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:
@@ -0,0 +1,169 @@
|
||||
# ADR-0062: Lazy `tl.load` — non-blocking HBM load with auto-wait on first use
|
||||
|
||||
## Status
|
||||
|
||||
Accepted
|
||||
|
||||
> 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).
|
||||
Reference in New Issue
Block a user