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 plan for the two GQA FlashAttention kernels (decode=reduce, prefill=ring) on AHBM in kernbench. Audience: reviewer resuming with "GQA 검토 시작하자". Read ADR-0060 first — it is now the authoritative design record (TL;DR has full pseudocode for both kernels). This DDD is the how to build it: file plan, phase plan, helper signatures, tests. It does not re-derive the design; it points to ADR-0060 sections.

Alignment note (this revision). ADR-0060 moved to a composite hybrid + hierarchical CUBE-Group SP design with two kernels. This DDD is rewritten to match. The earlier DDD (greenlet-primitive, tl.load_async, single unified kernel) is superseded.

---

1. Goal and success criteria

Goal: two efficient GQA fused-attention kernels on AHBM in kernbench — decode+SP (head-replicated, KV static shard, 2-level reduce) and prefill+SP (1 Q head per CUBE, Ring KV) — at realistic 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 — levers observable in op_log (ADR-0060 §11): GQA K/V-load amortisation; decode 2-level reduce (⌈log₂P⌉ + center-mesh over C, not C·P−1); prefill ring (no reduce, KV rotates); load/compute overlap (lazy tl.load + composite streaming); composite GEMM offload.
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-SIP head split (a head stays in one SIP, ADR-0060 §0); cycle-accurate microarchitecture (SPEC §5); a bespoke "flash-composite" command kind (ADR-0060 §8 item 4).

---

2. Where the design sits in the current code

runtime_api (host)
  RuntimeContext.zeros/empty/launch    context.py  (tensor deploy, DPPolicy(+cube_start), launch)
        │  KernelLaunchMsg
        ▼
sim_engine
  GraphEngine(enable_data=…)           engine.py
  DataExecutor (Phase 2)               data_executor.py     ← _execute_broadcast (ADR-0061, optional)
  MemoryStore                          memory_store.py
        │
        ▼
components / PE pipeline (per PE)
  PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM}
  greenlet kernel ↔ SimPy              kernel_runner.py, pe_cpu.py
  CompositeCmd → tile plan → engines   pe_scheduler.py (ADR-0014 D6)
        │
        ▼
tl programming model                   tl_context.py
  load (→ lazy, ADR-0062) / store / composite / dot / max/sum/exp/maximum
  send/recv/recv_async / scratch_scope (ADR-0063) / broadcast (ADR-0061, opt)

Both kernels are bench-side Python (src/kernbench/benches/) run as greenlet kernels emitting tl.* ops; the GEMMs are tl.composite (scheduler-managed tiling + K/V DMA streaming, existing CompositeCmd — no new command kind), the softmax merge + reduction/ring stay kernel-level (ADR-0060 §1). The remote work already added building blocks: DPPolicy.cube_start, _attention_mesh_mlo_2d.py (2D cube reduce — currently AllReduce, to become reduce-to-root), topologies/llama70b_4sip.yaml.

---

3. File plan

3.1 New / modified production files

File	Change	ADR
src/kernbench/triton_emu/tl_context.py	tl.load → lazy (non-blocking + auto-wait on first use); scratch_scope ctx-mgr; broadcast (optional)	0062/0063/0061
src/kernbench/components/builtin/pe_dma.py	non-blocking read path: post DMA, resolve on a scheduled event (the recv_async pattern for loads)	0062
src/kernbench/triton_emu/kernel_runner.py	auto-wait: yield a handle's pending load event at first consume; lazy-load dispatch arm	0062
src/kernbench/sim_engine/data_executor.py	_execute_broadcast (optional)	0061
src/kernbench/components/builtin/pe_cpu.py + cost table	per-op-type CPU issue cost (replaces uniform dispatch_cycles)	0064 (separate)

No new GEMM command kind: the two GEMMs use the existing CompositeCmd.

3.2 Kernel + bench files

File	Role
src/kernbench/benches/_gqa_decode.py	decode+SP kernel — evolve _attention_mesh_mlo/_mlo_2d to composite-hybrid + 2-level reduce-to-root (Level-2 PE tree + Level-1 center-mesh; currently 2D AllReduce)
src/kernbench/benches/_gqa_prefill.py	prefill+SP kernel — evolve _attention_mesh_kv to head-parallel (1 Q head/CUBE) Ring KV + composite-hybrid
src/kernbench/benches/_gqa_helpers.py	head_of_group, valid_len_2level, init_running, tree/center-mesh topology, hierarchical_reduce_and_store, ring helpers, causal block_* predicates + mask builders
extend src/kernbench/benches/milestone_gqa_llama70b.py	headline panels: real G, C=G=8, contiguous KV, 4-SIP
topologies/llama70b_4sip.yaml	exists (remote) — 4 SIP × 16-CUBE 4×4 × 8 PE

3.3 New tests

File	Covers
tests/attention/test_gqa_decode.py	decode reduce: GQA reuse (dma_read ⟂ G), 2-level reduce step count, root-only output, long context
tests/attention/test_gqa_prefill.py	prefill ring: no (m,ℓ,O) reduce, C KV rotations, causal step-skip, per-CUBE distributed output
tests/test_tl_lazy_load.py	ADR-0062 unit (overlap, auto-wait correctness, op_log parity)
tests/test_tl_scratch_scope.py	ADR-0063 unit
tests/test_tl_broadcast.py	ADR-0061 unit (optional)
tests/test_dppolicy_cube_start.py	exists (remote) — sub-mesh placement

---

4. Data model and launch contract

4.1 Tensor placement — CUBE Group, 2-level SP, contiguous KV

A CUBE Group is C CUBEs in one SIP owning one KV head (ADR-0060 §0). Placement via DPPolicy(..., cube_start=...):

Tensor	Decode+SP	Prefill+SP (C=G)
Q	replicate (all G heads, M-folded) over C×P ranks	head-parallel: CUBE i holds Q head i
K,V	sequence-sharded two axes (cube row_wise over C, pe row_wise over P)	sequence-sharded over C CUBEs (the ring blocks)
O	reduced to CUBE-Group root (all G heads)	per-CUBE (one head each), distributed

KV cache is contiguous C×P position blocks (rank r = [r·B,(r+1)·B)), shared by both kernels — prefill writes it (via qkv_rope), decode reads

• extends it, no reshard. Contiguous is required for prefill causal skip (ADR-0060 §2.1). cube_start offsets the group's CUBE sub-mesh within the 4×4 SIP mesh.

Driver SP-enable threshold (fallback). Contiguous sharding under-utilizes ranks for short/early decode (only frontier ranks hold data). So the driver enables SP only past a context-length threshold where KV-sweep time dominates reduction + placement imbalance; below it, fall back to a smaller C (or C=P=1, single rank, no reduce). The threshold is a sweep item (P7 / §9), not a fixed constant.

--device enumerates SIPs; one SIP-device runs its 2 CUBE Groups (2 KV heads); the head is picked by CUBE coordinate (head_of_group), not an in-kernel loop (ADR-0060 §0).

4.2 Launch signatures

# decode+SP
ctx.launch(f"{panel}_gqa_decode", gqa_decode_sp,
           q, k, v, o, counter, start_pe, start_cube, C, P, softmax_scale)
# prefill+SP
ctx.launch(f"{panel}_gqa_prefill", gqa_prefill_sp,
           q, k, v, o, T_q, S_kv_local, d, C, softmax_scale, q_block, cube_start)

Full per-kernel I/O contract: ADR-0060 §0.5.3/§0.5.4 (note the output head distribution differs — decode root-gathered, prefill distributed).

4.3 GQA reuse — fold G into the matmul M dim (decode)

Decode replicates Q and M-folds the G query heads into the GEMM row dim so one Q·Kᵀ serves all G heads sharing one K (ADR-0060 §0, §5.2): Q group [G, T_q, d] → [G·T_q, d] (byte-conserving _view); the composite carries m = G·T_q (timing counts all G rows). Prefill does not M-fold (1 head/CUBE). Caveat: tl.trans is reshape-not-transpose → store K pre-transposed [d, S/(C·P)] (ADR-0060 §3, §B).

---

5. Kernel design

The full pseudocode for both kernels (and the 3 decode CPU-pipelining variants) lives in ADR-0060 TL;DR + §5. Implementation notes only here.

5.1 Decode (reduce) — _gqa_decode.py

• Inner tile = §3 composite-hybrid: Sj = composite(q_g, Kⱼ)·scale → softmax MATH → Oj = composite(P, Vⱼ) → running merge (ADR-0060 §3).
• Ship opt3 (software pipelining: issue next tile's Q·Kᵀ before this tile's softmax; Sj in a persistent double buffer) — removes the GEMM-engine bubble, no new command kind (ADR-0060 §5.6). opt1 is the naïve baseline; opt2 (ex_composite) waits on ADR-0064.
• Combine = hierarchical_reduce_and_store (ADR-0060 §4): Level-2 PE reduce-to-root tree (intra-CUBE, the cube's KV slice further split across its P PEs — decision (a)) → Level-1 center-root CUBE-mesh reduce (intra-CUBE-Group). Data-driven, level-pipelined, root-only.

5.2 Prefill (ring) — _gqa_prefill.py

• Head-parallel: CUBE i owns Q head i + KV slice i. No reduce (each CUBE outputs a distinct head). Ring rotates KV blocks; GQA reuse via the rotation (ADR-0060 §5.5).
• K/V are tl.load'd TCM handles (not tl.ref) so the ring can tl.send/recv_async them; K pre-transposed [d, S/C]. Receive buffers ping-pong (persistent arena) — recv into a separate buffer while computing/sending the current one (do not clobber).
• Causal step-skip + boundary mask; recv_async overlaps next block's receive with current compute.
• Within a CUBE: tile the Q head's T_q rows across P PEs (disjoint output rows → no intra-CUBE reduce); fall back to KV-block split only if T_q < P (ADR-0060 §B decode-split item 3).

5.3 Shared helpers — _gqa_helpers.py

head_of_group(cube_id) = cube_id // C; init_running(...) → persistent (m=-inf, ℓ=0, O=0); valid_len_2level(...) (contiguous block extent); hierarchical_reduce_and_store(...); tree/center-mesh child/parent dirs; block_all_future/block_partial + causal_mask.

---

6. Supporting primitives — integration points

6.1 ADR-0062 lazy tl.load (efficiency — overlap)

• tl.load issues DmaReadCmd non-blocking, returns a handle with a pending event; the runtime auto-inserts the wait at first consume (generalises the recv_async/wait pattern, kernel_runner.py:248-285).
• pe_dma.py: non-blocking read; DMA occupies the (capacity-1) read channel, overlaps compute on PE_GEMM/PE_MATH.
• Global semantics change → existing goldens regenerate (ADR-0062 D3).
• op_log unchanged (memory/dma_read) ⇒ dma_read_count stable.

6.2 ADR-0063 tl.scratch_scope (scale)

• ctx-mgr saving/restoring _scratch_cursor. Persistent arena for (m,ℓ,O) (+ decode opt3 Sj double buffer, + prefill ring ping-pong buffers, + in-flight lazy-load buffers) lives outside the scope. Removes S=16.

6.3 ADR-0061 tl.broadcast (optional)

• Not on the GQA critical path (M-fold gives reuse). Convenience for additive mask construction. Lowest priority.

6.4 ADR-0064 per-op-type CPU issue cost (separate ADR)

• Replaces uniform dispatch_cycles=0. Makes the hybrid's CPU-saturation win (and the decode opt2 ex_composite fewer-issues win) measurable. Land before claiming hybrid latency wins (ADR-0060 §9).

---

7. Phased implementation plan (test-first per CLAUDE.md)

Each phase is an independent Phase-1→Phase-2 cycle; each keeps the baseline green.

Phase	Deliverable	Gate
P1	Real GQA, decode, composite-hybrid M-fold; one-shot, no tiling. No new primitive (existing CompositeCmd).	data-mode completes h_q=G·h_kv; K/V dma_read_count ⟂ G; composite tile plan m=G·T_q
P2	Decode 2-level reduce-to-root (Level-2 PE tree + Level-1 center-mesh), replacing the 2D AllReduce	reduce rounds = ⌈log₂P⌉+center-mesh; root-only output; Level-1 on CUBE NOC, Level-2 on PE IPCQ
P3	ADR-0063 scratch_scope + tiled sweep	long-S decode completes (today caps S=16); scratch O(1)
P4	ADR-0062 lazy tl.load + composite K/V streaming	tiled-sweep latency < serial Σ(load+compute); goldens regenerated
P5	Decode opt3 software pipelining	GEMM engine non-idle across tiles (structural / op_log)
P6	Prefill head-parallel Ring KV + causal step-skip	C KV rotations, no (m,ℓ,O) reduce, per-CUBE distributed O, causal skip ≈½
P7	Contiguous shared KV layout + 4-SIP topology + headline milestone panels (real G, C=G)	shared layout (no reshard); headline sweep.json rows
P8 (opt)	ADR-0064 cost model + decode opt2 ex_composite; ADR-0061 broadcast	cost-model goldens; opt2 fewer-issues measurable

P1 delivers "GQA runs" on the composite path. P3 is the key scale feature; P4 the overlap lever; P2/P6 the SP algorithms.

---

8. Verification plan (concrete)

Mirrors ADR-0060 §11; grounded in SPEC R2/R5, ADR-0023/0025, ADR-0046, ADR-0054.

Runs + numeric (enable_data=True):

• decode & prefill complete for (G∈{1,8}, T_q∈{1,16}, S, C∈{1,8}, P∈{1,8}) without byte-conservation error (the G=8 cases baseline cannot express).
• Numeric parity (secondary): O ≈ numpy FlashAttention for symmetric/identity inputs (reshape-as-transpose exact); full asymmetric parity gated on a real tl.transpose (deferred).

Levers (op_log):

• GQA amortisation: K/V dma_read_count ⟂ G; compute scales with G.
• Decode reduce: rounds = ⌈log₂P⌉ + center-mesh (not C·P−1); result at one rank; Level-1=CUBE NOC, Level-2=PE IPCQ.
• Prefill ring: C KV rotations, no reduce, each CUBE writes a distinct head.
• Causal skip: GEMM/step count = lower-triangular.
• Overlap: tiled-sweep latency < Σ(load+compute).
• Scratch: peak cursor bounded by one tile.

Invariants: determinism (identical op_log + latency); every routed request latency > 0 (SPEC §0.1).

---

9. Performance model (expected)

Per-rank decode (KV-load-bound), composite-streamed + overlapped:

t_decode_rank ≈ max( DMA(K+V over S/(C·P) tiles),  compute(QKᵀ+PV, ×G amortised) )
              + t_reduce( ⌈log₂P⌉ intra-CUBE  +  center-mesh over C inter-CUBE )

Prefill (ring): t ≈ C · max( recv(KV block), compute(QKᵀ+PV over block) ) with causal step-skip removing ≈half the steps; no reduce term.

Levers: GQA reuse → KV bytes H_kv·S·d not H_q·S·d; overlap → max not sum; decode reduce ⌈log₂P⌉+center-mesh vs baseline C·P−1; prefill ring trades the reduce for KV rotation (good when O is big). The CPU-saturation win (composite offload) is measurable only with ADR-0064.

---

10. Open items (status)

Most original "Open Decisions" are now resolved in ADR-0060; the live review items live in ADR-0060 §B (three groups: hybrid pivot, hierarchical CUBE-Group SP, decode/prefill split). Key resolved choices:

Was open	Now
greenlet vs composite	composite hybrid (GEMMs→composite, merge→kernel)
one kernel vs two	two kernels (decode reduce / prefill ring)
GQA matmul shape	M-fold (decode); head-parallel (prefill)
reduction topology	decode 2-level reduce-to-root; prefill no reduce (ring)
tl.load_async	lazy tl.load (ADR-0062, redefined)
KV placement	contiguous C×P, shared prefill/decode (ADR-0060 §2.1)

Still to confirm (ADR-0060 §B): DDD↔impl reconcile (_attention_mesh_mlo_2d AllReduce → reduce-to-root); head_of_group sub-mesh partition; C=G coupling; short-context KV balance; ADR-0064 calibration. Pre-existing: ghost ADRs 0055–0059 (separate backfill); bf16→f16 proxy.

---

11. Risks

Risk	Likelihood	Mitigation
Numeric parity blocked by trans-as-reshape	med	structural-first verify; pre-transpose K (ADR-0060 §B item 2)
2-level reduce pairs not 1-hop on the mesh	med	verify SFR neighbour table; center-root sub-mesh partition (§B)
scratch_scope use-after-scope (Sj / ring / lazy-load buffers)	med	persistent-arena discipline; ping-pong buffers (§5.2)
prefill ring buffer clobbered while in flight	med	recv into the other ping-pong buffer (§5.2)
decode opt2 acc read before composites drain	med	tl.wait() after the loop (ADR-0060 TL;DR opt2)
impl drift (_attention_mesh_mlo_2d is AllReduce, Q-replicated)	med	reconcile to reduce-to-root; add prefill ring (§B)
CPU-saturation win invisible	exp	ADR-0064 cost model (P8)

---

12. Glossary & references

• GQA / G — G = H_q/H_kv query heads share one KV head. Llama3-70B: G=8.
• CUBE Group — C CUBEs in one SIP owning one KV head (ADR-0060 §0).
• 2-level SP — Level-1 inter-CUBE (over C) × Level-2 intra-CUBE PE (over P); ranks = C·P.
• Composite hybrid — GEMMs via tl.composite (scheduler), softmax merge
  • reduction in the kernel (ADR-0060 §1).
• M-fold — stack the G query heads into the GEMM M (row) dim so one Q·Kᵀ does all heads sharing one K.
• FlashAttention / FlashDecoding / Ring Attention — ADR-0060 lineage.

Source anchors: tl_context.py (tl API), pe_scheduler.py:104-143 (composite tile plan), pe_dma.py:45,89 (read channel, drain), pe_cpu.py (dispatch_cycles), lrab_hierarchical_allreduce.py (center-root pattern), _attention_mesh_mlo_2d.py / _attention_mesh_kv.py (impl to evolve), DPPolicy.cube_start, topologies/llama70b_4sip.yaml.

ADRs: 0060 (this design), 0062 (lazy load), 0063 (scratch scope), 0061 (broadcast, optional), 0064 (CPU issue cost model); related accepted 0014/0017/0020/0023/0025/0046/0054.