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

PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see
D6). Calling its index `g`:

| OpSpec position | scope | Placement in plan |
|---|---|---|
| `0 .. g-1` | KERNEL | pre-loop stages (run once before tile loop) |
| `g` | KERNEL | head GEMM (drives tile loop with `extra["m"], extra["k"], extra["n"]`) |
| `g+1 .. ` | K_TILE | per K-tile epilogue (inside accumulation loop) |
| `g+1 .. ` | OUTPUT_TILE | per (m, n) tile epilogue (after K accumulation) |
| `g+1 .. ` | KERNEL | post-loop stages (run once after tile loop) |

A composite with no GEMM op (e.g., MATH-only composite) treats all ops
as KERNEL-scope sequential.

### D4. PE_SCHEDULER auto-inserts DMAs from operand `space`

PE_SCHEDULER scans the GEMM's `operands` and `out`. For each:
- `space == "hbm"` → emit DMA_READ Stage (or DMA_WRITE for `out`) before
  the FETCH/GEMM Stages
- `space == "tcm"` → no DMA Stage, operand consumed in place

This is already partially modelled today via `a_pinned`/`b_pinned`; D4
makes the rule uniform and visible: **the kernel never emits explicit
DMA cmds for composite operands; PE_SCHEDULER infers them from
handles**. `tl.ref(addr, shape)` returns a handle with `space="hbm"`;
`tl.zeros` / `tl.full` / `tl.load` outputs return `space="tcm"`.

Non-GEMM ops (MATH chain in prologue / epilogue) have TCM-resident
operands (kernel allocates via `tl.zeros` / `tl.full`); no DMA inserted
for them.

### 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 head GEMM's `operands["a"] = P_handle`.
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.

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

## 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 `space`.** A GEMM composite with
   `b=tl.ref(...)` (space=hbm) emits a DMA_READ Stage; with `b` from
   `tl.zeros(...)` (space=tcm) it does not. Output handle with
   `space=hbm` emits DMA_WRITE; `space=tcm` does not.
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).** opt3 vs opt2 per-tile
   PE_CPU dispatch cycles ratio ≈ 2.4× at default calibration.
8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a
   validation error at TLContext emit.

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

Land order:
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**.