IV.4: One command buffer per token
II.5 put the big matmuls on the GPU and II.6 taught the GPU to read quantized weights; together they made --backend metal the fastest prefill in the engine. Decode is another story, and it's time to tell it straight: the metal backend has never decoded faster than the plain CPU. On the 511-token prompt it decodes at 16.1 tokens/sec where single-core simd gets 24.3, and where IV.3's pool gets 34.7.
The embarrassing part is that during decode the metal backend doesn't use the GPU at all. Every decode matmul has one output row, m == 1, and II.5's MIN_M_FOR_GPU_MATMUL threshold sends it to the CPU fallback. That threshold exists for a good reason: II.5's dispatch builds a command buffer, submits it, and blocks in wait_until_completed for every single operation. A decode step needs 197 matvecs (7 per layer, 28 layers, plus the lm_head), so decoding on the GPU that way would mean about 200 dispatch-and-wait round trips per token, each costing a fraction of a millisecond before any arithmetic happens, with every non-matmul op (norms, RoPE, softmax, residual adds) running on the CPU in between. II.6 measured what per-row dispatching does to prefill: 23 seconds instead of 1.5. So decode stayed home, and "GPU decode" meant CPU decode wearing a Metal jacket. (Slower CPU decode, in fact: the fallback path composes attention from primitive ops through the Backend trait's default fused_attention, not the NEON kernel IV.1 gave SimdCpu, which is the 16.1-versus-24.3 gap at long context.)
This is the same disease IV.3 just cured on the CPU: per-job overhead comparable to the job. The GPU's version is two orders of magnitude worse, and the cure is the same shape. Stop asking per operation. Ask once per token.
What llama.cpp does, and what we do instead
llama.cpp's Metal backend does not dispatch a matmul and wait. It walks the model's compute graph and encodes every operation of the forward pass (every matmul, norm, RoPE, attention, residual add, all of it) into a single command buffer, submits that, and waits once. The GPU chews through hundreds of kernel launches back to back with no CPU in the loop; launch-to-launch cost inside a committed command buffer is measured in microseconds, not the fraction of a millisecond a synchronous round trip costs.
llama.cpp needs a general graph representation for this, because it runs hundreds of architectures. We run one. Qwen3's decode step is a fixed sequence of operations that we already wrote out by hand in forward_decode_with_kv_cache, so instead of building a graph engine we will hand-encode that fixed sequence: one readable Rust function that encodes 507 kernel dispatches (18 per layer times 28 layers, plus a final norm, the lm_head, and an argmax) into one command buffer, commits, waits once, and reads back a single integer, the next token id. That is the honest tutorial trade: less general than llama.cpp, and much easier to see through.
Getting there needs three new files under src/backend/metal/: decode_shaders.rs (eight small kernels covering every op decode performs), decoder.rs (the encoder and the resident GPU state), and gpu.rs (a sixth backend, gpu, that is metal-for-prefill plus the new decoder). Plus a thin seam through the model so a backend can claim a whole decode step.
For this to be worth anything, two more round trips have to die alongside the per-op waits: weights and KV cache. The weights get uploaded to GPU buffers once at startup and stay resident. The KV cache gets GPU-resident buffers that the GPU itself appends to each step, so no cache bytes cross per token. Apple Silicon's unified memory makes all of this cheap to arrange (the "upload" is a memcpy within one physical RAM pool, and readback of the token id is a pointer dereference), but the design would be the same on a discrete card; it would just cost more to get wrong.
Eight kernels
src/backend/metal/decode_shaders.rs holds the MSL source as a Rust string constant, same convention as II.5's shaders.rs. The preamble:
pub const DECODE_SHADERS: &str = r#"
#include <metal_stdlib>
using namespace metal;
constant uint BLOCK_QLEN = 32u;
constant uint BLOCK_BYTES = 34u;
constant uint ATTN_MAX_SEQ = 4096u;
inline float decode_scale(device const uchar* base) {
ushort bits = ushort(base[0]) | (ushort(base[1]) << 8);
return float(as_type<half>(bits));
}BLOCK_QLEN and BLOCK_BYTES are the Q8_0 block geometry from II.6 (32 quants, 2-byte f16 scale), and decode_scale is the same two-byte-to-half trick II.6's prefill kernel uses. ATTN_MAX_SEQ caps how long a context the attention kernel handles; it comes back below.
Every kernel here is decode-shaped: single token, single row, matrix-vector work. That's why II.5's tiled SIMD-group matmul kernel is the wrong tool (its 32x32 output tiles would be 97% padding on a one-row output) and these are the right ones.
q8_matvec is the workhorse, dispatched 197 times per step (seven projections per layer times 28 layers, plus the lm_head):
kernel void q8_matvec(
device const float* x [[buffer(0)]],
device const uchar* w [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& n [[buffer(3)]],
constant uint& p [[buffer(4)]],
uint row [[threadgroup_position_in_grid]],
uint lane [[thread_index_in_threadgroup]]
) {
if (row >= p) return;
const uint bpr = n / BLOCK_QLEN;
device const uchar* wrow = w + row * bpr * BLOCK_BYTES;
float sum = 0.0f;
for (uint b = lane; b < bpr; b += 32u) {
device const uchar* blk = wrow + b * BLOCK_BYTES;
device const float* xb = x + b * BLOCK_QLEN;
float acc = 0.0f;
for (uint j = 0; j < BLOCK_QLEN; j++) {
acc += float(int(char(blk[2u + j]))) * xb[j];
}
sum += decode_scale(blk) * acc;
}
sum = simd_sum(sum);
if (lane == 0) out[row] = sum;
}One threadgroup per output element, and each threadgroup is exactly one SIMD-group: 32 threads in lockstep (II.5 defined SIMD-groups; this kernel is nothing but one). Lane l handles blocks l, l+32, l+64, ... of the weight row: dequantize-in-the-dot, accumulating quant * x within a block and scaling once per block, exactly the fused shape II.6 established on the CPU. simd_sum then folds the 32 lane totals into one number without touching memory (cross-lane reduction is what SIMD-group hardware is for), and lane 0 writes the output element. Note what this kernel does not share with the CPU's sdot kernel: the CPU quantizes the activations and multiplies integers; this multiplies floats. Same mathematical dot product, different rounding and different summation order. Hold that thought for the correctness section.
rms_norm_row covers all three RMSNorm flavors in the model (I.5's layer norms and the per-head QK-norm):
kernel void rms_norm_row(
device const float* x [[buffer(0)]],
device const float* wgt [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& dim [[buffer(3)]],
constant float& eps [[buffer(4)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]
) {
threadgroup float partial[4];
device const float* xr = x + row * dim;
float ss = 0.0f;
for (uint i = tid; i < dim; i += 128u) {
float v = xr[i];
ss += v * v;
}
ss = simd_sum(ss);
if ((tid & 31u) == 0u) partial[tid >> 5] = ss;
threadgroup_barrier(mem_flags::mem_threadgroup);
float total = partial[0] + partial[1] + partial[2] + partial[3];
float inv = rsqrt(total / float(dim) + eps);
for (uint i = tid; i < dim; i += 128u) {
out[row * dim + i] = xr[i] * inv * wgt[i];
}
}One threadgroup of 128 threads (four SIMD-groups) per row. Each thread strides through the row accumulating squares, simd_sum reduces within each SIMD-group, the four per-group results land in threadgroup memory, a barrier, and then every thread computes the same total and inv and scales its strided share of the row. Launched with a grid of 1 for the layer norms (one 1024-wide row), or a grid of 16 or 8 rows of head_dim for the QK-norms: same kernel, three shapes.
rope_row is I.5's rotation, one threadgroup per head, one thread per rotated pair:
kernel void rope_row(
device float* x [[buffer(0)]],
constant uint& head_dim [[buffer(1)]],
constant uint& pos [[buffer(2)]],
constant float& theta [[buffer(3)]],
uint head [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]
) {
const uint half_dim = head_dim / 2u;
if (tid >= half_dim) return;
const uint base = head * head_dim;
float freq = 1.0f / precise::pow(theta, float(tid) / float(half_dim));
float ang = float(pos) * freq;
float c = precise::cos(ang);
float s = precise::sin(ang);
float a = x[base + tid];
float b = x[base + tid + half_dim];
x[base + tid] = a * c - b * s;
x[base + tid + half_dim] = a * s + b * c;
}The precise:: prefixes matter: Metal's default transcendental functions are fast approximations, and RoPE angles at position 500 are large enough for the approximation error to show up in the KV cache. precise:: selects the slower, correctly-rounded versions, keeping GPU rotations within float noise of the CPU's.
append_kv is cache.push_row for the GPU-resident cache: copy the freshly computed K and V rows into slot pos of the cache buffers.
kernel void append_kv(
device const float* k [[buffer(0)]],
device const float* v [[buffer(1)]],
device float* kc [[buffer(2)]],
device float* vc [[buffer(3)]],
constant uint& width [[buffer(4)]],
constant uint& pos [[buffer(5)]],
uint i [[thread_position_in_grid]]
) {
if (i >= width) return;
kc[pos * width + i] = k[i];
vc[pos * width + i] = v[i];
}This one-line-per-thread kernel is load-bearing: because the GPU appends to its own cache, no KV bytes cross between CPU and GPU during steady-state decode. Ever.
attention_decode is the GPU twin of IV.1's fused CPU attention, one threadgroup of 128 threads per query head:
kernel void attention_decode(
device const float* q [[buffer(0)]],
device const float* kc [[buffer(1)]],
device const float* vc [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& seq [[buffer(4)]],
constant uint& kv_width [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
constant uint& kv_group [[buffer(7)]],
uint head [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]
) {
threadgroup float scores[ATTN_MAX_SEQ];
threadgroup float red[4];
const uint kv_off = (head / kv_group) * head_dim;
device const float* qh = q + head * head_dim;
const float scale = 1.0f / precise::sqrt(float(head_dim));
for (uint t = tid; t < seq; t += 128u) {
device const float* krow = kc + t * kv_width + kv_off;
float d = 0.0f;
for (uint j = 0; j < head_dim; j++) d += qh[j] * krow[j];
scores[t] = d * scale;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float m = -INFINITY;
for (uint t = tid; t < seq; t += 128u) m = max(m, scores[t]);
m = simd_max(m);
if ((tid & 31u) == 0u) red[tid >> 5] = m;
threadgroup_barrier(mem_flags::mem_threadgroup);
m = max(max(red[0], red[1]), max(red[2], red[3]));
threadgroup_barrier(mem_flags::mem_threadgroup);
float s = 0.0f;
for (uint t = tid; t < seq; t += 128u) {
float e = exp(scores[t] - m);
scores[t] = e;
s += e;
}
s = simd_sum(s);
if ((tid & 31u) == 0u) red[tid >> 5] = s;
threadgroup_barrier(mem_flags::mem_threadgroup);
s = red[0] + red[1] + red[2] + red[3];
const float inv = 1.0f / s;
if (tid < head_dim) {
float acc = 0.0f;
for (uint t = 0; t < seq; t++) acc += scores[t] * vc[t * kv_width + kv_off + tid];
out[head * head_dim + tid] = acc * inv;
}
}kv_off implements GQA (I.5): query head h reads KV head h / kv_group. Then three phases, separated by barriers, all inside one threadgroup. Phase one: the 128 threads stride over the seq cached positions, each computing full query-key dot products into the shared scores array. Phase two: softmax, done as two passes over scores, first the max (the subtract-the-max trick from I.4, reduced with simd_max per SIMD-group and then across the four groups via the tiny red array), then the exponentials and their sum, with the same two-level reduction. Phase three: the first head_dim threads each own one output element and walk the whole value cache accumulating scores[t] * v, scaling by 1/s at the end.
The threadgroup float scores[ATTN_MAX_SEQ] array is why the 4096 cap exists: 4096 floats is 16 KB of threadgroup memory, and the whole score row must fit on-chip for the two softmax passes to reread it cheaply. Contexts longer than 4096 fall back to CPU decode. (Holding the entire score row at once is exactly what flash attention exists to avoid; that is the act's last chapter, not this one.)
The last three kernels are the small fry, elementwise ops that exist so nothing in the step needs the CPU:
kernel void add_inplace(
device float* x [[buffer(0)]],
device const float* y [[buffer(1)]],
constant uint& n [[buffer(2)]],
uint i [[thread_position_in_grid]]
) {
if (i < n) x[i] += y[i];
}
kernel void silu_hadamard(
device const float* gate [[buffer(0)]],
device const float* up [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint i [[thread_position_in_grid]]
) {
if (i >= n) return;
float g = gate[i];
out[i] = (g / (1.0f + exp(-g))) * up[i];
}add_inplace is the residual add; silu_hadamard fuses I.4's silu and Hadamard product into one pass, the SwiGLU gate in a single kernel.
And argmax_row, which exists so the only thing read back per token is four bytes:
kernel void argmax_row(
device const float* x [[buffer(0)]],
device uint* out [[buffer(1)]],
constant uint& n [[buffer(2)]],
uint tid [[thread_index_in_threadgroup]]
) {
threadgroup float bval[8];
threadgroup uint bidx[8];
float best = -INFINITY;
uint besti = 0xFFFFFFFFu;
for (uint i = tid; i < n; i += 256u) {
float v = x[i];
if (v > best || (v == best && i < besti)) {
best = v;
besti = i;
}
}
for (uint off = 16u; off > 0u; off >>= 1) {
float ov = simd_shuffle_down(best, off);
uint oi = simd_shuffle_down(besti, off);
if (ov > best || (ov == best && oi < besti)) {
best = ov;
besti = oi;
}
}
if ((tid & 31u) == 0u) {
bval[tid >> 5] = best;
bidx[tid >> 5] = besti;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid == 0u) {
for (uint s = 1u; s < 8u; s++) {
if (bval[s] > best || (bval[s] == best && bidx[s] < besti)) {
best = bval[s];
besti = bidx[s];
}
}
out[0] = besti;
}
}
"#;256 threads scan the 151,936 logits in strides, each keeping its best (value, index) pair; simd_shuffle_down (pass a value directly from a higher lane to a lower one, no memory) folds each SIMD-group in five steps; eight per-group champions meet in threadgroup memory and thread 0 picks the winner. Every comparison breaks ties toward the lower index, because that's what the CPU's argmax_with_prob does; greedy decoding (I.6) must not depend on which processor ran it.
The decoder: kernels and helpers
src/backend/metal/decoder.rs starts with the plumbing:
use std::collections::HashMap;
use mtl_gpu::device as mtl_device;
use mtl_gpu::{
Buffer, CommandQueue, ComputeCommandEncoder, ComputePipelineState, Device, ResourceOptions,
Size,
};
use super::decode_shaders::DECODE_SHADERS;
use crate::backend::backend_trait::GpuDecodePlan;
use crate::cache::KvCache;
use crate::tensor::q8_0::blocks_packed_byte_len;
use crate::tensor::{Tensor, TensorData};
const DECODE_KERNELS: &[&str] = &[
"q8_matvec",
"rms_norm_row",
"rope_row",
"append_kv",
"attention_decode",
"add_inplace",
"silu_hadamard",
"argmax_row",
];
const ATTN_MAX_SEQ: usize = 4096;
const KV_HEADROOM_ROWS: usize = 1024;
#[derive(Clone, Copy)]
pub(crate) enum Scalar {
U(u32),
F(f32),
}
pub(crate) struct Kernels {
device: Device,
queue: CommandQueue,
pipelines: HashMap<&'static str, ComputePipelineState>,
}Kernels is this module's MetalContext: device, queue, and a compiled compute pipeline per kernel name. The Scalar enum is a small upgrade over II.5's u32-only set_u32; some kernels take floats now (eps, theta), so scalar arguments carry their type. ATTN_MAX_SEQ mirrors the MSL constant, and KV_HEADROOM_ROWS sizes cache over-allocation, both used shortly.
impl Kernels {
pub fn try_new() -> Option<Self> {
let device = mtl_device::system_default()?;
let library = match device.new_library_with_source(DECODE_SHADERS, None) {
Ok(lib) => lib,
Err(e) => {
eprintln!("gpu decode shaders failed to compile: {e:?}");
return None;
}
};
let mut pipelines = HashMap::new();
for name in DECODE_KERNELS {
let func = library.new_function_with_name(name)?;
let pipeline = device
.new_compute_pipeline_state_with_function(&func)
.ok()?;
pipelines.insert(*name, pipeline);
}
let queue = device.new_command_queue()?;
Some(Self {
device,
queue,
pipelines,
})
}Unlike MetalContext::new, which unwraps its way through setup, try_new returns an Option. Every failure path (no GPU, shader compile error, missing kernel) becomes None, and None will mean "decode on the CPU instead." The GPU decoder is an offer, not a requirement; the engine still runs on a machine with no Metal at all.
fn alloc(&self, bytes: usize) -> Buffer {
self.device
.new_buffer(bytes.max(4), ResourceOptions::STORAGE_MODE_SHARED)
.expect("Metal buffer allocation failed")
}
fn upload(&self, bytes: &[u8]) -> Buffer {
self.device
.new_buffer_with_bytes(bytes, ResourceOptions::STORAGE_MODE_SHARED)
.expect("Metal buffer upload failed")
}
fn upload_f32(&self, data: &[f32]) -> Buffer {
let bytes =
unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4) };
self.upload(bytes)
}
fn write_f32(&self, buf: &Buffer, data: &[f32]) {
let dst = buf.contents().expect("buffer has no CPU mapping") as *mut f32;
unsafe { std::ptr::copy_nonoverlapping(data.as_ptr(), dst, data.len()) };
}
fn read_f32(&self, buf: &Buffer, len: usize) -> Vec<f32> {
let src = buf.contents().expect("buffer has no CPU mapping") as *const f32;
let mut out = vec![0.0f32; len];
unsafe { std::ptr::copy_nonoverlapping(src, out.as_mut_ptr(), len) };
out
}
fn read_u32(&self, buf: &Buffer) -> u32 {
let src = buf.contents().expect("buffer has no CPU mapping") as *const u32;
unsafe { *src }
}Buffer traffic, all through STORAGE_MODE_SHARED unified memory (II.5). alloc makes an empty GPU-owned buffer; upload makes one initialized with a copy of the given bytes. Note the shift from II.5: there we wrapped short-lived CPU vectors with newBufferWithBytesNoCopy because the data changed every call. Here the interesting data lives for the whole process, so we copy it into GPU-owned buffers once and never think about Rust lifetimes versus GPU lifetimes again. write_f32, read_f32, and read_u32 go through contents(), the buffer's CPU-visible pointer; on unified memory that's a plain memcpy or, for read_u32, literally one load. That one load is the entire per-token readback.
fn encode(
&self,
enc: &ComputeCommandEncoder,
kernel: &str,
bufs: &[&Buffer],
scalars: &[Scalar],
grid: Size,
threads: Size,
) {
enc.set_compute_pipeline_state(&self.pipelines[kernel]);
for (i, buf) in bufs.iter().enumerate() {
enc.set_buffer(buf, 0, i);
}
for (i, s) in scalars.iter().enumerate() {
let raw = match s {
Scalar::U(v) => v.to_ne_bytes(),
Scalar::F(v) => v.to_ne_bytes(),
};
enc.set_bytes(&raw, bufs.len() + i);
}
enc.dispatch_threadgroups(grid, threads);
}
fn run_once(
&self,
kernel: &str,
bufs: &[&Buffer],
scalars: &[Scalar],
grid: Size,
threads: Size,
) {
let cmd = self.queue.command_buffer().unwrap();
let enc =
unsafe { ComputeCommandEncoder::from_raw(cmd.compute_command_encoder()) }.unwrap();
self.encode(&enc, kernel, bufs, scalars, grid, threads);
enc.end_encoding();
cmd.commit();
cmd.wait_until_completed();
}
}Compare encode to II.5's dispatch and you're looking at the whole chapter in two function signatures. dispatch created a command buffer, encoded one kernel, committed, and waited. encode takes an encoder someone else owns and just adds one dispatch to it: no command buffer, no commit, no wait. The caller decides when to pay the round trip. Successive dispatches in one command buffer run in order with each one's writes visible to the next, which is exactly the dependency chain a forward pass needs. run_once is dispatch rebuilt from encode (create, encode one, commit, wait); the decode path never calls it, but the kernel unit tests below live on it.
Resident state
struct LayerBufs {
attn_norm: Buffer,
q_proj: Buffer,
k_proj: Buffer,
v_proj: Buffer,
o_proj: Buffer,
q_norm: Buffer,
k_norm: Buffer,
ffn_norm: Buffer,
gate_proj: Buffer,
up_proj: Buffer,
down_proj: Buffer,
}
pub(crate) struct GpuDecoder {
kernels: Kernels,
layers: Vec<LayerBufs>,
out_norm: Buffer,
lm_head: Buffer,
x: Buffer,
xn: Buffer,
q: Buffer,
k: Buffer,
v: Buffer,
attn: Buffer,
resid: Buffer,
gate: Buffer,
up: Buffer,
ffn_h: Buffer,
logits: Buffer,
amax: Buffer,
kcache: Vec<Buffer>,
vcache: Vec<Buffer>,
cache_cap: usize,
next_pos: usize,
hidden: usize,
q_width: usize,
kv_width: usize,
inter: usize,
vocab: usize,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
eps: f32,
theta: f32,
}GpuDecoder is everything that stays on the GPU between tokens, in three groups. Weights: one LayerBufs per layer (eleven buffers, matching the eleven weight tensors a Qwen3 layer owns) plus out_norm and lm_head, uploaded once. Activations: x is the residual stream and xn its normed partner (the two ping-pong through the step: norms read x and write xn, projections read xn, adds write back into x), then q/k/v/attn/resid/gate/up/ffn_h/logits/amax, each a scratch buffer sized once and reused by all 28 layers and every subsequent token; the step allocates nothing. Cache: kcache/vcache, one K and one V buffer per layer, with cache_cap rows allocated and next_pos tracking how many are filled.
Two free helpers handle the data conversions at the boundary:
fn q8_bytes<'a>(t: &'a Tensor, rows: usize, cols: usize) -> Option<&'a [u8]> {
match t.as_data() {
TensorData::Q8_0(blocks) => {
if t.shape() != [rows, cols] {
return None;
}
Some(unsafe {
std::slice::from_raw_parts(
blocks.as_ptr() as *const u8,
blocks_packed_byte_len(blocks),
)
})
}
TensorData::Fp32(_) => None,
}
}
fn embed_row(embed: &Tensor, token: usize, hidden: usize) -> Vec<f32> {
match embed.as_data() {
TensorData::Fp32(data) => data[token * hidden..(token + 1) * hidden].to_vec(),
TensorData::Q8_0(blocks) => {
let bpr = hidden / 32;
let mut out = Vec::with_capacity(hidden);
for block in &blocks[token * bpr..(token + 1) * bpr] {
let scale = block.scale_f32();
for &q in block.quants() {
out.push(scale * q as f32);
}
}
out
}
}
}q8_bytes reinterprets a Q8_0 tensor's blocks as the raw byte stream the kernel reads (sound because Q8_0Block is #[repr(C, packed)], 34 bytes exactly, the same reinterpretation II.6 used), returning None if the tensor has an unexpected type or shape; that None propagates all the way to "use the CPU," so a checkpoint we don't recognize degrades instead of crashing. embed_row is the one op we deliberately keep on the CPU: the embedding lookup is a single-row gather, 4 KB of output, cheaper to compute and upload than to dispatch.
The constructor uploads the model:
impl GpuDecoder {
pub fn try_new(plan: &GpuDecodePlan) -> Option<Self> {
let kernels = Kernels::try_new()?;
let h = plan.hidden;
let qw = plan.num_heads * plan.head_dim;
let kvw = plan.num_kv_heads * plan.head_dim;
let inter = plan.intermediate;
let vocab = plan.vocab;
let mut layers = Vec::with_capacity(plan.layers.len());
for l in &plan.layers {
layers.push(LayerBufs {
attn_norm: kernels.upload_f32(l.attn_norm.as_f32_slice()),
q_proj: kernels.upload(q8_bytes(l.q_proj, qw, h)?),
k_proj: kernels.upload(q8_bytes(l.k_proj, kvw, h)?),
v_proj: kernels.upload(q8_bytes(l.v_proj, kvw, h)?),
o_proj: kernels.upload(q8_bytes(l.o_proj, h, qw)?),
q_norm: kernels.upload_f32(l.q_norm.as_f32_slice()),
k_norm: kernels.upload_f32(l.k_norm.as_f32_slice()),
ffn_norm: kernels.upload_f32(l.ffn_norm.as_f32_slice()),
gate_proj: kernels.upload(q8_bytes(l.gate_proj, inter, h)?),
up_proj: kernels.upload(q8_bytes(l.up_proj, inter, h)?),
down_proj: kernels.upload(q8_bytes(l.down_proj, h, inter)?),
});
}
let out_norm = kernels.upload_f32(plan.out_norm.as_f32_slice());
let lm_head = kernels.upload(q8_bytes(plan.lm_head, vocab, h)?);
let x = kernels.alloc(h * 4);
let xn = kernels.alloc(h * 4);
let q = kernels.alloc(qw * 4);
let k = kernels.alloc(kvw * 4);
let v = kernels.alloc(kvw * 4);
let attn = kernels.alloc(qw * 4);
let resid = kernels.alloc(h * 4);
let gate = kernels.alloc(inter * 4);
let up = kernels.alloc(inter * 4);
let ffn_h = kernels.alloc(inter * 4);
let logits = kernels.alloc(vocab * 4);
let amax = kernels.alloc(4);
Some(Self {
kernels,
layers,
out_norm,
lm_head,
x,
xn,
q,
k,
v,
attn,
resid,
gate,
up,
ffn_h,
logits,
amax,
kcache: Vec::new(),
vcache: Vec::new(),
cache_cap: 0,
next_pos: usize::MAX,
hidden: h,
q_width: qw,
kv_width: kvw,
inter,
vocab,
n_heads: plan.num_heads,
n_kv_heads: plan.num_kv_heads,
head_dim: plan.head_dim,
eps: plan.rms_eps,
theta: plan.rope_theta,
})
}The GpuDecodePlan argument is defined in the seam section below; for now, read it as "borrowed references to every weight tensor plus the model's hyperparameters." try_new walks it, uploading each Q8 projection as raw bytes and each FP32 norm vector as floats: ~600 MB of weights crossing into GPU-owned buffers exactly once, at first use. Every ? is another graceful-degradation exit. The cache vectors start empty and next_pos starts at a sentinel, because the KV cache can't be uploaded yet; at construction time it doesn't exist.
The cache comes over once
Prefill still runs on the hybrid CPU/GPU path from Act 2, filling the CPU-side KV cache. So when the first decode step arrives, the GPU's cache buffers must be populated from the CPU's, and after that the GPU keeps its own copy current via append_kv:
fn sync_cache(&mut self, cache: &dyn KvCache, pos: usize) -> Option<()> {
let cap = pos + KV_HEADROOM_ROWS;
self.kcache.clear();
self.vcache.clear();
for li in 0..self.layers.len() {
let (k, v) = cache.kv_view(li).map(|(k, v)| (k.clone(), v.clone()))?;
if k.shape() != [pos, self.kv_width] || v.shape() != [pos, self.kv_width] {
return None;
}
let kb = self.kernels.alloc(cap * self.kv_width * 4);
let vb = self.kernels.alloc(cap * self.kv_width * 4);
self.kernels.write_f32(&kb, k.as_f32_slice());
self.kernels.write_f32(&vb, v.as_f32_slice());
self.kcache.push(kb);
self.vcache.push(vb);
}
self.cache_cap = cap;
self.next_pos = pos;
Some(())
}
fn grow_cache(&mut self) {
let new_cap = self.cache_cap + KV_HEADROOM_ROWS;
let used = self.next_pos * self.kv_width;
for li in 0..self.kcache.len() {
for old in [&mut self.kcache[li], &mut self.vcache[li]] {
let nb = self.kernels.alloc(new_cap * self.kv_width * 4);
let data = self.kernels.read_f32(old, used);
self.kernels.write_f32(&nb, &data);
*old = nb;
}
}
self.cache_cap = new_cap;
}sync_cache reads each layer's K and V through kv_view, the zero-copy window IV.2 added (a cache that can't provide the view returns None, and decode stays on the CPU), verifies the shapes, and writes the contents into fresh GPU buffers sized pos + 1024 rows. The headroom is the trick that makes this a one-time cost: the next 1024 decode steps append in place on the GPU without any reallocation, and if a generation outruns the headroom, grow_cache does one bigger-alloc-and-copy per thousand tokens. next_pos records how many rows are valid; decode_step checks it each call, so the upload happens only when the GPU's copy doesn't already match (the first step, or a new sequence reusing the decoder).
One command buffer
Now the function this chapter is named after:
pub fn decode_step(
&mut self,
plan: &GpuDecodePlan,
token_id: usize,
pos: usize,
cache: &dyn KvCache,
) -> Option<usize> {
let seq = pos + 1;
if seq > ATTN_MAX_SEQ {
return None;
}
if self.next_pos != pos {
self.sync_cache(cache, pos)?;
}
if seq > self.cache_cap {
self.grow_cache();
}
let x_row = embed_row(plan.embed, token_id, self.hidden);
self.kernels.write_f32(&self.x, &x_row);
let (h, qw, kvw, inter) = (
self.hidden as u32,
self.q_width as u32,
self.kv_width as u32,
self.inter as u32,
);
use Scalar::{F, U};
let ks = &self.kernels;
let cmd = ks.queue.command_buffer().unwrap();
let enc =
unsafe { ComputeCommandEncoder::from_raw(cmd.compute_command_encoder()) }.unwrap();
let tg128 = Size::new(128, 1, 1);
let tg32 = Size::new(32, 1, 1);
let g1 = Size::new(1, 1, 1);Guards first: contexts past 4096 hand back None (CPU decode takes over), a stale next_pos triggers the one-time cache sync, a full cache grows. Then the embedding row is computed on the CPU and written into x, one command buffer and one encoder are created, and the encoding begins. Everything from here to commit is the CPU describing work, not doing it; nothing executes yet.
for (li, l) in self.layers.iter().enumerate() {
ks.encode(
&enc,
"rms_norm_row",
&[&self.x, &l.attn_norm, &self.xn],
&[U(h), F(self.eps)],
g1,
tg128,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &l.q_proj, &self.q],
&[U(h), U(qw)],
Size::new(self.q_width, 1, 1),
tg32,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &l.k_proj, &self.k],
&[U(h), U(kvw)],
Size::new(self.kv_width, 1, 1),
tg32,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &l.v_proj, &self.v],
&[U(h), U(kvw)],
Size::new(self.kv_width, 1, 1),
tg32,
);
ks.encode(
&enc,
"rms_norm_row",
&[&self.q, &l.q_norm, &self.q],
&[U(self.head_dim as u32), F(self.eps)],
Size::new(self.n_heads, 1, 1),
tg128,
);
ks.encode(
&enc,
"rms_norm_row",
&[&self.k, &l.k_norm, &self.k],
&[U(self.head_dim as u32), F(self.eps)],
Size::new(self.n_kv_heads, 1, 1),
tg128,
);
ks.encode(
&enc,
"rope_row",
&[&self.q],
&[U(self.head_dim as u32), U(pos as u32), F(self.theta)],
Size::new(self.n_heads, 1, 1),
Size::new(64, 1, 1),
);
ks.encode(
&enc,
"rope_row",
&[&self.k],
&[U(self.head_dim as u32), U(pos as u32), F(self.theta)],
Size::new(self.n_kv_heads, 1, 1),
Size::new(64, 1, 1),
);
ks.encode(
&enc,
"append_kv",
&[&self.k, &self.v, &self.kcache[li], &self.vcache[li]],
&[U(kvw), U(pos as u32)],
Size::new(self.kv_width.div_ceil(128), 1, 1),
tg128,
);
ks.encode(
&enc,
"attention_decode",
&[&self.q, &self.kcache[li], &self.vcache[li], &self.attn],
&[
U(seq as u32),
U(kvw),
U(self.head_dim as u32),
U((self.n_heads / self.n_kv_heads) as u32),
],
Size::new(self.n_heads, 1, 1),
tg128,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.attn, &l.o_proj, &self.resid],
&[U(qw), U(h)],
Size::new(self.hidden, 1, 1),
tg32,
);
ks.encode(
&enc,
"add_inplace",
&[&self.x, &self.resid],
&[U(h)],
Size::new(self.hidden.div_ceil(128), 1, 1),
tg128,
);
ks.encode(
&enc,
"rms_norm_row",
&[&self.x, &l.ffn_norm, &self.xn],
&[U(h), F(self.eps)],
g1,
tg128,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &l.gate_proj, &self.gate],
&[U(h), U(inter)],
Size::new(self.inter, 1, 1),
tg32,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &l.up_proj, &self.up],
&[U(h), U(inter)],
Size::new(self.inter, 1, 1),
tg32,
);
ks.encode(
&enc,
"silu_hadamard",
&[&self.gate, &self.up, &self.ffn_h],
&[U(inter)],
Size::new(self.inter.div_ceil(128), 1, 1),
tg128,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.ffn_h, &l.down_proj, &self.resid],
&[U(inter), U(h)],
Size::new(self.hidden, 1, 1),
tg32,
);
ks.encode(
&enc,
"add_inplace",
&[&self.x, &self.resid],
&[U(h)],
Size::new(self.hidden.div_ceil(128), 1, 1),
tg128,
);
}If you've internalized I.5's forward pass, this reads like its table of contents: norm, Q/K/V projections, QK-norm, RoPE on both, append to cache, attention, output projection, residual add, norm, gate and up projections, SwiGLU, down projection, residual add. Eighteen dispatches per layer, each one naming its input buffers, its output buffer, and its grid. The buffers make the dataflow legible: x in, xn normed, q/k/v projected, attn out of attention, resid carrying whatever is about to be added back into x. It is forward_decode_with_kv_cache transcribed into dispatch descriptions, which is precisely the claim "we don't need a graph, we have one model" cashes out to.
ks.encode(
&enc,
"rms_norm_row",
&[&self.x, &self.out_norm, &self.xn],
&[U(h), F(self.eps)],
g1,
tg128,
);
ks.encode(
&enc,
"q8_matvec",
&[&self.xn, &self.lm_head, &self.logits],
&[U(h), U(self.vocab as u32)],
Size::new(self.vocab, 1, 1),
tg32,
);
ks.encode(
&enc,
"argmax_row",
&[&self.logits, &self.amax],
&[U(self.vocab as u32)],
g1,
Size::new(256, 1, 1),
);
enc.end_encoding();
cmd.commit();
cmd.wait_until_completed();
self.next_pos = pos + 1;
Some(ks.read_u32(&self.amax) as usize)
}
pub fn read_logits(&self) -> Vec<f32> {
self.kernels.read_f32(&self.logits, self.vocab)
}
}Final norm, the lm_head matvec (151,936 threadgroups, the single biggest dispatch of the step), argmax. Then the three lines that were the point all along: end_encoding, commit, wait_until_completed. One wait. 507 dispatches went in; four bytes come out, read_u32(&self.amax), the argmax'd token id. Compare the ledger against Act 2's design: roughly 200 synchronous round trips per token if decode had used the old dispatch, versus exactly one here; a 151,936-float logits readback per step on the CPU path, versus one integer. read_logits survives as a debugging and testing hook, off the hot path.
The seam
The decoder needs to be reachable from a Model that only knows the Backend trait, without the model growing Metal-specific code. Three small additions make the seam.
First, src/backend/backend_trait.rs defines the plan structs (borrowed references, no copies: a map of the model's weights, not a container of them) and gives Backend a defaulted method that backends may override to claim an entire decode step:
use crate::cache::KvCache;
use crate::tensor::Tensor;
pub struct GpuLayerPlan<'a> {
pub attn_norm: &'a Tensor,
pub q_proj: &'a Tensor,
pub k_proj: &'a Tensor,
pub v_proj: &'a Tensor,
pub o_proj: &'a Tensor,
pub q_norm: &'a Tensor,
pub k_norm: &'a Tensor,
pub ffn_norm: &'a Tensor,
pub gate_proj: &'a Tensor,
pub up_proj: &'a Tensor,
pub down_proj: &'a Tensor,
}
pub struct GpuDecodePlan<'a> {
pub embed: &'a Tensor,
pub layers: Vec<GpuLayerPlan<'a>>,
pub out_norm: &'a Tensor,
pub lm_head: &'a Tensor,
pub num_heads: usize,
pub num_kv_heads: usize,
pub head_dim: usize,
pub hidden: usize,
pub intermediate: usize,
pub vocab: usize,
pub rms_eps: f32,
pub rope_theta: f32,
}
pub trait Backend: Send + Sync {
fn name(&self) -> String;
fn gpu_decode_step(
&self,
plan: &GpuDecodePlan,
token_id: usize,
pos: usize,
cache: &dyn KvCache,
) -> Option<usize> {
let _ = (plan, token_id, pos, cache);
None
}
// ...The default returns None: "I don't do whole steps, run the ops one at a time like always." Every existing backend keeps that behavior for free.
Second, the new backend itself, src/backend/metal/gpu.rs. Gpu wraps Metal, so prefill and any fallback op run exactly the II.5/II.6 hybrid path, and adds the lazily-built decoder:
use std::sync::Mutex;
use crate::backend::Backend;
use crate::backend::backend_trait::GpuDecodePlan;
use crate::cache::KvCache;
use crate::tensor::Tensor;
use super::Metal;
use super::decoder::GpuDecoder;
enum DecoderSlot {
Uninit,
Ready(GpuDecoder),
Failed,
}
pub struct Gpu<B: Backend> {
inner: Metal<B>,
decoder: Mutex<DecoderSlot>,
}
impl<B: Backend> Gpu<B> {
pub fn new(fallback: B) -> Self {
Self {
inner: Metal::new(fallback),
decoder: Mutex::new(DecoderSlot::Uninit),
}
}
}
impl<B: Backend> Backend for Gpu<B> {
fn name(&self) -> String {
"gpu".to_string()
}
fn gpu_decode_step(
&self,
plan: &GpuDecodePlan,
token_id: usize,
pos: usize,
cache: &dyn KvCache,
) -> Option<usize> {
let mut slot = self.decoder.lock().unwrap();
if matches!(*slot, DecoderSlot::Uninit) {
*slot = match GpuDecoder::try_new(plan) {
Some(d) => DecoderSlot::Ready(d),
None => {
eprintln!("gpu backend: decoder init failed, decode falls back to CPU");
DecoderSlot::Failed
}
};
}
match &mut *slot {
DecoderSlot::Ready(d) => d.decode_step(plan, token_id, pos, cache),
_ => None,
}
}The Mutex<DecoderSlot> does two jobs: it gives the shared, immutably-borrowed backend a place to keep mutable decoder state (the resident buffers, next_pos), and its three states make failure a remembered decision. The first call builds the decoder (that's when the ~600 MB weight upload happens); if the build fails, Failed is recorded and every later call returns None immediately instead of retrying a hopeless init 63 more times.
The rest of gpu.rs is the standard decorator delegation, everything to inner:
fn matmul(&self, a: &Tensor, b: &Tensor) -> Tensor {
self.inner.matmul(a, b)
}
fn sum_squares_axis(&self, x: &Tensor, axis: usize) -> Tensor {
self.inner.sum_squares_axis(x, axis)
}
fn add(&self, a: &Tensor, b: &Tensor) -> Tensor {
self.inner.add(a, b)
}
fn hadamard(&self, a: &Tensor, b: &Tensor) -> Tensor {
self.inner.hadamard(a, b)
}
fn scale(&self, x: &Tensor, s: f32) -> Tensor {
self.inner.scale(x, s)
}
fn silu(&self, x: &Tensor) -> Tensor {
self.inner.silu(x)
}
fn add_scalar(&self, a: &Tensor, s: f32) -> Tensor {
self.inner.add_scalar(a, s)
}
fn rsqrt_elem(&self, x: &Tensor) -> Tensor {
self.inner.rsqrt_elem(x)
}
fn broadcast_row_scalars(&self, t: &Tensor, d: usize) -> Tensor {
self.inner.broadcast_row_scalars(t, d)
}
fn transpose_2d(&self, a: &Tensor) -> Tensor {
self.inner.transpose_2d(a)
}
fn softmax_rows(&self, x: &Tensor) -> Tensor {
self.inner.softmax_rows(x)
}
fn gather_rows(&self, table: &Tensor, row_indices: &[usize]) -> Tensor {
self.inner.gather_rows(table, row_indices)
}
fn reshape_data(&self, x: &Tensor, shape: Vec<usize>) -> Tensor {
self.inner.reshape_data(x, shape)
}
fn fill_strict_upper_tri(&self, x: &Tensor, value: f32) -> Tensor {
self.inner.fill_strict_upper_tri(x, value)
}
fn copy_2d_from_cols(&self, src: &Tensor, w: usize, col_offset: usize) -> Tensor {
self.inner.copy_2d_from_cols(src, w, col_offset)
}
fn copy_2d_into_cols(&self, dst: &mut [f32], dst_cols: usize, src: &Tensor, col_offset: usize) {
self.inner.copy_2d_into_cols(dst, dst_cols, src, col_offset)
}
fn repeat_row_as_matrix(&self, weight: &Tensor, rows: usize) -> Tensor {
self.inner.repeat_row_as_matrix(weight, rows)
}
fn apply_rope(&self, x: &Tensor, head_dim: usize, rope_theta: f32) -> Tensor {
self.inner.apply_rope(x, head_dim, rope_theta)
}
fn apply_rope_single_row(
&self,
x: &Tensor,
position: usize,
head_dim: usize,
rope_theta: f32,
) -> Tensor {
self.inner
.apply_rope_single_row(x, position, head_dim, rope_theta)
}
fn concat_dim0(&self, a: &Tensor, b: &Tensor) -> Tensor {
self.inner.concat_dim0(a, b)
}
fn copy_row_2d(&self, x: &Tensor, row: usize) -> Tensor {
self.inner.copy_row_2d(x, row)
}
fn copy_contiguous_into(&self, x: &Tensor, start: usize, dst: &mut [f32]) {
self.inner.copy_contiguous_into(x, start, dst)
}
fn argmax_with_prob(&self, x: &Tensor) -> (usize, f32) {
self.inner.argmax_with_prob(x)
}
fn fused_attention(
&self,
q: &Tensor,
k: &Tensor,
v: &Tensor,
n_heads: usize,
n_kv_heads: usize,
head_dim: usize,
) -> Tensor {
self.inner
.fused_attention(q, k, v, n_heads, n_kv_heads, head_dim)
}
}Third, the model side. Qwen3Model learns to describe itself as a plan, and to offer a backend the whole step. In src/model/qwen3/forward.rs:
pub(crate) fn gpu_decode_plan(&self) -> GpuDecodePlan<'_> {
let cfg = &self.config;
GpuDecodePlan {
embed: &self.embed,
layers: self
.layers
.iter()
.map(|l| GpuLayerPlan {
attn_norm: &l.input_layernorm,
q_proj: &l.q_proj,
k_proj: &l.k_proj,
v_proj: &l.v_proj,
o_proj: &l.o_proj,
q_norm: &l.attn_q_norm,
k_norm: &l.attn_k_norm,
ffn_norm: &l.post_attention_layernorm,
gate_proj: &l.gate_proj,
up_proj: &l.up_proj,
down_proj: &l.down_proj,
})
.collect(),
out_norm: &self.norm,
lm_head: &self.lm_head,
num_heads: cfg.num_attention_heads,
num_kv_heads: cfg.num_key_value_heads,
head_dim: cfg.head_dim,
hidden: cfg.hidden_size,
intermediate: cfg.intermediate_size,
vocab: cfg.vocab_size,
rms_eps: cfg.rms_norm_eps,
rope_theta: cfg.rope_theta,
}
}
pub fn decode_step_token(
&self,
token_id: usize,
pos: usize,
cache: &mut dyn KvCache,
) -> Option<usize> {
let plan = self.gpu_decode_plan();
self.cpu_backend
.gpu_decode_step(&plan, token_id, pos, cache)
}with the matching import line at the top of the file:
use crate::backend::{GpuDecodePlan, GpuLayerPlan};gpu_decode_plan is pure bookkeeping: borrow every weight, copy the hyperparameters. decode_step_token builds the plan and asks the backend; on simd, pool, or metal the default gpu_decode_step answers None and nothing changes, while on gpu the decoder takes the step. The Model trait in src/model/model_trait.rs exposes it with a None default so other model implementations never hear about any of this:
fn decode_step_token(
&self,
token_id: usize,
pos: usize,
cache: &mut dyn KvCache,
) -> Option<usize> {
let _ = (token_id, pos, cache);
None
}and Qwen3Model forwards its trait method to the inherent one:
fn decode_step_token(
&self,
token_id: usize,
pos: usize,
cache: &mut dyn KvCache,
) -> Option<usize> {
Qwen3Model::decode_step_token(self, token_id, pos, cache)
}The last consumer is the greedy loop from I.6, which now distinguishes "the model gave me logits" from "the model gave me the token itself." src/decode/greedy.rs, the whole function:
pub fn greedy_generate(
model: &dyn Model,
ops: &dyn Backend,
tokenizer: &dyn Tokenizer,
prompt_ids: &[usize],
max_new_tokens: usize,
eos_token_id: usize,
metrics: &mut Metrics,
cache: &mut Option<Box<dyn KvCache>>,
mut on_token: impl FnMut(usize),
) -> Vec<usize> {
let mut ids = prompt_ids.to_vec();
enum StepOut {
Logits(crate::tensor::Tensor),
Token(usize),
}
for step in 0..max_new_tokens {
let out = metrics.record_timed(|| match cache.as_mut() {
Some(c) if step == 0 => {
StepOut::Logits(model.forward_prefill_with_kv_cache(prompt_ids, c.as_mut()))
}
Some(c) => {
let last = *ids.last().unwrap();
let pos = ids.len() - 1;
match model.decode_step_token(last, pos, c.as_mut()) {
Some(id) => StepOut::Token(id),
None => StepOut::Logits(model.forward_decode_with_kv_cache(last, pos, c.as_mut())),
}
}
None => StepOut::Logits(model.forward(&ids)),
});
let (next_id, prob) = match out {
StepOut::Token(id) => (id, 1.0),
StepOut::Logits(logits) => next_token_id_from_logits(ops, &logits),
};
ids.push(next_id);
info!(
step = step + 1,
token_id = next_id,
prob = prob * 100.0,
token = %tokenizer.decode(&[next_id]),
phase = if step == 0 { "prefill" } else { "decode" },
);
if next_id == eos_token_id {
info!("[EOS reached]");
break;
}
on_token(next_id);
}
ids
}Prefill still takes the logits path (step 0). Every later step tries decode_step_token first and falls back seamlessly, which is what makes the over-4096-context and decoder-init-failure cases safe: the same loop just resumes computing logits on the CPU. One wrinkle: a GPU step yields a token with no probability attached (the logits never left the GPU), so the log line reports 1.0. An honest small cost of not reading back 151,936 floats.
Wiring completes the same way as every backend. TracingBackend forwards and counts the new method in src/backend/tracing.rs:
fn gpu_decode_step(
&self,
plan: &super::backend_trait::GpuDecodePlan,
token_id: usize,
pos: usize,
cache: &dyn crate::cache::KvCache,
) -> Option<usize> {
self.bump("gpu_decode_step");
self.inner.gpu_decode_step(plan, token_id, pos, cache)
}the factory gains "gpu":
"gpu" => Ok(wrap_gpu(enable_tracing)),
other => Err(format!(
"unknown backend {other:?} (supported: scalar, simd, parallel, pool, metal, gpu)"
)),
// ...
fn wrap_gpu(enable_tracing: bool) -> Arc<dyn Backend> {
let gpu = Gpu::new(SimdCpu::new(CpuBackend));
if enable_tracing {
Arc::new(TracingBackend::new(gpu))
} else {
Arc::new(gpu)
}
}and the module files export the new pieces:
mod backend;
mod context;
mod decode_shaders;
mod decoder;
mod gpu;
mod shaders;
pub use backend::Metal;
pub use gpu::Gpu;pub use backend_trait::{Backend, GpuDecodePlan, GpuLayerPlan};
// ...
pub(crate) use metal::{Gpu, Metal};Two files also pick up #[cfg(test)] exports so the whole-model test below can name internals:
#[cfg(test)]
pub(crate) use basic::BasicKvCache;#[cfg(test)]
pub(crate) use qwen3::Qwen3Model;(The commit's remaining touches, in src/cache/paged/layer.rs, src/cache/paged/paged_cache.rs, and src/model/load.rs, are cargo fmt reflows and import reorderings with no behavior change; we won't reprint them.)
The correctness ladder
Here's the part that deserves as much care as the speed. IV.1 set a hard bar for CPU optimizations: byte-identical logits, verified. A GPU port cannot meet that bar, and it's better to say so precisely than to hide it. Floating-point addition is not associative; q8_matvec sums each dot product in 32 lanes then folds them, where the CPU's sdot kernel accumulates in a different order with integer-quantized activations. Same math, different rounding at every one of the millions of adds per token. So we climb a ladder of weaker-but-verifiable guarantees instead, each rung a test in decoder.rs.
Rung one: every kernel against its CPU counterpart. The test module bootstraps with a tiny deterministic RNG, a skip-if-no-GPU helper, an element-wise comparator, and a Q8 weight fabricator:
#[cfg(test)]
mod gpu_kernel_tests {
use super::*;
use crate::backend::Backend;
use crate::backend::cpu::CpuBackend;
use crate::backend::simd_cpu::{SimdCpu, matvec_q8_0};
use crate::tensor::Tensor;
use crate::tensor::q8_0::{Q8_0_BLOCK_BYTES, Q8_0_BLOCK_SIZE, Q8_0Block};
fn rng_f32(state: &mut u64) -> f32 {
*state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((*state >> 33) as f32 / (1u64 << 31) as f32) - 1.0
}
fn kernels_or_skip() -> Option<Kernels> {
let k = Kernels::try_new();
if k.is_none() {
eprintln!("no Metal device; skipping GPU kernel test");
}
k
}
fn assert_close(gpu: &[f32], cpu: &[f32], tol: f32, label: &str) {
assert_eq!(gpu.len(), cpu.len(), "{label}: length mismatch");
let mut max_diff = 0.0f32;
for (a, b) in gpu.iter().zip(cpu) {
max_diff = max_diff.max((a - b).abs());
}
assert!(max_diff < tol, "{label}: max abs diff {max_diff} >= {tol}");
eprintln!("{label}: max abs diff {max_diff:e}");
}
fn random_q8_bytes(rows: usize, cols: usize, seed: u64) -> Vec<u8> {
let mut s = seed;
let bpr = cols / Q8_0_BLOCK_SIZE;
let mut bytes = Vec::with_capacity(rows * bpr * Q8_0_BLOCK_BYTES);
for _ in 0..rows * bpr {
bytes.extend_from_slice(&[0x00, 0x2c]);
for _ in 0..Q8_0_BLOCK_SIZE {
bytes.push((rng_f32(&mut s) * 127.0) as i8 as u8);
}
}
bytes
}(0x00, 0x2c is the f16 bit pattern for a realistic small block scale.) Then one test per kernel, each running the GPU kernel via run_once and comparing against the trusted CPU implementation, the same technique II.6 used to validate its integer kernels. The matvec:
#[test]
fn gpu_q8_matvec_matches_cpu() {
let Some(ks) = kernels_or_skip() else { return };
let (p, n) = (50, 96);
let wbytes = random_q8_bytes(p, n, 7);
let blocks: Vec<Q8_0Block> = wbytes
.chunks(Q8_0_BLOCK_BYTES)
.map(Q8_0Block::from_bytes)
.collect();
let mut s = 42u64;
let x: Vec<f32> = (0..n).map(|_| rng_f32(&mut s)).collect();
let bpr = n / Q8_0_BLOCK_SIZE;
let expected: Vec<f32> = (0..p)
.map(|r| {
let mut sum = 0.0f32;
for b in 0..bpr {
let block = &blocks[r * bpr + b];
let scale = block.scale_f32();
for j in 0..Q8_0_BLOCK_SIZE {
sum += scale * block.quants()[j] as f32 * x[b * Q8_0_BLOCK_SIZE + j];
}
}
sum
})
.collect();
let simd_ref = matvec_q8_0(&blocks, &x, p, n).unwrap();
let quant_err = expected
.iter()
.zip(&simd_ref)
.map(|(a, b)| (a - b).abs())
.fold(0.0f32, f32::max);
eprintln!("q8_matvec: cpu simd activation-quant error vs float ref: {quant_err:e}");
let xb = ks.upload_f32(&x);
let wb = ks.upload(&wbytes);
let ob = ks.alloc(p * 4);
ks.run_once(
"q8_matvec",
&[&xb, &wb, &ob],
&[Scalar::U(n as u32), Scalar::U(p as u32)],
Size::new(p, 1, 1),
Size::new(32, 1, 1),
);
assert_close(&ks.read_f32(&ob, p), &expected, 5e-4, "q8_matvec");
}Note the honesty baked into the reference: expected is the exact float dot product, and the test also prints how far the CPU's own matvec_q8_0 sits from that reference (it quantizes activations, remember). The GPU kernel must land within 5e-4 of exact; the printout shows the CPU kernel's activation-quantization error is of similar size. Neither processor is "the truth"; both approximate the same float math, from different directions.
The norm, RoPE, and attention tests, each against CPU code that has been in service since Act 1 or IV.1:
#[test]
fn gpu_rms_norm_matches_cpu() {
let Some(ks) = kernels_or_skip() else { return };
for (rows, dim) in [(1usize, 1024usize), (16, 128)] {
let mut s = 3u64;
let x: Vec<f32> = (0..rows * dim).map(|_| rng_f32(&mut s)).collect();
let w: Vec<f32> = (0..dim).map(|_| rng_f32(&mut s) + 1.5).collect();
let eps = 1e-6f32;
let mut expected = vec![0.0f32; rows * dim];
for r in 0..rows {
let row = &x[r * dim..(r + 1) * dim];
let ss: f32 = row.iter().map(|v| v * v).sum();
let inv = 1.0 / (ss / dim as f32 + eps).sqrt();
for i in 0..dim {
expected[r * dim + i] = row[i] * inv * w[i];
}
}
let xb = ks.upload_f32(&x);
let wb = ks.upload_f32(&w);
let ob = ks.alloc(rows * dim * 4);
ks.run_once(
"rms_norm_row",
&[&xb, &wb, &ob],
&[Scalar::U(dim as u32), Scalar::F(eps)],
Size::new(rows, 1, 1),
Size::new(128, 1, 1),
);
assert_close(
&ks.read_f32(&ob, rows * dim),
&expected,
1e-4,
&format!("rms_norm rows={rows} dim={dim}"),
);
}
}
#[test]
fn gpu_rope_matches_cpu() {
let Some(ks) = kernels_or_skip() else { return };
let (heads, hd, pos, theta) = (4usize, 16usize, 137usize, 1_000_000.0f32);
let mut s = 11u64;
let x: Vec<f32> = (0..heads * hd).map(|_| rng_f32(&mut s)).collect();
let cpu = CpuBackend;
let expected = cpu
.apply_rope_single_row(&Tensor::new(x.clone(), vec![1, heads * hd]), pos, hd, theta)
.as_f32_slice()
.to_vec();
let xb = ks.upload_f32(&x);
ks.run_once(
"rope_row",
&[&xb],
&[
Scalar::U(hd as u32),
Scalar::U(pos as u32),
Scalar::F(theta),
],
Size::new(heads, 1, 1),
Size::new(64, 1, 1),
);
assert_close(&ks.read_f32(&xb, heads * hd), &expected, 1e-4, "rope_row");
}
#[test]
fn gpu_attention_decode_matches_cpu() {
let Some(ks) = kernels_or_skip() else { return };
let (n_heads, n_kv, hd, seq) = (4usize, 2usize, 16usize, 9usize);
let kvw = n_kv * hd;
let qw = n_heads * hd;
let mut s = 23u64;
let q: Vec<f32> = (0..qw).map(|_| rng_f32(&mut s)).collect();
let k: Vec<f32> = (0..seq * kvw).map(|_| rng_f32(&mut s)).collect();
let v: Vec<f32> = (0..seq * kvw).map(|_| rng_f32(&mut s)).collect();
let simd = SimdCpu::new(CpuBackend);
let expected = simd
.fused_attention(
&Tensor::new(q.clone(), vec![1, qw]),
&Tensor::new(k.clone(), vec![seq, kvw]),
&Tensor::new(v.clone(), vec![seq, kvw]),
n_heads,
n_kv,
hd,
)
.as_f32_slice()
.to_vec();
let qb = ks.upload_f32(&q);
let kb = ks.upload_f32(&k);
let vb = ks.upload_f32(&v);
let ob = ks.alloc(qw * 4);
ks.run_once(
"attention_decode",
&[&qb, &kb, &vb, &ob],
&[
Scalar::U(seq as u32),
Scalar::U(kvw as u32),
Scalar::U(hd as u32),
Scalar::U((n_heads / n_kv) as u32),
],
Size::new(n_heads, 1, 1),
Size::new(128, 1, 1),
);
assert_close(&ks.read_f32(&ob, qw), &expected, 1e-4, "attention_decode");
}And the three small kernels plus the argmax, whose test plants a tie at indices 140011 and 140012 to pin down the tie-breaking rule:
#[test]
fn gpu_elementwise_kernels_match_cpu() {
let Some(ks) = kernels_or_skip() else { return };
let n = 3072usize;
let mut s = 5u64;
let a: Vec<f32> = (0..n).map(|_| rng_f32(&mut s)).collect();
let b: Vec<f32> = (0..n).map(|_| rng_f32(&mut s)).collect();
let ab = ks.upload_f32(&a);
let bb = ks.upload_f32(&b);
ks.run_once(
"add_inplace",
&[&ab, &bb],
&[Scalar::U(n as u32)],
Size::new(n.div_ceil(128), 1, 1),
Size::new(128, 1, 1),
);
let expected_add: Vec<f32> = a.iter().zip(&b).map(|(x, y)| x + y).collect();
assert_close(&ks.read_f32(&ab, n), &expected_add, 1e-6, "add_inplace");
let gb = ks.upload_f32(&a);
let ub = ks.upload_f32(&b);
let ob = ks.alloc(n * 4);
ks.run_once(
"silu_hadamard",
&[&gb, &ub, &ob],
&[Scalar::U(n as u32)],
Size::new(n.div_ceil(128), 1, 1),
Size::new(128, 1, 1),
);
let expected_sh: Vec<f32> = a
.iter()
.zip(&b)
.map(|(g, u)| g / (1.0 + (-g).exp()) * u)
.collect();
assert_close(&ks.read_f32(&ob, n), &expected_sh, 1e-5, "silu_hadamard");
}
#[test]
fn gpu_append_kv_writes_row() {
let Some(ks) = kernels_or_skip() else { return };
let (w, pos, cap) = (32usize, 3usize, 8usize);
let mut s = 9u64;
let k: Vec<f32> = (0..w).map(|_| rng_f32(&mut s)).collect();
let v: Vec<f32> = (0..w).map(|_| rng_f32(&mut s)).collect();
let kb = ks.upload_f32(&k);
let vb = ks.upload_f32(&v);
let kc = ks.alloc(cap * w * 4);
let vc = ks.alloc(cap * w * 4);
ks.run_once(
"append_kv",
&[&kb, &vb, &kc, &vc],
&[Scalar::U(w as u32), Scalar::U(pos as u32)],
Size::new(w.div_ceil(128), 1, 1),
Size::new(128, 1, 1),
);
let kc_all = ks.read_f32(&kc, cap * w);
let vc_all = ks.read_f32(&vc, cap * w);
assert_close(&kc_all[pos * w..(pos + 1) * w], &k, 1e-7, "append_kv k row");
assert_close(&vc_all[pos * w..(pos + 1) * w], &v, 1e-7, "append_kv v row");
}
#[test]
fn gpu_argmax_matches_cpu() {
let Some(ks) = kernels_or_skip() else { return };
let n = 151936usize;
let mut s = 77u64;
let mut x: Vec<f32> = (0..n).map(|_| rng_f32(&mut s)).collect();
x[140011] = 9.5;
x[140012] = 9.5;
let expected = x
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap().then(b.0.cmp(&a.0)))
.unwrap()
.0;
let xb = ks.upload_f32(&x);
let ob = ks.alloc(4);
ks.run_once(
"argmax_row",
&[&xb, &ob],
&[Scalar::U(n as u32)],
Size::new(1, 1, 1),
Size::new(256, 1, 1),
);
assert_eq!(ks.read_u32(&ob) as usize, expected, "argmax_row");
assert_eq!(expected, 140011, "tie should pick lowest index");
}Every kernel sits within 1e-4 of its CPU counterpart (5e-4 for the matvec, where quantization is in play; down to 1e-6 for the pure adds). That is rung one: each piece is right.
Rung two: the whole step. Ignored by default (it needs the model file on disk; run it with cargo test -- --ignored), this test prefills the real model on the CPU, then runs the same first decode step both ways and compares all 151,936 logits:
#[test]
#[ignore]
fn gpu_first_decode_logits_match_simd() {
let Some(_) = kernels_or_skip() else { return };
let path = std::path::Path::new(
"/Users/kothari/projects/inference/archive/Qwen3-0.6B-Q8_0-GGUF/qwen3-0.6b-q8_0.gguf",
);
if !path.exists() {
eprintln!("model file missing; skipping logits gate");
return;
}
let simd = crate::backend::create_backend("scalar", false).unwrap();
let mut gguf = crate::gguf::GGUF::parse(path);
let model = crate::model::Qwen3Model::new(&mut gguf, simd.clone());
let (layers, kvw) = model.kv_cache_layers_and_width();
let prompt = [
12522usize, 5193, 264, 882, 11, 304, 264, 2613, 14126, 553, 279, 9396,
];
let mut cache = crate::cache::BasicKvCache::new(layers, kvw, simd.clone());
let logits = model.forward_prefill_with_kv_cache(&prompt, &mut cache);
let (next, _) = simd.argmax_with_prob(&logits);
let mut cache_cpu = cache.clone();
let cpu_logits = model.forward_decode_with_kv_cache(next, prompt.len(), &mut cache_cpu);
let cpu = cpu_logits.as_f32_slice();
let plan = model.gpu_decode_plan();
let mut dec = GpuDecoder::try_new(&plan).expect("GpuDecoder init failed");
let gpu_tok = dec
.decode_step(&plan, next, prompt.len(), &cache)
.expect("gpu decode step failed");
let gpu = dec.read_logits();
let mut max_diff = 0.0f32;
for (a, b) in gpu.iter().zip(cpu) {
max_diff = max_diff.max((a - b).abs());
}
let (cpu_tok, _) = simd.argmax_with_prob(&cpu_logits);
eprintln!(
"first-decode-step logits: max abs diff {max_diff:e}; gpu argmax {gpu_tok} cpu argmax {cpu_tok}"
);
assert!(max_diff < 2e-2, "logits max abs diff {max_diff} >= 2e-2");
assert_eq!(gpu_tok, cpu_tok, "argmax token mismatch");
}
}Measured on this machine: max absolute difference of 1.2e-2 across the 151,936 logits, with the argmax token identical. Per-kernel error was under 1e-4; after 28 layers of compounding (each layer's slightly-different output feeds the next layer's slightly-different input), it grows three orders of magnitude, and that is normal. Against logits whose interesting values span tens of units, a 0.012 worst-case wobble almost never reorders the top of the distribution.
Rung three, where the ladder honestly ends: transcripts. Almost never is not never. Run the 64-token bakery prompt on gpu and on simd and the outputs agree token for token through 27 generated tokens, then diverge: somewhere in there, two logits land close enough that a 1e-2 wobble flips the argmax, and greedy decoding never looks back. The CPU continues "...decided to bake a special cake for his friends"; the GPU continues "...decided to bake a special cake. He wanted to make a cake that would be as delicious as possible." Both are the model, both are coherent; they are two valid readings of the same weights under different rounding. This is the documented trade of the chapter: within the CPU, IV.1 engineered byte-identity and we should demand it; across processors, no engine gets it (llama.cpp's Metal decode diverges from its own CPU decode the same way), and the correct move is to measure the divergence, bound it, test the bound, and say it out loud.
Running it
cargo run --release --bin model-generate -- --kv --backend metal path/to/qwen3-0.6b-q8_0.gguf "Once upon a time, in a small village by the sea, there lived a baker" 64
cargo run --release --bin model-generate -- --kv --backend gpu path/to/qwen3-0.6b-q8_0.gguf "Once upon a time, in a small village by the sea, there lived a baker" 64The metal baseline on the short prompt:
MetalBackend: device=Apple M2 Pro unified_memory=true
backend: metal
kv cache: basic
metrics:
time_to_first_token_ms: 300.193
decode_tokens_per_second: 42.167
per_forward_ms: min 21.186 max 300.193 mean 28.035 (n=64)And gpu:
MetalBackend: device=Apple M2 Pro unified_memory=true
backend: gpu
kv cache: basic
metrics:
time_to_first_token_ms: 288.795
decode_tokens_per_second: 108.178
per_forward_ms: min 7.432 max 288.795 mean 13.612 (n=64)42.2 to 108.2 tokens/sec, 2.6x, from the same silicon that sat idle through every decode step of Act 2. TTFT is unchanged (~0.29 s versus ~0.30 s) because prefill runs the identical hybrid path in both backends; only decode changed hands. The mean per-forward is dragged up by the first decode step, which pays the one-time weight upload and cache sync; steady-state steps run at the min, ~7.4 ms.
Long context is where the story completes, because that's where CPU decode sagged and where the GPU's parallelism has the most cached history to chew on. The 511-token prompt:
MetalBackend: device=Apple M2 Pro unified_memory=true
backend: metal
kv cache: basic
metrics:
time_to_first_token_ms: 1533.498
decode_tokens_per_second: 16.097
per_forward_ms: min 58.860 max 1533.498 mean 85.112 (n=64)MetalBackend: device=Apple M2 Pro unified_memory=true
backend: gpu
kv cache: basic
metrics:
time_to_first_token_ms: 1526.944
decode_tokens_per_second: 65.159
per_forward_ms: min 13.450 max 1526.944 mean 38.966 (n=64)16.1 to 65.2 tokens/sec, 4x, and TTFT again identical (~1.53 s). Put the whole ladder side by side at 511 tokens of context: simd 24.3, pool on four cores 34.7, gpu 65.2. For the first time in this series, the GPU is the fastest way to decode, and by nearly 2x over the best CPU configuration.
Read the two changes that produced the 4x, because neither is a kernel. The arithmetic per token is the same as it ever was; the kernels were never the problem. What changed is that the ~200 synchronous dispatch-and-wait round trips a per-op GPU decode would have paid became exactly one, and the weights plus KV cache stopped commuting: resident on the GPU, appended in place by append_kv, with four bytes of token id as the entire per-step traffic. The wall was never compute. It was the asking.
Where this leaves us
--backend gpu is the sixth backend: II.5's hybrid engine for prefill, and for decode a hand-encoded forward pass, eight kernels dispatched 507 times into one command buffer per token, weights and KV cache resident, one wait, one integer back. Decode reaches 108 tokens/sec on the short prompt and 65 at 511 tokens of context, with correctness held by a ladder of kernel tests, a bounded whole-step comparison, and a documented transcript divergence.
The gap between 108 and 65 points at the next bottleneck. Every kernel in this chapter scales gently with context except one: attention_decode walks the whole cached history through a 16 KB threadgroup-memory score array, one threadgroup per head, and it both slows down as context grows and caps out entirely at 4096 tokens. The act's last chapter replaces it with the algorithm built for exactly this shape of problem: flash attention.