gqa(adr-0065): P2 — softmax_merge recipe + TLContext prologue lowering

New triton_emu/tl_recipes.py: RecipeDescriptor/EngineOp/PrimaryOutSpec + RECIPE_DESCRIPTORS['softmax_merge'] (8-step engine_seq). PE_SCHEDULER does not import it (ADR-0065 D5 boundary). TLContext.composite(): add prologue=[...] + out=TensorHandle kwargs, a optional. _expand_prologue lowers a recipe into flat MATH OpSpecs (scope=KERNEL), allocates TCM scratch, derives the primary-out slot 'P' and auto-binds it into the head GEMM a (D6.6 conflict check); rw_handles=(m,l,O). decode-opt2 #2 lowers to 10 ops [rmax,max_elem,exp_diff,exp_diff,rsum,fma,mul_bcast,copy,gemm,add]. D6.7 (MATH operand TCM-only) scoped to prologue recipe ops only — the head op (gemm or math) keeps existing DMA-staged-from-HBM behavior. D6.1 (GEMM count <=1) on the whole composite. Host-side lowering only; PE_SCHEDULER position-scan is P3. Also commit the ADR-0064 Rev2 promotion content that the prior commit's git mv dropped: Status Proposed->Accepted + D7 amended to hard-cap ValueError (no segmentation), EN+KO. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-10 19:48:29 -07:00
parent 47e2c78c66
commit 184f654295
6 changed files with 508 additions and 85 deletions
@@ -0,0 +1,87 @@
+"""Phase 1 spec tests for ADR-0065 P2 — composite emit-time validation.
+
+Covers the two reachable invariants enforced by the TLContext lowering pass
+(ADR-0065 D6.6 / D6.7):
+
+- **Auto-bind conflict (D6.6).** If a prologue recipe declares a
+  `primary_out` (which auto-binds into the head GEMM's `a`) AND the kernel
+  also passes `a` explicitly, lowering raises — ambiguous which value wins.
+- **MATH operand TCM-only (D6.7).** Every operand of a *prologue recipe*
+  MATH op must be TCM-resident; passing an HBM handle as a recipe operand
+  raises at emit time. (The head op — gemm *or* math — keeps the existing
+  DMA-staged-from-HBM behavior; D6.7 does not apply to it.)
+
+(D6.1 GEMM-count ≤ 1 is implemented as a defensive guard but is not
+reachable through the public API while `softmax_merge` is MATH-only and the
+head carries a single GEMM — its dedicated test is deferred until a
+GEMM-emitting recipe exists.)
+
+Phase 1 (this commit): tests only. FAIL until P2.
+"""
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from kernbench.common.pe_commands import TensorHandle
+from kernbench.triton_emu.tl_context import TLContext
+
+G, TILE, D = 8, 64, 128
+
+
+def _tcm(addr: int, shape: tuple[int, ...]) -> TensorHandle:
+    return TensorHandle(
+        id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
+        nbytes=2 * math.prod(shape), space="tcm",
+    )
+
+
+def _hbm(addr: int, shape: tuple[int, ...]) -> TensorHandle:
+    return TensorHandle(
+        id=f"hb{addr:x}", addr=addr, shape=shape, dtype="f16",
+        nbytes=2 * math.prod(shape), space="hbm",
+    )
+
+
+def test_autobind_conflict_raises(monkeypatch=None):
+    """D6.6 — explicit `a` + a primary_out recipe is ambiguous."""
+    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
+    A = _tcm(0x9000, (G, TILE))
+    s = _tcm(0x1000, (G, TILE))
+    m = _tcm(0x2000, (G,))
+    l = _tcm(0x3000, (G,))
+    O = _tcm(0x4000, (G, D))
+    V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
+    with pytest.raises(ValueError, match="auto-bind"):
+        tl.composite(
+            op="gemm", a=A, b=V, out=O,
+            prologue=[{"op": "softmax_merge", "s": s, "m": m, "l": l, "O": O}],
+        )
+
+
+def test_math_operand_must_be_tcm_raises():
+    """D6.7 — an HBM handle reaching a prologue recipe MATH op is rejected.
+
+    softmax_merge's first op (rmax) reads ``s``; an HBM ``s`` makes the
+    expanded MATH op reference HBM, which the recipe prohibits.
+    """
+    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
+    s_hbm = _hbm(0x1000, (G, TILE))     # Sj as HBM — illegal recipe operand
+    m = _tcm(0x2000, (G,))
+    l = _tcm(0x3000, (G,))
+    O = _tcm(0x4000, (G, D))
+    V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
+    with pytest.raises(ValueError, match="TCM"):
+        tl.composite(
+            op="gemm", b=V, out=O,
+            prologue=[{"op": "softmax_merge", "s": s_hbm, "m": m, "l": l, "O": O}],
+        )
+
+
+def test_math_head_from_hbm_is_allowed():
+    """The head op (op="math") keeps the existing HBM-streamed behavior —
+    D6.7 does not apply to it (only to prologue recipe ops)."""
+    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
+    a_hbm = _hbm(0x1000, (G, TILE))
+    tl.composite(op="math", a=a_hbm, math_op="exp", out_ptr=0x6000)  # no raise