III.7: Batched decode

The scheduler from III.6 runs concurrent requests by interleaving them: decode_tick walks every active slot and runs each one's forward pass, one after another. Four active slots means four separate forward passes per tick.

That works, but it leaves performance on the table, and to see why we have to look at the shape of a decode-step forward pass.

A decode step processes exactly one token. Every projection in the model (the Q/K/V projections, the attention output projection, the MLP's three matrices) is a matmul of that one token's hidden vector against a weight matrix. A vector times a matrix. In matmul terms, the "batch" dimension is 1.

That is the worst case for matmul hardware. A CPU's SIMD units and a GPU's cores are built to chew through matrix times matrix, a tall left operand, many rows processed per pass. Hand them a single-row left operand and most of the machine sits idle: the weight matrix still has to be streamed from memory, but it's used for just one row instead of many. The matmul is memory-bound, and the compute units starve.

Now the key observation. When four slots all take a decode step in the same tick, that's four separate single-row matmuls against the same weight matrices. Stack those four hidden vectors into a 4-row matrix and you get one matmul with a batch dimension of 4: the weight matrix streamed once, used four times. Same arithmetic, a quarter of the weight-memory traffic. This is batched decode, and it's this chapter.

What batches and what doesn't

Not every part of the forward pass batches the same way. There are two kinds of computation in a transformer layer:

Position-independent projections. Q/K/V projections, the attention output projection, the MLP. Each is hidden_vector × weight_matrix, and the weight matrix is shared by every sequence. Stack the hidden vectors → one bigger matmul. These batch perfectly.
Attention itself. Each sequence attends over its own KV history, with different lengths, different contents, and a different KV cache per slot. There's no shared matrix to stack against. Attention stays per-sequence: a loop over the batch.

So a batched decode step looks like: gather the batch's input tokens, do the shared projections as fused matmuls, loop per-sequence for the attention (each slot reads and appends to its own paged KV cache from III.4), then do the shared MLP as a fused matmul. The expensive matmuls (projections and MLP, the bulk of the FLOPs) are batched; only attention, which is comparatively cheap at decode time, stays in a loop.

This chapter adds forward_decode_batch to the model, teaches the scheduler worker to call it when two or more slots are active, and ships a GuideLLM A/B benchmark to measure the gain.

The trait method

src/model/model_trait.rs gains forward_decode_batch, and it ships with a default implementation so any model that doesn't override it still works:

src/model/model_trait.rsRUST

    fn forward_decode_batch(
        &self,
        token_ids: &[usize],
        positions: &[usize],
        caches: &mut [&mut dyn KvCache],
    ) -> Result<Vec<Tensor>, String> {
        assert!(token_ids.len() == positions.len() && token_ids.len() == caches.len());
        let mut out = Vec::with_capacity(token_ids.len());
        for i in 0..token_ids.len() {
            out.push(self.forward_decode_with_kv_cache(
                token_ids[i],
                positions[i],
                caches[i],
            ));
        }
        Ok(out)
    }

The signature: a slice of token ids (one per batched slot), a slice of positions, and a slice of KV cache references (&mut [&mut dyn KvCache], one mutable cache borrow per slot). It returns one logits tensor per slot.

The default just loops, calling the single-sequence forward_decode_with_kv_cache once per slot, exactly what III.6's decode_tick did. So the trait method changes nothing on its own; the speedup comes from a model that overrides it with a genuinely fused pass. Qwen3 does.