gqa(adr-0065): N4 — opt2 numeric parity vs opt3 in data mode

composite() now awaits its prologue recipe operands (e.g. the score tile from a prior #1 composite) so the cross-composite RAW (#1 writes scores -> #2 reads them) is ordered: #2 waits for #1 before reading. With N1-N3, opt2 (the two-composite recipe kernel) now computes a correct attention output end-to-end in data mode and matches opt3 within fp tolerance. Scope note: the e2e opt2<->opt3 parity holds in the near-uniform-attention regime (small inputs). kernbench's K reshape-as-transpose (opt3 kernel docstring: 'correct for zero/symmetric inputs') makes opt2's per-sub-tile K interpretation differ from opt3's single-tile one for high-contrast inputs — a pre-existing kernbench limitation shared by both kernels, not the recipe. The EXACT recipe math (online-softmax merge m/l + O = O*corr + P@V) is verified non-trivially and exactly by N2/N3 (test_recipe_data_mode.py). Suite 817 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-10 23:13:16 -07:00
parent 5bbe293bea
commit e9d62b908c
2 changed files with 61 additions and 0 deletions
@@ -179,3 +179,57 @@ def test_opt2_bench_completes_oplog_mode():
                                                 enable_data=False),
    )
    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]}"