ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(adr): pivot ADR-0060 to composite hybrid + lazy tl.load; add ADR-0064 cost model

Browse Source 
- ADR-0060: GEMMs (Q.Kt, P.V) via existing tl.composite (scheduler-managed
  tiling + K/V DMA streaming); softmax merge + IPCQ tree reduction stay in
  kernel. Front TL;DR pseudocode of the final composite kernel; new section
  B lists open design items (DDD sync, K pre-transpose, dma_read lever,
  kernel-vs-scheduler tiling, ring path).
- ADR-0062: redefined from a new load_async op to global lazy tl.load
  (non-blocking + auto-wait on first use; API unchanged; goldens regenerate).
- ADR-0064 (new): per-op-type CPU issue cost model (composite ~40ns >>
  primitive) so the hybrid's CPU-saturation win becomes measurable
  (currently dispatch_cycles=0 hides it). Cost-model impl deferred.
- KO mirrors for ADR-0060/0062/0064 (-ko suffix, adr-proposed).

Rationale: non-blocking CompositeCmd offloads tiling to PE_SCHEDULER,
decoupling CPU issue-rate from execution so the CPU can saturate the
engines; the prior 'composite = no latency benefit' claim was an artifact
of dispatch_cycles=0. Docs only; no production code changed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-04 10:19:37 -07:00
parent 6d24b9306f
commit 7a9d4ec47b
6 changed files with 1675 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/adr-proposed/ADR-0062-prog-tl-async-load.md
+126 -105
View File
@@ -1,148 +1,169 @@
# ADR-0062: `tl.load_async` — non-blocking HBM tile load for KV prefetch
# 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**; the dominant lever is
> overlapping the next K/V tile's DMA with the current tile's compute.
> Today the only async primitive is for IPCQ comms, not for HBM loads.
> 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 K/V tiles: while the MATH/GEMM engine works on
tile `j` (Q·Kⱼᵀ → softmax → P·V), the DMA engine should already be
pulling tile `j+1` (and `j+2` at prefetch depth 2). With that overlap a
KV-load-bound kernel runs at roughly `max(compute, dma)` per tile instead
of `compute + dma`.
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
(`src/kernbench/triton_emu/tl_context.py:177-203`; greenlet drive in
`kernel_runner.py`). No two `tl.load`s can be in flight from one
kernel.
- `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` + `tl.wait(future)`
(`tl_context.py:543-560`, `tl_context.py:660-693`). It proves the
machinery — a non-blocking command that returns a future, resolved
later by `tl.wait` — works in the greenlet model.
- The DMA engine already models a read channel as a SimPy resource
separate from the write channel
(`src/kernbench/components/builtin/pe_dma.py`), so concurrent in-flight
reads are representable at the component layer; only the *kernel-facing
API* serialises them.
`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.
There is **no** `tl.load_async`. Prefetch/double-buffering cannot be
expressed.
`tl.load` cannot today overlap with the compute that follows it.
## Decision
Add a non-blocking HBM load that mirrors the existing `recv_async`/`wait`
contract.
**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
### D1. `tl` surface — unchanged
```python
def load_async(self, ptr: int, shape: tuple[int, ...],
dtype: str = "f16") -> LoadFuture:
"""Issue a DMA read and return immediately. Resolve with tl.wait(fut)
-> TensorHandle. Multiple loads may be in flight; ordering of
resolution is by tl.wait calls, not issue order."""
```
`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.
`tl.wait` is extended to accept a `LoadFuture` (it already dispatches on
`CompletionHandle` vs `RecvFuture` — add a third arm), returning the
loaded `TensorHandle`.
### D2. Mechanism — non-blocking issue + auto-wait on use
### D2. Command
- `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.
Reuse the existing `DmaReadCmd` with a `blocking=False` flag (mirroring
`IpcqRecvCmd.blocking`), or a thin `DmaReadAsyncCmd` sibling — whichever
keeps PE_DMA's handler simplest. The op_log entry is unchanged
(`memory/dma_read`); asynchrony is a scheduling property, not a new op
kind, so existing op_log consumers and the `dma_read_count` metric
(milestone bench) keep working.
### D3. Scope — global
### D3. Latency / overlap semantics
`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).
- `load_async` charges the **issue** (descriptor push) only; the kernel
proceeds.
- The DMA transfer occupies the read channel for its modelled duration
in parallel with whatever compute the kernel issues next.
- `tl.wait(fut)` blocks only if the transfer has not finished; if it has,
it returns immediately (same fast-path as `recv_async`/`wait`).
- Determinism is preserved: completion is a scheduled event on the
modelled DMA channel/link (SPEC §0.1, R8). No magic — the overlap is
real modelled concurrency, not a hand-waved latency subtraction.
### D4. Latency / overlap semantics
### D4. Double-buffer usage (the intended pattern)
- `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.
```python
# Prime: issue first two tile loads.
f0 = tl.load_async(K_base + 0*tile_bytes, (TILE, d))
f1 = tl.load_async(K_base + 1*tile_bytes, (TILE, d))
for j in range(n_tiles):
Kj = tl.wait(f_cur) # resolves the j-th load
if j + 2 < n_tiles: # keep depth-2 pipeline full
f_next2 = tl.load_async(K_base + (j+2)*tile_bytes, (TILE, d))
s = tl.dot(q, tl.trans(Kj)) # compute overlaps f_{j+1}'s DMA
...
```
(V tiles double-buffer the same way; ADR-0060 §3 schedules the V load so
it overlaps the Q·Kᵀ + softmax of the same tile.)
> **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. Use a tiled `tl.composite` instead
### A1. Separate `tl.load_async` op (earlier proposal)
`tl.composite` already pipelines tiles inside PE_SCHEDULER (DMA of tile
`j+1` overlaps compute of tile `j`) and recycles per-tile scratch — so
the composite path gets prefetch "for free." Rejected as the *sole*
mechanism because today's composite is a single GEMM head + math
epilogues (`pe_commands.py:144-162`); it cannot express the
two-GEMM-with-running-`(m,l,O)` flash inner loop. Building a bespoke
"flash composite" kind is a much larger change than `load_async`, which
is a small, general primitive that also helps non-composite kernels.
`load_async` and a future flash-composite are not mutually exclusive
(ADR-0060 §8).
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. Rely on `recv_async` only (no HBM async)
### A2. Per-kernel opt-in lazy load
Only works for the Ring-Attention path where KV arrives over IPCQ
(comms, already overlappable). The decode and no-SP paths load KV from
**local HBM**, which has no async path today. Insufficient.
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
- Enables the headline KV-load-bound optimisation for decode and
long-context, in the proven greenlet model.
- Symmetric with `recv_async`/`wait` — low conceptual surface area.
- General: any bandwidth-bound kernel can prefetch.
- 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
- Kernels must manage futures and buffer lifetimes explicitly (the
double-buffer dance). Mitigated by keeping the pattern in a shared
helper used by all four ADR-0060 cases.
- Interacts with scratch lifetime (ADR-0063): in-flight buffers must not
be recycled before their `tl.wait`. The recycling scope must exclude
live prefetch buffers.
- 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 `load_async` then a GEMM of
comparable duration completes in ≈`max(load, gemm)`, not `load+gemm`
(assert end-to-end latency strictly below the serial sum).
2. **Correctness**: `tl.wait(load_async(...))` returns the same bytes as
blocking `tl.load(...)` for the same address (Phase 2).
3. **Two in flight**: two outstanding `load_async`s to distinct
addresses both resolve to correct, independent tensors.
4. **op_log compatibility**: each `load_async` still logs exactly one
`memory/dma_read`, so `dma_read_count` metrics are unchanged vs the
blocking version.
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).