perf(cuda): head_dim-128 arm for the warp-split-KV decode kernel, gated OFF

[filipsajdak]

PagedAttentionDecodeOptKernel is reachable only at d == 32 * kDecEpl (256), so head_dim 128 -- the width Qwen3-dense, Llama and Mistral actually use -- never reaches it and falls to the generic block kernel. This adds the EPL=4 instantiation behind VT_ATTN_DECODE_D128, default OFF.

Refs #382.

What changes

The decode-opt 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 == 32 * kDecEpl (256) branch, so at head_dim 128 it can never run; excluding qpk == kDecGqaQG would strand exactly those models (e.g. Qwen3-32B, hq/num_kv_heads == 8) on the block kernel this arm exists to replace.

Why it is gated OFF

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. Default-OFF keeps every existing golden byte-identical.

This mirrors how the FA2 decode GQA group-swap landed -- perf(fa2): decode GQA group-swap port, gated OFF - correctness-complete (#48), flipped ON in #49 against the full gate. The flip here is a separate change and owes the near-tie razor, the distributional gate, and regen under the ratified-tie rule (dense_nvfp4_gemm.h:119-120). That decision is yours -- we cannot run your SACRED gates and are not asking you to take a golden shift on this evidence.

The flag is also the same-binary A/B your perf protocol requires: one binary, VT_ATTN_DECODE_D128=0/1, no rebuild between arms. We cannot run that comparison ourselves under our own operational rules (below), but it exists so you can produce the flip-ON evidence on your rig under your rules.

Motivation is sharper on sm_110 than the general case

cmake/CudaArchFeatures.cmake:349 reads fa2|8.0,8.6,8.7,8.9,12.0a,12.1a -- 11.0 is absent, so VLLM_CPP_FLASH_ATTN is not defined for sm_110 and the whole FA2 block, including the d128 fa2_decode_qwen3 arm, compiles to false. On Thor every d128 decode therefore lands on the generic block kernel; there is no FA2 fallback underneath it. docs/BUILD.md:56-58 describes this as attention falling back to "the portable path, which is slower. Nothing fails and no test goes red."

Whether that fa2 cell should gain 11.0 is a separate question from this PR and we have deliberately not bundled it. We raise it only because it is why this arm matters more on sm_110 than the head_dim gate alone suggests.

Evidence

sm_110 / Jetson AGX Thor, Qwen3-1.7B-NVFP4A16, W4A16 Marlin build (the one #326 enabled).

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 context lengths are bit-for-bit identical. The worst case is 0.5% relative error on small-magnitude outputs at ctx 1000 -- stated in those terms rather than as "6.1e-05", because 0.5% is the magnitude at which a near-tie flips, which is the whole reason your ratified-tie rule exists.

Comparator self-test. Seeding 1-ULP corruption into 3 of the 2048 elements of the ctx=31 row (offsets 5, 1024, 2047 -- interior, midpoint, last element) moves that row from 100.0000% bitwise / max_abs 0 to 99.8535% / max_abs 4.883e-04, and no other row changes. So the comparison is demonstrably able to fail, detects a single-ULP deviation, and localizes correctly. To be precise about what this is: it is a comparator self-test, not a kernel negative control -- it proves the measuring instrument works, not that the kernel comparison could have caught a bad kernel.

Serving A/B, same harness / model / config, c=1,2,4,8, 128 tokens, temp 0, 3 warmups:

	c=1	c=2	c=4	c=8
OFF	81.630	162.738	316.255	477.070
ON	131.023	258.953	508.139	936.048
ratio	1.61x	1.59x	1.61x	1.96x

Build. The changed TU compiles clean on the device: ninja CMakeFiles/vllm.dir/src/vt/cuda/cuda_paged_attn.cu.o, nvcc, sm_110, Release, exit 0, no warnings.

Provenance of the numbers above

The recorded tensors and tok/s come from an on-device experiment, not from this branch's binary. The generator script that produced that experiment has since been found defective -- it emits a self-referential default template argument (int EPL = EPL) -- so the script does not reproduce what was built. The source tree that was actually compiled on the device does not contain that defect; it carries the correct int EPL = kDecEpl. We diffed this branch against that compiled tree: the only functional difference is the VT_ATTN_DECODE_D128 gate itself. With the flag on, this branch is behaviourally the code that produced the measurements.

We are flagging this rather than leaving it for you to find, because "the numbers came from a slightly different tree" is the kind of thing that should come from the submitter.

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. A token-exact run was performed during the experiment but its output was never written to disk, so we will not quote it. Re-running means loading models on a box that serves live production traffic, which we will not do outside a scheduled maintenance window. We would rather submit this with the gap named than with a number we cannot show you the artifact for.
• 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.

Gates

Rebased onto dbd0d51c immediately before submitting. Run with upstream/main as base:

scripts/check-commit-trailers.py --range upstream/main..HEAD   OK: commit trailer contract
scripts/check-pr-size.py --base <upstream/main> --head <HEAD>  OK: every explicit path class is within its review budget
scripts/check-env-doc.py                                       OK: all 312 production env vars documented or classified
scripts/check-doc-checkpoint.py                                OK: public documents match the claims this change makes
scripts/check-agent-record.py                                  OK: ENGINE=149 MODEL=362 QUANT=82 KERNEL=51 BACKEND=80

scripts/agent-preflight.sh was run on this branch and on a pristine upstream/main checkout as a baseline, because the first run's failures pointed at commits that were not ours. On a fork clone the committed-range gates resolve origin/main, which is not your main, so they judge the wrong commits. Baseline fails 7; this branch fails those same 7 plus exactly one -- check-role-discipline, see below. The delta is one gate, and it is a branch-naming question rather than a defect.

VT_ATTN_DECODE_D128 was added to scripts/env-doc-allowlist.txt next to VT_ATTN_DECODE_OPT / VT_ATTN_DECODE_GQA, matching how the sibling kernel switches are classified; check-env-doc.py fails without it.

Two open questions

1. Row claim / branch name. check-role-discipline.py fails because this branch is not row/<ID>: "repository change reached main without arriving on a task branch." Same question as build(cuda): enable marlin-nvfp4 for sm_110 (Jetson Thor) - 8.0x-29.0x faster NVFP4 W4A16 decode #326, which landed without a board row. Tell us how you want an outside contribution to satisfy this and we will redo it -- we did not want to claim an internal roadmap row unasked.
2. The flip. If you want this on by default, say so and we will open it as its own change -- but the near-tie razor, the distributional gate and the golden regen need your hardware and your goldens.

Not done

No test was added. VT_ATTN_DECODE_OPT appears in no test today, so there is no decode-dispatch coverage to extend, and a RED-first test for this arm needs a device gate we cannot run. Happy to write one against whatever harness you would accept.

[filipsajdak]

Closing the gap this PR named in its own Limitations: the model-level token-exact comparison has now been run, as a same-binary A/B. One binary, VT_ATTN_DECODE_D128=0 vs =1, no rebuild between arms -- which is the comparison the perf protocol asks for and which the original two-build experiment could not provide.

sm_110 / Jetson AGX Thor, qwen3-1.7b-nvfp4a16, greedy temperature=0, max_tokens=512, KV pool bounded explicitly at --num-blocks 256.

prompt	OFF sha256	ON sha256	chars
tiny	a5b893b2263e13ba40..	d843fbdea6ace041cf..	2431 / 2410	DIVERGES
just_under_page	a83423cd49941b24e1..	a83423cd49941b24e1..	2560 / 2560	identical
on_page_boundary	a83423cd49941b24e1..	a83423cd49941b24e1..	2560 / 2560	identical
just_over_page	a83423cd49941b24e1..	a83423cd49941b24e1..	2560 / 2560	identical
multi_page	fe6548d5525959898c..	fe6548d5525959898c..	2823 / 2823	identical
long_prefill	68feca229cf73c818a..	68feca229cf73c818a..	2833 / 2833	identical

5 of 6 prompts are byte-identical, including the two longest contexts (multi_page, long_prefill). The prompts were chosen to straddle the 32-token page boundary.

The one divergence, tiny (prompt "Hi."), first differs at character offset 1456 of 2431/2410:

shared prefix ...'to tie this to their own problems? Maybe suggest they ask themselves questions that would '
OFF           'be different from their usual approach. Like, "What if I approached th'
ON            'change their approach. Like, "What if I approached this from a differe'

On a separate 400-token run of the same prompt, reading token ids out of logprobs.tokens, the first differing token is index 316 (387 vs 2297), after 316 identical tokens. Only 1 of the following 83 tokens matches -- complete decorrelation after the flip, which is the expected signature of a single argmax flip under greedy decoding rather than of accumulating numerical error.

What this does NOT show, and why

The near-tie razor could not be run. /v1/completions with {"logprobs": 5} returned logprobs.tokens but top_logprobs: null at every position, on this build, model and configuration. The request was valid -- the cap check at input_processor.cpp:137-150 rejects an over-cap request rather than dropping it, and 5 is far below vocab size.

I can point at where it degrades, though not at why. All 400 returned tokens are the literal string token_id:<N> (token_id:358, token_id:2776, ...), which is emitted at exactly one place: serving_utils.cpp:131, the fallback arm of BuildCompletionLogProbs taken when step == nullptr || step->empty() || step_token == nullptr, which also pushes std::nullopt for both token_logprobs and top_logprobs. So the serving layer behaved exactly as documented -- it is the "position with no logprobs" path, not a bug in that function. And output.logprobs.has_value() must have been true, or serving_completion.cpp:366 would have skipped the payload entirely and the response would carry no logprobs object at all.

That narrows it to: the engine supplied a logprobs container whose per-position entries were empty, or did not contain the sampled token id. I have not determined why, and I am not calling it a defect -- it may be a sampler path or configuration I have not found. If there is a supported way to get top-k logprobs out of this server, I would like to know it, because that is what would let a contributor run your near-tie methodology without needing the oracle.

The consequence for this PR is concrete: the flip at token 316 is consistent with a reduction-order near-tie -- 5/6 identical, the divergence appearing only deep in a long generation, the kernel-level numerics 6/8 contexts bit-identical -- but it has not been demonstrated to be one, and no logprob gap can be quoted. Distinguishing "benign near-tie" from "real numerical error" still needs either top-k logprobs or your oracle.

Limitations, updated from measurement

The body's bullet "No model-level token-exact gate was run against this arm" is now superseded by the above. Two of the other limitations are also narrowed:

• The A/B is no longer cross-build: the arms above are one binary toggled by the env var.
• Still one model, one head_dim, one arch; still no vLLM-oracle comparison, so this is a self-consistency result (new arm vs the block kernel it replaces), not a parity result against vLLM.

Nothing here changes the recommendation that this ships gated OFF. If anything it sharpens the reason: a default flip would move the greedy anchor on at least some prompts, so it owes the near-tie razor and a golden regen under the ratified-tie rule -- on your hardware, under your rules. The flag exists so that comparison is a same-binary A/B when you come to make it.

Raw artifacts (per-prompt SHA256 both arms, full generated text, token id sequences, the run report) were written to disk this time and can be attached in whatever form is useful.

[filipsajdak]

One correction to the provenance of the token-exact numbers above, which I would rather state than have you ask for.

The tree those two arms were built from was not a pristine checkout of main. It is upstream e3cc4f6 -- roughly 150 commits behind current main -- carrying 1,111 lines of local modification across 13 files. Most of that is unrelated experimental work of ours: ~634 lines in src/vllm/tokenizer/ (pretokenizer / unicode_data), plus a local reproduction of the marlin-nvfp4 sm_110 change that has since landed upstream as 9875931f, plus this PR's own diff.

Why the comparison is nevertheless valid: both arms are the same binary, selected by VT_ATTN_DECODE_D128=0 vs =1. Every one of those local modifications, and the stale base, are present identically on both sides and cancel in the diff. What the 5/6 result measures is the EPL=4 arm against the block kernel it replaces, within one build -- and that is exactly the claim being made.

What it does mean: the SHA256 values are not reproducible against a clean main build, and should not be read as absolute reference outputs. They are meaningful only as an equality test between the two arms.

Separately, and mentioned here only because it was found on the same tree: that build also failed every response_format: json_schema request. I did not report it from this tree, because our local tokenizer changes were a live alternative explanation -- a grammar backend keys on the token vocabulary. I rebuilt from a pristine --depth 1 clone of main at 0eb049f7 and the failure is byte-for-byte identical, so our modifications are not the cause. That is now #398, with its own reproduction; it has no bearing on the numbers in this PR.

[localai-bot]

Superseded by #425, which lands this work on current main with your authorship preserved — git log on that branch shows author=Filip Sajdak for the kernel commit, unchanged. Thank you for it; the change is careful and the reasoning in your commit message is why it went in as-is.

Neither CI failure here was yours. This branch sat 46 commits behind, and both reds were stale-base:

cuda-fat-build failed the gencode audit on seven files, every one of them Marlin:

- src/vt/cuda/cuda_marlin_dense.cu: gencode ['110','120a','121a'] != expected ['120a','121a']
- src/vt/cuda/cuda_moe_marlin.cu:   gencode ['110','120a','121a'] != expected ['120a','121a']
  … 5 more marlin files

Not one is cuda_paged_attn.cu, the only source you touch. That is exactly the defect #407 fixed — "Marlin gencode drifted from the feature table" — which landed after your branch was cut.

agent-record failed on audit-live-rows, which reads live branch state at run time rather than anything in the diff. I reproduced it at your tip: audit-live-rows.py --check returns rc=0, both before and after merging main. It was transient record state on main (an ACTIVE row whose branch had been deleted), since repaired.

One thing I added, in a separate commit that is mine, not folded into yours: #382 was never in the roadmap intake table. AGENTS.md wants the issue number in three places that must agree — the issue table, the row's spec, and the PR body — and only two were true. It is now filed under KERNEL-ATTN-PAGED, the row that anchors src/vt/cuda/cuda_paged_attn.cu.

Nothing else changed. The lever stays default-OFF, and flipping it ON remains the separate change your commit message describes — owing the near-tie razor and the distributional gate.

If #425 does not land for any reason, I will reopen this rather than leave the work stranded.