From ab324c155c92e0db7672b2205932dd8a3bbbf859 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Fri, 5 Jun 2026 00:52:38 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr):=20review=20fixes=20=E2=80=94=20unify?= =?UTF-8?q?=20KV=20placement=20to=20contiguous,=20ring/reduce=20cost=20mod?= =?UTF-8?q?el,=20opt3=20wording,=20rank=20def?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review: - Placement unified to *contiguous* C×P blocks throughout (S0/S2.1/S2.2/S0.5); round-robin demoted to the rejected alternative. Adds driver SP-enable threshold fallback (smaller C / C=P=1 for short/early decode). - Ring-vs-reduce cost model in S5.5: reduce ~ G*T_q*log(C*P) (O dominant; m,l scalars), ring ~ 2*S (total K+V bytes a CUBE injects over C-1 rotations; recv_async pipelines so latency ~ max-step) -> ring wins when T_q > 2S/(G*log(C*P)). - opt3 'removes the bubble' -> 'hides (subject to scheduler+engine balance)' everywhere; table 'hidden*' with footnote. - 'rank' defined (SP participant = a PE in a CUBE, KV shard in its HBM->TCM). - out-proj handoff contract (S0.5.4); S11 gate-type note (absolute latency deferred to ADR-0064; structural + relative-to-baseline gates now). - greenlet-as-contrast tightened to 'primitive-op (tl.dot) path' (S1, SB). KO mirror synced; DDD gets the SP-enable threshold fallback. Docs only. Co-Authored-By: Claude Opus 4.8 (1M context) --- ...R-0060-algo-gqa-fused-attention-ahbm-ko.md | 141 ++++++++++------ .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 155 ++++++++++++------ ...060-gqa-fused-attention-detailed-design.md | 7 + 3 files changed, 208 insertions(+), 95 deletions(-) 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 index 2db6016..c2e6edf 100644 --- 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 @@ -76,7 +76,7 @@ def gqa_decode_v1(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl) ``` ```python -# OPTION 3 — software pipelining (current primitives; bubble removed) ← SHIP THIS +# OPTION 3 — software pipelining (current primitives; bubble hidden, not provably removed) ← SHIP THIS def gqa_decode_v3(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl): cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0) kv = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d)) @@ -112,8 +112,12 @@ def gqa_decode_v2(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl) | | new HW cmd | GEMM bubble | CPU intra-tile wait | issues | now | |---|---|---|---|---|---| | **opt1 current** | no | **yes** | yes | `O(tiles·ops)` | ✓ | -| **opt3 sw-pipe** | no | no | yes (reordered) | `O(tiles·ops)` | ✓ | -| **opt2 ex_composite** | **#2 only** | no | **no** | `O(tiles)` | ✗ (build #2) | +| **opt3 sw-pipe** | no | hidden\* | yes (reordered) | `O(tiles·ops)` | ✓ | +| **opt2 ex_composite** | **#2 only** | hidden\* | **no** | `O(tiles)` | ✗ (build #2) | + +\* "hidden"이지 "removed" 아님: softmax 도는 동안 GEMM 엔진을 채우는 건 +스케줄러/엔진 균형(softmax-vs-GEMM 시간, K DMA 준비, 큐 깊이)에 달림 — 버블을 +*줄이거나 가린다*는 것이지 증명적으로 제거하는 게 아님. (`#1` = 기존 GEMM composite + `scale`; `#2` = softmax + P·V + stateful online-softmax accumulator만 신규. MATH 엔진엔 max/sum/exp가 이미 있음 — 신규는 @@ -239,7 +243,10 @@ sequence-parallelism(SP)으로 샤딩된다: 분할. - **Level-2 (intra-CUBE, PE 간):** 각 CUBE의 조각을 다시 `P` PE에 분할. -따라서 한 KV head는 **`C × P` rank**에 매핑되며 모두 한 SIP 내부다. **그 `G` query +따라서 한 KV head는 **`C × P` rank**에 매핑되며 모두 한 SIP 내부다. **rank**은 +메모리가 아니라 한 **SP 참가자** = 특정 **CUBE의 특정 PE**(`rank = +cube_local·P + pe_local`)이다; 그 rank의 KV shard는 그 PE가 소유한 HBM 영역 +(`K_base[rank]`)에 *상주*하고 계산 시 그 PE의 TCM으로 DMA된다. **그 `G` query head가 rank에 어떻게 매핑되는지는 case별로 다르다**(두 커널, TL;DR / §5): - **Decode** (§4): **Q는 replicate** — 모든 rank가 전 `G` query head를 가지며, @@ -266,10 +273,12 @@ CUBE NOC(Level-1) + PE IPCQ(Level-2)를 쓴다 — UCIe는 결코 아님. (출 Group(2 KV head)을 구동 — head는 CUBE 좌표로 공간적으로 선택, 커널 내 `for kv` 루프 **아님**. CLI가 4 SIP-device를 논리 병렬로 돌려 8 head 전부를 커버. -**Token → rank 배치** (CUBE Group 내, round-robin SP, ≤1 token 균형): KV token -`i`는 Level-1 rank `(start + i) mod C`, 이어 Level-2 rank -`((i // C) + start_pe) mod P`에 안착 — 계층적 round-robin이라 `C × P` rank 각각이 -dense·연속 로컬 slice를 소유. +**Token → rank 배치**는 CUBE Group 내 **contiguous 블록 샤딩**이다: rank `r`이 +위치 블록 `[r·B, (r+1)·B)`(`B = ⌈S/(C·P)⌉`)을 소유(Level-1은 CUBE, Level-2는 그 +안의 PE 선택). contiguous 블록이라야 prefill이 전체-미래 블록을 causal-skip할 수 +있고 prefill+decode가 한 KV cache를 공유할 수 있다(§2.1). (Round-robin — +interleaved `token i → rank (start+i) mod (C·P)` — 는 짧은-context decode 균형엔 +낫지만 causal-skip 불가·prefill 공유 불가라 **기각**된 대안, §2.1.) Reduction은 **계층적·intra-SIP**: 한 KV head는 group의 `C` CUBE에 **걸친다** (Level-1), CUBE NOC로 reduce; 이는 이전 "query head는 CUBE를 가로지르지 않는다" @@ -302,8 +311,8 @@ SIP에 머문다). 대해 **pure read**. - **P4.** V는 회전 안 함; V-cache는 raw projected V 보유. - **P5.** upstream RoPE 위치는 token의 **절대 global 위치** - `global_idx = local_slot·(C·P) + rank` (여기서 `rank = cube_local·P + pe_local`, - §2.1). Round-robin placement ≠ RoPE 위치. + `global_idx = rank·B + local_slot` (`B = ⌈S/(C·P)⌉`, contiguous 블록, + `rank = cube_local·P + pe_local`, §2.1). 배치 ≠ RoPE 위치. ### 0.5.2 Shape 기호 `T_q` = 이번 launch의 query 길이(decode: 1; prefill/chunk: chunk 폭). @@ -316,7 +325,7 @@ SIP에 머문다). | Input | Shape (per KV head) | Location | Notes | |---|---|---|---| | `Q` | `[G, T_q, d]` | per-rank HBM (loaded to TCM) | post-RoPE (P1). group의 `G` query 행 batched. | -| `K_cache` | `[S/(C·P), d]` | per-rank HBM, base `K_base[rank]`, contiguous | post-RoPE (P2). Read-only. 로컬 dense, global index는 `C·P`로 strided (§2.1). | +| `K_cache` | `[S/(C·P), d]` | per-rank HBM, base `K_base[rank]`, contiguous | post-RoPE (P2). Read-only. **연속** global 위치 블록 `[r·B,(r+1)·B)` (§2.1). | | `V_cache` | `[S/(C·P), d]` | per-rank HBM, base `V_base[rank]` | raw V (P4). Read-only. | | `global_token_counter` | scalar | launch arg | 커널이 로컬 len, slot↔global, causal bound 도출. | | `start_cube`, `start_pe` (= `f(request_id)`) | scalars | launch arg | Level-1/Level-2 회전. | @@ -339,6 +348,13 @@ V_block`을 ping-pong 버퍼에(post-RoPE), 그리고 causal step-skip용 | **Decode** (§4 reduce) | `O = [G, 1, d]` (all heads) | `O_base` at the **CUBE-Group root** | head-replicated → 2-level reduce 후 전 `G` head가 한 rank에. | | **Prefill** (§5.5 ring) | `O_i = [1, T_q, d]` (one head each) | per-CUBE `O_base[i]`, **distributed** | head-parallel → CUBE `i`가 head `i`의 행을 제자리에 기록; reduce 없음. | +**out-projection(downstream 별도 커널)으로의 handoff 계약.** 본 ADR은 위 두 +레이아웃을 *보장*한다; out-proj 설계는 범위 밖이나 handoff는 맞아야 한다: +**decode**는 전 `G` head를 Group root에 co-located로 전달(out-proj이 로컬로 읽거나 +자기 SP로 scatter); **prefill**은 head `i`를 CUBE `i`에 남김(out-proj이 자체 +head-parallel이거나 먼저 `G` head를 **gather**해야 함). 이를 `out_proj` launch +계약에 커널별로 고정(§B decode-split 항목 2). + **Decode 중간값** (non-root rank, IPCQ 상 — 커널-가시 아님): `(m_i, ℓ_i, O_i)`이 2-level tree를 따라 부모(Level-2 PE 부모, 이어 Level-1 CUBE 부모)에 `tl.send`로 push(§4). `O_i`가 무거운 부분. **Prefill**은 그런 partial이 @@ -401,7 +417,9 @@ per tile j: 허용하는 바 그대로다(`kernel_runner.py`, ADR-0020 D3). > **이것이 무엇을 대체하나.** 이전 iteration은 "composite는 latency 이점 없음"이라는 -> 논거로 순수 greenlet primitive 경로(전부 `tl.dot`, composite 없음)를 제안했다. +> 논거로 순수 primitive-op 경로(전부 `tl.dot`, composite 없음 — 둘 다 greenlet에서 +> 발행되므로 대비 축은 "greenlet vs composite"가 아니라 *scheduler-managed composite* +> vs *kernel-issued ops*다)를 제안했다. > 그것은 **오직** 시뮬레이터가 현재 per-op CPU 발행 비용을 **0**으로 청구하기 > 때문에만(`dispatch_cycles=0`, `pe_cpu.py`) 성립한다 — descriptor offload가 숨기려 > 존재하는 바로 그 CPU issue-rate / DMA-program 비용을 모델이 지워버린 것이다. @@ -433,24 +451,28 @@ intra-CUBE PE, §0). kernbench에서는 sequence 차원을 `cube` 축(`C`에 대 `DPPolicy(pe="row_wise")`가 (cube, pe) 두 축 `row_wise`가 된다. - per-rank KV 버퍼는 `⌈max_context / (C·P)⌉ × d × dtype`로 K와 V 각각 sizing. -- rank 안에서 할당된 token은 **연속 append**(slot 0,1,2,…); per-rank 버퍼는 dense - ⇒ global index가 `C·P`로 strided여도 DMA는 연속 유지. -- Slot → global (계층적 round-robin): `rank = cube_local·P + pe_local` +- 각 rank의 **연속 위치 블록**을 dense하게 저장(`B = ⌈max_context/(C·P)⌉` slot) + ⇒ DMA는 연속 유지. +- Slot → global (**contiguous 블록**): `rank = cube_local·P + pe_local` (여기서 `cube_local = (cube_id − start_cube) mod C`, `pe_local = (pe_id − start_pe) mod P`); 이후 - `global_idx = local_slot·(C·P) + rank`. + `global_idx = rank·B + local_slot`. -> **prefill/decode 공유 레이아웃 — round-robin이 아니라 *contiguous* 블록.** -> prefill이 KV cache를 *쓰고*(`qkv_rope`, upstream) decode가 같은 cache를 *읽고 -> 연장*하므로, 두 커널은 하나의 물리 레이아웃을 공유한다 — prefill→decode 경계에 -> reshard 없음. prefill의 **causal skip은 연속 위치 블록을 요구**한다(늦은 연속 -> 블록은 통째로 미래라 skip 가능; round-robin 블록은 전 위치에 걸쳐 결코 skip -> 불가). 그래서 공유 레이아웃은 각 KV head를 **contiguous `C × P` 위치 블록**으로 -> 샤딩한다(rank `r`이 `[r·B, (r+1)·B)` 소유, `B = ⌈max_context/(C·P)⌉`). Decode는 -> 자기 블록을 읽고 reduce; prefill은 `C` CUBE-레벨 블록을 ring. *Caveat:* -> contiguous는 **짧은** context에서 rank를 덜 씀(frontier rank만 데이터 보유) — -> long-context 타깃에선 수용 가능; 짧은-context 균형은 별도 연구(§B). (위 round-robin -> 공식은 decode-only 균형용 대안; contiguous 블록 형태가 두 커널을 다 서비스.) +> **배치가 *contiguous* 블록인 이유 (rationale).** prefill이 KV cache를 +> *쓰고*(`qkv_rope`, upstream) decode가 같은 cache를 *읽고 연장*하므로, 두 커널은 +> 하나의 물리 레이아웃을 공유한다 — prefill→decode 경계에 reshard 없음. prefill의 +> **causal skip은 연속 위치 블록을 요구**한다(늦은 연속 블록은 통째로 미래라 skip +> 가능; round-robin 블록은 전 위치에 걸쳐 결코 skip 불가). 그래서 공유 레이아웃은 +> 각 KV head를 **contiguous `C × P` 위치 블록**으로 샤딩한다(rank `r`이 +> `[r·B, (r+1)·B)` 소유, `B = ⌈max_context/(C·P)⌉`). Decode는 자기 블록을 읽고 +> reduce; prefill은 `C` CUBE-레벨 블록을 ring. *Caveat:* contiguous는 **짧은** +> context에서 rank를 덜 씀(frontier rank만 데이터 보유) — long-context 타깃에선 수용 +> 가능; 짧은-context 균형은 별도 연구(§B). **Fallback:** contiguous 샤딩이 rank를 +> 덜 쓰는 짧은/초기 decode에선 드라이버가 더 작은 `C`(또는 `C=P=1`)를 고를 수 있다; +> SP는 KV-sweep 시간이 reduction·배치 불균형을 지배하는 context-길이 임계 이후에만 +> 켠다. **Rejected: round-robin**(interleaved `token i → rank (start+i) mod (C·P)`) +> — 짧은-context decode 균형만 더 낫지만 round-robin 블록은 전 위치에 걸쳐 causal-skip +> 불가·prefill 공유 불가. ### 2.2 드라이버 launch당 의무 (최소) 드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를 @@ -465,17 +487,20 @@ intra-CUBE PE, §0). kernbench에서는 sequence 차원을 `cube` 축(`C`에 대 | `cube_id`, `pe_id` (`tl.program_id` 1/0), `C`, `P` | CUBE-Group geometry | | `q_block_meta` | prefill/SP query block 시작 + 길이 | -`R = C·P`, `cube_local = (cube_id − start_cube) mod C`, +`R = C·P`, `B = ⌈max_context/R⌉`, `cube_local = (cube_id − start_cube) mod C`, `pe_local = (pe_id − start_pe) mod P`, 이 rank의 인덱스 -`rank = cube_local·P + pe_local`라 하자. 커널 도출(평범한 산술): -- **이번 step에 내가 쓸 차례:** `(start_rank + counter) mod R == rank` -- **내 valid 길이:** `base = counter // R; rem = counter % R; - my_len = base + (1 if rank < rem else 0)` +`rank = cube_local·P + pe_local`라 하자. 커널 도출(평범한 산술, `counter` = 채워진 +token 수): + +- **내 valid 길이**(연속 블록 `[rank·B, (rank+1)·B)`): + `my_len = clamp(counter − rank·B, 0, B)`. +- **이번 step에 내가 쓸 차례:** 새 decode token(위치 `counter`)은 rank `counter // B` + (그 연속 블록)에 안착. - **read 범위:** tile `0 .. ⌈my_len / TILE⌉`. - **causal bound / per-tile skip:** query block과 각 tile의 global 위치에서. -드라이버 = base + counter + rotation. 두 SP 축에 대한 단일 공식 -`(request_id + token_idx) mod (C·P)`이 placement policy 전부다. +드라이버 = base + counter + rotation. placement policy는 단일 contiguous-block 맵 +`token t → rank ⌊t/B⌋`(`start_*`로 회전). --- @@ -669,9 +694,9 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, ### 5.3 DECODE, with SP / KV-parallel (`C × P` rank) -- 한 request의 KV가 CUBE Group의 `C·P` rank에 계층적-round-robin(Level-1은 `C` - CUBE, Level-2는 `P` PE); 각 rank가 ≈`my_len` token 소유하고 자기 로컬 tile에 - 대해 `G=8` query 행을 GQA-batch. +- 한 request의 KV가 CUBE Group의 `C·P` rank에 **contiguous-block 샤딩**(Level-1은 + `C` CUBE, Level-2는 `P` PE); 각 rank가 자기 `B`-token 블록을 소유하고 자기 로컬 + tile에 대해 `G=8` query 행을 GQA-batch. - `T_q=1` ⇒ 짧은 per-rank sweep ⇒ **reduction 지배** ⇒ §4 2-level reduce (`⌈log₂ P⌉` + `C`에 대한 center-mesh)가 구조적으로 중요한 부분. Reduction latency는 batch될 때 다른 동시 decode token으로, 긴 single-stream context의 긴 per-rank @@ -694,6 +719,22 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, 이동; 대신 **head를 shard**하고 (역시 큰) **KV를 옮긴다** — 어차피 각 head가 전체 KV를 필요로 하므로. +> **ring이 reduce를 이기는 조건 (cost model — "O가 크다"만이 아님).** **총 +> cross-rank 바이트**를 비교(head dim `d` 소거): +> +> - **reduce** ≈ `G·T_q · ⌈log₂(C·P)⌉` — payload는 `(m,ℓ,O)`지만 **`O`가 지배** +> (`m,ℓ`은 `[G·T_q]` scalar, `O`는 `[G·T_q, d]`) → `O`-바이트 × tree hop; +> head-replicated라 `G` 인자 존재. +> - **ring** ≈ `2·S` — 한 CUBE가 `C−1` 회전 동안 **주입하는 총 K+V 바이트** +> (`2·(S/C)·(C−1) ≈ 2·S`); `recv_async`가 회전을 **파이프라인**하므로 *latency*는 +> 가장 느린 한 step ≈, *traffic*은 합. **`G` 인자·log 없음** — 한 KV slice가 회전을 +> 통해 전 `C` head에 *공유*되므로. +> +> 그래서 **총 바이트 기준 ring < reduce 조건: `T_q > 2S / (G·⌈log₂(C·P)⌉)`** — +> decode(`T_q=1`) → **reduce**, prefill(`T_q≈S`, 우변 `≈ 2/(G·log) ≪ 1`) → **ring**. +> KV는 크지만 reduce의 `O`가 `G·log` 인자를 더 지므로 prefill에선 ring이 이긴다. +> (임계 sweep: §9.) + - **Head-parallel 배치:** CUBE Group 내에서 **CUBE `i`가 query head `i` 하나만 소유**(`G` query head → `C=G` CUBE, CUBE당 Q head 1개)하고 KV slice `i`. 각 CUBE가 **자기 Q head 하나의** full attention 계산. 각 CUBE가 *서로 다른* head를 생산하므로 @@ -718,11 +759,11 @@ head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추 세 decode 변형은 **TL;DR에 풀 코드**로 있다(Kernel 1): **opt1** 현재 `CompositeCmd`(CPU가 `Sj`를 auto-wait하는 동안 GEMM 엔진 버블), **opt3** software pipelining(다음 타일 `Q·Kᵀ`를 이번 타일 softmax 전에 발행, `Sj`는 persistent 이중 -버퍼 — 버블 제거, 새 커맨드 종류 없음), **opt2** `ex_composite`(분할: `#1` Q·Kᵀ = +버퍼 — 스케줄러+엔진 균형에 따라 버블을 *가림*, 새 커맨드 종류 없음), **opt2** `ex_composite`(분할: `#1` Q·Kᵀ = K를 먼저 읽는 기존 composite, `#2` softmax+P·V+accumulator = V를 나중에 읽는 유일한 신규 flash-epilogue 기계장치). -**권고:** 지금은 **opt3** 출시(새 기계장치 0, 버블 제거); cost model(ADR-0064)이 +**권고:** 지금은 **opt3** 출시(새 기계장치 0, 스케줄러+엔진 균형에 따라 버블 가림); cost model(ADR-0064)이 적은-CPU-발행 이득을 측정 가능하게 만들면 **opt2** 재검토(§8 항목 4). MATH 엔진엔 max/sum/exp가 이미 있음 — opt2의 유일한 진짜 신규는 composite의 **stateful `(m,ℓ,O)` accumulator**이지 ops가 아님. @@ -781,7 +822,7 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다. data-driven·레벨 파이프라인, `tl.send`/`tl.recv` 상의 순수 커널 제어 흐름. - causal tile skip + additive boundary mask(§3/§5.4) — 커널 `if` + `+`로 더한 mask 텐서. -- 2-level round-robin KV placement / valid-length 산술(§0, §2) — launch-arg 산술 + +- 2-level contiguous-block KV placement / valid-length 산술(§0, §2) — launch-arg 산술 + `cube`와 `pe` 양쪽 `row_wise`인 `DPPolicy`. **efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):** @@ -873,6 +914,13 @@ ADR-0023/0025 (IPCQ), ADR-0046 (`tl` 계약), ADR-0054 (eval bench). exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`은 `f16`으로 모델된다. 따라서 검증은 **구조/타이밍 우선**, 수치 parity는 제한된 2차 체크. +**Gate 종류.** *절대* latency target은 연기 — ADR-0064 cost model이 있어야만 +의미(오늘 `dispatch_cycles=0`). 아래 gate는 (a) **구조적**(op_log 형태)과 (b) 지금 +*테스트 가능한* **baseline-상대** 비율이며 반드시 통과해야 한다: reduce 라운드 +`⌈log₂P⌉+center-mesh` **vs baseline `C·P−1`**; KV 바이트 `H_kv·S·d` **vs `H_q·S·d` +(G×=8× 적음)**; causal GEMM 수 = 하삼각 **vs full grid**; prefill ring traffic = +`C` KV 회전 **vs** `[G·T_q,d]` reduce. + 1. **data 모드 실행(`enable_data=True`):** GQA 커널(`h_q = G·h_kv`)이 baseline head-packing이 부딪히는 byte-conservation 에러 없이 — 네 case 모두 — 완료. (이는 §5.2 재구조화가 필요하지, 신규 primitive가 *아님*.) @@ -908,7 +956,7 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` ## B. 하이브리드 전환에서 나온 Open 설계 항목 (추후 검토) -이들은 결정이 순수 greenlet primitive 경로에서 **composite hybrid + lazy +이들은 결정이 순수 primitive-op(`tl.dot`) 경로에서 **composite hybrid + lazy `tl.load`**(이번 개정)로 옮겨갈 때 생겼다. 설계를 막는 것은 없으며; 각각 구현 중 검증 패스가 필요하다. (묻지 않고) 작업 합의에 따라 여기 기록한다 — 권고는 내가 예측한 기본값이며; 검토 시 수정. @@ -985,10 +1033,11 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` 이 커널 범위에선 발견 안 됨, 단 `configure_sfr_*` 이웃 wiring이 Level-1용 CUBE-NOC `N/S/E/W`를 PE IPCQ뿐 아니라 노출해야 함. -5. **`valid_len_2level` / 계층적 round-robin 정확성.** 두-축 `(cube_local·P + - pe_local)` 배치(§2.1)는 시퀀스를 gap/overlap 없이 타일링하고 각 rank의 로컬 버퍼를 - dense하게 유지해야 함. **권고:** 커널 전에 index 산술을 unit-test(모든 global - token이 정확히 한 rank에 매핑; per-rank slice 연속). +5. **`valid_len_2level` / contiguous-block 정확성.** 두-축 + `rank = cube_local·P + pe_local`, 블록 `[rank·B, (rank+1)·B)` 배치(§2.1)는 + 시퀀스를 gap/overlap 없이 타일링하고 각 rank의 로컬 버퍼를 dense하게 유지해야 함. + **권고:** 커널 전에 index 산술을 unit-test(모든 global token이 정확히 한 rank에 + 매핑; per-rank 블록 연속). ### decode-reduce / prefill-ring 분리에서 나온 항목 (이번 개정) @@ -1027,5 +1076,5 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` 7. **출시할 decode CPU-pipelining 변형 (§5.6).** 세 decode 변형 존재(opt1 current / opt3 software-pipelining / opt2 ex_composite). **권고:** **opt3**(software pipelining: 다음 Q·Kᵀ를 이번 타일 softmax 전에 발행, `Sj`는 persistent 이중 버퍼) - 구현 — 새 커맨드 종류 없이 GEMM 엔진 버블 제거. **opt2**(2-composite `ex_composite`, + 구현 — 새 커맨드 종류 없이 GEMM 엔진 버블 가림(스케줄러+엔진 균형에 따라). **opt2**(2-composite `ex_composite`, `#2`만 신규)는 ADR-0064 cost model이 적은-발행 이득을 측정 가능하게 만들 때까지 연기. 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 675d87d..f8f2d5c 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 @@ -82,7 +82,7 @@ def gqa_decode_v1(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl) ``` ```python -# OPTION 3 — software pipelining (current primitives; bubble removed) ← SHIP THIS +# OPTION 3 — software pipelining (current primitives; bubble hidden, not provably removed) ← SHIP THIS def gqa_decode_v3(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl): cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0) kv = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d)) @@ -118,8 +118,12 @@ def gqa_decode_v2(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl) | | new HW cmd | GEMM bubble | CPU intra-tile wait | issues | now | |---|---|---|---|---|---| | **opt1 current** | no | **yes** | yes | `O(tiles·ops)` | ✓ | -| **opt3 sw-pipe** | no | no | yes (reordered) | `O(tiles·ops)` | ✓ | -| **opt2 ex_composite** | **#2 only** | no | **no** | `O(tiles)` | ✗ (build #2) | +| **opt3 sw-pipe** | no | hidden\* | yes (reordered) | `O(tiles·ops)` | ✓ | +| **opt2 ex_composite** | **#2 only** | hidden\* | **no** | `O(tiles)` | ✗ (build #2) | + +\* "hidden" not "removed": filling the GEMM engine during softmax depends on +the scheduler/engine balance (softmax-vs-GEMM time, K DMA readiness, queue +depth) — it *reduces/hides* the bubble, not provably eliminates it. (`#1` = the *existing* GEMM composite + `scale`; only `#2` = softmax + P·V + the stateful online-softmax accumulator is new. MATH engine already has @@ -257,9 +261,12 @@ two levels** of sequence-parallelism (SP): - **Level-2 (intra-CUBE, across PEs):** each CUBE's slice is split again across its `P` PEs. -So one KV head maps to **`C × P` ranks**, all within one SIP. **How the -`G` query heads map onto those ranks differs by case** (the two kernels, -TL;DR / §5): +So one KV head maps to **`C × P` ranks**, all within one SIP. A **rank** is +not memory — it is one **SP participant** = a specific **PE in a specific +CUBE** (`rank = cube_local·P + pe_local`); its KV shard *resides in* that +PE's owned HBM region (`K_base[rank]`) and is DMA'd to the PE's TCM for +compute. **How the `G` query heads map onto those ranks differs by case** +(the two kernels, TL;DR / §5): - **Decode** (§4): **Q is replicated** — every rank holds all `G` query heads, **M-folded** (stacked into the matmul M / row dimension: `Q` for @@ -290,10 +297,14 @@ drives its 2 CUBE Groups (2 KV heads) — the head is selected spatially by CUBE coordinate, **not** an in-kernel `for kv` loop. The CLI runs the 4 SIP-devices logically in parallel to cover all 8 heads. -**Token → rank placement** within a CUBE Group (round-robin SP, balanced -to ≤1 token): KV token `i` lands on Level-1 rank `(start + i) mod C`, then -Level-2 rank `((i // C) + start_pe) mod P` — a hierarchical round-robin so -each of the `C × P` ranks owns a dense, contiguous local slice. +**Token → rank placement** within a CUBE Group is **contiguous block +sharding**: rank `r` owns the position block `[r·B, (r+1)·B)`, +`B = ⌈S/(C·P)⌉` (Level-1 picks the CUBE, Level-2 the PE within it). A +contiguous block is **required** so prefill can causal-skip whole future +blocks **and** so prefill + decode share one KV cache (§2.1). (Round-robin — +interleaved `token i → rank (start+i) mod (C·P)` — gives better +short-context decode balance but cannot causal-skip and cannot be shared +with prefill; it is the **rejected** alternative, §2.1.) Reduction is **hierarchical and intra-SIP**: a KV head **does** span the `C` CUBEs of its group (Level-1), reduced over the CUBE NOC; this @@ -326,8 +337,8 @@ existed only because the `H_kv=1` baseline never needed it. Crossing launches.** ⇒ this kernel is **pure read** on the KV cache. - **P4.** V is not rotated; V-cache holds raw projected V. - **P5.** RoPE position upstream is the token's **absolute global - position** `global_idx = local_slot·(C·P) + rank` where `rank = - cube_local·P + pe_local` (§2.1). Round-robin placement ≠ RoPE position. + position** `global_idx = rank·B + local_slot` (`B = ⌈S/(C·P)⌉`, contiguous + block, `rank = cube_local·P + pe_local`, §2.1). Placement ≠ RoPE position. ### 0.5.2 Shape symbols `T_q` = query length this launch (decode: 1; prefill/chunk: chunk width). @@ -341,7 +352,7 @@ matters. | Input | Shape (per KV head) | Location | Notes | |---|---|---|---| | `Q` | `[G, T_q, d]` | per-rank HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. | -| `K_cache` | `[S/(C·P), d]` | per-rank HBM, base `K_base[rank]`, contiguous | post-RoPE (P2). Read-only. Dense locally, strided by `C·P` in global index (§2.1). | +| `K_cache` | `[S/(C·P), d]` | per-rank HBM, base `K_base[rank]`, contiguous | post-RoPE (P2). Read-only. A **contiguous** global position block `[r·B,(r+1)·B)` (§2.1). | | `V_cache` | `[S/(C·P), d]` | per-rank HBM, base `V_base[rank]` | raw V (P4). Read-only. | | `global_token_counter` | scalar | launch arg | kernel derives local len, slot↔global, causal bounds. | | `start_cube`, `start_pe` (= `f(request_id)`) | scalars | launch arg | Level-1/Level-2 rotation. | @@ -364,6 +375,14 @@ out-projection must consume them accordingly: | **Decode** (§4 reduce) | `O = [G, 1, d]` (all heads) | `O_base` at the **CUBE-Group root** | head-replicated → all `G` heads end on one rank after the 2-level reduce. | | **Prefill** (§5.5 ring) | `O_i = [1, T_q, d]` (one head each) | per-CUBE `O_base[i]`, **distributed** | head-parallel → CUBE `i` writes head `i`'s rows in place; no reduce. | +**Handoff contract to out-projection (downstream, separate kernel).** This +ADR *guarantees* the two layouts above; designing out-proj is out of scope, +but the handoff must match: **decode** delivers all `G` heads co-located at +the Group root (out-proj reads them locally, or scatters per its own SP); +**prefill** leaves head `i` on CUBE `i` (out-proj is itself head-parallel, +or must **gather** the `G` heads first). Pin this per-kernel in the +`out_proj` launch contract (§B decode-split item 2). + **Decode intermediate** (non-root rank, on IPCQ — not kernel-visible): `(m_i, ℓ_i, O_i)` pushed up the 2-level tree to the parent (Level-2 PE parent, then Level-1 CUBE parent) via `tl.send` (§4). `O_i` is the heavy @@ -433,8 +452,10 @@ 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). -> **What this supersedes.** An earlier iteration proposed a pure greenlet -> primitive path (all `tl.dot`, no composite) on the argument that +> **What this supersedes.** An earlier iteration proposed a pure +> primitive-op path (all `tl.dot`, no composite — both are issued from the +> greenlet, so the contrast is *scheduler-managed composite* vs +> *kernel-issued ops*, not "greenlet vs 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 @@ -472,16 +493,15 @@ One KV head's sequence is sharded across `C × P` ranks (Level-1 inter-CUBE - Per-rank KV buffers sized to `⌈max_context / (C·P)⌉ × d × dtype`, K and V separately. -- Within a rank, assigned tokens are **appended contiguously** (slot - 0,1,2,…); the per-rank buffer is dense ⇒ DMA stays contiguous, even - though global indices are strided by `C·P`. -- Slot → global (hierarchical round-robin): `rank = cube_local·P + pe_local` +- Each rank's **contiguous position block** is stored dense (`B = + ⌈max_context/(C·P)⌉` slots) ⇒ DMA stays contiguous. +- Slot → global (**contiguous block**): `rank = cube_local·P + pe_local` where `cube_local = (cube_id − start_cube) mod C`, `pe_local = (pe_id − start_pe) mod P`; then - `global_idx = local_slot·(C·P) + rank`. + `global_idx = rank·B + local_slot`. -> **Shared prefill/decode layout — use *contiguous* blocks, not -> round-robin.** Because prefill *writes* the KV cache (`qkv_rope`, +> **Why the placement is *contiguous* blocks (rationale).** Because +> prefill *writes* the KV cache (`qkv_rope`, > upstream) and decode *reads + extends* the same cache, both kernels share > one physical layout — no reshard at the prefill→decode boundary. Prefill's > **causal skip needs contiguous position blocks** (a late contiguous block @@ -492,8 +512,14 @@ One KV head's sequence is sharded across `C × P` ranks (Level-1 inter-CUBE > the `C` CUBE-level blocks. *Caveat:* contiguous under-uses ranks for > **short** context (only frontier ranks hold data) — acceptable given the > long-context target; short-context balance is a separate study (§B). -> (The round-robin formula above is the alternative for decode-only balance; -> the contiguous block form is the one that serves both kernels.) +> **Fallback:** for short/early decode where contiguous-block sharding +> under-utilizes ranks, the driver may select a smaller `C` (or `C=P=1`); SP +> is enabled only past the context-length threshold where KV-sweep time +> dominates reduction and placement imbalance. +> **Rejected: round-robin** (interleaved `token i → rank (start+i) mod (C·P)`) +> — better short-context decode balance *alone*, but a round-robin block +> spans all positions so it cannot causal-skip and cannot be shared with +> prefill. ### 2.2 Driver per-launch duties (minimal) The driver supplies, per launch, the bases + counter + rotation; the @@ -508,19 +534,21 @@ kernel derives the rest: | `cube_id`, `pe_id` (`tl.program_id` 1/0), `C`, `P` | CUBE-Group geometry | | `q_block_meta` | prefill/SP query block start + length | -Let `R = C·P`, `cube_local = (cube_id − start_cube) mod C`, -`pe_local = (pe_id − start_pe) mod P`, and this rank's index -`rank = cube_local·P + pe_local`. Kernel-derived (plain arithmetic): -- **my turn to write this step:** `(start_rank + counter) mod R == rank` -- **my valid length:** `base = counter // R; rem = counter % R; - my_len = base + (1 if rank < rem else 0)` +Let `R = C·P`, `B = ⌈max_context/R⌉`, `cube_local = (cube_id − start_cube) +mod C`, `pe_local = (pe_id − start_pe) mod P`, and this rank's index +`rank = cube_local·P + pe_local`. Kernel-derived (plain arithmetic, with +`counter` = # filled tokens): + +- **my valid length** (contiguous block `[rank·B, (rank+1)·B)`): + `my_len = clamp(counter − rank·B, 0, B)`. +- **my turn to write this step:** the new decode token (position `counter`) + lands on rank `counter // B` (its contiguous block). - **read range:** tiles `0 .. ⌈my_len / TILE⌉`. - **causal bounds / per-tile skip:** from global positions of the query block vs each tile. -Driver = bases + counter + rotation. The single formula -`(request_id + token_idx) mod (C·P)` over the two SP axes is the entire -placement policy. +Driver = bases + counter + rotation. The placement policy is the single +contiguous-block map `token t → rank ⌊t/B⌋` (rotated by `start_*`). --- @@ -728,10 +756,10 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, ### 5.3 DECODE, with SP / KV-parallel (`C × P` ranks) -- One request's KV is hierarchical-round-robin across the `C·P` ranks of +- One request's KV is **contiguous-block sharded** across the `C·P` ranks of the CUBE Group (Level-1 over `C` CUBEs, Level-2 over `P` PEs); each rank - owns ≈`my_len` tokens and GQA-batches its `G=8` query rows over its local - tiles. + owns its `B`-token block and GQA-batches its `G=8` query rows over its + local tiles. - `T_q=1` ⇒ short per-rank sweep ⇒ **reduction dominates** ⇒ the §4 2-level reduce (`⌈log₂ P⌉` + center-mesh over `C`) is the structurally important part. Reduction latency is hidden by other concurrent decode tokens when @@ -757,6 +785,23 @@ head is **big**, so reducing it across ranks would move `[T_q,d]` per rank; instead we **shard the heads** and **move the (also big) KV**, which each head needs in full anyway. +> **When ring beats reduce (cost model — the choice is *not* just "O is +> big").** Compare **total cross-rank bytes**, head dim `d` factored out: +> +> - **reduce** ≈ `G·T_q · ⌈log₂(C·P)⌉` — the payload is `(m,ℓ,O)` but **`O` +> dominates** (`m,ℓ` are `[G·T_q]` scalars vs `O` = `[G·T_q, d]`), so this +> is `O`-bytes × tree hops; head-replicated, so the `G` factor is present. +> - **ring** ≈ `2·S` — the **total K+V bytes one CUBE injects** over the +> `C−1` rotations (`2·(S/C)·(C−1) ≈ 2·S`); `recv_async` **pipelines** the +> rotations so *latency* ≈ the slowest single step, but *traffic* is the +> sum. **No `G` factor, no log** — one KV slice is *shared* across all `C` +> heads by the rotation. +> +> So **ring < reduce (by total bytes) when `T_q > 2S / (G·⌈log₂(C·P)⌉)`** — +> decode (`T_q=1`) → **reduce**, prefill (`T_q≈S`, RHS `≈ 2/(G·log) ≪ 1`) → +> **ring**. KV is big, but the reduce's `O` carries the extra `G·log` factor, +> so ring still wins for prefill. (Threshold sweep: §9.) + - **Head-parallel placement:** within a CUBE Group, **CUBE `i` owns exactly one query head `i`** (the `G` query heads → the `C=G` CUBEs, one Q head per CUBE) and KV slice `i`. Each CUBE computes **its one Q head's** full @@ -786,12 +831,14 @@ The three decode variants are shown **in full in the TL;DR** (Kernel 1): **opt1** current `CompositeCmd` (has a GEMM-engine bubble while the CPU auto-waits on `Sj`), **opt3** software pipelining (issue the next tile's `Q·Kᵀ` before this tile's softmax, `Sj` in a persistent double buffer — -removes the bubble, no new command kind), **opt2** `ex_composite` (split +*hides* the bubble subject to scheduler+engine balance, no new command +kind), **opt2** `ex_composite` (split into `#1` Q·Kᵀ = existing composite reading K first, `#2` softmax+P·V+ accumulator = the only new flash-epilogue machinery, reading V later). -**Recommend:** ship **opt3** now (no new machinery, removes the bubble); -revisit **opt2** once the cost model (ADR-0064) makes the fewer-CPU-issues +**Recommend:** ship **opt3** now (no new machinery, hides the bubble subject +to scheduler+engine balance); revisit **opt2** once the cost model (ADR-0064) +makes the fewer-CPU-issues win measurable (§8 item 4). The MATH engine already has max/sum/exp — the only genuinely new part of opt2 is the composite's **stateful `(m,ℓ,O)` accumulator**, not the ops. @@ -859,7 +906,7 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head, pure kernel control flow over `tl.send`/`tl.recv`. - Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` + a mask tensor added with `+`. -- 2-level round-robin KV placement / valid-length arithmetic (§0, §2) — +- 2-level contiguous-block KV placement / valid-length arithmetic (§0, §2) — launch-arg arithmetic + `DPPolicy` `row_wise` over both `cube` and `pe`. **New primitives for efficiency / scale (each has a supporting ADR):** @@ -964,6 +1011,14 @@ reshape-not-transpose, and `bf16` is modelled as `f16`. Verification is therefore **structural/timing-first**, with numeric parity as a bounded secondary check. +**Gate type.** *Absolute* latency targets are deferred — they are only +meaningful once ADR-0064's cost model lands (`dispatch_cycles=0` today). The +gates below are (a) **structural** (op_log shape) and (b) **relative-to- +baseline** ratios that *are* testable now and must pass: reduce rounds +`⌈log₂P⌉+center-mesh` **vs baseline `C·P−1`**; KV bytes `H_kv·S·d` **vs +`H_q·S·d` (G×=8× less)**; causal GEMM count = lower-triangular **vs full +grid**; prefill ring traffic = `C` KV rotations **vs** a `[G·T_q,d]` reduce. + 1. **Runs in data mode (`enable_data=True`):** the GQA kernel (`h_q = G·h_kv`) completes without the byte-conservation error that the baseline head-packing hits — for all four cases. (This needs the §5.2 @@ -1005,8 +1060,8 @@ secondary check. ## 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 +These arose when the decision moved from a pure primitive-op (`tl.dot`) 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. @@ -1095,11 +1150,12 @@ predicted default; revise on review. none found in this kernel's scope, but the `configure_sfr_*` neighbour wiring must expose CUBE-NOC `N/S/E/W` for Level-1, not just PE IPCQ. -5. **`valid_len_2level` / hierarchical round-robin correctness.** The - two-axis `(cube_local·P + pe_local)` placement (§2.1) must tile the - sequence with no gaps/overlaps and keep each rank's local buffer dense. - **Recommend:** unit-test the index math (every global token maps to - exactly one rank; per-rank slices are contiguous) before the kernel. +5. **`valid_len_2level` / contiguous-block correctness.** The two-axis + `rank = cube_local·P + pe_local`, block `[rank·B, (rank+1)·B)` placement + (§2.1) must tile the sequence with no gaps/overlaps and keep each rank's + local buffer dense. **Recommend:** unit-test the index math (every global + token maps to exactly one rank; per-rank blocks are contiguous) before + the kernel. ### Items from the decode-reduce / prefill-ring split (this revision) @@ -1147,7 +1203,8 @@ predicted default; revise on review. 7. **Decode CPU-pipelining variant to ship (§5.6).** Three decode variants exist (opt1 current / opt3 software-pipelining / opt2 ex_composite). **Recommend:** implement **opt3** (software pipelining: issue next Q·Kᵀ - before this tile's softmax, `Sj` in a persistent double buffer) — removes - the GEMM-engine bubble with no new command kind. Defer **opt2** (the + before this tile's softmax, `Sj` in a persistent double buffer) — hides + the GEMM-engine bubble (subject to scheduler+engine balance) with no new + command kind. Defer **opt2** (the two-composite `ex_composite`, only `#2` is new) until ADR-0064's cost model makes its fewer-issues win measurable. diff --git a/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md b/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md index 83a8454..0ffe3fd 100644 --- a/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md +++ b/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md @@ -130,6 +130,13 @@ Placement via `DPPolicy(..., cube_start=...)`: (ADR-0060 §2.1). `cube_start` offsets the group's CUBE sub-mesh within the 4×4 SIP mesh. +**Driver SP-enable threshold (fallback).** Contiguous sharding under-utilizes +ranks for short/early decode (only frontier ranks hold data). So the driver +**enables SP only past a context-length threshold** where KV-sweep time +dominates reduction + placement imbalance; below it, fall back to a smaller +`C` (or `C=P=1`, single rank, no reduce). The threshold is a sweep item +(P7 / §9), not a fixed constant. + `--device` enumerates **SIPs**; one SIP-device runs its 2 CUBE Groups (2 KV heads); the head is picked by CUBE coordinate (`head_of_group`), not an in-kernel loop (ADR-0060 §0).