79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Owned, dtype-grouped snapshots for asynchronous runner outputs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch.utils._pytree import tree_flatten, tree_unflatten
|
|
|
|
|
|
class PackedOutputSnapshot(dict):
|
|
"""A normal payload mapping with an internal batched-copy plan.
|
|
|
|
The runner's snapshot-slot event protects the device slabs until D2H is
|
|
complete. Host slabs are allocated per output, so downstream consumers can
|
|
retain their views without depending on reuse of the device ring slot.
|
|
"""
|
|
|
|
def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None:
|
|
super().__init__(tree_unflatten(leaves, spec))
|
|
self._leaves = leaves
|
|
self._spec = spec
|
|
self._groups = groups
|
|
|
|
def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]:
|
|
leaves = list(self._leaves)
|
|
for indices, slab in self._groups:
|
|
cpu = copy_tensor(slab)
|
|
offset = 0
|
|
for index in indices:
|
|
value = leaves[index]
|
|
size = value.numel()
|
|
leaves[index] = cpu.narrow(0, offset, size).view(value.shape)
|
|
offset += size
|
|
return tree_unflatten(leaves, self._spec)
|
|
|
|
|
|
def pack_output_snapshot(
|
|
payload: dict[str, Any], slot: dict[tuple[Any, ...], torch.Tensor], *, max_buckets: int
|
|
) -> PackedOutputSnapshot | None:
|
|
"""Copy tensor leaves once per dtype/device, retaining their exact values.
|
|
|
|
No CUDA synchronization is introduced here. The caller must select the
|
|
producer stream and wait for the previous consumer before reusing a slot.
|
|
"""
|
|
leaves, spec = tree_flatten(payload)
|
|
grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list)
|
|
for index, value in enumerate(leaves):
|
|
if isinstance(value, torch.Tensor):
|
|
# Preserve the existing per-tensor path for unusual layouts.
|
|
if value.layout != torch.strided or value.is_quantized:
|
|
return None
|
|
grouped[(value.dtype, value.device)].append(index)
|
|
groups = []
|
|
for (dtype, device), indices in grouped.items():
|
|
values = [leaves[index].detach().reshape(-1) for index in indices]
|
|
size = sum(value.numel() for value in values)
|
|
key = ("packed-output", dtype, device, size)
|
|
slab = slot.get(key)
|
|
if slab is None:
|
|
slab = torch.empty(size, dtype=dtype, device=device)
|
|
if len(slot) < max_buckets:
|
|
slot[key] = slab
|
|
if len(values) == 1:
|
|
slab.copy_(values[0])
|
|
else:
|
|
torch.cat(values, out=slab)
|
|
offset = 0
|
|
for index, value in zip(indices, values):
|
|
shape = leaves[index].shape
|
|
leaves[index] = slab.narrow(0, offset, value.numel()).view(shape)
|
|
offset += value.numel()
|
|
groups.append((indices, slab))
|
|
return PackedOutputSnapshot(leaves, spec, groups)
|