Description
Gemma 4's ten global layers attend at head_dim 512. The vendored FlashInfer FA2 path serves that shape by time-sharing K and V in one shared-memory buffer, V loaded only after Q·K^T, with a 64-key CTA tile — its own comment says so, to fit head_dim 512 on the 99 KB SKUs (SM86/89/120/121). There is no architecture condition on that choice, so on a Hopper card with 227 KB of shared memory the kernel pays twice for a budget that is not its own.
Measured on one GH200 at the chunked prefill's real shapes (one 8192-row chunk against everything before it, 32 query heads over 4 KV heads, bf16, causal) with the roofline taken in the same process (686 TFLOP/s sustained bf16 GEMM): FlashAttention 4 holds 44–49% of that roofline across the range, the kernel we serve 19–28%, 1.46× to 2.22× slower than FA4 over the same page table. At 31B the global attention is the largest single account of a long prefill, and this one kernel accounts for essentially the whole of the long-prompt gap against vLLM.
Proposed Plan
- A TileLang-generated hd512 varlen paged prefill kernel for sm_90, generated at build time next to the crate that serves it, through the same
PEGAINFER_GEMMA4_TILELANG_PYTHON door the routed lane already uses; the AOT section of the kernels build takes a family table so a second family costs a row, not a branch.
- The global pool paged at the attention key block, so a key tile is one copy.
- A runtime knob,
PEGAINFER_GLOBAL_ATTN=tilelang, selecting the generated kernel for the global family's prompt rows; unset or off keeps the incumbent kernel and byte-identical serving.
- A paired harness timing both kernels on one pool and one plan, and an oracle holding the generated kernel against the incumbent on the checkpoint.
Acceptance Criteria
- The generated kernel is bit-identical to its JIT twin on every AOT case, including hostile padding past the plan, and within 0.002 of an fp32 reference.
- The oracle on the checkpoint holds the generated prefill against the incumbent within its line; the split state's oracle numbers do not move.
- A same-host, same-binary A/B (the knob is the only difference) shows the prefill account moving and nothing else: TPOT and peak memory unchanged.
Description
Gemma 4's ten global layers attend at head_dim 512. The vendored FlashInfer FA2 path serves that shape by time-sharing K and V in one shared-memory buffer, V loaded only after Q·K^T, with a 64-key CTA tile — its own comment says so, to fit head_dim 512 on the 99 KB SKUs (SM86/89/120/121). There is no architecture condition on that choice, so on a Hopper card with 227 KB of shared memory the kernel pays twice for a budget that is not its own.
Measured on one GH200 at the chunked prefill's real shapes (one 8192-row chunk against everything before it, 32 query heads over 4 KV heads, bf16, causal) with the roofline taken in the same process (686 TFLOP/s sustained bf16 GEMM): FlashAttention 4 holds 44–49% of that roofline across the range, the kernel we serve 19–28%, 1.46× to 2.22× slower than FA4 over the same page table. At 31B the global attention is the largest single account of a long prefill, and this one kernel accounts for essentially the whole of the long-prompt gap against vLLM.
Proposed Plan
PEGAINFER_GEMMA4_TILELANG_PYTHONdoor the routed lane already uses; the AOT section of the kernels build takes a family table so a second family costs a row, not a branch.PEGAINFER_GLOBAL_ATTN=tilelang, selecting the generated kernel for the global family's prompt rows; unset oroffkeeps the incumbent kernel and byte-identical serving.Acceptance Criteria