[Feature] Support manually enabling the cumem allocator (#33648)

Signed-off-by: Kebe <mail@kebe7jun.com>

docs/features/nixl_connector_usage.md

@@ -375,6 +375,10 @@ The `kv_load_failure_policy` setting controls how the system handles failures wh
 !!! warning
     Using `kv_load_failure_policy="recompute"` can lead to performance degradation in production deployments. When KV loads fail, the decode instance will execute prefill work with decode-optimized configurations, which is inefficient and defeats the purpose of disaggregated prefilling. This also increases tail latency for other ongoing decode requests.
 
+### For NVIDIA GB-series GPUs
+
+GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVCache must be registered as VMM during KVCache registration. To enable this feature, you need to set `--enable-cumem-allocator` or `--enable-sleep-mode` flags, and set `UCX_CUDA_IPC_ENABLE_MNNVL: 'y'` env. Otherwise, NIXL can only use RDMA/TCP for cross-node KVCache transfers.
+
 ## Experimental Feature
 
 ### Heterogeneous KV Layout support

tests/v1/kv_connector/unit/test_config.py

@@ -66,7 +66,10 @@ def test_kv_connector(
 
 
 def _build_config(
-    *, kv_connector: str | None, enable_sleep_mode: bool = False
+    *,
+    kv_connector: str | None,
+    enable_sleep_mode: bool = False,
+    enable_cumem_allocator: bool = False,
 ) -> VllmConfig:
     """Build a VllmConfig that exercises _verify_kv_transfer_compat without
     requiring a real model (avoids HF downloads in CI)."""
@@ -79,7 +82,10 @@ def _build_config(
     )
     cfg = VllmConfig.__new__(VllmConfig)
     cfg.kv_transfer_config = kv_transfer_config
-    cfg.model_config = SimpleNamespace(enable_sleep_mode=enable_sleep_mode)
+    cfg.model_config = SimpleNamespace(
+        enable_sleep_mode=enable_sleep_mode,
+        enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode),
+    )
     cfg._verify_kv_transfer_compat()
     return cfg
 
@@ -105,6 +111,14 @@ def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch):
     _build_config(kv_connector="NixlConnector", enable_sleep_mode=True)
 
 
+def test_kv_connector_allows_expandable_segments_with_cumem_allocator(
+    monkeypatch,
+):
+    """Manual CuMem allocation must also bypass expandable_segments."""
+    monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+    _build_config(kv_connector="NixlConnector", enable_cumem_allocator=True)
+
+
 def test_kv_connector_allows_other_alloc_conf(monkeypatch):
     """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected."""
     monkeypatch.setenv(

vllm/config/model.py

@@ -80,6 +80,16 @@ else:
 
 logger = init_logger(__name__)
 
+
+def is_cumem_allocator_available() -> bool:
+    try:
+        from vllm.device_allocator.cumem import cumem_available
+    except ImportError:
+        return False
+
+    return cumem_available
+
+
 RunnerOption = Literal["auto", RunnerType]
 ConvertType = Literal["none", "embed", "classify"]
 ConvertOption = Literal["auto", ConvertType]
@@ -295,6 +305,13 @@ class ModelConfig:
     enable_sleep_mode: bool = False
     """Enable sleep mode for the engine (only cuda and
     hip platforms are supported)."""
+    enable_cumem_allocator: bool = False
+    """Enable the custom cumem allocator to leverage advanced GPU memory
+    allocation features such as multi-node NVLink support.
+
+    Sleep mode automatically enables this allocator. Only cuda and hip
+    platforms are supported.
+    """
     model_impl: str | ModelImpl = "auto"
     """Which implementation of the model to use:
 
@@ -510,8 +527,16 @@ class ModelConfig:
                 stacklevel=2,
             )
 
-        if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():
-            raise ValueError("Sleep mode is not supported on current platform.")
+        if self.enable_sleep_mode:
+            if not current_platform.is_sleep_mode_available():
+                raise ValueError("Sleep mode is not supported on current platform.")
+            if not self.enable_cumem_allocator:
+                logger.info_once(
+                    "Enabling cumem allocator because sleep mode requires it."
+                )
+                self.enable_cumem_allocator = True
+        if self.enable_cumem_allocator and not is_cumem_allocator_available():
+            raise ValueError("cumem allocator is not supported on current platform.")
 
         hf_config = get_config(
             self.hf_config_path or self.model,

vllm/config/vllm.py

@@ -786,7 +786,7 @@ class VllmConfig:
         # pins memory, so we conservatively reject the combination whenever
         # any KV connector is configured.
         #
-        # Sleep mode is exempt: CuMemAllocator.use_memory_pool toggles
+        # CuMem allocator is exempt: CuMemAllocator.use_memory_pool toggles
         # expandable_segments off around its pool (see #40812), so the KV
         # cache allocated within that context lands on stable physical pages
         # even when the env var is set.
@@ -794,17 +794,18 @@ class VllmConfig:
             "PYTORCH_CUDA_ALLOC_CONF", ""
         ):
             return
-        if self.model_config is not None and self.model_config.enable_sleep_mode:
+        if self.model_config is not None and (self.model_config.enable_cumem_allocator):
             return
 
         raise ValueError(
             f"KV connector {self.kv_transfer_config.kv_connector} is "
             "incompatible with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True "
-            "unless enable_sleep_mode is also enabled. PyTorch's CUDA VMM "
+            "unless enable_cumem_allocator is also enabled. PyTorch's CUDA VMM "
             "allocator can remap KV cache virtual addresses to different "
             "physical pages, invalidating any pinned/registered KV memory "
             "(e.g. IB memory regions registered by NIXL or Mooncake). Either "
-            "unset expandable_segments:True or enable sleep mode (which "
+            "unset expandable_segments:True or enable the cumem allocator "
+            "(sleep mode does this automatically and also "
             "routes KV allocations through CuMemAllocator's pool, where "
             "expandable_segments is automatically disabled)."
         )

vllm/engine/arg_utils.py

@@ -655,6 +655,7 @@ class EngineArgs:
 
     generation_config: str = ModelConfig.generation_config
     enable_sleep_mode: bool = ModelConfig.enable_sleep_mode
+    enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator
     override_generation_config: dict[str, Any] = get_field(
         ModelConfig, "override_generation_config"
     )
@@ -839,6 +840,9 @@ class EngineArgs:
         model_group.add_argument(
             "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"]
         )
+        model_group.add_argument(
+            "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"]
+        )
         model_group.add_argument("--model-impl", **model_kwargs["model_impl"])
         model_group.add_argument(
             "--override-attention-dtype", **model_kwargs["override_attention_dtype"]
@@ -1590,6 +1594,7 @@ class EngineArgs:
             generation_config=self.generation_config,
             override_generation_config=self.override_generation_config,
             enable_sleep_mode=self.enable_sleep_mode,
+            enable_cumem_allocator=self.enable_cumem_allocator,
             model_impl=self.model_impl,
             override_attention_dtype=self.override_attention_dtype,
             logits_processors=self.logits_processors,

vllm/v1/worker/gpu_worker.py

@@ -199,7 +199,7 @@ class Worker(WorkerBase):
             self.model_runner.post_kv_cache_wake_up()
 
     def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
-        if not self.vllm_config.model_config.enable_sleep_mode:
+        if not self.vllm_config.model_config.enable_cumem_allocator:
             return nullcontext()
 
         from vllm.device_allocator.cumem import CuMemAllocator
@@ -207,7 +207,7 @@ class Worker(WorkerBase):
         allocator = CuMemAllocator.get_instance()
         if tag == "weights":
             assert allocator.get_current_usage() == 0, (
-                "Sleep mode can only be used for one instance per process."
+                "CuMem allocator can only be used for one instance per process."
             )
         return allocator.use_memory_pool(tag=tag)
 
@@ -550,13 +550,7 @@ class Worker(WorkerBase):
         # related to kv cache connector (e.g. kv cache sharing layers).
         ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)
 
-        if self.vllm_config.model_config.enable_sleep_mode:
-            from vllm.device_allocator.cumem import CuMemAllocator
-
-            allocator = CuMemAllocator.get_instance()
-            with allocator.use_memory_pool(tag="kv_cache"):
-                self.model_runner.initialize_kv_cache(kv_cache_config)
-        else:
+        with self._maybe_get_memory_pool_context(tag="kv_cache"):
             self.model_runner.initialize_kv_cache(kv_cache_config)
 
         if self.model_config.enable_return_routed_experts: