diff --git a/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md b/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md index 41d90e7..83a8454 100644 --- a/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md +++ b/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md @@ -1,39 +1,42 @@ # Detailed Design Document — AHBM GQA Fused Attention (ADR-0060) **Status:** Draft (companion to ADR-0060, Proposed) -**Scope:** implementation-ready design for an efficient GQA FlashAttention -kernel on AHBM in the kernbench simulator. -**Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read ADR-0060 -first (decision record), then this (the how), then the *Open Decisions* -section (§10) for the choices made autonomously that need your sign-off. +**Scope:** implementation-ready plan for the **two** GQA FlashAttention +kernels (decode=reduce, prefill=ring) on AHBM in kernbench. +**Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read **ADR-0060 +first** — it is now the authoritative design record (TL;DR has full +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 -> design↔evaluation iterations against the real kernbench source -> (`src/kernbench/triton_emu/tl_context.py`, `sim_engine/*`, -> `components/builtin/pe_*`, the `_attention_mesh_*` kernels, and the GQA -> 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). +> **Alignment note (this revision).** ADR-0060 moved to a **composite +> hybrid + hierarchical CUBE-Group SP** design with **two kernels**. This +> DDD is rewritten to match. The earlier DDD (greenlet-primitive, +> `tl.load_async`, single unified kernel) is superseded. --- ## 1. Goal and success criteria -**Goal:** an *efficient* GQA fused attention kernel that runs on AHBM in -kernbench, across the four operating points decode/prefill × -single-/multi-user, at realistic (not just validation) scale, in both +**Goal:** two *efficient* GQA fused-attention kernels on AHBM in kernbench — +**decode+SP** (head-replicated, KV static shard, 2-level reduce) and +**prefill+SP** (1 Q head per CUBE, Ring KV) — at realistic scale, in both timing (`enable_data=False`) and data (`enable_data=True`) modes. **Success criteria:** -1. Correct: kernel output matches a numpy FlashAttention reference within - fp 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: - GQA K/V-load amortisation, KV prefetch overlap, `⌈log₂ N⌉` reduction. +1. Correct: kernel output matches a numpy FlashAttention reference within fp + tolerance, with **real GQA** (`H_q=64, H_kv=8, G=8`). +2. Efficient — levers observable in op_log (ADR-0060 §11): + 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. 4. Deterministic and SPEC-compliant (all latency from modelled events). **Non-goals:** RoPE / QKV projection / KV-cache writes (upstream -`qkv_rope`); output projection (downstream); cross-CUBE query-head split; -cycle-accurate microarchitecture (SPEC §5). +`qkv_rope`); output projection (downstream); cross-**SIP** head split +(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) - RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy, kernel launch) + RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy(+cube_start), launch) │ KernelLaunchMsg ▼ sim_engine - GraphEngine(enable_data=…) engine.py - DataExecutor (Phase 2) data_executor.py ← ADR-0061 adds _execute_broadcast - MemoryStore (nbytes check) memory_store.py ← the GQA blocker lives here + GraphEngine(enable_data=…) engine.py + DataExecutor (Phase 2) data_executor.py ← _execute_broadcast (ADR-0061, optional) + MemoryStore memory_store.py │ ▼ components / PE pipeline (per PE) - PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,FETCH_STORE,TCM,IPCQ} - greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py + PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM} + 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 - load/store/dot/softmax/max/sum/exp/maximum/send/recv/trans/… - ← ADR-0061 adds broadcast/repeat - ← ADR-0062 adds load_async (+ wait arm) - ← ADR-0063 adds scratch_scope +tl programming model tl_context.py + load (→ lazy, ADR-0062) / store / composite / dot / max/sum/exp/maximum + send/recv/recv_async / scratch_scope (ADR-0063) / broadcast (ADR-0061, opt) ``` -The attention kernels are **bench-side Python** (`src/kernbench/benches/`) -invoked via `ctx.launch(name, kernel_fn, *args)`; they run as greenlet -kernels that emit `tl.*` ops. No component (device-side) code changes are -needed for the algorithm — only the three `tl`/sim_engine primitives. +Both kernels are **bench-side Python** (`src/kernbench/benches/`) run as +greenlet kernels emitting `tl.*` ops; the **GEMMs are `tl.composite`** +(scheduler-managed tiling + K/V DMA streaming, existing `CompositeCmd` — no +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 | |---|---|---| -| `src/kernbench/triton_emu/tl_context.py` | add `broadcast`/`repeat`; `load_async` + `wait` arm; `scratch_scope` ctx-mgr | 0061/0062/0063 | -| `src/kernbench/common/pe_commands.py` | add `BroadcastCmd`; `DmaReadCmd.blocking` flag (or `DmaReadAsyncCmd`) | 0061/0062 | -| `src/kernbench/sim_engine/data_executor.py` | add `_execute_broadcast`; handle async dma_read resolution | 0061/0062 | -| `src/kernbench/components/builtin/pe_dma.py` | non-blocking read path returns a future-resolving event | 0062 | -| `src/kernbench/sim_engine/event_log.py` / op_log | recognise `broadcast` op_name (math) | 0061 | +| `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/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/triton_emu/kernel_runner.py` | auto-wait: yield a handle's pending load event at first consume; lazy-load dispatch arm | 0062 | +| `src/kernbench/sim_engine/data_executor.py` | `_execute_broadcast` (optional) | 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 | |---|---| -| `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_helpers.py` | `valid_len`, `tile_*` predicates, mask builders, tree topology (children/parent dirs), GQA-batched Q load | -| extend `src/kernbench/benches/milestone_gqa_llama70b.py` | add headline-scale panels with real `G`, prefetch, tiling | +| `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_prefill.py` | prefill+SP kernel — evolve `_attention_mesh_kv` to **head-parallel (1 Q head/CUBE) Ring KV** + composite-hybrid | +| `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 | File | Covers | |---|---| -| `tests/attention/test_gqa_correctness.py` | numpy reference parity, all 4 cases, real GQA | -| `tests/attention/test_gqa_levers.py` | GQA load amortisation, tree step count, causal skip, prefetch overlap | -| `tests/test_tl_broadcast.py` | ADR-0061 unit | -| `tests/test_tl_load_async.py` | ADR-0062 unit | +| `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_prefill.py` | prefill ring: no `(m,ℓ,O)` reduce, `C` KV rotations, causal step-skip, per-CUBE distributed output | +| `tests/test_tl_lazy_load.py` | ADR-0062 unit (overlap, auto-wait correctness, op_log parity) | | `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.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` | -| `K`, `V` (sequence-sharded) | `cube=replicate, pe=row_wise, num_pes=N` | `cube=row_wise, pe=replicate, num_cubes=N` | -| `O` (group, replicated) | as `Q` | as `Q` | +| `Q` | `replicate` (all `G` heads, M-folded) over `C×P` ranks | **head-parallel**: CUBE `i` holds Q head `i` | +| `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` | reduced to CUBE-Group root (all `G` heads) | per-CUBE (one head each), distributed | -This matches the baseline `_PANEL_DISPATCH` -(`milestone_gqa_llama70b.py:64-81`). `row_wise` over the sequence (M) axis -realises the round-robin shard (§2.1 of ADR-0060); the contiguous per-PE -buffer is the row-wise slice. +**KV cache is *contiguous* `C×P` position blocks** (rank `r` = `[r·B,(r+1)·B)`), +**shared by both kernels** — prefill writes it (via `qkv_rope`), decode reads ++ extends it, no reshard. Contiguous is required for prefill causal skip +(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 -single-user (PE id in cube), `1` for multi-user (cube level; only PE 0 of -each cube participates) — kept from the baseline so the SFR installs -(`configure_sfr_intracube_pe_ring` / `configure_sfr_intercube_multisip`) -apply unchanged. +`--device` enumerates **SIPs**; one SIP-device runs its 2 CUBE Groups +(2 KV heads); the head is picked by CUBE coordinate (`head_of_group`), not an +in-kernel loop (ADR-0060 §0). -### 4.2 Launch signature +### 4.2 Launch signatures ```python -ctx.launch( - f"{panel}_gqa", gqa_attention_kernel, - q, k, v, o, # Tensors (TensorArg) - S_q, S_kv_per_rank, H_q, H_kv, D_HEAD, n_ranks, # scalars - rank_axis, mode, counter, start_pe, # scalars (mode: 0=decode,1=prefill) - softmax_scale, # float scalar - _auto_dim_remap=False, # keep d_head from colliding (baseline note) -) +# decode+SP +ctx.launch(f"{panel}_gqa_decode", gqa_decode_sp, + q, k, v, o, counter, start_pe, start_cube, C, P, softmax_scale) +# prefill+SP +ctx.launch(f"{panel}_gqa_prefill", gqa_prefill_sp, + q, k, v, o, T_q, S_kv_local, d, C, softmax_scale, q_block, cube_start) ``` -`H_q > H_kv` is the change from the baseline's `H_q=H_kv=1`; it requires -ADR-0061 to survive Phase 2. +Full per-kernel I/O contract: ADR-0060 §0.5.3/§0.5.4 (note the **output head +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 -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 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). +Decode replicates Q and **M-folds** the `G` query heads into the GEMM row +dim so one `Q·Kᵀ` serves all `G` heads sharing one `K` (ADR-0060 §0, §5.2): +`Q` group `[G, T_q, d]` → `[G·T_q, d]` (byte-conserving `_view`); the +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 → +store `K` pre-transposed `[d, S/(C·P)]` (ADR-0060 §3, §B). --- -## 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 -def flash_tile_sweep(q_g, k_base, v_base, n_tiles, TILE, d, scale, - q_block, causal, tl, PREFETCH=2): - """One PE's FlashAttention sweep over its local KV tiles. - Returns running (m, l, O) in the persistent arena.""" - m, l, O = _init_running(q_g.shape, d, tl) # persistent arena - f_k, f_v = {}, {} - for j in range(min(PREFETCH, n_tiles)): # prime (ADR-0062) - if causal and tile_all_future(j, q_block): continue - f_k[j] = tl.load_async(k_base + j*TILE*d, (TILE, d)) - f_v[j] = tl.load_async(v_base + j*TILE*d, (TILE, d)) - 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 -``` +### 5.1 Decode (reduce) — `_gqa_decode.py` +- Inner tile = §3 composite-hybrid: `Sj = composite(q_g, Kⱼ)·scale` → softmax + MATH → `Oj = composite(P, Vⱼ)` → running merge (ADR-0060 §3). +- **Ship opt3** (software pipelining: issue next tile's `Q·Kᵀ` before this + tile's softmax; `Sj` in a persistent double buffer) — removes the + GEMM-engine bubble, no new command kind (ADR-0060 §5.6). opt1 is the naïve + baseline; opt2 (`ex_composite`) waits on ADR-0064. +- Combine = **`hierarchical_reduce_and_store`** (ADR-0060 §4): Level-2 PE + reduce-to-root tree (intra-CUBE, the cube's KV slice further split across + its `P` PEs — decision (a)) → Level-1 center-root CUBE-mesh reduce + (intra-CUBE-Group). Data-driven, level-pipelined, root-only. -> Note: `m, l, O` are reassigned inside the `with` but their **storage** -> must be the persistent arena, not the recycled scope. Concretely the -> merge ops write to pre-allocated persistent handles (double-buffered); -> ADR-0063 §D3 describes the persistent-vs-scoped arena split. This is the -> one subtlety the implementer must get right — see §10.5. +### 5.2 Prefill (ring) — `_gqa_prefill.py` +- Head-parallel: CUBE `i` owns Q head `i` + KV slice `i`. **No reduce** (each + CUBE outputs a distinct head). Ring rotates KV blocks; GQA reuse via the + rotation (ADR-0060 §5.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 - -```python -def tree_reduce_and_store(m, l, O, pe_id, N, o_ptr, tl): - for cdir in tree_children_dirs(pe_id, N): # static tree (§4 ADR-0060) - 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). +### 5.3 Shared helpers — `_gqa_helpers.py` +`head_of_group(cube_id) = cube_id // C`; `init_running(...)` → persistent +`(m=-inf, ℓ=0, O=0)`; `valid_len_2level(...)` (contiguous block extent); +`hierarchical_reduce_and_store(...)`; tree/center-mesh child/parent dirs; +`block_all_future`/`block_partial` + `causal_mask`. --- -## 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) -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`. -- Lowest priority; deliver after the efficiency/scale features land. +### 6.1 ADR-0062 lazy `tl.load` (efficiency — overlap) +- `tl.load` issues `DmaReadCmd` non-blocking, returns a handle with a pending + event; the runtime **auto-inserts the wait at first consume** (generalises + the `recv_async`/wait pattern, `kernel_runner.py:248-285`). +- `pe_dma.py`: non-blocking read; DMA occupies the (capacity-1) read channel, + overlaps compute on PE_GEMM/PE_MATH. +- **Global** semantics change → existing goldens regenerate (ADR-0062 D3). +- op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` stable. -### 6.2 ADR-0062 `tl.load_async` -- `tl_context.py`: `load_async(ptr, shape, dtype) → LoadFuture`; extend - `wait` to resolve `LoadFuture`. -- `pe_dma.py`: non-blocking read path; transfer occupies the read channel - in parallel; completion is a scheduled event. -- op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` metric stable. +### 6.2 ADR-0063 `tl.scratch_scope` (scale) +- ctx-mgr saving/restoring `_scratch_cursor`. Persistent arena for `(m,ℓ,O)` + (+ decode opt3 `Sj` double buffer, + prefill ring ping-pong buffers, + + in-flight lazy-load buffers) lives **outside** the scope. Removes `S=16`. -### 6.3 ADR-0063 `tl.scratch_scope` -- `tl_context.py`: context manager saving/restoring `_scratch_cursor`. -- Persistent arena: `m,l,O` (+ double-buffer) and live `LoadFuture` - buffers allocated outside the scope. -- Removes the `S=16` cap (`test_milestone_gqa_llama70b.py:123-148`). +### 6.3 ADR-0061 `tl.broadcast` (optional) +- Not on the GQA critical path (M-fold gives reuse). Convenience for additive + mask construction. Lowest priority. + +### 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) -Each phase is an independent Phase-1→Phase-2 cycle. Ordered so each builds -on a green predecessor; every phase keeps the existing baseline green. +Each phase is an independent Phase-1→Phase-2 cycle; each keeps the baseline +green. | 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` | -| **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 | +| **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** | 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** `scratch_scope` + tiled sweep | long-`S` decode completes (today caps S=16); scratch O(1) | +| **P4** | **ADR-0062** lazy `tl.load` + composite K/V streaming | tiled-sweep latency < serial `Σ(load+compute)`; goldens regenerated | +| **P5** | Decode **opt3** software pipelining | GEMM engine non-idle across tiles (structural / op_log) | +| **P6** | Prefill **head-parallel Ring KV** + causal step-skip | `C` KV rotations, no `(m,ℓ,O)` reduce, per-CUBE distributed `O`, causal skip ≈½ | +| **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**, -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. +P1 delivers "GQA runs" on the composite path. P3 is the key *scale* feature; +P4 the overlap lever; P2/P6 the SP algorithms. --- ## 8. Verification plan (concrete) -Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (`tl`), -ADR-0054 (eval bench). Mirrors ADR-0060 §11. +Mirrors ADR-0060 §11; grounded in SPEC R2/R5, ADR-0023/0025, ADR-0046, ADR-0054. -**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. +**Runs + numeric (`enable_data=True`):** +- decode & prefill complete for `(G∈{1,8}, T_q∈{1,16}, S, C∈{1,8}, P∈{1,8})` + without byte-conservation error (the `G=8` cases baseline cannot express). - *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). + symmetric/identity inputs (reshape-as-transpose exact); full asymmetric + parity gated on a real `tl.transpose` (deferred). -**Levers (op_log assertions):** -- GQA amortisation: K/V `dma_read_count` constant as `G` grows 1→8; - `gemm` work grows with `G`. -- Tree: decode-SP `ipcq_send`/`recv` counts = tree edges, reduction - rounds = `⌈log₂ N⌉` (not `N−1`). -- Causal skip: prefill `gemm_count` = lower-triangular tile count. -- Prefetch overlap: tiled-sweep latency strictly < `Σ(load+compute)`. -- Scratch: peak `_scratch_cursor` bounded by one tile, independent of - `n_tiles`. +**Levers (op_log):** +- GQA amortisation: K/V `dma_read_count` ⟂ `G`; compute scales with `G`. +- Decode reduce: rounds = `⌈log₂P⌉` + center-mesh (not `C·P−1`); result at + one rank; Level-1=CUBE NOC, Level-2=PE IPCQ. +- Prefill ring: `C` KV rotations, **no** reduce, each CUBE writes a distinct + head. +- Causal skip: GEMM/step count = lower-triangular. +- Overlap: tiled-sweep latency < `Σ(load+compute)`. +- Scratch: peak cursor bounded by one tile. -**Invariants:** -- Determinism: identical inputs → identical op_log + latency (SPEC §0.1). -- Every routed request latency > 0 (SPEC §0.1). +**Invariants:** determinism (identical op_log + latency); every routed +request latency > 0 (SPEC §0.1). --- -## 9. Performance model (what we expect to see) - -Rough per-PE decode cost (KV-load-bound), with prefetch overlap: +## 9. Performance model (expected) +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_tree_reduce(⌈log₂ N⌉ hops, payload O[G,d]) +t_decode_rank ≈ max( DMA(K+V over S/(C·P) tiles), compute(QKᵀ+PV, ×G amortised) ) + + 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: -- **GQA reuse** → KV bytes moved are `H_kv·S·d` not `H_q·S·d` (8× less), - with compute still `H_q`-scaled. Decode being KV-load-bound, this is the - dominant decode win. -- **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). +Levers: GQA reuse → KV bytes `H_kv·S·d` not `H_q·S·d`; overlap → `max` not +sum; decode reduce `⌈log₂P⌉+center-mesh` vs baseline `C·P−1`; prefill ring +trades the reduce for KV rotation (good when `O` is big). The CPU-saturation +win (composite offload) is *measurable* only with ADR-0064. --- -## 10. Open Decisions / Assumptions made autonomously (REVIEW THESE) +## 10. Open items (status) -Choices I made to keep progressing without your input. Each is reversible; -none is yet implemented in production code (docs only). Marked -**[recommend]** = my default if you don't object. +Most original "Open Decisions" are now **resolved** in ADR-0060; the live +review items live in **ADR-0060 §B** (three groups: hybrid pivot, +hierarchical CUBE-Group SP, decode/prefill split). Key resolved choices: -### 10.1 Greenlet `tl` model over composites **[recommend: greenlet]** -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 -build than a flash-composite. **Risk:** if you specifically want the -*composite* abstraction exercised for the eval narrative (e.g. to study -scheduler-managed tiling), we'd instead invest in a flash-composite -command kind. I judged that out of proportion to the goal. *Confirm the -greenlet direction.* +| Was open | Now | +|---|---| +| greenlet vs composite | **composite hybrid** (GEMMs→composite, merge→kernel) | +| one kernel vs two | **two kernels** (decode reduce / prefill ring) | +| GQA matmul shape | **M-fold** (decode); head-parallel (prefill) | +| reduction topology | decode **2-level reduce-to-root**; prefill **no reduce (ring)** | +| `tl.load_async` | **lazy `tl.load`** (ADR-0062, redefined) | +| 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]** -ADR-0060 §4 fixes a binary tree (N=8: root=7). The mapping of tree edges -to physical mesh directions depends on the SFR install's neighbour -wiring. I assumed `configure_sfr_intracube_pe_ring` gives a ring on which -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.* +Still to confirm (ADR-0060 §B): DDD↔impl reconcile (`_attention_mesh_mlo_2d` +AllReduce → reduce-to-root); `head_of_group` sub-mesh partition; `C=G` +coupling; short-context KV balance; ADR-0064 calibration. Pre-existing: +ghost ADRs 0055–0059 (separate backfill); `bf16→f16` proxy. --- @@ -500,33 +315,34 @@ critical path: *GQA itself is not gated on any new feature; scale | Risk | Likelihood | Mitigation | |---|---|---| -| 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) | -| GQA broadcast `.copy()` memory blow-up | low | bounded by TILE in tiled design | -| Effort underestimated (6 phases) | med | P0–P1 alone deliver the core unblock; later phases independent | +| Numeric parity blocked by trans-as-reshape | med | structural-first verify; pre-transpose K (ADR-0060 §B item 2) | +| 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 (Sj / ring / lazy-load buffers) | med | persistent-arena discipline; ping-pong buffers (§5.2) | +| prefill ring buffer clobbered while in flight | med | recv into the *other* ping-pong buffer (§5.2) | +| decode opt2 `acc` read before composites drain | med | `tl.wait()` after the loop (ADR-0060 TL;DR opt2) | +| 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 -- **GQA / G** — grouped-query attention; `G = H_q/H_kv` query heads share - one KV head. Llama3-70B: `G=8`. -- **FlashAttention / FlashDecoding / Ring Attention** — see ADR-0060 - lineage note. -- **IPCQ** — inter-PE queues; `tl.send`/`tl.recv` over `N/S/E/W` - (ADR-0023/0025/0032). -- **Greenlet `tl` model** — kernel interleaves with SimPy; ops emit - commands scheduled on PE engines (ADR-0020, ADR-0046). -- **Persistent vs scoped arena** — ADR-0063 scratch split. +- **GQA / G** — `G = H_q/H_kv` query heads share one KV head. Llama3-70B: `G=8`. +- **CUBE Group** — `C` CUBEs in one SIP owning one KV head (ADR-0060 §0). +- **2-level SP** — Level-1 inter-CUBE (over `C`) × Level-2 intra-CUBE PE + (over `P`); ranks = `C·P`. +- **Composite hybrid** — GEMMs via `tl.composite` (scheduler), softmax merge + + reduction in the kernel (ADR-0060 §1). +- **M-fold** — stack the `G` query heads into the GEMM M (row) dim so one + `Q·Kᵀ` does all heads sharing one `K`. +- **FlashAttention / FlashDecoding / Ring Attention** — ADR-0060 lineage. -Key source anchors: `tl_context.py` (tl API), `memory_store.py:67-73` -(nbytes check), `data_executor.py` (Phase 2), `pe_dma.py`/`pe_gemm.py`/ -`pe_math.py` (latency), `_attention_mesh_kv.py` / `_attention_mesh_mlo.py` -(baseline kernels), `milestone_gqa_llama70b.py` (eval bench), -`tests/attention/test_milestone_gqa_llama70b.py` (current limits). +Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143` +(composite tile plan), `pe_dma.py:45,89` (read channel, drain), `pe_cpu.py` +(dispatch_cycles), `lrab_hierarchical_allreduce.py` (center-root pattern), +`_attention_mesh_mlo_2d.py` / `_attention_mesh_kv.py` (impl to evolve), +`DPPolicy.cube_start`, `topologies/llama70b_4sip.yaml`. -ADRs: **0060** (this design), **0061** (broadcast), **0062** (async load), -**0063** (scratch scope); related accepted **0014/0020/0023/0025/0042/ -0045/0046/0052/0054**. +ADRs: **0060** (this design), **0062** (lazy load), **0063** (scratch scope), +**0061** (broadcast, optional), **0064** (CPU issue cost model); related +accepted **0014/0017/0020/0023/0025/0046/0054**.