[ROCm] Enable the cuMem CUDA-graph pool offload on ROCm (#59523)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: indianspeedster <cspandey016@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: indianspeedster <cspandey016@gmail.com>

csrc/cumem_allocator.cpp

@@ -226,28 +226,6 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
     return;
   }
 
-#ifdef USE_ROCM
-  // ROCm-only: zero the freshly (re)mapped region.
-  //
-  // On some amdgpu driver versions the physical VRAM backing a buffer is not
-  // scrubbed on release, so the pages handed back here (e.g. on wake_up after a
-  // sleep that discarded a tag) can still contain stale data from a previously
-  // freed allocation -- possibly from another process. Callers of the device
-  // allocator expect newly mapped memory to be well defined; in particular
-  // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent
-  // state cache, and a new sequence that reads an unzeroed slot will propagate
-  // garbage/NaN, collapsing generation to token id 0 (decoded as "!").
-  //
-  // This is scoped to the ROCm path on purpose: the CUDA path is unchanged.
-  // Use CUDA_CHECK so a failure sets the global error_code; the Python caller
-  // (python_create_and_map) only inspects error_code, so returning without
-  // setting it would make a failed memset look successful and still hand back
-  // unzeroed memory.
-  CUDA_CHECK((CUresult)hipMemset(reinterpret_cast<void*>(d_mem), 0, size));
-  if (error_code != 0) {
-    return;
-  }
-#endif
   // std::cout << "create_and_map: device=" << device << ", size=" << size << ",
   // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl;
 }

docs/features/sleep_mode.md

@@ -107,7 +107,7 @@ By default, CUDA graph memory stays on the GPU while asleep. With
 `sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured
 into a cuMem pool that sleep backs up to CPU memory at both levels and any wake
 restores in place, so graphs are replayed, not recaptured. It needs the default
-`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect.
+`cumem` backend, CUDA or ROCm, and CUDA graphs; otherwise it has no effect.
 
 ```python
 llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True)

tests/basic_correctness/memory/cumem/test_cumem.py

@@ -359,7 +359,7 @@ def test_runtime_state_survives_sleep(level):
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
 @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
 def test_level2_discards_ordinary_tensor_with_weights_tag():
-    """Reproduce the level-2 variant for an ordinary tensor in weights."""
+    """Discarded weights-tag memory is remapped; ROCm zeroes it over stale pages."""
     allocator = get_mem_allocator_instance()
 
     with allocator.use_memory_pool("weights"):
@@ -374,8 +374,9 @@ def test_level2_discards_ordinary_tensor_with_weights_tag():
     torch.accelerator.synchronize()
 
     assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers
-    assert torch.all(fake_weight == 0xA5)
-    assert torch.all(ordinary_tensor == 0xA5)
+    expected = 0 if current_platform.is_rocm() else 0xA5
+    assert torch.all(fake_weight == expected)
+    assert torch.all(ordinary_tensor == expected)
 
 
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
@@ -427,7 +428,9 @@ def test_cumem_with_cudagraph():
 
 @pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"])
 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool")
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool"
+)
 def test_cudagraph_pool_sleep(level):
     """Routing, and the graph pool backed up at both sleep levels and restored
     in place by any wake."""

tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py

@@ -209,8 +209,20 @@ def _cudagraph_bytes(worker) -> int:
 def _custom_ar_active(worker) -> bool:
     from vllm.distributed.parallel_state import get_tp_group
 
-    ca_comm = get_tp_group().device_communicator.ca_comm
-    return ca_comm is not None and not ca_comm.disabled
+    comm = get_tp_group().device_communicator
+    return any(
+        c is not None and not c.disabled for c in (comm.ca_comm, comm.aiter_ar_comm)
+    )
+
+
+def _capture_registers_buffers(worker) -> bool:
+    from vllm.distributed.parallel_state import get_tp_group
+
+    comm = get_tp_group().device_communicator
+    ca, aiter = comm.ca_comm, comm.aiter_ar_comm
+    return (ca is not None and ca._capture_registered) or (
+        aiter is not None and aiter.aiter_ca.enable_register_for_capturing
+    )
 
 
 @pytest.mark.parametrize(
@@ -227,7 +239,9 @@ def _custom_ar_active(worker) -> bool:
     ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"],
 )
 @create_new_process_for_each_test()
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool")
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool"
+)
 def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload):
     """Each capture site uses the pool, exact across sleeps; TP=2 needs
     unregistered custom AR; off changes nothing."""
@@ -255,8 +269,10 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload):
     assert all(graph_bytes) if offload else not any(graph_bytes)
     if not offload:
         return
-    if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)):
-        pytest.skip("Custom allreduce is unavailable on these GPUs")
+    if tp > 1:
+        if not all(llm.collective_rpc(_custom_ar_active)):
+            pytest.skip("Custom allreduce is unavailable on these GPUs")
+        assert not any(llm.collective_rpc(_capture_registers_buffers))
     prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10)
     expected = llm.generate(prompt, params)[0].outputs[0].text
     for level in (1, 2):
@@ -268,6 +284,42 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload):
         assert llm.generate(prompt, params)[0].outputs[0].text == expected
 
 
+@multi_gpu_test(num_gpus=2)
+@create_new_process_for_each_test()
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="AITER fused allreduce")
+def test_sleep_cudagraph_pool_aiter_fused_ar_group_quant(monkeypatch):
+    """The fused AITER AR + RMSNorm + per-group FP8 quant captures into the
+    cuMem pool without IPC-registering its input, across sleeps."""
+    for name, value in [
+        ("VLLM_USE_V2_MODEL_RUNNER", "1"),
+        ("VLLM_ROCM_USE_AITER", "1"),
+        ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"),
+    ]:
+        monkeypatch.setenv(name, value)
+    monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False)
+    llm = LLM(
+        "Qwen/Qwen3-0.6B-FP8",
+        enable_sleep_mode=True,
+        sleep_mode_offload_cudagraph=True,
+        tensor_parallel_size=2,
+        max_model_len=1024,
+        compilation_config={"pass_config": {"fuse_allreduce_rms": True}},
+    )
+    if not all(llm.collective_rpc(_custom_ar_active)):
+        pytest.skip("AITER custom allreduce is unavailable on these GPUs")
+    assert not any(llm.collective_rpc(_capture_registers_buffers))
+    assert all(llm.collective_rpc(_cudagraph_bytes))
+    prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10)
+    assert llm.generate(prompt, params)[0].outputs[0].text
+    for level in (1, 2):
+        llm.sleep(level=level)
+        llm.wake_up(tags=["weights"])
+        if level == 2:
+            llm.collective_rpc("reload_weights")
+        llm.wake_up(tags=["kv_cache"])
+        assert llm.generate(prompt, params)[0].outputs[0].text
+
+
 @create_new_process_for_each_test()
 def test_deep_sleep_async(gpu_memory_cleared):
     async def test():

vllm/_aiter_ops.py

@@ -1167,7 +1167,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
         w=weight,
         eps=epsilon,
         group_size=group_size,
-        registered=torch.cuda.is_current_stream_capturing(),
+        registered=torch.cuda.is_current_stream_capturing()
+        and ca.enable_register_for_capturing,
         use_1stage=use_1stage,
     )
     assert result is not None
@@ -1219,7 +1220,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
         w=weight,
         eps=epsilon,
         group_size=group_size,
-        registered=torch.cuda.is_current_stream_capturing(),
+        registered=torch.cuda.is_current_stream_capturing()
+        and ca.enable_register_for_capturing,
         use_1stage=use_1stage,
         emit_bf16=True,
     )

vllm/config/vllm.py

@@ -702,7 +702,7 @@ class VllmConfig:
             and model_config.enable_sleep_mode
             and model_config.sleep_mode_backend == "cumem"
             and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
-            and current_platform.is_cuda()
+            and current_platform.is_cuda_alike()
         )
 
     @property

vllm/device_allocator/cumem.py

@@ -98,7 +98,7 @@ class CuMemAllocator:
     offloaded to CPU memory, and the rest of the tensors will be discarded.
     When we call `wake_up`, all tensors that are previously offloaded
     will be loaded back to GPU memory, and the rest of the tensors will
-    have empty memory.
+    have empty memory (zeroed on ROCm).
 
     Why it needs to be a singleton?
     When allocated tensors are garbage collected, PyTorch will call
@@ -331,7 +331,7 @@ class CuMemAllocator:
     def wake_up(self, tags: list[str] | None = None) -> None:
         """Wake up the allocator from sleep mode.
         All data that is previously offloaded will be loaded back to GPU
-        memory, and the rest of the data will have empty memory.
+        memory, and the rest will have empty memory (zeroed on ROCm).
 
         Args:
             tags: The deferrable tags (weights, kv_cache) to load back to GPU
@@ -358,6 +358,10 @@ class CuMemAllocator:
                         cpu_ptr = cpu_backup_tensor.data_ptr()
                         libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)
                         data.cpu_backup_tensor = None
+                elif current_platform.is_rocm():
+                    # amdgpu <6.14 may return stale VRAM (#44972); drop with its support
+                    libcudart.cudaMemset(ptr, 0, handle[1])
+                    torch.accelerator.synchronize()
 
     @contextmanager
     def use_memory_pool(self, tag: str | None = None):

vllm/distributed/device_communicators/aiter_custom_all_reduce.py

@@ -30,6 +30,8 @@ class AiterCustomAllreduce:
         group: ProcessGroup,
         device: int | str | torch.device,
         max_size: int | None = None,
+        *,
+        register_graph_buffers: bool = True,
     ):
         from aiter.dist.device_communicators.custom_all_reduce import (
             CustomAllreduce as _AiterCustomAllreduce,
@@ -38,7 +40,10 @@ class AiterCustomAllreduce:
         if max_size is None:
             max_size = self.MAX_SIZE
 
-        self._impl = _AiterCustomAllreduce(group, device, max_size=max_size)
+        kwargs = {}
+        if not register_graph_buffers:
+            kwargs["enable_register_for_capturing"] = False
+        self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs)
 
     @property
     def aiter_ca(self):

vllm/distributed/device_communicators/cuda_communicator.py

@@ -151,6 +151,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             self.aiter_ar_comm = AiterCustomAllreduce(
                 group=self.cpu_group,
                 device=self.device,
+                register_graph_buffers=register,
             )
 
         if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1:
@@ -174,6 +175,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             self.aiter_ar_comm = AiterCustomAllreduce(
                 group=self.cpu_group,
                 device=self.device,
+                register_graph_buffers=register,
             )
             if self.aiter_ar_comm.disabled:
                 self.aiter_ar_comm = None