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:
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,148 +0,0 @@
|
||||
# ADR-0062: Lazy `tl.load` — 첫 사용 시점 auto-wait를 갖는 non-blocking HBM load
|
||||
|
||||
## Status
|
||||
|
||||
Proposed
|
||||
|
||||
> **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 즉시 발동).
|
||||
@@ -1,169 +0,0 @@
|
||||
# ADR-0062: Lazy `tl.load` — non-blocking HBM load with auto-wait on first use
|
||||
|
||||
## Status
|
||||
|
||||
Proposed
|
||||
|
||||
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and
|
||||
> long-context attention are **KV-load-bound**; load/compute overlap is a
|
||||
> dominant lever. This ADR makes `tl.load` itself **lazy** (non-blocking,
|
||||
> with the wait automatically inserted at first use) rather than adding a
|
||||
> separate `load_async` op. Today the only async primitive is for IPCQ
|
||||
> comms, not for HBM loads.
|
||||
|
||||
## Context
|
||||
|
||||
### What overlap requires
|
||||
|
||||
FlashAttention streams operands: while the GEMM/MATH engine works on the
|
||||
current tile, the DMA engine should already be pulling the next operand.
|
||||
With that overlap a bandwidth-bound kernel runs at roughly
|
||||
`max(compute, dma)` per tile instead of `compute + dma`.
|
||||
|
||||
In ADR-0060's hybrid design the two GEMMs are issued as
|
||||
`tl.composite(op="gemm")` whose K/V operands are `tl.ref` (HBM-resident,
|
||||
streamed per tile by PE_SCHEDULER), so the **per-tile K/V prefetch is
|
||||
handled by the composite scheduler**. Lazy `tl.load` covers the remaining
|
||||
explicit loads (the Q group, and any non-composite kernel) so those also
|
||||
overlap the compute that follows instead of stalling the greenlet.
|
||||
|
||||
### What exists
|
||||
|
||||
- `tl.load(ptr, shape, dtype)` is **blocking**: it emits `DmaReadCmd` and
|
||||
the greenlet kernel suspends until PE_DMA signals completion
|
||||
(`tl_context.py:177-203`; greenlet drive `kernel_runner.py:146-153`).
|
||||
No two `tl.load`s can be in flight from one kernel.
|
||||
- An async pattern **does** exist, but only for IPCQ:
|
||||
`tl.recv_async(dir, ...) -> RecvFuture` + a deferred wait
|
||||
(`kernel_runner.py:248-285`). It proves the machinery — a non-blocking
|
||||
command that returns a future, resolved later by a wait check
|
||||
(`if not future.event.triggered: yield future.event`) — works in the
|
||||
greenlet model.
|
||||
- The DMA engine models a read channel as a SimPy resource (capacity 1,
|
||||
`pe_dma.py:45`) separate from the write channel, so in-flight reads are
|
||||
representable; they serialise on the single read channel while
|
||||
overlapping compute on PE_GEMM/PE_MATH. Only the *kernel-facing API*
|
||||
currently serialises load-vs-compute.
|
||||
|
||||
`tl.load` cannot today overlap with the compute that follows it.
|
||||
|
||||
## Decision
|
||||
|
||||
**Make `tl.load` lazy: it issues the `DmaReadCmd` and returns a handle
|
||||
immediately (non-blocking); the runtime auto-inserts the wait at the
|
||||
first point the loaded data is actually consumed.** The kernel-facing API
|
||||
is unchanged — authors keep writing `tl.load` — so this is a *semantics*
|
||||
change, not a new op. It generalises the existing `recv_async`/wait
|
||||
machinery (1) to the HBM-load path and (2) from an explicit `tl.wait`
|
||||
call to an implicit, dependency-driven wait at first use.
|
||||
|
||||
### D1. `tl` surface — unchanged
|
||||
|
||||
`tl.load(ptr, shape, dtype)` keeps its signature and `TensorHandle`
|
||||
return. What changes is *when* it blocks: never at issue, only implicitly
|
||||
when its result is first read by a consuming op.
|
||||
|
||||
### D2. Mechanism — non-blocking issue + auto-wait on use
|
||||
|
||||
- `tl.load` posts the `DmaReadCmd` to PE_DMA without yielding its
|
||||
completion event, and records the pending event on the returned handle
|
||||
(the `recv_async` pattern, applied to loads).
|
||||
- When a consuming op (`tl.dot`, a MATH op, `tl.store`, a `tl.composite`
|
||||
operand, …) is dispatched, the runtime checks each input handle for a
|
||||
pending load event and yields it first if not yet triggered — i.e. the
|
||||
wait is inserted automatically at the latest correct point (first use).
|
||||
- The op_log entry is unchanged (`memory/dma_read`): asynchrony is a
|
||||
scheduling property, not a new op kind, so `dma_read_count` and existing
|
||||
op_log consumers keep working.
|
||||
|
||||
### D3. Scope — global
|
||||
|
||||
`tl.load` is lazy **everywhere**, not behind an opt-in flag. This is the
|
||||
faithful model: blocking on every load is a property of a *naive* kernel;
|
||||
an efficient kernel (and a real compiler) hoists the load and waits only
|
||||
at use. Consequence: existing kernels that have independent work between a
|
||||
`tl.load` and its first use see **lower (faster) latency** — a
|
||||
correctness improvement of the model, not a behaviour regression. Golden
|
||||
latencies that change must be **regenerated**; kernels that load then
|
||||
immediately use see no change (the auto-wait fires at once, identical to
|
||||
blocking).
|
||||
|
||||
### D4. Latency / overlap semantics
|
||||
|
||||
- `tl.load` charges the **issue** (descriptor push) only; the kernel
|
||||
proceeds. (Per-op issue cost is `dispatch_cycles`, currently 0; an
|
||||
op-type-differentiated issue cost is tracked separately — ADR-0060 §1,
|
||||
§9.)
|
||||
- The DMA transfer occupies the read channel (capacity 1) for its
|
||||
modelled duration in parallel with whatever compute the kernel issues
|
||||
next; multiple in-flight loads serialise on the channel but overlap
|
||||
compute.
|
||||
- The auto-inserted wait blocks only if the transfer has not finished.
|
||||
- Determinism is preserved: the wait point is fixed by program order
|
||||
(first use), and completion is a scheduled event on the modelled DMA
|
||||
channel (SPEC §0.1, R8). No latency subtraction — the overlap is real
|
||||
modelled concurrency.
|
||||
|
||||
> **Modelling assumption.** Auto-wait-at-first-use models a *well-
|
||||
> scheduled* kernel (the compiler places the wait at the latest correct
|
||||
> point). Real compilers approximate this; some loads cannot be hoisted
|
||||
> (register pressure, aliasing). For a performance simulator (SPEC §0)
|
||||
> modelling the well-scheduled case is the intended behaviour.
|
||||
|
||||
## Alternatives
|
||||
|
||||
### A1. Separate `tl.load_async` op (earlier proposal)
|
||||
|
||||
Add an explicit `load_async`/`wait` pair the kernel calls by hand (the
|
||||
double-buffer dance). Rejected: it enlarges the kernel-facing API and
|
||||
pushes buffer-lifetime bookkeeping onto every kernel author, when lazy
|
||||
`tl.load` gives the same overlap with **no** API-surface change and a
|
||||
compiler-style auto-wait.
|
||||
|
||||
### A2. Per-kernel opt-in lazy load
|
||||
|
||||
Keep `tl.load` blocking by default; make new kernels opt in. Rejected:
|
||||
splits `tl.load` semantics into two variants and hides the model
|
||||
improvement from existing benches; global lazy is cleaner and the
|
||||
golden-regeneration cost is one-time (D3).
|
||||
|
||||
### A3. Rely on composite streaming only (no lazy load)
|
||||
|
||||
ADR-0060's composites stream their `tl.ref` K/V operands, so the GEMM
|
||||
operand DMA already overlaps. But explicit loads (the Q group, and any
|
||||
non-composite kernel) still stall without lazy `tl.load`. Insufficient on
|
||||
its own.
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
- Load/compute overlap with **zero** kernel-facing API change; authors
|
||||
keep writing `tl.load`.
|
||||
- Symmetric with `recv_async`/wait — low conceptual surface area.
|
||||
- General: any bandwidth-bound kernel prefetches automatically.
|
||||
|
||||
### Negative
|
||||
- Existing golden latencies shift (faster) for kernels with independent
|
||||
work between load and use → one-time regeneration (D3).
|
||||
- Auto-wait requires the runtime to track per-handle pending events and
|
||||
check them at consuming-op dispatch (data-dependency tracking).
|
||||
- Interacts with scratch lifetime (ADR-0063): an in-flight load's target
|
||||
buffer must not be recycled before its auto-wait fires. The recycling
|
||||
scope must exclude live (un-waited) load buffers.
|
||||
|
||||
## Test Requirements
|
||||
|
||||
1. **Overlap is real**: a kernel that issues `tl.load` then an
|
||||
independent GEMM of comparable duration completes in ≈`max(load,
|
||||
gemm)`, not `load+gemm` (end-to-end latency strictly below the serial
|
||||
sum).
|
||||
2. **Auto-wait correctness**: the loaded `TensorHandle`, when first
|
||||
consumed, carries the same bytes as today's blocking `tl.load` for the
|
||||
same address (Phase 2).
|
||||
3. **Two in flight**: two `tl.load`s to distinct addresses, consumed
|
||||
later, both resolve to correct, independent tensors; their DMAs
|
||||
serialise on the read channel.
|
||||
4. **op_log compatibility**: each `tl.load` still logs exactly one
|
||||
`memory/dma_read`; `dma_read_count` unchanged vs the blocking version.
|
||||
5. **No-overlap no-change**: a kernel that loads then immediately uses has
|
||||
identical latency to the blocking model (auto-wait fires at once).
|
||||
@@ -1,224 +0,0 @@
|
||||
# ADR-0063: `tl.scratch_scope` — per-tile scratch recycling for long context
|
||||
|
||||
## Status
|
||||
|
||||
Proposed
|
||||
|
||||
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Long-context
|
||||
> attention sweeps many K/V tiles; each tile's intermediates
|
||||
> (`scores`, `P`, `exp`, partial `O`, …) currently allocate fresh
|
||||
> scratch that is never freed within a kernel invocation. The 1 MiB
|
||||
> per-PE scratch budget is exhausted long before a realistic context
|
||||
> length.
|
||||
|
||||
## Context
|
||||
|
||||
### The bump allocator
|
||||
|
||||
`TLContext` allocates every math/compute output handle from a per-PE
|
||||
scratch pool with a **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
|
||||
```
|
||||
|
||||
The docstring states the cursor **"resets on every kernel invocation"** —
|
||||
i.e. only at kernel entry, never *within* the kernel body. Every
|
||||
`tl.dot`, `tl.softmax`, `tl.exp`, `tl.sum`, `a - b`, `a * b`, … grabs a
|
||||
fresh slice and nothing is reclaimed until the kernel returns.
|
||||
|
||||
### Why this bites the GQA kernel
|
||||
|
||||
The current milestone bench keeps `S_q = S_kv_per_rank = 16` **explicitly
|
||||
because** of this limit
|
||||
(`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."
|
||||
|
||||
A FlashAttention sweep over `n_tiles` tiles allocates O(`n_tiles` ×
|
||||
per-tile-intermediates). For any realistic context this overflows. The
|
||||
*math* of flash attention needs only **O(1)** live scratch — the running
|
||||
`(m, l, O)` plus the current tile's working set — because each tile's
|
||||
temporaries are dead once that tile is folded into the running state. The
|
||||
allocator just doesn't know they're dead.
|
||||
|
||||
## Decision
|
||||
|
||||
Add a **scratch scope** that lets a kernel mark a region of allocations
|
||||
as reclaimable, so per-tile temporaries are recycled while live running
|
||||
state (and in-flight prefetch buffers) are preserved.
|
||||
|
||||
### D1. `tl` surface — context manager
|
||||
|
||||
```python
|
||||
with tl.scratch_scope():
|
||||
s = tl.dot(q, k_t) # all handles allocated inside the `with`
|
||||
p = tl.softmax(s) # share a region that is rewound on exit
|
||||
o_j = tl.dot(p, v)
|
||||
# ... fold o_j into running (m,l,O) which live OUTSIDE the scope ...
|
||||
# on __exit__: cursor rewinds to its value at __enter__
|
||||
```
|
||||
|
||||
Semantics:
|
||||
- `__enter__` records the current `_scratch_cursor` as a save-point.
|
||||
- `__exit__` restores the cursor to the save-point, freeing everything
|
||||
allocated inside.
|
||||
- Handles allocated **outside** the scope (running `m,l,O`, prefetch
|
||||
buffers held by `LoadFuture` from ADR-0062) keep their addresses —
|
||||
they were allocated before the save-point or in an enclosing scope.
|
||||
|
||||
### D2. Safety contract
|
||||
|
||||
A handle allocated inside a scope **must not** be read after the scope
|
||||
exits — its bytes may be overwritten by the next scope's allocations.
|
||||
The flash loop respects this naturally: the only values that survive a
|
||||
tile iteration are the running accumulators, which are allocated outside
|
||||
the per-tile scope and updated by ops *inside* it writing to outside
|
||||
addresses (the merge writes new running state — see D3).
|
||||
|
||||
### D3. Interaction with running accumulators
|
||||
|
||||
The online-softmax merge reads the old running `(m, l, O)` and the
|
||||
current tile's `(m_j, l_j, O_j)`, producing new running values. To keep
|
||||
the new running values outside the recycled region, the merge writes them
|
||||
to **stable scratch** allocated once before the loop (a small fixed
|
||||
"running-state" arena, distinct from the per-tile scope). Concretely the
|
||||
kernel keeps two arenas:
|
||||
|
||||
- **persistent arena** (allocated once): `m, l, O` (and their
|
||||
double-buffer if needed for the merge).
|
||||
- **scoped arena** (rewound each tile): `scores, P, exp, O_j, scale_*`.
|
||||
|
||||
This mirrors how real flash-attention SRAM budgeting works: a small
|
||||
persistent accumulator region + a recycled tile working set.
|
||||
|
||||
#### D3.1 The persistent-arena write mechanism: `tl.copy_to(dst, src)`
|
||||
|
||||
The merge ops (`tl.maximum`, `tl.exp`, binary `*` / `+`) all call
|
||||
`_make_compute_out(...)` which allocates from the bump cursor (D1).
|
||||
Inside a `scratch_scope`, their result handles therefore live **inside**
|
||||
the scope and vanish on `__exit__`. To realise D3's two-arena split, the
|
||||
kernel needs a way to **write a scoped result's bytes back to a
|
||||
persistent address**.
|
||||
|
||||
The primitive that closes this gap:
|
||||
|
||||
```python
|
||||
def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
|
||||
"""Copy ``src``'s bytes into ``dst``'s address (both TCM).
|
||||
|
||||
Shapes and dtypes must match. ``dst`` is typically a handle
|
||||
allocated outside any active ``scratch_scope`` (the persistent
|
||||
arena); ``src`` is a scoped handle whose bytes must outlive
|
||||
scope ``__exit__``.
|
||||
"""
|
||||
```
|
||||
|
||||
Symmetric to `tl.store` (the HBM-side byte copy), kept TCM-only here so
|
||||
the running-state writeback doesn't pollute op_log with spurious DMA.
|
||||
|
||||
**Mechanics:**
|
||||
- New `CopyCmd(src, dst, nbytes, data_op=True)` command.
|
||||
- op_log: `op_kind="math"`, `op_name="copy"` — runs on the vector engine.
|
||||
- Latency: `pe_math._compute_ns(prod(shape))` — models on-chip register
|
||||
writeback, not HBM transfer.
|
||||
- Emit-time validation: `dst.shape == src.shape`, `dst.dtype == src.dtype`,
|
||||
`dst.space == "tcm"`, `src.space == "tcm"`. Authoring errors surface in
|
||||
Phase 1, not deep in Phase 2 data execution.
|
||||
|
||||
**Call-site pattern:**
|
||||
|
||||
```python
|
||||
m, l, O = init_running(...) # persistent (outside scope)
|
||||
for j in range(n_tiles):
|
||||
with tl.scratch_scope():
|
||||
... # per-tile work (recycled)
|
||||
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) # ← persist new running state
|
||||
tl.copy_to(l, l_new)
|
||||
tl.copy_to(O, O_new)
|
||||
# exit: scoped m_new/l_new/O_new gone; their bytes live in persistent m/l/O
|
||||
```
|
||||
|
||||
The copy happens **before** `__exit__`, so the read of `src` (scoped) is
|
||||
valid; after exit only `dst` (persistent) is read, satisfying D2's
|
||||
no-read-after-exit safety contract.
|
||||
|
||||
**Why a dedicated primitive rather than `dst=` kwargs on every math op**
|
||||
(considered, rejected): adding `dst=` to `tl.maximum`, `tl.exp`,
|
||||
`_binary_math`, `_unary_math`, `_reduction` is ~25 LOC across 5
|
||||
op-families and breaks the uniform "call returns a fresh handle"
|
||||
pattern. `tl.copy_to` is one primitive, one command, one executor
|
||||
handler — minimal surface area for the same effect.
|
||||
|
||||
### D4. Nesting
|
||||
|
||||
Scopes nest (stack of save-points). Inner scope exit rewinds to the inner
|
||||
save-point; outer exit rewinds further. This supports an outer
|
||||
"per-query-block" scope around an inner "per-KV-tile" scope for prefill.
|
||||
|
||||
## Alternatives
|
||||
|
||||
### A1. Round-trip temporaries through HBM
|
||||
|
||||
Store intermediates to HBM and reload to "free" TCM scratch. Rejected:
|
||||
turns a TCM-resident streaming kernel into an HBM-bandwidth-bound one —
|
||||
the opposite of the goal, and it pollutes op_log with spurious DMA.
|
||||
|
||||
### A2. Tiled `tl.composite` (scheduler-managed scratch)
|
||||
|
||||
`tl.composite` recycles per-tile scratch inside PE_SCHEDULER
|
||||
automatically. As in ADR-0062 A1, this is attractive but blocked on a
|
||||
flash-capable composite kind (two GEMMs + carried `(m,l,O)`), a much
|
||||
larger change. `scratch_scope` gives the same memory behaviour for the
|
||||
greenlet kernel with a tiny, general primitive. The two can coexist.
|
||||
|
||||
### A3. Grow the scratch budget
|
||||
|
||||
Bump `pe_tcm.kernel_scratch_mb`. Rejected as a fix: it only pushes the
|
||||
context-length ceiling out linearly while still leaking O(`n_tiles`)
|
||||
scratch, and it misrepresents the hardware (real TCM is small; the point
|
||||
of flash attention is O(1) working set). Useful only as a coarse knob,
|
||||
not a substitute for recycling.
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
- Removes the artificial `S = 16` validation-scale ceiling; enables
|
||||
realistic context lengths in both timing and data modes.
|
||||
- Faithfully models flash attention's O(1) working-set property.
|
||||
- Small, general primitive (any tiled kernel benefits).
|
||||
|
||||
### Negative
|
||||
- A use-after-scope bug silently reads stale bytes. Mitigated by the D2
|
||||
contract, by keeping the scope discipline inside a shared attention
|
||||
helper, and (optionally) by a debug build that poisons rewound regions.
|
||||
- Must coordinate with ADR-0062: prefetch buffers are live across tile
|
||||
iterations, so they belong to the persistent arena, not the scoped one.
|
||||
|
||||
## Test Requirements
|
||||
|
||||
1. **Recycling**: a loop of `N` tiles inside `tl.scratch_scope()` keeps
|
||||
peak `_scratch_cursor` bounded by one tile's footprint, independent of
|
||||
`N` (today it grows linearly and overflows).
|
||||
2. **Correctness**: a flash sweep with scopes produces the same `O` as
|
||||
the same sweep without scopes at a small `N` that fits without
|
||||
recycling (Phase 2).
|
||||
3. **Long context**: a sweep at an `N` that would overflow 1 MiB without
|
||||
scopes completes (the exact failure the `S=16` cap avoids today).
|
||||
4. **Persistent-vs-scoped isolation**: running `(m,l,O)` allocated
|
||||
outside the scope retains correct values across `__exit__`.
|
||||
5. **Nesting**: nested scopes rewind to the correct save-points.
|
||||
Reference in New Issue
Block a user