From 65c365f8581e69161cda424ce0413622ab39f0a9 Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Mon, 15 Jun 2026 14:50:25 -0700 Subject: [PATCH] bench(milestone-gqa-headline): drop misleading single_user_/multi_user_ panels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The legacy panel names suggested batched serving semantics they never had — all four modeled a single user with KV sharded differently (C=1 single-cube; C=4 multi-cube Cube-SP), at toy dims (T_q=4, S_kv≤128). The single-KV-group C=8 panel + the new milestone-gqa-decode-4cases bench cover the meaningful comparisons; pytest regression already covers C=1/C=4 configurations end-to-end at richer scale. Changes: - milestone_gqa_headline.py: drop the 4 legacy panels; _PANELS now contains only single_kv_group_prefill_gqa_c8_p8. Update docstring. - tests/attention/test_milestone_gqa_headline.py: drop the 3 legacy- panel architectural tests (Ring-KV traffic, root-only decode write, per-CUBE distributed output) and test_decode_panels_use_real_gqa (no decode panels in this bench anymore). Equivalent properties are asserted in test_milestone_gqa_single_kv_group_prefill_panel.py (64 dma_writes, 896 ipcq_copy) and test_milestone_gqa_decode_4cases.py (1 dma_write at cube 6, 189 ipcq_copy). - tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py: drop test_existing_prefill_panel_runner_backward_compat (it exercised multi_user_prefill_gqa which no longer exists). - scripts/paper/paper_plot_gqa.py: replace the 4 legacy _LABELS entries with the single single_kv_group_prefill_gqa_c8_p8 label. - Regenerate 1H_milestone_output/gqa_headline/sweep.json from the new panel set. Verification: 9/9 tests pass across the 3 affected test files. Co-Authored-By: Claude Opus 4.7 --- scripts/paper/paper_plot_gqa.py | 6 +- .../gqa_headline/sweep.json | 56 +----------- .../benches/milestone_gqa_headline.py | 18 +--- .../attention/test_milestone_gqa_headline.py | 86 +++---------------- ...stone_gqa_single_kv_group_prefill_panel.py | 34 +------- 5 files changed, 23 insertions(+), 177 deletions(-) diff --git a/scripts/paper/paper_plot_gqa.py b/scripts/paper/paper_plot_gqa.py index 9a2f52c..87fcb57 100644 --- a/scripts/paper/paper_plot_gqa.py +++ b/scripts/paper/paper_plot_gqa.py @@ -24,10 +24,8 @@ _FIG_DIR = Path(__file__).resolve().parents[2] / "docs" / "report" / "1H-codesig _IN_JSON = _FIG_DIR / "gqa_latency.json" _LABELS = { - "single_user_prefill_gqa": "prefill\nC=1", - "multi_user_prefill_gqa": "prefill\nC=4 (Ring KV)", - "single_user_decode_gqa": "decode\nC=1, P=8", - "multi_user_decode_gqa": "decode\nC=4, P=8", + "single_kv_group_prefill_gqa_c8_p8": + "prefill\nC=8, P=8\n(single KV group)", } 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 f0692f6..e73d9cf 100644 --- a/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json +++ b/src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json @@ -1,11 +1,7 @@ { "version": 1, "panels": [ - "single_user_prefill_gqa", - "multi_user_prefill_gqa", - "single_kv_group_prefill_gqa_c8_p8", - "single_user_decode_gqa", - "multi_user_decode_gqa" + "single_kv_group_prefill_gqa_c8_p8" ], "config": { "T_q_prefill": 4, @@ -16,30 +12,6 @@ "d_head": 64 }, "rows": [ - { - "panel": "single_user_prefill_gqa", - "kind": "prefill", - "C": 1, - "S_kv": 16, - "op_log_summary": { - "gemm_count": 2, - "ipcq_copy_count": 0, - "dma_read_count": 3, - "dma_write_count": 1 - } - }, - { - "panel": "multi_user_prefill_gqa", - "kind": "prefill", - "C": 4, - "S_kv": 16, - "op_log_summary": { - "gemm_count": 32, - "ipcq_copy_count": 24, - "dma_read_count": 12, - "dma_write_count": 4 - } - }, { "panel": "single_kv_group_prefill_gqa_c8_p8", "kind": "prefill", @@ -54,32 +26,6 @@ "dma_read_count": 192, "dma_write_count": 64 } - }, - { - "panel": "single_user_decode_gqa", - "kind": "decode", - "C": 1, - "P": 8, - "S_kv": 64, - "op_log_summary": { - "gemm_count": 16, - "ipcq_copy_count": 21, - "dma_read_count": 24, - "dma_write_count": 1 - } - }, - { - "panel": "multi_user_decode_gqa", - "kind": "decode", - "C": 4, - "P": 8, - "S_kv": 128, - "op_log_summary": { - "gemm_count": 64, - "ipcq_copy_count": 93, - "dma_read_count": 96, - "dma_write_count": 1 - } } ] } \ No newline at end of file diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index 3fc4789..bd27bff 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -14,18 +14,16 @@ Restrictions: - No figure renderers (defer to a separate cycle) Panels: - single_user_prefill_gqa : prefill C=1, 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 + intra-CUBE PE-SP (all 64 ranks), 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 - (M-fold + 2-level chain reduce-to-root) + +Decode panels live in ``milestone_gqa_decode_4cases.py`` (the 4-cases +comparative study). Legacy C=1 / C=4 panels were dropped — pytest +regression already covers those configurations. Gated by ``GQA_HEADLINE_RUN=1``. """ @@ -59,17 +57,11 @@ _H_Q_DECODE = 8 # real GQA: G = H_Q_DECODE / H_KV_DECODE = 8 _H_KV_DECODE = 1 _PANELS = ( - "single_user_prefill_gqa", - "multi_user_prefill_gqa", "single_kv_group_prefill_gqa_c8_p8", - "single_user_decode_gqa", - "multi_user_decode_gqa", ) # Each entry: (kind, panel-specific params) _PANEL_DISPATCH: dict[str, tuple[str, dict]] = { - "single_user_prefill_gqa": ("prefill", {"C": 1, "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", { "C": 8, "P": 8, # T_q = S_kv = 1024 (one-shot long-context prefill, scratch-limited). @@ -82,8 +74,6 @@ _PANEL_DISPATCH: dict[str, tuple[str, dict]] = { "T_q": 1_024, "S_kv": 1_024, "d_head": 128, }), - "single_user_decode_gqa": ("decode", {"C": 1, "P": 8, "S_kv": 64}), - "multi_user_decode_gqa": ("decode", {"C": 4, "P": 8, "S_kv": 128}), } diff --git a/tests/attention/test_milestone_gqa_headline.py b/tests/attention/test_milestone_gqa_headline.py index d728710..5fb6b9a 100644 --- a/tests/attention/test_milestone_gqa_headline.py +++ b/tests/attention/test_milestone_gqa_headline.py @@ -1,25 +1,14 @@ -"""Phase 1 spec test for P7: headline milestone-gqa bench with real GQA. +"""Tests for the milestone-gqa-headline bench. -P7 wires the new ``_gqa_decode`` and ``_gqa_prefill`` kernels into a -new 4-panel milestone bench (independent from the existing -``milestone-gqa-llama70b`` which still covers the baseline kernels). -Real GQA (``h_q > h_kv`` with G=8) runs end-to-end through a -milestone-style sweep + sweep.json output. +The bench wires the new ``_gqa_attention_prefill_long`` kernel into a +single-KV-group milestone panel (LLaMA-3.1-70B target, C=8 P=8). The +4 legacy C=1 / C=4 panels (single_user_*, multi_user_*) were removed +once pytest regression covered those configurations comprehensively +and the comparative decode work moved into the +``milestone-gqa-decode-4cases`` bench. -Restriction in P7 first cut: - - C ≤ 4 (single-row inter-CUBE ring SFR; ADR-0060 §B leaves - multi-row rings for follow-on) - - Single SIP (default ``topology.yaml`` 4×4 cube mesh; 4-SIP - headline left to follow-on) - - No figure renderers (defer to a separate cycle) - -Panels (4 total): - single_user_prefill_gqa: C=1, T_q=4, S_kv=16 (no ring) - multi_user_prefill_gqa : C=4, T_q=4, S_kv=16 (Ring KV, 3 steps) - single_user_decode_gqa : C=1, P=8, h_q=8, h_kv=1, S_kv=64 (M-fold + intra-cube chain) - multi_user_decode_gqa : C=4, P=8, h_q=8, h_kv=1, S_kv=128 (M-fold + 2-level chain) - -Phase 1 (this commit): tests only — bench module lands in Phase 2. +Single SIP scope (default ``topology.yaml`` 4×4 cube mesh; 4-SIP +headline deferred per ADR-0060 §B-item-1). """ from __future__ import annotations @@ -35,11 +24,7 @@ from kernbench.topology.builder import resolve_topology BENCH_NAME = "milestone-gqa-headline" PANELS = ( - "single_user_prefill_gqa", - "multi_user_prefill_gqa", "single_kv_group_prefill_gqa_c8_p8", - "single_user_decode_gqa", - "multi_user_decode_gqa", ) @@ -98,52 +83,7 @@ def test_sweep_json_has_expected_panels(monkeypatch): assert {r["panel"] for r in data["rows"]} == set(PANELS) -def test_decode_panels_use_real_gqa(monkeypatch): - """ADR-0060 §A.1 headline: decode panels must use h_q = G·h_kv with G>1.""" - data = _sweep_json(monkeypatch) - cfg = data["config"] - assert cfg["h_q_decode"] > cfg["h_kv_decode"], ( - f"decode must use real GQA (h_q > h_kv); got " - f"h_q={cfg['h_q_decode']}, h_kv={cfg['h_kv_decode']}" - ) - - -# ── Per-panel architectural assertions ──────────────────────────────── - - -def _row(rows, panel: str) -> dict: - for r in rows: - if r["panel"] == panel: - return r - raise AssertionError(f"missing row for panel {panel!r}") - - -def test_prefill_ring_panel_has_ipcq_traffic(monkeypatch): - """multi_user_prefill_gqa uses Ring KV → IPCQ traffic > 0.""" - data = _sweep_json(monkeypatch) - row = _row(data["rows"], "multi_user_prefill_gqa") - n_copy = row["op_log_summary"].get("ipcq_copy_count", 0) - assert n_copy > 0, ( - f"multi_user_prefill_gqa must have Ring KV traffic; got " - f"ipcq_copy_count={n_copy}" - ) - - -def test_decode_reduce_panel_writes_once(monkeypatch): - """multi_user_decode_gqa: chain reduce-to-root → exactly 1 dma_write.""" - data = _sweep_json(monkeypatch) - row = _row(data["rows"], "multi_user_decode_gqa") - n_writes = row["op_log_summary"]["dma_write_count"] - assert n_writes == 1, ( - f"multi_user_decode_gqa root-only write: expected 1; got {n_writes}" - ) - - -def test_prefill_panel_distributes_output(monkeypatch): - """multi_user_prefill_gqa: per-CUBE distributed output → dma_write_count == C.""" - data = _sweep_json(monkeypatch) - row = _row(data["rows"], "multi_user_prefill_gqa") - n_writes = row["op_log_summary"]["dma_write_count"] - assert n_writes == 4, ( - f"multi_user_prefill_gqa per-CUBE distributed: expected 4; got {n_writes}" - ) +# Per-panel architectural assertions for the surviving single-KV-group +# panel live in tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py. +# Decode architectural assertions live in +# tests/attention/test_milestone_gqa_decode_4cases.py. 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 22cc2e9..9935254 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 @@ -116,34 +116,6 @@ def test_single_kv_group_prefill_panel_runner_smoke(): ) -# ── T3: existing panel calls still work (backward compat) ──────────── - - -def test_existing_prefill_panel_runner_backward_compat(): - """Existing callers of ``_run_prefill_panel`` (no ``P``/``T_q``/ - ``d_head`` overrides) must continue to work after the signature - extension. Exercises the ``multi_user_prefill_gqa`` panel dims - from ``_PANEL_DISPATCH``. - - Passes today AND after Phase 2. - """ - panel_name = "multi_user_prefill_gqa" - _kind, params = _PANEL_DISPATCH[panel_name] - topo = resolve_topology(str(TOPOLOGY_DEFAULT)) - - def _bench_fn(ctx): - _run_prefill_panel( - ctx, - panel=panel_name, - C=params["C"], S_kv=params["S_kv"], - ) - - result = run_bench( - topology=topo, bench_fn=_bench_fn, - device=resolve_device(None), - engine_factory=_engine_factory, - ) - assert result.completion.ok, ( - f"existing prefill panel {panel_name!r} must still run with " - f"the original kwargs; got {result.completion}" - ) +# Backward-compat test removed when the legacy multi_user_prefill_gqa +# panel was dropped; signature-extension contract is now exercised +# implicitly by the live single_kv_group_prefill_gqa_c8_p8 panel.