mukesh
d282144339
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.
ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
(m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.
ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
owns whole KV heads; PE-parallel heads with intra-group chain
reduce. Prefill has no Ring KV (each head fully resident).
ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).
ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.
ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
consults table with dispatch_cycles fallback (ADR-0046 §D6
back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
measurable end-to-end.
P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).
Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
_attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
(EN + KO kept in sync).
Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
202 lines
8.1 KiB
Markdown
202 lines
8.1 KiB
Markdown
# ADR-0061: `tl.broadcast` / `tl.repeat` — data-faithful GQA head reuse
|
||
|
||
## Status
|
||
|
||
Proposed
|
||
|
||
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention).
|
||
>
|
||
> **⚠ Re-evaluated during design iteration — this is now OPTIONAL, not a
|
||
> GQA blocker.** The original framing (below) held that GQA `h_q > h_kv`
|
||
> is blocked by the MemoryStore byte-conservation check on a broadcast
|
||
> view. That is true *only of the baseline's head-packing hack*
|
||
> (`_view(K, (h_q·d, S_kv))`). The correct GQA design (ADR-0060 §5.2)
|
||
> processes **one KV head at a time** and folds the `G` group rows into
|
||
> the matmul **M** dimension — using only byte-conserving reshapes — so it
|
||
> needs **no broadcast op at all**, and `np.matmul` already broadcasts a
|
||
> shared operand in Phase-2 data mode. This ADR therefore drops to a
|
||
> *convenience* primitive: it cleans up additive-mask construction across
|
||
> the `G·T_q` rows and serves general kernels. **Lowest priority of the
|
||
> three supporting ADRs.** The rest of this document is retained as the
|
||
> design for that convenience op.
|
||
|
||
## 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
|
||
|
||
Pre-ADR-0060 mesh kernels reshaped tensors with a metadata-only `_view`
|
||
helper that rewrote `shape` but kept the original `nbytes`:
|
||
|
||
```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 pre-ADR-0060 baseline documented this as a deliberate limitation:
|
||
|
||
> "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-headline` 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.
|