mukesh
65c365f858
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 <noreply@anthropic.com>
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
"""Tests for the milestone-gqa-headline bench.
|
||
|
||
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.
|
||
|
||
Single SIP scope (default ``topology.yaml`` 4×4 cube mesh; 4-SIP
|
||
headline deferred per ADR-0060 §B-item-1).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
|
||
import kernbench.benches.milestone_gqa_headline as bench_mod # noqa: F401 (Phase 2)
|
||
from kernbench.benches.registry import resolve
|
||
from kernbench.runtime_api.bench_runner import run_bench
|
||
from kernbench.runtime_api.types import resolve_device
|
||
from kernbench.sim_engine.engine import GraphEngine
|
||
from kernbench.topology.builder import resolve_topology
|
||
|
||
BENCH_NAME = "milestone-gqa-headline"
|
||
|
||
PANELS = (
|
||
"single_kv_group_prefill_gqa_c8_p8",
|
||
)
|
||
|
||
|
||
def _run_validation():
|
||
topo = resolve_topology("topology.yaml")
|
||
return run_bench(
|
||
topology=topo,
|
||
bench_fn=resolve(BENCH_NAME).run,
|
||
device=resolve_device(None),
|
||
engine_factory=lambda t, d: GraphEngine(
|
||
getattr(t, "topology_obj", t), enable_data=True,
|
||
),
|
||
)
|
||
|
||
|
||
def _sweep_json(monkeypatch) -> dict:
|
||
monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
|
||
out = bench_mod._OUTPUT_DIR / "sweep.json"
|
||
if not out.exists():
|
||
result = _run_validation()
|
||
assert result.completion.ok, result.completion
|
||
assert out.exists(), f"missing {out}"
|
||
return json.loads(out.read_text())
|
||
|
||
|
||
# ── Registration ───────────────────────────────────────────────────────
|
||
|
||
|
||
def test_bench_registered():
|
||
spec = resolve(BENCH_NAME)
|
||
assert spec.name == BENCH_NAME
|
||
assert callable(spec.run)
|
||
assert spec.description.strip(), "description must be non-empty"
|
||
|
||
|
||
# ── Validation run completes ──────────────────────────────────────────
|
||
|
||
|
||
def test_validation_run_completes_ok(monkeypatch):
|
||
monkeypatch.setenv("GQA_HEADLINE_RUN", "1")
|
||
result = _run_validation()
|
||
assert result.completion.ok, (
|
||
f"headline validation run failed: {result.completion}"
|
||
)
|
||
|
||
|
||
# ── sweep.json shape ──────────────────────────────────────────────────
|
||
|
||
|
||
def test_sweep_json_has_expected_panels(monkeypatch):
|
||
data = _sweep_json(monkeypatch)
|
||
assert set(data["panels"]) == set(PANELS), (
|
||
f"panels mismatch: expected {set(PANELS)}, got {set(data['panels'])}"
|
||
)
|
||
assert len(data["rows"]) == len(PANELS)
|
||
assert {r["panel"] for r in data["rows"]} == set(PANELS)
|
||
|
||
|
||
# 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.
|