diff --git a/docs/adr-proposed/ADR-0070-algo-gqa-short-context-attention.md b/docs/adr-proposed/ADR-0070-algo-gqa-short-context-attention.md new file mode 100644 index 0000000..ce0665c --- /dev/null +++ b/docs/adr-proposed/ADR-0070-algo-gqa-short-context-attention.md @@ -0,0 +1,243 @@ +# ADR-0070: GQA Short-Context Attention — Unified A1/A2/A4/B Mapping for Prefill and Decode + +## Status + +Proposed — short-context (single tile to a handful of KV tiles) GQA +attention benchmark covering **both prefill and decode** in **four +mapping modes** (A1/A2/A4/B) and **three composite tiers** (without +composite, GEMM-only composite, composite + softmax_merge fused). +Compute-only attention — weight DMA, ITL batching, and end-to-end +LLM driver scope are out of scope (separate ADRs). + +## Context + +A GQA layer at the LLaMA-3.1-70B headline shape (`h_q=64`, `h_kv=8`, +`d_head=128`, GQA group `G = h_q/h_kv = 8`) on a 4×4 cube SIP can be +mapped to AHBM in several ways that trade KV-cache memory, HBM read +volume per cube, and intra-cube IPCQ traffic against PE parallelism. + +> **Bench-vs-headline shape.** The kernels and tests use a +> `d_head = 64` proxy (`tests/attention/test_gqa_short_context.py` +> `D_HEAD = 64`) so the per-tile working set fits comfortably in +> scratch across all four modes. The mapping decisions are +> `d_head`-independent — only the per-PE GEMM tile bytes scale. + +The four candidate mappings (per *GQA-mapping-on-AHBM-short-context* note §1): + +``` +Mode kv_per_cube C group_size Reduce / broadcast topology +---- ----------- ------- ---------- ---------------------------------- +A1 1 h_kv P (=8) row chain + col bridge (2×4 mesh) +A2 2 h_kv/2 P/2 (=4) row chain only +A4 4 h_kv/4 P/4 (=2) single intra_W / intra_E hop +B 8 1 1 no broadcast / no reduce (single PE) +``` + +Within each cube the 8 PEs split into `kv_per_cube` groups of +`group_size = P/kv_per_cube` PEs; each group owns one KV head. + +The kernel must be able to compare these mappings at multiple context +lengths (8K..64K KV tokens) and across the three composite tiers so +the team can size the per-cube HBM / IPCQ / GEMM tradeoff empirically. + +## Decision + +### 1. Mapping (unified A1/A2/A4/B) + +Two unified kernels select mode at launch via `kv_per_cube ∈ {1,2,4,8}`: + +- **Prefill** (`gqa_attention_prefill_short_kernel`): + - Q-tile split — `T_q` rows floor-balanced across the group's + `group_size` PEs. + - FA2 head fusion — `G` Q heads fused into the M dimension of one + batched GEMM per PE per tile. + - **IPCQ KV broadcast** — group root (`pe_in_group == 0`) loads K/V + from HBM and IPCQ-broadcasts each tile to the rest of its group; + q-tiles are independent so no intra-group reduce. + +- **Decode** (`gqa_attention_decode_short_kernel`): + - Sequence-shard — each PE owns `S_local = S_kv/group_size` tokens + of the group's KV head. + - FA2 head fusion — same as prefill. + - **IPCQ chain reduce** — per-PE partial `(m, ℓ, O)` chain-reduce + up to group root (PE 0), which normalizes and stores. + +Geometry (2×4 mesh): + +``` +group_size = 8 : 2 row × 4 col (A1) +group_size = 4 : 1 row × 4 col (A2) +group_size = 2 : 1 row × 2 col (A4) +group_size = 1 : single PE (B) +``` + +### 2. Shard addressing (ADR-0011 D-VA1 contract) + +Deploy is per-(sip, cube, pe) but `tl.load` receives a single global +VA per tensor. The kernel computes its own shard base offset from +`tl.program_id(axis=0)` (PE id) and `tl.program_id(axis=1)` (cube id): + +```python +# Decode shape (pe=row_wise): K is split across PEs by sequence shard. +cube_K_base = cube_id * kv_per_cube * K_HEAD_BYTES +head_K_base = group_id_in_cube * K_HEAD_BYTES +pe_K_seq_offset = pe_in_group * n_tiles_per_pe * K_TILE_BYTES +k_shard_base = k_ptr + cube_K_base + head_K_base + pe_K_seq_offset +``` + +For **prefill** the K/V dp is `pe=replicate` (group root reads and +broadcasts) so the `pe_in_group` term drops out: +`k_head_shard_base = k_ptr + cube_K_base + head_K_base`. + +Skipping `cube_id` collapses all cubes onto cube 0's HBM region +(observed pre-fix as an 11.5× per-cube DMA imbalance). + +### 3. Three composite tiers + +The same mapping is exercised against three GEMM/MATH dispatch styles +so the contribution of the composite API vs the recipe-driven fusion +can be isolated: + +``` +(1) without composite primitives only (tl.dot, tl.exp, …) +(2) with composite (GEMM-only) tl.composite(op="gemm") +(3) with composite + softmax_merge tl.composite(prologue=[softmax_merge], + op="gemm", out=O, + epilogue=[add]) +``` + +File layout (`src/kernbench/benches/gqa_helpers/short_ctx/`): + +``` +_gqa_attention_{prefill,decode}_short.py (1) +_gqa_attention_{prefill,decode}_short_composite.py (2) +_gqa_attention_{prefill,decode}_short_composite_fused.py (3) +``` + +Tier (2) is **GEMM-only by definition** — no recipe-driven fusion, +so both prefill's and decode's P·V stay `tl.dot`. Recipe-driven +fusion is exactly what tier (3) adds. + +Tier (3) on **multi-cube modes (A1/A2/A4) of prefill** relies on the +D4 supplement: any composite operand that is an IPCQ recv'd slot +(non-root PE's K_T in Q·Kᵀ, V in P·V) is pinned and read in place +instead of being DMA-streamed from HBM. Without the supplement every +recv slot fed to a composite — both Q·Kᵀ's `b=K_T` and P·V's `b=V` +— PageFaults on PA decode (PE scratch addresses overflow the 51-bit +PA range). + +### 4. Caller contract — `_validate_config` + +Each kernel module exposes a single `_validate_config(...)` helper +called by the bench wrapper before launch. The kernel itself is lean +(no inline `if` guards): the contract block is enforced caller-side, +sim cost zero, but catches every silent-shard-corruption foot-gun +(`kv_per_cube=3 → group_size=2` via integer division, non-integer +GQA group, mismatched cube count, partial-tile `S_kv`, …) before any +address arithmetic runs. + +### 5. Tensor layouts (host-side, mode-invariant byte totals) + +- **Q**: `(kv_per_cube·T_q, h_q·d_head/kv_per_cube)` + `dp=(cube=column_wise, pe=replicate)` over `C = h_kv/kv_per_cube`. +- **K**: `(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)` tile-major. + Decode: `dp=(cube=row_wise, pe=row_wise)`. + Prefill: `dp=(cube=row_wise, pe=replicate)` (broadcast model). +- **V**: `(h_kv·S_kv, d_head)` native. Same dp as K. +- **O**: same dp as Q. + +Caller pre-scales Q by `1/√d_head`. + +## Verification + +### Smoke + regression (`tests/attention/test_gqa_short_context.py`) + +- Prefill / decode smoke (all 4 modes, single-tile S_kv). +- Multi-tile coverage (all 4 modes, S_kv chosen so each PE owns + ≥2 tiles). +- Op-count invariants — prefill `dma_writes = group_size·kv_per_cube·C`, + `dma_reads = P·C + 2·kv_per_cube·C·n_tiles`; decode + `dma_writes = kv_per_cube·C`. +- IPCQ topology — short-context kernel never emits inter-CUBE E/W + IPCQ. +- **ADR-0011 D-VA1 regression** — `per_cube_disjoint_src_addrs`, + `per_cube_dma_balanced` (max/min DMA busy ratio < 1.1×) per mode + per phase. + +### Composite smoke + +- `test_{prefill,decode}_composite_smoke[A1/A2/A4/B]` — tier (2), + 4 modes × 2 phases. +- `test_{prefill,decode}_composite_fused_smoke[A1/A2/A4/B]` — tier + (3), 4 modes × 2 phases. +- `test_decode_composite_fused_multitile[A1/A2/A4/B]` — + `n_tiles_per_pe == 2` per mode so the recipe-fused tile loop + actually executes. + +### Mode × context sweep (`tests/attention/test_gqa_short_context_sweep_*.py`) + +Six sweep files (3 tiers × {prefill, decode}) measure each +`(mode, S_kv)` cell and dump the same metric set to CSV: + +``` +wall_us, n_pe, gemm_util, math_count, math_pipeline_us, hbm_bw_util, +hbm_read_mb, hbm_write_kb, ipcq_kb, kv_cache_per_cube_mb +``` + +S_kv ∈ {8K, 16K, 32K, 64K}; outputs under `docs/sweeps/`. + +Composite/fused sweeps count both `gemm_f16` (non-pipeline) **and** +`TileToken/GEMM` (pipeline composite path); fused additionally counts +`TileToken/MATH`. Baseline sweeps count `gemm_f16` only. + +## Known limitations + +1. **Causal mask** — kernel is non-causal. Adding causal masking is + orthogonal to mapping and is scheduled as a separate ADR. +2. **f16 online-softmax accumulator** — `(m, ℓ, O)` are f16 throughout. + Long-context numerical drift will need an f32 accumulator before + correctness-grade use. +3. **End-to-end numeric validation** — current tests are op-count / + topology / per-cube-DMA invariants. No full-kernel + `np.allclose` against a torch/numpy reference yet. +4. **Skinny-M GEMM underfill** — decode (`T_q=1`) gives `M = G = 8`, + well below `PE_SCHEDULER` supertile `TILE_M = 32`. The two GEMM + paths handle this differently: + - **Primitive `tl.dot`** dispatches the GEMM at the actual + `m = 8` (no padding); GEMM time reflects the real 8-row work. + - **`tl.composite(op="gemm")`** tiles at `TILE_M = 32`, padding + `M` 4× with zeros; GEMM time reflects the padded 32-row work. + Consequence: when sweep CSVs show tier (2)/(3) `gemm_util` + higher than tier (1), most of that gap is supertile padding + overhead, **not** extra useful work or fusion savings — read it as + "composite tile shape doesn't match decode-skinny shape" rather + than "composite/fusion is more compute-intensive." This is the + **intended decode shape** for this benchmark; lifting it requires + batched-M (multi-request inference, `B_sys > 1`), tracked as a + separate batched variant — see "Future work". +5. **Multi-cube prefill composite without the ADR-0065 D4 supplement** + PageFaults. The supplement — IPCQ recv'd slots are pinned and read + in place as composite operands — ships alongside this ADR. + +## Future work + +- Batched-M variant (`B_sys ∈ {1,2,4,8,16}`) so composite/fused + pipeline overlap shows up in wall_clock rather than only in + per-engine utilization. +- Long-context (S_kv ≥ 128K) sweep extension for the headline LLaMA + decode target. +- Causal-mask + f32 accumulator promotion. +- 3-variant comparison plots from the sweep CSVs (wall, gemm util, + hbm bw util, ipcq, kv cache) — generated under `docs/diagrams/`. + +## References + +- ADR-0011 — PhysAddr + VA contract (D-VA1: kernel-computed shard + offset). +- ADR-0060 — GQA fused attention on AHBM (precursor of the unified + mapping). +- ADR-0064 — CPU issue-cost model (composite supertile motivation). +- ADR-0065 — Flat ops, composite, softmax_merge recipe (with the D4 + supplement applied here). +- ADR-0070 supersedes ADR-0060 §B.split.2 short-prefill clause. +- *GQA-mapping-on-AHBM-short-context* note (research design rationale).