From f4f55b2c1c4519d244865947bbb868893bfff49a Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 3 Jun 2026 17:38:55 -0700 Subject: [PATCH] gqa(adr): add supporting feature ADRs 0061-0063 for GQA fused attention Proposed prerequisites surfaced while evaluating the GQA fused-attention ADR against the actual kernbench tl/sim_engine implementation: - ADR-0061 tl.broadcast: data-faithful GQA head reuse (fixes the MemoryStore nbytes check that forces h_q==h_kv==1 today). - ADR-0062 tl.load_async: non-blocking HBM tile load for KV prefetch (KV-load-bound decode/long-context overlap). - ADR-0063 tl.scratch_scope: per-tile scratch recycling (removes the 1 MiB bump-allocator ceiling that caps context at S=16). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ADR-0061-prog-tl-broadcast.md | 193 ++++++++++++++++++ .../ADR-0062-prog-tl-async-load.md | 148 ++++++++++++++ .../ADR-0063-prog-tl-scratch-scope.md | 162 +++++++++++++++ 3 files changed, 503 insertions(+) create mode 100644 docs/adr-proposed/ADR-0061-prog-tl-broadcast.md create mode 100644 docs/adr-proposed/ADR-0062-prog-tl-async-load.md create mode 100644 docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md diff --git a/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md b/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md new file mode 100644 index 0000000..23be02e --- /dev/null +++ b/docs/adr-proposed/ADR-0061-prog-tl-broadcast.md @@ -0,0 +1,193 @@ +# ADR-0061: `tl.broadcast` / `tl.repeat` — data-faithful GQA head reuse + +## Status + +Proposed + +> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). This ADR +> removes the single hard blocker that keeps the existing attention +> kernels at `h_q == h_kv == 1`: there is no way to expand a KV head to +> the `G` query heads that share it without violating the MemoryStore +> byte-conservation check. + +## Context + +### The GQA reuse lever + +Grouped-Query Attention amortises K/V loads: `G = H_q / H_kv` query heads +share one KV head. For Llama3-70B `G = 8`. In a kernel this means a KV +tile of shape `[S_kv, d]` (for one KV head) must be matmul'd against the +`G` query rows of its group. Loading K/V **once** and reusing it across +the group is the decode-time efficiency win (ADR-0060 §5.2). + +### Why it does not work today + +The existing mesh kernels reshape tensors with a metadata-only helper +`_view` (`src/kernbench/benches/_attention_mesh_kv.py:25-36`, +`_attention_mesh_mlo.py:25-36`): + +```python +def _view(handle, new_shape): + return TensorHandle(..., shape=new_shape, nbytes=handle.nbytes, ...) + # ^^^ new shape ^^^ OLD nbytes (unchanged) +``` + +This rewrites `shape` but keeps the original `nbytes`. It is purely +symbolic. It survives **Phase 1 (timing)** because Phase 1 never reads +tensor bytes. It fails **Phase 2 (`enable_data=True`)** because the +data executor materialises every op through `MemoryStore.read`, which +enforces byte conservation +(`src/kernbench/sim_engine/memory_store.py:67-73`): + +```python +if shape is not None and arr.shape != shape: + if arr.nbytes != np.prod(shape) * arr.dtype.itemsize: + raise ValueError(f"Shape mismatch: stored {arr.shape} ({arr.nbytes}B) " + f"vs requested {shape} ({...}B)") + arr = arr.reshape(shape) +``` + +A reshape conserves bytes; a **broadcast does not** (`[S_kv, 1, d]` → +`[S_kv, G, d]` multiplies the element count by `G`). So the moment a +kernel tries to view a `h_kv`-headed K as if it had `h_q` heads, the +stored array has `1/G` of the requested bytes and `read` raises. + +The current test suite documents this as a deliberate limitation +(`tests/attention/test_milestone_gqa_llama70b.py:137-142`): + +> "v1 uses `h_q == h_kv == 1` to avoid … GQA broadcast view (which is +> symbolic and does not survive MemoryStore's nbytes check under +> simulator data execution). Real GQA (`h_q > h_kv`) is deferred." + +There is **no** broadcast / repeat / expand primitive in the `tl` API +today (`src/kernbench/triton_emu/tl_context.py` exposes `trans`, +`zeros`, `full`, `arange` as metadata-only ops; none replicate data). + +## Decision + +Add a **data-faithful broadcast op** to the `tl` API and the data +executor. Unlike `_view`, it updates `nbytes` and, in Phase 2, actually +replicates the underlying ndarray. The latency it charges models the +on-chip register/SRAM fan-out, not an HBM re-load (the KV tile is read +from HBM once; broadcast is an on-chip replication). + +### D1. `tl` surface + +```python +def broadcast(self, x: TensorHandle, new_shape: tuple[int, ...]) -> TensorHandle: + """Replicate x along size-1 axes to new_shape (numpy broadcast rules). + + new_shape must be broadcast-compatible with x.shape: every axis is + either equal to x's, or x's is 1 (the axis being expanded). nbytes of + the result reflect new_shape (unlike the metadata-only _view). + """ +``` + +`tl.repeat(x, axis, n)` is offered as a thin convenience wrapper over +`broadcast` for the common "insert a group axis and expand it to G" +pattern, but `broadcast` is the primitive. + +### D2. Command + op_log + +Emit a new `BroadcastCmd` (sibling of `MathCmd`): + +```python +@dataclass(frozen=True) +class BroadcastCmd: + src: TensorHandle + out: TensorHandle # out.nbytes reflects new_shape + new_shape: tuple[int, ...] + data_op: bool = True +``` + +op_log entry: `op_kind="math"`, `op_name="broadcast"` (it runs on the +MATH/vector engine, like the other element-wise ops). The output handle +is allocated from PE scratch via the existing `_make_compute_out` +(`tl_context.py:148-156`), so `out.nbytes` is computed from `new_shape` +and the byte-conservation check passes in Phase 2. + +### D3. Phase 2 data execution + +Add `_execute_broadcast` to the data executor +(`src/kernbench/sim_engine/data_executor.py`): + +```python +def _execute_broadcast(self, op): + p = op.params + src = self._resolve_read(p["src_space"], p["src_addr"], + p["src_shape"], p["dtype"], op.t_start) + result = np.broadcast_to(src, p["new_shape"]).copy() # .copy() → writable, owns bytes + self.store.write(p["dst_space"], p["dst_addr"], result) +``` + +`.copy()` is required: `np.broadcast_to` returns a zero-stride view, and +MemoryStore stores references — a later consumer's nbytes check must see +a real `prod(new_shape)`-element array. + +### D4. Latency model + +Broadcast latency = vector-engine cost of writing `prod(new_shape)` +elements (reuse `pe_math._compute_ns(num_elements)`, +`src/kernbench/components/builtin/pe_math.py:46-51`). This models the +on-chip replication into the register file / SRAM that real GQA hardware +performs. It is **not** an HBM transaction — no `dma_read` is logged, +consistent with "K/V tile loaded once" (ADR-0060 §5.2). + +## Alternatives + +### A1. Indexed per-head dot (no broadcast op) + +Loop the `G` query heads, each doing its own `tl.dot(Q_g, Kᵀ)` against +the *same* (re-`ref`'d, not re-loaded) KV tile: + +```python +for g in range(G): + s_g = tl.dot(Q[g], K_T) # G separate GEMMs, KV reused via tl.ref +``` + +- Pro: no new primitive; works in Phase 2 today. +- Con: `G×` more GEMM commands (op_log inflation), no batched-GEMM + modelling, and the per-call overhead is charged `G` times. It also + obscures the "single batched matmul over the group" the hardware does. + Rejected as the primary path; acceptable as a fallback if the batched + 2D dot's shape handling proves awkward. + +### A2. Keep `_view`, relax the nbytes check + +Make `MemoryStore.read` tolerate broadcast (stride-0) shapes. Rejected: +it silently breaks byte conservation for *all* readers and would mask +genuine shape bugs elsewhere. The check is a correctness invariant +(ADR-0052), not an obstacle to route around. + +### A3. Materialise the broadcast at deploy time + +Store K/V already expanded to `h_q` heads in HBM. Rejected: defeats the +entire point of GQA (it `G×`'s KV-cache HBM footprint and bandwidth). + +## Consequences + +### Positive +- Unblocks real GQA (`h_q > h_kv`) in Phase 2 — the headline blocker for + ADR-0060 and the deferred milestone work. +- A general, reusable primitive (any kernel needing numpy-style + broadcast benefits, not just attention). +- Byte conservation stays intact; no weakening of MemoryStore. + +### Negative +- One new command type + one data-executor handler + one `tl` method. +- Phase 2 `.copy()` allocates the expanded array in the store; for very + large groups this is real memory. Bounded by tile size in the tiled + design (ADR-0060), so acceptable. + +## Test Requirements + +1. **Phase 1 timing**: `tl.broadcast([S,1,d] → [S,G,d])` logs one + `math/broadcast` op with `num_elements == S*G*d`; no `dma_read`. +2. **Phase 2 data**: broadcasting a known array yields the numpy + `broadcast_to(...).copy()` result; a downstream `tl.dot` consuming it + passes the MemoryStore nbytes check (the exact case that fails today). +3. **GQA end-to-end**: re-run a reduced `milestone-gqa-llama70b` panel + with `h_q=8, h_kv=1` under `enable_data=True`; it must complete (today + it raises `ValueError: Shape mismatch`). +4. **Broadcast-incompatible shape** (e.g. axis size 3 → 8) raises a + clear error at emit time, not deep in Phase 2. diff --git a/docs/adr-proposed/ADR-0062-prog-tl-async-load.md b/docs/adr-proposed/ADR-0062-prog-tl-async-load.md new file mode 100644 index 0000000..e12310d --- /dev/null +++ b/docs/adr-proposed/ADR-0062-prog-tl-async-load.md @@ -0,0 +1,148 @@ +# ADR-0062: `tl.load_async` — non-blocking HBM tile load for KV prefetch + +## Status + +Proposed + +> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and +> long-context attention are **KV-load-bound**; the dominant lever is +> overlapping the next K/V tile's DMA with the current tile's compute. +> Today the only async primitive is for IPCQ comms, not for HBM loads. + +## Context + +### What overlap requires + +FlashAttention streams K/V tiles: while the MATH/GEMM engine works on +tile `j` (Q·Kⱼᵀ → softmax → P·V), the DMA engine should already be +pulling tile `j+1` (and `j+2` at prefetch depth 2). With that overlap a +KV-load-bound kernel runs at roughly `max(compute, dma)` per tile instead +of `compute + dma`. + +### What exists + +- `tl.load(ptr, shape, dtype)` is **blocking**: it emits `DmaReadCmd` + and the greenlet kernel suspends until PE_DMA signals completion + (`src/kernbench/triton_emu/tl_context.py:177-203`; greenlet drive in + `kernel_runner.py`). No two `tl.load`s can be in flight from one + kernel. +- An async pattern **does** exist, but only for IPCQ: + `tl.recv_async(dir, ...) -> RecvFuture` + `tl.wait(future)` + (`tl_context.py:543-560`, `tl_context.py:660-693`). It proves the + machinery — a non-blocking command that returns a future, resolved + later by `tl.wait` — works in the greenlet model. +- The DMA engine already models a read channel as a SimPy resource + separate from the write channel + (`src/kernbench/components/builtin/pe_dma.py`), so concurrent in-flight + reads are representable at the component layer; only the *kernel-facing + API* serialises them. + +There is **no** `tl.load_async`. Prefetch/double-buffering cannot be +expressed. + +## Decision + +Add a non-blocking HBM load that mirrors the existing `recv_async`/`wait` +contract. + +### D1. `tl` surface + +```python +def load_async(self, ptr: int, shape: tuple[int, ...], + dtype: str = "f16") -> LoadFuture: + """Issue a DMA read and return immediately. Resolve with tl.wait(fut) + -> TensorHandle. Multiple loads may be in flight; ordering of + resolution is by tl.wait calls, not issue order.""" +``` + +`tl.wait` is extended to accept a `LoadFuture` (it already dispatches on +`CompletionHandle` vs `RecvFuture` — add a third arm), returning the +loaded `TensorHandle`. + +### D2. Command + +Reuse the existing `DmaReadCmd` with a `blocking=False` flag (mirroring +`IpcqRecvCmd.blocking`), or a thin `DmaReadAsyncCmd` sibling — whichever +keeps PE_DMA's handler simplest. The op_log entry is unchanged +(`memory/dma_read`); asynchrony is a scheduling property, not a new op +kind, so existing op_log consumers and the `dma_read_count` metric +(milestone bench) keep working. + +### D3. Latency / overlap semantics + +- `load_async` charges the **issue** (descriptor push) only; the kernel + proceeds. +- The DMA transfer occupies the read channel for its modelled duration + in parallel with whatever compute the kernel issues next. +- `tl.wait(fut)` blocks only if the transfer has not finished; if it has, + it returns immediately (same fast-path as `recv_async`/`wait`). +- Determinism is preserved: completion is a scheduled event on the + modelled DMA channel/link (SPEC §0.1, R8). No magic — the overlap is + real modelled concurrency, not a hand-waved latency subtraction. + +### D4. Double-buffer usage (the intended pattern) + +```python +# Prime: issue first two tile loads. +f0 = tl.load_async(K_base + 0*tile_bytes, (TILE, d)) +f1 = tl.load_async(K_base + 1*tile_bytes, (TILE, d)) +for j in range(n_tiles): + Kj = tl.wait(f_cur) # resolves the j-th load + if j + 2 < n_tiles: # keep depth-2 pipeline full + f_next2 = tl.load_async(K_base + (j+2)*tile_bytes, (TILE, d)) + s = tl.dot(q, tl.trans(Kj)) # compute overlaps f_{j+1}'s DMA + ... +``` + +(V tiles double-buffer the same way; ADR-0060 §3 schedules the V load so +it overlaps the Q·Kᵀ + softmax of the same tile.) + +## Alternatives + +### A1. Use a tiled `tl.composite` instead + +`tl.composite` already pipelines tiles inside PE_SCHEDULER (DMA of tile +`j+1` overlaps compute of tile `j`) and recycles per-tile scratch — so +the composite path gets prefetch "for free." Rejected as the *sole* +mechanism because today's composite is a single GEMM head + math +epilogues (`pe_commands.py:144-162`); it cannot express the +two-GEMM-with-running-`(m,l,O)` flash inner loop. Building a bespoke +"flash composite" kind is a much larger change than `load_async`, which +is a small, general primitive that also helps non-composite kernels. +`load_async` and a future flash-composite are not mutually exclusive +(ADR-0060 §8). + +### A2. Rely on `recv_async` only (no HBM async) + +Only works for the Ring-Attention path where KV arrives over IPCQ +(comms, already overlappable). The decode and no-SP paths load KV from +**local HBM**, which has no async path today. Insufficient. + +## Consequences + +### Positive +- Enables the headline KV-load-bound optimisation for decode and + long-context, in the proven greenlet model. +- Symmetric with `recv_async`/`wait` — low conceptual surface area. +- General: any bandwidth-bound kernel can prefetch. + +### Negative +- Kernels must manage futures and buffer lifetimes explicitly (the + double-buffer dance). Mitigated by keeping the pattern in a shared + helper used by all four ADR-0060 cases. +- Interacts with scratch lifetime (ADR-0063): in-flight buffers must not + be recycled before their `tl.wait`. The recycling scope must exclude + live prefetch buffers. + +## Test Requirements + +1. **Overlap is real**: a kernel that issues `load_async` then a GEMM of + comparable duration completes in ≈`max(load, gemm)`, not `load+gemm` + (assert end-to-end latency strictly below the serial sum). +2. **Correctness**: `tl.wait(load_async(...))` returns the same bytes as + blocking `tl.load(...)` for the same address (Phase 2). +3. **Two in flight**: two outstanding `load_async`s to distinct + addresses both resolve to correct, independent tensors. +4. **op_log compatibility**: each `load_async` still logs exactly one + `memory/dma_read`, so `dma_read_count` metrics are unchanged vs the + blocking version. diff --git a/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md b/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md new file mode 100644 index 0000000..0fc178f --- /dev/null +++ b/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md @@ -0,0 +1,162 @@ +# ADR-0063: `tl.scratch_scope` — per-tile scratch recycling for long context + +## Status + +Proposed + +> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Long-context +> attention sweeps many K/V tiles; each tile's intermediates +> (`scores`, `P`, `exp`, partial `O`, …) currently allocate fresh +> scratch that is never freed within a kernel invocation. The 1 MiB +> per-PE scratch budget is exhausted long before a realistic context +> length. + +## Context + +### The bump allocator + +`TLContext` allocates every math/compute output handle from a per-PE +scratch pool with a **linear bump cursor** +(`src/kernbench/triton_emu/tl_context.py`; `_scratch_alloc`): + +```python +def _scratch_alloc(self, nbytes): + aligned = (nbytes + 15) & ~15 + addr = self._scratch_base + self._scratch_cursor + self._scratch_cursor += aligned + if self._scratch_cursor > self._scratch_size: # default 1 << 20 = 1 MiB + raise RuntimeError("TLContext scratch overflow: ...") + return addr +``` + +The docstring states the cursor **"resets on every kernel invocation"** — +i.e. only at kernel entry, never *within* the kernel body. Every +`tl.dot`, `tl.softmax`, `tl.exp`, `tl.sum`, `a - b`, `a * b`, … grabs a +fresh slice and nothing is reclaimed until the kernel returns. + +### Why this bites the GQA kernel + +The current milestone bench keeps `S_q = S_kv_per_rank = 16` **explicitly +because** of this limit +(`tests/attention/test_milestone_gqa_llama70b.py:123-148`): + +> "S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the +> simulator's 1 MB per-PE TCM kernel scratch is not exhausted by the +> bump-allocated handle outputs of softmax/exp/dot/sum chains over +> n_ranks ring steps." + +A FlashAttention sweep over `n_tiles` tiles allocates O(`n_tiles` × +per-tile-intermediates). For any realistic context this overflows. The +*math* of flash attention needs only **O(1)** live scratch — the running +`(m, l, O)` plus the current tile's working set — because each tile's +temporaries are dead once that tile is folded into the running state. The +allocator just doesn't know they're dead. + +## Decision + +Add a **scratch scope** that lets a kernel mark a region of allocations +as reclaimable, so per-tile temporaries are recycled while live running +state (and in-flight prefetch buffers) are preserved. + +### D1. `tl` surface — context manager + +```python +with tl.scratch_scope(): + s = tl.dot(q, k_t) # all handles allocated inside the `with` + p = tl.softmax(s) # share a region that is rewound on exit + o_j = tl.dot(p, v) + # ... fold o_j into running (m,l,O) which live OUTSIDE the scope ... +# on __exit__: cursor rewinds to its value at __enter__ +``` + +Semantics: +- `__enter__` records the current `_scratch_cursor` as a save-point. +- `__exit__` restores the cursor to the save-point, freeing everything + allocated inside. +- Handles allocated **outside** the scope (running `m,l,O`, prefetch + buffers held by `LoadFuture` from ADR-0062) keep their addresses — + they were allocated before the save-point or in an enclosing scope. + +### D2. Safety contract + +A handle allocated inside a scope **must not** be read after the scope +exits — its bytes may be overwritten by the next scope's allocations. +The flash loop respects this naturally: the only values that survive a +tile iteration are the running accumulators, which are allocated outside +the per-tile scope and updated by ops *inside* it writing to outside +addresses (the merge writes new running state — see D3). + +### D3. Interaction with running accumulators + +The online-softmax merge reads the old running `(m, l, O)` and the +current tile's `(m_j, l_j, O_j)`, producing new running values. To keep +the new running values outside the recycled region, the merge writes them +to **stable scratch** allocated once before the loop (a small fixed +"running-state" arena, distinct from the per-tile scope). Concretely the +kernel keeps two arenas: + +- **persistent arena** (allocated once): `m, l, O` (and their + double-buffer if needed for the merge). +- **scoped arena** (rewound each tile): `scores, P, exp, O_j, scale_*`. + +This mirrors how real flash-attention SRAM budgeting works: a small +persistent accumulator region + a recycled tile working set. + +### D4. Nesting + +Scopes nest (stack of save-points). Inner scope exit rewinds to the inner +save-point; outer exit rewinds further. This supports an outer +"per-query-block" scope around an inner "per-KV-tile" scope for prefill. + +## Alternatives + +### A1. Round-trip temporaries through HBM + +Store intermediates to HBM and reload to "free" TCM scratch. Rejected: +turns a TCM-resident streaming kernel into an HBM-bandwidth-bound one — +the opposite of the goal, and it pollutes op_log with spurious DMA. + +### A2. Tiled `tl.composite` (scheduler-managed scratch) + +`tl.composite` recycles per-tile scratch inside PE_SCHEDULER +automatically. As in ADR-0062 A1, this is attractive but blocked on a +flash-capable composite kind (two GEMMs + carried `(m,l,O)`), a much +larger change. `scratch_scope` gives the same memory behaviour for the +greenlet kernel with a tiny, general primitive. The two can coexist. + +### A3. Grow the scratch budget + +Bump `pe_tcm.kernel_scratch_mb`. Rejected as a fix: it only pushes the +context-length ceiling out linearly while still leaking O(`n_tiles`) +scratch, and it misrepresents the hardware (real TCM is small; the point +of flash attention is O(1) working set). Useful only as a coarse knob, +not a substitute for recycling. + +## Consequences + +### Positive +- Removes the artificial `S = 16` validation-scale ceiling; enables + realistic context lengths in both timing and data modes. +- Faithfully models flash attention's O(1) working-set property. +- Small, general primitive (any tiled kernel benefits). + +### Negative +- A use-after-scope bug silently reads stale bytes. Mitigated by the D2 + contract, by keeping the scope discipline inside a shared attention + helper, and (optionally) by a debug build that poisons rewound regions. +- Must coordinate with ADR-0062: prefetch buffers are live across tile + iterations, so they belong to the persistent arena, not the scoped one. + +## Test Requirements + +1. **Recycling**: a loop of `N` tiles inside `tl.scratch_scope()` keeps + peak `_scratch_cursor` bounded by one tile's footprint, independent of + `N` (today it grows linearly and overflows). +2. **Correctness**: a flash sweep with scopes produces the same `O` as + the same sweep without scopes at a small `N` that fits without + recycling (Phase 2). +3. **Long context**: a sweep at an `N` that would overflow 1 MiB without + scopes completes (the exact failure the `S=16` cap avoids today). +4. **Persistent-vs-scoped isolation**: running `(m,l,O)` allocated + outside the scope retains correct values across `__exit__`. +5. **Nesting**: nested scopes rewind to the correct save-points.