adr-0070: gqa short context attention prefill/decode + 3 composite tiers
This commit is contained in:
eusonice
@@ -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).
|
||||||
Reference in New Issue
Block a user