From 1742600641916043778e0ba26486d16773f42023 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 3 Jun 2026 17:42:00 -0700 Subject: [PATCH] gqa(adr): rewrite GQA fused-attention ADR as ADR-0060, aligned to kernbench Reframes the proposed GQA FlashAttention design onto kernbench's actual execution model (greenlet tl API + IPCQ), replacing the composite-centric mechanism that does not match the simulator: - Records relationship to existing baseline kernels (_attention_mesh_kv/mlo, milestone-gqa-llama70b) and their 3 deliberate limitations. - Mechanism is greenlet tl (per-op latency; no fusion benefit), not composites; running (m,l,O) is Python handles; reduction is tl.send/recv. - Tree reduction (log N) replaces baseline all-to-all fan-out (N-1). - Pseudocode rewritten in real tl.* signatures; depends on ADR-0061/62/63. - Rejects composite-IPCQ-push + composite-carried-state + flash-composite with documented efficient alternatives. - Adds verification plan. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 576 ++++++++++++++++++ .../ADR-XXXX-gqa-fused-attention-ahbm.md | 397 ------------ 2 files changed, 576 insertions(+), 397 deletions(-) create mode 100644 docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md delete mode 100644 docs/adr-proposed/ADR-XXXX-gqa-fused-attention-ahbm.md diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md new file mode 100644 index 0000000..ed7cfd7 --- /dev/null +++ b/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -0,0 +1,576 @@ +# ADR-0060: AHBM GQA Fused Attention Kernel (Llama3-70B) + +## Status + +Proposed + +**Context model:** Llama3-70B. +**Decision drivers:** agentic workload → low batch, long context; +KV-load-bound decode; sequence-parallel (Ring KV) for long-context prefill. + +**Supersedes / extends:** the existing mesh-native attention kernels +`_attention_mesh_kv` (prefill) and `_attention_mesh_mlo` (decode) and the +`milestone-gqa-llama70b` eval bench. See *§A. Relationship to existing +kernbench work* — that code is the baseline this ADR upgrades to a real +GQA, causal, long-context kernel. + +**Supporting ADRs (prerequisites):** +- **ADR-0061** `tl.broadcast` — data-faithful GQA head reuse. +- **ADR-0062** `tl.load_async` — non-blocking HBM tile load (KV prefetch). +- **ADR-0063** `tl.scratch_scope` — per-tile scratch recycling. + +**Algorithm lineage.** This kernel is **FlashAttention** (tiling + +online/streaming softmax with fused P·V — no full score matrix +materialised). The KV-parallel split-and-combine in §4 is +**FlashDecoding** (split-KV with log-sum-exp merge). The Ring path in +§5.5 is **Ring Attention** (KV blocks rotated around the mesh, folded by +the same online softmax). No new math is introduced; this ADR maps those +known algorithms onto the kernbench **greenlet `tl` programming model** +(ADR-0020, ADR-0046) and the **IPCQ** PE↔PE collective (ADR-0023/0025). + +--- + +## A. Relationship to existing kernbench work (read first) + +kernbench **already runs FlashAttention with an online-softmax `(m, ℓ, O)` +merge over IPCQ today.** Two kernels exist: + +| File | Role | Mechanism | +|---|---|---| +| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold | +| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge | + +Both are driven by `milestone-gqa-llama70b` (4 panels: +`{single,multi}_user × {prefill,decode}`, +`src/kernbench/benches/milestone_gqa_llama70b.py`) and tested in +`tests/attention/test_milestone_gqa_llama70b.py`. + +**They are written in the greenlet `tl` API — not composites:** `tl.load`, +`tl.dot`, `tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, +and Python `-`/`*`/`/` on `TensorHandle` (each emits a `MathCmd`). The +running `(m, ℓ, O)` is just Python `TensorHandle`s threaded through the +loop. This **matters for this ADR's design** (see §1). + +**Three deliberate limitations of the baseline** — exactly what an +*efficient GQA* kernel must lift: + +1. **No GQA reuse.** `h_q == h_kv == 1` + (`test_milestone_gqa_llama70b.py:137-142`). Real GQA (`h_q = G·h_kv`) + is blocked by the MemoryStore byte-conservation check on the symbolic + broadcast view → fixed by **ADR-0061**. +2. **O(N) reduction.** The baseline does an all-to-all bidirectional + fan-out so *every* rank ends with the full answer (`n_ranks − 1` + steps). Attention only needs `O` at the query owner → a **tree + reduction** to root is `⌈log₂ N⌉` steps (§4). +3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump + allocator leaks per-tile temporaries + (`test_milestone_gqa_llama70b.py:123-148`) and there is no causal + masking / tiling → fixed by **ADR-0063** (recycling) + §5 (tiling, + causal skip) + **ADR-0062** (prefetch). + +**Documentation debt (out of scope but recorded):** the baseline cites +ADR-0055/0056/0057/0058/0059, **none of which exist as files** — they are +ghost references. This ADR does not retro-write them; see the Detailed +Design Document's *Open Decisions* for the recommendation. + +--- + +## 0. Reference dimensions (Llama3-70B) + +| Symbol | Meaning | Value | +|---|---|---| +| `H_q` | query heads | 64 | +| `H_kv` | KV heads | 8 | +| `G` | GQA group size = `H_q / H_kv` | 8 | +| `d` | head dim | 128 | +| `L` | layers | 80 | +| `D` | model dim | 8192 | + +Hardware recap: +- **AHBM (chip)** = set of **CUBE**s (memory cubes, each with a logic die + containing **PE**s) + **IO die** (ADR-0003). +- **IPCQ**: PE↔PE queues, 4 mesh-direction queue-pairs per PE + (`N/S/E/W`, ADR-0023 D3; `global_*` for inter-SIP, ADR-0032). Kernel + API: `tl.send(dir, src)` / `tl.recv(dir, shape, dtype)` + (`tl_context.py:402-499`). +- **Composite command** (`CompositeCmd`, `pe_commands.py:144-162`): a + single GEMM (or MATH) *head* plus element-wise *epilogue* stages + (`bias/relu/scale/add/...`). It is **not** a general multi-op DAG: it + cannot chain two GEMMs, cannot carry register state across instances, + and cannot pop/wait on IPCQ. This ADR therefore does **not** require + composites for the attention inner loop (see §1, §8). +- **Allocation policy (SP):** one query head per CUBE in the multi-user + panels; the `G` query heads of a KV group map within a CUBE's PEs. + +**Two orthogonal mapping layers** (round-robin placement): +- **Layer 1 (KV-parallel, intra-request):** KV token `i` lands on PE + `(start_pe + i) mod N` — round-robin so one long request is split + across `N` PEs, balanced to ≤1 token. +- **Layer 2 (inter-request):** `start_pe = request_id mod N` rotates the + "PE-0 role" per request. +- Combined: `pe = (request_id + global_token_idx) mod N`. + +`N` = number of PEs in the KV-parallel reduction group for one +(query head, request). Reduction stays **intra-CUBE** (a query head never +spans CUBEs — explicit non-goal). + +--- + +## 0.5 Kernel boundary, preconditions, I/O contract + +### 0.5.1 Position in the decoder layer + +``` +1. RMSNorm +2. QKV projection (GEMM) ─┐ qkv_rope kernel (SEPARATE, upstream) +3. RoPE on Q and K ─┤ +4. write new K,V → KV cache ─┘ +5. ===== THIS KERNEL: FlashAttention ===== (post-RoPE Q, K-cache, V-cache → O) +6. Output projection (GEMM) out_proj kernel (SEPARATE, downstream) +7. residual add → FFN ... +``` + +**Preconditions (upstream `qkv_rope`, NOT this kernel):** +- **P1.** Q is already RoPE-rotated. No rotation here. +- **P2.** K-cache stores **post-RoPE** K (never re-rotated at attention + time — the reason for post-RoPE caching). +- **P3.** For decode, the new step's K row is RoPE-rotated and appended + to its owning PE's K-cache slot **by `qkv_rope` before this kernel + launches.** ⇒ this kernel is **pure read** on the KV cache. +- **P4.** V is not rotated; V-cache holds raw projected V. +- **P5.** RoPE position upstream is the token's **absolute global + position** `global_idx = local_slot·N + ((pe_id − start_pe) mod N)` + (§2.1). Round-robin placement ≠ RoPE position. + +### 0.5.2 Shape symbols +`T_q` = query length this launch (decode: 1; prefill/chunk: chunk width). +`S` = total context length. `S_pe` = keys owned by this PE ≈ `⌈S/N⌉`. +Storage `bf16` (numpy proxy `f16`, `memory_store.py:16`); `m,ℓ,O` +accumulators `f32` where precision matters. + +### 0.5.3 INPUTS (per kernel launch) + +| Input | Shape (per KV head) | Location | Notes | +|---|---|---|---| +| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. | +| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. Dense locally, strided in global index (§2.1). | +| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. | +| `global_token_counter` | scalar | launch arg | kernel derives `S_pe`, slot↔global, causal bounds. | +| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 rotation. | +| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. | +| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. | +| `O_base` | address | launch arg | where final O is written (root PE only). | +| `softmax_scale` | scalar | launch arg | `1/√d`. | + +For **Ring Attention (§5.5)** add per ring step: incoming `K_block, +V_block` via IPCQ into ping-pong buffers (post-RoPE), plus +`step_kv_global_range` for the causal step-skip. + +### 0.5.4 OUTPUTS + +| Output | Shape | Location | Notes | +|---|---|---|---| +| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE only** | normalised `O_acc / ℓ_acc`, cast to storage dtype. | + +**Intermediate (non-root PE, on IPCQ — not a kernel-visible output):** +`(m_i, ℓ_i, O_i)` (`m,ℓ`: `[G, T_q]`; `O_i`: `[G, T_q, d]` unnormalised) +pushed to `tree_parent` via `tl.send` (§4). `O_i` is the heavy part. + +**No-SP (`N=1`):** no IPCQ partials; the single PE's running `(m,ℓ,O)` is +normalised in place and written to `O_base`. + +**Explicitly NOT outputs:** KV-cache writes (upstream, P3), output +projection (downstream), score `S` and probs `P` (never materialised). + +--- + +## 1. Decision (mechanism) + +**Implement the kernel in the greenlet `tl` programming model, not as +composites.** Rationale grounded in kernbench's execution + latency model: + +- The simulator charges latency **per op on its modelled component** + (GEMM on PE_GEMM, vector ops on PE_MATH, DMA on PE_DMA — `pe_gemm.py`, + `pe_math.py`, `pe_dma.py`). "Fusing" several ops into one + `CompositeCmd` does **not** reduce modelled latency; the stages still + run on the same engines. The win a real fused kernel gets — *overlap* + (prefetch) and *small working set* (scratch recycling) — is obtained + here by **ADR-0062** and **ADR-0063**, which are smaller and reusable. +- The running `(m, ℓ, O)` flash state is naturally a set of Python + `TensorHandle`s threaded through the loop (the baseline already does + this). No composite-carried register state is needed. +- Cross-PE combination uses kernel-level `tl.send`/`tl.recv` (the + baseline already does this and it is tested). No composite-driven IPCQ + push is needed. + +So the per-tile inner pipeline is the op sequence + +``` +K_j load → Q·Kⱼᵀ → (+ causal mask) → online-softmax update → V_j load → P·Vⱼ → running (m,ℓ,O) update +``` + +issued as ordinary `tl.*` ops, with the **next** tile's K/V issued via +`tl.load_async` (ADR-0062) so its DMA overlaps the current tile's +compute, and each tile's temporaries wrapped in `tl.scratch_scope` +(ADR-0063) so scratch stays O(1). + +Cross-PE combination (KV-parallel / SP) is a **log-sum-exp tree +reduction** over `(m, ℓ, O)` after P·V, flowed through IPCQ with +kernel-level `tl.send`/`tl.recv` (§4). + +Control flow (tile skip, mask generation, reduction scheduling, address +arithmetic) lives in the **kernel** (plain Python `if`/arithmetic in the +greenlet body). This is exactly what kernbench's greenlet model already +permits (`kernel_runner.py`, ADR-0020 D3). + +> **Why this is the efficient choice.** It reuses the proven baseline +> kernels and the proven IPCQ collective; the only genuinely new +> machinery is three small, general primitives (ADR-0061/0062/0063). The +> originally-proposed "everything-in-one-composite + composite IPCQ push +> + composite-carried state" would require a bespoke flash-composite +> command type, a register-lifetime model across composites, and an +> IPCQ-push epilogue — large, special-purpose, and with no latency +> benefit over the greenlet path. Those are **rejected**; see §8. + +--- + +## 2. Memory layout & driver responsibilities + +### 2.1 KV cache allocation +- Per-PE KV buffers sized to per-PE max context `⌈max_context / N⌉ × d × + dtype`, for K and V separately. In kernbench these are deployed with a + `DPPolicy` whose `pe` (or `cube`) axis is `row_wise` over the sequence + dimension (`policy/placement/dp.py`; the baseline uses + `DPPolicy(pe="row_wise")` for K/V and `replicate` for Q). +- Within a PE, assigned tokens are **appended contiguously** (slot + 0,1,2,…). Global indices are strided (`i, i+N, …`) but the per-PE + buffer is dense ⇒ DMA stays contiguous. +- Slot → global: `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`. + +### 2.2 Driver per-launch duties (minimal) +The driver supplies, per launch, the bases + counter + rotation; the +kernel derives the rest: + +| Launch arg | Purpose | +|---|---| +| `K_base[pe]`, `V_base[pe]` | per-PE KV buffer bases (from tensor VA) | +| `O_base` | attention output destination | +| `global_token_counter` | current sequence position | +| `start_pe = request_id mod N` | Layer-2 rotation | +| `pe_id` (`tl.program_id`), `N` | geometry | +| `q_block_meta` | prefill/SP query block start + length | + +Kernel-derived (plain arithmetic): +- **my turn to write this step:** `(start_pe + counter) mod N == pe_id` +- **my valid length:** `base = counter // N; rem = counter % N; + my_len = base + (1 if ((pe_id − start_pe) mod N) < rem else 0)` +- **read range:** tiles `0 .. ⌈my_len / TILE⌉`. +- **causal bounds / per-tile skip:** from global positions of the query + block vs each tile. + +Driver = bases + counter + rotation. The single formula +`(request_id + token_idx) mod N` is the entire placement policy. + +--- + +## 3. Per-tile op sequence (greenlet `tl`) + +One iteration = one KV tile on one PE. Real `tl` names (`tl_context.py`), +with `tl.load_async` (ADR-0062), `tl.broadcast` (ADR-0061), +`tl.scratch_scope` (ADR-0063): + +```python +# running state (persistent arena — allocated once, outside the scope) +# m: [G, T_q] l: [G, T_q] O: [G, T_q, d] + +with tl.scratch_scope(): # per-tile temporaries recycled + Kj = tl.wait(f_k[j]) # prefetched (ADR-0062) + Sj = tl.dot(q_g, tl.trans(Kj)) * softmax_scale # [G, TILE]; q_g is GQA-batched + if mask_j is not None: + Sj = Sj + mask_j # additive causal mask (boundary tile) + m_j = tl.max(Sj, axis=-1) + m_new = tl.maximum(m, m_j) + P = tl.exp(Sj - m_new) # no full-matrix softmax; streaming + corr = tl.exp(m - m_new) # rescale factor for old accumulators + l = l * corr + tl.sum(P, axis=-1) + Vj = tl.wait(f_v[j]) + O = O * corr + tl.dot(P, Vj) # P·V folded into running O + m = m_new + if j + PREFETCH < n_tiles: # keep the pipeline full + f_k[j + PREFETCH] = tl.load_async(K_base + (j+PREFETCH)*TILE*d, (TILE, d)) + f_v[j + PREFETCH] = tl.load_async(V_base + (j+PREFETCH)*TILE*d, (TILE, d)) +``` + +Notes: +- `q_g` is the GQA-batched query: `[G, d]` (decode) or `[G, T_q, d]` + (prefill). One K/V tile load serves all `G` rows — the GQA reuse lever. + Where the matmul needs the KV head expanded to `G`, use + `tl.broadcast` (ADR-0061), never the symbolic `_view`. +- The V load is issued (prefetched) *before* it is needed so its DMA + overlaps the Q·Kᵀ + softmax of the same/earlier tile. +- Masking: the kernel builds the boundary-tile mask from query/KV global + offsets and adds it (`Sj + mask_j`); full-past tiles pass `None`; + full-future tiles are **skipped** (the `if` never enqueues them). +- No `CompositeCmd` is used. (A future `tl.composite` flash kind is an + *optional* optimisation — §8 item 4 — not a requirement.) + +--- + +## 4. Reduction (KV-parallel / SP combine) — tree to root + +After its tile sweep each PE holds `(m_i, ℓ_i, O_i)` (unnormalised). +Combine via the associative/commutative log-sum-exp merge (identical math +to the baseline's fold, `_attention_mesh_mlo.py:117-122`): + +```python +def merge(m_a, l_a, O_a, m_b, l_b, O_b): + m = tl.maximum(m_a, m_b) + sa, sb = tl.exp(m_a - m), tl.exp(m_b - m) + return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb +# final: O = O_root / l_root # normalise once at the tree root +``` + +- **Timing:** exchange is **after P·V** (each PE has its final `O_i`). +- **Topology:** mesh tree, depth `⌈log₂ N⌉`, over physical neighbours. + For `N=8`: level-0 `(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2 + `(3↔7)`, root = `7`. Each tree pair must be a physical `N/S/E/W` + neighbour so `tl.send(dir)`/`tl.recv(dir)` use a real link + (tuning item §9; the SFR install that wires neighbours is + `configure_sfr_intracube_pe_ring`, used by the baseline). +- **Why a tree, not the baseline fan-out:** attention needs `O` at one + place (the PE that writes `O_base`). A tree is `⌈log₂ N⌉` steps + (3 for `N=8`) vs the baseline all-to-all's `N−1` (7). For decode, where + the per-PE sweep is short, the reduction is the dominant cost, so this + is a real win. (The baseline's fan-out replicates the answer to all + ranks — unnecessary here.) +- **Payload:** `O_i` (`d`-vector) is heavy; `m,ℓ` are scalars per `(g, + T_q)`. Sent as separate `tl.send`s (the baseline sends the triplet as + three sends — same pattern). + +Kernel structure (greenlet): + +```python +# leaf / internal node: fold children that send to me, then forward up +for child_dir in tree_children_dirs(pe_id, N): + m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape) + O_c = tl.recv(child_dir, O.shape) + m, l, O = merge(m, l, O, m_c, l_c, O_c) +if not is_root(pe_id): + tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O) +else: + tl.store(O_base, O / l) +``` + +`tl.recv` blocks until data arrives (`tl_context.py:446-499`); the merge +order is fixed by the static tree, so no data-dependent `pop` is needed — +which is why **no hardware/composite `pop`-as-dependency change is +required** (§6). + +--- + +## 5. Per-case kernels + +All four cases share §3 (tile sweep) + §4 (reduction); they differ only +in `N`, query-block width, and which `tile_*` predicates fire. + +### 5.1 Common skeleton + +```python +def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N, + q_block, softmax_scale, *, tl): + pe_id = tl.program_id(axis=rank_axis) + my_len = valid_len(counter, start_pe, pe_id, N) + n_tiles = ceil(my_len / TILE) + q_g = load_Q_group(q_ptr) # [G,d] decode / [G,T_q,d] prefill; tl.broadcast for GQA + m, l, O = init_running() # persistent arena: -inf, 0, zeros + prime_prefetch(k_ptr, v_ptr, n_tiles) # tl.load_async first PREFETCH tiles + for j in range(n_tiles): + if tile_all_future(j, q_block): # causal skip (kernel if) + continue + run_tile(j, mask_or_null(j, q_block)) # §3 op sequence, in tl.scratch_scope + if N == 1: + tl.store(o_ptr, O / l) # no reduction + else: + tree_reduce_and_store(pe_id, N, o_ptr) # §4 +``` + +### 5.2 DECODE, no SP (`N=1`, one PE owns the head's KV) +- `T_q = 1`, attends to all past KV ⇒ no future tiles, mask only on the + final ragged tile. +- **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8` + query rows of the KV head are batched (`q_g = [G, d]`), so each K/V + tile is loaded once and reused by 8 rows (Q·Kᵀ is an `8×TILE` batched + GEMV, P·V an `8×d`). `tl.broadcast` expands the KV head where needed. +- If `S_pe` fits in scratch (small/medium context) this degenerates to a + **one-shot** partial attention (one `tl.dot` for Q·Kᵀ, one softmax, one + `tl.dot` for P·V) — exactly the baseline `_attention_mesh_mlo` + `_partial_attention`, just GQA-batched. Tiling (§3) only kicks in when + `S_pe` exceeds the scratch scope's tile budget. + +### 5.3 DECODE, with SP / KV-parallel (`N=8`) +- One request's KV is round-robin across 8 PEs; each owns ≈`my_len` + tokens. Each PE GQA-batches its `G=8` query rows over its local tiles. +- `T_q=1` ⇒ short per-PE sweep ⇒ **reduction dominates** ⇒ the §4 tree + (3 steps) is the structurally important part. Reduction latency is + hidden by other concurrent decode tokens when batched, or by the long + per-PE tile sweep for long single-stream context. + +### 5.4 PREFILL, no SP +- Whole prompt resident; query is a block of `T_q` tokens (chunked). +- Causality is real, per query block `[qs, qe)` vs KV tile `[ks, ke)`: + - `ke ≤ qs` → `tile_all_past` → `mask=None`, full compute. + - `ks ≥ qe` → `tile_all_future` → **skip** (kernel `if`). + - overlap → `tile_partial` → kernel builds a triangular additive mask, + the tile op sequence adds it (`Sj + mask_j`). +- GQA batches the group's query heads along the same axis as decode. + +### 5.5 PREFILL, with SP (Ring KV) +- KV sharded across `N` PEs (the ring); each step delivers a peer's KV + block over IPCQ. Running `(m,ℓ,O)` is carried **across ring steps** + (Python handles, exactly as `_attention_mesh_kv` does today). +- IPCQ overlaps next step's KV receive with current step's compute via + `tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists). +- **Causal ring skip:** if the incoming step's KV block is entirely + *after* the local query block, skip its compute (don't recv-consume / + don't fold) — eliminates ≈half the steps for causal attention. +- Receive buffers ping-pong (persistent arena, not recycled). +- After the ring, `(m,ℓ,O)` is final per query-owning PE; if KV-parallel + splits remain, §4 tree-merges them, then normalise + store. + +The baseline `_attention_mesh_kv` already implements the ring fold; this +ADR adds GQA reuse, causal step-skip, and `recv_async` overlap. + +--- + +## 6. Why no hardware / composite change is needed + +- The only place a HW/composite `pop`-as-dependency would help is a lone + reduction with no other work to hide a poll — i.e. batch=1 **and** short + context. The target is **agentic = low batch, long context** ⇒ each PE + has many KV tiles; the §4 tree's `tl.recv` blocks are covered by the + sweep work of concurrent tokens / long context. +- The §4 reduction uses a **static** tree, so the recv order is fixed at + compile time — no data-dependent pop. `tl.recv` (blocking) suffices. +- **Decision: ship with the greenlet `tl.send`/`tl.recv` collective.** + Revisit a HW `pop` only if a short-context, single-stream, + latency-critical target emerges. + +--- + +## 7. Control vs execution split (the load-bearing principle) + +| Concern | Owner | +|---|---| +| tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) | +| address / offset / valid-length arithmetic | **kernel** (from counter) | +| reduction scheduling, IPCQ send/recv ordering | **kernel** (static tree) | +| K/V load, Q·Kᵀ, mask add, softmax math, P·V | **`tl` ops** on PE engines | + +The kernel decides; the `tl` ops execute already-decided work. This is +exactly the greenlet model kernbench already supports — no new control +abstraction. + +--- + +## 8. Required kernbench changes + +**Genuinely new (each has a supporting ADR):** + +1. **GQA head broadcast in data mode** — **ADR-0061** (`tl.broadcast`). + *The* blocker for `h_q > h_kv`. Without it the GQA reuse lever cannot + run under `enable_data=True`. +2. **Async HBM tile load / KV prefetch** — **ADR-0062** (`tl.load_async` + + `tl.wait`). The KV-load-bound overlap lever for decode/long-context. +3. **Per-tile scratch recycling** — **ADR-0063** (`tl.scratch_scope`). + Removes the `S=16` ceiling so realistic context lengths run. + +**Algorithm work in the kernel (no new primitive; existing `tl` API):** + +- GQA Q-axis batching (`G` rows share each K/V load) — `tl.broadcast` + + 2-D `tl.dot`. +- Tree reduction to root (§4) replacing the baseline all-to-all fan-out — + pure kernel control flow over `tl.send`/`tl.recv`. +- Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` + + a mask tensor added with `+`. +- Round-robin KV placement / valid-length arithmetic (§2) — launch-arg + arithmetic + `DPPolicy(pe="row_wise")`. + +**Explicitly REJECTED (efficient alternative chosen):** + +4. ~~Composite that chains DMA→MM→VEC→DMA→MM→VEC with carried `(m,ℓ,O)` + register state and a tail IPCQ push.~~ Replaced by the greenlet path: + per-op latency is identical; overlap comes from ADR-0062; small + working set from ADR-0063; running state is Python handles; the IPCQ + push is `tl.send`. Building a bespoke flash-composite command type + + cross-composite register lifetime + an IPCQ-push epilogue is large, + special-purpose, and yields no latency benefit. **A tiled + `tl.composite` "flash" kind remains a possible *future* optimisation** + (it would fold prefetch+recycling into the scheduler), but it is not + required for an efficient kernel and is out of scope here. +5. ~~Hardware `pop`-as-dependency.~~ Out of scope (§6). +6. ~~RoPE / QKV projection / KV-cache write inside this kernel.~~ Upstream + `qkv_rope` (P1–P5). Folding RoPE in would force re-rotating past tiles + every decode step and break Ring Attention's post-RoPE pass-through. + +--- + +## 9. Open tuning items (measured in kernbench, not blocking) + +1. **KV tile prefetch depth** (ADR-0062) — dominant for KV-load-bound; + sweep 2→4. +2. **PE↔mesh-neighbour mapping for the reduction tree** — ensure each + depth-`⌈log₂ N⌉` pair is a physical `N/S/E/W` neighbour; bad mapping + adds hops. Verify against the SFR install. +3. **TILE size** — balance scratch residency (`K_tile + S/P + V_tile + + O_acc + G-way GQA`) against DMA efficiency; interacts with ADR-0063. +4. **Ring buffer ping-pong vs `recv_async` depth** in §5.5. +5. **One-shot vs tiled crossover for decode** (§5.2) — the `S_pe` + threshold where tiling beats a single `tl.dot`. + +--- + +## 10. Coverage summary + +| Case | KV placement | Inner loop | Reduction | Masking | +|---|---|---|---|---| +| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only | +| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only | +| Prefill, no SP | resident | tiled per q-block | none | skip future / triangular boundary | +| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary | + +**I/O per case** (full contract §0.5): + +| Case | Inputs | Output | IPCQ partials | +|---|---|---|---| +| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none | +| Decode, SP | `Q[G,1,d]`, per-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,ℓ,O_i)` per PE → tree | +| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none | +| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,ℓ,O)` + tree | + +In all cases: KV-cache writes and RoPE happen **upstream**; output +projection **downstream**; score `S` and probs `P` are never materialised. + +--- + +## 11. Verification plan (Phase 1 test outline) + +SPEC/ADR coverage: R5 (PE↔PE IPCQ, PE↔HBM), R2 (latency by traversal), +ADR-0023/0025 (IPCQ), ADR-0046 (`tl` contract), ADR-0054 (eval bench). + +1. **Correctness vs reference (Phase 2, `enable_data=True`):** for a + small `(G, T_q, S, N)`, kernel `O` matches a numpy FlashAttention + reference within fp tolerance — for all four cases. This is the test + that *requires* ADR-0061 (GQA) to even run with `h_q>h_kv`. +2. **GQA reuse:** with `h_q = G·h_kv`, the K/V `dma_read_count` is + independent of `G` (one load per tile, reused across the group), while + GEMM work scales with `G`. Asserts the lever actually fires. +3. **Tree reduction step count:** decode-SP issues `⌈log₂ N⌉` reduction + rounds (not `N−1`); op_log `ipcq_send`/`recv` counts match the tree. +4. **Causal skip:** prefill skips all `tile_all_future` tiles — GEMM + count equals the lower-triangular tile count, not the full grid. +5. **Long context (ADR-0063):** a sweep at `S` that overflows 1 MiB + without scopes completes and matches the reference. +6. **Prefetch overlap (ADR-0062):** end-to-end latency of the tiled sweep + is below the serial `Σ(load+compute)` (overlap is real). +7. **Determinism:** identical inputs → identical op_log + latency + (SPEC §0.1). diff --git a/docs/adr-proposed/ADR-XXXX-gqa-fused-attention-ahbm.md b/docs/adr-proposed/ADR-XXXX-gqa-fused-attention-ahbm.md deleted file mode 100644 index f0cbac4..0000000 --- a/docs/adr-proposed/ADR-XXXX-gqa-fused-attention-ahbm.md +++ /dev/null @@ -1,397 +0,0 @@ -# ADR-001: AHBM GQA Fused Attention Kernel (Llama3-70B) - -**Status:** Proposed -**Context model:** Llama3-70B -**Decision drivers:** agentic workload → low batch, long context; KV-load-bound decode; SP (Ring KV) for long-context prefill. - -**Algorithm lineage.** This kernel is **FlashAttention** (tiling + online/streaming softmax with fused P·V — no full score matrix materialized). The KV-parallel split-and-combine in §4 is **FlashDecoding** (split-KV with log-sum-exp merge). The SP path in §5.5 is **Ring Attention** (block-wise KV rotated around the ring, accumulated by the same online softmax). No new math is introduced; this ADR maps those known algorithms onto AHBM composite commands + IPCQ. - ---- - -## 0. Reference dimensions (Llama3-70B) - -| Symbol | Meaning | Value | -|---|---|---| -| `H_q` | query heads | 64 | -| `H_kv` | KV heads | 8 | -| `G` | GQA group size = `H_q / H_kv` | 8 | -| `d` | head dim | 128 | -| `L` | layers | 80 | -| `D` | model dim | 8192 | - -Hardware recap: -- **AHBM (chip)** = set of **CUBE**s (memory cubes, each with a logic die containing **PE**s) + **IO die**. -- **IPCQ**: PE↔PE queues. Each PE holds 4 neighbor queue-pairs (one per mesh direction). -- **Composite command**: a fused multi-stage command (load → matmul → vector op → ...) executed as one unit. *Cannot* express a data-dependent `pop` as an internal dependency (see §7). It *can* take a precomputed mask tile as an input and apply it as a stage. -- **Allocation policy (SP):** one query head per CUBE. For GQA-8, one KV head's 8 query heads map to 8 CUBEs. - -**Two orthogonal mapping layers** (decided in design discussion): -- **Layer 1 (KV-parallel, intra-request):** within a request, KV token `i` lands on PE `(start_pe + i) mod N` — round-robin so a single long request is split across `N` PEs and all PEs stay balanced (chunk length differs by ≤1 token). -- **Layer 2 (inter-request):** `start_pe = request_id mod N` rotates the "PE-0 role" per request to spread write/start load. -- Combined destination PE for a token: `pe = (request_id + global_token_idx) mod N`. - -Here `N` = number of PEs participating in the KV-parallel reduction group for one (query head, request). Within a CUBE the reduction group is the set of PEs splitting that head's sequence; reduction stays **intra-CUBE** (query head never spans CUBEs — explicit non-goal). - ---- - -## 0.5 Kernel boundary, preconditions, and I/O contract - -### 0.5.1 Where this kernel sits in the decoder layer - -``` -1. RMSNorm -2. QKV projection (GEMM) ─┐ qkv_rope kernel (SEPARATE, upstream) -3. RoPE on Q and K ─┤ -4. write new K,V → KV cache ─┘ -5. ===== THIS KERNEL: FlashAttention ===== (post-RoPE Q, post-RoPE K-cache, V-cache → O) -6. Output projection (GEMM) out_proj kernel (SEPARATE, downstream) -7. residual add → FFN ... -``` - -**Preconditions (responsibilities of the upstream `qkv_rope` kernel, NOT this kernel):** -- **P1.** Q is already RoPE-rotated. This kernel does **no** rotation on Q. -- **P2.** K-cache stores **post-RoPE** K. Past tiles are never re-rotated at attention time (this is the reason for post-RoPE caching). -- **P3.** For decode, the new step's K row is RoPE-rotated **and appended to its owning PE's K-cache slot by `qkv_rope` before this kernel launches.** For prefill, the whole query block's K is post-RoPE in cache. ⇒ this kernel is **pure read** on the KV cache; it never writes KV. -- **P4.** V is **not** rotated (RoPE applies to Q,K only). V-cache holds raw projected V. -- **P5.** RoPE position used upstream is the token's **absolute global position**, i.e. `global_idx = local_slot*N + ((pe_id - start_pe) mod N)` (§2.1). Round-robin placement must not be confused with RoPE position — the angle depends on global index, never on the local slot. - -(If a fused megakernel is ever desired, stages 2–4 would prepend to §5.1; the chosen design keeps them separate so KV cache is the clean interface and Ring Attention only ever passes post-RoPE tensors.) - -### 0.5.2 Symbols for shapes -`T_q` = query length this launch (decode: 1; prefill/chunk: chunk width). `S` = total context length (keys seen). `S_pe` = keys owned by this PE = `valid_len` (§2.2), `≈ ceil(S/N)`. `bf16` storage, `fp32` for `m,l,O` accumulators. - -### 0.5.3 INPUTS (per kernel launch) - -| Input | Shape (per head) | Layout / location | Notes | -|---|---|---|---| -| `Q` | `[G, T_q, d]` | on-chip regs/SRAM of the PE | post-RoPE (P1). `G=8` query rows of the GQA group batched together. | -| `K_cache` | `[S_pe, d]` | per-PE HBM buffer, **contiguous**, base = `K_base[pe]` | post-RoPE (P2). Read-only. Strided in global index, dense locally (§2.1). | -| `V_cache` | `[S_pe, d]` | per-PE HBM buffer, base = `V_base[pe]` | raw V (P4). Read-only. | -| `global_token_counter` | scalar | launch arg | kernel derives `S_pe`, slot↔global, causal bounds. | -| `start_pe` | scalar = `request_id mod N` | launch arg | Layer-2 rotation. | -| `pe_id`, `N` | scalars | known to PE / launch | reduction-group geometry. | -| `q_block_meta` | `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. | -| `O_base` | address | launch arg | where final O is written (root PE only). | -| `softmax_scale` | scalar | launch arg | `1/sqrt(d)` (=1/√128). | - -For **Ring Attention (§5.5)** add: incoming `K_block, V_block` per ring step arrive via IPCQ/comms into ping-pong receive buffers (post-RoPE), plus `step_kv_global_range` so the kernel can apply the causal step-skip. - -### 0.5.4 OUTPUTS - -| Output | Shape (per head) | Location | Notes | -|---|---|---|---| -| `O` (final) | `[G, T_q, d]` | `O_base`, **written by tree-root PE only** | normalized: `O_acc / l_acc`. fp32→bf16 cast on store. | - -**Intermediate (per non-root PE, on IPCQ, not a kernel-visible output):** - -| Partial | Shape | Channel | Notes | -|---|---|---|---| -| `(m_i, l_i, O_i)` | `m,l`: `[G, T_q]` scalars; `O_i`: `[G, T_q, d]` unnormalized | pushed to `tree_parent` via neighbor IPCQ | log-sum-exp merge payload (§4). `O_i` is the heavy part. | - -**No-SP cases (`N=1`):** there are no IPCQ partials; the single PE's running `(m,l,O)` is normalized in place and written to `O_base`. The kernel's only output is `O`. - -**Explicitly NOT outputs of this kernel:** KV cache writes (done upstream, P3), output projection (downstream stage 6), the score matrix `S` and probabilities `P` (never materialized — FlashAttention). - ---- - -## 1. Decision - -Fuse the attention inner pipeline into a single composite per KV tile: - -``` -K_j load → Q·Kjᵀ → (+ mask stage) → online-softmax update → V_j load → P·V → running (m,l,O) update → downstream IPCQ push -``` - -Cross-PE combination (KV-parallel / SP) is a **log-sum-exp tree reduction** over `(m, l, O)` performed **after P·V**, flowed through IPCQ with a **kernel-driven async-submit + bounded-polling loop** (no hardware `pop`-as-dependency change required for the target workload). - -Control flow (tile skip, mask generation, reduction scheduling, address arithmetic) lives in the **kernel**; the **composite** only executes already-decided work. - ---- - -## 2. Memory layout & driver responsibilities - -### 2.1 KV cache allocation -- Pre-allocate per-PE KV buffers sized to **per-PE max context** = `ceil(max_context / N)` tokens × `d` × dtype, for K and V separately. -- Within a PE, tokens assigned to it are **appended contiguously** (slot 0,1,2,...). Even though global indices are strided (`i, i+N, i+2N, ...`), the per-PE buffer is dense ⇒ DMA stays contiguous, bandwidth not fragmented. -- Slot → global index is implicit: `global_idx = local_slot * N + ((pe_id - start_pe) mod N)`. - -### 2.2 Driver per-step duties (kept minimal) -The driver does **not** compute write slots or read ranges. It supplies, per kernel launch: - -| Launch arg | Purpose | -|---|---| -| `K_base[pe]`, `V_base[pe]` | base address of each PE's KV buffer | -| `O_base` | attention output destination (separate from KV write) | -| `global_token_counter` | current sequence position; kernel derives everything from it | -| `start_pe` (= `request_id mod N`) | Layer-2 rotation | -| `pe_id` | known to PE; used with counter for `mod`/`div` | -| `q_block_meta` | for prefill/SP: query block start + length | - -From `global_token_counter`, `start_pe`, `pe_id` the kernel derives: -- **"is it my turn to write this step":** `(start_pe + counter) mod N == pe_id` -- **my valid length:** `base = counter / N; rem = counter % N; my_len = base + ( ((pe_id - start_pe) mod N) < rem ? 1 : 0 )` -- **write offset:** `my_len` (next free slot) — kernel computes `K_base[pe] + my_len*d`. -- **read range:** tiles `0 .. ceil(my_len / TILE)`. -- **causal bounds & per-tile skip:** from global positions of query block vs each tile. - -Driver = address bases + counter + rotation. One formula (`(request_id + token_idx) mod N`) is the entire placement policy. - ---- - -## 3. Composite structure - -One composite instance = one KV tile's worth of work on one PE. - -``` -COMPOSITE attn_tile(q_reg, K_base, V_base, tile_idx, mask_tile_or_null, - m_reg, l_reg, O_reg, push_target): - S1: DMA K_tile ← K_base + tile_idx*TILE*d # contiguous load - S2: MM S = q_reg · K_tileᵀ (scale 1/sqrt(d)) - S3: VEC if mask_tile != null: S += mask_tile # masked boundary tile - S4: VEC online softmax: - m_new = max(m_reg, rowmax(S)) - P = exp(S - m_new) - l_reg = l_reg*exp(m_reg - m_new) + rowsum(P) - O_reg = O_reg*exp(m_reg - m_new) # rescale accumulator - m_reg = m_new - S5: DMA V_tile ← V_base + tile_idx*TILE*d # issued early, hidden behind S2-S4 - S6: MM O_reg += P · V_tile - S7: (tail) PUSH (m_reg,l_reg,O_reg) → push_target # only on reduction-producing composite -``` - -Notes: -- `m_reg, l_reg, O_reg` are **carried across composites** (running state), so consecutive tile composites form an in-order chain on one PE. -- S5 (V load) is scheduled so its latency overlaps S2–S4; with prefetch depth ≥ 2 the next tile's K/V are already in flight. -- S7 push is only on the **last** composite of a PE's tile sweep (when its `O_local` is final). Push is a tail *action*, not a dependency — composites can express push, not pop. -- Mask handling: composite only **applies** a mask tile it is *given*. The kernel generates the boundary-tile mask from query/KV offsets and passes it in. Full-past tiles get `null`; full-future tiles are never enqueued (skipped in kernel). - ---- - -## 4. Reduction (KV-parallel / SP combine) - -After every PE finishes its tile sweep it holds `(m_i, l_i, O_i)` (unnormalized). Combine via associative/commutative log-sum-exp merge: - -``` -merge((m_a,l_a,O_a),(m_b,l_b,O_b)): - m = max(m_a,m_b) - l = l_a*exp(m_a-m) + l_b*exp(m_b-m) - O = O_a*exp(m_a-m) + O_b*exp(m_b-m) - return (m,l,O) -final O = O_root / l_root # normalize once at tree root -``` - -- **Timing: exchange is AFTER P·V** (each PE must finish its `O_local`). -- **Topology:** mesh tree, depth `ceil(log2(N))`. For `N=8`: level-0 pairs (0↔1,2↔3,4↔5,6↔7), level-1 (1↔3,5↔7), level-2 (3↔7), root = PE matching last-in-group. PE↔mesh-neighbor numbering must be chosen so each tree pair is a physical 4-direction neighbor (tuning item, §9). -- **Payload:** `O` is the heavy part (`d`-vector = 128 elems); `m,l` are scalars appended. - ---- - -## 5. Kernel pseudocode - -### 5.1 Common loop skeleton (all four cases share this) - -```python -def attn_kernel(pe_id, start_pe, counter, K_base, V_base, O_base, q_block): - # ---- derive geometry ---- - my_len = valid_len(counter, start_pe, pe_id, N) - n_tiles = ceil(my_len / TILE) - q_reg = load_Q(q_block) # decode: 1 row; prefill: q_block rows - m,l,O = -inf, 0, zeros(d) - - PREFETCH = 2 - issued = 0 - # prime prefetch - for j in range(min(PREFETCH, n_tiles)): - if not tile_all_future(j, q_block): # kernel-level if (§7) - enqueue(build_composite(j, q_block)); issued += 1 - - # main sweep - for j in range(n_tiles): - wait_complete(j) # in-order; running m,l,O updated - nxt = j + PREFETCH - if nxt < n_tiles and not tile_all_future(nxt, q_block): - enqueue(build_composite(nxt, q_block)) - # opportunistic reduction progress (see 5.5) — only when this PE - # already produced O_local AND has spare queue depth - try_drain_reduction() - - # ---- reduction phase ---- - push_local((m,l,O), to=tree_parent(pe_id)) - if is_tree_internal_or_root(pe_id): - reduce_tree_node(pe_id) # async-submit + bounded poll - if is_root(pe_id): - O_final = O_acc / l_acc - store(O_base, O_final) - -def build_composite(j, q_block): - if tile_partial(j, q_block): # boundary / diagonal tile - mask = make_mask_tile(j, q_block) # kernel computes from offsets - else: # tile_all_past - mask = null - return composite.attn_tile(q_reg, K_base, V_base, j, mask, m,l,O, push_target) -``` - -`tile_all_future / tile_partial / tile_all_past` are pure position comparisons (query global pos vs tile's KV global-pos range). Decode degenerates them (see below). - -### 5.2 DECODE, **no SP** (single PE owns the whole head's KV) - -- `N = 1` for the reduction group → no cross-PE merge, no tree. -- Query length = 1; it attends to **all** past KV ⇒ `tile_all_future` never true, `tile_partial` only on the final ragged tile, no real masking. -- GQA reuse: the `G=8` query heads sharing one KV head are processed in the **same** PE/composite stream so each K/V tile is loaded once and reused by 8 query rows (Q·Kᵀ becomes an 8×TILE GEMV-batch, P·V an 8×d). This is the main lever; decode is KV-load-bound, so amortizing K/V load over the group is where the win is. -- Bottleneck = K/V tile DMA. Tune `PREFETCH` (≥2) so the queue never starves. - -``` -load q_reg = [8 query rows for this KV head] # GQA group batched -for j in 0..n_tiles: - composite.attn_tile(...) # 8×TILE Q·Kᵀ, mask=null (except last) -store O[8 rows] # no reduction -``` - -### 5.3 DECODE, **with SP / KV-parallel** (N=8, head split across 8 PEs) - -- One request's KV is round-robin across 8 PEs; each PE owns ~`my_len` tokens (≤1 apart). -- Each PE still batches its `G=8` query rows (GQA) over its local tiles. -- Query length 1 ⇒ no future tiles, minimal mask ⇒ tile sweep is short per PE; **reduction is the structurally interesting part**, done once per token at sweep end. -- Reduction latency for a single token is hidden by the remaining tiles of *other* concurrent decode tokens if batched; for strict batch=1 long-context, hidden by the fact that each PE still has many tiles (long context) — sweep work >> reduction handshake. - -``` -per PE: - my tiles sweep (short, GQA-batched) → (m_i,l_i,O_i) - push to tree_parent - tree merge (depth 3) → root normalizes → O[8 rows] for this KV head -``` - -### 5.4 PREFILL, **no SP** - -- Whole prompt resident on the PE(s) for the head; query is a block of `T` tokens (chunked). -- Now causality is real: for query block `[qs, qe)` and KV tile covering `[ks, ke)`: - - `ke <= qs` → `tile_all_past` → mask=null, full compute. - - `ks >= qe` → `tile_all_future`→ **skip enqueue** (kernel `if`). - - overlap → `tile_partial` → kernel builds triangular mask tile, composite applies it. -- Chunked prefill walks query blocks; for each query block runs the tile sweep with the skip/mask decisions above. -- GQA: 8 query heads of a group share K/V tiles within the block ⇒ batched along the same axis as decode. - -``` -for q_block in chunks(prompt): - m,l,O = reset - for j in tiles_up_to(q_block.end): - if tile_all_future(j,q_block): continue # skip - composite.attn_tile(..., mask=mask_or_null(j,q_block)) - store O[q_block] -``` - -### 5.5 PREFILL, **with SP (Ring KV)** - -- KV is sharded across `N` PEs (the SP ring); each ring step delivers a different KV block from a peer. -- **Composite is instantiated per ring step's KV block:** `K load → Q·Kᵀ → mask → online-softmax update → V → P·V → running (m,l,O)`. Running `(m,l,O)` is carried **across ring steps**. -- IPCQ overlaps **next step's KV receive (comms)** with **current step's compute**. KV and V of a step arrive together (it's comms, not local DMA) ⇒ V is already present when P·V runs. -- **Causal ring optimization via kernel `if`:** if the incoming step's KV block is entirely *after* the local query block, the whole step's compute is **skipped** (don't enqueue the composite) — can eliminate ~half the ring steps for causal attention. -- Receive buffers ping-pong; buffer swap aligned to composite boundary. -- After the ring completes, `(m,l,O)` is final per PE → same log-sum-exp tree reduction to produce O. (If the ring already linearly accumulated the full sequence on each query-owning PE, the "reduction" is just the running state; if KV-parallel splits remain, tree-merge them.) - -``` -init m,l,O -recv KV_block[0] -for step in 0..N-1: - issue_recv(KV_block[step+1]) # overlap comms - if step_all_future(step, q_block): # causal ring skip - continue - mask = step_mask(step, q_block) # null if fully past - composite.attn_tile(q_reg, Kbuf[step%2], Vbuf[step%2], 0, mask, m,l,O, _) -swap buffers each step -# m,l,O now final for this PE's query block -reduce_tree → normalize → store O -``` - -### 5.6 Reduction drain loop (the no-hardware-change mechanism) - -Because composites cannot `pop` as a dependency, the kernel does it explicitly. The poll never sits on the critical path because there is always other tile work to enqueue (long context): - -```python -def reduce_tree_node(pe_id): - children = tree_children(pe_id) - acc = local_state - pending = set(children) - while pending: - for c in list(pending): - if ipcq_nonempty(dir_to(c)): # poll ONLY the child direction - part = ipcq_pop(dir_to(c)) - acc = merge(acc, part) - pending.remove(c) - if pending: - # miss → do useful work instead of busy-wait - if has_unissued_tiles(): enqueue(next_tile_composite()) - else: spin_short() # rare: only batch=1 + short ctx - if not is_root(pe_id): - ipcq_push(dir_to(parent), acc) - else: - global_acc = acc -``` - ---- - -## 6. Why no hardware change for the target workload - -- The only place a HW `pop`-as-dependency helps is a **lone reduction with no other work to hide the poll** — i.e. batch=1 *and* short context. -- Target is **agentic = low batch but long context** ⇒ each PE has many KV tiles; a reduction poll-miss is always covered by issuing the next tile composite. The handshake cost is hidden, not on the critical path. -- HW `pop` would also **remove the kernel's ability to choose "do another tile instead of waiting"** and complicate buffer-lifetime/deadlock reasoning (a blocked composite holds PE resources). For long-context that flexibility is worth more than the saved handshake. -- **Decision: ship with software async-submit + bounded poll. Revisit HW `pop` only if a short-context, single-stream, latency-critical target emerges.** - ---- - -## 7. Control vs execution split (the load-bearing principle) - -| Concern | Owner | -|---|---| -| tile skip (future), mask generation, causal bounds | **kernel** (`if`, arithmetic) | -| address/offset/valid-length arithmetic | **kernel** (from counter) | -| reduction scheduling, IPCQ poll/push timing | **kernel** | -| K/V load, Q·Kᵀ, mask-apply, softmax, P·V, push | **composite** (no internal branch) | - -Composites never branch on data and never `pop`. They receive fully-decided inputs (which tile, which mask, where to push). This is what lets the design avoid HW changes. - ---- - -## 8. Required changes to the composite command implementation - -1. **Multi-stage fusion across DMA + 2×MM + VEC in one command.** Current composites must support the chain `DMA(K) → MM → VEC(mask+softmax) → DMA(V) → MM(P·V) → VEC(accumulate)` as a single fused unit, with the two DMAs schedulable at *different* dependency points (K early, V mid) rather than both up front. -2. **Carried register state across composite instances.** `(m, l, O)` accumulators must persist between consecutive composites on a PE (running flash state) instead of being command-local. Needs a stable register/SRAM binding the kernel passes in and out. -3. **Mask tile as an input operand.** Composite must accept an optional mask tile and apply it as a VEC stage (`S += mask`). No mask generation inside the composite. -4. **Tail push action to a named IPCQ direction.** Composite must be able to, as a terminal action, write `(m,l,O)` to a neighbor queue-pair identified by direction. Push only — **pop remains a kernel operation** (no internal pop-dependency). -5. **Early/decoupled V DMA scheduling.** The V load stage must be issuable so its latency overlaps the preceding MM/VEC stages (prefetch depth ≥ 2 across instances), not serialized after softmax. -6. **GQA batching on the Q axis.** Q·Kᵀ and P·V stages must accept a batched Q of `G=8` rows so one K/V tile load serves the whole query group (single KV-head reuse). -7. **(NOT required) internal data-dependent branch / pop-as-dependency.** Explicitly out of scope; kernel handles all branching. Listed so reviewers don't add it speculatively. -8. **(NOT required here) RoPE / QKV projection.** Handled by the upstream `qkv_rope` kernel (§0.5). This kernel assumes post-RoPE Q and post-RoPE K-cache and performs **no rotation and no KV-cache write**. Listed so reviewers don't fold RoPE into the attention composite (doing so would force re-rotating past tiles every decode step and break Ring Attention's post-RoPE pass-through). - ---- - -## 9. Open tuning items (not blocking; measured in Kernbench) - -1. **KV tile prefetch depth** — dominant for KV-load-bound decode/prefill; sweep 2→4. -2. **PE↔mesh-neighbor mapping for the reduction tree** — ensure each depth-3 tree pair is a physical 4-direction neighbor; bad mapping adds hops. -3. **TILE size** — balance SRAM residency (K_tile + S/P + V_tile + O_acc + 8-way GQA) against DMA efficiency. -4. **Ring buffer ping-pong vs prefetch depth** interaction in §5.5. - ---- - -## 10. Coverage summary - -| Case | KV placement | Composite unit | Reduction | Masking | -|---|---|---|---|---| -| Decode, no SP | 1 PE, all KV | per tile, GQA-batched | none | last tile only | -| Decode, SP | round-robin N PEs | per tile, GQA-batched | tree (depth 3) | last tile only | -| Prefill, no SP | resident | per tile per q-block | none | skip future / triangular boundary | -| Prefill, SP (Ring) | ring-sharded | per ring step | running state + tree | causal step-skip + boundary | - -All four share the §5.1 skeleton and the §3 composite; they differ only in `N`, query-block width, and which `tile_*` predicates fire. - -**I/O per case** (see §0.5 for full contract): - -| Case | Inputs | Output | Partials on IPCQ | -|---|---|---|---| -| Decode, no SP | `Q[G,1,d]` post-RoPE, full `K_cache[S,d]`, `V_cache[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none | -| Decode, SP | `Q[G,1,d]`, per-PE `K_cache[S_pe,d]`, `V_cache[S_pe,d]` | `O[G,1,d]` (root PE) | `(m,l,O_i)` per PE → tree | -| Prefill, no SP | `Q[G,T_q,d]` post-RoPE, `K_cache[≤end,d]`, `V_cache` | `O[G,T_q,d]` at `O_base` | none | -| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K_block,V_block` (post-RoPE) per step | `O[G,T_q,d]` (root PE) | running `(m,l,O)` + tree partials | - -In all cases: KV-cache writes and RoPE happen **upstream**; output projection happens **downstream**; score `S` and probs `P` are never materialized.