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:
@@ -3,6 +3,7 @@
|
|||||||
"panels": [
|
"panels": [
|
||||||
"single_user_prefill_gqa",
|
"single_user_prefill_gqa",
|
||||||
"multi_user_prefill_gqa",
|
"multi_user_prefill_gqa",
|
||||||
|
"single_kv_group_prefill_gqa_c8_p8",
|
||||||
"single_user_decode_gqa",
|
"single_user_decode_gqa",
|
||||||
"multi_user_decode_gqa"
|
"multi_user_decode_gqa"
|
||||||
],
|
],
|
||||||
@@ -39,6 +40,21 @@
|
|||||||
"dma_write_count": 4
|
"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",
|
"panel": "single_user_decode_gqa",
|
||||||
"kind": "decode",
|
"kind": "decode",
|
||||||
|
|||||||
@@ -18,8 +18,10 @@ Panels:
|
|||||||
multi_user_prefill_gqa : prefill C=4 Ring KV, T_q=4, S_kv=16
|
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 +
|
single_kv_group_prefill_gqa_c8_p8: prefill C=8 snake Ring KV +
|
||||||
intra-CUBE PE-SP (all 64 ranks),
|
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
|
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
|
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)
|
(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
|
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}),
|
"multi_user_prefill_gqa": ("prefill", {"C": 4, "S_kv": _S_KV_PREFILL}),
|
||||||
"single_kv_group_prefill_gqa_c8_p8": ("prefill", {
|
"single_kv_group_prefill_gqa_c8_p8": ("prefill", {
|
||||||
"C": 8, "P": 8,
|
"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,
|
"d_head": 128,
|
||||||
}),
|
}),
|
||||||
"single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}),
|
"single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}),
|
||||||
|
|||||||
@@ -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:
|
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)
|
- C=8 head-parallel + Ring KV via snake-mapped 2×4 sub-mesh (Inc 1, 3)
|
||||||
- P=8 intra-CUBE PE-SP (Inc 4)
|
- 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:
|
This file verifies the bench-config wiring:
|
||||||
T1 the new ``single_kv_group_prefill_gqa_c8_p8`` panel is registered
|
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
|
"""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.
|
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)
|
C = 8 (head-parallel, snake sub-mesh)
|
||||||
P = 8 (intra-CUBE PE-SP)
|
P = 8 (intra-CUBE PE-SP)
|
||||||
T_q = 32_768 (one-shot long-context prefill)
|
T_q = 1_024 (one-shot long-context prefill)
|
||||||
S_kv = 32_768
|
S_kv = 1_024
|
||||||
|
(scratch-limited; LLaMA 32K headline awaits Q-axis kernel tiling)
|
||||||
d_head = 128 (LLaMA-3.1-70B)
|
d_head = 128 (LLaMA-3.1-70B)
|
||||||
"""
|
"""
|
||||||
panel_name = "single_kv_group_prefill_gqa_c8_p8"
|
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 kind == "prefill", f"kind={kind!r}, expected 'prefill'"
|
||||||
assert params.get("C") == 8, f"C={params.get('C')}, expected 8"
|
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("P") == 8, f"P={params.get('P')}, expected 8"
|
||||||
assert params.get("T_q") == 32_768, (
|
assert params.get("T_q") == 1_024, (
|
||||||
f"T_q={params.get('T_q')}, expected 32_768"
|
f"T_q={params.get('T_q')}, expected 1_024"
|
||||||
)
|
)
|
||||||
assert params.get("S_kv") == 32_768, (
|
assert params.get("S_kv") == 1_024, (
|
||||||
f"S_kv={params.get('S_kv')}, expected 32_768"
|
f"S_kv={params.get('S_kv')}, expected 1_024"
|
||||||
)
|
)
|
||||||
assert params.get("d_head") == 128, (
|
assert params.get("d_head") == 128, (
|
||||||
f"d_head={params.get('d_head')}, expected 128"
|
f"d_head={params.get('d_head')}, expected 128"
|
||||||
|
|||||||
Reference in New Issue
Block a user