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 + tree reduction → kernel; tl.load is lazy. Per-KV-head, G folded into the matmul M dim (real GQA, no broadcast). This is the reference shape; the prose sections elaborate each piece.

def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
                  counter, start_pe, N, q_block, scale, *, tl):
    # ---- geometry (kernel arithmetic, §2) ----
    pe_id   = tl.program_id(axis=rank_axis)
    G       = H_q // H_kv                       # query heads per KV head (=8)
    causal  = q_block.is_prefill

    for kv in range(H_kv):                       # one KV head per iteration (§5.2)
        # Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
        # issue now, auto-wait at first use inside the first composite.
        q_g     = tl.load(q_base(kv), (G * q_block.T_q, d))      # [G·T_q, d]
        my_len  = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv
        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

        # ---- cross-PE combine (§4): log-sum-exp tree to root, or store ----
        if N == 1:
            tl.store(o_base(kv), O / l)
        else:
            tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv))  # tl.send/tl.recv

Why this shape: the two tl.composite calls offload tiling + K/V DMA streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse descriptors and runs ahead → engines stay saturated, §1); the softmax merge and the IPCQ tree reduction stay in the kernel because the existing CompositeCmd cannot carry cross-tile state. 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 tree reduction to root is ⌈log₂ N⌉ steps (§4).
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).
• Allocation policy (SP): one query head per CUBE in the multi-user panels; the G query heads of a KV group map within a CUBE's PEs.

Two orthogonal mapping layers (round-robin placement):

• Layer 1 (KV-parallel, intra-request): KV token i lands on PE (start_pe + i) mod N — round-robin so one long request is split across N PEs, balanced to ≤1 token.
• Layer 2 (inter-request): start_pe = request_id mod N rotates the "PE-0 role" per request.
• Combined: pe = (request_id + global_token_idx) mod N.

N = number of PEs in the KV-parallel reduction group for one (query head, request). Reduction stays intra-CUBE (a query head never spans CUBEs — explicit non-goal).

---

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·N + ((pe_id − start_pe) mod N) (§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. S_pe = keys owned by this PE ≈ ⌈S/N⌉. 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-PE HBM (loaded to TCM)	post-RoPE (P1). G query rows of the group batched.
K_cache	[S_pe, d]	per-PE HBM, base K_base[pe], contiguous	post-RoPE (P2). Read-only. Dense locally, strided in global index (§2.1).
V_cache	[S_pe, d]	per-PE HBM, base V_base[pe]	raw V (P4). Read-only.
global_token_counter	scalar	launch arg	kernel derives S_pe, slot↔global, causal bounds.
start_pe (= request_id mod N)	scalar	launch arg	Layer-2 rotation.
pe_id, N	scalars	launch / tl.program_id	reduction-group geometry.
q_block_meta {q_start, T_q}	launch arg	prefill/SP causal masking & skip.
O_base	address	launch arg	where final O is written (root PE only).
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, root PE only	normalised O_acc / ℓ_acc, cast to storage dtype.

Intermediate (non-root PE, on IPCQ — not a kernel-visible output): (m_i, ℓ_i, O_i) (m,ℓ: [G, T_q]; O_i: [G, T_q, d] unnormalised) pushed to tree_parent via tl.send (§4). O_i is the heavy part.

No-SP (N=1): no IPCQ partials; the single PE'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-PE 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

• Per-PE KV buffers sized to per-PE max context ⌈max_context / N⌉ × d × dtype, for K and V separately. In kernbench these are deployed with a DPPolicy whose pe (or cube) axis is row_wise over the sequence dimension (policy/placement/dp.py; the baseline uses DPPolicy(pe="row_wise") for K/V and replicate for Q).
• Within a PE, assigned tokens are appended contiguously (slot 0,1,2,…). Global indices are strided (i, i+N, …) but the per-PE buffer is dense ⇒ DMA stays contiguous.
• Slot → global: global_idx = local_slot·N + ((pe_id − start_pe) mod N).

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[pe], V_base[pe]	per-PE KV buffer bases (from tensor VA)
O_base	attention output destination
global_token_counter	current sequence position
start_pe = request_id mod N	Layer-2 rotation
pe_id (tl.program_id), N	geometry
q_block_meta	prefill/SP query block start + length

Kernel-derived (plain arithmetic):

• my turn to write this step: (start_pe + counter) mod N == pe_id
• my valid length: base = counter // N; rem = counter % N; my_len = base + (1 if ((pe_id − start_pe) mod N) < 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 N 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) — tree to root

After its tile sweep each PE holds (m_i, ℓ_i, O_i) (unnormalised). 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 tree root

• Timing: exchange is after P·V (each PE has its final O_i).
• Topology: mesh tree, depth ⌈log₂ N⌉, over physical neighbours. For N=8: level-0 (0↔1,2↔3,4↔5,6↔7), level-1 (1↔3,5↔7), level-2 (3↔7), root = 7. Each tree pair must be a physical N/S/E/W neighbour so tl.send(dir)/tl.recv(dir) use a real link (tuning item §9; the SFR install that wires neighbours is configure_sfr_intracube_pe_ring, used by the baseline).
• Why a tree, not the baseline fan-out: attention needs O at one place (the PE that writes O_base). A tree is ⌈log₂ N⌉ steps (3 for N=8) vs the baseline all-to-all's N−1 (7). For decode, where the per-PE sweep is short, the reduction is the dominant cost, so this is a real win. (The baseline's fan-out replicates the answer to all ranks — unnecessary here.)
• Payload: O_i (d-vector) is heavy; m,ℓ are scalars per (g, T_q). Sent as separate tl.sends (the baseline sends the triplet as three sends — same pattern).

Kernel structure (greenlet):

# leaf / internal node: fold children that send to me, then forward up
for child_dir in tree_children_dirs(pe_id, N):
    m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape)
    O_c = tl.recv(child_dir, O.shape)
    m, l, O = merge(m, l, O, m_c, l_c, O_c)
if not is_root(pe_id):
    tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O)
else:
    tl.store(O_base, O / l)

tl.recv blocks until data arrives (tl_context.py:446-499); the merge order is fixed by the static tree, so no data-dependent pop is needed — which is why 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 N, 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, N,
                q_block, softmax_scale, *, tl):
    pe_id   = tl.program_id(axis=rank_axis)
    my_len  = valid_len(counter, start_pe, pe_id, N)
    n_tiles = ceil(my_len / TILE)
    q_g     = load_Q_group(q_ptr)            # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
    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 N == 1:
        tl.store(o_ptr, O / l)               # no reduction
    else:
        tree_reduce_and_store(pe_id, N, o_ptr)  # §4

5.2 DECODE, no SP (N=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_pe 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_pe exceeds the scratch scope's tile budget.

5.3 DECODE, with SP / KV-parallel (N=8)

• One request's KV is round-robin across 8 PEs; each owns ≈my_len tokens. Each PE GQA-batches its G=8 query rows over its local tiles.
• T_q=1 ⇒ short per-PE sweep ⇒ reduction dominates ⇒ the §4 tree (3 steps) is the structurally important part. Reduction latency is hidden by other concurrent decode tokens when batched, or by the long per-PE tile sweep for long single-stream context.

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 N PEs (the ring); 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).
• 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 query-owning PE; if KV-parallel splits remain, §4 tree-merges them, then normalise + store.

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 PE has many KV tiles; the §4 tree's tl.recv blocks are covered by the sweep work of concurrent tokens / long context.
• The §4 reduction uses a static tree, 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 tree)
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).
• Tree reduction to root (§4) replacing the baseline all-to-all fan-out — 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 +.
• Round-robin KV placement / valid-length arithmetic (§2) — launch-arg arithmetic + DPPolicy(pe="row_wise").

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. PE↔mesh-neighbour mapping for the reduction tree — ensure each depth-⌈log₂ N⌉ pair is a physical N/S/E/W neighbour; bad mapping adds hops. 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_pe 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 PE, all KV	one-shot or tiled, GQA-batched	none	last tile only
Decode, SP	round-robin N PEs	tiled, GQA-batched	§4 tree (⌈log₂ N⌉)	last tile only
Prefill, no SP	resident	tiled per q-block	none	skip future / triangular boundary
Prefill, SP (Ring)	ring-sharded	per ring step (recv_async overlap)	running state + §4 tree	causal step-skip + boundary

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 PE	O[G,1,d] at O_base	none
Decode, SP	Q[G,1,d], per-PE K/V[S_pe,d]	O[G,1,d] (root)	(m,ℓ,O_i) per PE → tree
Prefill, no SP	Q[G,T_q,d], K/V[≤end,d]	O[G,T_q,d] at O_base	none
Prefill, SP (Ring)	Q[G,T_q,d], ring-delivered K/V blocks	O[G,T_q,d] (root)	running (m,ℓ,O) + tree

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. Tree reduction step count: decode-SP issues ⌈log₂ N⌉ reduction rounds (not N−1); op_log ipcq_send/recv counts match the tree.
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_pe, 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_pe] 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.