Repository navigation
[Bugfix][ROCm] Fix startup OOM in AITER MLA FP8 prefill workspace sizing - #57923
Merged
dllehr-amd merged 2 commits intoSep 23, 2026
Merged
Conversation
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
simondanielsson
marked this pull request as ready for review
September 21, 2026 11:48
simondanielsson
requested review from
AndreasKaratzas and
tjtanaa
as code owners
September 21, 2026 11:48
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
7 of 8 tasks
AndreasKaratzas
added
the
ready
|
✅ @simondanielsson, CI is now available for this PR.
|
Contributor
|
/ci run |
|
❌ @okorzh-amd, Only reviewers with write access can use CI commands before CI is delegated to the PR author. |
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #90611 for commit |
dllehr-amd
self-requested a review
September 23, 2026 14:51
dllehr-amd
approved these changes
Sep 23, 2026
dllehr-amd
left a comment
dllehr-amd
left a comment
Contributor
There was a problem hiding this comment.
Looks good. Thanks @simondanielsson
Contributor
|
@simondanielsson on the Aiter side, can we add a benchmark to check the workspace sizes for future cases where we may explode the allocations. Thanks @jamesETsmith for the recommendation. |
Contributor
Author
Added regression test in the aiter PR: ROCm/aiter@d22abd7: |
fululi12
added a commit
to fululi12/vllm
that referenced
this pull request
Sep 28, 2026
…unks `_expand_page_indices_kernel` is launched with a grid of one program per request, and each program walks that request's tokens in a serial loop. The parallelism is therefore entirely over the batch, which degenerates at low concurrency: with num_reqs == 1 a single workgroup expands the whole sequence while the rest of the GPU idles. On a long-context decode workload (MI355X, TP8, ~97k-token context, concurrency 1) this measured 61.6 us per call and 7 calls per model step, 0.431 ms or 1.23% of a 34.9 ms step. For scale, 97k int32 stores is about 388 KB -- well under a microsecond of HBM bandwidth. The kernel is not bandwidth-bound, it is serialised. Add a second grid dimension over token chunks and drop the loop, so each program handles exactly one BLOCK_SIZE chunk. The arithmetic, the masks and the stores are unchanged; only the iteration space is mapped onto the GPU rather than onto one workgroup. The chunk count is derived host-side from the block table width, an upper bound on tokens per request, so no device synchronisation is introduced; over-launched programs mask out entirely. This affects any MLA model on ROCm at low concurrency, not one model in particular. Test plan --------- Bit-exactness against the existing serial kernel, both built from source and compared on device over 8 shapes: conc=1 long context, kbs=1 97000 indices identical conc=1 long context, kbs=16 97000 indices identical exact multiple of 1024 4096 indices identical 1024+1 1025 indices identical 1024-1 1023 indices identical ragged 4 reqs 38011 indices identical kbs=16 ragged tail 4212 indices identical single token 1 indices identical End-to-end, same workload, 1200 s per arm against a shared baseline: interactivity p90 110.28 -> 112.06 (+1.61%), tpot p90 0.00907 -> 0.00892. Accuracy is unaffected by construction -- the page indices are bit- identical, so the model reads exactly the same KV -- and this was confirmed by GSM8k over the full 1319 questions remaining at 0.9659. AI assistance ------------- AI assistance was used (Claude). The measurements above come from runs on real hardware, not estimates. The submitter has reviewed every changed line. Not a duplicate: no open PR touches `expand_page_indices`; the nearby ROCm MLA PRs (vllm-project#57085, vllm-project#53492, vllm-project#57923, vllm-project#46172, vllm-project#54899) address DCP, the sparse indexer, prefill workspace sizing and head padding respectively. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: lifulu <fululi12@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fixes a startup OOM on Kimi K2.5 (or similar) when using the AITER ASM MLA prefill kernel, caused by overly pessimistic bounds on the size of the required persistent metadata buffers. These sizes were subsequently passed to the workspace manager which triggered the OOM.
Triggered by: The AITER bump to v0.1.22post1 picked up ROCm/aiter#3606 which slightly changed the logic for workspace sizing.
Long term fix in AITER: ROCm/aiter#5759
Why
For causal FP8 MLA the AITER backend uses a persistent kernel. It needs buffers to hold the work items to be completed by each work group. The metadata scheduler is responsible for sizing these buffers so that they are sufficiently large to handle any requests that might come in.
The MLA metadata scheduler now emits a very loose bound. The bound is based on the assumption that every request can use max_qlen query tokens. (The actual upper bound is that the sum of query tokens equals max_num_batched_tokens). If one uses these to size the logits/attn_lse partial buffers we get huge KV reservations:
Short term fix: Size with assumption that the total query tokens is bounded by the token budget.
Long term fix: AITER should give us this budget-aware max buffer size. Ticket: ROCm/aiter#5729
References:
Test Result
Run DSv3 with FP8 KV on MI350 with max num batched tokens 16k and model len 160k.
vllm serve deepseek-ai/DeepSeek-V3-0324 \ --tensor-parallel-size 8 \ --enable-expert-parallel \ --kv-cache-dtype fp8 \ --gpu-memory-utilization 0.9 \ --max-model-len 160kBefore:
After:
Also long context bench serve runs without crashing
GSM8k DSv3 fp8 kv
This branch
Main branch before bumping to aiter v0.1.22post1
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.