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>
This commit is contained in:
@@ -140,24 +140,44 @@ ctx.launch(
|
||||
`H_q > H_kv` is the change from the baseline's `H_q=H_kv=1`; it requires
|
||||
ADR-0061 to survive Phase 2.
|
||||
|
||||
### 4.3 GQA-batched Q load
|
||||
### 4.3 GQA-batched Q load — fold `G` into the matmul M dim (no broadcast)
|
||||
|
||||
For a KV head `kv`, its group is query heads `[kv·G, (kv+1)·G)`. The
|
||||
kernel loads those `G` rows and, where the matmul needs the KV head
|
||||
expanded to `G`, broadcasts the KV head — **not** the symbolic `_view`:
|
||||
For a KV head `kv`, its group is query heads `[kv·G, (kv+1)·G)`. Load
|
||||
those `G` rows and **reshape the group into the matmul M dimension** —
|
||||
this is the key to GQA reuse and it needs only a byte-conserving reshape,
|
||||
**not** a broadcast:
|
||||
|
||||
```python
|
||||
def load_q_group(q_ptr, kv, G, T_q, d, tl):
|
||||
# Q laid out [H_q, T_q, d]; the group is contiguous in head axis.
|
||||
return tl.load(q_ptr + kv*G*T_q*d, shape=(G, T_q, d), dtype="f16")
|
||||
|
||||
def kv_for_group(k_tile, G, tl): # ADR-0061
|
||||
# k_tile: [TILE, d] (one KV head). Expand to [G, TILE, d] for batched dot.
|
||||
return tl.broadcast(k_tile[None, :, :], (G, k_tile.shape[0], k_tile.shape[1]))
|
||||
# Q laid out [H_q, T_q, d]; the group is contiguous in the head axis.
|
||||
qg = tl.load(q_ptr + kv*G*T_q*d, shape=(G, T_q, d), dtype="f16")
|
||||
return _view(qg, (G*T_q, d)) # byte-conserving reshape → M = G*T_q
|
||||
```
|
||||
|
||||
(Whether the batched `tl.dot` consumes `[G,T_q,d]·[G,d,TILE]` directly or
|
||||
the kernel loops `g` is an Open Decision — §10.4.)
|
||||
Then per KV tile (`Kj`, `Vj` are `[TILE, d]`, one KV head, loaded once):
|
||||
|
||||
```python
|
||||
Sj = tl.dot(q_g, tl.trans(Kj)) * scale # [G*T_q, d]·[d, TILE] -> [G*T_q, TILE]
|
||||
Oj = tl.dot(P, Vj) # [G*T_q, TILE]·[TILE, d] -> [G*T_q, d]
|
||||
```
|
||||
|
||||
K/V are the **shared operands** (the matmul's `b`), loaded once and reused
|
||||
across all `G·T_q` M-rows automatically. The emitted `GemmCmd` carries
|
||||
`m = G·T_q`, so Phase-1 timing counts all `G` rows' work correctly.
|
||||
|
||||
> **Why not a leading batch axis** (`[G,T_q,d]·[G,d,TILE]`)? `tl.dot`
|
||||
> derives `m,k,n` from the **last two dims only** (`tl_context.py:220`)
|
||||
> and `GemmCmd` carries no batch dim — so a leading `G` axis is computed
|
||||
> correctly by `np.matmul` in Phase 2 (it broadcasts) but **under-counted
|
||||
> by `G×` in Phase-1 timing**. Folding into M fixes both. This supersedes
|
||||
> the earlier batched-vs-loop question.
|
||||
|
||||
> **Transpose caveat.** `tl.trans` is metadata-only and `MemoryStore.read`
|
||||
> *reshapes* (not transposes) (`memory_store.py:73`,
|
||||
> `data_executor.py:153-168` does `np.matmul` on the reshaped operands).
|
||||
> Fine for timing/structure and for the baseline's zero data; for
|
||||
> non-trivial numeric parity, store K pre-transposed as `[d, S_pe]` or add
|
||||
> a real `tl.transpose` (§10.10).
|
||||
|
||||
---
|
||||
|
||||
@@ -273,16 +293,17 @@ inner update — running `(m,l,O)` carried across steps as Python handles
|
||||
|
||||
## 6. The three new primitives — integration points
|
||||
|
||||
### 6.1 ADR-0061 `tl.broadcast`
|
||||
### 6.1 ADR-0061 `tl.broadcast` (OPTIONAL — not on the GQA critical path)
|
||||
GQA reuse is achieved by folding `G` into M (§4.3), so broadcast is **not**
|
||||
required to unblock `h_q > h_kv`. It remains a clean primitive for additive
|
||||
causal-mask construction across the `G·T_q` rows and for general kernels.
|
||||
- `tl_context.py`: new `broadcast(x, new_shape)` → emits `BroadcastCmd`,
|
||||
output handle from `_make_compute_out` (so `nbytes` is correct).
|
||||
- `data_executor.py`: `_execute_broadcast` = `np.broadcast_to(src,
|
||||
shape).copy()` → `store.write`.
|
||||
- Latency: `pe_math._compute_ns(prod(new_shape))`; logs `math/broadcast`,
|
||||
**no** `dma_read`.
|
||||
- Unblocks: the `tl.dot(q_g[G,…], kv_for_group(...)[G,…])` consuming a
|
||||
broadcast KV head passes the MemoryStore nbytes check
|
||||
(`memory_store.py:67-73`) that fails today.
|
||||
- Lowest priority; deliver after the efficiency/scale features land.
|
||||
|
||||
### 6.2 ADR-0062 `tl.load_async`
|
||||
- `tl_context.py`: `load_async(ptr, shape, dtype) → LoadFuture`; extend
|
||||
@@ -306,16 +327,18 @@ on a green predecessor; every phase keeps the existing baseline green.
|
||||
|
||||
| Phase | Deliverable | Gate |
|
||||
|---|---|---|
|
||||
| **P0** | ADR-0061 `tl.broadcast` + unit test | broadcast Phase 2 parity; downstream dot passes nbytes check |
|
||||
| **P1** | Real GQA in baseline decode kernel (`h_q=G·h_kv`), one-shot, no tiling, using broadcast | numpy parity at small `S`; K/V `dma_read_count` independent of `G` |
|
||||
| **P2** | Tree reduction replacing fan-out (decode SP) | `⌈log₂ N⌉` reduction rounds; parity preserved |
|
||||
| **P3** | ADR-0063 `tl.scratch_scope` + tiled sweep | long-`S` decode completes; parity |
|
||||
| **P4** | ADR-0062 `tl.load_async` + prefetch in sweep | latency < serial Σ(load+compute); parity |
|
||||
| **P5** | Causal masking + tile/step skip (prefill, ring) | GEMM count = triangular tile count; parity |
|
||||
| **P1** | Real GQA by restructuring (per KV head, `G`→M fold); one-shot, no tiling. **No new primitive.** | data-mode run completes for `h_q=G·h_kv`; K/V `dma_read_count` independent of `G`; `gemm` m=`G·T_q` |
|
||||
| **P2** | Tree reduction replacing fan-out (decode SP) | `⌈log₂ N⌉` reduction rounds; structure preserved |
|
||||
| **P3** | ADR-0063 `tl.scratch_scope` + tiled sweep | long-`S` decode completes (today caps at S=16) |
|
||||
| **P4** | ADR-0062 `tl.load_async` + prefetch in sweep | latency < serial Σ(load+compute) |
|
||||
| **P5** | Causal masking + tile/step skip (prefill, ring) | GEMM count = triangular tile count |
|
||||
| **P6** | Headline-scale milestone panels (real `G`, tiling, prefetch) + figures | sweep.json headline rows; figures render |
|
||||
| **P7** *(opt)* | ADR-0061 `tl.broadcast` (mask convenience); optional `tl.transpose` for numeric parity | broadcast unit test; asymmetric numeric parity |
|
||||
|
||||
P0–P1 alone deliver "GQA actually runs in data mode" — the single most
|
||||
important unblock. P2–P5 deliver efficiency. P6 is the eval payoff.
|
||||
**P1 alone delivers "GQA actually runs" — and it needs no new feature**,
|
||||
only the §4.3/§5 restructuring (the iteration-2 finding). P3 is the most
|
||||
important *new* feature (removes the S=16 scale ceiling). P4 is the
|
||||
efficiency lever. P2/P5 are algorithmic. P7 is optional polish.
|
||||
|
||||
---
|
||||
|
||||
@@ -324,10 +347,14 @@ important unblock. P2–P5 deliver efficiency. P6 is the eval payoff.
|
||||
Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (`tl`),
|
||||
ADR-0054 (eval bench). Mirrors ADR-0060 §11.
|
||||
|
||||
**Correctness (Phase 2, `enable_data=True`):**
|
||||
- `test_gqa_correctness.py`: for `(G∈{1,8}, T_q∈{1,16}, S, N∈{1,8})`,
|
||||
kernel `O` ≈ numpy FlashAttention reference. The `G=8` cases are the
|
||||
ones that cannot run today (ADR-0061 gate).
|
||||
**Runs + numeric (Phase 2, `enable_data=True`):**
|
||||
- `test_gqa_correctness.py`: for `(G∈{1,8}, T_q∈{1,16}, S, N∈{1,8})` the
|
||||
GQA kernel **completes** (no byte-conservation error) — the `G=8` cases
|
||||
that the baseline head-packing cannot express. This is gated on the
|
||||
§4.3 restructuring, **not** on a new primitive.
|
||||
- *Numeric parity (secondary):* `O ≈ numpy FlashAttention` for
|
||||
symmetric/identity inputs where reshape-as-transpose is exact; full
|
||||
asymmetric parity is gated on an optional `tl.transpose` (§10.10).
|
||||
|
||||
**Levers (op_log assertions):**
|
||||
- GQA amortisation: K/V `dma_read_count` constant as `G` grows 1→8;
|
||||
@@ -397,12 +424,14 @@ on root only. Attention needs root-only, so tree wins. If any downstream
|
||||
consumer needs O replicated (none found), we'd keep fan-out. *Assuming
|
||||
root-only is fine.*
|
||||
|
||||
### 10.4 Batched 2-D dot vs per-`g` loop for GQA **[recommend: try batched first]**
|
||||
`tl.dot` semantics for a leading group axis (`[G,T_q,d]·[G,d,TILE]`) need
|
||||
confirming in `tl_context.dot` (it computes `m,k,n` from the last two
|
||||
dims; `tl_context.py:213-229`). If batched dot over a leading axis isn't
|
||||
supported, fall back to a `for g in range(G)` loop (ADR-0061 A1) — correct
|
||||
but `G×` GEMM ops. *I assumed batched works or is a small extension.*
|
||||
### 10.4 GQA matmul shape — RESOLVED in iteration 2 **[fold G into M]**
|
||||
Earlier this asked "batched leading axis vs per-`g` loop." Resolved: do
|
||||
**neither** — fold `G·T_q` into the matmul M dimension (§4.3). A leading
|
||||
`G` axis would be computed by `np.matmul` in data mode but **under-counted
|
||||
`G×` in Phase-1 timing** (`GemmCmd` carries no batch dim,
|
||||
`tl_context.py:220`). Folding into M gives correct timing *and* correct
|
||||
data with a single byte-conserving reshape and one shared K/V operand. No
|
||||
new primitive, no per-`g` loop.
|
||||
|
||||
### 10.5 Persistent-vs-scoped arena split **[recommend: two explicit arenas]**
|
||||
ADR-0063 requires `m,l,O` to live outside the recycled scope. The cleanest
|
||||
@@ -438,13 +467,40 @@ parity tests must use tolerances that accommodate this, and we should not
|
||||
claim bf16 *accuracy* results — only timing/structure. Pre-existing
|
||||
simulator limitation; flagged so results aren't over-interpreted.
|
||||
|
||||
### 10.10 `tl.trans` is reshape-not-transpose **[decision: structural-first verify]**
|
||||
Discovered in iteration 2: `tl.trans` only swaps shape metadata
|
||||
(`tl_context.py:390`) and `MemoryStore.read` *reshapes* the bytes
|
||||
(`memory_store.py:73`); `data_executor._execute_gemm` then `np.matmul`s
|
||||
the reshaped operands (`data_executor.py:153-168`). So for **non-trivial
|
||||
numeric data**, `Q·Kᵀ` via `tl.trans(K)` is mathematically a reshape, not
|
||||
a transpose — wrong numbers (the baseline never notices: it runs on
|
||||
all-zeros and only asserts op_log structure). **Decision:** treat the
|
||||
simulator as a *performance* model (SPEC §0) — verify structure + timing +
|
||||
determinism first; treat numeric parity as a bounded secondary check
|
||||
(symmetric/identity inputs, or store K pre-transposed). A real
|
||||
data-materialising `tl.transpose` is a *candidate* further primitive if
|
||||
true numeric validation is ever required — **recommend deferring it**; it
|
||||
is not needed for the performance-modeling goal. *Confirm you're content
|
||||
with structural-first verification.*
|
||||
|
||||
### 10.11 GQA reuse = fold `G` into M, NOT a broadcast op **[finding, iteration 2]**
|
||||
The headline mechanism finding. The deferred-GQA limitation
|
||||
(`test_milestone_gqa_llama70b.py:137-142`) was blamed on a missing
|
||||
broadcast / the MemoryStore nbytes check. Iteration 2 shows the real
|
||||
cause is the baseline's head-packing reshape, and the fix is kernel
|
||||
restructuring (§4.3), which needs **no new primitive**. This demoted
|
||||
ADR-0061 from "the blocker" to "optional convenience" and reordered the
|
||||
phase plan (§7). Recorded explicitly because it changes the project's
|
||||
critical path: *GQA itself is not gated on any new feature; scale
|
||||
(ADR-0063) and efficiency (ADR-0062) are the features that matter.*
|
||||
|
||||
---
|
||||
|
||||
## 11. Risks
|
||||
|
||||
| Risk | Likelihood | Mitigation |
|
||||
|---|---|---|
|
||||
| Batched dot over leading axis unsupported | med | per-`g` fallback (§10.4); still correct |
|
||||
| Numeric parity blocked by trans-as-reshape | med | structural-first verify (§10.10); pre-transpose K or defer `tl.transpose` |
|
||||
| Tree pairs not 1-hop on installed mesh | med | verify SFR neighbour table (§10.2); relabel PEs |
|
||||
| scratch_scope use-after-scope bug | med | two-arena discipline in shared helper; debug poison |
|
||||
| prefetch buffer recycled while in flight | med | live `LoadFuture` buffers in persistent arena (§10.5) |
|
||||
|
||||
Reference in New Issue
Block a user