ywkang
7a9d4ec47b
- ADR-0060: GEMMs (Q.Kt, P.V) via existing tl.composite (scheduler-managed tiling + K/V DMA streaming); softmax merge + IPCQ tree reduction stay in kernel. Front TL;DR pseudocode of the final composite kernel; new section B lists open design items (DDD sync, K pre-transpose, dma_read lever, kernel-vs-scheduler tiling, ring path). - ADR-0062: redefined from a new load_async op to global lazy tl.load (non-blocking + auto-wait on first use; API unchanged; goldens regenerate). - ADR-0064 (new): per-op-type CPU issue cost model (composite ~40ns >> primitive) so the hybrid's CPU-saturation win becomes measurable (currently dispatch_cycles=0 hides it). Cost-model impl deferred. - KO mirrors for ADR-0060/0062/0064 (-ko suffix, adr-proposed). Rationale: non-blocking CompositeCmd offloads tiling to PE_SCHEDULER, decoupling CPU issue-rate from execution so the CPU can saturate the engines; the prior 'composite = no latency benefit' claim was an artifact of dispatch_cycles=0. Docs only; no production code changed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
818 lines
42 KiB
Markdown
818 lines
42 KiB
Markdown
# 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** (efficiency / scale enablers — *not* GQA blockers;
|
||
see §8 correction): **ADR-0063** `tl.scratch_scope` (per-tile scratch
|
||
recycling — required for realistic context length), **ADR-0062** lazy
|
||
`tl.load` (non-blocking load with auto-wait on first use → load/compute
|
||
overlap), **ADR-0061** `tl.broadcast` (optional mask/general
|
||
convenience). The two GEMMs (Q·Kᵀ, P·V) are issued as scheduler-managed
|
||
`tl.composite` commands (the existing `CompositeCmd`; no new command
|
||
kind); real GQA itself needs only kernel restructuring (§5.2).
|
||
|
||
**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).
|
||
|
||
---
|
||
|
||
## TL;DR — final GQA kernel (composite hybrid) pseudocode
|
||
|
||
The decision (§1) in one place: **GEMMs → `tl.composite`; softmax merge +
|
||
tree reduction → kernel; `tl.load` is lazy.** Per-KV-head, `G` folded into
|
||
the matmul M dim (real GQA, no broadcast). This is the reference shape; the
|
||
prose sections elaborate each piece.
|
||
|
||
```python
|
||
def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
|
||
counter, start_pe, N, q_block, scale, *, tl):
|
||
# ---- geometry (kernel arithmetic, §2) ----
|
||
pe_id = tl.program_id(axis=rank_axis)
|
||
G = H_q // H_kv # query heads per KV head (=8)
|
||
causal = q_block.is_prefill
|
||
|
||
for kv in range(H_kv): # one KV head per iteration (§5.2)
|
||
# Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
|
||
# issue now, auto-wait at first use inside the first composite.
|
||
q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d]
|
||
my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv
|
||
n_tiles = ceil(my_len / TILE)
|
||
|
||
# persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
|
||
m, l, O = init_running(G * q_block.T_q, d)
|
||
|
||
for j in range(n_tiles):
|
||
if causal and tile_all_future(j, q_block):
|
||
continue # causal skip (kernel if)
|
||
with tl.scratch_scope(): # per-tile temporaries (ADR-0063)
|
||
# --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite ---
|
||
Sj = tl.composite("gemm", a=q_g,
|
||
b=tl.ref(k_tile(kv, j), (d, TILE))) * scale # Kᵀ pre-stored [d,TILE]
|
||
if causal and tile_partial(j, q_block):
|
||
Sj = Sj + causal_mask(j, q_block) # additive boundary mask
|
||
# --- online softmax merge in the kernel (MATH ops) ---
|
||
m_new = tl.maximum(m, tl.max(Sj, axis=-1))
|
||
P = tl.exp(Sj - m_new)
|
||
corr = tl.exp(m - m_new)
|
||
l = l * corr + tl.sum(P, axis=-1)
|
||
# --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite ---
|
||
Oj = tl.composite("gemm", a=P,
|
||
b=tl.ref(v_tile(kv, j), (TILE, d)))
|
||
O = O * corr + Oj # running merge
|
||
m = m_new
|
||
|
||
# ---- cross-PE combine (§4): log-sum-exp tree to root, or store ----
|
||
if N == 1:
|
||
tl.store(o_base(kv), O / l)
|
||
else:
|
||
tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv
|
||
```
|
||
|
||
> **Why this shape:** the two `tl.composite` calls offload tiling + K/V DMA
|
||
> streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse
|
||
> descriptors and runs ahead → engines stay saturated, §1); the softmax
|
||
> merge and the IPCQ tree reduction stay in the kernel because the existing
|
||
> `CompositeCmd` cannot carry cross-tile state. `K` is pre-stored
|
||
> transposed `[d, TILE]` to sidestep the reshape-not-transpose caveat
|
||
> (§3, §B). A bespoke "flash-composite" command kind is **not** introduced
|
||
> (§8 item 4).
|
||
|
||
---
|
||
|
||
## 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:** `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 GEMMs as
|
||
**blocking `tl.dot`**, not composites. The running `(m, ℓ, O)` is just
|
||
Python `TensorHandle`s threaded through the loop. This ADR keeps the
|
||
running state and the softmax merge in the kernel but **moves the two
|
||
GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this
|
||
**matters for the design**.
|
||
|
||
**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`). The test attributes this to
|
||
a MemoryStore byte-conservation failure on a *broadcast view*, but that
|
||
failure is a property of the baseline's **head-packing hack**
|
||
(`_view(K, (h_q·d, S_kv))`, which conflates all heads into one matmul
|
||
dim and only conserves bytes when `h_q == h_kv`). The correct fix is
|
||
**kernel restructuring**, not a broadcast op: process **one KV head at
|
||
a time** and fold the `G` group rows into the matmul **M** dimension
|
||
(§5.2). That uses only byte-conserving reshapes, so real GQA
|
||
(`h_q = G·h_kv`) runs with **no new primitive** — see §8.
|
||
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) + composite K/V streaming (§3) + **ADR-0062** (lazy load
|
||
overlap).
|
||
|
||
**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/...`), issued **non-blocking** to PE_SCHEDULER,
|
||
which generates a tile plan and streams DMA→GEMM→write per tile
|
||
(ADR-0014 D6; `pe_scheduler.py:104-143`). 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 issues
|
||
**each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and
|
||
keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel —
|
||
it does **not** need a new "flash-composite" command kind (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 as a *hybrid*: issue the two GEMMs (Q·Kᵀ and P·V)
|
||
as scheduler-managed `tl.composite(op="gemm")` commands, and keep the
|
||
online-softmax merge and the cross-PE reduction as kernel-level `tl`
|
||
ops.** `tl.load` is **lazy** (non-blocking; the wait is auto-inserted at
|
||
first use of the loaded data — ADR-0062), so explicit HBM loads overlap
|
||
the compute that follows. Rationale grounded in kernbench's execution +
|
||
latency model:
|
||
|
||
- **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is
|
||
non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`):
|
||
the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole
|
||
per-PE tile sweep) and the scheduler generates the tile plan and streams
|
||
DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the
|
||
scheduler streams from HBM, so per-tile **K/V prefetch is the
|
||
scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed
|
||
to issue the **next** composite while the current one runs, so the
|
||
scheduler keeps the GEMM engine saturated across tiles.
|
||
- **This reflects the hardware** and decouples CPU issue-rate from
|
||
execution-rate. The blocking per-op `tl.dot` path, by contrast, stalls
|
||
the CPU on every GEMM and leaves GEMM-engine **bubbles** during the
|
||
interleaved softmax MATH ops; it is realistic only if the CPU can keep
|
||
up with fine-grained per-tile issue.
|
||
- **The running `(m, ℓ, O)` flash state stays Python `TensorHandle`s**
|
||
threaded through the loop (the baseline already does this); the softmax
|
||
merge (max/exp/sum/rescale) is kernel-level `tl` MATH **between** the two
|
||
GEMM composites. The existing `CompositeCmd` cannot chain two GEMMs or
|
||
carry cross-tile register state (§0), so the merge necessarily lives in
|
||
the kernel — this is the hybrid split, not a limitation worked around.
|
||
- **Cross-PE combination** is a log-sum-exp **tree** over `(m, ℓ, O)`
|
||
after P·V, via kernel-level `tl.send`/`tl.recv` (§4) — unchanged.
|
||
|
||
So the per-tile inner pipeline is:
|
||
|
||
```
|
||
q_g = tl.load(Q group) # lazy; auto-wait at first use
|
||
per tile j:
|
||
Sⱼ = tl.composite("gemm", a=q_g, b=tl.ref(Kⱼ)) → Sⱼ # scheduler streams Kⱼ DMA + GEMM
|
||
Sⱼ += maskⱼ # kernel MATH, boundary tile only
|
||
online-softmax: mⱼ, m_new, P, corr, ℓ # kernel MATH
|
||
Oⱼ = tl.composite("gemm", a=P, b=tl.ref(Vⱼ)) → Oⱼ # scheduler streams Vⱼ DMA + GEMM
|
||
O = O*corr + Oⱼ; m = m_new # kernel MATH (running merge)
|
||
```
|
||
|
||
with each tile's MATH temporaries wrapped in `tl.scratch_scope`
|
||
(ADR-0063) so scratch stays O(1), and the next tile's composites issued
|
||
before the current tile's results are waited on (non-blocking handles) so
|
||
the scheduler pipelines across tiles.
|
||
|
||
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).
|
||
|
||
> **What this supersedes.** An earlier iteration proposed a pure greenlet
|
||
> primitive path (all `tl.dot`, no composite) on the argument that
|
||
> "composite yields no latency benefit." That holds **only because** the
|
||
> simulator currently charges **zero** per-op CPU issue cost
|
||
> (`dispatch_cycles=0`, `pe_cpu.py`) — it models away exactly the CPU
|
||
> issue-rate / DMA-program cost that descriptor offload exists to hide.
|
||
> The hybrid is the faithful representation of an efficient kernel. The
|
||
> **measurable** size of the win (can the CPU saturate the engines for
|
||
> many tiles?) is gated on modelling an op-type-differentiated issue cost,
|
||
> tracked as future work (cost model; §9). Even at `dispatch_cycles=0` the
|
||
> non-blocking composite path fills the GEMM-engine bubbles the blocking
|
||
> `tl.dot` path leaves.
|
||
>
|
||
> **Why this is the efficient choice.** GEMM tiling + DMA streaming +
|
||
> cross-tile pipelining are offloaded to the proven `CompositeCmd`
|
||
> scheduler path; the softmax merge and the proven IPCQ collective stay in
|
||
> the kernel. The only genuinely new machinery is two small, general
|
||
> primitives (ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); the
|
||
> reduction reuses `tl.send`/`tl.recv`. A bespoke "flash-composite"
|
||
> command (one kind internalising the softmax merge + carried register
|
||
> state + an IPCQ-push epilogue) is **not** built — large, special-purpose,
|
||
> and its only delta over this hybrid (full softmax offload) is not
|
||
> justified at the current modelling fidelity; 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. The two GEMMs are
|
||
`tl.composite(op="gemm")` (scheduler-managed tiling + K/V DMA streaming);
|
||
the softmax merge is kernel `tl` MATH between them. Real `tl` names
|
||
(`tl_context.py`), with lazy `tl.load` (ADR-0062), `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]
|
||
q_g = tl.load(Q_group_ptr, (G*T_q, d)) # lazy; auto-wait at first use (ADR-0062)
|
||
|
||
with tl.scratch_scope(): # per-tile MATH temporaries recycled
|
||
Sj = tl.composite("gemm", a=q_g, # [G·T_q, TILE]; scheduler streams Kⱼ DMA
|
||
b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale
|
||
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)
|
||
Oj = tl.composite("gemm", a=P, # [G·T_q, d]; scheduler streams Vⱼ DMA
|
||
b=tl.ref(V_base + j*TILE*d, (TILE, d)))
|
||
O = O * corr + Oj # running merge (kernel MATH)
|
||
m = m_new
|
||
```
|
||
|
||
Notes:
|
||
- `q_g` is the GQA-batched query reshaped to `[G·T_q, d]` (the `G` group
|
||
rows folded into the matmul M dim; byte-conserving). One K/V tile serves
|
||
all `G·T_q` rows — the GQA reuse lever — with no broadcast.
|
||
- **K/V are `tl.ref` operands** the composite scheduler streams from HBM
|
||
per tile (`pe_scheduler.py:104-143`): that *is* the prefetch/pipeline,
|
||
so there is no explicit prefetch op. Issuing tile `j+1`'s composites
|
||
before waiting on tile `j` (non-blocking handles) keeps the scheduler
|
||
pipelined across tiles and the GEMM engine saturated.
|
||
- `tl.trans` is **metadata-only** in kernbench (`tl_context.py:390`) and
|
||
`MemoryStore.read` *reshapes* rather than transposes
|
||
(`memory_store.py:73`). For zero/structural runs this is harmless; for
|
||
non-trivial numeric data it yields a reshape-not-transpose, so Q·Kᵀ via
|
||
a transposed K needs care (§11) — store K pre-transposed `[d, TILE]`, or
|
||
add a real `tl.transpose` (a candidate further primitive, likely
|
||
unnecessary given the simulator's performance-modeling purpose).
|
||
- 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).
|
||
- A **new** "flash-composite" command kind (one that internalises the
|
||
softmax merge + carried `(m,ℓ,O)`) is **not** used; the existing
|
||
`CompositeCmd` covers each GEMM and the merge stays in the kernel
|
||
(§1, §8 item 4).
|
||
|
||
---
|
||
|
||
## 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·T_q, d]; lazy tl.load, G folded into M (no broadcast)
|
||
m, l, O = init_running() # persistent arena: -inf, 0, zeros
|
||
for j in range(n_tiles): # K/V streamed per tile by the composite scheduler (§3)
|
||
if tile_all_future(j, q_block): # causal skip (kernel if)
|
||
continue
|
||
run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, 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 folded into the matmul **M** dimension
|
||
(`q_g` reshaped `[G, T_q, d] → [G·T_q, d]`, a byte-conserving reshape).
|
||
Then `Q·Kᵀ` is `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]` and
|
||
`P·V` is `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. The KV tile
|
||
(`[TILE, d]`) is the shared `K`/`V` operand — **streamed once, reused by
|
||
all `G·T_q` rows automatically** because they are the M rows of the
|
||
GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile
|
||
plan also makes the timing count all `G` rows' work correctly (a leading
|
||
batch axis would *not* be counted — see §8).
|
||
- If `S_pe` fits in scratch (small/medium context) this degenerates to a
|
||
**one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one
|
||
composite for P·V) — exactly the baseline `_attention_mesh_mlo`
|
||
`_partial_attention`, just GQA-batched and on the composite path. 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) |
|
||
| Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER |
|
||
| Q load, mask add, softmax math, running `(m,ℓ,O)` merge | **`tl` ops** on PE engines (kernel-issued) |
|
||
|
||
The kernel decides; the GEMMs are offloaded to the scheduler as
|
||
composites; the remaining `tl` ops execute already-decided work. This is
|
||
exactly the greenlet + composite model kernbench already supports — no new
|
||
control abstraction.
|
||
|
||
---
|
||
|
||
## 8. Required kernbench changes
|
||
|
||
**Correction from design iteration:** real GQA (`h_q > h_kv`) needs **no
|
||
new primitive** — only the kernel restructuring in §5.2 (per KV head,
|
||
`G` folded into M, byte-conserving reshapes). The supporting ADRs are
|
||
*efficiency / scale* enablers, not GQA blockers.
|
||
|
||
**Algorithm work in the kernel (no new primitive; existing `tl` API):**
|
||
|
||
- **GQA Q-axis batching** (the reuse lever) — fold `G·T_q` into the matmul
|
||
M dim per KV head (§5.2); `_view`-style byte-conserving reshape; the
|
||
GEMM is a `tl.composite(op="gemm")` with M = `G·T_q`. Runs today in both
|
||
timing and data mode.
|
||
- **GEMMs via composite** (§1/§3) — Q·Kᵀ and P·V each issued as a
|
||
non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and
|
||
streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new
|
||
command kind).
|
||
- 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")`.
|
||
|
||
**New primitives for efficiency / scale (each has a supporting ADR):**
|
||
|
||
1. **Per-tile scratch recycling** — **ADR-0063** (`tl.scratch_scope`).
|
||
*Required for scale*: removes the `S=16` ceiling (1 MiB bump
|
||
allocator) so realistic context lengths run. Highest-value of the
|
||
three.
|
||
2. **Lazy `tl.load`** — **ADR-0062** (non-blocking load + auto-wait on
|
||
first use; API surface unchanged). *Efficiency*: overlaps explicit
|
||
loads (the Q group, non-composite kernels) with following compute. The
|
||
per-tile **K/V** prefetch is handled by the composite scheduler (§1),
|
||
so this covers the remaining explicit loads. Global semantics change →
|
||
existing goldens regenerate (ADR-0062 D3).
|
||
3. **GQA head / mask broadcast** — **ADR-0061** (`tl.broadcast`).
|
||
*Optional convenience*, not a GQA blocker (see correction above).
|
||
Useful for additive-mask construction across the `G·T_q` rows and for
|
||
general kernels; `np.matmul` already broadcasts in data mode, so it is
|
||
not needed for correctness. Lowest priority.
|
||
|
||
**Explicitly REJECTED (efficient alternative chosen):**
|
||
|
||
4. ~~A bespoke "flash-composite" command kind that internalises the whole
|
||
inner loop — DMA→MM→VEC→DMA→MM→VEC with carried `(m,ℓ,O)` register
|
||
state and a tail IPCQ push.~~ The two GEMMs **do** use the existing
|
||
`CompositeCmd` (§1/§3) — that gives scheduler-managed tiling, K/V DMA
|
||
streaming, and cross-tile pipelining. What is rejected is a **new**
|
||
command kind that also absorbs the softmax merge + cross-tile register
|
||
lifetime + an IPCQ-push epilogue: it is large and special-purpose, and
|
||
its only delta over the hybrid (full softmax offload) is not justified
|
||
at the current modelling fidelity (per-op CPU issue cost = 0; see §1).
|
||
Revisit if the cost model (§9) makes full offload measurably worthwhile.
|
||
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. **Composite tile-pipeline depth** — dominant for KV-load-bound; how far
|
||
ahead the kernel issues non-blocking composites before waiting, and the
|
||
scheduler's per-tile streaming depth.
|
||
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 (`S/P tiles + O_acc + G-way
|
||
GQA`) against DMA efficiency; interacts with ADR-0063 and the
|
||
scheduler's `TILE_M/K/N` (`pe_scheduler.py`).
|
||
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 composite.
|
||
6. **Per-op CPU issue cost (cost model)** — currently `dispatch_cycles=0`
|
||
(`pe_cpu.py`), so composite-vs-primitive issue overhead is invisible.
|
||
An op-type-differentiated issue cost (a `tl.composite` descriptor push
|
||
≫ a primitive op) is what makes the hybrid's CPU-saturation win
|
||
**measurable** (§1). Specified in **ADR-0064**; tracked as separate
|
||
future work.
|
||
|
||
---
|
||
|
||
## 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).
|
||
|
||
The simulator's contract is **latency by traversal + determinism +
|
||
structural correctness** (SPEC §0, §0.1), not bit-exact numerics — Phase 2
|
||
data exists mainly to exercise the data path, `tl.trans` is
|
||
reshape-not-transpose, and `bf16` is modelled as `f16`. Verification is
|
||
therefore **structural/timing-first**, with numeric parity as a bounded
|
||
secondary check.
|
||
|
||
1. **Runs in data mode (`enable_data=True`):** the GQA kernel
|
||
(`h_q = G·h_kv`) completes without the byte-conservation error that the
|
||
baseline head-packing hits — for all four cases. (This needs the §5.2
|
||
restructuring, *not* a new primitive.)
|
||
**Numeric parity (secondary):** for symmetric/identity inputs where
|
||
reshape-as-transpose is exact, kernel `O` matches a numpy
|
||
FlashAttention reference within fp tolerance. Full asymmetric parity is
|
||
gated on a real `tl.transpose` (out of scope; flagged).
|
||
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. **Load/compute overlap:** end-to-end latency of the tiled sweep is
|
||
below the serial `Σ(load+compute)` — from the composite scheduler
|
||
streaming K/V per tile (§1/§3) and lazy `tl.load` (ADR-0062) overlapping
|
||
the Q load. (Overlap is real modelled concurrency, not a subtraction.)
|
||
7. **Composite GEMM offload (structural):** each tile's Q·Kᵀ and P·V emit a
|
||
`CompositeCmd` (non-blocking) to PE_SCHEDULER, not a blocking `tl.dot`;
|
||
op_log shows the composite tile plan and the kernel issues the next
|
||
tile's composites before waiting (cross-tile pipelining).
|
||
8. **Determinism:** identical inputs → identical op_log + latency
|
||
(SPEC §0.1).
|
||
|
||
---
|
||
|
||
## B. Open design items from the hybrid pivot (review later)
|
||
|
||
These arose when the decision moved from a pure greenlet primitive path to
|
||
the **composite hybrid + lazy `tl.load`** (this revision). None blocks the
|
||
design; each needs a verification pass during implementation. Recorded here
|
||
(rather than asked) per the working agreement — the recommendation is my
|
||
predicted default; revise on review.
|
||
|
||
1. **DDD-0060 is not yet synced.** The Detailed Design Document still
|
||
describes the old `tl.load_async` double-buffer path and primitive
|
||
`tl.dot` inner loop (its §4.3/§5/§10). It must be updated to the hybrid
|
||
(composite GEMMs, lazy load, K pre-transposed). *Left for review*
|
||
because the DDD is a derived how-to and a large rewrite; the ADR is now
|
||
the authoritative record. **Recommend:** sync DDD as a follow-up before
|
||
implementation starts.
|
||
|
||
2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b =
|
||
[d, TILE]`, but the KV cache stores K as `[S_pe, d]`. `tl.trans` is
|
||
metadata-only and `MemoryStore.read` reshapes, not transposes
|
||
(`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial
|
||
data. **Recommend:** store K **pre-transposed** `[d, S_pe]` in the cache
|
||
(the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))`
|
||
a contiguous slice. Verify the upstream `qkv_rope` write layout supports
|
||
this, or add a real `tl.transpose` (heavier; deferred).
|
||
|
||
3. **Composite output buffer vs `tl.scratch_scope`.** Each Q·Kᵀ composite
|
||
writes `Sj` to an `out_addr`; the kernel then reads it for the softmax
|
||
MATH. That output buffer, and the in-flight composites' targets, must
|
||
live where the per-tile `scratch_scope` (ADR-0063) will **not** recycle
|
||
them before they are consumed — same discipline as in-flight lazy loads
|
||
(ADR-0062 D-Negative). **Recommend:** composite outputs for the *current*
|
||
tile live in the scoped arena (consumed same iteration); the persistent
|
||
`(m,ℓ,O)` stays outside. Verify no use-after-recycle when the next
|
||
tile's composites are issued early (cross-tile pipelining).
|
||
|
||
4. **GQA `dma_read_count` lever under composite streaming.** The lever
|
||
(§11.2: K/V `dma_read_count` independent of `G`) assumes the composite
|
||
emits **one** K/V tile DMA reused across all `G·T_q` M-rows. The
|
||
scheduler's `generate_gemm_plan` tiles by `TILE_M/K/N`
|
||
(`pe_scheduler.py:35-37`, 32/64/32) — confirm the M-tiling over `G·T_q`
|
||
does **not** re-issue the shared K/V tile DMA per M-tile (i.e. operand
|
||
DMA is shared across M-tiles, or the lever weakens). **Recommend:**
|
||
assert it in the levers test; if violated, the GQA win is in compute
|
||
only, not DMA — still correct, but the headline changes.
|
||
|
||
5. **Kernel TILE vs scheduler `TILE_M/K/N`.** The kernel reasons about a
|
||
logical KV `TILE`; the scheduler re-tiles internally at fixed
|
||
`TILE_M/K/N`. Two tiling layers interact (scratch residency, pipeline
|
||
depth). **Recommend:** treat the kernel TILE as the K/V streaming
|
||
granularity and let the scheduler sub-tile the GEMM; document the
|
||
relationship in the DDD and sweep both (§9 items 1, 3).
|
||
|
||
6. **Cost model is a separate ADR.** The hybrid's CPU-saturation benefit is
|
||
invisible while `dispatch_cycles=0`. The per-op-type issue-cost model is
|
||
specified in **ADR-0064**; this ADR's §1/§9 depend on it for the
|
||
*measurable* (not just structural) win. **Recommend:** land ADR-0064's
|
||
model before claiming hybrid latency wins in the eval.
|
||
|
||
7. **Ring path (§5.5) GEMMs.** §5.5 still describes the ring fold with
|
||
primitive ops + `recv_async`. For consistency the ring's per-step Q·Kᵀ /
|
||
P·V should also be composites; the IPCQ `recv_async` overlap is
|
||
orthogonal and stays. **Recommend:** apply the same hybrid shape to the
|
||
ring step during implementation; low risk, mirrors §3.
|