gqa: fold decode_long into the Case 4 kernel; drop 1D-chain dead code

Option B fold: the lrab math previously in
``_gqa_attention_decode_long.py`` (used only as a thin re-export by
the Case 4 wrapper and as the import source for ``_merge_running``
in Cases 1 / 3) now lives inline in the Case 4 kernel file. ``sub_w``
is hardcoded to 4 (the 4×2 cube sub-mesh geometry; root cube 6);
the legacy ``sub_w=0`` 1D-chain backward-compat path is removed
along with its dedicated tests — the dropped ``single_user_*`` /
``multi_user_*`` panels that exercised it are already gone.

Production changes:
  - inline lrab math into _gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
    (drops sub_w param; _ROOT_CUBE=6 baked in)
  - inline _merge_running into Cases 1 (cube_sp_pe_tp) and 3
    (cube_repl_pe_sp) so they no longer depend on the deleted file
  - delete src/kernbench/benches/_gqa_attention_decode_long.py
  - remove dead _run_decode_panel / _DECODE_* constants /
    decode-side _PANEL_DISPATCH / _make_bench_fn decode branch
    from milestone_gqa_headline.py (only the prefill panel remains)
  - update _gqa_attention_decode_opt2.py docstring reference

Test changes:
  - delete 7 legacy test_gqa_*.py files that pre-dated the 4-cases
    architectural split (coverage now subsumed by the 16 4-cases tests)
  - remove test_opt2_matches_opt3_data_mode + _run_decode_data helper
    from test_gqa_decode_opt2.py (the parity check required sub_w=0
    which no longer exists; opt2's other 4 tests preserved)

20/20 tests pass (16 4-cases + 4 opt2 smoke/dispatch).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 16:32:37 -07:00
parent ddee28a499
commit 756680f4e6
14 changed files with 315 additions and 1810 deletions
-52
View File
@@ -181,55 +181,3 @@ def test_opt2_bench_completes_oplog_mode():
assert result.completion.ok, f"opt2 decode failed: {result.completion}"
# ── opt2 ↔ opt3 numeric parity in data mode (ADR-0065 N4) ─────────────
def _run_decode_data(kernel, name, q_d, k_d, v_d, *, S_kv=16):
"""Run a decode kernel (C=P=1) in data mode with seeded Q/K/V; return the
HBM output array."""
from pathlib import Path
import numpy as np
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.data_executor import DataExecutor
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology
topo = resolve_topology(str(Path(__file__).resolve().parents[2] / "topology.yaml"))
cap: dict = {}
def bench_fn(ctx):
dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=1)
q = ctx.zeros((1, 8 * D), dtype="f16", dp=dp, name=name + "q")
k = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "k")
v = ctx.zeros((S_kv, D), dtype="f16", dp=dp, name=name + "v")
o = ctx.empty((1, 8 * D), dtype="f16", dp=dp, name=name + "o")
q.copy_(ctx.from_numpy(q_d)); k.copy_(ctx.from_numpy(k_d)); v.copy_(ctx.from_numpy(v_d))
cap["o"] = o
ctx.launch(name, kernel, q, k, v, o, 1, S_kv, 8, 1, D, 1, 1, _auto_dim_remap=False)
r = run_bench(topology=topo, bench_fn=bench_fn, device=resolve_device(None),
engine_factory=lambda t, d: GraphEngine(getattr(t, "topology_obj", t),
enable_data=True))
assert r.completion.ok
DataExecutor(r.engine.op_log, r.engine.memory_store).run()
o = cap["o"]
return r.engine.memory_store.read("hbm", o._handle.va_base, shape=(1, 8 * D), dtype="f16")
def test_opt2_matches_opt3_data_mode():
import numpy as np
from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel
from kernbench.benches._gqa_attention_decode_opt2 import gqa_attention_decode_opt2_kernel
rng = np.random.default_rng(0)
q_d = (0.1 * rng.standard_normal((1, 8 * D))).astype(np.float16)
k_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
v_d = (0.1 * rng.standard_normal((16, D))).astype(np.float16)
o3 = _run_decode_data(gqa_attention_decode_long_kernel, "p3", q_d, k_d, v_d)
o2 = _run_decode_data(gqa_attention_decode_opt2_kernel, "p2", q_d, k_d, v_d)
assert np.allclose(np.asarray(o2, np.float32), np.asarray(o3, np.float32),
rtol=5e-2, atol=5e-2), f"opt2 {np.asarray(o2).ravel()[:4]} vs opt3 {np.asarray(o3).ravel()[:4]}"