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:
2026-06-10 16:17:32 -07:00
parent a8c50238c6
commit 7fad0371c5
22 changed files with 667 additions and 152 deletions
File diff suppressed because it is too large Load Diff
+148
View File
@@ -0,0 +1,148 @@
# ADR-0062: Lazy `tl.load` — 첫 사용 시점 auto-wait를 갖는 non-blocking HBM load
## Status
Accepted
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context
> 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본
> ADR은 별도의 `load_async` op을 추가하는 대신 `tl.load` 자체를 **lazy**
> (non-blocking, wait를 첫 사용 시점에 자동 삽입)로 만든다. 오늘날 유일한
> async primitive는 IPCQ 통신용이지 HBM load용이 아니다.
## Context
### 오버랩이 요구하는 것
FlashAttention은 operand를 스트리밍한다: GEMM/MATH 엔진이 현재 타일을 처리하는
동안 DMA 엔진은 이미 다음 operand를 당기고 있어야 한다. 그 오버랩이 있으면
대역폭-bound 커널은 타일당 `compute + dma` 대신 대략 `max(compute, dma)`
돈다.
ADR-0060의 하이브리드 설계에서 두 GEMM은 `tl.composite(op="gemm")`로 발행되고,
그 K/V operand는 `tl.ref`(HBM 상주, PE_SCHEDULER가 타일당 스트리밍)이므로 **타일당
K/V prefetch는 composite scheduler가 처리**한다. lazy `tl.load`는 나머지 명시적
load(Q group, 그리고 비-composite 커널)를 담당하여 그것들도 greenlet을 멈추는
대신 뒤따르는 compute와 오버랩되게 한다.
### 현재 존재하는 것
- `tl.load(ptr, shape, dtype)`은 **blocking**이다: `DmaReadCmd`를 emit하고
greenlet 커널은 PE_DMA가 완료를 알릴 때까지 suspend된다
(`tl_context.py:177-203`; greenlet 구동 `kernel_runner.py:146-153`). 한
커널에서 두 `tl.load`가 in-flight일 수 없다.
- async 패턴이 **존재**하긴 하나 IPCQ 전용이다:
`tl.recv_async(dir, ...) -> RecvFuture` + 지연 wait
(`kernel_runner.py:248-285`). 이는 메커니즘을 입증한다 — future를 반환하고
나중에 wait 체크(`if not future.event.triggered: yield future.event`)로
resolve되는 non-blocking 커맨드가 greenlet 모델에서 동작한다.
- DMA 엔진은 read 채널을 write 채널과 분리된 SimPy resource(capacity 1,
`pe_dma.py:45`)로 모델한다. 따라서 in-flight read는 표현 가능하며, 단일 read
채널에서 직렬화되면서 PE_GEMM/PE_MATH의 compute와 오버랩된다. 오직
*커널-대면 API*만이 현재 load-vs-compute를 직렬화한다.
`tl.load`는 오늘날 뒤따르는 compute와 오버랩될 수 없다.
## Decision
**`tl.load`를 lazy로 만든다: `DmaReadCmd`를 발행하고 핸들을 즉시 반환
(non-blocking); 런타임은 load된 데이터가 실제로 처음 소비되는 지점에 wait를
자동 삽입한다.** 커널-대면 API는 불변이다 — 작성자는 계속 `tl.load`를 쓴다 —
따라서 이는 새 op이 아니라 *시맨틱* 변경이다. 기존 `recv_async`/wait 기계장치를
(1) HBM-load 경로로, (2) 명시적 `tl.wait` 호출에서 암묵적·의존성 기반 first-use
wait로 일반화한다.
### D1. `tl` 표면 — 불변
`tl.load(ptr, shape, dtype)`은 시그니처와 `TensorHandle` 반환을 유지한다.
바뀌는 것은 *언제* blocking하느냐다: 발행 시점에는 결코 아니고, 결과가 소비 op에
의해 처음 읽힐 때만 암묵적으로.
### D2. 메커니즘 — non-blocking 발행 + 사용 시점 auto-wait
- `tl.load`는 완료 event를 yield하지 않고 `DmaReadCmd`를 PE_DMA에 post하며,
반환 핸들에 pending event를 기록한다(`recv_async` 패턴을 load에 적용).
- 소비 op(`tl.dot`, MATH op, `tl.store`, `tl.composite` operand, …)이
dispatch될 때, 런타임은 각 입력 핸들의 pending load event를 확인하고 아직
triggered가 아니면 먼저 yield한다 — 즉 wait가 가장 늦은 올바른 지점(첫 사용)에
자동 삽입된다.
- op_log 엔트리는 불변(`memory/dma_read`): asynchrony는 스케줄링 속성이지 새 op
종류가 아니므로 `dma_read_count` 및 기존 op_log 소비자가 계속 동작한다.
### D3. 범위 — 전역
`tl.load`는 opt-in 플래그 뒤가 아니라 **모든 곳에서** lazy다. 이것이 충실한
모델이다: 매 load마다 blocking하는 것은 *naive* 커널의 속성이고, 효율 커널(그리고
실제 컴파일러)은 load를 hoist하고 사용 시점에만 wait한다. 결과: `tl.load`와 첫
사용 사이에 독립 작업이 있는 기존 커널은 **더 낮은(빠른) latency**를 본다 —
동작의 회귀가 아니라 모델의 정확도 개선. 바뀌는 골든 latency는 **재생성**해야
하며; load 직후 곧바로 사용하는 커널은 변화가 없다(auto-wait이 즉시 발동하여
blocking과 동일).
### D4. Latency / 오버랩 시맨틱
- `tl.load`**발행**(descriptor push)만 청구하고 커널은 진행한다. (per-op
발행 비용은 `dispatch_cycles`, 현재 0; op 종류별 차등 발행 비용은 별도로
추적 — ADR-0060 §1, §9, ADR-0064.)
- DMA 전송은 read 채널(capacity 1)을 모델 지속시간만큼 점유하며 커널이 다음에
내는 compute와 병렬; 다중 in-flight load는 채널에서 직렬화되나 compute와
오버랩.
- 자동 삽입된 wait는 전송이 끝나지 않았을 때만 blocking.
- 결정성 보존: wait 지점은 프로그램 순서(첫 사용)로 고정되고, 완료는 modeled DMA
채널의 scheduled event다(SPEC §0.1, R8). latency 빼기 없음 — 오버랩은 실제
모델 동시성.
> **모델링 가정.** auto-wait-at-first-use는 *잘 스케줄된* 커널(컴파일러가
> wait를 가장 늦은 올바른 지점에 둠)을 모델한다. 실제 컴파일러는 이를 근사하며,
> 일부 load는 hoist 불가(register pressure, aliasing). 성능 시뮬레이터(SPEC
> §0)에서 잘 스케줄된 경우를 모델하는 것이 의도된 동작이다.
## Alternatives
### A1. 별도 `tl.load_async` op (이전 제안)
커널이 손으로 호출하는 명시적 `load_async`/`wait` 쌍 추가(더블버퍼 댄스). 기각:
커널-대면 API를 키우고 버퍼 수명 부기를 모든 커널 작성자에게 떠넘긴다. lazy
`tl.load`**API 표면 변경 없이** 동일 오버랩과 컴파일러식 auto-wait를 준다.
### A2. 커널별 opt-in lazy load
`tl.load`를 기본 blocking으로 두고 새 커널만 opt-in. 기각: `tl.load` 시맨틱을 두
변종으로 쪼개고 기존 bench로부터 모델 개선을 숨긴다; 전역 lazy가 더 깔끔하고
골든 재생성 비용은 1회성(D3).
### A3. composite 스트리밍에만 의존 (lazy load 없이)
ADR-0060의 composite는 `tl.ref` K/V operand를 스트리밍하므로 GEMM operand DMA는
이미 오버랩된다. 그러나 명시적 load(Q group, 비-composite 커널)는 lazy `tl.load`
없이 여전히 stall. 단독으로는 불충분.
## Consequences
### Positive
- 커널-대면 API 변경 **제로**로 load/compute 오버랩; 작성자는 계속 `tl.load`.
- `recv_async`/wait와 대칭 — 개념적 표면적 작음.
- 일반적: 어떤 대역폭-bound 커널도 자동 prefetch.
### Negative
- load와 사용 사이에 독립 작업이 있는 커널의 기존 골든 latency가 이동(빨라짐)
→ 1회성 재생성(D3).
- auto-wait는 런타임이 핸들별 pending event를 추적하고 소비-op dispatch에서
확인하도록 요구(데이터 의존성 추적).
- scratch 수명(ADR-0063)과 상호작용: in-flight load의 대상 버퍼는 auto-wait이
발동하기 전에 recycle되면 안 된다. recycling scope는 살아있는(미-wait된) load
버퍼를 제외해야 한다.
## Test Requirements
1. **오버랩이 실재**: `tl.load` 후 비슷한 지속시간의 독립 GEMM을 내는 커널이
`load+gemm`이 아니라 ≈`max(load, gemm)`으로 완료(end-to-end latency가 직렬
합보다 엄격히 작음).
2. **auto-wait 정확성**: load된 `TensorHandle`이 처음 소비될 때, 같은 주소에
대해 오늘날 blocking `tl.load`와 동일한 바이트를 보유(Phase 2).
3. **둘 in-flight**: 서로 다른 주소로의 두 `tl.load`가 나중에 소비되어 둘 다
올바른 독립 텐서로 resolve; 그 DMA는 read 채널에서 직렬화.
4. **op_log 호환성**: 각 `tl.load`은 여전히 정확히 하나의 `memory/dma_read`
로그; `dma_read_count`는 blocking 버전 대비 불변.
5. **무-오버랩 무-변화**: load 직후 곧바로 사용하는 커널은 blocking 모델과 동일한
latency(auto-wait 즉시 발동).
@@ -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.