perf(gemma4): serve the sliding family through generated kernels and fuse the projections under the generated states

[FeathBow]

Description

Closes #1056

What was left was four accounts, none of them attention at depth. With the global family on the generated kernels, the step was accounted kernel by kernel on both engines with one instrument: the weight-streaming GEMMs equal and at the roofline, and a length-independent constant left in the sliding decode read, the projection launches, the elementwise preps and the sliding prompt read.

The sliding decode reads through a generated windowed split-KV kernel. Under tilelang and tilelang640, a pure-decode row of the sliding family reads its resident window through a generated hd256 partial-and-merge kernel with a window lower bound; mixed steps keep the windowed prefill read. The resident pages hold up to a page more than the window, and both reads mask those keys the same way.

The projections are one weight each, and the split state's bytes do not move. The loader stacks q, k and v rows and gate and up rows in place (fused_rows), no transient copy, no extra device memory. The split state projects through row ranges of the stack, the same cuBLAS shapes and alignments as the separate matrices, so its output is bit-identical, which the oracle's fingerprints check rather than assume. The generated states issue one GEMM per stack; a fused shape may draw a different algorithm and does, so the fusion lives only where the numbers already differ from the incumbent. Columns names a column band of a fused row, the four prep wrappers and the activation take one, and every kernel carries a row stride.

Eight tokens a block through the preps, eight elements a thread through the activation. Thread d still owns element d, so the RMS reduction tree is unchanged and the split state stays bit-identical.

The sliding prompt rows read through a generated windowed prefill. Under the generated states the sliding family's prompt rows attend through a generated hd256 varlen kernel with a window lower bound and a page-64 key tile that is one copy; the sliding pool is paged at 64 for every state. The CTA's owner walk is unrolled into register arrays: the rolled walk cost a third of the kernel.

Test Env

• One GH200, CUDA 12.6 toolchain for the kernels crate, TileLang 0.1.12; Gemma 4 31B bf16 for the serving A/B and the 31B oracles, the pinned 12B checkpoint and fixtures for the maintained gate suite.

Verification

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

• The split state is byte-identical across the range: the serving oracles' incumbent fingerprints (FNV-1a over the bf16 bits of the prefill row and of every one of 17 decode steps) taken on the tree before each change equal the ones after, for the fused projections, the eight-token preps and the page-64 sliding pool alike.

• AOT gates: the windowed decode's four cases (window plus one page with the oldest page masked, ragged lengths, an offset row, short rows) and the windowed prefill's four cases (crossing the window, ragged, continuation, tiny) bit-identical between AOT and JIT, |Δ| ≤ 0.0043 and ≤ 0.011 against fp32, hostile padding unwritten, every refusal row refused.

• Parity on a pool and plan built the serving way: the generated windowed decode within 0.0039 of the incumbent windowed read at residencies of 1040, 1024, 528 and 1 tokens (a single token bit-identical); the generated windowed prefill within 0.0039 of it over three ragged requests and over sixteen (a plan at the slot ceiling), and the generated global prefill within 0.0039 of the incumbent over one request and over sixteen. The fused-band test holds the preps and the activation bit-identical between a fused row and separate buffers and refuses a band past the row.

• Kernel accounts at a 10.6K prompt: the hd256 prep 84 → 34 ms and the gelu-mul 81 → 29 ms, the hd512 prep 20.1 → 8.7 ms; the sliding prompt read over 100 launches 87–94 → 56.3 ms; the fused projections about −0.85 ms per decode step end to end (a bare GEMM bench predicted 1.78; the rest was launch gaps other kernels already fill).

• 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%, every paired interval off zero; vLLM at the 165,888 ceiling it negotiated with an 8192-token chunk, KV bf16 both sides, prefix caching off): tilelang640 against vLLM at 10,602 / 40,002 / 81,653 / 163,336 tokens, E2EL −4.1 / −6.4 / −10.2 / −14.9%, TTFT −3.3 .. −2.0 / −8.7 .. −8.1 / −12.4 .. −12.1 / −16.3 .. −16.1%, TPOT −4.7 .. −4.5 / −4.3 .. −3.7 / −4.6 .. −3.3 / −5.0 .. −3.7% (19.73 / 20.35 / 21.08 / 22.49 ms against 20.69 / 21.22 / 22.00 / 23.63). Twelve of twelve cells lead outside the spread; before this range the TTFT at 10.6K and the TPOT at every length trailed. Peak device memory 77.9 GB against 97.1 for vLLM at its utilization budget.

• The split state's four cells are within the round spread of the arm measured before this range (E2EL 7.18 / 13.84 / 28.82 / 76.84 s against 7.22 / 13.91 / 28.90 / 76.95), so paging the sliding pool at 64 cost it nothing.

• 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 the known environmental lt_algo_store case.