From 7fad0371c578ca157eff2b84a7cb820767a2c749 Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Wed, 10 Jun 2026 16:17:32 -0700 Subject: [PATCH] =?UTF-8?q?gqa:=20tile-granular=20Ring=20KV=20(P3c)=20+=20?= =?UTF-8?q?rename=20to=20gqa=5Fattention=5F*=20+=20ADR-0060/62/63/64=20?= =?UTF-8?q?=E2=86=92=20Accepted?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three logically distinct changes, bundled for atomic test green: 1. **P3c — prefill_long tile-granular Ring KV** (ADR-0060 §5.5.1 amendment). Convert the ring from slice-granular (one full ``(d_head, S_local)`` KV slice per step) to tile-granular (``n_tiles`` tiles of ``TILE_S_KV`` per step). Nested loop with outer tile, inner ring step: each tile propagates through all C ring positions before the next tile starts, so IPCQ in-flight depth stays at 1 per direction. Bootstrap at ``(t=0, k=0)`` outside the scratch_scope establishes the persistent ``(m, ℓ, O)``; every other iteration scope-wraps + persists via ``copy_to``. Per-rank persistent scratch shrinks to ~1 KB; per-tile scope bounded by TILE_S_KV regardless of S_local. Headline: prefill_long now completes at S_kv=128K (previously overflowed). New: ``tests/attention/test_gqa_prefill_long_tile_ring.py`` (3 tests — ceiling-lift + tile-granular ipcq_copy count + per-CUBE distributed output regression guard). 2. **Rename ``gqa_*`` → ``gqa_attention_*``** across kernel files, function names, and importers. The "attention" name makes the role explicit (GQA is grouped-query attention) and matches upstream Triton FlashAttention naming conventions. Renames: _gqa_decode_long.py -> _gqa_attention_decode_long.py _gqa_decode_short.py -> _gqa_attention_decode_short.py _gqa_prefill_long.py -> _gqa_attention_prefill_long.py _gqa_prefill_short.py -> _gqa_attention_prefill_short.py And function names ``gqa___kernel`` → ``gqa_attention___kernel``. Updated 1 bench file (milestone_gqa_headline.py) and 10 test files. 3. **ADR-0060 / 0062 / 0063 / 0064: Proposed → Accepted**. All four are reflected in production code and covered by tests: - ADR-0060 (GQA fused attention): 4 kernels deployed; §5.5.1 amendment added for the tile-granular Ring KV introduced by P3c (EN + KO mirror). - ADR-0062 (lazy tl.load): LoadFuture + _await_pending live in tl_context.py. - ADR-0063 (tl.scratch_scope + tl.copy_to): used in every chain reduce + tile sweep + ring step. EN-only previously; KO translation authored as part of this commit (CLAUDE.md bidirectional rule). - ADR-0064 (per-op-type CPU issue cost): cpu_issue_cost.py + issue_cost_table wiring in tl_context.py (Phase E). Files git mv'd from docs/adr-proposed/ to docs/adr/ (EN) and docs/adr-ko/ (KO). ADR-0061 (tl.broadcast) stays Proposed — no implementation; documented as optional convenience primitive in the ADR itself. Tests: 88/88 focused regression green (tests/attention/ + Phase E + TL discipline). ADR pair verification: ``python tools/verify_adr_lang_pairs.py`` OK. Co-Authored-By: Claude Opus 4.7 --- ...ADR-0060-algo-gqa-fused-attention-ahbm.md} | 42 +++- .../ADR-0062-prog-tl-async-load.md} | 2 +- docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md | 222 ++++++++++++++++++ .../ADR-0060-algo-gqa-fused-attention-ahbm.md | 50 +++- .../ADR-0062-prog-tl-async-load.md | 2 +- .../ADR-0063-prog-tl-scratch-scope.md | 2 +- ..._long.py => _gqa_attention_decode_long.py} | 2 +- ...hort.py => _gqa_attention_decode_short.py} | 2 +- .../benches/_gqa_attention_prefill_long.py | 151 ++++++++++++ ...ort.py => _gqa_attention_prefill_short.py} | 2 +- src/kernbench/benches/_gqa_prefill_long.py | 104 -------- .../benches/milestone_gqa_headline.py | 8 +- tests/attention/test_gqa_decode.py | 4 +- tests/attention/test_gqa_decode_mc.py | 4 +- tests/attention/test_gqa_decode_sp.py | 4 +- tests/attention/test_gqa_long_context.py | 8 +- tests/attention/test_gqa_prefill.py | 6 +- .../test_gqa_prefill_long_tile_ring.py | 162 +++++++++++++ tests/attention/test_gqa_prefill_ring.py | 4 +- tests/attention/test_gqa_scoped_writeback.py | 8 +- tests/attention/test_gqa_short_context.py | 12 +- tests/attention/test_gqa_tile_sweep.py | 18 +- 22 files changed, 667 insertions(+), 152 deletions(-) rename docs/{adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md => adr-ko/ADR-0060-algo-gqa-fused-attention-ahbm.md} (96%) rename docs/{adr-proposed/ADR-0062-prog-tl-async-load-ko.md => adr-ko/ADR-0062-prog-tl-async-load.md} (99%) create mode 100644 docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md rename docs/{adr-proposed => adr}/ADR-0060-algo-gqa-fused-attention-ahbm.md (96%) rename docs/{adr-proposed => adr}/ADR-0062-prog-tl-async-load.md (99%) rename docs/{adr-proposed => adr}/ADR-0063-prog-tl-scratch-scope.md (99%) rename src/kernbench/benches/{_gqa_decode_long.py => _gqa_attention_decode_long.py} (99%) rename src/kernbench/benches/{_gqa_decode_short.py => _gqa_attention_decode_short.py} (99%) create mode 100644 src/kernbench/benches/_gqa_attention_prefill_long.py rename src/kernbench/benches/{_gqa_prefill_short.py => _gqa_attention_prefill_short.py} (99%) delete mode 100644 src/kernbench/benches/_gqa_prefill_long.py create mode 100644 tests/attention/test_gqa_prefill_long_tile_ring.py diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md b/docs/adr-ko/ADR-0060-algo-gqa-fused-attention-ahbm.md similarity index 96% rename from docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md rename to docs/adr-ko/ADR-0060-algo-gqa-fused-attention-ahbm.md index 74ae334..70982c3 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md +++ b/docs/adr-ko/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -2,7 +2,7 @@ ## Status -Proposed +Accepted **Context model:** Llama3-70B. **Decision drivers:** agentic workload → 낮은 batch, 긴 context; @@ -748,6 +748,46 @@ KV를 필요로 하므로. (현재 제거된) baseline이 이미 ring fold를 구현; 본 ADR은 GQA 재사용, head-parallel 배치, causal step-skip, composite-hybrid inner tile(§3)을 추가. +#### 5.5.1 Tile-단위 Ring KV (수정안) + +위 §5.5 설명은 매 ring step마다 *전체 `(d_head, S_local)` KV 슬라이스* 를 +전송한다. `S_local`이 작을 때는 동작하지만, rank당 scratch가 `Kc + Vc` +(≈ `2·d·S_local·2` 바이트)에 종속되며, 이 값이 score-stack을 압도하여 +`S_local = 32K` 한참 이전에 1 MiB pool을 초과한다. + +구현된 ring은 **tile-단위(tile-granular)** 다 (ADR-0063 §A.2 + 본 ADR): +각 ring step은 전체 슬라이스가 아닌 `(d_head, TILE_S_KV)` K tile 1개와 +그에 대응하는 V tile 1개만 전송한다. 커널 루프는 **중첩(nested)** 구조 — +`for t in range(n_tiles): for k in range(C): ...` — 다음 tile이 시작되기 +전에 현재 tile이 ring의 `C`개 위치를 모두 통과한다. Rank당 persistent +scratch는 `(m, ℓ, O)` (≈ 1 KB) 로 축소되고, tile당 in-scope scratch는 +`S_local`에 무관하게 `TILE_S_KV`로 제한된다. + +Send 횟수 영향 (CUBE당, 방향당): + +- **슬라이스-단위 ring** (위 §5.5 baseline): `2·(C−1)` — + `C−1` ring step × 2 handle (K, V). +- **Tile-단위 ring** (구현): `2·n_tiles·(C−1)` — + `n_tiles · (C−1)` send/recv pair × 2 handle. + +총 IPCQ 바이트 수는 동일하다 (같은 데이터가 더 잘게 쪼개져 순환). IPCQ +command 수는 `n_tiles` 배로 증가. `n_tiles = 1` (작은 `S_local`) 에서는 +두 공식이 일치하므로 기존 `test_prefill_ring_c_*` (`S_kv ∈ {16, 32}`) +테스트들은 원래 count를 그대로 만족한다. + +루프 중첩 순서가 IPCQ buffer depth에 중요하다: **외부 tile, 내부 ring +step** 구조에서는 한 tile이 sender가 다음 tile을 만들기 전에 다음 CUBE가 +소비하므로 in-flight depth가 방향당 1로 유지된다. 반대 중첩 순서 (외부 +ring step, 내부 tile) 는 `k = 0`에서 한 CUBE가 `recv` 없이 `2·n_tiles`개의 +send를 누적시켜 IPCQ slot pool을 초과하면 deadlock된다. + +구현 위치는 `src/kernbench/benches/_gqa_attention_prefill_long.py` +(부트스트랩은 `(t=0, k=0)` 에서 peel off; 나머지 iteration은 +`tl.scratch_scope` 으로 감싸고 `tl.copy_to`로 `(m, ℓ, O)` 를 persist). +검증: `tests/attention/test_gqa_prefill_long_tile_ring.py` +(`S_kv = 128K` ceiling-lift, tile-단위 ipcq_copy count, CUBE별 출력 쓰기 +횟수 regression guard). + ### 5.6 Decode CPU-pipelining 변형 (opt1 / opt3 / opt2) 세 decode 변형은 **TL;DR에 풀 코드**로 있다(Kernel 1): **opt1** 현재 diff --git a/docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md b/docs/adr-ko/ADR-0062-prog-tl-async-load.md similarity index 99% rename from docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md rename to docs/adr-ko/ADR-0062-prog-tl-async-load.md index 83d12f2..e6463a1 100644 --- a/docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md +++ b/docs/adr-ko/ADR-0062-prog-tl-async-load.md @@ -2,7 +2,7 @@ ## Status -Proposed +Accepted > **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context > 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본 diff --git a/docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md b/docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md new file mode 100644 index 0000000..1ebf342 --- /dev/null +++ b/docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md @@ -0,0 +1,222 @@ +# ADR-0063: `tl.scratch_scope` — long-context를 위한 tile 단위 scratch 재활용 + +## Status + +Accepted + +> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Long-context attention +> 은 수많은 K/V tile을 sweep하며, 매 tile의 중간 결과 +> (`scores`, `P`, `exp`, partial `O`, …) 는 현재 kernel 호출 동안 회수되지 +> 않는 fresh scratch를 새로 할당한다. PE당 1 MiB scratch 예산은 실제로 +> 의미 있는 context length에 도달하기 훨씬 전에 소진된다. + +## Context + +### Bump allocator + +`TLContext` 는 모든 math/compute output handle을 PE-local scratch pool +에서 **linear bump cursor** 로 할당한다 +(`src/kernbench/triton_emu/tl_context.py`; `_scratch_alloc`): + +```python +def _scratch_alloc(self, nbytes): + aligned = (nbytes + 15) & ~15 + addr = self._scratch_base + self._scratch_cursor + self._scratch_cursor += aligned + if self._scratch_cursor > self._scratch_size: # default 1 << 20 = 1 MiB + raise RuntimeError("TLContext scratch overflow: ...") + return addr +``` + +Docstring은 cursor가 **"매 kernel 호출마다 reset된다"** 고 명시한다 — +즉 kernel entry에서만 reset, kernel body 내에서는 절대 reset되지 않는다. +모든 `tl.dot`, `tl.softmax`, `tl.exp`, `tl.sum`, `a - b`, `a * b`, … 는 +fresh slice를 가져가고 kernel이 return될 때까지 아무것도 회수되지 않는다. + +### 왜 GQA kernel에서 문제가 되는가 + +현재 milestone bench는 이 한계 **때문에** `S_q = S_kv_per_rank = 16` 으로 +명시적으로 고정해 두었다 +(`tests/attention/test_milestone_gqa_llama70b.py:123-148`): + +> "S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the +> simulator's 1 MB per-PE TCM kernel scratch is not exhausted by the +> bump-allocated handle outputs of softmax/exp/dot/sum chains over +> n_ranks ring steps." + +`n_tiles`개의 tile sweep에 대해 FlashAttention은 O(`n_tiles` × +tile당 중간값) 만큼 할당한다. 의미 있는 context에서는 overflow된다. +Flash attention의 *수학* 자체는 **O(1)** 의 live scratch만 필요로 한다 — +running `(m, l, O)` + 현재 tile의 working set — 매 tile의 임시값은 +running state에 fold된 직후 dead가 되기 때문이다. 다만 allocator가 그 +사실을 모를 뿐이다. + +## Decision + +Kernel이 할당 영역을 **재활용 가능(reclaimable)** 으로 표시할 수 있는 +**scratch scope** 을 추가한다. 이를 통해 tile당 임시값은 재활용되고, +live running state(및 in-flight prefetch buffer)는 보존된다. + +### D1. `tl` 표면 — context manager + +```python +with tl.scratch_scope(): + s = tl.dot(q, k_t) # `with` 내부의 모든 handle은 + p = tl.softmax(s) # 종료 시 rewind되는 영역을 공유 + o_j = tl.dot(p, v) + # ... o_j를 외부 영역에 사는 running (m,l,O) 에 fold ... +# __exit__: cursor가 __enter__ 시점 값으로 rewind +``` + +Semantics: +- `__enter__` 는 현재 `_scratch_cursor` 를 save-point로 기록. +- `__exit__` 는 cursor를 save-point로 복원, 내부에서 할당된 모든 것을 + free. +- Scope **외부**에서 할당된 handle (running `m,l,O`, ADR-0062의 + `LoadFuture` 가 보유한 prefetch buffer) 은 주소를 유지 — save-point + 이전에 또는 enclosing scope에서 할당되었기 때문. + +### D2. 안전 계약 + +Scope 내부에서 할당된 handle은 scope exit 이후 **읽으면 안 된다** — +그 바이트는 다음 scope의 할당으로 덮어쓰여질 수 있다. Flash loop은 +이를 자연스럽게 지킨다: tile iteration 사이에 살아남는 값은 running +accumulator 뿐이고, 이들은 per-tile scope **외부**에서 할당되며 +*내부*의 op들이 외부 주소로 쓰면서 갱신된다(merge가 새 running state +를 쓰는 방식 — D3 참조). + +### D3. Running accumulator와의 상호작용 + +Online-softmax merge는 이전 running `(m, l, O)` 와 현재 tile의 +`(m_j, l_j, O_j)` 를 읽어 새 running 값을 만든다. 새 running 값을 +재활용 영역 바깥에 두기 위해, merge는 결과를 loop 시작 전 1회 할당된 +**stable scratch** (per-tile scope과는 별개의 작은 고정 "running-state" +arena) 로 쓴다. 구체적으로 kernel은 두 개의 arena를 둔다: + +- **persistent arena** (1회 할당): `m, l, O` (merge용 double-buffer가 + 필요한 경우 그것까지). +- **scoped arena** (매 tile rewind): `scores, P, exp, O_j, scale_*`. + +이는 실제 flash-attention SRAM budgeting의 모습과 같다: 작은 persistent +accumulator 영역 + 재활용되는 tile working set. + +#### D3.1 Persistent-arena 쓰기 mechanism: `tl.copy_to(dst, src)` + +Merge op (`tl.maximum`, `tl.exp`, binary `*` / `+`) 모두 D1의 bump +cursor로부터 할당하는 `_make_compute_out(...)` 을 호출한다. 따라서 +`scratch_scope` 내부에서 그 결과 handle들은 scope **내부**에 살고 +`__exit__` 시점에 사라진다. D3의 two-arena 분리를 실현하려면 kernel은 +**scoped 결과의 바이트를 persistent 주소로 쓰는** 방법이 필요하다. + +이 gap을 메우는 primitive: + +```python +def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None: + """``src`` 의 바이트를 ``dst`` 의 주소로 복사한다 (양쪽 모두 TCM). + + Shape, dtype 일치 필요. ``dst`` 는 통상 어떤 활성 ``scratch_scope`` 도 + 바깥에서 할당된 handle (persistent arena); ``src`` 는 scope ``__exit__`` + 이후에도 바이트가 살아 있어야 하는 scoped handle. + """ +``` + +`tl.store` (HBM 측 바이트 복사) 와 대칭이지만, running-state writeback이 +op_log를 불필요한 DMA로 오염시키지 않도록 TCM 전용으로 유지한다. + +**Mechanics:** +- 신규 `CopyCmd(src, dst, nbytes, data_op=True)` command. +- op_log: `op_kind="math"`, `op_name="copy"` — vector engine에서 실행. +- Latency: `pe_math._compute_ns(prod(shape))` — HBM 전송이 아니라 + on-chip register writeback을 모델링. +- Emit-time 검증: `dst.shape == src.shape`, `dst.dtype == src.dtype`, + `dst.space == "tcm"`, `src.space == "tcm"`. 작성자 오류는 Phase 2 + data execution 깊숙한 곳이 아니라 Phase 1에서 노출. + +**호출 패턴:** + +```python +m, l, O = init_running(...) # persistent (scope 외부) +for j in range(n_tiles): + with tl.scratch_scope(): + ... # tile 작업 (재활용) + m_new = tl.maximum(m, mj) # scoped scratch + l_new = l * scale_old + l_step * scale_step + O_new = O * scale_old + O_step * scale_step + + tl.copy_to(m, m_new) # ← 새 running state를 persist + tl.copy_to(l, l_new) + tl.copy_to(O, O_new) + # exit: scoped m_new/l_new/O_new 는 사라지고, 그 바이트는 + # persistent m/l/O 에 살아있다 +``` + +`copy_to` 는 `__exit__` **이전** 에 실행되므로 `src` (scoped) 의 read는 +유효하며, exit 이후에는 `dst` (persistent) 만 read되어 D2의 +"exit 이후에는 읽지 말 것" 안전 계약을 만족한다. + +**왜 모든 math op에 `dst=` kwarg를 추가하지 않고 전용 primitive 인가** +(검토 후 기각): `tl.maximum`, `tl.exp`, `_binary_math`, `_unary_math`, +`_reduction` 에 `dst=` 를 추가하는 것은 5개 op-family 에 걸쳐 ~25 LOC +이며, "호출은 fresh handle을 return한다" 라는 일관된 패턴을 깬다. +`tl.copy_to` 는 단일 primitive, 단일 command, 단일 executor handler — +같은 효과에 비해 최소한의 surface area. + +### D4. Nesting + +Scope는 nest 가능 (save-point의 stack). Inner scope exit은 inner +save-point로 rewind, outer exit은 더 멀리 rewind한다. Prefill의 +"per-query-block" 외부 scope을 둘러싼 "per-KV-tile" 내부 scope을 +지원한다. + +## Alternatives + +### A1. Temporary를 HBM에 round-trip + +중간값을 HBM에 store하고 reload하여 TCM scratch를 "free" 한다. 기각: +TCM-resident streaming kernel을 HBM-bandwidth-bound로 만들어 목적에 +정반대이며, op_log를 불필요한 DMA로 오염시킨다. + +### A2. Tiled `tl.composite` (scheduler-관리 scratch) + +`tl.composite` 는 PE_SCHEDULER 내부에서 tile당 scratch를 자동 재활용 +한다. ADR-0062 A1과 마찬가지로 매력적이지만 flash-capable composite +kind (두 GEMM + carry되는 `(m,l,O)`) 가 필요하므로 훨씬 큰 변경에 +종속된다. `scratch_scope` 는 작고 일반적인 primitive로 greenlet kernel +에 동일한 memory 거동을 제공한다. 두 방식은 공존 가능. + +### A3. Scratch 예산 증가 + +`pe_tcm.kernel_scratch_mb` 를 증가. 해결책으로는 기각: O(`n_tiles`) +scratch leak이 여전히 존재하면서 context-length ceiling만 선형적으로 +밀어내며, hardware 실제 모습 (실 TCM은 작고, flash attention의 핵심은 +O(1) working set) 을 잘못 표현한다. 거친 조절 knob으로만 유용, 재활용의 +대체재는 아니다. + +## Consequences + +### Positive +- 인위적인 `S = 16` validation-scale ceiling 제거; 시간/데이터 mode + 양쪽에서 의미 있는 context length 가능. +- Flash attention의 O(1) working-set 특성을 충실히 모델링. +- 작고 일반적인 primitive (모든 tiled kernel이 이점). + +### Negative +- Use-after-scope 버그는 stale 바이트를 조용히 읽는다. D2 계약, attention + helper 안에서 scope discipline 공유, (선택적으로) rewound 영역을 + poison시키는 debug build로 완화. +- ADR-0062와 조정 필요: prefetch buffer는 tile iteration 사이에 살아 + 있어야 하므로 scoped arena가 아니라 persistent arena에 속한다. + +## Test Requirements + +1. **Recycling**: `tl.scratch_scope()` 내부의 `N` tile loop에서 peak + `_scratch_cursor` 가 `N`과 무관하게 한 tile의 footprint로 bound됨 + (오늘날은 선형 증가 후 overflow). +2. **Correctness**: scope이 있는 flash sweep이 scope이 없는 동일 sweep + 과 (재활용 없이도 수렴하는 작은 `N` 에서) 동일한 `O` 를 산출 + (Phase 2). +3. **Long context**: scope 없이 1 MiB를 초과하는 `N` 의 sweep이 완료 + (오늘날 `S=16` cap이 회피하는 그 실패). +4. **Persistent-vs-scoped isolation**: scope 외부에서 할당된 running + `(m,l,O)` 가 `__exit__` 이후에도 올바른 값을 유지. +5. **Nesting**: 중첩된 scope이 올바른 save-point로 rewind. diff --git a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md b/docs/adr/ADR-0060-algo-gqa-fused-attention-ahbm.md similarity index 96% rename from docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md rename to docs/adr/ADR-0060-algo-gqa-fused-attention-ahbm.md index bf0db71..600b27e 100644 --- a/docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md +++ b/docs/adr/ADR-0060-algo-gqa-fused-attention-ahbm.md @@ -2,7 +2,7 @@ ## Status -Proposed +Accepted **Context model:** Llama3-70B. **Decision drivers:** agentic workload → low batch, long context; @@ -823,6 +823,49 @@ The (now-removed) baseline already implemented the ring fold; this ADR adds GQA reuse, the head-parallel placement, causal step-skip, and the composite-hybrid inner tile (§3). +#### 5.5.1 Tile-granular Ring KV (amendment) + +The §5.5 description above ships *full `(d_head, S_local)` KV slices* +per ring step. That works at small `S_local` but ties per-rank scratch +to `Kc + Vc` (≈ `2·d·S_local·2` bytes), which dominates the score-stack +and overflows the 1 MiB pool well below `S_local = 32K`. + +The implemented ring is **tile-granular** (ADR-0063 §A.2 + this ADR): +each ring step transmits an `(d_head, TILE_S_KV)` K tile and its V +counterpart, not a full slice. The kernel loop is **nested** — +`for t in range(n_tiles): for k in range(C): ...` — so each tile +propagates through all `C` ring positions before the next tile starts. +Per-rank persistent scratch shrinks to `(m, ℓ, O)` (≈ 1 KB); per-tile +in-scope scratch is bounded by `TILE_S_KV` regardless of `S_local`. + +Send-count consequence (per CUBE per direction): + +- **Slice-granular ring** (the §5.5 baseline above): `2·(C−1)` — + `C−1` ring steps × 2 handles (K, V). +- **Tile-granular ring** (the implementation): `2·n_tiles·(C−1)` — + `n_tiles · (C−1)` send/recv pairs × 2 handles. + +Total IPCQ bytes are unchanged (same data circulates, just chunked). +The number of IPCQ commands grows by `n_tiles`. At `n_tiles = 1` +(small `S_local`) the two formulas coincide, so all existing +`test_prefill_ring_c_*` tests at `S_kv ∈ {16, 32}` continue to assert +the original count. + +Loop nesting matters for IPCQ buffer depth: with **outer tile, +inner ring step**, each tile is consumed by the next CUBE before the +sender produces another, so in-flight depth stays at 1 per direction. +The opposite nesting (outer ring step, inner tile) piles `2·n_tiles` +sends in flight per CUBE at `k = 0` before any `recv` drains them, which +deadlocks past the IPCQ slot pool. + +Implementation lives in +`src/kernbench/benches/_gqa_attention_prefill_long.py` (bootstrap +peeled off at `(t=0, k=0)`; remaining iterations wrap in +`tl.scratch_scope` and persist `(m, ℓ, O)` via `tl.copy_to`). +Verification: `tests/attention/test_gqa_prefill_long_tile_ring.py` +(ceiling-lift at `S_kv = 128K`, tile-granular ipcq_copy count, +per-CUBE output write count regression guard). + ### 5.6 Decode CPU-pipelining variants (opt1 / opt3 / opt2) The three decode variants are shown **in full in the TL;DR** (Kernel 1): @@ -1249,7 +1292,8 @@ separate study"). This subsection pins the two open numbers. Level-1 (inter-CUBE) collapses to a no-op. The output stays distributed per CUBE (each CUBE writes its owned heads' `O` slice). - **Recommend:** ship as a separate kernel file (`_gqa_decode_short.py` / - `_gqa_prefill_short.py`) and a separate bench dispatcher that picks + **Recommend:** ship as a separate kernel file + (`_gqa_attention_decode_short.py` / `_gqa_attention_prefill_short.py`) + and a separate bench dispatcher that picks short vs long by `S_kv ⋚ 256K`. Keep `kv_per_cube` as a kernel arg so the topology cost trade-off can be measured per workload. diff --git a/docs/adr-proposed/ADR-0062-prog-tl-async-load.md b/docs/adr/ADR-0062-prog-tl-async-load.md similarity index 99% rename from docs/adr-proposed/ADR-0062-prog-tl-async-load.md rename to docs/adr/ADR-0062-prog-tl-async-load.md index b200583..b3a111d 100644 --- a/docs/adr-proposed/ADR-0062-prog-tl-async-load.md +++ b/docs/adr/ADR-0062-prog-tl-async-load.md @@ -2,7 +2,7 @@ ## Status -Proposed +Accepted > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and > long-context attention are **KV-load-bound**; load/compute overlap is a diff --git a/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md b/docs/adr/ADR-0063-prog-tl-scratch-scope.md similarity index 99% rename from docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md rename to docs/adr/ADR-0063-prog-tl-scratch-scope.md index d1d31d3..1ee9ea6 100644 --- a/docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md +++ b/docs/adr/ADR-0063-prog-tl-scratch-scope.md @@ -2,7 +2,7 @@ ## Status -Proposed +Accepted > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Long-context > attention sweeps many K/V tiles; each tile's intermediates diff --git a/src/kernbench/benches/_gqa_decode_long.py b/src/kernbench/benches/_gqa_attention_decode_long.py similarity index 99% rename from src/kernbench/benches/_gqa_decode_long.py rename to src/kernbench/benches/_gqa_attention_decode_long.py index 25a73b0..5f3c9c3 100644 --- a/src/kernbench/benches/_gqa_decode_long.py +++ b/src/kernbench/benches/_gqa_attention_decode_long.py @@ -35,7 +35,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): return m_new, l_new, O_new -def gqa_decode_long_kernel( +def gqa_attention_decode_long_kernel( q_ptr: int, k_ptr: int, v_ptr: int, diff --git a/src/kernbench/benches/_gqa_decode_short.py b/src/kernbench/benches/_gqa_attention_decode_short.py similarity index 99% rename from src/kernbench/benches/_gqa_decode_short.py rename to src/kernbench/benches/_gqa_attention_decode_short.py index afd9b62..3e244c1 100644 --- a/src/kernbench/benches/_gqa_decode_short.py +++ b/src/kernbench/benches/_gqa_attention_decode_short.py @@ -44,7 +44,7 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): return m_new, l_new, O_new -def gqa_decode_short_kernel( +def gqa_attention_decode_short_kernel( q_ptr: int, k_ptr: int, v_ptr: int, diff --git a/src/kernbench/benches/_gqa_attention_prefill_long.py b/src/kernbench/benches/_gqa_attention_prefill_long.py new file mode 100644 index 0000000..b2cd21f --- /dev/null +++ b/src/kernbench/benches/_gqa_attention_prefill_long.py @@ -0,0 +1,151 @@ +"""GQA fused-attention prefill kernel — long context (ADR-0060 §5.5). + +Head-parallel: each CUBE owns one Q head and one KV slice. Over C ring +steps the KV blocks rotate around the C CUBEs (W-send / E-recv) so every +CUBE sees every block; the online-softmax merge folds each step into the +running ``(m, ℓ, O)``. No inter-CUBE reduce — each CUBE writes its own +head's output. + +The ring is **tile-granular** (ADR-0060 §5.5 + ADR-0063 §A.2): each +ring step transmits ``n_tiles`` tiles of size ``TILE_S_KV`` rather than +one full ``S_local`` slice. The kernel loop is nested over +``(ring_step, tile_idx)``, with the bootstrap at ``(k=0, t=0)`` +establishing the persistent ``(m, ℓ, O)`` and every subsequent +iteration folding its tile in inside a ``tl.scratch_scope``. Per-rank +persistent scratch is ``(m, ℓ, O)`` only (~1 KB); per-tile in-scope +scratch is bounded by ``TILE_S_KV`` regardless of ``S_local``. + +Topology / SFR: + - Requires ``configure_sfr_intercube_ring(ring_size=C)`` (1D ring of + C CUBEs with wrap at the CUBE level). + - Only PE 0 of each CUBE participates (head-parallel; intra-CUBE PE + parallelism is a separate phase). + +Layout caveats: + - GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``). + - K loaded as ``(d_head, tile_s)`` via byte-conserving reshape of + the deployed ``(tile_s, d_head)`` shard (ADR-0060 §3 + reshape-not-transpose caveat). + - No causal masking / step-skip; blocking ``tl.recv``. + - ``TILE_S_KV`` is assumed to divide ``S_local`` evenly; last-tile + padding is not modelled. +""" +from __future__ import annotations + + +TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width). + + +def _merge_running(m, l, O, m_step, l_step, O_step, *, tl): + """Online-softmax merge of two partial ``(m, ℓ, O)`` triples.""" + m_new = tl.maximum(m, m_step) + scale_old = tl.exp(m - m_new) + scale_step = tl.exp(m_step - m_new) + l_new = l * scale_old + l_step * scale_step + O_new = O * scale_old + O_step * scale_step + return m_new, l_new, O_new + + +def gqa_attention_prefill_long_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + d_head: int, + C: int, + *, + tl, +) -> None: + """Head-parallel prefill attention with tile-granular Ring KV (ADR-0060 §5.5). + + Tensor layout: + Q : (T_q, d_head) one head per CUBE; replicated. + K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns + (S_local, d_head). Tiles loaded as (d_head, tile_s) via + byte-conserving reshape. + V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns + (S_local, d_head). Tiles loaded as (tile_s, d_head). + O : (T_q * C, d_head) sharded cube_row_wise → each CUBE + writes its own (T_q, d_head) slice. NO reduce. + + Algorithm: nested loop over (ring_step k, tile_idx t). At k=0 each + CUBE loads its own block's tiles from HBM; at k > 0 it receives + tiles from its E neighbour. Tiles are forwarded W to the next + ring step. Each tile's partial is folded into the running + (m, ℓ, O) via online-softmax merge. + """ + pe_id = tl.program_id(axis=0) + # Head-parallel: only PE 0 of each CUBE participates. + if pe_id != 0: + return + + S_local = S_kv // C + n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV + KV_ROW_BYTES = d_head * 2 # f16 + Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16") + + # ── Bootstrap: (t=0, k=0) — load my own tile 0, establish (m, ℓ, O) ── + # Outside ``scratch_scope`` so the persistent running state survives + # subsequent ring iterations. + tile_s = min(TILE_S_KV, S_local) + K_t = tl.load(k_ptr, shape=(d_head, tile_s), dtype="f16") + V_t = tl.load(v_ptr, shape=(tile_s, d_head), dtype="f16") + if C > 1: + tl.send(dir="W", src=K_t) + tl.send(dir="W", src=V_t) + scores = tl.dot(Q, K_t) + m = tl.max(scores, axis=-1) + exp_scores = tl.exp(scores - m) + l = tl.sum(exp_scores, axis=-1) + O = tl.dot(exp_scores, V_t) + + # ── Nested loop: outer tile, inner ring step ── + # Each tile propagates all the way through the ring before the next + # tile starts; IPCQ in-flight depth stays at 1 per direction. + # + # Per outer ``t``: + # k=0 : load my own tile ``t`` from HBM + # k=1..C-1 : receive tile ``t`` of a peer's block from E + # (sent by my E neighbour during their previous k step) + # k None: - """Head-parallel prefill attention with Ring KV (C>1) — ADR-0060 §5.5. - - Tensor layout consumed by this kernel: - Q : (T_q, d_head) one head per CUBE; replicated. - K : (S_kv, d_head) sharded cube_row_wise → each CUBE owns - (S_kv/C, d_head); kernel loads as (d_head, S_local) via - byte-conserving reshape (reshape-not-transpose caveat). - V : (S_kv, d_head) sharded cube_row_wise → each CUBE owns - (S_local, d_head). - O : (T_q * C, d_head) sharded cube_row_wise → each CUBE - writes its own (T_q, d_head) slice. NO reduce. - - Algorithm: each CUBE computes a local partial against its current - KV block, then over C-1 ring steps the K and V blocks rotate W - while online-softmax merges each step into running (m, ℓ, O). - """ - pe_id = tl.program_id(axis=0) - # Head-parallel: only PE 0 of each CUBE participates. - if pe_id != 0: - return - - S_local = S_kv // C - Q = tl.load(q_ptr, shape=(T_q, d_head), dtype="f16") - Kc = tl.load(k_ptr, shape=(d_head, S_local), dtype="f16") - Vc = tl.load(v_ptr, shape=(S_local, d_head), dtype="f16") - - # ── Local attention: initial partial against own KV block ── - # Establishes the persistent (m, ℓ, O) running state. - scores = tl.dot(Q, Kc) - m = tl.max(scores, axis=-1) - exp_scores = tl.exp(scores - m) - l = tl.sum(exp_scores, axis=-1) - O = tl.dot(exp_scores, Vc) - - # ── Communication: Ring KV rotation + online-softmax merge ── - # Each step sends K, V to W and receives from E. Per-step - # intermediates are scope-recycled; the merged (m, ℓ, O) is - # persisted via tl.copy_to. Triton port: drop the scope and - # replace each copy_to with a Python rebind. - for _ in range(1, C): - tl.send(dir="W", src=Kc) - tl.send(dir="W", src=Vc) - Kc = tl.recv(dir="E", shape=(d_head, S_local), dtype="f16") - Vc = tl.recv(dir="E", shape=(S_local, d_head), dtype="f16") - - with tl.scratch_scope(): - scores = tl.dot(Q, Kc) - m_step = tl.max(scores, axis=-1) - exp_scores = tl.exp(scores - m_step) - l_step = tl.sum(exp_scores, axis=-1) - O_step = tl.dot(exp_scores, Vc) - - m_new = tl.maximum(m, m_step) - scale_old = tl.exp(m - m_new) - scale_step = tl.exp(m_step - m_new) - l_new = l * scale_old + l_step * scale_step - O_new = O * scale_old + O_step * scale_step - - tl.copy_to(m, m_new) - tl.copy_to(l, l_new) - tl.copy_to(O, O_new) - - # ── Final normalise + store (each CUBE writes its own head) ── - O_final = O / l - tl.store(o_ptr, O_final) diff --git a/src/kernbench/benches/milestone_gqa_headline.py b/src/kernbench/benches/milestone_gqa_headline.py index 461deee..94e55d1 100644 --- a/src/kernbench/benches/milestone_gqa_headline.py +++ b/src/kernbench/benches/milestone_gqa_headline.py @@ -28,8 +28,8 @@ import json import os from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel -from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel from kernbench.benches.registry import bench from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import ( @@ -96,7 +96,7 @@ def _run_prefill_panel(ctx, *, panel: str, C: int, S_kv: int) -> None: o = ctx.empty((_T_Q_PREFILL * C, _D_HEAD), dtype=_DTYPE, dp=dp_o, name=f"{panel}_o") ctx.launch( - panel, gqa_prefill_long_kernel, + panel, gqa_attention_prefill_long_kernel, q, k, v, o, _T_Q_PREFILL, S_kv, _D_HEAD, C, _auto_dim_remap=False, @@ -118,7 +118,7 @@ def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None: o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD), dtype=_DTYPE, dp=dp_full, name=f"{panel}_o") ctx.launch( - panel, gqa_decode_long_kernel, + panel, gqa_attention_decode_long_kernel, q, k, v, o, _T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_decode.py b/tests/attention/test_gqa_decode.py index 5bbcf51..5f8895f 100644 --- a/tests/attention/test_gqa_decode.py +++ b/tests/attention/test_gqa_decode.py @@ -23,7 +23,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 (Phase 2 deliverable) +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 (Phase 2 deliverable) from kernbench.policy.placement.dp import DPPolicy from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device @@ -72,7 +72,7 @@ def _run_decode_p1(*, h_q: int, h_kv: int): dtype=DTYPE, dp=dp, name=f"o_h{h_q}_kv{h_kv}") ctx.launch( f"gqa_decode_p1_h{h_q}_kv{h_kv}", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, T_Q, S_KV, h_q, h_kv, D_HEAD, 1, 1, # C=1, P=1 (no SP, degenerate) diff --git a/tests/attention/test_gqa_decode_mc.py b/tests/attention/test_gqa_decode_mc.py index c2d16f6..ce2f03d 100644 --- a/tests/attention/test_gqa_decode_mc.py +++ b/tests/attention/test_gqa_decode_mc.py @@ -31,7 +31,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -82,7 +82,7 @@ def _run_decode_mc(*, h_q: int, h_kv: int, C: int, P: int, S_kv: int): name=f"o_h{h_q}_kv{h_kv}_c{C}_p{P}") ctx.launch( f"gqa_decode_mc_h{h_q}_kv{h_kv}_c{C}_p{P}", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, T_Q, S_kv, h_q, h_kv, D_HEAD, C, P, diff --git a/tests/attention/test_gqa_decode_sp.py b/tests/attention/test_gqa_decode_sp.py index 21e4e26..5ecbafd 100644 --- a/tests/attention/test_gqa_decode_sp.py +++ b/tests/attention/test_gqa_decode_sp.py @@ -22,7 +22,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -70,7 +70,7 @@ def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): dtype=DTYPE, dp=dp_full, name=f"o_h{h_q}_kv{h_kv}_p{P}") ctx.launch( f"gqa_decode_sp_h{h_q}_kv{h_kv}_p{P}", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, T_Q, S_kv, h_q, h_kv, D_HEAD, 1, P, # C=1, P=P (single-CUBE SP) diff --git a/tests/attention/test_gqa_long_context.py b/tests/attention/test_gqa_long_context.py index f3baf6e..d5354c7 100644 --- a/tests/attention/test_gqa_long_context.py +++ b/tests/attention/test_gqa_long_context.py @@ -23,8 +23,8 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import ( configure_sfr_intercube_multisip, @@ -84,7 +84,7 @@ def test_decode_long_context_32k_completes(): dtype=DTYPE, dp=dp_full, name="q_long_dec_2") ctx.launch( "gqa_decode_long_32k", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, 1, S_kv, 8, 1, D_HEAD, 1, P, @@ -141,7 +141,7 @@ def test_prefill_long_context_completes_after_scope_discipline(): dtype=DTYPE, dp=dp_o, name="o_long_pre") ctx.launch( "gqa_prefill_long_64k", - gqa_prefill_long_kernel, + gqa_attention_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_prefill.py b/tests/attention/test_gqa_prefill.py index ac9ca2e..090b479 100644 --- a/tests/attention/test_gqa_prefill.py +++ b/tests/attention/test_gqa_prefill.py @@ -1,6 +1,6 @@ """Phase 1 spec test for P6a GQA prefill kernel (head-parallel, C=1 baseline). -P6a introduces ``_gqa_prefill_long.py`` with the head-parallel structure (one +P6a introduces ``_gqa_attention_prefill_long.py`` with the head-parallel structure (one Q head per CUBE, per-CUBE distributed output, no reduce). C=1 is the degenerate case — no Ring KV, no IPCQ traffic. Validates kernel structure and T_q > 1 attention. @@ -21,7 +21,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 from kernbench.policy.placement.dp import DPPolicy from kernbench.runtime_api.bench_runner import run_bench from kernbench.runtime_api.types import resolve_device @@ -62,7 +62,7 @@ def _run_prefill(*, T_q: int, S_kv: int, C: int = 1): name=f"o_t{T_q}_c{C}") ctx.launch( f"gqa_prefill_p6a_t{T_q}_s{S_kv}_c{C}", - gqa_prefill_long_kernel, + gqa_attention_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_prefill_long_tile_ring.py b/tests/attention/test_gqa_prefill_long_tile_ring.py new file mode 100644 index 0000000..93833a6 --- /dev/null +++ b/tests/attention/test_gqa_prefill_long_tile_ring.py @@ -0,0 +1,162 @@ +"""Phase 1 spec test for P3c: tile-granular Ring KV in prefill_long +(ADR-0060 §5.5 amendment + ADR-0063 §A.2). + +Today's prefill_long ring sends and receives full ``(d_head, S_local)`` +KV slices per step. Step 0's local-attention intermediates also live +outside any ``scratch_scope`` (they're loaded as full-slice ``Kc``, +``Vc`` and feed ``scores``, ``exp_scores`` as persistent allocations). +At larger ``S_local`` both the step-0 leak and the ring step's +in-scope intermediates grow linearly with ``S_local``, and at +``S_local = 32K`` (S_kv=128K, C=4, T_q=4) the peak overflows the +1 MiB pool. + +P3c converts the ring to **tile-granular**: a nested loop +``for k in range(C): for t in range(n_tiles): ...`` where each +iteration sends/recvs one ``(d_head, TILE_S_KV)`` tile (and its V +counterpart). The persistent state shrinks to ``(m, ℓ, O)`` only +(~1 KB); per-tile in-scope scratch is bounded by +``TILE_S_KV`` regardless of ``S_local``. Ceiling lifted. + +Trade-off: the per-CUBE send count grows from ``2·(C-1)`` to +``2·n_tiles·(C-1)``. At ``n_tiles=1`` (small ``S_local``) the count is +unchanged, so the existing ``test_prefill_ring_c_*`` tests at +``S_kv ∈ {16, 32}`` still pass. + +Phase 1 (this commit): tests only — production code lands in Phase 2. +T1 fails today with a ``TLContext scratch overflow``; T2 fails today +with the slice-granular ipcq_copy count; T5 passes today and is a +regression guard for the per-CUBE output write count. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_ring +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml" + +D_HEAD = 64 +DTYPE = "f16" + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _count(op_log, name: str) -> int: + return sum(1 for r in op_log if r.op_name == name) + + +def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): + """Head-parallel prefill with Ring KV across C CUBEs.""" + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_ring( + ctx.engine, ctx.spec, _ccl_cfg(), ring_size=C, + ) + dp_q = DPPolicy(cube="replicate", pe="replicate", + num_cubes=C, num_pes=1) + dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + dp_o = DPPolicy(cube="row_wise" if C > 1 else "replicate", + pe="replicate", num_cubes=C, num_pes=1) + q = ctx.zeros((T_q, D_HEAD), dtype=DTYPE, dp=dp_q, + name=f"q_tlr_t{T_q}_c{C}_s{S_kv}") + k = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"k_tlr_t{T_q}_c{C}_s{S_kv}") + v = ctx.zeros((S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv, + name=f"v_tlr_t{T_q}_c{C}_s{S_kv}") + o = ctx.empty((T_q * C, D_HEAD), dtype=DTYPE, dp=dp_o, + name=f"o_tlr_t{T_q}_c{C}_s{S_kv}") + ctx.launch( + f"gqa_prefill_long_tile_ring_t{T_q}_c{C}_s{S_kv}", + gqa_attention_prefill_long_kernel, + q, k, v, o, + T_q, S_kv, D_HEAD, C, + _auto_dim_remap=False, + ) + + return run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), + engine_factory=_engine_factory, + ) + + +# ── T1: 128K ceiling lift ──────────────────────────────────────────── + + +def test_prefill_long_context_128k_completes(): + """ADR-0063 §A.2 headline ceiling lift. At S_kv=128K with C=4, + S_local=32K. Today: step-0 score-stack (~768 KB persistent) + ring + scope (~768 KB) → 1.5 MB peak → TLContext scratch overflow. After + P3c the persistent state is just ``(m, ℓ, O)`` (≈ 1 KB) and the + per-tile in-scope scratch is bounded by TILE_S_KV. + """ + result = _run_prefill_ring(T_q=4, S_kv=131_072, C=4) + assert result.completion.ok, ( + f"prefill_long at S_kv=128K must complete after tile-granular " + f"ring lands; got {result.completion}" + ) + + +# ── T2: tile-granular ipcq_copy count ──────────────────────────────── + + +def test_prefill_long_tile_granular_ipcq_count(): + """ADR-0060 §5.5 amendment: with tile-granular sends, the per-CUBE + send count grows from ``2·(C-1)`` to ``2·n_tiles·(C-1)``. Aggregated + across all C CUBEs the total ipcq_copy becomes + ``2·n_tiles·(C-1)·C``. + + Config: T_q=4, S_kv=4096, C=2 → S_local=2048, n_tiles=2 (with + TILE_S_KV=1024). Today: slice-granular total = + ``(C-1)·2·C = 4``. After P3c: tile-granular total = + ``(C-1)·n_tiles·2·C = 8``. + """ + C = 2 + n_tiles = 2 # S_local=2048 / TILE_S_KV=1024 + result = _run_prefill_ring(T_q=4, S_kv=4096, C=C) + assert result.completion.ok, ( + f"prefill_long multi-tile ring must complete; got {result.completion}" + ) + n_copy = _count(result.engine.op_log, "ipcq_copy") + expected = (C - 1) * n_tiles * 2 * C + assert n_copy == expected, ( + f"tile-granular ring: expected {expected} ipcq_copy " + f"((C-1)·n_tiles·2·C = {C - 1}·{n_tiles}·2·{C}); got {n_copy}" + ) + + +# ── T5: per-CUBE distributed output unchanged (regression guard) ───── + + +def test_prefill_long_tile_ring_dma_write_count(): + """ADR-0060 §5.5: per-CUBE distributed output must hold under the + tile-granular ring rewrite. Each CUBE still writes its own head's + output (no inter-CUBE reduce); dma_write_count == C. + + Today passes; must continue to pass after P3c. + """ + C = 4 + result = _run_prefill_ring(T_q=4, S_kv=4096, C=C) + assert result.completion.ok + n_writes = _count(result.engine.op_log, "dma_write") + assert n_writes == C, ( + f"per-CUBE distributed output: expected {C} dma_writes (one per " + f"CUBE); got {n_writes}" + ) diff --git a/tests/attention/test_gqa_prefill_ring.py b/tests/attention/test_gqa_prefill_ring.py index 16be9c6..1ce602c 100644 --- a/tests/attention/test_gqa_prefill_ring.py +++ b/tests/attention/test_gqa_prefill_ring.py @@ -34,7 +34,7 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_ring # noqa: F401 (Phase 2) from kernbench.policy.placement.dp import DPPolicy @@ -89,7 +89,7 @@ def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): name=f"o_t{T_q}_c{C}_ring") ctx.launch( f"gqa_prefill_ring_t{T_q}_s{S_kv}_c{C}", - gqa_prefill_long_kernel, + gqa_attention_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_scoped_writeback.py b/tests/attention/test_gqa_scoped_writeback.py index 825cf2f..d6541d0 100644 --- a/tests/attention/test_gqa_scoped_writeback.py +++ b/tests/attention/test_gqa_scoped_writeback.py @@ -23,8 +23,8 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_prefill_long import gqa_prefill_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_prefill_long import gqa_attention_prefill_long_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import ( configure_sfr_intercube_multisip, @@ -79,7 +79,7 @@ def _run_decode_sp(*, h_q: int, h_kv: int, P: int, S_kv: int): dtype=DTYPE, dp=dp_full, name=f"o_sc_{P}") ctx.launch( f"gqa_decode_scoped_{P}", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, 1, S_kv, h_q, h_kv, D_HEAD, 1, P, @@ -138,7 +138,7 @@ def _run_prefill_ring(*, T_q: int, S_kv: int, C: int): dtype=DTYPE, dp=dp_o, name=f"o_ring_{C}") ctx.launch( f"gqa_prefill_scoped_{C}", - gqa_prefill_long_kernel, + gqa_attention_prefill_long_kernel, q, k, v, o, T_q, S_kv, D_HEAD, C, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_short_context.py b/tests/attention/test_gqa_short_context.py index ca99277..b1191fd 100644 --- a/tests/attention/test_gqa_short_context.py +++ b/tests/attention/test_gqa_short_context.py @@ -24,8 +24,8 @@ After chain reduce-to-group-root, the group's root PE writes its owned head's output. No inter-CUBE reduce. Phase 1 (this commit): tests only — production code lands in Phase 2. -All tests fail because the short kernels (``_gqa_decode_short.py`` and -``_gqa_prefill_short.py``) do not exist yet. +All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and +``_gqa_attention_prefill_short.py``) do not exist yet. """ from __future__ import annotations @@ -76,7 +76,7 @@ def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int, O: replicated; each group root writes the full byte-conserving (h_q·T_q, D_HEAD) result. """ - from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # Phase 2 + from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2 topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -98,7 +98,7 @@ def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int, dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}") ctx.launch( f"gqa_decode_short_{kv_per_cube}_{C}", - gqa_decode_short_kernel, + gqa_attention_decode_short_kernel, q, k, v, o, 1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube, _auto_dim_remap=False, @@ -196,7 +196,7 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int, Layout: same head-stacked K/V scheme as decode short, with Q/O replicated and byte-conserving reshape inside the kernel. """ - from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel # Phase 2 + from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2 topo = resolve_topology(str(TOPOLOGY_DEFAULT)) @@ -218,7 +218,7 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int, dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}") ctx.launch( f"gqa_prefill_short_{kv_per_cube}_{C}", - gqa_prefill_short_kernel, + gqa_attention_prefill_short_kernel, q, k, v, o, T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube, _auto_dim_remap=False, diff --git a/tests/attention/test_gqa_tile_sweep.py b/tests/attention/test_gqa_tile_sweep.py index 6853a9e..322cd49 100644 --- a/tests/attention/test_gqa_tile_sweep.py +++ b/tests/attention/test_gqa_tile_sweep.py @@ -1,8 +1,8 @@ """Phase 1 spec test for P3b: S_kv tile sweep in the GQA kernels (ADR-0063 §A.2 + ADR-0060 §B "long/short context split"). -The local one-shot partial in three kernels — ``_gqa_decode_long.py``, -``_gqa_decode_short.py``, ``_gqa_prefill_short.py`` — loads its rank's +The local one-shot partial in three kernels — ``_gqa_attention_decode_long.py``, +``_gqa_attention_decode_short.py``, ``_gqa_attention_prefill_short.py`` — loads its rank's entire ``(d_head, S_local)`` / ``(S_local, d_head)`` KV slice in one shot, so per-rank scratch grows linearly with ``S_local``. The ``_attention_*`` baselines hit a TCM ceiling once ``S_local`` exceeds @@ -20,7 +20,7 @@ When ``S_local ≤ TILE_S_KV`` the loop body never runs and the op_log is structurally identical to today's one-shot path — every existing validation-scale test continues to pass. -``_gqa_prefill_long.py`` is **out of scope** for P3b. Its inner step +``_gqa_attention_prefill_long.py`` is **out of scope** for P3b. Its inner step is IPCQ partial-recv (Ring KV rotation), not an HBM load; tiling it intersects the ring-step structure and is a separate phase. @@ -34,9 +34,9 @@ from __future__ import annotations from pathlib import Path -from kernbench.benches._gqa_decode_long import gqa_decode_long_kernel # noqa: F401 -from kernbench.benches._gqa_decode_short import gqa_decode_short_kernel # noqa: F401 -from kernbench.benches._gqa_prefill_short import gqa_prefill_short_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_long import gqa_attention_decode_long_kernel # noqa: F401 +from kernbench.benches._gqa_attention_decode_short import gqa_attention_decode_short_kernel # noqa: F401 +from kernbench.benches._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # noqa: F401 from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip from kernbench.policy.placement.dp import DPPolicy @@ -94,7 +94,7 @@ def _run_decode_long(*, C: int, P: int, S_kv: int, h_q: int = 1, h_kv: int = 1): dtype=DTYPE, dp=dp_full, name=f"o_tl_c{C}_p{P}_s{S_kv}") ctx.launch( f"gqa_decode_long_tile_c{C}_p{P}_s{S_kv}", - gqa_decode_long_kernel, + gqa_attention_decode_long_kernel, q, k, v, o, 1, S_kv, h_q, h_kv, D_HEAD, C, P, _auto_dim_remap=False, @@ -130,7 +130,7 @@ def _run_decode_short(*, kv_per_cube: int, C: int, P: int, S_kv: int, dtype=DTYPE, dp=dp_full, name=f"o_dsh_s{S_kv}") ctx.launch( f"gqa_decode_short_tile_s{S_kv}", - gqa_decode_short_kernel, + gqa_attention_decode_short_kernel, q, k, v, o, 1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube, _auto_dim_remap=False, @@ -162,7 +162,7 @@ def _run_prefill_short(*, kv_per_cube: int, C: int, P: int, dtype=DTYPE, dp=dp_full, name=f"o_psh_s{S_kv}") ctx.launch( f"gqa_prefill_short_tile_s{S_kv}", - gqa_prefill_short_kernel, + gqa_attention_prefill_short_kernel, q, k, v, o, T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube, _auto_dim_remap=False,