IV.1: Fused attention
Act 2 ended with fast kernels: an integer dot product for the weights, a SIMD-group matrix kernel for the GPU. The act intro showed that llama.cpp beats us anyway, on the same file and the same laptop. If the kernels are fine, the next place to look is what the forward pass does with them, and for that we need a new measurement. Not how long each op takes, which the profiler from II.3 already prints, but how often each op gets called.
Teaching the profiler to count
The TracingBackend from II.3 wraps every Backend method and emits one timed trace line per call. Counting is a small addition to the same file. src/backend/tracing.rs gains a mutex-guarded map from op name to count:
use super::Backend;
use crate::tensor::Tensor;
use std::cell::Cell;
use std::collections::BTreeMap;
use std::sync::Mutex;
use std::time::Instant;pub struct TracingBackend<B: Backend> {
inner: B,
name: String,
counts: Mutex<BTreeMap<&'static str, u64>>,
}
impl<B: Backend> TracingBackend<B> {
pub fn new(inner: B) -> Self {
let name = inner.name().to_uppercase();
Self {
inner,
name,
counts: Mutex::new(BTreeMap::new()),
}
}
fn bump(&self, op: &'static str) {
*self.counts.lock().unwrap().entry(op).or_insert(0) += 1;
}
}
impl<B: Backend> Drop for TracingBackend<B> {
fn drop(&mut self) {
let counts = self.counts.lock().unwrap();
let total: u64 = counts.values().sum();
tracing::trace!("[{}] op invocation counts (total {total}):", self.name);
for (op, n) in counts.iter() {
tracing::trace!(" {op}: {n}");
}
}
}bump increments one op's counter; the Drop impl prints the whole histogram when the backend is torn down at process exit, sorted by name (that's the BTreeMap). The trace_fn! macro that generates every wrapper method gets one new line, so every op call bumps its counter before running:
// ...
self.bump(stringify!($fn));
let t0 = Instant::now();
let result = self.inner.$fn($($arg),*);
// ...Like the rest of the profiler, this only attaches under RUST_LOG=trace, so normal runs pay nothing.
What one token costs
The counter itself lands in this chapter's commit, but the per-call trace lines have been there since II.3, so we can tally the old code by counting lines. Run the Act 3 tip with RUST_LOG=trace, a one-token prompt, and two generated tokens, then slice out the trace of the second forward pass (the lone decode step) and count each op's lines:
2016 copy_2d_from_cols
1486 hadamard
1458 reshape_data
1177 scale
1120 copy_2d_into_cols
1093 matmul
729 sum_squares_axis
729 rsqrt_elem
729 repeat_row_as_matrix
729 broadcast_row_scalars
729 add_scalar
449 softmax_rows
448 transpose_2d
56 concat_dim0
56 apply_rope_single_row
56 add
28 silu
1 gather_rows
1 copy_row_2d
1 argmax_with_probThat is 13,091 backend calls to produce one token. Look at the matmul line: 1,093 calls. The model has 28 layers, and each layer's real work is seven weight matmuls (q, k, v, and output projections, plus three in the MLP), which with the final output head accounts for 197. The other 896 are the per-head attention loop from II.2: 16 heads x 2 matmuls (scores, then context) x 28 layers. Attention is also where the 448 transpose_2d calls come from (one per head per layer), 448 of the 449 softmax_rows, 448 of the 1,177 scale calls, 1,344 of the copy_2d_from_cols (slicing each head's q, k, and v out of the full matrices), and 448 of the copy_2d_into_cols (writing each head's result back). Add it up and the decode attention loop is about 4,000 backend calls per token.
Now look at the size of that work. A per-head score matmul at short context is a [1, 128] query against a few dozen cached keys: on the order of a few thousand multiplies. The trace timings show it finishing in single-digit microseconds. Meanwhile every one of those calls pays the same fixed costs as a real matmul: allocate a fresh Tensor (a heap Vec), check shapes, dispatch through the trait, and in a traced run, format its arguments. The overhead is charged per call, and we are making roughly 4,000 calls for work that is a handful of numbers deep.
The fix: one op for all of attention
The per-head loop exists because II.2 built decode attention out of the ops the Backend trait already had: slice a head, transpose, matmul, scale, softmax, matmul, copy back. The fix is to make attention itself a backend op. One call, fused_attention, takes the query row, the whole cached K, the whole cached V, and the head geometry, and returns the merged attention output for all heads. This is kernel fusion: replacing a chain of small ops, each allocating and re-reading memory, with one op that does the whole computation in place.
The trait gets the new method, and the important design choice is a default implementation: the old per-head loop, verbatim, built from the same primitive ops. src/backend/backend_trait.rs:
fn matmul(&self, a: &Tensor, b: &Tensor) -> Tensor;
fn fused_attention(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
) -> Tensor {
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 out = vec![0.0f32; qw];
for h in 0..n_heads {
let kv_h = h / kv_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);
let attn = self.softmax_rows(&scores);
let ctx = self.matmul(&attn, &v_h);
self.copy_2d_into_cols(&mut out, qw, &ctx, h * head_dim);
}
Tensor::new(out, vec![1, qw])
}
fn gather_rows(&self, table: &Tensor, row_indices: &[usize]) -> Tensor;kv_group = n_heads / n_kv_heads is the GQA mapping: Qwen3 has 16 query heads sharing 8 key/value heads, so two query heads read each cached K/V head, and kv_h = h / kv_group picks the right one. Because the default reproduces the old math with the old ops, every backend (scalar, parallel, Metal) stays correct the moment the call site switches over. Backends that care can override it. One does.
The SIMD kernel
SimdCpu overrides fused_attention with a flat-buffer version: no Tensor allocations in the loop, no trait dispatch, just three scratch Vecs reused across all 16 heads and the NEON saxpy_row_f32 kernel that II.3 built for matmul. src/backend/simd_cpu.rs:
fn fused_attention(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
) -> Tensor {
let seq = k.shape()[0];
let qw = n_heads * head_dim;
let kvw = n_kv_heads * head_dim;
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; qw];
let mut scores = vec![0.0f32; seq];
let mut k_t = vec![0.0f32; head_dim * seq];
for h in 0..n_heads {
let kv_h = h / kv_group;
let q_h = &qd[h * head_dim..h * head_dim + head_dim];
let kv_off = kv_h * head_dim;
for t in 0..seq {
let k_row = &kd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
for d in 0..head_dim {
k_t[d * seq + t] = k_row[d];
}
}
scores.iter_mut().for_each(|s| *s = 0.0);
for d in 0..head_dim {
Self::saxpy_row_f32(scores.as_mut_ptr(), k_t[d * seq..].as_ptr(), q_h[d], seq);
}
for s in &mut scores {
*s *= scale;
}
let mut max = f32::NEG_INFINITY;
for &s in &scores {
if s > max {
max = s;
}
}
let mut sum = 0.0f32;
for s in &mut scores {
let e = (*s - max).exp();
*s = e;
sum += e;
}
for s in &mut scores {
*s /= sum;
}
let dst = &mut out[h * head_dim..h * head_dim + head_dim];
for t in 0..seq {
let v_row = &vd[t * kvw + kv_off..t * kvw + kv_off + head_dim];
Self::saxpy_row_f32(dst.as_mut_ptr(), v_row.as_ptr(), scores[t], head_dim);
}
}
Tensor::new(out, vec![1, qw])
}Read it top to bottom and it is the same four steps as before, per head: scores, scale, softmax, weighted sum of values. The differences are all about allocation and layout. q_h is a slice into the query row rather than a copied tensor. The K transpose is written into the reused k_t scratch buffer, so the score computation can run as head_dim SAXPY passes over rows of the transposed keys, the exact loop structure the SIMD matmul uses. The softmax is inlined with the same subtract-the-max trick as the backend's softmax_rows. And the context accumulates directly into this head's 128-float slice of out, one saxpy_row_f32 per cached token. Sixteen heads, zero intermediate tensors.
The call site collapses. gqa_attention_decode_one_query in src/model/common/attention.rs loses its head loop (and the helper gqa_attention_context_one_head_decode, which is deleted outright) and becomes a wrapper: one fused call, then the output projection:
pub(crate) fn gqa_attention_decode_one_query(
ops: &dyn Backend,
q: &Tensor,
k_cache: &Tensor,
v_cache: &Tensor,
o_proj: &Tensor,
num_attention_heads: usize,
num_key_value_heads: usize,
head_dim: usize,
) -> Tensor {
let merged = ops.fused_attention(
q,
k_cache,
v_cache,
num_attention_heads,
num_key_value_heads,
head_dim,
);
ops.matmul(&merged, o_proj)
}Every caller (the single-token decode path from II.2, the batched decode path from III.7) picks this up without changing a line. Finally the profiler learns to trace the new op, one more trace_fn! in the list:
trace_fn!(fn matmul (&self, a: &Tensor, b: &Tensor) -> Tensor);
trace_fn!(fn fused_attention (&self, q: &Tensor, k: &Tensor, v: &Tensor, n_heads: usize, n_kv_heads: usize, head_dim: usize) -> Tensor);The output must not change
The kernel above is the second version, and the difference between it and the first is worth a section, because the first one was wrong in a way no test caught.
The first version computed each query-key score the obvious fast way: NEON fused multiply-adds into four accumulator lanes, then a horizontal add, the dot-product idiom from II.3. Cleaner and slightly quicker than transposing K into a scratch buffer. It also changed the model's output: on the standard test prompt the generated text matched the baseline for 47 tokens and then diverged at token 48.
Nothing was broken in the usual sense. Floating-point addition is not associative: (a + b) + c and a + (b + c) can differ in the last bit, and a lane-wise tree reduction adds the same 128 products in a different order than the old path's sequential SAXPY accumulation. The logits differed from the baseline by something like one part in ten million, which is far below any tolerance a numerical test would use. But greedy decoding takes an argmax over 151,936 logits, and at token 48 the top two happened to be nearly tied. A last-bit wiggle flipped which one won, and every token after that followed the new leader.
The shipped kernel therefore does the transpose-then-SAXPY dance you see above, not because it is the fastest possible dot product but because it reproduces the exact float summation order of the code it replaces: the old path computed scores as a matmul against a transposed K, accumulating over the 128 dimensions in sequence, and the new one does the same additions in the same order. The result is byte-identical logits and byte-identical output, which is what lets the benchmark below claim a pure speedup rather than a speedup plus a mystery. Determinism across a refactor is not something floats give you for free; it is a property you choose and then engineer for.
Running it
Same harness, same prompts as Act 2. First the profiler, to confirm the calls are actually gone. The histogram from a traced run (one-token prompt, two generated tokens, so one prefill pass plus one decode pass) now prints at exit:
[SIMD] op invocation counts (total 22570):
add: 112
add_scalar: 1458
apply_rope: 56
apply_rope_single_row: 56
argmax_with_prob: 2
broadcast_row_scalars: 1458
concat_dim0: 56
copy_2d_from_cols: 2688
copy_2d_into_cols: 1792
copy_row_2d: 2
fill_strict_upper_tri: 448
fused_attention: 28
gather_rows: 2
hadamard: 2972
matmul: 1290
repeat_row_as_matrix: 1458
reshape_data: 2916
rsqrt_elem: 1458
scale: 1906
silu: 56
softmax_rows: 450
sum_squares_axis: 1458
transpose_2d: 448Two passes are mixed together here, so read it carefully. The prefill pass still runs the per-head loop with the causal mask (that is where all 448 transpose_2d and 448 fill_strict_upper_tri calls now live; this act returns to prefill in IV.5). The decode pass is the transformed one: slice out its ops the same way as before and it makes 197 matmul calls plus 28 fused_attention, and 9,087 backend calls in total, down from 13,091. Per decode step, 1,093 matmuls became 197, and roughly 4,000 attention bookkeeping calls became 28.
Now the clock. Short prompt (17 tokens, 64 generated), before:
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 277.890
decode_tokens_per_second: 41.286
per_forward_ms: min 22.048 max 277.890 mean 28.185 (n=64)After:
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 263.908
decode_tokens_per_second: 44.135
per_forward_ms: min 21.503 max 263.908 mean 26.428 (n=64)Decode goes from 41.3 to 44.1 tokens/sec, +7%. (Time to first token is prefill's number and this chapter didn't touch prefill; the difference is run-to-run noise.) And the generated text is token-for-token identical, which after the determinism section you know was the hard part.
A +7% payoff for deleting 82% of the matmul calls deserves an explanation, and the long-context run provides it. Same measurement with the 511-token prompt from II.4, before:
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 9973.194
decode_tokens_per_second: 15.900
per_forward_ms: min 56.443 max 9973.194 mean 217.740 (n=64)After:
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 10030.769
decode_tokens_per_second: 20.589
per_forward_ms: min 42.389 max 10030.769 mean 204.541 (n=64)15.9 to 20.6 tokens/sec, +30%.
Call count is not time share
The two numbers together are the lesson. At short context, 82% of the matmul calls were attention, but the seven big weight matmuls per layer dominate wall time: each one reads megabytes of Q8_0 weights while a per-head attention matmul touches a few kilobytes of cache. Deleting thousands of cheap calls only trims the margin. At 511-token context the picture shifts: every head's slice-transpose-matmul-softmax chain now walks 500-plus cached rows through tensor allocations and copies, so the bookkeeping is real work, and killing it buys 30%. The overhead you delete scales with context length, which is exactly why the fused kernel matters more the longer the conversation gets.
The histogram has one more thing to say, though. Look back at it: concat_dim0: 56. Every decode step still calls it 56 times, and each call rebuilds an entire cache tensor just to append one row. That is the next chapter.