From d747bd052ef9bae97459b57b522762e524f62d68 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 8 Aug 2026 16:00:58 +0800 Subject: [PATCH] feat(vlm): auto-select cuda vmm on multi-node mnnvl (#33936) --- .../autoregressive/Moonshotai/Kimi-K3.mdx | 10 +- python/sglang/srt/managers/schedule_batch.py | 20 ++- python/sglang/srt/model_loader/utils.py | 5 + python/sglang/srt/models/kimi_k25.py | 2 + python/sglang/srt/models/kimi_k3.py | 2 + python/sglang/srt/models/qwen3_vl.py | 2 + .../srt/multimodal/processors/kimi_k25.py | 4 +- .../srt/multimodal/processors/kimi_k3.py | 2 +- python/sglang/srt/server_args.py | 84 ++++++++---- test/registered/unit/models/test_kimi_k25.py | 23 ++++ .../unit/server_args/test_server_args.py | 120 ++++++++++++++++++ 11 files changed, 245 insertions(+), 29 deletions(-) diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index c3ac3d27f..32c1d402d 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -309,7 +309,13 @@ sglang serve \ --port 30000 ``` -- `--mm-feature-transport cuda_ipc` — single-node only: skips the CPU round trip, bounded pool (per-tensor CPU fallback when full), reserves up to `SGLANG_MM_FEATURE_CACHE_MB` on the base GPU. Multi-node recipes use CPU transport. +- `--mm-feature-transport cuda_ipc` — single-node only: skips the CPU round + trip and uses a bounded GPU pool. +- On multi-node GB200/GB300, SGLang automatically uses CUDA VMM with CUDA + FABRIC handles when an IMEX channel is available; otherwise it uses CPU. + Pass `--mm-feature-transport cpu` to opt out. +- CUDA IPC and CUDA VMM share the `SGLANG_MM_FEATURE_CACHE_MB` HBM budget + (1 GiB by default) and fall back to CPU per tensor when the pool is full. - 2 processor / 16 I/O workers are the measured defaults; more adds contention. - Leave `--mm-attention-backend` unset — auto-selected, with a correctness fallback. - Don't add `--mm-enable-dp-encoder`; K3 already shards images across TP ranks. @@ -321,7 +327,7 @@ sglang serve \ | PD | Supported. Image processing and ViT run on prefill; the PD transfer then moves both paged MLA KV and KDA recurrent state as described in [PD disaggregation](#pd-disaggregation). | | EPD | Supported on the public `kimi-k3` branch. Use an `--encoder-only` vision role and a `--language-only` prefill role; add the normal decode role for full EPD. See the [EPD guide](../../../docs/advanced_features/epd_disaggregation). | | MM encoder DP | Built in. K3 shards complete images across TP ranks, so leave `--mm-enable-dp-encoder` unset in unified, PD-prefill, and encoder-only roles. | -| CUDA IPC | Compatible with the local processor-to-scheduler path on a single-node unified or PD-prefill role. It does not replace `--encoder-transfer-backend` for EPD or the PD KV/KDA transfer, and its bounded pool consumes HBM. | +| MM feature transport | CUDA IPC is used on a single node. CUDA VMM with FABRIC handles is auto-selected for multi-node GB200/GB300 when IMEX is available. Both use a bounded HBM pool and are separate from EPD's `--encoder-transfer-backend` and the PD KV/KDA transfer. | | ViT BCG | Compatible with unified and encoder-only roles, but recommended only for repeated encoder shapes after measuring the HBM trade-off below. | #### Should ViT BCG be enabled? diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f414e60f4..02ff636cd 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -431,11 +431,14 @@ class MultimodalDataItem: def reconstruct(self, target_device: int, ipc_consumer_count: int = 1): """materialize cuda ipc proxy tensors in-place on target_device""" if isinstance(self.feature, CudaIpcTensorTransportProxy): - if ipc_consumer_count == 1: + consumer_count = self._resolve_transport_consumer_count( + self.feature, ipc_consumer_count + ) + if consumer_count == 1: self.feature = self.feature.reconstruct_on_target_device(target_device) else: self.feature = self.feature.reconstruct_on_target_device( - target_device, consumer_count=ipc_consumer_count + target_device, consumer_count=consumer_count ) if isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy): self.precomputed_embeddings = ( @@ -474,8 +477,21 @@ class MultimodalDataItem: def acknowledge_deferred_cuda_ipc_feature(self, consumer_count: int = 1): """Release a lazy IPC feature when an embedding-cache hit skips ViT.""" if isinstance(self.feature, CudaIpcTensorTransportProxy): + consumer_count = self._resolve_transport_consumer_count( + self.feature, consumer_count + ) self.feature.acknowledge_consumption(consumer_count) + @staticmethod + def _resolve_transport_consumer_count(proxy, requested_count: int) -> int: + """Clamp a group acknowledgement to the proxy's actual consumer set.""" + proxy_count = getattr( + proxy, + "total_consumer_count", + getattr(proxy, "consumer_count", requested_count), + ) + return min(requested_count, proxy_count) + @dataclasses.dataclass class MultimodalProcessorOutput: diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 1611b8fcc..8ba2f94de 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -235,6 +235,11 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module], return model_cls, resolved_arch +def supports_cuda_vmm_feature_transport(model_config: ModelConfig) -> bool: + model_cls, _ = get_model_architecture(model_config) + return bool(getattr(model_cls, "supports_cuda_vmm_feature_transport", False)) + + def get_resolved_model_impl(model_config: ModelConfig) -> ModelImpl: resolved_model_impl = getattr(model_config, "_resolved_model_impl", None) if resolved_model_impl is not None: diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 36d1d18d5..80a1ae4b8 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -637,6 +637,8 @@ def mm_projection_auto( class KimiK25ForConditionalGeneration(nn.Module): + supports_cuda_vmm_feature_transport = True + # Support nvidia/Kimi-K2.5-NVFP4 naming: language_model.layers.*. # Ref: HF config.json for nvidia/Kimi-K2.5-NVFP4 # https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/config.json diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 1d07b16a0..98a213b9a 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -2933,6 +2933,8 @@ class KimiK3LinearForCausalLM(nn.Module): class KimiK3ForConditionalGeneration(nn.Module): """K3 multimodal wrapper: MoonViT3d tower + KimiK3LinearForCausalLM.""" + supports_cuda_vmm_feature_transport = True + # Raw HF checkpoint prefixes, before hf_to_sglang_mapper is applied. encoder_only_safetensors_weight_prefixes = ( "vision_tower.", diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 83697f140..20a468836 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -1210,6 +1210,8 @@ class Qwen3LLMModel(Qwen3Model): class Qwen3VLForConditionalGeneration(nn.Module): + supports_cuda_vmm_feature_transport = True + # To ensure correct weight loading and mapping. hf_to_sglang_mapper = WeightsMapper( orig_to_new_substr={ diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py index 34816acb7..7c84069f5 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k25.py +++ b/python/sglang/srt/multimodal/processors/kimi_k25.py @@ -592,10 +592,10 @@ class KimiK2_5VLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): ) # K2.5/K2.7 encoder-DP assigns an image to exactly one TP rank. Keep - # its IPC proxy lazy until that assignment is known, avoiding a full + # its GPU transport proxy lazy until that assignment is known, avoiding a full # image copy to every rank. The scheduler only honors this marker once # the processor has already set the item's hash and pad value. - if self.use_cuda_ipc and self.server_args.mm_enable_dp_encoder: + if self.keep_mm_features_on_device and self.server_args.mm_enable_dp_encoder: for item in mm_items: item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = ( True diff --git a/python/sglang/srt/multimodal/processors/kimi_k3.py b/python/sglang/srt/multimodal/processors/kimi_k3.py index cf9df0563..0e6c67631 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k3.py +++ b/python/sglang/srt/multimodal/processors/kimi_k3.py @@ -389,7 +389,7 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): # that assignment is known: one tokenizer/scheduler crossing per # image instead of one per rank. K2.5 gates this on # --mm-enable-dp-encoder; K3 needs no flag. - if getattr(self, "use_cuda_ipc", False): + if self.keep_mm_features_on_device: for item in mm_items: item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = ( True diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e0c95ad2a..5db8727a8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2747,13 +2747,13 @@ class ServerArgs: mm_feature_transport: A[ 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. " + "pool, or a bounded CUDA VMM pool. " "Unset resolves automatically: multimodal models on single-node CUDA " - "deployments (without disaggregation) use cuda_ipc, everything else uses " - "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.", + "deployments (without disaggregation) use cuda_ipc; validated multi-node " + "GB200/GB300 MNNVL models use cuda_vmm when an IMEX channel is available; " + "all other deployments use cpu. GPU transports reserve " + "SGLANG_MM_FEATURE_CACHE_MB (default 1024 MiB) on the base GPU and fall " + "back to CPU transport when the pool is full.", NS("mm"), ] = None keep_mm_feature_on_device: A[ @@ -7591,10 +7591,10 @@ class ServerArgs: def _handle_multimodal_feature_transport(self): """Resolve multimodal feature transport before tokenizer workers start. - CUDA IPC is deliberately opt-in: its fixed pool lives on ``base_gpu_id`` - and reduces the memory left for model/KV-cache allocations. The legacy - flag and environment variable remain supported so existing deployments - continue to work, but both map to this single policy. + GPU transports use a fixed pool on ``base_gpu_id`` and therefore reduce + the memory left for model/KV-cache allocations. The legacy CUDA IPC flag + and environment variable remain supported so existing deployments map + to this single policy. """ requested_transport = self.mm_feature_transport legacy_ipc_is_set = envs.SGLANG_USE_CUDA_IPC_TRANSPORT.is_set() @@ -7631,20 +7631,48 @@ class ServerArgs: elif ( self.get_model_config().is_multimodal and is_cuda() - and self.nnodes == 1 and self.disaggregation_mode == "null" ): - # Auto policy: single-node CUDA serving defaults to the bounded - # CUDA-IPC pool for multimodal models. Text-only deployments do - # not need feature transport. Multi-node (IPC handles are - # intra-node) and PD-disaggregated deployments keep CPU transport. - # A full pool degrades to CPU transport per tensor. - requested_transport = "cuda_ipc" - logger.info( - "Multimodal feature transport auto-resolved to cuda_ipc " - "(single-node CUDA). Pass --mm-feature-transport=cpu to " - "opt out." - ) + # A full GPU pool always degrades to CPU transport per tensor. + # CUDA IPC is intra-node; multi-node auto-selection is limited + # to GB200/GB300 systems where the runtime already enables the + # MNNVL/IMEX communication stack. + if self.nnodes == 1: + requested_transport = "cuda_ipc" + logger.info( + "Multimodal feature transport auto-resolved to cuda_ipc " + "(single-node CUDA). Pass --mm-feature-transport=cpu to " + "opt out." + ) + elif is_mnnvl_fabric_device() and os.path.exists( + "/dev/nvidia-caps-imex-channels/channel0" + ): + from sglang.srt.model_loader.utils import ( + supports_cuda_vmm_feature_transport, + ) + + if supports_cuda_vmm_feature_transport(self.get_model_config()): + requested_transport = "cuda_vmm" + logger.info( + "Multimodal feature transport auto-resolved to " + "cuda_vmm (multi-node GB200/GB300 MNNVL). Pass " + "--mm-feature-transport=cpu to opt out." + ) + else: + requested_transport = "cpu" + logger.info( + "Multimodal feature transport auto-resolved to cpu: " + "the model has not opted into CUDA VMM transport." + ) + else: + requested_transport = "cpu" + if is_mnnvl_fabric_device(): + logger.info( + "Multimodal feature transport auto-resolved to cpu: " + "GB200/GB300 was detected but no IMEX channel is " + "mounted. Configure the MNNVL compute domain or pass " + "--mm-feature-transport=cuda_vmm after doing so." + ) else: requested_transport = "cpu" elif legacy_ipc_is_set and legacy_ipc_enabled != ( @@ -7681,6 +7709,18 @@ class ServerArgs: "--mm-feature-transport=cuda_vmm is not supported with " "SGLANG_RUST_SERVER." ) + pool_budget_mb = envs.SGLANG_MM_FEATURE_CACHE_MB.get() + handle_kind = "CUDA FABRIC" if self.nnodes > 1 else "POSIX FD" + logger.info( + "Using CUDA VMM for multimodal features with %s sharing: " + "reserving up to %d MiB on base GPU %d across %d tokenizer " + "worker(s). This reduces KV cache headroom; a full pool falls " + "back to inline CPU transport.", + handle_kind, + pool_budget_mb, + self.base_gpu_id, + self.tokenizer_worker_num, + ) if requested_transport == "cuda_ipc": if not is_cuda(): diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index 6f5f57ad7..2bc9dfeca 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/test/registered/unit/models/test_kimi_k25.py @@ -524,6 +524,28 @@ def test_kimi_lazy_ipc_feature_acknowledges_all_tp_consumers(): proxy.reconstruct_on_target_device.assert_called_once_with(0, consumer_count=8) +def test_kimi_lazy_vmm_feature_uses_proxy_consumer_count(): + proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy) + proxy.consumer_count = 2 + proxy.reconstruct_on_target_device = Mock(return_value=torch.randn(1, 2)) + item = MultimodalDataItem(modality=Modality.IMAGE, feature=proxy) + + item.reconstruct(0, ipc_consumer_count=8) + + proxy.reconstruct_on_target_device.assert_called_once_with(0, consumer_count=2) + + +def test_kimi_lazy_vmm_cache_hit_uses_proxy_consumer_count(): + proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy) + proxy.consumer_count = 2 + proxy.acknowledge_consumption = Mock() + item = MultimodalDataItem(modality=Modality.IMAGE, feature=proxy) + + item.acknowledge_deferred_cuda_ipc_feature(consumer_count=8) + + proxy.acknowledge_consumption.assert_called_once_with(2) + + class _Tokenizer: def encode(self, text, allowed_special=None): tokens = { @@ -675,6 +697,7 @@ def test_kimi_k3_rejects_silently_dropped_images(): def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries(): processor = object.__new__(KimiK3ImageProcessor) + processor.mm_feature_transport = "cpu" processor.mm_tokens = SimpleNamespace(image_token_id=99) processor.fast_load_mm_data = AsyncMock( return_value=SimpleNamespace( diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index e4a000c10..97b4e278d 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -181,6 +181,16 @@ class TestMultimodalFeatureTransport(CustomTestCase): self.assertIn("deprecated", logs.output[0]) + def test_legacy_keep_flag_rejects_explicit_cuda_vmm(self): + server_args = ServerArgs( + model_path="dummy", + keep_mm_feature_on_device=True, + mm_feature_transport="cuda_vmm", + ) + + with self.assertRaisesRegex(ValueError, "conflicts.*cuda_vmm"): + server_args._handle_multimodal_feature_transport() + @patch("sglang.srt.server_args.is_cuda", return_value=True) def test_explicit_cpu_overrides_legacy_environment(self, _mock_is_cuda): server_args = ServerArgs(model_path="dummy", mm_feature_transport="cpu") @@ -231,6 +241,91 @@ class TestMultimodalFeatureTransport(CustomTestCase): self.assertIn("auto-resolved to cuda_ipc", "\n".join(logs.output)) + @patch("sglang.srt.server_args.os.path.exists", return_value=True) + @patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + @patch( + "sglang.srt.model_loader.utils.supports_cuda_vmm_feature_transport", + return_value=True, + ) + def test_default_transport_is_cuda_vmm_for_supported_multinode_mnnvl( + self, + _mock_supports_cuda_vmm, + _mock_is_cuda, + _mock_is_mnnvl, + _mock_path_exists, + ): + server_args = ServerArgs(model_path="dummy", nnodes=2) + self._set_model_type(server_args, is_multimodal=True) + + with patch.dict(os.environ, {}, clear=False): + envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear() + with self.assertLogs(server_args_module.logger, level="INFO") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cuda_vmm") + self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + output = "\n".join(logs.output) + self.assertIn("auto-resolved to cuda_vmm", output) + self.assertIn("CUDA FABRIC", output) + + @patch("sglang.srt.server_args.os.path.exists", return_value=True) + @patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + @patch( + "sglang.srt.model_loader.utils.supports_cuda_vmm_feature_transport", + return_value=False, + ) + def test_default_transport_is_cpu_for_unsupported_multinode_model( + self, + _mock_supports_cuda_vmm, + _mock_is_cuda, + _mock_is_mnnvl, + _mock_path_exists, + ): + server_args = ServerArgs(model_path="dummy", nnodes=2) + self._set_model_type(server_args, is_multimodal=True) + + with self.assertLogs(server_args_module.logger, level="INFO") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cpu") + self.assertIn("has not opted into CUDA VMM", "\n".join(logs.output)) + + @patch("sglang.srt.server_args.os.path.exists", return_value=False) + @patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_default_transport_is_cpu_without_imex_channel( + self, _mock_is_cuda, _mock_is_mnnvl, _mock_path_exists + ): + server_args = ServerArgs(model_path="dummy", nnodes=2) + self._set_model_type(server_args, is_multimodal=True) + + with patch.dict(os.environ, {}, clear=False): + envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear() + with self.assertLogs(server_args_module.logger, level="INFO") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cpu") + + self.assertIn("no IMEX channel", "\n".join(logs.output)) + + @patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=False) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_default_transport_is_cpu_for_multinode_non_mnnvl( + self, _mock_is_cuda, _mock_is_mnnvl + ): + server_args = ServerArgs(model_path="dummy", nnodes=2) + self._set_model_type(server_args, is_multimodal=True) + + with patch.dict(os.environ, {}, clear=False): + envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear() + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cpu") + self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + @patch("sglang.srt.server_args.is_cuda", return_value=True) def test_default_transport_is_cuda_ipc_for_language_only_model(self, _mock_is_cuda): server_args = ServerArgs(model_path="dummy", language_only=True) @@ -259,6 +354,31 @@ class TestMultimodalFeatureTransport(CustomTestCase): with self.assertRaisesRegex(ValueError, "single node"): server_args._handle_multimodal_feature_transport() + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_cuda_vmm_is_explicit_and_uses_shared_budget(self, _mock_is_cuda): + server_args = ServerArgs( + model_path="dummy", + mm_feature_transport="cuda_vmm", + nnodes=2, + tokenizer_worker_num=2, + ) + + with ( + patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "1"}), + envs.SGLANG_MM_FEATURE_CACHE_MB.override(256), + ): + with self.assertLogs(server_args_module.logger, level="INFO") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cuda_vmm") + self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + output = "\n".join(logs.output) + self.assertIn("CUDA FABRIC", output) + self.assertIn("256 MiB", output) + self.assertIn("2 tokenizer worker", output) + self.assertIn("falls back to inline CPU", output) + @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")