diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 112378567..c273912b0 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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: diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 31e78d992..cd81cf5d4 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -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: diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index d1495c9ce..9c9817c42 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -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: diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 22a8dbe62..afd7659f4 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -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 diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index d92e32e79..fba2fb06d 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -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): diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py index 7c84069f5..023f1fc13 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k25.py +++ b/python/sglang/srt/multimodal/processors/kimi_k25.py @@ -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, ) diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py index 7b1bc6f59..6b1170757 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k3.py +++ b/python/sglang/srt/multimodal/processors/kimi_k3.py @@ -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]: diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index 6ad157b26..cac51f22b 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -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 + ) diff --git a/python/sglang/srt/multimodal/transport/__init__.py b/python/sglang/srt/multimodal/transport/__init__.py new file mode 100644 index 000000000..659092775 --- /dev/null +++ b/python/sglang/srt/multimodal/transport/__init__.py @@ -0,0 +1 @@ +"""GPU transports for multimodal feature tensors.""" diff --git a/python/sglang/srt/multimodal/transport/cuda_ipc.py b/python/sglang/srt/multimodal/transport/cuda_ipc.py new file mode 100644 index 000000000..d1d30af04 --- /dev/null +++ b/python/sglang/srt/multimodal/transport/cuda_ipc.py @@ -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 diff --git a/python/sglang/srt/multimodal/transport/memory_pool.py b/python/sglang/srt/multimodal/transport/memory_pool.py new file mode 100644 index 000000000..c9df875b9 --- /dev/null +++ b/python/sglang/srt/multimodal/transport/memory_pool.py @@ -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() diff --git a/python/sglang/srt/utils/cuda_ipc_transport_utils.py b/python/sglang/srt/utils/cuda_ipc_transport_utils.py index acc30e1ed..4108f1e68 100644 --- a/python/sglang/srt/utils/cuda_ipc_transport_utils.py +++ b/python/sglang/srt/utils/cuda_ipc_transport_utils.py @@ -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", +] diff --git a/test/registered/unit/managers/test_mm_process_config.py b/test/registered/unit/managers/test_mm_process_config.py index 36f198861..7f9c8b18c 100644 --- a/test/registered/unit/managers/test_mm_process_config.py +++ b/test/registered/unit/managers/test_mm_process_config.py @@ -174,6 +174,7 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase): mm_io_worker_num=0, tokenizer_worker_num=1, base_gpu_id=2, + tp_size=8, ) @staticmethod @@ -263,6 +264,93 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase): memory_pool.assert_not_called() +class TestStreamOrderedMmFeaturePool(CustomTestCase): + def test_consumer_slot_uses_global_tp_rank(self): + from sglang.srt.multimodal.transport.memory_pool import resolve_consumer_rank + + parallel = SimpleNamespace(tp_rank=6, attn_tp_rank=2) + with patch("sglang.srt.runtime_context.get_parallel", return_value=parallel): + self.assertEqual(resolve_consumer_rank(8), 6) + + def test_complete_group_acknowledges_each_consumer_slot(self): + from sglang.srt.multimodal.transport import memory_pool + + consumer = memory_pool.StreamOrderedPoolConsumerMixin() + consumer._init_stream_ordered_consumer( + ready_byte_offset=64, + ack_byte_offset=68, + generation=3, + total_consumer_count=4, + transport_name="test", + ) + with patch.object(memory_pool, "stream_write_value32") as write: + consumer._acknowledge_on_stream(1000, 0, consumer_count=4) + + self.assertEqual( + [call.args[1] for call in write.call_args_list], + [1068, 1072, 1076, 1080], + ) + self.assertEqual([call.args[2] for call in write.call_args_list], [3] * 4) + + def test_reused_pool_slot_gets_new_generation(self): + from sglang.srt.multimodal.transport.memory_pool import ( + StreamOrderedMmFeaturePool, + ) + + pool = object.__new__(StreamOrderedMmFeaturePool) + pool._available_ranges = [(256, 4096)] + pool._available_slots = [0] + pool._slot_generations = [0] + pool._occupied = {} + pool.control_words_per_slot = 2 + pool.transport_name = "test" + + first = pool._allocate_locked(512) + pool._release_locked(first) + pool._merge_ranges_locked() + second = pool._allocate_locked(512) + + self.assertEqual(first.generation, 1) + self.assertEqual(second.generation, 2) + + def test_pool_rejects_duplicate_release(self): + from sglang.srt.multimodal.transport.memory_pool import ( + StreamOrderedMmFeaturePool, + ) + + pool = object.__new__(StreamOrderedMmFeaturePool) + pool._available_ranges = [(256, 4096)] + pool._available_slots = [0] + pool._slot_generations = [0] + pool._occupied = {} + pool.control_words_per_slot = 2 + pool.transport_name = "test" + + lease = pool._allocate_locked(512) + pool._release_locked(lease) + + with self.assertRaisesRegex(RuntimeError, "inactive test pool lease"): + pool._release_locked(lease) + + def test_pool_shutdown_wakes_recycler_before_returning(self): + from sglang.srt.multimodal.transport.memory_pool import ( + StreamOrderedMmFeaturePool, + ) + + pool = object.__new__(StreamOrderedMmFeaturePool) + pool._recycler_stop_event = threading.Event() + pool._recycle_thread = threading.Thread( + target=pool._recycler_stop_event.wait, + args=(60,), + daemon=True, + ) + pool._recycle_thread.start() + + pool.shutdown() + + self.assertFalse(pool._recycle_thread.is_alive()) + + class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase): @staticmethod def _processor(enabled): diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index feec9e90a..c8d2694de 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/test/registered/unit/models/test_kimi_k25.py @@ -39,11 +39,11 @@ from sglang.srt.multimodal.processors.kimi_k25 import ( _resize_bicubic_if_needed, _resize_images_by_source_shape, ) -from sglang.srt.runtime_context import get_context, get_parallel -from sglang.srt.utils.cuda_ipc_transport_utils import ( +from sglang.srt.multimodal.transport.cuda_ipc import ( DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, CudaIpcTensorTransportProxy, ) +from sglang.srt.runtime_context import get_context, get_parallel from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") diff --git a/test/registered/unit/models/test_qwen3_vl_feature_materialization.py b/test/registered/unit/models/test_qwen3_vl_feature_materialization.py index bb9eaf0cd..091cfda41 100644 --- a/test/registered/unit/models/test_qwen3_vl_feature_materialization.py +++ b/test/registered/unit/models/test_qwen3_vl_feature_materialization.py @@ -2,10 +2,16 @@ import unittest from types import SimpleNamespace +from unittest.mock import Mock, patch import torch +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration +from sglang.srt.multimodal.processors.qwen_vl import QwenVLImageProcessor +from sglang.srt.multimodal.transport.cuda_ipc import ( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -27,9 +33,61 @@ class _RecordingVisual: class TestQwen3VLFeatureMaterialization(CustomTestCase): + @staticmethod + def _model(visual, *, use_data_parallel): + model = Qwen3VLForConditionalGeneration.__new__(Qwen3VLForConditionalGeneration) + torch.nn.Module.__init__(model) + model.visual = visual + model.use_data_parallel = use_data_parallel + return model + + def test_processor_defers_gpu_transport_for_encoder_dp(self): + for transport in ("cuda_ipc", "cuda_vmm"): + with self.subTest(transport=transport): + processor = QwenVLImageProcessor.__new__(QwenVLImageProcessor) + processor.mm_feature_transport = transport + processor.server_args = SimpleNamespace(mm_enable_dp_encoder=True) + processor.model_type = "qwen3_vl" + items = [ + MultimodalDataItem(modality=Modality.IMAGE), + MultimodalDataItem(modality=Modality.VIDEO), + MultimodalDataItem(modality=Modality.AUDIO), + ] + + processor._mark_dp_encoder_features_for_deferred_reconstruction(items) + + self.assertTrue( + items[0].model_specific_data[ + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY + ] + ) + self.assertTrue( + items[1].model_specific_data[ + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY + ] + ) + self.assertNotIn( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, + items[2].model_specific_data, + ) + + def test_processor_does_not_defer_cpu_transport(self): + processor = QwenVLImageProcessor.__new__(QwenVLImageProcessor) + processor.mm_feature_transport = "cpu" + processor.server_args = SimpleNamespace(mm_enable_dp_encoder=True) + processor.model_type = "qwen3_vl" + item = MultimodalDataItem(modality=Modality.IMAGE) + + processor._mark_dp_encoder_features_for_deferred_reconstruction([item]) + + self.assertNotIn( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, + item.model_specific_data, + ) + def test_image_features_are_packed_on_the_visual_device(self): visual = _RecordingVisual() - model = SimpleNamespace(visual=visual, use_data_parallel=False) + model = self._model(visual, use_data_parallel=False) items = [ SimpleNamespace( feature=torch.ones(2, 3), @@ -40,7 +98,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase): image_grid_thw=torch.tensor([[1, 1, 1]]), ), ] - output = Qwen3VLForConditionalGeneration.get_image_feature(model, items) + output = model.get_image_feature(items) self.assertIs(visual.pixel_values, output) self.assertEqual(output.shape, (3, 3)) @@ -49,7 +107,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase): def test_video_features_are_packed_on_the_visual_device(self): visual = _RecordingVisual() - model = SimpleNamespace(visual=visual, use_data_parallel=False) + model = self._model(visual, use_data_parallel=False) items = [ SimpleNamespace( feature=torch.ones(3, 4), @@ -60,7 +118,7 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase): video_grid_thw=torch.tensor([[1, 1, 2]]), ), ] - output = Qwen3VLForConditionalGeneration.get_video_feature(model, items) + output = model.get_video_feature(items) self.assertIs(visual.pixel_values, output) self.assertEqual(output.shape, (5, 4)) @@ -70,6 +128,57 @@ class TestQwen3VLFeatureMaterialization(CustomTestCase): torch.equal(visual.grid_thw, torch.tensor([[1, 1, 3], [1, 1, 2]])) ) + def test_encoder_dp_materializes_only_locally_assigned_visual_items(self): + visual = SimpleNamespace(device=torch.device("cuda:0"), dtype=torch.bfloat16) + model = self._model(visual, use_data_parallel=True) + + for modality, feature_method, grid_attribute in ( + ("image", model.get_image_feature, "image_grid_thw"), + ("video", model.get_video_feature, "video_grid_thw"), + ): + with self.subTest(modality=modality): + items = [ + SimpleNamespace( + feature=torch.ones(2, 3), + reconstruct=Mock(), + **{grid_attribute: torch.tensor([[1, 1, 2]])}, + ), + SimpleNamespace( + feature=torch.ones(1, 3), + reconstruct=Mock(), + **{grid_attribute: torch.tensor([[1, 1, 1]])}, + ), + ] + local_features = object() + encoded = object() + + def run_dp(_visual, pixel_values, grid_thw, **kwargs): + self.assertIsNone(pixel_values) + self.assertEqual(grid_thw, [[1, 1, 2], [1, 1, 1]]) + self.assertIs( + kwargs["load_local_pixel_values"]([1]), local_features + ) + return encoded + + with patch( + "sglang.srt.models.qwen3_vl.run_dp_sharded_mrope_vision_model", + side_effect=run_dp, + ), patch( + "sglang.srt.models.qwen3_vl.materialize_multimodal_features", + return_value=local_features, + ) as materialize, patch( + "sglang.srt.models.qwen3_vl.get_parallel", + return_value=SimpleNamespace(tp_size=8), + ): + output = feature_method(items) + + self.assertIs(output, encoded) + items[0].reconstruct.assert_not_called() + items[1].reconstruct.assert_called_once_with(0, ipc_consumer_count=8) + materialize.assert_called_once_with( + [items[1].feature], device=visual.device, dtype=visual.dtype + ) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py b/test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py index 535314c76..b09b3bb54 100644 --- a/test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py +++ b/test/registered/unit/multimodal/test_cuda_ipc_pool_budget.py @@ -2,7 +2,7 @@ import unittest -from sglang.srt.utils.cuda_ipc_transport_utils import ( +from sglang.srt.multimodal.transport.cuda_ipc import ( get_mm_feature_pool_size_per_worker, ) from sglang.test.ci.ci_register import register_cpu_ci diff --git a/test/registered/unit/multimodal/test_cuda_ipc_transport.py b/test/registered/unit/multimodal/test_cuda_ipc_transport.py index aa1723d6c..9d64ac8ab 100644 --- a/test/registered/unit/multimodal/test_cuda_ipc_transport.py +++ b/test/registered/unit/multimodal/test_cuda_ipc_transport.py @@ -8,11 +8,13 @@ CPU-only policy tests intentionally cannot exercise this cross-process handle. import gc import multiprocessing as mp import queue +import time import unittest +from unittest.mock import Mock, patch import torch -from sglang.srt.utils.cuda_ipc_transport_utils import ( +from sglang.srt.multimodal.transport.cuda_ipc import ( CudaIpcTensorTransportProxy, MmItemMemoryPool, _pool_handle_cache_clear, @@ -25,39 +27,40 @@ register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large") def _produce_pooled_tensor(proxy_queue, consumer_done, result_queue): """Create a tokenizer-worker-like CUDA IPC pool in a spawned producer.""" - pool = source = pool_slice = proxy = None + pool = source = proxy = None try: torch.cuda.set_device(0) pool = MmItemMemoryPool( memory_size=1 << 20, - recycle_interval=60, + recycle_interval=0.01, base_gpu_id=0, + consumer_count=1, ) source = torch.arange(35, dtype=torch.float32, device="cuda").reshape(5, 7) - expected = source.cpu().tolist() - sync_meta, pool_slice, byte_offset = pool.return_a_slice_tensor_with_flag( - source + expected = torch.arange(35, dtype=torch.float32).reshape(5, 7).tolist() + proxy = pool.wrap_tensor( + source, + use_pool_handle_cache=True, ) - if pool_slice is None: + if proxy is None: raise RuntimeError("test tensor did not fit in the CUDA IPC pool") - pool_slice.copy_(source.view(torch.int8).view(-1), non_blocking=True) - torch.cuda.synchronize() - proxy = CudaIpcTensorTransportProxy( - data=pool_slice, - info_data=source, - sync_buffer_meta=sync_meta, - pool_ipc_handle=pool._pool_ipc_handle, - pool_byte_offset=byte_offset, - pool_device_index=pool._pool_device_index, - ) + # Intentionally do not synchronize the producer. The consumer stream + # wait must order its copy after the producer-ready write. proxy_queue.put((proxy, expected)) if not consumer_done.wait(timeout=60): raise TimeoutError("consumer did not release the CUDA IPC tensor") + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + if pool.active_lease_count == 0: + break + time.sleep(0.01) + else: + raise TimeoutError("pool did not observe the stream-ordered consumer ack") except Exception as exc: # pragma: no cover - returned to the parent result_queue.put(("error", repr(exc))) return finally: - del proxy, pool_slice, source + del proxy, source if pool is not None: pool.shutdown() del pool @@ -119,6 +122,21 @@ class TestCudaIpcTransport(CustomTestCase): producer.join(timeout=10) self.assertEqual(producer.exitcode, 0) + def test_uncached_mapping_waits_before_proxy_release(self): + proxy = object.__new__(CudaIpcTensorTransportProxy) + proxy.proxy_state = {"ipc_extra": {"use_pool_handle_cache": False}} + proxy._pool_storage = None + stream = Mock() + + with patch( + "sglang.srt.multimodal.transport.cuda_ipc.torch.cuda.current_stream", + return_value=stream, + ): + proxy._retain_storage_until_stream_completes(object(), 0) + + stream.synchronize.assert_called_once_with() + self.assertIsNone(proxy._pool_storage) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/unit/test_global_config_read_ratchet.py b/test/registered/unit/test_global_config_read_ratchet.py index e6e01729f..474cfe77d 100644 --- a/test/registered/unit/test_global_config_read_ratchet.py +++ b/test/registered/unit/test_global_config_read_ratchet.py @@ -29,8 +29,6 @@ Where the remaining reads live (``runtime_context.py``, exempt by module): the short circuit is the point: with PP off the group is never touched, which is what lets the ``Indexer`` be constructed before distributed init. The live property would demand the group either way. - - ``cuda_ipc_transport_utils.tp_size`` runs in the tokenizer process, which has - no groups at all (the call site already guards for "not published yet"). - ``dp_attention.attn_cp_size`` / ``moe_dp_size``: the configuration the predicate detects (``attn_cp_size > moe_dp_size``) is the one where ``initialize_model_parallel`` aliases ``_MOE_DP`` to ``_ATTN_CP``, so the live @@ -102,14 +100,10 @@ _CONFIGURED_SIZE_CALL_SITES = { "the same dict already carries the live moe_dp_size under 'dp'; this entry " "is the configured intent" ), - ("srt/utils/cuda_ipc_transport_utils.py", "configured_tp_size"): ( - "runs in the tokenizer process, which has no parallel groups at all" - ), ("srt/models/kimi_k25.py", "configured_tp_size"): ( - "the IPC refcount has to name the same number the recycler waits on, and " - "that waiter (MmItemMemoryPool.try_to_recycle) reads configured_tp_size() " - "because it runs in the tokenizer process; a refcount taken from the live " - "attention subgroup would strand items in the bounded pool" + "the IPC refcount must match the configured TP consumer count captured " + "when the tokenizer creates MmItemMemoryPool; a live attention subgroup " + "size could strand leases in the bounded pool" ), ("srt/models/kimi_k3.py", "configured_tp_size"): ( "same as kimi_k25: the IPC refcount must agree with the recycler's waiter"