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

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

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

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

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

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