Use cumem allocator for the KV cache by default

Enable the cumem (CUDA VMM) allocator by default on CUDA and ROCm
platforms so users get stable physical pages for KV cache without
needing to set --enable-cumem-allocator explicitly. This is required
for MNNVL KV transfers.

Changes:
- ModelConfig.enable_cumem_allocator defaults to None, resolved to
  True when the cumem C extension is available
- Cached fabric handle probe in csrc/cumem_allocator.cpp uses a real
  cuMemCreate probe instead of trusting cuDeviceGetAttribute, with
  fallback to POSIX FD when fabric handles aren't available
- Safety-net fallback: if fabric was probed as available but a real
  allocation still fails, updates the cache and retries with POSIX FD
- Weight loading bypasses cumem pool when sleep mode is off (cumem is
  only needed for KV cache stability in that case)
- Guard against stale engine allocations corrupting the singleton

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

csrc/cumem_allocator.cpp

@@ -1,6 +1,7 @@
 // A CUDAPluggableAllocator based on cumem* APIs.
 // Important: allocation size, CUdeviceptr and CUmemGenericAllocationHandle*
 // need to be unsigned long long
+#include <atomic>
 #include <iostream>
 
 #include "cumem_allocator_compat.h"
@@ -116,6 +117,59 @@ void ensure_context(unsigned long long device) {
   }
 }
 
+// ---------------------------------------------------------------------------
+// Cached fabric handle probe (CUDA 12.4+, NVIDIA only):
+
+#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040
+// Per-device cache: 0 = not probed, 1 = supported, 2 = not supported
+static constexpr int MAX_DEVICES = 32;
+static std::atomic<int> fabric_support[MAX_DEVICES] = {};
+
+static bool probe_fabric_support(unsigned long long device) {
+  if (device >= MAX_DEVICES) return false;
+  int cached = fabric_support[device].load(std::memory_order_acquire);
+  if (cached != 0) return cached == 1;
+
+  int fab_flag = 0;
+  CUresult r = cuDeviceGetAttribute(
+      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
+  if (r != CUDA_SUCCESS || !fab_flag) {
+    fabric_support[device].store(2, std::memory_order_release);
+    return false;
+  }
+
+  // Attribute says supported — verify with a real allocation.
+  // cuDeviceGetAttribute can report supported even when IMEX is not
+  // configured, so we need a real probe.
+  CUmemAllocationProp probe_prop = {};
+  probe_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+  probe_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+  probe_prop.location.id = device;
+  probe_prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
+
+  size_t granularity;
+  r = cuMemGetAllocationGranularity(&granularity, &probe_prop,
+                                    CU_MEM_ALLOC_GRANULARITY_MINIMUM);
+  if (r != CUDA_SUCCESS) {
+    fabric_support[device].store(2, std::memory_order_release);
+    return false;
+  }
+
+  CUmemGenericAllocationHandle test_handle;
+  r = cuMemCreate(&test_handle, granularity, &probe_prop, 0);
+  if (r == CUDA_SUCCESS) {
+    cuMemRelease(test_handle);
+    fabric_support[device].store(1, std::memory_order_release);
+    return true;
+  }
+
+  fabric_support[device].store(2, std::memory_order_release);
+  return false;
+}
+#endif
+
+// ---------------------------------------------------------------------------
+
 void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
 #ifndef USE_ROCM
                     CUmemGenericAllocationHandle* p_memHandle) {
@@ -136,32 +190,40 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
   CUresult rdma_result = cuDeviceGetAttribute(
       &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
       device);
-  if (rdma_result == CUDA_SUCCESS &&
-      flag) {  // support GPUDirect RDMA if possible
+  if (rdma_result == CUDA_SUCCESS && flag) {
     prop.allocFlags.gpuDirectRDMACapable = 1;
   }
-  int fab_flag = 0;
-  CUresult fab_result = cuDeviceGetAttribute(
-      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
-  if (fab_result == CUDA_SUCCESS &&
-      fab_flag) {  // support fabric handle if possible
+
+  #if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
+  if (probe_fabric_support(device)) {
     prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
+  } else {
+    prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
   }
+  #else
+  prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+  #endif
 #endif
 
 #ifndef USE_ROCM
   // Allocate memory using cuMemCreate
   CUresult ret = (CUresult)cuMemCreate(p_memHandle, size, &prop, 0);
   if (ret) {
-    if (fab_flag &&
+  #if defined(CUDA_VERSION) && CUDA_VERSION >= 12040
+    // Safety net: if fabric was probed as available but this allocation
+    // still fails, fall back to POSIX FD and update the cache.
+    if (device < MAX_DEVICES &&
+        fabric_support[device].load(std::memory_order_acquire) == 1 &&
         (ret == CUDA_ERROR_NOT_PERMITTED || ret == CUDA_ERROR_NOT_SUPPORTED)) {
-      // Fabric allocation may fail without multi-node nvlink,
-      // fallback to POSIX file descriptor
+      fabric_support[device].store(2, std::memory_order_release);
       prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
       CUDA_CHECK(cuMemCreate(p_memHandle, size, &prop, 0));
     } else {
       CUDA_CHECK(ret);
     }
+  #else
+    CUDA_CHECK(ret);
+  #endif
   }
   if (error_code != 0) {
     return;
@@ -326,14 +388,13 @@ void* my_malloc(ssize_t size, int device, CUstream stream) {
   // first allocation, align the size, and reserve an address, and also allocate
   // a CUmemGenericAllocationHandle
 
-  // Define memory allocation properties
+  // No handle type here; create_and_map sets fabric/POSIX as needed.
   CUmemAllocationProp prop = {};
   prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
   prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
   prop.location.id = device;
   prop.allocFlags.compressionType = CU_MEM_ALLOCATION_COMP_NONE;
 
-  // Check if the allocation is supported
   size_t granularity;
   CUDA_CHECK(cuMemGetAllocationGranularity(&granularity, &prop,
                                            CU_MEM_ALLOC_GRANULARITY_MINIMUM));

tests/basic_correctness/test_mem.py

@@ -248,6 +248,32 @@ def test_deep_sleep_async():
     asyncio.run(test())
 
 
+@create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn")
+def test_cumem_without_sleep_mode():
+    """Verify cumem allocator works independently of sleep mode."""
+    llm = LLM("hmellor/tiny-random-LlamaForCausalLM", enable_cumem_allocator=True)
+    prompt = "How are you?"
+    sampling_params = SamplingParams(temperature=0, max_tokens=10)
+    output = llm.generate(prompt, sampling_params)
+    assert output[0].outputs[0].text
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Sleep mode requires CUDA or ROCm",
+)
+def test_sleep_mode_auto_enables_cumem():
+    """Verify sleep mode automatically enables cumem allocator."""
+    from vllm.config.model import ModelConfig
+
+    cfg = ModelConfig(
+        "hmellor/tiny-random-LlamaForCausalLM",
+        enable_sleep_mode=True,
+        enable_cumem_allocator=False,
+    )
+    assert cfg.enable_cumem_allocator is True
+
+
 @requires_fp8
 def test_deep_sleep_fp8_kvcache():
     model = "Qwen/Qwen2-0.5B"

tests/v1/kv_connector/unit/test_config.py

@@ -101,7 +101,6 @@ def test_kv_connector(
 def _build_config(
     *,
     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
@@ -115,10 +114,7 @@ def _build_config(
     )
     cfg = VllmConfig.__new__(VllmConfig)
     cfg.kv_transfer_config = kv_transfer_config
-    cfg.model_config = SimpleNamespace(
-        enable_sleep_mode=enable_sleep_mode,
-        enable_cumem_allocator=(enable_cumem_allocator or enable_sleep_mode),
-    )
+    cfg.model_config = SimpleNamespace(enable_cumem_allocator=enable_cumem_allocator)
     cfg._verify_kv_transfer_compat()
     return cfg
 
@@ -137,13 +133,6 @@ def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector):
         _build_config(kv_connector=kv_connector)
 
 
-def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch):
-    """Sleep mode routes KV allocations through CuMemAllocator's pool, which
-    auto-disables expandable_segments (see #40812)."""
-    monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-    _build_config(kv_connector="NixlConnector", enable_sleep_mode=True)
-
-
 def test_kv_connector_allows_expandable_segments_with_cumem_allocator(
     monkeypatch,
 ):

vllm/config/model.py

@@ -297,13 +297,11 @@ 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.
-    """
+    enable_cumem_allocator: bool | None = None
+    """Enable the cumem allocator for GPU memory management.
+    Automatically enabled when a KV connector is configured (for
+    stable physical pages required by MNNVL/RDMA transfers) or
+    when sleep mode is active. Only CUDA and ROCm are supported."""
     model_impl: str | ModelImpl = "auto"
     """Which implementation of the model to use:
 
@@ -525,6 +523,8 @@ class ModelConfig:
                 stacklevel=2,
             )
 
+        if self.enable_cumem_allocator is None:
+            self.enable_cumem_allocator = False
         if self.enable_sleep_mode:
             if not current_platform.is_sleep_mode_available():
                 raise ValueError("Sleep mode is not supported on current platform.")

vllm/config/vllm.py

@@ -820,28 +820,28 @@ class VllmConfig:
         # pins memory, so we conservatively reject the combination whenever
         # any KV connector is configured.
         #
-        # 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.
+        # The 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.
         if "expandable_segments:True" not in os.environ.get(
             "PYTORCH_CUDA_ALLOC_CONF", ""
         ):
             return
-        if self.model_config is not None and (self.model_config.enable_cumem_allocator):
+        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_cumem_allocator is also enabled. PyTorch's CUDA VMM "
+            "unless the cumem allocator is 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 the cumem allocator "
-            "(sleep mode does this automatically and also "
+            "(--enable-cumem-allocator) which "
             "routes KV allocations through CuMemAllocator's pool, where "
-            "expandable_segments is automatically disabled)."
+            "expandable_segments is automatically disabled."
         )
 
     def __post_init__(self):
@@ -1538,6 +1538,18 @@ class VllmConfig:
             if "-quant_fp8" not in custom_ops:
                 custom_ops.append("+quant_fp8")
 
+        # Auto-enable the cumem allocator when a KV connector is configured.
+        # KV connectors that pin memory (NIXL, Mooncake, etc.) need stable
+        # physical pages — cumem provides that via cuMemCreate.
+        if (
+            self.kv_transfer_config is not None
+            and self.kv_transfer_config.kv_connector is not None
+            and self.model_config is not None
+            and not self.model_config.enable_cumem_allocator
+            and current_platform.is_cumem_allocator_available()
+        ):
+            self.model_config.enable_cumem_allocator = True
+
         self._verify_kv_transfer_compat()
         # Log the custom passes that are enabled
         self.compilation_config.pass_config.log_enabled_passes()

vllm/engine/arg_utils.py

@@ -659,7 +659,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
+    enable_cumem_allocator: bool | None = ModelConfig.enable_cumem_allocator
     override_generation_config: dict[str, Any] = get_field(
         ModelConfig, "override_generation_config"
     )

vllm/v1/worker/gpu_worker.py

@@ -215,11 +215,20 @@ class Worker(WorkerBase):
         if current_platform.is_cpu():
             return nullcontext()
 
+        # Only route weight loading through the cumem pool when sleep mode
+        # needs it for weight offloading.  Without sleep mode the primary
+        # purpose of cumem is stable physical pages for KV cache (MNNVL),
+        # and weight loading through the regular allocator avoids both the
+        # singleton constraint and OOM on smaller GPUs.
+        if tag == "weights" and not self.vllm_config.model_config.enable_sleep_mode:
+            return nullcontext()
+
         allocator = get_mem_allocator_instance()
         if tag == "weights":
             assert allocator.get_current_usage() == 0, (
                 "CuMem allocator can only be used for one instance per process."
             )
+
         return allocator.use_memory_pool(tag=tag)
 
     @contextmanager