feat(vlm): auto-select cuda vmm on multi-node mnnvl (#33936)

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?

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:

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:

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

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.",

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={

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

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

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 "
+        "pool, or a bounded CUDA VMM pool. "
-        "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 "
+        "deployments (without disaggregation) use cuda_ipc; validated multi-node "
-        "cpu. Both CUDA transports reserve SGLANG_MM_FEATURE_CACHE_MB (default "
+        "GB200/GB300 MNNVL models use cuda_vmm when an IMEX channel is available; "
-        "1024 MiB) on the base GPU across tokenizer workers and fall back to CPU "
+        "all other deployments use cpu. GPU transports reserve "
-        "transport per tensor when full.",
+        "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``
+        GPU transports use a fixed pool on ``base_gpu_id`` and therefore reduce
-        and reduces the memory left for model/KV-cache allocations.  The legacy
+        the memory left for model/KV-cache allocations. The legacy CUDA IPC flag
-        flag and environment variable remain supported so existing deployments
+        and environment variable remain supported so existing deployments map
-        continue to work, but both map to this single policy.
+        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
+                # A full GPU pool always degrades to CPU transport per tensor.
-                # CUDA-IPC pool for multimodal models. Text-only deployments do
+                # CUDA IPC is intra-node; multi-node auto-selection is limited
-                # not need feature transport. Multi-node (IPC handles are
+                # to GB200/GB300 systems where the runtime already enables the
-                # intra-node) and PD-disaggregated deployments keep CPU transport.
+                # MNNVL/IMEX communication stack.
-                # A full pool degrades to CPU transport per tensor.
+                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():

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(

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")