perf(cuda): head_dim-128 arm for the warp-split-KV decode kernel, gated OFF (#382) - #425
Merged
Conversation
…ed OFF
PagedAttentionDecodeOptKernel is reachable only at `d == 32 * kDecEpl` (256),
so head_dim 128 -- what Qwen3-dense, Llama and Mistral actually use -- falls to
the generic block kernel. This adds the EPL=4 instantiation behind
VT_ATTN_DECODE_D128, DEFAULT OFF.
The kernel is templated on elements-per-lane. EPL == 8 forwards through
LoadRowN<8> to the existing LoadRow8, so the d256 path keeps the same
transactions, the same order and the same accumulation -- the d256 launcher
still instantiates the default EPL and is byte-for-byte what it was. EPL == 4
adds one 64-bit load per lane for bf16 and one 128-bit for f32; 32 lanes * 4
elems == 128, so a warp still covers exactly one head-dim row.
No GQA carve-out: PagedAttentionDecodeGqaKernel sits inside the d == 256 branch,
so at head_dim 128 it can never run, and excluding qpk == kDecGqaQG would strand
exactly those models (e.g. Qwen3-32B) on the block kernel this arm replaces.
Gated OFF because it is correctness-complete but NOT byte-exact against the
block kernel it replaces: the two reduce the KV sequence in a different ORDER
(warp-strided online softmax vs the block kernel's per-tile loop), so a greedy
anchor can move at an exact bf16 tie. OFF keeps every existing golden
byte-identical. This mirrors how the FA2 decode GQA group-swap landed --
correctness-complete and gated OFF -- before it was flipped ON against the full
gate. The flip is a separate change and owes the near-tie razor, the
distributional gate, and regen under the ratified-tie rule.
The env var is also the same-binary A/B your perf protocol requires: one binary,
VT_ATTN_DECODE_D128=0/1, no rebuild between arms.
EVIDENCE (sm_110 / Jetson AGX Thor, Qwen3-1.7B-NVFP4A16, W4A16 Marlin build):
Kernel numerics, EPL=4 vs the generic block kernel, 8 context lengths
straddling the 32-token page boundary, 2048 elements each:
ctx max_abs max_rel bitwise%
1 0 0 100.0000
31 0 0 100.0000
32 0 0 100.0000
33 0 0 100.0000
64 0 0 100.0000
100 3.725e-09 0 99.9512
685 0 0 100.0000
1000 6.104e-05 5.076e-03 99.9512
Six of eight are bit-for-bit identical. The worst case is 0.5% relative on
small-magnitude outputs at ctx 1000 -- named plainly because that is the
magnitude at which a near-tie flips.
Comparator self-test: seeding 1-ULP corruption into 3 of the 2048 elements of
the ctx=31 case (offsets 5, 1024, 2047) moves that row from 100.0000% to
99.8535% bitwise with max_abs 4.883e-04, and no other row changes. So the
comparison is demonstrably able to fail and localizes correctly. This is a
COMPARATOR self-test, not a kernel negative control.
Serving A/B, same harness / model / config, c=1,2,4,8, 128 tokens, temp 0:
OFF 81.630 / 162.738 / 316.255 / 477.070 tok/s
ON 131.023 / 258.953 / 508.139 / 936.048 tok/s
ratio 1.61x / 1.59x / 1.61x / 1.96x
Build: the single TU compiles clean on the device (nvcc, sm_110, Release,
`ninja CMakeFiles/vllm.dir/src/vt/cuda/cuda_paged_attn.cu.o`, exit 0).
LIMITATIONS, stated because they bound what the numbers above prove:
- ONE model, ONE head_dim, ONE arch. Not measured anywhere but sm_110.
- The serving A/B is a SINGLE repetition per point (~1 s per point) and was
taken across two builds, not through the new flag.
- No model-level token-exact gate was run against this arm.
- LoadRowN<4, float> and the windowed (HasWindow=true) EPL=4 instantiation
are compiled but unexercised by the evidence above.
- The GQA-fused kernel remains d256-only; this change does not widen it.
Refs #382.
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:claude-opus-5 [ClaudeCode]
The cherry-picked change (Filip Sajdak, authorship preserved) refs #382, but that issue was never in the roadmap intake table. AGENTS.md requires the number in three places that must agree -- the issue table, the row's spec, and the PR body -- and only two of the three were true. Placed under `KERNEL-ATTN-PAGED`, which is the row that anchors `src/vt/cuda/cuda_paged_attn.cu` and therefore owns the decode-opt kernel this arm extends. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code]
Keeps the branch current while CI is queue-blocked (#274). No conflicts. Amended: git merge had already created this commit, so the message that was meant to accompany it never applied and it carried git's default -- which has no trailers at all and reds the trailer gate. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code]
Keeps Filip Sajdak's d128 decode arm current while CI is queue-blocked (#274). No conflicts; the kernel commit is untouched and still authored by him. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code]
localai-bot
force-pushed
the
row/KERNEL-ATTN-DECODE-D128
branch
from
August 11, 2026 21:13
2a65376 to
cb5f151
Compare
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.
Lands @filipsajdak's #383 on current
main, with authorship preserved (git logshowsauthor=Filip Sajdak). Issue #382. Supersedes #383, which is red for reasons that are not its own.The work is @filipsajdak's, unchanged
PagedAttentionDecodeOptKernelis reachable only atd == 32 * kDecEpl(256), so head_dim 128 — what Qwen3-dense, Llama and Mistral actually use — never reaches it and falls to the generic block kernel. This adds theEPL=4instantiation behindVT_ATTN_DECODE_D128, default OFF.EPL == 8forwards throughLoadRowN<8>to the existingLoadRow8, so the d256 path keeps the same transactions, order and accumulation — byte-for-byte what it was.Gated OFF for an honest reason: it is correctness-complete but not byte-exact against the block kernel, because the two reduce the KV sequence in a different order, so a greedy anchor can move at an exact bf16 tie. That mirrors how the FA2 decode GQA group-swap landed (#48) before being flipped ON (#49).
Why #383 was red, and why none of it was the PR
Both failures were stale-base — it sat 46 commits behind:
cuda-fat-buildfailed the gencode audit on seven files, every one of them Marlin:Not one is
cuda_paged_attn.cu, the only source the PR touches. That is precisely the defect #407 fixed — "Marlin gencode drifted from the feature table" — which landed after this branch was cut.agent-recordfailed onaudit-live-rows, which reads live branch state at run time. Reproduced at the PR tip:--checkreturns rc=0, both before and after merging main. It was transient record state onmain(an ACTIVE row whose branch had gone), since repaired.One thing added
#382was never in the roadmap intake table.AGENTS.mdwants the number in three places that must agree — the issue table, the row's spec, the PR body — and only two were true. Linked underKERNEL-ATTN-PAGED, the row that anchorssrc/vt/cuda/cuda_paged_attn.cu.That record commit is mine; the kernel commit is Filip's and is untouched.
Evidence
preflight rc=0 on the merged tree;
check-agent-recordOK. The original commit already carried compliant trailers and passescheck-commit-trailerson its own range.The lever is default-OFF, so no golden moves. Flipping it ON is a separate change that owes the near-tie razor and the distributional gate, as Filip's own commit message states.
🤖 Generated with Claude Code