7fad0371c5
Three logically distinct changes, bundled for atomic test green:
1. **P3c — prefill_long tile-granular Ring KV** (ADR-0060 §5.5.1 amendment).
Convert the ring from slice-granular (one full ``(d_head, S_local)``
KV slice per step) to tile-granular (``n_tiles`` tiles of
``TILE_S_KV`` per step). Nested loop with outer tile, inner ring step:
each tile propagates through all C ring positions before the next
tile starts, so IPCQ in-flight depth stays at 1 per direction.
Bootstrap at ``(t=0, k=0)`` outside the scratch_scope establishes the
persistent ``(m, ℓ, O)``; every other iteration scope-wraps + persists
via ``copy_to``. Per-rank persistent scratch shrinks to ~1 KB; per-tile
scope bounded by TILE_S_KV regardless of S_local. Headline:
prefill_long now completes at S_kv=128K (previously overflowed).
New: ``tests/attention/test_gqa_prefill_long_tile_ring.py``
(3 tests — ceiling-lift + tile-granular ipcq_copy count +
per-CUBE distributed output regression guard).
2. **Rename ``gqa_*`` → ``gqa_attention_*``** across kernel files,
function names, and importers. The "attention" name makes the role
explicit (GQA is grouped-query attention) and matches upstream Triton
FlashAttention naming conventions. Renames:
_gqa_decode_long.py -> _gqa_attention_decode_long.py
_gqa_decode_short.py -> _gqa_attention_decode_short.py
_gqa_prefill_long.py -> _gqa_attention_prefill_long.py
_gqa_prefill_short.py -> _gqa_attention_prefill_short.py
And function names ``gqa_<phase>_<context>_kernel`` →
``gqa_attention_<phase>_<context>_kernel``. Updated 1 bench file
(milestone_gqa_headline.py) and 10 test files.
3. **ADR-0060 / 0062 / 0063 / 0064: Proposed → Accepted**.
All four are reflected in production code and covered by tests:
- ADR-0060 (GQA fused attention): 4 kernels deployed; §5.5.1
amendment added for the tile-granular Ring KV introduced by P3c
(EN + KO mirror).
- ADR-0062 (lazy tl.load): LoadFuture + _await_pending live in
tl_context.py.
- ADR-0063 (tl.scratch_scope + tl.copy_to): used in every chain
reduce + tile sweep + ring step. EN-only previously; KO
translation authored as part of this commit (CLAUDE.md
bidirectional rule).
- ADR-0064 (per-op-type CPU issue cost): cpu_issue_cost.py +
issue_cost_table wiring in tl_context.py (Phase E).
Files git mv'd from docs/adr-proposed/ to docs/adr/ (EN) and
docs/adr-ko/ (KO). ADR-0061 (tl.broadcast) stays Proposed — no
implementation; documented as optional convenience primitive in
the ADR itself.
Tests: 88/88 focused regression green
(tests/attention/ + Phase E + TL discipline).
ADR pair verification: ``python tools/verify_adr_lang_pairs.py`` OK.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1115 lines
69 KiB
Markdown
1115 lines
69 KiB
Markdown
# ADR-0060: AHBM GQA Fused Attention 커널 (Llama3-70B)
|
||
|
||
## Status
|
||
|
||
Accepted
|
||
|
||
**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 어텐션 커널과 그 `milestone-gqa-llama70b`
|
||
eval bench(본 ADR의 커널이 도입되면서 제거됨). *§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 — 두 개의 SP 커널 (decode = reduce, prefill = ring)
|
||
|
||
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
|
||
# 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
|
||
```
|
||
|
||
```python
|
||
# 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))
|
||
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))
|
||
```
|
||
|
||
```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 | 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가 이미 있음 — 신규는
|
||
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`로 오버랩한다.
|
||
|
||
---
|
||
|
||
## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것)
|
||
|
||
kernbench는 본 ADR 이전에 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로
|
||
FlashAttention을 돌리는 두 mesh 커널이 있었다(현재 제거됨):
|
||
|
||
| Role | Mechanism |
|
||
|---|---|
|
||
| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
|
||
| 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}`)이 구동했다.
|
||
|
||
**이들은 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 재사용 없음.** baseline은 `h_q == h_kv == 1`로 제한되어 있었다.
|
||
해당 baseline 테스트는 이를 *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`가
|
||
필요 → **2-level reduce-to-root**(intra-CUBE tree + intra-CUBE-Group
|
||
center-mesh, §4)는 `⌈log₂ P⌉` + center-mesh-over-`C` 단계.
|
||
3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile
|
||
임시값을 누수하고(baseline 테스트는 그래서 S=16으로 제한되어 있었음) 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 참조).
|
||
|
||
### CUBE Group — 배치 단위 (계층적 SP)
|
||
|
||
**`CUBE Group`은 한 KV head와 그 `G` query head를 공동 소유하는 `C`개의 CUBE
|
||
(한 SIP 내부)이다.** 한 KV head의 KV 시퀀스는 **두 레벨**의
|
||
sequence-parallelism(SP)으로 샤딩된다:
|
||
- **Level-1 (inter-CUBE, CUBE Group 내):** head의 시퀀스를 group의 `C` CUBE에
|
||
분할.
|
||
- **Level-2 (intra-CUBE, PE 간):** 각 CUBE의 조각을 다시 `P` PE에 분할.
|
||
|
||
따라서 한 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를 가지며,
|
||
**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`
|
||
이면 SIP은 **2 CUBE Group = 2 KV head**를 담고; 따라서 전체 `H_kv = 8` 모델은
|
||
**4 SIP**(8/2)에 걸친다. 각 CUBE Group은 **intra-SIP**이므로 한 head의 reduction은
|
||
CUBE NOC(Level-1) + PE IPCQ(Level-2)를 쓴다 — UCIe는 결코 아님. (출하된
|
||
`topology.yaml`은 `sips: 2`; full scale은 4-SIP config 필요 — §B.)
|
||
|
||
**`--device`는 SIP을 열거한다**(CLI 의미): 한 SIP-device 벤치가 자기 2 CUBE
|
||
Group(2 KV head)을 구동 — head는 CUBE 좌표로 공간적으로 선택, 커널 내 `for kv`
|
||
루프 **아님**. CLI가 4 SIP-device를 논리 병렬로 돌려 8 head 전부를 커버.
|
||
|
||
**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를 가로지르지 않는다"
|
||
non-goal을 **뒤집는데**, 그 non-goal은 `H_kv=1` baseline이 필요로 한 적이 없어서
|
||
존재했을 뿐이다. 한 head가 **SIP**을 가로지르는 것은 여전히 non-goal(head는 한
|
||
SIP에 머문다).
|
||
|
||
---
|
||
|
||
## 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 = 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 폭).
|
||
`S` = 전체 context 길이. `R = C·P` = CUBE Group당 rank 수;
|
||
`S_rank` = 이 rank 소유 key ≈ `⌈S/R⌉`. 저장 `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-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. **연속** 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 회전. |
|
||
| `cube_id`, `pe_id`, `C`, `P` | scalars | launch / `tl.program_id` 1/0 | CUBE-Group reduction geometry (`R = C·P`). |
|
||
| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
|
||
| `O_base` | address | launch arg | 최종 O가 쓰이는 곳 (CUBE-Group root만). |
|
||
| `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
|
||
|
||
**출력 head 분포가 커널별로 다르다** — downstream out-projection이 그에 맞게
|
||
소비해야 한다:
|
||
|
||
| 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 없음. |
|
||
|
||
**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이
|
||
없음(reduce 없음); 대신 KV 블록이 순환(§5.5).
|
||
|
||
**No-SP (`C=P=1`):** IPCQ partial 없음; 단일 rank의 실행 중 `(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`, per-rank 전체 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 이점 없음"이라는
|
||
> 논거로 순수 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 비용을 모델이 지워버린 것이다.
|
||
> 하이브리드가 효율 커널의 충실한 표현이다. 이점의 **측정 가능한** 크기(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 할당 (2-level SP)
|
||
|
||
한 KV head의 시퀀스는 `C × P` rank에 샤딩된다(Level-1 inter-CUBE × Level-2
|
||
intra-CUBE PE, §0). kernbench에서는 sequence 차원을 `cube` 축(`C`에 대해
|
||
`row_wise`)과 `pe` 축(`P`에 대해 `row_wise`) **양쪽**으로 샤딩하고 Q는 `replicate`인
|
||
`DPPolicy`다(`policy/placement/dp.py`). baseline의 단일-축
|
||
`DPPolicy(pe="row_wise")`가 (cube, pe) 두 축 `row_wise`가 된다.
|
||
|
||
- per-rank KV 버퍼는 `⌈max_context / (C·P)⌉ × d × dtype`로 K와 V 각각 sizing.
|
||
- 각 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 = rank·B + local_slot`.
|
||
|
||
> **배치가 *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을 공급하고; 커널이 나머지를
|
||
도출한다:
|
||
|
||
| Launch arg | Purpose |
|
||
|---|---|
|
||
| `K_base[rank]`, `V_base[rank]` | per-rank KV 버퍼 base (tensor VA에서) |
|
||
| `O_base` | 어텐션 출력 목적지 |
|
||
| `global_token_counter` | 현재 sequence 위치 |
|
||
| `start_cube`, `start_pe = f(request_id)` | Level-1/Level-2 회전 |
|
||
| `cube_id`, `pe_id` (`tl.program_id` 1/0), `C`, `P` | CUBE-Group geometry |
|
||
| `q_block_meta` | prefill/SP query block 시작 + 길이 |
|
||
|
||
`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`라 하자. 커널 도출(평범한 산술, `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. placement policy는 단일 contiguous-block 맵
|
||
`token t → rank ⌊t/B⌋`(`start_*`로 회전).
|
||
|
||
---
|
||
|
||
## 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) — 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와 동일 수학):
|
||
|
||
```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 CUBE-Group root
|
||
```
|
||
|
||
머지가 결합적 **이고** 교환적이므로, combine은 **reduce-to-root**(어텐션은 한
|
||
곳 — `O_base`를 쓰는 rank — 에서만 `O`가 필요, all-reduce 아님)이며 **두 레벨**로
|
||
진행한다:
|
||
|
||
### 4.1 Level-2 — intra-CUBE (P PE → 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)
|
||
|
||
- group의 `C` CUBE에 대한 **center-root bidirectional CUBE-mesh reduce**
|
||
(PE-root-only, `program_id(axis=1)` = `cube_id`), CUBE NOC 2D mesh(ADR-0017) 상.
|
||
이는 **`lrab_hierarchical_allreduce.py`의 검증된 inter-CUBE 패턴(Phase 1–2:
|
||
row-then-col이 center CUBE로 수렴)을 차용**하되 — **reduce-only**(broadcast-back
|
||
Phase 4–5 제거; 어텐션은 root-only) + plain `+` 대신 **log-sum-exp `merge`**.
|
||
center-root는 corner root 대비 critical path 절반
|
||
(`lrab_hierarchical_allreduce.py:113-116`). inter-SIP phase 없음 — CUBE Group은
|
||
intra-SIP(§0).
|
||
|
||
### 4.3 Data-driven 오버랩, global barrier 없음
|
||
|
||
- rank는 자기 로컬 P·V가 끝나는 **즉시** tree 위로 send — sibling rank를 **기다리지
|
||
않음**; 내부 노드는 자식이 도착하는 대로 `merge`. 따라서 reduction은 **느린
|
||
rank의 compute와 오버랩**(causal prefill 불균형, batched decode token).
|
||
- 두 레벨은 **파이프라인**: 어떤 CUBE의 Level-2 reduction이 끝나면 그 즉시 Level-1에
|
||
기여, 다른 CUBE가 아직 Level-2 중이어도 — **레벨 사이 barrier 없음**. (rank는
|
||
*자기* 로컬 sweep 완료 전엔 send 불가; pre-final 부분값 전송은 무거운 `O` payload를
|
||
여러 번 보내야 해서 가치 없음.)
|
||
- **왜 baseline fan-out / all-reduce가 아니라 reduce-to-root인가:** baseline은 답을
|
||
모든 rank에 복제(`R−1` 단계); 어텐션은 한 번만 필요. reduce-to-root는
|
||
`⌈log₂ P⌉ + (C에 대한 center-mesh depth)` — decode(짧은 sweep, reduction 지배)엔
|
||
이것이 지배적 이득.
|
||
|
||
### 4.4 Configurable 토폴로지
|
||
|
||
각 레벨의 collective는 **선택 가능**(튜닝 항목, §9): 위 기본값(Level-2 tree,
|
||
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`(decode는 `[G·1, d]`)가 무거운 부분; `m,ℓ`은 가볍다. 별도
|
||
`tl.send`로 전송(baseline 패턴).
|
||
|
||
커널 구조(greenlet), 두 레벨이 `merge` 재사용:
|
||
|
||
```python
|
||
def hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base):
|
||
# ---- Level-2: PE tree within the CUBE → CUBE root PE ----
|
||
for child in tree_children_dirs(pe_id, P): # PE IPCQ N/S/E/W
|
||
m, l, O = merge(m, l, O, *recv_triplet(child))
|
||
if not is_cube_root(pe_id, P):
|
||
send_triplet(parent_dir(pe_id, P), m, l, O); return
|
||
# ---- Level-1: CUBE-mesh center-root reduce → Group root (cube roots only) ----
|
||
for child in mesh_children_dirs(cube_id, C): # CUBE NOC, center-root
|
||
m, l, O = merge(m, l, O, *recv_triplet(child))
|
||
if is_group_root(cube_id, C):
|
||
tl.store(o_base, O / l) # normalise once
|
||
else:
|
||
send_triplet(parent_dir_mesh(cube_id, C), m, l, O)
|
||
```
|
||
|
||
reduction tree/mesh는 **정적**(컴파일 타임 고정)이라 recv 순서가 data-independent
|
||
— data-dependent `pop` 없음, plain blocking `tl.recv`로 충분, **하드웨어/composite
|
||
`pop`-as-dependency 변경이 필요 없는** 이유(§6).
|
||
|
||
---
|
||
|
||
## 5. Case별 커널
|
||
|
||
모든 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 Decode skeleton (head-replicated, 정적 shard, reduce)
|
||
|
||
```python
|
||
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_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·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): # 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 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.
|
||
- **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_rank`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial
|
||
attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) —
|
||
바로 baseline의 `_partial_attention` 구조, 단지 GQA-batched에 composite 경로.
|
||
Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만 발동.
|
||
|
||
### 5.3 DECODE, with SP / KV-parallel (`C × P` rank)
|
||
|
||
- 한 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
|
||
sweep으로, §4.3 레벨 파이프라이닝으로 숨겨진다. 짧은 decode엔 **작은 `C`**(inter-CUBE
|
||
reduction 적음) 선호; context 길이가 추가 KV-parallelism을 요구할 때만 `C`를 올린다.
|
||
|
||
### 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) — head-parallel, reduce 없음
|
||
|
||
이것은 **두 번째 커널**(TL;DR의 Kernel 2)이며 decode reduce 경로(§4)와 구조적으로
|
||
다르다. head당 출력 `O = [T_q, d]`가 **크므로** rank 간 reduce하면 rank당 `[T_q,d]`
|
||
이동; 대신 **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를 생산하므로
|
||
**`(m,ℓ,O)` reduce 없음** — 각 CUBE가 자기 head의 행을 정규화·기록.
|
||
- **Ring KV:** `C` KV slice가 CUBE ring을 **회전**; 각 CUBE가 들어오는 블록을 자기
|
||
head의 실행 중 `(m,ℓ,O)`에 fold(online-softmax, ring step 가로질러 Python 핸들로
|
||
carry). `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` — 이미 존재); 수신 버퍼는 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이 이미 ring fold를 구현; 본 ADR은 GQA 재사용,
|
||
head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가.
|
||
|
||
#### 5.5.1 Tile-단위 Ring KV (수정안)
|
||
|
||
위 §5.5 설명은 매 ring step마다 *전체 `(d_head, S_local)` KV 슬라이스* 를
|
||
전송한다. `S_local`이 작을 때는 동작하지만, rank당 scratch가 `Kc + Vc`
|
||
(≈ `2·d·S_local·2` 바이트)에 종속되며, 이 값이 score-stack을 압도하여
|
||
`S_local = 32K` 한참 이전에 1 MiB pool을 초과한다.
|
||
|
||
구현된 ring은 **tile-단위(tile-granular)** 다 (ADR-0063 §A.2 + 본 ADR):
|
||
각 ring step은 전체 슬라이스가 아닌 `(d_head, TILE_S_KV)` K tile 1개와
|
||
그에 대응하는 V tile 1개만 전송한다. 커널 루프는 **중첩(nested)** 구조 —
|
||
`for t in range(n_tiles): for k in range(C): ...` — 다음 tile이 시작되기
|
||
전에 현재 tile이 ring의 `C`개 위치를 모두 통과한다. Rank당 persistent
|
||
scratch는 `(m, ℓ, O)` (≈ 1 KB) 로 축소되고, tile당 in-scope scratch는
|
||
`S_local`에 무관하게 `TILE_S_KV`로 제한된다.
|
||
|
||
Send 횟수 영향 (CUBE당, 방향당):
|
||
|
||
- **슬라이스-단위 ring** (위 §5.5 baseline): `2·(C−1)` —
|
||
`C−1` ring step × 2 handle (K, V).
|
||
- **Tile-단위 ring** (구현): `2·n_tiles·(C−1)` —
|
||
`n_tiles · (C−1)` send/recv pair × 2 handle.
|
||
|
||
총 IPCQ 바이트 수는 동일하다 (같은 데이터가 더 잘게 쪼개져 순환). IPCQ
|
||
command 수는 `n_tiles` 배로 증가. `n_tiles = 1` (작은 `S_local`) 에서는
|
||
두 공식이 일치하므로 기존 `test_prefill_ring_c_*` (`S_kv ∈ {16, 32}`)
|
||
테스트들은 원래 count를 그대로 만족한다.
|
||
|
||
루프 중첩 순서가 IPCQ buffer depth에 중요하다: **외부 tile, 내부 ring
|
||
step** 구조에서는 한 tile이 sender가 다음 tile을 만들기 전에 다음 CUBE가
|
||
소비하므로 in-flight depth가 방향당 1로 유지된다. 반대 중첩 순서 (외부
|
||
ring step, 내부 tile) 는 `k = 0`에서 한 CUBE가 `recv` 없이 `2·n_tiles`개의
|
||
send를 누적시켜 IPCQ slot pool을 초과하면 deadlock된다.
|
||
|
||
구현 위치는 `src/kernbench/benches/_gqa_attention_prefill_long.py`
|
||
(부트스트랩은 `(t=0, k=0)` 에서 peel off; 나머지 iteration은
|
||
`tl.scratch_scope` 으로 감싸고 `tl.copy_to`로 `(m, ℓ, O)` 를 persist).
|
||
검증: `tests/attention/test_gqa_prefill_long_tile_ring.py`
|
||
(`S_kv = 128K` ceiling-lift, tile-단위 ipcq_copy count, CUBE별 출력 쓰기
|
||
횟수 regression guard).
|
||
|
||
### 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`로 오버랩한다.
|
||
|
||
---
|
||
|
||
## 6. 왜 하드웨어 / composite 변경이 불필요한가
|
||
|
||
- HW/composite `pop`-as-dependency가 도움 될 유일한 곳은 poll을 숨길 다른 작업이
|
||
없는 외로운 reduction — 즉 batch=1 **그리고** 짧은 context. 타깃은 **agentic =
|
||
낮은 batch, 긴 context** ⇒ 각 rank가 많은 KV tile을 가짐; §4 reduce의 `tl.recv`
|
||
blocking은 동시 token / 긴 context / §4.3 레벨 파이프라인의 sweep 작업으로 가려진다.
|
||
- §4 reduction은 **정적** 2-level tree/mesh를 쓰므로 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 순서 | **커널** (정적 2-level tree/mesh) |
|
||
| 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`; 새 커맨드 종류 없음).
|
||
- **2-level reduce-to-root**(§4)가 baseline all-to-all fan-out 대체 — Level-2 PE
|
||
tree(intra-CUBE) + Level-1 center-root CUBE-mesh reduce(intra-CUBE-Group,
|
||
`lrab_hierarchical_allreduce` inter-CUBE 패턴을 reduce-only + log-sum-exp로 차용),
|
||
data-driven·레벨 파이프라인, `tl.send`/`tl.recv` 상의 순수 커널 제어 흐름.
|
||
- causal tile skip + additive boundary mask(§3/§5.4) — 커널 `if` + `+`로 더한
|
||
mask 텐서.
|
||
- 2-level contiguous-block KV placement / valid-length 산술(§0, §2) — launch-arg 산술 +
|
||
`cube`와 `pe` 양쪽 `row_wise`인 `DPPolicy`.
|
||
|
||
**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 수명까지 흡수하는 **새** 커맨드 종류다. **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
|
||
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 토폴로지 (§4.4)** — Level-2(intra-CUBE PE)와 Level-1
|
||
(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**.
|
||
5. **decode의 one-shot vs tiled 교차점**(§5.2) — tiling이 단일 composite를 이기는
|
||
`S_rank` 임계값.
|
||
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 | Head map | KV strategy | Cross-rank comm | Masking |
|
||
|---|---|---|---|---|
|
||
| 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 | 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]` (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[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`는 미생성.
|
||
|
||
---
|
||
|
||
## 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차 체크.
|
||
|
||
**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가 *아님*.)
|
||
**수치 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. **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가 직렬
|
||
`Σ(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 설계 항목 (추후 검토)
|
||
|
||
이들은 결정이 순수 primitive-op(`tl.dot`) 경로에서 **composite hybrid + lazy
|
||
`tl.load`**(이번 개정)로 옮겨갈 때 생겼다. 설계를 막는 것은 없으며; 각각 구현 중
|
||
검증 패스가 필요하다. (묻지 않고) 작업 합의에 따라 여기 기록한다 — 권고는 내가
|
||
예측한 기본값이며; 검토 시 수정.
|
||
|
||
1. **DDD-0060 동기화 완료 (이번 개정).** Detailed Design Document를 두 커널
|
||
composite-hybrid + CUBE-Group SP 설계에 맞춰 재작성함(file plan, placement,
|
||
phase plan P1–P8, verification, risks). ADR이 여전히 권위 있는 설계 기록이고;
|
||
DDD는 구현 how-to로서 rationale는 여기(및 §B)를 참조한다.
|
||
|
||
2. **composite GEMM의 K operand 방향.** Q·Kᵀ는 `b = [d, TILE]`가 필요하나 KV
|
||
cache는 K를 `[S_rank, d]`로 저장한다. `tl.trans`는 메타데이터 전용이고
|
||
`MemoryStore.read`는 transpose가 아니라 reshape(`memory_store.py:73`) — 비자명
|
||
데이터엔 런타임 transpose가 틀리다. **권고:** cache에 K를 **사전 transpose**
|
||
`[d, S_rank]`로 저장(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 미러.
|
||
|
||
### 계층적 CUBE-Group SP 전환에서 나온 항목 (이번 개정)
|
||
|
||
1. **full scale엔 4-SIP 토폴로지 config 필요.** 출하된 `topology.yaml`은
|
||
`sips: count: 2`; `C=8`의 full `H_kv=8` 모델은 **4 SIP**(각 2 KV head, §0) 필요.
|
||
**권고:** headline eval용 4-SIP 토폴로지 config(SIP당 16-CUBE 4×4 mesh, 8 PE/CUBE
|
||
— 이미 per-SIP 형태) 추가; validation-스케일 실행은 더 적은 CUBE/SIP 사용 가능.
|
||
벤치는 launch당 single-device(한 SIP) 유지해야 함.
|
||
|
||
2. **`head_of_cube` / 4×4 mesh의 CUBE-Group 분할.** `C=8`이면 각 SIP의 16-CUBE
|
||
mesh가 8짜리 CUBE Group 2개로 분할. 정확한 sub-mesh 형태(4×2 vs 2×4)가 Level-1
|
||
center-root mesh hop 수와 어느 CUBE가 1-hop 이웃인지를 정한다. **권고:** 각 CUBE
|
||
Group이 진짜 center CUBE를 갖는 연속 직사각 sub-mesh가 되도록 분할 선택(그래야
|
||
`lrab` 식 center-root 적용); CUBE NOC 라우팅(ADR-0017) 대조 검증.
|
||
|
||
3. **operating-point별 노브로서의 `C`.** `C`는 case별로 달라야 함(reduction 지배인
|
||
decode엔 작게, long-context prefill엔 크게). **권고:** `C`(및 레벨별 토폴로지,
|
||
§4.4)를 launch/bench config로 노출, milestone 벤치에서 sweep하고 decode-vs-prefill
|
||
최적값을 보고. 기본 `C`는 측정으로 TBD.
|
||
|
||
4. **§0 "never spans CUBEs" non-goal을 뒤집음 — 의존하는 downstream 가정이 없는지
|
||
확인.** 이전 텍스트와 어쩌면 다른 ADR이 intra-CUBE-only reduction을 invariant로
|
||
취급. **권고:** 구현 전 그 가정(SFR install, address policy, diagram)을 grep;
|
||
이 커널 범위에선 발견 안 됨, 단 `configure_sfr_*` 이웃 wiring이 Level-1용 CUBE-NOC
|
||
`N/S/E/W`를 PE IPCQ뿐 아니라 노출해야 함.
|
||
|
||
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 분리에서 나온 항목 (이번 개정)
|
||
|
||
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(현재 제거됨)과 정합.** 그 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이 적은-발행 이득을 측정 가능하게 만들 때까지 연기.
|