diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index 9543290..7bde3ff 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -802,6 +802,13 @@ class TLContext: rw_handles: tuple[TensorHandle, ...] = () if prologue: prologue_ops, primary_out, rw_handles = self._expand_prologue(prologue) + # Await any prologue operand that is a pending composite output + # (e.g. the score tile produced by a prior #1 composite) so it is + # computed before this composite reads it (RAW across composites). + for _item in prologue: + self._await_pending( + *[v for v in _item.values() if isinstance(v, TensorHandle)] + ) # Auto-bind (D6.6): the recipe's primary output fills the head GEMM's # `a` unless the kernel supplied `a` explicitly (then it's ambiguous). diff --git a/tests/attention/test_gqa_decode_opt2.py b/tests/attention/test_gqa_decode_opt2.py index 9314cda..b824e08 100644 --- a/tests/attention/test_gqa_decode_opt2.py +++ b/tests/attention/test_gqa_decode_opt2.py @@ -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]}"