diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 763c96422..9a83da8a1 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1191,26 +1191,32 @@ class Engine(EngineScoreMixin, EngineBase): tokenizer_manager = MultiTokenizerRouter(server_args, port_args) template_manager = None - # Wait for the model to finish loading - scheduler_init_result.wait_for_ready() + startup_complete = False + try: + # Wait for the model to finish loading + scheduler_init_result.wait_for_ready() - cls._set_startup_time(tokenizer_manager, scheduler_init_result, startup_tic) + cls._set_startup_time(tokenizer_manager, scheduler_init_result, startup_tic) - # Get back some info from scheduler to tokenizer_manager - tokenizer_manager.max_req_input_len = scheduler_init_result.scheduler_infos[0][ - "max_req_input_len" - ] + # Get back some info from scheduler to tokenizer_manager + tokenizer_manager.max_req_input_len = scheduler_init_result.scheduler_infos[ + 0 + ]["max_req_input_len"] - # Set up subprocess liveness watchdog to detect crashes - # Note: RayEngine returns scheduler_procs=None as it uses Ray actors instead of mp.Process - processes = list(scheduler_procs or []) - names = [f"scheduler_{i}" for i in range(len(processes))] - processes.extend(detoken_procs) - names.extend(detoken_names) - subprocess_watchdog = SubprocessWatchdog( - processes=processes, process_names=names - ) - subprocess_watchdog.start() + # Set up subprocess liveness watchdog to detect crashes + # Note: RayEngine returns scheduler_procs=None as it uses Ray actors instead of mp.Process + processes = list(scheduler_procs or []) + names = [f"scheduler_{i}" for i in range(len(processes))] + processes.extend(detoken_procs) + names.extend(detoken_names) + subprocess_watchdog = SubprocessWatchdog( + processes=processes, process_names=names + ) + subprocess_watchdog.start() + startup_complete = True + finally: + if not startup_complete and isinstance(tokenizer_manager, TokenizerManager): + tokenizer_manager.cuda_vmm_feature_transport.shutdown() return ( tokenizer_manager, @@ -1225,26 +1231,30 @@ class Engine(EngineScoreMixin, EngineBase): """Shutdown the engine; block until the scheduler subprocess releases its GPU context so the caller can immediately reallocate on the same device.""" - if ( - self.tokenizer_manager is not None - and self.tokenizer_manager._subprocess_watchdog is not None - ): - self.tokenizer_manager._subprocess_watchdog.stop() + try: + if ( + self.tokenizer_manager is not None + and self.tokenizer_manager._subprocess_watchdog is not None + ): + self.tokenizer_manager._subprocess_watchdog.stop() - send_to_rpc = getattr(self, "send_to_rpc", None) - if send_to_rpc is not None: - send_to_rpc.close(linger=0) - self.send_to_rpc = None + send_to_rpc = getattr(self, "send_to_rpc", None) + if send_to_rpc is not None: + send_to_rpc.close(linger=0) + self.send_to_rpc = None - # Gracefully stop weight cache daemons *before* the blanket - # kill_process_tree below, so their SIGTERM handlers can unlink the - # .sock/.ready files instead of being SIGKILLed and leaving stale state. - daemon_procs = getattr(self, "_weight_cache_daemon_procs", None) - if daemon_procs: - self._terminate_weight_cache_daemons(daemon_procs) - self._weight_cache_daemon_procs = [] + # Gracefully stop weight cache daemons *before* the blanket + # kill_process_tree below, so their SIGTERM handlers can unlink the + # .sock/.ready files instead of being SIGKILLed and leaving stale state. + daemon_procs = getattr(self, "_weight_cache_daemon_procs", None) + if daemon_procs: + self._terminate_weight_cache_daemons(daemon_procs) + self._weight_cache_daemon_procs = [] - kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60) + kill_process_tree(os.getpid(), include_parent=False, wait_timeout=60) + finally: + if isinstance(self.tokenizer_manager, TokenizerManager): + self.tokenizer_manager.cuda_vmm_feature_transport.shutdown() def __enter__(self): return self diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c8bdceb57..4ec1e169e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1834,6 +1834,10 @@ class Scheduler( def process_input_requests(self, recv_reqs: List): now = time.monotonic() self.session_controller.maybe_reap(now) + if self.server_args.mm_feature_transport == "cuda_vmm": + for recv_req in recv_reqs: + self._materialize_cuda_vmm_inputs(recv_req) + for recv_req in recv_reqs: # Skip health check when server is busy — ongoing requests already carry health info. if is_health_check_generate_req(recv_req) and not self.is_fully_idle( @@ -1861,6 +1865,28 @@ class Scheduler( if self.external_corpus_manager is not None: self.external_corpus_manager.check_pending_load() + def _materialize_cuda_vmm_inputs(self, recv_req): + """Release VMM slices before request handling can reject the request.""" + if isinstance( + recv_req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput) + ): + tokenized_reqs = (recv_req,) + elif isinstance( + recv_req, + (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput), + ): + tokenized_reqs = recv_req + else: + return + + for tokenized_req in tokenized_reqs: + if tokenized_req.mm_inputs is not None and not isinstance( + tokenized_req.mm_inputs, MultimodalInputs + ): + tokenized_req.mm_inputs = MultimodalInputs.from_processor_output( + tokenized_req.mm_inputs + ) + def init_profiler(self) -> None: self.profiler_manager = SchedulerProfilerManager( ps=self.ps, @@ -2210,11 +2236,13 @@ class Scheduler( return image_inputs - def _get_multimodal_inputs(self, mm_inputs_dict): + def _get_multimodal_inputs(self, mm_inputs): + if isinstance(mm_inputs, MultimodalInputs): + return mm_inputs + if get_mm().enable_broadcast_mm_inputs_process: - return self._process_and_broadcast_mm_inputs(mm_inputs_dict) - else: - return MultimodalInputs.from_processor_output(mm_inputs_dict) + return self._process_and_broadcast_mm_inputs(mm_inputs) + return MultimodalInputs.from_processor_output(mm_inputs) @staticmethod def _try_apply_padded_mm_input_ids(recv_req, req, image_inputs) -> bool: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index f3525b2ef..36c639f84 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -135,6 +135,7 @@ from sglang.srt.utils import ( kill_process_tree, ) from sglang.srt.utils.aio_rwlock import RWLock +from sglang.srt.utils.cuda_vmm_transport_utils import CudaVmmFeatureTransport from sglang.srt.utils.cudacore_pyspy_dump_utils import ( collect_scheduler_processes, pyspy_dump_schedulers, @@ -416,6 +417,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # Init model config self.init_model_config() + self._validate_cuda_vmm_feature_transport_support() # Initialize tokenizer and multimodalprocessor self.init_tokenizer_and_processor() @@ -444,6 +446,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): # Init request dispatcher self.init_request_dispatcher() + # Construct this last so later initialization failures cannot orphan + # the transport's recycler thread. + self.cuda_vmm_feature_transport = CudaVmmFeatureTransport( + self.server_args, self.mm_processor + ) + def init_model_config(self): server_args = self.server_args model_config_class = getattr(self, "model_config_class", ModelConfig) @@ -516,6 +524,19 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): else: self.async_dynamic_batch_tokenizer = None + def _validate_cuda_vmm_feature_transport_support(self) -> None: + if self.server_args.mm_feature_transport != "cuda_vmm": + return + + from sglang.srt.model_loader.utils import get_model_architecture + + model_class, _ = get_model_architecture(self.model_config) + if not getattr(model_class, "supports_cuda_vmm_feature_transport", False): + raise ValueError( + "--mm-feature-transport=cuda_vmm is not supported by model class " + f"{model_class.__name__}" + ) + def init_ipc_channels(self, port_args: PortArgs): context = zmq.asyncio.Context(2) self.recv_from_detokenizer = get_zmq_socket( @@ -1541,16 +1562,26 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self, tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput], ): - tokenized_obj.time_stats.set_api_server_dispatch_time() - tokenized_obj = wrap_shm_features(tokenized_obj) - time_stats = tokenized_obj.time_stats - tokenized_obj.wrap_pickle_fields() - self._dispatch_to_scheduler(tokenized_obj) - state = self.rid_to_state.get(tokenized_obj.rid) - if state is not None: - state.dispatched = True - tokenized_obj.time_stats = time_stats - tokenized_obj.time_stats.set_api_server_dispatch_finish_time() + prepared_mm_items = [] + dispatched = False + try: + prepared_mm_items = self.cuda_vmm_feature_transport.prepare_for_dispatch( + (tokenized_obj.mm_inputs,) + ) + tokenized_obj.time_stats.set_api_server_dispatch_time() + tokenized_obj = wrap_shm_features(tokenized_obj) + time_stats = tokenized_obj.time_stats + tokenized_obj.wrap_pickle_fields() + self._dispatch_to_scheduler(tokenized_obj) + dispatched = True + state = self.rid_to_state.get(tokenized_obj.rid) + if state is not None: + state.dispatched = True + tokenized_obj.time_stats = time_stats + tokenized_obj.time_stats.set_api_server_dispatch_finish_time() + finally: + if not dispatched: + self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items) def _send_batch_request( self, @@ -1559,24 +1590,35 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ], ): """Send a batch of tokenized requests as a single batched request to the scheduler.""" - set_time_batch(tokenized_objs, "set_api_server_dispatch_time") - time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs] - for tokenized_obj in tokenized_objs: - tokenized_obj.wrap_pickle_fields() + prepared_mm_items = [] + dispatched = False + try: + prepared_mm_items = self.cuda_vmm_feature_transport.prepare_for_dispatch( + tokenized_obj.mm_inputs for tokenized_obj in tokenized_objs + ) - if isinstance(tokenized_objs[0], TokenizedGenerateReqInput): - batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs) - else: - batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) + set_time_batch(tokenized_objs, "set_api_server_dispatch_time") + time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs] + for tokenized_obj in tokenized_objs: + tokenized_obj.wrap_pickle_fields() - self._dispatch_to_scheduler(batch_req) - for tokenized_obj in tokenized_objs: - state = self.rid_to_state.get(tokenized_obj.rid) - if state is not None: - state.dispatched = True - for tokenized_obj, time_stat in zip(tokenized_objs, time_stats): - tokenized_obj.time_stats = time_stat - set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") + if isinstance(tokenized_objs[0], TokenizedGenerateReqInput): + batch_req = BatchTokenizedGenerateReqInput(batch=tokenized_objs) + else: + batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) + + self._dispatch_to_scheduler(batch_req) + dispatched = True + for tokenized_obj in tokenized_objs: + state = self.rid_to_state.get(tokenized_obj.rid) + if state is not None: + state.dispatched = True + for tokenized_obj, time_stat in zip(tokenized_objs, time_stats): + tokenized_obj.time_stats = time_stat + set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time") + finally: + if not dispatched: + self.cuda_vmm_feature_transport.cancel_for_dispatch(prepared_mm_items) def _coalesce_streaming_chunks( self, diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index b4e37c721..d92e32e79 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -200,7 +200,7 @@ class BaseMultimodalProcessor(ABC): ) self.mm_feature_transport = ( configured_mm_feature_transport - if configured_mm_feature_transport in ("cpu", "cuda_ipc") + if configured_mm_feature_transport in ("cpu", "cuda_ipc", "cuda_vmm") else "cpu" ) self.use_cuda_ipc = self.mm_feature_transport == "cuda_ipc" @@ -289,8 +289,11 @@ class BaseMultimodalProcessor(ABC): self.mm_processor_worker_num, "auto" if requested_mm_processor_worker_num == 0 else "explicit", ) + cpu_worker_start_method = ( + "spawn" if self.mm_feature_transport == "cuda_vmm" else "fork" + ) self.cpu_executor = concurrent.futures.ProcessPoolExecutor( - mp_context=mp.get_context("fork"), + mp_context=mp.get_context(cpu_worker_start_method), max_workers=int(os.environ.get("SGLANG_CPU_WORKERS", os.cpu_count())), ) @@ -363,6 +366,10 @@ class BaseMultimodalProcessor(ABC): self.server_args.base_gpu_id, ) + @property + def keep_mm_features_on_device(self) -> bool: + return self.mm_feature_transport in ("cuda_ipc", "cuda_vmm") + def compute_mrope_positions(self, input_ids, mm_items): """Compute M-RoPE positions from expanded input_ids and multimodal items. @@ -588,7 +595,10 @@ class BaseMultimodalProcessor(ABC): ) # Deferred: the hash is computed on the GPU tensor first, and # _precompute_hashes_before_cpu_transfer moves it down afterwards. - if not self.use_cuda_ipc and not self.precompute_hash_before_cpu_transfer: + if ( + not self.keep_mm_features_on_device + and not self.precompute_hash_before_cpu_transfer + ): # move feature tensors to cpu for feature_name in self.FEATURE_NAMES: if feature_name in result and isinstance( @@ -1395,7 +1405,7 @@ class BaseMultimodalProcessor(ABC): for item in mm_items: item.set_pad_value() - if not self.use_cuda_ipc: + if not self.keep_mm_features_on_device: item.feature = self._move_feature_to_cpu(item.feature) item.precomputed_embeddings = self._move_feature_to_cpu( item.precomputed_embeddings diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f1b9a225a..50900c44e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2746,13 +2746,15 @@ class ServerArgs: bool, "Adopt base image processor instead of fast image processor.", NS("mm") ] = False mm_feature_transport: A[ - Optional[Literal["cpu", "cuda_ipc"]], - "Transport multimodal features through CPU memory or a bounded CUDA IPC pool. " + Optional[Literal["cpu", "cuda_ipc", "cuda_vmm"]], + "Transport multimodal features through CPU memory, a bounded CUDA IPC " + "pool, or a bounded CUDA VMM pool. CUDA VMM must be selected explicitly " + "and is available only to models that opt in. " "Unset resolves automatically: multimodal models on single-node CUDA " "deployments (without disaggregation) use cuda_ipc, everything else uses " - "cpu. CUDA IPC reserves SGLANG_MM_FEATURE_CACHE_MB (default 1024 MiB) on " - "the base GPU and falls back to CPU transport per tensor when the pool is " - "full.", + "cpu. Both CUDA transports reserve SGLANG_MM_FEATURE_CACHE_MB (default " + "1024 MiB) on the base GPU across tokenizer workers and fall back to CPU " + "transport per tensor when full.", NS("mm"), ] = None keep_mm_feature_on_device: A[ @@ -7582,10 +7584,10 @@ class ServerArgs: legacy_ipc_enabled = envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get() if self.keep_mm_feature_on_device: - if requested_transport == "cpu": + if requested_transport not in (None, "cuda_ipc"): raise ValueError( "--keep-mm-feature-on-device conflicts with " - "--mm-feature-transport=cpu. Use only " + f"--mm-feature-transport={requested_transport}. Use only " "--mm-feature-transport=cuda_ipc." ) requested_transport = "cuda_ipc" @@ -7638,14 +7640,31 @@ class ServerArgs: int(legacy_ipc_enabled), ) - if self.encoder_only and requested_transport == "cuda_ipc": + if self.encoder_only and requested_transport in ("cuda_ipc", "cuda_vmm"): logger.warning( - "--mm-feature-transport=cuda_ipc does not control encoder-only " + "--mm-feature-transport=%s does not control encoder-only " "output transfer; using cpu for this inactive transport. Select " - "--encoder-transfer-backend for encoder outputs." + "--encoder-transfer-backend for encoder outputs.", + requested_transport, ) requested_transport = "cpu" + if requested_transport == "cuda_vmm": + if not is_cuda(): + raise ValueError( + "--mm-feature-transport=cuda_vmm requires NVIDIA CUDA." + ) + if self.pp_size != 1: + raise ValueError( + "--mm-feature-transport=cuda_vmm does not support pipeline " + "parallelism." + ) + if envs.SGLANG_RUST_SERVER.get(): + raise ValueError( + "--mm-feature-transport=cuda_vmm is not supported with " + "SGLANG_RUST_SERVER." + ) + if requested_transport == "cuda_ipc": if not is_cuda(): raise ValueError( diff --git a/python/sglang/srt/utils/cuda_vmm_transport_utils.py b/python/sglang/srt/utils/cuda_vmm_transport_utils.py new file mode 100644 index 000000000..b7250a3cd --- /dev/null +++ b/python/sglang/srt/utils/cuda_vmm_transport_utils.py @@ -0,0 +1,1095 @@ +from __future__ import annotations + +import logging +import os +import secrets +import socket +import threading +from collections.abc import Iterable, Sequence +from dataclasses import dataclass + +import torch + +from sglang.srt.distributed.device_communicators.vmm_utils import ( + _FD_SEND_TIMEOUT_S, + _get_cuda_driver, + _recv_fd, + _send_fd, + check_drv, + import_and_map_alloc, + make_rw_access_desc, + release_mappings, +) +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, +) +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils.cuda_ipc_transport_utils import ( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, + MM_FEATURE_CACHE_SIZE, + MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL, + CudaIpcTensorTransportProxy, + get_mm_feature_pool_size_per_worker, +) + +logger = logging.getLogger(__name__) + +_CONTROL_ALIGNMENT = 256 +_CONTROL_WORD_BYTES = 4 + + +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _tensor_from_pointer(pointer: int, size: int, device_index: int) -> torch.Tensor: + device = torch.device(f"cuda:{device_index}") + storage = torch._C._construct_storage_from_data_pointer(pointer, device, size) + return torch.empty(0, dtype=torch.uint8, device=device).set_( + storage, 0, (size,), (1,) + ) + + +class _PosixFdBroker: + """Serve one exported CUDA allocation FD to local consumer processes.""" + + def __init__(self, fd: int) -> None: + self.fd = fd + self._stop = threading.Event() + self._error: Exception | None = None + self.socket_path = f"\0sgl_mm_vmm_{secrets.token_hex(16)}" + self._server = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) + self._server.bind(self.socket_path) + self._server.listen() + self._server.settimeout(0.1) + self._thread = threading.Thread( + target=self._serve, + name="CudaVmmPosixFdBroker", + daemon=True, + ) + try: + self._thread.start() + except BaseException: + self._server.close() + raise + + def _serve(self) -> None: + while not self._stop.is_set(): + try: + conn, _ = self._server.accept() + except TimeoutError: + continue + except OSError as error: + if self._stop.is_set(): + return + self._error = error + logger.exception("CUDA VMM POSIX FD broker failed") + return + + try: + with conn: + _send_fd(conn, self.fd, src_rank=0, base_idx=0) + except Exception as error: + self._error = error + logger.exception("CUDA VMM POSIX FD broker failed") + return + + def raise_if_failed(self) -> None: + if self._error is not None: + raise RuntimeError("CUDA VMM POSIX FD broker failed") from self._error + + def close(self) -> None: + self._stop.set() + self._server.close() + self._thread.join(timeout=1.0) + if self._thread.is_alive(): + raise RuntimeError("CUDA VMM POSIX FD broker did not stop") + + +def _receive_posix_fd(socket_path: str) -> int: + with socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET) as sock: + sock.settimeout(_FD_SEND_TIMEOUT_S) + sock.connect(socket_path) + packet = _recv_fd(sock) + if packet is None: + raise RuntimeError("CUDA VMM POSIX FD broker returned no file descriptor") + _src_rank, _base_idx, fd = packet + return fd + + +@dataclass +class _CudaVmmMemoryChunk: + start: int + end: int + + @property + def size(self) -> int: + return self.end - self.start + + +@dataclass(frozen=True) +class _CudaVmmPackedTensorLayout: + relative_offset: int + data_nbytes: int + shape: torch.Size + dtype: torch.dtype + + +def _build_packed_tensor_layout( + tensors: Sequence[torch.Tensor], +) -> tuple[list[_CudaVmmPackedTensorLayout], int]: + layouts = [] + next_offset = 0 + for tensor in tensors: + next_offset = _align_up(next_offset, tensor.element_size()) + data_nbytes = tensor.numel() * tensor.element_size() + layouts.append( + _CudaVmmPackedTensorLayout( + relative_offset=next_offset, + data_nbytes=data_nbytes, + shape=tensor.shape, + dtype=tensor.dtype, + ) + ) + next_offset += data_nbytes + return layouts, next_offset + + +def _contains_tensor_container(value) -> bool: + return isinstance(value, (list, tuple)) and any( + isinstance(item, torch.Tensor) or _contains_tensor_container(item) + for item in value + ) + + +def get_vmm_feature_consumer_count(server_args) -> int: + if server_args.enable_dp_attention: + return server_args.tp_size // server_args.dp_size + return server_args.tp_size + + +class CudaVmmMemoryPool: + """Bounded CUDA VMM pool shared through FABRIC or a local POSIX FD.""" + + def __init__( + self, + memory_size: int, + recycle_interval: float, + base_gpu_id: int, + consumer_count: int, + allow_posix_fallback: bool = False, + ) -> None: + if memory_size <= 0: + raise ValueError("memory_size must be positive") + if consumer_count <= 0: + raise ValueError("consumer_count must be positive") + if recycle_interval <= 0: + raise ValueError("recycle_interval must be positive") + + self.device_index = int(base_gpu_id) + self.consumer_count = int(consumer_count) + self._recycle_interval = float(recycle_interval) + self._lock = threading.Lock() + self._publisher_condition = threading.Condition(self._lock) + self._shutdown_lock = threading.Lock() + self._active_publishers = 0 + self._closing = False + self._pool_full_warned = False + self._stop_recycler = threading.Event() + self._pool_error: BaseException | None = None + self._closed = False + + self._allocation_handle = None + self._pool_pointer = None + self._allocation_mapped = False + self.allocation_size = 0 + self.shareable_handle = None + self.memory_pool = None + self._fd_broker: _PosixFdBroker | None = None + self.posix_socket_path: str | None = None + self._recycle_stream = None + self._recycle_thread = None + + self.use_fabric = True + try: + self._allocate(memory_size) + except RuntimeError as error: + if not allow_posix_fallback: + raise + logger.warning( + "CUDA FABRIC VMM allocation is unavailable; falling back to " + "a POSIX FD handle: %s", + error, + ) + self.use_fabric = False + self._allocate(memory_size) + try: + if not self.use_fabric: + self._fd_broker = _PosixFdBroker(self.shareable_handle) + self.posix_socket_path = self._fd_broker.socket_path + + self.available_chunks = [_CudaVmmMemoryChunk(0, self.allocation_size)] + self.occupied_chunks = [] + self._recycle_stream = torch.cuda.Stream(device=self.device_index) + self._recycle_thread = threading.Thread( + target=self._recycle_loop, + name="CudaVmmMemoryPoolRecycler", + daemon=True, + ) + self._recycle_thread.start() + except BaseException as error: + cleanup_errors = [] + self._stop_recycler.set() + if self._recycle_thread is not None and self._recycle_thread.is_alive(): + self._recycle_thread.join(timeout=1.0) + if self._recycle_thread.is_alive(): + cleanup_errors.append( + RuntimeError("CUDA VMM recycler did not stop during rollback") + ) + if self._fd_broker is not None: + try: + self._fd_broker.close() + except BaseException as cleanup_error: + cleanup_errors.append(cleanup_error) + try: + self._release_allocation() + except BaseException as cleanup_error: + cleanup_errors.append(cleanup_error) + if cleanup_errors: + error.add_note( + f"{len(cleanup_errors)} CUDA VMM initialization rollback " + "operation(s) also failed" + ) + raise error from cleanup_errors[0] + raise + + @property + def fabric_handle(self) -> bytes | None: + return self.shareable_handle if self.use_fabric else None + + def _allocate(self, memory_size: int) -> None: + drv = _get_cuda_driver() + handle_type = ( + drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + if self.use_fabric + else drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ) + prop = drv.CUmemAllocationProp() + prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = self.device_index + prop.requestedHandleTypes = handle_type + if self.use_fabric: + prop.allocFlags.gpuDirectRDMACapable = 1 + + recommended = ( + drv.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED + ) + with torch.cuda.device(self.device_index): + check_drv(drv.cuInit(0), "cuInit") + granularity = int( + check_drv( + drv.cuMemGetAllocationGranularity(prop, recommended), + "cuMemGetAllocationGranularity(VMM transport)", + ) + ) + allocation_size = memory_size // granularity * granularity + if allocation_size == 0: + raise ValueError( + f"memory_size={memory_size} is smaller than CUDA VMM " + f"granularity={granularity}" + ) + + handle = pointer = exported = None + mapped = False + try: + handle = check_drv( + drv.cuMemCreate(allocation_size, prop, 0), + "cuMemCreate(VMM transport)", + ) + pointer = int( + check_drv( + drv.cuMemAddressReserve(allocation_size, granularity, 0, 0), + "cuMemAddressReserve(VMM transport)", + ) + ) + check_drv( + drv.cuMemMap(pointer, allocation_size, 0, handle, 0), + "cuMemMap(VMM transport)", + ) + mapped = True + access = make_rw_access_desc(self.device_index) + check_drv( + drv.cuMemSetAccess(pointer, allocation_size, [access], 1), + "cuMemSetAccess(VMM transport)", + ) + exported = check_drv( + drv.cuMemExportToShareableHandle(handle, handle_type, 0), + "cuMemExportToShareableHandle(VMM transport)", + ) + memory_pool = _tensor_from_pointer( + pointer, allocation_size, self.device_index + ) + except BaseException: + if mapped: + drv.cuMemUnmap(pointer, allocation_size) + if pointer is not None: + drv.cuMemAddressFree(pointer, allocation_size) + if handle is not None: + drv.cuMemRelease(handle) + if not self.use_fabric and exported is not None: + os.close(int(exported)) + raise + + self._allocation_handle = handle + self._pool_pointer = pointer + self._allocation_mapped = True + self.allocation_size = allocation_size + self.shareable_handle = ( + bytes(exported.data) if self.use_fabric else int(exported) + ) + self.memory_pool = memory_pool + + @property + def control_size(self) -> int: + return _align_up(self.consumer_count * _CONTROL_WORD_BYTES, _CONTROL_ALIGNMENT) + + def _raise_if_failed(self) -> None: + if self._pool_error is not None: + raise RuntimeError("CUDA VMM multimodal pool failed") from self._pool_error + if self._fd_broker is not None: + self._fd_broker.raise_if_failed() + + def _reserve_for_publish(self, required_size: int) -> _CudaVmmMemoryChunk | None: + with self._publisher_condition: + self._raise_if_failed() + if self._closing or self._closed: + raise RuntimeError("CUDA VMM multimodal pool is closing") + chunk = self._reserve_chunk(required_size) + if chunk is not None: + self._active_publishers += 1 + return chunk + + def _finish_publish(self) -> None: + with self._publisher_condition: + self._active_publishers -= 1 + if self._active_publishers == 0: + self._publisher_condition.notify_all() + + def wrap_tensor(self, tensor: torch.Tensor): + self._raise_if_failed() + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + data_nbytes = tensor.numel() * tensor.element_size() + required_size = _align_up(self.control_size + data_nbytes, _CONTROL_ALIGNMENT) + source_bytes = tensor.reshape(-1).view(torch.uint8) + + chunk = self._reserve_for_publish(required_size) + if chunk is None: + self._warn_pool_full_once(data_nbytes) + return tensor.cpu() + + producer_stream = None + copy_synchronized = False + try: + with torch.cuda.device(self.device_index): + producer_stream = torch.cuda.current_stream(self.device_index) + control_offset = chunk.start + data_offset = control_offset + self.control_size + control_end = control_offset + self.control_size + self.memory_pool[control_offset:control_end].zero_() + self.memory_pool[data_offset : data_offset + data_nbytes].copy_( + source_bytes, non_blocking=True + ) + # Imported VMM memory does not support cuStreamWaitValue32 on + # current GB-class drivers, so publish only after this copy. + producer_stream.synchronize() + copy_synchronized = True + + proxy = CudaVmmTensorTransportProxy( + fabric_handle=self.fabric_handle, + posix_socket_path=self.posix_socket_path, + allocation_size=self.allocation_size, + data_offset=data_offset, + data_nbytes=data_nbytes, + control_offset=control_offset, + consumer_count=self.consumer_count, + shape=tensor.shape, + dtype=tensor.dtype, + ) + with self._lock: + self.occupied_chunks.append(chunk) + return proxy + except BaseException: + safe_to_release = copy_synchronized or producer_stream is None + if not safe_to_release: + try: + producer_stream.synchronize() + safe_to_release = True + except BaseException as cleanup_error: + self._pool_error = cleanup_error + if safe_to_release: + with self._lock: + self._release_reserved_chunk(chunk) + raise + finally: + self._finish_publish() + + def wrap_tensors( + self, + tensors: Sequence[torch.Tensor], + ) -> list[CudaVmmPackedTensorTransportProxy] | None: + """Publish tensors through one shared VMM chunk. + + Every tensor must have the same dispatch and consumer lifetime because + reconstructing one child copies and acknowledges the full packed chunk. + ``None`` means that no contiguous pool chunk was available. No tensor is + published in that case, so the caller can use another transport path. + """ + self._raise_if_failed() + tensors = list(tensors) + if not tensors: + return [] + + layouts, packed_data_nbytes = _build_packed_tensor_layout(tensors) + required_size = _align_up( + self.control_size + packed_data_nbytes, _CONTROL_ALIGNMENT + ) + chunk = self._reserve_for_publish(required_size) + if chunk is None: + return None + + producer_stream = None + copy_synchronized = False + try: + contiguous_tensors = [ + tensor if tensor.is_contiguous() else tensor.contiguous() + for tensor in tensors + ] + with torch.cuda.device(self.device_index): + producer_stream = torch.cuda.current_stream(self.device_index) + control_offset = chunk.start + data_offset = control_offset + self.control_size + control_end = control_offset + self.control_size + self.memory_pool[control_offset:control_end].zero_() + for tensor, layout in zip(contiguous_tensors, layouts): + data_start = data_offset + layout.relative_offset + self.memory_pool[ + data_start : data_start + layout.data_nbytes + ].copy_(tensor.reshape(-1).view(torch.uint8), non_blocking=True) + # A single synchronization publishes every child together. + producer_stream.synchronize() + copy_synchronized = True + + owner = _CudaVmmPackedTransportOwner( + fabric_handle=self.fabric_handle, + posix_socket_path=self.posix_socket_path, + allocation_size=self.allocation_size, + data_offset=data_offset, + data_nbytes=packed_data_nbytes, + control_offset=control_offset, + consumer_count=self.consumer_count, + ) + proxies = [ + CudaVmmPackedTensorTransportProxy( + owner=owner, + layout=layout, + ) + for layout in layouts + ] + with self._lock: + self.occupied_chunks.append(chunk) + return proxies + except BaseException: + safe_to_release = copy_synchronized or producer_stream is None + if not safe_to_release: + try: + producer_stream.synchronize() + safe_to_release = True + except BaseException as cleanup_error: + self._pool_error = cleanup_error + if safe_to_release: + with self._lock: + self._release_reserved_chunk(chunk) + raise + finally: + self._finish_publish() + + def _reserve_chunk(self, required_size: int) -> _CudaVmmMemoryChunk | None: + candidates = [ + chunk for chunk in self.available_chunks if chunk.size >= required_size + ] + if not candidates: + return None + + available = min(candidates, key=lambda chunk: chunk.size) + self.available_chunks.remove(available) + occupied = _CudaVmmMemoryChunk( + start=available.start, + end=available.start + required_size, + ) + if occupied.end < available.end: + self.available_chunks.append( + _CudaVmmMemoryChunk(occupied.end, available.end) + ) + return occupied + + def _release_reserved_chunk(self, chunk: _CudaVmmMemoryChunk) -> None: + if chunk in self.occupied_chunks: + self.occupied_chunks.remove(chunk) + self.available_chunks.append(_CudaVmmMemoryChunk(chunk.start, chunk.end)) + self._merge_chunks() + + def _cancel_control_offset(self, control_offset: int) -> None: + with self._lock: + chunk = next( + ( + chunk + for chunk in self.occupied_chunks + if chunk.start == control_offset + ), + None, + ) + if chunk is None: + raise RuntimeError( + "CUDA VMM pool has no occupied slice at control offset " + f"{control_offset}" + ) + self._release_reserved_chunk(chunk) + + def cancel_proxy(self, proxy: CudaVmmTensorTransportProxy) -> None: + """Return a published slice when its request was never dispatched.""" + if isinstance(proxy, CudaVmmPackedTensorTransportProxy): + proxy._packed_owner.cancel_from_pool(self) + return + self._cancel_control_offset(proxy.control_offset) + + def _warn_pool_full_once(self, data_nbytes: int) -> None: + if self._pool_full_warned: + return + self._pool_full_warned = True + logger.warning( + "CUDA VMM multimodal pool has no free chunk for a %.2f MiB tensor " + "(pool size: %.2f MiB); falling back to CPU transport. Increase " + "SGLANG_MM_FEATURE_CACHE_MB to avoid inline request broadcasts.", + data_nbytes / (1024 * 1024), + self.allocation_size / (1024 * 1024), + ) + + def _recycle_loop(self) -> None: + while not self._stop_recycler.wait(self._recycle_interval): + try: + with self._lock: + self._recycle_chunks() + self._merge_chunks() + except Exception as error: + logger.exception("CUDA VMM multimodal pool recycle failed") + self._pool_error = error + self._stop_recycler.set() + + def _recycle_chunks(self) -> None: + remaining = [] + recycled = [] + with ( + torch.cuda.device(self.device_index), + torch.cuda.stream(self._recycle_stream), + ): + for chunk in self.occupied_chunks: + ack_start = chunk.start + ack_end = ack_start + self.consumer_count * _CONTROL_WORD_BYTES + ack_count = int( + torch.count_nonzero( + self.memory_pool[ack_start:ack_end].view(torch.int32) + ).item() + ) + if ack_count == self.consumer_count: + recycled.append(_CudaVmmMemoryChunk(chunk.start, chunk.end)) + else: + remaining.append(chunk) + + self.available_chunks.extend(recycled) + self.occupied_chunks = remaining + + def _merge_chunks(self) -> None: + merged = [] + for chunk in sorted(self.available_chunks, key=lambda item: item.start): + if merged and merged[-1].end == chunk.start: + merged[-1].end = chunk.end + else: + merged.append(chunk) + self.available_chunks = merged + + def _release_allocation(self) -> None: + self.memory_pool = None + if not self.use_fabric and self.shareable_handle is not None: + os.close(self.shareable_handle) + self.shareable_handle = None + if self._pool_pointer is None and self._allocation_handle is None: + return + drv = _get_cuda_driver() + with torch.cuda.device(self.device_index): + if self._allocation_mapped: + check_drv( + drv.cuMemUnmap(self._pool_pointer, self.allocation_size), + "cuMemUnmap(VMM transport pool)", + ) + self._allocation_mapped = False + if self._pool_pointer is not None: + check_drv( + drv.cuMemAddressFree(self._pool_pointer, self.allocation_size), + "cuMemAddressFree(VMM transport pool)", + ) + self._pool_pointer = None + if self._allocation_handle is not None: + check_drv( + drv.cuMemRelease(self._allocation_handle), + "cuMemRelease(VMM transport pool)", + ) + self._allocation_handle = None + + def shutdown(self) -> None: + with self._shutdown_lock: + if self._closed: + return + with self._publisher_condition: + self._closing = True + while self._active_publishers: + self._publisher_condition.wait() + + self._stop_recycler.set() + self._recycle_thread.join(timeout=1.0) + if self._recycle_thread.is_alive(): + raise RuntimeError("CUDA VMM recycler did not stop") + if self._fd_broker is not None: + self._fd_broker.close() + self._fd_broker = None + self._release_allocation() + self._closed = True + + +@dataclass +class _ImportedCudaVmmPool: + pointer: int + allocation_size: int + memory: torch.Tensor | None + + def close(self) -> None: + self.memory = None + release_mappings( + [ + ( + self.pointer, + self.allocation_size, + [(0, self.allocation_size)], + ) + ] + ) + + +_imported_pool_cache: dict[tuple, _ImportedCudaVmmPool] = {} +_imported_pool_cache_lock = threading.Lock() + + +def _get_imported_pool( + *, + fabric_handle: bytes | None, + posix_socket_path: str | None, + allocation_size: int, + device_index: int, +) -> _ImportedCudaVmmPool: + use_fabric = fabric_handle is not None + transport_handle = fabric_handle if use_fabric else posix_socket_path + if transport_handle is None: + raise RuntimeError("CUDA VMM proxy has no shareable handle") + key = (device_index, allocation_size, transport_handle) + pool = _imported_pool_cache.get(key) + if pool is not None: + return pool + + with _imported_pool_cache_lock: + pool = _imported_pool_cache.get(key) + if pool is not None: + return pool + + fd = None + try: + if not use_fabric: + fd = _receive_posix_fd(posix_socket_path) + with torch.cuda.device(device_index): + pointer = import_and_map_alloc( + fabric_handle, + fd, + allocation_size, + device_index, + use_fabric=use_fabric, + peer_rank=-1, + ) + try: + memory = _tensor_from_pointer( + pointer, allocation_size, device_index + ) + except Exception: + release_mappings( + [(pointer, allocation_size, [(0, allocation_size)])] + ) + raise + finally: + if fd is not None: + os.close(fd) + + pool = _ImportedCudaVmmPool( + pointer=pointer, + allocation_size=allocation_size, + memory=memory, + ) + _imported_pool_cache[key] = pool + return pool + + +def _imported_pool_cache_clear() -> None: + with _imported_pool_cache_lock: + pools = list(_imported_pool_cache.values()) + _imported_pool_cache.clear() + for pool in pools: + pool.close() + + +class CudaVmmTensorTransportProxy(CudaIpcTensorTransportProxy): + """Multimodal tensor proxy backed by a shared CUDA VMM pool.""" + + def __init__( + self, + *, + fabric_handle: bytes | None, + posix_socket_path: str | None, + allocation_size: int, + data_offset: int, + data_nbytes: int, + control_offset: int, + consumer_count: int, + shape, + dtype, + ) -> None: + self.fabric_handle = fabric_handle + self.posix_socket_path = posix_socket_path + self.allocation_size = allocation_size + self.data_offset = data_offset + self.data_nbytes = data_nbytes + self.control_offset = control_offset + self.consumer_count = consumer_count + self.shape = shape + self.dtype = dtype + self.reconstruct_tensor = None + self._consumer_acknowledged = False + + def _pool(self, device_index: int) -> _ImportedCudaVmmPool: + return _get_imported_pool( + fabric_handle=self.fabric_handle, + posix_socket_path=self.posix_socket_path, + allocation_size=self.allocation_size, + device_index=device_index, + ) + + def _acknowledgement_range(self, consumer_count: int) -> tuple[int, int]: + if consumer_count <= 0: + raise ValueError("consumer_count must be positive") + if consumer_count == self.consumer_count: + return 0, self.consumer_count + + parallel = get_parallel() + group_start = parallel.attn_cp_rank * parallel.attn_tp_size + group_end = group_start + parallel.attn_tp_size + if not 0 <= group_start < group_end <= self.consumer_count: + raise ValueError( + "attention group range " + f"[{group_start}, {group_end}) is outside " + f"consumer_count={self.consumer_count}" + ) + if consumer_count == 1: + slot = group_start + parallel.attn_tp_rank + return slot, slot + 1 + if consumer_count == parallel.attn_tp_size: + return group_start, group_end + raise ValueError( + "consumer_count must be 1, the attention TP size, or the full " + f"consumer count ({self.consumer_count}); got {consumer_count}" + ) + + def _resolve_consumer_count(self, consumer_count: int | None) -> int: + return 1 if consumer_count is None else consumer_count + + def _acknowledge_consumption(self, device_index: int, consumer_count: int) -> None: + if self._consumer_acknowledged: + return + pool = self._pool(device_index) + ack_start = self.control_offset + ack_end = ack_start + self.consumer_count * _CONTROL_WORD_BYTES + ack_words = pool.memory[ack_start:ack_end].view(torch.int32) + slot_start, slot_end = self._acknowledgement_range(consumer_count) + # This kernel is ordered after the remote read on the consumer stream; + # observing the flag therefore means the pool slice is safe to reuse. + ack_words[slot_start:slot_end].fill_(1) + self._consumer_acknowledged = True + + def acknowledge_consumption(self, consumer_count: int | None = None) -> None: + consumer_count = self._resolve_consumer_count(consumer_count) + device_index = torch.cuda.current_device() + with torch.cuda.device(device_index): + self._acknowledge_consumption(device_index, consumer_count) + + def reconstruct_on_target_device( + self, rebuild_device_idx, consumer_count: int | None = None + ): + consumer_count = self._resolve_consumer_count(consumer_count) + 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._consumer_acknowledged: + raise RuntimeError("CUDA VMM tensor has already released its pool slice") + + pool = self._pool(rebuild_device_idx) + try: + with torch.cuda.device(rebuild_device): + source = pool.memory[ + self.data_offset : self.data_offset + self.data_nbytes + ] + reconstructed = torch.empty( + self.shape, dtype=self.dtype, device=rebuild_device + ).contiguous() + reconstructed.reshape(-1).view(torch.uint8).copy_( + source, non_blocking=True + ) + self._acknowledge_consumption(rebuild_device_idx, consumer_count) + except BaseException as error: + try: + with torch.cuda.device(rebuild_device): + self._acknowledge_consumption(rebuild_device_idx, consumer_count) + except BaseException as cleanup_error: + error.add_note( + "CUDA VMM reconstruction cleanup also failed; the pool " + "slice was not acknowledged" + ) + raise error from cleanup_error + raise + + self.reconstruct_tensor = reconstructed + return reconstructed + + +class _CudaVmmPackedTransportOwner(CudaVmmTensorTransportProxy): + def __init__( + self, + *, + fabric_handle: bytes | None, + posix_socket_path: str | None, + allocation_size: int, + data_offset: int, + data_nbytes: int, + control_offset: int, + consumer_count: int, + ) -> None: + super().__init__( + fabric_handle=fabric_handle, + posix_socket_path=posix_socket_path, + allocation_size=allocation_size, + data_offset=data_offset, + data_nbytes=data_nbytes, + control_offset=control_offset, + consumer_count=consumer_count, + shape=(data_nbytes,), + dtype=torch.uint8, + ) + self._producer_cancelled = False + + def cancel_from_pool(self, pool: CudaVmmMemoryPool) -> None: + if self._producer_cancelled: + return + pool._cancel_control_offset(self.control_offset) + self._producer_cancelled = True + + +class CudaVmmPackedTensorTransportProxy(CudaVmmTensorTransportProxy): + """One typed view within a packed CUDA VMM transfer.""" + + def __init__( + self, + *, + owner: _CudaVmmPackedTransportOwner, + layout: _CudaVmmPackedTensorLayout, + ) -> None: + super().__init__( + fabric_handle=owner.fabric_handle, + posix_socket_path=owner.posix_socket_path, + allocation_size=owner.allocation_size, + data_offset=owner.data_offset + layout.relative_offset, + data_nbytes=layout.data_nbytes, + control_offset=owner.control_offset, + consumer_count=owner.consumer_count, + shape=layout.shape, + dtype=layout.dtype, + ) + self._packed_owner = owner + self._packed_relative_offset = layout.relative_offset + + def acknowledge_consumption(self, consumer_count: int | None = None) -> None: + raise RuntimeError( + "Packed CUDA VMM features must be reconstructed before release" + ) + + def reconstruct_on_target_device( + self, rebuild_device_idx, consumer_count: int | None = 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 + if self._consumer_acknowledged: + raise RuntimeError("CUDA VMM tensor has already released its pool slice") + + packed_buffer = self._packed_owner.reconstruct_on_target_device( + rebuild_device_idx, consumer_count=consumer_count + ) + tensor_bytes = packed_buffer[ + self._packed_relative_offset : self._packed_relative_offset + + self.data_nbytes + ] + reconstructed = tensor_bytes.view(self.dtype).reshape(self.shape) + self.reconstruct_tensor = reconstructed + self._consumer_acknowledged = True + return reconstructed + + +class CudaVmmFeatureTransport: + """Tokenizer-owned VMM transport for one tokenizer worker.""" + + def __init__(self, server_args, mm_processor) -> None: + self.pool: CudaVmmMemoryPool | None = None + if server_args.mm_feature_transport != "cuda_vmm": + return + if mm_processor is None: + raise RuntimeError( + "A CUDA VMM-enabled model must provide a multimodal processor" + ) + + per_worker_pool_size = get_mm_feature_pool_size_per_worker( + MM_FEATURE_CACHE_SIZE, server_args.tokenizer_worker_num + ) + self.pool = CudaVmmMemoryPool( + memory_size=per_worker_pool_size, + recycle_interval=MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL, + base_gpu_id=server_args.base_gpu_id, + consumer_count=get_vmm_feature_consumer_count(server_args), + allow_posix_fallback=server_args.nnodes == 1, + ) + + def prepare_for_dispatch( + self, + mm_inputs_batch: Iterable[MultimodalProcessorOutput | None], + ) -> list[MultimodalDataItem]: + if self.pool is None: + return [] + + prepared_mm_items = [] + preparation_complete = False + try: + for mm_inputs in mm_inputs_batch: + if mm_inputs is None or not mm_inputs.mm_items: + continue + mm_items = mm_inputs.mm_items + self.wrap_items(mm_items) + prepared_mm_items.extend(mm_items) + preparation_complete = True + return prepared_mm_items + finally: + if not preparation_complete: + self.cancel_for_dispatch(prepared_mm_items) + + def wrap_items(self, mm_items: list[MultimodalDataItem]) -> None: + if self.pool is None: + return + + updates = [] + try: + pack_candidates = [ + (item, item.feature) + for item in mm_items + if item.modality == Modality.IMAGE + and isinstance(item.feature, torch.Tensor) + and item.feature.numel() > 0 + and not item.model_specific_data.get( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, False + ) + ] + if len(pack_candidates) >= 2: + packed = self.pool.wrap_tensors( + [tensor for _, tensor in pack_candidates] + ) + if packed is not None: + for (item, tensor), proxy in zip( + pack_candidates, packed, strict=True + ): + item.feature = proxy + updates.append((item, "feature", tensor, proxy)) + + for item in mm_items: + for field in ("feature", "precomputed_embeddings"): + tensor = getattr(item, field) + if _contains_tensor_container(tensor): + raise TypeError( + "CUDA VMM feature transport requires each feature " + "field to contain a single tensor" + ) + if not isinstance(tensor, torch.Tensor): + continue + wrapped = self.pool.wrap_tensor(tensor) + setattr(item, field, wrapped) + updates.append((item, field, tensor, wrapped)) + except BaseException as error: + rollback_errors = [] + for item, field, tensor, wrapped in reversed(updates): + try: + if isinstance(wrapped, CudaVmmTensorTransportProxy): + self.pool.cancel_proxy(wrapped) + except BaseException as rollback_error: + rollback_errors.append(rollback_error) + finally: + setattr(item, field, tensor) + if rollback_errors: + error.add_note( + f"{len(rollback_errors)} VMM rollback operation(s) also failed" + ) + raise error from rollback_errors[0] + raise + + def cancel_for_dispatch(self, mm_items: list[MultimodalDataItem]) -> None: + if self.pool is None or not mm_items: + return + + errors = [] + for item in mm_items: + for field in ("feature", "precomputed_embeddings"): + proxy = getattr(item, field) + if not isinstance(proxy, CudaVmmTensorTransportProxy): + continue + try: + self.pool.cancel_proxy(proxy) + except BaseException as error: + errors.append(error) + finally: + setattr(item, field, None) + if errors: + raise RuntimeError( + f"Failed to cancel {len(errors)} VMM transport slice(s)" + ) from errors[0] + + def shutdown(self) -> None: + if self.pool is None: + return + self.pool.shutdown() diff --git a/test/registered/unit/managers/test_mm_process_config.py b/test/registered/unit/managers/test_mm_process_config.py index 67e6a053f..36f198861 100644 --- a/test/registered/unit/managers/test_mm_process_config.py +++ b/test/registered/unit/managers/test_mm_process_config.py @@ -237,6 +237,31 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase): self.assertFalse(processor.use_ipc_pool_handle_cache) memory_pool.assert_not_called() + def test_cuda_vmm_keeps_features_on_device_without_ipc_pool(self): + from sglang.srt.multimodal.processors import base_processor + + hf_processor = self._processor() + feature = torch.empty(1, device="meta") + hf_processor.return_value = {"pixel_values": feature} + with patch.object( + base_processor.BaseMultimodalProcessor, "__abstractmethods__", set() + ), patch.object(base_processor, "MmItemMemoryPool") as memory_pool: + processor = base_processor.BaseMultimodalProcessor( + hf_config=MagicMock(), + server_args=self._server_args("cuda_vmm"), + _processor=hf_processor, + transport_mode=None, + ) + + result = processor.process_mm_data("test") + + self.assertEqual(processor.mm_feature_transport, "cuda_vmm") + self.assertFalse(processor.use_cuda_ipc) + self.assertTrue(processor.keep_mm_features_on_device) + self.assertEqual(processor.cpu_executor._mp_context.get_start_method(), "spawn") + self.assertIs(result["pixel_values"], feature) + memory_pool.assert_not_called() + class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase): @staticmethod @@ -251,6 +276,7 @@ class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase): processor = BaseMultimodalProcessor() processor.precompute_hash_before_cpu_transfer = enabled processor.use_cuda_ipc = False + processor.mm_feature_transport = "cpu" return processor def test_enabled_path_sets_hash_and_pad_value(self): diff --git a/test/registered/unit/multimodal/test_cuda_vmm_transport.py b/test/registered/unit/multimodal/test_cuda_vmm_transport.py new file mode 100644 index 000000000..83becd597 --- /dev/null +++ b/test/registered/unit/multimodal/test_cuda_vmm_transport.py @@ -0,0 +1,490 @@ +"""CUDA VMM multimodal feature transport regression tests.""" + +from __future__ import annotations + +import gc +import multiprocessing as mp +import os +import pickle +import queue +import threading +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmMemoryPool, + CudaVmmPackedTensorTransportProxy, + _imported_pool_cache_clear, + _PosixFdBroker, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="base-c", runner_config="4-gpu-gb300") + + +class _FabricUnavailableCudaVmmMemoryPool(CudaVmmMemoryPool): + def _allocate(self, memory_size: int) -> None: + if self.use_fabric: + raise RuntimeError("forced FABRIC allocation failure") + super()._allocate(memory_size) + + +def _produce_vmm_tensor(proxy_queue, consumer_done, result_queue, mode): + pool = source = proxy = None + try: + torch.cuda.set_device(0) + pool_cls = ( + _FabricUnavailableCudaVmmMemoryPool + if mode == "posix_fallback" + else CudaVmmMemoryPool + ) + pool = pool_cls( + memory_size=4 << 20, + recycle_interval=60, + base_gpu_id=0, + consumer_count=2, + allow_posix_fallback=True, + ) + source = torch.arange(35, dtype=torch.float32, device="cuda").reshape(5, 7) + expected = source.cpu().tolist() + proxy = pool.wrap_tensor(source) + proxy_queue.put((proxy, expected)) + if not consumer_done.wait(timeout=60): + raise TimeoutError("consumers did not release the CUDA VMM tensor") + with pool._lock: + pool._recycle_chunks() + pool._merge_chunks() + if pool.occupied_chunks: + raise RuntimeError( + "consumer acknowledgements did not recycle the slice" + ) + except Exception as exc: # noqa: BLE001 # pragma: no cover + result_queue.put(("error", repr(exc))) + return + finally: + del proxy, source + if pool is not None: + pool.shutdown() + del pool + gc.collect() + result_queue.put(("ok", None)) + + +class TestCudaVmmTransport(CustomTestCase): + @classmethod + def setUpClass(cls): + if ( + not torch.cuda.is_available() + or torch.version.cuda is None + or torch.cuda.device_count() < 3 + ): + raise unittest.SkipTest("At least three NVIDIA CUDA GPUs are required") + + def _run_round_trip(self, mode: str): + consumer_devices = (1, 2) + torch.cuda.set_device(consumer_devices[0]) + ctx = mp.get_context("spawn") + proxy_queue = ctx.Queue() + producer_results = ctx.Queue() + consumer_done = ctx.Event() + producer = ctx.Process( + target=_produce_vmm_tensor, + args=(proxy_queue, consumer_done, producer_results, mode), + ) + producer.start() + proxy = second_proxy = None + reconstructed = [] + producer_result = None + try: + try: + proxy, expected = proxy_queue.get(timeout=60) + except queue.Empty: + producer_result = producer_results.get(timeout=5) + _status, payload = producer_result + self.fail( + f"CUDA VMM producer failed before sending its proxy: {payload}" + ) + + if mode == "posix_fallback": + self.assertIsNone(proxy.fabric_handle) + self.assertIsNotNone(proxy.posix_socket_path) + else: + self.assertIsNotNone(proxy.fabric_handle) + self.assertIsNone(proxy.posix_socket_path) + second_proxy = pickle.loads(pickle.dumps(proxy)) + for tp_rank, (consumer_proxy, device) in enumerate( + zip((proxy, second_proxy), consumer_devices) + ): + torch.cuda.set_device(device) + with get_parallel().override( + attn_tp_size=2, + attn_tp_rank=tp_rank, + attn_cp_size=1, + attn_cp_rank=0, + ): + tensor = consumer_proxy.reconstruct_on_target_device( + device, consumer_count=1 + ) + torch.cuda.synchronize(device) + self.assertEqual(tensor.cpu().tolist(), expected) + reconstructed.append(tensor) + finally: + del reconstructed, second_proxy, proxy + _imported_pool_cache_clear() + gc.collect() + consumer_done.set() + producer.join(timeout=60) + try: + if producer_result is None: + producer_result = producer_results.get(timeout=5) + status, payload = producer_result + self.assertEqual(status, "ok", payload) + finally: + if producer.is_alive(): + producer.terminate() + producer.join(timeout=10) + torch.cuda.set_device(0) + self.assertEqual(producer.exitcode, 0) + + def test_posix_fd_fallback_tensor_round_trip_and_recycling(self): + self._run_round_trip(mode="posix_fallback") + + def test_auto_prefers_fabric_tensor_round_trip_and_recycling(self): + self._run_round_trip(mode="auto") + + def test_reused_chunk_clears_acknowledgements(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 2, allow_posix_fallback=True) + try: + old = pool.wrap_tensor(torch.ones(1024, dtype=torch.uint8, device="cuda:0")) + pool.memory_pool[old.control_offset : old.control_offset + 8].view( + torch.int32 + ).fill_(1) + torch.cuda.synchronize(0) + with pool._lock: + pool._recycle_chunks() + pool._merge_chunks() + + pool.wrap_tensor(torch.ones(100, dtype=torch.uint8, device="cuda:0")) + live = pool.wrap_tensor(torch.ones(256, dtype=torch.uint8, device="cuda:0")) + control = pool.memory_pool[ + live.control_offset : live.control_offset + 8 + ].view(torch.int32) + self.assertTrue(torch.equal(control, torch.zeros_like(control))) + + with pool._lock: + pool._recycle_chunks() + self.assertIn( + live.control_offset, + [chunk.start for chunk in pool.occupied_chunks], + ) + finally: + pool.shutdown() + + def test_packed_tensors_round_trip_through_one_shared_buffer(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + sources = [ + torch.arange(24, dtype=torch.float32, device="cuda:0") + .reshape(4, 6) + .transpose(0, 1), + torch.arange(7, dtype=torch.bfloat16), + torch.arange(5, dtype=torch.int64, device="cuda:0"), + ] + expected = [source.contiguous().cpu() for source in sources] + proxies = reconstructed = None + try: + stream = MagicMock(wraps=torch.cuda.current_stream(0)) + with patch("torch.cuda.current_stream", return_value=stream): + proxies = pool.wrap_tensors(sources) + + self.assertIsNotNone(proxies) + self.assertEqual(stream.synchronize.call_count, 1) + self.assertEqual(len(pool.occupied_chunks), 1) + self.assertTrue( + all( + isinstance(proxy, CudaVmmPackedTensorTransportProxy) + for proxy in proxies + ) + ) + self.assertEqual(len({proxy.control_offset for proxy in proxies}), 1) + + proxies = pickle.loads(pickle.dumps(proxies)) + self.assertIs(proxies[0]._packed_owner, proxies[-1]._packed_owner) + with get_parallel().override( + attn_tp_size=1, + attn_tp_rank=0, + attn_cp_size=1, + attn_cp_rank=0, + ): + reconstructed = [ + proxy.reconstruct_on_target_device(0, consumer_count=1) + for proxy in proxies + ] + torch.cuda.synchronize(0) + + for actual, wanted in zip(reconstructed, expected): + self.assertTrue(torch.equal(actual.cpu(), wanted)) + packed_storage = proxies[ + 0 + ]._packed_owner.reconstruct_tensor.untyped_storage() + self.assertTrue( + all( + tensor.untyped_storage().data_ptr() == packed_storage.data_ptr() + for tensor in reconstructed + ) + ) + with pool._lock: + pool._recycle_chunks() + self.assertFalse(pool.occupied_chunks) + finally: + del reconstructed, proxies, expected, sources + _imported_pool_cache_clear() + pool.shutdown() + + def test_packed_cancel_is_shared_and_idempotent(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + proxies = pool.wrap_tensors( + [ + torch.ones(8, dtype=torch.float32, device="cuda:0"), + torch.ones(8, dtype=torch.float32, device="cuda:0"), + ] + ) + self.assertIsNotNone(proxies) + + pool.cancel_proxy(proxies[0]) + pool.cancel_proxy(proxies[1]) + + self.assertFalse(pool.occupied_chunks) + self.assertEqual( + sum(chunk.size for chunk in pool.available_chunks), + pool.allocation_size, + ) + finally: + pool.shutdown() + + def test_packed_reservation_failure_returns_fallback_signal(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + source = torch.empty( + pool.allocation_size, dtype=torch.uint8, device="cuda:0" + ) + + self.assertIsNone(pool.wrap_tensors([source])) + self.assertFalse(pool.occupied_chunks) + self.assertEqual( + sum(chunk.size for chunk in pool.available_chunks), + pool.allocation_size, + ) + finally: + pool.shutdown() + + def test_oversized_tensor_falls_back_to_cpu(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + source = torch.empty( + pool.allocation_size, dtype=torch.uint8, device="cuda:0" + ) + + fallback = pool.wrap_tensor(source) + + self.assertTrue(fallback.is_cpu) + self.assertFalse(pool.occupied_chunks) + self.assertEqual( + sum(chunk.size for chunk in pool.available_chunks), + pool.allocation_size, + ) + finally: + pool.shutdown() + + def test_failed_copy_rolls_back_reservation(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + with self.assertRaisesRegex(NotImplementedError, "meta tensor"): + pool.wrap_tensor(torch.ones(16, device="meta")) + with self.assertRaisesRegex(NotImplementedError, "meta tensor"): + pool.wrap_tensors( + [ + torch.ones(16, device="cuda:0"), + torch.ones(16, device="meta"), + ] + ) + self.assertFalse(pool.occupied_chunks) + self.assertEqual( + sum(chunk.size for chunk in pool.available_chunks), + pool.allocation_size, + ) + finally: + pool.shutdown() + + def test_undispatched_proxy_can_be_cancelled_immediately(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + proxy = pool.wrap_tensor(torch.ones(16, device="cuda:0")) + + pool.cancel_proxy(proxy) + + self.assertFalse(pool.occupied_chunks) + self.assertEqual( + sum(chunk.size for chunk in pool.available_chunks), + pool.allocation_size, + ) + finally: + pool.shutdown() + + def test_failed_cleanup_sync_quarantines_pool(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + stream = MagicMock() + stream.synchronize.side_effect = RuntimeError("forced sync failure") + try: + with ( + patch("torch.cuda.current_stream", return_value=stream), + self.assertRaisesRegex(RuntimeError, "forced sync failure"), + ): + pool.wrap_tensor(torch.ones(16, device="cuda:0")) + torch.cuda.synchronize(0) + self.assertIsNotNone(pool._pool_error) + with self.assertRaisesRegex(RuntimeError, "pool failed"): + pool.wrap_tensor(torch.ones(16, device="cuda:0")) + finally: + pool.shutdown() + + def test_shutdown_waits_for_active_publisher(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + publisher_entered = threading.Event() + allow_publisher_to_finish = threading.Event() + shutdown_entered = threading.Event() + shutdown_finished = threading.Event() + errors = [] + real_stream = torch.cuda.current_stream(0) + stream = MagicMock(wraps=real_stream) + + def synchronize(): + publisher_entered.set() + if not allow_publisher_to_finish.wait(timeout=10): + raise TimeoutError("publisher was not released") + real_stream.synchronize() + + stream.synchronize.side_effect = synchronize + + def publish(): + try: + with patch("torch.cuda.current_stream", return_value=stream): + pool.wrap_tensor(torch.ones(16, device="cuda:0")) + except Exception as error: # pragma: no cover + errors.append(error) + + def shutdown(): + shutdown_entered.set() + try: + pool.shutdown() + except Exception as error: # pragma: no cover + errors.append(error) + finally: + shutdown_finished.set() + + publisher = threading.Thread(target=publish) + shutdown_thread = threading.Thread(target=shutdown) + try: + publisher.start() + self.assertTrue(publisher_entered.wait(timeout=10)) + shutdown_thread.start() + self.assertTrue(shutdown_entered.wait(timeout=10)) + self.assertFalse(shutdown_finished.wait(timeout=0.1)) + finally: + allow_publisher_to_finish.set() + publisher.join(timeout=10) + shutdown_thread.join(timeout=10) + if not pool._closed: + pool.shutdown() + + self.assertFalse(publisher.is_alive()) + self.assertFalse(shutdown_thread.is_alive()) + self.assertFalse(errors) + + def test_consumer_copy_failure_releases_slice_without_allowing_retry(self): + pool = CudaVmmMemoryPool(4 << 20, 60, 0, 1, allow_posix_fallback=True) + try: + proxy = pool.wrap_tensor(torch.ones(1, device="cuda:0")) + proxy.shape = (2,) + with ( + get_parallel().override( + attn_tp_size=1, + attn_tp_rank=0, + attn_cp_size=1, + attn_cp_rank=0, + ), + self.assertRaises(RuntimeError), + ): + proxy.reconstruct_on_target_device(0, consumer_count=1) + torch.cuda.synchronize(0) + with pool._lock: + pool._recycle_chunks() + self.assertFalse(pool.occupied_chunks) + + with ( + get_parallel().override( + attn_tp_size=1, + attn_tp_rank=0, + attn_cp_size=1, + attn_cp_rank=0, + ), + self.assertRaisesRegex(RuntimeError, "already released"), + ): + proxy.reconstruct_on_target_device(0, consumer_count=1) + finally: + _imported_pool_cache_clear() + pool.shutdown() + + def test_posix_export_fd_closes_when_allocation_setup_fails(self): + with ( + patch( + "sglang.srt.utils.cuda_vmm_transport_utils._tensor_from_pointer", + side_effect=RuntimeError("forced storage failure"), + ), + patch( + "sglang.srt.utils.cuda_vmm_transport_utils.os.close", + wraps=os.close, + ) as close_fd, + self.assertRaisesRegex(RuntimeError, "forced storage failure"), + ): + _FabricUnavailableCudaVmmMemoryPool( + 4 << 20, 60, 0, 1, allow_posix_fallback=True + ) + close_fd.assert_called_once() + + def test_stream_setup_failure_releases_pool_and_posix_broker(self): + release_allocation = CudaVmmMemoryPool._release_allocation + close_broker = _PosixFdBroker.close + with ( + patch( + "sglang.srt.utils.cuda_vmm_transport_utils.torch.cuda.Stream", + side_effect=RuntimeError("forced stream failure"), + ), + patch.object( + CudaVmmMemoryPool, + "_release_allocation", + autospec=True, + side_effect=release_allocation, + ) as release_pool, + patch.object( + _PosixFdBroker, + "close", + autospec=True, + side_effect=close_broker, + ) as close_fd_broker, + self.assertRaisesRegex(RuntimeError, "forced stream failure"), + ): + _FabricUnavailableCudaVmmMemoryPool( + 4 << 20, 60, 0, 1, allow_posix_fallback=True + ) + + close_fd_broker.assert_called_once() + release_pool.assert_called_once() + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/multimodal/test_gpu_feature_transport.py b/test/registered/unit/multimodal/test_gpu_feature_transport.py new file mode 100644 index 000000000..6e802c71d --- /dev/null +++ b/test/registered/unit/multimodal/test_gpu_feature_transport.py @@ -0,0 +1,666 @@ +import unittest +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, call, patch + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestCudaVmmFeatureTransport(unittest.TestCase): + def test_partial_pool_release_can_be_retried(self): + from sglang.srt.utils import cuda_vmm_transport_utils as vmm + + pool = object.__new__(vmm.CudaVmmMemoryPool) + pool.memory_pool = object() + pool.use_fabric = True + pool.shareable_handle = b"handle" + pool._pool_pointer = 123 + pool._allocation_handle = 456 + pool._allocation_mapped = True + pool.allocation_size = 4096 + pool.device_index = 0 + driver = MagicMock() + driver.cuMemUnmap.return_value = "unmap" + driver.cuMemAddressFree.return_value = "address_free" + driver.cuMemRelease.return_value = "release" + failed_once = False + + def check_driver(result, _operation): + nonlocal failed_once + if result == "address_free" and not failed_once: + failed_once = True + raise RuntimeError("forced address-free failure") + return result + + with ( + patch.object(vmm, "_get_cuda_driver", return_value=driver), + patch.object(vmm.torch.cuda, "device", return_value=nullcontext()), + patch.object(vmm, "check_drv", side_effect=check_driver), + self.assertRaisesRegex(RuntimeError, "forced address-free failure"), + ): + pool._release_allocation() + + self.assertFalse(pool._allocation_mapped) + self.assertEqual(pool._pool_pointer, 123) + self.assertEqual(pool._allocation_handle, 456) + + with ( + patch.object(vmm, "_get_cuda_driver", return_value=driver), + patch.object(vmm.torch.cuda, "device", return_value=nullcontext()), + patch.object(vmm, "check_drv", side_effect=lambda result, _: result), + ): + pool._release_allocation() + + self.assertIsNone(pool._pool_pointer) + self.assertIsNone(pool._allocation_handle) + self.assertEqual(driver.cuMemUnmap.call_count, 1) + self.assertEqual(driver.cuMemAddressFree.call_count, 2) + self.assertEqual(driver.cuMemRelease.call_count, 1) + + def test_model_class_controls_cuda_vmm_opt_in(self): + from sglang.srt.managers.tokenizer_manager import TokenizerManager + + class SupportedModel: + supports_cuda_vmm_feature_transport = True + + class UnsupportedModel: + pass + + manager = object.__new__(TokenizerManager) + manager.server_args = SimpleNamespace(mm_feature_transport="cuda_vmm") + manager.model_config = object() + + with patch( + "sglang.srt.model_loader.utils.get_model_architecture", + return_value=(SupportedModel, "supported"), + ): + manager._validate_cuda_vmm_feature_transport_support() + + with ( + patch( + "sglang.srt.model_loader.utils.get_model_architecture", + return_value=(UnsupportedModel, "unsupported"), + ), + self.assertRaisesRegex(ValueError, "UnsupportedModel"), + ): + manager._validate_cuda_vmm_feature_transport_support() + + def test_cpu_transport_skips_model_opt_in_lookup(self): + from sglang.srt.managers.tokenizer_manager import TokenizerManager + + manager = object.__new__(TokenizerManager) + manager.server_args = SimpleNamespace(mm_feature_transport="cpu") + manager.model_config = object() + + with patch( + "sglang.srt.model_loader.utils.get_model_architecture" + ) as get_model_architecture: + manager._validate_cuda_vmm_feature_transport_support() + + get_model_architecture.assert_not_called() + + def test_vmm_transport_initializes_pool(self): + from sglang.srt.utils import cuda_vmm_transport_utils as vmm + + server_args = SimpleNamespace( + mm_feature_transport="cuda_vmm", + tokenizer_worker_num=2, + base_gpu_id=3, + enable_dp_attention=False, + tp_size=4, + nnodes=1, + ) + pool = object() + with ( + patch.object(vmm, "get_mm_feature_pool_size_per_worker", return_value=123), + patch.object(vmm, "CudaVmmMemoryPool", return_value=pool) as pool_class, + ): + transport = vmm.CudaVmmFeatureTransport(server_args, SimpleNamespace()) + + self.assertIs(transport.pool, pool) + pool_class.assert_called_once_with( + memory_size=123, + recycle_interval=vmm.MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL, + base_gpu_id=3, + consumer_count=4, + allow_posix_fallback=True, + ) + + def test_disabled_transport_is_a_noop(self): + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = CudaVmmFeatureTransport( + SimpleNamespace(mm_feature_transport="cpu"), None + ) + + self.assertEqual(transport.prepare_for_dispatch([None]), []) + transport.cancel_for_dispatch([]) + transport.shutdown() + self.assertIsNone(transport.pool) + + def test_vmm_transport_requires_processor(self): + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + with self.assertRaisesRegex(RuntimeError, "multimodal processor"): + CudaVmmFeatureTransport( + SimpleNamespace(mm_feature_transport="cuda_vmm"), None + ) + + def test_image_features_are_packed_per_request(self): + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = MagicMock() + features = [torch.arange(4), torch.arange(4, 8)] + proxies = [object(), object()] + transport.pool.wrap_tensors.return_value = proxies + items = [ + MultimodalDataItem(modality=Modality.IMAGE, feature=feature) + for feature in features + ] + + transport.wrap_items(items) + + transport.pool.wrap_tensors.assert_called_once_with(features) + transport.pool.wrap_tensor.assert_not_called() + self.assertEqual([item.feature for item in items], proxies) + + def test_deferred_features_are_not_packed(self): + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.utils.cuda_ipc_transport_utils import ( + DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, + ) + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = MagicMock() + features = [torch.arange(4), torch.arange(4, 8)] + proxies = [object(), object()] + transport.pool.wrap_tensor.side_effect = proxies + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + feature=feature, + model_specific_data={DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY: True}, + ) + for feature in features + ] + + transport.wrap_items(items) + + transport.pool.wrap_tensors.assert_not_called() + self.assertEqual( + transport.pool.wrap_tensor.call_args_list, + [call(feature) for feature in features], + ) + self.assertEqual([item.feature for item in items], proxies) + + def test_tensor_containers_fail_closed(self): + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = MagicMock() + item = MultimodalDataItem( + modality=Modality.IMAGE, + feature=[torch.arange(4), torch.arange(4, 8)], + ) + + with self.assertRaisesRegex(TypeError, "single tensor"): + transport.wrap_items([item]) + + transport.pool.wrap_tensor.assert_not_called() + transport.pool.wrap_tensors.assert_not_called() + + def test_partial_failure_restores_tensors_and_cancels_packed_chunk_once(self): + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + CudaVmmMemoryPool, + CudaVmmPackedTensorTransportProxy, + _CudaVmmPackedTransportOwner, + ) + + owner = object.__new__(_CudaVmmPackedTransportOwner) + owner.control_offset = 64 + owner._producer_cancelled = False + proxies = [object.__new__(CudaVmmPackedTensorTransportProxy) for _ in range(2)] + for proxy in proxies: + proxy._packed_owner = owner + + pool = object.__new__(CudaVmmMemoryPool) + pool.wrap_tensors = MagicMock(return_value=proxies) + pool.wrap_tensor = MagicMock(side_effect=RuntimeError("copy failed")) + pool._cancel_control_offset = MagicMock() + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = pool + + features = [torch.arange(4), torch.arange(4, 8)] + embedding = torch.arange(2) + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + feature=features[0], + precomputed_embeddings=embedding, + ), + MultimodalDataItem(modality=Modality.IMAGE, feature=features[1]), + ] + + with self.assertRaisesRegex(RuntimeError, "copy failed"): + transport.wrap_items(items) + + for item, feature in zip(items, features, strict=True): + self.assertIs(item.feature, feature) + self.assertIs(items[0].precomputed_embeddings, embedding) + pool._cancel_control_offset.assert_called_once_with(owner.control_offset) + + def test_text_request_uses_base_send_path(self): + from sglang.srt.managers import tokenizer_manager + from sglang.srt.managers.tokenizer_manager import TokenizerManager + + manager = object.__new__(TokenizerManager) + transport = MagicMock() + transport.prepare_for_dispatch.return_value = [] + manager.cuda_vmm_feature_transport = transport + manager._dispatch_to_scheduler = MagicMock() + state = SimpleNamespace(dispatched=False) + manager.rid_to_state = {"test-request": state} + tokenized_obj = SimpleNamespace( + rid="test-request", + mm_inputs=None, + time_stats=MagicMock(), + wrap_pickle_fields=MagicMock(), + ) + + with patch.object(tokenizer_manager, "wrap_shm_features", lambda obj: obj): + manager._send_one_request(tokenized_obj) + + manager._dispatch_to_scheduler.assert_called_once_with(tokenized_obj) + transport.prepare_for_dispatch.assert_called_once_with((None,)) + transport.cancel_for_dispatch.assert_not_called() + self.assertTrue(state.dispatched) + + def test_failed_dispatch_cancels_published_items(self): + from sglang.srt.managers import tokenizer_manager + from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, + ) + + manager = object.__new__(tokenizer_manager.TokenizerManager) + transport = MagicMock() + manager._dispatch_to_scheduler = MagicMock( + side_effect=RuntimeError("send failed") + ) + state = SimpleNamespace(dispatched=False) + manager.rid_to_state = {"test-request": state} + items = [MultimodalDataItem(modality=Modality.IMAGE, feature=torch.arange(2))] + tokenized_obj = SimpleNamespace( + rid="test-request", + mm_inputs=MultimodalProcessorOutput(input_ids=[1], mm_items=items), + time_stats=MagicMock(), + wrap_pickle_fields=MagicMock(), + ) + transport.prepare_for_dispatch.return_value = items + manager.cuda_vmm_feature_transport = transport + + with ( + patch.object(tokenizer_manager, "wrap_shm_features", lambda obj: obj), + self.assertRaisesRegex(RuntimeError, "send failed"), + ): + manager._send_one_request(tokenized_obj) + + transport.prepare_for_dispatch.assert_called_once_with( + (tokenized_obj.mm_inputs,) + ) + transport.cancel_for_dispatch.assert_called_once_with(items) + self.assertFalse(state.dispatched) + + def test_post_dispatch_failure_does_not_cancel_published_items(self): + from sglang.srt.managers import tokenizer_manager + from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, + ) + + manager = object.__new__(tokenizer_manager.TokenizerManager) + transport = MagicMock() + manager._dispatch_to_scheduler = MagicMock() + state = SimpleNamespace(dispatched=False) + manager.rid_to_state = {"test-request": state} + time_stats = MagicMock() + time_stats.set_api_server_dispatch_finish_time.side_effect = RuntimeError( + "bookkeeping failed" + ) + items = [MultimodalDataItem(modality=Modality.IMAGE, feature=torch.arange(2))] + tokenized_obj = SimpleNamespace( + rid="test-request", + mm_inputs=MultimodalProcessorOutput(input_ids=[1], mm_items=items), + time_stats=time_stats, + wrap_pickle_fields=MagicMock(), + ) + transport.prepare_for_dispatch.return_value = items + manager.cuda_vmm_feature_transport = transport + + with ( + patch.object(tokenizer_manager, "wrap_shm_features", lambda obj: obj), + self.assertRaisesRegex(RuntimeError, "bookkeeping failed"), + ): + manager._send_one_request(tokenized_obj) + + manager._dispatch_to_scheduler.assert_called_once_with(tokenized_obj) + transport.cancel_for_dispatch.assert_not_called() + self.assertTrue(state.dispatched) + + def test_prepare_batch_cancels_prior_groups_on_failure(self): + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = MagicMock() + transport.wrap_items = MagicMock( + side_effect=[None, RuntimeError("wrap failed")] + ) + transport.cancel_for_dispatch = MagicMock() + item_groups = [[object()], [object()]] + mm_inputs_batch = [SimpleNamespace(mm_items=items) for items in item_groups] + + with self.assertRaisesRegex(RuntimeError, "wrap failed"): + transport.prepare_for_dispatch(mm_inputs_batch) + + self.assertEqual( + transport.wrap_items.call_args_list, + [call(item_groups[0]), call(item_groups[1])], + ) + transport.cancel_for_dispatch.assert_called_once_with(item_groups[0]) + + def test_prepare_batch_returns_flattened_items(self): + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + transport.pool = MagicMock() + transport.wrap_items = MagicMock() + item_groups = [[object()], [object(), object()]] + + prepared = transport.prepare_for_dispatch( + [ + None, + SimpleNamespace(mm_items=[]), + *(SimpleNamespace(mm_items=items) for items in item_groups), + ] + ) + + self.assertEqual(prepared, item_groups[0] + item_groups[1]) + self.assertEqual( + transport.wrap_items.call_args_list, + [call(items) for items in item_groups], + ) + + def test_engine_shutdown_is_idempotent(self): + from sglang.srt.entrypoints import engine as engine_module + from sglang.srt.entrypoints.engine import Engine + from sglang.srt.managers.tokenizer_manager import TokenizerManager + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + manager = object.__new__(TokenizerManager) + transport = object.__new__(CudaVmmFeatureTransport) + pool = MagicMock() + transport.pool = pool + manager.cuda_vmm_feature_transport = transport + manager._subprocess_watchdog = None + engine = object.__new__(Engine) + engine.tokenizer_manager = manager + + with patch.object( + engine_module, + "kill_process_tree", + side_effect=RuntimeError("base failed"), + ): + for _ in range(2): + with self.assertRaisesRegex(RuntimeError, "base failed"): + engine.shutdown() + + self.assertEqual(pool.shutdown.call_count, 2) + self.assertIs(transport.pool, pool) + + def test_engine_startup_failure_releases_parent_pool(self): + from sglang.srt.entrypoints import engine as engine_module + from sglang.srt.entrypoints.engine import Engine + from sglang.srt.managers.tokenizer_manager import TokenizerManager + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + manager = object.__new__(TokenizerManager) + transport = object.__new__(CudaVmmFeatureTransport) + pool = MagicMock() + transport.pool = pool + manager.cuda_vmm_feature_transport = transport + server_args = SimpleNamespace( + remote_instance_weight_loader_start_seed_via_transfer_engine=False, + reasoning_parser=None, + tool_call_parser=None, + weight_cache_mode=None, + enable_elastic_expert_backup=False, + elastic_ep_backend=None, + node_rank=0, + tokenizer_worker_num=1, + check_server_args=MagicMock(), + ) + scheduler_init_result = SimpleNamespace( + all_child_pids=[], + scheduler_infos=[], + wait_for_ready=MagicMock(side_effect=RuntimeError("startup failed")), + engine_info_bootstrap_server=None, + ) + + with ( + patch.object(engine_module, "configure_logger"), + patch.object(engine_module, "_set_envs_and_config"), + patch.object(engine_module, "load_plugins"), + patch.object( + Engine, + "_launch_scheduler_processes", + return_value=(scheduler_init_result, []), + ), + patch.object( + Engine, "_launch_detokenizer_subprocesses", return_value=([], []) + ), + self.assertRaisesRegex(RuntimeError, "startup failed"), + ): + Engine._launch_subprocesses( + server_args=server_args, + init_tokenizer_manager_func=MagicMock(return_value=(manager, object())), + run_scheduler_process_func=MagicMock(), + run_detokenizer_process_func=MagicMock(), + port_args=SimpleNamespace(), + ) + + pool.shutdown.assert_called_once_with() + self.assertIs(transport.pool, pool) + + def test_failed_pool_shutdown_remains_retryable(self): + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmFeatureTransport, + ) + + transport = object.__new__(CudaVmmFeatureTransport) + pool = MagicMock() + pool.shutdown.side_effect = [RuntimeError("shutdown failed"), None] + transport.pool = pool + + with self.assertRaisesRegex(RuntimeError, "shutdown failed"): + transport.shutdown() + self.assertIs(transport.pool, pool) + + transport.shutdown() + self.assertIs(transport.pool, pool) + + +class TestSchedulerMmTransportBoundary(unittest.TestCase): + @staticmethod + def _prepare_scheduler(scheduler): + scheduler.session_controller = SimpleNamespace(maybe_reap=MagicMock()) + scheduler._request_dispatcher = MagicMock(return_value=None) + scheduler.flush_wrapper = SimpleNamespace(check_pending=MagicMock()) + scheduler.external_corpus_manager = None + + def test_materializes_inputs_directly_before_base_dispatch(self): + from sglang.srt.managers import scheduler as scheduler_module + + scheduler = object.__new__(scheduler_module.Scheduler) + scheduler.server_args = SimpleNamespace( + mm_feature_transport="cuda_vmm", + enable_broadcast_mm_inputs_process=True, + ) + self._prepare_scheduler(scheduler) + raw_inputs = object() + materialized = object() + request = SimpleNamespace(mm_inputs=raw_inputs) + + with ( + patch.object( + scheduler_module, "TokenizedGenerateReqInput", SimpleNamespace + ), + patch.object( + scheduler_module.MultimodalInputs, + "from_processor_output", + return_value=materialized, + ) as build_inputs, + patch.object( + scheduler, "_process_and_broadcast_mm_inputs" + ) as cpu_broadcast, + patch.object( + scheduler_module, "is_health_check_generate_req", return_value=False + ), + ): + scheduler.process_input_requests([request]) + + build_inputs.assert_called_once_with(raw_inputs) + self.assertIs(request.mm_inputs, materialized) + scheduler._request_dispatcher.assert_called_once_with(request) + cpu_broadcast.assert_not_called() + + def test_materializes_batched_inputs_before_dispatch(self): + from sglang.srt.managers import scheduler as scheduler_module + + class TokenizedRequest: + def __init__(self, mm_inputs): + self.mm_inputs = mm_inputs + + class BatchRequest: + def __init__(self, batch): + self.batch = batch + + def __iter__(self): + return iter(self.batch) + + scheduler = object.__new__(scheduler_module.Scheduler) + scheduler.server_args = SimpleNamespace(mm_feature_transport="cuda_vmm") + self._prepare_scheduler(scheduler) + raw_inputs = [object(), object()] + materialized = [object(), object()] + inner_requests = [TokenizedRequest(value) for value in raw_inputs] + request = BatchRequest(inner_requests) + + with ( + patch.object( + scheduler_module, "TokenizedGenerateReqInput", TokenizedRequest + ), + patch.object( + scheduler_module, "TokenizedEmbeddingReqInput", TokenizedRequest + ), + patch.object( + scheduler_module, "BatchTokenizedGenerateReqInput", BatchRequest + ), + patch.object( + scheduler_module, "BatchTokenizedEmbeddingReqInput", BatchRequest + ), + patch.object( + scheduler_module.MultimodalInputs, + "from_processor_output", + side_effect=materialized, + ) as build_inputs, + patch.object( + scheduler_module, "is_health_check_generate_req", return_value=False + ), + ): + scheduler.process_input_requests([request]) + + self.assertEqual( + build_inputs.call_args_list, + [call(value) for value in raw_inputs], + ) + self.assertEqual( + [inner.mm_inputs for inner in inner_requests], + materialized, + ) + scheduler._request_dispatcher.assert_called_once_with(request) + + def test_already_materialized_inputs_are_reused(self): + from sglang.srt.managers.schedule_batch import MultimodalInputs + from sglang.srt.managers.scheduler import Scheduler + + scheduler = object.__new__(Scheduler) + mm_inputs = MultimodalInputs(mm_items=[]) + + with patch.object( + scheduler, "_process_and_broadcast_mm_inputs" + ) as process_and_broadcast: + self.assertIs(scheduler._get_multimodal_inputs(mm_inputs), mm_inputs) + + process_and_broadcast.assert_not_called() + + +class TestVmmConsumerCount(unittest.TestCase): + def test_proxy_defaults_to_one_consumer(self): + from sglang.srt.utils import cuda_vmm_transport_utils as vmm + + proxy = object.__new__(vmm.CudaVmmTensorTransportProxy) + proxy.consumer_count = 4 + self.assertEqual(proxy._resolve_consumer_count(None), 1) + self.assertEqual(proxy._resolve_consumer_count(2), 2) + + def test_acknowledgement_ranges_include_cp_rank(self): + from sglang.srt.runtime_context import get_parallel + from sglang.srt.utils.cuda_vmm_transport_utils import ( + CudaVmmTensorTransportProxy, + ) + + proxy = object.__new__(CudaVmmTensorTransportProxy) + proxy.consumer_count = 4 + with get_parallel().override( + attn_tp_size=2, + attn_tp_rank=1, + attn_cp_size=2, + attn_cp_rank=1, + ): + self.assertEqual(proxy._acknowledgement_range(1), (3, 4)) + self.assertEqual(proxy._acknowledgement_range(2), (2, 4)) + self.assertEqual(proxy._acknowledgement_range(4), (0, 4)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 6fbdbeee9..e4a000c10 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -259,6 +259,32 @@ class TestMultimodalFeatureTransport(CustomTestCase): with self.assertRaisesRegex(ValueError, "single node"): server_args._handle_multimodal_feature_transport() + @patch("sglang.srt.server_args.is_cuda", return_value=False) + def test_cuda_vmm_rejects_non_nvidia_platforms(self, _mock_is_cuda): + server_args = ServerArgs(model_path="dummy", mm_feature_transport="cuda_vmm") + + with self.assertRaisesRegex(ValueError, "requires NVIDIA CUDA"): + server_args._handle_multimodal_feature_transport() + + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_cuda_vmm_rejects_rust_server(self, _mock_is_cuda): + server_args = ServerArgs(model_path="dummy", mm_feature_transport="cuda_vmm") + + with ( + envs.SGLANG_RUST_SERVER.override(True), + self.assertRaisesRegex(ValueError, "SGLANG_RUST_SERVER"), + ): + server_args._handle_multimodal_feature_transport() + + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_cuda_vmm_rejects_pipeline_parallelism(self, _mock_is_cuda): + server_args = ServerArgs( + model_path="dummy", mm_feature_transport="cuda_vmm", pp_size=2 + ) + + with self.assertRaisesRegex(ValueError, "pipeline parallelism"): + server_args._handle_multimodal_feature_transport() + class TestMambaCacheStochasticRounding(unittest.TestCase): def test_rejects_fp32_ssm_cache(self):