gqa(adr/ddd): iteration-2 corrections — GQA reuse via M-fold, transpose caveat

Verified against sim_engine data path (memory_store, data_executor):

- GQA reuse does NOT need a broadcast op. The baseline's h_q==h_kv limit
  is its head-packing reshape hack, not a missing primitive. Correct fix:
  per-KV-head loop with G folded into matmul M dim (byte-conserving
  reshape) — runs today, timing correct (m=G*T_q), data mode runs.
- ADR-0061 broadcast demoted from 'the blocker' to optional convenience.
- Surfaced tl.trans = reshape-not-transpose (memory_store reshapes;
  data_executor np.matmul on reshaped operands) -> numeric parity is
  bounded; verification is structural/timing/determinism-first (matches
  SPEC perf-model purpose). Optional tl.transpose deferred.
- Reordered DDD phase plan (P1 GQA needs no new feature; P3 scratch_scope
  is the key scale feature); added open decisions 10.10 (transpose) and
  10.11 (GQA-via-M-fold finding).

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

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

@@ -14,10 +14,12 @@ KV-load-bound decode; sequence-parallel (Ring KV) for long-context prefill.
 kernbench work* — that code is the baseline this ADR upgrades to a real
 GQA, causal, long-context kernel.
 
-**Supporting ADRs (prerequisites):**
-- **ADR-0061** `tl.broadcast` — data-faithful GQA head reuse.
-- **ADR-0062** `tl.load_async` — non-blocking HBM tile load (KV prefetch).
-- **ADR-0063** `tl.scratch_scope` — per-tile scratch recycling.
+**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
@@ -55,9 +57,15 @@ loop. This **matters for this ADR's design** (see §1).
 *efficient GQA* kernel must lift:
 
 1. **No GQA reuse.** `h_q == h_kv == 1`
-   (`test_milestone_gqa_llama70b.py:137-142`). Real GQA (`h_q = G·h_kv`)
-   is blocked by the MemoryStore byte-conservation check on the symbolic
-   broadcast view → fixed by **ADR-0061**.
+   (`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
@@ -302,10 +310,16 @@ with tl.scratch_scope():                              # per-tile temporaries rec
 ```
 
 Notes:
-- `q_g` is the GQA-batched query: `[G, d]` (decode) or `[G, T_q, d]`
-  (prefill). One K/V tile load serves all `G` rows — the GQA reuse lever.
-  Where the matmul needs the KV head expanded to `G`, use
-  `tl.broadcast` (ADR-0061), never the symbolic `_view`.
+- `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
@@ -398,9 +412,15 @@ def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
 - `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 batched (`q_g = [G, d]`), so each K/V
-  tile is loaded once and reused by 8 rows (Q·Kᵀ is an `8×TILE` batched
-  GEMV, P·V an `8×d`). `tl.broadcast` expands the KV head where needed.
+  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`
@@ -474,20 +494,16 @@ abstraction.
 
 ## 8. Required kernbench changes
 
-**Genuinely new (each has a supporting ADR):**
-
-1. **GQA head broadcast in data mode** — **ADR-0061** (`tl.broadcast`).
-   *The* blocker for `h_q > h_kv`. Without it the GQA reuse lever cannot
-   run under `enable_data=True`.
-2. **Async HBM tile load / KV prefetch** — **ADR-0062** (`tl.load_async`
-   + `tl.wait`). The KV-load-bound overlap lever for decode/long-context.
-3. **Per-tile scratch recycling** — **ADR-0063** (`tl.scratch_scope`).
-   Removes the `S=16` ceiling so realistic context lengths run.
+**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 (`G` rows share each K/V load) — `tl.broadcast` +
-  2-D `tl.dot`.
+- **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` +
@@ -495,6 +511,21 @@ abstraction.
 - 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)`
@@ -557,10 +588,21 @@ projection **downstream**; score `S` and probs `P` are never materialised.
 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).
 
-1. **Correctness vs reference (Phase 2, `enable_data=True`):** for a
-   small `(G, T_q, S, N)`, kernel `O` matches a numpy FlashAttention
-   reference within fp tolerance — for all four cases. This is the test
-   that *requires* ADR-0061 (GQA) to even run with `h_q>h_kv`.
+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.