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 `TensorHandle`s 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 **CUBE**s (memory cubes, each with a logic die
  containing **PE**s) + **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
  `TensorHandle`s 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):

```python
# 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`):

```python
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.send`s (the baseline sends the triplet as
  three sends — same pattern).

Kernel structure (greenlet):

```python
# 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

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