bench(milestone-gqa-headline): scale single_kv_group prefill panel T_q=S_kv=1K (scratch budget)

The headline T_q=S_kv=32K target overflows the 1 MB per-PE scratch
pool: at T_q_local=4K and tile_s=1024 the scores matrix alone is 8 MB.
The prefill kernel's bootstrap section also leaves K_t/V_t/scores/
exp_scores persistent (outside tl.scratch_scope), inflating baseline.

Scale-down to T_q=S_kv=1K (T_q_local=128, fits comfortably) preserves
the C=8 + P=8 architecture demonstration; the true LLaMA 32K headline
awaits a future increment to add Q-axis tiling and tighten bootstrap
scratch discipline.

Verified end-to-end: kernbench run --bench milestone-gqa-headline now
produces sweep.json with all 5 panels. The new panel shows
ipcq_copy=896 (matches (C-1)·n_tiles·2·C·P = 7·1·2·8·8) and
dma_write=64 (one per PE, head-parallel + intra-CUBE PE-SP).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 14:00:30 -07:00
parent 9e1242039b
commit 5b4d9cb597
3 changed files with 38 additions and 10 deletions
@@ -18,8 +18,10 @@ Panels:
multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16
single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
intra-CUBE PE-SP (all 64 ranks),
T_q=S_kv=32K, d_head=128 — the
T_q=S_kv=1K, d_head=128 — the
LLaMA-3.1-70B single-KV-group target
(scratch-limited; 32K headline awaits
Q-axis kernel tiling)
single_user_decode_gqa : decode C=1, P=8, h_q=8, h_kv=1, S_kv=64
(M-fold + intra-cube row-then-col chain)
multi_user_decode_gqa : decode C=4, P=8, h_q=8, h_kv=1, S_kv=128
@@ -70,7 +72,14 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = {
"multi_user_prefill_gqa": ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
"single_kv_group_prefill_gqa_c8_p8": ("prefill", {
"C": 8, "P": 8,
"T_q": 32_768, "S_kv": 32_768,
# T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited).
# The bootstrap section of the prefill kernel leaves K_t/V_t/scores/
# exp_scores persistent (outside tl.scratch_scope), inflating the
# baseline. At T_q=2048 the peak hits ~1.05 MB vs 1.0 MB budget; at
# 1024 it fits comfortably. The true LLaMA 32K headline awaits a
# future increment to (a) add Q-axis tiling and (b) move the
# bootstrap into scratch discipline.
"T_q": 1_024, "S_kv": 1_024,
"d_head": 128,
}),
"single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}),