Files

T

ywkang 973a452b77 gqa(adr): sync ADR-0060 KO mirror to current EN (two-kernel design)

Regenerate the KO mirror to match the current EN: two SP kernels
(decode=reduce / prefill=ring), TL;DR full code for the 3 decode variants
+ prefill, 'Q replicated / M-fold' and '1 Q head per CUBE' terminology,
contiguous shared KV layout, opt2 tl.wait, and all SB items. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

2026-06-04 23:02:49 -07:00

64 KiB

Raw Blame History

ADR-0060: AHBM GQA Fused Attention 커널 (Llama3-70B)

Status

Proposed

Context model: Llama3-70B. Decision drivers: agentic workload → 낮은 batch, 긴 context; KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV).

Supersedes / extends: 기존 mesh-native 어텐션 커널 _attention_mesh_kv(prefill)와 _attention_mesh_mlo(decode), 그리고 milestone-gqa-llama70b eval bench. §A. 기존 kernbench 작업과의 관계 참조 — 그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는 baseline이다.

Supporting ADRs (efficiency / scale enabler — GQA blocker 아님; §8 정정 참조): ADR-0063 tl.scratch_scope(per-tile scratch 재활용 — 현실적 context 길이에 필요), ADR-0062 lazy tl.load(첫 사용 시점 auto-wait를 갖는 non-blocking load → load/compute 오버랩), ADR-0061 tl.broadcast(선택적 mask/범용 편의). 두 GEMM(Q·Kᵀ, P·V)은 scheduler가 관리하는 tl.composite 커맨드로 발행된다(기존 CompositeCmd; 새 커맨드 종류 없음); 진짜 GQA 자체는 커널 재구조화만 필요(§5.2).

Algorithm lineage. 이 커널은 FlashAttention(tiling + online/streaming softmax, P·V fused — full score matrix 미생성)이다. §4의 KV-parallel split-and-combine는 FlashDecoding(split-KV + log-sum-exp merge). §5.5의 Ring 경로는 Ring Attention(KV 블록을 mesh 주위로 회전, 같은 online softmax로 fold). 새 수학은 도입하지 않으며, 본 ADR은 이 알려진 알고리즘을 kernbench의 greenlet tl 프로그래밍 모델(ADR-0020, ADR-0046)과 IPCQ PE↔PE collective(ADR-0023/0025)에 매핑한다.

TL;DR — 두 개의 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):

# 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

# OPTION 3 — software pipelining (current primitives; bubble removed) ← SHIP THIS
def gqa_decode_v3(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv  = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d))
    m, l, O = init_running(G, d); n = ceil(S_kv_local / TILE)
    Sb = double_buffer()                               # 2 persistent Sj buffers (outside scratch_scope)
    h  = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(0), (d, TILE)), out=Sb[0], epi=[scale])   # prime
    for j in range(n):
        Sj = Sb[j % 2]
        if j + 1 < n:                                  # ← issue NEXT Q·Kᵀ before softmax → fills GEMM engine
            h = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j+1), (d, TILE)), out=Sb[(j+1)%2], epi=[scale])
        with tl.scratch_scope():
            m2 = tl.maximum(m, tl.max(Sj, -1)); P = tl.exp(Sj - m2); corr = tl.exp(m - m2)
            l  = l * corr + tl.sum(P, -1)
            O  = O * corr + tl.composite("gemm", a=P, b=tl.ref(v_tile(j), (TILE, d))); m = m2
    hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))

# OPTION 2 — ex_composite, 2-split (NEW flash-epilogue cmd; gives K-before-V DMA priority)
def gqa_decode_v2(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv  = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d))
    acc = init_running(G, d)                           # (m,l,O): scheduler-updated flash accumulator
    for j in range(ceil(S_kv_local / TILE)):
        Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale])   # #1: reads K (priority)
        tl.ex_composite("softmax_pv", s=Sj, v=tl.ref(v_tile(j), (TILE, d)), acc=acc, scale=scale)  # #2: reads V
        # ↑ both non-blocking; CPU never waits intra-tile → max run-ahead, fewest issues
    tl.wait()                                          # ← drain ALL composites: acc is updated async (no auto-wait,
                                                       #   #2 is a serial chain through acc); only final after the last #2
    hierarchical_reduce_and_store(*acc, cube_id, pe_id, C, P, o_base(kv))

	new HW cmd	GEMM bubble	CPU intra-tile wait	issues	now
opt1 current	no	yes	yes	`O(tiles·ops)`	✓
opt3 sw-pipe	no	no	yes (reordered)	`O(tiles·ops)`	✓
opt2 ex_composite	#2 only	no	no	`O(tiles)`	✗ (build #2)

(#1 = 기존 GEMM composite + scale; #2 = softmax + P·V + stateful online-softmax accumulator만 신규. MATH 엔진엔 max/sum/exp가 이미 있음 — 신규는 flash accumulator이지 ops가 아님.)

Kernel 2 — PREFILL + SP (CUBE당 Q head 1개, head-parallel; Ring KV, reduce 없음)

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는 오늘날 이미 IPCQ 상에서 online-softmax (m, ℓ, O) 머지로 FlashAttention을 돌린다. 두 커널이 존재한다:

File	Role	Mechanism
`src/kernbench/benches/_attention_mesh_kv.py`	prefill (Ring K/V)	per-rank partial attention, bidirectional K/V fan-out, online-softmax fold
`src/kernbench/benches/_attention_mesh_mlo.py`	decode (split-KV)	per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge

둘 다 milestone-gqa-llama70b(4 패널: {single,multi}_user × {prefill,decode}, src/kernbench/benches/milestone_gqa_llama70b.py)이 구동하며 tests/attention/test_milestone_gqa_llama70b.py에서 테스트된다.

이들은 greenlet tl API로 작성되었다: tl.load, tl.dot, tl.softmax/tl.max/tl.sum/tl.exp, tl.send/tl.recv, 그리고 TensorHandle에 대한 Python -/*//(각각 MathCmd emit) — GEMM은 composite가 아니라 **blocking tl.dot**로. 실행 중 (m, ℓ, O)는 루프를 관통하는 Python TensorHandle일 뿐이다. 본 ADR은 실행 상태와 softmax 머지는 커널에 유지하되 두 GEMM은 scheduler가 관리하는 tl.composite 경로로 옮긴다(§1 참조) — 이것이 설계에 중요하다.

baseline의 의도된 세 한계 — 효율적 GQA 커널이 정확히 들어내야 하는 것:

GQA 재사용 없음. h_q == h_kv == 1 (test_milestone_gqa_llama70b.py:137-142). 테스트는 이를 broadcast view에 대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의 head-packing 핵의 속성이다(_view(K, (h_q·d, S_kv))는 모든 head를 하나의 matmul 차원에 뭉치고 h_q == h_kv일 때만 byte를 보존). 올바른 수정은 broadcast op이 아니라 커널 재구조화다: 한 번에 한 KV head를 처리하고 G group 행을 matmul M 차원에 fold(§5.2). 이는 byte-보존 reshape만 쓰므로 진짜 GQA(h_q = G·h_kv)가 신규 primitive 없이 돈다 — §8 참조.
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 단계.
검증 스케일만. S = 16인 이유는 1 MiB scratch bump allocator가 per-tile 임시값을 누수하고(test_milestone_gqa_llama70b.py:123-148) causal masking / tiling이 없기 때문 → ADR-0063(재활용) + §5(tiling, causal skip) + composite K/V 스트리밍(§3) + ADR-0062(lazy load 오버랩)으로 해결.

문서 부채(범위 밖이나 기록): baseline은 ADR-0055/0056/0057/0058/0059를 인용하나 파일로 존재하지 않는다 — ghost 참조다. 본 ADR은 이를 소급 작성하지 않으며; 권고는 Detailed Design Document의 Open Decisions 참조.

0. 참조 차원 (Llama3-70B)

Symbol	Meaning	Value
`H_q`	query heads	64
`H_kv`	KV heads	8
`G`	GQA group size = `H_q / H_kv`	8
`d`	head dim	128
`L`	layers	80
`D`	model dim	8192

하드웨어 요약:

AHBM (chip) = CUBE(메모리 cube, 각각 PE를 담은 logic die) 집합 + IO die(ADR-0003).
IPCQ: PE↔PE 큐, PE당 4 mesh-방향 queue-pair (N/S/E/W, ADR-0023 D3; inter-SIP용 global_*, ADR-0032). 커널 API: tl.send(dir, src) / tl.recv(dir, shape, dtype) (tl_context.py:402-499).
Composite command (CompositeCmd, pe_commands.py:144-162): 단일 GEMM(또는 MATH) head + element-wise epilogue 단계 (bias/relu/scale/add/...), PE_SCHEDULER에 non-blocking으로 발행되며, scheduler가 tile plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다 (ADR-0014 D6; pe_scheduler.py:104-143). 이는 일반적 multi-op DAG가 아니다: 두 GEMM을 chain할 수 없고, 인스턴스 간 register 상태를 못 들며, IPCQ를 pop/wait 못 한다. 따라서 본 ADR은 두 어텐션 GEMM(Q·Kᵀ, P·V)을 각각 자체 composite로 발행하고 cross-GEMM softmax 머지 + IPCQ reduction은 커널에 유지한다 — 새 "flash-composite" 커맨드 종류는 필요 없다(§1, §8 참조).

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 내부다. 그 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 내, 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에 머문다).

0.5 커널 경계, 전제조건, I/O 계약

0.5.1 디코더 레이어 내 위치

1. RMSNorm
2. QKV projection (GEMM)        ─┐  qkv_rope kernel (SEPARATE, upstream)
3. RoPE on Q and K              ─┤
4. write new K,V → KV cache     ─┘
5. ===== THIS KERNEL: FlashAttention =====   (post-RoPE Q, K-cache, V-cache → O)
6. Output projection (GEMM)        out_proj kernel (SEPARATE, downstream)
7. residual add → FFN ...

전제조건 (upstream qkv_rope, 본 커널 아님):

P1. Q는 이미 RoPE-회전됨. 여기서 회전 없음.
P2. K-cache는 post-RoPE K 저장(어텐션 시점에 재회전 안 함 — post-RoPE 캐싱의 이유).
P3. decode의 경우, 새 step의 K 행은 RoPE-회전되어 그 소유 PE의 K-cache 슬롯에 본 커널 launch 전에 qkv_rope가 추가한다. ⇒ 본 커널은 KV cache에 대해 pure read.
P4. V는 회전 안 함; V-cache는 raw projected V 보유.
P5. upstream RoPE 위치는 token의 절대 global 위치 global_idx = local_slot·(C·P) + rank (여기서 rank = cube_local·P + pe_local, §2.1). Round-robin placement ≠ 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. 로컬 dense, global index는 `C·P`로 strided (§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 없음.

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 이점 없음"이라는 논거로 순수 greenlet primitive 경로(전부 tl.dot, composite 없음)를 제안했다. 그것은 오직 시뮬레이터가 현재 per-op CPU 발행 비용을 0으로 청구하기 때문에만(dispatch_cycles=0, pe_cpu.py) 성립한다 — descriptor offload가 숨기려 존재하는 바로 그 CPU issue-rate / DMA-program 비용을 모델이 지워버린 것이다. 하이브리드가 효율 커널의 충실한 표현이다. 이점의 측정 가능한 크기(CPU가 많은 타일에 대해 엔진을 saturate할 수 있는가?)는 op 종류별 발행 비용 모델링에 좌우되며, future work로 추적된다(cost model; §9). dispatch_cycles=0에서도 non-blocking composite 경로는 blocking tl.dot 경로가 남기는 GEMM 엔진 버블을 채운다.

이것이 효율적 선택인 이유. GEMM tiling + DMA 스트리밍 + cross-tile 파이프라이닝은 검증된 CompositeCmd scheduler 경로에 offload되고; softmax 머지와 검증된 IPCQ collective는 커널에 남는다. 진짜 새 기계장치는 작고 범용적인 두 primitive뿐이다(ADR-0062 lazy tl.load, ADR-0063 tl.scratch_scope); reduction은 tl.send/tl.recv를 재사용한다. 전용 "flash-composite" 커맨드(softmax 머지 + carried register 상태 + IPCQ-push epilogue를 내재화하는 한 종류)는 만들지 않는다 — 크고 특수목적이며, 하이브리드 대비 유일한 delta(완전 softmax offload)가 현재 모델링 충실도에서 정당화되지 않는다; §8 참조.

2. 메모리 레이아웃 & 드라이버 책임

2.1 KV cache 할당 (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 안에서 할당된 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.

prefill/decode 공유 레이아웃 — round-robin이 아니라 contiguous 블록. prefill이 KV cache를 쓰고(qkv_rope, upstream) decode가 같은 cache를 읽고 연장하므로, 두 커널은 하나의 물리 레이아웃을 공유한다 — prefill→decode 경계에 reshard 없음. prefill의 causal skip은 연속 위치 블록을 요구한다(늦은 연속 블록은 통째로 미래라 skip 가능; round-robin 블록은 전 위치에 걸쳐 결코 skip 불가). 그래서 공유 레이아웃은 각 KV head를 contiguous C × P 위치 블록으로 샤딩한다(rank r이 [r·B, (r+1)·B) 소유, B = ⌈max_context/(C·P)⌉). Decode는 자기 블록을 읽고 reduce; prefill은 C CUBE-레벨 블록을 ring. Caveat: contiguous는 짧은 context에서 rank를 덜 씀(frontier rank만 데이터 보유) — long-context 타깃에선 수용 가능; 짧은-context 균형은 별도 연구(§B). (위 round-robin 공식은 decode-only 균형용 대안; contiguous 블록 형태가 두 커널을 다 서비스.)

2.2 드라이버 launch당 의무 (최소)

드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를 도출한다:

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, cube_local = (cube_id − start_cube) mod C, pe_local = (pe_id − start_pe) mod P, 이 rank의 인덱스 rank = cube_local·P + pe_local라 하자. 커널 도출(평범한 산술):

이번 step에 내가 쓸 차례: (start_rank + counter) mod R == rank
내 valid 길이: base = counter // R; rem = counter % R; my_len = base + (1 if rank < rem else 0)
read 범위: tile 0 .. ⌈my_len / TILE⌉.
causal bound / per-tile skip: query block과 각 tile의 global 위치에서.

드라이버 = base + counter + rotation. 두 SP 축에 대한 단일 공식 (request_id + token_idx) mod (C·P)이 placement policy 전부다.

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) 사용:

# 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와 동일 수학, _attention_mesh_mlo.py:117-122):

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 재사용:

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)

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 _attention_mesh_mlo의 _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에 계층적-round-robin(Level-1은 C CUBE, Level-2는 P PE); 각 rank가 ≈my_len 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를 필요로 하므로.

Head-parallel 배치: CUBE Group 내에서 CUBE i가 query head i 하나만 소유(G query head → C=G CUBE, CUBE당 Q head 1개)하고 KV slice i. 각 CUBE가 자기 Q head 하나의 full attention 계산. 각 CUBE가 서로 다른 head를 생산하므로 (m,ℓ,O) reduce 없음 — 각 CUBE가 자기 head의 행을 정규화·기록.
Ring KV: C KV slice가 CUBE ring을 회전; 각 CUBE가 들어오는 블록을 자기 head의 실행 중 (m,ℓ,O)에 fold(online-softmax, ring step 가로질러 Python 핸들로 carry, _attention_mesh_kv가 오늘 하듯). C step 후 모든 head가 전 KV를 봄. GQA 재사용은 회전에서 — slice j가 전 C CUBE를 방문하며 전 G head를 서비스.
IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 tl.recv_async/tl.wait로 오버랩(tl_context.py:543-560 — 이미 존재); 수신 버퍼는 ping-pong(persistent arena, 재활용 안 함).
Causal ring skip: 들어오는 KV 블록이 전부 이 head의 query block 이후면 그 compute를 skip — 약 절반 step 제거.
CUBE 내(P PE): head의 query 행 [T_q, d] 및/또는 현재 KV 블록을 P PE에 타일(여기엔 decode와 달리 query 축이 존재); 상세는 §B.

baseline _attention_mesh_kv가 이미 ring fold를 구현; 본 ADR은 GQA 재사용, head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가.

5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2)

세 decode 변형은 TL;DR에 풀 코드로 있다(Kernel 1): opt1 현재 CompositeCmd(CPU가 Sj를 auto-wait하는 동안 GEMM 엔진 버블), opt3 software pipelining(다음 타일 Q·Kᵀ를 이번 타일 softmax 전에 발행, Sj는 persistent 이중 버퍼 — 버블 제거, 새 커맨드 종류 없음), opt2 ex_composite(분할: #1 Q·Kᵀ = K를 먼저 읽는 기존 composite, #2 softmax+P·V+accumulator = V를 나중에 읽는 유일한 신규 flash-epilogue 기계장치).

권고: 지금은 opt3 출시(새 기계장치 0, 버블 제거); cost model(ADR-0064)이 적은-CPU-발행 이득을 측정 가능하게 만들면 opt2 재검토(§8 항목 4). MATH 엔진엔 max/sum/exp가 이미 있음 — opt2의 유일한 진짜 신규는 composite의 stateful (m,ℓ,O) accumulator이지 ops가 아님.

이 변형들은 decode 관심사다: prefill(§5.5)에선 causal if가 composite에 못 들어가는 커널 제어흐름이고, ring은 이미 recv_async로 오버랩한다.

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 round-robin KV placement / valid-length 산술(§0, §2) — launch-arg 산술 + cube와 pe 양쪽 row_wise인 DPPolicy.

efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):

per-tile scratch 재활용 — ADR-0063(tl.scratch_scope). 스케일에 필요: S=16 상한(1 MiB bump allocator)을 제거해 현실적 context 길이가 돌게 함. 셋 중 최고 가치.
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).
GQA head / mask broadcast — ADR-0061(tl.broadcast). 선택적 편의, GQA blocker 아님(위 정정 참조). G·T_q 행에 걸친 additive-mask 구성과 범용 커널에 유용; np.matmul이 data 모드에서 이미 broadcast하므로 정확성엔 불필요. 최저 우선순위.

명시적 REJECTED (효율적 대안 선택):

~~내부 루프 전체를 내재화하는 전용 "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 없음) 출시.
~~하드웨어 pop-as-dependency.~~ 범위 밖(§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 아님)

Composite tile-pipeline depth — KV-load-bound에 지배적; 커널이 wait 전에 non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당 스트리밍 depth.
레벨별 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 대조 검증.
TILE 크기 — scratch 상주(S/P tile + O_acc + G-way GQA)와 DMA 효율 균형; ADR-0063 및 scheduler의 TILE_M/K/N(pe_scheduler.py)과 상호작용.
§5.5의 Ring 버퍼 ping-pong vs recv_async depth.
decode의 one-shot vs tiled 교차점(§5.2) — tiling이 단일 composite를 이기는 S_rank 임계값.
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차 체크.

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에 좌우됨(범위 밖; 플래그됨).
GQA 재사용: h_q = G·h_kv에서 K/V dma_read_count가 G와 무관(타일당 한 load, group에 걸쳐 재사용), 반면 GEMM 작업은 G로 스케일. 레버가 실제 발동함을 assert.
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에 소비됨으로.
Causal skip: prefill이 모든 tile_all_future tile(및 ring의 전체-미래 KV 블록)을 skip — GEMM 수가 full grid가 아니라 하삼각 tile/step 수와 일치.
Long context (ADR-0063): scope 없이 1 MiB를 넘기는 S sweep이 완료하고 reference와 일치.
Load/compute 오버랩: tiled sweep의 end-to-end latency가 직렬 Σ(load+compute)보다 낮음 — composite scheduler의 타일당 K/V 스트리밍(§1/§3)과 Q load를 오버랩하는 lazy tl.load(ADR-0062)에서. (오버랩은 실제 모델 동시성, 빼기가 아님.)
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 파이프라이닝).
결정성: 동일 입력 → 동일 op_log + latency (SPEC §0.1).

B. 하이브리드 전환에서 나온 Open 설계 항목 (추후 검토)

이들은 결정이 순수 greenlet primitive 경로에서 composite hybrid + lazy tl.load(이번 개정)로 옮겨갈 때 생겼다. 설계를 막는 것은 없으며; 각각 구현 중 검증 패스가 필요하다. (묻지 않고) 작업 합의에 따라 여기 기록한다 — 권고는 내가 예측한 기본값이며; 검토 시 수정.

DDD-0060이 아직 미동기화. Detailed Design Document는 여전히 옛 tl.load_async 더블버퍼 경로와 primitive tl.dot 내부 루프를 기술한다(그 §4.3/§5/§10). 하이브리드(composite GEMM, lazy load, K 사전 transpose)로 갱신해야 한다. DDD가 파생 how-to이고 큰 재작성이라 검토용으로 남김; ADR이 이제 권위 있는 기록이다. 권고: 구현 시작 전 후속으로 DDD 동기화.
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 추가(더 무거움; 연기).
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이 없음을 검증.
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에만 — 여전히 정확하나 헤드라인이 바뀜.
커널 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).
cost model은 별도 ADR. 하이브리드의 CPU-saturation 이점은 dispatch_cycles=0인 동안 보이지 않는다. op 종류별 issue-cost 모델은 ADR-0064에 명세; 본 ADR의 §1/§9는 측정 가능한(구조적뿐 아닌) 이점을 위해 그것에 의존. 권고: eval에서 하이브리드 latency 이점을 주장하기 전에 ADR-0064 모델을 도입.
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 전환에서 나온 항목 (이번 개정)

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) 유지해야 함.
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) 대조 검증.
operating-point별 노브로서의 C. C는 case별로 달라야 함(reduction 지배인 decode엔 작게, long-context prefill엔 크게). 권고: C(및 레벨별 토폴로지, §4.4)를 launch/bench config로 노출, milestone 벤치에서 sweep하고 decode-vs-prefill 최적값을 보고. 기본 C는 측정으로 TBD.
§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뿐 아니라 노출해야 함.
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 연속).

decode-reduce / prefill-ring 분리에서 나온 항목 (이번 개정)

두 커널, 두 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를 대신 이동. 권고: 두 커널로 유지; 재병합 금지.
출력 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).
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.
prefill의 C = G 결합. head-parallel 매핑은 C = G = 8(CUBE당 Q head 1개)을 가정. C ≠ G면 매핑 재검토 필요(CUBE당 다중 head, 또는 head가 부분 ring에 걸침). 권고: headline 스케일에서 prefill 커널은 C = G 고정; C ≠ G는 별도 연구.
_attention_mesh_mlo_2d(현재 impl)와 정합. 원격 impl의 2D 커널은 Q replicated로 cube에 걸친 AllReduce — 즉 decode-reduce 계열이나 reduce-to-root 아닌 all-reduce(broadcast-back), 그리고 prefill head-parallel ring은 아직 없음. 권고: (a) 그 2D AllReduce → reduce-to-root(broadcast-back 제거)로 decode 커널; (b) §5.5 head-parallel Ring-KV 커널을 prefill용으로 추가.
공유 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 균형은 별도 연구).
출시할 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이 적은-발행 이득을 측정 가능하게 만들 때까지 연기.

64 KiB Raw Blame History Unescape Escape