gqa(adr): ADR-0060 hierarchical CUBE-Group SP — 2-level reduce-to-root

Placement pivot (user-approved): a CUBE Group = C CUBEs within one SIP that
jointly own one KV head + its G=8 query heads. KV sequence sharded 2 levels
(Level-1 inter-CUBE over C, Level-2 intra-CUBE PE over P=8 = C*P ranks); G
folded into matmul M. C is a knob (8/4; C=1 = single-CUBE). Reverses the old
'a query head never spans CUBEs' non-goal (held only because the baseline
was H_kv=1). device=SIP; the for-kv loop is gone (head picked by CUBE coord).

Reduction is a 2-level reduce-to-root (not all-reduce): Level-2 PE tree ->
Level-1 center-root CUBE-mesh reduce, adapting lrab_hierarchical_allreduce's
inter-CUBE pattern as reduce-only + log-sum-exp. Data-driven (send on local
P.V completion, no global barrier) + level-pipelined; per-level topology
configurable (tree for decode, ring for long prefill).

Rewrites SS0/2/4/5/0.5/8-11, pseudocode, and adds SSB items (4-SIP config,
mesh partition, C knob, invariant-reversal check, index-math test). KO mirror
updated. Topology grounded in topology.yaml (4x4 CUBE mesh/SIP, 8 PE/cube).
Docs only; no production code changed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-04 17:35:43 -07:00
parent f009ba7b1f
commit a61fed3ce0
2 changed files with 604 additions and 342 deletions
@@ -34,23 +34,27 @@ collective(ADR-0023/0025)에 매핑한다.
## TL;DR — 최종 GQA 커널 (composite hybrid) pseudocode ## TL;DR — 최종 GQA 커널 (composite hybrid) pseudocode
결정(§1)을 한 곳에: **GEMM → `tl.composite`; softmax 머지 + tree reduction 결정(§1)을 한 곳에: **GEMM → `tl.composite`; softmax 머지 + 2-level reduction
커널; `tl.load`은 lazy.** KV head별로, `G`를 matmul M 차원에 fold(진짜 GQA, 커널; `tl.load`은 lazy.** KV head는 **CUBE 좌표**로 선택(head당 CUBE Group
broadcast 없음). 이것이 reference 형태이며, 산문 섹션이 각 부분을 부연한다. 하나, §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 형태:
```python ```python
def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr, def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
counter, start_pe, N, q_block, scale, *, tl): counter, start_pe, q_block, scale, *, tl):
# ---- geometry (kernel arithmetic, §2) ---- # ---- geometry: this rank within its CUBE Group (§0, §2) ----
pe_id = tl.program_id(axis=rank_axis) 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) G = H_q // H_kv # query heads per KV head (=8)
causal = q_block.is_prefill causal = q_block.is_prefill
for kv in range(H_kv): # one KV head per iteration (§5.2) # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062).
# Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
# issue now, auto-wait at first use inside the first composite.
q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d] q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d]
my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv 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) n_tiles = ceil(my_len / TILE)
# persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros # persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
@@ -76,19 +80,19 @@ def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
O = O * corr + Oj # running merge O = O * corr + Oj # running merge
m = m_new m = m_new
# ---- cross-PE combine (§4): log-sum-exp tree to root, or store ---- # ---- 2-level combine to the CUBE Group root (§4); data-driven, no barrier ----
if N == 1: hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))
tl.store(o_base(kv), O / l)
else:
tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv
``` ```
> **이 형태인 이유:** 두 `tl.composite` 호출이 tiling + K/V DMA 스트리밍 + > **이 형태인 이유:** 두 `tl.composite` 호출이 tiling + K/V DMA 스트리밍 +
> cross-tile 파이프라이닝을 PE_SCHEDULER에 offload한다(CPU는 coarse descriptor를 > cross-tile 파이프라이닝을 PE_SCHEDULER에 offload한다(CPU는 coarse descriptor를
> 내고 앞서 나가 → 엔진이 saturate 유지, §1); softmax 머지와 IPCQ tree reduction은 > 내고 앞서 나가 → 엔진이 saturate 유지, §1); softmax 머지와 IPCQ reduction은
> 커널에 남는다(기존 `CompositeCmd`가 cross-tile 상태를 못 들기 때문). `K`는 > 커널에 남는다(기존 `CompositeCmd`가 cross-tile 상태를 못 들기 때문). reduction은
> reshape-not-transpose caveat를 피하려 transpose된 `[d, TILE]`로 미리 저장된다 > **2-level**(CUBE 내 Level-2 PE tree → CUBE Group 내 Level-1 CUBE-mesh reduce),
> (§3, §B). 전용 "flash-composite" 커맨드 종류는 도입하지 **않는다**(§8 항목 4). > 둘 다 reduce-to-root·log-sum-exp이며, 각 rank가 자기 로컬 sweep을 마치는 대로
> 전송(§4). `K`는 reshape-not-transpose caveat를 피하려 transpose된 `[d, TILE]`로
> 미리 저장(§3, §B). 전용 "flash-composite" 커맨드 종류는 도입하지 **않는다**
> (§8 항목 4).
--- ---
@@ -127,7 +131,8 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
진짜 GQA(`h_q = G·h_kv`)가 **신규 primitive 없이** 돈다 — §8 참조. 진짜 GQA(`h_q = G·h_kv`)가 **신규 primitive 없이** 돈다 — §8 참조.
2. **O(N) reduction.** baseline은 all-to-all bidirectional fan-out을 하여 *모든* 2. **O(N) reduction.** baseline은 all-to-all bidirectional fan-out을 하여 *모든*
rank가 full 답을 갖는다(`n_ranks 1` 단계). 어텐션은 query owner에서만 `O` rank가 full 답을 갖는다(`n_ranks 1` 단계). 어텐션은 query owner에서만 `O`
필요 → root로의 **tree reduction**은 `⌈log₂ N⌉` 단계(§4). 필요 → **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 3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile
임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking / 임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking /
tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite
@@ -166,20 +171,42 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
pop/wait 못 한다. 따라서 본 ADR은 두 어텐션 GEMM(Q·Kᵀ, P·V)을 **각각** 자체 pop/wait 못 한다. 따라서 본 ADR은 두 어텐션 GEMM(Q·Kᵀ, P·V)을 **각각** 자체
composite로 발행하고 cross-GEMM softmax 머지 + IPCQ reduction은 커널에 유지한다 composite로 발행하고 cross-GEMM softmax 머지 + IPCQ reduction은 커널에 유지한다
— 새 "flash-composite" 커맨드 종류는 **필요 없다**(§1, §8 참조). — 새 "flash-composite" 커맨드 종류는 **필요 없다**(§1, §8 참조).
- **Allocation policy (SP):** multi-user 패널에서 CUBE당 query head 하나; KV
group의 `G` query head는 한 CUBE의 PE 안에 매핑.
**직교하는 두 매핑 레이어** (round-robin placement): ### CUBE Group — 배치 단위 (계층적 SP)
- **Layer 1 (KV-parallel, intra-request):** KV token `i`는 PE
`(start_pe + i) mod N`에 안착 — round-robin이라 한 긴 request가 `N` PE에 걸쳐
분할, ≤1 token으로 균형.
- **Layer 2 (inter-request):** `start_pe = request_id mod N`이 request별 "PE-0
역할"을 회전.
- 결합: `pe = (request_id + global_token_idx) mod N`.
`N` = 하나의 (query head, request)에 대한 KV-parallel reduction group의 PE 수. **`CUBE Group`은 한 KV head와 그 `G` query head를 공동 소유하는 `C`개의 CUBE
Reduction은 **intra-CUBE** 유지(query head는 CUBE를 가로지르지 않음 — 명시적 (한 SIP 내부)이다.** 한 KV head의 KV 시퀀스는 **두 레벨**의
non-goal). sequence-parallelism(SP)으로 샤딩된다:
- **Level-1 (inter-CUBE, CUBE Group 내):** head의 시퀀스를 group의 `C` CUBE에
분할.
- **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만).
**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 내, 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를 소유.
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에 머문다).
--- ---
@@ -206,27 +233,27 @@ non-goal).
대해 **pure read**. 대해 **pure read**.
- **P4.** V는 회전 안 함; V-cache는 raw projected V 보유. - **P4.** V는 회전 안 함; V-cache는 raw projected V 보유.
- **P5.** upstream RoPE 위치는 token의 **절대 global 위치** - **P5.** upstream RoPE 위치는 token의 **절대 global 위치**
`global_idx = local_slot·N + ((pe_id start_pe) mod N)`(§2.1). Round-robin `global_idx = local_slot·(C·P) + rank` (여기서 `rank = cube_local·P + pe_local`,
placement ≠ RoPE 위치. §2.1). Round-robin placement ≠ RoPE 위치.
### 0.5.2 Shape 기호 ### 0.5.2 Shape 기호
`T_q` = 이번 launch의 query 길이(decode: 1; prefill/chunk: chunk 폭). `T_q` = 이번 launch의 query 길이(decode: 1; prefill/chunk: chunk 폭).
`S` = 전체 context 길이. `S_pe` = 이 PE 소유 key ≈ `⌈S/N⌉`. `S` = 전체 context 길이. `R = C·P` = CUBE Group당 rank 수;
저장 `bf16`(numpy proxy `f16`, `memory_store.py:16`); `m,,O` 누산기는 정밀도가 `S_rank` = 이 rank 소유 key ≈ `⌈S/R⌉`. 저장 `bf16`(numpy proxy `f16`,
중요한 곳에서 `f32`. `memory_store.py:16`); `m,,O` 누산기는 정밀도가 중요한 곳에서 `f32`.
### 0.5.3 INPUTS (커널 launch당) ### 0.5.3 INPUTS (커널 launch당)
| Input | Shape (per KV head) | Location | Notes | | Input | Shape (per KV head) | Location | Notes |
|---|---|---|---| |---|---|---|---|
| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). group의 `G` query 행 batched. | | `Q` | `[G, T_q, d]` | per-rank HBM (loaded to TCM) | post-RoPE (P1). group의 `G` query 행 batched. |
| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. 로컬 dense, global index는 strided (§2.1). | | `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). |
| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. | | `V_cache` | `[S/(C·P), d]` | per-rank HBM, base `V_base[rank]` | raw V (P4). Read-only. |
| `global_token_counter` | scalar | launch arg | 커널이 `S_pe`, slot↔global, causal bound 도출. | | `global_token_counter` | scalar | launch arg | 커널이 로컬 len, slot↔global, causal bound 도출. |
| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 회전. | | `start_cube`, `start_pe` (= `f(request_id)`) | scalars | launch arg | Level-1/Level-2 회전. |
| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. | | `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. | | `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
| `O_base` | address | launch arg | 최종 O가 쓰이는 곳 (root PE만). | | `O_base` | address | launch arg | 최종 O가 쓰이는 곳 (CUBE-Group root만). |
| `softmax_scale` | scalar | launch arg | `1/√d`. | | `softmax_scale` | scalar | launch arg | `1/√d`. |
**Ring Attention (§5.5)**에는 ring step마다 추가: IPCQ로 들어오는 `K_block, **Ring Attention (§5.5)**에는 ring step마다 추가: IPCQ로 들어오는 `K_block,
@@ -237,14 +264,15 @@ V_block`을 ping-pong 버퍼에(post-RoPE), 그리고 causal step-skip용
| Output | Shape | Location | Notes | | Output | Shape | Location | Notes |
|---|---|---|---| |---|---|---|---|
| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE** | 정규화 `O_acc / _acc`, 저장 dtype으로 cast. | | `O` (final) | `[G, T_q, d]` | `O_base`, **CUBE-Group root만** | 정규화 `O_acc / _acc`, 저장 dtype으로 cast. |
**중간값 (non-root PE, IPCQ 상 — 커널-가시 출력 아님):** **중간값 (non-root rank, IPCQ 상 — 커널-가시 출력 아님):**
`(m_i, _i, O_i)` (`m,`: `[G, T_q]`; `O_i`: `[G, T_q, d]` 비정규화)이 `(m_i, _i, O_i)` (`m,`: `[G, T_q]`; `O_i`: `[G, T_q, d]` 비정규화)이 2-level
`tl.send``tree_parent`에 push(§4). `O_i`가 무거운 부분. tree를 따라 부모(Level-2 PE 부모, 이어 Level-1 CUBE 부모)에 `tl.send`로 push(§4).
`O_i`가 무거운 부분.
**No-SP (`N=1`):** IPCQ partial 없음; 단일 PE의 실행 중 `(m,,O)` 제자리에서 **No-SP (`C=P=1`):** IPCQ partial 없음; 단일 rank의 실행 중 `(m,,O)`
정규화하여 `O_base`에 기록. 제자리에서 정규화하여 `O_base`에 기록.
**명시적 비-출력:** KV-cache write(upstream, P3), output projection(downstream), **명시적 비-출력:** KV-cache write(upstream, P3), output projection(downstream),
score `S` 및 prob `P`(미생성). score `S` 및 prob `P`(미생성).
@@ -262,7 +290,7 @@ load가 뒤따르는 compute와 오버랩된다. kernbench의 실행 + latency
- **GEMM tiling이 PE_SCHEDULER에 offload된다.** `CompositeCmd`는 non-blocking - **GEMM tiling이 PE_SCHEDULER에 offload된다.** `CompositeCmd`는 non-blocking
(`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): 커널이 **하나의 coarse (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): 커널이 **하나의 coarse
descriptor**(M = `G·T_q`, PE당 전체 tile sweep)를 push하면 scheduler가 tile descriptor**(M = `G·T_q`, per-rank 전체 tile sweep)를 push하면 scheduler가 tile
plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다(ADR-0014 D6). K/V는 plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다(ADR-0014 D6). K/V는
scheduler가 HBM에서 스트리밍하는 `tl.ref` operand이므로, 타일당 **K/V prefetch는 scheduler가 HBM에서 스트리밍하는 `tl.ref` operand이므로, 타일당 **K/V prefetch는
scheduler의 일**이다 — 명시적 prefetch op 없음. CPU(greenlet)는 현재 composite가 scheduler의 일**이다 — 명시적 prefetch op 없음. CPU(greenlet)는 현재 composite가
@@ -324,14 +352,21 @@ per tile j:
## 2. 메모리 레이아웃 & 드라이버 책임 ## 2. 메모리 레이아웃 & 드라이버 책임
### 2.1 KV cache 할당 ### 2.1 KV cache 할당 (2-level SP)
- per-PE KV 버퍼는 per-PE 최대 context `⌈max_context / N⌉ × d × dtype`로 K와 V를
각각 sizing. kernbench에서는 sequence 차원에 대해 `pe`(또는 `cube`) 축이 한 KV head의 시퀀스는 `C × P` rank에 샤딩된다(Level-1 inter-CUBE × Level-2
`row_wise``DPPolicy`로 배치된다(`policy/placement/dp.py`; baseline은 K/V에 intra-CUBE PE, §0). kernbench에서는 sequence 차원을 `cube` 축(`C`에 대해
`DPPolicy(pe="row_wise")`, Q `replicate`). `row_wise`)과 `pe` 축(`P`에 대해 `row_wise`) **양쪽**으로 샤딩하고 Q `replicate`
- PE 안에서 할당된 token은 **연속 append**된다(slot 0,1,2,…). global index는 `DPPolicy`다(`policy/placement/dp.py`). baseline의 단일-축
strided(`i, i+N, …`)이나 per-PE 버퍼는 dense ⇒ DMA가 연속 유지. `DPPolicy(pe="row_wise")`가 (cube, pe) 두 축 `row_wise`가 된다.
- Slot → global: `global_idx = local_slot·N + ((pe_id start_pe) mod N)`.
- 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`
(여기서 `cube_local = (cube_id start_cube) mod C`,
`pe_local = (pe_id start_pe) mod P`); 이후
`global_idx = local_slot·(C·P) + rank`.
### 2.2 드라이버 launch당 의무 (최소) ### 2.2 드라이버 launch당 의무 (최소)
드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를 드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를
@@ -339,22 +374,24 @@ per tile j:
| Launch arg | Purpose | | Launch arg | Purpose |
|---|---| |---|---|
| `K_base[pe]`, `V_base[pe]` | per-PE KV 버퍼 base (tensor VA에서) | | `K_base[rank]`, `V_base[rank]` | per-rank KV 버퍼 base (tensor VA에서) |
| `O_base` | 어텐션 출력 목적지 | | `O_base` | 어텐션 출력 목적지 |
| `global_token_counter` | 현재 sequence 위치 | | `global_token_counter` | 현재 sequence 위치 |
| `start_pe = request_id mod N` | Layer-2 회전 | | `start_cube`, `start_pe = f(request_id)` | Level-1/Level-2 회전 |
| `pe_id` (`tl.program_id`), `N` | geometry | | `cube_id`, `pe_id` (`tl.program_id` 1/0), `C`, `P` | CUBE-Group geometry |
| `q_block_meta` | prefill/SP query block 시작 + 길이 | | `q_block_meta` | prefill/SP query block 시작 + 길이 |
커널 도출(평범한 산술): `R = C·P`, `cube_local = (cube_id start_cube) mod C`,
- **이번 step에 내가 쓸 차례:** `(start_pe + counter) mod N == pe_id` `pe_local = (pe_id start_pe) mod P`, 이 rank의 인덱스
- **내 valid 길이:** `base = counter // N; rem = counter % N; `rank = cube_local·P + pe_local`라 하자. 커널 도출(평범한 산술):
my_len = base + (1 if ((pe_id start_pe) mod N) < rem else 0)` - **이번 step에 내가 쓸 차례:** `(start_rank + counter) mod R == rank`
- **내 valid 길이:** `base = counter // R; rem = counter % R;
my_len = base + (1 if rank < rem else 0)`
- **read 범위:** tile `0 .. ⌈my_len / TILE⌉`. - **read 범위:** tile `0 .. ⌈my_len / TILE⌉`.
- **causal bound / per-tile skip:** query block과 각 tile의 global 위치에서. - **causal bound / per-tile skip:** query block과 각 tile의 global 위치에서.
드라이버 = base + counter + rotation. 단일 공식 드라이버 = base + counter + rotation. 두 SP 축에 대한 단일 공식
`(request_id + token_idx) mod N`이 placement policy 전부다. `(request_id + token_idx) mod (C·P)`이 placement policy 전부다.
--- ---
@@ -410,78 +447,117 @@ Notes:
--- ---
## 4. Reduction (KV-parallel / SP combine) — root로의 tree ## 4. Reduction (KV-parallel / SP combine) — 2-level reduce-to-root
tile sweep 후 각 PE는 `(m_i, _i, O_i)`(비정규화)를 가진다. 결합/교환 가능한 tile sweep 후 각 rank는 자기 `1/(C·P)` shard에 대해 `(m_i, _i, O_i)`(비정규화)를
log-sum-exp 머지로 결합(baseline의 fold와 동일 수학, `_attention_mesh_mlo.py:117-122`): 가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학,
`_attention_mesh_mlo.py:117-122`):
```python ```python
def merge(m_a, l_a, O_a, m_b, l_b, O_b): def merge(m_a, l_a, O_a, m_b, l_b, O_b):
m = tl.maximum(m_a, m_b) m = tl.maximum(m_a, m_b)
sa, sb = tl.exp(m_a - m), tl.exp(m_b - m) 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 return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb
# final: O = O_root / l_root # normalise once at the tree root # final: O = O_root / l_root # normalise once at the CUBE-Group root
``` ```
- **타이밍:** 교환은 **P·V 후**(각 PE가 최종 `O_i`를 가짐). 머지가 결합적 **이고** 교환적이므로, combine은 **reduce-to-root**(어텐션은 한
- **토폴로지:** mesh tree, depth `⌈log₂ N⌉`, 물리 이웃 상. `N=8`의 경우: level-0 곳 — `O_base`를 쓰는 rank — 에서만 `O`가 필요, all-reduce 아님)이며 **두 레벨**로
`(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2 `(3↔7)`, root = `7`. 각 tree 진행한다:
pair는 물리 `N/S/E/W` 이웃이어야 `tl.send(dir)`/`tl.recv(dir)`이 실제 link를
쓴다(튜닝 항목 §9; 이웃을 wiring하는 SFR install은
`configure_sfr_intracube_pe_ring`, baseline이 사용).
- **왜 baseline fan-out이 아니라 tree인가:** 어텐션은 한 곳(`O_base`를 쓰는 PE)에서
`O`가 필요하다. tree는 `⌈log₂ N⌉` 단계(`N=8`에 3) vs baseline all-to-all의 `N1`
(7). per-PE sweep이 짧은 decode에서 reduction이 지배적 비용이므로 실제 이득이다.
(baseline fan-out은 답을 모든 rank에 복제 — 여기선 불필요.)
- **Payload:** `O_i`(`d`-vector)가 무겁고; `m,`은 `(g, T_q)`당 scalar. 별도
`tl.send`로 전송(baseline은 triplet을 세 send로 — 같은 패턴).
커널 구조(greenlet): ### 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에 안착.
### 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 12:
row-then-col이 center CUBE로 수렴)을 차용**하되 — **reduce-only**(broadcast-back
Phase 45 제거; 어텐션은 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에 복제(`R1` 단계); 어텐션은 한 번만 필요. 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에 적합; **ring** 옵션
(chunked `O`)은 큰 `T_q·d` payload의 대역폭-bound long-context prefill에 적합.
`C=1`이면 Level-2만으로 축약(single-CUBE SP).
**Payload:** `O_i`(`[G·T_q, d]`)가 무겁고; `m,`은 가볍다. 별도 `tl.send`로 전송
(baseline 패턴).
커널 구조(greenlet), 두 레벨이 `merge` 재사용:
```python ```python
# leaf / internal node: fold children that send to me, then forward up def hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base):
for child_dir in tree_children_dirs(pe_id, N): # ---- Level-2: PE tree within the CUBE → CUBE root PE ----
m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape) for child in tree_children_dirs(pe_id, P): # PE IPCQ N/S/E/W
O_c = tl.recv(child_dir, O.shape) m, l, O = merge(m, l, O, *recv_triplet(child))
m, l, O = merge(m, l, O, m_c, l_c, O_c) if not is_cube_root(pe_id, P):
if not is_root(pe_id): send_triplet(parent_dir(pe_id, P), m, l, O); return
tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O) # ---- 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: else:
tl.store(O_base, O / l) send_triplet(parent_dir_mesh(cube_id, C), m, l, O)
``` ```
`tl.recv`는 데이터 도착까지 blocking(`tl_context.py:446-499`); 머지 순서가 정적 reduction tree/mesh는 **정적**(컴파일 타임 고정)이라 recv 순서가 data-independent
tree로 고정되므로 data-dependent `pop`이 불필요 — **하드웨어/composite data-dependent `pop` 없음, plain blocking `tl.recv`로 충분, **하드웨어/composite
`pop`-as-dependency 변경이 필요 없는** 이유(§6). `pop`-as-dependency 변경이 필요 없는** 이유(§6).
--- ---
## 5. Case별 커널 ## 5. Case별 커널
네 case 모두 §3(tile sweep) + §4(reduction)를 공유하며; `N`, query-block 폭, 네 case 모두 §3(tile sweep) + §4(reduction)를 공유하며; SP extent `C·P`,
어느 `tile_*` predicate가 발동하는지에서만 다르다. query-block 폭, 어느 `tile_*` predicate가 발동하는지에서만 다르다.
### 5.1 공통 skeleton ### 5.1 공통 skeleton
```python ```python
def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
q_block, softmax_scale, *, tl): C, P, q_block, softmax_scale, *, tl):
pe_id = tl.program_id(axis=rank_axis) cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
my_len = valid_len(counter, start_pe, pe_id, N) kv = head_of_cube(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) n_tiles = ceil(my_len / TILE)
q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast) q_g = load_Q_group(q_ptr, kv) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
m, l, O = init_running() # persistent arena: -inf, 0, zeros 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) 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) if tile_all_future(j, q_block): # causal skip (kernel if)
continue continue
run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope
if N == 1: if C == 1 and P == 1:
tl.store(o_ptr, O / l) # no reduction tl.store(o_base(kv), O / l) # no reduction (single rank)
else: else:
tree_reduce_and_store(pe_id, N, o_ptr) # §4 hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) # §4
``` ```
### 5.2 DECODE, no SP (`N=1`, 한 PE가 head의 KV 소유) ### 5.2 DECODE, no SP (`C=P=1`, 한 PE가 head의 KV 소유)
- `T_q = 1`, 모든 과거 KV에 attend ⇒ future tile 없음, 마지막 ragged tile만 mask. - `T_q = 1`, 모든 과거 KV에 attend ⇒ future tile 없음, 마지막 ragged tile만 mask.
- **GQA 재사용이 전부**(decode는 KV-load-bound): KV head의 `G=8` query 행을 matmul - **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-보존). **M** 차원에 fold(`q_g`를 `[G, T_q, d] → [G·T_q, d]`로 reshape, byte-보존).
@@ -491,18 +567,22 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
그것들이 GEMM의 M 행이기 때문. K/V broadcast 불필요; composite의 tile plan 내 그것들이 GEMM의 M 행이기 때문. K/V broadcast 불필요; composite의 tile plan 내
`m = G·T_q`가 timing이 모든 `G` 행의 작업을 올바르게 세게 한다(leading batch `m = G·T_q`가 timing이 모든 `G` 행의 작업을 올바르게 세게 한다(leading batch
축은 세어지지 *않음* — §8 참조). 축은 세어지지 *않음* — §8 참조).
- `S_pe`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial - `S_rank`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial
attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) — attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) —
바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에 바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에
composite 경로. Tiling(§3)은 `S_pe`가 scratch scope의 tile 예산을 초과할 때만 composite 경로. Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만
발동. 발동.
### 5.3 DECODE, with SP / KV-parallel (`N=8`) ### 5.3 DECODE, with SP / KV-parallel (`C × P` rank)
- 한 request의 KV가 8 PE에 round-robin; 각각 ≈`my_len` token 소유. 각 PE가 자기
로컬 tile에 대해 `G=8` query 행을 GQA-batch. - 한 request의 KV가 CUBE Group의 `C·P` rank에 계층적-round-robin(Level-1은 `C`
- `T_q=1` ⇒ 짧은 per-PE sweep ⇒ **reduction 지배** ⇒ §4 tree(3단계)가 구조적으로 CUBE, Level-2는 `P` PE); 각 rank가 ≈`my_len` token 소유하고 자기 로컬 tile에
중요한 부분. Reduction latency는 batch될 때 다른 동시 decode token으로, 혹은 긴 대해 `G=8` query 행을 GQA-batch.
single-stream context의 긴 per-PE tile sweep으로 숨겨진다. - `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 ### 5.4 PREFILL, no SP
- 전체 prompt 상주; query는 `T_q` token 블록(chunk). - 전체 prompt 상주; query는 `T_q` token 블록(chunk).
@@ -514,17 +594,19 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
- GQA는 group의 query head를 decode와 같은 축으로 batch. - GQA는 group의 query head를 decode와 같은 축으로 batch.
### 5.5 PREFILL, with SP (Ring KV) ### 5.5 PREFILL, with SP (Ring KV)
- KV가 `N` PE(ring)에 sharded; 각 step이 peer의 KV 블록을 IPCQ로 전달. 실행 중
`(m,,O)`는 **ring step을 가로질러** carry(Python 핸들, `_attention_mesh_kv`가 - 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 아님)이 자연스러운 선택.
- IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로 - IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로
오버랩(`tl_context.py:543-560` — 이미 존재). 오버랩(`tl_context.py:543-560` — 이미 존재).
- **Causal ring skip:** 들어오는 step의 KV 블록이 전부 로컬 query block *이후*면 그 - **Causal ring skip:** 들어오는 step의 KV 블록이 전부 로컬 query block *이후*면 그
compute를 skip(recv-consume 안 함 / fold 안 함) — causal 어텐션에서 약 절반 step compute를 skip(recv-consume 안 함 / fold 안 함) — causal 어텐션에서 약 절반 step
제거. 제거.
- 수신 버퍼는 ping-pong(persistent arena, 재활용 안 함). - 수신 버퍼는 ping-pong(persistent arena, 재활용 안 함).
- ring 후, `(m,,O)`는 query-소유 PE별 최종; KV-parallel 분할이 남으면 §4가 - ring 후, `(m,,O)`는 rank별 최종; 남은 KV-parallel 분할은 §4 2-level reduce로
tree-merge한 뒤 정규화 + store. 결합한 뒤 Group root에서 정규화 + store.
baseline `_attention_mesh_kv`가 이미 ring fold를 구현한다; 본 ADR은 GQA 재사용, baseline `_attention_mesh_kv`가 이미 ring fold를 구현한다; 본 ADR은 GQA 재사용,
causal step-skip, `recv_async` 오버랩을 추가한다. causal step-skip, `recv_async` 오버랩을 추가한다.
@@ -535,10 +617,10 @@ causal step-skip, `recv_async` 오버랩을 추가한다.
- HW/composite `pop`-as-dependency가 도움 될 유일한 곳은 poll을 숨길 다른 작업이 - HW/composite `pop`-as-dependency가 도움 될 유일한 곳은 poll을 숨길 다른 작업이
없는 외로운 reduction — 즉 batch=1 **그리고** 짧은 context. 타깃은 **agentic = 없는 외로운 reduction — 즉 batch=1 **그리고** 짧은 context. 타깃은 **agentic =
낮은 batch, 긴 context** ⇒ 각 PE가 많은 KV tile을 가짐; §4 tree의 `tl.recv` 낮은 batch, 긴 context** ⇒ 각 rank가 많은 KV tile을 가짐; §4 reduce의 `tl.recv`
blocking은 동시 token / 긴 context의 sweep 작업으로 가려진다. blocking은 동시 token / 긴 context / §4.3 레벨 파이프라인의 sweep 작업으로 가려진다.
- §4 reduction은 **정적** tree를 쓰므로 recv 순서가 컴파일 타임에 고정 - §4 reduction은 **정적** 2-level tree/mesh를 쓰므로 recv 순서가 컴파일 타임에 고정
data-dependent pop 없음. `tl.recv`(blocking)로 충분. data-dependent pop 없음. `tl.recv`(blocking)로 충분.
- **결정: greenlet `tl.send`/`tl.recv` collective로 출시.** 짧은-context, - **결정: greenlet `tl.send`/`tl.recv` collective로 출시.** 짧은-context,
single-stream, latency-critical 타깃이 나타날 때만 HW `pop` 재검토. single-stream, latency-critical 타깃이 나타날 때만 HW `pop` 재검토.
@@ -550,7 +632,7 @@ causal step-skip, `recv_async` 오버랩을 추가한다.
|---|---| |---|---|
| tile skip (future), mask 생성, causal bound | **커널** (greenlet 본문의 Python `if` + 산술) | | tile skip (future), mask 생성, causal bound | **커널** (greenlet 본문의 Python `if` + 산술) |
| 주소 / offset / valid-length 산술 | **커널** (counter에서) | | 주소 / offset / valid-length 산술 | **커널** (counter에서) |
| reduction 스케줄링, IPCQ send/recv 순서 | **커널** (정적 tree) | | reduction 스케줄링, IPCQ send/recv 순서 | **커널** (정적 2-level tree/mesh) |
| Q·Kᵀ, P·V (per-tile K/V DMA 스트리밍 + tiling 포함) | **`tl.composite`** → PE_SCHEDULER | | 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 엔진 상 (커널 발행) | | Q load, mask add, softmax math, 실행 중 `(m,,O)` 머지 | **`tl` op** PE 엔진 상 (커널 발행) |
@@ -574,12 +656,14 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다.
- **composite를 통한 GEMM**(§1/§3) — Q·Kᵀ와 P·V를 각각 non-blocking - **composite를 통한 GEMM**(§1/§3) — Q·Kᵀ와 P·V를 각각 non-blocking
`tl.composite(op="gemm")`로 발행; PE_SCHEDULER가 tiling하고 `tl.ref` K/V `tl.composite(op="gemm")`로 발행; PE_SCHEDULER가 tiling하고 `tl.ref` K/V
operand의 DMA를 스트리밍(기존 `CompositeCmd`; 새 커맨드 종류 없음). operand의 DMA를 스트리밍(기존 `CompositeCmd`; 새 커맨드 종류 없음).
- root로의 tree reduction(§4) baseline all-to-all fan-out 대체 — `tl.send`/ - **2-level reduce-to-root**(§4) baseline all-to-all fan-out 대체 — Level-2 PE
`tl.recv` 상의 순수 커널 제어 흐름. 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` + `+`로 더한 - causal tile skip + additive boundary mask(§3/§5.4) — 커널 `if` + `+`로 더한
mask 텐서. mask 텐서.
- round-robin KV placement / valid-length 산술(§2) — launch-arg 산술 + - 2-level round-robin KV placement / valid-length 산술(§0, §2) — launch-arg 산술 +
`DPPolicy(pe="row_wise")`. `cube`와 `pe` 양쪽 `row_wise`인 `DPPolicy`.
**efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):** **efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):**
@@ -617,13 +701,16 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다.
1. **Composite tile-pipeline depth** — KV-load-bound에 지배적; 커널이 wait 전에 1. **Composite tile-pipeline depth** — KV-load-bound에 지배적; 커널이 wait 전에
non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당 non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당
스트리밍 depth. 스트리밍 depth.
2. **reduction tree의 PE↔mesh-이웃 매핑** — 각 depth-`⌈log₂ N⌉` pair가 물리 2. **레벨별 reduction 토폴로지 (§4.4)** — Level-2(intra-CUBE PE)와 Level-1
`N/S/E/W` 이웃인지 보장; 나쁜 매핑은 hop 추가. SFR install 대조 검증. (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 대조 검증.
3. **TILE 크기** — scratch 상주(`S/P tile + O_acc + G-way GQA`)와 DMA 효율 균형; 3. **TILE 크기** — scratch 상주(`S/P tile + O_acc + G-way GQA`)와 DMA 효율 균형;
ADR-0063 및 scheduler의 `TILE_M/K/N`(`pe_scheduler.py`)과 상호작용. ADR-0063 및 scheduler의 `TILE_M/K/N`(`pe_scheduler.py`)과 상호작용.
4. §5.5의 **Ring 버퍼 ping-pong vs `recv_async` depth**. 4. §5.5의 **Ring 버퍼 ping-pong vs `recv_async` depth**.
5. **decode의 one-shot vs tiled 교차점**(§5.2) — tiling이 단일 composite를 이기는 5. **decode의 one-shot vs tiled 교차점**(§5.2) — tiling이 단일 composite를 이기는
`S_pe` 임계값. `S_rank` 임계값.
6. **per-op CPU 발행 비용 (cost model)** — 현재 `dispatch_cycles=0`(`pe_cpu.py`)이라 6. **per-op CPU 발행 비용 (cost model)** — 현재 `dispatch_cycles=0`(`pe_cpu.py`)이라
composite-vs-primitive 발행 오버헤드가 보이지 않음. op 종류별 차등 발행 비용 composite-vs-primitive 발행 오버헤드가 보이지 않음. op 종류별 차등 발행 비용
(`tl.composite` descriptor push ≫ primitive op)이 하이브리드의 CPU-saturation (`tl.composite` descriptor push ≫ primitive op)이 하이브리드의 CPU-saturation
@@ -635,19 +722,19 @@ ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다.
| Case | KV placement | Inner loop | Reduction | Masking | | Case | KV placement | Inner loop | Reduction | Masking |
|---|---|---|---|---| |---|---|---|---|---|
| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only | | Decode, no SP | 1 rank (`C=P=1`), all KV | one-shot or tiled, GQA-batched | none | last tile only |
| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only | | 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, no SP | resident | tiled per q-block | none | skip future / triangular boundary |
| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary | | Prefill, SP (Ring) | ring over `C·P` ranks | per ring step (recv_async overlap) | running state + §4 2-level | causal step-skip + boundary |
**Case별 I/O** (전체 계약 §0.5): **Case별 I/O** (전체 계약 §0.5):
| Case | Inputs | Output | IPCQ partials | | Case | Inputs | Output | IPCQ partials |
|---|---|---|---| |---|---|---|---|
| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none | | Decode, 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-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,,O_i)` per PE → tree | | 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 |
| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none | | Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none |
| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,,O)` + tree | | 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 |
모든 case에서: KV-cache write와 RoPE는 **upstream**에서; output projection은 모든 case에서: KV-cache write와 RoPE는 **upstream**에서; output projection은
**downstream**에서; score `S`와 prob `P`는 미생성. **downstream**에서; score `S`와 prob `P`는 미생성.
@@ -673,8 +760,11 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`
2. **GQA 재사용:** `h_q = G·h_kv`에서 K/V `dma_read_count`가 `G`와 무관(타일당 한 2. **GQA 재사용:** `h_q = G·h_kv`에서 K/V `dma_read_count`가 `G`와 무관(타일당 한
load, group에 걸쳐 재사용), 반면 GEMM 작업은 `G`로 스케일. 레버가 실제 발동함을 load, group에 걸쳐 재사용), 반면 GEMM 작업은 `G`로 스케일. 레버가 실제 발동함을
assert. assert.
3. **tree reduction 단계 수:** decode-SP가 `⌈log₂ N` reduction 라운드(`N1`이 3. **2-level reduction 단계 수:** decode-SP가 `⌈log₂ P`(Level-2, intra-CUBE) +
아님)를 발행; op_log `ipcq_send`/`recv` 수가 tree와 일치. `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 4. **Causal skip:** prefill이 모든 `tile_all_future` tile을 skip — GEMM 수가 full
grid가 아니라 하삼각 tile 수와 일치. grid가 아니라 하삼각 tile 수와 일치.
5. **Long context (ADR-0063):** scope 없이 1 MiB를 넘기는 `S` sweep이 완료하고 5. **Long context (ADR-0063):** scope 없이 1 MiB를 넘기는 `S` sweep이 완료하고
@@ -705,10 +795,10 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`
권위 있는 기록이다. **권고:** 구현 시작 전 후속으로 DDD 동기화. 권위 있는 기록이다. **권고:** 구현 시작 전 후속으로 DDD 동기화.
2. **composite GEMM의 K operand 방향.** Q·Kᵀ는 `b = [d, TILE]`가 필요하나 KV 2. **composite GEMM의 K operand 방향.** Q·Kᵀ는 `b = [d, TILE]`가 필요하나 KV
cache는 K를 `[S_pe, d]`로 저장한다. `tl.trans`는 메타데이터 전용이고 cache는 K를 `[S_rank, d]`로 저장한다. `tl.trans`는 메타데이터 전용이고
`MemoryStore.read`는 transpose가 아니라 reshape(`memory_store.py:73`) — 비자명 `MemoryStore.read`는 transpose가 아니라 reshape(`memory_store.py:73`) — 비자명
데이터엔 런타임 transpose가 틀리다. **권고:** cache에 K를 **사전 transpose** 데이터엔 런타임 transpose가 틀리다. **권고:** cache에 K를 **사전 transpose**
`[d, S_pe]`로 저장(pseudocode와 §3이 이를 가정)하여 `tl.ref(k_tile, (d, TILE))`이 `[d, S_rank]`로 저장(pseudocode와 §3이 이를 가정)하여 `tl.ref(k_tile, (d, TILE))`이
연속 slice가 되게. upstream `qkv_rope` write 레이아웃이 이를 지원하는지 검증, 연속 slice가 되게. upstream `qkv_rope` write 레이아웃이 이를 지원하는지 검증,
아니면 실제 `tl.transpose` 추가(더 무거움; 연기). 아니면 실제 `tl.transpose` 추가(더 무거움; 연기).
@@ -745,3 +835,33 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`
fold를 기술한다. 일관성을 위해 ring의 per-step Q·Kᵀ / P·V도 composite여야 하고; fold를 기술한다. 일관성을 위해 ring의 per-step Q·Kᵀ / P·V도 composite여야 하고;
IPCQ `recv_async` 오버랩은 직교하며 유지. **권고:** 구현 중 ring step에 같은 IPCQ `recv_async` 오버랩은 직교하며 유지. **권고:** 구현 중 ring step에 같은
하이브리드 형태 적용; 낮은 리스크, §3 미러. 하이브리드 형태 적용; 낮은 리스크, §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` / 계층적 round-robin 정확성.** 두-축 `(cube_local·P +
pe_local)` 배치(§2.1)는 시퀀스를 gap/overlap 없이 타일링하고 각 rank의 로컬 버퍼를
dense하게 유지해야 함. **권고:** 커널 전에 index 산술을 unit-test(모든 global
token이 정확히 한 rank에 매핑; per-rank slice 연속).
@@ -37,23 +37,27 @@ known algorithms onto the kernbench **greenlet `tl` programming model**
## TL;DR — final GQA kernel (composite hybrid) pseudocode ## TL;DR — final GQA kernel (composite hybrid) pseudocode
The decision (§1) in one place: **GEMMs → `tl.composite`; softmax merge + The decision (§1) in one place: **GEMMs → `tl.composite`; softmax merge +
tree reduction → kernel; `tl.load` is lazy.** Per-KV-head, `G` folded into 2-level reduction → kernel; `tl.load` is lazy.** The KV head is selected by
the matmul M dim (real GQA, no broadcast). This is the reference shape; the **CUBE coordinate** (one CUBE Group per head, §0) — **no `for kv` loop**.
prose sections elaborate each piece. `G` is folded into the matmul M dim (real GQA, no broadcast). The KV
sequence is sharded `C × P` ranks (Level-1 inter-CUBE × Level-2 intra-CUBE
PE); the combine is a 2-level reduce-to-root (§4). Reference shape:
```python ```python
def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr, def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
counter, start_pe, N, q_block, scale, *, tl): counter, start_pe, q_block, scale, *, tl):
# ---- geometry (kernel arithmetic, §2) ---- # ---- geometry: this rank within its CUBE Group (§0, §2) ----
pe_id = tl.program_id(axis=rank_axis) 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) G = H_q // H_kv # query heads per KV head (=8)
causal = q_block.is_prefill causal = q_block.is_prefill
for kv in range(H_kv): # one KV head per iteration (§5.2) # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062).
# Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
# issue now, auto-wait at first use inside the first composite.
q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d] q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d]
my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv 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) n_tiles = ceil(my_len / TILE)
# persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros # persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
@@ -79,21 +83,20 @@ def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
O = O * corr + Oj # running merge O = O * corr + Oj # running merge
m = m_new m = m_new
# ---- cross-PE combine (§4): log-sum-exp tree to root, or store ---- # ---- 2-level combine to the CUBE Group root (§4); data-driven, no barrier ----
if N == 1: hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))
tl.store(o_base(kv), O / l)
else:
tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv
``` ```
> **Why this shape:** the two `tl.composite` calls offload tiling + K/V DMA > **Why this shape:** the two `tl.composite` calls offload tiling + K/V DMA
> streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse > streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse
> descriptors and runs ahead → engines stay saturated, §1); the softmax > descriptors and runs ahead → engines stay saturated, §1); the softmax
> merge and the IPCQ tree reduction stay in the kernel because the existing > merge and the IPCQ reduction stay in the kernel because the existing
> `CompositeCmd` cannot carry cross-tile state. `K` is pre-stored > `CompositeCmd` cannot carry cross-tile state. The reduction is **2-level**
> transposed `[d, TILE]` to sidestep the reshape-not-transpose caveat > (Level-2 PE tree within a CUBE → Level-1 CUBE-mesh reduce within the CUBE
> (§3, §B). A bespoke "flash-composite" command kind is **not** introduced > Group), both reduce-to-root, log-sum-exp, sent as each rank finishes its
> (§8 item 4). > local sweep (§4). `K` is pre-stored transposed `[d, TILE]` to sidestep the
> reshape-not-transpose caveat (§3, §B). A bespoke "flash-composite" command
> kind is **not** introduced (§8 item 4).
--- ---
@@ -136,8 +139,9 @@ GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this
(`h_q = G·h_kv`) runs with **no new primitive** — see §8. (`h_q = G·h_kv`) runs with **no new primitive** — see §8.
2. **O(N) reduction.** The baseline does an all-to-all bidirectional 2. **O(N) reduction.** The baseline does an all-to-all bidirectional
fan-out so *every* rank ends with the full answer (`n_ranks 1` fan-out so *every* rank ends with the full answer (`n_ranks 1`
steps). Attention only needs `O` at the query owner → a **tree steps). Attention only needs `O` at the query owner → a **2-level
reduction** to root is `⌈log₂ N⌉` steps (§4). reduce-to-root** (intra-CUBE tree + intra-CUBE-Group center-mesh, §4) is
`⌈log₂ P⌉` + center-mesh-over-`C` steps.
3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump 3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump
allocator leaks per-tile temporaries allocator leaks per-tile temporaries
(`test_milestone_gqa_llama70b.py:123-148`) and there is no causal (`test_milestone_gqa_llama70b.py:123-148`) and there is no causal
@@ -180,20 +184,48 @@ Hardware recap:
**each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and **each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and
keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel — keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel —
it does **not** need a new "flash-composite" command kind (see §1, §8). it does **not** need a new "flash-composite" command kind (see §1, §8).
- **Allocation policy (SP):** one query head per CUBE in the multi-user
panels; the `G` query heads of a KV group map within a CUBE's PEs.
**Two orthogonal mapping layers** (round-robin placement): ### CUBE Group — the placement unit (hierarchical SP)
- **Layer 1 (KV-parallel, intra-request):** KV token `i` lands on PE
`(start_pe + i) mod N` — round-robin so one long request is split
across `N` PEs, balanced to ≤1 token.
- **Layer 2 (inter-request):** `start_pe = request_id mod N` rotates the
"PE-0 role" per request.
- Combined: `pe = (request_id + global_token_idx) mod N`.
`N` = number of PEs in the KV-parallel reduction group for one **A `CUBE Group` is `C` CUBEs (within one SIP) that jointly own one KV
(query head, request). Reduction stays **intra-CUBE** (a query head never head and its `G` query heads.** One KV head's KV sequence is **sharded
spans CUBEs — explicit non-goal). two levels** of sequence-parallelism (SP):
- **Level-1 (inter-CUBE, within the CUBE Group):** the head's sequence is
split across the `C` CUBEs of the group.
- **Level-2 (intra-CUBE, across PEs):** each CUBE's slice is split again
across its `P` PEs.
So the KV-parallel reduction group for one (KV head, request) is **`C × P`
ranks**, all within one SIP. The `G` query heads are folded into the
matmul M dim (replicated across all ranks). `C` is a **tuning knob**
(e.g. 8 or 4): large `C` for long-context prefill, small `C` (→ 1 = a
single CUBE, intra-CUBE SP only) for short decode where reduction
dominates.
**Topology grounding** (`topology.yaml`): a SIP is a `4×4` CUBE mesh
(16 CUBEs, ADR-0017 NOC); a CUBE has `P = 8` PEs
(`hbm_pseudo_channels/hbm_channels_per_pe = 64/8`). With `C = 8` a SIP
holds **2 CUBE Groups = 2 KV heads**; the full `H_kv = 8` model therefore
spans **4 SIPs** (8/2). Each CUBE Group is **intra-SIP**, so one head's
reduction uses the CUBE NOC (Level-1) + PE IPCQ (Level-2) — never UCIe.
(The shipped `topology.yaml` has `sips: 2`; full scale needs a 4-SIP
config — §B.)
**`--device` enumerates SIPs** (CLI semantics): one SIP-device bench
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.
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
**reverses** the earlier "a query head never spans CUBEs" non-goal, which
existed only because the `H_kv=1` baseline never needed it. Crossing
**SIPs** for one head remains a non-goal (a head stays in one SIP).
--- ---
@@ -220,27 +252,28 @@ spans CUBEs — explicit non-goal).
launches.** ⇒ this kernel is **pure read** on the KV cache. launches.** ⇒ this kernel is **pure read** on the KV cache.
- **P4.** V is not rotated; V-cache holds raw projected V. - **P4.** V is not rotated; V-cache holds raw projected V.
- **P5.** RoPE position upstream is the token's **absolute global - **P5.** RoPE position upstream is the token's **absolute global
position** `global_idx = local_slot·N + ((pe_id start_pe) mod N)` position** `global_idx = local_slot·(C·P) + rank` where `rank =
(§2.1). Round-robin placement ≠ RoPE position. cube_local·P + pe_local` (§2.1). Round-robin placement ≠ RoPE position.
### 0.5.2 Shape symbols ### 0.5.2 Shape symbols
`T_q` = query length this launch (decode: 1; prefill/chunk: chunk width). `T_q` = query length this launch (decode: 1; prefill/chunk: chunk width).
`S` = total context length. `S_pe` = keys owned by this PE ≈ `⌈S/N⌉`. `S` = total context length. `R = C·P` = ranks per CUBE Group;
Storage `bf16` (numpy proxy `f16`, `memory_store.py:16`); `m,,O` `S_rank` = keys owned by this rank ≈ `⌈S/R⌉`. Storage `bf16` (numpy proxy
accumulators `f32` where precision matters. `f16`, `memory_store.py:16`); `m,,O` accumulators `f32` where precision
matters.
### 0.5.3 INPUTS (per kernel launch) ### 0.5.3 INPUTS (per kernel launch)
| Input | Shape (per KV head) | Location | Notes | | Input | Shape (per KV head) | Location | Notes |
|---|---|---|---| |---|---|---|---|
| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. | | `Q` | `[G, T_q, d]` | per-rank HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. |
| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. Dense locally, strided in global index (§2.1). | | `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). |
| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. | | `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 `S_pe`, slot↔global, causal bounds. | | `global_token_counter` | scalar | launch arg | kernel derives local len, slot↔global, causal bounds. |
| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 rotation. | | `start_cube`, `start_pe` (= `f(request_id)`) | scalars | launch arg | Level-1/Level-2 rotation. |
| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. | | `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. | | `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
| `O_base` | address | launch arg | where final O is written (root PE only). | | `O_base` | address | launch arg | where final O is written (CUBE-Group root only). |
| `softmax_scale` | scalar | launch arg | `1/√d`. | | `softmax_scale` | scalar | launch arg | `1/√d`. |
For **Ring Attention (§5.5)** add per ring step: incoming `K_block, For **Ring Attention (§5.5)** add per ring step: incoming `K_block,
@@ -251,14 +284,15 @@ V_block` via IPCQ into ping-pong buffers (post-RoPE), plus
| Output | Shape | Location | Notes | | Output | Shape | Location | Notes |
|---|---|---|---| |---|---|---|---|
| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE only** | normalised `O_acc / _acc`, cast to storage dtype. | | `O` (final) | `[G, T_q, d]` | `O_base`, **CUBE-Group root only** | normalised `O_acc / _acc`, cast to storage dtype. |
**Intermediate (non-root PE, on IPCQ — not a kernel-visible output):** **Intermediate (non-root rank, on IPCQ — not a kernel-visible output):**
`(m_i, _i, O_i)` (`m,`: `[G, T_q]`; `O_i`: `[G, T_q, d]` unnormalised) `(m_i, _i, O_i)` (`m,`: `[G, T_q]`; `O_i`: `[G, T_q, d]` unnormalised)
pushed to `tree_parent` via `tl.send` (§4). `O_i` is the heavy part. 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 part.
**No-SP (`N=1`):** no IPCQ partials; the single PE's running `(m,,O)` is **No-SP (`C=P=1`):** no IPCQ partials; the single rank's running `(m,,O)`
normalised in place and written to `O_base`. is normalised in place and written to `O_base`.
**Explicitly NOT outputs:** KV-cache writes (upstream, P3), output **Explicitly NOT outputs:** KV-cache writes (upstream, P3), output
projection (downstream), score `S` and probs `P` (never materialised). projection (downstream), score `S` and probs `P` (never materialised).
@@ -278,7 +312,7 @@ latency model:
- **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is - **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is
non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`):
the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole
per-PE tile sweep) and the scheduler generates the tile plan and streams per-rank tile sweep) and the scheduler generates the tile plan and streams
DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the
scheduler streams from HBM, so per-tile **K/V prefetch is the scheduler streams from HBM, so per-tile **K/V prefetch is the
scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed
@@ -348,16 +382,24 @@ permits (`kernel_runner.py`, ADR-0020 D3).
## 2. Memory layout & driver responsibilities ## 2. Memory layout & driver responsibilities
### 2.1 KV cache allocation ### 2.1 KV cache allocation (2-level SP)
- Per-PE KV buffers sized to per-PE max context `⌈max_context / N⌉ × d ×
dtype`, for K and V separately. In kernbench these are deployed with a One KV head's sequence is sharded across `C × P` ranks (Level-1 inter-CUBE
`DPPolicy` whose `pe` (or `cube`) axis is `row_wise` over the sequence × Level-2 intra-CUBE PE, §0). In kernbench this is a `DPPolicy` whose
dimension (`policy/placement/dp.py`; the baseline uses **both** the `cube` axis (`row_wise` over `C`) **and** the `pe` axis
`DPPolicy(pe="row_wise")` for K/V and `replicate` for Q). (`row_wise` over `P`) shard the sequence dimension, with Q `replicate`
- Within a PE, assigned tokens are **appended contiguously** (slot (`policy/placement/dp.py`). The baseline's single-axis
0,1,2,…). Global indices are strided (`i, i+N, …`) but the per-PE `DPPolicy(pe="row_wise")` becomes a two-axis `row_wise` over (cube, pe).
buffer is dense ⇒ DMA stays contiguous.
- Slot → global: `global_idx = local_slot·N + ((pe_id start_pe) mod N)`. - 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`
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`.
### 2.2 Driver per-launch duties (minimal) ### 2.2 Driver per-launch duties (minimal)
The driver supplies, per launch, the bases + counter + rotation; the The driver supplies, per launch, the bases + counter + rotation; the
@@ -365,23 +407,26 @@ kernel derives the rest:
| Launch arg | Purpose | | Launch arg | Purpose |
|---|---| |---|---|
| `K_base[pe]`, `V_base[pe]` | per-PE KV buffer bases (from tensor VA) | | `K_base[rank]`, `V_base[rank]` | per-rank KV buffer bases (from tensor VA) |
| `O_base` | attention output destination | | `O_base` | attention output destination |
| `global_token_counter` | current sequence position | | `global_token_counter` | current sequence position |
| `start_pe = request_id mod N` | Layer-2 rotation | | `start_cube`, `start_pe = f(request_id)` | Level-1/Level-2 rotation |
| `pe_id` (`tl.program_id`), `N` | geometry | | `cube_id`, `pe_id` (`tl.program_id` 1/0), `C`, `P` | CUBE-Group geometry |
| `q_block_meta` | prefill/SP query block start + length | | `q_block_meta` | prefill/SP query block start + length |
Kernel-derived (plain arithmetic): Let `R = C·P`, `cube_local = (cube_id start_cube) mod C`,
- **my turn to write this step:** `(start_pe + counter) mod N == pe_id` `pe_local = (pe_id start_pe) mod P`, and this rank's index
- **my valid length:** `base = counter // N; rem = counter % N; `rank = cube_local·P + pe_local`. Kernel-derived (plain arithmetic):
my_len = base + (1 if ((pe_id start_pe) mod N) < rem else 0)` - **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)`
- **read range:** tiles `0 .. ⌈my_len / TILE⌉`. - **read range:** tiles `0 .. ⌈my_len / TILE⌉`.
- **causal bounds / per-tile skip:** from global positions of the query - **causal bounds / per-tile skip:** from global positions of the query
block vs each tile. block vs each tile.
Driver = bases + counter + rotation. The single formula Driver = bases + counter + rotation. The single formula
`(request_id + token_idx) mod N` is the entire placement policy. `(request_id + token_idx) mod (C·P)` over the two SP axes is the entire
placement policy.
--- ---
@@ -440,84 +485,124 @@ Notes:
--- ---
## 4. Reduction (KV-parallel / SP combine) — tree to root ## 4. Reduction (KV-parallel / SP combine) — 2-level reduce-to-root
After its tile sweep each PE holds `(m_i, _i, O_i)` (unnormalised). After its tile sweep each rank holds `(m_i, _i, O_i)` (unnormalised) over
Combine via the associative/commutative log-sum-exp merge (identical math its `1/(C·P)` shard. Combine via the associative/commutative log-sum-exp
to the baseline's fold, `_attention_mesh_mlo.py:117-122`): merge (identical math to the baseline's fold, `_attention_mesh_mlo.py:117-122`):
```python ```python
def merge(m_a, l_a, O_a, m_b, l_b, O_b): def merge(m_a, l_a, O_a, m_b, l_b, O_b):
m = tl.maximum(m_a, m_b) m = tl.maximum(m_a, m_b)
sa, sb = tl.exp(m_a - m), tl.exp(m_b - m) 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 return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb
# final: O = O_root / l_root # normalise once at the tree root # final: O = O_root / l_root # normalise once at the CUBE-Group root
``` ```
- **Timing:** exchange is **after P·V** (each PE has its final `O_i`). Because the merge is associative **and** commutative, the combine is
- **Topology:** mesh tree, depth `⌈log₂ N⌉`, over physical neighbours. **reduce-to-root** (attention needs `O` at one place — the rank that writes
For `N=8`: level-0 `(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2 `O_base` — not an all-reduce) and proceeds in **two levels**:
`(3↔7)`, root = `7`. Each tree pair must be a physical `N/S/E/W`
neighbour so `tl.send(dir)`/`tl.recv(dir)` use a real link
(tuning item §9; the SFR install that wires neighbours is
`configure_sfr_intracube_pe_ring`, used by the baseline).
- **Why a tree, not the baseline fan-out:** attention needs `O` at one
place (the PE that writes `O_base`). A tree is `⌈log₂ N⌉` steps
(3 for `N=8`) vs the baseline all-to-all's `N1` (7). For decode, where
the per-PE sweep is short, the reduction is the dominant cost, so this
is a real win. (The baseline's fan-out replicates the answer to all
ranks — unnecessary here.)
- **Payload:** `O_i` (`d`-vector) is heavy; `m,` are scalars per `(g,
T_q)`. Sent as separate `tl.send`s (the baseline sends the triplet as
three sends — same pattern).
Kernel structure (greenlet): ### 4.1 Level-2 — intra-CUBE (P PEs → CUBE root PE)
- A **reduce-to-root tree** over the `P` PEs of a CUBE, depth `⌈log₂ P⌉`
(3 for `P=8`), over the PE IPCQ `N/S/E/W` mesh
(`configure_sfr_intracube_pe_ring`). Each tree pair is a 1-hop physical
neighbour. Result lands on the CUBE's root PE.
### 4.2 Level-1 — intra-CUBE-Group (C CUBE roots → Group root)
- A **center-root bidirectional CUBE-mesh reduce** over the `C` CUBEs of
the group (PE-root-only, `program_id(axis=1)` = `cube_id`), over the
CUBE NOC 2D mesh (ADR-0017). This **adapts the proven inter-CUBE pattern
of `lrab_hierarchical_allreduce.py` (Phases 12: row-then-col converge at
the center CUBE)** — but **reduce-only** (drop its broadcast-back Phases
45; attention needs root-only) and with the **log-sum-exp `merge`**
replacing its plain `+`. Center-root halves the critical path vs a corner
root (`lrab_hierarchical_allreduce.py:113-116`). No inter-SIP phase — the
CUBE Group is intra-SIP (§0).
### 4.3 Data-driven overlap, no global barrier
- A rank sends up the tree the **instant its own local P·V finishes** — it
does **not** wait for sibling ranks; internal nodes `merge` children as
they arrive. So the reduction **overlaps the compute of slower ranks**
(causal prefill imbalance, batched decode tokens).
- The two levels **pipeline**: a CUBE whose Level-2 reduction is done feeds
Level-1 immediately, while other CUBEs are still doing Level-2 — **no
barrier between levels**. (A rank cannot send before its *own* local
sweep completes; sending pre-final partials would multiply the heavy `O`
payload and is not worth it.)
- **Why reduce-to-root, not the baseline fan-out / an all-reduce:** the
baseline replicates the answer to every rank (`R1` steps); attention
needs it once. Reduce-to-root is `⌈log₂ P⌉ + (center-mesh depth over C)`
— for decode (short sweep, reduction-dominated) this is the dominant win.
### 4.4 Configurable topology
Each level's collective is **selectable** (a tuning item, §9): the defaults
above (Level-2 tree, Level-1 center-root mesh) suit the latency-bound,
small-`O` decode case; a **ring** option (chunked `O`) suits bandwidth-bound
long-context prefill with large `T_q·d` payload. `C=1` degenerates to
Level-2 only (single-CUBE SP).
**Payload:** `O_i` (`[G·T_q, d]`) is heavy; `m,` are light. Sent as
separate `tl.send`s (baseline pattern).
Kernel structure (greenlet), both levels reuse `merge`:
```python ```python
# leaf / internal node: fold children that send to me, then forward up def hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base):
for child_dir in tree_children_dirs(pe_id, N): # ---- Level-2: PE tree within the CUBE → CUBE root PE ----
m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape) for child in tree_children_dirs(pe_id, P): # PE IPCQ N/S/E/W
O_c = tl.recv(child_dir, O.shape) m, l, O = merge(m, l, O, *recv_triplet(child))
m, l, O = merge(m, l, O, m_c, l_c, O_c) if not is_cube_root(pe_id, P):
if not is_root(pe_id): send_triplet(parent_dir(pe_id, P), m, l, O); return
tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O) # ---- 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: else:
tl.store(O_base, O / l) send_triplet(parent_dir_mesh(cube_id, C), m, l, O)
``` ```
`tl.recv` blocks until data arrives (`tl_context.py:446-499`); the merge The reduction tree/mesh is **static** (fixed at compile time), so the recv
order is fixed by the static tree, so no data-dependent `pop` is needed — order is data-independent — no data-dependent `pop`, plain blocking
which is why **no hardware/composite `pop`-as-dependency change is `tl.recv` suffices, and **no hardware/composite `pop`-as-dependency change
required** (§6). is required** (§6).
--- ---
## 5. Per-case kernels ## 5. Per-case kernels
All four cases share §3 (tile sweep) + §4 (reduction); they differ only All four cases share §3 (tile sweep) + §4 (reduction); they differ only
in `N`, query-block width, and which `tile_*` predicates fire. in the SP extents `C·P`, query-block width, and which `tile_*` predicates
fire.
### 5.1 Common skeleton ### 5.1 Common skeleton
```python ```python
def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
q_block, softmax_scale, *, tl): C, P, q_block, softmax_scale, *, tl):
pe_id = tl.program_id(axis=rank_axis) cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
my_len = valid_len(counter, start_pe, pe_id, N) kv = head_of_cube(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) n_tiles = ceil(my_len / TILE)
q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast) q_g = load_Q_group(q_ptr, kv) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
m, l, O = init_running() # persistent arena: -inf, 0, zeros 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) 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) if tile_all_future(j, q_block): # causal skip (kernel if)
continue continue
run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope
if N == 1: if C == 1 and P == 1:
tl.store(o_ptr, O / l) # no reduction tl.store(o_base(kv), O / l) # no reduction (single rank)
else: else:
tree_reduce_and_store(pe_id, N, o_ptr) # §4 hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv)) # §4
``` ```
### 5.2 DECODE, no SP (`N=1`, one PE owns the head's KV) ### 5.2 DECODE, no SP (`C=P=1`, one PE owns the head's KV)
- `T_q = 1`, attends to all past KV ⇒ no future tiles, mask only on the - `T_q = 1`, attends to all past KV ⇒ no future tiles, mask only on the
final ragged tile. final ragged tile.
- **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8` - **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8`
@@ -530,19 +615,25 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile
plan also makes the timing count all `G` rows' work correctly (a leading plan also makes the timing count all `G` rows' work correctly (a leading
batch axis would *not* be counted — see §8). batch axis would *not* be counted — see §8).
- If `S_pe` fits in scratch (small/medium context) this degenerates to a - If `S_rank` fits in scratch (small/medium context) this degenerates to a
**one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one **one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one
composite for P·V) — exactly the baseline `_attention_mesh_mlo` composite for P·V) — exactly the baseline `_attention_mesh_mlo`
`_partial_attention`, just GQA-batched and on the composite path. Tiling `_partial_attention`, just GQA-batched and on the composite path. Tiling
(§3) only kicks in when `S_pe` exceeds the scratch scope's tile budget. (§3) only kicks in when `S_rank` exceeds the scratch scope's tile budget.
### 5.3 DECODE, with SP / KV-parallel (`N=8`) ### 5.3 DECODE, with SP / KV-parallel (`C × P` ranks)
- One request's KV is round-robin across 8 PEs; each owns ≈`my_len`
tokens. Each PE GQA-batches its `G=8` query rows over its local tiles. - One request's KV is hierarchical-round-robin across the `C·P` ranks of
- `T_q=1` ⇒ short per-PE sweep ⇒ **reduction dominates** ⇒ the §4 tree the CUBE Group (Level-1 over `C` CUBEs, Level-2 over `P` PEs); each rank
(3 steps) is the structurally important part. Reduction latency is owns ≈`my_len` tokens and GQA-batches its `G=8` query rows over its local
hidden by other concurrent decode tokens when batched, or by the long tiles.
per-PE tile sweep for long single-stream context. - `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
batched, by the long per-rank sweep for long single-stream context, and
by the §4.3 level-pipelining. For short decode prefer a **small `C`**
(less inter-CUBE reduction); push `C` up only when context length demands
the extra KV-parallelism.
### 5.4 PREFILL, no SP ### 5.4 PREFILL, no SP
- Whole prompt resident; query is a block of `T_q` tokens (chunked). - Whole prompt resident; query is a block of `T_q` tokens (chunked).
@@ -554,17 +645,21 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
- GQA batches the group's query heads along the same axis as decode. - GQA batches the group's query heads along the same axis as decode.
### 5.5 PREFILL, with SP (Ring KV) ### 5.5 PREFILL, with SP (Ring KV)
- KV sharded across `N` PEs (the ring); each step delivers a peer's KV
block over IPCQ. Running `(m,,O)` is carried **across ring steps** - KV sharded across the `C·P` ranks in **ring order**; each step delivers a
(Python handles, exactly as `_attention_mesh_kv` does today). peer's KV block over IPCQ. Running `(m,,O)` is carried **across ring
steps** (Python handles, exactly as `_attention_mesh_kv` does today). The
ring is the long-context, bandwidth-bound case where §4.4's **ring**
collective option (not the tree) is the natural choice.
- IPCQ overlaps next step's KV receive with current step's compute via - IPCQ overlaps next step's KV receive with current step's compute via
`tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists). `tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists).
- **Causal ring skip:** if the incoming step's KV block is entirely - **Causal ring skip:** if the incoming step's KV block is entirely
*after* the local query block, skip its compute (don't recv-consume / *after* the local query block, skip its compute (don't recv-consume /
don't fold) — eliminates ≈half the steps for causal attention. don't fold) — eliminates ≈half the steps for causal attention.
- Receive buffers ping-pong (persistent arena, not recycled). - Receive buffers ping-pong (persistent arena, not recycled).
- After the ring, `(m,,O)` is final per query-owning PE; if KV-parallel - After the ring, `(m,,O)` is final per rank; remaining KV-parallel splits
splits remain, §4 tree-merges them, then normalise + store. combine via the §4 2-level reduce, then normalise + store at the Group
root.
The baseline `_attention_mesh_kv` already implements the ring fold; this The baseline `_attention_mesh_kv` already implements the ring fold; this
ADR adds GQA reuse, causal step-skip, and `recv_async` overlap. ADR adds GQA reuse, causal step-skip, and `recv_async` overlap.
@@ -575,11 +670,12 @@ ADR adds GQA reuse, causal step-skip, and `recv_async` overlap.
- The only place a HW/composite `pop`-as-dependency would help is a lone - The only place a HW/composite `pop`-as-dependency would help is a lone
reduction with no other work to hide a poll — i.e. batch=1 **and** short reduction with no other work to hide a poll — i.e. batch=1 **and** short
context. The target is **agentic = low batch, long context** ⇒ each PE context. The target is **agentic = low batch, long context** ⇒ each rank
has many KV tiles; the §4 tree's `tl.recv` blocks are covered by the has many KV tiles; the §4 reduce's `tl.recv` blocks are covered by the
sweep work of concurrent tokens / long context. sweep work of concurrent tokens / long context / the §4.3 level-pipeline.
- The §4 reduction uses a **static** tree, so the recv order is fixed at - The §4 reduction uses a **static** 2-level tree/mesh, so the recv order is
compile time — no data-dependent pop. `tl.recv` (blocking) suffices. fixed at compile time — no data-dependent pop. `tl.recv` (blocking)
suffices.
- **Decision: ship with the greenlet `tl.send`/`tl.recv` collective.** - **Decision: ship with the greenlet `tl.send`/`tl.recv` collective.**
Revisit a HW `pop` only if a short-context, single-stream, Revisit a HW `pop` only if a short-context, single-stream,
latency-critical target emerges. latency-critical target emerges.
@@ -592,7 +688,7 @@ ADR adds GQA reuse, causal step-skip, and `recv_async` overlap.
|---|---| |---|---|
| tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) | | tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) |
| address / offset / valid-length arithmetic | **kernel** (from counter) | | address / offset / valid-length arithmetic | **kernel** (from counter) |
| reduction scheduling, IPCQ send/recv ordering | **kernel** (static tree) | | reduction scheduling, IPCQ send/recv ordering | **kernel** (static 2-level tree/mesh) |
| Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER | | Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER |
| Q load, mask add, softmax math, running `(m,,O)` merge | **`tl` ops** on PE engines (kernel-issued) | | Q load, mask add, softmax math, running `(m,,O)` merge | **`tl` ops** on PE engines (kernel-issued) |
@@ -620,12 +716,15 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head,
non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and
streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new
command kind). command kind).
- Tree reduction to root (§4) replacing the baseline all-to-all fan-out - **2-level reduce-to-root** (§4) replacing the baseline all-to-all fan-out
— Level-2 PE tree (intra-CUBE) + Level-1 center-root CUBE-mesh reduce
(intra-CUBE-Group, adapting the `lrab_hierarchical_allreduce` inter-CUBE
pattern as reduce-only + log-sum-exp), data-driven and level-pipelined,
pure kernel control flow over `tl.send`/`tl.recv`. pure kernel control flow over `tl.send`/`tl.recv`.
- Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` + - Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` +
a mask tensor added with `+`. a mask tensor added with `+`.
- Round-robin KV placement / valid-length arithmetic (§2) — launch-arg - 2-level round-robin KV placement / valid-length arithmetic (§0, §2) —
arithmetic + `DPPolicy(pe="row_wise")`. launch-arg arithmetic + `DPPolicy` `row_wise` over both `cube` and `pe`.
**New primitives for efficiency / scale (each has a supporting ADR):** **New primitives for efficiency / scale (each has a supporting ADR):**
@@ -669,14 +768,17 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head,
1. **Composite tile-pipeline depth** — dominant for KV-load-bound; how far 1. **Composite tile-pipeline depth** — dominant for KV-load-bound; how far
ahead the kernel issues non-blocking composites before waiting, and the ahead the kernel issues non-blocking composites before waiting, and the
scheduler's per-tile streaming depth. scheduler's per-tile streaming depth.
2. **PE↔mesh-neighbour mapping for the reduction tree** — ensure each 2. **Reduction topology per level (§4.4)** — Level-2 (intra-CUBE PE) and
depth-`⌈log₂ N⌉` pair is a physical `N/S/E/W` neighbour; bad mapping Level-1 (intra-CUBE-Group CUBE-mesh) each selectable tree/ring/center-mesh;
adds hops. Verify against the SFR install. sweep tree (latency, small-`O` decode) vs ring (bandwidth, large-`O`
prefill). Also the CUBE-Group size `C` knob (small for decode, large for
long context). Ensure each tree pair is a 1-hop physical neighbour; verify
against the SFR install.
3. **TILE size** — balance scratch residency (`S/P tiles + O_acc + G-way 3. **TILE size** — balance scratch residency (`S/P tiles + O_acc + G-way
GQA`) against DMA efficiency; interacts with ADR-0063 and the GQA`) against DMA efficiency; interacts with ADR-0063 and the
scheduler's `TILE_M/K/N` (`pe_scheduler.py`). scheduler's `TILE_M/K/N` (`pe_scheduler.py`).
4. **Ring buffer ping-pong vs `recv_async` depth** in §5.5. 4. **Ring buffer ping-pong vs `recv_async` depth** in §5.5.
5. **One-shot vs tiled crossover for decode** (§5.2) — the `S_pe` 5. **One-shot vs tiled crossover for decode** (§5.2) — the `S_rank`
threshold where tiling beats a single composite. threshold where tiling beats a single composite.
6. **Per-op CPU issue cost (cost model)** — currently `dispatch_cycles=0` 6. **Per-op CPU issue cost (cost model)** — currently `dispatch_cycles=0`
(`pe_cpu.py`), so composite-vs-primitive issue overhead is invisible. (`pe_cpu.py`), so composite-vs-primitive issue overhead is invisible.
@@ -691,19 +793,19 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head,
| Case | KV placement | Inner loop | Reduction | Masking | | Case | KV placement | Inner loop | Reduction | Masking |
|---|---|---|---|---| |---|---|---|---|---|
| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only | | Decode, no SP | 1 rank (`C=P=1`), all KV | one-shot or tiled, GQA-batched | none | last tile only |
| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only | | 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, no SP | resident | tiled per q-block | none | skip future / triangular boundary |
| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary | | Prefill, SP (Ring) | ring over `C·P` ranks | per ring step (recv_async overlap) | running state + §4 2-level | causal step-skip + boundary |
**I/O per case** (full contract §0.5): **I/O per case** (full contract §0.5):
| Case | Inputs | Output | IPCQ partials | | Case | Inputs | Output | IPCQ partials |
|---|---|---|---| |---|---|---|---|
| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none | | Decode, 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-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,,O_i)` per PE → tree | | 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 |
| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none | | Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none |
| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,,O)` + tree | | 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 |
In all cases: KV-cache writes and RoPE happen **upstream**; output In all cases: KV-cache writes and RoPE happen **upstream**; output
projection **downstream**; score `S` and probs `P` are never materialised. projection **downstream**; score `S` and probs `P` are never materialised.
@@ -733,8 +835,12 @@ secondary check.
2. **GQA reuse:** with `h_q = G·h_kv`, the K/V `dma_read_count` is 2. **GQA reuse:** with `h_q = G·h_kv`, the K/V `dma_read_count` is
independent of `G` (one load per tile, reused across the group), while independent of `G` (one load per tile, reused across the group), while
GEMM work scales with `G`. Asserts the lever actually fires. GEMM work scales with `G`. Asserts the lever actually fires.
3. **Tree reduction step count:** decode-SP issues `⌈log₂ N` reduction 3. **2-level reduction step count:** decode-SP issues `⌈log₂ P` (Level-2,
rounds (not `N1`); op_log `ipcq_send`/`recv` counts match the tree. intra-CUBE) + center-mesh-depth-over-`C` (Level-1) reduction rounds — not
the baseline's `C·P 1` all-to-all; op_log `ipcq_send`/`recv` counts
match the static 2-level tree/mesh, and the result lands at exactly one
rank (CUBE-Group root). A separate assertion checks Level-1 sends use the
CUBE NOC and Level-2 the PE IPCQ.
4. **Causal skip:** prefill skips all `tile_all_future` tiles — GEMM 4. **Causal skip:** prefill skips all `tile_all_future` tiles — GEMM
count equals the lower-triangular tile count, not the full grid. count equals the lower-triangular tile count, not the full grid.
5. **Long context (ADR-0063):** a sweep at `S` that overflows 1 MiB 5. **Long context (ADR-0063):** a sweep at `S` that overflows 1 MiB
@@ -769,10 +875,10 @@ predicted default; revise on review.
implementation starts. implementation starts.
2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b = 2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b =
[d, TILE]`, but the KV cache stores K as `[S_pe, d]`. `tl.trans` is [d, TILE]`, but the KV cache stores K as `[S_rank, d]`. `tl.trans` is
metadata-only and `MemoryStore.read` reshapes, not transposes metadata-only and `MemoryStore.read` reshapes, not transposes
(`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial (`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial
data. **Recommend:** store K **pre-transposed** `[d, S_pe]` in the cache data. **Recommend:** store K **pre-transposed** `[d, S_rank]` in the cache
(the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))` (the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))`
a contiguous slice. Verify the upstream `qkv_rope` write layout supports a contiguous slice. Verify the upstream `qkv_rope` write layout supports
this, or add a real `tl.transpose` (heavier; deferred). this, or add a real `tl.transpose` (heavier; deferred).
@@ -815,3 +921,39 @@ predicted default; revise on review.
P·V should also be composites; the IPCQ `recv_async` overlap is P·V should also be composites; the IPCQ `recv_async` overlap is
orthogonal and stays. **Recommend:** apply the same hybrid shape to the orthogonal and stays. **Recommend:** apply the same hybrid shape to the
ring step during implementation; low risk, mirrors §3. ring step during implementation; low risk, mirrors §3.
### Items from the hierarchical CUBE-Group SP pivot (this revision)
1. **A 4-SIP topology config is needed for full scale.** The shipped
`topology.yaml` has `sips: count: 2`; the full `H_kv=8` model at `C=8`
needs **4 SIPs** (2 KV heads each, §0). **Recommend:** add a 4-SIP
topology config (16-CUBE 4×4 mesh per SIP, 8 PE/CUBE — already the per-SIP
shape) for the headline eval; validation-scale runs can use fewer
CUBEs/SIPs. The bench must remain single-device per launch (one SIP).
2. **`head_of_cube` / CUBE-Group partition of the 4×4 mesh.** With `C=8`,
each SIP's 16-CUBE mesh splits into 2 CUBE Groups of 8. The exact sub-mesh
shape (4×2 vs 2×4) sets the Level-1 center-root mesh hop counts and which
CUBEs are 1-hop neighbours. **Recommend:** pick the partition that keeps
each CUBE Group a contiguous rectangular sub-mesh with a true center CUBE
(so `lrab`-style center-root applies); verify against the CUBE NOC routing
(ADR-0017).
3. **`C` as a per-operating-point knob.** `C` should differ by case (small
for decode where reduction dominates, large for long-context prefill).
**Recommend:** expose `C` (and per-level topology, §4.4) as launch/bench
config, sweep it in the milestone bench, and report the decode-vs-prefill
optimum. Default `C` TBD by measurement.
4. **Reverses the §0 "never spans CUBEs" non-goal — confirm no downstream
assumption depends on it.** Earlier text and possibly other ADRs treated
intra-CUBE-only reduction as invariant. **Recommend:** grep for that
assumption (SFR installs, address policy, diagrams) before implementation;
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.