gqa: tile-granular Ring KV (P3c) + rename to gqa_attention_* + ADR-0060/62/63/64 → Accepted
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_<phase>_<context>_kernel`` →
``gqa_attention_<phase>_<context>_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 <noreply@anthropic.com>
This commit is contained in:
+41
-1
@@ -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** 현재
|
||||
+1
-1
@@ -2,7 +2,7 @@
|
||||
|
||||
## Status
|
||||
|
||||
Proposed
|
||||
Accepted
|
||||
|
||||
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context
|
||||
> 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본
|
||||
@@ -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.
|
||||
+47
-3
@@ -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.
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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,
|
||||
+1
-1
@@ -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,
|
||||
@@ -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<C-1 : forward the tile to W
|
||||
#
|
||||
# The ``(t=0, k=0)`` iteration is the bootstrap above; the loop skips
|
||||
# it via ``k_start``. Every other iteration uses ``scratch_scope`` +
|
||||
# ``copy_to`` to recycle per-tile intermediates.
|
||||
#
|
||||
# Triton port: drop ``with tl.scratch_scope():`` and replace each
|
||||
# ``copy_to`` with a Python rebind (e.g. ``m = m_new``).
|
||||
for t in range(n_tiles):
|
||||
tile_start = t * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
k_start = 1 if t == 0 else 0
|
||||
for k in range(k_start, C):
|
||||
with tl.scratch_scope():
|
||||
if k == 0:
|
||||
K_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
else:
|
||||
K_t = tl.recv(dir="E", shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.recv(dir="E", shape=(tile_s, d_head), dtype="f16")
|
||||
if k < C - 1:
|
||||
tl.send(dir="W", src=K_t)
|
||||
tl.send(dir="W", src=V_t)
|
||||
scores = tl.dot(Q, K_t)
|
||||
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, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m, l, O, m_step, l_step, O_step, tl=tl,
|
||||
)
|
||||
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)
|
||||
+1
-1
@@ -27,7 +27,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_prefill_short_kernel(
|
||||
def gqa_attention_prefill_short_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
@@ -1,104 +0,0 @@
|
||||
"""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.
|
||||
|
||||
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, S_local)`` via byte-conserving reshape of
|
||||
the deployed ``(S_local, d_head)`` slice (ADR-0060 §3
|
||||
reshape-not-transpose caveat).
|
||||
- No causal masking / step-skip; blocking ``tl.recv``.
|
||||
- Step-0 local partial is NOT tile-swept — its scratch bound is set
|
||||
by the ring's full-slice ``Kc``/``Vc`` carry, which dominates the
|
||||
score-stack regardless. Tile-granular ring is a separate phase.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def gqa_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 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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user