kernbench2/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md

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

Status

Proposed

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

Supersedes / extends: the existing mesh-native attention kernels _attention_mesh_kv (prefill) and _attention_mesh_mlo (decode) and the milestone-gqa-llama70b eval bench. See §A. Relationship to existing kernbench work — that code is the baseline this ADR upgrades to a real GQA, causal, long-context kernel.

Supporting ADRs (efficiency / scale enablers — not GQA blockers; see §8 correction): ADR-0063 tl.scratch_scope (per-tile scratch recycling — required for realistic context length), ADR-0062 lazy tl.load (non-blocking load with auto-wait on first use → load/compute overlap), ADR-0061 tl.broadcast (optional mask/general convenience). The two GEMMs (Q·Kᵀ, P·V) are issued as scheduler-managed tl.composite commands (the existing CompositeCmd; no new command kind); real GQA itself needs only kernel restructuring (§5.2).

Algorithm lineage. This kernel is FlashAttention (tiling + online/streaming softmax with fused P·V — no full score matrix materialised). The KV-parallel split-and-combine in §4 is FlashDecoding (split-KV with log-sum-exp merge). The Ring path in §5.5 is Ring Attention (KV blocks rotated around the mesh, folded by the same online softmax). No new math is introduced; this ADR maps those known algorithms onto the kernbench greenlet tl programming model (ADR-0020, ADR-0046) and the IPCQ PE↔PE collective (ADR-0023/0025).

---

TL;DR — two SP kernels (decode = reduce, prefill = ring)

Decode-SP and prefill-SP are structurally different and are two kernels — one kernel cannot do both. The principle is move the smaller thing:

• Decode (T_q=1): the output O = [G, d] is tiny, the KV cache is big. → keep KV statically sharded & resident, do the local sweep, and move only the small (m,ℓ,O) via a 2-level reduce (§4). Q is replicated — all G query heads stacked into the GEMM M-dim (M-fold), so one Q·Kᵀ does all heads sharing one K. No KV movement.
• Prefill (T_q=S): the output O = [S, d] per head is big — reducing it would move [S,d] per rank. → instead make each CUBE own one Q head (head-parallel, C=G) and rotate the KV so each head sees all KV (Ring KV, §5.5). No (m,ℓ,O) reduce (each CUBE outputs a different head). Only KV blocks move.

Memory has room, so replicate Q/weights to avoid communicating them; only the irreducible data moves — (m,ℓ,O) (decode) or KV blocks (prefill). Both share §3's composite-hybrid inner tile (GEMMs → tl.composite, softmax merge in kernel, lazy tl.load).

Kernel 1 — DECODE + SP (head-replicated, KV static shard, 2-level reduce; NO ring)

The decode inner tile has a hard chain Q·Kᵀ → softmax → P·V: softmax (tl.max(Sj)) waits for the first GEMM, so a naïve loop stalls the CPU on Sj and leaves the GEMM engine idle during softmax (a bubble). Three variants trade CPU/HW complexity against that bubble (ship opt3 now; opt2 needs a new command kind — revisit with the 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
    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 = the existing GEMM composite + scale; only #2 = softmax + P·V + the stateful online-softmax accumulator is new. MATH engine already has max/sum/exp — the new part is the flash accumulator, not the ops.)

Kernel 2 — PREFILL + SP (1 Q head per CUBE, head-parallel; Ring KV, NO 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

Why two shapes. Decode keeps KV put and moves the tiny (m,ℓ,O) (§4 2-level reduce); prefill moves KV (Ring, §5.5) and keeps each head's big output local. Both use §3's composite-hybrid tile. K is pre-stored transposed to sidestep the reshape-not-transpose caveat (§3, §B); no bespoke "flash-composite" kind on the decode critical path (§8 item 4). Output head distribution differs — decode lands all G heads at the CUBE-Group root; prefill leaves one Q head per CUBE (§0.5.4). The opt2/opt3 variants are a decode concern: in prefill the causal if is kernel control flow that cannot enter a composite, and the ring already overlaps via recv_async.

---

A. Relationship to existing kernbench work (read first)

kernbench already runs FlashAttention with an online-softmax (m, ℓ, O) merge over IPCQ today. Two kernels exist:

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

Both are driven by milestone-gqa-llama70b (4 panels: {single,multi}_user × {prefill,decode}, src/kernbench/benches/milestone_gqa_llama70b.py) and tested in tests/attention/test_milestone_gqa_llama70b.py.

They are written in the greenlet tl API: tl.load, tl.dot, tl.softmax/tl.max/tl.sum/tl.exp, tl.send/tl.recv, and Python -/*// on TensorHandle (each emits a MathCmd) — the GEMMs as blocking tl.dot, not composites. The running (m, ℓ, O) is just Python TensorHandles threaded through the loop. This ADR keeps the running state and the softmax merge in the kernel but moves the two GEMMs onto the scheduler-managed tl.composite path (see §1) — this matters for the design.

Three deliberate limitations of the baseline — exactly what an efficient GQA kernel must lift:

1. No GQA reuse. h_q == h_kv == 1 (test_milestone_gqa_llama70b.py:137-142). The test attributes this to a MemoryStore byte-conservation failure on a broadcast view, but that failure is a property of the baseline's head-packing hack (_view(K, (h_q·d, S_kv)), which conflates all heads into one matmul dim and only conserves bytes when h_q == h_kv). The correct fix is kernel restructuring, not a broadcast op: process one KV head at a time and fold the G group rows into the matmul M dimension (§5.2). That uses only byte-conserving reshapes, so real GQA (h_q = G·h_kv) runs with no new primitive — see §8.
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 steps). Attention only needs O at the query owner → a 2-level 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 allocator leaks per-tile temporaries (test_milestone_gqa_llama70b.py:123-148) and there is no causal masking / tiling → fixed by ADR-0063 (recycling) + §5 (tiling, causal skip) + composite K/V streaming (§3) + ADR-0062 (lazy load overlap).

Documentation debt (out of scope but recorded): the baseline cites ADR-0055/0056/0057/0058/0059, none of which exist as files — they are ghost references. This ADR does not retro-write them; see the Detailed Design Document's Open Decisions for the recommendation.

---

0. Reference dimensions (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

Hardware recap:

• AHBM (chip) = set of CUBEs (memory cubes, each with a logic die containing PEs) + IO die (ADR-0003).
• IPCQ: PE↔PE queues, 4 mesh-direction queue-pairs per PE (N/S/E/W, ADR-0023 D3; global_* for inter-SIP, ADR-0032). Kernel API: tl.send(dir, src) / tl.recv(dir, shape, dtype) (tl_context.py:402-499).
• Composite command (CompositeCmd, pe_commands.py:144-162): a single GEMM (or MATH) head plus element-wise epilogue stages (bias/relu/scale/add/...), issued non-blocking to PE_SCHEDULER, which generates a tile plan and streams DMA→GEMM→write per tile (ADR-0014 D6; pe_scheduler.py:104-143). It is not a general multi-op DAG: it cannot chain two GEMMs, cannot carry register state across instances, and cannot pop/wait on IPCQ. This ADR therefore issues 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 — it does not need a new "flash-composite" command kind (see §1, §8).

CUBE Group — the placement unit (hierarchical SP)

A CUBE Group is C CUBEs (within one SIP) that jointly own one KV head and its G query heads. One KV head's KV sequence is sharded 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 one KV head maps to C × P ranks, all within one SIP. How the G query heads map onto those ranks differs by case (the two kernels, TL;DR / §5):

• Decode (§4): Q is replicated — every rank holds all G query heads, M-folded (stacked into the matmul M / row dimension: Q for the group [G, T_q, d] → [G·T_q, d], so one Q·Kᵀ GEMM computes all G heads while sharing the single K — the GQA reuse). KV is sequence-sharded C × P ways; outputs reduce.
• Prefill (§5.5): Q is head-parallel — with C = G, CUBE i owns exactly one query head i; KV rotates (Ring); no reduce.

Q is small, so replicating it (decode) or distributing it one-head-per-CUBE (prefill) is cheap — only the irreducible data moves (decode: (m,ℓ,O); prefill: KV blocks). C is a tuning knob: for prefill set C = G (one Q head per CUBE); for decode C trades inter-CUBE reduction against KV-parallel breadth (small C, even C = 1 single-CUBE, for short context 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).

---

0.5 Kernel boundary, preconditions, I/O contract

0.5.1 Position in the decoder layer

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 ...

Preconditions (upstream qkv_rope, NOT this kernel):

• P1. Q is already RoPE-rotated. No rotation here.
• P2. K-cache stores post-RoPE K (never re-rotated at attention time — the reason for post-RoPE caching).
• P3. For decode, the new step's K row is RoPE-rotated and appended to its owning PE's K-cache slot by qkv_rope before this kernel launches. ⇒ this kernel is pure read on the KV cache.
• P4. V is not rotated; V-cache holds raw projected V.
• P5. RoPE position upstream is the token's absolute global position global_idx = local_slot·(C·P) + rank where rank = cube_local·P + pe_local (§2.1). Round-robin placement ≠ RoPE position.

0.5.2 Shape symbols

T_q = query length this launch (decode: 1; prefill/chunk: chunk width). S = total context length. R = C·P = ranks per CUBE Group; S_rank = keys owned by this rank ≈ ⌈S/R⌉. Storage bf16 (numpy proxy f16, memory_store.py:16); m,ℓ,O accumulators f32 where precision matters.

0.5.3 INPUTS (per kernel launch)

Input	Shape (per KV head)	Location	Notes
Q	[G, T_q, d]	per-rank HBM (loaded to TCM)	post-RoPE (P1). G query rows of the group batched.
K_cache	[S/(C·P), d]	per-rank HBM, base K_base[rank], contiguous	post-RoPE (P2). Read-only. Dense locally, strided by C·P in global index (§2.1).
V_cache	[S/(C·P), d]	per-rank HBM, base V_base[rank]	raw V (P4). Read-only.
global_token_counter	scalar	launch arg	kernel derives local len, slot↔global, causal bounds.
start_cube, start_pe (= f(request_id))	scalars	launch arg	Level-1/Level-2 rotation.
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	where final O is written (CUBE-Group root only).
softmax_scale	scalar	launch arg	1/√d.

For Ring Attention (§5.5) add per ring step: incoming K_block, V_block via IPCQ into ping-pong buffers (post-RoPE), plus step_kv_global_range for the causal step-skip.

0.5.4 OUTPUTS

The output head distribution differs by kernel — downstream out-projection must consume them accordingly:

Kernel	Output	Location	Notes
Decode (§4 reduce)	O = [G, 1, d] (all heads)	O_base at the CUBE-Group root	head-replicated → all G heads end on one rank after the 2-level reduce.
Prefill (§5.5 ring)	O_i = [1, T_q, d] (one head each)	per-CUBE O_base[i], distributed	head-parallel → CUBE i writes head i's rows in place; no reduce.

Decode intermediate (non-root rank, on IPCQ — not kernel-visible): (m_i, ℓ_i, O_i) pushed up the 2-level tree to the parent (Level-2 PE parent, then Level-1 CUBE parent) via tl.send (§4). O_i is the heavy part. Prefill has no such partials (no reduce); instead KV blocks circulate (§5.5).

No-SP (C=P=1): no IPCQ partials; the single rank's running (m,ℓ,O) is normalised in place and written to O_base.

Explicitly NOT outputs: KV-cache writes (upstream, P3), output projection (downstream), score S and probs P (never materialised).

---

1. Decision (mechanism)

Implement the kernel as a hybrid: issue the two GEMMs (Q·Kᵀ and P·V) as scheduler-managed tl.composite(op="gemm") commands, and keep the online-softmax merge and the cross-PE reduction as kernel-level tl ops. tl.load is lazy (non-blocking; the wait is auto-inserted at first use of the loaded data — ADR-0062), so explicit HBM loads overlap the compute that follows. Rationale grounded in kernbench's execution + latency model:

• GEMM tiling is offloaded to PE_SCHEDULER. A CompositeCmd is non-blocking (kernel_runner.py:182-191, pe_scheduler.py:104-121): the kernel pushes one coarse descriptor (M = G·T_q, the whole 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 scheduler streams from HBM, so per-tile K/V prefetch is the scheduler's job — no explicit prefetch op. The CPU (greenlet) is freed to issue the next composite while the current one runs, so the scheduler keeps the GEMM engine saturated across tiles.
• This reflects the hardware and decouples CPU issue-rate from execution-rate. The blocking per-op tl.dot path, by contrast, stalls the CPU on every GEMM and leaves GEMM-engine bubbles during the interleaved softmax MATH ops; it is realistic only if the CPU can keep up with fine-grained per-tile issue.
• The running (m, ℓ, O) flash state stays Python TensorHandles threaded through the loop (the baseline already does this); the softmax merge (max/exp/sum/rescale) is kernel-level tl MATH between the two GEMM composites. The existing CompositeCmd cannot chain two GEMMs or carry cross-tile register state (§0), so the merge necessarily lives in the kernel — this is the hybrid split, not a limitation worked around.
• Cross-PE combination is a log-sum-exp tree over (m, ℓ, O) after P·V, via kernel-level tl.send/tl.recv (§4) — unchanged.

So the per-tile inner pipeline is:

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)

with each tile's MATH temporaries wrapped in tl.scratch_scope (ADR-0063) so scratch stays O(1), and the next tile's composites issued before the current tile's results are waited on (non-blocking handles) so the scheduler pipelines across tiles.

Control flow (tile skip, mask generation, reduction scheduling, address arithmetic) lives in the kernel (plain Python if/arithmetic in the greenlet body). This is exactly what kernbench's greenlet model already permits (kernel_runner.py, ADR-0020 D3).

What this supersedes. An earlier iteration proposed a pure greenlet primitive path (all tl.dot, no composite) on the argument that "composite yields no latency benefit." That holds only because the simulator currently charges zero per-op CPU issue cost (dispatch_cycles=0, pe_cpu.py) — it models away exactly the CPU issue-rate / DMA-program cost that descriptor offload exists to hide. The hybrid is the faithful representation of an efficient kernel. The measurable size of the win (can the CPU saturate the engines for many tiles?) is gated on modelling an op-type-differentiated issue cost, tracked as future work (cost model; §9). Even at dispatch_cycles=0 the non-blocking composite path fills the GEMM-engine bubbles the blocking tl.dot path leaves.

Why this is the efficient choice. GEMM tiling + DMA streaming + cross-tile pipelining are offloaded to the proven CompositeCmd scheduler path; the softmax merge and the proven IPCQ collective stay in the kernel. The only genuinely new machinery is two small, general primitives (ADR-0062 lazy tl.load, ADR-0063 tl.scratch_scope); the reduction reuses tl.send/tl.recv. A bespoke "flash-composite" command (one kind internalising the softmax merge + carried register state + an IPCQ-push epilogue) is not built — large, special-purpose, and its only delta over this hybrid (full softmax offload) is not justified at the current modelling fidelity; see §8.

---

2. Memory layout & driver responsibilities

2.1 KV cache allocation (2-level SP)

One KV head's sequence is sharded across C × P ranks (Level-1 inter-CUBE × Level-2 intra-CUBE PE, §0). In kernbench this is a DPPolicy whose both the cube axis (row_wise over C) and the pe axis (row_wise over P) shard the sequence dimension, with Q replicate (policy/placement/dp.py). The baseline's single-axis DPPolicy(pe="row_wise") becomes a two-axis row_wise over (cube, pe).

• 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.

Shared prefill/decode layout — use contiguous blocks, not round-robin. Because prefill writes the KV cache (qkv_rope, upstream) and decode reads + extends the same cache, both kernels share one physical layout — no reshard at the prefill→decode boundary. Prefill's causal skip needs contiguous position blocks (a late contiguous block can be wholly future and skipped; a round-robin block spans all positions and can never be skipped). So the shared layout shards each KV head into contiguous C × P position blocks (rank r owns [r·B, (r+1)·B), B = ⌈max_context/(C·P)⌉). Decode reads its block + reduces; prefill rings the C CUBE-level blocks. Caveat: contiguous under-uses ranks for short context (only frontier ranks hold data) — acceptable given the long-context target; short-context balance is a separate study (§B). (The round-robin formula above is the alternative for decode-only balance; the contiguous block form is the one that serves both kernels.)

2.2 Driver per-launch duties (minimal)

The driver supplies, per launch, the bases + counter + rotation; the kernel derives the rest:

Launch arg	Purpose
K_base[rank], V_base[rank]	per-rank KV buffer bases (from tensor VA)
O_base	attention output destination
global_token_counter	current sequence position
start_cube, start_pe = f(request_id)	Level-1/Level-2 rotation
cube_id, pe_id (tl.program_id 1/0), C, P	CUBE-Group geometry
q_block_meta	prefill/SP query block start + length

Let R = C·P, cube_local = (cube_id − start_cube) mod C, pe_local = (pe_id − start_pe) mod P, and this rank's index rank = cube_local·P + pe_local. Kernel-derived (plain arithmetic):

• my turn to write this step: (start_rank + counter) mod R == rank
• my valid length: base = counter // R; rem = counter % R; my_len = base + (1 if rank < rem else 0)
• read range: tiles 0 .. ⌈my_len / TILE⌉.
• causal bounds / per-tile skip: from global positions of the query block vs each tile.

Driver = bases + counter + rotation. The single formula (request_id + token_idx) mod (C·P) over the two SP axes is the entire placement policy.

---

3. Per-tile op sequence (greenlet tl)

One iteration = one KV tile on one PE. The two GEMMs are tl.composite(op="gemm") (scheduler-managed tiling + K/V DMA streaming); the softmax merge is kernel tl MATH between them. Real tl names (tl_context.py), with 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 is the GQA-batched query reshaped to [G·T_q, d] (the G group rows folded into the matmul M dim; byte-conserving). One K/V tile serves all G·T_q rows — the GQA reuse lever — with no broadcast.
• K/V are tl.ref operands the composite scheduler streams from HBM per tile (pe_scheduler.py:104-143): that is the prefetch/pipeline, so there is no explicit prefetch op. Issuing tile j+1's composites before waiting on tile j (non-blocking handles) keeps the scheduler pipelined across tiles and the GEMM engine saturated.
• tl.trans is metadata-only in kernbench (tl_context.py:390) and MemoryStore.read reshapes rather than transposes (memory_store.py:73). For zero/structural runs this is harmless; for non-trivial numeric data it yields a reshape-not-transpose, so Q·Kᵀ via a transposed K needs care (§11) — store K pre-transposed [d, TILE], or add a real tl.transpose (a candidate further primitive, likely unnecessary given the simulator's performance-modeling purpose).
• Masking: the kernel builds the boundary-tile mask from query/KV global offsets and adds it (Sj + mask_j); full-past tiles pass None; full-future tiles are skipped (the if never enqueues them).
• A new "flash-composite" command kind (one that internalises the softmax merge + carried (m,ℓ,O)) is not used; the existing CompositeCmd covers each GEMM and the merge stays in the kernel (§1, §8 item 4).

---

4. Reduction (KV-parallel / SP combine) — 2-level reduce-to-root

Scope: this is the DECODE-SP path (head-replicated, KV statically sharded, small O). Prefill-SP does not reduce — it rotates KV (Ring, §5.5). Reduce is chosen for decode because O = [G, d] is tiny, so moving (m,ℓ,O) beats moving the resident KV.

After its tile sweep each rank holds (m_i, ℓ_i, O_i) (unnormalised) over its 1/(C·P) shard. Combine via the associative/commutative log-sum-exp merge (identical math to the baseline's 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

Because the merge is associative and commutative, the combine is reduce-to-root (attention needs O at one place — the rank that writes O_base — not an all-reduce) and proceeds in two levels:

4.1 Level-2 — intra-CUBE (P PEs → CUBE root PE)

• The CUBE's KV slice is itself sequence-SP-split across its P PEs (decision (a): the only way decode, T_q=1, can use all P PEs — there is no query axis to split). Each PE sweeps its 1/(C·P) sub-shard, then a reduce-to-root tree over the P PEs, 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 1–2: row-then-col converge at the center CUBE) — but reduce-only (drop its broadcast-back Phases 4–5; 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 (R−1 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 variant (chunked) is available if a specific reduction proves bandwidth-bound. (Note: this is the reduce collective for decode; prefill does not reduce at all — it uses the Ring-KV kernel of §5.5.) C=1 degenerates to Level-2 only (single-CUBE SP).

Payload: O_i ([G·1, d] for decode) is the heavy part; m,ℓ are light. Sent as separate tl.sends (baseline pattern).

Kernel structure (greenlet), both levels reuse 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)

The reduction tree/mesh is static (fixed at compile time), so the recv order is data-independent — no data-dependent pop, plain blocking tl.recv suffices, and no hardware/composite pop-as-dependency change is required (§6).

---

5. Per-case kernels

All cases share §3's composite-hybrid inner tile, but split into two kernels (TL;DR): the decode-reduce skeleton below (§5.2/§5.3) and the prefill-ring kernel (§5.5). They differ in head mapping, KV strategy, and cross-rank communication — see §10.

5.1 Decode skeleton (head-replicated, static 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 is this same shape with C=P=1, T_q>1, and causal masking; prefill-with-SP is the separate Ring kernel of §5.5, not this skeleton.)

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 final ragged tile.
• GQA reuse is the whole game (decode is KV-load-bound): the G=8 query rows of the KV head are folded into the matmul M dimension (q_g reshaped [G, T_q, d] → [G·T_q, d], a byte-conserving reshape). Then Q·Kᵀ is composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE] and P·V is composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]. The KV tile ([TILE, d]) is the shared K/V operand — streamed once, reused by all G·T_q rows automatically because they are the M rows of the 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 batch axis would not be counted — see §8).
• 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 composite for P·V) — exactly the baseline _attention_mesh_mlo _partial_attention, just GQA-batched and on the composite path. Tiling (§3) only kicks in when S_rank exceeds the scratch scope's tile budget.

5.3 DECODE, with SP / KV-parallel (C × P ranks)

• One request's KV is hierarchical-round-robin across the C·P ranks of the CUBE Group (Level-1 over C CUBEs, Level-2 over P PEs); each rank owns ≈my_len tokens and GQA-batches its G=8 query rows over its local tiles.
• 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

• Whole prompt resident; query is a block of T_q tokens (chunked).
• Causality is real, per query block [qs, qe) vs KV tile [ks, ke):
  • ke ≤ qs → tile_all_past → mask=None, full compute.
  • ks ≥ qe → tile_all_future → skip (kernel if).
  • overlap → tile_partial → kernel builds a triangular additive mask, the tile op sequence adds it (Sj + mask_j).
• GQA batches the group's query heads along the same axis as decode.

5.5 PREFILL, with SP (Ring KV) — head-parallel, NO reduce

This is the second kernel (Kernel 2 in the TL;DR), structurally distinct from the decode reduce path (§4). The output O = [T_q, d] per head is big, so reducing it across ranks would move [T_q,d] per rank; instead we shard the heads and move the (also big) KV, which each head needs in full anyway.

• Head-parallel placement: within a CUBE Group, CUBE i owns exactly one query head i (the G query heads → the C=G CUBEs, one Q head per CUBE) and KV slice i. Each CUBE computes its one Q head's full attention. Because each CUBE produces a different head, there is no (m,ℓ,O) reduce — each CUBE normalises and writes its own head's rows.
• Ring KV: the C KV slices rotate around the CUBE ring; each CUBE folds the incoming block into its head's running (m,ℓ,O) (online-softmax, carried across ring steps as Python handles, as _attention_mesh_kv does today). After C steps every head has seen all KV. GQA reuse comes from the rotation — slice j visits all C CUBEs, serving all G heads.
• IPCQ overlaps the next step's KV receive with the current step's compute via tl.recv_async/tl.wait (tl_context.py:543-560 — already exists); receive buffers ping-pong (persistent arena, not recycled).
• Causal ring skip: if the incoming KV block is entirely after this head's query block, skip its compute — eliminates ≈half the steps.
• Within a CUBE (P PEs): the head's query rows [T_q, d] and/or the current KV block tile across the P PEs (a query-axis split exists here, unlike decode); details in §B.

The baseline _attention_mesh_kv already implements the ring fold; this ADR adds GQA reuse, the head-parallel placement, causal step-skip, and the composite-hybrid inner tile (§3).

5.6 Decode CPU-pipelining variants (opt1 / opt3 / opt2)

The three decode variants are shown in full in the TL;DR (Kernel 1): opt1 current CompositeCmd (has a GEMM-engine bubble while the CPU auto-waits on Sj), opt3 software pipelining (issue the next tile's Q·Kᵀ before this tile's softmax, Sj in a persistent double buffer — removes the bubble, no new command kind), opt2 ex_composite (split into #1 Q·Kᵀ = existing composite reading K first, #2 softmax+P·V+ accumulator = the only new flash-epilogue machinery, reading V later).

Recommend: ship opt3 now (no new machinery, removes the bubble); revisit opt2 once the cost model (ADR-0064) makes the fewer-CPU-issues win measurable (§8 item 4). The MATH engine already has max/sum/exp — the only genuinely new part of opt2 is the composite's stateful (m,ℓ,O) accumulator, not the ops.

These variants are a decode concern: in prefill (§5.5) the causal if is kernel control flow that cannot enter a composite, and the ring already overlaps via recv_async.

---

6. Why no hardware / composite change is needed

• 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 context. The target is agentic = low batch, long context ⇒ each rank has many KV tiles; the §4 reduce's tl.recv blocks are covered by the sweep work of concurrent tokens / long context / the §4.3 level-pipeline.
• The §4 reduction uses a static 2-level tree/mesh, so the recv order is fixed at compile time — no data-dependent pop. tl.recv (blocking) suffices.
• Decision: ship with the greenlet tl.send/tl.recv collective. Revisit a HW pop only if a short-context, single-stream, latency-critical target emerges.

---

7. Control vs execution split (the load-bearing principle)

Concern	Owner
tile skip (future), mask generation, causal bounds	kernel (Python if + arithmetic in greenlet body)
address / offset / valid-length arithmetic	kernel (from counter)
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 load, mask add, softmax math, running (m,ℓ,O) merge	tl ops on PE engines (kernel-issued)

The kernel decides; the GEMMs are offloaded to the scheduler as composites; the remaining tl ops execute already-decided work. This is exactly the greenlet + composite model kernbench already supports — no new control abstraction.

---

8. Required kernbench changes

Correction from design iteration: real GQA (h_q > h_kv) needs no new primitive — only the kernel restructuring in §5.2 (per KV head, G folded into M, byte-conserving reshapes). The supporting ADRs are efficiency / scale enablers, not GQA blockers.

Algorithm work in the kernel (no new primitive; existing tl API):

• GQA Q-axis batching (the reuse lever) — fold G·T_q into the matmul M dim per KV head (§5.2); _view-style byte-conserving reshape; the GEMM is a tl.composite(op="gemm") with M = G·T_q. Runs today in both timing and data mode.
• GEMMs via composite (§1/§3) — Q·Kᵀ and P·V each issued as a non-blocking tl.composite(op="gemm"); PE_SCHEDULER tiles them and streams the tl.ref K/V operands' DMA (existing CompositeCmd; no new command kind).
• 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.
• Causal tile skip + additive boundary mask (§3/§5.4) — kernel if + a mask tensor added with +.
• 2-level round-robin KV placement / valid-length arithmetic (§0, §2) — launch-arg arithmetic + DPPolicy row_wise over both cube and pe.

New primitives for efficiency / scale (each has a supporting ADR):

1. Per-tile scratch recycling — ADR-0063 (tl.scratch_scope). Required for scale: removes the S=16 ceiling (1 MiB bump allocator) so realistic context lengths run. Highest-value of the three.
2. Lazy tl.load — ADR-0062 (non-blocking load + auto-wait on first use; API surface unchanged). Efficiency: overlaps explicit loads (the Q group, non-composite kernels) with following compute. The per-tile K/V prefetch is handled by the composite scheduler (§1), so this covers the remaining explicit loads. Global semantics change → existing goldens regenerate (ADR-0062 D3).
3. GQA head / mask broadcast — ADR-0061 (tl.broadcast). Optional convenience, not a GQA blocker (see correction above). Useful for additive-mask construction across the G·T_q rows and for general kernels; np.matmul already broadcasts in data mode, so it is not needed for correctness. Lowest priority.

Explicitly REJECTED (efficient alternative chosen):

4. A bespoke "flash-composite" command kind that internalises the whole inner loop — DMA→MM→VEC→DMA→MM→VEC with carried (m,ℓ,O) register state and a tail IPCQ push. The two GEMMs do use the existing CompositeCmd (§1/§3) — that gives scheduler-managed tiling, K/V DMA streaming, and cross-tile pipelining. What is rejected is a new command kind that also absorbs the softmax merge + cross-tile register lifetime. Sizing note (§5.6 opt2): if revisited, split it into two composites — #1 = Q·Kᵀ (the existing composite + scale; lets DMA prioritise K), #2 = softmax + P·V + the online-softmax accumulator merge (the only genuinely new machinery: reduction epilogues + a stateful (m,ℓ,O) accumulator). MATH engine already has max/sum/exp — the new part is the composite's stateful flash accumulator, not the ops. Revisit when the cost model (ADR-0064) makes the fewer-CPU-issues win measurable (§5.6); until then ship §5.6 opt3 (software pipelining, no new cmd).
5. Hardware pop-as-dependency. Out of scope (§6).
6. RoPE / QKV projection / KV-cache write inside this kernel. Upstream qkv_rope (P1–P5). Folding RoPE in would force re-rotating past tiles every decode step and break Ring Attention's post-RoPE pass-through.

---

9. Open tuning items (measured in kernbench, not blocking)

1. Composite tile-pipeline depth — dominant for KV-load-bound; how far ahead the kernel issues non-blocking composites before waiting, and the scheduler's per-tile streaming depth.
2. Reduction topology per level (§4.4) — Level-2 (intra-CUBE PE) and Level-1 (intra-CUBE-Group CUBE-mesh) each selectable tree/center-mesh (and a ring variant) — decode reduce only; prefill uses the §5.5 Ring-KV kernel, not a reduce. Also the decode C knob (small C for short context where reduction dominates). 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 GQA) against DMA efficiency; interacts with ADR-0063 and the scheduler's TILE_M/K/N (pe_scheduler.py).
4. Ring buffer ping-pong vs recv_async depth in §5.5.
5. One-shot vs tiled crossover for decode (§5.2) — the S_rank threshold where tiling beats a single composite.
6. Per-op CPU issue cost (cost model) — currently dispatch_cycles=0 (pe_cpu.py), so composite-vs-primitive issue overhead is invisible. An op-type-differentiated issue cost (a tl.composite descriptor push ≫ a primitive op) is what makes the hybrid's CPU-saturation win measurable (§1). Specified in ADR-0064; tracked as separate future work.

---

10. Coverage summary

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 (all G query heads stacked into the 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

I/O per case (full contract §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)

In all cases: KV-cache writes and RoPE happen upstream; output projection downstream; score S and probs P are never materialised.

---

11. Verification plan (Phase 1 test outline)

SPEC/ADR coverage: R5 (PE↔PE IPCQ, PE↔HBM), R2 (latency by traversal), ADR-0023/0025 (IPCQ), ADR-0046 (tl contract), ADR-0054 (eval bench).

The simulator's contract is latency by traversal + determinism + structural correctness (SPEC §0, §0.1), not bit-exact numerics — Phase 2 data exists mainly to exercise the data path, tl.trans is reshape-not-transpose, and bf16 is modelled as f16. Verification is therefore structural/timing-first, with numeric parity as a bounded secondary check.

1. Runs in data mode (enable_data=True): the GQA kernel (h_q = G·h_kv) completes without the byte-conservation error that the baseline head-packing hits — for all four cases. (This needs the §5.2 restructuring, not a new primitive.) Numeric parity (secondary): for symmetric/identity inputs where reshape-as-transpose is exact, kernel O matches a numpy FlashAttention reference within fp tolerance. Full asymmetric parity is gated on a real tl.transpose (out of scope; flagged).
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 GEMM work scales with G. Asserts the lever actually fires.
3. SP cross-rank traffic — both kernels:
   • Decode (reduce): issues ⌈log₂ P⌉ (Level-2, 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 match the static 2-level tree/mesh; result lands at exactly one rank (CUBE-Group root); Level-1 sends use the CUBE NOC, Level-2 the PE IPCQ.
   • Prefill (ring): no (m,ℓ,O) reduce — instead C KV-block rotations (ipcq_send/recv of K,V per step); each CUBE writes a distinct head's O in place (no Group root); GQA reuse shows as each KV slice consumed by all C heads over the ring.
4. Causal skip: prefill skips all tile_all_future tiles (and whole future KV blocks in the ring) — GEMM count equals the lower-triangular tile/step count, not the full grid.
5. Long context (ADR-0063): a sweep at S that overflows 1 MiB without scopes completes and matches the reference.
6. Load/compute overlap: end-to-end latency of the tiled sweep is below the serial Σ(load+compute) — from the composite scheduler streaming K/V per tile (§1/§3) and lazy tl.load (ADR-0062) overlapping the Q load. (Overlap is real modelled concurrency, not a subtraction.)
7. Composite GEMM offload (structural): each tile's Q·Kᵀ and P·V emit a CompositeCmd (non-blocking) to PE_SCHEDULER, not a blocking tl.dot; op_log shows the composite tile plan and the kernel issues the next tile's composites before waiting (cross-tile pipelining).
8. Determinism: identical inputs → identical op_log + latency (SPEC §0.1).

---

B. Open design items from the hybrid pivot (review later)

These arose when the decision moved from a pure greenlet primitive path to the composite hybrid + lazy tl.load (this revision). None blocks the design; each needs a verification pass during implementation. Recorded here (rather than asked) per the working agreement — the recommendation is my predicted default; revise on review.

1. DDD-0060 is not yet synced. The Detailed Design Document still describes the old tl.load_async double-buffer path and primitive tl.dot inner loop (its §4.3/§5/§10). It must be updated to the hybrid (composite GEMMs, lazy load, K pre-transposed). Left for review because the DDD is a derived how-to and a large rewrite; the ADR is now the authoritative record. Recommend: sync DDD as a follow-up before implementation starts.

2. K operand orientation for the composite GEMM. Q·Kᵀ needs b = [d, TILE], but the KV cache stores K as [S_rank, d]. tl.trans is metadata-only and MemoryStore.read reshapes, not transposes (memory_store.py:73) — so a runtime transpose is wrong for non-trivial 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)) a contiguous slice. Verify the upstream qkv_rope write layout supports this, or add a real tl.transpose (heavier; deferred).

3. Composite output buffer vs tl.scratch_scope. Each Q·Kᵀ composite writes Sj to an out_addr; the kernel then reads it for the softmax MATH. That output buffer, and the in-flight composites' targets, must live where the per-tile scratch_scope (ADR-0063) will not recycle them before they are consumed — same discipline as in-flight lazy loads (ADR-0062 D-Negative). Recommend: composite outputs for the current tile live in the scoped arena (consumed same iteration); the persistent (m,ℓ,O) stays outside. Verify no use-after-recycle when the next tile's composites are issued early (cross-tile pipelining).

4. GQA dma_read_count lever under composite streaming. The lever (§11.2: K/V dma_read_count independent of G) assumes the composite emits one K/V tile DMA reused across all G·T_q M-rows. The scheduler's generate_gemm_plan tiles by TILE_M/K/N (pe_scheduler.py:35-37, 32/64/32) — confirm the M-tiling over G·T_q does not re-issue the shared K/V tile DMA per M-tile (i.e. operand DMA is shared across M-tiles, or the lever weakens). Recommend: assert it in the levers test; if violated, the GQA win is in compute only, not DMA — still correct, but the headline changes.

5. Kernel TILE vs scheduler TILE_M/K/N. The kernel reasons about a logical KV TILE; the scheduler re-tiles internally at fixed TILE_M/K/N. Two tiling layers interact (scratch residency, pipeline depth). Recommend: treat the kernel TILE as the K/V streaming granularity and let the scheduler sub-tile the GEMM; document the relationship in the DDD and sweep both (§9 items 1, 3).

6. Cost model is a separate ADR. The hybrid's CPU-saturation benefit is invisible while dispatch_cycles=0. The per-op-type issue-cost model is specified in ADR-0064; this ADR's §1/§9 depend on it for the measurable (not just structural) win. Recommend: land ADR-0064's model before claiming hybrid latency wins in the eval.

7. Ring path (§5.5) GEMMs. §5.5 still describes the ring fold with primitive ops + recv_async. For consistency the ring's per-step Q·Kᵀ / P·V should also be composites; the IPCQ recv_async overlap is orthogonal and stays. Recommend: apply the same hybrid shape to the 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.

Items from the decode-reduce / prefill-ring split (this revision)

1. Two kernels, two head mappings. Decode-SP = head-replicated + static KV shard + 2-level reduce (§4); Prefill-SP = head-parallel (1 Q head/CUBE, C=G) + Ring KV + no reduce (§5.5). The principle is move the smaller thing — decode's O is tiny (reduce it), prefill's O is big (move KV instead). Recommend: keep them as two kernels; do not re-merge.

2. Output head distribution differs (downstream impact). Decode lands all G heads at the CUBE-Group root; prefill leaves one Q head per CUBE (distributed). The downstream out-projection must consume each layout (gather for decode-root vs in-place per-CUBE for prefill). Recommend: pin the O layout per kernel in the qkv_rope/out_proj contract before implementation (§0.5.4).

3. Prefill within-CUBE P PEs. §5.5 splits the head's query rows [T_q, d] and/or the current KV block across the P PEs (a query axis exists, unlike decode). Recommend: default to tiling the query rows across PEs (no intra-CUBE reduce needed — disjoint output rows); fall back to KV-block split + intra-CUBE reduce only if T_q < P.

4. C = G coupling for prefill. The head-parallel mapping assumes C = G = 8 (one Q head per CUBE). If C ≠ G, the mapping needs revisiting (multiple heads per CUBE, or heads spanning a partial ring). Recommend: fix C = G for the prefill kernel at headline scale; treat C ≠ G as a separate study.

5. Reconcile with _attention_mesh_mlo_2d (current impl). The remote impl's 2D kernel is an AllReduce over cubes with Q replicated — i.e. the decode-reduce family, but all-reduce (broadcast-back) not reduce-to-root, and not yet the prefill head-parallel ring. Recommend: (a) move its 2D AllReduce → reduce-to-root (drop broadcast-back) for the decode kernel; (b) add the §5.5 head-parallel Ring-KV kernel for prefill.

6. Shared KV cache layout (prefill writes, decode reads+extends). Both kernels share one physical KV cache, so it must use contiguous C×P position blocks (not round-robin) — prefill's causal skip needs contiguous blocks, and a shared layout avoids a prefill→decode reshard (§2.1). Recommend: standardise on the contiguous block layout in the qkv_rope write contract; flag that short-context decode under-uses ranks under contiguous (acceptable for the long-context target; a short-context balance scheme is a separate study).

7. Decode CPU-pipelining variant to ship (§5.6). Three decode variants exist (opt1 current / opt3 software-pipelining / opt2 ex_composite). Recommend: implement opt3 (software pipelining: issue next Q·Kᵀ before this tile's softmax, Sj in a persistent double buffer) — removes the GEMM-engine bubble with no new command kind. Defer opt2 (the two-composite ex_composite, only #2 is new) until ADR-0064's cost model makes its fewer-issues win measurable.