# Detailed Design Document — Flat-ops `CompositeCmd` + `softmax_merge` recipe (ADR-0065) **Status:** Accepted (companion to ADR-0065, Accepted) — reconciled to the as-built implementation. **Scope:** the `CompositeCmd` flat-ops restructure, the `softmax_merge` recipe, the decode opt2 variant in `_gqa_attention_decode_opt2.py`, and the data-mode numeric path that makes opt2's recipe compute. **Audience:** reviewer / implementer. Read **ADR-0065** and **ADR-0064 Revision 2** first. --- ## 0. As-built reconciliation (read first) This DDD was authored in the design round; the implementation diverged in several places. The sections below have been corrected, but the headline deltas are collected here. (ADR-0065 itself was kept in sync; see commits `47e2c78`,`184f654`,`55f025c`,`17fb940`,`35453cc`,`4a55ae5`,`e8d6c28`, `4089e18`,`3d4a3d4`,`5bbe293`,`e9d62b9`,`f0c5b9c`.) | Topic | Original DDD | As-built | | --- | --- | --- | | **Phase order** | P0 (ADR-0064 Rev2) lands first | **P1 (flat-ops) first**, then P0 — `logical_bytes` is naturally defined on the flat shape, so the refactor goes first | | **Operand-input DMA** | from `handle.space` (`space=="hbm"`→DMA_READ) | **from the `pinned` flag** (already-in-TCM→skip). The `space` flip was avoided (it would break byte-equality); `tl.ref`→hbm, but the *decision* stays pinned-based | | **Output DMA / handle** | (not addressed; `out=O`) | **ADR-0065 D8** — `tl.composite` returns the output `TensorHandle`; `out`=handle (HBM via `tl.ref`) \| `out_ptr` (HBM shorthand) \| omitted→auto-TCM; **DMA_WRITE gated on `out.space=="hbm"`** (TCM output stays on-chip, chainable) | | **Size cap (D7)** | deterministic segmentation | **hard cap → `ValueError`** at emit (no segmentation) | | **Prologue execution** | feeder yields `prologue_stages` first | feeder feeds each prologue stage as a **standalone 1-stage tile** (folding broke PE_MATH→PE_DMA routing) | | **D6.7 (MATH TCM-only)** | every non-GEMM op | **prologue recipe ops only**; the head op (gemm *or* math) keeps DMA-staged-from-HBM | | **Numeric parity** | "just works" via the existing data executor | composites did **not** compute in data mode at all — built in **N1–N4** (below): composite-as-gemm record, recipe MATH compute, GEMM accumulate, opt2 ordering | | **Dispatch ratio** | model ≈ 4.0× | **measured 5.22×** per-tile (gate `>2×`); R-sweep 4.20×/5.22×/5.45× | | **opt2↔opt3 parity** | exact in data mode | exact recipe math verified by N2/N3 (isolated, non-trivial); e2e opt2↔opt3 holds in the **near-uniform regime** (kernbench K reshape-as-transpose caveat, shared by both kernels) | **Data-mode numerics (N1–N4), absent from the original DDD — see §15.** --- ## 1. Goal and success criteria **Goal:** make decode opt2 *runnable*, *numerically equivalent to opt3*, and *measurably cheaper in PE_CPU dispatch* under the ADR-0064 Rev2 cost model. Existing benches keep their (post-ADR-0064 Rev2) goldens unchanged. **Success criteria:** 1. **Refactor safety.** `CompositeCmd` flat-ops restructure preserves the op_log byte-for-byte for every existing bench (`tests/` all green, no goldens diff). 2. **Recipe correctness.** opt2 final `(m, l, O)` matches opt3 within fp tolerance in data mode (`enable_data=True`). 3. **Boundary preservation.** PE_SCHEDULER imports nothing recipe-related; PE_MATH/PE_GEMM/PE_DMA see only `Stage.params["op_kind"]` (no new engine code). 4. **Dispatch ratio (robust).** Per-tile PE_CPU dispatch cycles satisfy `opt3 > 2 × opt2` (qualitative invariant, calibration-independent). Default ADR-0064 Rev2 calibration (FIXED=40 cycles, R=0.0625 cycles/byte) gives model-expected ≈ 4.0× — informative only, not the test gate. 5. **K-before-V invariant.** op_log of decode opt2 shows zero V-related DMA during #2's MATH prologue. **Non-goals:** prefill (causal `if` in kernel control flow, separate work); RW-aware scheduler reorder (strict FIFO ships); user-defined recipes; new engine micro-ops in PE_MATH (it stays op_kind-opaque). --- ## 2. Where the design sits in the current code ``` runtime_api (host) RuntimeContext.zeros/empty/launch context.py │ KernelLaunchMsg ▼ sim_engine GraphEngine(enable_data=…) engine.py MemoryStore memory_store.py │ ▼ components / PE pipeline (per PE) PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM} greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py CompositeCmd → tile plan → engines pe_scheduler.py (ADR-0014 D6) │ ▼ tl programming model tl_context.py composite / dot / load / store / scratch_scope ── ADR-0065 NEW ────────────────────────── tl_recipes.py RECIPE_DESCRIPTORS, EngineOp, RecipeDescriptor, PrimaryOutSpec ``` The new module `tl_recipes.py` sits beside `tl_context.py`. PE_SCHEDULER (`pe_scheduler.py`) and `tiling.py` extend to consume the flat ops list without importing `tl_recipes`. --- ## 3. File plan ### 3.1 Modified production files | File | Change | ADR ref | |---|---|---| | `src/kernbench/common/pe_commands.py` | `CompositeCmd` flat-ops dataclass; `OpSpec.operands: dict`; new `Scope` value semantics (KERNEL position-based); `logical_bytes` property on every `PeCommand`. | 0065 D1/D2, 0064 D2 | | `src/kernbench/triton_emu/tl_context.py` | `composite(prologue=[...], out=TensorHandle, …)`; recipe lowering pass (RECIPE_DESCRIPTORS lookup, scratch alloc, OpSpec expansion, primary_out auto-bind); legacy API preserved (`op="gemm", a, b, epilogue=[...]` still works, internally lowers to flat ops). | 0065 D1/D5 | | `src/kernbench/triton_emu/tl_recipes.py` (NEW) | `RecipeDescriptor`, `PrimaryOutSpec`, `EngineOp` dataclasses; `RECIPE_DESCRIPTORS["softmax_merge"]` entry; expansion helper. | 0065 D5 | | `src/kernbench/components/builtin/pe_scheduler.py` | `_generate_plan` rewritten for flat ops: scan for GEMM (≤1), position-based KERNEL placement, OUTPUT_TILE/K_TILE epilogue, strict-FIFO RW tracker keyed by handle id. | 0065 D3/D6 | | `src/kernbench/components/builtin/tiling.py` | `generate_plan_from_ops(ops, tile_*, bpe, pe_prefix)` wraps `generate_gemm_plan`/`generate_math_plan`. Input DMA_READ from the **`pinned`** flag; **output DMA_WRITE gated on `out.space=="hbm"`** (D8). `_math_stage` carries operand+output addrs on recipe prologue stages (N2). | 0065 D3/D4/D8 | | `src/kernbench/common/pe_cost_model.py` (NEW) | `PeCostModel` (FIXED/R/cap) + `from_node_attrs`; replaces the removed Rev1 `cpu_issue_cost.py`. | 0064 D1/D4 | | `src/kernbench/components/builtin/pe_cpu.py` | reads `pe_cost_model:` + `clock_freq_ghz` from node attrs, threads to TLContext (legacy + greenlet KernelRunner). | 0064 D1/D4 | | `src/kernbench/sim_engine/op_log.py` | `_extract_op_info(CompositeCmd)` scans for the GEMM op + `accumulate` flag (N1/N3); `record_end` promotes an address-carrying MATH stage to op_kind="math" (N2). | 0065 N1-N3 | | `src/kernbench/sim_engine/data_executor.py` | `_compute_math` +5 recipe ops; `_execute_gemm` best-effort skip + accumulate (N1-N3). | 0065 N1-N3 | | `src/kernbench/benches/_gqa_attention_decode_opt2.py` (NEW) | opt2 variant — tile 0 establishes (m,l,O) with primitives; the merge tile uses `#1 scores = tl.composite(op="gemm", a=Q, b=K_T1)` + `#2 tl.composite(prologue=[softmax_merge], op="gemm", b=V_ref, out=O, epilogue=[add])`. | 0065 §Migration P3 | ### 3.2 New tests | File | Covers | |---|---| | `tests/test_composite_flat_ops_refactor.py` | P1 flat-ops lowering shape (legacy + math + epilogue) | | `tests/test_pe_cost_model.py` | ADR-0064 Rev2 — `logical_bytes`/formula, override, `PeCpuOverheadCmd` bypass, D7 cap→error, live wiring/path-parity | | `tests/test_tl_recipe_softmax_merge.py` | recipe lowering — 10 ops in order + `rw_handles == (m, l, O)` + size cap | | `tests/test_tl_composite_validation.py` | auto-bind conflict (D6.6); prologue MATH operand TCM-only (D6.7) | | `tests/test_pe_scheduler_flat_ops_plan.py` | position-scan; pinned-based DMA; prologue stages; MATH-only | | `tests/test_pe_scheduler_strict_fifo.py` | `_RwHazardTracker` admit/retire (RAW/WAW/disjoint/legacy) | | `tests/test_composite_output_handle.py` | D8 — composite returns handle (auto-TCM); chaining; output-space DMA; `tl.ref`→hbm | | `tests/test_recipe_data_mode.py` | **N2/N3** — recipe MATH compute m/l/O-rescale + full accumulate (O = O·corr + P@V) vs numpy | | `tests/test_e2e_pipeline.py` | **N1** — composite GEMM computes a@b in data mode | | `tests/attention/test_gqa_decode_opt2.py` | dispatch ratio (>2×) + R-sweep; K-before-V; **N4** opt2↔opt3 parity | --- ## 4. Data model — `CompositeCmd` flat-ops ### 4.1 New shape ```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: return ( 4 # framing + 1 + sum(op.logical_bytes for op in self.ops) + 1 + 8 * len(self.rw_handles) ) ``` The previous `op`, `a`, `b`, `out_addr`, `out_nbytes`, `math_op` fields are **removed**. All info lives in `ops[i]`. ### 4.2 `OpSpec` extended ```python @dataclass(frozen=True) class OpSpec: kind: str # "gemm" | "rmax" | "exp_diff" | ... scope: Scope # KERNEL | K_TILE | OUTPUT_TILE operands: dict[str, TensorHandle] extra: dict[str, Any] = field(default_factory=dict) out: TensorHandle | None = None @property def logical_bytes(self) -> int: return ( 1 + 1 # opcode + scope enum + 1 + 8 * len(self.operands) # named handles + (8 if self.out is not None else 0) + 1 + 4 * len(self.extra) # scalars ) ``` Backward compat: existing `operands: tuple[Any, ...]` use sites (epilogue `bias`, `scale`, `add`, etc.) migrate to `operands={"other": h}` / `operands={"src": h}` per a small lookup in `pe_commands.EPILOGUE_OPS`. ### 4.3 `Scope` semantics — phase vs repetition (ADR-0065 D3) `Scope` keeps its three values (`KERNEL`, `K_TILE`, `OUTPUT_TILE`). What changes is the **semantics split**: - **Position** determines the *phase* (where in the tile-loop lifecycle the op fires — pre-loop / in-loop / post-loop). - **Scope** determines the *repetition* within the phase (KERNEL = once per `CompositeCmd` invocation; K_TILE = per K-tile iteration; OUTPUT_TILE = per output tile). `KERNEL` does *not* mean "once per kernel launch" or "once anywhere". It means "once per `CompositeCmd` invocation, at this op's position". A composite with multiple KERNEL-scope ops (e.g., softmax_merge's 8 prologue ops) executes each exactly once, in sequence order, before the GEMM tile loop begins. Concretely: - KERNEL ops *before* GEMM → pre-tile-loop (single shot, in order) - KERNEL ops *after* GEMM → post-tile-loop (single shot, in order) - K_TILE / OUTPUT_TILE ops *after* GEMM → inside tile loop as today --- ## 5. Recipe: `softmax_merge` ### 5.1 `tl_recipes.py` definitions ```python @dataclass(frozen=True) class PrimaryOutSpec: 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: engine: Literal["MATH", "GEMM", "DMA"] op_kind: str # PE_MATH op_kind label (opaque to engine) # Slot names: any of "src", "src_a", "src_b", "src_c", "dst" — these # are RecipeDescriptor.operands names or internal scratch names. 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 spec tile_alignment: Literal["single_shot", "tile_aligned"] internal_scratch_bytes_fn: Callable[..., int] engine_seq: tuple[EngineOp, ...] RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = { ... } ``` ### 5.2 Engine sequence (8 MATH ops) | # | engine | op_kind | reads (slots) | writes (slots) | extras | |---|---|---|---|---|---| | 1 | MATH | `rmax` | s | m_loc | `reduce_axis=-1` | | 2 | MATH | `max_elem` | m, m_loc | m_new | — | | 3 | MATH | `exp_diff` | m, m_new | corr | — | | 4 | MATH | `exp_diff` | s, m_new | P (primary_out)| `bcast_axis=0` | | 5 | MATH | `rsum` | P | l_loc | `reduce_axis=-1` | | 6 | MATH | `fma` | l, corr, l_loc | l | — | | 7 | MATH | `mul_bcast` | O, corr | O | `bcast_axis=1` | | 8 | MATH | `copy` | m_new | m | — | Intermediates `m_loc, m_new, corr, P, l_loc` are scratch slots allocated by TLContext before lowering. `P` is the recipe's primary output; auto- bound to the head GEMM's `operands["a"]`. ### 5.3 PE_MATH op_kind additions PE_MATH is op_kind-opaque ([pe_math.py:90-99](src/kernbench/components/builtin/pe_math.py#L90-L99)). Latency = `ceil(num_elements / vector_width) / clock`. Adding new op_kind labels (rmax, max_elem, exp_diff, rsum, fma, mul_bcast, copy) requires **no PE_MATH code changes** — they pass through as string labels in `Stage.params["op_kind"]`. Latency stays consistent for the same `num_elements`. (Reductions like `rmax` / `rsum` have log-depth final stages in real HW not modelled by `ceil(num_elements / vector_width)`; this is an existing approximation, not introduced by ADR-0065.) --- ## 6. TLContext lowering pipeline Input (user-facing API, preserved): ```python tl.composite( op="gemm", b=tl.ref(v_tile(j), (TILE, d)), out=O, prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}], epilogue=[{"op": "add", "other": O}], ) ``` Lowering steps in `TLContext.composite()`: 1. **Validate prologue items.** For each, lookup `RECIPE_DESCRIPTORS[item["op"]]`. Check all operands present, type tags (R / RW) match handle states (e.g., RW must not be `tl.ref`). 2. **Allocate internal scratch.** Call `recipe.internal_scratch_bytes_fn(G, TILE, d, bpe)`; allocate via `scratch_scope` helper (ADR-0063). Get slot addresses for `m_loc, m_new, corr, P, l_loc`. 3. **Derive primary_out handle.** From `recipe.primary_out`: `P` handle = `TensorHandle(shape=Sj.shape, dtype=Sj.dtype, addr=

, space="tcm")`. 4. **Expand `engine_seq` → flat OpSpecs.** For each `EngineOp`: - Resolve slot names (operands or scratch) to `TensorHandle`s. - Build `OpSpec(kind=op_kind, scope=Scope.KERNEL, operands={...}, extra={...}, out=...)`. - `out` is the handle for `dst` slot; `operands` is `{src*: handle}`. 5. **Auto-bind head GEMM (with conflict check, ADR-0065 D6.6).** GEMM OpSpec built from `op="gemm"`, `operands={"a": P_handle, "b": V_handle}`, `out=`, `extra={"m": G, "k": TILE, "n": d}`. **Validation:** if the kernel passed `a=` explicitly AND there is a prologue recipe with `primary_out` that would auto-bind into `a`, raise a validation error (auto-bind conflict). The kernel must omit `a` (let recipe bind) or omit the prologue / use a `primary_out=None` recipe (no auto-bind). 6. **Build epilogue OpSpecs.** Each `epilogue=[{...}]` entry maps to `OpSpec(kind=..., scope=Scope.OUTPUT_TILE, operands={...}, ...)`. The `add` epilogue specifically gets `operands={"other": O_handle}` so the Stage later carries `other_handle` (GAP 9 fix). 7. **Compute `rw_handles`.** Collect all handles tagged RW in prologue recipes + the head GEMM's `out` (if accumulating) + any RW epilogue targets. For softmax_merge: `rw_handles = (m, l, O)`. 8. **Operand validation (ADR-0065 D6.7, as-built scope).** Checked on the **prologue recipe ops only**: every operand handle must be `space == "tcm"` (HBM recipe operand → `ValueError`). The head op (gemm *or* math) is **exempt** — it may stream from HBM. softmax_merge's 8 MATH ops only reference TCM-resident handles → the invariant naturally holds. Also: GEMM-count ≤ 1 (D6.1). 9. **Check composite size cap (ADR-0064 D7, as-built).** If `cmd.logical_bytes > max_composite_logical_bytes` (default 1024), raise `ValueError` at emit — **hard cap, no auto-segmentation**. softmax_merge's 10-op composite (~322 bytes) is well under the cap. 10. **Resolve the output handle (ADR-0065 D8) + emit.** `out`=handle (HBM `tl.ref` \| in-place TCM) \| `out_ptr` (HBM shorthand) \| omitted→auto-TCM scratch. Emit `CompositeCmd(ops=(8 MATH + 1 GEMM + 1 MATH(add)), rw_handles=(m,l,O))`, attach a `CompositeFuture` (carrying the completion) to the output handle, and **return that handle** so the result chains like `tl.dot`'s (downstream auto-awaits via `_await_pending`). **Legacy lowering.** Pre-existing `tl.composite(op="gemm", a=A, b=B, epilogue=[{"op":"bias","bias":h}])` goes through the same emit path: no prologue, head OpSpec from `op="gemm"`, epilogue OpSpecs built from `epilogue=[...]`. Output `CompositeCmd` has flat ops list; semantics preserved. --- ## 7. PE_SCHEDULER tile plan generation (flat-ops) ```python def _generate_plan(cmd: CompositeCmd) -> PipelinePlan: pp = self._pe_prefix bpe = 2 # 1. Identify GEMM op (≤1). gemm_idx = next( (i for i, o in enumerate(cmd.ops) if o.kind == "gemm"), None, ) # 2. MATH-only composite (no GEMM): single tile of all ops in order. if gemm_idx is None: return _math_only_plan(cmd.ops, pp) gemm = cmd.ops[gemm_idx] M, K, N = gemm.extra["m"], gemm.extra["k"], gemm.extra["n"] # 3. Partition ops by position. pre_loop_ops = cmd.ops[:gemm_idx] # KERNEL scope after_ops = cmd.ops[gemm_idx+1:] # mixed scopes # 4. Pre-loop MATH Stages (single-shot, before tile loop). pre_stages = [_math_stage_from_op(op, pp) for op in pre_loop_ops] # 5. Tile loop (existing generate_gemm_plan pattern, but: # - DMA_READ inserted only when operand.space == "hbm" # - K_TILE epilogue from after_ops with scope == K_TILE # - OUTPUT_TILE epilogue from after_ops with scope == OUTPUT_TILE # - DMA_WRITE inserted only when out.space == "hbm") tile_plans = _generate_tile_loop( gemm, M, K, N, after_ops, pp, bpe, ) # 6. Post-loop KERNEL Stages (after tile loop completes). post_loop_ops = [op for op in after_ops if op.scope == Scope.KERNEL] post_stages = [_math_stage_from_op(op, pp) for op in post_loop_ops] return PipelinePlan( prologue_stages=tuple(pre_stages), tiles=tile_plans, epilogue_stages=tuple(post_stages), m_tiles=..., k_tiles=..., n_tiles=..., ) ``` `PipelinePlan` gets two new fields: ```python @dataclass(frozen=True) class PipelinePlan: tiles: list[TilePlan] m_tiles: int k_tiles: int n_tiles: int prologue_stages: tuple[Stage, ...] = () epilogue_stages: tuple[Stage, ...] = () ``` Feeder change (`_feed_loop`): yield `prologue_stages` Stages first, then tile loop as today, then `epilogue_stages`. For the single-Stage case (no prologue/epilogue), behaviour identical to today. --- ## 8. Strict-FIFO RW hazard tracker ```python class _RwHazardTracker: def __init__(self): self._inflight: list[tuple[CompletionHandle, set[str]]] = [] # (completion, set of handle ids written or RW) def admit(self, env, cmd: CompositeCmd) -> Generator: """Block until no in-flight composite has overlapping RW.""" rw_ids = {h.id for h in cmd.rw_handles} # Strict FIFO: wait for ALL prior composites whose rw_set # intersects, in order. while True: blocking = [c for c, s in self._inflight if s & rw_ids] if not blocking: break # Wait on the *first* blocking composite's completion. yield blocking[0].done_event self._inflight.append((cmd.completion, rw_ids)) def retire(self, completion: CompletionHandle) -> None: self._inflight = [ (c, s) for c, s in self._inflight if c != completion ] ``` `PE_SCHEDULER` invokes `admit()` before queueing a CompositeCmd's tile plan and `retire()` upon `done_event.succeed()`. Strict FIFO: even when overlap is partial, the new composite waits for all prior overlapping composites to *fully* drain — simple, easy to reason about, easy to test. (RW-aware reorder deferred per ADR-0065 A4.) **Conservatism trade-off (ADR-0065 Consequences/Negative).** This rule is safe but may under-expose overlap. Concretely for decode opt2: next-tile's #1 (Q·Kᵀ) does not touch `(m, l, O)` and *could* in principle run while this-tile's #2 is still computing on those handles. Under strict FIFO, next-tile's #1 still gets dispatched (no rw overlap with prior in-flight) — so #1 cross-tile pipelining is preserved. However, if a future bench composed multiple #2-like recipes that share handles, FIFO would serialize them even when their compute could overlap. Acceptable for Phase 1; revisit when a real workload shows the gap. --- ## 9. Phased implementation plan (test-first per CLAUDE.md) Each phase is an independent Phase-1 → Phase-2 cycle per CLAUDE.md Part 1. | Phase | Deliverable | Gate | |---|---|---| | **P0 (= ADR-0064 Rev2)** | `logical_bytes` property + dispatch formula + topology override + size cap; one-time goldens regeneration | typical 1-OpSpec composite ≈ 43 ns; dispatch ratio test infrastructure in place | | **P1** | `CompositeCmd` flat-ops dataclass refactor + `OpSpec.operands: dict` + legacy `tl.composite` lowering preserved | every existing bench's op_log byte-equal pre/post | | **P2** | `tl_recipes.py` + `RECIPE_DESCRIPTORS["softmax_merge"]` + TLContext lowering for `prologue=[...]` | recipe lowering test (10 ops, rw_handles, addresses) passes | | **P3** | PE_SCHEDULER `_generate_plan` flat-ops + `PipelinePlan.prologue_stages` + DMA auto-insertion + `_feed_loop` extension | all existing bench Stage sequences unchanged; new MATH chain visible in op_log for opt2 mock | | **P4** | Strict-FIFO RW hazard tracker + `_RwHazardTracker` integration | two-composite RW conflict test serializes correctly | | **P5** | `_gqa_decode_long.py` opt2 variant; data-mode numeric parity check; K-before-V invariant check | opt2 matches opt3 within fp tolerance; no V DMA during prologue | | **P6** | Dispatch-ratio measurement: opt3 vs opt2 per-tile PE_CPU cycles + R sensitivity sweep | `opt3 > 2 × opt2` at default calibration (gate); model-expected ≈ 4.0× (informative); `opt2 < opt3` at all `R ∈ {0.25, 0.0625, 0.03125}` (ADR-0064 Test #9) | **As-built ordering:** the implementation reversed P0↔P1 — **P1 (flat-ops refactor) landed first**, then P0 (ADR-0064 Rev2 cost model), because `logical_bytes` is naturally defined on the flat shape (doing P1 first avoids a throwaway transitional definition). P1 is a pure, byte-equal refactor; P0 then swaps in the structural formula and regenerates goldens (in practice no golden churn occurred). P5 was split: **P5(B)** shipped the opt2 kernel + op_log dispatch measurement first (numeric parity deferred), then **D8 + N1–N4** (§15) built the data-mode numeric path. P6 closed the measurement. --- ## 10. Verification plan (concrete) Mirrors ADR-0065 §Test Requirements; grounded in SPEC R2/R5, ADR-0023/ 0025, ADR-0046, ADR-0054. **Refactor safety (P1):** - Parametrize over `tests/.py` baseline op_logs. - Pre/post refactor diff must be empty for *every* bench. - Failure → revert P1 atomically; do not land partial. **Recipe lowering (P2):** - Given `Sj=(8,64), m=(8,), l=(8,), O=(8,128)`, lower `tl.composite(prologue=[{"op":"softmax_merge", "s":Sj,...}], op="gemm", b=V_ref, out=O, epilogue=[{"op":"add", "other":O}])`. - Assert `len(cmd.ops) == 10`, `cmd.ops[0..7].kind == [rmax, max_elem, exp_diff, exp_diff, rsum, fma, mul_bcast, copy]`, `cmd.ops[8].kind == "gemm"`, `cmd.ops[9].kind == "add"`. - Assert `cmd.rw_handles == (m, l, O)`. - Assert all OpSpec addresses are concrete (no None). **Scheduler plan (P3):** - For the above CompositeCmd, assert PipelinePlan has 8 prologue_stages (MATH), tile loop with FETCH/GEMM/MATH(add), no epilogue_stages. - Assert no DMA_READ for prologue MATH stages (operands TCM-resident). - Assert one DMA_READ for V (`b.space == "hbm"`). - Assert no DMA_WRITE (`out.space == "tcm"`). **RW hazard (P4):** - Submit composite A (rw_handles=(O,)) then composite B (operands reading O). Assert B's Stages start only after A's done_event. **Numeric parity (P5):** - In data mode, run opt2 and opt3 for `(G=8, T_q=1, S_kv=64, C=1, P=1)`. - Assert `O_opt2 ≈ O_opt3` within `fp16` tolerance. **K-before-V invariant (P5):** - op_log inspection: for each tile's #2 CompositeCmd, no `memory/dma_read` events between #2 start and the GEMM compute begin that are not the V tile DMA itself. **Dispatch ratio (P6):** - For `S_kv=64, n_tiles=16`, measure total PE_CPU dispatch cycles for opt3 vs opt2 paths. - Assert `opt3 > 2 × opt2` (robust gate — calibration-independent). - Record observed ratio; the default-calibration model expects ≈ 4.0× (DDD §11). The recorded number is informative for performance tracking, not a test gate. - **Sensitivity sweep.** Repeat with `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte; assert opt2 < opt3 in all three. (ADR-0064 Test #9.) **Invariant guards (P2):** - Auto-bind conflict: `tl.composite(op="gemm", a=A_handle, prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}], ...)` → `pytest.raises(ValidationError, match="auto-bind conflict")`. - MATH operand TCM-only: `tl.composite(op="math", a=tl.ref(addr, shape), ...)` → `pytest.raises(ValidationError, match="MATH operand must be TCM")`. - Per-op logical_bytes summation: a synthetic composite with the same handle `O` in 3 OpSpecs + in `rw_handles` reports `logical_bytes` higher than a deduplicated equivalent by exactly `3 × 8` (operand references) + `8` (rw_handles entry) = 32 bytes. **Determinism (all phases):** - Identical inputs → identical op_log + latency (SPEC §0.1). --- ## 11. Performance model (expected) Per-tile PE_CPU dispatch cycles, ADR-0064 Rev2 default calibration (FIXED=40 cycles, R=0.0625 cycles/byte, 1 cycle = 1 ns at 1 GHz): ``` opt3: 10 cmds × 40 + total_bytes(≈232) × 0.0625 = 400 + 14.5 ≈ 414.5 ns opt2: 2 cmds × 40 + total_bytes(≈380) × 0.0625 = 80 + 23.75 ≈ 103.75 ns ratio = 414.5 / 103.75 ≈ 4.0× ``` The FIXED term dominates (~96% of the differential), reflecting ADR-0064's framing: the primary signal is command-count reduction. The byte term is a secondary refinement (visible: opt2 carries more total bytes than opt3, so its byte cost is higher; this is more than offset by 8× fewer FIXED payments). Per-tile engine timing dominates total latency; dispatch is one component. The ratio above is the *dispatch-only* relative cost. Total latency improvement depends on whether dispatch was on the critical path — measured per workload in P6. **R sensitivity** (ADR-0064 Test #9 — sanity at neighbouring R): | R (cycles/byte) | opt3 ns | opt2 ns | ratio | |---|---|---|---| | 0.25 (4 B/cy) | 458 | 175 | 2.6× | | 0.0625 (16 B/cy) ✓| 414.5 | 103.75 | 4.0× | | 0.03125 (32 B/cy) | 407.25 | 91.875 | 4.4× | `opt2 < opt3` at all three; ratio is monotonic in `1/R` (more FIXED-dominated as R decreases). **As-built (measured, `test_gqa_decode_opt2.py`).** The per-tile emitters give opt3=960 ns, opt2=184 ns → **5.22×** at the default calibration (higher than the model's ≈4.0× because the opt3 emitter includes the full online merge); R-sweep ratios **4.20× / 5.22× / 5.45×** for R∈{0.25,0.0625,0.03125}, monotonically increasing as R decreases. The gate is the calibration-robust `opt3 > 2 × opt2`. --- ## 12. Open items (status) Most ADR-0065 §Open review items map to specific implementation choices already documented in §§5–8. Live remaining items: | Item | Resolution venue | |---|---| | Non-identity `PrimaryOutSpec.transform` (e.g., `"trans"`) | Future recipe addition; current code uses string dispatch in step 3 of §6 | | `_gqa_decode_long.py` opt2 — TILE size sweep | Open (existing §B-3, §9 of ADR-0060) | | Larger composite descriptors approaching ADR-0064 D7 cap (1024 bytes) | Hard cap → `ValueError` at emit (no segmentation); kernel must split the recipe | | `scratch_scope` budget vs softmax_merge intermediates (G·TILE for P) | ADR-0063 sizing check at S_kv = 256K | | Scheduler plan-gen cost for large composites (ADR-0064 Open Item 2) | If P6 measurements show scheduler-bound behaviour near the cap, expose per-op-count `overhead_ns` on PE_SCHEDULER | --- ## 13. Risks | Risk | Likelihood | Mitigation | |---|---|---| | `CompositeCmd` refactor (P1) changes op_log for some bench | med | Parametrized regression over all benches; revert atomically | | `OpSpec.operands` migration (positional → dict) breaks epilogue site | med | Per-EPILOGUE_OPS migration table; existing epilogue tests cover | | `_RwHazardTracker` false-positive serializes too aggressively | low | Strict FIFO is the safe direction; RW-aware reorder is a future ADR | | ADR-0064 Rev2 and ADR-0065 P1 land in wrong order → stale goldens | med | PR ordering enforced; CI gate requires both to be green | | `softmax_merge` scratch budget exceeds `scratch_scope` at long context | med | ADR-0063 sizing check; if needed, tile_alignment promotion (future) | | PE_MATH op_kind label collision with new strings | low | Labels are opaque strings — no enum conflict; just add to allowed set if validated | --- ## 14. Glossary & references - **flat ops list** — `CompositeCmd.ops: tuple[OpSpec, ...]` with no head/epilogue field; placement determined by op position + scope. - **prologue (user API)** — `tl.composite(prologue=[...])` kwarg for recipe ops that lower into pre-GEMM KERNEL-scope OpSpecs. *HW command does not carry this concept.* - **recipe** — a named macro op (e.g., `softmax_merge`) that TLContext expands to multiple flat OpSpecs at emit time. Registry = `RECIPE_DESCRIPTORS`. - **primary_out** — recipe's implicit output (e.g., `P` for softmax_merge). Auto-bound to head GEMM's first matrix operand. - **handle-identity strict FIFO** — cross-composite serialization rule: a composite waits for all prior composites whose `rw_handles` intersect with its own to complete. - **logical_bytes** — HW-logical byte size of a PE command, derived from a structural rule (ADR-0064 D2), not from Python's `sys.getsizeof`. Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143` (composite tile plan), `pe_dma.py:45,89` (read channel), `pe_math.py` (op_kind opaque latency), `pe_cpu.py` (dispatch hook), `pe_commands.py` (CompositeCmd, OpSpec, EPILOGUE_OPS). ADRs: **0065** (this design), **0064 Revision 2** (dispatch cost), **0060** (GQA fused attention, §5.6 / §8 item 4 carve-out), **0063** (scratch_scope), **0046** (tl_context contract), **0042** (tile plan generators), **0014** (PE pipeline execution model). --- ## 15. Data-mode numerics (D8 + N1–N4) — as-built The original DDD assumed opt2's numeric parity "just works" via the existing DataExecutor. It does not: a composite in the builtin tiled path is recorded **only as latency TileTokens** (`op_kind="unknown"`, no addresses) — its matmul/MATH is never computed; the output address is never written. Making opt2's recipe compute required building the composite/recipe data-mode path. **D8 — output handle + output-space DMA (commit `e8d6c28`).** Prerequisite: a composite's output must be a real, addressable, chainable handle, and a TCM-resident output must NOT be DMA_WRITTEN (its high-bit scratch address crashes the DMA PA decoder). `tl.composite` returns the output `TensorHandle` (auto-TCM if `out` omitted); the tile loop's DMA_WRITE is gated on `out.space=="hbm"`. This is what lets opt2 *run* in data mode at all. **N1 — composite GEMM record (commit `4089e18`).** `op_log._extract_op_info (CompositeCmd)` scans `ops` for the GEMM (not `ops[0]` — a recipe's GEMM sits after the prologue) and emits one `composite_gemm` record (a/b/out addrs+spaces). `PE_SCHEDULER` records the composite's numeric op (no-op without an op_logger). `DataExecutor._execute_gemm` is best-effort: skips on (KeyError, ValueError) reading inputs (torch.empty / sharded operands). **N2 — recipe MATH ops (commit `3d4a3d4`).** `_compute_math` gains `rmax`/`rsum` (keepdims) + `max_elem`/`exp_diff`/`mul_bcast` (binary; numpy broadcasting covers `bcast_axis`). `tiling._math_stage` carries operand+output addrs/shapes/spaces+axis on the **recipe prologue** stages only; `op_log.record_end` promotes a MATH stage to `op_kind="math"` *only* when it carries `input_addrs` (legacy epilogue stages stay opaque → byte-equal). A latent P2 bug was fixed: reductions keep the reduced axis as size 1 (keepdims) so they broadcast against `m=(G,1)`. Verified: the 8 MATH ops produce m, l (full online-softmax update) and O (rescaled by corr) vs a numpy reference. **N3 — GEMM-epilogue accumulate (commit `5bbe293`).** The composite GEMM record is marked `accumulate=True` when an `add` epilogue targets the GEMM's own output; `_execute_gemm` then does `O = O_existing + a@b`. The composite's numeric op is recorded at **completion** (`_retire_on_done`), not dispatch, so its `t_start` sorts AFTER its prologue MATH ops → the GEMM reads the computed P and the rescaled O. Verified: full recipe → `O = O_old·corr + P@V` vs numpy. **N4 — opt2 parity + ordering (commit `e9d62b9`).** `composite()` awaits its prologue recipe operands (the cross-composite RAW: #2's `s` is #1's score tile) so #2 waits for #1. With N1–N3, opt2 computes a correct attention output e2e in data mode and matches opt3 within fp tolerance. **Scope:** the e2e opt2↔opt3 parity holds in the near-uniform-attention regime — kernbench's K reshape-as-transpose (opt3 docstring: correct for zero/symmetric inputs) makes opt2's per-sub-tile K interpretation differ from opt3's single-tile one for high-contrast inputs (a pre-existing kernbench limitation shared by both kernels, not the recipe). The **exact** recipe math is verified non-trivially by N2/N3 (`test_recipe_data_mode.py`).