# ADR-0065: Flat-ops `CompositeCmd` + first stateful recipe (`softmax_merge`) ## Status Proposed (verification gated on ADR-0064 Revision 2 land) > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Implements > the §5.6 / §8 item 4 carve-out exactly: opt2 of decode = (#1 existing > GEMM composite for Q·Kᵀ) + (#2 a single composite carrying softmax + > P·V + online-softmax merge). The "ex_composite" / "flash_pv_merge" name > referenced in ADR-0060 §5.6 is *retired* — this ADR fixes the canonical > shape using the existing `tl.composite` entry with two structural > additions: (a) a flat `ops` tuple, and (b) a first stateful recipe > `softmax_merge` that lowers to a sequence of MATH micro-ops. ## Context ### What ADR-0060 left to this ADR ADR-0060 §5.6 shipped opt3 (software pipelining) as the recommended now, and explicitly carved out opt2 — fewer per-tile dispatches, K-before-V DMA priority, **measurable only after ADR-0064**'s cost model — as the follow-up. §8 item 4 sized the carve-out: "if revisited, split it into **two** composites — `#1` = Q·Kᵀ (existing composite + `scale`), `#2` = softmax + P·V + online-softmax accumulator merge". This ADR implements `#2`. ### Why composite needs a structural change Today's `CompositeCmd` has the shape `(op, a, b, out_addr, ops=(head, *epi))` with implicit conventions: - `op` selects the head engine (gemm | math) - `ops[0]` is the head OpSpec, `ops[1:]` are epilogue OpSpecs that run *after* the head per-output-tile or per-K-tile (scope-determined) For decode opt2 we need MATH ops to run **before** the head GEMM (the online-softmax `m, l, O` update + emit `P` for the GEMM input). The current shape has no slot for that. Two failed paths considered earlier: - Embedding softmax_merge as an *epilogue* of #1 (Q·Kᵀ) — wrong: it needs to read `Sj` (= #1's output) and writes a fresh `P` consumed by the *next* GEMM (#2's P·V); circular. - A single mixed-engine recipe `flash_pv_merge` absorbing MATH + GEMM + MATH — violates the engine boundary, complicates PE_SCHEDULER. The clean shape is: **drop the head/epilogue distinction from the command format**. Make `ops` a flat ordered tuple; let each op's `scope` + position determine its tile-loop placement. The "prologue" concept stays in the *user-facing* `tl.composite(...)` API as ergonomics — but the **HW command never sees it**. ### Why a recipe (not a generic op list) Decode opt2's `#2` MATH chain is 8 steps with a fixed structure (online softmax merge). Letting the kernel author wire those 8 steps would be both verbose and a hazard surface (correctness depends on step order). A *recipe* — a single named op (`softmax_merge`) that TLContext expands to the 8 steps with addresses pre-filled — keeps the kernel concise and the structure under one source of truth. The recipe table lives in TLContext (the compiler analog); PE_SCHEDULER never reads it. See boundary in D5. ## Decision ### D1. `CompositeCmd` is a flat ordered op list ```python @dataclass(frozen=True) class CompositeCmd: completion: CompletionHandle ops: tuple[OpSpec, ...] # ordered MATH/GEMM ops rw_handles: tuple[TensorHandle, ...] = () # cross-composite hazard data_op: bool = True @property def logical_bytes(self) -> int: ... # ADR-0064 D2 ``` The previous `(op, a, b, out_addr, out_nbytes, math_op)` fields are absorbed into `ops`. The first OpSpec in `ops` (if `kind == "gemm"`) is the head GEMM; preceding OpSpecs are pre-loop ops, following OpSpecs are post-loop / per-tile epilogue ops depending on `scope`. ### D2. `OpSpec` carries named operands ```python @dataclass(frozen=True) class OpSpec: kind: str # "gemm" | "rmax" | ... scope: Scope # KERNEL | K_TILE | OUTPUT_TILE operands: dict[str, TensorHandle] # named — see D5 extra: dict[str, Any] # scalars, axes, m/k/n, etc. out: TensorHandle | None = None # explicit write-back handle @property def logical_bytes(self) -> int: ... ``` Named operands let PE_SCHEDULER route inputs to engine ports (e.g., GEMM takes `operands["a"]`, `operands["b"]`) without positional convention. ### D3. Position + scope determines tile-loop placement **Semantics split.** *Position* determines the **phase** (where in the tile-loop lifecycle the op fires); *scope* determines the **repetition** within that phase. `Scope.KERNEL` means "**once per `CompositeCmd` invocation**", not "once per kernel launch" — its single firing happens at the op's position in the sequence. The same KERNEL value carries different *phases* depending on whether the OpSpec sits before, at, or after the GEMM op. PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see D6). Calling its index `g`: | OpSpec position | scope | Phase | Placement in plan | |---|---|---|---| | `0 .. g-1` | KERNEL | pre-tile-loop | run once before tile loop | | `g` | KERNEL | head GEMM | drives tile loop with `extra["m"], extra["k"], extra["n"]` | | `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (inside K accumulation loop) | | `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (after K accumulation) | | `g+1 .. ` | KERNEL | post-tile-loop | run once after tile loop | A composite with no GEMM op (e.g., MATH-only composite) treats all ops as KERNEL-scope sequential — phase collapses to "single-shot in order". ### D4. PE_SCHEDULER auto-inserts DMAs from the operand `pinned` flag PE_SCHEDULER scans the GEMM's `operands`. For each: - **not `pinned`** (data lives in HBM, must be streamed) → emit DMA_READ Stage before the FETCH/GEMM Stages - **`pinned`** (already staged in TCM via a prior `tl.load`, or a recipe's TCM scratch / primary-out `P`) → no DMA Stage, consumed in place The kernel never emits explicit DMA cmds for composite operands; PE_SCHEDULER infers them from handles. `tl.ref` operands are not pinned (DMA_READ); `tl.load` results and recipe scratch / primary-out are pinned (in place). > **As-built note (P3).** An earlier draft of D4 derived the DMA decision > from a `space` tag (`hbm`→DMA, `tcm`→in place) and would have flipped the > current `space` values (today `tl.load`=`hbm`, `tl.ref`=`tcm`). The > implementation **keeps the existing `pinned` flag** as the DMA signal > instead — `pinned` already encodes "already in TCM, skip DMA", flipping > `space` would change existing bench Stage sequences for no functional > gain. Unifying `space` and retiring `pinned` is a deferred low-priority > cleanup. **Prologue recipe ops are TCM-only (Phase 1 limit, enforced at TLContext emit).** Every operand of a *prologue recipe* (non-GEMM) OpSpec **must** be TCM-resident; passing an HBM handle as a recipe operand raises a validation error at emit time. This bounds DMA staging to the head op's operands in Phase 1. The **head op itself** (gemm *or* math) keeps the existing DMA-staged-from-HBM behavior — D4's TCM-only rule does **not** apply to it. This invariant is also restated in D6 #7. ### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py` Recipes live in a new TLContext-adjacent module (`src/kernbench/triton_emu/tl_recipes.py`). PE_SCHEDULER does **not** import it. The first recipe: ```python @dataclass(frozen=True) class RecipeDescriptor: # Operand contract: name → "R" (read-only) | "RW" (read-write) operands: dict[str, str] # Type / shape rule for the implicit primary output (auto-bound to # head's first matrix operand). None means recipe emits no value. primary_out: PrimaryOutSpec | None # Single-shot prologue (full reduction) vs tile-aligned execution. tile_alignment: Literal["single_shot", "tile_aligned"] # Internal scratch budget (bytes), function of operand dims. internal_scratch_bytes_fn: Callable[..., int] # Engine micro-op sequence: TLContext expands this into flat OpSpecs. engine_seq: tuple[EngineOp, ...] RECIPE_DESCRIPTORS["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"), ), ) ``` TLContext, on `tl.composite(prologue=[{"op": "softmax_merge", ...}], op="gemm", ...)`: 1. Validates operand R/RW against `RECIPE_DESCRIPTORS["softmax_merge"]`. 2. Allocates scratch (m_loc, m_new, corr, P, l_loc) via the scratch_scope helper (ADR-0063). 3. Derives the primary output P's shape from `s.shape` (identity). 4. Expands `engine_seq` into 8 flat MATH OpSpecs with all addresses / sizes filled. Each OpSpec gets `scope=KERNEL`. 5. **Auto-binds (with conflict check, D6.6).** If the head GEMM does *not* explicitly provide `operands["a"]`, auto-bind `operands["a"] = P_handle`. If the kernel *did* provide `a` explicitly, emit a validation error — ambiguous primary-output replacement. 6. Emits `CompositeCmd(ops=(8 MATH + 1 GEMM + epilogue), rw_handles= (m, l, O))`. **RECIPE_DESCRIPTORS appears nowhere in the HW path**. PE_SCHEDULER sees only the flat ops list. ### D6. Invariants 1. **GEMM count.** `0 ≤ count(op.kind == "gemm" in cmd.ops) ≤ 1`. A composite with 0 GEMMs is MATH-only. 2. **softmax_merge has no V operand.** `RECIPE_DESCRIPTORS["softmax_merge" ].operands` does not include any HBM-resident handle → no DMA inserted during the prologue → V DMA can only happen during the GEMM head → natural K-before-V DMA priority across #1 / #2 of decode opt2. 3. **Strict FIFO across composites.** PE_SCHEDULER tracks in-flight `rw_handles`. A new composite whose `rw_handles` (or any of its op inputs) intersect with an in-flight composite's `rw_handles` waits for all earlier composites to complete — no out-of-order reorder. 4. **No backwards-incompatibility for legacy callers.** The user-facing API `tl.composite(op="gemm", a, b, epilogue=[...])` is preserved. TLContext internally lowers it to a flat-ops `CompositeCmd`. 5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER does not import RECIPE_DESCRIPTORS, does not know that "softmax_merge" exists, and has no recipe-specific code branches. 6. **No implicit operand replacement (auto-bind conflict).** If a prologue recipe declares a `primary_out` AND the head op's auto-bind target operand (e.g., GEMM `a`) is *also* provided by the kernel explicitly, TLContext raises a validation error at emit time. The kernel must either let the recipe bind (omit the operand) **or** specify it explicitly (omit the prologue, or use a recipe without `primary_out`). This prevents the ambiguity of "which value wins". 7. **Prologue MATH operands TCM-only.** Restatement of D4: every operand of a *prologue recipe* (non-GEMM) OpSpec must be TCM-resident. The head op (gemm or math) is exempt — it may stream from HBM. Phase 1 enforced at TLContext emit. ### D7. Boundary summary (compiler vs scheduler vs engine) ``` HOST DEVICE TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher) - validate against RECIPE_DESC. - scan flat ops list - derive primary_out shape - identify GEMM (≤1) - allocate scratch - position-based stage placement - expand recipe → flat OpSpecs - auto-insert DMA stages from space - compute rw_handles - strict-FIFO RW hazard tracking - emit CompositeCmd - emit Stage list imports: RECIPE_DESCRIPTORS imports: nothing recipe-related ENGINES (PE_MATH / PE_GEMM / PE_DMA) - read Stage.params (op_kind, addresses, n_elements) - existing latency models unchanged imports: nothing recipe-related ``` ### D8. Composite output handle + output-space DMA `tl.composite(...)` returns the **output `TensorHandle`** (not a bare `CompletionHandle`), so a composite's result can feed a downstream op the same way `tl.dot`'s result does. The handle's `pending` field references the composite's completion; downstream consumers auto-await it (`_await_pending`, ADR-0062 lazy pattern) and `tl.wait(handle)` works. **`out` is always a `TensorHandle`** (never a raw address) — its location is the kernel's choice, never guessed: - **HBM output** → `out=tl.ref(addr, shape)`. `tl.ref` returns a handle with `space="hbm"`, so the composite writes it back via DMA_WRITE. The STORE + DMA_WRITE stay **inside** the composite pipeline — they are part of the fused op, *not* a separate `tl.store`. - **In-place TCM** → `out=` (e.g. the recipe accumulator `O`). `space="tcm"` → STORE only, result stays in TCM (chainable). - **Omitted** → TLContext **auto-allocates a TCM scratch** (like `tl.dot`) and returns it. For ergonomics, `out_ptr: int` is accepted as **shorthand for an HBM output** — equivalent to `out=tl.ref(out_ptr, )` (the compiler wraps it into an `space="hbm"` handle). When no output is specified the compiler owns the scratch address, never the kernel. (`tl.ref` returns `space="hbm"` — it references HBM data. The operand-**input** DMA decision stays `pinned`-based per D4, so this label does not affect input streaming.) The output handle's **`space` drives the tile loop's write-back** (the output analog of the operand `pinned` rule in D4): | `out.space` | tile-loop write-back | | --- | --- | | `"tcm"` | STORE only — result stays in TCM (chainable; **no DMA_WRITE**) | | `"hbm"` | STORE + DMA_WRITE — final HBM output | This makes a TCM-resident composite output (e.g. a recipe accumulator) legal **and** keeps a TCM-scratch address out of the DMA engine's PA decoder (which would otherwise reject the high-bit scratch address). Existing benches wrap their HBM output via `tl.ref` → write-back unchanged (byte-equal). ## Alternatives ### A1. Three composites per tile (no prologue concept) Split #2 further into `softmax_merge` + `gemm(P·V)` + `add(O)` as three separate CompositeCmds. Pros: no flat-ops restructuring (`add` reuses existing epilogue path). Cons: 3 dispatches per tile instead of 2 — diverges from ADR-0060 §8 item 4 sizing ("split into **two**"); loses ~32 ns × N_tiles of fixed-cost savings. **Rejected.** ### A2. Mixed-engine recipe (`flash_pv_merge`) One recipe absorbing MATH + GEMM + MATH (ADR-0060 §5.6 original phrasing). Pros: tightest dispatch (1 composite). Cons: PE_SCHEDULER must know engine-boundary recipes, violates the "engines see only Stage.op_kind" boundary, recipe table becomes ISA-like and pollutes HW interface. **Rejected.** ### A3. Generic op-list DAG (`tl.ex_composite([...])`) User-defined sequence of any ops with named slots and dataflow analysis. Pros: maximally flexible. Cons: a single user (softmax_merge) does not justify a generic DAG language; PE_SCHEDULER would need dataflow analysis; reintroduces ADR-0060 §8 item 4's "absorb the inner loop" trap. **Rejected (Simplicity First).** ### A4. RW-aware scheduler reorder (not strict FIFO) PE_SCHEDULER could allow a later CompositeCmd to overtake an in-flight one when their `rw_handles` don't intersect (e.g., next-tile #1 could overlap with this-tile #2 since #1 doesn't touch m/l/O). Pros: more overlap. Cons: more scheduler state, gains are bounded (next #1 waits for K DMA anyway). **Deferred** — strict FIFO ships first; RW-aware reorder may be a future follow-up ADR. ### A5. Embedding softmax_merge as an epilogue of #1 (Q·Kᵀ) Wrong: softmax_merge's output `P` is the input of #2's P·V GEMM, which runs *after* #1's epilogue. Embedding it would force #2 to read its own input from #1's epilogue scratch — circular tile-loop dependency. **Rejected (incorrect).** ## Consequences ### Positive - ADR-0060 §5.6 / §8 item 4 carve-out implemented exactly: opt2 = 2 composites per tile, K-before-V DMA priority natural. - Composite contract becomes uniform: flat ops list with scope-based placement. No structural distinction between "head" and "epilogue" in the HW cmd. - DMAs auto-inferred from operand `space`; kernel surface simpler. - HW interface (`CompositeCmd` shape) stays small even as recipes grow — recipe expansion happens host-side. - New fused patterns (linear+gelu+norm, rmsnorm+linear, …) can be added by inserting one `RECIPE_DESCRIPTORS` entry; no PE_SCHEDULER changes. ### Negative - CompositeCmd's struct shape changes — a refactor for all current callers. Semantics preserved (D6.4); existing goldens unchanged once ADR-0064 Revision 2 lands. - `OpSpec.operands` changes from positional `tuple[Any, ...]` to `dict[str, TensorHandle]` — touches the existing epilogue lowering. - New TLContext code: recipe expansion + scratch-slot allocation + primary_out auto-bind. ~80–120 LOC. - PE_SCHEDULER's `_generate_plan` gets a position-based scan + DMA auto-insertion + strict-FIFO RW tracker. ~60–100 LOC. - **Strict-FIFO RW hazard tracking is safe but conservative.** When two composites share any `rw_handle`, the new one waits for *all* prior overlapping composites to fully complete — even when their compute could in principle overlap (e.g., next-tile #1 = Q·Kᵀ does not touch m/l/O and could overlap with this-tile's #2, but FIFO still serializes any path that touches the same handles down the line). This **may under-expose composite-level overlap** relative to a smarter scheduler. RW-aware reorder (A4) is the deferred improvement. ## Open review items 1. **Recipe shape derivation rule beyond identity.** Future recipes may need non-identity transforms (e.g., transposed primary_out). The `PrimaryOutSpec.transform` field is forward-compat; add new transforms as needed. 2. **Recipe scratch allocator** integration with ADR-0063's `scratch_scope`. Initial wiring: TLContext allocates inside the current `scratch_scope` context if active, else in a kernel-scope persistent slot. 3. **`tl_recipes.py` location.** New module under `triton_emu/`. Keeps recipe knowledge co-located with the compiler analog (TLContext). 4. **GEMM `out=TensorHandle`.** New form alongside legacy `out_ptr: int`. Both accepted; handle form preferred for new code (enables strict-FIFO RW tracking by handle identity). ## Test Requirements 1. **CompositeCmd flat-ops refactor — meaning preservation.** Every existing bench's op_log is byte-equal pre/post refactor (legacy API `tl.composite(op="gemm", a, b, epilogue=[...])` lowers to the same Stage sequence). 2. **softmax_merge recipe lowering.** TLContext call `tl.composite( prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}], op="gemm", b=V_ref, out=O, epilogue=[{"op": "add", "other": O}])` produces a `CompositeCmd` with exactly 10 ops in the expected order (8 MATH + 1 GEMM + 1 MATH(add)) and `rw_handles == (m, l, O)`. 3. **K-before-V DMA priority invariant.** op_log of decode opt2 shows no V-related DMA during #2's MATH prologue (#2's V DMA begins only when the GEMM head starts). 4. **Strict-FIFO RW serialization.** Two consecutive composites sharing any `rw_handle` complete in dispatch order; their stages do not interleave. 5. **DMA auto-insertion from `pinned`.** A GEMM composite with a not-`pinned` operand (e.g. `b=tl.ref(...)`) emits a DMA_READ Stage; a `pinned` operand (e.g. a recipe's TCM scratch / primary-out, or a `tl.load` result) does not. (As-built D4 note: the DMA decision uses the existing `pinned` flag, not a `space` tag.) 6. **opt2 numeric parity with opt3.** In data mode, opt2's final `(m, l, O)` matches opt3's within fp tolerance. 7. **opt2 dispatch ratio (after ADR-0064 Rev2) — robust.** opt3 vs opt2 per-tile PE_CPU dispatch cycles satisfies `opt3 > 2 × opt2` (qualitative gate; calibration-independent). The default-calibration model expectation is ≈ 4.0× (FIXED=40 cycles, R=0.0625 cycles/byte) — see DDD-0065 §11 for the model-derived numbers; only the loose `> 2×` bound is the test gate so it survives calibration changes. Ratio is FIXED-dominated, reflecting command-count reduction as the primary signal. 8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a validation error at TLContext emit. 9. **Within composite size cap (ADR-0064 D7).** Decode opt2's `#2` composite (10 ops, ~322 logical bytes — per-op summation of operand handles) fits comfortably under the default `MAX_COMPOSITE_LOGICAL_BYTES=1024` — no segmentation. The recipe lowering test asserts `cmd.logical_bytes < 1024`. 10. **Auto-bind conflict.** `tl.composite(op="gemm", a=A_handle, prologue=[{"op": "softmax_merge", ...}], ...)` where softmax_merge declares `primary_out` raises a validation error at emit time (D6.6). 11. **MATH operand TCM invariant.** `tl.composite(op="math", a=tl.ref(addr, shape), ...)` (passing an HBM-resident handle as a MATH operand) raises a validation error at emit time (D6.7). ## Dependencies - **ADR-0060** §5.6, §8 item 4 — the carve-out this ADR implements. - **ADR-0064 Revision 2** — dispatch cost model; lands first, makes opt2's fewer-issues win measurable. - **ADR-0063** `scratch_scope` — recipe intermediates allocated inside; `m, l, O` allocated outside as persistent arena (per DDD-0060 §6.2). - **ADR-0046** `tl_context` contract — extended by `prologue=[...]` kwarg, `out=TensorHandle`, RECIPE_DESCRIPTORS lookup. - **ADR-0042** tile plan generators — extended (not redesigned) to process the flat ops list and auto-insert DMAs. - **ADR-0014** PE pipeline — boundary preserved: scheduler stays recipe-free, engines stay op_kind-opaque. ## Migration > **Implementation ordering note.** The ADR-0064-first ordering below was > the original plan. The actual implementation **reversed steps 1 and 2** — > ADR-0065 Phase 1 (flat-ops `CompositeCmd`) landed *before* ADR-0064 Rev2. > Reason: `logical_bytes` (ADR-0064 D2) is naturally defined on the flat-ops > shape, so doing the refactor first avoids a throwaway transitional > `logical_bytes`. Each step is independently safe: P1 is a meaning- > preserving refactor (op_log byte-equal) under the pre-existing cost path; > ADR-0064 Rev2 then swaps in the structural formula on the already-flat > shape. The end state is identical either way. Land order (original plan; see note above for the as-built reversal of 1↔2): 1. **ADR-0064 Revision 2** (separate PR): structural dispatch cost + `logical_bytes` property + topology config override + one-time goldens regeneration. 2. **ADR-0065 Phase 1** (this ADR, PR a): `CompositeCmd` flat-ops restructure + legacy-API lowering. Refactor only; goldens unchanged. 3. **ADR-0065 Phase 2** (this ADR, PR b): `tl_recipes.py` + `softmax_merge` recipe + PE_SCHEDULER position-scan + DMA auto-insertion + strict-FIFO RW tracker. 4. **ADR-0065 Phase 3** (this ADR, PR c): `_gqa_decode_long.py` opt2 variant + numeric parity test + dispatch-ratio measurement against opt3 baseline. Detailed implementation plan: see **DDD-0065**.