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 — final GQA kernel (composite hybrid) pseudocode

The decision (§1) in one place: GEMMs → tl.composite; softmax merge + 2-level reduction → kernel; tl.load is lazy. The KV head is selected by CUBE coordinate (one CUBE Group per head, §0) — no for kv loop. G is folded into the matmul M dim (real GQA, no broadcast). The KV sequence is sharded C × P ranks (Level-1 inter-CUBE × Level-2 intra-CUBE PE); the combine is a 2-level reduce-to-root (§4). Reference shape:

def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
                  counter, start_pe, q_block, scale, *, tl):
    # ---- geometry: this rank within its CUBE Group (§0, §2) ----
    cube_id = tl.program_id(axis=1)              # which CUBE in the SIP mesh
    pe_id   = tl.program_id(axis=0)              # which PE in the CUBE
    kv      = head_of_cube(cube_id)              # CUBE coord → KV head (no for-kv loop)
    C, P    = cube_group_size, pes_per_cube      # Level-1 × Level-2 SP extents
    G       = H_q // H_kv                         # query heads per KV head (=8)
    causal  = q_block.is_prefill

    # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062).
    q_g     = tl.load(q_base(kv), (G * q_block.T_q, d))          # [G·T_q, d]
    my_len  = valid_len_2level(counter, start_pe, cube_id, pe_id, C, P) \
                  if not causal else S_kv_local
    n_tiles = ceil(my_len / TILE)

    # persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
    m, l, O = init_running(G * q_block.T_q, d)

    for j in range(n_tiles):
        if causal and tile_all_future(j, q_block):
            continue                                             # causal skip (kernel if)
        with tl.scratch_scope():                                 # per-tile temporaries (ADR-0063)
            # --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite ---
            Sj = tl.composite("gemm", a=q_g,
                              b=tl.ref(k_tile(kv, j), (d, TILE))) * scale   # Kᵀ pre-stored [d,TILE]
            if causal and tile_partial(j, q_block):
                Sj = Sj + causal_mask(j, q_block)                # additive boundary mask
            # --- online softmax merge in the kernel (MATH ops) ---
            m_new = tl.maximum(m, tl.max(Sj, axis=-1))
            P     = tl.exp(Sj - m_new)
            corr  = tl.exp(m - m_new)
            l     = l * corr + tl.sum(P, axis=-1)
            # --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite ---
            Oj    = tl.composite("gemm", a=P,
                                 b=tl.ref(v_tile(kv, j), (TILE, d)))
            O     = O * corr + Oj                                # running merge
            m     = m_new

    # ---- 2-level combine to the CUBE Group root (§4); data-driven, no barrier ----
    hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))

Why this shape: the two tl.composite calls offload tiling + K/V DMA streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse descriptors and runs ahead → engines stay saturated, §1); the softmax merge and the IPCQ reduction stay in the kernel because the existing CompositeCmd cannot carry cross-tile state. The reduction is 2-level (Level-2 PE tree within a CUBE → Level-1 CUBE-mesh reduce within the CUBE Group), both reduce-to-root, log-sum-exp, sent as each rank finishes its local sweep (§4). K is pre-stored transposed [d, TILE] to sidestep the reshape-not-transpose caveat (§3, §B). A bespoke "flash-composite" command kind is not introduced (§8 item 4).

---

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 the KV-parallel reduction group for one (KV head, request) is C × P ranks, all within one SIP. The G query heads are folded into the matmul M dim (replicated across all ranks). C is a tuning knob (e.g. 8 or 4): large C for long-context prefill, small C (→ 1 = a single CUBE, intra-CUBE SP only) for short decode where reduction dominates.

Topology grounding (topology.yaml): a SIP is a 4×4 CUBE mesh (16 CUBEs, ADR-0017 NOC); a CUBE has P = 8 PEs (hbm_pseudo_channels/hbm_channels_per_pe = 64/8). With C = 8 a SIP holds 2 CUBE Groups = 2 KV heads; the full H_kv = 8 model therefore spans 4 SIPs (8/2). Each CUBE Group is intra-SIP, so one head's reduction uses the CUBE NOC (Level-1) + PE IPCQ (Level-2) — never UCIe. (The shipped topology.yaml has sips: 2; full scale needs a 4-SIP config — §B.)

--device enumerates SIPs (CLI semantics): one SIP-device bench drives its 2 CUBE Groups (2 KV heads) — the head is selected spatially by CUBE coordinate, not an in-kernel for kv loop. The CLI runs the 4 SIP-devices logically in parallel to cover all 8 heads.

Token → rank placement within a CUBE Group (round-robin SP, balanced to ≤1 token): KV token i lands on Level-1 rank (start + i) mod C, then Level-2 rank ((i // C) + start_pe) mod P — a hierarchical round-robin so each of the C × P ranks owns a dense, contiguous local slice.

Reduction is hierarchical and intra-SIP: a KV head does span the C CUBEs of its group (Level-1), reduced over the CUBE NOC; this reverses the earlier "a query head never spans CUBEs" non-goal, which existed only because the H_kv=1 baseline never needed it. Crossing SIPs for one head remains a non-goal (a head stays in one SIP).

---

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

Output	Shape	Location	Notes
O (final)	[G, T_q, d]	O_base, CUBE-Group root only	normalised O_acc / ℓ_acc, cast to storage dtype.

Intermediate (non-root rank, on IPCQ — not a kernel-visible output): (m_i, ℓ_i, O_i) (m,ℓ: [G, T_q]; O_i: [G, T_q, d] unnormalised) pushed up the 2-level tree to the parent (Level-2 PE parent, then Level-1 CUBE parent) via tl.send (§4). O_i is the heavy part.

No-SP (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.

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

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)

• A reduce-to-root tree over the P PEs of a CUBE, depth ⌈log₂ P⌉ (3 for P=8), over the PE IPCQ N/S/E/W mesh (configure_sfr_intracube_pe_ring). Each tree pair is a 1-hop physical neighbour. Result lands on the CUBE's root PE.

4.2 Level-1 — intra-CUBE-Group (C CUBE roots → Group root)

• A center-root bidirectional CUBE-mesh reduce over the C CUBEs of the group (PE-root-only, program_id(axis=1) = cube_id), over the CUBE NOC 2D mesh (ADR-0017). This adapts the proven inter-CUBE pattern of lrab_hierarchical_allreduce.py (Phases 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 option (chunked O) suits bandwidth-bound long-context prefill with large T_q·d payload. C=1 degenerates to Level-2 only (single-CUBE SP).

Payload: O_i ([G·T_q, d]) is heavy; m,ℓ are light. Sent as separate tl.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 four cases share §3 (tile sweep) + §4 (reduction); they differ only in the SP extents C·P, query-block width, and which tile_* predicates fire.

5.1 Common skeleton

def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
                C, P, q_block, softmax_scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv      = head_of_cube(cube_id)          # CUBE coord → KV head (no for-kv loop)
    my_len  = valid_len_2level(counter, start_cube, start_pe, cube_id, pe_id, C, P)
    n_tiles = ceil(my_len / TILE)
    q_g     = load_Q_group(q_ptr, kv)        # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
    m, l, O = init_running()                 # persistent arena: -inf, 0, zeros
    for j in range(n_tiles):                 # K/V streamed per tile by the composite scheduler (§3)
        if tile_all_future(j, q_block):      # causal skip (kernel if)
            continue
        run_tile(j, mask_or_null(j, q_block))   # §3: 2 composites + softmax MATH, in tl.scratch_scope
    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

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)

• KV sharded across the C·P ranks in ring order; each step delivers a peer's KV block over IPCQ. Running (m,ℓ,O) is carried across ring steps (Python handles, exactly as _attention_mesh_kv does today). The ring is the long-context, bandwidth-bound case where §4.4's ring collective option (not the tree) is the natural choice.
• IPCQ overlaps next step's KV receive with current step's compute via tl.recv_async/tl.wait (tl_context.py:543-560 — already exists).
• Causal ring skip: if the incoming step's KV block is entirely after the local query block, skip its compute (don't recv-consume / don't fold) — eliminates ≈half the steps for causal attention.
• Receive buffers ping-pong (persistent arena, not recycled).
• After the ring, (m,ℓ,O) is final per rank; remaining KV-parallel splits combine via the §4 2-level reduce, then normalise + store at the Group root.

The baseline _attention_mesh_kv already implements the ring fold; this ADR adds GQA reuse, causal step-skip, and recv_async overlap.

---

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 + an IPCQ-push epilogue: it is large and special-purpose, and its only delta over the hybrid (full softmax offload) is not justified at the current modelling fidelity (per-op CPU issue cost = 0; see §1). Revisit if the cost model (§9) makes full offload measurably worthwhile.
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/ring/center-mesh; sweep tree (latency, small-O decode) vs ring (bandwidth, large-O prefill). Also the CUBE-Group size C knob (small for decode, large for long context). Ensure each tree pair is a 1-hop physical neighbour; verify against the SFR install.
3. TILE size — balance scratch residency (S/P tiles + O_acc + G-way 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	KV placement	Inner loop	Reduction	Masking
Decode, no SP	1 rank (C=P=1), all KV	one-shot or tiled, GQA-batched	none	last tile only
Decode, SP	2-level RR over C·P ranks	tiled, GQA-batched	§4 2-level (⌈log₂P⌉+mesh over C)	last tile only
Prefill, no SP	resident	tiled per q-block	none	skip future / triangular boundary
Prefill, SP (Ring)	ring over C·P ranks	per ring step (recv_async overlap)	running state + §4 2-level	causal step-skip + boundary

I/O per case (full contract §0.5):

Case	Inputs	Output	IPCQ partials
Decode, no SP	Q[G,1,d], full K/V[S,d] on 1 rank	O[G,1,d] at O_base	none
Decode, SP	Q[G,1,d], per-rank K/V[S/(C·P),d]	O[G,1,d] (Group root)	(m,ℓ,O_i) per rank → 2-level
Prefill, no SP	Q[G,T_q,d], K/V[≤end,d]	O[G,T_q,d] at O_base	none
Prefill, SP (Ring)	Q[G,T_q,d], ring-delivered K/V blocks	O[G,T_q,d] (Group root)	running (m,ℓ,O) + 2-level

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. 2-level reduction step count: decode-SP 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 counts match the static 2-level tree/mesh, and the result lands at exactly one rank (CUBE-Group root). A separate assertion checks Level-1 sends use the CUBE NOC and Level-2 the PE IPCQ.
4. Causal skip: prefill skips all tile_all_future tiles — GEMM count equals the lower-triangular tile 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.