gqa(adr): ADR-0064 Rev2 (structural dispatch cost) + ADR-0065/DDD-0065 (flat-ops composite + softmax_merge recipe)
ADR-0064 Revision 2: replace per-op-type calibration table with a structural formula `FIXED_PER_CMD + cmd.logical_bytes × R`, defaults anchored at typical composite ≈ 45 ns. Topology yaml override under `pe_cost_model:` block; logical_bytes property per PE command. ADR-0065: implement ADR-0060 §5.6 / §8 item 4 carve-out as a flat-ops CompositeCmd (no head/epilogue structural fields — position + scope drives placement) + first stateful recipe `softmax_merge` (MATH-only 8-step). RECIPE_DESCRIPTORS lives in TLContext-adjacent module only; PE_SCHEDULER stays recipe-free and auto-inserts DMAs from operand `space`. Strict-FIFO RW hazard tracker; ≤1 GEMM per composite invariant. User-facing `tl.composite(prologue=[...], op=, epilogue=[...])` API preserved; existing benches unchanged (meaning-preserving refactor). DDD-0065: implementation-ready phased plan (P0=ADR-0064 Rev2 first, P1-P6 for ADR-0065), file plan, recipe engine sequence, scheduler plan-gen algorithm, RW tracker design, test matrix. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,166 +1,242 @@
|
|||||||
# ADR-0064: op 종류별 CPU 발행 비용 모델 (command construct + dispatch)
|
# ADR-0064: 구조적 CPU dispatch cost 모델 (`logical_bytes` + FIXED + R)
|
||||||
|
|
||||||
## Status
|
## Status
|
||||||
|
|
||||||
Proposed
|
Proposed (Revision 2)
|
||||||
|
|
||||||
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. 거기서의 하이브리드
|
> **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops
|
||||||
> 결정(GEMM은 `tl.composite`, softmax 머지는 커널)이 이기는 이유는
|
> composite + 첫 stateful recipe) 의 보조 ADR. 그 hybrid 의 핵심 — "GEMM 은
|
||||||
> **타일링을 PE_SCHEDULER에 offload하여 CPU가 coarse descriptor를 내리고
|
> `tl.composite` 로, softmax merge 는 커널에서" — 은 **PE_SCHEDULER 로
|
||||||
> 앞서 나가 엔진을 saturate**하기 때문이다. 그 이점이 현재 시뮬레이터에서는
|
> tiling 을 offload 해서 CPU 가 굵은 descriptor 만 issue 하고 앞서가며
|
||||||
> per-op CPU 발행 비용이 0이라 **보이지 않는다**. 본 ADR은 발행 비용을
|
> 엔진을 포화시킴** 으로 이깁니다. 그 win 이 현재 시뮬레이터에선 보이지
|
||||||
> 실재화하고 **op 종류별로 차등**화하여, composite-vs-primitive 트레이드오프
|
> 않습니다 (per-op CPU issue cost = 0).
|
||||||
> (그리고 CPU saturation)가 측정 가능하도록 한다.
|
>
|
||||||
|
> Revision 2 는 원안의 **op-type calibration 표** 를 **logical_bytes 기반
|
||||||
|
> 구조 공식** 으로 대체합니다 — op-type 별 calibration 불필요, 새 op_kind
|
||||||
|
> 도 자동 적용.
|
||||||
|
|
||||||
## Context
|
## Context
|
||||||
|
|
||||||
### 현재 상태
|
### 현재 코드
|
||||||
|
|
||||||
- 모든 `tl.*` op은 커맨드를 emit하기 전에 `_emit_dispatch_overhead()`를
|
- 모든 `tl.*` op 가 cmd emit 전에 `_emit_dispatch_overhead()` 호출
|
||||||
호출하며(`tl_context.py:123-125, 190, 227, 235, 612, …`), 이는
|
(`tl_context.py:196-212`) → `dispatch_cycles > 0` 일 때만
|
||||||
`dispatch_cycles > 0`일 때**만** `PeCpuOverheadCmd(cycles=dispatch_cycles)`를
|
`PeCpuOverheadCmd(cycles=dispatch_cycles)` 발행. knob 은 op kind 무관
|
||||||
emit한다. 즉 발행 비용은 `tl.load`, `tl.dot`, MATH op, `tl.composite`에
|
**uniform** 이며 두 실행 경로 모두 **0** 하드코딩
|
||||||
동일하게 적용되는 **단일 균일 노브**다.
|
(`pe_cpu.py:101` greenlet, `:195` replay).
|
||||||
- 그 노브는 두 실행 경로 모두에서 **0**으로 하드코딩되어 있다
|
- ⇒ 명령 issue (descriptor 구성 + scheduler 큐 push) 가 PE_CPU 에서
|
||||||
(`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ 커맨드 발행
|
**0 ns** 비용.
|
||||||
— *descriptor를 construct하고 scheduler 큐에 push* — 이 현재 PE_CPU에서
|
- `PeCpuOverheadCmd` 는 PE_CPU 에서 `yield env.timeout(cmd.cycles)` 로
|
||||||
**0 ns**다.
|
소비 (`kernel_runner.py:131-132`).
|
||||||
- `PeCpuOverheadCmd`는 PE_CPU(`kernel_runner.py:131-132`)와
|
|
||||||
PE_SCHEDULER(`pe_scheduler.py:97-100`)에서 `yield env.timeout(cmd.cycles)`로
|
|
||||||
소비된다. 수동 `tl.cycles(n)`도 존재한다(`tl_context.py:695`).
|
|
||||||
|
|
||||||
### 본 ADR을 규정하는 두 가지 발견 (ADR-0060 검토에서)
|
### uniform-zero 가 hybrid 에 부적합한 이유
|
||||||
|
|
||||||
- **Q1 — composite 발행 비용.** `tl.composite`를 construct + push 하는 비용은
|
ADR-0060 §1 의 핵심: composite 1 개가 `N_tiles` 분량의 GEMM tiling 을
|
||||||
CPU 시간 기준 **~40 ns** 정도(descriptor 빌드 + 큐 push)로 예상된다. 현재는
|
offload → CPU 는 `O(1)` 개의 굵은 cmd 만 issue, primitive 경로는
|
||||||
**0**이다. 훅은 존재하며, 값(그리고 op 종류별 차등)만 빠져 있다.
|
`O(N_tiles × ops/tile)`. issue cost = 0 이면 모델은:
|
||||||
- **Q2 — scheduler dispatch vs DMA latency.** PE_SCHEDULER의 composite
|
- primitive 경로의 *CPU 가 push 를 못 따라가 엔진 idle* 을 못 보임
|
||||||
dispatch는 **non-blocking**이다: `_dispatch_composite`는 tile plan을 생성하고
|
- composite 의 *구성 비용 ≫ primitive 1 개이지만 ≪ 대체된 primitive 들의 합* 을 못 보임
|
||||||
feeder에 enqueue한 뒤 즉시 반환한다(`pe_scheduler.py:104-121`). 실제 **DMA
|
|
||||||
latency는 타일이 흐를 때 PE_DMA에서 부과**되며(`drain_ns =
|
|
||||||
compute_drain_ns(path, nbytes)`, `pe_dma.py:89`), scheduler dispatch에
|
|
||||||
덩어리로 들어가지 **않는다** ⇒ **중복 계상 없음**, DMA는 modeled component에
|
|
||||||
머문다(SPEC §0.1). scheduler의 plan 생성 자체는 현재 sim 시간 **0**이다.
|
|
||||||
|
|
||||||
### 균일-그리고-0이 하이브리드에 틀린 이유
|
### Revision 1 의 op-type calibration 이 과한 이유
|
||||||
|
|
||||||
하이브리드의 논거 전체는 **하나의** composite descriptor가 `N_tiles`만큼의
|
원안은 `cost_table[kind]` (kind = `composite`, `load`, `dot`, `math`, …)
|
||||||
GEMM 타일링을 offload하므로, CPU가 `O(N_tiles × ops/tile)`개의 fine 커맨드
|
형태였습니다. 비용:
|
||||||
대신 `O(1)`개의 coarse 커맨드를 낸다는 것이다. 발행 비용 = 0(그리고 균일)이면
|
- kind 마다 값 필요 (calibration ≥ |kinds|)
|
||||||
모델은 다음을 보여줄 수 없다:
|
- 새 kind 추가 시마다 entry 추가
|
||||||
- CPU가 충분히 빨리 못 밀어넣을 때 primitive 경로가 엔진을 **saturate 못 할 수
|
- 그런데 잡으려던 *ratio* — "composite ≫ primitive, ≪ 대체된 primitive 들의 합" —
|
||||||
있음**(핵심 ADR-0060 §1 주장), 그리고
|
는 *op kind* 가 아니라 *cmd 가 들고 있는 필드 수* 의 함수.
|
||||||
- composite가 단일 primitive보다 construct 비용이 **더 크지만** 그것이 대체하는
|
|
||||||
다수의 primitive보다는 훨씬 작음.
|
|
||||||
|
|
||||||
단일 균일 `dispatch_cycles`로는 이를 표현할 수 없다: 실제 construct 비용에서
|
cmd 의 *byte 발자국* 이 자연 proxy: N OpSpec composite 는 1 OpSpec primitive 의 ~N 배 bytes.
|
||||||
`tl.load` ≪ `tl.composite`이기 때문이다.
|
*고정* 부분 (큐 tail 업데이트, completion 등록, MMIO-급 latency) 은 cmd 마다.
|
||||||
|
둘 합치면: `FIXED + bytes × R`.
|
||||||
|
|
||||||
## Decision
|
## Decision
|
||||||
|
|
||||||
### D1. PE_CPU의 op 종류별 발행 비용 테이블
|
### D1. 구조적 dispatch cost 공식
|
||||||
|
|
||||||
단일 `dispatch_cycles` 스칼라를 **커맨드 종류별 비용 테이블**로 교체하고,
|
PE_SCHEDULER 로 가는 모든 PE command 가 PE_CPU 에서 dispatch cycles 소모:
|
||||||
발행 시점(커맨드가 scheduler로 dispatch되기 전)에 PE_CPU에서 부과한다.
|
|
||||||
greenlet이 이 시간을 지불하며, 이것이 바로 CPU가 얼마나 빨리 work를 밀어넣을
|
|
||||||
수 있는지를 gate한다 — saturation 레버.
|
|
||||||
|
|
||||||
시작 추정치(추후 calibrate — 검토 항목 참조):
|
```
|
||||||
|
dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes × R
|
||||||
|
```
|
||||||
|
|
||||||
| 발행 op | CPU 발행 비용 (construct + push) |
|
- `FIXED_PER_CMD` (cycles/cmd): 큐 tail 업데이트, MMIO-급 RTT,
|
||||||
|
completion-event 등록 — cmd 크기와 무관 고정 비용.
|
||||||
|
- `R` (cycles/byte): scheduler 큐에 cmd 를 직렬화하는 큐-write 대역폭.
|
||||||
|
- `cmd.logical_bytes` (int): cmd 의 *HW-논리* byte 수 — D2 룰로 계산,
|
||||||
|
Python `sys.getsizeof` 가 아님.
|
||||||
|
|
||||||
|
PE_CPU 는 기존 hook 그대로 `PeCpuOverheadCmd(cycles=dispatch_cycles(cmd))`
|
||||||
|
를 dispatch 전에 발행 — 사이클 값만 변경.
|
||||||
|
|
||||||
|
### D2. `logical_bytes` 룰
|
||||||
|
|
||||||
|
각 PE command dataclass 가 `logical_bytes: int` (property) 노출. 룰
|
||||||
|
(HW 친화적, Python overhead 무시):
|
||||||
|
|
||||||
|
| 필드 종류 | bytes |
|
||||||
|---|---|
|
|---|---|
|
||||||
| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 추정) |
|
| cmd framing (cmd 종류 discriminator + completion id 참조) | 4 |
|
||||||
| `tl.load` / `tl.store` (DMA descriptor) | 작음 (수 ns) |
|
| Opcode (op kind enum) | 1 |
|
||||||
| `tl.dot` / MATH / IPCQ send/recv | 작음 (수 ns) |
|
| Enum (scope 등) | 1 |
|
||||||
|
| `TensorHandle` 참조 (address only — shape/dtype 은 descriptor table 가정) | 8 |
|
||||||
|
| Scalar (int/float) | 4 |
|
||||||
|
| Tuple length marker | 1 |
|
||||||
|
|
||||||
요점은 **비율**이다: composite construct ≫ primitive 발행, 그러나 그것이
|
`CompositeCmd` 는 `ops` 와 `rw_handles` 재귀 합산:
|
||||||
대체하는 primitive 발행들의 합 ≪. 정확한 ns는 설정 가능하다.
|
|
||||||
|
|
||||||
### D2. 실행 latency는 엔진에 유지 (Q2 — 변경 없음)
|
```python
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
4 # framing
|
||||||
|
+ 1 + sum(op.logical_bytes for op in self.ops)
|
||||||
|
+ 1 + 8 * len(self.rw_handles)
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
DMA/GEMM/MATH 실행 latency는 오늘처럼 PE_DMA / PE_GEMM / PE_MATH에서 부과된다.
|
`OpSpec`:
|
||||||
발행 비용(D1)은 **CPU 측에만** 추가되며, `drain_ns`를 건드리지 않으므로 중복
|
|
||||||
계상이 없다. 이는 SPEC §0.1(latency는 modeled component에서)을 보존한다.
|
|
||||||
|
|
||||||
### D3. scheduler plan 생성 비용 — 0에서 시작, 재검토
|
```python
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
1 + 1 # opcode + scope
|
||||||
|
+ 1 + 8 * len(self.operands) # named operand handles
|
||||||
|
+ (8 if self.out is not None else 0) # out handle
|
||||||
|
+ 1 + sum(4 for _ in self.extra.values()) # extra scalars
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
PE_SCHEDULER의 tile-plan 생성은 오늘 sim 시간 0이다. 일단 유지한다(지배적
|
(`DmaReadCmd`, `MathCmd` 등도 같은 룰 — dataclass 당 property 1 개,
|
||||||
레버는 CPU 발행 비용 D1); scheduler 측 비용이 유의미하다고 판명되면 기존
|
약 3 줄.)
|
||||||
`overhead_ns` node attr로 노출한다. 여기서 결정하지 않고 검토 항목으로 둔다.
|
|
||||||
|
|
||||||
### D4. 설정 가능한 값; 골든 재생성
|
### D3. Default — 일반 composite ≈ 45 ns 기준
|
||||||
|
|
||||||
비용 테이블은 설정 가능하다(per-topology / node attrs, 기본값은 D1 추정치).
|
anchor: 일반 1-op DMA→GEMM→DMA composite 의
|
||||||
발행 비용을 0이 아니게 하면 **모든** bench latency가 바뀌므로, 골든 latency를
|
`logical_bytes ≈ 52` (framing 4 + GEMM OpSpec 39 + rw_handles 9). dispatch
|
||||||
한 번 **재생성**한다 — ADR-0062의 전역 lazy-load 회귀와 동일한 자세의 모델
|
목표 45 ns. on-die producer→consumer 큐 4 bytes/cycle 가정.
|
||||||
충실도 개선이다.
|
|
||||||
|
```
|
||||||
|
clock = 1 GHz # 1 cycle = 1 ns
|
||||||
|
FIXED_PER_CMD = 32 cycles
|
||||||
|
R = 0.25 cycles/byte
|
||||||
|
```
|
||||||
|
|
||||||
|
검증: `32 + 52 × 0.25 = 45 cycles ≈ 45 ns` ✓
|
||||||
|
|
||||||
|
### D4. topology config override
|
||||||
|
|
||||||
|
default 는 `pe_cpu.py` 에 내장. topology yaml 의 PE 노드 attrs 아래
|
||||||
|
`pe_cost_model:` 절로 override:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
pe:
|
||||||
|
attrs:
|
||||||
|
pe_cost_model:
|
||||||
|
fixed_per_cmd_cycles: 32
|
||||||
|
byte_cycles_recip: 0.25
|
||||||
|
clock_freq_ghz: 1.0
|
||||||
|
```
|
||||||
|
|
||||||
|
누락된 키는 default. dispatch 공식은 PE_CPU init 시
|
||||||
|
`node.attrs["pe_cost_model"]` 에서 읽음.
|
||||||
|
|
||||||
|
### D5. Scope — 무엇이 cost 를 내고 무엇이 안 내나
|
||||||
|
|
||||||
|
| 경로 | dispatch cost? |
|
||||||
|
|---|---|
|
||||||
|
| PE_CPU → PE_SCHEDULER 의 모든 `PeCommand` | **예** |
|
||||||
|
| `PeCpuOverheadCmd` 자체 (이미 cycles 명시) | **아니오** (공식 우회) |
|
||||||
|
| PE_SCHEDULER 가 자동 생성한 Stage (DMA_READ/WRITE/FETCH/STORE) | **아니오** (PE_SCHEDULER 내부) |
|
||||||
|
| 엔진 compute latency (PE_DMA `drain_ns`, GEMM/MATH `_compute_ns`) | **변화 없음** — 엔진에 그대로 (SPEC §0.1) |
|
||||||
|
|
||||||
|
"latency 는 모델링된 컴포넌트에서" invariant 유지 — dispatch cost 는
|
||||||
|
*추가* CPU-side 시간, 엔진 시간에 fold 안 함.
|
||||||
|
|
||||||
|
### D6. 설정 가능 값; goldens 재생성
|
||||||
|
|
||||||
|
issue cost 를 0 → non-zero 로 바꾸면 **모든** bench 의 latency 변화.
|
||||||
|
이 ADR land 시 골든 latency **한 번** 재생성 — ADR-0062 D3 lazy-load 와
|
||||||
|
같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용.
|
||||||
|
|
||||||
## Alternatives
|
## Alternatives
|
||||||
|
|
||||||
### A1. 단일 균일 `dispatch_cycles > 0` 유지
|
### A1. Revision 1 의 op-type calibration 표 유지
|
||||||
|
|
||||||
기각: op construct 비용은 자릿수 단위로 다르다(`tl.load` vs `tl.composite`).
|
기각: calibration 비용이 |kinds| 에 비례, 그리고 표가 잡으려던 *ratio* 는
|
||||||
균일값은 primitive를 과대 청구하거나 composite를 과소 청구하며, 어느 쪽이든
|
구조적으로 cmd 크기의 함수. 구조 공식은 같은 정성적 동작을 N 개가 아닌
|
||||||
하이브리드가 의존하는 composite-vs-primitive 트레이드오프를 왜곡한다.
|
2 개의 calibratable 숫자로 달성.
|
||||||
|
|
||||||
### A2. 발행 비용을 PE_CPU 대신 PE_SCHEDULER에 부과
|
### A2. byte-only 공식 (FIXED 항 없음)
|
||||||
|
|
||||||
기각: saturation 질문은 *"CPU가 엔진을 바쁘게 유지할 만큼 descriptor를 빨리
|
기각. FIXED = 0 이면 opt2 (ADR-0065 Option Y) 가 opt3 를 **이기지 못함**
|
||||||
밀어넣을 수 있는가?"*이며 — 이는 **PE_CPU**의 issue-bandwidth 속성이다.
|
— per-tile dispatch *총 bytes* 가 비슷 (opt3 ≈ 232, opt2 ≈ 380); win 은
|
||||||
scheduler에 부과하면 CPU back-pressure를 모델하지 못한다.
|
전적으로 *per-cmd fixed cost 횟수 감소* 에서 옴. byte-only 는 모델이
|
||||||
|
보여야 할 신호를 지움.
|
||||||
|
|
||||||
### A3. DMA program/setup 시간을 별도의 고정 per-descriptor 비용으로 모델
|
### A3. PE_SCHEDULER 에 dispatch 비용 부과
|
||||||
|
|
||||||
실제 HW는 전송 시간과 별개로 DMA descriptor program 비용을 지불한다. 연기:
|
기각: saturation 질문은 *"CPU 가 descriptor 를 push 하는 속도가 엔진을
|
||||||
초기에는 descriptor-program 비용을 **발행 op**의 D1 비용에 합친다(DMA를
|
계속 바쁘게 할 수 있는가?"* — **PE_CPU** 의 issue-bandwidth 특성.
|
||||||
유발하는 `tl.load` / composite). calibration이 유의미함을 보이면 PE_DMA 고정
|
scheduler 에 부과하면 CPU back-pressure 모델링 안 됨.
|
||||||
setup으로 분리한다. 검토 항목.
|
|
||||||
|
### A4. DMA program/setup 시간을 별도 fixed per-descriptor 비용으로 모델
|
||||||
|
|
||||||
|
연기: 처음엔 descriptor-program 비용을 **issuing op 의** dispatch cost 에
|
||||||
|
fold. calibration 이 분리 필요성을 보이면 PE_DMA fixed setup 으로 분리.
|
||||||
|
|
||||||
## Consequences
|
## Consequences
|
||||||
|
|
||||||
### Positive
|
### Positive
|
||||||
- 하이브리드의 CPU-offload / saturation 이점(ADR-0060 §1)이 구조적일 뿐
|
- hybrid 의 CPU-offload / saturation win (ADR-0060 §1) 이 **측정 가능**,
|
||||||
아니라 **측정 가능**해진다.
|
구조적으로 정직한 모델 (calibration 표 없음).
|
||||||
- composite-vs-primitive 및 타일링 granularity 트레이드오프가 latency에
|
- 새 op_kind 추가 (예: ADR-0065 의 `softmax_merge` 8-step recipe) 가
|
||||||
보여 eval 서사를 가능케 한다.
|
*zero cost* — 같은 공식 자동 적용.
|
||||||
- 하드웨어에 더 충실(issue bandwidth는 실제 병목).
|
- 하드웨어에 더 충실 (queue-head MMIO RTT + queue-write 대역폭).
|
||||||
|
|
||||||
### Negative
|
### Negative
|
||||||
- **모든** bench 골든이 이동 → 한 번의 재생성(D4); CI 골든 fixture 갱신.
|
- **모든** bench goldens 변화 → 일회성 재생성 (D6); CI 골든 fixtures 업데이트.
|
||||||
- op 종류별 값은 calibration 필요; ~40 ns composite 수치는 추정이고
|
- 두 calibration knob (FIXED, R) 필요; default 는 문서화된 가정에 기반 —
|
||||||
primitive는 미지정 — 결과는 숫자만큼만 정확하다(calibrate 전까지 절대
|
HW reference 없는 동안엔 절대 latency 를 잠정 취급, **ratio** 만 방어.
|
||||||
latency를 과대 주장 금지).
|
- 각 PE command dataclass 에 작은 `logical_bytes` property 추가.
|
||||||
- 발행 경로에 비용 테이블 lookup 추가(런타임 영향 미미).
|
|
||||||
|
|
||||||
## Open review items (자율 결정; 검토 시 수정)
|
## Open review items
|
||||||
|
|
||||||
1. **op 종류별 값의 calibration 출처.** composite ≈ 40 ns는 작업 추정치;
|
1. **FIXED 와 R 의 calibration 출처.** default 는 "일반 composite = 45 ns
|
||||||
primitive 발행 비용은 placeholder. *권고:* 문서화된 가정(instruction-issue +
|
+ on-die 큐 4 bytes/cycle"; on-die producer→consumer 큐로 합리.
|
||||||
큐 push)에서 값을 택하고, 실제 reference가 생기기 전까지 절대 latency를
|
HW reference 등장 시 재방문.
|
||||||
잠정으로 취급; **비율**을 방어 가능하게 유지.
|
2. **Scheduler plan-gen 비용.** 0 유지 — D5 가 PE_SCHEDULER 의 plan
|
||||||
2. **scheduler plan-gen 비용 (D3).** *권고:* 초기에는 0 유지; scheduler-bound
|
생성을 dispatch 공식 밖에 둠. scheduler-bound 행동이 보이면 기존
|
||||||
동작을 보이는 workload가 나오면 `overhead_ns`로 노출.
|
`overhead_ns` 로 노출.
|
||||||
3. **DMA program 시간 (A3).** *권고:* 먼저 발행 op의 비용에 합치고; 필요 시
|
3. **Override 위치.** topology yaml 의 PE 노드 attrs 아래 `pe_cost_model:`
|
||||||
PE_DMA setup으로 분리.
|
block — 모든 knob 을 한 곳에, 리뷰 가능.
|
||||||
4. **테이블 위치.** *권고:* 흩어진 node attr이 아니라, 커맨드 종류로 키잉되고
|
4. **경로 parity.** greenlet (`_execute_legacy`, `kernel_runner`) 와
|
||||||
per-topology 오버라이드 가능한 작은 중앙 cost 모듈 — 값을 한 곳에서 검토
|
replay 둘 다 같은 cost model 읽어야 함. 검증.
|
||||||
가능하도록.
|
|
||||||
5. **회귀 롤아웃.** *권고:* lazy-load(ADR-0062)와 본 ADR의 cost model을 한
|
|
||||||
번의 골든 재생성 패스로 함께 도입하여 두 번의 churn을 피함.
|
|
||||||
6. **legacy replay 경로와의 상호작용**(`pe_cpu.py:_execute_legacy`) — greenlet과
|
|
||||||
replay 경로가 동일한 cost 테이블을 읽어 결과가 일치하도록 보장. 검증 요.
|
|
||||||
|
|
||||||
## Test Requirements
|
## Test Requirements
|
||||||
|
|
||||||
1. **composite는 한 번, primitive는 op마다 청구.** `N_tiles`에 걸쳐 하나의
|
1. **Anchor 보존.** 일반 DMA→GEMM→DMA composite (1 op,
|
||||||
`tl.composite`를 내는 커널은 PE_CPU에서 composite 발행 비용을 한 번 청구하고;
|
`logical_bytes ≈ 52`) 가 default 값에서 45 ns 에 dispatch.
|
||||||
동등한 `N_tiles × ops` primitive 커널은 op당 비용을 `N_tiles × ops`번 청구.
|
2. **구조적 ratio.** opt3 vs opt2 dispatch (ADR-0065 §verification 의):
|
||||||
PE_CPU busy time이 그에 따라 다름을 assert.
|
default calibration 에서 `opt3 / opt2 ≈ 2.4×`.
|
||||||
2. **DMA latency 불변 (Q2).** 고정된 전송에 대해 PE_DMA `drain_ns`는 ADR 이전
|
3. **Override 경로.** topology yaml 의 `pe_cost_model:` block 이 per-PE
|
||||||
값과 동일 — 발행 비용은 DMA에 합쳐지지 않고 PE_CPU에 가산(중복 계상 없음).
|
dispatch cost 변경; block 누락 시 default 복귀.
|
||||||
3. **saturation 관측 가능.** per-op 발행 비용이 0이 아니면, many-tile primitive
|
4. **`PeCpuOverheadCmd` 우회.** 수동 `tl.cycles(n)` 는 정확히 `n` cycles,
|
||||||
sweep는 GEMM 엔진 idle(CPU-bound 발행)을 보이고 composite sweep는 바쁘게
|
`n + dispatch_cycles(...)` 아님.
|
||||||
유지 — ADR-0060 §1 레버.
|
5. **double-count 없음.** PE_DMA `drain_ns`, PE_GEMM/MATH `_compute_ns` 가
|
||||||
4. **결정성:** 동일 입력 → 동일 op_log + latency (SPEC §0.1).
|
pre-ADR 값과 동일.
|
||||||
5. **경로 일치:** greenlet과 legacy-replay 경로가 동일 커널에 대해 동일한 발행
|
6. **결정성.** 동일 입력 → 동일 op_log + latency (SPEC §0.1).
|
||||||
비용 회계를 산출.
|
7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일
|
||||||
|
dispatch-cycle 회계.
|
||||||
|
|
||||||
|
## Migration
|
||||||
|
|
||||||
|
ADR-0064 Revision 2 는 단일 PR 로 land:
|
||||||
|
- 각 `PeCommand` dataclass 의 `logical_bytes` property
|
||||||
|
- `pe_cpu.py` dispatch 경로의 공식 적용
|
||||||
|
- PE_CPU init 의 `pe_cost_model:` override read
|
||||||
|
- 일회성 골든 재생성
|
||||||
|
|
||||||
|
land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음
|
||||||
|
(ADR-0065 는 기존 경로의 `CompositeCmd` 의미 보존 refactor; opt2 만 새 bench).
|
||||||
|
|||||||
@@ -1,179 +1,258 @@
|
|||||||
# ADR-0064: Per-op-type CPU issue cost model (command construct + dispatch)
|
# ADR-0064: Structural CPU dispatch cost model (`logical_bytes` + FIXED + R)
|
||||||
|
|
||||||
## Status
|
## Status
|
||||||
|
|
||||||
Proposed
|
Proposed (Revision 2)
|
||||||
|
|
||||||
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). The hybrid
|
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065**
|
||||||
> decision there (GEMMs via `tl.composite`, softmax merge in the kernel)
|
> (flat-ops composite + first stateful recipe). The hybrid decision there
|
||||||
> wins by **offloading tiling to PE_SCHEDULER so the CPU issues coarse
|
> (GEMMs via `tl.composite`, softmax merge in the kernel) wins by
|
||||||
> descriptors and runs ahead, keeping the engines saturated**. That win is
|
> **offloading tiling to PE_SCHEDULER so the CPU issues coarse descriptors
|
||||||
> currently **invisible in the simulator** because per-op CPU issue cost is
|
> and runs ahead, keeping the engines saturated**. That win is currently
|
||||||
> zero. This ADR makes the issue cost real and **op-type-differentiated**,
|
> **invisible in the simulator** because per-op CPU issue cost is zero.
|
||||||
> so composite-vs-primitive trade-offs (and CPU saturation) are measurable.
|
>
|
||||||
|
> Revision 2 replaces the **op-type calibration table** (the original
|
||||||
|
> proposal) with a **structural formula** derived from each command's
|
||||||
|
> `logical_bytes` — no per-op-type calibration needed; new op kinds are
|
||||||
|
> covered automatically.
|
||||||
|
|
||||||
## Context
|
## Context
|
||||||
|
|
||||||
### What exists today
|
### What exists today
|
||||||
|
|
||||||
- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its
|
- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its
|
||||||
command (`tl_context.py:123-125, 190, 227, 235, 612, …`), which emits
|
command (`tl_context.py:196-212`), which emits
|
||||||
`PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles
|
`PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles
|
||||||
> 0`. So the issue cost is a **single uniform knob** applied identically
|
> 0`. The knob is **uniform** across op kinds and hardcoded to **0**
|
||||||
to `tl.load`, `tl.dot`, a MATH op, and `tl.composite`.
|
in both live execution paths (`pe_cpu.py:101` greenlet, `:195` replay).
|
||||||
- That knob is hardcoded to **0** in both live execution paths
|
- ⇒ issuing a command — *constructing the descriptor and pushing it to
|
||||||
(`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ issuing a
|
the scheduler queue* — currently costs **0 ns** on PE_CPU.
|
||||||
command — *constructing the descriptor and pushing it to the scheduler
|
|
||||||
queue* — currently costs **0 ns** on PE_CPU.
|
|
||||||
- `PeCpuOverheadCmd` is consumed as `yield env.timeout(cmd.cycles)` on
|
- `PeCpuOverheadCmd` is consumed as `yield env.timeout(cmd.cycles)` on
|
||||||
PE_CPU (`kernel_runner.py:131-132`) and on PE_SCHEDULER
|
PE_CPU (`kernel_runner.py:131-132`).
|
||||||
(`pe_scheduler.py:97-100`); a manual `tl.cycles(n)` also exists
|
|
||||||
(`tl_context.py:695`).
|
|
||||||
|
|
||||||
### Two findings that frame this ADR (from ADR-0060 review)
|
|
||||||
|
|
||||||
- **Q1 — composite issue cost.** Constructing + pushing a `tl.composite`
|
|
||||||
is expected to cost on the order of **~40 ns** of CPU time (descriptor
|
|
||||||
build + queue push). Today it is **0**. The hook exists; only the value
|
|
||||||
(and its per-op-type differentiation) is missing.
|
|
||||||
- **Q2 — scheduler dispatch vs DMA latency.** `PE_SCHEDULER`'s composite
|
|
||||||
dispatch is **non-blocking**: `_dispatch_composite` generates the tile
|
|
||||||
plan and enqueues to the feeder, returning immediately
|
|
||||||
(`pe_scheduler.py:104-121`). The actual **DMA latency is charged on
|
|
||||||
PE_DMA** as each tile flows through (`drain_ns = compute_drain_ns(path,
|
|
||||||
nbytes)`, `pe_dma.py:89`), **not** lumped into the scheduler dispatch ⇒
|
|
||||||
**no double-counting**, and DMA stays on its modelled component
|
|
||||||
(SPEC §0.1). The scheduler's own plan-generation currently costs **0**
|
|
||||||
sim time.
|
|
||||||
|
|
||||||
### Why uniform-and-zero is wrong for the hybrid
|
### Why uniform-and-zero is wrong for the hybrid
|
||||||
|
|
||||||
The hybrid's whole argument is that **one** composite descriptor offloads
|
ADR-0060 §1's argument is that **one** composite descriptor offloads
|
||||||
`N_tiles` worth of GEMM tiling, so the CPU issues `O(1)` coarse commands
|
`N_tiles` worth of GEMM tiling, so the CPU issues `O(1)` coarse commands
|
||||||
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0 (and
|
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0, the
|
||||||
uniform), the model cannot show:
|
model cannot show:
|
||||||
- that the primitive path may **fail to saturate** the engines when the
|
- that the primitive path may **fail to saturate** the engines when the
|
||||||
CPU cannot push fast enough (the core ADR-0060 §1 claim), nor
|
CPU cannot push fast enough, nor
|
||||||
- that a composite **costs more to construct** than a single primitive but
|
- that a composite **costs more to construct** than a single primitive
|
||||||
far less than the many primitives it replaces.
|
but far less than the many primitives it replaces.
|
||||||
|
|
||||||
A single uniform `dispatch_cycles` cannot express this: `tl.load` ≪
|
### Why per-op-type calibration (Revision 1) was over-shaped
|
||||||
`tl.composite` in real construct cost.
|
|
||||||
|
The original proposal had a `cost_table[kind]` keyed by op kind
|
||||||
|
(`composite`, `load`, `dot`, `math`, …). That required:
|
||||||
|
- a value per kind (calibration cost ≥ |kinds|),
|
||||||
|
- a new entry every time a new kind appears,
|
||||||
|
- and yet the *ratio* it tried to capture — "composite ≫ primitive,
|
||||||
|
but ≪ the primitives it replaces" — is structurally a function of
|
||||||
|
**how many fields the command carries**, not of the op kind.
|
||||||
|
|
||||||
|
A command's *byte footprint* is the natural proxy: a composite carrying
|
||||||
|
N OpSpecs has ~N× the bytes of a primitive op with one OpSpec. The
|
||||||
|
*fixed* part (queue head update, completion register, MMIO-class
|
||||||
|
latency) is per-command. The two together compose: `FIXED + bytes × R`.
|
||||||
|
|
||||||
## Decision
|
## Decision
|
||||||
|
|
||||||
### D1. Per-op-type issue-cost table on PE_CPU
|
### D1. Structural dispatch cost formula
|
||||||
|
|
||||||
Replace the single `dispatch_cycles` scalar with a **cost table keyed by
|
Each PE command going to PE_SCHEDULER incurs PE_CPU dispatch cycles:
|
||||||
command type**, charged on PE_CPU at issue (before the command is
|
|
||||||
dispatched to the scheduler). The greenlet pays this time, which is
|
|
||||||
exactly what gates how fast the CPU can push work — the saturation lever.
|
|
||||||
|
|
||||||
Starting estimates (calibrate later — see review items):
|
```
|
||||||
|
dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes × R
|
||||||
|
```
|
||||||
|
|
||||||
| Issued op | CPU issue cost (construct + push) |
|
where:
|
||||||
|
- `FIXED_PER_CMD` (cycles per command) models queue-tail update, MMIO-class
|
||||||
|
RTT, completion-event registration — fixed per command regardless of size.
|
||||||
|
- `R` (cycles per byte) models the queue-write bandwidth — bytes of the
|
||||||
|
command serialized into the scheduler queue.
|
||||||
|
- `cmd.logical_bytes` (int) is each command's *HW-logical* byte size,
|
||||||
|
computed from D2 below — not Python's `sys.getsizeof`.
|
||||||
|
|
||||||
|
PE_CPU emits `PeCpuOverheadCmd(cycles=dispatch_cycles(cmd))` before
|
||||||
|
dispatching, exactly as the existing hook (`tl_context.py:_emit_dispatch_
|
||||||
|
overhead`) — only the cycle value changes.
|
||||||
|
|
||||||
|
### D2. `logical_bytes` rule
|
||||||
|
|
||||||
|
Each PE command dataclass exposes `logical_bytes: int` (property). The
|
||||||
|
counting rule (HW-friendly, ignores Python overhead):
|
||||||
|
|
||||||
|
| Field kind | Bytes |
|
||||||
|---|---|
|
|---|---|
|
||||||
| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 estimate) |
|
| Command framing (cmd-type discriminator + completion id ref) | 4 |
|
||||||
| `tl.load` / `tl.store` (DMA descriptor) | small (a few ns) |
|
| Opcode (op kind enum) | 1 |
|
||||||
| `tl.dot` / MATH / IPCQ send/recv | small (a few ns) |
|
| Enum (scope, etc.) | 1 |
|
||||||
|
| `TensorHandle` reference (address only — shape/dtype assumed in descriptor table) | 8 |
|
||||||
|
| Scalar (int/float) | 4 |
|
||||||
|
| Tuple length marker | 1 |
|
||||||
|
|
||||||
The point is the **ratio**: a composite construct ≫ a primitive issue, but
|
`CompositeCmd` recursively sums its `ops` and `rw_handles`:
|
||||||
≪ the sum of the primitive issues it replaces. Exact ns are configurable.
|
|
||||||
|
|
||||||
### D2. Execution latency stays on the engines (Q2 — no change)
|
```python
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
4 # framing
|
||||||
|
+ 1 + sum(op.logical_bytes for op in self.ops)
|
||||||
|
+ 1 + 8 * len(self.rw_handles)
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
DMA/GEMM/MATH execution latency remains charged on PE_DMA / PE_GEMM /
|
`OpSpec`:
|
||||||
PE_MATH as today. The issue cost (D1) is **CPU-side only** and additive; it
|
|
||||||
does **not** touch `drain_ns`, so there is no double-count. This preserves
|
|
||||||
SPEC §0.1 (latency on modelled components).
|
|
||||||
|
|
||||||
### D3. Scheduler plan-generation cost — start at 0, revisit
|
```python
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
1 + 1 # opcode + scope
|
||||||
|
+ 1 + 8 * len(self.operands) # named operand handles
|
||||||
|
+ (8 if self.out is not None else 0) # out handle
|
||||||
|
+ 1 + sum(4 for _ in self.extra.values()) # extra scalars
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
`PE_SCHEDULER`'s tile-plan generation costs 0 sim time today. Keep that for
|
(Identical rule for `DmaReadCmd`, `MathCmd`, etc. — one property per
|
||||||
now (the dominant lever is CPU issue cost, D1); expose it later via the
|
dataclass, ~3 lines each.)
|
||||||
existing `overhead_ns` node attr if a scheduler-side cost proves material.
|
|
||||||
Marked as a review item, not decided here.
|
|
||||||
|
|
||||||
### D4. Configurable values; goldens regenerate
|
### D3. Defaults — anchored on a typical composite ≈ 45 ns
|
||||||
|
|
||||||
The cost table is configurable (per-topology / node attrs, default to the
|
Anchor: a typical 1-op DMA→GEMM→DMA composite has `logical_bytes ≈ 52`
|
||||||
D1 estimates). Turning issue cost non-zero changes **every** bench's
|
(framing 4 + GEMM OpSpec 39 + rw_handles 9). Target dispatch = 45 ns.
|
||||||
latency, so golden latencies are **regenerated** once — a model-fidelity
|
On-die producer→consumer queue assumed at 4 bytes/cycle.
|
||||||
improvement, same posture as ADR-0062's global lazy-load regression.
|
|
||||||
|
```
|
||||||
|
clock = 1 GHz # 1 cycle = 1 ns
|
||||||
|
FIXED_PER_CMD = 32 cycles
|
||||||
|
R = 0.25 cycles/byte
|
||||||
|
```
|
||||||
|
|
||||||
|
Verification: `32 + 52 × 0.25 = 45 cycles ≈ 45 ns` ✓
|
||||||
|
|
||||||
|
### D4. Topology config override
|
||||||
|
|
||||||
|
Defaults are baked into `pe_cpu.py`. Topology yaml may override under a
|
||||||
|
`pe_cost_model:` section at the PE node attrs:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
pe:
|
||||||
|
attrs:
|
||||||
|
pe_cost_model:
|
||||||
|
fixed_per_cmd_cycles: 32
|
||||||
|
byte_cycles_recip: 0.25
|
||||||
|
clock_freq_ghz: 1.0
|
||||||
|
```
|
||||||
|
|
||||||
|
Missing keys fall back to defaults. The dispatch formula reads from
|
||||||
|
`node.attrs["pe_cost_model"]` at PE_CPU init.
|
||||||
|
|
||||||
|
### D5. Scope — what does and does not pay
|
||||||
|
|
||||||
|
| Path | Pays dispatch cost? |
|
||||||
|
|---|---|
|
||||||
|
| PE_CPU → PE_SCHEDULER for any `PeCommand` | **Yes** |
|
||||||
|
| `PeCpuOverheadCmd` itself (already cycles-explicit) | **No** (formula bypass) |
|
||||||
|
| Stages auto-generated by PE_SCHEDULER (DMA_READ/WRITE/FETCH/STORE) | **No** (PE_SCHEDULER-internal) |
|
||||||
|
| Engine compute latency (DMA `drain_ns`, GEMM/MATH `_compute_ns`) | **No change** — stays on engines (SPEC §0.1) |
|
||||||
|
|
||||||
|
This preserves the "latency on modelled components" invariant — dispatch
|
||||||
|
cost is *additional* CPU-side time, not folded into engine times.
|
||||||
|
|
||||||
|
### D6. Configurable values; goldens regenerate
|
||||||
|
|
||||||
|
Turning issue cost non-zero changes **every** bench's latency. Golden
|
||||||
|
latencies are **regenerated once** when this ADR lands — same posture as
|
||||||
|
ADR-0062 D3 lazy-load. After regeneration, the same calibration is in
|
||||||
|
effect for ADR-0065 opt2 measurement.
|
||||||
|
|
||||||
## Alternatives
|
## Alternatives
|
||||||
|
|
||||||
### A1. Keep a single uniform `dispatch_cycles > 0`
|
### A1. Keep Revision 1's op-type calibration table
|
||||||
|
|
||||||
Rejected: op construct costs differ by an order of magnitude (`tl.load` vs
|
Rejected: calibration cost scales with |kinds|, and the *ratio* the
|
||||||
`tl.composite`); a uniform value either over-charges primitives or
|
table tried to capture is structurally a function of cmd size. The
|
||||||
under-charges composites, and in both cases misrepresents the
|
structural formula reaches the same qualitative behaviour with two
|
||||||
composite-vs-primitive trade-off the hybrid depends on.
|
calibratable numbers instead of N.
|
||||||
|
|
||||||
### A2. Charge the issue cost on PE_SCHEDULER instead of PE_CPU
|
### A2. Byte-only formula (no FIXED term)
|
||||||
|
|
||||||
|
Rejected. With FIXED = 0, opt2 (Option Y per ADR-0065) does **not** win
|
||||||
|
over opt3 — the total *bytes* dispatched per tile are similar (opt3 ≈ 232,
|
||||||
|
opt2 ≈ 380); the win is entirely in *fewer per-cmd fixed costs*. A
|
||||||
|
byte-only formula erases the very signal the model needs to expose.
|
||||||
|
|
||||||
|
### A3. Charge dispatch on PE_SCHEDULER instead of PE_CPU
|
||||||
|
|
||||||
Rejected: the saturation question is *"can the CPU push descriptors fast
|
Rejected: the saturation question is *"can the CPU push descriptors fast
|
||||||
enough to keep the engines busy?"* — that is a **PE_CPU** issue-bandwidth
|
enough to keep the engines busy?"* — that is a **PE_CPU** issue-bandwidth
|
||||||
property. Charging it on the scheduler would not model CPU back-pressure.
|
property. Charging on the scheduler would not model CPU back-pressure.
|
||||||
|
|
||||||
### A3. Model DMA program/setup time as a separate fixed per-descriptor cost
|
### A4. Model DMA program/setup time as a separate fixed per-descriptor cost
|
||||||
|
|
||||||
Real HW pays a DMA descriptor program cost distinct from transfer time.
|
|
||||||
Deferred: initially fold the descriptor-program cost into the **issuing
|
Deferred: initially fold the descriptor-program cost into the **issuing
|
||||||
op's** D1 cost (a `tl.load` / composite that triggers DMA). Split it out to
|
op's** dispatch cost. Split it out to a PE_DMA fixed setup only if
|
||||||
a PE_DMA fixed setup only if calibration shows it matters. Review item.
|
calibration shows it matters.
|
||||||
|
|
||||||
## Consequences
|
## Consequences
|
||||||
|
|
||||||
### Positive
|
### Positive
|
||||||
- The hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
|
- Hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
|
||||||
**measurable**, not just structural.
|
**measurable**, with a structurally honest model (no calibration table).
|
||||||
- Composite-vs-primitive and tiling-granularity trade-offs are visible in
|
- Adding new op kinds (e.g., ADR-0065's `softmax_merge` 8-step recipe)
|
||||||
latency, enabling the eval narrative.
|
costs zero — they fit the same formula automatically.
|
||||||
- More faithful to hardware (issue bandwidth is a real bottleneck).
|
- More faithful to hardware (queue-head MMIO RTT + queue-write bandwidth).
|
||||||
|
|
||||||
### Negative
|
### Negative
|
||||||
- **All** bench goldens shift → one-time regeneration (D4); CI golden
|
- **All** bench goldens shift → one-time regeneration (D6); CI golden
|
||||||
fixtures update.
|
fixtures update.
|
||||||
- Per-op-type values need calibration; the ~40 ns composite figure is an
|
- Two calibration knobs (FIXED, R) need values; defaults are anchored on
|
||||||
estimate, primitives are unspecified — results are only as good as the
|
a documented assumption — treat absolute latencies as provisional until
|
||||||
numbers (do not over-claim absolute latencies until calibrated).
|
a reference exists; keep the **ratios** defensible.
|
||||||
- Adds a cost-table lookup on the issue path (negligible runtime).
|
- Adds a small `logical_bytes` property to each PE command dataclass.
|
||||||
|
|
||||||
## Open review items (decided autonomously; revise on review)
|
## Open review items
|
||||||
|
|
||||||
1. **Calibration source for the per-op-type values.** Composite ≈ 40 ns is
|
1. **Calibration source for FIXED and R.** Defaults from "typical
|
||||||
a working estimate; primitive issue costs are placeholders. *Recommend:*
|
composite = 45 ns + on-die queue 4 bytes/cycle"; reasonable for an
|
||||||
pick values from a documented assumption (instruction-issue + queue
|
on-die producer→consumer queue. Revisit when a HW reference appears.
|
||||||
push) and treat absolute latency as provisional until a real reference
|
2. **Scheduler plan-gen cost.** Stays 0 — D5 keeps PE_SCHEDULER's
|
||||||
exists; keep the **ratios** defensible.
|
plan-generation outside the dispatch formula. Expose via existing
|
||||||
2. **Scheduler plan-gen cost (D3).** *Recommend:* keep 0 initially; expose
|
`overhead_ns` if a workload shows scheduler-bound behaviour.
|
||||||
via `overhead_ns` if a workload shows scheduler-bound behaviour.
|
3. **Where the override lives.** `pe_cost_model:` block under PE node
|
||||||
3. **DMA program time (A3).** *Recommend:* fold into the issuing op's cost
|
attrs in topology yaml — keeps all knobs in one place, reviewable.
|
||||||
first; split to PE_DMA setup only if needed.
|
4. **Path parity.** Both greenlet (`_execute_legacy` and `kernel_runner`)
|
||||||
4. **Where the table lives.** *Recommend:* a small central cost module
|
and replay paths must read the same cost model. Verify.
|
||||||
keyed by command type, overridable per topology — not scattered node
|
|
||||||
attrs — so the values are reviewable in one place.
|
|
||||||
5. **Regression rollout.** *Recommend:* land lazy-load (ADR-0062) and this
|
|
||||||
ADR's cost model in one golden-regeneration pass to avoid two churns.
|
|
||||||
6. **Interaction with the legacy replay path** (`pe_cpu.py:_execute_legacy`)
|
|
||||||
— ensure both the greenlet and replay paths read the same cost table so
|
|
||||||
results match. Verify.
|
|
||||||
|
|
||||||
## Test Requirements
|
## Test Requirements
|
||||||
|
|
||||||
1. **Composite charges once; primitives charge per-op.** A kernel issuing
|
1. **Anchor preservation.** A typical DMA→GEMM→DMA composite (1 op,
|
||||||
one `tl.composite` over `N_tiles` charges one composite issue cost on
|
`logical_bytes ≈ 52`) dispatches in 45 ns at the default values.
|
||||||
PE_CPU; the equivalent `N_tiles × ops` primitive kernel charges the
|
2. **Structural ratio.** opt3 vs opt2 dispatch (per ADR-0065 §verification):
|
||||||
per-op cost `N_tiles × ops` times. Assert the PE_CPU busy time differs
|
`opt3 / opt2 ≈ 2.4×` at default calibration.
|
||||||
accordingly.
|
3. **Override path.** Topology yaml `pe_cost_model:` block changes the
|
||||||
2. **DMA latency unchanged (Q2).** For a fixed transfer, PE_DMA `drain_ns`
|
per-PE dispatch cost; default is recovered when block is missing.
|
||||||
is identical to the pre-ADR value — the issue cost is additive on
|
4. **`PeCpuOverheadCmd` bypass.** Manual `tl.cycles(n)` issues exactly `n`
|
||||||
PE_CPU, not folded into DMA (no double-count).
|
cycles, not `n + dispatch_cycles(...)`.
|
||||||
3. **Saturation is observable.** With non-zero per-op issue cost, a
|
5. **No double-count.** PE_DMA `drain_ns`, PE_GEMM/MATH `_compute_ns`
|
||||||
many-tile primitive sweep shows GEMM-engine idle (CPU-bound issue)
|
identical to pre-ADR values.
|
||||||
whereas the composite sweep keeps it busy — the ADR-0060 §1 lever.
|
6. **Determinism.** Identical inputs → identical op_log + latency
|
||||||
4. **Determinism:** identical inputs → identical op_log + latency
|
|
||||||
(SPEC §0.1).
|
(SPEC §0.1).
|
||||||
5. **Path parity:** greenlet and legacy-replay paths produce identical
|
7. **Path parity.** Greenlet and replay paths produce identical
|
||||||
issue-cost accounting for the same kernel.
|
dispatch-cycle accounting for the same kernel.
|
||||||
|
|
||||||
|
## Migration
|
||||||
|
|
||||||
|
ADR-0064 Revision 2 lands as a single PR with:
|
||||||
|
- `logical_bytes` property on each `PeCommand` dataclass
|
||||||
|
- formula application in `pe_cpu.py` dispatch path
|
||||||
|
- `pe_cost_model:` override read at PE_CPU init
|
||||||
|
- one-time goldens regeneration
|
||||||
|
|
||||||
|
After this lands, ADR-0065 builds on top with no further goldens churn
|
||||||
|
in existing benches (ADR-0065 is a meaning-preserving refactor of
|
||||||
|
`CompositeCmd` for the existing path; only opt2 is a new bench).
|
||||||
|
|||||||
@@ -0,0 +1,339 @@
|
|||||||
|
# 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 배치 결정
|
||||||
|
|
||||||
|
PE_SCHEDULER 가 `cmd.ops` 에서 GEMM op 검색 (composite 당 ≤ 1, D6 참조).
|
||||||
|
인덱스 `g`:
|
||||||
|
|
||||||
|
| OpSpec position | scope | plan 내 배치 |
|
||||||
|
|---|---|---|
|
||||||
|
| `0 .. g-1` | KERNEL | pre-loop stages (tile loop 진입 전 1 회 실행) |
|
||||||
|
| `g` | KERNEL | head GEMM (`extra["m"], extra["k"], extra["n"]` 로 tile loop drive) |
|
||||||
|
| `g+1 .. ` | K_TILE | per K-tile epilogue (누산 loop 안) |
|
||||||
|
| `g+1 .. ` | OUTPUT_TILE | per (m, n) tile epilogue (K 누산 후) |
|
||||||
|
| `g+1 .. ` | KERNEL | post-loop stages (tile loop 종료 후 1 회) |
|
||||||
|
|
||||||
|
GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope
|
||||||
|
직렬로 처리.
|
||||||
|
|
||||||
|
### D4. PE_SCHEDULER 가 operand `space` 보고 DMA 자동 삽입
|
||||||
|
|
||||||
|
PE_SCHEDULER 가 GEMM 의 `operands` 와 `out` 검사. 각각:
|
||||||
|
- `space == "hbm"` → FETCH/GEMM Stage 앞에 DMA_READ Stage 삽입 (`out` 은 DMA_WRITE)
|
||||||
|
- `space == "tcm"` → DMA Stage 없음, in-place 소비
|
||||||
|
|
||||||
|
오늘 이미 `a_pinned`/`b_pinned` 로 부분 모델링; D4 가 룰을 균일화 +
|
||||||
|
가시화: **커널은 composite operand 의 명시적 DMA cmd 를 절대 emit 안 함;
|
||||||
|
PE_SCHEDULER 가 핸들에서 추론**. `tl.ref(addr, shape)` → `space="hbm"`;
|
||||||
|
`tl.zeros` / `tl.full` / `tl.load` 결과 → `space="tcm"`.
|
||||||
|
|
||||||
|
비 GEMM op (prologue/epilogue 의 MATH chain) 은 TCM-resident operand
|
||||||
|
(커널이 `tl.zeros` / `tl.full` 로 할당); DMA 삽입 없음.
|
||||||
|
|
||||||
|
### 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. head GEMM 의 `operands["a"] = P_handle` auto-bind.
|
||||||
|
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 별 코드 분기 없음.
|
||||||
|
|
||||||
|
### 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 관련 없음
|
||||||
|
```
|
||||||
|
|
||||||
|
## 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.
|
||||||
|
|
||||||
|
## 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. **`space` 에서 DMA 자동 삽입.** `b=tl.ref(...)` (space=hbm) 의 GEMM
|
||||||
|
composite 가 DMA_READ Stage emit; `b` 가 `tl.zeros(...)` (space=tcm) 면
|
||||||
|
emit 안 함. 출력 핸들 `space=hbm` 이면 DMA_WRITE emit; `space=tcm` 이면
|
||||||
|
emit 안 함.
|
||||||
|
6. **opt2 가 opt3 와 수치 동등.** data mode 에서 opt2 의 최종
|
||||||
|
`(m, l, O)` 가 opt3 와 fp tolerance 안.
|
||||||
|
7. **opt2 dispatch ratio (ADR-0064 Rev2 이후).** opt3 vs opt2 per-tile
|
||||||
|
PE_CPU dispatch cycles ratio 가 default calibration 에서 ≈ 2.4×.
|
||||||
|
8. **GEMM-count invariant.** GEMM OpSpec 두 개 가진 composite 가
|
||||||
|
TLContext emit 시 validation error.
|
||||||
|
|
||||||
|
## 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
|
||||||
|
|
||||||
|
Land 순서:
|
||||||
|
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** 참조.
|
||||||
@@ -0,0 +1,374 @@
|
|||||||
|
# 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
|
||||||
|
|
||||||
|
PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see
|
||||||
|
D6). Calling its index `g`:
|
||||||
|
|
||||||
|
| OpSpec position | scope | Placement in plan |
|
||||||
|
|---|---|---|
|
||||||
|
| `0 .. g-1` | KERNEL | pre-loop stages (run once before tile loop) |
|
||||||
|
| `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 .. ` | OUTPUT_TILE | per (m, n) tile epilogue (after K accumulation) |
|
||||||
|
| `g+1 .. ` | KERNEL | post-loop stages (run once after tile loop) |
|
||||||
|
|
||||||
|
A composite with no GEMM op (e.g., MATH-only composite) treats all ops
|
||||||
|
as KERNEL-scope sequential.
|
||||||
|
|
||||||
|
### D4. PE_SCHEDULER auto-inserts DMAs from operand `space`
|
||||||
|
|
||||||
|
PE_SCHEDULER scans the GEMM's `operands` and `out`. For each:
|
||||||
|
- `space == "hbm"` → emit DMA_READ Stage (or DMA_WRITE for `out`) before
|
||||||
|
the FETCH/GEMM Stages
|
||||||
|
- `space == "tcm"` → no DMA Stage, operand consumed in place
|
||||||
|
|
||||||
|
This is already partially modelled today via `a_pinned`/`b_pinned`; D4
|
||||||
|
makes the rule uniform and visible: **the kernel never emits explicit
|
||||||
|
DMA cmds for composite operands; PE_SCHEDULER infers them from
|
||||||
|
handles**. `tl.ref(addr, shape)` returns a handle with `space="hbm"`;
|
||||||
|
`tl.zeros` / `tl.full` / `tl.load` outputs return `space="tcm"`.
|
||||||
|
|
||||||
|
Non-GEMM ops (MATH chain in prologue / epilogue) have TCM-resident
|
||||||
|
operands (kernel allocates via `tl.zeros` / `tl.full`); no DMA inserted
|
||||||
|
for them.
|
||||||
|
|
||||||
|
### 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 head GEMM's `operands["a"] = P_handle`.
|
||||||
|
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.
|
||||||
|
|
||||||
|
### 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
|
||||||
|
```
|
||||||
|
|
||||||
|
## 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.
|
||||||
|
|
||||||
|
## 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 `space`.** A GEMM composite with
|
||||||
|
`b=tl.ref(...)` (space=hbm) emits a DMA_READ Stage; with `b` from
|
||||||
|
`tl.zeros(...)` (space=tcm) it does not. Output handle with
|
||||||
|
`space=hbm` emits DMA_WRITE; `space=tcm` does not.
|
||||||
|
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).** opt3 vs opt2 per-tile
|
||||||
|
PE_CPU dispatch cycles ratio ≈ 2.4× at default calibration.
|
||||||
|
8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a
|
||||||
|
validation error at TLContext emit.
|
||||||
|
|
||||||
|
## 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
|
||||||
|
|
||||||
|
Land order:
|
||||||
|
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**.
|
||||||
@@ -0,0 +1,524 @@
|
|||||||
|
# Detailed Design Document — Flat-ops `CompositeCmd` + `softmax_merge` recipe (ADR-0065)
|
||||||
|
|
||||||
|
**Status:** Draft (companion to ADR-0065, Proposed)
|
||||||
|
**Scope:** implementation-ready plan for the `CompositeCmd` flat-ops
|
||||||
|
restructure, the `softmax_merge` recipe, and the decode opt2 variant in
|
||||||
|
`_gqa_decode_long.py`.
|
||||||
|
**Audience:** reviewer / implementer resuming after the ADR-0065 design
|
||||||
|
round. Read **ADR-0065 first** (boundary D5 + D7 set the layering), and
|
||||||
|
**ADR-0064 Revision 2** (dispatch cost model that opt2's win depends on).
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 1. Goal and success criteria
|
||||||
|
|
||||||
|
**Goal:** make decode opt2 *runnable*, *numerically equivalent to opt3*,
|
||||||
|
and *measurably cheaper in PE_CPU dispatch* under the ADR-0064 Rev2 cost
|
||||||
|
model. Existing benches keep their (post-ADR-0064 Rev2) goldens
|
||||||
|
unchanged.
|
||||||
|
|
||||||
|
**Success criteria:**
|
||||||
|
1. **Refactor safety.** `CompositeCmd` flat-ops restructure preserves the
|
||||||
|
op_log byte-for-byte for every existing bench
|
||||||
|
(`tests/<existing>` all green, no goldens diff).
|
||||||
|
2. **Recipe correctness.** opt2 final `(m, l, O)` matches opt3 within fp
|
||||||
|
tolerance in data mode (`enable_data=True`).
|
||||||
|
3. **Boundary preservation.** PE_SCHEDULER imports nothing recipe-related;
|
||||||
|
PE_MATH/PE_GEMM/PE_DMA see only `Stage.params["op_kind"]` (no new
|
||||||
|
engine code).
|
||||||
|
4. **Dispatch ratio.** Per-tile PE_CPU dispatch cycles
|
||||||
|
`ratio(opt3 / opt2) ≈ 2.4×` at default ADR-0064 calibration.
|
||||||
|
5. **K-before-V invariant.** op_log of decode opt2 shows zero V-related
|
||||||
|
DMA during #2's MATH prologue.
|
||||||
|
|
||||||
|
**Non-goals:** prefill (causal `if` in kernel control flow, separate
|
||||||
|
work); RW-aware scheduler reorder (strict FIFO ships); user-defined
|
||||||
|
recipes; new engine micro-ops in PE_MATH (it stays op_kind-opaque).
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 2. Where the design sits in the current code
|
||||||
|
|
||||||
|
```
|
||||||
|
runtime_api (host)
|
||||||
|
RuntimeContext.zeros/empty/launch context.py
|
||||||
|
│ KernelLaunchMsg
|
||||||
|
▼
|
||||||
|
sim_engine
|
||||||
|
GraphEngine(enable_data=…) engine.py
|
||||||
|
MemoryStore memory_store.py
|
||||||
|
│
|
||||||
|
▼
|
||||||
|
components / PE pipeline (per PE)
|
||||||
|
PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM}
|
||||||
|
greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py
|
||||||
|
CompositeCmd → tile plan → engines pe_scheduler.py (ADR-0014 D6)
|
||||||
|
│
|
||||||
|
▼
|
||||||
|
tl programming model tl_context.py
|
||||||
|
composite / dot / load / store / scratch_scope
|
||||||
|
── ADR-0065 NEW ──────────────────────────
|
||||||
|
tl_recipes.py RECIPE_DESCRIPTORS, EngineOp,
|
||||||
|
RecipeDescriptor, PrimaryOutSpec
|
||||||
|
```
|
||||||
|
|
||||||
|
The new module `tl_recipes.py` sits beside `tl_context.py`. PE_SCHEDULER
|
||||||
|
(`pe_scheduler.py`) and `tiling.py` extend to consume the flat ops list
|
||||||
|
without importing `tl_recipes`.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 3. File plan
|
||||||
|
|
||||||
|
### 3.1 Modified production files
|
||||||
|
|
||||||
|
| File | Change | ADR ref |
|
||||||
|
|---|---|---|
|
||||||
|
| `src/kernbench/common/pe_commands.py` | `CompositeCmd` flat-ops dataclass; `OpSpec.operands: dict`; new `Scope` value semantics (KERNEL position-based); `logical_bytes` property on every `PeCommand`. | 0065 D1/D2, 0064 D2 |
|
||||||
|
| `src/kernbench/triton_emu/tl_context.py` | `composite(prologue=[...], out=TensorHandle, …)`; recipe lowering pass (RECIPE_DESCRIPTORS lookup, scratch alloc, OpSpec expansion, primary_out auto-bind); legacy API preserved (`op="gemm", a, b, epilogue=[...]` still works, internally lowers to flat ops). | 0065 D1/D5 |
|
||||||
|
| `src/kernbench/triton_emu/tl_recipes.py` (NEW) | `RecipeDescriptor`, `PrimaryOutSpec`, `EngineOp` dataclasses; `RECIPE_DESCRIPTORS["softmax_merge"]` entry; expansion helper. | 0065 D5 |
|
||||||
|
| `src/kernbench/components/builtin/pe_scheduler.py` | `_generate_plan` rewritten for flat ops: scan for GEMM (≤1), position-based KERNEL placement, OUTPUT_TILE/K_TILE epilogue, strict-FIFO RW tracker keyed by handle id. | 0065 D3/D6 |
|
||||||
|
| `src/kernbench/components/builtin/tiling.py` | `generate_plan_from_ops(ops, rw_handles, pe_prefix)` replaces `generate_gemm_plan` / `generate_math_plan` (or wraps them). DMA auto-insertion from `handle.space`. Epilogue `add` stage emits `other_handle` in params (GAP 9 fix). | 0065 D3/D4 |
|
||||||
|
| `src/kernbench/components/builtin/pe_cpu.py` | dispatch cost from `cmd.logical_bytes` via ADR-0064 Rev2 formula; topology config override. | 0064 D1/D4 |
|
||||||
|
| topology yaml (e.g., `topologies/llama70b_4sip.yaml`) | optional `pe_cost_model:` block under PE node attrs. | 0064 D4 |
|
||||||
|
| `src/kernbench/benches/_gqa_decode_long.py` | opt2 variant — `tl.composite(prologue=[{"op":"softmax_merge", …}], op="gemm", b=V_ref, out=O, epilogue=[{"op":"add", "other":O}])`. | 0065 §Migration P3 |
|
||||||
|
|
||||||
|
### 3.2 New tests
|
||||||
|
|
||||||
|
| File | Covers |
|
||||||
|
|---|---|
|
||||||
|
| `tests/test_composite_flat_ops_refactor.py` | every existing bench's op_log byte-equal pre/post refactor (parametrized over bench list) |
|
||||||
|
| `tests/test_tl_recipe_softmax_merge.py` | recipe lowering — given inputs, asserts CompositeCmd has 10 ops in expected order + `rw_handles == (m, l, O)` |
|
||||||
|
| `tests/test_pe_scheduler_strict_fifo.py` | two consecutive composites sharing rw_handles complete in dispatch order; no stage interleave |
|
||||||
|
| `tests/test_pe_scheduler_auto_dma.py` | `space=hbm` operand → DMA_READ stage; `space=tcm` → none; `out.space=hbm` → DMA_WRITE |
|
||||||
|
| `tests/test_pe_cost_model.py` | ADR-0064 Rev2 — typical 1-GEMM composite at 45 ns; topology override path; `PeCpuOverheadCmd` bypass |
|
||||||
|
| `tests/attention/test_gqa_decode_opt2.py` | opt2 numeric parity with opt3; K-before-V DMA invariant; per-tile dispatch ratio measurement |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 4. Data model — `CompositeCmd` flat-ops
|
||||||
|
|
||||||
|
### 4.1 New shape
|
||||||
|
|
||||||
|
```python
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class CompositeCmd:
|
||||||
|
completion: CompletionHandle
|
||||||
|
ops: tuple[OpSpec, ...] # ordered MATH/GEMM ops
|
||||||
|
rw_handles: tuple[TensorHandle, ...] = () # cross-composite hazard
|
||||||
|
data_op: bool = True
|
||||||
|
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
4 # framing
|
||||||
|
+ 1 + sum(op.logical_bytes for op in self.ops)
|
||||||
|
+ 1 + 8 * len(self.rw_handles)
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
The previous `op`, `a`, `b`, `out_addr`, `out_nbytes`, `math_op` fields
|
||||||
|
are **removed**. All info lives in `ops[i]`.
|
||||||
|
|
||||||
|
### 4.2 `OpSpec` extended
|
||||||
|
|
||||||
|
```python
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class OpSpec:
|
||||||
|
kind: str # "gemm" | "rmax" | "exp_diff" | ...
|
||||||
|
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
|
||||||
|
operands: dict[str, TensorHandle]
|
||||||
|
extra: dict[str, Any] = field(default_factory=dict)
|
||||||
|
out: TensorHandle | None = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def logical_bytes(self) -> int:
|
||||||
|
return (
|
||||||
|
1 + 1 # opcode + scope enum
|
||||||
|
+ 1 + 8 * len(self.operands) # named handles
|
||||||
|
+ (8 if self.out is not None else 0)
|
||||||
|
+ 1 + 4 * len(self.extra) # scalars
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
Backward compat: existing `operands: tuple[Any, ...]` use sites
|
||||||
|
(epilogue `bias`, `scale`, `add`, etc.) migrate to `operands={"other":
|
||||||
|
h}` / `operands={"src": h}` per a small lookup in
|
||||||
|
`pe_commands.EPILOGUE_OPS`.
|
||||||
|
|
||||||
|
### 4.3 `Scope` semantics (position-based for KERNEL)
|
||||||
|
|
||||||
|
`Scope` keeps its three values (`KERNEL`, `K_TILE`, `OUTPUT_TILE`).
|
||||||
|
What changes: `KERNEL` no longer means "once anywhere"; it now means
|
||||||
|
"execute at this op's position in the sequence — once". Combined with
|
||||||
|
position relative to the GEMM op:
|
||||||
|
- KERNEL ops *before* GEMM → pre-tile-loop (single shot)
|
||||||
|
- KERNEL ops *after* GEMM → post-tile-loop (single shot)
|
||||||
|
- K_TILE / OUTPUT_TILE ops *after* GEMM → inside tile loop as before
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 5. Recipe: `softmax_merge`
|
||||||
|
|
||||||
|
### 5.1 `tl_recipes.py` definitions
|
||||||
|
|
||||||
|
```python
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PrimaryOutSpec:
|
||||||
|
from_shape: str # operand name to copy shape from
|
||||||
|
from_dtype: str # operand name to copy dtype from
|
||||||
|
transform: str # "identity" | "trans" (forward-compat)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class EngineOp:
|
||||||
|
engine: Literal["MATH", "GEMM", "DMA"]
|
||||||
|
op_kind: str # PE_MATH op_kind label (opaque to engine)
|
||||||
|
# Slot names: any of "src", "src_a", "src_b", "src_c", "dst" — these
|
||||||
|
# are RecipeDescriptor.operands names or internal scratch names.
|
||||||
|
src: str | None = None
|
||||||
|
src_a: str | None = None
|
||||||
|
src_b: str | None = None
|
||||||
|
src_c: str | None = None
|
||||||
|
dst: str | None = None
|
||||||
|
reduce_axis: int | None = None
|
||||||
|
bcast_axis: int | None = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class RecipeDescriptor:
|
||||||
|
operands: dict[str, str] # name → "R" | "RW"
|
||||||
|
primary_out: PrimaryOutSpec | None # implicit output spec
|
||||||
|
tile_alignment: Literal["single_shot", "tile_aligned"]
|
||||||
|
internal_scratch_bytes_fn: Callable[..., int]
|
||||||
|
engine_seq: tuple[EngineOp, ...]
|
||||||
|
|
||||||
|
|
||||||
|
RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = { ... }
|
||||||
|
```
|
||||||
|
|
||||||
|
### 5.2 Engine sequence (8 MATH ops)
|
||||||
|
|
||||||
|
| # | engine | op_kind | reads (slots) | writes (slots) | extras |
|
||||||
|
|---|---|---|---|---|---|
|
||||||
|
| 1 | MATH | `rmax` | s | m_loc | `reduce_axis=-1` |
|
||||||
|
| 2 | MATH | `max_elem` | m, m_loc | m_new | — |
|
||||||
|
| 3 | MATH | `exp_diff` | m, m_new | corr | — |
|
||||||
|
| 4 | MATH | `exp_diff` | s, m_new | P (primary_out)| `bcast_axis=0` |
|
||||||
|
| 5 | MATH | `rsum` | P | l_loc | `reduce_axis=-1` |
|
||||||
|
| 6 | MATH | `fma` | l, corr, l_loc | l | — |
|
||||||
|
| 7 | MATH | `mul_bcast` | O, corr | O | `bcast_axis=1` |
|
||||||
|
| 8 | MATH | `copy` | m_new | m | — |
|
||||||
|
|
||||||
|
Intermediates `m_loc, m_new, corr, P, l_loc` are scratch slots allocated
|
||||||
|
by TLContext before lowering. `P` is the recipe's primary output; auto-
|
||||||
|
bound to the head GEMM's `operands["a"]`.
|
||||||
|
|
||||||
|
### 5.3 PE_MATH op_kind additions
|
||||||
|
|
||||||
|
PE_MATH is op_kind-opaque ([pe_math.py:90-99](src/kernbench/components/builtin/pe_math.py#L90-L99)).
|
||||||
|
Latency = `ceil(num_elements / vector_width) / clock`. Adding new
|
||||||
|
op_kind labels (rmax, max_elem, exp_diff, rsum, fma, mul_bcast, copy)
|
||||||
|
requires **no PE_MATH code changes** — they pass through as string
|
||||||
|
labels in `Stage.params["op_kind"]`. Latency stays consistent for the
|
||||||
|
same `num_elements`.
|
||||||
|
|
||||||
|
(Reductions like `rmax` / `rsum` have log-depth final stages in real HW
|
||||||
|
not modelled by `ceil(num_elements / vector_width)`; this is an existing
|
||||||
|
approximation, not introduced by ADR-0065.)
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 6. TLContext lowering pipeline
|
||||||
|
|
||||||
|
Input (user-facing API, preserved):
|
||||||
|
```python
|
||||||
|
tl.composite(
|
||||||
|
op="gemm", b=tl.ref(v_tile(j), (TILE, d)), out=O,
|
||||||
|
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
|
||||||
|
epilogue=[{"op": "add", "other": O}],
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
Lowering steps in `TLContext.composite()`:
|
||||||
|
|
||||||
|
1. **Validate prologue items.** For each, lookup
|
||||||
|
`RECIPE_DESCRIPTORS[item["op"]]`. Check all operands present, type
|
||||||
|
tags (R / RW) match handle states (e.g., RW must not be `tl.ref`).
|
||||||
|
2. **Allocate internal scratch.** Call
|
||||||
|
`recipe.internal_scratch_bytes_fn(G, TILE, d, bpe)`; allocate via
|
||||||
|
`scratch_scope` helper (ADR-0063). Get slot addresses for `m_loc,
|
||||||
|
m_new, corr, P, l_loc`.
|
||||||
|
3. **Derive primary_out handle.** From `recipe.primary_out`: `P` handle =
|
||||||
|
`TensorHandle(shape=Sj.shape, dtype=Sj.dtype, addr=<P slot>,
|
||||||
|
space="tcm")`.
|
||||||
|
4. **Expand `engine_seq` → flat OpSpecs.** For each `EngineOp`:
|
||||||
|
- Resolve slot names (operands or scratch) to `TensorHandle`s.
|
||||||
|
- Build `OpSpec(kind=op_kind, scope=Scope.KERNEL, operands={...},
|
||||||
|
extra={...}, out=...)`.
|
||||||
|
- `out` is the handle for `dst` slot; `operands` is `{src*: handle}`.
|
||||||
|
5. **Auto-bind head GEMM.** GEMM OpSpec built from `op="gemm"`,
|
||||||
|
`operands={"a": P_handle, "b": V_handle}`, `out=<output handle>`,
|
||||||
|
`extra={"m": G, "k": TILE, "n": d}`.
|
||||||
|
6. **Build epilogue OpSpecs.** Each `epilogue=[{...}]` entry maps to
|
||||||
|
`OpSpec(kind=..., scope=Scope.OUTPUT_TILE, operands={...}, ...)`.
|
||||||
|
The `add` epilogue specifically gets `operands={"other": O_handle}`
|
||||||
|
so the Stage later carries `other_handle` (GAP 9 fix).
|
||||||
|
7. **Compute `rw_handles`.** Collect all handles tagged RW in prologue
|
||||||
|
recipes + the head GEMM's `out` (if accumulating) + any RW epilogue
|
||||||
|
targets. For softmax_merge: `rw_handles = (m, l, O)`.
|
||||||
|
8. **Emit `CompositeCmd`.** `ops = (8 MATH + 1 GEMM + 1 MATH(add))`,
|
||||||
|
`rw_handles = (m, l, O)`, `completion = new`.
|
||||||
|
|
||||||
|
**Legacy lowering.** Pre-existing
|
||||||
|
`tl.composite(op="gemm", a=A, b=B, epilogue=[{"op":"bias","bias":h}])`
|
||||||
|
goes through the same emit path: no prologue, head OpSpec from `op="gemm"`,
|
||||||
|
epilogue OpSpecs built from `epilogue=[...]`. Output `CompositeCmd` has
|
||||||
|
flat ops list; semantics preserved.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 7. PE_SCHEDULER tile plan generation (flat-ops)
|
||||||
|
|
||||||
|
```python
|
||||||
|
def _generate_plan(cmd: CompositeCmd) -> PipelinePlan:
|
||||||
|
pp = self._pe_prefix
|
||||||
|
bpe = 2
|
||||||
|
|
||||||
|
# 1. Identify GEMM op (≤1).
|
||||||
|
gemm_idx = next(
|
||||||
|
(i for i, o in enumerate(cmd.ops) if o.kind == "gemm"),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 2. MATH-only composite (no GEMM): single tile of all ops in order.
|
||||||
|
if gemm_idx is None:
|
||||||
|
return _math_only_plan(cmd.ops, pp)
|
||||||
|
|
||||||
|
gemm = cmd.ops[gemm_idx]
|
||||||
|
M, K, N = gemm.extra["m"], gemm.extra["k"], gemm.extra["n"]
|
||||||
|
|
||||||
|
# 3. Partition ops by position.
|
||||||
|
pre_loop_ops = cmd.ops[:gemm_idx] # KERNEL scope
|
||||||
|
after_ops = cmd.ops[gemm_idx+1:] # mixed scopes
|
||||||
|
|
||||||
|
# 4. Pre-loop MATH Stages (single-shot, before tile loop).
|
||||||
|
pre_stages = [_math_stage_from_op(op, pp) for op in pre_loop_ops]
|
||||||
|
|
||||||
|
# 5. Tile loop (existing generate_gemm_plan pattern, but:
|
||||||
|
# - DMA_READ inserted only when operand.space == "hbm"
|
||||||
|
# - K_TILE epilogue from after_ops with scope == K_TILE
|
||||||
|
# - OUTPUT_TILE epilogue from after_ops with scope == OUTPUT_TILE
|
||||||
|
# - DMA_WRITE inserted only when out.space == "hbm")
|
||||||
|
tile_plans = _generate_tile_loop(
|
||||||
|
gemm, M, K, N, after_ops, pp, bpe,
|
||||||
|
)
|
||||||
|
|
||||||
|
# 6. Post-loop KERNEL Stages (after tile loop completes).
|
||||||
|
post_loop_ops = [op for op in after_ops if op.scope == Scope.KERNEL]
|
||||||
|
post_stages = [_math_stage_from_op(op, pp) for op in post_loop_ops]
|
||||||
|
|
||||||
|
return PipelinePlan(
|
||||||
|
prologue_stages=tuple(pre_stages),
|
||||||
|
tiles=tile_plans,
|
||||||
|
epilogue_stages=tuple(post_stages),
|
||||||
|
m_tiles=..., k_tiles=..., n_tiles=...,
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
`PipelinePlan` gets two new fields:
|
||||||
|
|
||||||
|
```python
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PipelinePlan:
|
||||||
|
tiles: list[TilePlan]
|
||||||
|
m_tiles: int
|
||||||
|
k_tiles: int
|
||||||
|
n_tiles: int
|
||||||
|
prologue_stages: tuple[Stage, ...] = ()
|
||||||
|
epilogue_stages: tuple[Stage, ...] = ()
|
||||||
|
```
|
||||||
|
|
||||||
|
Feeder change (`_feed_loop`): yield `prologue_stages` Stages first, then
|
||||||
|
tile loop as today, then `epilogue_stages`. For the single-Stage case
|
||||||
|
(no prologue/epilogue), behaviour identical to today.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 8. Strict-FIFO RW hazard tracker
|
||||||
|
|
||||||
|
```python
|
||||||
|
class _RwHazardTracker:
|
||||||
|
def __init__(self):
|
||||||
|
self._inflight: list[tuple[CompletionHandle, set[str]]] = []
|
||||||
|
# (completion, set of handle ids written or RW)
|
||||||
|
|
||||||
|
def admit(self, env, cmd: CompositeCmd) -> Generator:
|
||||||
|
"""Block until no in-flight composite has overlapping RW."""
|
||||||
|
rw_ids = {h.id for h in cmd.rw_handles}
|
||||||
|
# Strict FIFO: wait for ALL prior composites whose rw_set
|
||||||
|
# intersects, in order.
|
||||||
|
while True:
|
||||||
|
blocking = [c for c, s in self._inflight if s & rw_ids]
|
||||||
|
if not blocking:
|
||||||
|
break
|
||||||
|
# Wait on the *first* blocking composite's completion.
|
||||||
|
yield blocking[0].done_event
|
||||||
|
self._inflight.append((cmd.completion, rw_ids))
|
||||||
|
|
||||||
|
def retire(self, completion: CompletionHandle) -> None:
|
||||||
|
self._inflight = [
|
||||||
|
(c, s) for c, s in self._inflight if c != completion
|
||||||
|
]
|
||||||
|
```
|
||||||
|
|
||||||
|
`PE_SCHEDULER` invokes `admit()` before queueing a CompositeCmd's tile
|
||||||
|
plan and `retire()` upon `done_event.succeed()`. Strict FIFO: even when
|
||||||
|
overlap is partial, the new composite waits for all prior overlapping
|
||||||
|
composites to *fully* drain — simple, easy to reason about, easy to
|
||||||
|
test. (RW-aware reorder deferred per ADR-0065 A4.)
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 9. Phased implementation plan (test-first per CLAUDE.md)
|
||||||
|
|
||||||
|
Each phase is an independent Phase-1 → Phase-2 cycle per CLAUDE.md Part 1.
|
||||||
|
|
||||||
|
| Phase | Deliverable | Gate |
|
||||||
|
|---|---|---|
|
||||||
|
| **P0 (= ADR-0064 Rev2)** | `logical_bytes` property + dispatch formula + topology override; one-time goldens regeneration | typical composite ≈ 45 ns; dispatch ratio test infrastructure in place |
|
||||||
|
| **P1** | `CompositeCmd` flat-ops dataclass refactor + `OpSpec.operands: dict` + legacy `tl.composite` lowering preserved | every existing bench's op_log byte-equal pre/post |
|
||||||
|
| **P2** | `tl_recipes.py` + `RECIPE_DESCRIPTORS["softmax_merge"]` + TLContext lowering for `prologue=[...]` | recipe lowering test (10 ops, rw_handles, addresses) passes |
|
||||||
|
| **P3** | PE_SCHEDULER `_generate_plan` flat-ops + `PipelinePlan.prologue_stages` + DMA auto-insertion + `_feed_loop` extension | all existing bench Stage sequences unchanged; new MATH chain visible in op_log for opt2 mock |
|
||||||
|
| **P4** | Strict-FIFO RW hazard tracker + `_RwHazardTracker` integration | two-composite RW conflict test serializes correctly |
|
||||||
|
| **P5** | `_gqa_decode_long.py` opt2 variant; data-mode numeric parity check; K-before-V invariant check | opt2 matches opt3 within fp tolerance; no V DMA during prologue |
|
||||||
|
| **P6** | Dispatch-ratio measurement: opt3 vs opt2 per-tile PE_CPU cycles | ratio ≈ 2.4× ± 10% at default calibration |
|
||||||
|
|
||||||
|
P0 must land first (separate ADR). P1 is pure refactor (safest). P2–P4
|
||||||
|
build the new path without touching opt3. P5 enables opt2. P6 closes
|
||||||
|
the measurement loop.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 10. Verification plan (concrete)
|
||||||
|
|
||||||
|
Mirrors ADR-0065 §Test Requirements; grounded in SPEC R2/R5, ADR-0023/
|
||||||
|
0025, ADR-0046, ADR-0054.
|
||||||
|
|
||||||
|
**Refactor safety (P1):**
|
||||||
|
- Parametrize over `tests/<existing bench>.py` baseline op_logs.
|
||||||
|
- Pre/post refactor diff must be empty for *every* bench.
|
||||||
|
- Failure → revert P1 atomically; do not land partial.
|
||||||
|
|
||||||
|
**Recipe lowering (P2):**
|
||||||
|
- Given `Sj=(8,64), m=(8,), l=(8,), O=(8,128)`, lower
|
||||||
|
`tl.composite(prologue=[{"op":"softmax_merge", "s":Sj,...}], op="gemm",
|
||||||
|
b=V_ref, out=O, epilogue=[{"op":"add", "other":O}])`.
|
||||||
|
- Assert `len(cmd.ops) == 10`, `cmd.ops[0..7].kind == [rmax, max_elem,
|
||||||
|
exp_diff, exp_diff, rsum, fma, mul_bcast, copy]`,
|
||||||
|
`cmd.ops[8].kind == "gemm"`, `cmd.ops[9].kind == "add"`.
|
||||||
|
- Assert `cmd.rw_handles == (m, l, O)`.
|
||||||
|
- Assert all OpSpec addresses are concrete (no None).
|
||||||
|
|
||||||
|
**Scheduler plan (P3):**
|
||||||
|
- For the above CompositeCmd, assert PipelinePlan has 8 prologue_stages
|
||||||
|
(MATH), tile loop with FETCH/GEMM/MATH(add), no epilogue_stages.
|
||||||
|
- Assert no DMA_READ for prologue MATH stages (operands TCM-resident).
|
||||||
|
- Assert one DMA_READ for V (`b.space == "hbm"`).
|
||||||
|
- Assert no DMA_WRITE (`out.space == "tcm"`).
|
||||||
|
|
||||||
|
**RW hazard (P4):**
|
||||||
|
- Submit composite A (rw_handles=(O,)) then composite B (operands
|
||||||
|
reading O). Assert B's Stages start only after A's done_event.
|
||||||
|
|
||||||
|
**Numeric parity (P5):**
|
||||||
|
- In data mode, run opt2 and opt3 for `(G=8, T_q=1, S_kv=64, C=1, P=1)`.
|
||||||
|
- Assert `O_opt2 ≈ O_opt3` within `fp16` tolerance.
|
||||||
|
|
||||||
|
**K-before-V invariant (P5):**
|
||||||
|
- op_log inspection: for each tile's #2 CompositeCmd, no
|
||||||
|
`memory/dma_read` events between #2 start and the GEMM compute begin
|
||||||
|
that are not the V tile DMA itself.
|
||||||
|
|
||||||
|
**Dispatch ratio (P6):**
|
||||||
|
- For `S_kv=64, n_tiles=16`, measure total PE_CPU dispatch cycles for
|
||||||
|
opt3 vs opt2 paths.
|
||||||
|
- Assert ratio in `[2.0, 2.8]` (centre 2.4×).
|
||||||
|
|
||||||
|
**Determinism (all phases):**
|
||||||
|
- Identical inputs → identical op_log + latency (SPEC §0.1).
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 11. Performance model (expected)
|
||||||
|
|
||||||
|
Per-tile PE_CPU dispatch cycles, ADR-0064 Rev2 default calibration:
|
||||||
|
|
||||||
|
```
|
||||||
|
opt3: 10 cmds × 32 + total_bytes(≈232) × 0.25 = 320 + 58 ≈ 378 ns
|
||||||
|
opt2: 2 cmds × 32 + total_bytes(≈380) × 0.25 = 64 + 95 ≈ 159 ns
|
||||||
|
ratio = 378 / 159 ≈ 2.4×
|
||||||
|
```
|
||||||
|
|
||||||
|
Per-tile engine timing dominates total latency; dispatch is one
|
||||||
|
component. The ratio above is the *dispatch-only* relative cost. Total
|
||||||
|
latency improvement depends on whether dispatch was on the critical
|
||||||
|
path — measured per workload in P6.
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 12. Open items (status)
|
||||||
|
|
||||||
|
Most ADR-0065 §Open review items map to specific implementation choices
|
||||||
|
already documented in §§5–8. Live remaining items:
|
||||||
|
|
||||||
|
| Item | Resolution venue |
|
||||||
|
|---|---|
|
||||||
|
| Non-identity `PrimaryOutSpec.transform` (e.g., `"trans"`) | Future recipe addition; current code uses string dispatch in step 3 of §6 |
|
||||||
|
| `_gqa_decode_long.py` opt2 — TILE size sweep | Open (existing §B-3, §9 of ADR-0060) |
|
||||||
|
| Larger composite descriptors (10+ OpSpecs) and `logical_bytes` growth | Monitored in P6; if `R` term dominates, calibration revisit |
|
||||||
|
| `scratch_scope` budget vs softmax_merge intermediates (G·TILE for P) | ADR-0063 sizing check at S_kv = 256K |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 13. Risks
|
||||||
|
|
||||||
|
| Risk | Likelihood | Mitigation |
|
||||||
|
|---|---|---|
|
||||||
|
| `CompositeCmd` refactor (P1) changes op_log for some bench | med | Parametrized regression over all benches; revert atomically |
|
||||||
|
| `OpSpec.operands` migration (positional → dict) breaks epilogue site | med | Per-EPILOGUE_OPS migration table; existing epilogue tests cover |
|
||||||
|
| `_RwHazardTracker` false-positive serializes too aggressively | low | Strict FIFO is the safe direction; RW-aware reorder is a future ADR |
|
||||||
|
| ADR-0064 Rev2 and ADR-0065 P1 land in wrong order → stale goldens | med | PR ordering enforced; CI gate requires both to be green |
|
||||||
|
| `softmax_merge` scratch budget exceeds `scratch_scope` at long context | med | ADR-0063 sizing check; if needed, tile_alignment promotion (future) |
|
||||||
|
| PE_MATH op_kind label collision with new strings | low | Labels are opaque strings — no enum conflict; just add to allowed set if validated |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 14. Glossary & references
|
||||||
|
|
||||||
|
- **flat ops list** — `CompositeCmd.ops: tuple[OpSpec, ...]` with no
|
||||||
|
head/epilogue field; placement determined by op position + scope.
|
||||||
|
- **prologue (user API)** — `tl.composite(prologue=[...])` kwarg for
|
||||||
|
recipe ops that lower into pre-GEMM KERNEL-scope OpSpecs. *HW command
|
||||||
|
does not carry this concept.*
|
||||||
|
- **recipe** — a named macro op (e.g., `softmax_merge`) that TLContext
|
||||||
|
expands to multiple flat OpSpecs at emit time. Registry =
|
||||||
|
`RECIPE_DESCRIPTORS`.
|
||||||
|
- **primary_out** — recipe's implicit output (e.g., `P` for
|
||||||
|
softmax_merge). Auto-bound to head GEMM's first matrix operand.
|
||||||
|
- **handle-identity strict FIFO** — cross-composite serialization rule:
|
||||||
|
a composite waits for all prior composites whose `rw_handles`
|
||||||
|
intersect with its own to complete.
|
||||||
|
- **logical_bytes** — HW-logical byte size of a PE command, derived
|
||||||
|
from a structural rule (ADR-0064 D2), not from Python's
|
||||||
|
`sys.getsizeof`.
|
||||||
|
|
||||||
|
Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143`
|
||||||
|
(composite tile plan), `pe_dma.py:45,89` (read channel), `pe_math.py`
|
||||||
|
(op_kind opaque latency), `pe_cpu.py` (dispatch hook),
|
||||||
|
`pe_commands.py` (CompositeCmd, OpSpec, EPILOGUE_OPS).
|
||||||
|
|
||||||
|
ADRs: **0065** (this design), **0064 Revision 2** (dispatch cost),
|
||||||
|
**0060** (GQA fused attention, §5.6 / §8 item 4 carve-out), **0063**
|
||||||
|
(scratch_scope), **0046** (tl_context contract), **0042** (tile plan
|
||||||
|
generators), **0014** (PE pipeline execution model).
|
||||||
Reference in New Issue
Block a user