@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 19cf334207

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 05:30 UTC

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

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

Kebe committed May 20, 2026 at 21:58 UTC 19cf334207ed81d3ed75a473acd1a95c785d9ed3

6 files changed +60 -17

docs/features/nixl_connector_usage.md

@@ -375,6 +375,10 @@ The `kv_load_failure_policy` setting controls how the system handles failures wh

375 !!! warning

376 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.

377

378 +### For NVIDIA GB-series GPUs

379 +

380 +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.

381 +

382 ## Experimental Feature

383

384 ### Heterogeneous KV Layout support

tests/v1/kv_connector/unit/test_config.py

+16 -2

@@ -66,7 +66,10 @@ def test_kv_connector(

68 def _build_config(

69 - *, kv_connector: str | None, enable_sleep_mode: bool = False

69 + *,

70 + kv_connector: str | None,

71 + enable_sleep_mode: bool = False,

72 + enable_cumem_allocator: bool = False,

73 ) -> VllmConfig:

74 """Build a VllmConfig that exercises _verify_kv_transfer_compat without

75 requiring a real model (avoids HF downloads in CI)."""

@@ -79,7 +82,10 @@ def _build_config(

82 )

83 cfg = VllmConfig.__new__(VllmConfig)

84 cfg.kv_transfer_config = kv_transfer_config

82 - cfg.model_config = SimpleNamespace(enable_sleep_mode=enable_sleep_mode)

85 + cfg.model_config = SimpleNamespace(

86 + enable_sleep_mode=enable_sleep_mode,

87 + enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode),

88 + )

89 cfg._verify_kv_transfer_compat()

90 return cfg

@@ -105,6 +111,14 @@ def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch):

111 _build_config(kv_connector="NixlConnector", enable_sleep_mode=True)

112

113

114 +def test_kv_connector_allows_expandable_segments_with_cumem_allocator(

115 + monkeypatch,

116 +):

117 + """Manual CuMem allocation must also bypass expandable_segments."""

118 + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

119 + _build_config(kv_connector="NixlConnector", enable_cumem_allocator=True)

120 +

121 +

122 def test_kv_connector_allows_other_alloc_conf(monkeypatch):

123 """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected."""

124 monkeypatch.setenv(

vllm/config/model.py

+27 -2

@@ -80,6 +80,16 @@ else:

81 logger = init_logger(__name__)

83 +

84 +def is_cumem_allocator_available() -> bool:

85 + try:

86 + from vllm.device_allocator.cumem import cumem_available

87 + except ImportError:

88 + return False

89 +

90 + return cumem_available

91 +

92 +

93 RunnerOption = Literal["auto", RunnerType]

94 ConvertType = Literal["none", "embed", "classify"]

95 ConvertOption = Literal["auto", ConvertType]

@@ -295,6 +305,13 @@ class ModelConfig:

305 enable_sleep_mode: bool = False

306 """Enable sleep mode for the engine (only cuda and

307 hip platforms are supported)."""

308 + enable_cumem_allocator: bool = False

309 + """Enable the custom cumem allocator to leverage advanced GPU memory

310 + allocation features such as multi-node NVLink support.

311 +

312 + Sleep mode automatically enables this allocator. Only cuda and hip

313 + platforms are supported.

314 + """

315 model_impl: str | ModelImpl = "auto"

316 """Which implementation of the model to use:

317

@@ -510,8 +527,16 @@ class ModelConfig:

527 stacklevel=2,

528 )

529

513 - if self.enable_sleep_mode and not current_platform.is_sleep_mode_available():

514 - raise ValueError("Sleep mode is not supported on current platform.")

530 + if self.enable_sleep_mode:

531 + if not current_platform.is_sleep_mode_available():

532 + raise ValueError("Sleep mode is not supported on current platform.")

533 + if not self.enable_cumem_allocator:

534 + logger.info_once(

535 + "Enabling cumem allocator because sleep mode requires it."

536 + )

537 + self.enable_cumem_allocator = True

538 + if self.enable_cumem_allocator and not is_cumem_allocator_available():

539 + raise ValueError("cumem allocator is not supported on current platform.")

540

541 hf_config = get_config(

542 self.hf_config_path or self.model,

vllm/config/vllm.py

+5 -4

@@ -786,7 +786,7 @@ class VllmConfig:

786 # pins memory, so we conservatively reject the combination whenever

787 # any KV connector is configured.

788 #

789 - # Sleep mode is exempt: CuMemAllocator.use_memory_pool toggles

789 + # CuMem allocator is exempt: CuMemAllocator.use_memory_pool toggles

790 # expandable_segments off around its pool (see #40812), so the KV

791 # cache allocated within that context lands on stable physical pages

792 # even when the env var is set.

@@ -794,17 +794,18 @@ class VllmConfig:

794 "PYTORCH_CUDA_ALLOC_CONF", ""

795 ):

796 return

797 - if self.model_config is not None and self.model_config.enable_sleep_mode:

797 + if self.model_config is not None and (self.model_config.enable_cumem_allocator):

798 return

799

800 raise ValueError(

801 f"KV connector {self.kv_transfer_config.kv_connector} is "

802 "incompatible with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True "

803 - "unless enable_sleep_mode is also enabled. PyTorch's CUDA VMM "

803 + "unless enable_cumem_allocator is also enabled. PyTorch's CUDA VMM "

804 "allocator can remap KV cache virtual addresses to different "

805 "physical pages, invalidating any pinned/registered KV memory "

806 "(e.g. IB memory regions registered by NIXL or Mooncake). Either "

807 - "unset expandable_segments:True or enable sleep mode (which "

807 + "unset expandable_segments:True or enable the cumem allocator "

808 + "(sleep mode does this automatically and also "

809 "routes KV allocations through CuMemAllocator's pool, where "

810 "expandable_segments is automatically disabled)."

811 )

vllm/engine/arg_utils.py

@@ -655,6 +655,7 @@ class EngineArgs:

655

656 generation_config: str = ModelConfig.generation_config

657 enable_sleep_mode: bool = ModelConfig.enable_sleep_mode

658 + enable_cumem_allocator: bool = ModelConfig.enable_cumem_allocator

659 override_generation_config: dict[str, Any] = get_field(

660 ModelConfig, "override_generation_config"

661 )

@@ -839,6 +840,9 @@ class EngineArgs:

840 model_group.add_argument(

841 "--enable-sleep-mode", **model_kwargs["enable_sleep_mode"]

842 )

843 + model_group.add_argument(

844 + "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"]

845 + )

846 model_group.add_argument("--model-impl", **model_kwargs["model_impl"])

847 model_group.add_argument(

848 "--override-attention-dtype", **model_kwargs["override_attention_dtype"]

@@ -1590,6 +1594,7 @@ class EngineArgs:

1594 generation_config=self.generation_config,

1595 override_generation_config=self.override_generation_config,

1596 enable_sleep_mode=self.enable_sleep_mode,

1597 + enable_cumem_allocator=self.enable_cumem_allocator,

1598 model_impl=self.model_impl,

1599 override_attention_dtype=self.override_attention_dtype,

1600 logits_processors=self.logits_processors,

vllm/v1/worker/gpu_worker.py

+3 -9

@@ -199,7 +199,7 @@ class Worker(WorkerBase):

199 self.model_runner.post_kv_cache_wake_up()

200

201 def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:

202 - if not self.vllm_config.model_config.enable_sleep_mode:

202 + if not self.vllm_config.model_config.enable_cumem_allocator:

203 return nullcontext()

204

205 from vllm.device_allocator.cumem import CuMemAllocator

@@ -207,7 +207,7 @@ class Worker(WorkerBase):

207 allocator = CuMemAllocator.get_instance()

208 if tag == "weights":

209 assert allocator.get_current_usage() == 0, (

210 - "Sleep mode can only be used for one instance per process."

210 + "CuMem allocator can only be used for one instance per process."

211 )

212 return allocator.use_memory_pool(tag=tag)

213

@@ -550,13 +550,7 @@ class Worker(WorkerBase):

550 # related to kv cache connector (e.g. kv cache sharing layers).

551 ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)

552

553 - if self.vllm_config.model_config.enable_sleep_mode:

554 - from vllm.device_allocator.cumem import CuMemAllocator

555 -

556 - allocator = CuMemAllocator.get_instance()

557 - with allocator.use_memory_pool(tag="kv_cache"):

558 - self.model_runner.initialize_kv_cache(kv_cache_config)

559 - else:

553 + with self._maybe_get_memory_pool_context(tag="kv_cache"):

554 self.model_runner.initialize_kv_cache(kv_cache_config)

555

556 if self.model_config.enable_return_routed_experts: