Files
kernbench2/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md
T
ywkang 1742600641 gqa(adr): rewrite GQA fused-attention ADR as ADR-0060, aligned to kernbench
Reframes the proposed GQA FlashAttention design onto kernbench's actual
execution model (greenlet tl API + IPCQ), replacing the composite-centric
mechanism that does not match the simulator:

- Records relationship to existing baseline kernels (_attention_mesh_kv/mlo,
  milestone-gqa-llama70b) and their 3 deliberate limitations.
- Mechanism is greenlet tl (per-op latency; no fusion benefit), not
  composites; running (m,l,O) is Python handles; reduction is tl.send/recv.
- Tree reduction (log N) replaces baseline all-to-all fan-out (N-1).
- Pseudocode rewritten in real tl.* signatures; depends on ADR-0061/62/63.
- Rejects composite-IPCQ-push + composite-carried-state + flash-composite
  with documented efficient alternatives.
- Adds verification plan.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-03 17:42:00 -07:00

27 KiB
Raw Blame History

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 (prerequisites):

  • ADR-0061 tl.broadcast — data-faithful GQA head reuse.
  • ADR-0062 tl.load_async — non-blocking HBM tile load (KV prefetch).
  • ADR-0063 tl.scratch_scope — per-tile scratch recycling.

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).


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 — not composites: 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 running (m, , O) is just Python TensorHandles threaded through the loop. This matters for this ADR's design (see §1).

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). Real GQA (h_q = G·h_kv) is blocked by the MemoryStore byte-conservation check on the symbolic broadcast view → fixed by ADR-0061.
  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) + ADR-0062 (prefetch).

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 CUBEs (memory cubes, each with a logic die containing PEs) + 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/...). 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 does not require composites for the attention inner loop (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 in the greenlet tl programming model, not as composites. Rationale grounded in kernbench's execution + latency model:

  • The simulator charges latency per op on its modelled component (GEMM on PE_GEMM, vector ops on PE_MATH, DMA on PE_DMA — pe_gemm.py, pe_math.py, pe_dma.py). "Fusing" several ops into one CompositeCmd does not reduce modelled latency; the stages still run on the same engines. The win a real fused kernel gets — overlap (prefetch) and small working set (scratch recycling) — is obtained here by ADR-0062 and ADR-0063, which are smaller and reusable.
  • The running (m, , O) flash state is naturally a set of Python TensorHandles threaded through the loop (the baseline already does this). No composite-carried register state is needed.
  • Cross-PE combination uses kernel-level tl.send/tl.recv (the baseline already does this and it is tested). No composite-driven IPCQ push is needed.

So the per-tile inner pipeline is the op sequence

K_j load → Q·Kⱼᵀ → (+ causal mask) → online-softmax update → V_j load → P·Vⱼ → running (m,,O) update

issued as ordinary tl.* ops, with the next tile's K/V issued via tl.load_async (ADR-0062) so its DMA overlaps the current tile's compute, and each tile's temporaries wrapped in tl.scratch_scope (ADR-0063) so scratch stays O(1).

Cross-PE combination (KV-parallel / SP) is a log-sum-exp tree reduction over (m, , O) after P·V, flowed through IPCQ with kernel-level tl.send/tl.recv (§4).

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).

Why this is the efficient choice. It reuses the proven baseline kernels and the proven IPCQ collective; the only genuinely new machinery is three small, general primitives (ADR-0061/0062/0063). The originally-proposed "everything-in-one-composite + composite IPCQ push

  • composite-carried state" would require a bespoke flash-composite command type, a register-lifetime model across composites, and an IPCQ-push epilogue — large, special-purpose, and with no latency benefit over the greenlet path. Those are rejected; 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. Real tl names (tl_context.py), with tl.load_async (ADR-0062), tl.broadcast (ADR-0061), tl.scratch_scope (ADR-0063):

# running state (persistent arena — allocated once, outside the scope)
# m: [G, T_q]  l: [G, T_q]  O: [G, T_q, d]

with tl.scratch_scope():                              # per-tile temporaries recycled
    Kj = tl.wait(f_k[j])                              # prefetched (ADR-0062)
    Sj = tl.dot(q_g, tl.trans(Kj)) * softmax_scale   # [G, TILE]; q_g is GQA-batched
    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)
    Vj    = tl.wait(f_v[j])
    O     = O * corr + tl.dot(P, Vj)                  # P·V folded into running O
    m     = m_new
    if j + PREFETCH < n_tiles:                        # keep the pipeline full
        f_k[j + PREFETCH] = tl.load_async(K_base + (j+PREFETCH)*TILE*d, (TILE, d))
        f_v[j + PREFETCH] = tl.load_async(V_base + (j+PREFETCH)*TILE*d, (TILE, d))

Notes:

  • q_g is the GQA-batched query: [G, d] (decode) or [G, T_q, d] (prefill). One K/V tile load serves all G rows — the GQA reuse lever. Where the matmul needs the KV head expanded to G, use tl.broadcast (ADR-0061), never the symbolic _view.
  • The V load is issued (prefetched) before it is needed so its DMA overlaps the Q·Kᵀ + softmax of the same/earlier tile.
  • 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).
  • No CompositeCmd is used. (A future tl.composite flash kind is an optional optimisation — §8 item 4 — not a requirement.)

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):

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 N1 (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.sends (the baseline sends the triplet as three sends — same pattern).

Kernel structure (greenlet):

# 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

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,d] decode / [G,T_q,d] prefill; tl.broadcast for GQA
    m, l, O = init_running()                 # persistent arena: -inf, 0, zeros
    prime_prefetch(k_ptr, v_ptr, n_tiles)    # tl.load_async first PREFETCH tiles
    for j in range(n_tiles):
        if tile_all_future(j, q_block):      # causal skip (kernel if)
            continue
        run_tile(j, mask_or_null(j, q_block))   # §3 op sequence, 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 batched (q_g = [G, d]), so each K/V tile is loaded once and reused by 8 rows (Q·Kᵀ is an 8×TILE batched GEMV, P·V an 8×d). tl.broadcast expands the KV head where needed.
  • If S_pe fits in scratch (small/medium context) this degenerates to a one-shot partial attention (one tl.dot for Q·Kᵀ, one softmax, one tl.dot for P·V) — exactly the baseline _attention_mesh_mlo _partial_attention, just GQA-batched. 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 ≤ qstile_all_pastmask=None, full compute.
    • ks ≥ qetile_all_futureskip (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)
K/V load, Q·Kᵀ, mask add, softmax math, P·V tl ops on PE engines

The kernel decides; the tl ops execute already-decided work. This is exactly the greenlet model kernbench already supports — no new control abstraction.


8. Required kernbench changes

Genuinely new (each has a supporting ADR):

  1. GQA head broadcast in data modeADR-0061 (tl.broadcast). The blocker for h_q > h_kv. Without it the GQA reuse lever cannot run under enable_data=True.
  2. Async HBM tile load / KV prefetchADR-0062 (tl.load_async
    • tl.wait). The KV-load-bound overlap lever for decode/long-context.
  3. Per-tile scratch recyclingADR-0063 (tl.scratch_scope). Removes the S=16 ceiling so realistic context lengths run.

Algorithm work in the kernel (no new primitive; existing tl API):

  • GQA Q-axis batching (G rows share each K/V load) — tl.broadcast + 2-D tl.dot.
  • 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").

Explicitly REJECTED (efficient alternative chosen):

  1. Composite that chains DMA→MM→VEC→DMA→MM→VEC with carried (m,,O) register state and a tail IPCQ push. Replaced by the greenlet path: per-op latency is identical; overlap comes from ADR-0062; small working set from ADR-0063; running state is Python handles; the IPCQ push is tl.send. Building a bespoke flash-composite command type + cross-composite register lifetime + an IPCQ-push epilogue is large, special-purpose, and yields no latency benefit. A tiled tl.composite "flash" kind remains a possible future optimisation (it would fold prefetch+recycling into the scheduler), but it is not required for an efficient kernel and is out of scope here.
  2. Hardware pop-as-dependency. Out of scope (§6).
  3. RoPE / QKV projection / KV-cache write inside this kernel. Upstream qkv_rope (P1P5). 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. KV tile prefetch depth (ADR-0062) — dominant for KV-load-bound; sweep 2→4.
  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 (K_tile + S/P + V_tile + O_acc + G-way GQA) against DMA efficiency; interacts with ADR-0063.
  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 tl.dot.

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).

  1. Correctness vs reference (Phase 2, enable_data=True): for a small (G, T_q, S, N), kernel O matches a numpy FlashAttention reference within fp tolerance — for all four cases. This is the test that requires ADR-0061 (GQA) to even run with h_q>h_kv.
  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 N1); 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. Prefetch overlap (ADR-0062): end-to-end latency of the tiled sweep is below the serial Σ(load+compute) (overlap is real).
  7. Determinism: identical inputs → identical op_log + latency (SPEC §0.1).