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.

`--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

```python
# 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**.