fix(vlm): stream-order cuda-ipc feature pool lifecycle and streamline multimodal transport module (#33949)
This commit is contained in:
Mick
@@ -110,6 +110,10 @@ from sglang.srt.model_executor.forward_batch_info import (
|
||||
ForwardBatch,
|
||||
ForwardMode,
|
||||
)
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
CudaIpcTensorTransportProxy,
|
||||
)
|
||||
from sglang.srt.observability.metrics_collector import (
|
||||
DPCooperationInfo,
|
||||
SchedulerMetricsCollector,
|
||||
@@ -124,10 +128,6 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import flatten_nested_list
|
||||
from sglang.srt.utils.cuda_ipc_transport_utils import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
CudaIpcTensorTransportProxy,
|
||||
)
|
||||
from sglang.srt.utils.token_sequence_matcher import TokenSequenceMatcher
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -730,9 +730,9 @@ class KimiK25ForConditionalGeneration(nn.Module):
|
||||
acknowledges the entire TP group so the bounded IPC pool remains
|
||||
recyclable.
|
||||
"""
|
||||
# Same source as MmItemMemoryPool.try_to_recycle(), which waits on
|
||||
# configured_tp_size(): the live world size agrees once dist is up,
|
||||
# but a refcount that disagrees with the waiter would strand items.
|
||||
# Match the configured TP consumer count captured when the
|
||||
# tokenizer creates MmItemMemoryPool. A live attention subgroup
|
||||
# size could leave acknowledgements missing and strand the lease.
|
||||
ipc_consumer_count = max(configured_tp_size(), 1)
|
||||
device_index = device.index
|
||||
if device.type == "cuda" and device_index is None:
|
||||
|
||||
@@ -3087,9 +3087,9 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
|
||||
def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
|
||||
"""Materialize only the images assigned to this vision-DP rank."""
|
||||
# Same source as MmItemMemoryPool.try_to_recycle(), which waits on
|
||||
# configured_tp_size(): the live world size agrees once dist is up,
|
||||
# but a refcount that disagrees with the waiter would strand items.
|
||||
# Match the configured TP consumer count captured when the
|
||||
# tokenizer creates MmItemMemoryPool. A live attention subgroup
|
||||
# size could leave acknowledgements missing and strand the lease.
|
||||
ipc_consumer_count = max(configured_tp_size(), 1)
|
||||
device_index = device.index
|
||||
if device.type == "cuda" and device_index is None:
|
||||
|
||||
@@ -1351,42 +1351,52 @@ class Qwen3VLForConditionalGeneration(nn.Module):
|
||||
|
||||
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||
_require_vision(self)
|
||||
pixel_values = materialize_multimodal_features(
|
||||
[item.feature for item in items],
|
||||
device=self.visual.device,
|
||||
dtype=self.visual.dtype,
|
||||
)
|
||||
image_grid_thw = torch.concat([item.image_grid_thw for item in items], dim=0)
|
||||
assert pixel_values.dim() == 2, pixel_values.dim()
|
||||
assert image_grid_thw.dim() == 2, image_grid_thw.dim()
|
||||
|
||||
if self.use_data_parallel:
|
||||
return run_dp_sharded_mrope_vision_model(
|
||||
self.visual,
|
||||
pixel_values,
|
||||
image_grid_thw.tolist(),
|
||||
rope_type="rope_3d",
|
||||
)
|
||||
else:
|
||||
return self.visual(pixel_values, grid_thw=image_grid_thw)
|
||||
return self._get_visual_feature(items, image_grid_thw)
|
||||
|
||||
def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||
_require_vision(self)
|
||||
pixel_values = materialize_multimodal_features(
|
||||
[item.feature for item in items],
|
||||
device=self.visual.device,
|
||||
dtype=self.visual.dtype,
|
||||
)
|
||||
video_grid_thw = torch.concat([item.video_grid_thw for item in items], dim=0)
|
||||
assert pixel_values.dim() == 2, pixel_values.dim()
|
||||
assert video_grid_thw.dim() == 2, video_grid_thw.dim()
|
||||
return self._get_visual_feature(items, video_grid_thw)
|
||||
|
||||
def _get_visual_feature(
|
||||
self, items: List[MultimodalDataItem], grid_thw: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
assert grid_thw.dim() == 2, grid_thw.dim()
|
||||
if self.use_data_parallel:
|
||||
return run_dp_sharded_mrope_vision_model(
|
||||
self.visual, pixel_values, video_grid_thw.tolist(), rope_type="rope_3d"
|
||||
self.visual,
|
||||
None,
|
||||
grid_thw.tolist(),
|
||||
rope_type="rope_3d",
|
||||
load_local_pixel_values=partial(self._materialize_visual_items, items),
|
||||
pixel_values_device=self.visual.device,
|
||||
pixel_values_dtype=self.visual.dtype,
|
||||
)
|
||||
else:
|
||||
video_embeds = self.visual(pixel_values, grid_thw=video_grid_thw)
|
||||
return video_embeds
|
||||
pixel_values = self._materialize_visual_items(items, range(len(items)))
|
||||
assert pixel_values.dim() == 2, pixel_values.dim()
|
||||
return self.visual(pixel_values, grid_thw=grid_thw)
|
||||
|
||||
def _materialize_visual_items(
|
||||
self, items: List[MultimodalDataItem], indices: Iterable[int]
|
||||
) -> torch.Tensor:
|
||||
device = self.visual.device
|
||||
device_index = device.index
|
||||
if device.type == "cuda" and device_index is None:
|
||||
device_index = torch.cuda.current_device()
|
||||
if device.type == "cuda":
|
||||
parallel = get_parallel()
|
||||
consumer_count = max(parallel.tp_size, 1)
|
||||
|
||||
features = []
|
||||
for index in indices:
|
||||
item = items[index]
|
||||
if device.type == "cuda":
|
||||
item.reconstruct(device_index, ipc_consumer_count=consumer_count)
|
||||
features.append(item.feature)
|
||||
return materialize_multimodal_features(
|
||||
features, device=device, dtype=self.visual.dtype
|
||||
)
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.model.embed_tokens
|
||||
|
||||
@@ -20,6 +20,12 @@ from sglang.srt.managers.schedule_batch import (
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.executor import MultimodalProcessorExecutor
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
MM_FEATURE_CACHE_SIZE,
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
||||
MmItemMemoryPool,
|
||||
get_mm_feature_pool_size_per_worker,
|
||||
)
|
||||
from sglang.srt.utils import (
|
||||
CLIENT_MEDIA_EXCEPTIONS,
|
||||
envs,
|
||||
@@ -31,13 +37,6 @@ from sglang.srt.utils import (
|
||||
load_video,
|
||||
logger,
|
||||
)
|
||||
from sglang.srt.utils.cuda_ipc_transport_utils import (
|
||||
MM_FEATURE_CACHE_SIZE,
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
||||
CudaIpcTensorTransportProxy,
|
||||
MmItemMemoryPool,
|
||||
get_mm_feature_pool_size_per_worker,
|
||||
)
|
||||
|
||||
_is_cpu = is_cpu()
|
||||
_is_npu = is_npu()
|
||||
@@ -364,6 +363,7 @@ class BaseMultimodalProcessor(ABC):
|
||||
per_worker_pool_size,
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
||||
self.server_args.base_gpu_id,
|
||||
self.server_args.tp_size,
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -1368,24 +1368,11 @@ class BaseMultimodalProcessor(ABC):
|
||||
if not tensor.is_cuda:
|
||||
return tensor
|
||||
|
||||
sync_flag, available_slice, byte_offset = (
|
||||
self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag(tensor)
|
||||
proxy = self.cudaipc_mmfeature_pool.wrap_tensor(
|
||||
tensor,
|
||||
use_pool_handle_cache=self.use_ipc_pool_handle_cache,
|
||||
)
|
||||
if isinstance(available_slice, torch.Tensor):
|
||||
available_slice.copy_(tensor.view(torch.int8).view(-1), non_blocking=True)
|
||||
return CudaIpcTensorTransportProxy(
|
||||
data=available_slice,
|
||||
info_data=tensor,
|
||||
sync_buffer_meta=sync_flag,
|
||||
pool_ipc_handle=(
|
||||
self.cudaipc_mmfeature_pool._pool_ipc_handle
|
||||
if self.use_ipc_pool_handle_cache
|
||||
else None
|
||||
),
|
||||
pool_byte_offset=byte_offset,
|
||||
pool_device_index=self.cudaipc_mmfeature_pool._pool_device_index,
|
||||
)
|
||||
return tensor.cpu()
|
||||
return proxy if proxy is not None else tensor.cpu()
|
||||
|
||||
@staticmethod
|
||||
def _move_feature_to_cpu(value):
|
||||
|
||||
@@ -20,7 +20,7 @@ from sglang.srt.multimodal.processors.base_processor import (
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
|
||||
from sglang.srt.utils.cuda_ipc_transport_utils import (
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
)
|
||||
|
||||
|
||||
@@ -44,10 +44,10 @@ from sglang.srt.multimodal.processors.kimi_k25 import (
|
||||
_grid_thw_from_resize_config,
|
||||
navit_resize_config,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda
|
||||
from sglang.srt.utils.cuda_ipc_transport_utils import (
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda
|
||||
|
||||
|
||||
def _encode_k3_special_tokens(tokenizer, text: str) -> list[int]:
|
||||
|
||||
@@ -37,6 +37,9 @@ from sglang.srt.multimodal.processors.base_processor import (
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
MultimodalSpecialTokens,
|
||||
)
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
)
|
||||
from sglang.srt.utils import cpu_has_amx_support, is_cpu
|
||||
from sglang.srt.utils.video_decoder import VideoDecoderWrapper
|
||||
from sglang.utils import logger
|
||||
@@ -774,6 +777,8 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
|
||||
base_output, self.mm_tokens, **processor_kwargs
|
||||
)
|
||||
|
||||
self._mark_dp_encoder_features_for_deferred_reconstruction(mm_items)
|
||||
|
||||
audio_feature_lengths = None
|
||||
|
||||
if self.model_type == "qwen3_omni_moe":
|
||||
@@ -885,3 +890,17 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
|
||||
mrope_positions=mrope_positions,
|
||||
mrope_position_delta=mrope_position_delta,
|
||||
)
|
||||
|
||||
def _mark_dp_encoder_features_for_deferred_reconstruction(self, mm_items):
|
||||
if not (
|
||||
self.keep_mm_features_on_device
|
||||
and self.server_args.mm_enable_dp_encoder
|
||||
and self.model_type
|
||||
in ("qwen3_vl", "qwen3_vl_moe", "qwen3_5", "qwen3_5_moe")
|
||||
):
|
||||
return
|
||||
for item in mm_items:
|
||||
if item.is_image() or item.is_video():
|
||||
item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = (
|
||||
True
|
||||
)
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""GPU transports for multimodal feature tensors."""
|
||||
@@ -0,0 +1,346 @@
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.multimodal.transport.memory_pool import (
|
||||
DEFAULT_MAX_INFLIGHT_SLICES,
|
||||
StreamOrderedMmFeaturePool,
|
||||
StreamOrderedPoolConsumerMixin,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MM_FEATURE_CACHE_SIZE = envs.SGLANG_MM_FEATURE_CACHE_MB.get() * 1024 * 1024
|
||||
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL = (
|
||||
envs.SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC.get()
|
||||
)
|
||||
|
||||
# Processors set this marker only when their encoder consumes each IPC feature
|
||||
# on a single TP rank. The scheduler then keeps the feature lazy until the
|
||||
# model has computed the data-parallel assignment.
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY = (
|
||||
"_sglang_defer_cuda_ipc_feature_reconstruction"
|
||||
)
|
||||
|
||||
|
||||
def get_mm_feature_pool_size_per_worker(
|
||||
total_pool_size: int, tokenizer_worker_num: int
|
||||
) -> int:
|
||||
"""Split the CUDA IPC feature-pool budget without exceeding it.
|
||||
|
||||
Each tokenizer worker owns a distinct CUDA allocation, even though all pools
|
||||
are created on ``base_gpu_id``. Therefore a minimum per-worker allocation
|
||||
would make the aggregate HBM reservation larger than the configured budget.
|
||||
Keep the configured value as a hard per-node cap and leave at most
|
||||
``tokenizer_worker_num - 1`` bytes unused when it is not evenly divisible.
|
||||
"""
|
||||
if total_pool_size <= 0:
|
||||
raise ValueError("total_pool_size must be positive")
|
||||
if tokenizer_worker_num <= 0:
|
||||
raise ValueError("tokenizer_worker_num must be positive")
|
||||
|
||||
return total_pool_size // tokenizer_worker_num
|
||||
|
||||
|
||||
# Cache for pool-level IPC handles on the consumer side.
|
||||
# Key: the pool CUDA IPC handle tuple. Value: opened UntypedStorage.
|
||||
_pool_storage_cache: dict = {}
|
||||
_pool_cache_lock = threading.Lock()
|
||||
|
||||
|
||||
def _normalize_pool_cache_key(pool_handle, device_index: int) -> tuple[Any, ...]:
|
||||
normalized_handle = (
|
||||
pool_handle if isinstance(pool_handle, tuple) else tuple(pool_handle)
|
||||
)
|
||||
return (device_index, normalized_handle)
|
||||
|
||||
|
||||
def _open_pooled_storage_uncached(pool_handle):
|
||||
return torch.UntypedStorage._new_shared_cuda(*pool_handle)
|
||||
|
||||
|
||||
def _pool_handle_cache_get_or_open(cache_key, pool_handle):
|
||||
storage = _pool_storage_cache.get(cache_key)
|
||||
if storage is None:
|
||||
with _pool_cache_lock:
|
||||
storage = _pool_storage_cache.get(cache_key)
|
||||
if storage is None:
|
||||
storage = _open_pooled_storage_uncached(pool_handle)
|
||||
_pool_storage_cache[cache_key] = storage
|
||||
return storage
|
||||
|
||||
|
||||
def _pool_handle_cache_set(cache_key, storage):
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache[cache_key] = storage
|
||||
|
||||
|
||||
def _pool_handle_cache_invalidate(cache_key):
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache.pop(cache_key, None)
|
||||
|
||||
|
||||
def _pool_handle_cache_clear():
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache.clear()
|
||||
|
||||
|
||||
class MmItemMemoryPool:
|
||||
def __init__(
|
||||
self,
|
||||
memory_size: int,
|
||||
recycle_interval: float,
|
||||
base_gpu_id: int,
|
||||
consumer_count: int,
|
||||
max_inflight_slices: int = DEFAULT_MAX_INFLIGHT_SLICES,
|
||||
):
|
||||
self.device_id = base_gpu_id
|
||||
self.consumer_count = consumer_count
|
||||
self.memory_pool = torch.empty(
|
||||
memory_size, dtype=torch.uint8, device=f"cuda:{base_gpu_id}"
|
||||
).contiguous()
|
||||
self._pool = StreamOrderedMmFeaturePool(
|
||||
memory_size=memory_size,
|
||||
byte_tensor=self.memory_pool,
|
||||
base_address=self.memory_pool.data_ptr(),
|
||||
device_id=base_gpu_id,
|
||||
consumer_count=consumer_count,
|
||||
recycle_interval=recycle_interval,
|
||||
transport_name="CUDA IPC",
|
||||
max_inflight_slices=max_inflight_slices,
|
||||
)
|
||||
storage = self.memory_pool.untyped_storage()
|
||||
self._pool_ipc_handle = storage._share_cuda_()
|
||||
self._pool_full_warned = False
|
||||
|
||||
logger.debug(
|
||||
f"[MmItemMemoryPool] init: memory_size={memory_size}, "
|
||||
f"recycle_interval={recycle_interval}s"
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
self._pool.shutdown()
|
||||
|
||||
@property
|
||||
def active_lease_count(self) -> int:
|
||||
return self._pool.active_lease_count
|
||||
|
||||
def wrap_tensor(
|
||||
self, tensor: torch.Tensor, *, use_pool_handle_cache: bool
|
||||
) -> Optional["CudaIpcTensorTransportProxy"]:
|
||||
lease, destination = self._pool.copy_tensor(tensor)
|
||||
if lease is None:
|
||||
nbytes = tensor.numel() * tensor.element_size()
|
||||
self._warn_pool_full_once(nbytes)
|
||||
return None
|
||||
|
||||
return CudaIpcTensorTransportProxy(
|
||||
data=destination,
|
||||
info_data=tensor,
|
||||
pool_ipc_handle=self._pool_ipc_handle,
|
||||
pool_byte_offset=lease.start,
|
||||
ready_byte_offset=lease.ready_byte_offset,
|
||||
ack_byte_offset=lease.ack_byte_offset,
|
||||
generation=lease.generation,
|
||||
total_consumer_count=self.consumer_count,
|
||||
use_pool_handle_cache=use_pool_handle_cache,
|
||||
)
|
||||
|
||||
def _warn_pool_full_once(self, nbytes: int):
|
||||
if self._pool_full_warned:
|
||||
return
|
||||
self._pool_full_warned = True
|
||||
pool_mb = (
|
||||
self.memory_pool.numel() * self.memory_pool.element_size() / (1024 * 1024)
|
||||
)
|
||||
need_mb = nbytes / (1024 * 1024)
|
||||
logger.warning(
|
||||
"MmItemMemoryPool has no free chunk large enough for a %.2f MiB tensor "
|
||||
"(pool size: %.2f MiB); falling back to non-IPC transport. "
|
||||
"Consider increasing SGLANG_MM_FEATURE_CACHE_MB.",
|
||||
need_mb,
|
||||
pool_mb,
|
||||
)
|
||||
|
||||
|
||||
class CudaIpcTensorTransportProxy(StreamOrderedPoolConsumerMixin):
|
||||
"""Serializable view of one tensor stored in a CUDA IPC memory pool.
|
||||
|
||||
The producer-ready word and one acknowledgement word per consumer live in
|
||||
the same CUDA allocation as the tensor. CUDA stream memory operations order
|
||||
the producer copy, consumer copy, and pool reuse without CPU shared memory
|
||||
or device-wide synchronization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data: torch.Tensor,
|
||||
info_data: torch.Tensor,
|
||||
pool_ipc_handle,
|
||||
pool_byte_offset: int,
|
||||
ready_byte_offset: int,
|
||||
ack_byte_offset: int,
|
||||
generation: int,
|
||||
total_consumer_count: int,
|
||||
use_pool_handle_cache: bool,
|
||||
):
|
||||
if (not isinstance(data, torch.Tensor)) or (
|
||||
not isinstance(info_data, torch.Tensor)
|
||||
):
|
||||
raise TypeError(
|
||||
f"Input 'data' must be a torch.Tensor, but got {type(data)}"
|
||||
)
|
||||
|
||||
self._init_stream_ordered_consumer(
|
||||
ready_byte_offset=ready_byte_offset,
|
||||
ack_byte_offset=ack_byte_offset,
|
||||
generation=generation,
|
||||
total_consumer_count=total_consumer_count,
|
||||
transport_name="CUDA IPC",
|
||||
)
|
||||
|
||||
self.proxy_state = {
|
||||
"ipc_extra": {
|
||||
"pool_handle": pool_ipc_handle,
|
||||
"pool_byte_offset": pool_byte_offset,
|
||||
"shape": data.shape,
|
||||
"dtype": data.dtype,
|
||||
"stride": data.stride(),
|
||||
"storage_offset": 0,
|
||||
"nbytes": data.numel() * data.element_size(),
|
||||
"recons_shape": info_data.shape,
|
||||
"recons_dtype": info_data.dtype,
|
||||
"use_pool_handle_cache": use_pool_handle_cache,
|
||||
},
|
||||
"tensor_data": None,
|
||||
}
|
||||
self.reconstruct_tensor = None
|
||||
# Keep uncached mappings alive until the work enqueued on the consumer
|
||||
# stream has completed.
|
||||
self._pool_storage = None
|
||||
|
||||
def _reconstruct_from_ipc_extra(
|
||||
self, ipc_extra, *, use_cache: bool, rebuild_device_idx: int
|
||||
):
|
||||
shape = ipc_extra["shape"]
|
||||
dtype = ipc_extra["dtype"]
|
||||
stride = ipc_extra["stride"]
|
||||
# Redirect handle[0] to the consumer's device so _new_shared_cuda's
|
||||
# CUDAGuard stays there; peer access handles the cross-GPU open.
|
||||
pool_handle = ipc_extra["pool_handle"]
|
||||
redirected_handle = (rebuild_device_idx,) + tuple(pool_handle)[1:]
|
||||
target_device = torch.device(f"cuda:{rebuild_device_idx}")
|
||||
cache_key = _normalize_pool_cache_key(pool_handle, rebuild_device_idx)
|
||||
|
||||
with torch.cuda.device(target_device):
|
||||
if use_cache:
|
||||
storage = _pool_handle_cache_get_or_open(cache_key, redirected_handle)
|
||||
else:
|
||||
storage = _open_pooled_storage_uncached(redirected_handle)
|
||||
slice_storage = storage[
|
||||
ipc_extra["pool_byte_offset"] : ipc_extra["pool_byte_offset"]
|
||||
+ ipc_extra["nbytes"]
|
||||
]
|
||||
slice_tensor = torch.empty(0, dtype=dtype, device=target_device).set_(
|
||||
slice_storage,
|
||||
storage_offset=ipc_extra["storage_offset"],
|
||||
size=shape,
|
||||
stride=stride,
|
||||
)
|
||||
|
||||
return slice_tensor, storage
|
||||
|
||||
def _open_pool_slice(self, rebuild_device_idx: int):
|
||||
ipc_extra = self.proxy_state["ipc_extra"]
|
||||
use_cache = ipc_extra["use_pool_handle_cache"]
|
||||
try:
|
||||
return self._reconstruct_from_ipc_extra(
|
||||
ipc_extra,
|
||||
use_cache=use_cache,
|
||||
rebuild_device_idx=rebuild_device_idx,
|
||||
)
|
||||
except Exception as exc:
|
||||
if not use_cache:
|
||||
raise
|
||||
cache_key = _normalize_pool_cache_key(
|
||||
ipc_extra["pool_handle"], rebuild_device_idx
|
||||
)
|
||||
logger.info(
|
||||
"Failed to deserialize from cached pooled CUDA IPC handle (%s). "
|
||||
"Invalidating cache entry and retrying uncached.",
|
||||
exc,
|
||||
)
|
||||
_pool_handle_cache_invalidate(cache_key)
|
||||
result = self._reconstruct_from_ipc_extra(
|
||||
ipc_extra,
|
||||
use_cache=False,
|
||||
rebuild_device_idx=rebuild_device_idx,
|
||||
)
|
||||
_pool_handle_cache_set(cache_key, result[1])
|
||||
return result
|
||||
|
||||
def _retain_storage_until_stream_completes(self, storage, device_id: int) -> None:
|
||||
if self.proxy_state["ipc_extra"]["use_pool_handle_cache"]:
|
||||
# The process-wide cache owns the mapping after this proxy is
|
||||
# replaced by its reconstructed tensor.
|
||||
self._pool_storage = storage
|
||||
else:
|
||||
# An uncached mapping is owned only by this proxy. The caller
|
||||
# replaces the proxy immediately, so finish the current stream
|
||||
# before allowing the mapping to close.
|
||||
torch.cuda.current_stream(device_id).synchronize()
|
||||
|
||||
def acknowledge_consumption(
|
||||
self, consumer_count: int = 1, consumer_rank: Optional[int] = None
|
||||
) -> None:
|
||||
"""Stream-order pool release when a cache hit needs no tensor copy."""
|
||||
if self._consumer_acknowledged:
|
||||
return
|
||||
device_id = torch.cuda.current_device()
|
||||
with torch.cuda.device(device_id):
|
||||
_, storage = self._open_pool_slice(device_id)
|
||||
base_address = storage.data_ptr()
|
||||
self._wait_until_ready(base_address, device_id)
|
||||
self._acknowledge_on_stream(
|
||||
base_address, device_id, consumer_count, consumer_rank
|
||||
)
|
||||
self._retain_storage_until_stream_completes(storage, device_id)
|
||||
|
||||
def reconstruct_on_target_device(
|
||||
self,
|
||||
rebuild_device_idx,
|
||||
consumer_count: int = 1,
|
||||
consumer_rank: Optional[int] = None,
|
||||
):
|
||||
rebuild_device = torch.device(f"cuda:{rebuild_device_idx}")
|
||||
if (
|
||||
isinstance(self.reconstruct_tensor, torch.Tensor)
|
||||
and self.reconstruct_tensor.device == rebuild_device
|
||||
):
|
||||
return self.reconstruct_tensor
|
||||
|
||||
ipc_extra = self.proxy_state["ipc_extra"]
|
||||
with torch.cuda.device(rebuild_device):
|
||||
slice_tensor, storage = self._open_pool_slice(rebuild_device_idx)
|
||||
base_address = storage.data_ptr()
|
||||
self._wait_until_ready(base_address, rebuild_device_idx)
|
||||
reconstructed_tensor = torch.empty(
|
||||
ipc_extra["recons_shape"],
|
||||
dtype=ipc_extra["recons_dtype"],
|
||||
device=rebuild_device,
|
||||
).contiguous()
|
||||
reconstructed_tensor.view(torch.uint8).reshape(-1).copy_(slice_tensor)
|
||||
self._acknowledge_on_stream(
|
||||
base_address,
|
||||
rebuild_device_idx,
|
||||
consumer_count,
|
||||
consumer_rank,
|
||||
)
|
||||
|
||||
self._retain_storage_until_stream_completes(storage, rebuild_device_idx)
|
||||
self.reconstruct_tensor = reconstructed_tensor
|
||||
return self.reconstruct_tensor
|
||||
@@ -0,0 +1,367 @@
|
||||
"""Shared stream-ordered lifecycle for GPU multimodal feature pools."""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CONTROL_WORD_BYTES = 4
|
||||
DATA_ALIGNMENT = 256
|
||||
DEFAULT_MAX_INFLIGHT_SLICES = 4096
|
||||
|
||||
|
||||
def align_up(value: int, alignment: int) -> int:
|
||||
return ((value + alignment - 1) // alignment) * alignment
|
||||
|
||||
|
||||
def _driver_modules():
|
||||
from cuda.bindings import driver as cuda
|
||||
|
||||
from sglang.srt.distributed.device_communicators.vmm_utils import check_drv
|
||||
|
||||
return cuda, check_drv
|
||||
|
||||
|
||||
def stream_wait_value32(
|
||||
device_id: int, address: int, value: int, transport_name: str
|
||||
) -> None:
|
||||
cuda, check_drv = _driver_modules()
|
||||
stream = torch.cuda.current_stream(device_id)
|
||||
check_drv(
|
||||
cuda.cuStreamWaitValue32(stream.cuda_stream, address, value, 0),
|
||||
f"cuStreamWaitValue32(mm {transport_name})",
|
||||
)
|
||||
|
||||
|
||||
def stream_write_value32(
|
||||
device_id: int, address: int, value: int, transport_name: str
|
||||
) -> None:
|
||||
cuda, check_drv = _driver_modules()
|
||||
stream = torch.cuda.current_stream(device_id)
|
||||
check_drv(
|
||||
cuda.cuStreamWriteValue32(stream.cuda_stream, address, value, 0),
|
||||
f"cuStreamWriteValue32(mm {transport_name})",
|
||||
)
|
||||
|
||||
|
||||
def resolve_consumer_rank(
|
||||
total_consumer_count: int,
|
||||
consumer_rank: Optional[int] = None,
|
||||
transport_name: str = "GPU",
|
||||
) -> int:
|
||||
if total_consumer_count == 1:
|
||||
return 0
|
||||
if consumer_rank is None:
|
||||
try:
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
|
||||
# Use the global TP rank. An attention/DCP subgroup rank can alias
|
||||
# another consumer's acknowledgement slot.
|
||||
rank = int(get_parallel().tp_rank)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"Cannot resolve the {transport_name} consumer rank before "
|
||||
"parallel state initialization"
|
||||
) from exc
|
||||
else:
|
||||
rank = int(consumer_rank)
|
||||
if not 0 <= rank < total_consumer_count:
|
||||
raise RuntimeError(
|
||||
f"{transport_name} consumer rank {rank} is outside "
|
||||
f"[0, {total_consumer_count})"
|
||||
)
|
||||
return rank
|
||||
|
||||
|
||||
class StreamOrderedPoolConsumerMixin:
|
||||
"""Ready/wait/ack protocol for stream-ordered GPU feature proxies."""
|
||||
|
||||
def _init_stream_ordered_consumer(
|
||||
self,
|
||||
*,
|
||||
ready_byte_offset: int,
|
||||
ack_byte_offset: int,
|
||||
generation: int,
|
||||
total_consumer_count: int,
|
||||
transport_name: str,
|
||||
) -> None:
|
||||
if total_consumer_count <= 0:
|
||||
raise ValueError("total_consumer_count must be positive")
|
||||
self.ready_byte_offset = ready_byte_offset
|
||||
self.ack_byte_offset = ack_byte_offset
|
||||
self.generation = generation
|
||||
self.total_consumer_count = total_consumer_count
|
||||
self.transport_name = transport_name
|
||||
self._consumer_acknowledged = False
|
||||
|
||||
def _wait_until_ready(self, base_address: int, device_id: int) -> None:
|
||||
stream_wait_value32(
|
||||
device_id,
|
||||
base_address + self.ready_byte_offset,
|
||||
self.generation,
|
||||
self.transport_name,
|
||||
)
|
||||
|
||||
def _acknowledge_on_stream(
|
||||
self,
|
||||
base_address: int,
|
||||
device_id: int,
|
||||
consumer_count: int,
|
||||
consumer_rank: Optional[int] = None,
|
||||
) -> None:
|
||||
if self._consumer_acknowledged:
|
||||
return
|
||||
if consumer_count == self.total_consumer_count:
|
||||
consumer_ranks = range(self.total_consumer_count)
|
||||
elif consumer_count == 1:
|
||||
consumer_ranks = (
|
||||
resolve_consumer_rank(
|
||||
self.total_consumer_count,
|
||||
consumer_rank,
|
||||
self.transport_name,
|
||||
),
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"{self.transport_name} acknowledgements support one consumer "
|
||||
"or the complete consumer group, got "
|
||||
f"{consumer_count}/{self.total_consumer_count}"
|
||||
)
|
||||
|
||||
for rank in consumer_ranks:
|
||||
stream_write_value32(
|
||||
device_id,
|
||||
base_address + self.ack_byte_offset + rank * CONTROL_WORD_BYTES,
|
||||
self.generation,
|
||||
self.transport_name,
|
||||
)
|
||||
self._consumer_acknowledged = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PoolLease:
|
||||
start: int
|
||||
end: int
|
||||
nbytes: int
|
||||
slot: int
|
||||
generation: int
|
||||
ready_byte_offset: int
|
||||
ack_byte_offset: int
|
||||
|
||||
|
||||
class StreamOrderedMmFeaturePool:
|
||||
"""Bounded GPU pool with generation-safe producer/consumer leases."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
memory_size: int,
|
||||
byte_tensor: torch.Tensor,
|
||||
base_address: int,
|
||||
device_id: int,
|
||||
consumer_count: int,
|
||||
recycle_interval: float,
|
||||
transport_name: str,
|
||||
max_inflight_slices: int = DEFAULT_MAX_INFLIGHT_SLICES,
|
||||
) -> None:
|
||||
if memory_size <= 0:
|
||||
raise ValueError("memory_size must be positive")
|
||||
if consumer_count <= 0:
|
||||
raise ValueError("consumer_count must be positive")
|
||||
if max_inflight_slices <= 0:
|
||||
raise ValueError("max_inflight_slices must be positive")
|
||||
if recycle_interval <= 0:
|
||||
raise ValueError("recycle_interval must be positive")
|
||||
if (
|
||||
not byte_tensor.is_cuda
|
||||
or byte_tensor.device.index != device_id
|
||||
or byte_tensor.dtype != torch.uint8
|
||||
or not byte_tensor.is_contiguous()
|
||||
or byte_tensor.numel() < memory_size
|
||||
):
|
||||
raise ValueError(
|
||||
"byte_tensor must be a sufficiently large contiguous uint8 tensor "
|
||||
f"on cuda:{device_id}"
|
||||
)
|
||||
|
||||
self.memory_size = memory_size
|
||||
self.byte_tensor = byte_tensor
|
||||
self.base_address = base_address
|
||||
self.device_id = device_id
|
||||
self.consumer_count = consumer_count
|
||||
self.control_words_per_slot = 1 + consumer_count
|
||||
self.max_inflight_slices = max_inflight_slices
|
||||
self.transport_name = transport_name
|
||||
control_bytes = (
|
||||
max_inflight_slices * self.control_words_per_slot * CONTROL_WORD_BYTES
|
||||
)
|
||||
self.data_start = align_up(control_bytes, DATA_ALIGNMENT)
|
||||
if memory_size <= self.data_start:
|
||||
raise ValueError(
|
||||
f"{transport_name} pool is too small after control metadata: "
|
||||
f"pool={memory_size}, control={self.data_start}"
|
||||
)
|
||||
|
||||
control_word_count = max_inflight_slices * self.control_words_per_slot
|
||||
self._control_words = (
|
||||
byte_tensor[: control_word_count * CONTROL_WORD_BYTES]
|
||||
.view(torch.int32)
|
||||
.view(max_inflight_slices, self.control_words_per_slot)
|
||||
)
|
||||
self._control_words.zero_()
|
||||
torch.cuda.synchronize(device_id)
|
||||
|
||||
self._available_ranges = [(self.data_start, memory_size)]
|
||||
self._available_slots = list(reversed(range(max_inflight_slices)))
|
||||
self._slot_generations = [0] * max_inflight_slices
|
||||
self._occupied: dict[int, PoolLease] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._recycle_interval = recycle_interval
|
||||
self._recycler_stop_event = threading.Event()
|
||||
self._recycle_thread = threading.Thread(
|
||||
target=self._recycle_loop,
|
||||
name=f"{transport_name}MmFeaturePoolRecycler",
|
||||
daemon=True,
|
||||
)
|
||||
self._recycle_thread.start()
|
||||
|
||||
@property
|
||||
def usable_size(self) -> int:
|
||||
return self.memory_size - self.data_start
|
||||
|
||||
@property
|
||||
def active_lease_count(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._occupied)
|
||||
|
||||
def _allocate_locked(self, nbytes: int) -> Optional[PoolLease]:
|
||||
allocation_bytes = align_up(nbytes, DATA_ALIGNMENT)
|
||||
candidates = [
|
||||
(end - start, index, start, end)
|
||||
for index, (start, end) in enumerate(self._available_ranges)
|
||||
if end - start >= allocation_bytes
|
||||
]
|
||||
if not candidates or not self._available_slots:
|
||||
return None
|
||||
_, index, start, end = min(candidates)
|
||||
self._available_ranges.pop(index)
|
||||
if start + allocation_bytes < end:
|
||||
self._available_ranges.append((start + allocation_bytes, end))
|
||||
slot = self._available_slots.pop()
|
||||
generation = self._slot_generations[slot] + 1
|
||||
if generation > 0x7FFFFFFF:
|
||||
raise RuntimeError(f"{self.transport_name} pool slot generation exhausted")
|
||||
self._slot_generations[slot] = generation
|
||||
ready_byte_offset = slot * self.control_words_per_slot * CONTROL_WORD_BYTES
|
||||
lease = PoolLease(
|
||||
start=start,
|
||||
end=start + allocation_bytes,
|
||||
nbytes=nbytes,
|
||||
slot=slot,
|
||||
generation=generation,
|
||||
ready_byte_offset=ready_byte_offset,
|
||||
ack_byte_offset=ready_byte_offset + CONTROL_WORD_BYTES,
|
||||
)
|
||||
self._occupied[slot] = lease
|
||||
return lease
|
||||
|
||||
def _release_locked(self, lease: PoolLease) -> None:
|
||||
active_lease = self._occupied.get(lease.slot)
|
||||
if active_lease != lease:
|
||||
raise RuntimeError(
|
||||
f"Cannot release inactive {self.transport_name} pool lease "
|
||||
f"(slot={lease.slot}, generation={lease.generation})"
|
||||
)
|
||||
del self._occupied[lease.slot]
|
||||
self._available_slots.append(lease.slot)
|
||||
self._available_ranges.append((lease.start, lease.end))
|
||||
|
||||
def _merge_ranges_locked(self) -> None:
|
||||
merged = []
|
||||
for start, end in sorted(self._available_ranges):
|
||||
if merged and merged[-1][1] == start:
|
||||
merged[-1] = (merged[-1][0], end)
|
||||
else:
|
||||
merged.append((start, end))
|
||||
self._available_ranges = merged
|
||||
|
||||
def _recycle_ready_leases_locked(self) -> None:
|
||||
if not self._occupied:
|
||||
return
|
||||
leases = list(self._occupied.values())
|
||||
slot_indices = torch.tensor(
|
||||
[lease.slot for lease in leases],
|
||||
dtype=torch.long,
|
||||
device=f"cuda:{self.device_id}",
|
||||
)
|
||||
expected = torch.tensor(
|
||||
[lease.generation for lease in leases],
|
||||
dtype=torch.int32,
|
||||
device=f"cuda:{self.device_id}",
|
||||
).unsqueeze(1)
|
||||
completed = (
|
||||
(self._control_words.index_select(0, slot_indices) == expected)
|
||||
.all(dim=1)
|
||||
.cpu()
|
||||
.tolist()
|
||||
)
|
||||
for lease, is_complete in zip(leases, completed):
|
||||
if is_complete:
|
||||
self._release_locked(lease)
|
||||
self._merge_ranges_locked()
|
||||
|
||||
def _recycle_loop(self) -> None:
|
||||
torch.cuda.set_device(self.device_id)
|
||||
while not self._recycler_stop_event.is_set():
|
||||
try:
|
||||
with self._lock, torch.cuda.device(self.device_id):
|
||||
self._recycle_ready_leases_locked()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"%s multimodal pool recycle failed",
|
||||
self.transport_name,
|
||||
exc_info=True,
|
||||
)
|
||||
self._recycler_stop_event.wait(self._recycle_interval)
|
||||
|
||||
def copy_tensor(
|
||||
self, tensor: torch.Tensor
|
||||
) -> tuple[Optional[PoolLease], Optional[torch.Tensor]]:
|
||||
if not tensor.is_cuda:
|
||||
raise ValueError(f"{self.transport_name} requires a CUDA tensor")
|
||||
source = tensor.contiguous()
|
||||
nbytes = source.numel() * source.element_size()
|
||||
if nbytes == 0:
|
||||
raise ValueError(f"{self.transport_name} cannot transport an empty tensor")
|
||||
with self._lock:
|
||||
lease = self._allocate_locked(nbytes)
|
||||
if lease is None:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
with torch.cuda.device(self.device_id):
|
||||
destination = self.byte_tensor[lease.start : lease.start + lease.nbytes]
|
||||
destination.copy_(
|
||||
source.view(torch.uint8).reshape(-1), non_blocking=True
|
||||
)
|
||||
stream_write_value32(
|
||||
self.device_id,
|
||||
self.base_address + lease.ready_byte_offset,
|
||||
lease.generation,
|
||||
self.transport_name,
|
||||
)
|
||||
except Exception:
|
||||
with self._lock:
|
||||
self._release_locked(lease)
|
||||
self._merge_ranges_locked()
|
||||
raise
|
||||
return lease, destination
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._recycler_stop_event.set()
|
||||
if self._recycle_thread.is_alive():
|
||||
self._recycle_thread.join()
|
||||
@@ -1,568 +1,24 @@
|
||||
import fcntl
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from multiprocessing import shared_memory
|
||||
from typing import Any, Tuple
|
||||
"""Compatibility imports for the multimodal CUDA IPC transport.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
New code should import from :mod:`sglang.srt.multimodal.transport.cuda_ipc`.
|
||||
"""
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.runtime_context import (
|
||||
configured_tp_size,
|
||||
)
|
||||
from sglang.srt.utils.stale_shm_cleanup import make_shm_name
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MM_FEATURE_CACHE_SIZE = envs.SGLANG_MM_FEATURE_CACHE_MB.get() * 1024 * 1024
|
||||
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL = (
|
||||
envs.SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC.get()
|
||||
from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
|
||||
MM_FEATURE_CACHE_SIZE,
|
||||
MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
|
||||
CudaIpcTensorTransportProxy,
|
||||
MmItemMemoryPool,
|
||||
_pool_handle_cache_clear,
|
||||
get_mm_feature_pool_size_per_worker,
|
||||
)
|
||||
|
||||
SHM_LOCK_FILE = "/tmp/shm_wr_lock.lock"
|
||||
|
||||
# Processors set this marker only when their encoder consumes each IPC feature
|
||||
# on a single TP rank. The scheduler then keeps the feature lazy until the
|
||||
# model has computed the data-parallel assignment.
|
||||
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY = (
|
||||
"_sglang_defer_cuda_ipc_feature_reconstruction"
|
||||
)
|
||||
|
||||
|
||||
def get_mm_feature_pool_size_per_worker(
|
||||
total_pool_size: int, tokenizer_worker_num: int
|
||||
) -> int:
|
||||
"""Split the CUDA IPC feature-pool budget without exceeding it.
|
||||
|
||||
Each tokenizer worker owns a distinct CUDA allocation, even though all pools
|
||||
are created on ``base_gpu_id``. Therefore a minimum per-worker allocation
|
||||
would make the aggregate HBM reservation larger than the configured budget.
|
||||
Keep the configured value as a hard per-node cap and leave at most
|
||||
``tokenizer_worker_num - 1`` bytes unused when it is not evenly divisible.
|
||||
"""
|
||||
if total_pool_size <= 0:
|
||||
raise ValueError("total_pool_size must be positive")
|
||||
if tokenizer_worker_num <= 0:
|
||||
raise ValueError("tokenizer_worker_num must be positive")
|
||||
|
||||
return total_pool_size // tokenizer_worker_num
|
||||
|
||||
|
||||
# Cache for pool-level IPC handles on the consumer side.
|
||||
# Key: the pool CUDA IPC handle tuple. Value: opened UntypedStorage.
|
||||
_pool_storage_cache: dict = {}
|
||||
_pool_cache_lock = threading.Lock()
|
||||
|
||||
|
||||
def _normalize_pool_cache_key(pool_handle, pool_device_index: int) -> tuple[Any, ...]:
|
||||
normalized_handle = (
|
||||
pool_handle if isinstance(pool_handle, tuple) else tuple(pool_handle)
|
||||
)
|
||||
return (pool_device_index, normalized_handle)
|
||||
|
||||
|
||||
def _open_pooled_storage_uncached(pool_handle):
|
||||
return torch.UntypedStorage._new_shared_cuda(*pool_handle)
|
||||
|
||||
|
||||
def _pool_handle_cache_get_or_open(cache_key, pool_handle):
|
||||
storage = _pool_storage_cache.get(cache_key)
|
||||
if storage is None:
|
||||
with _pool_cache_lock:
|
||||
storage = _pool_storage_cache.get(cache_key)
|
||||
if storage is None:
|
||||
storage = _open_pooled_storage_uncached(pool_handle)
|
||||
_pool_storage_cache[cache_key] = storage
|
||||
return storage
|
||||
|
||||
|
||||
def _pool_handle_cache_set(cache_key, storage):
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache[cache_key] = storage
|
||||
|
||||
|
||||
def _pool_handle_cache_invalidate(cache_key):
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache.pop(cache_key, None)
|
||||
|
||||
|
||||
def _pool_handle_cache_clear():
|
||||
with _pool_cache_lock:
|
||||
_pool_storage_cache.clear()
|
||||
|
||||
|
||||
class ShmSyncBuffer:
|
||||
def __init__(self, byte_size: int = 4):
|
||||
self.buffer = shared_memory.SharedMemory(
|
||||
create=True, size=byte_size, name=make_shm_name("sync")
|
||||
)
|
||||
self.buffer_wrapper = np.ndarray(1, dtype=np.float32, buffer=self.buffer.buf)
|
||||
self.buffer_wrapper *= 0
|
||||
self.meta_data = {
|
||||
"handle": self.buffer.name,
|
||||
"shape": self.buffer_wrapper.shape,
|
||||
"dtype": str(self.buffer_wrapper.dtype),
|
||||
}
|
||||
|
||||
def __del__(self):
|
||||
if isinstance(self.buffer, shared_memory.SharedMemory):
|
||||
self.buffer.close()
|
||||
self.buffer.unlink()
|
||||
|
||||
|
||||
class MmItemMemoryChunk:
|
||||
def __init__(self, area: Tuple, sync_buffer: ShmSyncBuffer):
|
||||
self.area = area
|
||||
self.sync_flag = sync_buffer
|
||||
|
||||
@property
|
||||
def mem_size(self):
|
||||
return self.area[1] - self.area[0]
|
||||
|
||||
@property
|
||||
def start(self):
|
||||
return self.area[0]
|
||||
|
||||
@property
|
||||
def end(self):
|
||||
return self.area[1]
|
||||
|
||||
def try_to_recycle(self) -> bool:
|
||||
try:
|
||||
tp_num = configured_tp_size()
|
||||
except Exception:
|
||||
logger.info(
|
||||
"server_args has not been published yet, skip this turn's recycle"
|
||||
)
|
||||
return False
|
||||
|
||||
val = float(self.sync_flag.buffer_wrapper.item())
|
||||
logger.debug(f"[try_to_recycle] area={self.area}, flag={val}, tp_size={tp_num}")
|
||||
|
||||
if val == float(tp_num):
|
||||
self.sync_flag.buffer_wrapper *= 0.0
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
class MmItemMemoryPool:
|
||||
def __init__(self, memory_size, recycle_interval, base_gpu_id):
|
||||
self.memory_pool = torch.empty(
|
||||
memory_size, dtype=torch.int8, device=f"cuda:{base_gpu_id}"
|
||||
).contiguous()
|
||||
storage = self.memory_pool.untyped_storage()
|
||||
self._pool_ipc_handle = storage._share_cuda_()
|
||||
self._pool_device_index = self.memory_pool.device.index
|
||||
|
||||
self.sync_flag_list = []
|
||||
|
||||
init_chunk = MmItemMemoryChunk((0, memory_size), self.pop_sync_buffer())
|
||||
self.available_chunks = [init_chunk]
|
||||
self.occupied_chunks = []
|
||||
|
||||
self._lock = threading.Lock()
|
||||
self._pool_full_warned = False
|
||||
|
||||
self._recycle_interval = recycle_interval
|
||||
self._stop_recycler = False
|
||||
self._recycle_thread = threading.Thread(
|
||||
target=self._recycle_loop, name="MmItemMemoryPoolRecycler", daemon=True
|
||||
)
|
||||
self._recycle_thread.start()
|
||||
|
||||
logger.debug(
|
||||
f"[MmItemMemoryPool] init: memory_size={memory_size}, "
|
||||
f"recycle_interval={recycle_interval}s"
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
self._stop_recycler = True
|
||||
if self._recycle_thread.is_alive():
|
||||
self._recycle_thread.join(timeout=1.0)
|
||||
|
||||
def _recycle_loop(self):
|
||||
while not self._stop_recycler:
|
||||
try:
|
||||
with self._lock:
|
||||
self.recycle_chunks()
|
||||
self.merge_chunks()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[MmItemMemoryPool] recycle loop error: {e}", exc_info=True
|
||||
)
|
||||
|
||||
time.sleep(self._recycle_interval)
|
||||
|
||||
def clear_sync_flag_list(self):
|
||||
# call each chunk's __del__
|
||||
self.sync_flag_list.clear()
|
||||
|
||||
def pop_sync_buffer(self):
|
||||
if len(self.sync_flag_list) == 0:
|
||||
try:
|
||||
new_sync_buffer = ShmSyncBuffer()
|
||||
return new_sync_buffer
|
||||
except:
|
||||
logger.info("allocate shm buffer failed")
|
||||
raise RuntimeError
|
||||
else:
|
||||
return self.sync_flag_list.pop()
|
||||
|
||||
def push_sync_buffer(self, sync_buffer):
|
||||
self.sync_flag_list.append(sync_buffer)
|
||||
|
||||
def get_available_chunk(self, src_tensor: torch.Tensor) -> MmItemMemoryChunk:
|
||||
# find currently available_chunks contain a available chunk or not
|
||||
# if not, return None
|
||||
src_tensor_size = src_tensor.numel() * src_tensor.element_size()
|
||||
min_size = self.memory_pool.numel() * self.memory_pool.element_size() + 1
|
||||
selected_chunk = None
|
||||
for chunk in self.available_chunks:
|
||||
if chunk.mem_size >= src_tensor_size:
|
||||
if chunk.mem_size < min_size:
|
||||
min_size = chunk.mem_size
|
||||
selected_chunk = chunk
|
||||
|
||||
if selected_chunk:
|
||||
occupied_chunk_area = (
|
||||
selected_chunk.start,
|
||||
selected_chunk.start + src_tensor_size,
|
||||
)
|
||||
occupied_chunk_sync_flag = selected_chunk.sync_flag
|
||||
new_occupied_chunk = MmItemMemoryChunk(
|
||||
occupied_chunk_area, occupied_chunk_sync_flag
|
||||
)
|
||||
|
||||
self.occupied_chunks.append(new_occupied_chunk)
|
||||
self.available_chunks.remove(selected_chunk)
|
||||
|
||||
available_split_chunk_area = (new_occupied_chunk.end, selected_chunk.end)
|
||||
# add a new chunk
|
||||
if available_split_chunk_area[0] != available_split_chunk_area[1]:
|
||||
split_available_chunk = MmItemMemoryChunk(
|
||||
available_split_chunk_area, self.pop_sync_buffer()
|
||||
)
|
||||
self.available_chunks.append(split_available_chunk)
|
||||
|
||||
return new_occupied_chunk
|
||||
|
||||
return None
|
||||
|
||||
def return_a_slice_tensor_with_flag(self, src_tensor: torch.Tensor):
|
||||
with self._lock:
|
||||
available_chunk = self.get_available_chunk(src_tensor)
|
||||
if available_chunk is not None:
|
||||
return (
|
||||
available_chunk.sync_flag.meta_data,
|
||||
self.memory_pool[available_chunk.start : available_chunk.end],
|
||||
available_chunk.start,
|
||||
)
|
||||
self._warn_pool_full_once(src_tensor)
|
||||
return None, None, None
|
||||
|
||||
def _warn_pool_full_once(self, src_tensor: torch.Tensor):
|
||||
if self._pool_full_warned:
|
||||
return
|
||||
self._pool_full_warned = True
|
||||
pool_mb = (
|
||||
self.memory_pool.numel() * self.memory_pool.element_size() / (1024 * 1024)
|
||||
)
|
||||
need_mb = src_tensor.numel() * src_tensor.element_size() / (1024 * 1024)
|
||||
logger.warning(
|
||||
"MmItemMemoryPool has no free chunk large enough for a %.2f MiB tensor "
|
||||
"(pool size: %.2f MiB); falling back to non-IPC transport. "
|
||||
"Consider increasing SGLANG_MM_FEATURE_CACHE_MB.",
|
||||
need_mb,
|
||||
pool_mb,
|
||||
)
|
||||
|
||||
def recycle_chunks(self):
|
||||
|
||||
new_occupied_chunks = []
|
||||
for chunk in self.occupied_chunks:
|
||||
if chunk.try_to_recycle():
|
||||
self.available_chunks.append(chunk)
|
||||
else:
|
||||
new_occupied_chunks.append(chunk)
|
||||
self.occupied_chunks = new_occupied_chunks
|
||||
|
||||
def merge_chunks(self):
|
||||
# merge_all_available_chunks
|
||||
merged_chunks = []
|
||||
for chunk in sorted(self.available_chunks, key=lambda x: x.start):
|
||||
if len(merged_chunks) == 0:
|
||||
merged_chunks.append(chunk)
|
||||
else:
|
||||
if chunk.start == merged_chunks[-1].end:
|
||||
to_merge_chunk = merged_chunks.pop()
|
||||
to_merge_chunk_sync = to_merge_chunk.sync_flag
|
||||
merged_chunk_area = (to_merge_chunk.start, chunk.end)
|
||||
merged_chunks.append(
|
||||
MmItemMemoryChunk(merged_chunk_area, to_merge_chunk_sync)
|
||||
)
|
||||
self.push_sync_buffer(chunk.sync_flag)
|
||||
else:
|
||||
merged_chunks.append(chunk)
|
||||
|
||||
self.available_chunks = merged_chunks
|
||||
|
||||
|
||||
class CudaIpcTensorTransportProxy:
|
||||
"""
|
||||
A torch.tensor's proxy used to do inter-process data-sharing
|
||||
including:
|
||||
|
||||
torch.tensor(on gpu)'s cuda-ipc-hande infos
|
||||
a shm sync buffer's meta data which is used to sync between different process
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data: torch.Tensor,
|
||||
info_data: torch.Tensor,
|
||||
sync_buffer_meta,
|
||||
pool_ipc_handle=None,
|
||||
pool_byte_offset: int = 0,
|
||||
pool_device_index: int = 0,
|
||||
):
|
||||
|
||||
if (not isinstance(data, torch.Tensor)) or (
|
||||
not isinstance(info_data, torch.Tensor)
|
||||
):
|
||||
raise TypeError(
|
||||
f"Input 'data' must be a torch.Tensor, but got {type(data)}"
|
||||
)
|
||||
|
||||
if pool_ipc_handle is not None:
|
||||
self.proxy_state = {
|
||||
"ipc_extra": {
|
||||
"pool_handle": pool_ipc_handle,
|
||||
"pool_byte_offset": pool_byte_offset,
|
||||
"pool_device_index": pool_device_index,
|
||||
"shape": data.shape,
|
||||
"dtype": data.dtype,
|
||||
"stride": data.stride(),
|
||||
"storage_offset": 0,
|
||||
"nbytes": data.numel() * data.element_size(),
|
||||
"recons_shape": info_data.shape,
|
||||
"recons_dtype": info_data.dtype,
|
||||
},
|
||||
"tensor_data": None,
|
||||
}
|
||||
else:
|
||||
self.proxy_state = self.get_proxy_state(data, info_data)
|
||||
self.reconstruct_tensor = None
|
||||
self.sync_data_meta = sync_buffer_meta
|
||||
self.sync_buffer = None
|
||||
self._consumer_acknowledged = False
|
||||
|
||||
@property
|
||||
def get_sync_flag(self):
|
||||
if not self.sync_buffer:
|
||||
shm_name = self.sync_data_meta["handle"]
|
||||
self.sync_buffer = shared_memory.SharedMemory(name=shm_name)
|
||||
|
||||
shape = self.sync_data_meta["shape"]
|
||||
dtype = self.sync_data_meta["dtype"]
|
||||
return np.ndarray(shape, dtype=dtype, buffer=self.sync_buffer.buf)
|
||||
|
||||
def close_shm(self):
|
||||
self.sync_buffer.close()
|
||||
self.sync_buffer = None
|
||||
|
||||
def get_proxy_state(self, data, info_data):
|
||||
# acquire all serialize metadata from _metadata
|
||||
state = {}
|
||||
|
||||
try:
|
||||
storage = data.untyped_storage()
|
||||
handle = storage._share_cuda_()
|
||||
|
||||
state["ipc_extra"] = {
|
||||
"handle": handle,
|
||||
"shape": data.shape,
|
||||
"dtype": data.dtype,
|
||||
"stride": data.stride(),
|
||||
"device_index": data.device.index,
|
||||
"storage_offset": data.storage_offset(),
|
||||
"recons_shape": info_data.shape,
|
||||
"recons_dtype": info_data.dtype,
|
||||
}
|
||||
state["tensor_data"] = None
|
||||
except Exception:
|
||||
# Failed to get CUDA IPC handle (possibly tp). Falling back to default transport.
|
||||
state["ipc_extra"] = None
|
||||
state["tensor_data"] = data
|
||||
|
||||
return state
|
||||
|
||||
def _reconstruct_from_ipc_extra(
|
||||
self, ipc_extra, *, use_cache: bool, rebuild_device_idx: int
|
||||
):
|
||||
shape = ipc_extra["shape"]
|
||||
dtype = ipc_extra["dtype"]
|
||||
stride = ipc_extra["stride"]
|
||||
# Redirect handle[0] to the consumer's device so _new_shared_cuda's
|
||||
# CUDAGuard stays there; peer access handles the cross-GPU open.
|
||||
pool_handle = ipc_extra["pool_handle"]
|
||||
redirected_handle = (rebuild_device_idx,) + tuple(pool_handle)[1:]
|
||||
target_device = torch.device(f"cuda:{rebuild_device_idx}")
|
||||
cache_key = _normalize_pool_cache_key(pool_handle, rebuild_device_idx)
|
||||
|
||||
with torch.cuda.device(target_device):
|
||||
if use_cache:
|
||||
storage = _pool_handle_cache_get_or_open(cache_key, redirected_handle)
|
||||
storage_to_cache = None
|
||||
else:
|
||||
storage = _open_pooled_storage_uncached(redirected_handle)
|
||||
storage_to_cache = storage
|
||||
slice_storage = storage[
|
||||
ipc_extra["pool_byte_offset"] : ipc_extra["pool_byte_offset"]
|
||||
+ ipc_extra["nbytes"]
|
||||
]
|
||||
slice_tensor = torch.empty(0, dtype=dtype, device=target_device).set_(
|
||||
slice_storage,
|
||||
storage_offset=ipc_extra["storage_offset"],
|
||||
size=shape,
|
||||
stride=stride,
|
||||
)
|
||||
|
||||
return slice_tensor, target_device, cache_key, storage_to_cache
|
||||
|
||||
def _acknowledge_consumption(self, consumer_count: int = 1):
|
||||
"""Mark this IPC feature as consumed without necessarily copying it.
|
||||
|
||||
A normal TP execution reconstructs a feature once per rank, so each
|
||||
consumer contributes one acknowledgement. Encoder-DP can instead
|
||||
route a feature to exactly one rank; that rank acknowledges all TP
|
||||
consumers after its copy completes. Keeping this acknowledgement
|
||||
idempotent is important for chunked-prefill cache hits, where the same
|
||||
proxy may be visited more than once.
|
||||
"""
|
||||
if getattr(self, "_consumer_acknowledged", False):
|
||||
return
|
||||
if consumer_count <= 0:
|
||||
raise ValueError("consumer_count must be positive")
|
||||
if self.sync_data_meta is not None:
|
||||
open(SHM_LOCK_FILE, "a").close()
|
||||
# Keep the counter update atomic across scheduler processes.
|
||||
with open(SHM_LOCK_FILE, "w+") as f:
|
||||
fcntl.flock(f, fcntl.LOCK_EX)
|
||||
sync_flag = self.get_sync_flag
|
||||
sync_flag += consumer_count
|
||||
fcntl.flock(f, fcntl.LOCK_UN)
|
||||
self.close_shm()
|
||||
self._consumer_acknowledged = True
|
||||
|
||||
def acknowledge_consumption(self, consumer_count: int = 1):
|
||||
"""Release an IPC-pool slice when a cache hit needs no tensor copy."""
|
||||
self._acknowledge_consumption(consumer_count)
|
||||
|
||||
def _copy_slice_tensor_to_target(
|
||||
self,
|
||||
slice_tensor: torch.Tensor,
|
||||
rebuild_device: torch.device,
|
||||
recons_shape,
|
||||
recons_dtype,
|
||||
consumer_count: int,
|
||||
):
|
||||
with torch.cuda.device(rebuild_device):
|
||||
reconstructed_tensor = torch.empty(
|
||||
recons_shape, dtype=recons_dtype, device=rebuild_device
|
||||
).contiguous()
|
||||
reconstructed_tensor.view(torch.int8).view(-1).copy_(slice_tensor)
|
||||
self._acknowledge_consumption(consumer_count)
|
||||
|
||||
return reconstructed_tensor
|
||||
|
||||
def reconstruct_on_target_device(self, rebuild_device_idx, consumer_count: int = 1):
|
||||
rebuild_device = torch.device(f"cuda:{rebuild_device_idx}")
|
||||
if (
|
||||
isinstance(self.reconstruct_tensor, torch.Tensor)
|
||||
and self.reconstruct_tensor.device == rebuild_device
|
||||
):
|
||||
return self.reconstruct_tensor
|
||||
|
||||
if self.proxy_state["ipc_extra"]:
|
||||
ipc_extra = self.proxy_state["ipc_extra"]
|
||||
recons_shape = ipc_extra["recons_shape"]
|
||||
recons_dtype = ipc_extra["recons_dtype"]
|
||||
|
||||
if "pool_handle" in ipc_extra:
|
||||
try:
|
||||
(
|
||||
slice_tensor,
|
||||
_target_device,
|
||||
cache_key,
|
||||
storage_to_cache,
|
||||
) = self._reconstruct_from_ipc_extra(
|
||||
ipc_extra,
|
||||
use_cache=True,
|
||||
rebuild_device_idx=rebuild_device_idx,
|
||||
)
|
||||
except Exception as e:
|
||||
cache_key = _normalize_pool_cache_key(
|
||||
ipc_extra["pool_handle"], rebuild_device_idx
|
||||
)
|
||||
logger.info(
|
||||
"Failed to deserialize from cached pooled CUDA IPC handle (%s). "
|
||||
"Invalidating cache entry and retrying uncached.",
|
||||
e,
|
||||
)
|
||||
_pool_handle_cache_invalidate(cache_key)
|
||||
(
|
||||
slice_tensor,
|
||||
_target_device,
|
||||
_cache_key,
|
||||
storage_to_cache,
|
||||
) = self._reconstruct_from_ipc_extra(
|
||||
ipc_extra,
|
||||
use_cache=False,
|
||||
rebuild_device_idx=rebuild_device_idx,
|
||||
)
|
||||
if storage_to_cache is not None:
|
||||
_pool_handle_cache_set(cache_key, storage_to_cache)
|
||||
else:
|
||||
# Non-pooled path: redirect handle[0] the same way as the pooled path.
|
||||
try:
|
||||
original_handle = ipc_extra["handle"]
|
||||
redirected_handle = (rebuild_device_idx,) + tuple(original_handle)[
|
||||
1:
|
||||
]
|
||||
target_device = torch.device(f"cuda:{rebuild_device_idx}")
|
||||
with torch.cuda.device(target_device):
|
||||
storage = torch.UntypedStorage._new_shared_cuda(
|
||||
*redirected_handle
|
||||
)
|
||||
slice_tensor = torch.empty(
|
||||
0, dtype=ipc_extra["dtype"], device=target_device
|
||||
).set_(
|
||||
storage,
|
||||
storage_offset=ipc_extra["storage_offset"],
|
||||
size=ipc_extra["shape"],
|
||||
stride=ipc_extra["stride"],
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info("Failed to deserialize from CUDA IPC handle (%s).", e)
|
||||
raise
|
||||
|
||||
reconstructed_tensor = self._copy_slice_tensor_to_target(
|
||||
slice_tensor,
|
||||
rebuild_device,
|
||||
recons_shape,
|
||||
recons_dtype,
|
||||
consumer_count,
|
||||
)
|
||||
elif isinstance(self.proxy_state["tensor_data"], torch.Tensor):
|
||||
reconstructed_tensor = self.proxy_state["tensor_data"].to(
|
||||
rebuild_device, non_blocking=True
|
||||
)
|
||||
else:
|
||||
raise TypeError("invalid proxy_state")
|
||||
|
||||
self.reconstruct_tensor = reconstructed_tensor
|
||||
return self.reconstruct_tensor
|
||||
__all__ = [
|
||||
"DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY",
|
||||
"MM_FEATURE_CACHE_SIZE",
|
||||
"MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL",
|
||||
"CudaIpcTensorTransportProxy",
|
||||
"MmItemMemoryPool",
|
||||
"_pool_handle_cache_clear",
|
||||
"get_mm_feature_pool_size_per_worker",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user