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:** pre-existing mesh-native attention kernels and
their `milestone-gqa-llama70b` eval bench (removed when this ADR's
kernels landed). See *§A. Relationship to existing kernbench work* — that
code was the baseline this ADR upgrades to a real GQA, causal,
long-context kernel.

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

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

---

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

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

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

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

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

The decode inner tile has a hard chain `Q·Kᵀ → softmax → P·V`: softmax
(`tl.max(Sj)`) waits for the first GEMM, so a naïve loop stalls the CPU on
`Sj` and **leaves the GEMM engine idle during softmax** (a bubble). Three
variants trade CPU/HW complexity against that bubble (ship **opt3** now;
**opt2** needs a new command kind — revisit with the cost model, §8/ADR-0064):

```python
# OPTION 1 — current CompositeCmd (today; HAS the GEMM-engine bubble)
def gqa_decode_v1(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv  = head_of_group(cube_id)                       # CUBE → its KV head
    q_g = tl.load(q_base(kv), (G * 1, d))              # T_q=1; Q replicated: G heads stacked into M (M-fold)
    m, l, O = init_running(G, d)                       # persistent arena (-inf, 0, 0)
    for j in range(ceil(S_kv_local / TILE)):           # sweep ONLY my KV shard (resident, no move)
        with tl.scratch_scope():
            Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale])
            # ↓ CPU auto-waits on Sj → GEMM engine IDLE during softmax (bubble)
            m2 = tl.maximum(m, tl.max(Sj, -1)); P = tl.exp(Sj - m2); corr = tl.exp(m - m2)
            l  = l * corr + tl.sum(P, -1)
            O  = O * corr + tl.composite("gemm", a=P, b=tl.ref(v_tile(j), (TILE, d))); m = m2
    hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))   # §4 2-level reduce
```

```python
# OPTION 3 — software pipelining (current primitives; bubble hidden, not provably removed) ← SHIP THIS
def gqa_decode_v3(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv  = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d))
    m, l, O = init_running(G, d); n = ceil(S_kv_local / TILE)
    Sb = double_buffer()                               # 2 persistent Sj buffers (outside scratch_scope)
    h  = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(0), (d, TILE)), out=Sb[0], epi=[scale])   # prime
    for j in range(n):
        Sj = Sb[j % 2]
        if j + 1 < n:                                  # ← issue NEXT Q·Kᵀ before softmax → fills GEMM engine
            h = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j+1), (d, TILE)), out=Sb[(j+1)%2], epi=[scale])
        with tl.scratch_scope():
            m2 = tl.maximum(m, tl.max(Sj, -1)); P = tl.exp(Sj - m2); corr = tl.exp(m - m2)
            l  = l * corr + tl.sum(P, -1)
            O  = O * corr + tl.composite("gemm", a=P, b=tl.ref(v_tile(j), (TILE, d))); m = m2
    hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))
```

```python
# OPTION 2 — ex_composite, 2-split (NEW flash-epilogue cmd; gives K-before-V DMA priority)
def gqa_decode_v2(q_ptr, k_ptr, v_ptr, o_ptr, S_kv_local, d, C, P, scale, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv  = head_of_group(cube_id); q_g = tl.load(q_base(kv), (G * 1, d))
    acc = init_running(G, d)                           # (m,l,O): scheduler-updated flash accumulator
    for j in range(ceil(S_kv_local / TILE)):
        Sj = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(j), (d, TILE)), epi=[scale])   # #1: reads K (priority)
        tl.ex_composite("softmax_pv", s=Sj, v=tl.ref(v_tile(j), (TILE, d)), acc=acc, scale=scale)  # #2: reads V
        # ↑ both non-blocking; CPU never waits intra-tile → max run-ahead, fewest issues
    tl.wait()                                          # ← drain ALL composites: acc is updated async (no auto-wait,
                                                       #   #2 is a serial chain through acc); only final after the last #2
    hierarchical_reduce_and_store(*acc, cube_id, pe_id, C, P, o_base(kv))
```

| | new HW cmd | GEMM bubble | CPU intra-tile wait | issues | now |
|---|---|---|---|---|---|
| **opt1 current** | no | **yes** | yes | `O(tiles·ops)` | ✓ |
| **opt3 sw-pipe** | no | hidden\* | yes (reordered) | `O(tiles·ops)` | ✓ |
| **opt2 ex_composite** | **#2 only** | hidden\* | **no** | `O(tiles)` | ✗ (build #2) |

\* "hidden" not "removed": filling the GEMM engine during softmax depends on
the scheduler/engine balance (softmax-vs-GEMM time, K DMA readiness, queue
depth) — it *reduces/hides* the bubble, not provably eliminates it.

(`#1` = the *existing* GEMM composite + `scale`; only `#2` = softmax + P·V +
the stateful online-softmax accumulator is new. MATH engine already has
max/sum/exp — the new part is the flash accumulator, not the ops.)

### Kernel 2 — PREFILL + SP (1 Q head per CUBE, head-parallel; Ring KV, NO reduce)

```python
def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_block, cube_start, *, tl):
    i   = tl.program_id(axis=1) - cube_start           # this CUBE's Q head = its start KV slice
    q   = tl.load(q_ptr, (T_q, d))                     # MY Q head's query rows (resident, TCM)
    Kc  = tl.load(k_ptr, (d, S_kv_local))              # my K slice, pre-stored transposed [d, S/C] → TCM
    Vc  = tl.load(v_ptr, (S_kv_local, d)); src = i     # my V slice [S/C, d] → TCM  (K/V in TCM so the ring can send them)
    m, l, O = init_running(T_q, d)
    for step in range(C):                              # ── Ring KV: rotate blocks around C CUBEs over IPCQ ──
        f = None
        if step < C - 1:                               # send current TCM block, recv next (overlap)
            tl.send("ring+", Kc); tl.send("ring+", Vc)
            f = (tl.recv_async("ring-", (d, S_kv_local)), tl.recv_async("ring-", (S_kv_local, d)))
        if not block_all_future(q_block, slice_pos(src)):         # causal skip whole-future blocks
            S = tl.composite("gemm", a=q, b=Kc, epi=[scale])      # [T_q,d]·[d,S/C] → [T_q,S/C]; Kc is TCM-resident
            if block_partial(q_block, slice_pos(src)): S = S + causal_mask(q_block, slice_pos(src))
            m2 = tl.maximum(m, tl.max(S, -1)); P = tl.exp(S - m2); corr = tl.exp(m - m2)
            l  = l * corr + tl.sum(P, -1)
            O  = O * corr + tl.composite("gemm", a=P, b=Vc); m = m2   # [T_q,S/C]·[S/C,d] → [T_q,d]
        if f: Kc, Vc = tl.wait(f[0]), tl.wait(f[1]); src = (src - 1) % C
    tl.store(o_ptr, O / l)                             # MY Q head's rows — NO reduce
```

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

---

## A. Relationship to existing kernbench work (read first)

kernbench previously ran FlashAttention with an online-softmax `(m, ℓ, O)`
merge over IPCQ via two pre-ADR-0060 mesh kernels (now removed):

| Role | Mechanism |
|---|---|
| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
| decode (split-KV)  | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |

Both were driven by `milestone-gqa-llama70b` (4 panels:
`{single,multi}_user × {prefill,decode}`).

**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 `TensorHandle`s 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.** Baseline was capped at `h_q == h_kv == 1`. The
   baseline test attributed this to
   a MemoryStore byte-conservation failure on a *broadcast view*, but that
   failure is a property of the baseline's **head-packing hack**
   (`_view(K, (h_q·d, S_kv))`, which conflates all heads into one matmul
   dim and only conserves bytes when `h_q == h_kv`). The correct fix is
   **kernel restructuring**, not a broadcast op: process **one KV head at
   a time** and fold the `G` group rows into the matmul **M** dimension
   (§5.2). That uses only byte-conserving reshapes, so real GQA
   (`h_q = G·h_kv`) runs with **no new primitive** — see §8.
2. **O(N) reduction.** The baseline does an all-to-all bidirectional
   fan-out so *every* rank ends with the full answer (`n_ranks − 1`
   steps). Attention only needs `O` at the query owner → a **2-level
   reduce-to-root** (intra-CUBE tree + intra-CUBE-Group center-mesh, §4) is
   `⌈log₂ P⌉` + center-mesh-over-`C` steps.
3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump
   allocator leaks per-tile temporaries (baseline test capped S at 16
   for this reason) 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 **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/...`), issued **non-blocking** to PE_SCHEDULER,
  which generates a tile plan and streams DMA→GEMM→write per tile
  (ADR-0014 D6; `pe_scheduler.py:104-143`). It is **not** a general
  multi-op DAG: it cannot chain two GEMMs, cannot carry register state
  across instances, and cannot pop/wait on IPCQ. This ADR therefore issues
  **each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and
  keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel —
  it does **not** need a new "flash-composite" command kind (see §1, §8).

### CUBE Group — the placement unit (hierarchical SP)

**A `CUBE Group` is `C` CUBEs (within one SIP) that jointly own one KV
head and its `G` query heads.** One KV head's KV sequence is **sharded
two levels** of sequence-parallelism (SP):
- **Level-1 (inter-CUBE, within the CUBE Group):** the head's sequence is
  split across the `C` CUBEs of the group.
- **Level-2 (intra-CUBE, across PEs):** each CUBE's slice is split again
  across its `P` PEs.

So one KV head maps to **`C × P` ranks**, all within one SIP. A **rank** is
not memory — it is one **SP participant** = a specific **PE in a specific
CUBE** (`rank = cube_local·P + pe_local`); its KV shard *resides in* that
PE's owned HBM region (`K_base[rank]`) and is DMA'd to the PE's TCM for
compute. **How the `G` query heads map onto those ranks differs by case**
(the two kernels, TL;DR / §5):

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

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

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

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

**Token → rank placement** within a CUBE Group is **contiguous block
sharding**: rank `r` owns the position block `[r·B, (r+1)·B)`,
`B = ⌈S/(C·P)⌉` (Level-1 picks the CUBE, Level-2 the PE within it). A
contiguous block is **required** so prefill can causal-skip whole future
blocks **and** so prefill + decode share one KV cache (§2.1). (Round-robin —
interleaved `token i → rank (start+i) mod (C·P)` — gives better
short-context decode balance but cannot causal-skip and cannot be shared
with prefill; it is the **rejected** alternative, §2.1.)

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

---

## 0.5 Kernel boundary, preconditions, I/O contract

### 0.5.1 Position in the decoder layer

```
1. RMSNorm
2. QKV projection (GEMM)        ─┐  qkv_rope kernel (SEPARATE, upstream)
3. RoPE on Q and K              ─┤
4. write new K,V → KV cache     ─┘
5. ===== THIS KERNEL: FlashAttention =====   (post-RoPE Q, K-cache, V-cache → O)
6. Output projection (GEMM)        out_proj kernel (SEPARATE, downstream)
7. residual add → FFN ...
```

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

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

### 0.5.3 INPUTS (per kernel launch)

| Input | Shape (per KV head) | Location | Notes |
|---|---|---|---|
| `Q` | `[G, T_q, d]` | per-rank HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. |
| `K_cache` | `[S/(C·P), d]` | per-rank HBM, base `K_base[rank]`, contiguous | post-RoPE (P2). Read-only. A **contiguous** global position block `[r·B,(r+1)·B)` (§2.1). |
| `V_cache` | `[S/(C·P), d]` | per-rank HBM, base `V_base[rank]` | raw V (P4). Read-only. |
| `global_token_counter` | scalar | launch arg | kernel derives local len, slot↔global, causal bounds. |
| `start_cube`, `start_pe` (= `f(request_id)`) | scalars | launch arg | Level-1/Level-2 rotation. |
| `cube_id`, `pe_id`, `C`, `P` | scalars | launch / `tl.program_id` 1/0 | CUBE-Group reduction geometry (`R = C·P`). |
| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
| `O_base` | address | launch arg | where final O is written (CUBE-Group root only). |
| `softmax_scale` | scalar | launch arg | `1/√d`. |

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

### 0.5.4 OUTPUTS

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

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

**Handoff contract to out-projection (downstream, separate kernel).** This
ADR *guarantees* the two layouts above; designing out-proj is out of scope,
but the handoff must match: **decode** delivers all `G` heads co-located at
the Group root (out-proj reads them locally, or scatters per its own SP);
**prefill** leaves head `i` on CUBE `i` (out-proj is itself head-parallel,
or must **gather** the `G` heads first). Pin this per-kernel in the
`out_proj` launch contract (§B decode-split item 2).

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

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

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

---

## 1. Decision (mechanism)

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

- **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is
  non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`):
  the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole
  per-rank tile sweep) and the scheduler generates the tile plan and streams
  DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the
  scheduler streams from HBM, so per-tile **K/V prefetch is the
  scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed
  to issue the **next** composite while the current one runs, so the
  scheduler keeps the GEMM engine saturated across tiles.
- **This reflects the hardware** and decouples CPU issue-rate from
  execution-rate. The blocking per-op `tl.dot` path, by contrast, stalls
  the CPU on every GEMM and leaves GEMM-engine **bubbles** during the
  interleaved softmax MATH ops; it is realistic only if the CPU can keep
  up with fine-grained per-tile issue.
- **The running `(m, ℓ, O)` flash state stays Python `TensorHandle`s**
  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
> primitive-op path (all `tl.dot`, no composite — both are issued from the
> greenlet, so the contrast is *scheduler-managed composite* vs
> *kernel-issued ops*, not "greenlet vs composite") on the argument that
> "composite yields no latency benefit." That holds **only because** the
> simulator currently charges **zero** per-op CPU issue cost
> (`dispatch_cycles=0`, `pe_cpu.py`) — it models away exactly the CPU
> issue-rate / DMA-program cost that descriptor offload exists to hide.
> The hybrid is the faithful representation of an efficient kernel. The
> **measurable** size of the win (can the CPU saturate the engines for
> many tiles?) is gated on modelling an op-type-differentiated issue cost,
> tracked as future work (cost model; §9). Even at `dispatch_cycles=0` the
> non-blocking composite path fills the GEMM-engine bubbles the blocking
> `tl.dot` path leaves.
>
> **Why this is the efficient choice.** GEMM tiling + DMA streaming +
> cross-tile pipelining are offloaded to the proven `CompositeCmd`
> scheduler path; the softmax merge and the proven IPCQ collective stay in
> the kernel. The only genuinely new machinery is two small, general
> primitives (ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); the
> reduction reuses `tl.send`/`tl.recv`. A bespoke "flash-composite"
> command (one kind internalising the softmax merge + carried register
> state + an IPCQ-push epilogue) is **not** built — large, special-purpose,
> and its only delta over this hybrid (full softmax offload) is not
> justified at the current modelling fidelity; see §8.

---

## 2. Memory layout & driver responsibilities

### 2.1 KV cache allocation (2-level SP)

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

- Per-rank KV buffers sized to `⌈max_context / (C·P)⌉ × d × dtype`, K and V
  separately.
- Each rank's **contiguous position block** is stored dense (`B =
  ⌈max_context/(C·P)⌉` slots) ⇒ DMA stays contiguous.
- Slot → global (**contiguous block**): `rank = cube_local·P + pe_local`
  where `cube_local = (cube_id − start_cube) mod C`,
  `pe_local = (pe_id − start_pe) mod P`; then
  `global_idx = rank·B + local_slot`.

> **Why the placement is *contiguous* blocks (rationale).** Because
> prefill *writes* the KV cache (`qkv_rope`,
> upstream) and decode *reads + extends* the same cache, both kernels share
> one physical layout — no reshard at the prefill→decode boundary. Prefill's
> **causal skip needs contiguous position blocks** (a late contiguous block
> can be wholly future and skipped; a round-robin block spans all positions
> and can never be skipped). So the shared layout shards each KV head into
> **contiguous `C × P` position blocks** (rank `r` owns `[r·B, (r+1)·B)`,
> `B = ⌈max_context/(C·P)⌉`). Decode reads its block + reduces; prefill rings
> the `C` CUBE-level blocks. *Caveat:* contiguous under-uses ranks for
> **short** context (only frontier ranks hold data) — acceptable given the
> long-context target; short-context balance is a separate study (§B).
> **Fallback:** for short/early decode where contiguous-block sharding
> under-utilizes ranks, the driver may select a smaller `C` (or `C=P=1`); SP
> is enabled only past the context-length threshold where KV-sweep time
> dominates reduction and placement imbalance.
> **Rejected: round-robin** (interleaved `token i → rank (start+i) mod (C·P)`)
> — better short-context decode balance *alone*, but a round-robin block
> spans all positions so it cannot causal-skip and cannot be shared with
> prefill.

### 2.2 Driver per-launch duties (minimal)
The driver supplies, per launch, the bases + counter + rotation; the
kernel derives the rest:

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

Let `R = C·P`, `B = ⌈max_context/R⌉`, `cube_local = (cube_id − start_cube)
mod C`, `pe_local = (pe_id − start_pe) mod P`, and this rank's index
`rank = cube_local·P + pe_local`. Kernel-derived (plain arithmetic, with
`counter` = # filled tokens):

- **my valid length** (contiguous block `[rank·B, (rank+1)·B)`):
  `my_len = clamp(counter − rank·B, 0, B)`.
- **my turn to write this step:** the new decode token (position `counter`)
  lands on rank `counter // B` (its contiguous block).
- **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 placement policy is the single
contiguous-block map `token t → rank ⌊t/B⌋` (rotated by `start_*`).

---

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

```python
# running state (persistent arena — allocated once, outside the scope)
# m: [G, T_q]  l: [G, T_q]  O: [G, T_q, d]
q_g = tl.load(Q_group_ptr, (G*T_q, d))               # lazy; auto-wait at first use (ADR-0062)

with tl.scratch_scope():                              # per-tile MATH temporaries recycled
    Sj    = tl.composite("gemm", a=q_g,                # [G·T_q, TILE]; scheduler streams Kⱼ DMA
                         b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale
    if mask_j is not None:
        Sj = Sj + mask_j                              # additive causal mask (boundary tile)
    m_j   = tl.max(Sj, axis=-1)
    m_new = tl.maximum(m, m_j)
    P     = tl.exp(Sj - m_new)                        # no full-matrix softmax; streaming
    corr  = tl.exp(m - m_new)                         # rescale factor for old accumulators
    l     = l * corr + tl.sum(P, axis=-1)
    Oj    = tl.composite("gemm", a=P,                  # [G·T_q, d]; scheduler streams Vⱼ DMA
                         b=tl.ref(V_base + j*TILE*d, (TILE, d)))
    O     = O * corr + Oj                             # running merge (kernel MATH)
    m     = m_new
```

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

---

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

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

After its tile sweep each rank holds `(m_i, ℓ_i, O_i)` (unnormalised) over
its `1/(C·P)` shard. Combine via the associative/commutative log-sum-exp
merge (identical math to the baseline's fold):

```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 CUBE-Group root
```

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

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

- The CUBE's KV slice is itself **sequence-SP-split across its `P` PEs**
  (decision (a): the only way decode, `T_q=1`, can use all `P` PEs — there
  is no query axis to split). Each PE sweeps its `1/(C·P)` sub-shard, then a
  **reduce-to-root tree** over the `P` PEs, depth `⌈log₂ P⌉` (3 for `P=8`),
  over the PE IPCQ `N/S/E/W` mesh (`configure_sfr_intracube_pe_ring`). Each
  tree pair is a 1-hop physical neighbour. Result lands on the CUBE's root
  PE.

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

- A **center-root bidirectional CUBE-mesh reduce** over the `C` CUBEs of
  the group (PE-root-only, `program_id(axis=1)` = `cube_id`), over the
  CUBE NOC 2D mesh (ADR-0017). This **adapts the proven inter-CUBE pattern
  of `lrab_hierarchical_allreduce.py` (Phases 1–2: row-then-col converge at
  the center CUBE)** — but **reduce-only** (drop its broadcast-back Phases
  4–5; attention needs root-only) and with the **log-sum-exp `merge`**
  replacing its plain `+`. Center-root halves the critical path vs a corner
  root (`lrab_hierarchical_allreduce.py:113-116`). No inter-SIP phase — the
  CUBE Group is intra-SIP (§0).

### 4.3 Data-driven overlap, no global barrier

- A rank sends up the tree the **instant its own local P·V finishes** — it
  does **not** wait for sibling ranks; internal nodes `merge` children as
  they arrive. So the reduction **overlaps the compute of slower ranks**
  (causal prefill imbalance, batched decode tokens).
- The two levels **pipeline**: a CUBE whose Level-2 reduction is done feeds
  Level-1 immediately, while other CUBEs are still doing Level-2 — **no
  barrier between levels**. (A rank cannot send before its *own* local
  sweep completes; sending pre-final partials would multiply the heavy `O`
  payload and is not worth it.)
- **Why reduce-to-root, not the baseline fan-out / an all-reduce:** the
  baseline replicates the answer to every rank (`R−1` steps); attention
  needs it once. Reduce-to-root is `⌈log₂ P⌉ + (center-mesh depth over C)`
  — for decode (short sweep, reduction-dominated) this is the dominant win.

### 4.4 Configurable topology

Each level's collective is **selectable** (a tuning item, §9): the defaults
above (Level-2 tree, Level-1 center-root mesh) suit the latency-bound,
small-`O` decode case; a **ring** variant (chunked) is available if a
specific reduction proves bandwidth-bound. (Note: this is the *reduce*
collective for **decode**; **prefill** does not reduce at all — it uses the
Ring-KV kernel of §5.5.) `C=1` degenerates to Level-2 only (single-CUBE SP).

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

Kernel structure (greenlet), both levels reuse `merge`:

```python
def hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base):
    # ---- Level-2: PE tree within the CUBE → CUBE root PE ----
    for child in tree_children_dirs(pe_id, P):          # PE IPCQ N/S/E/W
        m, l, O = merge(m, l, O, *recv_triplet(child))
    if not is_cube_root(pe_id, P):
        send_triplet(parent_dir(pe_id, P), m, l, O); return
    # ---- Level-1: CUBE-mesh center-root reduce → Group root (cube roots only) ----
    for child in mesh_children_dirs(cube_id, C):        # CUBE NOC, center-root
        m, l, O = merge(m, l, O, *recv_triplet(child))
    if is_group_root(cube_id, C):
        tl.store(o_base, O / l)                         # normalise once
    else:
        send_triplet(parent_dir_mesh(cube_id, C), m, l, O)
```

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

---

## 5. Per-case kernels

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

### 5.1 Decode skeleton (head-replicated, static shard, reduce)

```python
def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
                  C, P, *, tl):
    cube_id = tl.program_id(axis=1); pe_id = tl.program_id(axis=0)
    kv      = head_of_group(cube_id)         # CUBE coord → KV head (no for-kv loop)
    my_len  = valid_len_2level(counter, start_cube, start_pe, cube_id, pe_id, C, P)
    n_tiles = ceil(my_len / TILE)
    q_g     = load_Q_group(q_ptr, kv)        # [G·1, d]; lazy tl.load, G folded into M (replicated)
    m, l, O = init_running()                 # persistent arena: -inf, 0, zeros
    for j in range(n_tiles):                 # sweep ONLY my static KV shard (resident, no move)
        run_tile(j)                          # §3: 2 composites + softmax MATH, in tl.scratch_scope
    if C == 1 and P == 1:
        tl.store(o_base(kv), O / l)          # no reduction (single rank)
    else:
        hierarchical_reduce_and_store(m, l, O, cube_id, pe_id, C, P, o_base(kv))  # §4 reduce
```

(Prefill-no-SP is this same shape with `C=P=1`, `T_q>1`, and causal masking;
**prefill-with-SP is the separate Ring kernel of §5.5**, not this skeleton.)

### 5.2 DECODE, no SP (`C=P=1`, one PE owns the head's KV)

- `T_q = 1`, attends to all past KV ⇒ no future tiles, mask only on the
  final ragged tile.
- **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8`
  query rows of the KV head are folded into the matmul **M** dimension
  (`q_g` reshaped `[G, T_q, d] → [G·T_q, d]`, a byte-conserving reshape).
  Then `Q·Kᵀ` is `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]` and
  `P·V` is `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. The KV tile
  (`[TILE, d]`) is the shared `K`/`V` operand — **streamed once, reused by
  all `G·T_q` rows automatically** because they are the M rows of the
  GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile
  plan also makes the timing count all `G` rows' work correctly (a leading
  batch axis would *not* be counted — see §8).
- If `S_rank` fits in scratch (small/medium context) this degenerates to a
  **one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one
  composite for P·V) — exactly the baseline `_partial_attention` shape,
  just GQA-batched and on the composite path. Tiling
  (§3) only kicks in when `S_rank` exceeds the scratch scope's tile budget.

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

- One request's KV is **contiguous-block sharded** across the `C·P` ranks of
  the CUBE Group (Level-1 over `C` CUBEs, Level-2 over `P` PEs); each rank
  owns its `B`-token block and GQA-batches its `G=8` query rows over its
  local tiles.
- `T_q=1` ⇒ short per-rank sweep ⇒ **reduction dominates** ⇒ the §4 2-level
  reduce (`⌈log₂ P⌉` + center-mesh over `C`) is the structurally important
  part. Reduction latency is hidden by other concurrent decode tokens when
  batched, by the long per-rank sweep for long single-stream context, and
  by the §4.3 level-pipelining. For short decode prefer a **small `C`**
  (less inter-CUBE reduction); push `C` up only when context length demands
  the extra KV-parallelism.

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

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

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

> **When ring beats reduce (cost model — the choice is *not* just "O is
> big").** Compare **total cross-rank bytes**, head dim `d` factored out:
>
> - **reduce** ≈ `G·T_q · ⌈log₂(C·P)⌉` — the payload is `(m,ℓ,O)` but **`O`
>   dominates** (`m,ℓ` are `[G·T_q]` scalars vs `O` = `[G·T_q, d]`), so this
>   is `O`-bytes × tree hops; head-replicated, so the `G` factor is present.
> - **ring** ≈ `2·S` — the **total K+V bytes one CUBE injects** over the
>   `C−1` rotations (`2·(S/C)·(C−1) ≈ 2·S`); `recv_async` **pipelines** the
>   rotations so *latency* ≈ the slowest single step, but *traffic* is the
>   sum. **No `G` factor, no log** — one KV slice is *shared* across all `C`
>   heads by the rotation.
>
> So **ring < reduce (by total bytes) when `T_q > 2S / (G·⌈log₂(C·P)⌉)`** —
> decode (`T_q=1`) → **reduce**, prefill (`T_q≈S`, RHS `≈ 2/(G·log) ≪ 1`) →
> **ring**. KV is big, but the reduce's `O` carries the extra `G·log` factor,
> so ring still wins for prefill. (Threshold sweep: §9.)

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

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

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

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

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

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

---

## 6. Why no hardware / composite change is needed

- The only place a HW/composite `pop`-as-dependency would help is a lone
  reduction with no other work to hide a poll — i.e. batch=1 **and** short
  context. The target is **agentic = low batch, long context** ⇒ each rank
  has many KV tiles; the §4 reduce's `tl.recv` blocks are covered by the
  sweep work of concurrent tokens / long context / the §4.3 level-pipeline.
- The §4 reduction uses a **static** 2-level tree/mesh, so the recv order is
  fixed at compile time — no data-dependent pop. `tl.recv` (blocking)
  suffices.
- **Decision: ship with the greenlet `tl.send`/`tl.recv` collective.**
  Revisit a HW `pop` only if a short-context, single-stream,
  latency-critical target emerges.

---

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

| Concern | Owner |
|---|---|
| tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) |
| address / offset / valid-length arithmetic | **kernel** (from counter) |
| reduction scheduling, IPCQ send/recv ordering | **kernel** (static 2-level tree/mesh) |
| Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER |
| Q load, mask add, softmax math, running `(m,ℓ,O)` merge | **`tl` ops** on PE engines (kernel-issued) |

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

---

## 8. Required kernbench changes

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

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

- **GQA Q-axis batching** (the reuse lever) — fold `G·T_q` into the matmul
  M dim per KV head (§5.2); `_view`-style byte-conserving reshape; the
  GEMM is a `tl.composite(op="gemm")` with M = `G·T_q`. Runs today in both
  timing and data mode.
- **GEMMs via composite** (§1/§3) — Q·Kᵀ and P·V each issued as a
  non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and
  streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new
  command kind).
- **2-level reduce-to-root** (§4) replacing the baseline all-to-all fan-out
  — Level-2 PE tree (intra-CUBE) + Level-1 center-root CUBE-mesh reduce
  (intra-CUBE-Group, adapting the `lrab_hierarchical_allreduce` inter-CUBE
  pattern as reduce-only + log-sum-exp), data-driven and level-pipelined,
  pure kernel control flow over `tl.send`/`tl.recv`.
- Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` +
  a mask tensor added with `+`.
- 2-level contiguous-block KV placement / valid-length arithmetic (§0, §2) —
  launch-arg arithmetic + `DPPolicy` `row_wise` over both `cube` and `pe`.

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

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

**Explicitly REJECTED (efficient alternative chosen):**

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

---

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

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

---

## 10. Coverage summary

| Case | Head map | KV strategy | Cross-rank comm | Masking |
|---|---|---|---|---|
| Decode, no SP | `G` replicated, 1 rank | all KV resident | none | last tile only |
| **Decode, SP** | **Q replicated** (all `G` query heads stacked into the GEMM M-dim) | 2-level static shard `C·P` | **§4 2-level reduce** (small `O`) | last tile only |
| Prefill, no SP | `G` replicated, 1 rank | resident | none | triangular / skip future |
| **Prefill, SP (Ring)** | **1 Q head per CUBE** (`C=G`) | **Ring KV rotate** | **none** (KV blocks move, not `O`) | causal step-skip + boundary |

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

| Case | Inputs | Output | Cross-rank traffic |
|---|---|---|---|
| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 rank | `O[G,1,d]` at `O_base` | none |
| Decode, SP | `Q[G,1,d]` (replicated), per-rank `K/V[S/(C·P),d]` | `O[G,1,d]` at Group root | `(m,ℓ,O_i)` reduce (small) |
| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none |
| Prefill, SP (Ring) | `Q[1,T_q,d]` per CUBE, own `K/V[S/C,d]` | `O[1,T_q,d]` per CUBE (distributed) | KV blocks rotate (Ring) |

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

---

## 11. Verification plan (Phase 1 test outline)

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

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

**Gate type.** *Absolute* latency targets are deferred — they are only
meaningful once ADR-0064's cost model lands (`dispatch_cycles=0` today). The
gates below are (a) **structural** (op_log shape) and (b) **relative-to-
baseline** ratios that *are* testable now and must pass: reduce rounds
`⌈log₂P⌉+center-mesh` **vs baseline `C·P−1`**; KV bytes `H_kv·S·d` **vs
`H_q·S·d` (G×=8× less)**; causal GEMM count = lower-triangular **vs full
grid**; prefill ring traffic = `C` KV rotations **vs** a `[G·T_q,d]` reduce.

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

---

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

These arose when the decision moved from a pure primitive-op (`tl.dot`) 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 synced (this revision).** The Detailed Design Document has
   been rewritten to match the two-kernel composite-hybrid + CUBE-Group SP
   design (file plan, placement, phase plan P1–P8, verification, risks). The
   ADR remains the authoritative design record; the DDD is the
   implementation how-to and points back here (and to §B) for rationale.

2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b =
   [d, TILE]`, but the KV cache stores K as `[S_rank, d]`. `tl.trans` is
   metadata-only and `MemoryStore.read` reshapes, not transposes
   (`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial
   data. **Recommend:** store K **pre-transposed** `[d, S_rank]` in the cache
   (the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))`
   a contiguous slice. Verify the upstream `qkv_rope` write layout supports
   this, or add a real `tl.transpose` (heavier; deferred).

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

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

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

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

7. **Ring path (§5.5) GEMMs.** §5.5 still describes the ring fold with
   primitive ops + `recv_async`. For consistency the ring's per-step Q·Kᵀ /
   P·V should also be composites; the IPCQ `recv_async` overlap is
   orthogonal and stays. **Recommend:** apply the same hybrid shape to the
   ring step during implementation; low risk, mirrors §3.

### Items from the hierarchical CUBE-Group SP pivot (this revision)

1. **A 4-SIP topology config is needed for full scale.** The shipped
   `topology.yaml` has `sips: count: 2`; the full `H_kv=8` model at `C=8`
   needs **4 SIPs** (2 KV heads each, §0). **Recommend:** add a 4-SIP
   topology config (16-CUBE 4×4 mesh per SIP, 8 PE/CUBE — already the per-SIP
   shape) for the headline eval; validation-scale runs can use fewer
   CUBEs/SIPs. The bench must remain single-device per launch (one SIP).

2. **`head_of_cube` / CUBE-Group partition of the 4×4 mesh.** With `C=8`,
   each SIP's 16-CUBE mesh splits into 2 CUBE Groups of 8. The exact sub-mesh
   shape (4×2 vs 2×4) sets the Level-1 center-root mesh hop counts and which
   CUBEs are 1-hop neighbours. **Recommend:** pick the partition that keeps
   each CUBE Group a contiguous rectangular sub-mesh with a true center CUBE
   (so `lrab`-style center-root applies); verify against the CUBE NOC routing
   (ADR-0017).

3. **`C` as a per-operating-point knob.** `C` should differ by case (small
    for decode where reduction dominates, large for long-context prefill).
    **Recommend:** expose `C` (and per-level topology, §4.4) as launch/bench
    config, sweep it in the milestone bench, and report the decode-vs-prefill
    optimum. Default `C` TBD by measurement.

4. **Reverses the §0 "never spans CUBEs" non-goal — confirm no downstream
    assumption depends on it.** Earlier text and possibly other ADRs treated
    intra-CUBE-only reduction as invariant. **Recommend:** grep for that
    assumption (SFR installs, address policy, diagrams) before implementation;
    none found in this kernel's scope, but the `configure_sfr_*` neighbour
    wiring must expose CUBE-NOC `N/S/E/W` for Level-1, not just PE IPCQ.

5. **`valid_len_2level` / contiguous-block correctness.** The two-axis
    `rank = cube_local·P + pe_local`, block `[rank·B, (rank+1)·B)` placement
    (§2.1) must tile the sequence with no gaps/overlaps and keep each rank's
    local buffer dense. **Recommend:** unit-test the index math (every global
    token maps to exactly one rank; per-rank blocks are contiguous) before
    the kernel.

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

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

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

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

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

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

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

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

### Items from the long/short context split (this revision)

The split between the long-context kernels (§5.2 decode, §5.5 prefill —
both already specified) and a short-context variant is referenced
abstractly in items §B-1 and §B-6 above ("short-context balance is a
separate study"). This subsection pins the two open numbers.

1. **Short/long context threshold = 256 K tokens.** Below the threshold,
    the §5.5 Ring KV prefill pays C-1 ring rotations whose IPCQ cost is
    not amortised by enough per-rank compute; above it, the per-rank
    KV sweep dominates and ring is the right choice (§9 line 803).
    Likewise §5.2's per-rank `S_kv/C·P` shard is small enough at
    short context that the 2-level reduce-to-root hop count
    dominates the per-rank compute. The 256 K boundary matches the
    point at which `S_kv/C·P` (`C=4, P=8 → /32`) crosses the
    8 K tokens-per-rank mark above which the per-rank tile sweep
    (§3 / §B-3 §scratch_scope) becomes the limiting factor rather
    than collective overhead. **Recommend:** treat 256 K as the
    bench dispatch threshold; expose it as a launch knob for sweeps.

2. **Short-context kernel design — `kv_per_cube ∈ {1, 2, 4, 8}`.**
    At short context the §5.5 Ring KV motion is wasted work and the
    §5.2 cube-SP shard is too thin to feed the engines. The
    short-context variant therefore drops cube-SP entirely: each
    CUBE owns `kv_per_cube` *whole* KV heads (no `S_kv` sharding
    across CUBEs), and the existing PE-SP within a CUBE shards `S_kv`
    across the `P` PEs for each owned head.

    For `h_kv = 8` the natural distributions are:

    | `kv_per_cube` | participating CUBEs | head→CUBE map |
    |---|---|---|
    | 1 | `C = 8` | head `i` → CUBE `i` |
    | 2 | `C = 4` | heads `[2i, 2i+1]` → CUBE `i` |
    | 4 | `C = 2` | heads `[4i..4i+3]` → CUBE `i` |
    | 8 | `C = 1` | all heads → single CUBE |

    No inter-CUBE reduce within a head (each head fully owned by one
    CUBE), so the kernel runs Level-2 (PE) chain reduce-to-root only;
    Level-1 (inter-CUBE) collapses to a no-op. The output stays
    distributed per CUBE (each CUBE writes its owned heads' `O` slice).

    **Recommend:** ship as a separate kernel file (`_gqa_decode_short.py` /
    `_gqa_prefill_short.py`) and a separate bench dispatcher that picks
    short vs long by `S_kv ⋚ 256K`. Keep `kv_per_cube` as a kernel arg so
    the topology cost trade-off can be measured per workload.