30a0451481
Updated the detailed design to match the shipped implementation: an as-built reconciliation table (P1-first, pinned-based DMA, D7 hard-cap-error, D8 output handle, prologue 1-stage-tile feed, D6.7 prologue-scoped), corrected file/test lists, the measured 5.22x dispatch ratio, and a new section 15 documenting the data-mode numeric subsystem (D8 + N1-N4) the original DDD omitted. Status Draft->Accepted; moved docs/adr-proposed -> docs/ddd/. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
608 lines
28 KiB
Markdown
608 lines
28 KiB
Markdown
# Detailed Design Document — Flat-ops `CompositeCmd` + `softmax_merge` recipe (ADR-0065)
|
||
|
||
**Status:** Draft (companion to ADR-0065, Proposed)
|
||
**Scope:** implementation-ready plan for the `CompositeCmd` flat-ops
|
||
restructure, the `softmax_merge` recipe, and the decode opt2 variant in
|
||
`_gqa_decode_long.py`.
|
||
**Audience:** reviewer / implementer resuming after the ADR-0065 design
|
||
round. Read **ADR-0065 first** (boundary D5 + D7 set the layering), and
|
||
**ADR-0064 Revision 2** (dispatch cost model that opt2's win depends on).
|
||
|
||
---
|
||
|
||
## 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/<existing>` 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, rw_handles, pe_prefix)` replaces `generate_gemm_plan` / `generate_math_plan` (or wraps them). DMA auto-insertion from `handle.space`. Epilogue `add` stage emits `other_handle` in params (GAP 9 fix). | 0065 D3/D4 |
|
||
| `src/kernbench/components/builtin/pe_cpu.py` | dispatch cost from `cmd.logical_bytes` via ADR-0064 Rev2 formula; topology config override. | 0064 D1/D4 |
|
||
| topology yaml (e.g., `topologies/llama70b_4sip.yaml`) | optional `pe_cost_model:` block under PE node attrs. | 0064 D4 |
|
||
| `src/kernbench/benches/_gqa_decode_long.py` | opt2 variant — `tl.composite(prologue=[{"op":"softmax_merge", …}], op="gemm", b=V_ref, out=O, epilogue=[{"op":"add", "other":O}])`. | 0065 §Migration P3 |
|
||
|
||
### 3.2 New tests
|
||
|
||
| File | Covers |
|
||
|---|---|
|
||
| `tests/test_composite_flat_ops_refactor.py` | every existing bench's op_log byte-equal pre/post refactor (parametrized over bench list) |
|
||
| `tests/test_tl_recipe_softmax_merge.py` | recipe lowering — given inputs, asserts CompositeCmd has 10 ops in expected order + `rw_handles == (m, l, O)` |
|
||
| `tests/test_pe_scheduler_strict_fifo.py` | two consecutive composites sharing rw_handles complete in dispatch order; no stage interleave |
|
||
| `tests/test_pe_scheduler_auto_dma.py` | `space=hbm` operand → DMA_READ stage; `space=tcm` → none; `out.space=hbm` → DMA_WRITE |
|
||
| `tests/test_pe_cost_model.py` | ADR-0064 Rev2 — typical 1-GEMM composite at ~43 ns; topology override path; `PeCpuOverheadCmd` bypass; per-op summation rule (ADR-0064 D2) |
|
||
| `tests/test_tl_composite_validation.py` | auto-bind conflict (D6.6); MATH operand TCM-only (D6.7); GEMM-count > 1 (D6.1) |
|
||
| `tests/attention/test_gqa_decode_opt2.py` | opt2 numeric parity with opt3; K-before-V DMA invariant; per-tile dispatch ratio measurement |
|
||
|
||
---
|
||
|
||
## 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=<P slot>,
|
||
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=<output handle>`, `extra={"m": G, "k": TILE, "n":
|
||
d}`. **Validation:** if the kernel passed `a=<some handle>`
|
||
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 space validation (ADR-0065 D6.7).** For each generated
|
||
OpSpec in the about-to-emit `ops`:
|
||
- If `kind == "gemm"`: operands may be HBM (auto-DMA) or TCM. Output
|
||
handle may be HBM (auto-DMA_WRITE) or TCM.
|
||
- If `kind != "gemm"` (MATH): **every** `operands.values()` handle
|
||
must have `space == "tcm"`. HBM operand → validation error.
|
||
This catches kernel misuse (passing `tl.ref(...)` into MATH ops)
|
||
early. softmax_merge's 8 MATH ops only ever reference TCM-resident
|
||
handles (kernel-allocated `m, l, O` + scratch `m_loc, m_new, corr,
|
||
P, l_loc`) — invariant naturally holds.
|
||
9. **Check composite size cap (ADR-0064 D7).** Compute
|
||
`cmd.logical_bytes` using the per-op summation rule (ADR-0064 D2);
|
||
if > `MAX_COMPOSITE_LOGICAL_BYTES` (default 1024), segment
|
||
deterministically by op index. softmax_merge's 10-op composite
|
||
(~322 bytes per-op-summed) is well under the cap → single segment,
|
||
no split for decode opt2.
|
||
10. **Emit `CompositeCmd`.** `ops = (8 MATH + 1 GEMM + 1 MATH(add))`,
|
||
`rw_handles = (m, l, O)`, `completion = new`.
|
||
|
||
**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) |
|
||
|
||
P0 must land first (separate ADR). P1 is pure refactor (safest). P2–P4
|
||
build the new path without touching opt3. P5 enables opt2. P6 closes
|
||
the measurement loop.
|
||
|
||
---
|
||
|
||
## 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/<existing bench>.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).
|
||
|
||
---
|
||
|
||
## 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) | Monitored in P6; segmentation kicks in deterministically if exceeded |
|
||
| `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).
|