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>

docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md

@@ -2,7 +2,7 @@
 
 ## Status
 
-Proposed (Revision 2)
+Accepted (Revision 2)
 
 > **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops
 > composite + 첫 stateful recipe) 의 보조 ADR. 그 hybrid 의 핵심 — "GEMM 은
@@ -219,20 +219,14 @@ issue cost 를 0 → non-zero 로 바꾸면 **모든** bench 의 latency 변화.
 이 ADR land 시 골든 latency **한 번** 재생성 — ADR-0062 D3 lazy-load 와
 같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용.
 
-### D7. Composite 크기 cap (deterministic segmentation)
+### D7. Composite 크기 cap (하드 제한 — validation error)
 
-각 `CompositeCmd` 의 `logical_bytes` 가 **`MAX_COMPOSITE_LOGICAL_BYTES`**
+각 `CompositeCmd` 의 `logical_bytes` 가 **`max_composite_logical_bytes`**
-(default **1024 bytes**, D4 로 override) 로 제한. 초과하는 cmd 는
+(default **1024 bytes**, D4 로 override) 로 제한. `logical_bytes` 가 cap
-*emitter* (host-side TLContext, ADR-0065 D5) 가 N 개 연속 `CompositeCmd`
+을 초과하는 composite 는 host-side TLContext 가 emit 시점에 `ValueError`
-로 deterministic 하게 **segment** — 각 ≤ cap.
+로 **거부** — **자동 segmentation 없음**. 커널 작성자가 각 `CompositeCmd`
-
+가 descriptor capacity 에 맞도록 recipe 를 재구성(예: 여러 개의 더 작은
-- 각 segment 가 자기 `completion: CompletionHandle` 보유.
+composite 로 명시적으로 분할)해야 함.
-- 각 segment 가 자기 dispatch cost — `total = sum(FIXED + bytes_i × R) =
-  N × FIXED + total_bytes × R`. FIXED 항이 segment 당.
-- segment 가 `rw_handles` 공유 — **strict FIFO** (ADR-0065 D6.3) 가
-  write-after-write 순서 자동 보존.
-- segmentation 알고리즘 결정적 (emit 순서로 op 인덱스 greedy); host
-  emitter 의 일부, PE_SCHEDULER 가 아님.
 
 **cap 이 필요한 이유.** 실제 하드웨어는 descriptor queue entry 크기,
 scheduler parser buffer, command SRAM, firmware 입력의 하드 제한 보유.
@@ -248,18 +242,15 @@ composite (가장 큰 것이 decode opt2 의 `#2` ~322 bytes) 보다 훨씬 위
 이 default 의 역할은 cap 을 *원칙으로 존재시키는 것*, 특정 HW 에 맞추는
 것이 아님.
 
-**segment 간 순서 — strict FIFO 가 주체, `rw_handles` 아님.** segment 들은
+**자동 segmentation 대신 하드 에러를 택한 근거.** 초과 composite 를 자동
-emit 순서로 PE_CPU → PE_SCHEDULER 큐에 들어감. strict-FIFO dispatch
+분할(원래 Revision 2 제안)하는 것은 inter-segment 순서, `rw_handles`
-(ADR-0065 D6.3) 가 *순서의 주체* — segment 가 큐 진입 순서로 실행. 각
+공유, completion 체이닝 등 emitter 복잡도를 추가하면서, 실질적으로는
-segment 의 `rw_handles` block 은 cross-composite 해저드 추적기의 **의존성
+하드웨어가 허용하는 것보다 큰 descriptor 를 요청한 커널을 덮어주는
-메타데이터** — 그 자체로 inter-segment 순서를 보장하지 **않음**. 미래
+것일 뿐. 명시적 에러로 표면화하면 emitter 가 단순해지고, HW 제약이
-scheduler 가 FIFO 를 완화 (예: RW-aware reorder, ADR-0065 A4) 한다면
+작업을 어떻게 분할할지 가장 잘 아는 커널 작성자에게 보임.
-segmenter 는 (a) segment 를 단일 CompositeCmd 로 merge 또는
-(b) 명시적 completion-handle 의존 체인 도입 필요. Phase 1 에선 무관 —
-strict FIFO 가 유효.
 
 Decode opt2 의 `#2` composite (10 ops, ~322 bytes) 는 1024 cap 안에
-편안히 — GQA workload 에 segmentation 없음.
+편안히 — GQA workload 에 에러 없음.
 
 ## Alternatives
 
@@ -350,10 +341,8 @@ calibration 이 바뀌거나 OpSpec/CompositeCmd 필드가 추가되어도 유
 7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일
    dispatch-cycle 회계.
 8. **Composite 크기 cap (D7).** `logical_bytes >
-   MAX_COMPOSITE_LOGICAL_BYTES` 가 될 recipe 가 N 개 연속 `CompositeCmd`
+   max_composite_logical_bytes` 가 될 composite 는 emit 시점에 `ValueError`
-   로 segment; `sum(segment.logical_bytes) == original logical_bytes`;
+   (segmentation 없음). cap 이내 composite 는 정상 emit.
-   총 dispatch = segment dispatch 의 합 (FIXED 가 segment 당 지불);
-   inter-segment 순서 strict FIFO 로 보존 (`rw_handles` 단독 아님).
 9. **민감도 (정성적).** `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte 에서
    3 지점 모두 `opt3 > opt2`. 방향 (R 감소 시 ratio 단조 증가) 도 단언,
    단 절대 ratio 값은 *불필요*.
@@ -366,11 +355,11 @@ ADR-0064 Revision 2 는 단일 PR 로 land:
 - `pe_cpu.py` dispatch 경로의 공식 적용
 - PE_CPU init 의 `pe_cost_model:` override read (cycle-domain knob +
   `max_composite_logical_bytes`)
-- **composite 크기 cap (D7)** — TLContext-side segmentation 로직,
+- **composite 크기 cap (D7)** — TLContext-side 하드 cap,
-  `MAX_COMPOSITE_LOGICAL_BYTES` default 1024; 기존 bench 에 필요 없음
+  `max_composite_logical_bytes` default 1024: cap 초과 composite 는 emit
-  (현재 최대 composite 가 200 bytes 훨씬 아래), 하지만 ADR-0065 의 10-op
+  시점에 `ValueError` (자동 segmentation 없음). 기존 bench 는 트리거
-  decode-opt2 composite (~310 bytes) + 더 큰 recipe 가 나타날 때를 위해
+  안 됨 (현재 최대 composite 가 ~322 bytes), 하지만 recipe 가 커질 때
-  메커니즘 land.
+  descriptor-capacity 제한이 강제되도록 체크 land.
 - 일회성 골든 재생성
 
 land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음

docs/adr/ADR-0064-perf-cpu-issue-cost-model.md

@@ -2,7 +2,7 @@
 
 ## Status
 
-Proposed (Revision 2)
+Accepted (Revision 2)
 
 > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065**
 > (flat-ops composite + first stateful recipe). The hybrid decision there
@@ -234,22 +234,15 @@ latencies are **regenerated once** when this ADR lands — same posture as
 ADR-0062 D3 lazy-load. After regeneration, the same calibration is in
 effect for ADR-0065 opt2 measurement.
 
-### D7. Composite size cap (deterministic segmentation)
+### D7. Composite size cap (hard limit — validation error)
 
 Each `CompositeCmd`'s `logical_bytes` is capped at
-**`MAX_COMPOSITE_LOGICAL_BYTES`** (default **1024 bytes**, overridable
+**`max_composite_logical_bytes`** (default **1024 bytes**, overridable
-per D4). Oversized commands are deterministically **segmented** by the
+per D4). A composite whose `logical_bytes` exceeds the cap is **rejected
-*emitter* (host-side TLContext, ADR-0065 D5) into N consecutive
+at emit time** by the host-side TLContext with a `ValueError` — there is
-`CompositeCmd`s, each ≤ cap.
+**no automatic segmentation**. The kernel author must restructure the
-
+recipe (e.g., split it into multiple smaller composites explicitly) so
-- Each segment carries its own `completion: CompletionHandle`.
+each `CompositeCmd` fits within the descriptor capacity.
-- Each segment incurs its own dispatch cost — `total =
-  sum(FIXED + bytes_i × R) = N × FIXED + total_bytes × R`. The FIXED
-  term is paid per segment.
-- Segments share `rw_handles` where applicable; **strict FIFO**
-  (ADR-0065 D6.3) preserves write-after-write ordering automatically.
-- The segmentation algorithm is deterministic (greedy by op index in
-  emit order); it is part of the host emitter, not PE_SCHEDULER.
 
 **Why a cap is needed.** Real hardware imposes hard limits — descriptor
 queue entry size, scheduler parser buffer, command SRAM, firmware
@@ -266,20 +259,16 @@ overridable per topology (D4); when a real HW reference appears, the
 value should be recalibrated. The role of this default is to make the
 cap *exist as a discipline*, not to fit a specific HW.
 
-**Ordering of segments — driven by strict FIFO, not `rw_handles`.**
+**Rationale for a hard error over auto-segmentation.** Auto-splitting an
-Segments are emitted into the PE_CPU → PE_SCHEDULER queue in their
+oversized composite (the original Revision 2 proposal) added emitter
-emit order. Strict-FIFO dispatch (ADR-0065 D6.3) is the *ordering
+complexity — inter-segment ordering, shared `rw_handles`, completion
-source*: segments execute in the order they enter the queue. The
+chaining — to paper over what is, in practice, a kernel that asked for a
-`rw_handles` block on each segment is **dependency metadata** for the
+descriptor larger than the hardware allows. Surfacing it as an explicit
-cross-composite hazard tracker — it does **not** by itself guarantee
+error keeps the emitter simple and makes the HW constraint visible to the
-inter-segment ordering. If a future scheduler relaxed FIFO (e.g., to
+kernel author, who is best placed to decide how to split the work.
-RW-aware reorder, ADR-0065 A4), the segmenter would need to either
-(a) merge segments into a single CompositeCmd, or (b) introduce an
-explicit completion-handle dependency chain. For Phase 1 this does
-not arise: strict FIFO is in effect.
 
 Decode opt2's `#2` composite (10 ops, ~322 bytes) sits comfortably
-inside the 1024 cap — no segmentation for the GQA workload.
+inside the 1024 cap — no error for the GQA workload.
 
 ## Alternatives
 
@@ -379,12 +368,9 @@ OpSpec/CompositeCmd fields are added.
    (SPEC §0.1).
 7. **Path parity.** Greenlet and replay paths produce identical
    dispatch-cycle accounting for the same kernel.
-8. **Composite size cap (D7).** A recipe that would emit `logical_bytes
+8. **Composite size cap (D7).** A composite that would emit `logical_bytes
-   > MAX_COMPOSITE_LOGICAL_BYTES` is segmented into N consecutive
+   > max_composite_logical_bytes` raises a `ValueError` at emit time (no
-   `CompositeCmd`s; `sum(segment.logical_bytes) == original
+   segmentation). A composite within the cap emits normally.
-   logical_bytes`; total dispatch = sum of segment dispatches (FIXED
-   paid per segment); inter-segment ordering preserved by strict FIFO
-   (not by `rw_handles` alone).
 9. **Sensitivity (qualitative).** At `R ∈ {0.25, 0.0625, 0.03125}`
    cycles/byte, `opt3 > opt2` at *all* three points. Direction (ratio
    monotonically increases as R decreases) is also asserted, but
@@ -398,11 +384,12 @@ ADR-0064 Revision 2 lands as a single PR with:
 - formula application in `pe_cpu.py` dispatch path
 - `pe_cost_model:` override read at PE_CPU init (cycle-domain knobs +
   `max_composite_logical_bytes`)
-- **composite size cap (D7)** — TLContext-side segmentation logic with
+- **composite size cap (D7)** — TLContext-side hard cap with
-  `MAX_COMPOSITE_LOGICAL_BYTES` default 1024; not needed for any
+  `max_composite_logical_bytes` default 1024: a composite over the cap
-  existing bench (largest current composite is well under 200 bytes),
+  raises a `ValueError` at emit time (no auto-segmentation). Not
-  but the mechanism lands so it is in place when ADR-0065's
+  triggered by any existing bench (largest current composite is ~322
-  10-op decode-opt2 composite (~310 bytes) and larger recipes appear.
+  bytes), but the check lands so the descriptor-capacity limit is
+  enforced as recipes grow.
 - one-time goldens regeneration
 
 After this lands, ADR-0065 builds on top with no further goldens churn

src/kernbench/triton_emu/tl_context.py

@@ -735,28 +735,58 @@ class TLContext:
     def composite(
         self,
         op: Literal["gemm", "math"],
-        a: TensorHandle,
+        a: TensorHandle | None = None,
         b: TensorHandle | None = None,
         out_ptr: int = 0,
         math_op: str | None = None,
         *,
+        out: TensorHandle | None = None,
+        prologue: list[dict] | None = None,
         epilogue: list[dict] | None = None,
         acc_dtype: str | None = None,
         tile_shape: tuple[int, int, int] | None = None,
     ) -> CompletionHandle:
         """Submit a composite command (non-blocking, tiled pipeline).
 
-        Optional ``epilogue`` is an ordered list of dicts; each dict has a
+        ``prologue`` (ADR-0065 D5) is an ordered list of recipe dicts — each
-        required ``"op"`` key (one of ``EPILOGUE_OPS``) plus op-specific
+        with an ``"op"`` naming a ``RECIPE_DESCRIPTORS`` entry plus its
-        fields and an optional ``"scope"``. Validation happens here so
+        operands. TLContext expands each into flat MATH OpSpecs that run
-        typos fail before the command is emitted.
+        before the head op and auto-binds the recipe's primary output into
+        the head GEMM's ``a`` operand (D6.6). ``out`` (TensorHandle) is the
+        explicit write-back handle, preferred over ``out_ptr``. ``epilogue``
+        is an ordered list of ``EPILOGUE_OPS`` dicts.
 
         Returns CompletionHandle for use with wait().
         """
         # ADR-0062: composite operand DMA paths still need their inputs
         # to be resolved before the composite reads them via PE_SCHEDULER.
         self._await_pending(a, b)
-        # Compute output geometry based on op.
+
+        # Prologue recipes → flat MATH OpSpecs + primary-output handle (D5).
+        prologue_ops: tuple[OpSpec, ...] = ()
+        primary_out: TensorHandle | None = None
+        rw_handles: tuple[TensorHandle, ...] = ()
+        if prologue:
+            prologue_ops, primary_out, rw_handles = self._expand_prologue(prologue)
+
+        # 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).
+        if op == "gemm" and primary_out is not None:
+            if a is not None:
+                raise ValueError(
+                    "auto-bind conflict: head GEMM operand 'a' was provided "
+                    "explicitly but the prologue recipe declares a primary_out "
+                    "(ADR-0065 D6.6) — omit 'a' to let the recipe bind it"
+                )
+            a = primary_out
+
+        if a is None:
+            raise ValueError(
+                "composite requires operand 'a' (or a prologue recipe with a "
+                "primary_out)"
+            )
+
+        # Output geometry.
         if op == "gemm" and b is not None:
             m, k = a.shape[-2], a.shape[-1]
             n = b.shape[-1]
@@ -768,7 +798,10 @@ class TLContext:
             out_shape = a.shape
             out_nbytes = a.nbytes
 
-        # Head op's write-back handle carries the output address (ADR-0065 D1).
+        # Write-back handle: explicit `out` handle wins over `out_ptr`.
+        if out is not None:
+            out_handle = out
+        else:
             out_handle = TensorHandle(
                 id=self._next_handle_id(),
                 addr=out_ptr, shape=out_shape, dtype=out_dtype,
@@ -795,12 +828,144 @@ class TLContext:
         )
         epi_specs = tuple(self._build_epilogue_spec(e, i)
                           for i, e in enumerate(epilogue or []))
-        ops_tuple = (head_spec, *epi_specs)
+        ops_tuple = (*prologue_ops, head_spec, *epi_specs)
+
+        self._validate_composite_ops(ops_tuple, prologue_ops)
 
         completion = CompletionHandle(id=self._next_completion_id())
-        self._emit(CompositeCmd(completion=completion, ops=ops_tuple))
+        self._emit(CompositeCmd(
+            completion=completion, ops=ops_tuple, rw_handles=rw_handles,
+        ))
         return completion
 
+    def _expand_prologue(
+        self, prologue: list[dict],
+    ) -> tuple[tuple[OpSpec, ...], TensorHandle | None, tuple[TensorHandle, ...]]:
+        """Expand prologue recipe items into flat MATH OpSpecs (ADR-0065 D5).
+
+        Returns ``(math_ops, primary_out_handle, rw_handles)``. PE_SCHEDULER
+        never sees the recipe — only the flat ops.
+        """
+        from kernbench.triton_emu.tl_recipes import (
+            PRIMARY_OUT_SLOT, RECIPE_DESCRIPTORS,
+        )
+
+        math_ops: list[OpSpec] = []
+        primary_out: TensorHandle | None = None
+        rw_handles: list[TensorHandle] = []
+
+        for item in prologue:
+            if not isinstance(item, dict) or "op" not in item:
+                raise ValueError("prologue entry must be a dict with an 'op' key")
+            name = item["op"]
+            recipe = RECIPE_DESCRIPTORS.get(name)
+            if recipe is None:
+                known = ", ".join(sorted(RECIPE_DESCRIPTORS))
+                raise ValueError(
+                    f"unknown prologue recipe {name!r} (known: {known})"
+                )
+
+            # Bind recipe operands from the item.
+            slots: dict[str, TensorHandle] = {}
+            for opd_name in recipe.operands:
+                if opd_name not in item:
+                    raise ValueError(
+                        f"prologue recipe {name!r} missing operand {opd_name!r}"
+                    )
+                slots[opd_name] = item[opd_name]
+
+            # Expand engine_seq sequentially — later ops read earlier dsts.
+            for eop in recipe.engine_seq:
+                operands: dict[str, Any] = {}
+                for key in ("src", "src_a", "src_b", "src_c"):
+                    slot_name = getattr(eop, key)
+                    if slot_name is None:
+                        continue
+                    if slot_name not in slots:
+                        raise ValueError(
+                            f"recipe {name!r} op {eop.op_kind!r} reads unbound "
+                            f"slot {slot_name!r}"
+                        )
+                    operands[key] = slots[slot_name]
+
+                dst_handle = self._resolve_recipe_dst(eop, recipe, slots, operands)
+                extra: dict[str, Any] = {}
+                if eop.reduce_axis is not None:
+                    extra["reduce_axis"] = eop.reduce_axis
+                if eop.bcast_axis is not None:
+                    extra["bcast_axis"] = eop.bcast_axis
+
+                math_ops.append(OpSpec(
+                    kind=eop.op_kind, scope=Scope.KERNEL,
+                    operands=operands, extra=extra, out=dst_handle,
+                ))
+                slots[eop.dst] = dst_handle
+
+            if recipe.primary_out is not None:
+                primary_out = slots.get(PRIMARY_OUT_SLOT)
+            for opd_name, rw in recipe.operands.items():
+                if rw == "RW":
+                    rw_handles.append(slots[opd_name])
+
+        return tuple(math_ops), primary_out, tuple(rw_handles)
+
+    def _resolve_recipe_dst(
+        self, eop: Any, recipe: Any, slots: dict, operands: dict,
+    ) -> TensorHandle:
+        """Resolve an EngineOp's dst to a handle: reuse an existing slot
+        (RW operand / already-allocated), else allocate fresh TCM scratch
+        with an inferred shape."""
+        from kernbench.triton_emu.tl_recipes import PRIMARY_OUT_SLOT
+
+        dst_name = eop.dst
+        if dst_name in slots:
+            return slots[dst_name]            # write in place
+
+        if recipe.primary_out is not None and dst_name == PRIMARY_OUT_SLOT:
+            ref = slots[recipe.primary_out.from_shape]
+            shape, dtype = ref.shape, ref.dtype
+        elif eop.reduce_axis is not None:
+            ref = next(iter(operands.values()))
+            shape = ref.shape[:-1] if len(ref.shape) > 1 else ref.shape
+            dtype = ref.dtype
+        else:
+            ref = next(iter(operands.values()))
+            shape, dtype = ref.shape, ref.dtype
+
+        nbytes = self._nbytes(shape, dtype)
+        addr = self._scratch_alloc(nbytes)
+        return TensorHandle(
+            id=self._next_handle_id(),
+            addr=addr, shape=shape, dtype=dtype, nbytes=nbytes, space="tcm",
+        )
+
+    @staticmethod
+    def _validate_composite_ops(
+        ops: tuple[OpSpec, ...], prologue_ops: tuple[OpSpec, ...],
+    ) -> None:
+        """Emit-time invariants (ADR-0065 D6.1 / D6.7).
+
+        D6.1 (GEMM count ≤ 1) applies to the whole composite. D6.7 (MATH
+        operands TCM-only) applies to **prologue recipe ops only** — the
+        head op (gemm *or* math) and epilogue ops keep the existing
+        DMA-staged-from-HBM behavior (a math head's `a` is streamed by
+        PE_SCHEDULER, like a GEMM operand).
+        """
+        gemm_count = sum(1 for o in ops if o.kind == "gemm")
+        if gemm_count > 1:
+            raise ValueError(
+                f"composite must carry at most 1 GEMM op (ADR-0065 D6.1); "
+                f"got {gemm_count}"
+            )
+        for o in prologue_ops:
+            for h in o.operands.values():
+                if getattr(h, "space", "tcm") != "tcm":
+                    raise ValueError(
+                        f"MATH operand must be TCM-resident (ADR-0065 D6.7); "
+                        f"prologue op {o.kind!r} operand has space="
+                        f"{getattr(h, 'space', None)!r}"
+                    )
+
     @staticmethod
     def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec:
         if not isinstance(entry, dict) or "op" not in entry:

src/kernbench/triton_emu/tl_recipes.py

@@ -0,0 +1,84 @@
+"""TLContext recipe descriptors (ADR-0065 D5).
+
+A *recipe* is a named macro op (e.g. ``softmax_merge``) that TLContext
+expands, at composite-emit time, into a sequence of flat ``OpSpec``s with
+all addresses / sizes filled. The recipe table lives here, beside
+``tl_context.py`` (the compiler analog). **PE_SCHEDULER never imports this
+module** — it sees only the flat ops list (ADR-0065 D5 / D7 boundary).
+
+Slot names used in ``EngineOp`` are either recipe operand names
+(``RecipeDescriptor.operands``) or internal scratch slots produced by an
+earlier op in ``engine_seq``. The scratch slot named ``"P"`` is, by
+convention, the recipe's *primary output* — auto-bound into the head GEMM's
+first matrix operand by the lowering pass.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Literal
+
+PRIMARY_OUT_SLOT = "P"
+
+
+@dataclass(frozen=True)
+class PrimaryOutSpec:
+    """Spec for a recipe's implicit primary output (slot ``"P"``)."""
+
+    from_shape: str          # operand name to copy shape from
+    from_dtype: str          # operand name to copy dtype from
+    transform: str           # "identity" | "trans" (forward-compat)
+
+
+@dataclass(frozen=True)
+class EngineOp:
+    """One micro-op in a recipe's ``engine_seq``.
+
+    Slot fields (``src``, ``src_a``, ``src_b``, ``src_c``, ``dst``) name
+    either ``RecipeDescriptor.operands`` or internal scratch slots.
+    """
+
+    engine: Literal["MATH", "GEMM", "DMA"]
+    op_kind: str             # PE_MATH op_kind label (opaque to the engine)
+    src: str | None = None
+    src_a: str | None = None
+    src_b: str | None = None
+    src_c: str | None = None
+    dst: str | None = None
+    reduce_axis: int | None = None
+    bcast_axis: int | None = None
+
+
+@dataclass(frozen=True)
+class RecipeDescriptor:
+    operands: dict[str, str]                            # name → "R" | "RW"
+    primary_out: PrimaryOutSpec | None                  # implicit output (slot "P")
+    tile_alignment: Literal["single_shot", "tile_aligned"]
+    internal_scratch_bytes_fn: Callable[..., int]
+    engine_seq: tuple[EngineOp, ...]
+
+
+RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = {
+    # Online-softmax merge (flash-attention accumulator update). Reads the
+    # current score tile ``s`` and the running ``(m, l, O)``; emits the
+    # rescaled probabilities ``P`` (slot "P") for the downstream P·V GEMM.
+    "softmax_merge": RecipeDescriptor(
+        operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
+        primary_out=PrimaryOutSpec(
+            from_shape="s", from_dtype="s", transform="identity",
+        ),
+        tile_alignment="single_shot",
+        internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
+            G + G + G + G * TILE + G  # m_loc + m_new + corr + P + l_loc
+        ),
+        engine_seq=(
+            EngineOp("MATH", "rmax",      src="s",                   dst="m_loc", reduce_axis=-1),
+            EngineOp("MATH", "max_elem",  src_a="m",    src_b="m_loc", dst="m_new"),
+            EngineOp("MATH", "exp_diff",  src_a="m",    src_b="m_new", dst="corr"),
+            EngineOp("MATH", "exp_diff",  src_a="s",    src_b="m_new", dst="P", bcast_axis=0),
+            EngineOp("MATH", "rsum",      src="P",                   dst="l_loc", reduce_axis=-1),
+            EngineOp("MATH", "fma",       src_a="l",    src_b="corr", src_c="l_loc", dst="l"),
+            EngineOp("MATH", "mul_bcast", src_a="O",    src_b="corr", dst="O", bcast_axis=1),
+            EngineOp("MATH", "copy",      src="m_new",               dst="m"),
+        ),
+    ),
+}

tests/test_tl_composite_validation.py

@@ -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

tests/test_tl_recipe_softmax_merge.py

@@ -0,0 +1,111 @@
+"""Phase 1 spec tests for ADR-0065 P2 — `softmax_merge` recipe lowering.
+
+P2 adds `tl_recipes.py` (RECIPE_DESCRIPTORS) and a TLContext lowering pass
+so that
+
+    tl.composite(
+        op="gemm", b=V_ref, out=O,
+        prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
+        epilogue=[{"op": "add", "other": O}],
+    )
+
+expands the single `softmax_merge` recipe into 8 flat MATH OpSpecs, binds the
+recipe's primary output `P` into the head GEMM's `a` operand, and emits one
+`CompositeCmd` with 10 ops + `rw_handles == (m, l, O)` (ADR-0065 D5 / DDD §6).
+
+This is host-side lowering only — PE_SCHEDULER consumption of the pre-GEMM
+KERNEL-scope ops is P3, so these tests inspect the lowered CompositeCmd
+object and do NOT run it through the scheduler.
+
+Phase 1 (this commit): tests only. FAIL until P2:
+  - `kernbench.triton_emu.tl_recipes` does not exist,
+  - `TLContext.composite` has no `prologue=` / `out=TensorHandle` kwargs.
+"""
+from __future__ import annotations
+
+import math
+
+from kernbench.common.pe_commands import CompositeCmd, TensorHandle
+from kernbench.triton_emu.tl_context import TLContext
+
+G, TILE, D = 8, 64, 128
+
+_EXPECTED_KINDS = [
+    "rmax", "max_elem", "exp_diff", "exp_diff", "rsum", "fma",
+    "mul_bcast", "copy",            # 8 MATH (softmax_merge)
+    "gemm",                          # head GEMM (P · V)
+    "add",                           # epilogue
+]
+
+
+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 _lower():
+    """Lower decode-opt2 #2 and return (tl, cmd, handles)."""
+    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000,
+                   scratch_size=1 << 20)
+    s = _tcm(0x1000, (G, TILE))      # Sj — #1's output, TCM-resident
+    m = _tcm(0x2000, (G,))
+    l = _tcm(0x3000, (G,))
+    O = _tcm(0x4000, (G, D))
+    V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
+    tl.composite(
+        op="gemm", b=V, out=O,
+        prologue=[{"op": "softmax_merge", "s": s, "m": m, "l": l, "O": O}],
+        epilogue=[{"op": "add", "other": O}],
+    )
+    cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1]
+    return tl, cmd, dict(s=s, m=m, l=l, O=O, V=V)
+
+
+# ── recipe shape (ADR-0065 Test #2) ──────────────────────────────────
+
+
+def test_recipe_lowers_to_ten_ops_in_order():
+    _, cmd, _ = _lower()
+    assert [o.kind for o in cmd.ops] == _EXPECTED_KINDS, [o.kind for o in cmd.ops]
+
+
+def test_recipe_rw_handles_are_m_l_O():
+    _, cmd, h = _lower()
+    assert cmd.rw_handles == (h["m"], h["l"], h["O"]), cmd.rw_handles
+
+
+def test_head_gemm_autobinds_primary_out_and_geometry():
+    _, cmd, h = _lower()
+    gemm = cmd.ops[8]
+    assert gemm.kind == "gemm"
+    # b is V (the kernel-provided operand); a is the recipe's primary out P
+    assert gemm.operands["b"] == h["V"]
+    assert "a" in gemm.operands and gemm.operands["a"].shape == (G, TILE)
+    assert gemm.out is not None and gemm.out.addr == h["O"].addr
+    assert gemm.extra.get("m") == G
+    assert gemm.extra.get("k") == TILE
+    assert gemm.extra.get("n") == D
+
+
+def test_recipe_scratch_addresses_are_concrete_and_distinct():
+    _, cmd, _ = _lower()
+    # Collect every operand/out handle address across the 8 MATH ops.
+    addrs = set()
+    for op in cmd.ops[:8]:
+        for hd in op.operands.values():
+            assert hd.addr is not None
+            addrs.add(hd.addr)
+        if op.out is not None:
+            addrs.add(op.out.addr)
+    # The 5 scratch slots (m_loc, m_new, corr, P, l_loc) must have been
+    # allocated to distinct, non-zero addresses (scratch_base supplied).
+    scratch_addrs = {a for a in addrs if a >= 0x100000}
+    assert len(scratch_addrs) >= 5, scratch_addrs
+
+
+def test_recipe_composite_within_size_cap():
+    """ADR-0065 Test #9 — the 10-op composite fits the default 1024 cap."""
+    _, cmd, _ = _lower()
+    assert cmd.logical_bytes < 1024, cmd.logical_bytes