From 1a4ecb7aacba969c51ad7fb97c6d5ee0dd8fb571 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 23:27:43 -0700 Subject: [PATCH] gqa(ddd): DDD-0065 as-built content (prior commit only recorded the rename) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- ...composite-softmax-merge-detailed-design.md | 172 ++++++++++++++---- 1 file changed, 132 insertions(+), 40 deletions(-) diff --git a/docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md b/docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md index 33d349d..9c42e74 100644 --- a/docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md +++ b/docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md @@ -1,12 +1,36 @@ # 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). +**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.** --- @@ -81,22 +105,27 @@ without importing `tl_recipes`. | `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 | +| `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` | 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 | +| `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 | --- @@ -287,24 +316,23 @@ Lowering steps in `TLContext.composite()`: 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`. +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}])` @@ -441,9 +469,14 @@ Each phase is an independent Phase-1 → Phase-2 cycle per CLAUDE.md Part 1. | **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. +**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. --- @@ -547,6 +580,13 @@ path — measured per workload in P6. `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) @@ -558,7 +598,7 @@ already documented in §§5–8. Live remaining items: |---|---| | 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 | +| 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 | @@ -605,3 +645,55 @@ 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`).