diff --git a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json index b10be34..f0692f6 100644 --- a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json +++ b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json @@ -3,6 +3,7 @@ "panels": [ "single_user_prefill_gqa", "multi_user_prefill_gqa", + "single_kv_group_prefill_gqa_c8_p8", "single_user_decode_gqa", "multi_user_decode_gqa" ], @@ -39,6 +40,21 @@ "dma_write_count": 4 } }, + { + "panel": "single_kv_group_prefill_gqa_c8_p8", + "kind": "prefill", + "C": 8, + "P": 8, + "T_q": 1024, + "S_kv": 1024, + "d_head": 128, + "op_log_summary": { + "gemm_count": 1024, + "ipcq_copy_count": 896, + "dma_read_count": 192, + "dma_write_count": 64 + } + }, { "panel": "single_user_decode_gqa", "kind": "decode", diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index d8023b3..ae045a2 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -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}), diff --git a/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py b/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py index 0ba67d5..22cc2e9 100644 --- a/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py +++ b/tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py @@ -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"