From 30588e90f4aa90bc07773ac0683cb1ace391f21d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 5 Oct 2026 00:11:11 +0800 Subject: [PATCH 01/12] [ROCm] Enable the cuMem CUDA-graph pool offload on ROCm #59160 captures CUDA graphs into a cuMem pool that sleep mode offloads when sleep_mode_offload_cudagraph is on. This enables the same path on ROCm, behind the same option. - VllmConfig.use_cumem_cudagraph_pool requires is_cuda_alike() instead of is_cuda(). - cumem_allocator.cpp, ROCm path: the blocking hipMemset that zeroes every freshly mapped region joins the legacy stream, so the first capture into the cuMem graph pool fails with "operation would make the legacy stream depend on a capturing blocking stream" (hipErrorStreamCaptureImplicit). Zero with hipMemsetAsync on a private non-blocking stream, created once per device, then synchronize that stream. The CUDA path is untouched. - AiterCustomAllreduce passes enable_register_for_capturing=False when the cuMem graph pool is active: hipIpcGetMemHandle rejects VMM graph buffers, so capture must copy into the IPC buffer, as vLLM's own CustomAllreduce already does through _capture_registered. - Run the cuMem graph-pool tests on ROCm, and say CUDA or ROCm in the docs. Validated on MI355X (DeepSeek-V2-Lite TP1/TP2/DP2, DeepSeek-V4-Flash TP4): graph replay after wake is bit-exact through 10 sleep/wake cycles plus a level-2 reload; the cudagraph tag is 72-244 MiB per GPU. Signed-off-by: aoshen02 Co-Authored-By: Claude Fable 5.1 --- csrc/cumem_allocator.cpp | 31 ++++++++++++++++++- docs/features/sleep_mode.md | 2 +- .../memory/cumem/test_cumem.py | 4 ++- .../memory/sleep_mode/test_sleep_mode.py | 4 ++- vllm/config/vllm.py | 2 +- .../aiter_custom_all_reduce.py | 9 +++++- 6 files changed, 46 insertions(+), 6 deletions(-) diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 720c990450b6..595f416491a7 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -12,6 +12,7 @@ static const char* PYARGS_PARSE = "KKKK"; #include #include #include + #include // Default chunk size 256MB for ROCm. Can be overridden at runtime by the // environment variable VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE, specified in megabytes @@ -116,6 +117,24 @@ void ensure_context(unsigned long long device) { } } +#ifdef USE_ROCM +// A blocking hipMemset joins the legacy stream and breaks graph capture. +static hipStream_t get_zero_stream(unsigned long long device) { + static std::unordered_map streams; + auto it = streams.find(device); + if (it != streams.end()) { + return it->second; + } + hipStream_t stream = nullptr; + CUDA_CHECK((CUresult)hipStreamCreateWithFlags(&stream, hipStreamNonBlocking)); + if (error_code != 0) { + return nullptr; + } + streams[device] = stream; + return stream; +} +#endif + void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, #ifndef USE_ROCM CUmemGenericAllocationHandle* p_memHandle) { @@ -230,7 +249,17 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, // (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(d_mem), 0, size)); + // Zero on a private non-blocking stream so a capturing stream is unaffected. + hipStream_t zero_stream = get_zero_stream(device); + if (error_code != 0) { + return; + } + CUDA_CHECK((CUresult)hipMemsetAsync(reinterpret_cast(d_mem), 0, size, + zero_stream)); + if (error_code != 0) { + return; + } + CUDA_CHECK((CUresult)hipStreamSynchronize(zero_stream)); if (error_code != 0) { return; } diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index a8aa15af7b52..0db82de9d3c5 100644 --- a/docs/features/sleep_mode.md +++ b/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) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index df8c6048824c..19fa222d0a8d 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -427,7 +427,9 @@ def model(x): @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.""" diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 73af6e39103d..57f4eafc4c65 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -227,7 +227,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.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 32b49f42725e..fafae1806312 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -709,7 +709,7 @@ def use_cumem_cudagraph_pool(self) -> bool: 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 diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 6fbffcaaf73d..ba7b55bd2957 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -11,6 +11,8 @@ import torch from torch.distributed import ProcessGroup +from vllm.config import get_current_vllm_config_or_none +from vllm.device_allocator import cumem_cudagraph_pool_enabled from vllm.logger import init_logger logger = init_logger(__name__) @@ -38,7 +40,12 @@ def __init__( if max_size is None: max_size = self.MAX_SIZE - self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + kwargs = {} + config = get_current_vllm_config_or_none() + if config is not None and cumem_cudagraph_pool_enabled(config): + # cuMem graph buffers cannot be hipIpc-exported; capture copies them in. + kwargs["enable_register_for_capturing"] = False + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs) @property def aiter_ca(self): From 704c6ee709c0f0284bef449362c05dff7a4257da Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 1 Oct 2026 03:40:07 +0000 Subject: [PATCH 02/12] [ROCm] Pass register_graph_buffers to AiterCustomAllreduce AiterCustomAllreduce takes the same keyword-only register_graph_buffers flag as CustomAllreduce and no longer reads the global config or imports vllm.device_allocator. CudaCommunicator passes the value it already computes at both construction sites; enable_register_for_capturing=False is still only forwarded to AITER when registration is off. Co-Authored-By: Claude Fable 5.1 Signed-off-by: aoshen02 --- .../device_communicators/aiter_custom_all_reduce.py | 8 +++----- .../distributed/device_communicators/cuda_communicator.py | 2 ++ 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index ba7b55bd2957..9a5970c565d7 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -11,8 +11,6 @@ import torch from torch.distributed import ProcessGroup -from vllm.config import get_current_vllm_config_or_none -from vllm.device_allocator import cumem_cudagraph_pool_enabled from vllm.logger import init_logger logger = init_logger(__name__) @@ -32,6 +30,8 @@ def __init__( 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, @@ -41,9 +41,7 @@ def __init__( max_size = self.MAX_SIZE kwargs = {} - config = get_current_vllm_config_or_none() - if config is not None and cumem_cudagraph_pool_enabled(config): - # cuMem graph buffers cannot be hipIpc-exported; capture copies them in. + if not register_graph_buffers: kwargs["enable_register_for_capturing"] = False self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 99ebcd717ddf..ca1b5309d9f8 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -146,6 +146,7 @@ def __init__( 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: @@ -169,6 +170,7 @@ def __init__( 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 From 82a5b697484848eb7722e12c6da333bf3dc3f4e0 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Fri, 2 Oct 2026 16:15:36 +0000 Subject: [PATCH 03/12] [ROCm] Respect the graph-buffer policy in fused AITER all-reduce The fused all-reduce + RMSNorm + per-group FP8 quant ops, plain and emit_bf16, called AITER's fused_ar_rms_per_group_quant with registered=. That bypasses enable_register_for_capturing, so with the cuMem graph pool the capture input was IPC-registered and capture aborted (hipIpcGetMemHandle: invalid argument, reproduced on MI355X). Both now pass registered only when capturing and the communicator registers graph buffers. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/_aiter_ops.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index ca98e63b394f..ade322464561 100644 --- a/vllm/_aiter_ops.py +++ b/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, ) From 3e9dc1cd0d4cf6ba4419195bb27e829e93568041 Mon Sep 17 00:00:00 2001 From: indianspeedster Date: Thu, 8 Oct 2026 02:44:43 +0000 Subject: [PATCH 04/12] [ROCm][Test] Cover AITER allreduce in the sleep-mode graph-pool tests `_custom_ar_active` checked only `ca_comm`, which is None when AITER custom allreduce is selected, so the TP=2 graph-pool test skipped under VLLM_ROCM_USE_AITER=1. Also add a test for the fused AITER allreduce + RMSNorm + per-group FP8 quant op under the cuMem graph pool. Signed-off-by: indianspeedster --- .../memory/sleep_mode/test_sleep_mode.py | 41 ++++++++++++++++++- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 57f4eafc4c65..72cbf42ad7f3 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -209,8 +209,10 @@ 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) + ) @pytest.mark.parametrize( @@ -270,6 +272,41 @@ 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 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(): From 0e54b3fb74684797f27763c391ce01419a567861 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 11:28:48 +0800 Subject: [PATCH 05/12] [ROCm] Zero discarded allocations on wake instead of every mapping Address review: drop the ROCm zero-fill in create_and_map, so the CUDA and ROCm allocator paths are identical again and nothing touches graph capture. Stale VRAM from drivers that do not scrub on release (#44972) matters for discarded memory remapped on wake, so wake_up zeroes it there and synchronizes before returning. Assert that custom allreduce (vLLM or AITER) does not register graph buffers under the pool, and run the AMD sleep-mode CI step on 2 GPUs so the TP=2 cases run. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/basic_correctness.yaml | 2 +- csrc/cumem_allocator.cpp | 51 ------------------- .../memory/cumem/test_cumem.py | 8 +-- .../memory/sleep_mode/test_sleep_mode.py | 19 +++++-- vllm/device_allocator/cumem.py | 10 +++- 6 files changed, 30 insertions(+), 62 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0a050b6979f0..e55b96a45bb3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3049,7 +3049,7 @@ steps: timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false - agent_pool: mi300_1 + agent_pool: mi300_2 fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index d1d595ad0a98..4abda72e4b4c 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -52,7 +52,7 @@ steps: amd: label: ":amd: (MI300) Basic Correctness Sleep Mode" dind: false - device: mi300_1 + device: mi300_2 timeout_in_minutes: 60 depends_on: - image-build-amd diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 595f416491a7..840a88cbe507 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -12,7 +12,6 @@ static const char* PYARGS_PARSE = "KKKK"; #include #include #include - #include // Default chunk size 256MB for ROCm. Can be overridden at runtime by the // environment variable VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE, specified in megabytes @@ -117,24 +116,6 @@ void ensure_context(unsigned long long device) { } } -#ifdef USE_ROCM -// A blocking hipMemset joins the legacy stream and breaks graph capture. -static hipStream_t get_zero_stream(unsigned long long device) { - static std::unordered_map streams; - auto it = streams.find(device); - if (it != streams.end()) { - return it->second; - } - hipStream_t stream = nullptr; - CUDA_CHECK((CUresult)hipStreamCreateWithFlags(&stream, hipStreamNonBlocking)); - if (error_code != 0) { - return nullptr; - } - streams[device] = stream; - return stream; -} -#endif - void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, #ifndef USE_ROCM CUmemGenericAllocationHandle* p_memHandle) { @@ -232,38 +213,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. - // Zero on a private non-blocking stream so a capturing stream is unaffected. - hipStream_t zero_stream = get_zero_stream(device); - if (error_code != 0) { - return; - } - CUDA_CHECK((CUresult)hipMemsetAsync(reinterpret_cast(d_mem), 0, size, - zero_stream)); - if (error_code != 0) { - return; - } - CUDA_CHECK((CUresult)hipStreamSynchronize(zero_stream)); - 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; } diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index 19fa222d0a8d..2d2bac84fd62 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -32,7 +32,7 @@ def mapped_usage(allocator) -> int: def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None: - """Wake discarded allocations with deterministic nonzero contents.""" + """Wake up with every remapped page filled with nonzero bytes first.""" original_create_and_map = cumem.create_and_map def create_and_map_with_poison(handle) -> None: @@ -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 comes back zeroed, even over stale pages.""" allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): @@ -374,8 +374,8 @@ 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) + assert torch.all(fake_weight == 0) + assert torch.all(ordinary_tensor == 0) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 72cbf42ad7f3..12645baf8e02 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -215,6 +215,16 @@ def _custom_ar_active(worker) -> bool: ) +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( ("mode", "breakable", "tp", "offload"), [ @@ -259,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): @@ -295,6 +307,7 @@ def test_sleep_cudagraph_pool_aiter_fused_ar_group_quant(monkeypatch): ) 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 @@ -471,7 +484,7 @@ def create_and_map_with_poison(handle) -> None: monkeypatch.setattr(cumem, "create_and_map", create_and_map_with_poison) - # New requests must overwrite undefined remapped KV bytes before reading them. + # Remapped KV memory comes back zeroed even over stale pages. llm.wake_up(tags=["kv_cache"]) actual = llm.generate(prompt, sampling_params) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index b85be08d28e3..c5303d233b3b 100644 --- a/vllm/device_allocator/cumem.py +++ b/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. + be zeroed. Why it needs to be a singleton? When allocated tensors are garbage collected, PyTorch will call @@ -331,7 +331,7 @@ def discard(self, tags: tuple[str, ...] | str) -> None: 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 comes back zeroed. Args: tags: The deferrable tags (weights, kv_cache) to load back to GPU @@ -358,6 +358,12 @@ def wake_up(self, tags: list[str] | None = None) -> None: cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + else: + # Zero discarded memory: some amdgpu drivers hand back + # stale VRAM (#44972). + libcudart.cudaMemset(ptr, 0, handle[1]) + # cudaMemset is async; finish it before work on other streams. + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): From 9ddebaa8804c8235d8447f02fb9fbcd1475f1da6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 15:02:48 +0800 Subject: [PATCH 06/12] [ROCm] Zero discarded allocations on wake only on ROCm The zero-fill this replaces was ROCm-only (#44972); keep CUDA wake-up unchanged. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- .../basic_correctness/memory/cumem/test_cumem.py | 9 +++++---- .../memory/sleep_mode/test_sleep_mode.py | 2 +- vllm/device_allocator/cumem.py | 15 ++++++++------- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index 2d2bac84fd62..c7ebfd2609d6 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -32,7 +32,7 @@ def mapped_usage(allocator) -> int: def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None: - """Wake up with every remapped page filled with nonzero bytes first.""" + """Wake discarded allocations with deterministic nonzero contents.""" original_create_and_map = cumem.create_and_map def create_and_map_with_poison(handle) -> None: @@ -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(): - """Discarded weights-tag memory comes back zeroed, even over stale pages.""" + """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 == 0) - assert torch.all(ordinary_tensor == 0) + 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") diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index e0df42710e37..d897d710015b 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -483,7 +483,7 @@ def create_and_map_with_poison(handle) -> None: monkeypatch.setattr(cumem, "create_and_map", create_and_map_with_poison) - # Remapped KV memory comes back zeroed even over stale pages. + # New requests must overwrite undefined remapped KV bytes before reading them. llm.wake_up(tags=["kv_cache"]) actual = llm.generate(prompt, sampling_params) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index c5303d233b3b..861070ab5ab6 100644 --- a/vllm/device_allocator/cumem.py +++ b/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 - be zeroed. + 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 @@ def discard(self, tags: tuple[str, ...] | str) -> None: 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 comes back zeroed. + 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,12 +358,13 @@ def wake_up(self, tags: list[str] | None = None) -> None: cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None - else: - # Zero discarded memory: some amdgpu drivers hand back - # stale VRAM (#44972). + elif current_platform.is_rocm(): + # amdgpu <6.14 can hand back stale VRAM (#44972); newer + # drivers scrub on release, so this can go once they're the floor. libcudart.cudaMemset(ptr, 0, handle[1]) - # cudaMemset is async; finish it before work on other streams. - torch.accelerator.synchronize() + if current_platform.is_rocm(): + # The memset is async; finish it before work on other streams. + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): From 52c78f9d16e70f5d9f12bb53c1e2d21d30504867 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 15:49:45 +0800 Subject: [PATCH 07/12] [CI] Keep the AMD sleep-mode step on mi300_1 The MI300 distributed steps already run the TP=2 sleep tests. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/basic_correctness.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 6e67fbd9cb58..9e634ea02b4d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3017,7 +3017,7 @@ steps: timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false - agent_pool: mi300_2 + agent_pool: mi300_1 fast_check: true working_dir: "/vllm-workspace/tests" source_file_dependencies: diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index 4abda72e4b4c..d1d595ad0a98 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -52,7 +52,7 @@ steps: amd: label: ":amd: (MI300) Basic Correctness Sleep Mode" dind: false - device: mi300_2 + device: mi300_1 timeout_in_minutes: 60 depends_on: - image-build-amd From ea1f86f9bdd35636b69359a3dd1a4f8ecf4515d5 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 16:14:49 +0800 Subject: [PATCH 08/12] [ROCm] Drop the device sync after zeroing on wake Match main, which zeroes with hipMemset and no sync. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/device_allocator/cumem.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 861070ab5ab6..4e3712bc4608 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -362,9 +362,6 @@ def wake_up(self, tags: list[str] | None = None) -> None: # amdgpu <6.14 can hand back stale VRAM (#44972); newer # drivers scrub on release, so this can go once they're the floor. libcudart.cudaMemset(ptr, 0, handle[1]) - if current_platform.is_rocm(): - # The memset is async; finish it before work on other streams. - torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): From 64fdb6aaf267e444baed508bd19ca0916db875b6 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 19:49:33 +0800 Subject: [PATCH 09/12] [ROCm] Synchronize after zeroing discarded allocations on wake hipMemset at offset 0 returns before it finishes, and vLLM's non-blocking stream does not wait for the null stream, so a level-2 buffer restore can be overwritten by a late memset. Reproduced at the allocator level on MI355X; main has the same race. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/device_allocator/cumem.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 4e3712bc4608..4de447ab0ec5 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -362,6 +362,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: # amdgpu <6.14 can hand back stale VRAM (#44972); newer # drivers scrub on release, so this can go once they're the floor. libcudart.cudaMemset(ptr, 0, handle[1]) + if current_platform.is_rocm(): + torch.accelerator.synchronize() # hipMemset is async to other streams @contextmanager def use_memory_pool(self, tag: str | None = None): From daefeac4d05108faa2b3db5fadb1655436028a6e Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 19:51:00 +0800 Subject: [PATCH 10/12] [ROCm] Put the wake sync comment on its own line Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/device_allocator/cumem.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 4de447ab0ec5..f1d930d3d1f6 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -363,7 +363,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: # drivers scrub on release, so this can go once they're the floor. libcudart.cudaMemset(ptr, 0, handle[1]) if current_platform.is_rocm(): - torch.accelerator.synchronize() # hipMemset is async to other streams + # hipMemset is async to other streams; finish it before they run. + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): From c97885384333a190c8b1934b7796a4f9d987f0fa Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 8 Oct 2026 19:54:15 +0800 Subject: [PATCH 11/12] [ROCm] Zero discarded allocations on the current stream Use hipMemsetAsync on the current stream instead of hipMemset plus a device sync: work after wake is ordered behind the zero-fill by the stream. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/device_allocator/cumem.py | 9 +++------ .../distributed/device_communicators/cuda_wrapper.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index f1d930d3d1f6..9876ae32f208 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -359,12 +359,9 @@ def wake_up(self, tags: list[str] | None = None) -> None: libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None elif current_platform.is_rocm(): - # amdgpu <6.14 can hand back stale VRAM (#44972); newer - # drivers scrub on release, so this can go once they're the floor. - libcudart.cudaMemset(ptr, 0, handle[1]) - if current_platform.is_rocm(): - # hipMemset is async to other streams; finish it before they run. - torch.accelerator.synchronize() + # amdgpu <6.14 may return stale VRAM (#44972); drop with its support + stream = torch.cuda.current_stream().cuda_stream + libcudart.cudaMemsetAsync(ptr, 0, handle[1], stream) @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py index d27d84c03045..38250c1ff170 100644 --- a/vllm/distributed/device_communicators/cuda_wrapper.py +++ b/vllm/distributed/device_communicators/cuda_wrapper.py @@ -60,6 +60,12 @@ class CudaRTLibrary: Function( "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] ), + # cudaError_t cudaMemsetAsync ( void* devPtr, int value, size_t count, cudaStream_t stream ) # noqa + Function( + "cudaMemsetAsync", + cudaError_t, + [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t, ctypes.c_void_p], + ), # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa Function( "cudaMemcpy", @@ -99,6 +105,7 @@ class CudaRTLibrary: "cudaMalloc": "hipMalloc", "cudaFree": "hipFree", "cudaMemset": "hipMemset", + "cudaMemsetAsync": "hipMemsetAsync", "cudaMemcpy": "hipMemcpy", "cudaIpcGetMemHandle": "hipIpcGetMemHandle", "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle", @@ -175,6 +182,11 @@ def cudaFree(self, devPtr: ctypes.c_void_p) -> None: def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None: self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count)) + def cudaMemsetAsync( + self, devPtr: ctypes.c_void_p, value: int, count: int, stream: int + ) -> None: + self.CUDART_CHECK(self.funcs["cudaMemsetAsync"](devPtr, value, count, stream)) + def cudaMemcpy( self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int ) -> None: From 07b66ed928eb5fc01c244997bcfb50b3f1041fe4 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Fri, 9 Oct 2026 08:50:09 +0800 Subject: [PATCH 12/12] [ROCm] Synchronize right after each wake zero-fill Keep the existing hipMemset binding and sync after it, instead of adding a cudaMemsetAsync binding. Signed-off-by: aoshen02 Co-Authored-By: Claude Opus 5.5 --- vllm/device_allocator/cumem.py | 4 ++-- .../distributed/device_communicators/cuda_wrapper.py | 12 ------------ 2 files changed, 2 insertions(+), 14 deletions(-) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 9876ae32f208..a2e582f85b07 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -360,8 +360,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: data.cpu_backup_tensor = None elif current_platform.is_rocm(): # amdgpu <6.14 may return stale VRAM (#44972); drop with its support - stream = torch.cuda.current_stream().cuda_stream - libcudart.cudaMemsetAsync(ptr, 0, handle[1], stream) + libcudart.cudaMemset(ptr, 0, handle[1]) + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/distributed/device_communicators/cuda_wrapper.py b/vllm/distributed/device_communicators/cuda_wrapper.py index 38250c1ff170..d27d84c03045 100644 --- a/vllm/distributed/device_communicators/cuda_wrapper.py +++ b/vllm/distributed/device_communicators/cuda_wrapper.py @@ -60,12 +60,6 @@ class CudaRTLibrary: Function( "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t] ), - # cudaError_t cudaMemsetAsync ( void* devPtr, int value, size_t count, cudaStream_t stream ) # noqa - Function( - "cudaMemsetAsync", - cudaError_t, - [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t, ctypes.c_void_p], - ), # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa Function( "cudaMemcpy", @@ -105,7 +99,6 @@ class CudaRTLibrary: "cudaMalloc": "hipMalloc", "cudaFree": "hipFree", "cudaMemset": "hipMemset", - "cudaMemsetAsync": "hipMemsetAsync", "cudaMemcpy": "hipMemcpy", "cudaIpcGetMemHandle": "hipIpcGetMemHandle", "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle", @@ -182,11 +175,6 @@ def cudaFree(self, devPtr: ctypes.c_void_p) -> None: def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None: self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count)) - def cudaMemsetAsync( - self, devPtr: ctypes.c_void_p, value: int, count: int, stream: int - ) -> None: - self.CUDART_CHECK(self.funcs["cudaMemsetAsync"](devPtr, value, count, stream)) - def cudaMemcpy( self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int ) -> None: