From 7a9d4ec47b5cb9f7da487b1f9f9fca1d485fc430 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Thu, 4 Jun 2026 10:19:37 -0700 Subject: [PATCH] gqa(adr): pivot ADR-0060 to composite hybrid + lazy tl.load; add ADR-0064 cost model - ADR-0060: GEMMs (Q.Kt, P.V) via existing tl.composite (scheduler-managed tiling + K/V DMA streaming); softmax merge + IPCQ tree reduction stay in kernel. Front TL;DR pseudocode of the final composite kernel; new section B lists open design items (DDD sync, K pre-transpose, dma_read lever, kernel-vs-scheduler tiling, ring path). - ADR-0062: redefined from a new load_async op to global lazy tl.load (non-blocking + auto-wait on first use; API unchanged; goldens regenerate). - ADR-0064 (new): per-op-type CPU issue cost model (composite ~40ns >> primitive) so the hybrid's CPU-saturation win becomes measurable (currently dispatch_cycles=0 hides it). Cost-model impl deferred. - KO mirrors for ADR-0060/0062/0064 (-ko suffix, adr-proposed). Rationale: non-blocking CompositeCmd offloads tiling to PE_SCHEDULER, decoupling CPU issue-rate from execution so the CPU can saturate the engines; the prior 'composite = no latency benefit' claim was an artifact of dispatch_cycles=0. Docs only; no production code changed. Co-Authored-By: Claude Opus 4.8 (1M context) --- ...R-0060-algo-gqa-fused-attention-ahbm-ko.md | 747 ++++++++++++++++++ .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 419 +++++++--- .../ADR-0062-prog-tl-async-load-ko.md | 148 ++++ .../ADR-0062-prog-tl-async-load.md | 231 +++--- .../ADR-0064-perf-cpu-issue-cost-model-ko.md | 166 ++++ .../ADR-0064-perf-cpu-issue-cost-model.md | 179 +++++ 6 files changed, 1675 insertions(+), 215 deletions(-) create mode 100644 docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md create mode 100644 docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md create mode 100644 docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md create mode 100644 docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md new file mode 100644 index 0000000..adc1aef --- /dev/null +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md @@ -0,0 +1,747 @@ +# ADR-0060: AHBM GQA Fused Attention 커널 (Llama3-70B) + +## Status + +Proposed + +**Context model:** Llama3-70B. +**Decision drivers:** agentic workload → 낮은 batch, 긴 context; +KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV). + +**Supersedes / extends:** 기존 mesh-native 어텐션 커널 +`_attention_mesh_kv`(prefill)와 `_attention_mesh_mlo`(decode), 그리고 +`milestone-gqa-llama70b` eval bench. *§A. 기존 kernbench 작업과의 관계* 참조 — +그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는 +baseline이다. + +**Supporting ADRs** (efficiency / scale enabler — *GQA blocker 아님*; §8 정정 +참조): **ADR-0063** `tl.scratch_scope`(per-tile scratch 재활용 — 현실적 context +길이에 필요), **ADR-0062** lazy `tl.load`(첫 사용 시점 auto-wait를 갖는 +non-blocking load → load/compute 오버랩), **ADR-0061** `tl.broadcast`(선택적 +mask/범용 편의). 두 GEMM(Q·Kᵀ, P·V)은 scheduler가 관리하는 `tl.composite` +커맨드로 발행된다(기존 `CompositeCmd`; 새 커맨드 종류 없음); 진짜 GQA 자체는 +커널 재구조화만 필요(§5.2). + +**Algorithm lineage.** 이 커널은 **FlashAttention**(tiling + online/streaming +softmax, P·V fused — full score matrix 미생성)이다. §4의 KV-parallel +split-and-combine는 **FlashDecoding**(split-KV + log-sum-exp merge). §5.5의 Ring +경로는 **Ring Attention**(KV 블록을 mesh 주위로 회전, 같은 online softmax로 +fold). 새 수학은 도입하지 않으며, 본 ADR은 이 알려진 알고리즘을 kernbench의 +**greenlet `tl` 프로그래밍 모델**(ADR-0020, ADR-0046)과 **IPCQ** PE↔PE +collective(ADR-0023/0025)에 매핑한다. + +--- + +## TL;DR — 최종 GQA 커널 (composite hybrid) pseudocode + +결정(§1)을 한 곳에: **GEMM → `tl.composite`; softmax 머지 + tree reduction → +커널; `tl.load`은 lazy.** KV head별로, `G`를 matmul M 차원에 fold(진짜 GQA, +broadcast 없음). 이것이 reference 형태이며, 산문 섹션이 각 부분을 부연한다. + +```python +def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr, + counter, start_pe, N, q_block, scale, *, tl): + # ---- geometry (kernel arithmetic, §2) ---- + pe_id = tl.program_id(axis=rank_axis) + G = H_q // H_kv # query heads per KV head (=8) + causal = q_block.is_prefill + + for kv in range(H_kv): # one KV head per iteration (§5.2) + # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062): + # issue now, auto-wait at first use inside the first composite. + q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d] + my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv + n_tiles = ceil(my_len / TILE) + + # persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros + m, l, O = init_running(G * q_block.T_q, d) + + for j in range(n_tiles): + if causal and tile_all_future(j, q_block): + continue # causal skip (kernel if) + with tl.scratch_scope(): # per-tile temporaries (ADR-0063) + # --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite --- + Sj = tl.composite("gemm", a=q_g, + b=tl.ref(k_tile(kv, j), (d, TILE))) * scale # Kᵀ pre-stored [d,TILE] + if causal and tile_partial(j, q_block): + Sj = Sj + causal_mask(j, q_block) # additive boundary mask + # --- online softmax merge in the kernel (MATH ops) --- + m_new = tl.maximum(m, tl.max(Sj, axis=-1)) + P = tl.exp(Sj - m_new) + corr = tl.exp(m - m_new) + l = l * corr + tl.sum(P, axis=-1) + # --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite --- + Oj = tl.composite("gemm", a=P, + b=tl.ref(v_tile(kv, j), (TILE, d))) + O = O * corr + Oj # running merge + m = m_new + + # ---- cross-PE combine (§4): log-sum-exp tree to root, or store ---- + if N == 1: + tl.store(o_base(kv), O / l) + else: + tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv +``` + +> **이 형태인 이유:** 두 `tl.composite` 호출이 tiling + K/V DMA 스트리밍 + +> cross-tile 파이프라이닝을 PE_SCHEDULER에 offload한다(CPU는 coarse descriptor를 +> 내고 앞서 나가 → 엔진이 saturate 유지, §1); softmax 머지와 IPCQ tree reduction은 +> 커널에 남는다(기존 `CompositeCmd`가 cross-tile 상태를 못 들기 때문). `K`는 +> reshape-not-transpose caveat를 피하려 transpose된 `[d, TILE]`로 미리 저장된다 +> (§3, §B). 전용 "flash-composite" 커맨드 종류는 도입하지 **않는다**(§8 항목 4). + +--- + +## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것) + +kernbench는 **오늘날 이미 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로 +FlashAttention을 돌린다.** 두 커널이 존재한다: + +| File | Role | Mechanism | +|---|---|---| +| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | +| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | + +둘 다 `milestone-gqa-llama70b`(4 패널: +`{single,multi}_user × {prefill,decode}`, +`src/kernbench/benches/milestone_gqa_llama70b.py`)이 구동하며 +`tests/attention/test_milestone_gqa_llama70b.py`에서 테스트된다. + +**이들은 greenlet `tl` API로 작성되었다:** `tl.load`, `tl.dot`, +`tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, 그리고 +`TensorHandle`에 대한 Python `-`/`*`/`/`(각각 `MathCmd` emit) — GEMM은 composite가 +아니라 **blocking `tl.dot`**로. 실행 중 `(m, ℓ, O)`는 루프를 관통하는 Python +`TensorHandle`일 뿐이다. 본 ADR은 실행 상태와 softmax 머지는 커널에 유지하되 +**두 GEMM은 scheduler가 관리하는 `tl.composite` 경로로 옮긴다**(§1 참조) — 이것이 +**설계에 중요**하다. + +**baseline의 의도된 세 한계** — 효율적 GQA 커널이 정확히 들어내야 하는 것: + +1. **GQA 재사용 없음.** `h_q == h_kv == 1` + (`test_milestone_gqa_llama70b.py:137-142`). 테스트는 이를 *broadcast view*에 + 대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의 + **head-packing 핵**의 속성이다(`_view(K, (h_q·d, S_kv))`는 모든 head를 하나의 + matmul 차원에 뭉치고 `h_q == h_kv`일 때만 byte를 보존). 올바른 수정은 + broadcast op이 아니라 **커널 재구조화**다: **한 번에 한 KV head**를 처리하고 + `G` group 행을 matmul **M** 차원에 fold(§5.2). 이는 byte-보존 reshape만 쓰므로 + 진짜 GQA(`h_q = G·h_kv`)가 **신규 primitive 없이** 돈다 — §8 참조. +2. **O(N) reduction.** baseline은 all-to-all bidirectional fan-out을 하여 *모든* + rank가 full 답을 갖는다(`n_ranks − 1` 단계). 어텐션은 query owner에서만 `O`가 + 필요 → root로의 **tree reduction**은 `⌈log₂ N⌉` 단계(§4). +3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile + 임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking / + tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite + K/V 스트리밍(§3) + **ADR-0062**(lazy load 오버랩)으로 해결. + +**문서 부채(범위 밖이나 기록):** baseline은 ADR-0055/0056/0057/0058/0059를 +인용하나 **파일로 존재하지 않는다** — ghost 참조다. 본 ADR은 이를 소급 +작성하지 않으며; 권고는 Detailed Design Document의 *Open Decisions* 참조. + +--- + +## 0. 참조 차원 (Llama3-70B) + +| Symbol | Meaning | Value | +|---|---|---| +| `H_q` | query heads | 64 | +| `H_kv` | KV heads | 8 | +| `G` | GQA group size = `H_q / H_kv` | 8 | +| `d` | head dim | 128 | +| `L` | layers | 80 | +| `D` | model dim | 8192 | + +하드웨어 요약: +- **AHBM (chip)** = **CUBE**(메모리 cube, 각각 **PE**를 담은 logic die) 집합 + + **IO die**(ADR-0003). +- **IPCQ**: PE↔PE 큐, PE당 4 mesh-방향 queue-pair + (`N/S/E/W`, ADR-0023 D3; inter-SIP용 `global_*`, ADR-0032). 커널 + API: `tl.send(dir, src)` / `tl.recv(dir, shape, dtype)` + (`tl_context.py:402-499`). +- **Composite command** (`CompositeCmd`, `pe_commands.py:144-162`): 단일 + GEMM(또는 MATH) *head* + element-wise *epilogue* 단계 + (`bias/relu/scale/add/...`), PE_SCHEDULER에 **non-blocking**으로 발행되며, + scheduler가 tile plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다 + (ADR-0014 D6; `pe_scheduler.py:104-143`). 이는 일반적 multi-op DAG가 **아니다**: + 두 GEMM을 chain할 수 없고, 인스턴스 간 register 상태를 못 들며, IPCQ를 + pop/wait 못 한다. 따라서 본 ADR은 두 어텐션 GEMM(Q·Kᵀ, P·V)을 **각각** 자체 + composite로 발행하고 cross-GEMM softmax 머지 + IPCQ reduction은 커널에 유지한다 + — 새 "flash-composite" 커맨드 종류는 **필요 없다**(§1, §8 참조). +- **Allocation policy (SP):** multi-user 패널에서 CUBE당 query head 하나; KV + group의 `G` query head는 한 CUBE의 PE 안에 매핑. + +**직교하는 두 매핑 레이어** (round-robin placement): +- **Layer 1 (KV-parallel, intra-request):** KV token `i`는 PE + `(start_pe + i) mod N`에 안착 — round-robin이라 한 긴 request가 `N` PE에 걸쳐 + 분할, ≤1 token으로 균형. +- **Layer 2 (inter-request):** `start_pe = request_id mod N`이 request별 "PE-0 + 역할"을 회전. +- 결합: `pe = (request_id + global_token_idx) mod N`. + +`N` = 하나의 (query head, request)에 대한 KV-parallel reduction group의 PE 수. +Reduction은 **intra-CUBE** 유지(query head는 CUBE를 가로지르지 않음 — 명시적 +non-goal). + +--- + +## 0.5 커널 경계, 전제조건, I/O 계약 + +### 0.5.1 디코더 레이어 내 위치 + +``` +1. RMSNorm +2. QKV projection (GEMM) ─┐ qkv_rope kernel (SEPARATE, upstream) +3. RoPE on Q and K ─┤ +4. write new K,V → KV cache ─┘ +5. ===== THIS KERNEL: FlashAttention ===== (post-RoPE Q, K-cache, V-cache → O) +6. Output projection (GEMM) out_proj kernel (SEPARATE, downstream) +7. residual add → FFN ... +``` + +**전제조건 (upstream `qkv_rope`, 본 커널 아님):** +- **P1.** Q는 이미 RoPE-회전됨. 여기서 회전 없음. +- **P2.** K-cache는 **post-RoPE** K 저장(어텐션 시점에 재회전 안 함 — post-RoPE + 캐싱의 이유). +- **P3.** decode의 경우, 새 step의 K 행은 RoPE-회전되어 그 소유 PE의 K-cache + 슬롯에 **본 커널 launch 전에 `qkv_rope`가** 추가한다. ⇒ 본 커널은 KV cache에 + 대해 **pure read**. +- **P4.** V는 회전 안 함; V-cache는 raw projected V 보유. +- **P5.** upstream RoPE 위치는 token의 **절대 global 위치** + `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`(§2.1). Round-robin + placement ≠ RoPE 위치. + +### 0.5.2 Shape 기호 +`T_q` = 이번 launch의 query 길이(decode: 1; prefill/chunk: chunk 폭). +`S` = 전체 context 길이. `S_pe` = 이 PE 소유 key ≈ `⌈S/N⌉`. +저장 `bf16`(numpy proxy `f16`, `memory_store.py:16`); `m,ℓ,O` 누산기는 정밀도가 +중요한 곳에서 `f32`. + +### 0.5.3 INPUTS (커널 launch당) + +| Input | Shape (per KV head) | Location | Notes | +|---|---|---|---| +| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). group의 `G` query 행 batched. | +| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. 로컬은 dense, global index는 strided (§2.1). | +| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. | +| `global_token_counter` | scalar | launch arg | 커널이 `S_pe`, slot↔global, causal bound 도출. | +| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 회전. | +| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. | +| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. | +| `O_base` | address | launch arg | 최종 O가 쓰이는 곳 (root PE만). | +| `softmax_scale` | scalar | launch arg | `1/√d`. | + +**Ring Attention (§5.5)**에는 ring step마다 추가: IPCQ로 들어오는 `K_block, +V_block`을 ping-pong 버퍼에(post-RoPE), 그리고 causal step-skip용 +`step_kv_global_range`. + +### 0.5.4 OUTPUTS + +| Output | Shape | Location | Notes | +|---|---|---|---| +| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE만** | 정규화 `O_acc / ℓ_acc`, 저장 dtype으로 cast. | + +**중간값 (non-root PE, IPCQ 상 — 커널-가시 출력 아님):** +`(m_i, ℓ_i, O_i)` (`m,ℓ`: `[G, T_q]`; `O_i`: `[G, T_q, d]` 비정규화)이 +`tl.send`로 `tree_parent`에 push(§4). `O_i`가 무거운 부분. + +**No-SP (`N=1`):** IPCQ partial 없음; 단일 PE의 실행 중 `(m,ℓ,O)`를 제자리에서 +정규화하여 `O_base`에 기록. + +**명시적 비-출력:** KV-cache write(upstream, P3), output projection(downstream), +score `S` 및 prob `P`(미생성). + +--- + +## 1. 결정 (메커니즘) + +**커널을 *하이브리드*로 구현한다: 두 GEMM(Q·Kᵀ, P·V)을 scheduler가 관리하는 +`tl.composite(op="gemm")` 커맨드로 발행하고, online-softmax 머지와 cross-PE +reduction은 커널 수준 `tl` op으로 유지한다.** `tl.load`은 **lazy**다(non-blocking; +wait는 load된 데이터의 첫 사용 시점에 자동 삽입 — ADR-0062). 따라서 명시적 HBM +load가 뒤따르는 compute와 오버랩된다. kernbench의 실행 + latency 모델에 근거한 +근거: + +- **GEMM tiling이 PE_SCHEDULER에 offload된다.** `CompositeCmd`는 non-blocking + (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): 커널이 **하나의 coarse + descriptor**(M = `G·T_q`, PE당 전체 tile sweep)를 push하면 scheduler가 tile + plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다(ADR-0014 D6). K/V는 + scheduler가 HBM에서 스트리밍하는 `tl.ref` operand이므로, 타일당 **K/V prefetch는 + scheduler의 일**이다 — 명시적 prefetch op 없음. CPU(greenlet)는 현재 composite가 + 도는 동안 **다음** composite를 낼 수 있게 풀려나, scheduler가 타일을 가로질러 + GEMM 엔진을 saturate 유지한다. +- **이는 하드웨어를 반영**하며 CPU issue-rate를 execution-rate로부터 decouple한다. + 대조적으로 blocking per-op `tl.dot` 경로는 매 GEMM마다 CPU를 멈추고, 끼어드는 + softmax MATH op 동안 GEMM 엔진 **버블**을 남긴다; CPU가 fine-grained per-tile + 발행을 따라잡을 수 있을 때만 현실적이다. +- **실행 중 `(m, ℓ, O)` flash 상태는 Python `TensorHandle`로** 루프를 관통한다 + (baseline이 이미 그렇게 함); softmax 머지(max/exp/sum/rescale)는 두 GEMM + composite **사이**의 커널 수준 `tl` MATH다. 기존 `CompositeCmd`는 두 GEMM을 + chain하거나 cross-tile register 상태를 못 들므로(§0), 머지는 필연적으로 커널에 + 산다 — 이것이 하이브리드 분할이지, 우회한 한계가 아니다. +- **Cross-PE 결합**은 P·V 후 `(m, ℓ, O)`에 대한 log-sum-exp **tree**로, 커널 수준 + `tl.send`/`tl.recv`(§4)를 통해 — 불변. + +그래서 per-tile 내부 파이프라인은: + +``` +q_g = tl.load(Q group) # lazy; auto-wait at first use +per tile j: + Sⱼ = tl.composite("gemm", a=q_g, b=tl.ref(Kⱼ)) → Sⱼ # scheduler streams Kⱼ DMA + GEMM + Sⱼ += maskⱼ # kernel MATH, boundary tile only + online-softmax: mⱼ, m_new, P, corr, ℓ # kernel MATH + Oⱼ = tl.composite("gemm", a=P, b=tl.ref(Vⱼ)) → Oⱼ # scheduler streams Vⱼ DMA + GEMM + O = O*corr + Oⱼ; m = m_new # kernel MATH (running merge) +``` + +각 타일의 MATH 임시값을 `tl.scratch_scope`(ADR-0063)로 감싸 scratch를 O(1)로 +유지하고, 다음 타일의 composite를 현재 타일 결과를 wait하기 전에 발행 +(non-blocking 핸들)하여 scheduler가 타일을 가로질러 파이프라인되게 한다. + +제어 흐름(tile skip, mask 생성, reduction 스케줄링, 주소 산술)은 **커널**에 산다 +(greenlet 본문의 평범한 Python `if`/산술). 이는 kernbench의 greenlet 모델이 이미 +허용하는 바 그대로다(`kernel_runner.py`, ADR-0020 D3). + +> **이것이 무엇을 대체하나.** 이전 iteration은 "composite는 latency 이점 없음"이라는 +> 논거로 순수 greenlet primitive 경로(전부 `tl.dot`, composite 없음)를 제안했다. +> 그것은 **오직** 시뮬레이터가 현재 per-op CPU 발행 비용을 **0**으로 청구하기 +> 때문에만(`dispatch_cycles=0`, `pe_cpu.py`) 성립한다 — descriptor offload가 숨기려 +> 존재하는 바로 그 CPU issue-rate / DMA-program 비용을 모델이 지워버린 것이다. +> 하이브리드가 효율 커널의 충실한 표현이다. 이점의 **측정 가능한** 크기(CPU가 많은 +> 타일에 대해 엔진을 saturate할 수 있는가?)는 op 종류별 발행 비용 모델링에 +> 좌우되며, future work로 추적된다(cost model; §9). `dispatch_cycles=0`에서도 +> non-blocking composite 경로는 blocking `tl.dot` 경로가 남기는 GEMM 엔진 버블을 +> 채운다. +> +> **이것이 효율적 선택인 이유.** GEMM tiling + DMA 스트리밍 + cross-tile +> 파이프라이닝은 검증된 `CompositeCmd` scheduler 경로에 offload되고; softmax 머지와 +> 검증된 IPCQ collective는 커널에 남는다. 진짜 새 기계장치는 작고 범용적인 두 +> primitive뿐이다(ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); reduction은 +> `tl.send`/`tl.recv`를 재사용한다. 전용 "flash-composite" 커맨드(softmax 머지 + +> carried register 상태 + IPCQ-push epilogue를 내재화하는 한 종류)는 만들지 +> **않는다** — 크고 특수목적이며, 하이브리드 대비 유일한 delta(완전 softmax +> offload)가 현재 모델링 충실도에서 정당화되지 않는다; §8 참조. + +--- + +## 2. 메모리 레이아웃 & 드라이버 책임 + +### 2.1 KV cache 할당 +- per-PE KV 버퍼는 per-PE 최대 context `⌈max_context / N⌉ × d × dtype`로 K와 V를 + 각각 sizing. kernbench에서는 sequence 차원에 대해 `pe`(또는 `cube`) 축이 + `row_wise`인 `DPPolicy`로 배치된다(`policy/placement/dp.py`; baseline은 K/V에 + `DPPolicy(pe="row_wise")`, Q에 `replicate`). +- PE 안에서 할당된 token은 **연속 append**된다(slot 0,1,2,…). global index는 + strided(`i, i+N, …`)이나 per-PE 버퍼는 dense ⇒ DMA가 연속 유지. +- Slot → global: `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`. + +### 2.2 드라이버 launch당 의무 (최소) +드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를 +도출한다: + +| Launch arg | Purpose | +|---|---| +| `K_base[pe]`, `V_base[pe]` | per-PE KV 버퍼 base (tensor VA에서) | +| `O_base` | 어텐션 출력 목적지 | +| `global_token_counter` | 현재 sequence 위치 | +| `start_pe = request_id mod N` | Layer-2 회전 | +| `pe_id` (`tl.program_id`), `N` | geometry | +| `q_block_meta` | prefill/SP query block 시작 + 길이 | + +커널 도출(평범한 산술): +- **이번 step에 내가 쓸 차례:** `(start_pe + counter) mod N == pe_id` +- **내 valid 길이:** `base = counter // N; rem = counter % N; + my_len = base + (1 if ((pe_id − start_pe) mod N) < rem else 0)` +- **read 범위:** tile `0 .. ⌈my_len / TILE⌉`. +- **causal bound / per-tile skip:** query block과 각 tile의 global 위치에서. + +드라이버 = base + counter + rotation. 단일 공식 +`(request_id + token_idx) mod N`이 placement policy 전부다. + +--- + +## 3. Per-tile op 시퀀스 (greenlet `tl`) + +한 iteration = 한 PE의 한 KV tile. 두 GEMM은 `tl.composite(op="gemm")` +(scheduler 관리 tiling + K/V DMA 스트리밍); softmax 머지는 그 사이의 커널 `tl` +MATH다. 실제 `tl` 이름(`tl_context.py`), lazy `tl.load`(ADR-0062), +`tl.scratch_scope`(ADR-0063) 사용: + +```python +# running state (persistent arena — allocated once, outside the scope) +# m: [G, T_q] l: [G, T_q] O: [G, T_q, d] +q_g = tl.load(Q_group_ptr, (G*T_q, d)) # lazy; auto-wait at first use (ADR-0062) + +with tl.scratch_scope(): # per-tile MATH temporaries recycled + Sj = tl.composite("gemm", a=q_g, # [G·T_q, TILE]; scheduler streams Kⱼ DMA + b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale + if mask_j is not None: + Sj = Sj + mask_j # additive causal mask (boundary tile) + m_j = tl.max(Sj, axis=-1) + m_new = tl.maximum(m, m_j) + P = tl.exp(Sj - m_new) # no full-matrix softmax; streaming + corr = tl.exp(m - m_new) # rescale factor for old accumulators + l = l * corr + tl.sum(P, axis=-1) + Oj = tl.composite("gemm", a=P, # [G·T_q, d]; scheduler streams Vⱼ DMA + b=tl.ref(V_base + j*TILE*d, (TILE, d))) + O = O * corr + Oj # running merge (kernel MATH) + m = m_new +``` + +Notes: +- `q_g`는 `[G·T_q, d]`로 reshape된 GQA-batched query(`G` group 행을 matmul M + 차원에 fold; byte-보존). 한 K/V tile이 모든 `G·T_q` 행을 서비스 — GQA 재사용 + 레버 — broadcast 없이. +- **K/V는 `tl.ref` operand**로 composite scheduler가 HBM에서 타일당 스트리밍한다 + (`pe_scheduler.py:104-143`): 그것이 *바로* prefetch/파이프라인이므로 명시적 + prefetch op이 없다. 타일 `j+1`의 composite를 타일 `j`를 wait하기 전에 발행 + (non-blocking 핸들)하면 scheduler가 타일을 가로질러 파이프라인되고 GEMM 엔진이 + saturate 유지된다. +- `tl.trans`는 kernbench에서 **메타데이터 전용**이고(`tl_context.py:390`) + `MemoryStore.read`는 transpose가 아니라 *reshape*한다(`memory_store.py:73`). + zero/structural 실행엔 무해하나; 비자명 수치 데이터에선 reshape-not-transpose가 + 되어 transpose된 K로의 Q·Kᵀ는 주의가 필요하다(§11) — K를 transpose해 `[d, TILE]`로 + 미리 저장하거나, 실제 `tl.transpose`를 추가(후보 primitive, 시뮬레이터의 + 성능-모델링 목적상 아마 불필요). +- Masking: 커널이 query/KV global offset에서 boundary-tile mask를 만들어 더한다 + (`Sj + mask_j`); full-past tile은 `None` 전달; full-future tile은 **skip**(`if`이 + enqueue 안 함). +- **새** "flash-composite" 커맨드 종류(softmax 머지 + carried `(m,ℓ,O)`를 + 내재화하는 것)는 쓰지 **않는다**; 기존 `CompositeCmd`가 각 GEMM을 담당하고 머지는 + 커널에 남는다(§1, §8 항목 4). + +--- + +## 4. Reduction (KV-parallel / SP combine) — root로의 tree + +tile sweep 후 각 PE는 `(m_i, ℓ_i, O_i)`(비정규화)를 가진다. 결합/교환 가능한 +log-sum-exp 머지로 결합(baseline의 fold와 동일 수학, `_attention_mesh_mlo.py:117-122`): + +```python +def merge(m_a, l_a, O_a, m_b, l_b, O_b): + m = tl.maximum(m_a, m_b) + sa, sb = tl.exp(m_a - m), tl.exp(m_b - m) + return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb +# final: O = O_root / l_root # normalise once at the tree root +``` + +- **타이밍:** 교환은 **P·V 후**(각 PE가 최종 `O_i`를 가짐). +- **토폴로지:** mesh tree, depth `⌈log₂ N⌉`, 물리 이웃 상. `N=8`의 경우: level-0 + `(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2 `(3↔7)`, root = `7`. 각 tree + pair는 물리 `N/S/E/W` 이웃이어야 `tl.send(dir)`/`tl.recv(dir)`이 실제 link를 + 쓴다(튜닝 항목 §9; 이웃을 wiring하는 SFR install은 + `configure_sfr_intracube_pe_ring`, baseline이 사용). +- **왜 baseline fan-out이 아니라 tree인가:** 어텐션은 한 곳(`O_base`를 쓰는 PE)에서 + `O`가 필요하다. tree는 `⌈log₂ N⌉` 단계(`N=8`에 3) vs baseline all-to-all의 `N−1` + (7). per-PE sweep이 짧은 decode에서 reduction이 지배적 비용이므로 실제 이득이다. + (baseline fan-out은 답을 모든 rank에 복제 — 여기선 불필요.) +- **Payload:** `O_i`(`d`-vector)가 무겁고; `m,ℓ`은 `(g, T_q)`당 scalar. 별도 + `tl.send`로 전송(baseline은 triplet을 세 send로 — 같은 패턴). + +커널 구조(greenlet): + +```python +# leaf / internal node: fold children that send to me, then forward up +for child_dir in tree_children_dirs(pe_id, N): + m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape) + O_c = tl.recv(child_dir, O.shape) + m, l, O = merge(m, l, O, m_c, l_c, O_c) +if not is_root(pe_id): + tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O) +else: + tl.store(O_base, O / l) +``` + +`tl.recv`는 데이터 도착까지 blocking(`tl_context.py:446-499`); 머지 순서가 정적 +tree로 고정되므로 data-dependent `pop`이 불필요 — **하드웨어/composite +`pop`-as-dependency 변경이 필요 없는** 이유(§6). + +--- + +## 5. Case별 커널 + +네 case 모두 §3(tile sweep) + §4(reduction)를 공유하며; `N`, query-block 폭, +어느 `tile_*` predicate가 발동하는지에서만 다르다. + +### 5.1 공통 skeleton + +```python +def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, + q_block, softmax_scale, *, tl): + pe_id = tl.program_id(axis=rank_axis) + my_len = valid_len(counter, start_pe, pe_id, N) + n_tiles = ceil(my_len / TILE) + q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast) + m, l, O = init_running() # persistent arena: -inf, 0, zeros + for j in range(n_tiles): # K/V streamed per tile by the composite scheduler (§3) + if tile_all_future(j, q_block): # causal skip (kernel if) + continue + run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope + if N == 1: + tl.store(o_ptr, O / l) # no reduction + else: + tree_reduce_and_store(pe_id, N, o_ptr) # §4 +``` + +### 5.2 DECODE, no SP (`N=1`, 한 PE가 head의 KV 소유) +- `T_q = 1`, 모든 과거 KV에 attend ⇒ future tile 없음, 마지막 ragged tile만 mask. +- **GQA 재사용이 전부**(decode는 KV-load-bound): KV head의 `G=8` query 행을 matmul + **M** 차원에 fold(`q_g`를 `[G, T_q, d] → [G·T_q, d]`로 reshape, byte-보존). + 그러면 `Q·Kᵀ`는 `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]`이고 + `P·V`는 `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. KV tile(`[TILE, d]`)이 + 공유 `K`/`V` operand — **한 번 스트리밍되어 모든 `G·T_q` 행이 자동 재사용** — + 그것들이 GEMM의 M 행이기 때문. K/V broadcast 불필요; composite의 tile plan 내 + `m = G·T_q`가 timing이 모든 `G` 행의 작업을 올바르게 세게 한다(leading batch + 축은 세어지지 *않음* — §8 참조). +- `S_pe`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial + attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) — + 바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에 + composite 경로. Tiling(§3)은 `S_pe`가 scratch scope의 tile 예산을 초과할 때만 + 발동. + +### 5.3 DECODE, with SP / KV-parallel (`N=8`) +- 한 request의 KV가 8 PE에 round-robin; 각각 ≈`my_len` token 소유. 각 PE가 자기 + 로컬 tile에 대해 `G=8` query 행을 GQA-batch. +- `T_q=1` ⇒ 짧은 per-PE sweep ⇒ **reduction 지배** ⇒ §4 tree(3단계)가 구조적으로 + 중요한 부분. Reduction latency는 batch될 때 다른 동시 decode token으로, 혹은 긴 + single-stream context의 긴 per-PE tile sweep으로 숨겨진다. + +### 5.4 PREFILL, no SP +- 전체 prompt 상주; query는 `T_q` token 블록(chunk). +- causality가 실재, query block `[qs, qe)` vs KV tile `[ks, ke)`: + - `ke ≤ qs` → `tile_all_past` → `mask=None`, full compute. + - `ks ≥ qe` → `tile_all_future` → **skip**(커널 `if`). + - overlap → `tile_partial` → 커널이 삼각 additive mask 생성, tile op 시퀀스가 + 더함(`Sj + mask_j`). +- GQA는 group의 query head를 decode와 같은 축으로 batch. + +### 5.5 PREFILL, with SP (Ring KV) +- KV가 `N` PE(ring)에 sharded; 각 step이 peer의 KV 블록을 IPCQ로 전달. 실행 중 + `(m,ℓ,O)`는 **ring step을 가로질러** carry(Python 핸들, `_attention_mesh_kv`가 + 오늘 하는 그대로). +- IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로 + 오버랩(`tl_context.py:543-560` — 이미 존재). +- **Causal ring skip:** 들어오는 step의 KV 블록이 전부 로컬 query block *이후*면 그 + compute를 skip(recv-consume 안 함 / fold 안 함) — causal 어텐션에서 약 절반 step + 제거. +- 수신 버퍼는 ping-pong(persistent arena, 재활용 안 함). +- ring 후, `(m,ℓ,O)`는 query-소유 PE별 최종; KV-parallel 분할이 남으면 §4가 + tree-merge한 뒤 정규화 + store. + +baseline `_attention_mesh_kv`가 이미 ring fold를 구현한다; 본 ADR은 GQA 재사용, +causal step-skip, `recv_async` 오버랩을 추가한다. + +--- + +## 6. 왜 하드웨어 / composite 변경이 불필요한가 + +- HW/composite `pop`-as-dependency가 도움 될 유일한 곳은 poll을 숨길 다른 작업이 + 없는 외로운 reduction — 즉 batch=1 **그리고** 짧은 context. 타깃은 **agentic = + 낮은 batch, 긴 context** ⇒ 각 PE가 많은 KV tile을 가짐; §4 tree의 `tl.recv` + blocking은 동시 token / 긴 context의 sweep 작업으로 가려진다. +- §4 reduction은 **정적** tree를 쓰므로 recv 순서가 컴파일 타임에 고정 — + data-dependent pop 없음. `tl.recv`(blocking)로 충분. +- **결정: greenlet `tl.send`/`tl.recv` collective로 출시.** 짧은-context, + single-stream, latency-critical 타깃이 나타날 때만 HW `pop` 재검토. + +--- + +## 7. 제어 vs 실행 분할 (load-bearing 원칙) + +| Concern | Owner | +|---|---| +| tile skip (future), mask 생성, causal bound | **커널** (greenlet 본문의 Python `if` + 산술) | +| 주소 / offset / valid-length 산술 | **커널** (counter에서) | +| reduction 스케줄링, IPCQ send/recv 순서 | **커널** (정적 tree) | +| Q·Kᵀ, P·V (per-tile K/V DMA 스트리밍 + tiling 포함) | **`tl.composite`** → PE_SCHEDULER | +| Q load, mask add, softmax math, 실행 중 `(m,ℓ,O)` 머지 | **`tl` op** PE 엔진 상 (커널 발행) | + +커널이 결정하고; GEMM은 composite로 scheduler에 offload되며; 나머지 `tl` op은 +이미 결정된 작업을 실행한다. 이는 kernbench가 이미 지원하는 greenlet + composite +모델 그대로 — 새 제어 추상화 없음. + +--- + +## 8. 필요한 kernbench 변경 + +**설계 iteration의 정정:** 진짜 GQA(`h_q > h_kv`)는 **신규 primitive 불필요** — +§5.2의 커널 재구조화(KV head별, `G`를 M에 fold, byte-보존 reshape)만 필요. 보조 +ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다. + +**커널 내 알고리즘 작업 (신규 primitive 없음; 기존 `tl` API):** + +- **GQA Q축 batching**(재사용 레버) — KV head별로 `G·T_q`를 matmul M 차원에 fold + (§5.2); `_view` 식 byte-보존 reshape; GEMM은 M = `G·T_q`인 + `tl.composite(op="gemm")`. 오늘 timing/data 모드 모두 동작. +- **composite를 통한 GEMM**(§1/§3) — Q·Kᵀ와 P·V를 각각 non-blocking + `tl.composite(op="gemm")`로 발행; PE_SCHEDULER가 tiling하고 `tl.ref` K/V + operand의 DMA를 스트리밍(기존 `CompositeCmd`; 새 커맨드 종류 없음). +- root로의 tree reduction(§4)이 baseline all-to-all fan-out 대체 — `tl.send`/ + `tl.recv` 상의 순수 커널 제어 흐름. +- causal tile skip + additive boundary mask(§3/§5.4) — 커널 `if` + `+`로 더한 + mask 텐서. +- round-robin KV placement / valid-length 산술(§2) — launch-arg 산술 + + `DPPolicy(pe="row_wise")`. + +**efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):** + +1. **per-tile scratch 재활용** — **ADR-0063**(`tl.scratch_scope`). *스케일에 필요*: + `S=16` 상한(1 MiB bump allocator)을 제거해 현실적 context 길이가 돌게 함. 셋 중 + 최고 가치. +2. **Lazy `tl.load`** — **ADR-0062**(non-blocking load + 첫 사용 시점 auto-wait; API + 표면 불변). *효율*: 명시적 load(Q group, 비-composite 커널)를 뒤따르는 compute와 + 오버랩. 타일당 **K/V** prefetch는 composite scheduler가 처리(§1)하므로, 이것은 + 나머지 명시적 load를 담당. 전역 시맨틱 변경 → 기존 골든 재생성(ADR-0062 D3). +3. **GQA head / mask broadcast** — **ADR-0061**(`tl.broadcast`). *선택적 편의*, GQA + blocker 아님(위 정정 참조). `G·T_q` 행에 걸친 additive-mask 구성과 범용 커널에 + 유용; `np.matmul`이 data 모드에서 이미 broadcast하므로 정확성엔 불필요. 최저 + 우선순위. + +**명시적 REJECTED (효율적 대안 선택):** + +4. ~~내부 루프 전체를 내재화하는 전용 "flash-composite" 커맨드 종류 — + DMA→MM→VEC→DMA→MM→VEC + carried `(m,ℓ,O)` register 상태 + 꼬리 IPCQ push.~~ 두 + GEMM은 기존 `CompositeCmd`를 **쓴다**(§1/§3) — 그것이 scheduler 관리 tiling, K/V + DMA 스트리밍, cross-tile 파이프라이닝을 준다. 기각되는 것은 softmax 머지 + + cross-tile register 수명 + IPCQ-push epilogue까지 흡수하는 **새** 커맨드 종류다: + 크고 특수목적이며, 하이브리드 대비 유일한 delta(완전 softmax offload)가 현재 + 모델링 충실도(per-op CPU 발행 비용 = 0; §1 참조)에서 정당화되지 않는다. cost + model(§9)이 완전 offload를 측정 가능하게 가치 있게 만들면 재검토. +5. ~~하드웨어 `pop`-as-dependency.~~ 범위 밖(§6). +6. ~~본 커널 내 RoPE / QKV projection / KV-cache write.~~ Upstream `qkv_rope` + (P1–P5). RoPE를 fold하면 매 decode step마다 과거 tile을 재회전해야 하고 Ring + Attention의 post-RoPE pass-through를 깬다. + +--- + +## 9. Open 튜닝 항목 (kernbench에서 측정, blocking 아님) + +1. **Composite tile-pipeline depth** — KV-load-bound에 지배적; 커널이 wait 전에 + non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당 + 스트리밍 depth. +2. **reduction tree의 PE↔mesh-이웃 매핑** — 각 depth-`⌈log₂ N⌉` pair가 물리 + `N/S/E/W` 이웃인지 보장; 나쁜 매핑은 hop 추가. SFR install 대조 검증. +3. **TILE 크기** — scratch 상주(`S/P tile + O_acc + G-way GQA`)와 DMA 효율 균형; + ADR-0063 및 scheduler의 `TILE_M/K/N`(`pe_scheduler.py`)과 상호작용. +4. §5.5의 **Ring 버퍼 ping-pong vs `recv_async` depth**. +5. **decode의 one-shot vs tiled 교차점**(§5.2) — tiling이 단일 composite를 이기는 + `S_pe` 임계값. +6. **per-op CPU 발행 비용 (cost model)** — 현재 `dispatch_cycles=0`(`pe_cpu.py`)이라 + composite-vs-primitive 발행 오버헤드가 보이지 않음. op 종류별 차등 발행 비용 + (`tl.composite` descriptor push ≫ primitive op)이 하이브리드의 CPU-saturation + 이점을 **측정 가능**하게 한다(§1). **ADR-0064**에 명세; 별도 future work로 추적. + +--- + +## 10. Coverage 요약 + +| Case | KV placement | Inner loop | Reduction | Masking | +|---|---|---|---|---| +| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only | +| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only | +| Prefill, no SP | resident | tiled per q-block | none | skip future / triangular boundary | +| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary | + +**Case별 I/O** (전체 계약 §0.5): + +| Case | Inputs | Output | IPCQ partials | +|---|---|---|---| +| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none | +| Decode, SP | `Q[G,1,d]`, per-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,ℓ,O_i)` per PE → tree | +| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none | +| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,ℓ,O)` + tree | + +모든 case에서: KV-cache write와 RoPE는 **upstream**에서; output projection은 +**downstream**에서; score `S`와 prob `P`는 미생성. + +--- + +## 11. 검증 계획 (Phase 1 테스트 개요) + +SPEC/ADR coverage: R5 (PE↔PE IPCQ, PE↔HBM), R2 (traversal에 의한 latency), +ADR-0023/0025 (IPCQ), ADR-0046 (`tl` 계약), ADR-0054 (eval bench). + +시뮬레이터의 계약은 bit-exact 수치가 아니라 **traversal에 의한 latency + 결정성 + +구조적 정확성**(SPEC §0, §0.1)이다 — Phase 2 데이터는 주로 data path를 +exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`은 `f16`으로 +모델된다. 따라서 검증은 **구조/타이밍 우선**, 수치 parity는 제한된 2차 체크. + +1. **data 모드 실행(`enable_data=True`):** GQA 커널(`h_q = G·h_kv`)이 baseline + head-packing이 부딪히는 byte-conservation 에러 없이 — 네 case 모두 — 완료. (이는 + §5.2 재구조화가 필요하지, 신규 primitive가 *아님*.) + **수치 parity(2차):** reshape-as-transpose가 정확한 symmetric/identity 입력에 + 대해, 커널 `O`가 numpy FlashAttention reference와 fp tolerance 내 일치. 완전 + asymmetric parity는 실제 `tl.transpose`에 좌우됨(범위 밖; 플래그됨). +2. **GQA 재사용:** `h_q = G·h_kv`에서 K/V `dma_read_count`가 `G`와 무관(타일당 한 + load, group에 걸쳐 재사용), 반면 GEMM 작업은 `G`로 스케일. 레버가 실제 발동함을 + assert. +3. **tree reduction 단계 수:** decode-SP가 `⌈log₂ N⌉` reduction 라운드(`N−1`이 + 아님)를 발행; op_log `ipcq_send`/`recv` 수가 tree와 일치. +4. **Causal skip:** prefill이 모든 `tile_all_future` tile을 skip — GEMM 수가 full + grid가 아니라 하삼각 tile 수와 일치. +5. **Long context (ADR-0063):** scope 없이 1 MiB를 넘기는 `S` sweep이 완료하고 + reference와 일치. +6. **Load/compute 오버랩:** tiled sweep의 end-to-end latency가 직렬 + `Σ(load+compute)`보다 낮음 — composite scheduler의 타일당 K/V 스트리밍(§1/§3)과 + Q load를 오버랩하는 lazy `tl.load`(ADR-0062)에서. (오버랩은 실제 모델 동시성, + 빼기가 아님.) +7. **Composite GEMM offload (구조적):** 각 tile의 Q·Kᵀ와 P·V가 blocking `tl.dot`이 + 아니라 `CompositeCmd`(non-blocking)를 PE_SCHEDULER에 emit; op_log가 composite + tile plan을 보이고 커널이 다음 tile의 composite를 wait 전에 발행(cross-tile + 파이프라이닝). +8. **결정성:** 동일 입력 → 동일 op_log + latency (SPEC §0.1). + +--- + +## B. 하이브리드 전환에서 나온 Open 설계 항목 (추후 검토) + +이들은 결정이 순수 greenlet primitive 경로에서 **composite hybrid + lazy +`tl.load`**(이번 개정)로 옮겨갈 때 생겼다. 설계를 막는 것은 없으며; 각각 구현 중 +검증 패스가 필요하다. (묻지 않고) 작업 합의에 따라 여기 기록한다 — 권고는 내가 +예측한 기본값이며; 검토 시 수정. + +1. **DDD-0060이 아직 미동기화.** Detailed Design Document는 여전히 옛 + `tl.load_async` 더블버퍼 경로와 primitive `tl.dot` 내부 루프를 기술한다(그 + §4.3/§5/§10). 하이브리드(composite GEMM, lazy load, K 사전 transpose)로 + 갱신해야 한다. DDD가 파생 how-to이고 큰 재작성이라 *검토용으로 남김*; ADR이 이제 + 권위 있는 기록이다. **권고:** 구현 시작 전 후속으로 DDD 동기화. + +2. **composite GEMM의 K operand 방향.** Q·Kᵀ는 `b = [d, TILE]`가 필요하나 KV + cache는 K를 `[S_pe, d]`로 저장한다. `tl.trans`는 메타데이터 전용이고 + `MemoryStore.read`는 transpose가 아니라 reshape(`memory_store.py:73`) — 비자명 + 데이터엔 런타임 transpose가 틀리다. **권고:** cache에 K를 **사전 transpose** + `[d, S_pe]`로 저장(pseudocode와 §3이 이를 가정)하여 `tl.ref(k_tile, (d, TILE))`이 + 연속 slice가 되게. upstream `qkv_rope` write 레이아웃이 이를 지원하는지 검증, + 아니면 실제 `tl.transpose` 추가(더 무거움; 연기). + +3. **composite 출력 버퍼 vs `tl.scratch_scope`.** 각 Q·Kᵀ composite는 `Sj`를 + `out_addr`에 쓰고; 커널이 softmax MATH용으로 그것을 읽는다. 그 출력 버퍼와 + in-flight composite의 타깃은 per-tile `scratch_scope`(ADR-0063)가 소비 전에 + 재활용하지 **않는** 곳에 살아야 한다 — in-flight lazy load와 같은 규율(ADR-0062 + D-Negative). **권고:** *현재* tile의 composite 출력은 scoped arena에(같은 + iteration에 소비); persistent `(m,ℓ,O)`는 바깥에. 다음 tile의 composite를 일찍 + 발행할 때(cross-tile 파이프라이닝) use-after-recycle이 없음을 검증. + +4. **composite 스트리밍 하의 GQA `dma_read_count` 레버.** 레버(§11.2: K/V + `dma_read_count`가 `G`와 무관)는 composite가 모든 `G·T_q` M-행에 재사용되는 + **하나의** K/V tile DMA를 emit한다고 가정한다. scheduler의 `generate_gemm_plan`은 + `TILE_M/K/N`(`pe_scheduler.py:35-37`, 32/64/32)으로 tiling한다 — `G·T_q`에 대한 + M-tiling이 M-tile마다 공유 K/V tile DMA를 재발행하지 **않음**을 확인(즉 operand + DMA가 M-tile에 걸쳐 공유되거나, 아니면 레버가 약해짐). **권고:** levers 테스트에서 + assert; 위반 시 GQA 이득은 DMA가 아니라 compute에만 — 여전히 정확하나 헤드라인이 + 바뀜. + +5. **커널 TILE vs scheduler `TILE_M/K/N`.** 커널은 논리적 KV `TILE`을 다루고; + scheduler는 고정 `TILE_M/K/N`으로 내부 재tiling한다. 두 tiling 레이어가 + 상호작용(scratch 상주, 파이프라인 depth). **권고:** 커널 TILE을 K/V 스트리밍 + granularity로 보고 scheduler가 GEMM을 sub-tile하게; DDD에 관계를 문서화하고 둘 다 + sweep(§9 항목 1, 3). + +6. **cost model은 별도 ADR.** 하이브리드의 CPU-saturation 이점은 + `dispatch_cycles=0`인 동안 보이지 않는다. op 종류별 issue-cost 모델은 + **ADR-0064**에 명세; 본 ADR의 §1/§9는 *측정 가능한*(구조적뿐 아닌) 이점을 위해 + 그것에 의존. **권고:** eval에서 하이브리드 latency 이점을 주장하기 전에 ADR-0064 + 모델을 도입. + +7. **Ring 경로(§5.5) GEMM.** §5.5는 여전히 primitive op + `recv_async`로 ring + fold를 기술한다. 일관성을 위해 ring의 per-step Q·Kᵀ / P·V도 composite여야 하고; + IPCQ `recv_async` 오버랩은 직교하며 유지. **권고:** 구현 중 ring step에 같은 + 하이브리드 형태 적용; 낮은 리스크, §3 미러. diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md index 12e17c2..4f82995 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -16,10 +16,12 @@ GQA, causal, long-context kernel. **Supporting ADRs** (efficiency / scale enablers — *not* GQA blockers; see §8 correction): **ADR-0063** `tl.scratch_scope` (per-tile scratch -recycling — required for realistic context length), **ADR-0062** -`tl.load_async` (KV prefetch overlap — efficiency), **ADR-0061** -`tl.broadcast` (optional mask/general convenience). Real GQA itself needs -only kernel restructuring (§5.2). +recycling — required for realistic context length), **ADR-0062** lazy +`tl.load` (non-blocking load with auto-wait on first use → load/compute +overlap), **ADR-0061** `tl.broadcast` (optional mask/general +convenience). The two GEMMs (Q·Kᵀ, P·V) are issued as scheduler-managed +`tl.composite` commands (the existing `CompositeCmd`; no new command +kind); real GQA itself needs only kernel restructuring (§5.2). **Algorithm lineage.** This kernel is **FlashAttention** (tiling + online/streaming softmax with fused P·V — no full score matrix @@ -32,6 +34,69 @@ known algorithms onto the kernbench **greenlet `tl` programming model** --- +## TL;DR — final GQA kernel (composite hybrid) pseudocode + +The decision (§1) in one place: **GEMMs → `tl.composite`; softmax merge + +tree reduction → kernel; `tl.load` is lazy.** Per-KV-head, `G` folded into +the matmul M dim (real GQA, no broadcast). This is the reference shape; the +prose sections elaborate each piece. + +```python +def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr, + counter, start_pe, N, q_block, scale, *, tl): + # ---- geometry (kernel arithmetic, §2) ---- + pe_id = tl.program_id(axis=rank_axis) + G = H_q // H_kv # query heads per KV head (=8) + causal = q_block.is_prefill + + for kv in range(H_kv): # one KV head per iteration (§5.2) + # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062): + # issue now, auto-wait at first use inside the first composite. + q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d] + my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv + n_tiles = ceil(my_len / TILE) + + # persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros + m, l, O = init_running(G * q_block.T_q, d) + + for j in range(n_tiles): + if causal and tile_all_future(j, q_block): + continue # causal skip (kernel if) + with tl.scratch_scope(): # per-tile temporaries (ADR-0063) + # --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite --- + Sj = tl.composite("gemm", a=q_g, + b=tl.ref(k_tile(kv, j), (d, TILE))) * scale # Kᵀ pre-stored [d,TILE] + if causal and tile_partial(j, q_block): + Sj = Sj + causal_mask(j, q_block) # additive boundary mask + # --- online softmax merge in the kernel (MATH ops) --- + m_new = tl.maximum(m, tl.max(Sj, axis=-1)) + P = tl.exp(Sj - m_new) + corr = tl.exp(m - m_new) + l = l * corr + tl.sum(P, axis=-1) + # --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite --- + Oj = tl.composite("gemm", a=P, + b=tl.ref(v_tile(kv, j), (TILE, d))) + O = O * corr + Oj # running merge + m = m_new + + # ---- cross-PE combine (§4): log-sum-exp tree to root, or store ---- + if N == 1: + tl.store(o_base(kv), O / l) + else: + tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv +``` + +> **Why this shape:** the two `tl.composite` calls offload tiling + K/V DMA +> streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse +> descriptors and runs ahead → engines stay saturated, §1); the softmax +> merge and the IPCQ tree reduction stay in the kernel because the existing +> `CompositeCmd` cannot carry cross-tile state. `K` is pre-stored +> transposed `[d, TILE]` to sidestep the reshape-not-transpose caveat +> (§3, §B). A bespoke "flash-composite" command kind is **not** introduced +> (§8 item 4). + +--- + ## A. Relationship to existing kernbench work (read first) kernbench **already runs FlashAttention with an online-softmax `(m, ℓ, O)` @@ -47,11 +112,14 @@ Both are driven by `milestone-gqa-llama70b` (4 panels: `src/kernbench/benches/milestone_gqa_llama70b.py`) and tested in `tests/attention/test_milestone_gqa_llama70b.py`. -**They are written in the greenlet `tl` API — not composites:** `tl.load`, -`tl.dot`, `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, -and Python `-`/`*`/`/` on `TensorHandle` (each emits a `MathCmd`). The -running `(m, ℓ, O)` is just Python `TensorHandle`s threaded through the -loop. This **matters for this ADR's design** (see §1). +**They are written in the greenlet `tl` API:** `tl.load`, `tl.dot`, +`tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, and Python +`-`/`*`/`/` on `TensorHandle` (each emits a `MathCmd`) — the GEMMs as +**blocking `tl.dot`**, not composites. The running `(m, ℓ, O)` is just +Python `TensorHandle`s threaded through the loop. This ADR keeps the +running state and the softmax merge in the kernel but **moves the two +GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this +**matters for the design**. **Three deliberate limitations of the baseline** — exactly what an *efficient GQA* kernel must lift: @@ -74,7 +142,8 @@ loop. This **matters for this ADR's design** (see §1). allocator leaks per-tile temporaries (`test_milestone_gqa_llama70b.py:123-148`) and there is no causal masking / tiling → fixed by **ADR-0063** (recycling) + §5 (tiling, - causal skip) + **ADR-0062** (prefetch). + causal skip) + composite K/V streaming (§3) + **ADR-0062** (lazy load + overlap). **Documentation debt (out of scope but recorded):** the baseline cites ADR-0055/0056/0057/0058/0059, **none of which exist as files** — they are @@ -103,10 +172,14 @@ Hardware recap: (`tl_context.py:402-499`). - **Composite command** (`CompositeCmd`, `pe_commands.py:144-162`): a single GEMM (or MATH) *head* plus element-wise *epilogue* stages - (`bias/relu/scale/add/...`). It is **not** a general multi-op DAG: it - cannot chain two GEMMs, cannot carry register state across instances, - and cannot pop/wait on IPCQ. This ADR therefore does **not** require - composites for the attention inner loop (see §1, §8). + (`bias/relu/scale/add/...`), issued **non-blocking** to PE_SCHEDULER, + which generates a tile plan and streams DMA→GEMM→write per tile + (ADR-0014 D6; `pe_scheduler.py:104-143`). It is **not** a general + multi-op DAG: it cannot chain two GEMMs, cannot carry register state + across instances, and cannot pop/wait on IPCQ. This ADR therefore issues + **each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and + keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel — + it does **not** need a new "flash-composite" command kind (see §1, §8). - **Allocation policy (SP):** one query head per CUBE in the multi-user panels; the `G` query heads of a KV group map within a CUBE's PEs. @@ -194,51 +267,82 @@ projection (downstream), score `S` and probs `P` (never materialised). ## 1. Decision (mechanism) -**Implement the kernel in the greenlet `tl` programming model, not as -composites.** Rationale grounded in kernbench's execution + latency model: +**Implement the kernel as a *hybrid*: issue the two GEMMs (Q·Kᵀ and P·V) +as scheduler-managed `tl.composite(op="gemm")` commands, and keep the +online-softmax merge and the cross-PE reduction as kernel-level `tl` +ops.** `tl.load` is **lazy** (non-blocking; the wait is auto-inserted at +first use of the loaded data — ADR-0062), so explicit HBM loads overlap +the compute that follows. Rationale grounded in kernbench's execution + +latency model: -- The simulator charges latency **per op on its modelled component** - (GEMM on PE_GEMM, vector ops on PE_MATH, DMA on PE_DMA — `pe_gemm.py`, - `pe_math.py`, `pe_dma.py`). "Fusing" several ops into one - `CompositeCmd` does **not** reduce modelled latency; the stages still - run on the same engines. The win a real fused kernel gets — *overlap* - (prefetch) and *small working set* (scratch recycling) — is obtained - here by **ADR-0062** and **ADR-0063**, which are smaller and reusable. -- The running `(m, ℓ, O)` flash state is naturally a set of Python - `TensorHandle`s threaded through the loop (the baseline already does - this). No composite-carried register state is needed. -- Cross-PE combination uses kernel-level `tl.send`/`tl.recv` (the - baseline already does this and it is tested). No composite-driven IPCQ - push is needed. +- **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is + non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): + the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole + per-PE tile sweep) and the scheduler generates the tile plan and streams + DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the + scheduler streams from HBM, so per-tile **K/V prefetch is the + scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed + to issue the **next** composite while the current one runs, so the + scheduler keeps the GEMM engine saturated across tiles. +- **This reflects the hardware** and decouples CPU issue-rate from + execution-rate. The blocking per-op `tl.dot` path, by contrast, stalls + the CPU on every GEMM and leaves GEMM-engine **bubbles** during the + interleaved softmax MATH ops; it is realistic only if the CPU can keep + up with fine-grained per-tile issue. +- **The running `(m, ℓ, O)` flash state stays Python `TensorHandle`s** + threaded through the loop (the baseline already does this); the softmax + merge (max/exp/sum/rescale) is kernel-level `tl` MATH **between** the two + GEMM composites. The existing `CompositeCmd` cannot chain two GEMMs or + carry cross-tile register state (§0), so the merge necessarily lives in + the kernel — this is the hybrid split, not a limitation worked around. +- **Cross-PE combination** is a log-sum-exp **tree** over `(m, ℓ, O)` + after P·V, via kernel-level `tl.send`/`tl.recv` (§4) — unchanged. -So the per-tile inner pipeline is the op sequence +So the per-tile inner pipeline is: ``` -K_j load → Q·Kⱼᵀ → (+ causal mask) → online-softmax update → V_j load → P·Vⱼ → running (m,ℓ,O) update +q_g = tl.load(Q group) # lazy; auto-wait at first use +per tile j: + Sⱼ = tl.composite("gemm", a=q_g, b=tl.ref(Kⱼ)) → Sⱼ # scheduler streams Kⱼ DMA + GEMM + Sⱼ += maskⱼ # kernel MATH, boundary tile only + online-softmax: mⱼ, m_new, P, corr, ℓ # kernel MATH + Oⱼ = tl.composite("gemm", a=P, b=tl.ref(Vⱼ)) → Oⱼ # scheduler streams Vⱼ DMA + GEMM + O = O*corr + Oⱼ; m = m_new # kernel MATH (running merge) ``` -issued as ordinary `tl.*` ops, with the **next** tile's K/V issued via -`tl.load_async` (ADR-0062) so its DMA overlaps the current tile's -compute, and each tile's temporaries wrapped in `tl.scratch_scope` -(ADR-0063) so scratch stays O(1). - -Cross-PE combination (KV-parallel / SP) is a **log-sum-exp tree -reduction** over `(m, ℓ, O)` after P·V, flowed through IPCQ with -kernel-level `tl.send`/`tl.recv` (§4). +with each tile's MATH temporaries wrapped in `tl.scratch_scope` +(ADR-0063) so scratch stays O(1), and the next tile's composites issued +before the current tile's results are waited on (non-blocking handles) so +the scheduler pipelines across tiles. Control flow (tile skip, mask generation, reduction scheduling, address arithmetic) lives in the **kernel** (plain Python `if`/arithmetic in the greenlet body). This is exactly what kernbench's greenlet model already permits (`kernel_runner.py`, ADR-0020 D3). -> **Why this is the efficient choice.** It reuses the proven baseline -> kernels and the proven IPCQ collective; the only genuinely new -> machinery is three small, general primitives (ADR-0061/0062/0063). The -> originally-proposed "everything-in-one-composite + composite IPCQ push -> + composite-carried state" would require a bespoke flash-composite -> command type, a register-lifetime model across composites, and an -> IPCQ-push epilogue — large, special-purpose, and with no latency -> benefit over the greenlet path. Those are **rejected**; see §8. +> **What this supersedes.** An earlier iteration proposed a pure greenlet +> primitive path (all `tl.dot`, no composite) on the argument that +> "composite yields no latency benefit." That holds **only because** the +> simulator currently charges **zero** per-op CPU issue cost +> (`dispatch_cycles=0`, `pe_cpu.py`) — it models away exactly the CPU +> issue-rate / DMA-program cost that descriptor offload exists to hide. +> The hybrid is the faithful representation of an efficient kernel. The +> **measurable** size of the win (can the CPU saturate the engines for +> many tiles?) is gated on modelling an op-type-differentiated issue cost, +> tracked as future work (cost model; §9). Even at `dispatch_cycles=0` the +> non-blocking composite path fills the GEMM-engine bubbles the blocking +> `tl.dot` path leaves. +> +> **Why this is the efficient choice.** GEMM tiling + DMA streaming + +> cross-tile pipelining are offloaded to the proven `CompositeCmd` +> scheduler path; the softmax merge and the proven IPCQ collective stay in +> the kernel. The only genuinely new machinery is two small, general +> primitives (ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); the +> reduction reuses `tl.send`/`tl.recv`. A bespoke "flash-composite" +> command (one kind internalising the softmax merge + carried register +> state + an IPCQ-push epilogue) is **not** built — large, special-purpose, +> and its only delta over this hybrid (full softmax offload) is not +> justified at the current modelling fidelity; see §8. --- @@ -283,17 +387,20 @@ Driver = bases + counter + rotation. The single formula ## 3. Per-tile op sequence (greenlet `tl`) -One iteration = one KV tile on one PE. Real `tl` names (`tl_context.py`), -with `tl.load_async` (ADR-0062), `tl.broadcast` (ADR-0061), -`tl.scratch_scope` (ADR-0063): +One iteration = one KV tile on one PE. The two GEMMs are +`tl.composite(op="gemm")` (scheduler-managed tiling + K/V DMA streaming); +the softmax merge is kernel `tl` MATH between them. Real `tl` names +(`tl_context.py`), with lazy `tl.load` (ADR-0062), `tl.scratch_scope` +(ADR-0063): ```python # running state (persistent arena — allocated once, outside the scope) # m: [G, T_q] l: [G, T_q] O: [G, T_q, d] +q_g = tl.load(Q_group_ptr, (G*T_q, d)) # lazy; auto-wait at first use (ADR-0062) -with tl.scratch_scope(): # per-tile temporaries recycled - Kj = tl.wait(f_k[j]) # prefetched (ADR-0062) - Sj = tl.dot(q_g, tl.trans(Kj)) * softmax_scale # [G, TILE]; q_g is GQA-batched +with tl.scratch_scope(): # per-tile MATH temporaries recycled + Sj = tl.composite("gemm", a=q_g, # [G·T_q, TILE]; scheduler streams Kⱼ DMA + b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale if mask_j is not None: Sj = Sj + mask_j # additive causal mask (boundary tile) m_j = tl.max(Sj, axis=-1) @@ -301,32 +408,35 @@ with tl.scratch_scope(): # per-tile temporaries rec P = tl.exp(Sj - m_new) # no full-matrix softmax; streaming corr = tl.exp(m - m_new) # rescale factor for old accumulators l = l * corr + tl.sum(P, axis=-1) - Vj = tl.wait(f_v[j]) - O = O * corr + tl.dot(P, Vj) # P·V folded into running O + Oj = tl.composite("gemm", a=P, # [G·T_q, d]; scheduler streams Vⱼ DMA + b=tl.ref(V_base + j*TILE*d, (TILE, d))) + O = O * corr + Oj # running merge (kernel MATH) m = m_new - if j + PREFETCH < n_tiles: # keep the pipeline full - f_k[j + PREFETCH] = tl.load_async(K_base + (j+PREFETCH)*TILE*d, (TILE, d)) - f_v[j + PREFETCH] = tl.load_async(V_base + (j+PREFETCH)*TILE*d, (TILE, d)) ``` Notes: - `q_g` is the GQA-batched query reshaped to `[G·T_q, d]` (the `G` group - rows folded into the matmul M dim; byte-conserving). One K/V tile load - serves all `G·T_q` rows — the GQA reuse lever — with no broadcast. -- `tl.trans(Kj)` is **metadata-only** in kernbench (`tl_context.py:390`), - and `MemoryStore.read` *reshapes* rather than transposes + rows folded into the matmul M dim; byte-conserving). One K/V tile serves + all `G·T_q` rows — the GQA reuse lever — with no broadcast. +- **K/V are `tl.ref` operands** the composite scheduler streams from HBM + per tile (`pe_scheduler.py:104-143`): that *is* the prefetch/pipeline, + so there is no explicit prefetch op. Issuing tile `j+1`'s composites + before waiting on tile `j` (non-blocking handles) keeps the scheduler + pipelined across tiles and the GEMM engine saturated. +- `tl.trans` is **metadata-only** in kernbench (`tl_context.py:390`) and + `MemoryStore.read` *reshapes* rather than transposes (`memory_store.py:73`). For zero/structural runs this is harmless; for - non-trivial numeric data it yields a reshape-not-transpose. Data-mode - *numeric* parity therefore needs care (§11) — store K pre-transposed, - or add a real `tl.transpose` (a candidate further primitive, likely + non-trivial numeric data it yields a reshape-not-transpose, so Q·Kᵀ via + a transposed K needs care (§11) — store K pre-transposed `[d, TILE]`, or + add a real `tl.transpose` (a candidate further primitive, likely unnecessary given the simulator's performance-modeling purpose). -- The V load is issued (prefetched) *before* it is needed so its DMA - overlaps the Q·Kᵀ + softmax of the same/earlier tile. - Masking: the kernel builds the boundary-tile mask from query/KV global offsets and adds it (`Sj + mask_j`); full-past tiles pass `None`; full-future tiles are **skipped** (the `if` never enqueues them). -- No `CompositeCmd` is used. (A future `tl.composite` flash kind is an - *optional* optimisation — §8 item 4 — not a requirement.) +- A **new** "flash-composite" command kind (one that internalises the + softmax merge + carried `(m,ℓ,O)`) is **not** used; the existing + `CompositeCmd` covers each GEMM and the merge stays in the kernel + (§1, §8 item 4). --- @@ -395,13 +505,12 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, pe_id = tl.program_id(axis=rank_axis) my_len = valid_len(counter, start_pe, pe_id, N) n_tiles = ceil(my_len / TILE) - q_g = load_Q_group(q_ptr) # [G,d] decode / [G,T_q,d] prefill; tl.broadcast for GQA + q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast) m, l, O = init_running() # persistent arena: -inf, 0, zeros - prime_prefetch(k_ptr, v_ptr, n_tiles) # tl.load_async first PREFETCH tiles - for j in range(n_tiles): + for j in range(n_tiles): # K/V streamed per tile by the composite scheduler (§3) if tile_all_future(j, q_block): # causal skip (kernel if) continue - run_tile(j, mask_or_null(j, q_block)) # §3 op sequence, in tl.scratch_scope + run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope if N == 1: tl.store(o_ptr, O / l) # no reduction else: @@ -414,18 +523,18 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, - **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8` query rows of the KV head are folded into the matmul **M** dimension (`q_g` reshaped `[G, T_q, d] → [G·T_q, d]`, a byte-conserving reshape). - Then `Q·Kᵀ` is `tl.dot([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]` and - `P·V` is `tl.dot([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. The KV tile - (`[TILE, d]`) is the shared `K`/`V` operand — **loaded once, reused by + Then `Q·Kᵀ` is `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]` and + `P·V` is `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. The KV tile + (`[TILE, d]`) is the shared `K`/`V` operand — **streamed once, reused by all `G·T_q` rows automatically** because they are the M rows of the - GEMM. No broadcast of K/V is needed; `m = G·T_q` in the emitted - `GemmCmd` also makes the Phase-1 timing count all `G` rows' work - correctly (a leading batch axis would *not* be counted — see §8). + GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile + plan also makes the timing count all `G` rows' work correctly (a leading + batch axis would *not* be counted — see §8). - If `S_pe` fits in scratch (small/medium context) this degenerates to a - **one-shot** partial attention (one `tl.dot` for Q·Kᵀ, one softmax, one - `tl.dot` for P·V) — exactly the baseline `_attention_mesh_mlo` - `_partial_attention`, just GQA-batched. Tiling (§3) only kicks in when - `S_pe` exceeds the scratch scope's tile budget. + **one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one + composite for P·V) — exactly the baseline `_attention_mesh_mlo` + `_partial_attention`, just GQA-batched and on the composite path. Tiling + (§3) only kicks in when `S_pe` exceeds the scratch scope's tile budget. ### 5.3 DECODE, with SP / KV-parallel (`N=8`) - One request's KV is round-robin across 8 PEs; each owns ≈`my_len` @@ -484,11 +593,13 @@ ADR adds GQA reuse, causal step-skip, and `recv_async` overlap. | tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) | | address / offset / valid-length arithmetic | **kernel** (from counter) | | reduction scheduling, IPCQ send/recv ordering | **kernel** (static tree) | -| K/V load, Q·Kᵀ, mask add, softmax math, P·V | **`tl` ops** on PE engines | +| Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER | +| Q load, mask add, softmax math, running `(m,ℓ,O)` merge | **`tl` ops** on PE engines (kernel-issued) | -The kernel decides; the `tl` ops execute already-decided work. This is -exactly the greenlet model kernbench already supports — no new control -abstraction. +The kernel decides; the GEMMs are offloaded to the scheduler as +composites; the remaining `tl` ops execute already-decided work. This is +exactly the greenlet + composite model kernbench already supports — no new +control abstraction. --- @@ -502,8 +613,13 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head, **Algorithm work in the kernel (no new primitive; existing `tl` API):** - **GQA Q-axis batching** (the reuse lever) — fold `G·T_q` into the matmul - M dim per KV head (§5.2); `_view`-style byte-conserving reshape + 2-D - `tl.dot`. Runs today in both timing and data mode. + M dim per KV head (§5.2); `_view`-style byte-conserving reshape; the + GEMM is a `tl.composite(op="gemm")` with M = `G·T_q`. Runs today in both + timing and data mode. +- **GEMMs via composite** (§1/§3) — Q·Kᵀ and P·V each issued as a + non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and + streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new + command kind). - Tree reduction to root (§4) replacing the baseline all-to-all fan-out — pure kernel control flow over `tl.send`/`tl.recv`. - Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` + @@ -517,9 +633,12 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head, *Required for scale*: removes the `S=16` ceiling (1 MiB bump allocator) so realistic context lengths run. Highest-value of the three. -2. **Async HBM tile load / KV prefetch** — **ADR-0062** (`tl.load_async` - + `tl.wait`). *Efficiency*: the KV-load-bound overlap lever for - decode/long-context. Without it the kernel is correct but serial. +2. **Lazy `tl.load`** — **ADR-0062** (non-blocking load + auto-wait on + first use; API surface unchanged). *Efficiency*: overlaps explicit + loads (the Q group, non-composite kernels) with following compute. The + per-tile **K/V** prefetch is handled by the composite scheduler (§1), + so this covers the remaining explicit loads. Global semantics change → + existing goldens regenerate (ADR-0062 D3). 3. **GQA head / mask broadcast** — **ADR-0061** (`tl.broadcast`). *Optional convenience*, not a GQA blocker (see correction above). Useful for additive-mask construction across the `G·T_q` rows and for @@ -528,16 +647,16 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head, **Explicitly REJECTED (efficient alternative chosen):** -4. ~~Composite that chains DMA→MM→VEC→DMA→MM→VEC with carried `(m,ℓ,O)` - register state and a tail IPCQ push.~~ Replaced by the greenlet path: - per-op latency is identical; overlap comes from ADR-0062; small - working set from ADR-0063; running state is Python handles; the IPCQ - push is `tl.send`. Building a bespoke flash-composite command type + - cross-composite register lifetime + an IPCQ-push epilogue is large, - special-purpose, and yields no latency benefit. **A tiled - `tl.composite` "flash" kind remains a possible *future* optimisation** - (it would fold prefetch+recycling into the scheduler), but it is not - required for an efficient kernel and is out of scope here. +4. ~~A bespoke "flash-composite" command kind that internalises the whole + inner loop — DMA→MM→VEC→DMA→MM→VEC with carried `(m,ℓ,O)` register + state and a tail IPCQ push.~~ The two GEMMs **do** use the existing + `CompositeCmd` (§1/§3) — that gives scheduler-managed tiling, K/V DMA + streaming, and cross-tile pipelining. What is rejected is a **new** + command kind that also absorbs the softmax merge + cross-tile register + lifetime + an IPCQ-push epilogue: it is large and special-purpose, and + its only delta over the hybrid (full softmax offload) is not justified + at the current modelling fidelity (per-op CPU issue cost = 0; see §1). + Revisit if the cost model (§9) makes full offload measurably worthwhile. 5. ~~Hardware `pop`-as-dependency.~~ Out of scope (§6). 6. ~~RoPE / QKV projection / KV-cache write inside this kernel.~~ Upstream `qkv_rope` (P1–P5). Folding RoPE in would force re-rotating past tiles @@ -547,16 +666,24 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head, ## 9. Open tuning items (measured in kernbench, not blocking) -1. **KV tile prefetch depth** (ADR-0062) — dominant for KV-load-bound; - sweep 2→4. +1. **Composite tile-pipeline depth** — dominant for KV-load-bound; how far + ahead the kernel issues non-blocking composites before waiting, and the + scheduler's per-tile streaming depth. 2. **PE↔mesh-neighbour mapping for the reduction tree** — ensure each depth-`⌈log₂ N⌉` pair is a physical `N/S/E/W` neighbour; bad mapping adds hops. Verify against the SFR install. -3. **TILE size** — balance scratch residency (`K_tile + S/P + V_tile + - O_acc + G-way GQA`) against DMA efficiency; interacts with ADR-0063. +3. **TILE size** — balance scratch residency (`S/P tiles + O_acc + G-way + GQA`) against DMA efficiency; interacts with ADR-0063 and the + scheduler's `TILE_M/K/N` (`pe_scheduler.py`). 4. **Ring buffer ping-pong vs `recv_async` depth** in §5.5. 5. **One-shot vs tiled crossover for decode** (§5.2) — the `S_pe` - threshold where tiling beats a single `tl.dot`. + threshold where tiling beats a single composite. +6. **Per-op CPU issue cost (cost model)** — currently `dispatch_cycles=0` + (`pe_cpu.py`), so composite-vs-primitive issue overhead is invisible. + An op-type-differentiated issue cost (a `tl.composite` descriptor push + ≫ a primitive op) is what makes the hybrid's CPU-saturation win + **measurable** (§1). Specified in **ADR-0064**; tracked as separate + future work. --- @@ -612,7 +739,79 @@ secondary check. count equals the lower-triangular tile count, not the full grid. 5. **Long context (ADR-0063):** a sweep at `S` that overflows 1 MiB without scopes completes and matches the reference. -6. **Prefetch overlap (ADR-0062):** end-to-end latency of the tiled sweep - is below the serial `Σ(load+compute)` (overlap is real). -7. **Determinism:** identical inputs → identical op_log + latency +6. **Load/compute overlap:** end-to-end latency of the tiled sweep is + below the serial `Σ(load+compute)` — from the composite scheduler + streaming K/V per tile (§1/§3) and lazy `tl.load` (ADR-0062) overlapping + the Q load. (Overlap is real modelled concurrency, not a subtraction.) +7. **Composite GEMM offload (structural):** each tile's Q·Kᵀ and P·V emit a + `CompositeCmd` (non-blocking) to PE_SCHEDULER, not a blocking `tl.dot`; + op_log shows the composite tile plan and the kernel issues the next + tile's composites before waiting (cross-tile pipelining). +8. **Determinism:** identical inputs → identical op_log + latency (SPEC §0.1). + +--- + +## B. Open design items from the hybrid pivot (review later) + +These arose when the decision moved from a pure greenlet primitive path to +the **composite hybrid + lazy `tl.load`** (this revision). None blocks the +design; each needs a verification pass during implementation. Recorded here +(rather than asked) per the working agreement — the recommendation is my +predicted default; revise on review. + +1. **DDD-0060 is not yet synced.** The Detailed Design Document still + describes the old `tl.load_async` double-buffer path and primitive + `tl.dot` inner loop (its §4.3/§5/§10). It must be updated to the hybrid + (composite GEMMs, lazy load, K pre-transposed). *Left for review* + because the DDD is a derived how-to and a large rewrite; the ADR is now + the authoritative record. **Recommend:** sync DDD as a follow-up before + implementation starts. + +2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b = + [d, TILE]`, but the KV cache stores K as `[S_pe, d]`. `tl.trans` is + metadata-only and `MemoryStore.read` reshapes, not transposes + (`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial + data. **Recommend:** store K **pre-transposed** `[d, S_pe]` in the cache + (the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))` + a contiguous slice. Verify the upstream `qkv_rope` write layout supports + this, or add a real `tl.transpose` (heavier; deferred). + +3. **Composite output buffer vs `tl.scratch_scope`.** Each Q·Kᵀ composite + writes `Sj` to an `out_addr`; the kernel then reads it for the softmax + MATH. That output buffer, and the in-flight composites' targets, must + live where the per-tile `scratch_scope` (ADR-0063) will **not** recycle + them before they are consumed — same discipline as in-flight lazy loads + (ADR-0062 D-Negative). **Recommend:** composite outputs for the *current* + tile live in the scoped arena (consumed same iteration); the persistent + `(m,ℓ,O)` stays outside. Verify no use-after-recycle when the next + tile's composites are issued early (cross-tile pipelining). + +4. **GQA `dma_read_count` lever under composite streaming.** The lever + (§11.2: K/V `dma_read_count` independent of `G`) assumes the composite + emits **one** K/V tile DMA reused across all `G·T_q` M-rows. The + scheduler's `generate_gemm_plan` tiles by `TILE_M/K/N` + (`pe_scheduler.py:35-37`, 32/64/32) — confirm the M-tiling over `G·T_q` + does **not** re-issue the shared K/V tile DMA per M-tile (i.e. operand + DMA is shared across M-tiles, or the lever weakens). **Recommend:** + assert it in the levers test; if violated, the GQA win is in compute + only, not DMA — still correct, but the headline changes. + +5. **Kernel TILE vs scheduler `TILE_M/K/N`.** The kernel reasons about a + logical KV `TILE`; the scheduler re-tiles internally at fixed + `TILE_M/K/N`. Two tiling layers interact (scratch residency, pipeline + depth). **Recommend:** treat the kernel TILE as the K/V streaming + granularity and let the scheduler sub-tile the GEMM; document the + relationship in the DDD and sweep both (§9 items 1, 3). + +6. **Cost model is a separate ADR.** The hybrid's CPU-saturation benefit is + invisible while `dispatch_cycles=0`. The per-op-type issue-cost model is + specified in **ADR-0064**; this ADR's §1/§9 depend on it for the + *measurable* (not just structural) win. **Recommend:** land ADR-0064's + model before claiming hybrid latency wins in the eval. + +7. **Ring path (§5.5) GEMMs.** §5.5 still describes the ring fold with + primitive ops + `recv_async`. For consistency the ring's per-step Q·Kᵀ / + P·V should also be composites; the IPCQ `recv_async` overlap is + orthogonal and stays. **Recommend:** apply the same hybrid shape to the + ring step during implementation; low risk, mirrors §3. diff --git a/docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md b/docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md new file mode 100644 index 0000000..83d12f2 --- /dev/null +++ b/docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md @@ -0,0 +1,148 @@ +# ADR-0062: Lazy `tl.load` — 첫 사용 시점 auto-wait를 갖는 non-blocking HBM load + +## Status + +Proposed + +> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context +> 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본 +> ADR은 별도의 `load_async` op을 추가하는 대신 `tl.load` 자체를 **lazy** +> (non-blocking, wait를 첫 사용 시점에 자동 삽입)로 만든다. 오늘날 유일한 +> async primitive는 IPCQ 통신용이지 HBM load용이 아니다. + +## Context + +### 오버랩이 요구하는 것 + +FlashAttention은 operand를 스트리밍한다: GEMM/MATH 엔진이 현재 타일을 처리하는 +동안 DMA 엔진은 이미 다음 operand를 당기고 있어야 한다. 그 오버랩이 있으면 +대역폭-bound 커널은 타일당 `compute + dma` 대신 대략 `max(compute, dma)`로 +돈다. + +ADR-0060의 하이브리드 설계에서 두 GEMM은 `tl.composite(op="gemm")`로 발행되고, +그 K/V operand는 `tl.ref`(HBM 상주, PE_SCHEDULER가 타일당 스트리밍)이므로 **타일당 +K/V prefetch는 composite scheduler가 처리**한다. lazy `tl.load`는 나머지 명시적 +load(Q group, 그리고 비-composite 커널)를 담당하여 그것들도 greenlet을 멈추는 +대신 뒤따르는 compute와 오버랩되게 한다. + +### 현재 존재하는 것 + +- `tl.load(ptr, shape, dtype)`은 **blocking**이다: `DmaReadCmd`를 emit하고 + greenlet 커널은 PE_DMA가 완료를 알릴 때까지 suspend된다 + (`tl_context.py:177-203`; greenlet 구동 `kernel_runner.py:146-153`). 한 + 커널에서 두 `tl.load`가 in-flight일 수 없다. +- async 패턴이 **존재**하긴 하나 IPCQ 전용이다: + `tl.recv_async(dir, ...) -> RecvFuture` + 지연 wait + (`kernel_runner.py:248-285`). 이는 메커니즘을 입증한다 — future를 반환하고 + 나중에 wait 체크(`if not future.event.triggered: yield future.event`)로 + resolve되는 non-blocking 커맨드가 greenlet 모델에서 동작한다. +- DMA 엔진은 read 채널을 write 채널과 분리된 SimPy resource(capacity 1, + `pe_dma.py:45`)로 모델한다. 따라서 in-flight read는 표현 가능하며, 단일 read + 채널에서 직렬화되면서 PE_GEMM/PE_MATH의 compute와 오버랩된다. 오직 + *커널-대면 API*만이 현재 load-vs-compute를 직렬화한다. + +`tl.load`는 오늘날 뒤따르는 compute와 오버랩될 수 없다. + +## Decision + +**`tl.load`를 lazy로 만든다: `DmaReadCmd`를 발행하고 핸들을 즉시 반환 +(non-blocking); 런타임은 load된 데이터가 실제로 처음 소비되는 지점에 wait를 +자동 삽입한다.** 커널-대면 API는 불변이다 — 작성자는 계속 `tl.load`를 쓴다 — +따라서 이는 새 op이 아니라 *시맨틱* 변경이다. 기존 `recv_async`/wait 기계장치를 +(1) HBM-load 경로로, (2) 명시적 `tl.wait` 호출에서 암묵적·의존성 기반 first-use +wait로 일반화한다. + +### D1. `tl` 표면 — 불변 + +`tl.load(ptr, shape, dtype)`은 시그니처와 `TensorHandle` 반환을 유지한다. +바뀌는 것은 *언제* blocking하느냐다: 발행 시점에는 결코 아니고, 결과가 소비 op에 +의해 처음 읽힐 때만 암묵적으로. + +### D2. 메커니즘 — non-blocking 발행 + 사용 시점 auto-wait + +- `tl.load`는 완료 event를 yield하지 않고 `DmaReadCmd`를 PE_DMA에 post하며, + 반환 핸들에 pending event를 기록한다(`recv_async` 패턴을 load에 적용). +- 소비 op(`tl.dot`, MATH op, `tl.store`, `tl.composite` operand, …)이 + dispatch될 때, 런타임은 각 입력 핸들의 pending load event를 확인하고 아직 + triggered가 아니면 먼저 yield한다 — 즉 wait가 가장 늦은 올바른 지점(첫 사용)에 + 자동 삽입된다. +- op_log 엔트리는 불변(`memory/dma_read`): asynchrony는 스케줄링 속성이지 새 op + 종류가 아니므로 `dma_read_count` 및 기존 op_log 소비자가 계속 동작한다. + +### D3. 범위 — 전역 + +`tl.load`는 opt-in 플래그 뒤가 아니라 **모든 곳에서** lazy다. 이것이 충실한 +모델이다: 매 load마다 blocking하는 것은 *naive* 커널의 속성이고, 효율 커널(그리고 +실제 컴파일러)은 load를 hoist하고 사용 시점에만 wait한다. 결과: `tl.load`와 첫 +사용 사이에 독립 작업이 있는 기존 커널은 **더 낮은(빠른) latency**를 본다 — +동작의 회귀가 아니라 모델의 정확도 개선. 바뀌는 골든 latency는 **재생성**해야 +하며; load 직후 곧바로 사용하는 커널은 변화가 없다(auto-wait이 즉시 발동하여 +blocking과 동일). + +### D4. Latency / 오버랩 시맨틱 + +- `tl.load`는 **발행**(descriptor push)만 청구하고 커널은 진행한다. (per-op + 발행 비용은 `dispatch_cycles`, 현재 0; op 종류별 차등 발행 비용은 별도로 + 추적 — ADR-0060 §1, §9, ADR-0064.) +- DMA 전송은 read 채널(capacity 1)을 모델 지속시간만큼 점유하며 커널이 다음에 + 내는 compute와 병렬; 다중 in-flight load는 채널에서 직렬화되나 compute와 + 오버랩. +- 자동 삽입된 wait는 전송이 끝나지 않았을 때만 blocking. +- 결정성 보존: wait 지점은 프로그램 순서(첫 사용)로 고정되고, 완료는 modeled DMA + 채널의 scheduled event다(SPEC §0.1, R8). latency 빼기 없음 — 오버랩은 실제 + 모델 동시성. + +> **모델링 가정.** auto-wait-at-first-use는 *잘 스케줄된* 커널(컴파일러가 +> wait를 가장 늦은 올바른 지점에 둠)을 모델한다. 실제 컴파일러는 이를 근사하며, +> 일부 load는 hoist 불가(register pressure, aliasing). 성능 시뮬레이터(SPEC +> §0)에서 잘 스케줄된 경우를 모델하는 것이 의도된 동작이다. + +## Alternatives + +### A1. 별도 `tl.load_async` op (이전 제안) + +커널이 손으로 호출하는 명시적 `load_async`/`wait` 쌍 추가(더블버퍼 댄스). 기각: +커널-대면 API를 키우고 버퍼 수명 부기를 모든 커널 작성자에게 떠넘긴다. lazy +`tl.load`는 **API 표면 변경 없이** 동일 오버랩과 컴파일러식 auto-wait를 준다. + +### A2. 커널별 opt-in lazy load + +`tl.load`를 기본 blocking으로 두고 새 커널만 opt-in. 기각: `tl.load` 시맨틱을 두 +변종으로 쪼개고 기존 bench로부터 모델 개선을 숨긴다; 전역 lazy가 더 깔끔하고 +골든 재생성 비용은 1회성(D3). + +### A3. composite 스트리밍에만 의존 (lazy load 없이) + +ADR-0060의 composite는 `tl.ref` K/V operand를 스트리밍하므로 GEMM operand DMA는 +이미 오버랩된다. 그러나 명시적 load(Q group, 비-composite 커널)는 lazy `tl.load` +없이 여전히 stall. 단독으로는 불충분. + +## Consequences + +### Positive +- 커널-대면 API 변경 **제로**로 load/compute 오버랩; 작성자는 계속 `tl.load`. +- `recv_async`/wait와 대칭 — 개념적 표면적 작음. +- 일반적: 어떤 대역폭-bound 커널도 자동 prefetch. + +### Negative +- load와 사용 사이에 독립 작업이 있는 커널의 기존 골든 latency가 이동(빨라짐) + → 1회성 재생성(D3). +- auto-wait는 런타임이 핸들별 pending event를 추적하고 소비-op dispatch에서 + 확인하도록 요구(데이터 의존성 추적). +- scratch 수명(ADR-0063)과 상호작용: in-flight load의 대상 버퍼는 auto-wait이 + 발동하기 전에 recycle되면 안 된다. recycling scope는 살아있는(미-wait된) load + 버퍼를 제외해야 한다. + +## Test Requirements + +1. **오버랩이 실재**: `tl.load` 후 비슷한 지속시간의 독립 GEMM을 내는 커널이 + `load+gemm`이 아니라 ≈`max(load, gemm)`으로 완료(end-to-end latency가 직렬 + 합보다 엄격히 작음). +2. **auto-wait 정확성**: load된 `TensorHandle`이 처음 소비될 때, 같은 주소에 + 대해 오늘날 blocking `tl.load`와 동일한 바이트를 보유(Phase 2). +3. **둘 in-flight**: 서로 다른 주소로의 두 `tl.load`가 나중에 소비되어 둘 다 + 올바른 독립 텐서로 resolve; 그 DMA는 read 채널에서 직렬화. +4. **op_log 호환성**: 각 `tl.load`은 여전히 정확히 하나의 `memory/dma_read`를 + 로그; `dma_read_count`는 blocking 버전 대비 불변. +5. **무-오버랩 무-변화**: load 직후 곧바로 사용하는 커널은 blocking 모델과 동일한 + latency(auto-wait 즉시 발동). diff --git a/docs/adr-proposed/ADR-0062-prog-tl-async-load.md b/docs/adr-proposed/ADR-0062-prog-tl-async-load.md index e12310d..b200583 100644 --- a/docs/adr-proposed/ADR-0062-prog-tl-async-load.md +++ b/docs/adr-proposed/ADR-0062-prog-tl-async-load.md @@ -1,148 +1,169 @@ -# ADR-0062: `tl.load_async` — non-blocking HBM tile load for KV prefetch +# ADR-0062: Lazy `tl.load` — non-blocking HBM load with auto-wait on first use ## Status Proposed > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and -> long-context attention are **KV-load-bound**; the dominant lever is -> overlapping the next K/V tile's DMA with the current tile's compute. -> Today the only async primitive is for IPCQ comms, not for HBM loads. +> long-context attention are **KV-load-bound**; load/compute overlap is a +> dominant lever. This ADR makes `tl.load` itself **lazy** (non-blocking, +> with the wait automatically inserted at first use) rather than adding a +> separate `load_async` op. Today the only async primitive is for IPCQ +> comms, not for HBM loads. ## Context ### What overlap requires -FlashAttention streams K/V tiles: while the MATH/GEMM engine works on -tile `j` (Q·Kⱼᵀ → softmax → P·V), the DMA engine should already be -pulling tile `j+1` (and `j+2` at prefetch depth 2). With that overlap a -KV-load-bound kernel runs at roughly `max(compute, dma)` per tile instead -of `compute + dma`. +FlashAttention streams operands: while the GEMM/MATH engine works on the +current tile, the DMA engine should already be pulling the next operand. +With that overlap a bandwidth-bound kernel runs at roughly +`max(compute, dma)` per tile instead of `compute + dma`. + +In ADR-0060's hybrid design the two GEMMs are issued as +`tl.composite(op="gemm")` whose K/V operands are `tl.ref` (HBM-resident, +streamed per tile by PE_SCHEDULER), so the **per-tile K/V prefetch is +handled by the composite scheduler**. Lazy `tl.load` covers the remaining +explicit loads (the Q group, and any non-composite kernel) so those also +overlap the compute that follows instead of stalling the greenlet. ### What exists -- `tl.load(ptr, shape, dtype)` is **blocking**: it emits `DmaReadCmd` - and the greenlet kernel suspends until PE_DMA signals completion - (`src/kernbench/triton_emu/tl_context.py:177-203`; greenlet drive in - `kernel_runner.py`). No two `tl.load`s can be in flight from one - kernel. +- `tl.load(ptr, shape, dtype)` is **blocking**: it emits `DmaReadCmd` and + the greenlet kernel suspends until PE_DMA signals completion + (`tl_context.py:177-203`; greenlet drive `kernel_runner.py:146-153`). + No two `tl.load`s can be in flight from one kernel. - An async pattern **does** exist, but only for IPCQ: - `tl.recv_async(dir, ...) -> RecvFuture` + `tl.wait(future)` - (`tl_context.py:543-560`, `tl_context.py:660-693`). It proves the - machinery — a non-blocking command that returns a future, resolved - later by `tl.wait` — works in the greenlet model. -- The DMA engine already models a read channel as a SimPy resource - separate from the write channel - (`src/kernbench/components/builtin/pe_dma.py`), so concurrent in-flight - reads are representable at the component layer; only the *kernel-facing - API* serialises them. + `tl.recv_async(dir, ...) -> RecvFuture` + a deferred wait + (`kernel_runner.py:248-285`). It proves the machinery — a non-blocking + command that returns a future, resolved later by a wait check + (`if not future.event.triggered: yield future.event`) — works in the + greenlet model. +- The DMA engine models a read channel as a SimPy resource (capacity 1, + `pe_dma.py:45`) separate from the write channel, so in-flight reads are + representable; they serialise on the single read channel while + overlapping compute on PE_GEMM/PE_MATH. Only the *kernel-facing API* + currently serialises load-vs-compute. -There is **no** `tl.load_async`. Prefetch/double-buffering cannot be -expressed. +`tl.load` cannot today overlap with the compute that follows it. ## Decision -Add a non-blocking HBM load that mirrors the existing `recv_async`/`wait` -contract. +**Make `tl.load` lazy: it issues the `DmaReadCmd` and returns a handle +immediately (non-blocking); the runtime auto-inserts the wait at the +first point the loaded data is actually consumed.** The kernel-facing API +is unchanged — authors keep writing `tl.load` — so this is a *semantics* +change, not a new op. It generalises the existing `recv_async`/wait +machinery (1) to the HBM-load path and (2) from an explicit `tl.wait` +call to an implicit, dependency-driven wait at first use. -### D1. `tl` surface +### D1. `tl` surface — unchanged -```python -def load_async(self, ptr: int, shape: tuple[int, ...], - dtype: str = "f16") -> LoadFuture: - """Issue a DMA read and return immediately. Resolve with tl.wait(fut) - -> TensorHandle. Multiple loads may be in flight; ordering of - resolution is by tl.wait calls, not issue order.""" -``` +`tl.load(ptr, shape, dtype)` keeps its signature and `TensorHandle` +return. What changes is *when* it blocks: never at issue, only implicitly +when its result is first read by a consuming op. -`tl.wait` is extended to accept a `LoadFuture` (it already dispatches on -`CompletionHandle` vs `RecvFuture` — add a third arm), returning the -loaded `TensorHandle`. +### D2. Mechanism — non-blocking issue + auto-wait on use -### D2. Command +- `tl.load` posts the `DmaReadCmd` to PE_DMA without yielding its + completion event, and records the pending event on the returned handle + (the `recv_async` pattern, applied to loads). +- When a consuming op (`tl.dot`, a MATH op, `tl.store`, a `tl.composite` + operand, …) is dispatched, the runtime checks each input handle for a + pending load event and yields it first if not yet triggered — i.e. the + wait is inserted automatically at the latest correct point (first use). +- The op_log entry is unchanged (`memory/dma_read`): asynchrony is a + scheduling property, not a new op kind, so `dma_read_count` and existing + op_log consumers keep working. -Reuse the existing `DmaReadCmd` with a `blocking=False` flag (mirroring -`IpcqRecvCmd.blocking`), or a thin `DmaReadAsyncCmd` sibling — whichever -keeps PE_DMA's handler simplest. The op_log entry is unchanged -(`memory/dma_read`); asynchrony is a scheduling property, not a new op -kind, so existing op_log consumers and the `dma_read_count` metric -(milestone bench) keep working. +### D3. Scope — global -### D3. Latency / overlap semantics +`tl.load` is lazy **everywhere**, not behind an opt-in flag. This is the +faithful model: blocking on every load is a property of a *naive* kernel; +an efficient kernel (and a real compiler) hoists the load and waits only +at use. Consequence: existing kernels that have independent work between a +`tl.load` and its first use see **lower (faster) latency** — a +correctness improvement of the model, not a behaviour regression. Golden +latencies that change must be **regenerated**; kernels that load then +immediately use see no change (the auto-wait fires at once, identical to +blocking). -- `load_async` charges the **issue** (descriptor push) only; the kernel - proceeds. -- The DMA transfer occupies the read channel for its modelled duration - in parallel with whatever compute the kernel issues next. -- `tl.wait(fut)` blocks only if the transfer has not finished; if it has, - it returns immediately (same fast-path as `recv_async`/`wait`). -- Determinism is preserved: completion is a scheduled event on the - modelled DMA channel/link (SPEC §0.1, R8). No magic — the overlap is - real modelled concurrency, not a hand-waved latency subtraction. +### D4. Latency / overlap semantics -### D4. Double-buffer usage (the intended pattern) +- `tl.load` charges the **issue** (descriptor push) only; the kernel + proceeds. (Per-op issue cost is `dispatch_cycles`, currently 0; an + op-type-differentiated issue cost is tracked separately — ADR-0060 §1, + §9.) +- The DMA transfer occupies the read channel (capacity 1) for its + modelled duration in parallel with whatever compute the kernel issues + next; multiple in-flight loads serialise on the channel but overlap + compute. +- The auto-inserted wait blocks only if the transfer has not finished. +- Determinism is preserved: the wait point is fixed by program order + (first use), and completion is a scheduled event on the modelled DMA + channel (SPEC §0.1, R8). No latency subtraction — the overlap is real + modelled concurrency. -```python -# Prime: issue first two tile loads. -f0 = tl.load_async(K_base + 0*tile_bytes, (TILE, d)) -f1 = tl.load_async(K_base + 1*tile_bytes, (TILE, d)) -for j in range(n_tiles): - Kj = tl.wait(f_cur) # resolves the j-th load - if j + 2 < n_tiles: # keep depth-2 pipeline full - f_next2 = tl.load_async(K_base + (j+2)*tile_bytes, (TILE, d)) - s = tl.dot(q, tl.trans(Kj)) # compute overlaps f_{j+1}'s DMA - ... -``` - -(V tiles double-buffer the same way; ADR-0060 §3 schedules the V load so -it overlaps the Q·Kᵀ + softmax of the same tile.) +> **Modelling assumption.** Auto-wait-at-first-use models a *well- +> scheduled* kernel (the compiler places the wait at the latest correct +> point). Real compilers approximate this; some loads cannot be hoisted +> (register pressure, aliasing). For a performance simulator (SPEC §0) +> modelling the well-scheduled case is the intended behaviour. ## Alternatives -### A1. Use a tiled `tl.composite` instead +### A1. Separate `tl.load_async` op (earlier proposal) -`tl.composite` already pipelines tiles inside PE_SCHEDULER (DMA of tile -`j+1` overlaps compute of tile `j`) and recycles per-tile scratch — so -the composite path gets prefetch "for free." Rejected as the *sole* -mechanism because today's composite is a single GEMM head + math -epilogues (`pe_commands.py:144-162`); it cannot express the -two-GEMM-with-running-`(m,l,O)` flash inner loop. Building a bespoke -"flash composite" kind is a much larger change than `load_async`, which -is a small, general primitive that also helps non-composite kernels. -`load_async` and a future flash-composite are not mutually exclusive -(ADR-0060 §8). +Add an explicit `load_async`/`wait` pair the kernel calls by hand (the +double-buffer dance). Rejected: it enlarges the kernel-facing API and +pushes buffer-lifetime bookkeeping onto every kernel author, when lazy +`tl.load` gives the same overlap with **no** API-surface change and a +compiler-style auto-wait. -### A2. Rely on `recv_async` only (no HBM async) +### A2. Per-kernel opt-in lazy load -Only works for the Ring-Attention path where KV arrives over IPCQ -(comms, already overlappable). The decode and no-SP paths load KV from -**local HBM**, which has no async path today. Insufficient. +Keep `tl.load` blocking by default; make new kernels opt in. Rejected: +splits `tl.load` semantics into two variants and hides the model +improvement from existing benches; global lazy is cleaner and the +golden-regeneration cost is one-time (D3). + +### A3. Rely on composite streaming only (no lazy load) + +ADR-0060's composites stream their `tl.ref` K/V operands, so the GEMM +operand DMA already overlaps. But explicit loads (the Q group, and any +non-composite kernel) still stall without lazy `tl.load`. Insufficient on +its own. ## Consequences ### Positive -- Enables the headline KV-load-bound optimisation for decode and - long-context, in the proven greenlet model. -- Symmetric with `recv_async`/`wait` — low conceptual surface area. -- General: any bandwidth-bound kernel can prefetch. +- Load/compute overlap with **zero** kernel-facing API change; authors + keep writing `tl.load`. +- Symmetric with `recv_async`/wait — low conceptual surface area. +- General: any bandwidth-bound kernel prefetches automatically. ### Negative -- Kernels must manage futures and buffer lifetimes explicitly (the - double-buffer dance). Mitigated by keeping the pattern in a shared - helper used by all four ADR-0060 cases. -- Interacts with scratch lifetime (ADR-0063): in-flight buffers must not - be recycled before their `tl.wait`. The recycling scope must exclude - live prefetch buffers. +- Existing golden latencies shift (faster) for kernels with independent + work between load and use → one-time regeneration (D3). +- Auto-wait requires the runtime to track per-handle pending events and + check them at consuming-op dispatch (data-dependency tracking). +- Interacts with scratch lifetime (ADR-0063): an in-flight load's target + buffer must not be recycled before its auto-wait fires. The recycling + scope must exclude live (un-waited) load buffers. ## Test Requirements -1. **Overlap is real**: a kernel that issues `load_async` then a GEMM of - comparable duration completes in ≈`max(load, gemm)`, not `load+gemm` - (assert end-to-end latency strictly below the serial sum). -2. **Correctness**: `tl.wait(load_async(...))` returns the same bytes as - blocking `tl.load(...)` for the same address (Phase 2). -3. **Two in flight**: two outstanding `load_async`s to distinct - addresses both resolve to correct, independent tensors. -4. **op_log compatibility**: each `load_async` still logs exactly one - `memory/dma_read`, so `dma_read_count` metrics are unchanged vs the - blocking version. +1. **Overlap is real**: a kernel that issues `tl.load` then an + independent GEMM of comparable duration completes in ≈`max(load, + gemm)`, not `load+gemm` (end-to-end latency strictly below the serial + sum). +2. **Auto-wait correctness**: the loaded `TensorHandle`, when first + consumed, carries the same bytes as today's blocking `tl.load` for the + same address (Phase 2). +3. **Two in flight**: two `tl.load`s to distinct addresses, consumed + later, both resolve to correct, independent tensors; their DMAs + serialise on the read channel. +4. **op_log compatibility**: each `tl.load` still logs exactly one + `memory/dma_read`; `dma_read_count` unchanged vs the blocking version. +5. **No-overlap no-change**: a kernel that loads then immediately uses has + identical latency to the blocking model (auto-wait fires at once). diff --git a/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md b/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md new file mode 100644 index 0000000..ac2bc5f --- /dev/null +++ b/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md @@ -0,0 +1,166 @@ +# ADR-0064: op 종류별 CPU 발행 비용 모델 (command construct + dispatch) + +## Status + +Proposed + +> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. 거기서의 하이브리드 +> 결정(GEMM은 `tl.composite`, softmax 머지는 커널)이 이기는 이유는 +> **타일링을 PE_SCHEDULER에 offload하여 CPU가 coarse descriptor를 내리고 +> 앞서 나가 엔진을 saturate**하기 때문이다. 그 이점이 현재 시뮬레이터에서는 +> per-op CPU 발행 비용이 0이라 **보이지 않는다**. 본 ADR은 발행 비용을 +> 실재화하고 **op 종류별로 차등**화하여, composite-vs-primitive 트레이드오프 +> (그리고 CPU saturation)가 측정 가능하도록 한다. + +## Context + +### 현재 상태 + +- 모든 `tl.*` op은 커맨드를 emit하기 전에 `_emit_dispatch_overhead()`를 + 호출하며(`tl_context.py:123-125, 190, 227, 235, 612, …`), 이는 + `dispatch_cycles > 0`일 때**만** `PeCpuOverheadCmd(cycles=dispatch_cycles)`를 + emit한다. 즉 발행 비용은 `tl.load`, `tl.dot`, MATH op, `tl.composite`에 + 동일하게 적용되는 **단일 균일 노브**다. +- 그 노브는 두 실행 경로 모두에서 **0**으로 하드코딩되어 있다 + (`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ 커맨드 발행 + — *descriptor를 construct하고 scheduler 큐에 push* — 이 현재 PE_CPU에서 + **0 ns**다. +- `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 검토에서) + +- **Q1 — composite 발행 비용.** `tl.composite`를 construct + push 하는 비용은 + CPU 시간 기준 **~40 ns** 정도(descriptor 빌드 + 큐 push)로 예상된다. 현재는 + **0**이다. 훅은 존재하며, 값(그리고 op 종류별 차등)만 빠져 있다. +- **Q2 — scheduler dispatch vs DMA latency.** PE_SCHEDULER의 composite + dispatch는 **non-blocking**이다: `_dispatch_composite`는 tile plan을 생성하고 + 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이 하이브리드에 틀린 이유 + +하이브리드의 논거 전체는 **하나의** composite descriptor가 `N_tiles`만큼의 +GEMM 타일링을 offload하므로, CPU가 `O(N_tiles × ops/tile)`개의 fine 커맨드 +대신 `O(1)`개의 coarse 커맨드를 낸다는 것이다. 발행 비용 = 0(그리고 균일)이면 +모델은 다음을 보여줄 수 없다: +- CPU가 충분히 빨리 못 밀어넣을 때 primitive 경로가 엔진을 **saturate 못 할 수 + 있음**(핵심 ADR-0060 §1 주장), 그리고 +- composite가 단일 primitive보다 construct 비용이 **더 크지만** 그것이 대체하는 + 다수의 primitive보다는 훨씬 작음. + +단일 균일 `dispatch_cycles`로는 이를 표현할 수 없다: 실제 construct 비용에서 +`tl.load` ≪ `tl.composite`이기 때문이다. + +## Decision + +### D1. PE_CPU의 op 종류별 발행 비용 테이블 + +단일 `dispatch_cycles` 스칼라를 **커맨드 종류별 비용 테이블**로 교체하고, +발행 시점(커맨드가 scheduler로 dispatch되기 전)에 PE_CPU에서 부과한다. +greenlet이 이 시간을 지불하며, 이것이 바로 CPU가 얼마나 빨리 work를 밀어넣을 +수 있는지를 gate한다 — saturation 레버. + +시작 추정치(추후 calibrate — 검토 항목 참조): + +| 발행 op | CPU 발행 비용 (construct + push) | +|---|---| +| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 추정) | +| `tl.load` / `tl.store` (DMA descriptor) | 작음 (수 ns) | +| `tl.dot` / MATH / IPCQ send/recv | 작음 (수 ns) | + +요점은 **비율**이다: composite construct ≫ primitive 발행, 그러나 그것이 +대체하는 primitive 발행들의 합 ≪. 정확한 ns는 설정 가능하다. + +### D2. 실행 latency는 엔진에 유지 (Q2 — 변경 없음) + +DMA/GEMM/MATH 실행 latency는 오늘처럼 PE_DMA / PE_GEMM / PE_MATH에서 부과된다. +발행 비용(D1)은 **CPU 측에만** 추가되며, `drain_ns`를 건드리지 않으므로 중복 +계상이 없다. 이는 SPEC §0.1(latency는 modeled component에서)을 보존한다. + +### D3. scheduler plan 생성 비용 — 0에서 시작, 재검토 + +PE_SCHEDULER의 tile-plan 생성은 오늘 sim 시간 0이다. 일단 유지한다(지배적 +레버는 CPU 발행 비용 D1); scheduler 측 비용이 유의미하다고 판명되면 기존 +`overhead_ns` node attr로 노출한다. 여기서 결정하지 않고 검토 항목으로 둔다. + +### D4. 설정 가능한 값; 골든 재생성 + +비용 테이블은 설정 가능하다(per-topology / node attrs, 기본값은 D1 추정치). +발행 비용을 0이 아니게 하면 **모든** bench latency가 바뀌므로, 골든 latency를 +한 번 **재생성**한다 — ADR-0062의 전역 lazy-load 회귀와 동일한 자세의 모델 +충실도 개선이다. + +## Alternatives + +### A1. 단일 균일 `dispatch_cycles > 0` 유지 + +기각: op construct 비용은 자릿수 단위로 다르다(`tl.load` vs `tl.composite`). +균일값은 primitive를 과대 청구하거나 composite를 과소 청구하며, 어느 쪽이든 +하이브리드가 의존하는 composite-vs-primitive 트레이드오프를 왜곡한다. + +### A2. 발행 비용을 PE_CPU 대신 PE_SCHEDULER에 부과 + +기각: saturation 질문은 *"CPU가 엔진을 바쁘게 유지할 만큼 descriptor를 빨리 +밀어넣을 수 있는가?"*이며 — 이는 **PE_CPU**의 issue-bandwidth 속성이다. +scheduler에 부과하면 CPU back-pressure를 모델하지 못한다. + +### A3. DMA program/setup 시간을 별도의 고정 per-descriptor 비용으로 모델 + +실제 HW는 전송 시간과 별개로 DMA descriptor program 비용을 지불한다. 연기: +초기에는 descriptor-program 비용을 **발행 op**의 D1 비용에 합친다(DMA를 +유발하는 `tl.load` / composite). calibration이 유의미함을 보이면 PE_DMA 고정 +setup으로 분리한다. 검토 항목. + +## Consequences + +### Positive +- 하이브리드의 CPU-offload / saturation 이점(ADR-0060 §1)이 구조적일 뿐 + 아니라 **측정 가능**해진다. +- composite-vs-primitive 및 타일링 granularity 트레이드오프가 latency에 + 보여 eval 서사를 가능케 한다. +- 하드웨어에 더 충실(issue bandwidth는 실제 병목). + +### Negative +- **모든** bench 골든이 이동 → 한 번의 재생성(D4); CI 골든 fixture 갱신. +- op 종류별 값은 calibration 필요; ~40 ns composite 수치는 추정이고 + primitive는 미지정 — 결과는 숫자만큼만 정확하다(calibrate 전까지 절대 + latency를 과대 주장 금지). +- 발행 경로에 비용 테이블 lookup 추가(런타임 영향 미미). + +## Open review items (자율 결정; 검토 시 수정) + +1. **op 종류별 값의 calibration 출처.** composite ≈ 40 ns는 작업 추정치; + primitive 발행 비용은 placeholder. *권고:* 문서화된 가정(instruction-issue + + 큐 push)에서 값을 택하고, 실제 reference가 생기기 전까지 절대 latency를 + 잠정으로 취급; **비율**을 방어 가능하게 유지. +2. **scheduler plan-gen 비용 (D3).** *권고:* 초기에는 0 유지; scheduler-bound + 동작을 보이는 workload가 나오면 `overhead_ns`로 노출. +3. **DMA program 시간 (A3).** *권고:* 먼저 발행 op의 비용에 합치고; 필요 시 + PE_DMA setup으로 분리. +4. **테이블 위치.** *권고:* 흩어진 node attr이 아니라, 커맨드 종류로 키잉되고 + per-topology 오버라이드 가능한 작은 중앙 cost 모듈 — 값을 한 곳에서 검토 + 가능하도록. +5. **회귀 롤아웃.** *권고:* lazy-load(ADR-0062)와 본 ADR의 cost model을 한 + 번의 골든 재생성 패스로 함께 도입하여 두 번의 churn을 피함. +6. **legacy replay 경로와의 상호작용**(`pe_cpu.py:_execute_legacy`) — greenlet과 + replay 경로가 동일한 cost 테이블을 읽어 결과가 일치하도록 보장. 검증 요. + +## Test Requirements + +1. **composite는 한 번, primitive는 op마다 청구.** `N_tiles`에 걸쳐 하나의 + `tl.composite`를 내는 커널은 PE_CPU에서 composite 발행 비용을 한 번 청구하고; + 동등한 `N_tiles × ops` primitive 커널은 op당 비용을 `N_tiles × ops`번 청구. + PE_CPU busy time이 그에 따라 다름을 assert. +2. **DMA latency 불변 (Q2).** 고정된 전송에 대해 PE_DMA `drain_ns`는 ADR 이전 + 값과 동일 — 발행 비용은 DMA에 합쳐지지 않고 PE_CPU에 가산(중복 계상 없음). +3. **saturation 관측 가능.** per-op 발행 비용이 0이 아니면, many-tile primitive + sweep는 GEMM 엔진 idle(CPU-bound 발행)을 보이고 composite sweep는 바쁘게 + 유지 — ADR-0060 §1 레버. +4. **결정성:** 동일 입력 → 동일 op_log + latency (SPEC §0.1). +5. **경로 일치:** greenlet과 legacy-replay 경로가 동일 커널에 대해 동일한 발행 + 비용 회계를 산출. diff --git a/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md b/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md new file mode 100644 index 0000000..9b51b2d --- /dev/null +++ b/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md @@ -0,0 +1,179 @@ +# ADR-0064: Per-op-type CPU issue cost model (command construct + dispatch) + +## Status + +Proposed + +> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). The hybrid +> decision there (GEMMs via `tl.composite`, softmax merge in the kernel) +> wins by **offloading tiling to PE_SCHEDULER so the CPU issues coarse +> descriptors and runs ahead, keeping the engines saturated**. That win is +> currently **invisible in the simulator** because per-op CPU issue cost is +> zero. This ADR makes the issue cost real and **op-type-differentiated**, +> so composite-vs-primitive trade-offs (and CPU saturation) are measurable. + +## Context + +### What exists today + +- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its + command (`tl_context.py:123-125, 190, 227, 235, 612, …`), which emits + `PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles + > 0`. So the issue cost is a **single uniform knob** applied identically + to `tl.load`, `tl.dot`, a MATH op, and `tl.composite`. +- That knob is hardcoded to **0** in both live execution paths + (`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ issuing a + 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 + PE_CPU (`kernel_runner.py:131-132`) and on PE_SCHEDULER + (`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 + +The hybrid's whole argument is that **one** composite descriptor offloads +`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 +uniform), the model cannot show: +- that the primitive path may **fail to saturate** the engines when the + CPU cannot push fast enough (the core ADR-0060 §1 claim), nor +- that a composite **costs more to construct** than a single primitive but + far less than the many primitives it replaces. + +A single uniform `dispatch_cycles` cannot express this: `tl.load` ≪ +`tl.composite` in real construct cost. + +## Decision + +### D1. Per-op-type issue-cost table on PE_CPU + +Replace the single `dispatch_cycles` scalar with a **cost table keyed by +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): + +| Issued op | CPU issue cost (construct + push) | +|---|---| +| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 estimate) | +| `tl.load` / `tl.store` (DMA descriptor) | small (a few ns) | +| `tl.dot` / MATH / IPCQ send/recv | small (a few ns) | + +The point is the **ratio**: a composite construct ≫ a primitive issue, but +≪ the sum of the primitive issues it replaces. Exact ns are configurable. + +### D2. Execution latency stays on the engines (Q2 — no change) + +DMA/GEMM/MATH execution latency remains charged on PE_DMA / PE_GEMM / +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 + +`PE_SCHEDULER`'s tile-plan generation costs 0 sim time today. Keep that for +now (the dominant lever is CPU issue cost, D1); expose it later via the +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 + +The cost table is configurable (per-topology / node attrs, default to the +D1 estimates). Turning issue cost non-zero changes **every** bench's +latency, so golden latencies are **regenerated** once — a model-fidelity +improvement, same posture as ADR-0062's global lazy-load regression. + +## Alternatives + +### A1. Keep a single uniform `dispatch_cycles > 0` + +Rejected: op construct costs differ by an order of magnitude (`tl.load` vs +`tl.composite`); a uniform value either over-charges primitives or +under-charges composites, and in both cases misrepresents the +composite-vs-primitive trade-off the hybrid depends on. + +### A2. Charge the issue cost on PE_SCHEDULER instead of PE_CPU + +Rejected: the saturation question is *"can the CPU push descriptors fast +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. + +### A3. 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 +op's** D1 cost (a `tl.load` / composite that triggers DMA). Split it out to +a PE_DMA fixed setup only if calibration shows it matters. Review item. + +## Consequences + +### Positive +- The hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes + **measurable**, not just structural. +- Composite-vs-primitive and tiling-granularity trade-offs are visible in + latency, enabling the eval narrative. +- More faithful to hardware (issue bandwidth is a real bottleneck). + +### Negative +- **All** bench goldens shift → one-time regeneration (D4); CI golden + fixtures update. +- Per-op-type values need calibration; the ~40 ns composite figure is an + estimate, primitives are unspecified — results are only as good as the + numbers (do not over-claim absolute latencies until calibrated). +- Adds a cost-table lookup on the issue path (negligible runtime). + +## Open review items (decided autonomously; revise on review) + +1. **Calibration source for the per-op-type values.** Composite ≈ 40 ns is + a working estimate; primitive issue costs are placeholders. *Recommend:* + pick values from a documented assumption (instruction-issue + queue + push) and treat absolute latency as provisional until a real reference + exists; keep the **ratios** defensible. +2. **Scheduler plan-gen cost (D3).** *Recommend:* keep 0 initially; expose + via `overhead_ns` if a workload shows scheduler-bound behaviour. +3. **DMA program time (A3).** *Recommend:* fold into the issuing op's cost + first; split to PE_DMA setup only if needed. +4. **Where the table lives.** *Recommend:* a small central cost module + 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 + +1. **Composite charges once; primitives charge per-op.** A kernel issuing + one `tl.composite` over `N_tiles` charges one composite issue cost on + PE_CPU; the equivalent `N_tiles × ops` primitive kernel charges the + per-op cost `N_tiles × ops` times. Assert the PE_CPU busy time differs + accordingly. +2. **DMA latency unchanged (Q2).** For a fixed transfer, PE_DMA `drain_ns` + is identical to the pre-ADR value — the issue cost is additive on + PE_CPU, not folded into DMA (no double-count). +3. **Saturation is observable.** With non-zero per-op issue cost, a + many-tile primitive sweep shows GEMM-engine idle (CPU-bound issue) + whereas the composite sweep keeps it busy — the ADR-0060 §1 lever. +4. **Determinism:** identical inputs → identical op_log + latency + (SPEC §0.1). +5. **Path parity:** greenlet and legacy-replay paths produce identical + issue-cost accounting for the same kernel.