[ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations - #47567
Merged
AndreasKaratzas merged 1 commit intoJul 4, 2026
Conversation
…nuations The AITER persistent MLA work-stealing kernel (get_mla_metadata_v1 + mla_decode_fwd persistent path, enabled in vllm-project#41990) is numerically wrong for multi-token prefill batches. The per-token errors are small in isolation but compound across chunked-prefill continuation passes until long-context decode collapses into repetition/garbage (vllm-project#47042). The kernel is correct for pure decode and for fresh single-chunk prefills, so gate it on batch shape: fall back to the non-persistent split-KV path whenever any request is a chunked-prefill continuation (more than one query token this step with part of its context already computed in an earlier chunk). Decode and single-chunk prefills keep the persistent performance path. Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This was referenced Jul 3, 2026
AndreasKaratzas
added
the
ready
Rohan138
added a commit
to ROCm/vllm
that referenced
this pull request
Jul 6, 2026
…ext lengths The persistent sparse-MLA work-stealing metadata is cached and only recomputed when metadata_key changes. The key used min(seq_lens, topk), but the metadata also depends on the per-request context length: a chunked-prefill continuation whose context exceeds topk has a uniform per-token KV length, while the first chunk ramps 1..topk. Both clamp to the same seq_lens, so a >=3-chunk prefill lets a later chunk reuse an earlier chunk's work partition and corrupt long-context output (vllm-project#47042). Include the clamped per-request context length and per-request query lengths in metadata_key. This keeps the persistent path for chunked prefill (superseding the disable-persistent workaround in vllm-project#47567) with no decode-cache regression. Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
4 tasks
raviguptaamd
added a commit
to raviguptaamd/vllm
that referenced
this pull request
Jul 7, 2026
…ths (disagg-hardened) Port of vllm-project#47766: the sparse-MLA persistent work-stealing metadata cache was keyed on min(seq_lens, topk) alone, colliding a first prefill chunk (context ramps 1->topk) with a continuation chunk (uniform at topk). A >=3-chunk prefill then reuses the first chunk's work partition and corrupts output, compounding across layers into long-context collapse (vllm#47042). Fix: add clamped per-request context length + per-request query lengths to metadata_key. Keeps the persistent kernel ON for chunked prefill (no perf regression); supersedes the disable-persistent mitigation in vllm-project#47567. The persistent kernel itself is correct given correct metadata (ROCm/aiter#4076); this is purely a vLLM-side cache-key bug. Disagg hardening (ours, carried from the vllm-project#47567-era int64 fix): under cudagraph/DP dummy-batch padding, seg_lengths (np.diff(query_start_loc_cpu)) has length num_reqs_PADDED while seq_lens_cpu[:num_reqs] has length num_reqs; keying/subtracting them raw broadcast-mismatches on DP ranks where padded != real (numpy ValueError -> rank dies, peers see only a gloo reset). Slice seg_lengths[:num_reqs]. vllm-project#47766 was validated TP8 single-node only, so it does not carry this slice; required for WideEP disagg. Co-authored-by: Rohan138 <rohan.potdar@amd.com> Co-authored-by: Chuan Li <chuali@amd.com> Co-Authored-By: Claude <noreply@anthropic.com>
NickLucche
pushed a commit
to NickLucche/vllm
that referenced
this pull request
Jul 15, 2026
…nuations (vllm-project#47567) Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
philippesic
pushed a commit
to philippesic/vllm-semantic-cache
that referenced
this pull request
Jul 19, 2026
…nuations (vllm-project#47567) Signed-off-by: Rohan Potdar <rohan.potdar@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
MohitAMD
added a commit
to MohitAMD/vllm
that referenced
this pull request
Aug 14, 2026
vllm-project#49649) AITER has no non-persistent sparse-MLA decode kernel for gqa_ratio=64 fp8/fp8 (asm_mla.cu:949: "fp8/fp8 with gqa_ratio=64 only supports persistent mode"). A persistent-kernel gate that falls back to the non-persistent split-KV path for such groupings (e.g. chunked-prefill continuations, cf. vllm-project#47567) crashes the prefill worker deterministically for GLM-5.1-FP8 DSA disaggregated serving (DP=8/TP=1 => gqa_ratio=64). Adds `sparse_mla_requires_persistent(num_heads_per_head_k, kv_cache_dtype)` that encodes the invariant, and a fail-fast guard in the sparse-MLA decode path that raises a clear error (pointing to vllm-project#49649) if persistent metadata is ever absent for a persistent-only grouping, instead of the opaque asm_mla.cu:949 kernel crash. main is already always-persistent, so this is a safety net that keeps a future gate from silently breaking gqa_ratio=64 fp8. Includes a CPU-only reproducer/regression test for the decision function. Refs vllm-project#49649 (also vllm-project#47567, vllm-project#47766, ROCm/aiter#4076). Signed-off-by: Mohit Deopujari <Mohit.Deopujari@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
2 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.
[ROCm] Disable persistent sparse-MLA kernel for chunked-prefill continuations
Summary
Fixes #47042. GLM-4.6/4.5 / DeepSeek-V3.2-style sparse-MLA (DSA) models on ROCm
produce correct output up to ~20K tokens and then collapse into repetition /
garbage at longer contexts.
Root cause: the AITER persistent MLA work-stealing kernel
(
get_mla_metadata_v1+ thework_meta_datapersistent path inaiter.mla.mla_decode_fwd, enabled unconditionally for the sparse-MLA backendin #41990) is numerically wrong for multi-token (prefill) batches. The
per-token error is small in isolation, but chunked prefill runs a request
through several forward passes and the error compounds through the KV cache
across passes/layers until long-context decode collapses.
This is why the failure is gated on chunk count, not context length:
The persistent kernel is correct for pure decode (
qseqlen==1) and forfresh single-chunk prefills. This patch gates it on batch shape: it falls
back to the (correct) non-persistent split-KV path whenever any request in the
batch is a chunked-prefill continuation — i.e. it has more than one query
token this step and part of its context was already computed in an earlier
chunk (
seq_len > query_len). Decode and single-chunk prefills keep thepersistent performance path, so there is no decode throughput regression.
Why this is not a duplicate
same file but fixes a different bug: a stale-metadata race during CUDA-graph
replay. The bug here is a numerical correctness bug in the kernel itself —
it reproduces with
--enforce-eager/-cc.cudagraph_mode=NONE(no graphreplay), so metadata-ordering is not involved.
of this regression), not a fix.
Isolated (kernel-only) reproduction
A standalone script (no vLLM/model) builds a sparse-MLA batch, regenerates the
persistent metadata via
get_mla_metadata_v1, and compares the persistentkernel and the non-persistent kernel against a torch fp32 reference:
The non-persistent kernel matches the fp32 reference tightly everywhere; the
persistent kernel is correct for decode-shaped batches but produces large
errors for prefill-shaped batches. (The underlying kernel bug is also worth an
AITER-side fix; this PR is the vLLM-side mitigation.)
Testing (MI355X / gfx950, TP8, GLM-FP8, vLLM nightly + aiter 0.1.16.post2)
Needle-in-a-haystack via raw
/v1/completions, greedy,max-num-batched-tokens 8192(default small chunk), cudagraph/torch.compile enabled:No accuracy regression on short context:
lm_eval gsm8k --num_fewshot 20(concurrency 256): 0.944 flexible / 0.945strict exact_match, unchanged vs baseline (decode path is untouched).
Notes
glm_moe_dsa/DeepSeek-V3.2 indexer) backends.
author owns and has validated it.
pre-commit run --files vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.pyshould be run before submission.