IV.5: Flash attention
There was a promise hiding in IV.1's closing histogram. Two lines of it read fill_strict_upper_tri: 448 and transpose_2d: 448, and the text noted that every one of those calls lives in the prefill pass, then moved on. Decode attention got fused four chapters ago. Prefill still runs the textbook attention from I.5: for each of the model's 16 heads in each of its 28 layers, build the full score matrix (every query row dotted with every key row), mask the future, softmax every row, multiply by V.
That matrix is the last big intermediate the engine still materializes, and it is the wrong shape to be cheap. At 511 prompt tokens it is [511, 511]: about 261,000 floats, a megabyte per head. The engine writes it, rewrites it with the mask, rewrites it again as probabilities, then reads the whole thing back for the scores-times-V matmul. That happens 448 times per prefill (16 heads times 28 layers), roughly half a gigabyte of score traffic before the first token comes out. And it is the only quadratic term left in the pass: the Q/K/V projections, the MLP, the norms all grow linearly with prompt length, while score matrices grow with its square. At 2,044 tokens each head's matrix is 16.7 MB and the total crosses 7 GB.
This chapter deletes the matrix. Flash attention computes exactly the numbers that mask-softmax-matmul computes while never holding more than a 64-score tile at a time, and on the CPU backends it cuts a 511-token prefill by 13% on one core and by a quarter on the pool, stretching to 36% at 2k context where the quadratic term bites harder. It also hands us the most instructive negative result of the act: wired into the Metal backends, flash attention made prefill more than twice as slow, and the fix was for those backends to refuse it. Both halves of that story are the chapter.
Softmax without seeing the whole row
The obstacle to streaming attention is softmax. I.4 built it with the subtract-the-max trick: find the row's maximum M, exponentiate every score relative to it, divide by the sum. The max and the sum are properties of the entire row. Written that way, no output can be produced until every score exists, which is exactly why the code materializes them all first.
But look at what attention does with the softmaxed row. The output for one query is a weighted average of value rows, and if you substitute softmax's formula in, the whole computation for a row of scores s_0..s_n is:
out = ( exp(s_0 - M) * v_0 + exp(s_1 - M) * v_1 + ... + exp(s_n - M) * v_n )
/ ( exp(s_0 - M) + exp(s_1 - M) + ... + exp(s_n - M) )One vector-valued numerator, one scalar denominator, both plain sums over the keys. Sums can be accumulated in pieces. The only global thing in the formula is M, and M is only there to keep exp from overflowing; any value at least as large as the true max works. So keep three running quantities per query row:
m: the largest score seen so far,l: the running denominator, every term computed asexp(s_j - m),acc: the running numerator,head_dimfloats, every termexp(s_j - m) * v_j.
Walk the keys in tiles of 64. For each tile, compute its 64 scores and find the tile's max. If it does not beat m, exponentiate against m and add into l and acc as usual. If it does beat m, every term already accumulated is stale: it was computed as exp(s_j - m_old) when it should have been exp(s_j - m_new). The fix is one identity:
exp(s_j - m_new) = exp(s_j - m_old) * exp(m_old - m_new)The correction factor is the same for every accumulated term, so one multiply of acc and l by exp(m_old - m_new) repairs everything at once, and the walk continues with m = m_new. Because m only ever rises, that factor is at most 1 and every exponent ever taken is at most 0: overflow is impossible, the same guarantee subtract-the-max gave, maintained incrementally. When the keys run out, acc / l is the softmax-weighted average. Not an approximation of it: the same terms in the same ratio, grouped differently.
This incremental max-and-sum bookkeeping is called online softmax, and the algorithm that applies it to attention, tiling K and V so scores never touch slow memory, is flash attention (Dao et al., 2022). The paper's setting is a GPU, where the score matrix would live in HBM (the GPU's main memory) and the tiles live in on-chip SRAM hundreds of times faster; its point was that attention is bound by that memory traffic, not by its arithmetic. The same restructuring pays on a CPU for the same reason with different names: a materialized score matrix streams through RAM, a 64-score tile stays in L1 cache and registers.
One more thing falls out for free. The causal mask says query row i may only attend to keys 0..=i. In the materialized version that is a tensor-wide pass writing negative infinity into the upper triangle so softmax zeroes it. In the streaming version a masked score is a term we never accumulate, so the key walk for row i just stops at key i. The mask stops being data and becomes a loop bound.
The reference implementation
The engine grows one new Backend operation, and like fused_attention in IV.1 it ships as a default method on the trait itself in src/backend/backend_trait.rs, so every backend inherits a correct implementation and any backend can override it with a faster or more honest one. Both halves of that sentence matter this chapter. Here is the whole thing, walked in pieces:
/// Flash attention over a whole prefill: computes softmax(q k^T * scale) v
/// per head without ever materializing the [seq x seq] score matrix.
///
/// Keys/values are walked in tiles; each query row keeps a running max `m`
/// and running sum `l` (online softmax), rescaling its output accumulator
/// whenever a new tile raises the max. `q` is [seq, n_heads*head_dim];
/// `k`/`v` are [seq, n_kv_heads*head_dim]. With `causal`, query row i only
/// attends to keys 0..=i.
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
const TILE: usize = 64;
let seq_q = q.shape()[0];
let seq_k = k.shape()[0];
let qw = n_heads * head_dim;
let kvw = n_kv_heads * head_dim;
assert_eq!(q.shape(), &[seq_q, qw][..]);
assert_eq!(k.shape(), &[seq_k, kvw][..]);
assert_eq!(v.shape(), &[seq_k, kvw][..]);
if causal {
assert_eq!(
seq_q, seq_k,
"causal flash attention expects seq_q == seq_k"
);
}The signature is fused_attention's plus a causal flag, and the doc comment is the algorithm in miniature. Prefill always passes causal = true; the flag exists because the math does not require causality (the tests exercise it, and a future non-causal caller could too), but a causal call only makes sense when queries and keys cover the same positions, hence the extra assert.
let kv_group = n_heads / n_kv_heads;
let scale = 1.0 / (head_dim as f32).sqrt();
let qd = q.as_f32_slice();
let kd = k.as_f32_slice();
let vd = v.as_f32_slice();
let mut out = vec![0.0f32; seq_q * qw];
let mut s = [0.0f32; TILE];
let mut acc = vec![0.0f32; head_dim];Setup. kv_group is I.5's grouped-query attention bookkeeping: 16 query heads share 8 KV heads, two to one. The two small buffers are the entire memory story of this chapter: s holds one tile of scores (64 floats, 256 bytes) and acc holds one running numerator (head_dim floats, 512 bytes). Both are allocated once and reused for every row of every head. Against the megabyte-per-head matrix they replace, that is the whole pitch, visible in two lets.
for h in 0..n_heads {
let kv_off = (h / kv_group) * head_dim;
for i in 0..seq_q {
let q_row = &qd[i * qw + h * head_dim..i * qw + h * head_dim + head_dim];
let kv_end = if causal { i + 1 } else { seq_k };
let mut m = f32::NEG_INFINITY;
let mut l = 0.0f32;
acc.iter_mut().for_each(|a| *a = 0.0);
let mut t0 = 0;Per head, per query row: slice out the query, set the causal loop bound (kv_end = i + 1 is "stop at your own position"), and reset the online-softmax state: m to negative infinity, l to zero, acc to zeros.
while t0 < kv_end {
let t1 = (t0 + TILE).min(kv_end);
let mut tile_max = f32::NEG_INFINITY;
for t in t0..t1 {
let k_row = &kd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
let mut dot = 0.0f32;
for d in 0..head_dim {
dot += q_row[d] * k_row[d];
}
let sc = dot * scale;
s[t - t0] = sc;
if sc > tile_max {
tile_max = sc;
}
}The first pass over a tile: dot the query against each key row, scale, stash the score in s, and track the tile's max. This is the reference implementation, so the dot product is the plain scalar loop; the NEON version below swaps it for vector code.
if tile_max > m {
let corr = (m - tile_max).exp();
acc.iter_mut().for_each(|a| *a *= corr);
l *= corr;
m = tile_max;
}The rescale, the five lines the whole algorithm turns on. If the tile raised the max, corr = exp(m_old - m_new) repairs every already-accumulated term in acc and l at once, exactly the identity from the last section.
for t in t0..t1 {
let p = (s[t - t0] - m).exp();
l += p;
let v_row = &vd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
for d in 0..head_dim {
acc[d] += p * v_row[d];
}
}
t0 = t1;
}The second pass over the tile: exponentiate each stashed score against the (possibly just-raised) max, grow the denominator, and add p * v_row into the numerator. This is where V is consumed, tile by tile, without the probabilities ever existing as a matrix.
let inv = 1.0 / l;
let dst = &mut out[i * qw + h * head_dim..i * qw + h * head_dim + head_dim];
for d in 0..head_dim {
dst[d] = acc[d] * inv;
}
}
}
Tensor::new(out, vec![seq_q, qw])
}After the last tile, one divide by l turns the numerator into the softmax-weighted average, written into this head's columns of the output row. The function returns the same [seq, n_heads * head_dim] concatenated-heads tensor the old per-head loop produced, ready for the output projection.
The NEON version
The scalar reference spends nearly all its time in two inner loops: the query-key dot product and the p * v_row accumulation. II.3 already built vector versions of both shapes: dot_f32 (four independent FMA accumulators, 16 floats per iteration) and saxpy_row_f32 (fused multiply-add of a scaled row, 16 floats per iteration), and IV.1's fused decode attention was their second customer. Flash attention is their third. src/backend/simd_cpu.rs factors the whole per-head walk into a free function:
pub(crate) const FLASH_TILE: usize = 64;
/// One head of flash-attention prefill, NEON-accelerated: online softmax over
/// K/V tiles, never materializing the score matrix. Writes the [seq_q, head_dim]
/// context for this head into `out` at column `out_off` (row stride `out_cols`).
///
/// # Safety
/// `out` must be valid for `seq_q * out_cols` floats and no other thread may
/// write columns `out_off..out_off + head_dim` concurrently.
#[allow(clippy::too_many_arguments)]
pub(crate) unsafe fn flash_attention_head_neon(
qd: &[f32],
kd: &[f32],
vd: &[f32],
seq_q: usize,
seq_k: usize,
qw: usize,
kvw: usize,
q_off: usize,
kv_off: usize,
head_dim: usize,
scale: f32,
causal: bool,
out: *mut f32,
out_cols: usize,
out_off: usize,
) {The argument list is long because the function is deliberately dumb: raw slices, raw offsets, and a raw output pointer. The pointer is the interesting part. Each head writes a disjoint 128-column stripe of the shared output matrix, so several heads can run this function concurrently against the same allocation without ever touching the same float; the # Safety contract states exactly that, and it is why the function is unsafe and takes *mut f32 instead of &mut [f32] (a mutable slice would promise exclusive access to the whole matrix, which is precisely what parallel heads do not have).
let mut s = [0.0f32; FLASH_TILE];
let mut acc = vec![0.0f32; head_dim];
for i in 0..seq_q {
let q_row = &qd[i * qw + q_off..i * qw + q_off + head_dim];
let kv_end = if causal { i + 1 } else { seq_k };
let mut m = f32::NEG_INFINITY;
let mut l = 0.0f32;
acc.iter_mut().for_each(|a| *a = 0.0);
let mut t0 = 0;
while t0 < kv_end {
let t1 = (t0 + FLASH_TILE).min(kv_end);
let mut tile_max = f32::NEG_INFINITY;
for t in t0..t1 {
let k_row = &kd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
let sc = SimdCpu::<super::CpuBackend>::dot_f32(q_row, k_row) * scale;
s[t - t0] = sc;
if sc > tile_max {
tile_max = sc;
}
}
if tile_max > m {
let corr = (m - tile_max).exp();
acc.iter_mut().for_each(|a| *a *= corr);
l *= corr;
m = tile_max;
}
for t in t0..t1 {
let p = (s[t - t0] - m).exp();
l += p;
let v_row = &vd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
SimdCpu::<super::CpuBackend>::saxpy_row_f32(
acc.as_mut_ptr(),
v_row.as_ptr(),
p,
head_dim,
);
}
t0 = t1;
}
let inv = 1.0 / l;
unsafe {
let dst = out.add(i * out_cols + out_off);
for (d, a) in acc.iter().enumerate() {
*dst.add(d) = a * inv;
}
}
}
}Line for line this is the reference implementation with the two hot loops swapped for the NEON helpers: dot_f32 computes each score, saxpy_row_f32 does the acc += p * v_row accumulation. The online-softmax skeleton (tile walk, m/l/rescale) is untouched; it is bookkeeping, not bandwidth.
The SimdCpu backend override runs the heads one after another on its single core:
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
let seq_q = q.shape()[0];
let seq_k = k.shape()[0];
let qw = n_heads * head_dim;
let kvw = n_kv_heads * head_dim;
assert_eq!(q.shape(), &[seq_q, qw][..]);
assert_eq!(k.shape(), &[seq_k, kvw][..]);
assert_eq!(v.shape(), &[seq_k, kvw][..]);
if causal {
assert_eq!(
seq_q, seq_k,
"causal flash attention expects seq_q == seq_k"
);
}
let kv_group = n_heads / n_kv_heads;
let scale = 1.0 / (head_dim as f32).sqrt();
let qd = q.as_f32_slice();
let kd = k.as_f32_slice();
let vd = v.as_f32_slice();
let mut out = vec![0.0f32; seq_q * qw];
let out_ptr = out.as_mut_ptr();
for h in 0..n_heads {
let kv_off = (h / kv_group) * head_dim;
unsafe {
flash_attention_head_neon(
qd,
kd,
vd,
seq_q,
seq_k,
qw,
kvw,
h * head_dim,
kv_off,
head_dim,
scale,
causal,
out_ptr,
qw,
h * head_dim,
);
}
}
Tensor::new(out, vec![seq_q, qw])
}The PoolCpu backend from IV.3 has cores to spend, and the per-head function was shaped for exactly this moment: sixteen heads, sixteen independent jobs, disjoint output stripes.
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
use super::simd_cpu::flash_attention_head_neon;
let seq_q = q.shape()[0];
let seq_k = k.shape()[0];
let qw = n_heads * head_dim;
let kvw = n_kv_heads * head_dim;
assert_eq!(q.shape(), &[seq_q, qw][..]);
assert_eq!(k.shape(), &[seq_k, kvw][..]);
assert_eq!(v.shape(), &[seq_k, kvw][..]);
if causal {
assert_eq!(
seq_q, seq_k,
"causal flash attention expects seq_q == seq_k"
);
}
let kv_group = n_heads / n_kv_heads;
let scale = 1.0 / (head_dim as f32).sqrt();
let qd = q.as_f32_slice();
let kd = k.as_f32_slice();
let vd = v.as_f32_slice();
let mut out = vec![0.0f32; seq_q * qw];
let out_ptr = SendPtr(out.as_mut_ptr());
// Heads write disjoint column ranges of the shared output, so the
// parallel raw-pointer writes never overlap.
(0..n_heads).into_par_iter().for_each(|h| {
let dst = &out_ptr;
let kv_off = (h / kv_group) * head_dim;
unsafe {
flash_attention_head_neon(
qd,
kd,
vd,
seq_q,
seq_k,
qw,
kvw,
h * head_dim,
kv_off,
head_dim,
scale,
causal,
dst.0,
qw,
h * head_dim,
);
}
});
Tensor::new(out, vec![seq_q, qw])
}SendPtr is the same tell-the-compiler-I-know wrapper the pooled matvec already uses, and the heads fan out with rayon's parallel iterator from II.4 (the persistent pool's dispatch splits the rows of one matvec; sixteen chunky head-jobs are the shape rayon handles well as-is). The comment above the loop carries the safety argument: disjoint columns, no overlap.
The Parallel backend forwards to the NEON implementation of the SimdCpu it wraps:
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
self.inner
.flash_attention_prefill(q, k, v, n_heads, n_kv_heads, head_dim, causal)
}One call replaces the loop
Now the payoff in the model code. The prefill attention in src/model/common/attention.rs has carried a per-head loop since I.5: slice out the head, matmul against transposed K, scale, mask, softmax, matmul against V, copy the context into place. For the first time in this series, a chapter deletes model code instead of adding it. The loop, its two helper functions, and the mask import are gone, replaced by one line:
pub(crate) fn gqa_attention_forward_with_kv(
ops: &dyn Backend,
x: &Tensor,
q_proj: &Tensor,
k_proj: &Tensor,
v_proj: &Tensor,
o_proj: &Tensor,
attn_q_norm: &Tensor,
attn_k_norm: &Tensor,
num_attention_heads: usize,
num_key_value_heads: usize,
head_dim: usize,
rms_norm_eps: f32,
rope_theta: f32,
) -> (Tensor, Tensor, Tensor) {
let nh = num_attention_heads;
let nkv = num_key_value_heads;
let hd = head_dim;
let mut q = ops.matmul(x, q_proj);
let mut k = ops.matmul(x, k_proj);
let v = ops.matmul(x, v_proj);
q = headwise_rms_norm_weighted(ops, &q, nh, hd, attn_q_norm, rms_norm_eps);
k = headwise_rms_norm_weighted(ops, &k, nkv, hd, attn_k_norm, rms_norm_eps);
let q = ops.apply_rope(&q, hd, rope_theta);
let k_rope = ops.apply_rope(&k, hd, rope_theta);
let merged = ops.flash_attention_prefill(&q, &k_rope, &v, nh, nkv, hd, true);
let out = ops.matmul(&merged, o_proj);
(out, k_rope, v)
}The true is the causal flag. Everything else in the function (projections, QK-norm, RoPE) is untouched, k_rope and v still return to the caller for the KV cache, and gqa_attention_decode_one_query just below it still runs IV.1's fused_attention: decode was already streaming, and this chapter does not touch it.
With the loop gone, nothing in the engine calls causal_mask_upper_tri anymore, except the tests you will meet in a moment, which use it to rebuild the old path as a reference. src/model/common/mask.rs gets an attribute so non-test builds compile clean:
use crate::backend::Backend;
use crate::tensor::Tensor;
#[cfg_attr(not(test), allow(dead_code))]
pub fn causal_mask_upper_tri(ops: &dyn Backend, scores: &Tensor) -> Tensor {
ops.fill_strict_upper_tri(scores, f32::NEG_INFINITY)
}And the TracingBackend from II.3 learns to time the new op, one macro line among its siblings:
trace_fn!(fn fused_attention (&self, q: &Tensor, k: &Tensor, v: &Tensor, n_heads: usize, n_kv_heads: usize, head_dim: usize) -> Tensor);
trace_fn!(fn flash_attention_prefill (&self, q: &Tensor, k: &Tensor, v: &Tensor, n_heads: usize, n_kv_heads: usize, head_dim: usize, causal: bool) -> Tensor);
trace_fn!(fn concat_dim0 (&self, a: &Tensor, b: &Tensor) -> Tensor);That tracing line lets the profiler confirm the deletion. Rerun IV.1's traced histogram (simd backend, 17-token prompt, two generated tokens) and the prefill-side ops collapse: fill_strict_upper_tri falls from 448 calls to 0, transpose_2d from 448 to 0, softmax_rows from 450 to 2 (only the sampling path's calls remain), matmul from 1,290 to 394, and one new op appears: flash_attention_prefill, 28 calls, one per layer. The score matrix does not get built less often. It does not get built.
The backend that said no
The plan was to stop there: a default implementation for correctness, a NEON one for the SIMD backends, heads-in-parallel for the pool, and the Metal backends inheriting the default until someone writes a GPU version. Then the benchmark ran. With the trait default in place for everyone, the metal and gpu backends prefilled the 511-token prompt in about 3.6 seconds. They were at about 1.5 seconds before this chapter, and at 2,000 tokens the regression grew to 3x. The memory-optimal algorithm more than doubled the time to first token.
The explanation is in II.5. On those backends the per-head score matmuls never ran on a CPU core: a [511, 128] by [128, 511] matmul clears the m >= 8 threshold easily, so the score matrix and the scores-times-V product were computed by the GPU's SIMD-group matrix hardware, thousands of lanes doing 8x8 fragment multiply-accumulates, with the mask and softmax riding along between dispatches. The trait's flash loop runs on one CPU core. Adopting it saved the score matrix's memory traffic and paid for that saving by walking away from an order of magnitude more compute. Bandwidth arithmetic only wins when the machine you move to is at least as strong as the machine you left.
So the Metal backend overrides flash_attention_prefill to not flash: it keeps the materialized score path, per head, on the GPU. The old model-side loop does not die after all; it moves into the backend that still wants it, with the mask applied only when causal asks:
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
let seq = q.shape()[0];
let width = n_heads * head_dim;
let group = n_heads / n_kv_heads;
let scale_attn = 1.0 / (head_dim as f32).sqrt();
let mut concat = vec![0.0f32; seq * width];
for h in 0..n_heads {
let kv_h = h / group;
let q_h = self.copy_2d_from_cols(q, head_dim, h * head_dim);
let k_h = self.copy_2d_from_cols(k, head_dim, kv_h * head_dim);
let v_h = self.copy_2d_from_cols(v, head_dim, kv_h * head_dim);
let scores = self.matmul(&q_h, &self.transpose_2d(&k_h));
let scores = self.scale(&scores, scale_attn);
let scores = if causal {
self.fill_strict_upper_tri(&scores, f32::NEG_INFINITY)
} else {
scores
};
let attn = self.softmax_rows(&scores);
let ctx = self.matmul(&attn, &v_h);
self.copy_2d_into_cols(&mut concat, width, &ctx, h * head_dim);
}
Tensor::new(concat, vec![seq, width])
}Every op in that loop is one of Metal's own: the matmuls go to the SIMD-group kernel, and fill_strict_upper_tri, softmax_rows, scale fall back to the CPU exactly as they have since II.5. The gpu backend from IV.4 wraps a Metal and forwards, inheriting the same decision:
fn flash_attention_prefill(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
causal: bool,
) -> Tensor {
self.inner
.flash_attention_prefill(q, k, v, n_heads, n_kv_heads, head_dim, causal)
}This is the second thing the trait-default design bought us. Because flash_attention_prefill is an overridable method rather than a hardcoded call, "which attention algorithm" became a per-backend decision, and the honest benchmark got to make it. There is no universally right attention implementation; there is a right implementation per location, and the location that owns our prefill crown is a GPU whose strength is huge matmuls. The real fix, and the single biggest prefill lever this engine has left, is the one llama.cpp ships: flash attention as a Metal kernel, tiles staged in threadgroup memory, online softmax running in GPU registers, the SIMD-group hardware doing the tile products. That is future work, named honestly. Until it exists, flash runs where it wins (the CPU backends) and the GPU keeps its matmul path.
Pinned to the reference, then a flipped coin
Deleting the old prefill path raises an obvious question: how do we know the streaming math is right? The old path answers it from inside a test module. The chapter's tests reconstruct materialized attention out of backend ops, name it what it now is, the reference, and require every flash implementation to match it:
#[cfg(test)]
mod flash_attention_tests {
use crate::backend::{Backend, CpuBackend, PoolCpu, SimdCpu};
use crate::model::common::mask::causal_mask_upper_tri;
use crate::tensor::Tensor;
/// The pre-flash prefill attention: materialize the [seq x seq] score
/// matrix per head, causal-mask it, softmax, then scores @ V. Kept here as
/// the reference the flash implementations must match.
fn materialized_attention_reference(
ops: &dyn Backend,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
) -> Tensor {
let seq = q.shape()[0];
let qw = n_heads * head_dim;
let kv_group = n_heads / n_kv_heads;
let scale = 1.0 / (head_dim as f32).sqrt();
let mut concat = vec![0.0f32; seq * qw];
for h in 0..n_heads {
let kv_h = h / kv_group;
let q_h = ops.copy_2d_from_cols(q, head_dim, h * head_dim);
let k_h = ops.copy_2d_from_cols(k, head_dim, kv_h * head_dim);
let v_h = ops.copy_2d_from_cols(v, head_dim, kv_h * head_dim);
let scores = ops.matmul(&q_h, &ops.transpose_2d(&k_h));
let scores = ops.scale(&scores, scale);
let scores = causal_mask_upper_tri(ops, &scores);
let attn = ops.softmax_rows(&scores);
let ctx = ops.matmul(&attn, &v_h);
ops.copy_2d_into_cols(&mut concat, qw, &ctx, h * head_dim);
}
Tensor::new(concat, vec![seq, qw])
}(Two small helpers are elided here: pseudo_random, a deterministic bit-mixing generator so the tensors are reproducible without pulling in a crate, and max_abs_diff.) The check runs every implementation against the reference at Qwen3's real head geometry, across sequence lengths that cover the tile boundary cases: 1 (a single row, degenerate tiles), 17 (one partial tile), and 100 (a full tile plus a partial one):
fn check_backend_matches_reference(flash_ops: &dyn Backend, label: &str) {
let cpu = CpuBackend;
let n_heads = 16;
let n_kv_heads = 8;
let head_dim = 128;
for (si, &seq) in [1usize, 17, 100].iter().enumerate() {
let q = Tensor::new(
pseudo_random(11 + si as u32, seq * n_heads * head_dim),
vec![seq, n_heads * head_dim],
);
let k = Tensor::new(
pseudo_random(23 + si as u32, seq * n_kv_heads * head_dim),
vec![seq, n_kv_heads * head_dim],
);
let v = Tensor::new(
pseudo_random(37 + si as u32, seq * n_kv_heads * head_dim),
vec![seq, n_kv_heads * head_dim],
);
let reference =
materialized_attention_reference(&cpu, &q, &k, &v, n_heads, n_kv_heads, head_dim);
let flash =
flash_ops.flash_attention_prefill(&q, &k, &v, n_heads, n_kv_heads, head_dim, true);
let diff = max_abs_diff(&reference, &flash);
eprintln!("{label} flash vs reference at seq {seq}: max abs diff {diff:e}");
assert!(
diff < 1e-4,
"{label} flash attention diverges from materialized reference at seq {seq}: max abs diff {diff}"
);
}
}
#[test]
fn flash_default_matches_materialized_reference() {
check_backend_matches_reference(&CpuBackend, "scalar default");
}
#[test]
fn flash_simd_matches_materialized_reference() {
check_backend_matches_reference(&SimdCpu::new(CpuBackend), "simd");
}
#[test]
fn flash_pool_matches_materialized_reference() {
check_backend_matches_reference(&PoolCpu::new(SimdCpu::new(CpuBackend)), "pool");
}
}The tolerance is 1e-4. The measured differences are on the order of 2e-7 (1.79e-7 at seq 17, 2.38e-7 at seq 100): the two computations agree to within float rounding, which is all that reordered float arithmetic can promise.
And then the transcript changed anyway. The short-prompt benchmark, the baker story from Act 2, has opened the same way for four chapters: " named Mr. Smith. He had a small shop with a small kitchen." With flash attention, the simd backend writes " named Mr. Baker. He had a small shop with a small kitchen." A story about a baker now names its baker Mr. Baker, which is hard to argue with.
The mechanics are the same non-associativity lesson as IV.1, one notch further along. The first three generated tokens are identical with essentially identical probabilities (46.8% versus 46.4%, 5.1% versus 4.8%, 99.18% versus 99.18%). The fourth token is the surname, and there the model has no real opinion: the distribution is nearly flat, the old path's top pick is " Smith" at 5.70%, the flash path's is " Baker" at 6.28%. Two candidates separated by noise-level probability mass, an argmax that must pick one, and a summation order that changed underneath it: the coin lands on the other face, and greedy decoding follows the new leader for the rest of the story. The metal and gpu transcripts are byte-identical before and after, because those backends kept the old path; only the backends that actually changed their arithmetic flipped.
IV.1 met this exact phenomenon and engineered it away, restructuring the fused kernel to reproduce the old summation order bit for bit. That option does not exist here, and it is worth being precise about why: IV.1 changed how the same sums were computed, so there was an old order to imitate. Flash attention changes which sums exist; the rescale-as-you-go structure has no materialized-softmax order to reproduce. When the loop structure is the optimization, you keep the math pinned to a reference within float rounding, you keep the tests, and you accept that ties can land differently. This time, we accept it.
Running it
Same harness, same prompts. The 511-token pangram prompt on the single-core simd backend, IV.4's commit versus this one:
backend: simd
kv cache: basic
-- IV.4 (materialized scores) --
time_to_first_token_ms: 10044.634
decode_tokens_per_second: 21.177
-- IV.5 (flash attention) --
time_to_first_token_ms: 8747.114
decode_tokens_per_second: 21.304Prefill drops from 10.04 s to 8.75 s, 13% off, and decode is untouched (it never ran this code). On the pool backend, where the heads fan out across cores:
backend: pool
kv cache: basic
-- IV.4 (materialized scores) --
time_to_first_token_ms: 2392.510
decode_tokens_per_second: 27.840
-- IV.5 (flash attention) --
time_to_first_token_ms: 1792.267
decode_tokens_per_second: 27.9182.39 s to 1.79 s, a quarter of the prefill gone. And because the deleted term is the quadratic one, the win grows with context. The 2,044-token prompt, pool backend:
backend: pool
kv cache: basic
-- IV.4 (materialized scores) --
time_to_first_token_ms: 17210.213
decode_tokens_per_second: 5.969
-- IV.5 (flash attention) --
time_to_first_token_ms: 11074.357
decode_tokens_per_second: 5.91317.2 s to 11.1 s, 36% off. Single-core simd shows the same curve: 13% at 511 tokens becomes 26% at 2k (72.2 s down to 53.3 s). Extrapolate and the story is clear: the longer the prompt, the more of prefill is the score matrix, and flash's share of the win keeps climbing.
The gated backends, meanwhile, do exactly what the gate promises. Metal's 511-token prefill:
backend: metal
kv cache: basic
time_to_first_token_ms: 1517.821Within run noise of the 1.53 s it posted before this chapter, and still ahead of the pool's new best of 1.79 s: the GPU prefill crown does not move until flash becomes a kernel. And the regression check that matters most after IV.4, GPU decode on the short prompt:
backend: gpu
kv cache: basic
time_to_first_token_ms: 283.380
decode_tokens_per_second: 115.512115 tokens/sec, unaffected. Decode never touches flash_attention_prefill; the one-command-buffer path is exactly as fast as we left it.
The end of the act
Flash attention closes Act 4 the way the act opened: with a structural change, measured honestly, including the backend where it lost. The score matrix is gone from every path that was actually paying for it, the win compounds with context length exactly as the algebra says it should, and the one backend that kept the matrix kept it for the right reason, measured twice. The engine now decodes at 115 tokens per second on the GPU and 109 on four CPU cores, prefills a 511-token prompt in 1.5 seconds, and its remaining gaps have names and owners. The Act 4 recap takes stock of the whole ladder, tallies what we deliberately did not build, and puts the final numbers next to llama.cpp's.