feat(gemma4): decode the global family through a generated split-KV kernel over a folded 640-column pool

[FeathBow]

Description

Closes #1053.

The decode read was moving bytes it never used. Gemma 4's global layers rotate a quarter of the head, so a split K|V row of 1024 columns carries 384 columns of K that only the norm weight touches, and the incumbent decode read at 163K spent 4.5 ms per step where its bytes set a floor of 2.5 ms.

The page format is an axis of the layout. PagedKvLayout carries a KvFormat: Split is the row every existing kernel reads; Folded holds K only where the rotation touches it, V in full, and folds K's norm weight into the query at prep time, 640 columns per token per head. The hd512 prep writes either form; the format is the pool's for its whole life.

Two generated kernels, two formats each. Beside the generated prefill, a split-KV decode kernel (partial over key chunks, merge in two levels) is generated for the global family, both in split-row and folded-row variants; the folded row streams as one tile. PEGAINFER_GLOBAL_ATTN=tilelang serves prefill and decode through the generated kernels on the split pool, tilelang640 on the folded pool; unset or off is the incumbent kernel on the split pool, byte-identical to before. The server reports the global pool's bytes at start-up.

Gated on both formats. The AOT gates run every case in both formats against the JIT twin and an fp32 reference with hostile padding past the plan; the serving oracles hold the generated decode against the incumbent and the folded pool against the split one on the checkpoint.

Test Env

• One GH200, CUDA 13.1 toolchain for the kernels crate, TileLang 0.1.14; Gemma 4 31B bf16 for the serving A/B and the oracles.

Verification

Bound to the tip of this range; every run on the box above.

• AOT gates, both formats: prefill six cases and decode eight cases bit-identical between AOT and JIT, |Δ| ≤ 0.002 against fp32, hostile padding unwritten, refusals refused.
• Serving oracles on the 31B checkpoint, one gate per process, each arm's incumbent asked twice first so the floor is measured rather than chosen: the generated prefill within |Δ logit| 0.5 of the incumbent on the prompt row, the generated decode within 2.6875 over 16 steps, and the folded pool within 2.28125 of the split one over the same steps, argmax equal throughout. Every floor was 0.
• Paired kernel harness at the 31B shape, incumbent / split generated / folded generated: decode read at 10.6K 100.1 / 48.5 / 41.9 µs, at 40K 246.5 / 156.0 / 100.5, at 81.6K 458.9 / 287.2 / 176, at 163K 899.2 / 550.5 / 314.8 (2.86× the incumbent; 1855 GB/s effective on the 640-column row before the two-level merge, higher after).
• Four-arm A/B against vLLM on the same card (single request, four prompt lengths, four rounds, 12 kept requests per cell, round spread ≤ 0.9%; vLLM at the 165,888 ceiling it negotiated with an 8192-token chunk, KV bf16 both sides, prefix caching off): with this range on top of the generated prefill, tilelang640 against vLLM is E2EL +3.0% at 10,602 tokens, −0.5% at 40,002, −5.8% at 81,653 and −12.2% at 163,336; TPOT +2.4 .. +2.7%, +2.7 .. +2.9%, +2.1 .. +2.7% and +1.0 .. +1.6% (from +6.8% and +27.3% before the generated decode); TTFT −13.9 .. −13.2% at 163K. Peak device memory 77.8 GB on the folded pool against 82.7 GB on the split one and 97.1 GB for vLLM at its utilization budget.
• The split state's cells at every length are within the round spread of the arm measured before this range.
• cargo fmt --check, clippy -D warnings on the kernels and gemma4 crates including tests and examples, lib tests green; the one red in the kernels test targets is lt_algo_store, which this range does not touch: the assertion that fails is that test's own precondition that a write was cut short, and the RLIMIT_FSIZE it uses to cut one does not truncate on this filesystem.

[FeathBow]

decode_merge_groups folds a group into its own first slot, so the store lands on rows the walk just read, and every thread wrote the TmpS scalar with nothing between. T.sync_threads() now separates the two phases, outside the lo < hi branch so it is unconditional -- both bounds come from the block's own indices -- and the scalar has a single writer. In the emitted CUDA the two reads, the barrier and the threadIdx.x == 0 store sit in that order, and the TmpV store moved behind the barrier with them; the file differs by eleven lines from the version you reviewed.

A numeric comparison cannot hold this: the race needs the losing schedule, which is why your repro had to inject one. check_in_place_merge_barrier holds it on the emitted source instead, refusing to generate a merge whose store is not separated from the group's reads. It matches every decode_merge_groups kernel by name, so the sliding family reuses it in the next range. Red and green both exercised on the real artifact: it passes as generated and fails with the barrier removed.

Re-run on the fixed kernels: gates green, oracles 0.5 / 2.6875 / 2.28125, unchanged from before the fix. That is not evidence the fix works -- it says the race does not surface in natural execution here, which matches your own note. The barrier in the emitted CUDA and the generation gate are the evidence.

On the decode tolerance you are right that one prompt is not a basis, and measuring it changed the answer rather than the number. Over sixteen cells -- three prompts at 512, 1500 and 3000 tokens, plus a length sweep around 3000 -- the generated decode reaches 7.91 and the folded pool 7.06, so the 4.0 in the body does not hold at 3000. The incumbent's own gap under a difference already accepted as neutral, a walk against a single prefill of the tokens it walked, reaches 5.75 on the same rows. Neither quantity tracks prompt or length: lengths one page apart differ sevenfold (1.13 at 3072 against 7.91 at 3008) and the worst row lands anywhere from the prompt row to the last, so page alignment is not behind it either. A maximum over a walk of these is a draw from a heavy tail, and no constant fitted to it is a correctness bound.

So the gates now separate the two. The argmax is the correctness they carry, equal to the incumbent's on every row of every cell for both arms. The raw-logit gap became a drift line at 12.0, above the measured spread rather than at its edge. What bounds the kernels' own error stays the AOT gate, at 0.002 against an fp32 reference on hostile cases. Both decode gates sweep the nine cells and print each one's gap beside the neutral figure, and the runner's per-gate output window went from twenty lines to forty, since a nine-cell table was being truncated to five with nothing to show it had been.

The body gains the run entry points: PEGAINFER_REQUIRE_GPU=1 cargo test --release -p pegainfer-kernels --features gemma4 for the kernel gates, PEGAINFER_TEST_MODEL_PATH=<31B> PEGAINFER_GEMMA4_FIXTURE_TAG=31b scripts/gemma4_gates.sh the_replacement_global and the same with the_folded_pool for the oracles, and --example hd512_decode_bench for the paired harness.