perf(fa2): port decode GQA group-swap into d128 varlen launcher (gated OFF) - #48
Closed
localai-bot wants to merge 1 commit into
Closed
perf(fa2): port decode GQA group-swap into d128 varlen launcher (gated OFF)#48localai-bot wants to merge 1 commit into
localai-bot wants to merge 1 commit into
Conversation
…d OFF) Ports vLLM's FA2 seqlenq_ngroups_swapped decode optimization into LaunchDecodeVarlenFA2Bf16 (gate VT_FA2_DECODE_GQA_SWAP, default OFF), so the Qwen3-dense head_dim-128 varlen decode launches (batch, kv_heads) instead of (batch, hq): the ngroups query heads of a KV group pack into seqlen_q, KV is read once per group, and the num_splits heuristic sees batch*kv_heads. This halves the CTA count at batch>=2, the exact over-wave #47 localized as the MXFP4 c2-c8 gap (ours decode flash grid (batch,hq) vs vLLM's (batch,kv_heads); decode flash 63.7us c2 / 218us c8 vs vLLM 41.7/151). Presented WITHOUT a materialized transpose via kv-major-group-minor strides, a 1:1 mirror of the already-shipped d256 LaunchDecodeFA2Bf16 swap (the vendored flash_fwd_kernel.h get_lse_tile/combine already honor the flag in both the num_splits==1 direct-write and >1 combine paths). OFF path is byte-identical to the shipped plain-varlen reduction; ON is non-byte-exact only when num_splits>1 (split reduction order changes -> near-tie, toward vLLM's own numerics). Blast radius: the shared d128 varlen decode launcher (Qwen3-dense 0.6B/4B/8B + MXFP4-8B); the 27B/35B gate models decode on their own d256 path, untouched. Correctness (dgx GB10, sm_121a, CUDA 13.0): - RED-first op test: test_ops_paged_attn 5/5 varlen-d128 cases, 280/280 assn (both GQA ratios x batch{1,2,4,8} x short+long ctx; swap_launches==1 proves the swapped grid engaged; swap-vs-plain near-tie; MHA qpk==1 inert). RED proven: wrong swapped o_head_stride -> 26528 violations. Full binary 28/28 cases / 454679 assertions, no regression. - compute-sanitizer memcheck: 0 errors, 0 bytes leaked. - #44 MXFP4 e2e smoke (default async/graphed) swap-ON: 3/3 deterministic TOKEN-EXACT + coherent, byte-identical to swap-OFF (no token flip e2e). Default stays OFF: the c1-c8 x3 binding re-bench + default flip is the recorded next step (parity-enablers: flip only on a measured speed win with no regression). #47 projects flash alone closes ~28%@c2 / ~55%@c8 of the gap. Records: STATUS/BENCHMARKS/FEATURES + kernel/quantization matrices + state (KERNEL-FA2-GQA-SWAP) + NOW + benchmark-record + parity-ledger. FOLLOWING_AGENTS_PROTOCOL Assisted-by: Claude Code:claude-opus-4-8 [ClaudeCode]
localai-bot
marked this pull request as ready for review
August 6, 2026 03:06
Collaborator
Author
|
Landed on main as the squash (mudler authorship). 🤖 Generated with Claude Code |
mudler
added a commit
that referenced
this pull request
Aug 6, 2026
…te (#48) row/KERNEL-FA2-GQA-SWAP squash. LaunchDecodeVarlenFA2Bf16 gains the upstream mha_fwd_kvcache ngroups swap (logical q/o reshape via strides, no transpose; h=kv_heads, seqlen_q=ngroups; num_splits on swapped dims; 1:1 mirror of the shipped d256 LaunchDecodeFA2Bf16). Gate VT_FA2_DECODE_GQA_SWAP default OFF, OFF path untouched. Blast radius: Qwen3-dense d128 family only (0.6B/4B/8B); gate models already swap on d256; Voxtral MHA inert. Gates: op RED-first 5/5 (280 assn; wrong-stride injection fails loudly 26,528 violations), full binary 28/28 (454,679), memcheck 0/0, #44 smoke swap-ON 3/3 token-exact AND byte-identical to OFF (graphed). Flip conditioned on the c1-c8 re-bench + 0.6B/4B e2e (next campaign); #47 projects ~28% c2 / ~55% c8 gap closure. FOLLOWING_AGENTS_PROTOCOL Assisted-by: Claude Code:claude-opus-4-8 [ClaudeCode]
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.
Ports vLLM's FA2
seqlenq_ngroups_swappeddecode optimization intoLaunchDecodeVarlenFA2Bf16(gateVT_FA2_DECODE_GQA_SWAP, default OFF), the #47-localized MXFP4 c2-c8 lever. The Qwen3-dense head_dim-128 varlen decode now launches(batch, kv_heads)instead of(batch, hq), halving the CTA count at batch>=2 (the exact over-wave #47 measured). A 1:1 mirror of the already-shipped d256LaunchDecodeFA2Bf16swap (the vendored kernel's get_lse_tile/combine already honor the flag). OFF path byte-identical; ON is a split-reduction near-tie toward vLLM's numerics.Blast radius: the shared d128 varlen decode launcher (Qwen3-dense 0.6B/4B/8B + MXFP4-8B). 27B/35B gate models decode on their own d256 path, untouched.
Correctness (dgx GB10, sm_121a, CUDA 13.0):
test_ops_paged_attn5/5 varlen-d128 cases, 280/280 assn (both GQA ratios x batch{1,2,4,8} x short+long ctx;swap_launches==1proves the grid engaged; swap-vs-plain near-tie; MHA-inert). RED proven: wrong swapped stride -> 26528 violations. Full binary 28/28 / 454679 assn, no regression.Default stays OFF. The c1-c8 x3 binding re-bench + default flip is the recorded next step (parity-enablers). #47 projects flash alone closes ~28%@c2 / ~55%@c8 of the gap.
🤖 Generated with Claude Code