[ROCm] Re-enable cudagraph memory profiling, captured on the current stream - #48526
Conversation
…stream Re-enables cudagraph memory profiling on ROCm (is_cuda() -> is_cuda_alike()) that vllm-project#48440 disabled as a stopgap, and fixes the decode-throughput regression that motivated that revert. profile_cudagraph_memory() captures throwaway graphs via graph_capture(), which by default allocates a fresh side stream. torch's caching allocator pools free blocks per stream, so the side-stream forward strands a persistent aiter MLA-sparse scratch buffer (natively owned, not reclaimed by empty_cache) in a separate pool. Allocated before the real KV cache, it shifts the KV/activation buffers' physical placement, slowing bandwidth-bound decode ~20% (GLM-5.2-MXFP4, tp4, batch1: 83.4 -> 69.5 tok/s). Fix: on ROCm, capture on the current stream instead, so profiling allocations share the runtime pool and placement matches the no-profiling path. graph_capture() gains an optional graph_capture_context param for this. The graphs are discarded, so a side stream is unnecessary. Scoped to ROCm: on CUDA the current stream is the legacy default stream, on which capture cannot begin, so its side-stream path is unchanged. Restores baseline throughput with the estimate enabled. Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Fixing the issue #48453 |
|
Confirmed on a second model/arch — thanks for the fix. DeepSeek-R1-0528-MXFP4, MI350X (gfx950) x8, TP=8, ROCM_AITER_MLA, fp8 KV (
Profiling ran and produced a sane estimate — all 8 workers logged Independent root-causing landed on the same mechanism: the side-stream capture strands a natively-owned aiter scratch buffer in a separate per-stream allocator pool — invisible to (Confirmation done with Claude Code.) |
Apply ruff-format's parenthesized multi-line `with` form to the profiling graph_capture() call in profile_cudagraph_memory(), fixing the pre-commit failure on the prior commit. Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
| per_graph_estimate = {} | ||
| encoder_memory_estimate = 0 | ||
|
|
||
| # On ROCm, capture these throwaway profiling graphs on the current stream |
There was a problem hiding this comment.
This PR LGTM for now, but should we check if there's a way to tell the model runner to delete the AITER persistent kernel buffers after CG capture? I wonder why this issue isn't seen on CUDA cc @dllehr-amd @AndreasKaratzas @tjtanaa
There was a problem hiding this comment.
Thanks @Rohan138. I looked into the "free the AITER buffers after capture" idea — it doesn't recover the throughput, and I think the reason is worth spelling out because it's not a leak and not a graph-memory-release problem.
What the buffer is. The stranded block is a ~76 MiB (79,691,776 B) persistent scratch/workspace that AITER's sparse-MLA decode kernel allocates lazily on its first forward pass (dispatched from mla_decode_fwd at vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py:702). It's allocated inside AITER's compiled kernel code — a gc walk finds no owning Python torch.Tensor — so nothing on the vLLM side holds a handle to it. It's a legitimate runtime workspace: the decode kernels reuse it on every step, so it's meant to persist.
Why the profiling stage strands it. profile_cudagraph_memory() runs its throwaway forward+capture inside graph_capture(), which by default allocates a fresh side stream (parallel_state.py:600, torch.cuda.Stream(...)). torch's caching allocator pools free blocks per stream. So when that first ever sparse-MLA forward happens on the side stream, the 76 MiB workspace is created in the side stream's allocator pool rather than the default/runtime pool. Confirmed with memory_snapshot: the fast path has all segments on one (default) stream; the slow path has an extra segment (~0.07 GiB) owned by a second stream.
Why it degrades decode later. That side-stream block is allocated before the real KV cache. Its presence in a separate pool shifts the physical HBM placement of the KV/activation buffers allocated afterward. Every bandwidth-bound decode kernel then reads from a worse physical region — traces show ~+20–30% on bandwidth-bound kernels (allreduce, wvSplitK, concat_and_cache_mla) while compute-bound MFMA barely moves. Same kernels, same launch configs, same total reserved memory — purely a placement effect. Pre-#47366, ROCm never ran profiling, so this workspace was first allocated on the default stream during real warmup → good placement → full speed.
Why freeing after capture doesn't work (I benchmarked it, MI355 tp4, GLM-5.2-MXFP4):
| approach | tok/s |
|---|---|
| side-stream capture (broken) | 69.5 |
| side-stream capture + explicit gc.collect() + empty_cache() after capture | 69.3 |
| current-stream capture (this PR) | 83.3 |
Two reasons it can't help: (1) empty_cache only returns unreferenced blocks — this one is still live (AITER holds it and decode needs it), so it's not reclaimed; and (2) even a working free is too late — the placement damage is done when the KV cache is allocated after profiling, and the side-stream pool already exists at that point. prewarm_default and reset_ws (in my root-causing) failed for the same ordering reason. Capturing on the current stream is the only approach that prevents the separate pool from ever existing, so KV placement matches the no-profiling path.
Why CUDA doesn't hit this. Two independent reasons: (1) the triggering workspace is allocated by AITER, which is ROCm-only — CUDA uses FlashAttention/FlashInfer backends that don't create this particular persistent buffer; and (2) the fix itself is ROCm-specific anyway — on CUDA current_stream() at init is the legacy default stream, on which graph capture can't begin, so CUDA is forced onto a side stream regardless. That's why the change is scoped to is_rocm().
Net: it's not a graph-memory-release issue (a profiling run that captures no graph at all still strands the buffer and is still slow), and we don't change any AITER pool — we just ensure that first forward runs on the runtime stream so the workspace lands in the right pool. Happy to file a follow-up if AITER wants to expose a way to pre-touch/own that workspace explicitly.
There was a problem hiding this comment.
Please let me know if it resolves your question. It is really an aiter related issue, cuda does not have it.
There was a problem hiding this comment.
Well ... my question is that this workspace or its handles should likely be allocated by aiter.get_mla_metadata_v1 from vllm: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py#L550
So while the allocation itself does live inside AITER, why doesn't torch "see" this workspace and gcit after cudagraph capture in https://github.com/vllm-project/vllm/blob/85f50eb/vllm/v1/worker/gpu_model_runner.py#L5315?
There was a problem hiding this comment.
Good question — the key is that there are two separate buffer sets here, and the stranded block isn't the one get_mla_metadata_v1 touches.
The get_mla_metadata_v1 buffers are torch-visible — and they're not the problem. _mla_work_meta_data, _mla_work_info_set, _mla_work_indptr, mla_reduce* (allocated at rocm_aiter_mla_sparse.py:442-457) are plain torch.empty(...) tensors owned by the metadata builder. get_mla_metadata_v1 doesn't allocate them — it's passed them as out-params and just fills them in (:550-568). torch sees them perfectly. But they're persistent attributes on a live object held for the runner's lifetime, so gc/empty_cache correctly won't reclaim them — and they're small. These aren't what strands the 76 MiB.
The stranded 76 MiB comes from mla_decode_fwd, not get_mla_metadata_v1. The block is a scratch/workspace the AITER sparse-decode kernel allocates lazily inside compiled kernel code on its first forward (dispatched from mla_decode_fwd at :704). It's not wrapped by any Python torch.Tensor — I confirmed with a gc walk that no owning object exists. torch's caching allocator tracks the underlying segment (that's why it shows up in memory_snapshot), but there's no Python handle to it, so it's invisible to any Python-side cleanup.
So on your gpu_model_runner.py#L5315 question specifically — that's why the gc.collect() + empty_cache() there is a no-op on this block, for two independent reasons:
- empty_cache only returns free, unreferenced blocks to the OS. This block is live — AITER holds it internally and every decode step reuses it — so it's never on the free list to reclaim.
- Even if it were reclaimable, empty_cache doesn't relocate blocks or change which stream-pool owns them. The damage isn't that the buffer exists; it's that (during profiling) its first-ever allocation happened on the side stream, so it landed in the side-stream allocator pool, which shifts the physical placement of the KV cache allocated afterward. That placement is baked in before L5315 ever runs.
I did benchmark the explicit-free approach (in my earlier reply's table): side-stream capture + gc.collect() + empty_cache() after capture = 69.3 tok/s vs. 83.3 for capturing on the current stream. It doesn't recover throughput, for exactly these reasons.
Net: nothing to gc here — the visible buffers should persist, and the invisible one is native-owned + live. Capturing that first forward on the runtime stream is what keeps the workspace in the right pool in the first place. Happy to file an AITER follow-up if it'd help to expose a way to pre-own/pre-touch that decode workspace explicitly.
AndreasKaratzas
added
the
ready
| # ROCm is included: #44825 moved the profiler to | ||
| # torch.accelerator.get_memory_info (reliable on ROCm, as used by | ||
| # the AMD-CI mem tests), and graph_pool_handle resolves to the same | ||
| # torch.cuda handle the live capture path already uses on ROCm. | ||
| # XPU stays excluded (see #39977). |
There was a problem hiding this comment.
I think this can be deleted, we don't need now-irrelevant references to historical state.
Partial revert of vllm-project#48526: this removes only its `current_platform.is_rocm()` branch, so ROCm rejoins the side-stream path CUDA already uses. The rest of vllm-project#48526 -- the gc freeze, the capturing-enabled guard, the cleanup-only try -- is retained. With that branch gone nothing passes an explicit context, so the `graph_capture_context` parameter added to the module-level `graph_capture()` goes too; GroupCoordinator.graph_capture keeps its own long-standing one. The branch was added to fix a decode slowdown reported in vllm-project#48453. It offered two rationales, neither substantiated: that the slowdown came from stream synchronisation, and that torch's per-stream block pooling stranded a natively owned aiter scratch buffer and shifted the KV cache's physical placement. The second does not hold mechanically: per-stream block pooling governs allocations made through the torch allocator, not a natively-owned buffer, and a profile-run allocation is reclaimed before the KV cache is sized. It also predicts a memory symptom, while the report was a latency one. Measured on both models at the reported configuration, alternating arms: DeepSeek-V4-Pro TP8, conc 1/8/32/128: +0.16% / +0.32% / -0.46% / -0.01% GLM-5.2-MXFP4 TP4, conc 1/8/32/128: +0.30% / -0.55% / -0.79% / +1.64% Every delta is smaller than the same-arm run-to-run spread at that concurrency, with mixed sign, against a claimed ~20% effect. At concurrency 1 -- where vllm-project#48453's regression and vllm-project#48526's fix were both demonstrated -- the arms are indistinguishable. Peak VRAM is equal or lower without the divergence, and nothing is stranded. The branch is also MRV1-only; MRV2's capture path has no equivalent and has not regressed. Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
4 participants
Purpose
Re-enables cudagraph memory profiling on ROCm (
is_cuda()->is_cuda_alike()) that #48440 disabled as a stopgap, and fixes the decode-throughput regression that motivated that revert.profile_cudagraph_memory()captures throwaway graphs viagraph_capture(), which by default allocates a fresh side stream (torch.cuda.Stream(...)). torch's caching allocator pools free blocks per stream, so the side-stream forward strands a persistent aiter MLA-sparse scratch buffer (natively owned, not reclaimed byempty_cache) in a separate pool. Allocated before the real KV cache, it shifts the KV/activation buffers' physical placement and slows every bandwidth-bound decode kernel ~20%.Fix: on ROCm, capture on the current stream instead (
torch.cuda.current_stream(...)), so profiling allocations share the runtime pool and placement matches the no-profiling path.graph_capture()gains an optionalgraph_capture_contextparam to select the stream. The graphs are discarded, so a side stream is unnecessary. Scoped to ROCm ~@~T on CUDA the current stream is the legacy default stream, on which capture cannot begin, so its side-stream path is unchanged.Not a duplicate: follow-up to the merged stopgap #48440 and the profiling added in #47366; no open PR re-enables ROCm cudagraph profiling or changes the profiling capture stream (#46515 is an unrelated IPC-handle leak). AI assistance (Claude) was used to root-cause the regression and prepare this PR.
Test Plan
Serving throughput A/B on MI355 (gfx950), tp4,
amd/GLM-5.2-MXFP4, comparing profiling-disabledmainagainst this branch, both built clean:Correctness is covered by the existing ROCm spec-decode e2e suite exercising this path:
pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference".Test Result
Output token throughput, same config, MI355 gfx950 tp4:
main(#48440 stopgap)Restores baseline throughput (83.25 vs 83.42, within run-to-run noise) with the memory estimate enabled, closing the ~20% gap. CUDA is not exercised by this change (ROCm-scoped).
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.