diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index 09e5c12a59..cb5a3dca03 100644 --- a/docs/features/nixl_connector_usage.md +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_config.py b/tests/v1/kv_connector/unit/test_config.py index 8218074245..33c9abd09e 100644 --- a/tests/v1/kv_connector/unit/test_config.py +++ b/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( diff --git a/vllm/config/model.py b/vllm/config/model.py index 26737f313e..936409a01f 100644 --- a/vllm/config/model.py +++ b/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, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index b0e74d8403..6005250662 100644 --- a/vllm/config/vllm.py +++ b/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)." ) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 49189862eb..0e25a4a35c 100644 --- a/vllm/engine/arg_utils.py +++ b/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, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 53e62c9256..0dcf1614f4 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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: