[ROCm] Enable the cuMem CUDA-graph pool offload on ROCm

[aoshen02]

Purpose

#59160 captures CUDA graphs (both model runners) into a cuMem pool that sleep mode offloads, opt-in with sleep_mode_offload_cudagraph (off by default, like SGLang's SGLANG_MEMORY_SAVER_CUDA_GRAPH). Its policy (VllmConfig.use_cumem_cudagraph_pool) requires current_platform.is_cuda(), so on ROCm it is inert even when opted in: every capture site gets nullcontext(pool), no cudagraph tag appears, and asleep memory equals main with #59156. This PR enables the same path on ROCm.

Based on main, which now includes #59156, #59158 and #59160. RFC: #59055. It composes with #59909 (release the graph pool instead of backing it up); see Test Result.

What breaks with the Python gate alone

Relaxing only the policy's platform check makes the first PIECEWISE capture fail on MI355X with

CUDA Error: operation would make the legacy stream depend on a capturing blocking stream at csrc/cumem_allocator.cpp:233

(hipErrorStreamCaptureImplicit), followed by unknown error at :461 and an engine start failure. The cause is ROCm-only code at the PR base: create_and_map zeroes every freshly mapped region with a blocking hipMemset on the legacy stream (the stale-VRAM guard for GDN state from #43022/#46203). On CUDA there is no memset, so a cuMem allocation during capture is legal; on ROCm the graph pool's first allocation during capture triggers the memset, and the memset breaks the capture.

Changes

• vllm/config/vllm.py: VllmConfig.use_cumem_cudagraph_pool checks is_cuda_alike() instead of is_cuda().
• csrc/cumem_allocator.cpp: the ROCm-only zero-fill in create_and_map is removed, so the CUDA and ROCm paths are identical again and nothing touches graph capture. Instead, on ROCm only, CuMemAllocator.wake_up zeroes discarded allocations (those without a CPU backup) after remapping them, and synchronizes after each zero-fill. hipMemset returns before it finishes, and vLLM's non-blocking stream does not wait for it. Without the sync, a late memset can overwrite the level-2 buffer restore. Main has the same race. On MI355X it reproduces at the allocator level: 200/200 on main, 0/200 with the sync. That is the memory that can hold stale VRAM on drivers that don't scrub on release ([Test][V1] Add sleep/wake correctness regression test for hybrid GDN/… #44972); startup allocations don't need it (the KV cache is torch.zeros). Suggested by @mawong-amd. It stays because vLLM's ROCm floor (6.3) pairs with amdgpu 6.8–6.10, which don't scrub VRAM on release; amdgpu ≥6.14 does, so it can go once that is the floor. CUDA wake is unchanged. On ROCm it costs about 35 ms per wake for a 150 GiB KV cache.
• AiterCustomAllreduce: passes enable_register_for_capturing=False to AITER's CustomAllreduce when the cuMem graph pool is active. This is required: at capture AITER IPC-registers the graph input pointers, and hipIpcGetMemHandle rejects VMM memory just as cudaIpcGetMemHandle does on CUDA ([Feature] Release the CUDA graph pool on sleep #59160's custom-allreduce fix). Without it, capture with AITER custom allreduce at TP2 fails with hipIpcGetMemHandle ... invalid argument (custom_all_reduce.cuh:4080) and the engine doesn't start (gfx942, reproduced by @mawong-amd). AITER already forces copy-in for expandable-segment pointers for the same reason.
• vllm/_aiter_ops.py: the fused all-reduce + RMSNorm + per-group FP8 quant ops (plain and emit_bf16) passed registered=<capturing> to AITER's raw fused_ar_rms_per_group_quant, which bypasses enable_register_for_capturing, so under the cuMem graph pool the capture input was IPC-registered and capture aborted (hipIpcGetMemHandle ... invalid argument, reproduced on MI355X). Both now pass registered=capturing and ca.enable_register_for_capturing; with the fix capture succeeds and replay is bit-exact, also after sleep/wake of the cudagraph tag.
• Tests: the cuMem graph-pool tests (allocator-level sleep/release test and the engine-level capture-site test) run on ROCm. The TP=2 case asserts that neither vLLM's nor AITER's custom allreduce registers graph buffers, and the level-2 poisoned-remap test asserts that discarded memory comes back zeroed. On AMD CI the TP=2 cases already run in the MI300 distributed steps.
• AITER test coverage, by @indianspeedster (from [ROCm][Test] Cover AITER allreduce in the sleep-mode graph-pool tests #60525, commit authored by them): test_sleep_cudagraph_pool[custom-ar-tp2] now also runs when AITER custom allreduce is active (it skipped before because ca_comm is None then), and a new test covers the fused AITER allreduce + RMSNorm + per-group FP8 quant op under the cuMem graph pool. The new test fails without the _aiter_ops.py change above.

Confirmed and left alone:

• vLLM's own CustomAllreduce needs [Feature] Release the CUDA graph pool on sleep #59160's _capture_registered = False on ROCm too: forcing registered=True fails capture with custom_all_reduce_hip.cuh:167 'invalid argument'. The copy-in path works unchanged.
• [Feature] Release the CUDA graph pool on sleep #59160's NCCL_GRAPH_REGISTER default (VllmConfig._verify_cumem_cudagraph_pool: setdefault("0") when the pool is active, an explicit value kept with a warning, as in fix(nccl): disable graph buffer registration for pausable graph pools sgl-project/sglang#40648) keys on use_cumem_cudagraph_pool, so it applies to RCCL once this PR enables the gate on ROCm (RCCL 2.30.4 honours NCCL_GRAPH_REGISTER). Not duplicated. AITER's enable_register_for_capturing=False, like vLLM's custom allreduce (Support releasing CUDA graph memory when paused sgl-project/sglang#7873), applies only when the pool is active.
• QuickReduce copies into its own static IPC buffer and never registers graph buffers; bit-exact over 10 cycles with VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=FP.
• [Feature] Release the CUDA graph pool on sleep #59160 no longer requires Model Runner V2, so ROCm architectures that default to V1 (e.g. DeepSeek-V4-Flash) also use the pool when opted in. The MI355X numbers below were measured on V2 before that change; V1 on ROCm is not yet validated here.

Results on MI355X

8x MI355X, vLLM nightly e975732152 with #59156 (PR1), #59160 and this change overlaid; --gpu-memory-utilization 0.85 --max-model-len 4096, level-1 sleep. Values are rocm-smi "VRAM Total Used" in MiB per GPU, polled until stable (rocm-smi lags frees by 7–13 s).

Config	Unpatched	PR1	PR1 + this	PR1 + #59158 + this	cudagraph tag
DeepSeek-V2-Lite TP1	5677	4747 (−930)	4583 (−1094)	4479 (−1198)	92
DeepSeek-V2-Lite TP2	7946	7176 (−770)	7038 (−908)	6947 (−999)	72
DeepSeek-V2-Lite DP2	9430	7894 (−1536)	7783 (−1647)	–	94
DeepSeek-V4-Flash TP4 (VLLM_USE_V2_MODEL_RUNNER=1)	11777	–	10747 (−1030)	–	244

• The cudagraph tag is present, asleep while sleeping and restored on any wake, including weights-only.
• Correctness (8 greedy prompts, 48 tokens): TP2, DP2 and V4-Flash TP4 are bit-exact after full wake and after selective wake (weights, then kv_cache). TP2 stays bit-exact through a 10-cycle sleep/wake loop plus level 2 (sleep?level=2, wake weights, reload_weights, wake kv_cache), with and without [Feature] Offload KV-init runtime state on sleep #59158, and with QuickReduce. V4-Flash TP4 is bit-exact over 5 cycles. (DSv2-Lite TP1 with graphs is nondeterministic before any sleep, unpatched included, so exact match is not meaningful there; the graph-replay-after-wake property is covered by the unit test.)
• Wake latency: a wake that restores weights takes +0.2–0.3 s (TP2 0.37 s → 0.56–0.80 s; TP1 0.66 s → 0.95–0.99 s), the cost of copying the graph pool back from host memory. kv_cache-only wake is unchanged (0.07 s).

What this does not release

On ROCm the torch graph private pool is only 72–244 MiB per GPU, while graphs cost 948 / 1044 / 1543 MiB (DSv2-Lite TP1 / TP2 / DP2) and 1743 MiB (V4-Flash TP4) against --enforce-eager. The remaining 0.85–1.5 GB is HIP runtime graph memory held outside any torch pool, the ROCm counterpart of the graph-exec residue described in #59160. One difference matters: in a bare-process probe on MI355X, destroying graphs does return that memory (200 graphs: +4070 MiB captured, back to baseline after reset() + empty_cache()), whereas cuGraphExecDestroy returned 0 on GB300. So "destroy graphs on sleep, recapture on wake" would recover roughly the full graphs-vs-eager delta on ROCm, at the cost of recapture time (9–50 s here). That is a possible follow-up, not part of this PR.

Test Plan

• pre-commit run --files <changed files>.

• CUDA: csrc/cumem_allocator.cpp only loses code inside #ifdef USE_ROCM; run the cuMem and sleep-mode tests on CUDA.

• ROCm unit test (single GPU):

  pytest -v tests/basic_correctness/memory/cumem/test_cumem.py -k "cudagraph_pool_survives_sleep or plain_cudagraph_capture"

• ROCm end to end: DeepSeek-V2-Lite TP1/TP2/DP2 and DeepSeek-V4-Flash TP4 (V2 runner), sleep/wake cycles with greedy-output comparison, as above.

Test Result

• pre-commit: all hooks pass (ruff check/format, clang-format, mypy, typos, SPDX, forbidden imports, torch.cuda API check, test tethering).
• CUDA: g++ -std=c++17 against CUDA 12.9 headers builds the TU; the preprocessed output is identical to the base PR head apart from the __LINE__ values that CUDA_CHECK embeds in its messages. Not built with hipcc in this environment.
• ROCm (MI355X, numbers above): a port of test_cudagraph_pool_survives_sleep passes at levels 1 and 2 with the C++ change and fails at capture without it. The ROCm runs used a build that created the zero stream per create_and_map call; this PR creates it once per device.
• Re-run with the once-per-device build: pool_test.py passes at levels 1 and 2; DeepSeek-V4-Flash TP4 (V2 runner, with [Feature] Release WorkspaceManager scratch on sleep #59156 + [Feature] Offload KV-init runtime state on sleep #59158) 30/30 cycles bit-exact, flat at 10779 MiB asleep, wake 0.88 s (1.58 s with the per-call build); DeepSeek-V2-Lite TP2 10/10 cycles including a level-2 reload.
• Re-run on this head (1c6815d402; csrc/ unchanged since the build above): 25 cold startups with this build (DeepSeek-V2-Lite TP2 ×21, V4-Flash TP4 ×4, 5 nodes, every GPU pair) all came up, against 23 with the earlier per-call build. One earlier startup with this build had hung during KV-cache allocation (AMDKFD_IOC_ALLOC_MEMORY_OF_GPU); it did not reproduce and no common factor was found, so a rare driver-side issue cannot be fully excluded. Memory after wake vs before sleep: in-process −12 MiB (V4-Flash TP4) and 0 MiB (DSv2-Lite TP2); 30/30 cycles bit-exact on both, including a level-2 reload on TP2. test_cumem.py -k cudagraph 4 passed; test_sleep_cudagraph_custom_allreduce_tp2 passed.
• GSM8K (1319 questions, 5-shot, temperature 0), V4-Flash TP4 (V2 runner), all three PRs, this head: awake 0.951 / 0.945 → full wake 0.952, selective wake 0.948, +10 cycles 0.952 (per-call build: 0.943 / 0.945 → 0.948 / 0.948 / 0.945); unpatched V2-runner control 0.947 / 0.949 → 0.951 / 0.948 / 0.948.

Re-test after rebase onto main bdd31c31d (no conflicts)

AMD MI355X (gfx950), ROCm nightly 29468dde (ROCm 7.2.3), cumem allocator rebuilt from this branch:

• test_cumem.py: 21 passed, 3 runs (incl. test_cudagraph_pool_sleep at levels 1 and 2).
• test_sleep_mode.py -k test_sleep_cudagraph_pool, 2 GPUs: 5 passed, 3 runs (full, piecewise, breakable, custom-ar-tp2, off-by-default).
• test_sleep_mode_backend.py: 9 passed, 3 runs.
• DeepSeek-V2-Lite TP2 with --sleep-mode-offload-cudagraph, 5 runs of 3 cycles (level-1 full wake, level-1 selective wake, level 2 + reload_weights): outputs identical to awake in every cycle, with vLLM custom allreduce and with AITER custom allreduce. Asleep 9549 MiB/GPU vs 9696 MiB with the option off (the graph pool is 72 MiB/GPU on this model).
• The same suite with [Core] Release the CUDA graph pool on sleep instead of backing it up #59909 on top: all pass, 3 runs, outputs identical; the pool is no longer copied to host memory while asleep.

Re-test after review (0e54b3f: zero on wake instead of every mapping)

AMD MI355X, ROCm nightly 92044241a (already includes #59156, #59158, #59160), cumem allocator rebuilt from this branch. The .so no longer imports hipMemset, hipMemsetAsync or any hipStream* function.

Test	Result
test_cumem.py	21 passed ×3. Removing the wake-up memset makes only test_level2_discards_ordinary_tensor_with_weights_tag fail.
test_sleep_mode.py -k test_sleep_cudagraph_pool, 2 GPUs	5 passed ×3 (vLLM custom AR) and ×6 with VLLM_ROCM_USE_AITER=1 (AITER custom AR). No custom AR registers graph buffers.
fused AITER AR + group-quant test	Segfaults loading FP8 without #60367 (9/9); passes ×6 with it.
test_sleep_mode.py -k "undefined_remap or deep_sleep"	5 passed, 2 skipped ×3
test_sleep_mode_backend.py	9 passed ×3

• DeepSeek-V2-Lite TP2 with --sleep-mode-offload-cudagraph, 3 cycles (L1 full wake, L1 selective wake, L2 + reload_weights):
  • TRITON_MLA is bit-exact.
  • ROCM_AITER_MLA is bit-exact ×3 (AITER on, AITER MoE off, with [Bugfix][ROCm] Keep MLA derived weights in place on weight reload #60389 for L2). With every AITER kernel on, two identical requests already give different outputs before any sleep, on main too.
  • L1 wake takes 0.37–0.44 s, against 0.38–0.40 s before.
• CUDA (GB300): test_cumem.py 21 passed. test_sleep_mode.py -k "cudagraph_pool or undefined_remap" 6 passed. Main gives the same results.

AITER tests from #60525 (@indianspeedster)

8x MI355X (gfx950), ROCm 7.2.3, vllm/vllm-openai-rocm:nightly-bb87d227, cumem allocator rebuilt from this branch plus #60367:

Check	Result
custom-ar-tp2 with VLLM_ROCM_USE_AITER=1, before	SKIPPED
custom-ar-tp2 with VLLM_ROCM_USE_AITER=1, after	PASSED
New fused AITER test, with the _aiter_ops.py fix	PASSED
New fused AITER test, fix reverted	FAILED (hipIpcGetMemHandle ... invalid argument)

The new test needs #60367 to start on ROCm. CI's ROCm sleep-mode step runs on 1 GPU, so both TP=2 cases skip there.

🤖 Generated with Claude Code

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mawong-amd]

@aoshen02 on whether we need the zeroing, I dug into this more on MI300X (gfx942, amdgpu 6.16):

• You were right that the L2 NaNs in our earlier run came from [Bugfix][ROCm] Keep MLA derived weights in place on weight reload #60389, not KV.
• With AITER MoE off and [Bugfix][ROCm] Keep MLA derived weights in place on weight reload #60389 applied, ROCM_AITER_MLA is still sensitive to KV slots it hasn't written, but only to NaN bit patterns. With no sleep at all, overwriting the idle KV cache with finite junk (bf16 0.0115 or 3.0) changes nothing, while overwriting it with 0xFF changes all 8 prompts. TRITON_MLA doesn't care either way. So block reuse within a run is fine, and zeroing on wake only protects against unscrubbed VRAM.
• On a scrubbing driver, that can't happen: with no zeroing anywhere, discarded cuMem allocations from 2 MiB to 8 GiB all came back zeroed after remap.
• Cost: memset runs at about 4.4 TiB/s here, so about 35 ms per wake for a 150 GiB KV cache. That's roughly +50% on a kv_cache-only wake and a few percent on a weights wake.

So I agree it only matters if ≤6.10 is still supported. If it is, keeping it ROCm-only (or gated on the driver version) would avoid the cost elsewhere.

Separately, a correction to my earlier CI comment: the TP=2 sleep tests weren't actually skipped on AMD. pytest basic_correctness/ -m 'distributed(num_gpus=2)' in the existing MI300 Distributed steps already collects custom-ar-tp2, the new fused AITER test and test_deep_sleep_lora_tp2. I ran them the way those steps do and all three pass, with the fused per-group op in the compiled graph. Could you move the sleep-mode step back to mi300_1 in both test-amd.yaml and test_areas/basic_correctness.yaml?

[aoshen02]

@mawong-amd thanks for digging into this. The VRAM-scrub check and the cost numbers settle it for current drivers.

One question before I finalize: does vLLM still need to support amdgpu 6.10 and older, the drivers that don't scrub VRAM on release?

• If not, I'll drop the zeroing completely, so the ROCm wake path matches CUDA.
• If yes, I wonder in what case does zeroing help?

[mawong-amd]

Yes — vLLM's documented floor is ROCm 6.3, and ROCm 6.x stacks pair with amdgpu 6.8–6.10, the range that doesn't scrub VRAM on release. So those drivers are in the supported window.

The case where it helps: on those drivers a wake remap can return stale VRAM, and we showed ROCM_AITER_MLA reads KV slots it hasn't written — NaN bit patterns there change outputs. #44972's GDN case is the same shape (collapse to token 0). Both fail silently, with no error.

So I'd keep it: worst case ~35 ms per wake for a 150 GiB KV cache, and wake isn't on a known hot path. Cheap insurance for a supported driver range. Keeping it ROCm-only, as you proposed, also leaves CUDA untouched. Could be worth a one-line comment there noting it's unnecessary on amdgpu ≥6.14, which scrubs on release, so a future reader knows it can go once that floor moves.

[aoshen02]

Pushed 9ddebaa and 52c78f9 on top of the main merge:

• Zero-on-wake now runs on ROCm only, with a comment that amdgpu <6.14 can return stale VRAM and that the zeroing can go once newer drivers are the floor. CUDA wake is unchanged.
• The sleep-mode step is back on mi300_1 in both files.

Thanks for checking the driver range.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #93733 for commit 9b681ed61724.

[aoshen02]

/ci retry

[github-actions]

✅ Queued 4 failed job(s) for retry in Buildkite CI #93733.