gqa(adr): ADR-0060 — 1 Q head/CUBE term, decode 3 variants (S5.6), shared contiguous KV layout

- Terminology: 'Q replicated' (all G query heads stacked into the GEMM M-dim;
  M-fold explained) for decode; 'one Q head per CUBE' precise for prefill.
- New S5.6: three decode CPU-pipelining variants — opt1 current CompositeCmd
  (has GEMM-engine bubble), opt3 software pipelining (issue next Q.Kt before
  this tile's softmax; Sj in persistent double buffer; ships now, no new cmd),
  opt2 ex_composite split into two (#1 = existing GEMM+scale reads K first;
  #2 = softmax+P.V+accumulator merge, the only new flash-epilogue machinery,
  gives DMA K-before-V priority). MATH engine already has max/sum/exp — the
  new part is the stateful flash accumulator, not the ops.
- S2.1/SB: shared prefill/decode KV layout = contiguous CxP blocks (prefill
  causal-skip needs contiguous; avoids prefill->decode reshard; short-context
  under-use caveat). S8 item 4 sizing note for the two-composite split.

Prefill note: opt2/opt3 give little for prefill (causal if can't enter a
composite; recv_async already overlaps). Docs only; KO mirror deferred.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md

@@ -97,7 +97,7 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b
 > transposed to sidestep the reshape-not-transpose caveat (§3, §B); no
 > bespoke "flash-composite" kind (§8 item 4). **Output head distribution
 > differs** — decode lands all `G` heads at the CUBE-Group root; prefill
-> leaves one head per CUBE (§0.5.4).
+> leaves one Q head per CUBE (§0.5.4).
 
 ---
 
@@ -200,15 +200,20 @@ So one KV head maps to **`C × P` ranks**, all within one SIP. **How the
 `G` query heads map onto those ranks differs by case** (the two kernels,
 TL;DR / §5):
 
-- **Decode** (§4): `G` heads are **replicated** (folded into the matmul M
+- **Decode** (§4): **Q is replicated** — every rank holds all `G` query
-  dim) on every rank; KV is sequence-sharded `C × P` ways; outputs reduce.
+  heads, **M-folded** (stacked into the matmul M / row dimension: `Q` for
-- **Prefill** (§5.5): `G` heads are **head-parallel** — with `C = G`, CUBE
+  the group `[G, T_q, d]` → `[G·T_q, d]`, so one `Q·Kᵀ` GEMM computes all
-  `i` owns query head `i`; KV rotates (Ring); no reduce.
+  `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.
 
-`C` is a **tuning knob** (e.g. 8 or 4): for prefill set `C = G` (one head
+`Q` is small, so replicating it (decode) or distributing it one-head-per-CUBE
-per CUBE); for decode `C` trades inter-CUBE reduction against KV-parallel
+(prefill) is cheap — only the irreducible data moves (decode: `(m,ℓ,O)`;
-breadth (small `C`, even `C = 1` single-CUBE, for short context where
+prefill: KV blocks). `C` is a **tuning knob**: for prefill set `C = G`
-reduction dominates).
+(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
@@ -414,6 +419,21 @@ One KV head's sequence is sharded across `C × P` ranks (Level-1 inter-CUBE
   `pe_local = (pe_id − start_pe) mod P`; then
   `global_idx = local_slot·(C·P) + rank`.
 
+> **Shared prefill/decode layout — use *contiguous* blocks, not
+> round-robin.** 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).
+> (The round-robin formula above is the alternative for decode-only balance;
+> the contiguous block form is the one that serves both kernels.)
+
 ### 2.2 Driver per-launch duties (minimal)
 The driver supplies, per launch, the bases + counter + rotation; the
 kernel derives the rest:
@@ -676,11 +696,11 @@ 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.
 
-- **Head-parallel placement:** within a CUBE Group, **CUBE `i` owns query
+- **Head-parallel placement:** within a CUBE Group, **CUBE `i` owns exactly
-  head `i`** (the `G` heads → the `C=G` CUBEs, one per CUBE) and KV slice
+  one query head `i`** (the `G` query heads → the `C=G` CUBEs, one Q head per
-  `i`. Each CUBE computes **its one head's** full attention. Because each
+  CUBE) and KV slice `i`. Each CUBE computes **its one Q head's** full
-  CUBE produces a *different* head, there is **no `(m,ℓ,O)` reduce** — each
+  attention. Because each CUBE produces a *different* head, there is **no
-  CUBE normalises and writes its own head's rows.
+  `(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, as `_attention_mesh_kv` does
@@ -699,6 +719,75 @@ The baseline `_attention_mesh_kv` already implements 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 (3 kernels)
+
+The decode inner loop has a hard intra-tile chain `Q·Kᵀ → softmax →
+P·V`: the 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 the
+softmax** (a bubble). Three variants trade CPU/HW complexity against that
+bubble (assume a realistic non-zero per-op CPU issue cost — §9/ADR-0064):
+
+**Option 1 — current `CompositeCmd` (today; has the bubble):**
+
+```python
+for j in range(n_tiles):
+    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 while softmax runs (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
+```
+
+**Option 3 — software pipelining (current primitives; bubble removed):**
+issue the *next* tile's `Q·Kᵀ` **before** this tile's softmax, so the GEMM
+engine runs `Q·Kᵀ_{j+1}` during `softmax_j`. `Sj` lives in a persistent
+**double buffer** (outside `scratch_scope`, so the next composite does not
+clobber it).
+
+```python
+Sb = double_buffer()                                         # 2 persistent Sj buffers
+h  = tl.composite("gemm", a=q_g, b=tl.ref(k_tile(0),(d,TILE)), out=Sb[0], epi=[scale])
+for j in range(n_tiles):
+    Sj = Sb[j % 2]
+    if j+1 < n_tiles:                                        # ← 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
+```
+
+**Option 2 — extended composite (`ex_composite`; needs a new command kind):**
+split into **two** composites so DMA can prioritise **K first, V later**
+(V is only needed after softmax). `#1` is the *existing* composite (GEMM +
+`scale`); only `#2` (softmax + P·V + the online-softmax accumulator merge)
+needs the new flash-epilogue machinery (reduction epilogues + a stateful
+`(m,ℓ,O)` accumulator, §8 item 4). The CPU issues both non-blocking and
+**never waits intra-tile** → maximal run-ahead, fewest issues.
+
+```python
+for j in range(n_tiles):
+    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=(m,l,O), scale=scale)  # #2: reads V
+```
+
+| | new HW cmd | GEMM bubble | CPU intra-tile wait | issues | available now |
+|---|---|---|---|---|---|
+| **opt1 current** | no | **yes** | yes | `O(tiles·ops)` | ✓ |
+| **opt3 sw-pipe** | no | no | yes (reordered) | `O(tiles·ops)` | ✓ |
+| **opt2 ex_composite** | **#2 only** | no | **no** | `O(tiles)` | ✗ (build #2) |
+
+All three end with the §4 2-level reduce. **Recommend:** ship **opt3** now
+(no new machinery, removes the bubble); revisit **opt2** once the cost
+model (ADR-0064) makes the fewer-issues win measurable.
+
+> **Prefill note:** these variants are a **decode** concern. In prefill
+> (§5.5) the causal `if` (skip-future / partial-mask) is data-dependent
+> kernel control flow that **cannot enter a composite** and makes
+> pre-issuing speculative; prefill's overlap is the `recv_async` KV
+> prefetch, already present. So opt2/opt3 give little for prefill.
+
 ---
 
 ## 6. Why no hardware / composite change is needed
@@ -787,10 +876,14 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head,
    `CompositeCmd` (§1/§3) — that gives scheduler-managed tiling, K/V DMA
    streaming, and cross-tile pipelining. What is rejected is a **new**
    command kind that also absorbs the softmax merge + cross-tile register
-   lifetime + an IPCQ-push epilogue: it is large and special-purpose, and
+   lifetime. **Sizing note (§5.6 opt2):** if revisited, split it into **two**
-   its only delta over the hybrid (full softmax offload) is not justified
+   composites — `#1` = Q·Kᵀ (the *existing* composite + `scale`; lets DMA
-   at the current modelling fidelity (per-op CPU issue cost = 0; see §1).
+   prioritise K), `#2` = softmax + P·V + the online-softmax accumulator merge
-   Revisit if the cost model (§9) makes full offload measurably worthwhile.
+   (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
@@ -829,9 +922,9 @@ new primitive** — only the kernel restructuring in §5.2 (per KV head,
 | 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** | `G` replicated (M-fold) | 2-level static shard `C·P` | **§4 2-level reduce** (small `O`) | 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 head per CUBE** (`C=G`) | **Ring KV rotate** | **none** (KV blocks move, not `O`) | causal step-skip + boundary |
+| **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):
 
@@ -1001,13 +1094,13 @@ predicted default; revise on review.
 ### 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 head/CUBE,
+    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 head per CUBE
+    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
@@ -1020,7 +1113,7 @@ predicted default; revise on review.
     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 head per CUBE). If `C ≠ G`, the mapping needs revisiting
+    `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.
@@ -1031,3 +1124,20 @@ predicted default; revise on review.
     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) — removes
+    the GEMM-engine bubble 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.