261 lines
8.3 KiB
Python
261 lines
8.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""TransferTensorBuffer: memory staging area for disaggregated tensor transfer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
|
|
import torch
|
|
|
|
from sglang.multimodal_gen.runtime.disaggregation.transport.allocator import (
|
|
BuddyAllocator,
|
|
)
|
|
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import (
|
|
str_to_dtype,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SlotHandle:
|
|
request_id: str
|
|
offset: int # byte offset in the pool
|
|
size: int # allocated size in bytes
|
|
tensor_views: dict[str, torch.Tensor | list[torch.Tensor]] = field(
|
|
default_factory=dict
|
|
)
|
|
|
|
|
|
class TransferTensorBuffer:
|
|
"""Memory pool for staging tensor payloads between roles.
|
|
|
|
Wraps a contiguous block of memory (CPU pinned or GPU) with a BuddyAllocator.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool_size: int,
|
|
min_block_size: int = 1 << 20,
|
|
role_name: str = "unknown",
|
|
device: str = "cpu",
|
|
):
|
|
self._role_name = role_name
|
|
self._device = device
|
|
self._allocator = BuddyAllocator(pool_size, min_block_size)
|
|
actual_size = self._allocator.pool_size
|
|
|
|
if device == "cpu":
|
|
self._pool = torch.empty(actual_size, dtype=torch.uint8, pin_memory=True)
|
|
else:
|
|
self._pool = torch.empty(actual_size, dtype=torch.uint8, device=device)
|
|
self._pool_ptr = self._pool.data_ptr()
|
|
|
|
pool_location = "pinned CPU" if device == "cpu" else f"GPU ({device})"
|
|
logger.info(
|
|
"TransferTensorBuffer[%s]: allocated %d MiB %s memory "
|
|
"(min_block=%d KiB)",
|
|
role_name,
|
|
actual_size >> 20,
|
|
pool_location,
|
|
min_block_size >> 10,
|
|
)
|
|
|
|
@property
|
|
def pool_size(self) -> int:
|
|
return self._allocator.pool_size
|
|
|
|
@property
|
|
def device(self) -> str:
|
|
return self._device
|
|
|
|
@property
|
|
def pool_data_ptr(self) -> int:
|
|
return self._pool_ptr
|
|
|
|
def allocate(self, size: int, request_id: str) -> SlotHandle | None:
|
|
"""Allocate a slot. Returns None if pool is full."""
|
|
offset = self._allocator.allocate(size, request_id=request_id)
|
|
if offset is None:
|
|
logger.warning(
|
|
"TransferTensorBuffer[%s]: allocation failed for %s (%d bytes). "
|
|
"Pool stats: %s",
|
|
self._role_name,
|
|
request_id,
|
|
size,
|
|
self._allocator.get_stats(),
|
|
)
|
|
return None
|
|
|
|
block = self._allocator.get_block_info(offset)
|
|
return SlotHandle(
|
|
request_id=request_id,
|
|
offset=offset,
|
|
size=block.size if block else size,
|
|
)
|
|
|
|
def free(self, handle: SlotHandle) -> bool:
|
|
return self._allocator.free(handle.offset)
|
|
|
|
def write_tensor(
|
|
self,
|
|
handle: SlotHandle,
|
|
tensor: torch.Tensor,
|
|
byte_offset: int = 0,
|
|
stream: torch.Stream | None = None,
|
|
) -> int:
|
|
"""Copy a tensor into the pool slot. Returns bytes written."""
|
|
src_tensor = tensor.contiguous()
|
|
nbytes = src_tensor.numel() * src_tensor.element_size()
|
|
|
|
if byte_offset + nbytes > handle.size:
|
|
raise ValueError(
|
|
f"Write exceeds slot: offset={byte_offset}, nbytes={nbytes}, "
|
|
f"slot_size={handle.size}"
|
|
)
|
|
|
|
dst = self._pool[
|
|
handle.offset + byte_offset : handle.offset + byte_offset + nbytes
|
|
]
|
|
src_bytes = src_tensor.view(torch.uint8).reshape(-1)
|
|
|
|
if stream is not None:
|
|
with torch.get_device_module().stream(stream):
|
|
dst.copy_(src_bytes, non_blocking=True)
|
|
else:
|
|
dst.copy_(src_bytes, non_blocking=True)
|
|
|
|
return nbytes
|
|
|
|
def read_tensor(
|
|
self,
|
|
handle: SlotHandle,
|
|
shape: list[int],
|
|
dtype: torch.dtype,
|
|
byte_offset: int = 0,
|
|
device: torch.device | str = "cpu",
|
|
stream: torch.Stream | None = None,
|
|
) -> torch.Tensor:
|
|
"""Read a tensor from the pool slot. Returns a clone on target device."""
|
|
nbytes = 1
|
|
for s in shape:
|
|
nbytes *= s
|
|
nbytes *= torch.tensor([], dtype=dtype).element_size()
|
|
|
|
raw = self._pool[
|
|
handle.offset + byte_offset : handle.offset + byte_offset + nbytes
|
|
]
|
|
src = raw.view(dtype).reshape(shape)
|
|
|
|
pool_dev = str(self._pool.device)
|
|
target_dev = str(device)
|
|
|
|
same_device = pool_dev == target_dev
|
|
|
|
if same_device:
|
|
# Clone to decouple tensor lifetime from pool slot
|
|
if stream is not None:
|
|
with torch.get_device_module().stream(stream):
|
|
return src.clone()
|
|
return src.clone()
|
|
|
|
if stream is not None:
|
|
with torch.get_device_module().stream(stream):
|
|
return src.to(device, non_blocking=True)
|
|
return src.to(device, non_blocking=True)
|
|
|
|
def write_tensors_from_gpu(
|
|
self,
|
|
handle: SlotHandle,
|
|
tensors: dict[str, torch.Tensor | list[torch.Tensor] | None],
|
|
stream: torch.Stream | None = None,
|
|
) -> dict[str, list[dict]]:
|
|
"""Batch-write GPU tensors into a slot. Returns a manifest for later reads."""
|
|
manifest: dict[str, list[dict]] = {}
|
|
byte_offset = 0
|
|
|
|
# Ensure copy stream sees all prior compute kernels
|
|
if stream is not None:
|
|
stream.wait_stream(torch.get_device_module().current_stream())
|
|
|
|
for name, value in tensors.items():
|
|
if value is None:
|
|
continue
|
|
|
|
entries = []
|
|
if isinstance(value, torch.Tensor):
|
|
nbytes = self.write_tensor(handle, value, byte_offset, stream)
|
|
entries.append(
|
|
{
|
|
"offset": byte_offset,
|
|
"shape": list(value.shape),
|
|
"dtype": str(value.dtype).replace("torch.", ""),
|
|
}
|
|
)
|
|
byte_offset += nbytes
|
|
byte_offset = (byte_offset + 511) & ~511 # align to 512B
|
|
|
|
elif isinstance(value, list):
|
|
for i, t in enumerate(value):
|
|
if t is None:
|
|
continue
|
|
nbytes = self.write_tensor(handle, t, byte_offset, stream)
|
|
entries.append(
|
|
{
|
|
"offset": byte_offset,
|
|
"shape": list(t.shape),
|
|
"dtype": str(t.dtype).replace("torch.", ""),
|
|
"list_index": i,
|
|
}
|
|
)
|
|
byte_offset += nbytes
|
|
byte_offset = (byte_offset + 511) & ~511
|
|
|
|
if entries:
|
|
manifest[name] = entries
|
|
|
|
return manifest
|
|
|
|
def read_tensors_from_manifest(
|
|
self,
|
|
handle: SlotHandle,
|
|
manifest: dict[str, list[dict]],
|
|
device: torch.device | str = "cpu",
|
|
stream: torch.Stream | None = None,
|
|
) -> dict[str, torch.Tensor | list[torch.Tensor]]:
|
|
"""Batch-read tensors from a slot using a manifest."""
|
|
result: dict[str, torch.Tensor | list[torch.Tensor]] = {}
|
|
|
|
for name, entries in manifest.items():
|
|
if not entries:
|
|
continue
|
|
has_list_index = any("list_index" in e for e in entries)
|
|
|
|
if has_list_index:
|
|
max_idx = max(e.get("list_index", 0) for e in entries) + 1
|
|
tensors = [None] * max_idx
|
|
for entry in entries:
|
|
t = self.read_tensor(
|
|
handle,
|
|
entry["shape"],
|
|
str_to_dtype(entry["dtype"]),
|
|
entry["offset"],
|
|
device,
|
|
stream,
|
|
)
|
|
tensors[entry["list_index"]] = t
|
|
result[name] = tensors
|
|
else:
|
|
entry = entries[0]
|
|
result[name] = self.read_tensor(
|
|
handle,
|
|
entry["shape"],
|
|
str_to_dtype(entry["dtype"]),
|
|
entry["offset"],
|
|
device,
|
|
stream,
|
|
)
|
|
|
|
return result
|