gqa(adr): ADR-0065 Proposed -> Accepted
Flat-ops CompositeCmd + softmax_merge recipe is fully implemented and tested: structural restructure (P1-P4), dispatch-cost measurement (P5b/P6, 5.22x per-tile win), and data-mode numerics (D8 output handle/space + N1-N4 composite GEMM / recipe MATH / accumulate / opt2 parity). git mv EN -> docs/adr/, KO -> docs/adr-ko/ (verify_adr_lang_pairs OK). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,440 +0,0 @@
|
||||
# ADR-0065: 평평한 ops `CompositeCmd` + 첫 stateful recipe (`softmax_merge`)
|
||||
|
||||
## Status
|
||||
|
||||
Proposed (verification gated on ADR-0064 Revision 2 land)
|
||||
|
||||
> **ADR-0060** (AHBM GQA Fused Attention) 의 보조 ADR. §5.6 / §8 item 4 의
|
||||
> carve-out 을 정확히 구현: decode opt2 = (#1 기존 Q·Kᵀ GEMM composite) +
|
||||
> (#2 softmax + P·V + online-softmax merge 를 담은 단일 composite). ADR-0060
|
||||
> §5.6 의 "ex_composite" / "flash_pv_merge" 명칭은 *retire* — 본 ADR 이
|
||||
> 정식 모양을 확정: 기존 `tl.composite` 입구를 그대로 사용하면서 두 가지
|
||||
> 구조적 추가 — (a) 평평한 `ops` 튜플, (b) MATH micro-op 시퀀스로
|
||||
> 펼쳐지는 첫 stateful recipe `softmax_merge`.
|
||||
|
||||
## Context
|
||||
|
||||
### ADR-0060 가 이 ADR 에 남긴 것
|
||||
|
||||
ADR-0060 §5.6 는 opt3 (software pipelining) 를 *지금 ship* 으로 권고,
|
||||
opt2 는 — per-tile dispatch 적음, K-before-V DMA priority,
|
||||
**ADR-0064 의 cost model 이후에만 측정 가능** — 후속으로 carve out.
|
||||
|
||||
§8 item 4 가 carve-out 의 sizing: "재방문 시 **두 개** composite 로 분할 —
|
||||
`#1` = Q·Kᵀ (기존 composite + `scale`), `#2` = softmax + P·V + 온라인-softmax
|
||||
누산기 merge". 본 ADR 이 `#2` 를 구현.
|
||||
|
||||
### composite 가 구조 변경이 필요한 이유
|
||||
|
||||
현재 `CompositeCmd` 는 `(op, a, b, out_addr, ops=(head, *epi))` 모양 + 암묵 컨벤션:
|
||||
- `op` 가 head 엔진 선택 (gemm | math)
|
||||
- `ops[0]` 는 head OpSpec, `ops[1:]` 는 epilogue OpSpec 으로 head *이후*
|
||||
per-output-tile 또는 per-K-tile 에 실행 (scope 결정)
|
||||
|
||||
decode opt2 는 MATH op 이 head GEMM **이전** 에 실행 (`m, l, O` 의
|
||||
online-softmax 갱신 + GEMM 입력 `P` emit) 이 필요. 현 모양엔 자리 없음.
|
||||
|
||||
이전에 검토 후 기각된 경로:
|
||||
- softmax_merge 를 #1 (Q·Kᵀ) 의 *epilogue* 로 두기 — `Sj` 를 읽고 *다음*
|
||||
GEMM (#2 P·V) 가 소비할 `P` 를 쓰는 구조라 순환.
|
||||
- 단일 mixed-engine recipe `flash_pv_merge` 가 MATH + GEMM + MATH 흡수 —
|
||||
engine boundary 위반, PE_SCHEDULER 복잡화.
|
||||
|
||||
깔끔한 모양: **command 포맷에서 head/epilogue 구분 제거**. `ops` 를 평평한
|
||||
순서 튜플로, 각 op 의 `scope` + position 이 tile-loop 배치 결정. "prologue"
|
||||
개념은 *사용자 API* 의 `tl.composite(...)` 에 ergonomic 으로 남되 —
|
||||
**HW command 는 보지 않음**.
|
||||
|
||||
### recipe 가 (generic op list 가 아닌) 이유
|
||||
|
||||
decode opt2 의 `#2` MATH chain 은 8 단계 고정 구조 (online softmax merge).
|
||||
커널 작성자에게 8 단계 직접 와이어링은 verbose + correctness hazard
|
||||
(단계 순서 의존). *Recipe* — TLContext 가 8 단계로 (주소 미리 채워)
|
||||
펼치는 단일 명명 op (`softmax_merge`) — 가 커널을 간결히 + 구조를 단일
|
||||
source 에 보존. recipe 표는 TLContext (compiler analog) 에 위치;
|
||||
PE_SCHEDULER 는 안 읽음. D5 boundary 참조.
|
||||
|
||||
## Decision
|
||||
|
||||
### D1. `CompositeCmd` 는 평평한 ordered op 리스트
|
||||
|
||||
```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: ... # ADR-0064 D2
|
||||
```
|
||||
|
||||
이전의 `(op, a, b, out_addr, out_nbytes, math_op)` 필드는 `ops` 에 흡수.
|
||||
`ops` 의 첫 OpSpec 이 `kind == "gemm"` 이면 head GEMM; 앞의 OpSpec 들이
|
||||
pre-loop, 뒤의 OpSpec 들이 `scope` 에 따라 post-loop / per-tile epilogue.
|
||||
|
||||
### D2. `OpSpec` 이 named operand 보유
|
||||
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class OpSpec:
|
||||
kind: str # "gemm" | "rmax" | ...
|
||||
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
|
||||
operands: dict[str, TensorHandle] # named — D5 참조
|
||||
extra: dict[str, Any] # scalars, axes, m/k/n 등
|
||||
out: TensorHandle | None = None # 명시적 write-back handle
|
||||
|
||||
@property
|
||||
def logical_bytes(self) -> int: ...
|
||||
```
|
||||
|
||||
named operand 로 PE_SCHEDULER 가 (positional 컨벤션 없이) 엔진 port 에
|
||||
입력 라우팅 (예: GEMM 은 `operands["a"]`, `operands["b"]`).
|
||||
|
||||
### 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 참조).
|
||||
인덱스 `g`:
|
||||
|
||||
| OpSpec position | scope | Phase | plan 내 배치 |
|
||||
|---|---|---|---|
|
||||
| `0 .. g-1` | KERNEL | pre-tile-loop | tile loop 진입 전 1 회 |
|
||||
| `g` | KERNEL | head GEMM | `extra["m"], extra["k"], extra["n"]` 로 tile loop drive |
|
||||
| `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (K 누산 loop 안) |
|
||||
| `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (K 누산 후) |
|
||||
| `g+1 .. ` | KERNEL | post-tile-loop | tile loop 종료 후 1 회 |
|
||||
|
||||
GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope
|
||||
직렬로 처리 — phase 가 "순서대로 single-shot" 으로 collapse.
|
||||
|
||||
### D4. PE_SCHEDULER 가 operand `pinned` 플래그 보고 DMA 자동 삽입
|
||||
|
||||
PE_SCHEDULER 가 GEMM 의 `operands` 검사. 각각:
|
||||
- **not `pinned`** (데이터가 HBM 에 있어 스트리밍 필요) → FETCH/GEMM Stage
|
||||
앞에 DMA_READ Stage 삽입
|
||||
- **`pinned`** (이전 `tl.load` 로 이미 TCM 에 staged, 또는 recipe 의 TCM
|
||||
scratch / primary-out `P`) → DMA Stage 없음, in-place 소비
|
||||
|
||||
커널은 composite operand 의 명시적 DMA cmd 를 절대 emit 안 함;
|
||||
PE_SCHEDULER 가 핸들에서 추론. `tl.ref` operand 는 not pinned (DMA_READ);
|
||||
`tl.load` 결과와 recipe scratch / primary-out 은 pinned (in-place).
|
||||
|
||||
> **as-built 노트 (P3).** D4 의 초안은 DMA 결정을 `space` 태그
|
||||
> (`hbm`→DMA, `tcm`→in-place) 에서 도출하며 현재 `space` 값
|
||||
> (오늘 `tl.load`=`hbm`, `tl.ref`=`tcm`) 을 뒤집어야 했음. 구현은 기존
|
||||
> **`pinned` 플래그**를 DMA 신호로 그대로 유지 — `pinned` 이 이미
|
||||
> "이미 TCM 에 있음, DMA skip" 을 인코딩하고, `space` 뒤집기는 기존 bench
|
||||
> Stage 시퀀스를 바꾸면서 기능적 이득이 없음. `space` 통일 + `pinned`
|
||||
> 폐기는 후순위 cleanup 으로 연기.
|
||||
|
||||
**Prologue recipe op 은 TCM-only (Phase 1 한계, TLContext emit 시 강제).**
|
||||
*prologue recipe* (비 GEMM) OpSpec 의 모든 operand 는 **반드시** TCM 상주;
|
||||
HBM 핸들을 recipe operand 로 전달하면 emit 시 validation error. Phase 1
|
||||
에서 DMA staging 을 head op 의 operand 로 제한. **head op 자체** (gemm
|
||||
또는 math) 는 기존 DMA-staged-from-HBM 동작 유지 — D4 의 TCM-only 룰은
|
||||
head 에 적용 **안 됨**. 이 invariant 는 D6 #7 에 재기술.
|
||||
|
||||
### D5. RECIPE_DESCRIPTORS — TLContext 내부, `pe_commands.py` 가 아님
|
||||
|
||||
Recipe 는 새 TLContext-adjacent 모듈
|
||||
(`src/kernbench/triton_emu/tl_recipes.py`) 에 위치. PE_SCHEDULER 는
|
||||
import **안 함**. 첫 recipe:
|
||||
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class RecipeDescriptor:
|
||||
operands: dict[str, str] # name → "R" | "RW"
|
||||
primary_out: PrimaryOutSpec | None # 암묵 primary output 의 type/shape 룰
|
||||
tile_alignment: Literal["single_shot", "tile_aligned"]
|
||||
internal_scratch_bytes_fn: Callable[..., int]
|
||||
engine_seq: tuple[EngineOp, ...] # TLContext 가 평평한 OpSpec 으로 펼침
|
||||
|
||||
|
||||
RECIPE_DESCRIPTORS["softmax_merge"] = RecipeDescriptor(
|
||||
operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
|
||||
primary_out=PrimaryOutSpec(
|
||||
from_shape="s", from_dtype="s", transform="identity",
|
||||
),
|
||||
tile_alignment="single_shot",
|
||||
internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
|
||||
G + G + G + G * TILE + G # m_loc + m_new + corr + P + l_loc
|
||||
),
|
||||
engine_seq=(
|
||||
EngineOp("MATH", "rmax", src="s", dst="m_loc", reduce_axis=-1),
|
||||
EngineOp("MATH", "max_elem", src_a="m", src_b="m_loc", dst="m_new"),
|
||||
EngineOp("MATH", "exp_diff", src_a="m", src_b="m_new", dst="corr"),
|
||||
EngineOp("MATH", "exp_diff", src_a="s", src_b="m_new", dst="P", bcast_axis=0),
|
||||
EngineOp("MATH", "rsum", src="P", dst="l_loc", reduce_axis=-1),
|
||||
EngineOp("MATH", "fma", src_a="l", src_b="corr", src_c="l_loc", dst="l"),
|
||||
EngineOp("MATH", "mul_bcast", src_a="O", src_b="corr", dst="O", bcast_axis=1),
|
||||
EngineOp("MATH", "copy", src="m_new", dst="m"),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
TLContext 가 `tl.composite(prologue=[{"op": "softmax_merge", ...}],
|
||||
op="gemm", ...)` 에서:
|
||||
1. operand R/RW 를 `RECIPE_DESCRIPTORS["softmax_merge"]` 와 검증.
|
||||
2. scratch (m_loc, m_new, corr, P, l_loc) 를 scratch_scope helper 로 할당 (ADR-0063).
|
||||
3. primary output P 의 shape 을 `s.shape` 에서 derive (identity).
|
||||
4. `engine_seq` 를 8 개의 평평한 MATH OpSpec 으로 펼침 (모든 주소/크기 채움).
|
||||
각 OpSpec 의 `scope = KERNEL`.
|
||||
5. **Auto-bind (충돌 검사, D6.6).** head GEMM 이 `operands["a"]` 를
|
||||
*명시하지 않은* 경우에만 `operands["a"] = P_handle` auto-bind.
|
||||
커널이 `a` 를 명시적으로 제공했다면 validation error — 모호한
|
||||
primary-output 대체.
|
||||
6. `CompositeCmd(ops=(8 MATH + 1 GEMM + epilogue), rw_handles=(m, l, O))` emit.
|
||||
|
||||
**RECIPE_DESCRIPTORS 는 HW 경로 어디에도 안 나타남**. PE_SCHEDULER 는
|
||||
평평한 ops 리스트만 봄.
|
||||
|
||||
### D6. Invariants
|
||||
|
||||
1. **GEMM 수.** `0 ≤ count(op.kind == "gemm" in cmd.ops) ≤ 1`. GEMM 0 개면
|
||||
MATH-only composite.
|
||||
2. **softmax_merge 가 V operand 없음.** `RECIPE_DESCRIPTORS[
|
||||
"softmax_merge"].operands` 에 HBM-resident handle 없음 → prologue 동안
|
||||
DMA 삽입 없음 → V DMA 는 GEMM head 시작 시에만 발생 → decode opt2 의
|
||||
#1 / #2 간 K-before-V DMA priority 자연 강제.
|
||||
3. **Cross-composite strict FIFO.** PE_SCHEDULER 가 in-flight `rw_handles`
|
||||
추적. 신규 composite 의 `rw_handles` (또는 op 입력) 가 in-flight 와
|
||||
교차하면 모든 이전 composite 완료까지 대기 — out-of-order reorder 없음.
|
||||
4. **legacy caller 후방 호환.** 사용자 API `tl.composite(op="gemm", a, b,
|
||||
epilogue=[...])` 보존. TLContext 가 내부적으로 평평한 ops `CompositeCmd`
|
||||
로 lowering.
|
||||
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER 는 RECIPE_DESCRIPTORS import
|
||||
안 함, "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. **Prologue MATH operand TCM-only.** D4 의 재기술: *prologue recipe*
|
||||
(비 GEMM) OpSpec 의 모든 operand 는 TCM 상주. head op (gemm 또는 math)
|
||||
은 예외 — HBM 스트리밍 허용. Phase 1 에서 TLContext emit 시 강제.
|
||||
|
||||
### D7. Boundary 요약 (compiler vs scheduler vs engine)
|
||||
|
||||
```
|
||||
HOST DEVICE
|
||||
TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher)
|
||||
- RECIPE_DESC 와 검증 - 평평한 ops 리스트 스캔
|
||||
- primary_out shape derive - GEMM 식별 (≤ 1)
|
||||
- scratch 할당 - position-기반 stage 배치
|
||||
- recipe → 평평한 OpSpec 펼침 - space 보고 DMA stage 자동 삽입
|
||||
- rw_handles 계산 - strict-FIFO RW 해저드 추적
|
||||
- CompositeCmd emit - Stage 리스트 emit
|
||||
imports: RECIPE_DESCRIPTORS imports: recipe 관련 없음
|
||||
|
||||
ENGINES (PE_MATH / PE_GEMM / PE_DMA)
|
||||
- Stage.params 읽음 (op_kind, addresses, n_elements)
|
||||
- 기존 latency 모델 변화 없음
|
||||
imports: recipe 관련 없음
|
||||
```
|
||||
|
||||
### D8. Composite 출력 handle + 출력-space DMA
|
||||
|
||||
`tl.composite(...)` 는 **출력 `TensorHandle`** 을 반환(단순 `CompletionHandle`
|
||||
아님) — composite 결과를 `tl.dot` 결과처럼 downstream op 에 먹일 수 있음.
|
||||
handle 의 `pending` 필드가 composite 완료를 참조; downstream consumer 가
|
||||
auto-await(`_await_pending`, ADR-0062 lazy 패턴), `tl.wait(handle)` 동작.
|
||||
|
||||
**`out` 은 항상 `TensorHandle`** (raw 주소 아님) — 위치는 커널의 선택, 추측 안 함:
|
||||
|
||||
- **HBM 출력** → `out=tl.ref(addr, shape)`. `tl.ref` 가 `space="hbm"` handle
|
||||
반환 → composite 가 DMA_WRITE 로 write-back. STORE + DMA_WRITE 는 composite
|
||||
파이프라인 **내부**에 유지 — fused op 의 일부이지 별도 `tl.store` 아님.
|
||||
- **In-place TCM** → `out=<기존 TCM handle>` (예: recipe 누적기 `O`).
|
||||
`space="tcm"` → STORE 만, 결과 TCM 에 남음 (chainable).
|
||||
- **미지정** → TLContext 가 **TCM scratch 자동 할당** (`tl.dot` 과 동일) 후 반환.
|
||||
|
||||
편의상 `out_ptr: int` 는 **HBM 출력 shorthand** 로 허용 —
|
||||
`out=tl.ref(out_ptr, <출력 shape>)` 와 동치(컴파일러가 `space="hbm"` handle 로
|
||||
감쌈). 출력 미지정 시 scratch 주소는 커널이 아니라 컴파일러가 소유.
|
||||
(`tl.ref` 는 `space="hbm"` 반환 — HBM 데이터 참조이므로. operand-**입력** DMA
|
||||
결정은 D4 대로 `pinned` 기반이라 이 라벨이 입력 스트리밍에 영향 없음.)
|
||||
|
||||
출력 handle 의 **`space` 가 타일 루프의 write-back 을 결정** (D4 의 operand
|
||||
`pinned` 규칙의 출력판):
|
||||
|
||||
| `out.space` | 타일 루프 write-back |
|
||||
| --- | --- |
|
||||
| `"tcm"` | STORE 만 — TCM 에 남음 (chainable; **DMA_WRITE 없음**) |
|
||||
| `"hbm"` | STORE + DMA_WRITE — 최종 HBM 출력 |
|
||||
|
||||
이로써 출력이 TCM 에 남는 composite(예: recipe 누적기)가 합법이 되고,
|
||||
TCM-scratch 주소가 DMA 엔진의 PA 디코더(고비트 scratch 주소를 거부)로 가는
|
||||
일을 막음. 기존 벤치는 HBM 출력을 `tl.ref` 로 감싸므로 write-back
|
||||
불변(byte-equal).
|
||||
|
||||
## Alternatives
|
||||
|
||||
### A1. tile 당 3 composite (prologue 개념 없이)
|
||||
|
||||
#2 를 `softmax_merge` + `gemm(P·V)` + `add(O)` 의 3 개로 분할. 장점:
|
||||
flat-ops 재구조화 없음 (`add` 가 기존 epilogue 경로 재사용). 단점: tile
|
||||
당 dispatch 3 회 (대신 2) — ADR-0060 §8 item 4 sizing ("**두 개** 로
|
||||
분할") 과 어긋남; ~32 ns × N_tiles 의 fixed-cost 절감 소실. **기각.**
|
||||
|
||||
### A2. Mixed-engine recipe (`flash_pv_merge`)
|
||||
|
||||
MATH + GEMM + MATH 흡수하는 단일 recipe (ADR-0060 §5.6 원안 표현). 장점:
|
||||
가장 적은 dispatch (1 composite). 단점: PE_SCHEDULER 가 engine-경계
|
||||
recipe 알아야 함, "엔진은 Stage.op_kind 만" boundary 위반, recipe 표가
|
||||
ISA-급이 되어 HW 인터페이스 오염. **기각.**
|
||||
|
||||
### A3. Generic op-list DAG (`tl.ex_composite([...])`)
|
||||
|
||||
사용자 정의 임의 op 시퀀스 + 명명 슬롯 + 데이터플로 분석. 장점: 최대
|
||||
유연성. 단점: 단일 user (softmax_merge) 가 generic DAG 언어를 정당화 못함;
|
||||
PE_SCHEDULER 에 데이터플로 분석 필요; ADR-0060 §8 item 4 의 "inner loop
|
||||
흡수" 트랩 재등장. **기각 (Simplicity First).**
|
||||
|
||||
### A4. RW-aware scheduler reorder (strict FIFO 가 아님)
|
||||
|
||||
PE_SCHEDULER 가 `rw_handles` 교차 없으면 후속 composite 가 in-flight 를
|
||||
추월 허용 (예: 다음 tile 의 #1 가 현재 tile 의 #2 와 겹침 — #1 은
|
||||
m/l/O 안 만짐). 장점: 더 많은 overlap. 단점: scheduler state 증가, 이득
|
||||
한정 (다음 #1 은 어차피 K DMA 대기). **연기** — strict FIFO 먼저;
|
||||
RW-aware reorder 는 후속 ADR 가능.
|
||||
|
||||
### A5. softmax_merge 를 #1 (Q·Kᵀ) 의 epilogue 로 embedding
|
||||
|
||||
오답: softmax_merge 의 출력 `P` 는 #2 의 P·V GEMM 의 입력, #1 의
|
||||
epilogue *이후* 실행. embedding 하면 #2 가 자기 입력을 #1 의 epilogue
|
||||
scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect).**
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
- ADR-0060 §5.6 / §8 item 4 carve-out 정확 구현: opt2 = tile 당 2
|
||||
composite, K-before-V DMA priority 자연.
|
||||
- composite contract 균일화: scope-기반 배치의 평평한 ops 리스트. HW cmd
|
||||
에서 "head" 와 "epilogue" 구조 구분 없음.
|
||||
- DMA 가 operand `space` 에서 자동 추론; 커널 표면 단순화.
|
||||
- HW 인터페이스 (`CompositeCmd` 모양) 가 recipe 증가에도 작게 유지 —
|
||||
recipe 펼침이 host-side.
|
||||
- 새 fused 패턴 (linear+gelu+norm, rmsnorm+linear 등) 은 `RECIPE_DESCRIPTORS`
|
||||
entry 추가만으로 가능; PE_SCHEDULER 변경 없음.
|
||||
|
||||
### Negative
|
||||
- CompositeCmd 구조 변경 — 모든 현재 caller 의 refactor. 의미 보존
|
||||
(D6.4); ADR-0064 Revision 2 land 후 기존 골든 불변.
|
||||
- `OpSpec.operands` 가 positional `tuple[Any, ...]` 에서
|
||||
`dict[str, TensorHandle]` 로 변경 — 기존 epilogue lowering 접촉.
|
||||
- 새 TLContext 코드: recipe 펼침 + scratch slot 할당 + primary_out
|
||||
auto-bind. ~80–120 LOC.
|
||||
- PE_SCHEDULER 의 `_generate_plan` 에 position-기반 스캔 + DMA 자동 삽입
|
||||
+ strict-FIFO RW tracker. ~60–100 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
|
||||
|
||||
1. **identity 를 넘는 recipe shape derivation 룰.** 미래 recipe 가
|
||||
non-identity transform (예: 전치된 primary_out) 필요할 수 있음.
|
||||
`PrimaryOutSpec.transform` 필드가 forward-compat; 필요 시 새 transform 추가.
|
||||
2. **Recipe scratch allocator** 의 ADR-0063 `scratch_scope` 통합. 초기
|
||||
배선: TLContext 가 활성 `scratch_scope` 가 있으면 그 안에서 할당,
|
||||
없으면 kernel-scope persistent slot.
|
||||
3. **`tl_recipes.py` 위치.** `triton_emu/` 아래 새 모듈. recipe 지식을
|
||||
compiler analog (TLContext) 와 함께 위치.
|
||||
4. **GEMM `out=TensorHandle`.** 기존 `out_ptr: int` 와 함께 새 형태.
|
||||
둘 다 허용; handle 형태가 신규 코드에 권장 (handle identity 로
|
||||
strict-FIFO RW 추적 가능).
|
||||
|
||||
## Test Requirements
|
||||
|
||||
1. **CompositeCmd 평평화 refactor — 의미 보존.** 기존 모든 bench 의
|
||||
op_log 가 refactor 전후로 byte-equal (legacy API `tl.composite(
|
||||
op="gemm", a, b, epilogue=[...])` 가 같은 Stage 시퀀스로 lowering).
|
||||
2. **softmax_merge recipe lowering.** TLContext 호출 `tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
|
||||
op="gemm", b=V_ref, out=O, epilogue=[{"op": "add", "other": O}])` 가
|
||||
정확히 10 ops (8 MATH + 1 GEMM + 1 MATH(add)) + `rw_handles == (m, l, O)`
|
||||
의 `CompositeCmd` 생성.
|
||||
3. **K-before-V DMA priority invariant.** decode opt2 의 op_log 에서
|
||||
#2 의 MATH prologue 동안 V 관련 DMA 없음 (#2 의 V DMA 는 GEMM head
|
||||
시작 시에만).
|
||||
4. **Strict-FIFO RW 직렬화.** `rw_handle` 공유하는 두 연속 composite 가
|
||||
dispatch 순서대로 완료; stage 가 interleave 안 함.
|
||||
5. **`pinned` 에서 DMA 자동 삽입.** not-`pinned` operand (예: `b=tl.ref(...)`)
|
||||
의 GEMM composite 가 DMA_READ Stage emit; `pinned` operand (예: recipe
|
||||
의 TCM scratch / primary-out, 또는 `tl.load` 결과) 는 emit 안 함.
|
||||
(as-built D4 노트: DMA 결정은 `space` 태그가 아니라 기존 `pinned`
|
||||
플래그 사용.)
|
||||
6. **opt2 가 opt3 와 수치 동등.** data mode 에서 opt2 의 최종
|
||||
`(m, l, O)` 가 opt3 와 fp tolerance 안.
|
||||
7. **opt2 dispatch ratio (ADR-0064 Rev2 이후) — robust.** opt3 vs opt2
|
||||
per-tile PE_CPU dispatch cycles 가 `opt3 > 2 × opt2` 만족
|
||||
(정성적 gate, calibration-독립). Default-calibration 모델 기대치는
|
||||
≈ 4.0× (FIXED=40 cycles, R=0.0625 cycles/byte) — DDD-0065 §11 의
|
||||
모델-유도 숫자 참조; 테스트 gate 는 느슨한 `> 2×` 한계라
|
||||
calibration 변경에도 살아남음. Ratio 가 FIXED-dominated —
|
||||
command-count 감소가 1차 신호임을 반영.
|
||||
8. **GEMM-count invariant.** GEMM OpSpec 두 개 가진 composite 가
|
||||
TLContext emit 시 validation error.
|
||||
9. **Composite 크기 cap 안에 (ADR-0064 D7).** Decode opt2 의 `#2`
|
||||
composite (10 ops, ~322 logical bytes — operand 핸들의 per-op 합산)
|
||||
가 default `MAX_COMPOSITE_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
|
||||
|
||||
- **ADR-0060** §5.6, §8 item 4 — 본 ADR 이 구현하는 carve-out.
|
||||
- **ADR-0064 Revision 2** — dispatch cost model; 먼저 land, opt2 의
|
||||
fewer-issues win 을 측정 가능하게 함.
|
||||
- **ADR-0063** `scratch_scope` — recipe intermediate 가 안에 할당;
|
||||
`m, l, O` 는 persistent arena 로 밖에 할당 (DDD-0060 §6.2 참조).
|
||||
- **ADR-0046** `tl_context` contract — `prologue=[...]` kwarg,
|
||||
`out=TensorHandle`, RECIPE_DESCRIPTORS lookup 으로 확장.
|
||||
- **ADR-0042** tile plan generators — (재설계 아닌) 확장 — 평평한 ops
|
||||
리스트 처리 + DMA 자동 삽입.
|
||||
- **ADR-0014** PE pipeline — boundary 보존: scheduler 가 recipe-free,
|
||||
엔진이 op_kind-opaque.
|
||||
|
||||
## Migration
|
||||
|
||||
> **구현 순서 노트.** 아래의 ADR-0064-우선 순서가 원래 계획이었음.
|
||||
> 실제 구현은 **1↔2 단계를 뒤집어** — ADR-0065 Phase 1 (평평화
|
||||
> `CompositeCmd`) 이 ADR-0064 Rev2 *보다 먼저* land. 이유: `logical_bytes`
|
||||
> (ADR-0064 D2) 가 평평화 형태 위에서 자연스럽게 정의되므로, refactor 를
|
||||
> 먼저 하면 버려지는 transitional `logical_bytes` 를 피함. 각 단계는
|
||||
> 독립적으로 안전: P1 은 기존 cost 경로 하에서 의미 보존 refactor
|
||||
> (op_log byte-equal); ADR-0064 Rev2 가 이미 평평화된 형태 위에 구조적
|
||||
> 공식을 교체. 최종 상태는 어느 쪽이든 동일.
|
||||
|
||||
Land 순서 (원래 계획; as-built 1↔2 뒤집힘은 위 노트 참조):
|
||||
1. **ADR-0064 Revision 2** (별도 PR): 구조적 dispatch cost +
|
||||
`logical_bytes` property + topology config override + 일회성 골든 재생성.
|
||||
2. **ADR-0065 Phase 1** (본 ADR, PR a): `CompositeCmd` 평평화 refactor +
|
||||
legacy-API lowering. refactor only; 골든 불변.
|
||||
3. **ADR-0065 Phase 2** (본 ADR, PR b): `tl_recipes.py` + `softmax_merge`
|
||||
recipe + PE_SCHEDULER position-스캔 + DMA 자동 삽입 + strict-FIFO
|
||||
RW tracker.
|
||||
4. **ADR-0065 Phase 3** (본 ADR, PR c): `_gqa_decode_long.py` opt2
|
||||
variant + 수치 동등 테스트 + opt3 대비 dispatch-ratio 측정.
|
||||
|
||||
상세 구현 계획: **DDD-0065** 참조.
|
||||
@@ -1,488 +0,0 @@
|
||||
# ADR-0065: Flat-ops `CompositeCmd` + first stateful recipe (`softmax_merge`)
|
||||
|
||||
## Status
|
||||
|
||||
Proposed (verification gated on ADR-0064 Revision 2 land)
|
||||
|
||||
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Implements
|
||||
> the §5.6 / §8 item 4 carve-out exactly: opt2 of decode = (#1 existing
|
||||
> GEMM composite for Q·Kᵀ) + (#2 a single composite carrying softmax +
|
||||
> P·V + online-softmax merge). The "ex_composite" / "flash_pv_merge" name
|
||||
> referenced in ADR-0060 §5.6 is *retired* — this ADR fixes the canonical
|
||||
> shape using the existing `tl.composite` entry with two structural
|
||||
> additions: (a) a flat `ops` tuple, and (b) a first stateful recipe
|
||||
> `softmax_merge` that lowers to a sequence of MATH micro-ops.
|
||||
|
||||
## Context
|
||||
|
||||
### What ADR-0060 left to this ADR
|
||||
|
||||
ADR-0060 §5.6 shipped opt3 (software pipelining) as the recommended now,
|
||||
and explicitly carved out opt2 — fewer per-tile dispatches, K-before-V
|
||||
DMA priority, **measurable only after ADR-0064**'s cost model — as the
|
||||
follow-up.
|
||||
|
||||
§8 item 4 sized the carve-out: "if revisited, split it into **two**
|
||||
composites — `#1` = Q·Kᵀ (existing composite + `scale`), `#2` = softmax +
|
||||
P·V + online-softmax accumulator merge". This ADR implements `#2`.
|
||||
|
||||
### Why composite needs a structural change
|
||||
|
||||
Today's `CompositeCmd` has the shape `(op, a, b, out_addr, ops=(head,
|
||||
*epi))` with implicit conventions:
|
||||
- `op` selects the head engine (gemm | math)
|
||||
- `ops[0]` is the head OpSpec, `ops[1:]` are epilogue OpSpecs that run
|
||||
*after* the head per-output-tile or per-K-tile (scope-determined)
|
||||
|
||||
For decode opt2 we need MATH ops to run **before** the head GEMM (the
|
||||
online-softmax `m, l, O` update + emit `P` for the GEMM input). The
|
||||
current shape has no slot for that.
|
||||
|
||||
Two failed paths considered earlier:
|
||||
- Embedding softmax_merge as an *epilogue* of #1 (Q·Kᵀ) — wrong: it
|
||||
needs to read `Sj` (= #1's output) and writes a fresh `P` consumed by
|
||||
the *next* GEMM (#2's P·V); circular.
|
||||
- A single mixed-engine recipe `flash_pv_merge` absorbing MATH + GEMM +
|
||||
MATH — violates the engine boundary, complicates PE_SCHEDULER.
|
||||
|
||||
The clean shape is: **drop the head/epilogue distinction from the
|
||||
command format**. Make `ops` a flat ordered tuple; let each op's
|
||||
`scope` + position determine its tile-loop placement. The "prologue"
|
||||
concept stays in the *user-facing* `tl.composite(...)` API as
|
||||
ergonomics — but the **HW command never sees it**.
|
||||
|
||||
### Why a recipe (not a generic op list)
|
||||
|
||||
Decode opt2's `#2` MATH chain is 8 steps with a fixed structure (online
|
||||
softmax merge). Letting the kernel author wire those 8 steps would be
|
||||
both verbose and a hazard surface (correctness depends on step order).
|
||||
A *recipe* — a single named op (`softmax_merge`) that TLContext expands
|
||||
to the 8 steps with addresses pre-filled — keeps the kernel concise and
|
||||
the structure under one source of truth. The recipe table lives in
|
||||
TLContext (the compiler analog); PE_SCHEDULER never reads it. See
|
||||
boundary in D5.
|
||||
|
||||
## Decision
|
||||
|
||||
### D1. `CompositeCmd` is a flat ordered op list
|
||||
|
||||
```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: ... # ADR-0064 D2
|
||||
```
|
||||
|
||||
The previous `(op, a, b, out_addr, out_nbytes, math_op)` fields are
|
||||
absorbed into `ops`. The first OpSpec in `ops` (if `kind == "gemm"`) is
|
||||
the head GEMM; preceding OpSpecs are pre-loop ops, following OpSpecs
|
||||
are post-loop / per-tile epilogue ops depending on `scope`.
|
||||
|
||||
### D2. `OpSpec` carries named operands
|
||||
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class OpSpec:
|
||||
kind: str # "gemm" | "rmax" | ...
|
||||
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
|
||||
operands: dict[str, TensorHandle] # named — see D5
|
||||
extra: dict[str, Any] # scalars, axes, m/k/n, etc.
|
||||
out: TensorHandle | None = None # explicit write-back handle
|
||||
|
||||
@property
|
||||
def logical_bytes(self) -> int: ...
|
||||
```
|
||||
|
||||
Named operands let PE_SCHEDULER route inputs to engine ports (e.g., GEMM
|
||||
takes `operands["a"]`, `operands["b"]`) without positional convention.
|
||||
|
||||
### 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
|
||||
D6). Calling its index `g`:
|
||||
|
||||
| OpSpec position | scope | Phase | Placement in plan |
|
||||
|---|---|---|---|
|
||||
| `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+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (inside K accumulation loop) |
|
||||
| `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (after K accumulation) |
|
||||
| `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
|
||||
as KERNEL-scope sequential — phase collapses to "single-shot in order".
|
||||
|
||||
### D4. PE_SCHEDULER auto-inserts DMAs from the operand `pinned` flag
|
||||
|
||||
PE_SCHEDULER scans the GEMM's `operands`. For each:
|
||||
- **not `pinned`** (data lives in HBM, must be streamed) → emit DMA_READ
|
||||
Stage before the FETCH/GEMM Stages
|
||||
- **`pinned`** (already staged in TCM via a prior `tl.load`, or a recipe's
|
||||
TCM scratch / primary-out `P`) → no DMA Stage, consumed in place
|
||||
|
||||
The kernel never emits explicit DMA cmds for composite operands;
|
||||
PE_SCHEDULER infers them from handles. `tl.ref` operands are not pinned
|
||||
(DMA_READ); `tl.load` results and recipe scratch / primary-out are pinned
|
||||
(in place).
|
||||
|
||||
> **As-built note (P3).** An earlier draft of D4 derived the DMA decision
|
||||
> from a `space` tag (`hbm`→DMA, `tcm`→in place) and would have flipped the
|
||||
> current `space` values (today `tl.load`=`hbm`, `tl.ref`=`tcm`). The
|
||||
> implementation **keeps the existing `pinned` flag** as the DMA signal
|
||||
> instead — `pinned` already encodes "already in TCM, skip DMA", flipping
|
||||
> `space` would change existing bench Stage sequences for no functional
|
||||
> gain. Unifying `space` and retiring `pinned` is a deferred low-priority
|
||||
> cleanup.
|
||||
|
||||
**Prologue recipe ops are TCM-only (Phase 1 limit, enforced at TLContext
|
||||
emit).** Every operand of a *prologue recipe* (non-GEMM) OpSpec **must** be
|
||||
TCM-resident; passing an HBM handle as a recipe operand raises a validation
|
||||
error at emit time. This bounds DMA staging to the head op's operands in
|
||||
Phase 1. The **head op itself** (gemm *or* math) keeps the existing
|
||||
DMA-staged-from-HBM behavior — D4's TCM-only rule does **not** apply to it.
|
||||
This invariant is also restated in D6 #7.
|
||||
|
||||
### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py`
|
||||
|
||||
Recipes live in a new TLContext-adjacent module
|
||||
(`src/kernbench/triton_emu/tl_recipes.py`). PE_SCHEDULER does **not**
|
||||
import it. The first recipe:
|
||||
|
||||
```python
|
||||
@dataclass(frozen=True)
|
||||
class RecipeDescriptor:
|
||||
# Operand contract: name → "R" (read-only) | "RW" (read-write)
|
||||
operands: dict[str, str]
|
||||
# Type / shape rule for the implicit primary output (auto-bound to
|
||||
# head's first matrix operand). None means recipe emits no value.
|
||||
primary_out: PrimaryOutSpec | None
|
||||
# Single-shot prologue (full reduction) vs tile-aligned execution.
|
||||
tile_alignment: Literal["single_shot", "tile_aligned"]
|
||||
# Internal scratch budget (bytes), function of operand dims.
|
||||
internal_scratch_bytes_fn: Callable[..., int]
|
||||
# Engine micro-op sequence: TLContext expands this into flat OpSpecs.
|
||||
engine_seq: tuple[EngineOp, ...]
|
||||
|
||||
|
||||
RECIPE_DESCRIPTORS["softmax_merge"] = RecipeDescriptor(
|
||||
operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
|
||||
primary_out=PrimaryOutSpec(
|
||||
from_shape="s", from_dtype="s", transform="identity",
|
||||
),
|
||||
tile_alignment="single_shot",
|
||||
internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
|
||||
G + G + G + G * TILE + G # m_loc + m_new + corr + P + l_loc
|
||||
),
|
||||
engine_seq=(
|
||||
EngineOp("MATH", "rmax", src="s", dst="m_loc", reduce_axis=-1),
|
||||
EngineOp("MATH", "max_elem", src_a="m", src_b="m_loc", dst="m_new"),
|
||||
EngineOp("MATH", "exp_diff", src_a="m", src_b="m_new", dst="corr"),
|
||||
EngineOp("MATH", "exp_diff", src_a="s", src_b="m_new", dst="P", bcast_axis=0),
|
||||
EngineOp("MATH", "rsum", src="P", dst="l_loc", reduce_axis=-1),
|
||||
EngineOp("MATH", "fma", src_a="l", src_b="corr", src_c="l_loc", dst="l"),
|
||||
EngineOp("MATH", "mul_bcast", src_a="O", src_b="corr", dst="O", bcast_axis=1),
|
||||
EngineOp("MATH", "copy", src="m_new", dst="m"),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
TLContext, on `tl.composite(prologue=[{"op": "softmax_merge", ...}],
|
||||
op="gemm", ...)`:
|
||||
1. Validates operand R/RW against `RECIPE_DESCRIPTORS["softmax_merge"]`.
|
||||
2. Allocates scratch (m_loc, m_new, corr, P, l_loc) via the scratch_scope
|
||||
helper (ADR-0063).
|
||||
3. Derives the primary output P's shape from `s.shape` (identity).
|
||||
4. Expands `engine_seq` into 8 flat MATH OpSpecs with all addresses /
|
||||
sizes filled. Each OpSpec gets `scope=KERNEL`.
|
||||
5. **Auto-binds (with conflict check, D6.6).** If the head GEMM does
|
||||
*not* explicitly provide `operands["a"]`, auto-bind
|
||||
`operands["a"] = P_handle`. If the kernel *did* provide `a`
|
||||
explicitly, emit a validation error — ambiguous primary-output
|
||||
replacement.
|
||||
6. Emits `CompositeCmd(ops=(8 MATH + 1 GEMM + epilogue), rw_handles=
|
||||
(m, l, O))`.
|
||||
|
||||
**RECIPE_DESCRIPTORS appears nowhere in the HW path**. PE_SCHEDULER
|
||||
sees only the flat ops list.
|
||||
|
||||
### D6. Invariants
|
||||
|
||||
1. **GEMM count.** `0 ≤ count(op.kind == "gemm" in cmd.ops) ≤ 1`. A
|
||||
composite with 0 GEMMs is MATH-only.
|
||||
2. **softmax_merge has no V operand.** `RECIPE_DESCRIPTORS["softmax_merge"
|
||||
].operands` does not include any HBM-resident handle → no DMA inserted
|
||||
during the prologue → V DMA can only happen during the GEMM head →
|
||||
natural K-before-V DMA priority across #1 / #2 of decode opt2.
|
||||
3. **Strict FIFO across composites.** PE_SCHEDULER tracks in-flight
|
||||
`rw_handles`. A new composite whose `rw_handles` (or any of its op
|
||||
inputs) intersect with an in-flight composite's `rw_handles` waits
|
||||
for all earlier composites to complete — no out-of-order reorder.
|
||||
4. **No backwards-incompatibility for legacy callers.** The user-facing
|
||||
API `tl.composite(op="gemm", a, b, epilogue=[...])` is preserved.
|
||||
TLContext internally lowers it to a flat-ops `CompositeCmd`.
|
||||
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER does not import
|
||||
RECIPE_DESCRIPTORS, does not know that "softmax_merge" exists, and
|
||||
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. **Prologue MATH operands TCM-only.** Restatement of D4: every operand
|
||||
of a *prologue recipe* (non-GEMM) OpSpec must be TCM-resident. The head
|
||||
op (gemm or math) is exempt — it may stream from HBM. Phase 1 enforced
|
||||
at TLContext emit.
|
||||
|
||||
### D7. Boundary summary (compiler vs scheduler vs engine)
|
||||
|
||||
```
|
||||
HOST DEVICE
|
||||
TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher)
|
||||
- validate against RECIPE_DESC. - scan flat ops list
|
||||
- derive primary_out shape - identify GEMM (≤1)
|
||||
- allocate scratch - position-based stage placement
|
||||
- expand recipe → flat OpSpecs - auto-insert DMA stages from space
|
||||
- compute rw_handles - strict-FIFO RW hazard tracking
|
||||
- emit CompositeCmd - emit Stage list
|
||||
imports: RECIPE_DESCRIPTORS imports: nothing recipe-related
|
||||
|
||||
ENGINES (PE_MATH / PE_GEMM / PE_DMA)
|
||||
- read Stage.params (op_kind, addresses, n_elements)
|
||||
- existing latency models unchanged
|
||||
imports: nothing recipe-related
|
||||
```
|
||||
|
||||
### D8. Composite output handle + output-space DMA
|
||||
|
||||
`tl.composite(...)` returns the **output `TensorHandle`** (not a bare
|
||||
`CompletionHandle`), so a composite's result can feed a downstream op the
|
||||
same way `tl.dot`'s result does. The handle's `pending` field references
|
||||
the composite's completion; downstream consumers auto-await it
|
||||
(`_await_pending`, ADR-0062 lazy pattern) and `tl.wait(handle)` works.
|
||||
|
||||
**`out` is always a `TensorHandle`** (never a raw address) — its location is
|
||||
the kernel's choice, never guessed:
|
||||
|
||||
- **HBM output** → `out=tl.ref(addr, shape)`. `tl.ref` returns a handle with
|
||||
`space="hbm"`, so the composite writes it back via DMA_WRITE. The
|
||||
STORE + DMA_WRITE stay **inside** the composite pipeline — they are part of
|
||||
the fused op, *not* a separate `tl.store`.
|
||||
- **In-place TCM** → `out=<existing TCM handle>` (e.g. the recipe accumulator
|
||||
`O`). `space="tcm"` → STORE only, result stays in TCM (chainable).
|
||||
- **Omitted** → TLContext **auto-allocates a TCM scratch** (like `tl.dot`)
|
||||
and returns it.
|
||||
|
||||
For ergonomics, `out_ptr: int` is accepted as **shorthand for an HBM
|
||||
output** — equivalent to `out=tl.ref(out_ptr, <output shape>)` (the compiler
|
||||
wraps it into an `space="hbm"` handle). When no output is specified the
|
||||
compiler owns the scratch address, never the kernel. (`tl.ref` returns
|
||||
`space="hbm"` — it references HBM data. The operand-**input** DMA decision
|
||||
stays `pinned`-based per D4, so this label does not affect input streaming.)
|
||||
|
||||
The output handle's **`space` drives the tile loop's write-back** (the
|
||||
output analog of the operand `pinned` rule in D4):
|
||||
|
||||
| `out.space` | tile-loop write-back |
|
||||
| --- | --- |
|
||||
| `"tcm"` | STORE only — result stays in TCM (chainable; **no DMA_WRITE**) |
|
||||
| `"hbm"` | STORE + DMA_WRITE — final HBM output |
|
||||
|
||||
This makes a TCM-resident composite output (e.g. a recipe accumulator) legal
|
||||
**and** keeps a TCM-scratch address out of the DMA engine's PA decoder
|
||||
(which would otherwise reject the high-bit scratch address). Existing benches
|
||||
wrap their HBM output via `tl.ref` → write-back unchanged (byte-equal).
|
||||
|
||||
## Alternatives
|
||||
|
||||
### A1. Three composites per tile (no prologue concept)
|
||||
|
||||
Split #2 further into `softmax_merge` + `gemm(P·V)` + `add(O)` as three
|
||||
separate CompositeCmds. Pros: no flat-ops restructuring (`add` reuses
|
||||
existing epilogue path). Cons: 3 dispatches per tile instead of 2 —
|
||||
diverges from ADR-0060 §8 item 4 sizing ("split into **two**"); loses
|
||||
~32 ns × N_tiles of fixed-cost savings. **Rejected.**
|
||||
|
||||
### A2. Mixed-engine recipe (`flash_pv_merge`)
|
||||
|
||||
One recipe absorbing MATH + GEMM + MATH (ADR-0060 §5.6 original
|
||||
phrasing). Pros: tightest dispatch (1 composite). Cons: PE_SCHEDULER
|
||||
must know engine-boundary recipes, violates the "engines see only
|
||||
Stage.op_kind" boundary, recipe table becomes ISA-like and pollutes
|
||||
HW interface. **Rejected.**
|
||||
|
||||
### A3. Generic op-list DAG (`tl.ex_composite([...])`)
|
||||
|
||||
User-defined sequence of any ops with named slots and dataflow analysis.
|
||||
Pros: maximally flexible. Cons: a single user (softmax_merge) does not
|
||||
justify a generic DAG language; PE_SCHEDULER would need dataflow
|
||||
analysis; reintroduces ADR-0060 §8 item 4's "absorb the inner loop"
|
||||
trap. **Rejected (Simplicity First).**
|
||||
|
||||
### A4. RW-aware scheduler reorder (not strict FIFO)
|
||||
|
||||
PE_SCHEDULER could allow a later CompositeCmd to overtake an in-flight
|
||||
one when their `rw_handles` don't intersect (e.g., next-tile #1 could
|
||||
overlap with this-tile #2 since #1 doesn't touch m/l/O). Pros: more
|
||||
overlap. Cons: more scheduler state, gains are bounded (next #1 waits
|
||||
for K DMA anyway). **Deferred** — strict FIFO ships first; RW-aware
|
||||
reorder may be a future follow-up ADR.
|
||||
|
||||
### A5. Embedding softmax_merge as an epilogue of #1 (Q·Kᵀ)
|
||||
|
||||
Wrong: softmax_merge's output `P` is the input of #2's P·V GEMM, which
|
||||
runs *after* #1's epilogue. Embedding it would force #2 to read its own
|
||||
input from #1's epilogue scratch — circular tile-loop dependency.
|
||||
**Rejected (incorrect).**
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
- ADR-0060 §5.6 / §8 item 4 carve-out implemented exactly: opt2 = 2
|
||||
composites per tile, K-before-V DMA priority natural.
|
||||
- Composite contract becomes uniform: flat ops list with scope-based
|
||||
placement. No structural distinction between "head" and "epilogue"
|
||||
in the HW cmd.
|
||||
- DMAs auto-inferred from operand `space`; kernel surface simpler.
|
||||
- HW interface (`CompositeCmd` shape) stays small even as recipes grow
|
||||
— recipe expansion happens host-side.
|
||||
- New fused patterns (linear+gelu+norm, rmsnorm+linear, …) can be added
|
||||
by inserting one `RECIPE_DESCRIPTORS` entry; no PE_SCHEDULER changes.
|
||||
|
||||
### Negative
|
||||
- CompositeCmd's struct shape changes — a refactor for all current
|
||||
callers. Semantics preserved (D6.4); existing goldens unchanged once
|
||||
ADR-0064 Revision 2 lands.
|
||||
- `OpSpec.operands` changes from positional `tuple[Any, ...]` to
|
||||
`dict[str, TensorHandle]` — touches the existing epilogue lowering.
|
||||
- New TLContext code: recipe expansion + scratch-slot allocation +
|
||||
primary_out auto-bind. ~80–120 LOC.
|
||||
- PE_SCHEDULER's `_generate_plan` gets a position-based scan + DMA
|
||||
auto-insertion + strict-FIFO RW tracker. ~60–100 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
|
||||
|
||||
1. **Recipe shape derivation rule beyond identity.** Future recipes may
|
||||
need non-identity transforms (e.g., transposed primary_out). The
|
||||
`PrimaryOutSpec.transform` field is forward-compat; add new
|
||||
transforms as needed.
|
||||
2. **Recipe scratch allocator** integration with ADR-0063's
|
||||
`scratch_scope`. Initial wiring: TLContext allocates inside the
|
||||
current `scratch_scope` context if active, else in a kernel-scope
|
||||
persistent slot.
|
||||
3. **`tl_recipes.py` location.** New module under `triton_emu/`. Keeps
|
||||
recipe knowledge co-located with the compiler analog (TLContext).
|
||||
4. **GEMM `out=TensorHandle`.** New form alongside legacy `out_ptr:
|
||||
int`. Both accepted; handle form preferred for new code (enables
|
||||
strict-FIFO RW tracking by handle identity).
|
||||
|
||||
## Test Requirements
|
||||
|
||||
1. **CompositeCmd flat-ops refactor — meaning preservation.** Every
|
||||
existing bench's op_log is byte-equal pre/post refactor (legacy API
|
||||
`tl.composite(op="gemm", a, b, epilogue=[...])` lowers to the same
|
||||
Stage sequence).
|
||||
2. **softmax_merge recipe lowering.** TLContext call `tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
|
||||
op="gemm", b=V_ref, out=O, epilogue=[{"op": "add", "other": O}])`
|
||||
produces a `CompositeCmd` with exactly 10 ops in the expected order
|
||||
(8 MATH + 1 GEMM + 1 MATH(add)) and `rw_handles == (m, l, O)`.
|
||||
3. **K-before-V DMA priority invariant.** op_log of decode opt2 shows
|
||||
no V-related DMA during #2's MATH prologue (#2's V DMA begins only
|
||||
when the GEMM head starts).
|
||||
4. **Strict-FIFO RW serialization.** Two consecutive composites sharing
|
||||
any `rw_handle` complete in dispatch order; their stages do not
|
||||
interleave.
|
||||
5. **DMA auto-insertion from `pinned`.** A GEMM composite with a
|
||||
not-`pinned` operand (e.g. `b=tl.ref(...)`) emits a DMA_READ Stage; a
|
||||
`pinned` operand (e.g. a recipe's TCM scratch / primary-out, or a
|
||||
`tl.load` result) does not. (As-built D4 note: the DMA decision uses
|
||||
the existing `pinned` flag, not a `space` tag.)
|
||||
6. **opt2 numeric parity with opt3.** In data mode, opt2's final
|
||||
`(m, l, O)` matches opt3's within fp tolerance.
|
||||
7. **opt2 dispatch ratio (after ADR-0064 Rev2) — robust.** opt3 vs
|
||||
opt2 per-tile PE_CPU dispatch cycles satisfies `opt3 > 2 × opt2`
|
||||
(qualitative gate; calibration-independent). The default-calibration
|
||||
model expectation is ≈ 4.0× (FIXED=40 cycles, R=0.0625 cycles/byte)
|
||||
— see DDD-0065 §11 for the model-derived numbers; only the loose
|
||||
`> 2×` bound is the test gate so it survives calibration changes.
|
||||
Ratio is FIXED-dominated, reflecting command-count reduction as the
|
||||
primary signal.
|
||||
8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a
|
||||
validation error at TLContext emit.
|
||||
9. **Within composite size cap (ADR-0064 D7).** Decode opt2's `#2`
|
||||
composite (10 ops, ~322 logical bytes — per-op summation of operand
|
||||
handles) fits comfortably under the default
|
||||
`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
|
||||
|
||||
- **ADR-0060** §5.6, §8 item 4 — the carve-out this ADR implements.
|
||||
- **ADR-0064 Revision 2** — dispatch cost model; lands first, makes
|
||||
opt2's fewer-issues win measurable.
|
||||
- **ADR-0063** `scratch_scope` — recipe intermediates allocated inside;
|
||||
`m, l, O` allocated outside as persistent arena (per DDD-0060 §6.2).
|
||||
- **ADR-0046** `tl_context` contract — extended by `prologue=[...]`
|
||||
kwarg, `out=TensorHandle`, RECIPE_DESCRIPTORS lookup.
|
||||
- **ADR-0042** tile plan generators — extended (not redesigned) to
|
||||
process the flat ops list and auto-insert DMAs.
|
||||
- **ADR-0014** PE pipeline — boundary preserved: scheduler stays
|
||||
recipe-free, engines stay op_kind-opaque.
|
||||
|
||||
## Migration
|
||||
|
||||
> **Implementation ordering note.** The ADR-0064-first ordering below was
|
||||
> the original plan. The actual implementation **reversed steps 1 and 2** —
|
||||
> ADR-0065 Phase 1 (flat-ops `CompositeCmd`) landed *before* ADR-0064 Rev2.
|
||||
> Reason: `logical_bytes` (ADR-0064 D2) is naturally defined on the flat-ops
|
||||
> shape, so doing the refactor first avoids a throwaway transitional
|
||||
> `logical_bytes`. Each step is independently safe: P1 is a meaning-
|
||||
> preserving refactor (op_log byte-equal) under the pre-existing cost path;
|
||||
> ADR-0064 Rev2 then swaps in the structural formula on the already-flat
|
||||
> shape. The end state is identical either way.
|
||||
|
||||
Land order (original plan; see note above for the as-built reversal of 1↔2):
|
||||
1. **ADR-0064 Revision 2** (separate PR): structural dispatch cost +
|
||||
`logical_bytes` property + topology config override + one-time
|
||||
goldens regeneration.
|
||||
2. **ADR-0065 Phase 1** (this ADR, PR a): `CompositeCmd` flat-ops
|
||||
restructure + legacy-API lowering. Refactor only; goldens unchanged.
|
||||
3. **ADR-0065 Phase 2** (this ADR, PR b): `tl_recipes.py` +
|
||||
`softmax_merge` recipe + PE_SCHEDULER position-scan + DMA
|
||||
auto-insertion + strict-FIFO RW tracker.
|
||||
4. **ADR-0065 Phase 3** (this ADR, PR c): `_gqa_decode_long.py` opt2
|
||||
variant + numeric parity test + dispatch-ratio measurement against
|
||||
opt3 baseline.
|
||||
|
||||
Detailed implementation plan: see **DDD-0065**.
|
||||
Reference in New Issue
Block a user