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
|
||||
Reference in New Issue
Block a user