perf(cuda): full tensor-core (mma.sync) flash-attention prefill for d=512 — beat the half2 kernel (#141 follow-up)

[pekkah]

✅ RESOLVED by PR #148 (2026-06-06)

Tensor-core (mma.sync) flash-attention prefill shipped for the Gemma 4 head dims: #146 single-warp (+5%, occupancy-limited at d=512) and the #147 multi-warp / d-split register-resident-O variant (+27–40%, default-on) which beats the half2 kernel — this issue's goal. Parity 3e-4 vs the sequential AttentionBatched oracle. TC-flash test-coverage hardening is tracked in #151; the deeper MMQ-GEMM prefill lever in #152.

---

Follow-up to #141. PR #143 landed a half2 fp16x2 flash-attention prefill kernel (llm_flash_attn_prefill_f32): one-warp-per-query, shared K/V tile reuse, online softmax, half2 QK dot, scalar-fp32 PV. That took prefill 1564 → 2853 t/s and made attention (previously 67% of prefill, BW-bound) far cheaper.

Remaining gap to llama.cpp (~8475 t/s) is now attention compute. The next lever is a full int8/fp16 tensor-core (mma.sync) flash kernel — QK and PV on tensor cores — analogous to llama.cpp's fattn-mma-f16.cuh.

Why it was deferred from the #143 session (the half2 path was chosen instead):

• A study of llama.cpp's fattn-mma-f16.cuh showed that at Gemma 4's global head_dim = 512, the full-TC flash kernel runs occupancy-1 even in llama.cpp — the register/shared pressure of a 512-wide head dim is brutal. So the expected win over half2 was uncertain and the kernel is high-risk.
• half2 captured most of the QK speedup at a fraction of the implementation risk.

Scope (not an estimate, just the pieces):

• mma.sync m16n8k32 s8/s8 (or f16) fragments for QK^T and for P·V, online-softmax rescale between tiles.
• Handle the two Gemma 4 head dims (512 global / 256 SWA) and GQA 4:1.
• d=512 occupancy/register-pressure mitigation (split-D accumulation or 2-pass).
• Argmax-stable parity vs the sequential AttentionBatched oracle (fast path pinned off in bit-exact tests), same pattern as CudaFlashAttnTests / Gemma4_E4B_BatchedPrefill_FlashAttnMatchesSequential.

Gate behind a SHARPI_PREFILL_FLASH_TC-style flag until it beats the half2 kernel on the 4070 Ti.

[pekkah]

Tried the full tensor-core flash (#146). It works, +5% prefill, but is occupancy-limited.

Shipped on branch perf/gemma4-cuda-flash-tc-146 (PR pending):

• Validated the mma.sync m16n8k16 f32 fragment layouts in isolation first (maxAbs 4.77e-7) to de-risk the raw-PTX packing.
• Full TC flash kernel llm_flash_attn_prefill_tc: QK^T and P·V on the mma cores. The QK^T score C-fragment is reused directly as the P·V A-fragment (no transpose — the key index is QK^T's N-col and P·V's contraction, and the m16n8k16 C/A layouts coincide on (row, 2·tig)). O[16×512] is too big for registers (256 regs/lane), so it lives in shared fp32, rescaled in place; K/V time-share one fp16 buffer → fits 48 KB at d=512.
• Parity (CudaFlashAttnTcTests): maxAbs ~3-4e-4 vs the scalar kernels across GQA / 256-SWA & 512-global / windowing / partial tiles — more accurate than the half2 kernel. E2E FlashTcMatchesSequential green.

A/B vs the half2 baseline (all-GPU, warm, gemma-4-E4B Q8_0, 4070 Ti):

N	half2	TC	Δ
965	2585	2732	+5.7%
2054	2716	2849	+4.9%

The win is occupancy-limited: one warp/block + 48 KB shared → only ~2 warps/SM, so the tensor cores can't hide latency. This is why it's only +5% despite both matmuls being on the mma units. Opt-in (SHARPI_PREFILL_FLASH_TC=1) until the multi-warp version lands — see follow-up.