Attention and KV Cache
The original Transformer 1 attention mechanism computes
softmax(QK^T / sqrtd_k) * V. Computed naively, this materializes an [L x L] score matrix
that does not fit in GPU SRAM for long contexts. poot handles attention in three stages:
primitive decomposition (what the tracer records), flash attention rewrite (what the compiler
produces), and KV caching (how past context is stored).
Attention as primitive ops
The model tracer (ops::attention_masked) emits attention as a sequence of about ten
primitive eqns, plus the KV write and the GQA repeat:
# q: [1, H_q, 1, D]; the new k/v are written into the state buffers at position `pos`
k_cache = DynamicUpdateSlice(k_cache, k, pos) # state_in -> state_out
v_cache = DynamicUpdateSlice(v_cache, v, pos)
# GQA repeat: Reshape -> Broadcast -> Reshape (no-op if H_q == H_kv)
k_exp = repeat_kv(k_cache, n_rep)
v_exp = repeat_kv(v_cache, n_rep)
scores = MatMul(q, Transpose(k_exp)) # [1, H_q, 1, cap]
scores = Mul(scores, scale) # scale, usually 1/sqrt(D)
scores = Add(scores, mask) # large negative past `pos`; a graph computation over
# Slot::Pos
# Numerically stable softmax (max-subtract, then normalize); binary ops broadcast implicitly
m = Reduce(Max, scores, axis=-1, keepdim=True)
xs = Sub(scores, m)
ex = Unary(Exp, xs)
s = Reduce(Sum, ex, axis=-1, keepdim=True)
p = Div(ex, s)
out = MatMul(p, v_exp) # [1, H_q, 1, D]
No special attention opcode. The flash attention rewrite pass recognizes this chain and replaces it.
Multi-head, GQA, and MQA
Multi-query attention (MQA) 2 and grouped-query attention (GQA) 3 reduce the
number of K and V heads relative to Q. poot handles this with repeat_kv
(Reshape, Broadcast, Reshape) over the head axis, no special opcode. For a model with
32 Q heads and 8 KV heads, K and V are expanded 4x. After the flash rewrite the fused op
carries n_rep = 4 and the kernel reads kv_head = q_head / n_rep, so the same K/V rows
are shared by several Q heads without a copy.
The attention rewrite pass
Flash attention 4 5 avoids materializing the [Lq x Lk] score matrix by
streaming over K/V and maintaining a running (max, denominator, accumulator) triple in
workgroup memory. poot implements this as a compiler pass, not an opt-in mode.
compile's attention rewrite runs after canonicalization and RoPE recognition, using the
target grid limit. It recognizes supported attention subgraphs and replaces them with a single
Attention eqn. Decode (one query row against a capacity-long cache) and prefill (every
query row against the prompt) are the same op: q [B, Hq, Lq, Dk], k [B, Hkv, Lk, Dk],
v [B, Hkv, Lk, Dv] (the value width is independent of the key width), an additive
mask [B, Hm, Lq, Lk] whose head axis is 1 or Hq (ALiBi), and an optional attention-logit
softcap. Three more forms ride on the same op: a split key, a paged key and a live length (below).
A sparse-attention indexer uses a block-selected form instead: no mask operand, the selected block ids and
the query positions in its place, so each row visits only the (K + 1) * ratio positions of its K selected
blocks and its own unfinished block rather than all Lk. The tracer states the block selection and the mask
over the selected blocks as primitives, and the selection pass of compile recognizes them and forms the
BlockTopK and SelectedAttention eqns, the same way the attention rewrite forms Attention.
The planner turns the eqn into one generation request that names the shape, the score
transform (scale, or scale plus softcap), the mask kind, the key bound and the schedule. One
generator serves every phase, target and head width. The planner picks the schedule's lane
count from the device's workgroup memory and records the request with the plan. It refuses the
eqn with a typed capability error only when the device cannot hold the running state of even
one lane (AttentionLds, naming the value width and the device's bytes).
Models compose attention with ops::attention_masked or ops::attention_prefill. The compiler
selects the fused op when the shape and target satisfy the matcher checks. Unsupported patterns
keep their primitive decomposition.
The streaming schedule
The implemented schedule runs one workgroup of lanes threads per (batch row, query head,
query row): B * n_q_heads * Lq workgroups, laid over a 2-D grid when that count passes the
device's X grid limit. Each lane strides over the keys (t = lane, lane + lanes, ...) and
keeps its own running state. The lane count is the widest power of two, at most 64, whose
state fits the device's workgroup memory: lanes * (dv + 2) floats.
The online softmax algorithm 6 keeps three accumulators per lane in LDS (workgroup-local memory):
// Per lane, over its keys:
let score = softcap_or_identity(dot(q, k[t]) * scale) + mask[t];
let m_new = max(m, score);
let rescale = exp(m - m_new); // 0 while the lane has seen nothing
let weight = exp(score - m_new);
l = l * rescale + weight;
for d in 0..head_dim {
o[d] = o[d] * rescale + weight * v[t][d];
}
m = m_new;
// One barrier, then every lane folds the lanes' partials over a slice of the output channels:
// M = max of the lane maxima, out[c] = sum_i o_i[c] * exp(m_i - M) / sum_i l_i * exp(m_i - M)
The accumulators are in workgroup LDS, not private arrays: private arrays with dynamic indices crash some SPIR-V drivers (RADV). The same generated kernel runs on wgpu, raw Vulkan, ROCm and PTX.
A key whose score is -inf adds exactly zero weight, and a row with no visible key at all is
0 / 0, as the oracle's softmax gives. Peak memory is O(head_dim) per workgroup: the
[Lq x Lk] score matrix is never written.
Split keys, paged keys and the live length
Split keys. Multi-head latent attention 7 builds each key from a per-head part and one
rotary part shared by every head: concat(k_nope, broadcast(k_pe)). The rewrite recognizes that
chain and keeps the two parts as separate operands, k [B, Hkv, Lk, Dk - s] and
k_shared [B, 1, Lk, s]. The kernel reads the shared part once per key for every head, so the
broadcast key is never built. A value narrower than the key is part of the op's type.
Paged keys. A paged KV cache keeps K and V in per-layer pools of fixed slots (see
Paged KV). The paged key form takes the pools themselves, k [S, Hkv, Dk] and
v [S, Hkv, Dv], and an I32 read [B, Lk] operand that names the pool slot of each logical
position. The kernel addresses key t of row b and key head g as pool row
read[b, t] * Hkv + g, so the per-row [B, Hkv, Lk, D] view is never gathered or transposed. The
schedule, live length and mask handling are the ones the dense form uses; only the address
differs. The rewrite fuses a paged key only when the graph proves every word of read is a slot of
the pool (the map's declared range, which the bind enforces); otherwise the attention stays over the
gathered view. Attention that reads other cached values (latent, indexer and sliding-window caches)
keeps the gathered view.
Live length. A decode over a capacity-long cache normally visits every slot and lets the mask
hide the unused ones. The live form takes an I32 len [B] operand and visits only the first
len[b] keys of each batch row, so a decode step costs the live length, not the capacity. The grid
and the recorded launch do not depend on len, so capture and replay stay valid. The length is built
in the graph from the position (len = pos + 1).
ops::attention_live is the one definition. It requires 1 <= len[b] <= Lk and bounds the length in
the graph with an index guard whose witness is a validation output: an out-of-range length fails
through the validation channel and publishes nothing. It then selects the inactive key and value rows
to zero before either matmul (a true select, so a NaN in an unused row cannot reach the output;
multiplying by zero would keep it) and selects the inactive scores to -inf before the softmax. The
rewrite emits the live bound when the graph proves 1 <= len <= Lk from value ranges, and it holds
the clamp to the same contract: a length whose range comes from an index guard's clamp fuses only
when that guard's witness is a declared validation output, so a bad length never turns silently into
a valid-looking answer; a range that comes from a declared position slot, which every binder checks,
needs no witness.
A graph with only a score mask (the in-graph causal mask of the dense families), a guard on some
operands, or a length it cannot prove keeps its capacity form or its primitive decomposition,
and its results do not change. The live bound types at every float dtype, but the generated kernel is F32 only: a
BF16 or F16 chain (live or not) stays the primitive decomposition, with the same selects at its own dtype, and runs on
ROCm and PTX (a SPIR-V target refuses the two-byte select with a typed DtypeLowering). That primitive form runs on every device: the key positions it
compares against len are an exact I32 range, with no round trip through F32. A generated
kernel clamps len into 0..=Lk before it bounds the loop, so no length can read outside the
buffers.
Split-KV decode
A decode step has one query row per head, so the streaming schedule launches B * n_q_heads
workgroups. At a few heads on a device with tens of compute units, most of the device idles while
each workgroup walks the whole cache. The split-KV schedule (Flash-Decoding 8) cuts
the keys into splits contiguous chunks and runs two launches:
- Partials. One workgroup per (row, split) runs the streaming body over its own chunk of the
live range and writes the unnormalized state
(o[dv], m, l)to a scratch buffer at[(row * splits + split) * (dv + 2)]. A split whose chunk is past the live length, or whose keys are all masked, publishes the neutral(0, -inf, 0)by the same arithmetic as a lane that saw nothing, so the scratch needs no initialization and a stale slot cannot leak into a result. - Combine. One workgroup of
lanesthreads per row reduces the row's partials with the sameexp(m_i - M)weighting the lane combine uses, and writes the output.
The scratch is a plan-local buffer the equation declares (its storage and element count) beside its two dispatches, and the dispatches run in order. The buffer plan colors it with the values: it lives exactly as long as its equation, is never an export, and shares its arena slot with the next layer's scratch, so the arena does not grow with the layer count and the stage's memory budget already counts it.
The planner chooses the split from the device and the shape: only for a decode (Lq == 1) whose
rows do not already fill the device's compute units, at most 16 splits, never a chunk under 128
keys. It keeps the streaming schedule when the partial state would exceed the device's
single-buffer limit or the launch's 2-D fold would pass the device's grid; a plan that declares a
scratch the device cannot allocate anyway is refused with a typed ScratchBytes. The split count
is part of the request, so it is part of the plan key.
The combine reassociates the float sums (per-split partials, then across splits), so it matches the streaming schedule and the oracle within a tolerance, not bit for bit.
Every schedule shares the same score, mask and softmax-combine semantics; the planner chooses the schedule per shape.
KV cache
The attention kernel reads from K/V buffers accumulated over all past tokens. Without caching, K and V would be recomputed from scratch every decode step at O(L) cost.
poot supports two KV allocation schemes:
Fixed-capacity with causal masking: K and V are Storage::State buffers pre-allocated to
cap positions per layer. Per token, new K/V are written at position pos (a
DynamicUpdateSlice). The attention kernel sweeps the full [cap] allocation; a mask hides
positions after pos with a large negative value. That mask is built in-graph from Slot::Pos
against an Iota; no slot carries it. The grid is
constant every token, capturable as a CUDA graph. Memory usage = capacity regardless of actual
sequence length.
Paged KV 9: K/V space is divided into fixed-size blocks (16 tokens,
poot_llm::driver::block_table::BLOCK_SIZE) from a shared pool, inspired by virtual memory
management in operating systems. Each layer holds one [pool_slots, kv_heads, head_dim] pool
shared by every row of a step. The driver keeps a block table per sequence and binds two inputs
per step: read (Slot::SlotMap), an [rows, cap] i32 map from each row's logical position to
its flat physical slot, and write (Slot::WriteMap), a [rows * tokens] i32 map from each new
token to the pool slot it is written to (or -1, so a padding token writes nothing). A write map is
bound only as a checked ForwardSlotMap, which refuses a slot outside the pool or named twice. The
read map is declared in [0, pool_slots), so a word outside the pool is refused at the bind as a
typed range error before anything dispatches. The
paged graph writes new rows with ScatterUpdate through write, in place on the donated pool and
at a cost that follows the tokens, not the pool. Attention then reads each row's K/V from the
written pools in place through read (a paged key, see Paged keys), with no gather and no transpose, and
applies the same causal mask as the contiguous layout. Blocks are allocated on demand and freed
when the sequence completes. Memory scales with actual usage. Shared prefix blocks enable prefix
caching across requests.
Prefix caching
When a new request begins with a prefix that matches a previous request's K/V blocks (identified by a hash of token ids plus the effective model weights), those blocks are pinned in the pool and the new request's block table is pointed at them. The prefill for the matching prefix is skipped. Because the weights identity is part of the key, the base model and each LoRA adapter keep separate prefix entries even when their token ids match, and a hot-reloaded adapter gets a fresh identity so it never reuses the previous adapter's cached K/V. This is a block-allocation policy, the kernel itself does not change.
Implementation details
- Key and value widths. The value width is independent of the key width, and a split key (multi-head latent attention 7) is a first-class operand layout. The matcher declines a chain whose key or value the op's typing rule cannot type, such as a single key head the matmul merely broadcasts. See Crates and implementation details.
- Softcap. The softcap is part of the score transform on every phase: decode and prefill both run the generated kernel.
- Head width. There is no fixed head-width limit: the lane count follows the device's workgroup memory, so a 512-wide head runs on every target with fewer lanes than a 64-wide one.
More crate-level notes live on Crates and implementation details.