ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
23a2bc3fb3fe7b19ab8cda1cff5c6062f2610ca6
kernbench2/docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md
T
ywkang 1a4ecb7aac gqa(ddd): DDD-0065 as-built content (prior commit only recorded the rename) 
The reconciliation table, corrected file/test lists, measured 5.22x ratio, and the new section 15 (data-mode numerics D8 + N1-N4) — the content the git-mv in 30a0451 left uncommitted.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-10 23:27:43 -07:00

34 KiB
Raw Blame History

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/<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, 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

@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

@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

@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). 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):

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 TensorHandles.
    • 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 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)

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:

@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

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/<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).

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

Reference in New Issue View Git Blame Copy Permalink