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 2d03052..65226cb 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 @@ -32,67 +32,125 @@ collective(ADR-0023/0025)에 매핑한다. --- -## TL;DR — 최종 GQA 커널 (composite hybrid) pseudocode +## TL;DR — 두 개의 SP 커널 (decode = reduce, prefill = ring) -결정(§1)을 한 곳에: **GEMM → `tl.composite`; softmax 머지 + 2-level reduction -→ 커널; `tl.load`은 lazy.** KV head는 **CUBE 좌표**로 선택(head당 CUBE Group -하나, §0) — **`for kv` 루프 없음**. `G`는 matmul M 차원에 fold(진짜 GQA, -broadcast 없음). KV 시퀀스는 `C × P` rank로 샤딩(Level-1 inter-CUBE × Level-2 -intra-CUBE PE); combine은 2-level reduce-to-root(§4). Reference 형태: +Decode-SP와 prefill-SP는 **구조적으로 다르며 두 개의 커널**이다 — 하나로 둘 다 +못 한다. 원리는 *더 작은 것을 옮긴다*: + +- **Decode** (T_q=1): 출력 `O = [G, d]`가 작고, KV cache는 크다. → KV를 **정적 + 샤딩·상주**시키고 로컬 sweep 후 작은 `(m,ℓ,O)`만 **2-level reduce**(§4)로 옮긴다. + **Q는 replicate** — 전 `G` query head를 GEMM M-차원에 쌓음(**M-fold**), 한 + `Q·Kᵀ`가 K를 공유하며 전 head를 계산. KV 이동 없음. +- **Prefill** (T_q=S): head당 출력 `O = [S, d]`가 크다 — reduce하면 rank당 + `[S,d]` 이동. → 대신 **각 CUBE가 query head 하나만 소유**(head-parallel, `C=G`) + 하고 **KV를 회전**시켜 각 head가 전체 KV를 보게 한다(**Ring KV**, §5.5). + `(m,ℓ,O)` reduce 없음(각 CUBE가 서로 다른 head 출력). KV 블록만 이동. + +메모리 여유가 있으니 **Q/weight는 replicate해서 통신을 없애고**; 옮길 수밖에 없는 +것만 옮긴다 — `(m,ℓ,O)`(decode) 또는 KV 블록(prefill). 둘 다 §3의 composite-hybrid +inner tile(GEMM → `tl.composite`, softmax merge는 커널, lazy `tl.load`)을 공유. + +### Kernel 1 — DECODE + SP (head-replicated, KV 정적 shard, 2-level reduce; ring 없음) + +decode inner tile은 강한 체인 `Q·Kᵀ → softmax → P·V`를 가진다: softmax +(`tl.max(Sj)`)가 첫 GEMM을 기다리므로 naïve 루프는 CPU를 `Sj`에서 멈추고 **softmax +도는 동안 GEMM 엔진을 idle로 둔다(버블)**. 세 변형이 CPU/HW 복잡도와 그 버블을 +trade한다(지금은 **opt3** 출시; **opt2**는 새 커맨드 종류 필요 — cost model과 함께 +재검토, §8/ADR-0064): ```python -def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr, - counter, start_pe, q_block, scale, *, tl): - # ---- geometry: this rank within its CUBE Group (§0, §2) ---- - cube_id = tl.program_id(axis=1) # which CUBE in the SIP mesh - pe_id = tl.program_id(axis=0) # which PE in the CUBE - kv = head_of_cube(cube_id) # CUBE coord → KV head (no for-kv loop) - C, P = cube_group_size, pes_per_cube # Level-1 × Level-2 SP extents - G = H_q // H_kv # query heads per KV head (=8) - causal = q_block.is_prefill +# OPTION 1 — current CompositeCmd (today; HAS the GEMM-engine bubble) +def gqa_decode_v1(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) # CUBE → its KV head + q_g = tl.load(q_base(kv), (G * 1, d)) # T_q=1; Q replicated: G heads stacked into M (M-fold) + m, l, O = init_running(G, d) # persistent arena (-inf, 0, 0) + for j in range(ceil(S_kv_local / TILE)): # sweep ONLY my KV shard (resident, no move) + with tl.scratch_scope(): + Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale]) + # ↓ CPU auto-waits on Sj → GEMM engine IDLE during softmax (bubble) + m2 = tl.maximum(m, tl.max(Sj, -1)); P = tl.exp(Sj - m2); corr = tl.exp(m - m2) + l = l * corr + tl.sum(P, -1) + O = O * corr + tl.composite("gemm", a=P, b=tl.ref(v_tile(j), (TILE, d))); m = m2 + hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) # §4 2-level reduce +``` - # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062). - q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d] - my_len = valid_len_2level(counter, start_pe, cube_id, pe_id, C, P) \ - if not causal else S_kv_local - 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 - - # ---- 2-level combine to the CUBE Group root (§4); data-driven, no barrier ---- +```python +# OPTION 3 — software pipelining (current primitives; bubble 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)) + m, l, O = init_running(G, d); n = ceil(S_kv_local / TILE) + Sb = double_buffer() # 2 persistent Sj buffers (outside scratch_scope) + h = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(0), (d, TILE)), out=Sb[0], epi=[scale]) # prime + for j in range(n): + Sj = Sb[j % 2] + if j + 1 < n: # ← issue NEXT Q·Kᵀ before softmax → fills GEMM engine + h = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j+1), (d, TILE)), out=Sb[(j+1)%2], epi=[scale]) + with tl.scratch_scope(): + m2 = tl.maximum(m, tl.max(Sj, -1)); P = tl.exp(Sj - m2); corr = tl.exp(m - m2) + l = l * corr + tl.sum(P, -1) + O = O * corr + tl.composite("gemm", a=P, b=tl.ref(v_tile(j), (TILE, d))); m = m2 hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) ``` -> **이 형태인 이유:** 두 `tl.composite` 호출이 tiling + K/V DMA 스트리밍 + -> cross-tile 파이프라이닝을 PE_SCHEDULER에 offload한다(CPU는 coarse descriptor를 -> 내고 앞서 나가 → 엔진이 saturate 유지, §1); softmax 머지와 IPCQ reduction은 -> 커널에 남는다(기존 `CompositeCmd`가 cross-tile 상태를 못 들기 때문). reduction은 -> **2-level**(CUBE 내 Level-2 PE tree → CUBE Group 내 Level-1 CUBE-mesh reduce), -> 둘 다 reduce-to-root·log-sum-exp이며, 각 rank가 자기 로컬 sweep을 마치는 대로 -> 전송(§4). `K`는 reshape-not-transpose caveat를 피하려 transpose된 `[d, TILE]`로 -> 미리 저장(§3, §B). 전용 "flash-composite" 커맨드 종류는 도입하지 **않는다** -> (§8 항목 4). +```python +# OPTION 2 — ex_composite, 2-split (NEW flash-epilogue cmd; gives K-before-V DMA priority) +def gqa_decode_v2(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)) + acc = init_running(G, d) # (m,l,O): scheduler-updated flash accumulator + for j in range(ceil(S_kv_local / TILE)): + Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale]) # #1: reads K (priority) + tl.ex_composite("softmax_pv", s=Sj, v=tl.ref(v_tile(j), (TILE, d)), acc=acc, scale=scale) # #2: reads V + # ↑ both non-blocking; CPU never waits intra-tile → max run-ahead, fewest issues + tl.wait() # ← drain ALL composites: acc is updated async (no auto-wait, + # #2 is a serial chain through acc); only final after the last #2 + hierarchical_reduce_and_store(*acc, cube_id, pe_id, C, P, o_base(kv)) +``` + +| | 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) | + +(`#1` = 기존 GEMM composite + `scale`; `#2` = softmax + P·V + stateful +online-softmax accumulator만 신규. MATH 엔진엔 max/sum/exp가 이미 있음 — 신규는 +flash accumulator이지 ops가 아님.) + +### Kernel 2 — PREFILL + SP (CUBE당 Q head 1개, head-parallel; Ring KV, reduce 없음) + +```python +def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_block, cube_start, *, tl): + i = tl.program_id(axis=1) - cube_start # this CUBE's Q head = its start KV slice + q = tl.load(q_ptr, (T_q, d)) # MY Q head's query rows (resident, TCM) + Kc = tl.load(k_ptr, (d, S_kv_local)) # my K slice, pre-stored transposed [d, S/C] → TCM + Vc = tl.load(v_ptr, (S_kv_local, d)); src = i # my V slice [S/C, d] → TCM (K/V in TCM so the ring can send them) + m, l, O = init_running(T_q, d) + for step in range(C): # ── Ring KV: rotate blocks around C CUBEs over IPCQ ── + f = None + if step < C - 1: # send current TCM block, recv next (overlap) + tl.send("ring+", Kc); tl.send("ring+", Vc) + f = (tl.recv_async("ring-", (d, S_kv_local)), tl.recv_async("ring-", (S_kv_local, d))) + if not block_all_future(q_block, slice_pos(src)): # causal skip whole-future blocks + S = tl.composite("gemm", a=q, b=Kc, epi=[scale]) # [T_q,d]·[d,S/C] → [T_q,S/C]; Kc is TCM-resident + if block_partial(q_block, slice_pos(src)): S = S + causal_mask(q_block, slice_pos(src)) + m2 = tl.maximum(m, tl.max(S, -1)); P = tl.exp(S - m2); corr = tl.exp(m - m2) + l = l * corr + tl.sum(P, -1) + O = O * corr + tl.composite("gemm", a=P, b=Vc); m = m2 # [T_q,S/C]·[S/C,d] → [T_q,d] + if f: Kc, Vc = tl.wait(f[0]), tl.wait(f[1]); src = (src - 1) % C + tl.store(o_ptr, O / l) # MY Q head's rows — NO reduce +``` + +> **왜 두 형태인가.** Decode는 KV를 제자리에 두고 작은 `(m,ℓ,O)`만 옮긴다(§4 +> 2-level reduce); prefill은 KV를 옮기고(Ring, §5.5) 각 head의 큰 출력을 로컬에 +> 둔다. 둘 다 §3의 composite-hybrid tile을 쓴다. `K`는 reshape-not-transpose +> caveat를 피하려 transpose 저장(§3, §B); decode critical path에 전용 +> "flash-composite" 종류 없음(§8 항목 4). **출력 head 분포가 다르다** — decode는 +> 전 `G` head를 CUBE-Group root에, prefill은 CUBE마다 Q head 하나(§0.5.4). +> opt2/opt3 변형은 **decode** 관심사다: prefill에선 causal `if`가 composite에 못 +> 들어가는 커널 제어흐름이고, ring은 이미 `recv_async`로 오버랩한다. --- @@ -181,10 +239,21 @@ sequence-parallelism(SP)으로 샤딩된다: 분할. - **Level-2 (intra-CUBE, PE 간):** 각 CUBE의 조각을 다시 `P` PE에 분할. -따라서 한 (KV head, request)의 KV-parallel reduction group은 **`C × P` rank**이며 -모두 한 SIP 내부다. `G` query head는 matmul M 차원에 fold(전 rank에 replicate). -`C`는 **튜닝 노브**(예: 8 또는 4): long-context prefill엔 큰 `C`, reduction이 -지배적인 짧은 decode엔 작은 `C`(→ 1 = 단일 CUBE, intra-CUBE SP만). +따라서 한 KV head는 **`C × P` rank**에 매핑되며 모두 한 SIP 내부다. **그 `G` query +head가 rank에 어떻게 매핑되는지는 case별로 다르다**(두 커널, TL;DR / §5): + +- **Decode** (§4): **Q는 replicate** — 모든 rank가 전 `G` query head를 가지며, + **M-fold**(matmul M/행 차원에 쌓음: group의 `Q` `[G, T_q, d]` → `[G·T_q, d]`, + 한 `Q·Kᵀ` GEMM이 단일 `K`를 공유하며 전 `G` head 계산 — GQA 재사용). KV는 + 시퀀스 `C × P` 샤딩; 출력은 reduce. +- **Prefill** (§5.5): **Q는 head-parallel** — `C = G`이면 **CUBE `i`가 query head + `i` 하나만 소유**; KV는 회전(Ring); reduce 없음. + +`Q`는 작아서 replicate(decode)하거나 CUBE당 1 head로 분산(prefill)하는 게 쌈 — +옮길 수밖에 없는 것만 이동(decode: `(m,ℓ,O)`; prefill: KV 블록). `C`는 **튜닝 +노브**: prefill엔 `C = G`(CUBE당 Q head 1개); decode엔 `C`가 inter-CUBE reduction과 +KV-parallel breadth를 trade(짧은 context면 작은 `C`, 심지어 `C = 1` single-CUBE, +reduction 지배 시). **Topology grounding** (`topology.yaml`): SIP은 `4×4` CUBE mesh(16 CUBE, ADR-0017 NOC); CUBE는 `P = 8` PE(`hbm_pseudo_channels/hbm_channels_per_pe = 64/8`). `C = 8` @@ -262,14 +331,18 @@ V_block`을 ping-pong 버퍼에(post-RoPE), 그리고 causal step-skip용 ### 0.5.4 OUTPUTS -| Output | Shape | Location | Notes | -|---|---|---|---| -| `O` (final) | `[G, T_q, d]` | `O_base`, **CUBE-Group root만** | 정규화 `O_acc / ℓ_acc`, 저장 dtype으로 cast. | +**출력 head 분포가 커널별로 다르다** — downstream out-projection이 그에 맞게 +소비해야 한다: -**중간값 (non-root rank, IPCQ 상 — 커널-가시 출력 아님):** -`(m_i, ℓ_i, O_i)` (`m,ℓ`: `[G, T_q]`; `O_i`: `[G, T_q, d]` 비정규화)이 2-level -tree를 따라 부모(Level-2 PE 부모, 이어 Level-1 CUBE 부모)에 `tl.send`로 push(§4). -`O_i`가 무거운 부분. +| Kernel | Output | Location | Notes | +|---|---|---|---| +| **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 없음. | + +**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이 +없음(reduce 없음); 대신 KV 블록이 순환(§5.5). **No-SP (`C=P=1`):** IPCQ partial 없음; 단일 rank의 실행 중 `(m,ℓ,O)`를 제자리에서 정규화하여 `O_base`에 기록. @@ -285,8 +358,7 @@ score `S` 및 prob `P`(미생성). `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 모델에 근거한 -근거: +load가 뒤따르는 compute와 오버랩된다. kernbench의 실행 + latency 모델에 근거: - **GEMM tiling이 PE_SCHEDULER에 offload된다.** `CompositeCmd`는 non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): 커널이 **하나의 coarse @@ -368,6 +440,18 @@ intra-CUBE PE, §0). kernbench에서는 sequence 차원을 `cube` 축(`C`에 대 `pe_local = (pe_id − start_pe) mod P`); 이후 `global_idx = local_slot·(C·P) + rank`. +> **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 블록 형태가 두 커널을 다 서비스.) + ### 2.2 드라이버 launch당 의무 (최소) 드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를 도출한다: @@ -449,6 +533,10 @@ Notes: ## 4. Reduction (KV-parallel / SP combine) — 2-level reduce-to-root +**범위: 이것은 DECODE-SP 경로**(head-replicated, KV 정적 샤딩, 작은 `O`). Prefill-SP는 +reduce **안 함** — KV를 회전(Ring, §5.5). decode에 reduce를 택한 이유는 +`O = [G, d]`가 작아 `(m,ℓ,O)` 이동이 상주 KV 이동보다 싸기 때문. + tile sweep 후 각 rank는 자기 `1/(C·P)` shard에 대해 `(m_i, ℓ_i, O_i)`(비정규화)를 가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학, `_attention_mesh_mlo.py:117-122`): @@ -467,9 +555,12 @@ def merge(m_a, l_a, O_a, m_b, l_b, O_b): ### 4.1 Level-2 — intra-CUBE (P PE → CUBE root PE) -- CUBE의 `P` PE에 대한 **reduce-to-root tree**, depth `⌈log₂ P⌉`(`P=8`에 3), PE - IPCQ `N/S/E/W` mesh 상(`configure_sfr_intracube_pe_ring`). 각 tree pair는 1-hop - 물리 이웃. 결과는 CUBE의 root PE에 안착. +- CUBE의 KV slice 자체가 **그 `P` PE에 sequence-SP-split**된다(결정 (a): decode + `T_q=1`이 전 `P` PE를 쓰는 유일한 방법 — 쪼갤 query 축이 없음). 각 PE가 자기 + `1/(C·P)` 서브-shard를 sweep한 뒤, `P` PE에 대한 **reduce-to-root tree**, + depth `⌈log₂ P⌉`(`P=8`에 3), PE IPCQ `N/S/E/W` mesh 상 + (`configure_sfr_intracube_pe_ring`). 각 tree pair는 1-hop 물리 이웃. 결과는 + CUBE의 root PE에 안착. ### 4.2 Level-1 — intra-CUBE-Group (C CUBE root → Group root) @@ -499,12 +590,13 @@ def merge(m_a, l_a, O_a, m_b, l_b, O_b): ### 4.4 Configurable 토폴로지 각 레벨의 collective는 **선택 가능**(튜닝 항목, §9): 위 기본값(Level-2 tree, -Level-1 center-root mesh)은 latency-bound·작은-`O` decode에 적합; **ring** 옵션 -(chunked `O`)은 큰 `T_q·d` payload의 대역폭-bound long-context prefill에 적합. +Level-1 center-root mesh)은 latency-bound·작은-`O` decode에 적합; 특정 reduction이 +대역폭-bound로 판명되면 **ring** 변형(chunked)이 가능. (참고: 이것은 **decode**용 +*reduce* collective; **prefill**은 reduce를 전혀 안 함 — §5.5의 Ring-KV 커널 사용.) `C=1`이면 Level-2만으로 축약(single-CUBE SP). -**Payload:** `O_i`(`[G·T_q, d]`)가 무겁고; `m,ℓ`은 가볍다. 별도 `tl.send`로 전송 -(baseline 패턴). +**Payload:** `O_i`(decode는 `[G·1, d]`)가 무거운 부분; `m,ℓ`은 가볍다. 별도 +`tl.send`로 전송(baseline 패턴). 커널 구조(greenlet), 두 레벨이 `merge` 재사용: @@ -532,30 +624,32 @@ reduction tree/mesh는 **정적**(컴파일 타임 고정)이라 recv 순서가 ## 5. Case별 커널 -네 case 모두 §3(tile sweep) + §4(reduction)를 공유하며; SP extent `C·P`, -query-block 폭, 어느 `tile_*` predicate가 발동하는지에서만 다르다. +모든 case가 §3의 composite-hybrid inner tile을 공유하나 **두 커널**로 분리된다 +(TL;DR): 아래의 **decode-reduce** skeleton(§5.2/§5.3)과 **prefill-ring** +커널(§5.5). head 매핑, KV 전략, cross-rank 통신에서 다르다 — §10 참조. -### 5.1 공통 skeleton +### 5.1 Decode skeleton (head-replicated, 정적 shard, reduce) ```python -def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, - C, P, q_block, softmax_scale, *, tl): +def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, + C, P, *, tl): cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0) - kv = head_of_cube(cube_id) # CUBE coord → KV head (no for-kv loop) + kv = head_of_group(cube_id) # CUBE coord → KV head (no for-kv loop) my_len = valid_len_2level(counter, start_cube, start_pe, cube_id, pe_id, C, P) n_tiles = ceil(my_len / TILE) - q_g = load_Q_group(q_ptr, kv) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast) + q_g = load_Q_group(q_ptr, kv) # [G·1, d]; lazy tl.load, G folded into M (replicated) 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 + for j in range(n_tiles): # sweep ONLY my static KV shard (resident, no move) + run_tile(j) # §3: 2 composites + softmax MATH, in tl.scratch_scope if C == 1 and P == 1: tl.store(o_base(kv), O / l) # no reduction (single rank) else: - hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) # §4 + hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) # §4 reduce ``` +(Prefill-no-SP는 `C=P=1`, `T_q>1`, causal masking인 같은 형태; +**prefill-with-SP는 §5.5의 별도 Ring 커널**, 이 skeleton 아님.) + ### 5.2 DECODE, no SP (`C=P=1`, 한 PE가 head의 KV 소유) - `T_q = 1`, 모든 과거 KV에 attend ⇒ future tile 없음, 마지막 ragged tile만 mask. @@ -593,23 +687,48 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube, 더함(`Sj + mask_j`). - GQA는 group의 query head를 decode와 같은 축으로 batch. -### 5.5 PREFILL, with SP (Ring KV) +### 5.5 PREFILL, with SP (Ring KV) — head-parallel, reduce 없음 -- KV가 `C·P` rank에 **ring 순서**로 sharded; 각 step이 peer의 KV 블록을 IPCQ로 - 전달. 실행 중 `(m,ℓ,O)`는 **ring step을 가로질러** carry(Python 핸들, - `_attention_mesh_kv`가 오늘 하는 그대로). ring은 long-context·대역폭-bound 케이스로 - §4.4의 **ring** collective 옵션(tree 아님)이 자연스러운 선택. +이것은 **두 번째 커널**(TL;DR의 Kernel 2)이며 decode reduce 경로(§4)와 구조적으로 +다르다. head당 출력 `O = [T_q, d]`가 **크므로** rank 간 reduce하면 rank당 `[T_q,d]` +이동; 대신 **head를 shard**하고 (역시 큰) **KV를 옮긴다** — 어차피 각 head가 전체 +KV를 필요로 하므로. + +- **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를 생산하므로 + **`(m,ℓ,O)` reduce 없음** — 각 CUBE가 자기 head의 행을 정규화·기록. +- **Ring KV:** `C` KV slice가 CUBE ring을 **회전**; 각 CUBE가 들어오는 블록을 자기 + head의 실행 중 `(m,ℓ,O)`에 fold(online-softmax, ring step 가로질러 Python 핸들로 + carry, `_attention_mesh_kv`가 오늘 하듯). `C` step 후 모든 head가 전 KV를 봄. + **GQA 재사용은 회전에서** — slice `j`가 전 `C` CUBE를 방문하며 전 `G` head를 서비스. - 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)`는 rank별 최종; 남은 KV-parallel 분할은 §4 2-level reduce로 - 결합한 뒤 Group root에서 정규화 + store. + 오버랩(`tl_context.py:543-560` — 이미 존재); 수신 버퍼는 ping-pong(persistent + arena, 재활용 안 함). +- **Causal ring skip:** 들어오는 KV 블록이 전부 이 head의 query block *이후*면 그 + compute를 skip — 약 절반 step 제거. +- **CUBE 내(P PE):** head의 query 행 `[T_q, d]` 및/또는 현재 KV 블록을 `P` PE에 + 타일(여기엔 decode와 달리 query 축이 존재); 상세는 §B. -baseline `_attention_mesh_kv`가 이미 ring fold를 구현한다; 본 ADR은 GQA 재사용, -causal step-skip, `recv_async` 오버랩을 추가한다. +baseline `_attention_mesh_kv`가 이미 ring fold를 구현; 본 ADR은 GQA 재사용, +head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가. + +### 5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2) + +세 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ᵀ = +K를 먼저 읽는 기존 composite, `#2` softmax+P·V+accumulator = V를 나중에 읽는 유일한 +신규 flash-epilogue 기계장치). + +**권고:** 지금은 **opt3** 출시(새 기계장치 0, 버블 제거); cost model(ADR-0064)이 +적은-CPU-발행 이득을 측정 가능하게 만들면 **opt2** 재검토(§8 항목 4). MATH 엔진엔 +max/sum/exp가 이미 있음 — opt2의 유일한 진짜 신규는 composite의 **stateful +`(m,ℓ,O)` accumulator**이지 ops가 아님. + +이 변형들은 **decode** 관심사다: prefill(§5.5)에선 causal `if`가 composite에 못 +들어가는 커널 제어흐름이고, ring은 이미 `recv_async`로 오버랩한다. --- @@ -685,10 +804,13 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다. 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를 측정 가능하게 가치 있게 만들면 재검토. + cross-tile register 수명까지 흡수하는 **새** 커맨드 종류다. **Sizing note (§5.6 + opt2):** 재검토 시 **두** composite로 분할 — `#1` = Q·Kᵀ(기존 composite + `scale`; + DMA가 K를 우선시하게), `#2` = softmax + P·V + online-softmax accumulator 머지 + (*유일한* 진짜 신규: reduction epilogue + stateful `(m,ℓ,O)` accumulator). MATH + 엔진엔 max/sum/exp가 이미 있음 — 신규는 composite의 stateful flash accumulator이지 + ops가 아님. **cost model(ADR-0064)이 적은-CPU-발행 이득을 측정 가능하게 만들면 + 재검토**(§5.6); 그 전까지 §5.6 opt3(software pipelining, 새 cmd 없음) 출시. 5. ~~하드웨어 `pop`-as-dependency.~~ 범위 밖(§6). 6. ~~본 커널 내 RoPE / QKV projection / KV-cache write.~~ Upstream `qkv_rope` (P1–P5). RoPE를 fold하면 매 decode step마다 과거 tile을 재회전해야 하고 Ring @@ -702,10 +824,10 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다. non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당 스트리밍 depth. 2. **레벨별 reduction 토폴로지 (§4.4)** — Level-2(intra-CUBE PE)와 Level-1 - (intra-CUBE-Group CUBE-mesh) 각각 tree/ring/center-mesh 선택 가능; tree(latency, - 작은-`O` decode) vs ring(대역폭, 큰-`O` prefill) sweep. 그리고 CUBE-Group 크기 - `C` 노브(decode엔 작게, long context엔 크게). 각 tree pair가 1-hop 물리 이웃인지 - 보장; SFR install 대조 검증. + (intra-CUBE-Group CUBE-mesh) 각각 tree/center-mesh(및 ring 변형) 선택 가능 — + **decode reduce 전용**; prefill은 reduce가 아니라 §5.5 Ring-KV 커널 사용. 그리고 + decode `C` 노브(reduction 지배인 짧은 context엔 작은 `C`). 각 tree pair가 1-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**. @@ -720,21 +842,21 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다. ## 10. Coverage 요약 -| Case | KV placement | Inner loop | Reduction | Masking | +| Case | Head map | KV strategy | Cross-rank comm | Masking | |---|---|---|---|---| -| Decode, no SP | 1 rank (`C=P=1`), all KV | one-shot or tiled, GQA-batched | none | last tile only | -| Decode, SP | 2-level RR over `C·P` ranks | tiled, GQA-batched | §4 2-level (`⌈log₂P⌉`+mesh over `C`) | last tile only | -| Prefill, no SP | resident | tiled per q-block | none | skip future / triangular boundary | -| Prefill, SP (Ring) | ring over `C·P` ranks | per ring step (recv_async overlap) | running state + §4 2-level | causal step-skip + boundary | +| Decode, no SP | `G` replicated, 1 rank | all KV resident | none | last tile only | +| **Decode, SP** | **Q replicated** (전 `G` query head를 GEMM M-dim에 쌓음) | 2-level static shard `C·P` | **§4 2-level reduce** (small `O`) | last tile only | +| Prefill, no SP | `G` replicated, 1 rank | resident | none | triangular / skip future | +| **Prefill, SP (Ring)** | **1 Q head per CUBE** (`C=G`) | **Ring KV rotate** | **none** (KV blocks move, not `O`) | causal step-skip + boundary | **Case별 I/O** (전체 계약 §0.5): -| Case | Inputs | Output | IPCQ partials | +| Case | Inputs | Output | Cross-rank traffic | |---|---|---|---| | Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 rank | `O[G,1,d]` at `O_base` | none | -| Decode, SP | `Q[G,1,d]`, per-rank `K/V[S/(C·P),d]` | `O[G,1,d]` (Group root) | `(m,ℓ,O_i)` per rank → 2-level | +| Decode, SP | `Q[G,1,d]` (replicated), per-rank `K/V[S/(C·P),d]` | `O[G,1,d]` at Group root | `(m,ℓ,O_i)` reduce (small) | | 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]` (Group root) | running `(m,ℓ,O)` + 2-level | +| Prefill, SP (Ring) | `Q[1,T_q,d]` per CUBE, own `K/V[S/C,d]` | `O[1,T_q,d]` per CUBE (distributed) | KV blocks rotate (Ring) | 모든 case에서: KV-cache write와 RoPE는 **upstream**에서; output projection은 **downstream**에서; score `S`와 prob `P`는 미생성. @@ -760,13 +882,16 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` 2. **GQA 재사용:** `h_q = G·h_kv`에서 K/V `dma_read_count`가 `G`와 무관(타일당 한 load, group에 걸쳐 재사용), 반면 GEMM 작업은 `G`로 스케일. 레버가 실제 발동함을 assert. -3. **2-level reduction 단계 수:** decode-SP가 `⌈log₂ P⌉`(Level-2, intra-CUBE) + - `C`에-대한-center-mesh-depth(Level-1) reduction 라운드를 발행 — baseline의 - `C·P − 1` all-to-all이 아님; op_log `ipcq_send`/`recv` 수가 정적 2-level - tree/mesh와 일치하고, 결과가 정확히 한 rank(CUBE-Group root)에 안착. 별도 assert가 - Level-1 send는 CUBE NOC, Level-2는 PE IPCQ를 쓰는지 확인. -4. **Causal skip:** prefill이 모든 `tile_all_future` tile을 skip — GEMM 수가 full - grid가 아니라 하삼각 tile 수와 일치. +3. **SP cross-rank 트래픽 — 두 커널:** + - *Decode (reduce):* `⌈log₂ P⌉`(Level-2, intra-CUBE) + `C`에-대한-center-mesh-depth + (Level-1) reduction 라운드 발행 — baseline의 `C·P − 1` all-to-all이 아님; op_log + `ipcq_send`/`recv`가 정적 2-level tree/mesh와 일치; 결과가 정확히 한 rank + (CUBE-Group root)에 안착; Level-1 send는 CUBE NOC, Level-2는 PE IPCQ. + - *Prefill (ring):* `(m,ℓ,O)` reduce **없음** — 대신 `C` KV-블록 회전(step마다 + K,V의 `ipcq_send`/`recv`); 각 CUBE가 **서로 다른** head의 `O`를 제자리에 기록 + (Group root 없음); GQA 재사용은 각 KV slice가 ring을 통해 전 `C` head에 소비됨으로. +4. **Causal skip:** prefill이 모든 `tile_all_future` tile(및 ring의 전체-미래 KV + 블록)을 skip — GEMM 수가 full grid가 아니라 하삼각 tile/step 수와 일치. 5. **Long context (ADR-0063):** scope 없이 1 MiB를 넘기는 `S` sweep이 완료하고 reference와 일치. 6. **Load/compute 오버랩:** tiled sweep의 end-to-end latency가 직렬 @@ -865,3 +990,43 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16` pe_local)` 배치(§2.1)는 시퀀스를 gap/overlap 없이 타일링하고 각 rank의 로컬 버퍼를 dense하게 유지해야 함. **권고:** 커널 전에 index 산술을 unit-test(모든 global token이 정확히 한 rank에 매핑; per-rank slice 연속). + +### decode-reduce / prefill-ring 분리에서 나온 항목 (이번 개정) + +1. **두 커널, 두 head 매핑.** Decode-SP = head-replicated + 정적 KV shard + 2-level + reduce(§4); Prefill-SP = head-parallel(1 Q head/CUBE, `C=G`) + Ring KV + reduce + 없음(§5.5). 원리는 *더 작은 것을 옮긴다* — decode의 `O`는 작아서 reduce, prefill의 + `O`는 커서 KV를 대신 이동. **권고:** 두 커널로 유지; 재병합 금지. + +2. **출력 head 분포가 다름 (downstream 영향).** Decode는 전 `G` head를 CUBE-Group + root에, prefill은 CUBE당 Q head 1개(분산). downstream **out-projection**이 각 + 레이아웃을 소비해야 함(decode-root는 gather, prefill은 per-CUBE 제자리). **권고:** + 구현 전 `qkv_rope`/`out_proj` 계약에 커널별 O 레이아웃 고정(§0.5.4). + +3. **prefill CUBE 내 `P` PE.** §5.5는 head의 query 행 `[T_q, d]` 및/또는 현재 KV + 블록을 `P` PE에 분할(decode와 달리 query 축 존재). **권고:** 기본은 query 행을 PE에 + 타일(disjoint 출력 행 → intra-CUBE reduce 불필요); `T_q < P`일 때만 KV-블록 분할 + + intra-CUBE reduce로 fallback. + +4. **prefill의 `C = G` 결합.** head-parallel 매핑은 `C = G = 8`(CUBE당 Q head 1개)을 + 가정. `C ≠ G`면 매핑 재검토 필요(CUBE당 다중 head, 또는 head가 부분 ring에 걸침). + **권고:** headline 스케일에서 prefill 커널은 `C = G` 고정; `C ≠ G`는 별도 연구. + +5. **`_attention_mesh_mlo_2d`(현재 impl)와 정합.** 원격 impl의 2D 커널은 **Q + replicated**로 cube에 걸친 **AllReduce** — 즉 decode-reduce 계열이나 reduce-to-root + 아닌 all-reduce(broadcast-back), 그리고 prefill head-parallel ring은 아직 없음. + **권고:** (a) 그 2D AllReduce → reduce-to-root(broadcast-back 제거)로 decode 커널; + (b) §5.5 head-parallel Ring-KV 커널을 prefill용으로 추가. + +6. **공유 KV cache 레이아웃 (prefill 쓰고, decode 읽고+연장).** 두 커널이 하나의 물리 + KV cache를 공유하므로 **contiguous `C×P` 위치 블록**(round-robin 아님)이어야 함 — + prefill causal skip이 contiguous 블록을 요구하고, 공유 레이아웃이 prefill→decode + reshard를 피함(§2.1). **권고:** `qkv_rope` write 계약에서 contiguous 블록 레이아웃 + 표준화; **짧은-context** decode가 contiguous에서 rank를 덜 쓴다는 점 플래그 + (long-context 타깃엔 수용 가능; 짧은-context 균형은 별도 연구). + +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`, + `#2`만 신규)는 ADR-0064 cost model이 적은-발행 이득을 측정 가능하게 만들 때까지 연기.