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 tl.load_async (KV prefetch overlap — efficiency), ADR-0061 tl.broadcast (optional mask/general convenience). 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).

---

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 — not composites: 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 running (m, ℓ, O) is just Python TensorHandles threaded through the loop. This matters for this ADR's design (see §1).

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) + ADR-0062 (prefetch).

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/...). 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 does not require composites for the attention inner loop (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 in the greenlet tl programming model, not as composites. Rationale grounded in kernbench's execution + latency model:

• The simulator charges latency per op on its modelled component (GEMM on PE_GEMM, vector ops on PE_MATH, DMA on PE_DMA — pe_gemm.py, pe_math.py, pe_dma.py). "Fusing" several ops into one CompositeCmd does not reduce modelled latency; the stages still run on the same engines. The win a real fused kernel gets — overlap (prefetch) and small working set (scratch recycling) — is obtained here by ADR-0062 and ADR-0063, which are smaller and reusable.
• The running (m, ℓ, O) flash state is naturally a set of Python TensorHandles threaded through the loop (the baseline already does this). No composite-carried register state is needed.
• Cross-PE combination uses kernel-level tl.send/tl.recv (the baseline already does this and it is tested). No composite-driven IPCQ push is needed.

So the per-tile inner pipeline is the op sequence

K_j load → Q·Kⱼᵀ → (+ causal mask) → online-softmax update → V_j load → P·Vⱼ → running (m,ℓ,O) update

issued as ordinary tl.* ops, with the next tile's K/V issued via tl.load_async (ADR-0062) so its DMA overlaps the current tile's compute, and each tile's temporaries wrapped in tl.scratch_scope (ADR-0063) so scratch stays O(1).

Cross-PE combination (KV-parallel / SP) is a log-sum-exp tree reduction over (m, ℓ, O) after P·V, flowed through IPCQ with kernel-level tl.send/tl.recv (§4).

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

Why this is the efficient choice. It reuses the proven baseline kernels and the proven IPCQ collective; the only genuinely new machinery is three small, general primitives (ADR-0061/0062/0063). The originally-proposed "everything-in-one-composite + composite IPCQ push

• composite-carried state" would require a bespoke flash-composite command type, a register-lifetime model across composites, and an IPCQ-push epilogue — large, special-purpose, and with no latency benefit over the greenlet path. Those are rejected; 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. Real tl names (tl_context.py), with tl.load_async (ADR-0062), tl.broadcast (ADR-0061), 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]

with tl.scratch_scope():                              # per-tile temporaries recycled
    Kj = tl.wait(f_k[j])                              # prefetched (ADR-0062)
    Sj = tl.dot(q_g, tl.trans(Kj)) * softmax_scale   # [G, TILE]; q_g is GQA-batched
    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)
    Vj    = tl.wait(f_v[j])
    O     = O * corr + tl.dot(P, Vj)                  # P·V folded into running O
    m     = m_new
    if j + PREFETCH < n_tiles:                        # keep the pipeline full
        f_k[j + PREFETCH] = tl.load_async(K_base + (j+PREFETCH)*TILE*d, (TILE, d))
        f_v[j + PREFETCH] = tl.load_async(V_base + (j+PREFETCH)*TILE*d, (TILE, d))

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 load serves all G·T_q rows — the GQA reuse lever — with no broadcast.
• tl.trans(Kj) 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. Data-mode numeric parity therefore needs care (§11) — store K pre-transposed, or add a real tl.transpose (a candidate further primitive, likely unnecessary given the simulator's performance-modeling purpose).
• The V load is issued (prefetched) before it is needed so its DMA overlaps the Q·Kᵀ + softmax of the same/earlier tile.
• 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).
• No CompositeCmd is used. (A future tl.composite flash kind is an optional optimisation — §8 item 4 — not a requirement.)

---

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,d] decode / [G,T_q,d] prefill; tl.broadcast for GQA
    m, l, O = init_running()                 # persistent arena: -inf, 0, zeros
    prime_prefetch(k_ptr, v_ptr, n_tiles)    # tl.load_async first PREFETCH tiles
    for j in range(n_tiles):
        if tile_all_future(j, q_block):      # causal skip (kernel if)
            continue
        run_tile(j, mask_or_null(j, q_block))   # §3 op sequence, 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 tl.dot([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE] and P·V is tl.dot([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]. The KV tile ([TILE, d]) is the shared K/V operand — loaded 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 emitted GemmCmd also makes the Phase-1 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 tl.dot for Q·Kᵀ, one softmax, one tl.dot for P·V) — exactly the baseline _attention_mesh_mlo _partial_attention, just GQA-batched. 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)
K/V load, Q·Kᵀ, mask add, softmax math, P·V	tl ops on PE engines

The kernel decides; the tl ops execute already-decided work. This is exactly the greenlet 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 + 2-D tl.dot. Runs today in both timing and data mode.
• 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. Async HBM tile load / KV prefetch — ADR-0062 (tl.load_async
   • tl.wait). Efficiency: the KV-load-bound overlap lever for decode/long-context. Without it the kernel is correct but serial.
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. Composite that chains DMA→MM→VEC→DMA→MM→VEC with carried (m,ℓ,O) register state and a tail IPCQ push. Replaced by the greenlet path: per-op latency is identical; overlap comes from ADR-0062; small working set from ADR-0063; running state is Python handles; the IPCQ push is tl.send. Building a bespoke flash-composite command type + cross-composite register lifetime + an IPCQ-push epilogue is large, special-purpose, and yields no latency benefit. A tiled tl.composite "flash" kind remains a possible future optimisation (it would fold prefetch+recycling into the scheduler), but it is not required for an efficient kernel and is out of scope here.
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. KV tile prefetch depth (ADR-0062) — dominant for KV-load-bound; sweep 2→4.
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 (K_tile + S/P + V_tile + O_acc + G-way GQA) against DMA efficiency; interacts with ADR-0063.
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 tl.dot.

---

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. Prefetch overlap (ADR-0062): end-to-end latency of the tiled sweep is below the serial Σ(load+compute) (overlap is real).
7. Determinism: identical inputs → identical op_log + latency (SPEC §0.1).