Repository navigation
Conversation
On gfx1250 the aiter ASM MLA decode kernels need FP8 Q and a KV layout that sglang's KV pool does not use. With SGLANG_USE_AITER=1, MLA decode in the aiter backend calls aiter's paged Gluon mla_decode_fwd (bf16 Q, bf16 or FP8 KV) for a paged KV pool and plain causal attention with one query token. Other single-rank calls use decode_attention_fwd; DCP keeps its existing path. The paged kernel is used only when aiter carries ROCm/aiter#5935; before it, a batch with one segment per sequence writes softmax state over the output without an error. Co-authored-by: Cursor <cursoragent@cursor.com>
With SGLANG_USE_AITER=1 on gfx1250, aiter or triton MLA decode skips the standalone RoPE and uses aiter fused_qk_rope_cat_and_cache_mla, the same fused branch aiter MLA takes on gfx95, with a bf16 Q. sglang's RotaryEmbedding keeps cos and sin in one cos_sin_cache, so the kernel gets its two halves as strided views. Limited to plain rope modules whose table covers the whole rope head, without DSA or DCP, and to bf16 or fp8 e4m3 KV caches. Requires aiter with ROCm/aiter#5501. The layer-static terms of the gate are computed once in __init__. A CPU unit test pins down each term of the gate and the cos/sin order of the tables. Co-authored-by: Cursor <cursoragent@cursor.com>
mqhc2020
changed the title
mqhc2020
changed the title
1 task
mqhc2020
changed the title
This branch has not been deployed
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
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.
Motivation
On gfx1250, MLA decode runs the Triton kernel. AITER has a faster ASM decode kernel (fp8 Q, fp8 KV). This PR switches decode to it and folds its input preparation into the existing fused RoPE/KV-cache kernel, so no extra kernels are added per layer.
Changes
SGLANG_AITER_MLA_GFX1250_ASM=1with--decode-attention-backend aiter)Default behavior (flag off) is unchanged.
Dependencies
Results
Kimi-K2-Thinking MXFP4, 8× gfx1250, TP8, input/output 8192/1024, same server args apart from the decode backend:
GSM8K (1319 questions, 5-shot): Triton 0.934, ASM 0.929.
Test
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.