kernbench2/docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md

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:

1. Correct: kernel output matches a numpy FlashAttention reference within fp tolerance, with real GQA (H_q=64, H_kv=8, G=8).
2. Efficient: the three levers actually fire and are observable in op_log: GQA K/V-load amortisation, KV prefetch overlap, ⌈log₂ N⌉ reduction.
3. Scales: context length not capped by the 1 MiB scratch bump allocator.
4. 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

For a KV head kv, its group is query heads [kv·G, (kv+1)·G). The kernel loads those G rows and, where the matmul needs the KV head expanded to G, broadcasts the KV head — not the symbolic _view:

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 head axis.
    return tl.load(q_ptr + kv*G*T_q*d, shape=(G, T_q, d), dtype="f16")

def kv_for_group(k_tile, G, tl):           # ADR-0061
    # k_tile: [TILE, d] (one KV head). Expand to [G, TILE, d] for batched dot.
    return tl.broadcast(k_tile[None, :, :], (G, k_tile.shape[0], k_tile.shape[1]))

(Whether the batched tl.dot consumes [G,T_q,d]·[G,d,TILE] directly or the kernel loops g is an Open Decision — §10.4.)

---

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, O are reassigned inside the with but 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_async for the next step's K/V (overlap) — tl_context.py:543.
• causal step-skip: if step_all_future(step, q_block): continue before folding (skips ≈half the steps).
• GQA: q_g is the G-row group; tl.broadcast for 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

• tl_context.py: new broadcast(x, new_shape) → emits BroadcastCmd, output handle from _make_compute_out (so nbytes is correct).
• data_executor.py: _execute_broadcast = np.broadcast_to(src, shape).copy() → store.write.
• Latency: pe_math._compute_ns(prod(new_shape)); logs math/broadcast, no dma_read.
• Unblocks: the tl.dot(q_g[G,…], kv_for_group(...)[G,…]) consuming a broadcast KV head passes the MemoryStore nbytes check (memory_store.py:67-73) that fails today.

6.2 ADR-0062 tl.load_async

• tl_context.py: load_async(ptr, shape, dtype) → LoadFuture; extend wait to resolve LoadFuture.
• 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_count metric 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 live LoadFuture buffers allocated outside the scope.
• Removes the S=16 cap (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
P0	ADR-0061 tl.broadcast + unit test	broadcast Phase 2 parity; downstream dot passes nbytes check
P1	Real GQA in baseline decode kernel (h_q=G·h_kv), one-shot, no tiling, using broadcast	numpy parity at small S; K/V dma_read_count independent of G
P2	Tree reduction replacing fan-out (decode SP)	⌈log₂ N⌉ reduction rounds; parity preserved
P3	ADR-0063 tl.scratch_scope + tiled sweep	long-S decode completes; parity
P4	ADR-0062 tl.load_async + prefetch in sweep	latency < serial Σ(load+compute); parity
P5	Causal masking + tile/step skip (prefill, ring)	GEMM count = triangular tile count; parity
P6	Headline-scale milestone panels (real G, tiling, prefetch) + figures	sweep.json headline rows; figures render

P0–P1 alone deliver "GQA actually runs in data mode" — the single most important unblock. P2–P5 deliver efficiency. P6 is the eval payoff.

---

8. Verification plan (concrete)

Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (tl), ADR-0054 (eval bench). Mirrors ADR-0060 §11.

Correctness (Phase 2, enable_data=True):

• test_gqa_correctness.py: for (G∈{1,8}, T_q∈{1,16}, S, N∈{1,8}), kernel O ≈ numpy FlashAttention reference. The G=8 cases are the ones that cannot run today (ADR-0061 gate).

Levers (op_log assertions):

• GQA amortisation: K/V dma_read_count constant as G grows 1→8; gemm work grows with G.
• Tree: decode-SP ipcq_send/recv counts = tree edges, reduction rounds = ⌈log₂ N⌉ (not N−1).
• Causal skip: prefill gemm_count = lower-triangular tile count.
• Prefetch overlap: tiled-sweep latency strictly < Σ(load+compute).
• Scratch: peak _scratch_cursor bounded by one tile, independent of n_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·d not H_q·S·d (8× less), with compute still H_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 baseline N−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 Batched 2-D dot vs per-g loop for GQA [recommend: try batched first]

tl.dot semantics for a leading group axis ([G,T_q,d]·[G,d,TILE]) need confirming in tl_context.dot (it computes m,k,n from the last two dims; tl_context.py:213-229). If batched dot over a leading axis isn't supported, fall back to a for g in range(G) loop (ADR-0061 A1) — correct but G× GEMM ops. I assumed batched works or is a small extension.

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.

---

11. Risks

Risk	Likelihood	Mitigation
Batched dot over leading axis unsupported	med	per-g fallback (§10.4); still correct
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_kv query 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.recv over N/S/E/W (ADR-0023/0025/0032).
• Greenlet tl model — 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.