Repository navigation
[ROCm][Bugfix] Pass per-sequence context lengths to the AITER paged MQA-logits kernels - #59420
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Did you do any perf tests after this? Looks like this might introduce perf regression. Also do you think that the best place to fix it is here or in aiter? |
simondanielsson
left a comment
simondanielsson
left a comment
There was a problem hiding this comment.
I think this makes sense, as I understand it AITER expects 1d context lens (in contrast to Nv deepgemm expecting 2d) and then reconstructs the full one here https://github.com/ROCm/aiter/blob/4ee387457385e9e18c50c4c808b9b4d0300514d6/aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py#L513.
But I agree with Andreas, can we check that we don't hit this new contiguous call unless expected? e2e perf numbers with both mtp on and off would be good.
Thanks!
9e59266 to
71535a3
Compare
e132b17 to
33a8124
Compare
If we introduce the fix in aiter to accept a 2D context length array, then we might have to change other frameworks and other models that use this method. E.g. ATOM Deepseek v2/v32 passes the per batch context length - https://github.com/ROCm/ATOM/blob/868cb029cc154038fb748c24397afb35e7c2bc57/atom/model_ops/attentions/aiter_mla.py#L3031. Similar to ATOM, I could compute this in indexer metadata builder so that it removes that copy per step and instead turns it into a persistent buffer, and the AITER calls read it. @AndreasKaratzas which method do you recommend?
Currently there is no native decode for next_n > 2 - #59421. @simondanielsson I have now added the Performance number and accuracy numbers in the PR description. |
|
I still think this belongs in aiter. Preferably vllm should be mostly a frontend. Imagine that aiter is bumped later and this becomes a NOP and we don't even know about it |
For this, couldn't we make the aiter impl backwards compatible so you can pass either 1d or 2d? |
33a8124 to
464cc22
Compare
|
We should ship this as an accuracy and functionality fix, before we land the fixes in AITER. |
@AndreasKaratzas @simondanielsson Thanks for the comments. I have drafted an AITER PR that supports both 1D and 2D context_lens for MTP so that it is backward compatible for other frameworks. |
simondanielsson
added
the
verified
…s kernels Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
464cc22 to
0ac615a
Compare
dllehr-amd
left a comment
dllehr-amd
left a comment
There was a problem hiding this comment.
Looks good. Thanks for catching this @amd-sriram
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #93383 for commit |
github-actions
Bot
added
the
ready
CI selector (shadow): 4 test steps (7 jobs) instead of 69 (93 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (4)
Would skip (today's rules run them) (66)
Would add (today's rules do not run them) (1)
AMD mirrors: would skip (54)
AMD mirrors: would add (4)
2 changed files · base |
|
/ci run |
|
❌ This PR is 19 commits behind upstream |
|
/amd-ci retry |
|
✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/amd-ci/builds/14210 |
|
/ci run |
|
✅ Triggered Buildkite CI #93626 for commit |
|
/ci retry |
|
✅ Triggered Buildkite CI #93784 for commit |
aiter.paged_mqa_logits takes one length per sequence, like the deepgemm kernels vllm-project#59420 fixed. Route the (B, next_n) table through _per_sequence_context_lens on this path too. Signed-off-by: build <build@local>
|
I will revert this After ROCm/aiter#6451 (Triton/Gluon) and ROCm/aiter#6450 (FlyDSL). gets merged and lands in Vllm pin ROCm/aiter#6153 got split in to 2 PRs and is being closed. |
6 participants
Error
MTP with native DSA-indexer decode above
next_n=2crashes during warm-up on MI355X (seen atnext_n=3andnext_n=6). GLM-5.2-MXFP4, TP4,num_speculative_tokens=5(next_n=6), Gluon paged MQA logits:Where it does not crash, the indexer silently uses wrong causal bounds: in a direct call at
next_n=3, the 8192-token sequence loses its last 2 valid positions in every row (9/12 rows correct).Cause
With native decode (
next_n > 1),rocm_fp8_paged_mqa_logitsreceives a(B, next_n)table of per-row context lengths,[b, j] = L_b - next_n + j + 1, built inindexer.py#L1136-L1152. The AITER Gluondeepgemm_fp8_paged_mqa_logitsanddeepgemm_fp8_paged_mqa_logits_stage1kernels read the table as a flat[B]array and some sequences get another row's length.Example: 4 requests,
next_n=3, context lengths 8192 / 1500 / 3000 / 700:bj=0j=1j=2L_bThe kernel expects
[8192, 1500, 3000, 700]and computes each row's bound itself. Reading the(4, 3)table as[4], it gets the first four values in memory,[8190, 8191, 8192, 1498]: request 0 loses its last 2 positions in every row, and requests 1 to 3 get lengths past their own KV pages, the likely source of the fault. Columnj = next_n - 1equalsL_b, which is what the kernel needs.Fix
Pass the kernels the last column, which is the per-sequence length.
[B]and(B, 1)inputs, the C4A Triton path and the torch fallback are unchanged.On ROCm native decode is used for
next_n1 and 2 today, so MTP with one speculative token goes through this path. #59421 raises that cap, which needs this fix.Test Plan
pytest tests/v1/attention/test_rocm_paged_mqa_context_lens.py(new): the helper returns the per-sequence lengths for a(B, next_n)table and passes[B]/(B, 1)through unchanged.rocm_fp8_paged_mqa_logitson one MI355X against a torch reference:B=4,next_n=3, contexts 8192/1500/3000/700, block size 64.Test Result
With this fix, GLM-5.3 MTP at
next_n=3andnext_n=5completes the full throughput sweep, GPQA (0.889 / 0.904) and RULER niah_single_2 (1.0 / 1.0).Throughput per GPU:
The no-MTP run is the same setup without speculative decoding (decode paged MQA logits on the FlyDSL kernel instead of Gluon).