gqa(adr): ADR-0065 review fixes — KERNEL phase/repetition split, auto-bind conflict + MATH-TCM invariants, FIFO conservatism note, logical_bytes per-op summation

ADR-0065 review fixes:
- D3 semantics clarified: position determines *phase* (pre-loop / head /
  in-accumulation / per-output-tile / post-loop), scope determines
  *repetition*. KERNEL = "once per CompositeCmd invocation", not "once
  per kernel launch" — removes ambiguity around softmax_merge's KERNEL
  scope at different positions before/after GEMM.
- D4 MATH-TCM invariant: Phase 1 limits DMA auto-insertion to GEMM
  operands and head outputs; MATH operand with space="hbm" raises
  validation error at TLContext emit. Catches kernel misuse early; a
  MATH-with-HBM path is future work (explicit tl.load + MATH composite).
- D6 #6 NEW: auto-bind conflict — explicit GEMM `a=` + prologue
  primary_out simultaneously → validation error (prevents "which value
  wins" ambiguity).
- D6 #7 NEW: MATH operand TCM-only restatement.
- Consequences/Negative: strict-FIFO conservatism note — safe but may
  under-expose composite-level overlap when handle-sharing composites
  could in principle run in parallel.
- Test Requirements #10, #11: auto-bind conflict + MATH-HBM validation
  error tests.

ADR-0064 Revision 2 clarification:
- D2 counting rule: per-op summation, no deduplication. Same handle
  appearing in multiple OpSpecs or in rw_handles is counted
  independently — matches HW reality (each descriptor field is a
  distinct address slot, rw_handles is separate metadata block).
  Example walked through (mul_bcast with O in src_a + out + rw_handles).

DDD-0065 follow-on:
- §3.2: new test file test_tl_composite_validation.py covers auto-bind
  conflict, MATH-TCM, GEMM-count invariants.
- §4.3: phase vs repetition split aligned with ADR-0065 D3.
- §6 lowering pipeline: step 5 adds auto-bind conflict check; step 8
  NEW — operand space validation; step 9 (size cap); step 10 (emit).
- §8 strict-FIFO conservatism note: explains why next-tile #1 cross-tile
  pipelining still works under FIFO (no rw overlap), and where the
  conservatism would hurt (future handle-sharing recipes).
- §10 verification: invariant-guard tests (auto-bind conflict, MATH
  TCM, per-op summation accuracy).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-10 16:29:39 -07:00
parent 95cecccd8e
commit a8d04750e6
5 changed files with 197 additions and 46 deletions
@@ -139,6 +139,17 @@ def _extra_bytes(v) -> int:
(`DmaReadCmd`, `MathCmd` 등도 같은 룰 — dataclass 당 property 1 개, (`DmaReadCmd`, `MathCmd` 등도 같은 룰 — dataclass 당 property 1 개,
약 3 줄.) 약 3 줄.)
**카운팅 룰 — per-op 합산, dedup 없음.** `OpSpec` 안의 각 `TensorHandle`
참조는 독립적으로 카운트 — 같은 핸들이 여러 OpSpec 에 등장하거나
`rw_handles` 에도 있어도 마찬가지. `rw_handles` block 은 cross-composite
해저드 추적 (ADR-0065 D6.3) 의 메타데이터로 `ops` 의 operand 참조와
**별도** 카운트. 중복 제거 없음. HW 와 일치: descriptor 가 각 operand
slot 을 독립 address field 로 인코딩하고 dispatcher 가 `rw_handles`
별도 metadata block 으로 추적. 예: `OpSpec(kind="mul_bcast", operands=
{"src_a": O, "src_b": corr}, out=O)` 는 핸들 `O`*두 번* 카운트
(`src_a` 한 번, `out` 한 번); 만약 `O` 가 둘러싼 CompositeCmd 의
`rw_handles` 에도 있으면 세 번째 카운트.
### D3. Default — 일반 composite ≈ 43 ns 기준 ### D3. Default — 일반 composite ≈ 43 ns 기준
Anchor: **DMA-staged GEMM 경로용 단일-OpSpec composite** Anchor: **DMA-staged GEMM 경로용 단일-OpSpec composite**
@@ -150,6 +150,19 @@ def _extra_bytes(v) -> int:
(Identical rule for `DmaReadCmd`, `MathCmd`, etc. — one property per (Identical rule for `DmaReadCmd`, `MathCmd`, etc. — one property per
dataclass, ~3 lines each.) dataclass, ~3 lines each.)
**Counting rule — per-op summation, no deduplication.** Each
`TensorHandle` reference inside an `OpSpec` is counted independently,
even when the same handle appears in multiple OpSpecs or in
`rw_handles`. The `rw_handles` block is metadata for cross-composite
hazard tracking (ADR-0065 D6.3) and is counted separately from
operand references in `ops`. There is no deduplication. This matches
HW reality: the descriptor encodes each operand slot as an independent
address field, and the dispatcher tracks `rw_handles` as a distinct
metadata block. Example: `OpSpec(kind="mul_bcast", operands={"src_a":
O, "src_b": corr}, out=O)` counts handle `O` *twice* (once for
`src_a`, once for `out`); if `O` is also in the enclosing
CompositeCmd's `rw_handles`, it is counted a third time.
### D3. Defaults — anchored on a typical composite ≈ 43 ns ### D3. Defaults — anchored on a typical composite ≈ 43 ns
Anchor: a **single-OpSpec composite for a DMA-staged GEMM path** Anchor: a **single-OpSpec composite for a DMA-staged GEMM path**
@@ -94,19 +94,26 @@ named operand 로 PE_SCHEDULER 가 (positional 컨벤션 없이) 엔진 port 에
### D3. Position + scope 가 tile-loop 배치 결정 ### D3. Position + scope 가 tile-loop 배치 결정
**Semantics 분리.** *Position***phase** (tile-loop 수명주기 어디서
실행되는지) 결정; *scope* 가 그 phase 안에서의 **repetition** 결정.
`Scope.KERNEL` 의미는 "**`CompositeCmd` invocation 당 1 회**" — kernel
launch 당 1 회 아님. 그 1 회 실행이 op 의 시퀀스 내 *position* 에서
발생. 같은 KERNEL 값이 OpSpec 이 GEMM op 의 앞/위치/뒤에 있느냐에 따라
*phase* 가 달라짐.
PE_SCHEDULER 가 `cmd.ops` 에서 GEMM op 검색 (composite 당 ≤ 1, D6 참조). PE_SCHEDULER 가 `cmd.ops` 에서 GEMM op 검색 (composite 당 ≤ 1, D6 참조).
인덱스 `g`: 인덱스 `g`:
| OpSpec position | scope | plan 내 배치 | | OpSpec position | scope | Phase | plan 내 배치 |
|---|---|---| |---|---|---|---|
| `0 .. g-1` | KERNEL | pre-loop stages (tile loop 진입 전 1 회 실행) | | `0 .. g-1` | KERNEL | pre-tile-loop | tile loop 진입 전 1 회 |
| `g` | KERNEL | head GEMM (`extra["m"], extra["k"], extra["n"]` 로 tile loop drive) | | `g` | KERNEL | head GEMM | `extra["m"], extra["k"], extra["n"]` 로 tile loop drive |
| `g+1 .. ` | K_TILE | per K-tile epilogue (누산 loop 안) | | `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (K 누산 loop 안) |
| `g+1 .. ` | OUTPUT_TILE | per (m, n) tile epilogue (K 누산 후) | | `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (K 누산 후) |
| `g+1 .. ` | KERNEL | post-loop stages (tile loop 종료 후 1 회) | | `g+1 .. ` | KERNEL | post-tile-loop | tile loop 종료 후 1 회 |
GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope
직렬로 처리. 직렬로 처리 — phase 가 "순서대로 single-shot" 으로 collapse.
### D4. PE_SCHEDULER 가 operand `space` 보고 DMA 자동 삽입 ### D4. PE_SCHEDULER 가 operand `space` 보고 DMA 자동 삽입
@@ -119,8 +126,13 @@ PE_SCHEDULER 가 GEMM 의 `operands` 와 `out` 검사. 각각:
PE_SCHEDULER 가 핸들에서 추론**. `tl.ref(addr, shape)``space="hbm"`; PE_SCHEDULER 가 핸들에서 추론**. `tl.ref(addr, shape)``space="hbm"`;
`tl.zeros` / `tl.full` / `tl.load` 결과 → `space="tcm"`. `tl.zeros` / `tl.full` / `tl.load` 결과 → `space="tcm"`.
비 GEMM op (prologue/epilogue 의 MATH chain) 은 TCM-resident operand **비 GEMM op 은 TCM-only (Phase 1 한계, TLContext emit 시 강제).**
(커널이 `tl.zeros` / `tl.full` 로 할당); DMA 삽입 없음. 모든 non-GEMM (MATH) OpSpec 의 모든 operand 는 **반드시**
`space == "tcm"`. 커널이 `space == "hbm"` 핸들을 MATH operand 로 전달하면
TLContext 가 emit 시 validation error. Phase 1 에서 DMA 자동 삽입을
GEMM operand + head 출력 핸들로 제한; MATH-with-HBM 경로는 미래 확장
(명시적 `tl.load` + MATH composite 분리, 또는 새 "DMA-as-prologue" recipe
variant). 이 invariant 는 D6 #7 에 재기술.
### D5. RECIPE_DESCRIPTORS — TLContext 내부, `pe_commands.py` 가 아님 ### D5. RECIPE_DESCRIPTORS — TLContext 내부, `pe_commands.py` 가 아님
@@ -189,6 +201,14 @@ op="gemm", ...)` 에서:
로 lowering. 로 lowering.
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER 는 RECIPE_DESCRIPTORS import 5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER 는 RECIPE_DESCRIPTORS import
안 함, "softmax_merge" 의 존재 모름, recipe 별 코드 분기 없음. 안 함, "softmax_merge" 의 존재 모름, recipe 별 코드 분기 없음.
6. **암묵 operand 대체 없음 (auto-bind 충돌).** prologue recipe 가
`primary_out` 선언하고 head op 의 auto-bind 대상 operand (예: GEMM
`a`) 도 *커널이 명시적* 으로 제공하면, TLContext 가 emit 시 validation
error. 커널은 recipe 가 바인딩하게 두거나 (operand 생략) **혹은**
명시적으로 지정 (prologue 생략, 또는 `primary_out` 없는 recipe 사용)
해야 함. "어느 값이 이기느냐" 의 모호함 방지.
7. **MATH operand TCM-only.** D4 의 재기술: 모든 비 GEMM OpSpec 의 모든
operand 는 `space == "tcm"`. Phase 1 에서 TLContext emit 시 강제.
### D7. Boundary 요약 (compiler vs scheduler vs engine) ### D7. Boundary 요약 (compiler vs scheduler vs engine)
@@ -268,6 +288,13 @@ scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect
auto-bind. ~80120 LOC. auto-bind. ~80120 LOC.
- PE_SCHEDULER 의 `_generate_plan` 에 position-기반 스캔 + DMA 자동 삽입 - PE_SCHEDULER 의 `_generate_plan` 에 position-기반 스캔 + DMA 자동 삽입
+ strict-FIFO RW tracker. ~60100 LOC. + strict-FIFO RW tracker. ~60100 LOC.
- **strict-FIFO RW 해저드 추적이 안전하지만 보수적.** 두 composite 가
`rw_handle` 을 공유할 때 신규는 모든 이전 겹치는 composite 가 *완전히
완료* 될 때까지 대기 — 원칙적으로 overlap 가능한 경우에도
(예: 다음-tile #1 = Q·Kᵀ 가 m/l/O 안 만지므로 현재-tile 의 #2 와
overlap 가능하지만, 같은 핸들을 만지는 후속 path 도 FIFO 로 직렬화).
더 똑똑한 scheduler 대비 **composite-level overlap 을 과소 노출**.
RW-aware reorder (A4) 가 연기된 개선.
## Open review items ## Open review items
@@ -311,9 +338,17 @@ scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect
8. **GEMM-count invariant.** GEMM OpSpec 두 개 가진 composite 가 8. **GEMM-count invariant.** GEMM OpSpec 두 개 가진 composite 가
TLContext emit 시 validation error. TLContext emit 시 validation error.
9. **Composite 크기 cap 안에 (ADR-0064 D7).** Decode opt2 의 `#2` 9. **Composite 크기 cap 안에 (ADR-0064 D7).** Decode opt2 의 `#2`
composite (10 ops, ~310 logical bytes) 가 default composite (10 ops, ~322 logical bytes — operand 핸들의 per-op 합산)
`MAX_COMPOSITE_LOGICAL_BYTES=1024` 안에 편안히 — segmentation 없음. 가 default `MAX_COMPOSITE_LOGICAL_BYTES=1024` 안에 편안히 —
recipe lowering 테스트가 `cmd.logical_bytes < 1024` 단언. segmentation 없음. recipe lowering 테스트가 `cmd.logical_bytes
< 1024` 단언.
10. **Auto-bind 충돌.** `tl.composite(op="gemm", a=A_handle,
prologue=[{"op": "softmax_merge", ...}], ...)` 처럼 softmax_merge 가
`primary_out` 선언하는데 `a` 도 명시되면 emit 시 validation error
(D6.6).
11. **MATH operand TCM invariant.** `tl.composite(op="math",
a=tl.ref(addr, shape), ...)` (HBM-resident 핸들을 MATH operand 로
전달) 가 emit 시 validation error (D6.7).
## Dependencies ## Dependencies
@@ -103,19 +103,27 @@ takes `operands["a"]`, `operands["b"]`) without positional convention.
### D3. Position + scope determines tile-loop placement ### D3. Position + scope determines tile-loop placement
**Semantics split.** *Position* determines the **phase** (where in the
tile-loop lifecycle the op fires); *scope* determines the **repetition**
within that phase. `Scope.KERNEL` means "**once per `CompositeCmd`
invocation**", not "once per kernel launch" — its single firing happens
at the op's position in the sequence. The same KERNEL value carries
different *phases* depending on whether the OpSpec sits before, at, or
after the GEMM op.
PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see
D6). Calling its index `g`: D6). Calling its index `g`:
| OpSpec position | scope | Placement in plan | | OpSpec position | scope | Phase | Placement in plan |
|---|---|---| |---|---|---|---|
| `0 .. g-1` | KERNEL | pre-loop stages (run once before tile loop) | | `0 .. g-1` | KERNEL | pre-tile-loop | run once before tile loop |
| `g` | KERNEL | head GEMM (drives tile loop with `extra["m"], extra["k"], extra["n"]`) | | `g` | KERNEL | head GEMM | drives tile loop with `extra["m"], extra["k"], extra["n"]` |
| `g+1 .. ` | K_TILE | per K-tile epilogue (inside accumulation loop) | | `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (inside K accumulation loop) |
| `g+1 .. ` | OUTPUT_TILE | per (m, n) tile epilogue (after K accumulation) | | `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (after K accumulation) |
| `g+1 .. ` | KERNEL | post-loop stages (run once after tile loop) | | `g+1 .. ` | KERNEL | post-tile-loop | run once after tile loop |
A composite with no GEMM op (e.g., MATH-only composite) treats all ops A composite with no GEMM op (e.g., MATH-only composite) treats all ops
as KERNEL-scope sequential. as KERNEL-scope sequential — phase collapses to "single-shot in order".
### D4. PE_SCHEDULER auto-inserts DMAs from operand `space` ### D4. PE_SCHEDULER auto-inserts DMAs from operand `space`
@@ -130,9 +138,15 @@ DMA cmds for composite operands; PE_SCHEDULER infers them from
handles**. `tl.ref(addr, shape)` returns a handle with `space="hbm"`; handles**. `tl.ref(addr, shape)` returns a handle with `space="hbm"`;
`tl.zeros` / `tl.full` / `tl.load` outputs return `space="tcm"`. `tl.zeros` / `tl.full` / `tl.load` outputs return `space="tcm"`.
Non-GEMM ops (MATH chain in prologue / epilogue) have TCM-resident **Non-GEMM ops are TCM-only (Phase 1 limit, enforced at TLContext emit).**
operands (kernel allocates via `tl.zeros` / `tl.full`); no DMA inserted Every operand of every non-GEMM (MATH) OpSpec **must** have
for them. `space == "tcm"`. If a kernel passes a handle with `space == "hbm"` as a
MATH operand, TLContext raises a validation error at emit time. This
keeps DMA auto-insertion bounded to GEMM operands and head output
handles in Phase 1; a MATH-with-HBM path would require a future
extension (split into explicit `tl.load` + MATH composite, or a new
"DMA-as-prologue" recipe variant). This invariant is also restated in
D6 #7.
### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py` ### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py`
@@ -211,6 +225,16 @@ sees only the flat ops list.
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER does not import 5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER does not import
RECIPE_DESCRIPTORS, does not know that "softmax_merge" exists, and RECIPE_DESCRIPTORS, does not know that "softmax_merge" exists, and
has no recipe-specific code branches. has no recipe-specific code branches.
6. **No implicit operand replacement (auto-bind conflict).** If a
prologue recipe declares a `primary_out` AND the head op's auto-bind
target operand (e.g., GEMM `a`) is *also* provided by the kernel
explicitly, TLContext raises a validation error at emit time. The
kernel must either let the recipe bind (omit the operand) **or**
specify it explicitly (omit the prologue, or use a recipe without
`primary_out`). This prevents the ambiguity of "which value wins".
7. **MATH operands TCM-only.** Restatement of D4: every operand of a
non-GEMM OpSpec must have `space == "tcm"`. Phase 1 enforced at
TLContext emit.
### D7. Boundary summary (compiler vs scheduler vs engine) ### D7. Boundary summary (compiler vs scheduler vs engine)
@@ -297,6 +321,15 @@ input from #1's epilogue scratch — circular tile-loop dependency.
primary_out auto-bind. ~80120 LOC. primary_out auto-bind. ~80120 LOC.
- PE_SCHEDULER's `_generate_plan` gets a position-based scan + DMA - PE_SCHEDULER's `_generate_plan` gets a position-based scan + DMA
auto-insertion + strict-FIFO RW tracker. ~60100 LOC. auto-insertion + strict-FIFO RW tracker. ~60100 LOC.
- **Strict-FIFO RW hazard tracking is safe but conservative.** When
two composites share any `rw_handle`, the new one waits for *all*
prior overlapping composites to fully complete — even when their
compute could in principle overlap (e.g., next-tile #1 = Q·Kᵀ does
not touch m/l/O and could overlap with this-tile's #2, but FIFO
still serializes any path that touches the same handles down the
line). This **may under-expose composite-level overlap** relative to
a smarter scheduler. RW-aware reorder (A4) is the deferred
improvement.
## Open review items ## Open review items
@@ -344,9 +377,17 @@ input from #1's epilogue scratch — circular tile-loop dependency.
8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a 8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a
validation error at TLContext emit. validation error at TLContext emit.
9. **Within composite size cap (ADR-0064 D7).** Decode opt2's `#2` 9. **Within composite size cap (ADR-0064 D7).** Decode opt2's `#2`
composite (10 ops, ~310 logical bytes) fits comfortably under the composite (10 ops, ~322 logical bytes — per-op summation of operand
default `MAX_COMPOSITE_LOGICAL_BYTES=1024` — no segmentation. The handles) fits comfortably under the default
recipe lowering test asserts `cmd.logical_bytes < 1024`. `MAX_COMPOSITE_LOGICAL_BYTES=1024` — no segmentation. The recipe
lowering test asserts `cmd.logical_bytes < 1024`.
10. **Auto-bind conflict.** `tl.composite(op="gemm", a=A_handle,
prologue=[{"op": "softmax_merge", ...}], ...)` where softmax_merge
declares `primary_out` raises a validation error at emit time
(D6.6).
11. **MATH operand TCM invariant.** `tl.composite(op="math",
a=tl.ref(addr, shape), ...)` (passing an HBM-resident handle as a
MATH operand) raises a validation error at emit time (D6.7).
## Dependencies ## Dependencies
@@ -92,7 +92,8 @@ without importing `tl_recipes`.
| `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_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_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_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 45 ns; topology override path; `PeCpuOverheadCmd` bypass | | `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/attention/test_gqa_decode_opt2.py` | opt2 numeric parity with opt3; K-before-V DMA invariant; per-tile dispatch ratio measurement |
--- ---
@@ -147,15 +148,26 @@ Backward compat: existing `operands: tuple[Any, ...]` use sites
h}` / `operands={"src": h}` per a small lookup in h}` / `operands={"src": h}` per a small lookup in
`pe_commands.EPILOGUE_OPS`. `pe_commands.EPILOGUE_OPS`.
### 4.3 `Scope` semantics (position-based for KERNEL) ### 4.3 `Scope` semantics — phase vs repetition (ADR-0065 D3)
`Scope` keeps its three values (`KERNEL`, `K_TILE`, `OUTPUT_TILE`). `Scope` keeps its three values (`KERNEL`, `K_TILE`, `OUTPUT_TILE`).
What changes: `KERNEL` no longer means "once anywhere"; it now means What changes is the **semantics split**:
"execute at this op's position in the sequence — once". Combined with - **Position** determines the *phase* (where in the tile-loop lifecycle
position relative to the GEMM op: the op fires — pre-loop / in-loop / post-loop).
- KERNEL ops *before* GEMM → pre-tile-loop (single shot) - **Scope** determines the *repetition* within the phase
- KERNEL ops *after* GEMM → post-tile-loop (single shot) (KERNEL = once per `CompositeCmd` invocation; K_TILE = per K-tile
- K_TILE / OUTPUT_TILE ops *after* GEMM → inside tile loop as before 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
--- ---
@@ -258,9 +270,14 @@ Lowering steps in `TLContext.composite()`:
- Build `OpSpec(kind=op_kind, scope=Scope.KERNEL, operands={...}, - Build `OpSpec(kind=op_kind, scope=Scope.KERNEL, operands={...},
extra={...}, out=...)`. extra={...}, out=...)`.
- `out` is the handle for `dst` slot; `operands` is `{src*: handle}`. - `out` is the handle for `dst` slot; `operands` is `{src*: handle}`.
5. **Auto-bind head GEMM.** GEMM OpSpec built from `op="gemm"`, 5. **Auto-bind head GEMM (with conflict check, ADR-0065 D6.6).** GEMM
`operands={"a": P_handle, "b": V_handle}`, `out=<output handle>`, OpSpec built from `op="gemm"`, `operands={"a": P_handle, "b":
`extra={"m": G, "k": TILE, "n": d}`. 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 6. **Build epilogue OpSpecs.** Each `epilogue=[{...}]` entry maps to
`OpSpec(kind=..., scope=Scope.OUTPUT_TILE, operands={...}, ...)`. `OpSpec(kind=..., scope=Scope.OUTPUT_TILE, operands={...}, ...)`.
The `add` epilogue specifically gets `operands={"other": O_handle}` The `add` epilogue specifically gets `operands={"other": O_handle}`
@@ -268,12 +285,23 @@ Lowering steps in `TLContext.composite()`:
7. **Compute `rw_handles`.** Collect all handles tagged RW in prologue 7. **Compute `rw_handles`.** Collect all handles tagged RW in prologue
recipes + the head GEMM's `out` (if accumulating) + any RW epilogue recipes + the head GEMM's `out` (if accumulating) + any RW epilogue
targets. For softmax_merge: `rw_handles = (m, l, O)`. targets. For softmax_merge: `rw_handles = (m, l, O)`.
8. **Check composite size cap (ADR-0064 D7).** Compute 8. **Operand space validation (ADR-0065 D6.7).** For each generated
`cmd.logical_bytes`; if > `MAX_COMPOSITE_LOGICAL_BYTES` (default OpSpec in the about-to-emit `ops`:
1024), segment deterministically by op index. softmax_merge's 10-op - If `kind == "gemm"`: operands may be HBM (auto-DMA) or TCM. Output
composite (~310 bytes) is well under the cap → single segment, no handle may be HBM (auto-DMA_WRITE) or TCM.
split for decode opt2. - If `kind != "gemm"` (MATH): **every** `operands.values()` handle
9. **Emit `CompositeCmd`.** `ops = (8 MATH + 1 GEMM + 1 MATH(add))`, 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`. `rw_handles = (m, l, O)`, `completion = new`.
**Legacy lowering.** Pre-existing **Legacy lowering.** Pre-existing
@@ -384,6 +412,17 @@ overlap is partial, the new composite waits for all prior overlapping
composites to *fully* drain — simple, easy to reason about, easy to composites to *fully* drain — simple, easy to reason about, easy to
test. (RW-aware reorder deferred per ADR-0065 A4.) 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) ## 9. Phased implementation plan (test-first per CLAUDE.md)
@@ -392,7 +431,7 @@ Each phase is an independent Phase-1 → Phase-2 cycle per CLAUDE.md Part 1.
| Phase | Deliverable | Gate | | Phase | Deliverable | Gate |
|---|---|---| |---|---|---|
| **P0 (= ADR-0064 Rev2)** | `logical_bytes` property + dispatch formula + topology override; one-time goldens regeneration | typical composite ≈ 45 ns; dispatch ratio test infrastructure in place | | **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 | | **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 | | **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 | | **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 |
@@ -453,6 +492,18 @@ Mirrors ADR-0065 §Test Requirements; grounded in SPEC R2/R5, ADR-0023/
- **Sensitivity sweep.** Repeat with `R ∈ {0.25, 0.0625, 0.03125}` - **Sensitivity sweep.** Repeat with `R ∈ {0.25, 0.0625, 0.03125}`
cycles/byte; assert opt2 < opt3 in all three. (ADR-0064 Test #9.) 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):** **Determinism (all phases):**
- Identical inputs → identical op_log + latency (SPEC §0.1). - Identical inputs → identical op_log + latency (SPEC §0.1).