Verified against sim_engine data path (memory_store, data_executor): - GQA reuse does NOT need a broadcast op. The baseline's h_q==h_kv limit is its head-packing reshape hack, not a missing primitive. Correct fix: per-KV-head loop with G folded into matmul M dim (byte-conserving reshape) — runs today, timing correct (m=G*T_q), data mode runs. - ADR-0061 broadcast demoted from 'the blocker' to optional convenience. - Surfaced tl.trans = reshape-not-transpose (memory_store reshapes; data_executor np.matmul on reshaped operands) -> numeric parity is bounded; verification is structural/timing/determinism-first (matches SPEC perf-model purpose). Optional tl.transpose deferred. - Reordered DDD phase plan (P1 GQA needs no new feature; P3 scratch_scope is the key scale feature); added open decisions 10.10 (transpose) and 10.11 (GQA-via-M-fold finding). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
25 KiB
Detailed Design Document — AHBM GQA Fused Attention (ADR-0060)
Status: Draft (companion to ADR-0060, Proposed) Scope: implementation-ready design for an efficient GQA FlashAttention kernel on AHBM in the kernbench simulator. Audience: reviewer resuming with "GQA 검토 시작하자". Read ADR-0060 first (decision record), then this (the how), then the Open Decisions section (§10) for the choices made autonomously that need your sign-off.
How this document was produced. It is the result of several design↔evaluation iterations against the real kernbench source (
src/kernbench/triton_emu/tl_context.py,sim_engine/*,components/builtin/pe_*, the_attention_mesh_*kernels, and the GQA tests). Every claim about what exists is grounded in a file path; every claim about what is new is tied to a supporting ADR (0061/0062/0063).
1. Goal and success criteria
Goal: an efficient GQA fused attention kernel that runs on AHBM in
kernbench, across the four operating points decode/prefill ×
single-/multi-user, at realistic (not just validation) scale, in both
timing (enable_data=False) and data (enable_data=True) modes.
Success criteria:
- Correct: kernel output matches a numpy FlashAttention reference within
fp tolerance, with real GQA (
H_q=64, H_kv=8, G=8). - Efficient: the three levers actually fire and are observable in op_log:
GQA K/V-load amortisation, KV prefetch overlap,
⌈log₂ N⌉reduction. - Scales: context length not capped by the 1 MiB scratch bump allocator.
- Deterministic and SPEC-compliant (all latency from modelled events).
Non-goals: RoPE / QKV projection / KV-cache writes (upstream
qkv_rope); output projection (downstream); cross-CUBE query-head split;
cycle-accurate microarchitecture (SPEC §5).
2. Where the design sits in the current code
runtime_api (host)
RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy, kernel launch)
│ KernelLaunchMsg
▼
sim_engine
GraphEngine(enable_data=…) engine.py
DataExecutor (Phase 2) data_executor.py ← ADR-0061 adds _execute_broadcast
MemoryStore (nbytes check) memory_store.py ← the GQA blocker lives here
│
▼
components / PE pipeline (per PE)
PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,FETCH_STORE,TCM,IPCQ}
greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py
│
▼
tl programming model tl_context.py
load/store/dot/softmax/max/sum/exp/maximum/send/recv/trans/…
← ADR-0061 adds broadcast/repeat
← ADR-0062 adds load_async (+ wait arm)
← ADR-0063 adds scratch_scope
The attention kernels are bench-side Python (src/kernbench/benches/)
invoked via ctx.launch(name, kernel_fn, *args); they run as greenlet
kernels that emit tl.* ops. No component (device-side) code changes are
needed for the algorithm — only the three tl/sim_engine primitives.
3. File plan
3.1 New / modified production files
| File | Change | ADR |
|---|---|---|
src/kernbench/triton_emu/tl_context.py |
add broadcast/repeat; load_async + wait arm; scratch_scope ctx-mgr |
0061/0062/0063 |
src/kernbench/common/pe_commands.py |
add BroadcastCmd; DmaReadCmd.blocking flag (or DmaReadAsyncCmd) |
0061/0062 |
src/kernbench/sim_engine/data_executor.py |
add _execute_broadcast; handle async dma_read resolution |
0061/0062 |
src/kernbench/components/builtin/pe_dma.py |
non-blocking read path returns a future-resolving event | 0062 |
src/kernbench/sim_engine/event_log.py / op_log |
recognise broadcast op_name (math) |
0061 |
3.2 New kernel + bench files
| File | Role |
|---|---|
src/kernbench/benches/_gqa_attention.py |
the unified GQA kernel (4 cases via flags), with a shared _flash_tile_sweep helper and _tree_reduce helper |
src/kernbench/benches/_gqa_helpers.py |
valid_len, tile_* predicates, mask builders, tree topology (children/parent dirs), GQA-batched Q load |
extend src/kernbench/benches/milestone_gqa_llama70b.py |
add headline-scale panels with real G, prefetch, tiling |
3.3 New tests
| File | Covers |
|---|---|
tests/attention/test_gqa_correctness.py |
numpy reference parity, all 4 cases, real GQA |
tests/attention/test_gqa_levers.py |
GQA load amortisation, tree step count, causal skip, prefetch overlap |
tests/test_tl_broadcast.py |
ADR-0061 unit |
tests/test_tl_load_async.py |
ADR-0062 unit |
tests/test_tl_scratch_scope.py |
ADR-0063 unit |
4. Data model and launch contract
4.1 Tensor placement (DPPolicy)
Using DPPolicy(cube, pe, num_cubes, num_pes) (policy/placement/dp.py):
| Tensor | single-user (PE ring in 1 cube) | multi-user (cube ring) |
|---|---|---|
Q (group, replicated) |
cube=replicate, pe=replicate, num_pes=N |
cube=replicate, pe=replicate, num_cubes=N |
K, V (sequence-sharded) |
cube=replicate, pe=row_wise, num_pes=N |
cube=row_wise, pe=replicate, num_cubes=N |
O (group, replicated) |
as Q |
as Q |
This matches the baseline _PANEL_DISPATCH
(milestone_gqa_llama70b.py:64-81). row_wise over the sequence (M) axis
realises the round-robin shard (§2.1 of ADR-0060); the contiguous per-PE
buffer is the row-wise slice.
rank_axis selects which tl.program_id carries the ring rank: 0 for
single-user (PE id in cube), 1 for multi-user (cube level; only PE 0 of
each cube participates) — kept from the baseline so the SFR installs
(configure_sfr_intracube_pe_ring / configure_sfr_intercube_multisip)
apply unchanged.
4.2 Launch signature
ctx.launch(
f"{panel}_gqa", gqa_attention_kernel,
q, k, v, o, # Tensors (TensorArg)
S_q, S_kv_per_rank, H_q, H_kv, D_HEAD, n_ranks, # scalars
rank_axis, mode, counter, start_pe, # scalars (mode: 0=decode,1=prefill)
softmax_scale, # float scalar
_auto_dim_remap=False, # keep d_head from colliding (baseline note)
)
H_q > H_kv is the change from the baseline's H_q=H_kv=1; it requires
ADR-0061 to survive Phase 2.
4.3 GQA-batched Q load — fold G into the matmul M dim (no broadcast)
For a KV head kv, its group is query heads [kv·G, (kv+1)·G). Load
those G rows and reshape the group into the matmul M dimension —
this is the key to GQA reuse and it needs only a byte-conserving reshape,
not a broadcast:
def load_q_group(q_ptr, kv, G, T_q, d, tl):
# Q laid out [H_q, T_q, d]; the group is contiguous in the head axis.
qg = tl.load(q_ptr + kv*G*T_q*d, shape=(G, T_q, d), dtype="f16")
return _view(qg, (G*T_q, d)) # byte-conserving reshape → M = G*T_q
Then per KV tile (Kj, Vj are [TILE, d], one KV head, loaded once):
Sj = tl.dot(q_g, tl.trans(Kj)) * scale # [G*T_q, d]·[d, TILE] -> [G*T_q, TILE]
Oj = tl.dot(P, Vj) # [G*T_q, TILE]·[TILE, d] -> [G*T_q, d]
K/V are the shared operands (the matmul's b), loaded once and reused
across all G·T_q M-rows automatically. The emitted GemmCmd carries
m = G·T_q, so Phase-1 timing counts all G rows' work correctly.
Why not a leading batch axis (
[G,T_q,d]·[G,d,TILE])?tl.dotderivesm,k,nfrom the last two dims only (tl_context.py:220) andGemmCmdcarries no batch dim — so a leadingGaxis is computed correctly bynp.matmulin Phase 2 (it broadcasts) but under-counted byG×in Phase-1 timing. Folding into M fixes both. This supersedes the earlier batched-vs-loop question.
Transpose caveat.
tl.transis metadata-only andMemoryStore.readreshapes (not transposes) (memory_store.py:73,data_executor.py:153-168doesnp.matmulon the reshaped operands). Fine for timing/structure and for the baseline's zero data; for non-trivial numeric parity, store K pre-transposed as[d, S_pe]or add a realtl.transpose(§10.10).
5. Kernel design (concrete, real tl)
5.1 Shared tile sweep helper
def flash_tile_sweep(q_g, k_base, v_base, n_tiles, TILE, d, scale,
q_block, causal, tl, PREFETCH=2):
"""One PE's FlashAttention sweep over its local KV tiles.
Returns running (m, l, O) in the persistent arena."""
m, l, O = _init_running(q_g.shape, d, tl) # persistent arena
f_k, f_v = {}, {}
for j in range(min(PREFETCH, n_tiles)): # prime (ADR-0062)
if causal and tile_all_future(j, q_block): continue
f_k[j] = tl.load_async(k_base + j*TILE*d, (TILE, d))
f_v[j] = tl.load_async(v_base + j*TILE*d, (TILE, d))
for j in range(n_tiles):
if causal and tile_all_future(j, q_block): continue
with tl.scratch_scope(): # recycle temporaries (ADR-0063)
Kj = tl.wait(f_k[j])
Sj = tl.dot(q_g, tl.trans(Kj)) * scale
mask = make_mask_tile(j, q_block) if (causal and tile_partial(j, q_block)) else None
if mask is not None: Sj = Sj + mask
m_j = tl.max(Sj, axis=-1)
m_new = tl.maximum(m, m_j)
P = tl.exp(Sj - m_new)
corr = tl.exp(m - m_new)
l = l * corr + tl.sum(P, axis=-1)
Vj = tl.wait(f_v[j])
O = O * corr + tl.dot(P, Vj)
m = m_new
nxt = j + PREFETCH # refill pipeline
if nxt < n_tiles and not (causal and tile_all_future(nxt, q_block)):
f_k[nxt] = tl.load_async(k_base + nxt*TILE*d, (TILE, d))
f_v[nxt] = tl.load_async(v_base + nxt*TILE*d, (TILE, d))
return m, l, O
Note:
m, l, Oare reassigned inside thewithbut their storage must be the persistent arena, not the recycled scope. Concretely the merge ops write to pre-allocated persistent handles (double-buffered); ADR-0063 §D3 describes the persistent-vs-scoped arena split. This is the one subtlety the implementer must get right — see §10.5.
5.2 Tree reduction helper
def tree_reduce_and_store(m, l, O, pe_id, N, o_ptr, tl):
for cdir in tree_children_dirs(pe_id, N): # static tree (§4 ADR-0060)
m_c = tl.recv(cdir, m.shape); l_c = tl.recv(cdir, l.shape)
O_c = tl.recv(cdir, O.shape)
mn = tl.maximum(m, m_c); sa = tl.exp(m - mn); sb = tl.exp(m_c - mn)
l = l*sa + l_c*sb; O = O*sa + O_c*sb; m = mn
if is_root(pe_id, N):
tl.store(o_ptr, O / l)
else:
pdir = parent_dir(pe_id, N)
tl.send(pdir, m); tl.send(pdir, l); tl.send(pdir, O)
tree_children_dirs / parent_dir map the binary tree (level pairs in
ADR-0060 §4) onto physical N/S/E/W neighbours — this mapping is the
tuning item ADR-0060 §9.2 and an Open Decision (§10.2).
5.3 Top-level kernel (all four cases)
def gqa_attention_kernel(q_ptr,k_ptr,v_ptr,o_ptr,
S_q,S_kv_per_rank,H_q,H_kv,d,n_ranks,
rank_axis,mode,counter,start_pe,scale,*,tl):
if rank_axis != 0 and tl.program_id(axis=0) != 0: # multi-user: only PE0 per cube
return
pe_id = tl.program_id(axis=rank_axis)
G = H_q // H_kv
causal = (mode == 1) # prefill
q_block = (counter, S_q) if causal else (counter, 1)
# decode small-S fast path = one tile; otherwise tiled (TILE from config)
for kv in range(H_kv): # one KV head per iter (or PE-mapped)
q_g = load_q_group(q_ptr, kv, G, S_q, d, tl)
my_len = valid_len(counter, start_pe, pe_id, n_ranks) if mode==0 else S_kv_per_rank
n_tiles = ceil(my_len / TILE)
m,l,O = flash_tile_sweep(q_g, k_base(kv), v_base(kv), n_tiles, TILE, d,
scale, q_block, causal, tl)
if n_ranks == 1:
tl.store(o_ptr_for(kv), O / l)
elif mode == 1 and is_ring(...): # prefill ring already folded across steps
ring_fold_then_store(...) # §5.4
else:
tree_reduce_and_store(m,l,O, pe_id, n_ranks, o_ptr_for(kv), tl)
This collapses the baseline's two separate kernels (_attention_mesh_kv,
_attention_mesh_mlo) into one parameterised kernel; the mode/n_ranks
flags select the case. (Keeping them separate is also fine — Open
Decision §10.6.)
5.4 Ring (prefill SP) specifics
The Ring path keeps the baseline's structure (_attention_mesh_kv.py)
with three upgrades:
recv_asyncfor the next step's K/V (overlap) —tl_context.py:543.- causal step-skip:
if step_all_future(step, q_block): continuebefore folding (skips ≈half the steps). - GQA:
q_gis theG-row group;tl.broadcastfor the KV head.
The fold across ring steps is the same online-softmax merge as §5.2's
inner update — running (m,l,O) carried across steps as Python handles
(exactly as the baseline does today).
6. The three new primitives — integration points
6.1 ADR-0061 tl.broadcast (OPTIONAL — not on the GQA critical path)
GQA reuse is achieved by folding G into M (§4.3), so broadcast is not
required to unblock h_q > h_kv. It remains a clean primitive for additive
causal-mask construction across the G·T_q rows and for general kernels.
tl_context.py: newbroadcast(x, new_shape)→ emitsBroadcastCmd, output handle from_make_compute_out(sonbytesis correct).data_executor.py:_execute_broadcast=np.broadcast_to(src, shape).copy()→store.write.- Latency:
pe_math._compute_ns(prod(new_shape)); logsmath/broadcast, nodma_read. - Lowest priority; deliver after the efficiency/scale features land.
6.2 ADR-0062 tl.load_async
tl_context.py:load_async(ptr, shape, dtype) → LoadFuture; extendwaitto resolveLoadFuture.pe_dma.py: non-blocking read path; transfer occupies the read channel in parallel; completion is a scheduled event.- op_log unchanged (
memory/dma_read) ⇒dma_read_countmetric stable.
6.3 ADR-0063 tl.scratch_scope
tl_context.py: context manager saving/restoring_scratch_cursor.- Persistent arena:
m,l,O(+ double-buffer) and liveLoadFuturebuffers allocated outside the scope. - Removes the
S=16cap (test_milestone_gqa_llama70b.py:123-148).
7. Phased implementation plan (test-first per CLAUDE.md)
Each phase is an independent Phase-1→Phase-2 cycle. Ordered so each builds on a green predecessor; every phase keeps the existing baseline green.
| Phase | Deliverable | Gate |
|---|---|---|
| P1 | Real GQA by restructuring (per KV head, G→M fold); one-shot, no tiling. No new primitive. |
data-mode run completes for h_q=G·h_kv; K/V dma_read_count independent of G; gemm m=G·T_q |
| P2 | Tree reduction replacing fan-out (decode SP) | ⌈log₂ N⌉ reduction rounds; structure preserved |
| P3 | ADR-0063 tl.scratch_scope + tiled sweep |
long-S decode completes (today caps at S=16) |
| P4 | ADR-0062 tl.load_async + prefetch in sweep |
latency < serial Σ(load+compute) |
| P5 | Causal masking + tile/step skip (prefill, ring) | GEMM count = triangular tile count |
| P6 | Headline-scale milestone panels (real G, tiling, prefetch) + figures |
sweep.json headline rows; figures render |
| P7 (opt) | ADR-0061 tl.broadcast (mask convenience); optional tl.transpose for numeric parity |
broadcast unit test; asymmetric numeric parity |
P1 alone delivers "GQA actually runs" — and it needs no new feature, only the §4.3/§5 restructuring (the iteration-2 finding). P3 is the most important new feature (removes the S=16 scale ceiling). P4 is the efficiency lever. P2/P5 are algorithmic. P7 is optional polish.
8. Verification plan (concrete)
Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (tl),
ADR-0054 (eval bench). Mirrors ADR-0060 §11.
Runs + numeric (Phase 2, enable_data=True):
test_gqa_correctness.py: for(G∈{1,8}, T_q∈{1,16}, S, N∈{1,8})the GQA kernel completes (no byte-conservation error) — theG=8cases that the baseline head-packing cannot express. This is gated on the §4.3 restructuring, not on a new primitive.- Numeric parity (secondary):
O ≈ numpy FlashAttentionfor symmetric/identity inputs where reshape-as-transpose is exact; full asymmetric parity is gated on an optionaltl.transpose(§10.10).
Levers (op_log assertions):
- GQA amortisation: K/V
dma_read_countconstant asGgrows 1→8;gemmwork grows withG. - Tree: decode-SP
ipcq_send/recvcounts = tree edges, reduction rounds =⌈log₂ N⌉(notN−1). - Causal skip: prefill
gemm_count= lower-triangular tile count. - Prefetch overlap: tiled-sweep latency strictly <
Σ(load+compute). - Scratch: peak
_scratch_cursorbounded by one tile, independent ofn_tiles.
Invariants:
- Determinism: identical inputs → identical op_log + latency (SPEC §0.1).
- Every routed request latency > 0 (SPEC §0.1).
9. Performance model (what we expect to see)
Rough per-PE decode cost (KV-load-bound), with prefetch overlap:
t_decode_pe ≈ max( DMA(K+V over S_pe tiles), compute(QKᵀ + PV over S_pe, ×G amortised) )
+ t_tree_reduce(⌈log₂ N⌉ hops, payload O[G,d])
Levers and their expected signature in op_log / latency:
- GQA reuse → KV bytes moved are
H_kv·S·dnotH_q·S·d(8× less), with compute stillH_q-scaled. Decode being KV-load-bound, this is the dominant decode win. - Prefetch → DMA hidden behind compute ⇒
t ≈ max(...)not sum. - Tree → reduction
⌈log₂ N⌉(3 for N=8) vs baselineN−1(7). - Causal skip (prefill) → ≈½ the tile/step work.
These are the quantities the milestone bench's op_log summary + a new latency column should report (Open Decision §10.3 on adding latency to the sweep JSON).
10. Open Decisions / Assumptions made autonomously (REVIEW THESE)
Choices I made to keep progressing without your input. Each is reversible; none is yet implemented in production code (docs only). Marked [recommend] = my default if you don't object.
10.1 Greenlet tl model over composites [recommend: greenlet]
The biggest design call. ADR-0060 §1/§8 argues the greenlet path is equally efficient in kernbench's per-op latency model and far smaller to build than a flash-composite. Risk: if you specifically want the composite abstraction exercised for the eval narrative (e.g. to study scheduler-managed tiling), we'd instead invest in a flash-composite command kind. I judged that out of proportion to the goal. Confirm the greenlet direction.
10.2 Tree reduction topology / neighbour mapping [recommend: binary tree on N/S/E/W]
ADR-0060 §4 fixes a binary tree (N=8: root=7). The mapping of tree edges
to physical mesh directions depends on the SFR install's neighbour
wiring. I assumed configure_sfr_intracube_pe_ring gives a ring on which
the tree pairs are 1-hop neighbours. Needs verification against the
actual installed neighbour table; if pairs aren't 1-hop, either relabel
PEs or accept multi-hop sends. Confirm or hand me the intended mapping.
10.3 Keep the baseline all-to-all as an option? [recommend: replace with tree]
The baseline fan-out gives the answer on every rank; the tree gives it on root only. Attention needs root-only, so tree wins. If any downstream consumer needs O replicated (none found), we'd keep fan-out. Assuming root-only is fine.
10.4 GQA matmul shape — RESOLVED in iteration 2 [fold G into M]
Earlier this asked "batched leading axis vs per-g loop." Resolved: do
neither — fold G·T_q into the matmul M dimension (§4.3). A leading
G axis would be computed by np.matmul in data mode but under-counted
G× in Phase-1 timing (GemmCmd carries no batch dim,
tl_context.py:220). Folding into M gives correct timing and correct
data with a single byte-conserving reshape and one shared K/V operand. No
new primitive, no per-g loop.
10.5 Persistent-vs-scoped arena split [recommend: two explicit arenas]
ADR-0063 requires m,l,O to live outside the recycled scope. The cleanest
implementation reserves a small fixed persistent region at kernel start
and double-buffers the merge. I assumed this is acceptable extra
bookkeeping inside the shared helper. Alternative: make scratch_scope
take an explicit "keep-alive" handle list. I picked the two-arena model
for clarity.
10.6 One unified kernel vs two (decode/prefill) [recommend: unify]
§5.3 unifies into one parameterised kernel; the baseline has two. Unifying
reduces duplication but makes the kernel branch-heavy. Alternative: keep
_gqa_decode and _gqa_prefill separate, sharing helpers. Low stakes;
either is fine.
10.7 Ghost ADRs 0055–0059 [recommend: backfill as Draft, separate task]
The baseline cites ADR-0055/0056/0057/0058/0059 which do not exist
(grep finds refs in code/tests, no files; no git history of them). This
is pre-existing documentation debt, not created by this work. I did not
backfill them (out of scope, and reverse-engineering intent is risky).
Recommendation: a separate Draft ADR-backfill task to document the
existing mesh kernels + milestone bench + SFR install, so ADR-0060 has a
real predecessor to "supersede". Confirm you want this as follow-up.
10.8 Numbering [recommend: as chosen]
GQA ADR = 0060; supporting = 0061/0062/0063. I avoided 0055–0059
(claimed by ghost refs) to prevent collisions. The file was renamed from
ADR-XXXX-… to ADR-0060-algo-… (title was erroneously "ADR-001").
10.9 bf16 modelled as f16 [assumption, pre-existing]
memory_store.py:16 maps bf16→float16 (numpy has no bf16). Numeric
parity tests must use tolerances that accommodate this, and we should not
claim bf16 accuracy results — only timing/structure. Pre-existing
simulator limitation; flagged so results aren't over-interpreted.
10.10 tl.trans is reshape-not-transpose [decision: structural-first verify]
Discovered in iteration 2: tl.trans only swaps shape metadata
(tl_context.py:390) and MemoryStore.read reshapes the bytes
(memory_store.py:73); data_executor._execute_gemm then np.matmuls
the reshaped operands (data_executor.py:153-168). So for non-trivial
numeric data, Q·Kᵀ via tl.trans(K) is mathematically a reshape, not
a transpose — wrong numbers (the baseline never notices: it runs on
all-zeros and only asserts op_log structure). Decision: treat the
simulator as a performance model (SPEC §0) — verify structure + timing +
determinism first; treat numeric parity as a bounded secondary check
(symmetric/identity inputs, or store K pre-transposed). A real
data-materialising tl.transpose is a candidate further primitive if
true numeric validation is ever required — recommend deferring it; it
is not needed for the performance-modeling goal. Confirm you're content
with structural-first verification.
10.11 GQA reuse = fold G into M, NOT a broadcast op [finding, iteration 2]
The headline mechanism finding. The deferred-GQA limitation
(test_milestone_gqa_llama70b.py:137-142) was blamed on a missing
broadcast / the MemoryStore nbytes check. Iteration 2 shows the real
cause is the baseline's head-packing reshape, and the fix is kernel
restructuring (§4.3), which needs no new primitive. This demoted
ADR-0061 from "the blocker" to "optional convenience" and reordered the
phase plan (§7). Recorded explicitly because it changes the project's
critical path: GQA itself is not gated on any new feature; scale
(ADR-0063) and efficiency (ADR-0062) are the features that matter.
11. Risks
| Risk | Likelihood | Mitigation |
|---|---|---|
| Numeric parity blocked by trans-as-reshape | med | structural-first verify (§10.10); pre-transpose K or defer tl.transpose |
| Tree pairs not 1-hop on installed mesh | med | verify SFR neighbour table (§10.2); relabel PEs |
| scratch_scope use-after-scope bug | med | two-arena discipline in shared helper; debug poison |
| prefetch buffer recycled while in flight | med | live LoadFuture buffers in persistent arena (§10.5) |
GQA broadcast .copy() memory blow-up |
low | bounded by TILE in tiled design |
| Effort underestimated (6 phases) | med | P0–P1 alone deliver the core unblock; later phases independent |
12. Glossary & references
- GQA / G — grouped-query attention;
G = H_q/H_kvquery heads share one KV head. Llama3-70B:G=8. - FlashAttention / FlashDecoding / Ring Attention — see ADR-0060 lineage note.
- IPCQ — inter-PE queues;
tl.send/tl.recvoverN/S/E/W(ADR-0023/0025/0032). - Greenlet
tlmodel — kernel interleaves with SimPy; ops emit commands scheduled on PE engines (ADR-0020, ADR-0046). - Persistent vs scoped arena — ADR-0063 scratch split.
Key source anchors: tl_context.py (tl API), memory_store.py:67-73
(nbytes check), data_executor.py (Phase 2), pe_dma.py/pe_gemm.py/
pe_math.py (latency), _attention_mesh_kv.py / _attention_mesh_mlo.py
(baseline kernels), milestone_gqa_llama70b.py (eval bench),
tests/attention/test_milestone_gqa_llama70b.py (current limits).
ADRs: 0060 (this design), 0061 (broadcast), 0062 (async load), 0063 (scratch scope); related accepted 0014/0020/0023/0025/0042/ 0045/0046/0052/0054.