ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
Mukesh Garg
2026-06-15 14:00:30 -07:00
parent 9e1242039b
commit 5b4d9cb597
3 changed files with 38 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py
+11 -8
View File
@@ -4,7 +4,8 @@ The single-KV-group (LLaMA-3.1-70B target) prefill panel wires together
all Increment 1–4 work in a single bench panel:
- C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3)
- P=8 intra-CUBE PE-SP (Inc 4)
- d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 32K
- d_head=128 (LLaMA-3.1-70B), one-shot prefill T_q = S_kv = 1K
(scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling)
This file verifies the bench-config wiring:
T1 the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered
@@ -50,11 +51,13 @@ def test_single_kv_group_prefill_panel_registered():
"""The ``single_kv_group_prefill_gqa_c8_p8`` panel must be in ``_PANELS`` and
its dispatch entry must have the expected LLaMA-3.1-70B params.
Headline config:
Headline config (scratch-limited; LLaMA 32K headline awaits
Q-axis kernel tiling):
C = 8 (head-parallel, snake sub-mesh)
P = 8 (intra-CUBE PE-SP)
T_q = 32_768 (one-shot long-context prefill)
S_kv = 32_768
T_q = 1_024 (one-shot long-context prefill)
S_kv = 1_024
(scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling)
d_head = 128 (LLaMA-3.1-70B)
"""
panel_name = "single_kv_group_prefill_gqa_c8_p8"
@@ -68,11 +71,11 @@ def test_single_kv_group_prefill_panel_registered():
assert kind == "prefill", f"kind={kind!r}, expected 'prefill'"
assert params.get("C") == 8, f"C={params.get('C')}, expected 8"
assert params.get("P") == 8, f"P={params.get('P')}, expected 8"
assert params.get("T_q") == 32_768, (
f"T_q={params.get('T_q')}, expected 32_768"
assert params.get("T_q") == 1_024, (
f"T_q={params.get('T_q')}, expected 1_024"
)
assert params.get("S_kv") == 32_768, (
f"S_kv={params.get('S_kv')}, expected 32_768"
assert params.get("S_kv") == 1_024, (
f"S_kv={params.get('S_kv')}, expected 1_024"
)
assert params.get("d_head") == 128, (
f"d_head={params.get('d_head')}, expected 128"