kernbench2/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md

# 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
```

## 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**.