gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-06-09 18:15:59 -07:00
parent b3730a33eb
commit d282144339
52 changed files with 3952 additions and 2526 deletions
@@ -8,11 +8,9 @@ Proposed
**Decision drivers:** agentic workload → 낮은 batch, 긴 context;
KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV).
**Supersedes / extends:** 기존 mesh-native 어텐션 커널
`_attention_mesh_kv`(prefill)와 `_attention_mesh_mlo`(decode), 그리고
`milestone-gqa-llama70b` eval bench. *§A. 기존 kernbench 작업과의 관계* 참조 —
그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는
baseline이다.
**Supersedes / extends:** 이전 mesh-native 어텐션 커널과 그 `milestone-gqa-llama70b`
eval bench(본 ADR의 커널이 도입되면서 제거됨). *§A. 기존 kernbench 작업과의 관계* 참조 —
그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드한 baseline이다.
**Supporting ADRs** (efficiency / scale enabler — *GQA blocker 아님*; §8 정정
참조): **ADR-0063** `tl.scratch_scope`(per-tile scratch 재활용 — 현실적 context
@@ -160,18 +158,16 @@ def gqa_prefill_sp(q_ptr, k_ptr, v_ptr, o_ptr, T_q, S_kv_local, d, C, scale, q_b
## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것)
kernbench는 **오늘날 이미 IPCQ 상에서 online-softmax `(m, , O)` 머지로
FlashAttention을 돌린다.** 두 커널이 존재한다:
kernbench는 본 ADR 이전에 IPCQ 상에서 online-softmax `(m, , O)` 머지로
FlashAttention을 돌리는 두 mesh 커널이 있었다(현재 제거됨):
| File | Role | Mechanism |
|---|---|---|
| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,,O)` fan-out, log-sum-exp merge |
| Role | Mechanism |
|---|---|
| prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
| decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,,O)` fan-out, log-sum-exp merge |
둘 다 `milestone-gqa-llama70b`(4 패널:
`{single,multi}_user × {prefill,decode}`,
`src/kernbench/benches/milestone_gqa_llama70b.py`)이 구동하며
`tests/attention/test_milestone_gqa_llama70b.py`에서 테스트된다.
`{single,multi}_user × {prefill,decode}`)이 구동했다.
**이들은 greenlet `tl` API로 작성되었다:** `tl.load`, `tl.dot`,
`tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, 그리고
@@ -183,8 +179,8 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
**baseline의 의도된 세 한계** — 효율적 GQA 커널이 정확히 들어내야 하는 것:
1. **GQA 재사용 없음.** `h_q == h_kv == 1`
(`test_milestone_gqa_llama70b.py:137-142`). 테스트는 이를 *broadcast view*에
1. **GQA 재사용 없음.** baseline은 `h_q == h_kv == 1`로 제한되어 있었다.
해당 baseline 테스트는 이를 *broadcast view*에
대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의
**head-packing 핵**의 속성이다(`_view(K, (h_q·d, S_kv))`는 모든 head를 하나의
matmul 차원에 뭉치고 `h_q == h_kv`일 때만 byte를 보존). 올바른 수정은
@@ -196,7 +192,7 @@ FlashAttention을 돌린다.** 두 커널이 존재한다:
필요 → **2-level reduce-to-root**(intra-CUBE tree + intra-CUBE-Group
center-mesh, §4)는 `⌈log₂ P⌉` + center-mesh-over-`C` 단계.
3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile
임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking /
임시값을 누수하고(baseline 테스트는 그래서 S=16으로 제한되어 있었음) causal masking /
tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite
K/V 스트리밍(§3) + **ADR-0062**(lazy load 오버랩)으로 해결.
@@ -563,8 +559,7 @@ reduce **안 함** — KV를 회전(Ring, §5.5). decode에 reduce를 택한 이
`O = [G, d]`가 작아 `(m,,O)` 이동이 상주 KV 이동보다 싸기 때문.
tile sweep 후 각 rank는 자기 `1/(C·P)` shard에 대해 `(m_i, _i, O_i)`(비정규화)를
가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학,
`_attention_mesh_mlo.py:117-122`):
가진다. 결합/교환 가능한 log-sum-exp 머지로 결합(baseline의 fold와 동일 수학):
```python
def merge(m_a, l_a, O_a, m_b, l_b, O_b):
@@ -688,9 +683,8 @@ def gqa_decode_sp(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, start_cube,
축은 세어지지 *않음* — §8 참조).
- `S_rank`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial
attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) —
바로 baseline `_attention_mesh_mlo``_partial_attention`, 단지 GQA-batched에
composite 경로. Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만
발동.
바로 baseline의 `_partial_attention` 구조, 단지 GQA-batched에 composite 경로.
Tiling(§3)은 `S_rank`가 scratch scope의 tile 예산을 초과할 때만 발동.
### 5.3 DECODE, with SP / KV-parallel (`C × P` rank)
@@ -741,7 +735,7 @@ KV를 필요로 하므로.
**`(m,,O)` reduce 없음** — 각 CUBE가 자기 head의 행을 정규화·기록.
- **Ring KV:** `C` KV slice가 CUBE ring을 **회전**; 각 CUBE가 들어오는 블록을 자기
head의 실행 중 `(m,,O)`에 fold(online-softmax, ring step 가로질러 Python 핸들로
carry, `_attention_mesh_kv`가 오늘 하듯). `C` step 후 모든 head가 전 KV를 봄.
carry). `C` step 후 모든 head가 전 KV를 봄.
**GQA 재사용은 회전에서** — slice `j`가 전 `C` CUBE를 방문하며 전 `G` head를 서비스.
- IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`
오버랩(`tl_context.py:543-560` — 이미 존재); 수신 버퍼는 ping-pong(persistent
@@ -751,7 +745,7 @@ KV를 필요로 하므로.
- **CUBE 내(P PE):** head의 query 행 `[T_q, d]` 및/또는 현재 KV 블록을 `P` PE에
타일(여기엔 decode와 달리 query 축이 존재); 상세는 §B.
baseline `_attention_mesh_kv` 이미 ring fold를 구현; 본 ADR은 GQA 재사용,
(현재 제거된) baseline이 이미 ring fold를 구현; 본 ADR은 GQA 재사용,
head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가.
### 5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2)
@@ -1060,7 +1054,7 @@ exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`
가정. `C ≠ G`면 매핑 재검토 필요(CUBE당 다중 head, 또는 head가 부분 ring에 걸침).
**권고:** headline 스케일에서 prefill 커널은 `C = G` 고정; `C ≠ G`는 별도 연구.
5. **`_attention_mesh_mlo_2d`(현재 impl)와 정합.** 원격 impl의 2D 커널은 **Q
5. **이전 `_attention_mesh_mlo_2d` impl(현재 제거됨)과 정합.** 2D 커널은 **Q
replicated**로 cube에 걸친 **AllReduce** — 즉 decode-reduce 계열이나 reduce-to-root
아닌 all-reduce(broadcast-back), 그리고 prefill head-parallel ring은 아직 없음.
**권고:** (a) 그 2D AllReduce → reduce-to-root(broadcast-back 제거)로 decode 커널;