ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(ddd): rewrite DDD-0060 to match current ADR (two-kernel composite hybrid)

Browse Source 
Supersede the old greenlet-primitive/load_async/single-kernel DDD. Align to
ADR-0060's composite-hybrid + hierarchical CUBE-Group SP with two kernels:
decode+SP (head-replicated, contiguous C×P static shard, 2-level
reduce-to-root) and prefill+SP (1 Q head/CUBE, Ring KV, no reduce). Updated
file plan (lazy tl.load, scratch_scope, broadcast-opt; evolve
_attention_mesh_mlo_2d/_kv; DPPolicy.cube_start; llama70b_4sip.yaml),
placement (contiguous shared KV), phase plan (P1..P8), verification, perf
model, risks, glossary. Defers design rationale to ADR-0060 (now
authoritative); open items point to ADR-0060 §B. ADR-0064 cost model noted.
Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-04 23:08:02 -07:00
parent 973a452b77
commit c699e2ddfc
1 changed files with 221 additions and 405 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md
+221 -405
View File
@@ -1,39 +1,42 @@
# Detailed Design Document — AHBM GQA Fused Attention (ADR-0060) # Detailed Design Document — AHBM GQA Fused Attention (ADR-0060)
**Status:** Draft (companion to ADR-0060, Proposed) **Status:** Draft (companion to ADR-0060, Proposed)
**Scope:** implementation-ready design for an efficient GQA FlashAttention **Scope:** implementation-ready plan for the **two** GQA FlashAttention
kernel on AHBM in the kernbench simulator. kernels (decode=reduce, prefill=ring) on AHBM in kernbench.
**Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read ADR-0060 **Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read **ADR-0060
first (decision record), then this (the how), then the *Open Decisions* first** — it is now the authoritative design record (TL;DR has full
section (§10) for the choices made autonomously that need your sign-off. pseudocode for both kernels). This DDD is the *how to build it*: file plan,
phase plan, helper signatures, tests. It does **not** re-derive the design;
it points to ADR-0060 sections.
> **How this document was produced.** It is the result of several > **Alignment note (this revision).** ADR-0060 moved to a **composite
> design↔evaluation iterations against the real kernbench source > hybrid + hierarchical CUBE-Group SP** design with **two kernels**. This
> (`src/kernbench/triton_emu/tl_context.py`, `sim_engine/*`, > DDD is rewritten to match. The earlier DDD (greenlet-primitive,
> `components/builtin/pe_*`, the `_attention_mesh_*` kernels, and the GQA > `tl.load_async`, single unified kernel) is superseded.
> tests). Every claim about what exists is grounded in a file path; every
> claim about what is new is tied to a supporting ADR (0061/0062/0063).
--- ---
## 1. Goal and success criteria ## 1. Goal and success criteria
**Goal:** an *efficient* GQA fused attention kernel that runs on AHBM in **Goal:** two *efficient* GQA fused-attention kernels on AHBM in kernbench —
kernbench, across the four operating points decode/prefill × **decode+SP** (head-replicated, KV static shard, 2-level reduce) and
single-/multi-user, at realistic (not just validation) scale, in both **prefill+SP** (1 Q head per CUBE, Ring KV) — at realistic scale, in both
timing (`enable_data=False`) and data (`enable_data=True`) modes. timing (`enable_data=False`) and data (`enable_data=True`) modes.
**Success criteria:** **Success criteria:**
1. Correct: kernel output matches a numpy FlashAttention reference within 1. Correct: kernel output matches a numpy FlashAttention reference within fp
fp tolerance, with **real GQA** (`H_q=64, H_kv=8, G=8`). tolerance, with **real GQA** (`H_q=64, H_kv=8, G=8`).
2. Efficient: the three levers actually fire and are observable in op_log: 2. Efficient — levers observable in op_log (ADR-0060 §11):
GQA K/V-load amortisation, KV prefetch overlap, `⌈log₂ N⌉` reduction. GQA K/V-load amortisation; decode 2-level reduce (`⌈log₂P⌉` + center-mesh
over `C`, not `C·P−1`); prefill ring (no reduce, KV rotates); load/compute
overlap (lazy `tl.load` + composite streaming); composite GEMM offload.
3. Scales: context length not capped by the 1 MiB scratch bump allocator. 3. Scales: context length not capped by the 1 MiB scratch bump allocator.
4. Deterministic and SPEC-compliant (all latency from modelled events). 4. Deterministic and SPEC-compliant (all latency from modelled events).
**Non-goals:** RoPE / QKV projection / KV-cache writes (upstream **Non-goals:** RoPE / QKV projection / KV-cache writes (upstream
`qkv_rope`); output projection (downstream); cross-CUBE query-head split; `qkv_rope`); output projection (downstream); cross-**SIP** head split
cycle-accurate microarchitecture (SPEC §5). (a head stays in one SIP, ADR-0060 §0); cycle-accurate microarchitecture
(SPEC §5); a bespoke "flash-composite" command kind (ADR-0060 §8 item 4).
--- ---
@@ -41,31 +44,33 @@ cycle-accurate microarchitecture (SPEC §5).
``` ```
runtime_api (host) runtime_api (host)
RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy, kernel launch) RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy(+cube_start), launch)
│ KernelLaunchMsg │ KernelLaunchMsg
▼ ▼
sim_engine sim_engine
GraphEngine(enable_data=…) engine.py GraphEngine(enable_data=…) engine.py
DataExecutor (Phase 2) data_executor.py ← ADR-0061 adds _execute_broadcast DataExecutor (Phase 2) data_executor.py ← _execute_broadcast (ADR-0061, optional)
MemoryStore (nbytes check) memory_store.py ← the GQA blocker lives here MemoryStore memory_store.py
│ │
▼ ▼
components / PE pipeline (per PE) components / PE pipeline (per PE)
PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,FETCH_STORE,TCM,IPCQ} PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM}
greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py
CompositeCmd → tile plan → engines pe_scheduler.py (ADR-0014 D6)
│ │
▼ ▼
tl programming model tl_context.py tl programming model tl_context.py
load/store/dot/softmax/max/sum/exp/maximum/send/recv/trans/… load (→ lazy, ADR-0062) / store / composite / dot / max/sum/exp/maximum
← ADR-0061 adds broadcast/repeat send/recv/recv_async / scratch_scope (ADR-0063) / broadcast (ADR-0061, opt)
← ADR-0062 adds load_async (+ wait arm)
← ADR-0063 adds scratch_scope
``` ```
The attention kernels are **bench-side Python** (`src/kernbench/benches/`) Both kernels are **bench-side Python** (`src/kernbench/benches/`) run as
invoked via `ctx.launch(name, kernel_fn, *args)`; they run as greenlet greenlet kernels emitting `tl.*` ops; the **GEMMs are `tl.composite`**
kernels that emit `tl.*` ops. No component (device-side) code changes are (scheduler-managed tiling + K/V DMA streaming, existing `CompositeCmd` — no
needed for the algorithm — only the three `tl`/sim_engine primitives. new command kind), the softmax merge + reduction/ring stay kernel-level
(ADR-0060 §1). The remote work already added building blocks:
`DPPolicy.cube_start`, `_attention_mesh_mlo_2d.py` (2D cube reduce —
currently AllReduce, to become reduce-to-root), `topologies/llama70b_4sip.yaml`.
--- ---
@@ -75,424 +80,234 @@ needed for the algorithm — only the three `tl`/sim_engine primitives.
| File | Change | ADR | | File | Change | ADR |
|---|---|---| |---|---|---|
| `src/kernbench/triton_emu/tl_context.py` | add `broadcast`/`repeat`; `load_async` + `wait` arm; `scratch_scope` ctx-mgr | 0061/0062/0063 | | `src/kernbench/triton_emu/tl_context.py` | `tl.load` → **lazy** (non-blocking + auto-wait on first use); `scratch_scope` ctx-mgr; `broadcast` (optional) | 0062/0063/0061 |
| `src/kernbench/common/pe_commands.py` | add `BroadcastCmd`; `DmaReadCmd.blocking` flag (or `DmaReadAsyncCmd`) | 0061/0062 | | `src/kernbench/components/builtin/pe_dma.py` | non-blocking read path: post DMA, resolve on a scheduled event (the `recv_async` pattern for loads) | 0062 |
| `src/kernbench/sim_engine/data_executor.py` | add `_execute_broadcast`; handle async dma_read resolution | 0061/0062 | | `src/kernbench/triton_emu/kernel_runner.py` | auto-wait: yield a handle's pending load event at first consume; lazy-load dispatch arm | 0062 |
| `src/kernbench/components/builtin/pe_dma.py` | non-blocking read path returns a future-resolving event | 0062 | | `src/kernbench/sim_engine/data_executor.py` | `_execute_broadcast` (optional) | 0061 |
| `src/kernbench/sim_engine/event_log.py` / op_log | recognise `broadcast` op_name (math) | 0061 | | `src/kernbench/components/builtin/pe_cpu.py` + cost table | per-op-type CPU issue cost (replaces uniform `dispatch_cycles`) | **0064** (separate) |
### 3.2 New kernel + bench files No new GEMM command kind: the two GEMMs use the **existing** `CompositeCmd`.
### 3.2 Kernel + bench files
| File | Role | | File | Role |
|---|---| |---|---|
| `src/kernbench/benches/_gqa_attention.py` | the unified GQA kernel (4 cases via flags), with a shared `_flash_tile_sweep` helper and `_tree_reduce` helper | | `src/kernbench/benches/_gqa_decode.py` | decode+SP kernel — evolve `_attention_mesh_mlo`/`_mlo_2d` to composite-hybrid + **2-level reduce-to-root** (Level-2 PE tree + Level-1 center-mesh; currently 2D AllReduce) |
| `src/kernbench/benches/_gqa_helpers.py` | `valid_len`, `tile_*` predicates, mask builders, tree topology (children/parent dirs), GQA-batched Q load | | `src/kernbench/benches/_gqa_prefill.py` | prefill+SP kernel — evolve `_attention_mesh_kv` to **head-parallel (1 Q head/CUBE) Ring KV** + composite-hybrid |
| extend `src/kernbench/benches/milestone_gqa_llama70b.py` | add headline-scale panels with real `G`, prefetch, tiling | | `src/kernbench/benches/_gqa_helpers.py` | `head_of_group`, `valid_len_2level`, `init_running`, tree/center-mesh topology, `hierarchical_reduce_and_store`, ring helpers, causal `block_*` predicates + mask builders |
| extend `src/kernbench/benches/milestone_gqa_llama70b.py` | headline panels: real `G`, `C=G=8`, contiguous KV, 4-SIP |
| `topologies/llama70b_4sip.yaml` | **exists** (remote) — 4 SIP × 16-CUBE 4×4 × 8 PE |
### 3.3 New tests ### 3.3 New tests
| File | Covers | | File | Covers |
|---|---| |---|---|
| `tests/attention/test_gqa_correctness.py` | numpy reference parity, all 4 cases, real GQA | | `tests/attention/test_gqa_decode.py` | decode reduce: GQA reuse (dma_read ⟂ `G`), 2-level reduce step count, root-only output, long context |
| `tests/attention/test_gqa_levers.py` | GQA load amortisation, tree step count, causal skip, prefetch overlap | | `tests/attention/test_gqa_prefill.py` | prefill ring: no `(m,ℓ,O)` reduce, `C` KV rotations, causal step-skip, per-CUBE distributed output |
| `tests/test_tl_broadcast.py` | ADR-0061 unit | | `tests/test_tl_lazy_load.py` | ADR-0062 unit (overlap, auto-wait correctness, op_log parity) |
| `tests/test_tl_load_async.py` | ADR-0062 unit |
| `tests/test_tl_scratch_scope.py` | ADR-0063 unit | | `tests/test_tl_scratch_scope.py` | ADR-0063 unit |
| `tests/test_tl_broadcast.py` | ADR-0061 unit (optional) |
| `tests/test_dppolicy_cube_start.py` | **exists** (remote) — sub-mesh placement |
--- ---
## 4. Data model and launch contract ## 4. Data model and launch contract
### 4.1 Tensor placement (DPPolicy) ### 4.1 Tensor placement — CUBE Group, 2-level SP, contiguous KV
Using `DPPolicy(cube, pe, num_cubes, num_pes)` (`policy/placement/dp.py`): A `CUBE Group` is `C` CUBEs in one SIP owning one KV head (ADR-0060 §0).
Placement via `DPPolicy(..., cube_start=...)`:
| Tensor | single-user (PE ring in 1 cube) | multi-user (cube ring) | | Tensor | Decode+SP | Prefill+SP (`C=G`) |
|---|---|---| |---|---|---|
| `Q` (group, replicated) | `cube=replicate, pe=replicate, num_pes=N` | `cube=replicate, pe=replicate, num_cubes=N` | | `Q` | `replicate` (all `G` heads, M-folded) over `C×P` ranks | **head-parallel**: CUBE `i` holds Q head `i` |
| `K`, `V` (sequence-sharded) | `cube=replicate, pe=row_wise, num_pes=N` | `cube=row_wise, pe=replicate, num_cubes=N` | | `K`,`V` | sequence-sharded **two axes** (`cube` row_wise over `C`, `pe` row_wise over `P`) | sequence-sharded over `C` CUBEs (the ring blocks) |
| `O` (group, replicated) | as `Q` | as `Q` | | `O` | reduced to CUBE-Group root (all `G` heads) | per-CUBE (one head each), distributed |
This matches the baseline `_PANEL_DISPATCH` **KV cache is *contiguous* `C×P` position blocks** (rank `r` = `[r·B,(r+1)·B)`),
(`milestone_gqa_llama70b.py:64-81`). `row_wise` over the sequence (M) axis **shared by both kernels** — prefill writes it (via `qkv_rope`), decode reads
realises the round-robin shard (§2.1 of ADR-0060); the contiguous per-PE + extends it, no reshard. Contiguous is required for prefill causal skip
buffer is the row-wise slice. (ADR-0060 §2.1). `cube_start` offsets the group's CUBE sub-mesh within the
4×4 SIP mesh.
`rank_axis` selects which `tl.program_id` carries the ring rank: `0` for `--device` enumerates **SIPs**; one SIP-device runs its 2 CUBE Groups
single-user (PE id in cube), `1` for multi-user (cube level; only PE 0 of (2 KV heads); the head is picked by CUBE coordinate (`head_of_group`), not an
each cube participates) — kept from the baseline so the SFR installs in-kernel loop (ADR-0060 §0).
(`configure_sfr_intracube_pe_ring` / `configure_sfr_intercube_multisip`)
apply unchanged.
### 4.2 Launch signature ### 4.2 Launch signatures
```python ```python
ctx.launch( # decode+SP
f"{panel}_gqa", gqa_attention_kernel, ctx.launch(f"{panel}_gqa_decode", gqa_decode_sp,
q, k, v, o, # Tensors (TensorArg) q, k, v, o, counter, start_pe, start_cube, C, P, softmax_scale)
S_q, S_kv_per_rank, H_q, H_kv, D_HEAD, n_ranks, # scalars # prefill+SP
rank_axis, mode, counter, start_pe, # scalars (mode: 0=decode,1=prefill) ctx.launch(f"{panel}_gqa_prefill", gqa_prefill_sp,
softmax_scale, # float scalar q, k, v, o, T_q, S_kv_local, d, C, softmax_scale, q_block, cube_start)
_auto_dim_remap=False, # keep d_head from colliding (baseline note)
)
``` ```
`H_q > H_kv` is the change from the baseline's `H_q=H_kv=1`; it requires Full per-kernel I/O contract: ADR-0060 §0.5.3/§0.5.4 (note the **output head
ADR-0061 to survive Phase 2. distribution differs** — decode root-gathered, prefill distributed).
### 4.3 GQA-batched Q load — fold `G` into the matmul M dim (no broadcast) ### 4.3 GQA reuse — fold `G` into the matmul M dim (decode)
For a KV head `kv`, its group is query heads `[kv·G, (kv+1)·G)`. Load Decode replicates Q and **M-folds** the `G` query heads into the GEMM row
those `G` rows and **reshape the group into the matmul M dimension** — dim so one `Q·Kᵀ` serves all `G` heads sharing one `K` (ADR-0060 §0, §5.2):
this is the key to GQA reuse and it needs only a byte-conserving reshape, `Q` group `[G, T_q, d]` → `[G·T_q, d]` (byte-conserving `_view`); the
**not** a broadcast: composite carries `m = G·T_q` (timing counts all `G` rows). Prefill does
**not** M-fold (1 head/CUBE). Caveat: `tl.trans` is reshape-not-transpose →
```python store `K` pre-transposed `[d, S/(C·P)]` (ADR-0060 §3, §B).
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 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
```
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).
--- ---
## 5. Kernel design (concrete, real `tl`) ## 5. Kernel design
### 5.1 Shared tile sweep helper **The full pseudocode for both kernels (and the 3 decode CPU-pipelining
variants) lives in ADR-0060 TL;DR + §5.** Implementation notes only here.
```python ### 5.1 Decode (reduce) — `_gqa_decode.py`
def flash_tile_sweep(q_g, k_base, v_base, n_tiles, TILE, d, scale, - Inner tile = §3 composite-hybrid: `Sj = composite(q_g, Kⱼ)·scale` → softmax
q_block, causal, tl, PREFETCH=2): MATH → `Oj = composite(P, Vⱼ)` → running merge (ADR-0060 §3).
"""One PE's FlashAttention sweep over its local KV tiles. - **Ship opt3** (software pipelining: issue next tile's `Q·Kᵀ` before this
Returns running (m, l, O) in the persistent arena.""" tile's softmax; `Sj` in a persistent double buffer) — removes the
m, l, O = _init_running(q_g.shape, d, tl) # persistent arena GEMM-engine bubble, no new command kind (ADR-0060 §5.6). opt1 is the naïve
f_k, f_v = {}, {} baseline; opt2 (`ex_composite`) waits on ADR-0064.
for j in range(min(PREFETCH, n_tiles)): # prime (ADR-0062) - Combine = **`hierarchical_reduce_and_store`** (ADR-0060 §4): Level-2 PE
if causal and tile_all_future(j, q_block): continue reduce-to-root tree (intra-CUBE, the cube's KV slice further split across
f_k[j] = tl.load_async(k_base + j*TILE*d, (TILE, d)) its `P` PEs — decision (a)) → Level-1 center-root CUBE-mesh reduce
f_v[j] = tl.load_async(v_base + j*TILE*d, (TILE, d)) (intra-CUBE-Group). Data-driven, level-pipelined, root-only.
for j in range(n_tiles):
if causal and tile_all_future(j, q_block): continue
with tl.scratch_scope(): # recycle temporaries (ADR-0063)
Kj = tl.wait(f_k[j])
Sj = tl.dot(q_g, tl.trans(Kj)) * scale
mask = make_mask_tile(j, q_block) if (causal and tile_partial(j, q_block)) else None
if mask is not None: Sj = Sj + mask
m_j = tl.max(Sj, axis=-1)
m_new = tl.maximum(m, m_j)
P = tl.exp(Sj - m_new)
corr = tl.exp(m - m_new)
l = l * corr + tl.sum(P, axis=-1)
Vj = tl.wait(f_v[j])
O = O * corr + tl.dot(P, Vj)
m = m_new
nxt = j + PREFETCH # refill pipeline
if nxt < n_tiles and not (causal and tile_all_future(nxt, q_block)):
f_k[nxt] = tl.load_async(k_base + nxt*TILE*d, (TILE, d))
f_v[nxt] = tl.load_async(v_base + nxt*TILE*d, (TILE, d))
return m, l, O
```
> Note: `m, l, O` are reassigned inside the `with` but their **storage** ### 5.2 Prefill (ring) — `_gqa_prefill.py`
> must be the persistent arena, not the recycled scope. Concretely the - Head-parallel: CUBE `i` owns Q head `i` + KV slice `i`. **No reduce** (each
> merge ops write to pre-allocated persistent handles (double-buffered); CUBE outputs a distinct head). Ring rotates KV blocks; GQA reuse via the
> ADR-0063 §D3 describes the persistent-vs-scoped arena split. This is the rotation (ADR-0060 §5.5).
> one subtlety the implementer must get right — see §10.5. - **K/V are `tl.load`'d TCM handles** (not `tl.ref`) so the ring can
`tl.send`/`recv_async` them; `K` pre-transposed `[d, S/C]`. Receive buffers
**ping-pong** (persistent arena) — recv into a *separate* buffer while
computing/sending the current one (do not clobber).
- Causal step-skip + boundary mask; `recv_async` overlaps next block's
receive with current compute.
- Within a CUBE: tile the Q head's `T_q` rows across `P` PEs (disjoint output
rows → no intra-CUBE reduce); fall back to KV-block split only if `T_q < P`
(ADR-0060 §B decode-split item 3).
### 5.2 Tree reduction helper ### 5.3 Shared helpers — `_gqa_helpers.py`
`head_of_group(cube_id) = cube_id // C`; `init_running(...)` → persistent
```python `(m=-inf, ℓ=0, O=0)`; `valid_len_2level(...)` (contiguous block extent);
def tree_reduce_and_store(m, l, O, pe_id, N, o_ptr, tl): `hierarchical_reduce_and_store(...)`; tree/center-mesh child/parent dirs;
for cdir in tree_children_dirs(pe_id, N): # static tree (§4 ADR-0060) `block_all_future`/`block_partial` + `causal_mask`.
m_c = tl.recv(cdir, m.shape); l_c = tl.recv(cdir, l.shape)
O_c = tl.recv(cdir, O.shape)
mn = tl.maximum(m, m_c); sa = tl.exp(m - mn); sb = tl.exp(m_c - mn)
l = l*sa + l_c*sb; O = O*sa + O_c*sb; m = mn
if is_root(pe_id, N):
tl.store(o_ptr, O / l)
else:
pdir = parent_dir(pe_id, N)
tl.send(pdir, m); tl.send(pdir, l); tl.send(pdir, O)
```
`tree_children_dirs` / `parent_dir` map the binary tree (level pairs in
ADR-0060 §4) onto physical `N/S/E/W` neighbours — this mapping is the
tuning item ADR-0060 §9.2 and an Open Decision (§10.2).
### 5.3 Top-level kernel (all four cases)
```python
def gqa_attention_kernel(q_ptr,k_ptr,v_ptr,o_ptr,
S_q,S_kv_per_rank,H_q,H_kv,d,n_ranks,
rank_axis,mode,counter,start_pe,scale,*,tl):
if rank_axis != 0 and tl.program_id(axis=0) != 0: # multi-user: only PE0 per cube
return
pe_id = tl.program_id(axis=rank_axis)
G = H_q // H_kv
causal = (mode == 1) # prefill
q_block = (counter, S_q) if causal else (counter, 1)
# decode small-S fast path = one tile; otherwise tiled (TILE from config)
for kv in range(H_kv): # one KV head per iter (or PE-mapped)
q_g = load_q_group(q_ptr, kv, G, S_q, d, tl)
my_len = valid_len(counter, start_pe, pe_id, n_ranks) if mode==0 else S_kv_per_rank
n_tiles = ceil(my_len / TILE)
m,l,O = flash_tile_sweep(q_g, k_base(kv), v_base(kv), n_tiles, TILE, d,
scale, q_block, causal, tl)
if n_ranks == 1:
tl.store(o_ptr_for(kv), O / l)
elif mode == 1 and is_ring(...): # prefill ring already folded across steps
ring_fold_then_store(...) # §5.4
else:
tree_reduce_and_store(m,l,O, pe_id, n_ranks, o_ptr_for(kv), tl)
```
This collapses the baseline's two separate kernels (`_attention_mesh_kv`,
`_attention_mesh_mlo`) into one parameterised kernel; the `mode`/`n_ranks`
flags select the case. (Keeping them separate is also fine — Open
Decision §10.6.)
### 5.4 Ring (prefill SP) specifics
The Ring path keeps the baseline's structure (`_attention_mesh_kv.py`)
with three upgrades:
- `recv_async` for the next step's K/V (overlap) — `tl_context.py:543`.
- causal step-skip: `if step_all_future(step, q_block): continue` before
folding (skips ≈half the steps).
- GQA: `q_g` is the `G`-row group; `tl.broadcast` for the KV head.
The fold across ring steps is the same online-softmax merge as §5.2's
inner update — running `(m,l,O)` carried across steps as Python handles
(exactly as the baseline does today).
--- ---
## 6. The three new primitives — integration points ## 6. Supporting primitives — integration points
### 6.1 ADR-0061 `tl.broadcast` (OPTIONAL — not on the GQA critical path) ### 6.1 ADR-0062 lazy `tl.load` (efficiency — overlap)
GQA reuse is achieved by folding `G` into M (§4.3), so broadcast is **not** - `tl.load` issues `DmaReadCmd` non-blocking, returns a handle with a pending
required to unblock `h_q > h_kv`. It remains a clean primitive for additive event; the runtime **auto-inserts the wait at first consume** (generalises
causal-mask construction across the `G·T_q` rows and for general kernels. the `recv_async`/wait pattern, `kernel_runner.py:248-285`).
- `tl_context.py`: new `broadcast(x, new_shape)` → emits `BroadcastCmd`, - `pe_dma.py`: non-blocking read; DMA occupies the (capacity-1) read channel,
output handle from `_make_compute_out` (so `nbytes` is correct). overlaps compute on PE_GEMM/PE_MATH.
- `data_executor.py`: `_execute_broadcast` = `np.broadcast_to(src, - **Global** semantics change → existing goldens regenerate (ADR-0062 D3).
shape).copy()` → `store.write`. - op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` stable.
- Latency: `pe_math._compute_ns(prod(new_shape))`; logs `math/broadcast`,
**no** `dma_read`.
- Lowest priority; deliver after the efficiency/scale features land.
### 6.2 ADR-0062 `tl.load_async` ### 6.2 ADR-0063 `tl.scratch_scope` (scale)
- `tl_context.py`: `load_async(ptr, shape, dtype) → LoadFuture`; extend - ctx-mgr saving/restoring `_scratch_cursor`. Persistent arena for `(m,ℓ,O)`
`wait` to resolve `LoadFuture`. (+ decode opt3 `Sj` double buffer, + prefill ring ping-pong buffers, +
- `pe_dma.py`: non-blocking read path; transfer occupies the read channel in-flight lazy-load buffers) lives **outside** the scope. Removes `S=16`.
in parallel; completion is a scheduled event.
- op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` metric stable.
### 6.3 ADR-0063 `tl.scratch_scope` ### 6.3 ADR-0061 `tl.broadcast` (optional)
- `tl_context.py`: context manager saving/restoring `_scratch_cursor`. - Not on the GQA critical path (M-fold gives reuse). Convenience for additive
- Persistent arena: `m,l,O` (+ double-buffer) and live `LoadFuture` mask construction. Lowest priority.
buffers allocated outside the scope.
- Removes the `S=16` cap (`test_milestone_gqa_llama70b.py:123-148`). ### 6.4 ADR-0064 per-op-type CPU issue cost (separate ADR)
- Replaces uniform `dispatch_cycles=0`. Makes the hybrid's CPU-saturation
win (and the decode opt2 `ex_composite` fewer-issues win) **measurable**.
Land before claiming hybrid latency wins (ADR-0060 §9).
--- ---
## 7. Phased implementation plan (test-first per CLAUDE.md) ## 7. Phased implementation plan (test-first per CLAUDE.md)
Each phase is an independent Phase-1→Phase-2 cycle. Ordered so each builds Each phase is an independent Phase-1→Phase-2 cycle; each keeps the baseline
on a green predecessor; every phase keeps the existing baseline green. green.
| Phase | Deliverable | Gate | | Phase | Deliverable | Gate |
|---|---|---| |---|---|---|
| **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` | | **P1** | Real GQA, decode, composite-hybrid M-fold; one-shot, no tiling. **No new primitive** (existing `CompositeCmd`). | data-mode completes `h_q=G·h_kv`; K/V `dma_read_count` ⟂ `G`; composite tile plan `m=G·T_q` |
| **P2** | Tree reduction replacing fan-out (decode SP) | `⌈log₂ N⌉` reduction rounds; structure preserved | | **P2** | Decode **2-level reduce-to-root** (Level-2 PE tree + Level-1 center-mesh), replacing the 2D AllReduce | reduce rounds = `⌈log₂P⌉`+center-mesh; **root-only** output; Level-1 on CUBE NOC, Level-2 on PE IPCQ |
| **P3** | ADR-0063 `tl.scratch_scope` + tiled sweep | long-`S` decode completes (today caps at S=16) | | **P3** | **ADR-0063** `scratch_scope` + tiled sweep | long-`S` decode completes (today caps S=16); scratch O(1) |
| **P4** | ADR-0062 `tl.load_async` + prefetch in sweep | latency < serial Σ(load+compute) | | **P4** | **ADR-0062** lazy `tl.load` + composite K/V streaming | tiled-sweep latency < serial `Σ(load+compute)`; goldens regenerated |
| **P5** | Causal masking + tile/step skip (prefill, ring) | GEMM count = triangular tile count | | **P5** | Decode **opt3** software pipelining | GEMM engine non-idle across tiles (structural / op_log) |
| **P6** | Headline-scale milestone panels (real `G`, tiling, prefetch) + figures | sweep.json headline rows; figures render | | **P6** | Prefill **head-parallel Ring KV** + causal step-skip | `C` KV rotations, no `(m,ℓ,O)` reduce, per-CUBE distributed `O`, causal skip ≈½ |
| **P7** *(opt)* | ADR-0061 `tl.broadcast` (mask convenience); optional `tl.transpose` for numeric parity | broadcast unit test; asymmetric numeric parity | | **P7** | Contiguous shared KV layout + 4-SIP topology + headline milestone panels (real `G`, `C=G`) | shared layout (no reshard); headline sweep.json rows |
| **P8** *(opt)* | **ADR-0064** cost model + decode **opt2** `ex_composite`; **ADR-0061** broadcast | cost-model goldens; opt2 fewer-issues measurable |
**P1 alone delivers "GQA actually runs" — and it needs no new feature**, P1 delivers "GQA runs" on the composite path. P3 is the key *scale* feature;
only the §4.3/§5 restructuring (the iteration-2 finding). P3 is the most P4 the overlap lever; P2/P6 the SP algorithms.
important *new* feature (removes the S=16 scale ceiling). P4 is the
efficiency lever. P2/P5 are algorithmic. P7 is optional polish.
--- ---
## 8. Verification plan (concrete) ## 8. Verification plan (concrete)
Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (`tl`), Mirrors ADR-0060 §11; grounded in SPEC R2/R5, ADR-0023/0025, ADR-0046, ADR-0054.
ADR-0054 (eval bench). Mirrors ADR-0060 §11.
**Runs + numeric (Phase 2, `enable_data=True`):** **Runs + numeric (`enable_data=True`):**
- `test_gqa_correctness.py`: for `(G∈{1,8}, T_q∈{1,16}, S, N∈{1,8})` the - decode & prefill complete for `(G∈{1,8}, T_q∈{1,16}, S, C∈{1,8}, P∈{1,8})`
GQA kernel **completes** (no byte-conservation error) — the `G=8` cases without byte-conservation error (the `G=8` cases baseline cannot express).
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 - *Numeric parity (secondary):* `O ≈ numpy FlashAttention` for
symmetric/identity inputs where reshape-as-transpose is exact; full symmetric/identity inputs (reshape-as-transpose exact); full asymmetric
asymmetric parity is gated on an optional `tl.transpose` (§10.10). parity gated on a real `tl.transpose` (deferred).
**Levers (op_log assertions):** **Levers (op_log):**
- GQA amortisation: K/V `dma_read_count` constant as `G` grows 1→8; - GQA amortisation: K/V `dma_read_count` ⟂ `G`; compute scales with `G`.
`gemm` work grows with `G`. - Decode reduce: rounds = `⌈log₂P⌉` + center-mesh (not `C·P−1`); result at
- Tree: decode-SP `ipcq_send`/`recv` counts = tree edges, reduction one rank; Level-1=CUBE NOC, Level-2=PE IPCQ.
rounds = `⌈log₂ N⌉` (not `N−1`). - Prefill ring: `C` KV rotations, **no** reduce, each CUBE writes a distinct
- Causal skip: prefill `gemm_count` = lower-triangular tile count. head.
- Prefetch overlap: tiled-sweep latency strictly < `Σ(load+compute)`. - Causal skip: GEMM/step count = lower-triangular.
- Scratch: peak `_scratch_cursor` bounded by one tile, independent of - Overlap: tiled-sweep latency < `Σ(load+compute)`.
`n_tiles`. - Scratch: peak cursor bounded by one tile.
**Invariants:** **Invariants:** determinism (identical op_log + latency); every routed
- Determinism: identical inputs → identical op_log + latency (SPEC §0.1). request latency > 0 (SPEC §0.1).
- Every routed request latency > 0 (SPEC §0.1).
--- ---
## 9. Performance model (what we expect to see) ## 9. Performance model (expected)
Rough per-PE decode cost (KV-load-bound), with prefetch overlap:
Per-rank decode (KV-load-bound), composite-streamed + overlapped:
``` ```
t_decode_pe ≈ max( DMA(K+V over S_pe tiles), compute(QKᵀ + PV over S_pe, ×G amortised) ) t_decode_rank ≈ max( DMA(K+V over S/(C·P) tiles), compute(QKᵀ+PV, ×G amortised) )
+ t_tree_reduce(⌈log₂ N⌉ hops, payload O[G,d]) + t_reduce( ⌈log₂P⌉ intra-CUBE + center-mesh over C inter-CUBE )
``` ```
Prefill (ring): `t ≈ C · max( recv(KV block), compute(QKᵀ+PV over block) )`
with causal step-skip removing ≈half the steps; no reduce term.
Levers and their expected signature in op_log / latency: Levers: GQA reuse → KV bytes `H_kv·S·d` not `H_q·S·d`; overlap → `max` not
- **GQA reuse** → KV bytes moved are `H_kv·S·d` not `H_q·S·d` (8× less), sum; decode reduce `⌈log₂P⌉+center-mesh` vs baseline `C·P−1`; prefill ring
with compute still `H_q`-scaled. Decode being KV-load-bound, this is the trades the reduce for KV rotation (good when `O` is big). The CPU-saturation
dominant decode win. win (composite offload) is *measurable* only with ADR-0064.
- **Prefetch** → DMA hidden behind compute ⇒ `t ≈ max(...)` not sum.
- **Tree** → reduction `⌈log₂ N⌉` (3 for N=8) vs baseline `N−1` (7).
- **Causal skip** (prefill) → ≈½ the tile/step work.
These are the quantities the milestone bench's op_log summary + a new
latency column should report (Open Decision §10.3 on adding latency to the
sweep JSON).
--- ---
## 10. Open Decisions / Assumptions made autonomously (REVIEW THESE) ## 10. Open items (status)
Choices I made to keep progressing without your input. Each is reversible; Most original "Open Decisions" are now **resolved** in ADR-0060; the live
none is yet implemented in production code (docs only). Marked review items live in **ADR-0060 §B** (three groups: hybrid pivot,
**[recommend]** = my default if you don't object. hierarchical CUBE-Group SP, decode/prefill split). Key resolved choices:
### 10.1 Greenlet `tl` model over composites **[recommend: greenlet]** | Was open | Now |
The biggest design call. ADR-0060 §1/§8 argues the greenlet path is |---|---|
equally efficient in kernbench's per-op latency model and far smaller to | greenlet vs composite | **composite hybrid** (GEMMs→composite, merge→kernel) |
build than a flash-composite. **Risk:** if you specifically want the | one kernel vs two | **two kernels** (decode reduce / prefill ring) |
*composite* abstraction exercised for the eval narrative (e.g. to study | GQA matmul shape | **M-fold** (decode); head-parallel (prefill) |
scheduler-managed tiling), we'd instead invest in a flash-composite | reduction topology | decode **2-level reduce-to-root**; prefill **no reduce (ring)** |
command kind. I judged that out of proportion to the goal. *Confirm the | `tl.load_async` | **lazy `tl.load`** (ADR-0062, redefined) |
greenlet direction.* | KV placement | **contiguous `C×P`**, shared prefill/decode (ADR-0060 §2.1) |
### 10.2 Tree reduction topology / neighbour mapping **[recommend: binary tree on N/S/E/W]** Still to confirm (ADR-0060 §B): DDD↔impl reconcile (`_attention_mesh_mlo_2d`
ADR-0060 §4 fixes a binary tree (N=8: root=7). The mapping of tree edges AllReduce → reduce-to-root); `head_of_group` sub-mesh partition; `C=G`
to physical mesh directions depends on the SFR install's neighbour coupling; short-context KV balance; ADR-0064 calibration. Pre-existing:
wiring. I assumed `configure_sfr_intracube_pe_ring` gives a ring on which ghost ADRs 0055–0059 (separate backfill); `bf16→f16` proxy.
the tree pairs are 1-hop neighbours. **Needs verification** against the
actual installed neighbour table; if pairs aren't 1-hop, either relabel
PEs or accept multi-hop sends. *Confirm or hand me the intended mapping.*
### 10.3 Keep the baseline all-to-all as an option? **[recommend: replace with tree]**
The baseline fan-out gives the answer on *every* rank; the tree gives it
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 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
implementation reserves a small fixed persistent region at kernel start
and double-buffers the merge. I assumed this is acceptable extra
bookkeeping inside the shared helper. *Alternative:* make `scratch_scope`
take an explicit "keep-alive" handle list. I picked the two-arena model
for clarity.
### 10.6 One unified kernel vs two (decode/prefill) **[recommend: unify]**
§5.3 unifies into one parameterised kernel; the baseline has two. Unifying
reduces duplication but makes the kernel branch-heavy. *Alternative:* keep
`_gqa_decode` and `_gqa_prefill` separate, sharing helpers. Low stakes;
either is fine.
### 10.7 Ghost ADRs 0055–0059 **[recommend: backfill as Draft, separate task]**
The baseline cites ADR-0055/0056/0057/0058/0059 which **do not exist**
(`grep` finds refs in code/tests, no files; no git history of them). This
is pre-existing documentation debt, not created by this work. I did **not**
backfill them (out of scope, and reverse-engineering intent is risky).
*Recommendation:* a separate `Draft` ADR-backfill task to document the
existing mesh kernels + milestone bench + SFR install, so ADR-0060 has a
real predecessor to "supersede". *Confirm you want this as follow-up.*
### 10.8 Numbering **[recommend: as chosen]**
GQA ADR = **0060**; supporting = **0061/0062/0063**. I avoided 0055–0059
(claimed by ghost refs) to prevent collisions. The file was renamed from
`ADR-XXXX-…` to `ADR-0060-algo-…` (title was erroneously "ADR-001").
### 10.9 bf16 modelled as f16 **[assumption, pre-existing]**
`memory_store.py:16` maps `bf16→float16` (numpy has no bf16). Numeric
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.*
--- ---
@@ -500,33 +315,34 @@ critical path: *GQA itself is not gated on any new feature; scale
| Risk | Likelihood | Mitigation | | Risk | Likelihood | Mitigation |
|---|---|---| |---|---|---|
| Numeric parity blocked by trans-as-reshape | med | structural-first verify (§10.10); pre-transpose K or defer `tl.transpose` | | Numeric parity blocked by trans-as-reshape | med | structural-first verify; pre-transpose K (ADR-0060 §B item 2) |
| Tree pairs not 1-hop on installed mesh | med | verify SFR neighbour table (§10.2); relabel PEs | | 2-level reduce pairs not 1-hop on the mesh | med | verify SFR neighbour table; center-root sub-mesh partition (§B) |
| scratch_scope use-after-scope bug | med | two-arena discipline in shared helper; debug poison | | scratch_scope use-after-scope (Sj / ring / lazy-load buffers) | med | persistent-arena discipline; ping-pong buffers (§5.2) |
| prefetch buffer recycled while in flight | med | live `LoadFuture` buffers in persistent arena (§10.5) | | prefill ring buffer clobbered while in flight | med | recv into the *other* ping-pong buffer (§5.2) |
| GQA broadcast `.copy()` memory blow-up | low | bounded by TILE in tiled design | | decode opt2 `acc` read before composites drain | med | `tl.wait()` after the loop (ADR-0060 TL;DR opt2) |
| Effort underestimated (6 phases) | med | P0–P1 alone deliver the core unblock; later phases independent | | impl drift (`_attention_mesh_mlo_2d` is AllReduce, Q-replicated) | med | reconcile to reduce-to-root; add prefill ring (§B) |
| CPU-saturation win invisible | exp | ADR-0064 cost model (P8) |
--- ---
## 12. Glossary & references ## 12. Glossary & references
- **GQA / G** — grouped-query attention; `G = H_q/H_kv` query heads share - **GQA / G** — `G = H_q/H_kv` query heads share one KV head. Llama3-70B: `G=8`.
one KV head. Llama3-70B: `G=8`. - **CUBE Group** — `C` CUBEs in one SIP owning one KV head (ADR-0060 §0).
- **FlashAttention / FlashDecoding / Ring Attention** — see ADR-0060 - **2-level SP** — Level-1 inter-CUBE (over `C`) × Level-2 intra-CUBE PE
lineage note. (over `P`); ranks = `C·P`.
- **IPCQ** — inter-PE queues; `tl.send`/`tl.recv` over `N/S/E/W` - **Composite hybrid** — GEMMs via `tl.composite` (scheduler), softmax merge
(ADR-0023/0025/0032). + reduction in the kernel (ADR-0060 §1).
- **Greenlet `tl` model** — kernel interleaves with SimPy; ops emit - **M-fold** — stack the `G` query heads into the GEMM M (row) dim so one
commands scheduled on PE engines (ADR-0020, ADR-0046). `Q·Kᵀ` does all heads sharing one `K`.
- **Persistent vs scoped arena** — ADR-0063 scratch split. - **FlashAttention / FlashDecoding / Ring Attention** — ADR-0060 lineage.
Key source anchors: `tl_context.py` (tl API), `memory_store.py:67-73` Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143`
(nbytes check), `data_executor.py` (Phase 2), `pe_dma.py`/`pe_gemm.py`/ (composite tile plan), `pe_dma.py:45,89` (read channel, drain), `pe_cpu.py`
`pe_math.py` (latency), `_attention_mesh_kv.py` / `_attention_mesh_mlo.py` (dispatch_cycles), `lrab_hierarchical_allreduce.py` (center-root pattern),
(baseline kernels), `milestone_gqa_llama70b.py` (eval bench), `_attention_mesh_mlo_2d.py` / `_attention_mesh_kv.py` (impl to evolve),
`tests/attention/test_milestone_gqa_llama70b.py` (current limits). `DPPolicy.cube_start`, `topologies/llama70b_4sip.yaml`.
ADRs: **0060** (this design), **0061** (broadcast), **0062** (async load), ADRs: **0060** (this design), **0062** (lazy load), **0063** (scratch scope),
**0063** (scratch scope); related accepted **0014/0020/0023/0025/0042/ **0061** (broadcast, optional), **0064** (CPU issue cost model); related
0045/0046/0052/0054**. accepted **0014/0017/0020/0023/0025/0046/0054**.