IV.2: Zero-copy KV cache
IV.1 fused the per-head attention loop into one op and left a loose end in the histogram: concat_dim0: 56. With attention fused, that is the biggest remaining piece of per-token bookkeeping, and it is not a cheap piece. Those 56 calls are the KV cache appending one token: 28 layers times two tensors (K and V), every decode step, forever.
What push_row actually does
Recall II.2's BasicKvCache. Appending a token's K row is self.k = ops.concat_dim0(&self.k, k_row): allocate a brand-new tensor one row taller, copy the entire existing cache into it, copy the new row after it, drop the old tensor. II.2 admitted at the time that this was "not the most efficient layout." Now we can say precisely how inefficient. The copy grows with the sequence: at step t it copies t rows, so over a generation the copying is quadratic in sequence length, the same O(seq²) shape the KV cache was invented to remove from compute, quietly reintroduced as memcpy.
And that is only the append. The decode path then calls cache.materialize(li), which hands attention the full K and V by cloning both tensors, another two full copies per layer, every step. Put numbers on it: at the end of a 511-token-prompt run with 64 generated tokens, each layer's K is about 575 rows x 1024 floats, roughly 2.3 MB. Two concat copies plus two materialize clones is about 9.4 MB per layer per token; times 28 layers, the engine is copying roughly a quarter of a gigabyte of cache per generated token, none of which changes a single value. The fused attention kernel then reads those freshly copied tensors once and they are garbage.
The fix has two halves that are really one idea: make the cache a growable buffer that appends in place, and let attention read it through a borrow instead of a clone.
A Vec with capacity is already a preallocated buffer
There is a tempting heavyweight design here: preallocate the whole context up front and write rows into it. Do the arithmetic first. Qwen3 0.6B's maximum context is 40,960 tokens; a full-context cache would be 40,960 rows x 1024 floats x 4 bytes x 28 layers x 2 tensors, about 9.4 GB, allocated before the model generates a word. That is a bad trade for a laptop engine, and it is also unnecessary, because Rust's Vec already is a preallocated buffer with a movable end: reserve claims capacity without initializing it, extend_from_slice writes at the end without reallocating until capacity runs out, and when it does run out, Vec doubles, which keeps the amortized copying constant per token instead of linear. So the design is: reserve generous headroom once at prefill, append rows in place, and let doubling cover the rare overflow.
Tensor grows three methods to expose exactly that. src/tensor/tensor.rs:
pub(crate) fn empty_rows(cols: usize) -> Self {
Self {
data: TensorData::Fp32(Vec::new()),
shape: vec![0, cols],
}
}
pub(crate) fn reserve_rows(&mut self, rows: usize) {
let cols = self.shape[1];
match &mut self.data {
TensorData::Fp32(v) => v.reserve(rows * cols),
_ => panic!("reserve_rows: FP32 only"),
}
}
pub(crate) fn append_rows(&mut self, src: &Tensor) {
let cols = self.shape[1];
assert_eq!(src.shape[1], cols, "append_rows: column mismatch");
match (&mut self.data, &src.data) {
(TensorData::Fp32(dst), TensorData::Fp32(s)) => dst.extend_from_slice(s),
_ => panic!("append_rows: FP32 only"),
}
self.shape[0] += src.shape[0];
}
}empty_rows makes a zero-row tensor with a column width; reserve_rows claims capacity for that many more rows; append_rows copies rows onto the end and bumps shape[0]. The shape is the write offset: the data Vec's length always equals shape[0] * cols, and the reserved-but-unwritten capacity behind it is invisible to every other part of the engine. All three panic on quantized tensors, which is fine: KV data is always FP32 activations, never Q8_0 weights.
The cache appends in place
src/cache/basic.rs rewires LayerKvState around the new methods:
const KV_HEADROOM_ROWS: usize = 512;
impl LayerKvState {
pub fn empty(kv_width: usize) -> Self {
Self {
k: Tensor::empty_rows(kv_width),
v: Tensor::empty_rows(kv_width),
}
}
pub fn set_prefill(&mut self, k: &Tensor, v: &Tensor) {
let kw = k.shape()[1];
let cap = k.shape()[0] + KV_HEADROOM_ROWS;
let mut kb = Tensor::empty_rows(kw);
kb.reserve_rows(cap);
kb.append_rows(k);
let mut vb = Tensor::empty_rows(kw);
vb.reserve_rows(cap);
vb.append_rows(v);
self.k = kb;
self.v = vb;
}
pub fn push_row(&mut self, k_row: &Tensor, v_row: &Tensor) {
self.k.append_rows(k_row);
self.v.append_rows(v_row);
}set_prefill copies the prompt's K and V exactly once, into buffers reserved for the prompt plus 512 more rows of headroom, enough for any decode run in this book without a single reallocation. push_row is now two in-place appends: 1,024 floats each, no allocation, no copy of existing rows. It no longer needs a Backend at all (the old version took ops just to call concat_dim0, and II.2 spent a paragraph on the borrow-checker contortion that caused; both the parameter and the contortion are simply gone).
Reading without copying
The other half is materialize, the clone-per-read. The clean fix is to hand attention a borrow of the cache's tensors instead of copies. The KvCache trait gets a new method with a conservative default. src/cache/kv_cache_trait.rs:
fn set_prefill(&mut self, layer: usize, k: &Tensor, v: &Tensor);
fn push_row(&mut self, layer: usize, k_row: &Tensor, v_row: &Tensor);
fn materialize(&self, layer: usize) -> (Tensor, Tensor);
fn kv_view(&self, _layer: usize) -> Option<(&Tensor, &Tensor)> {
None
}
fn seq_len(&self, layer: usize) -> usize;
fn num_layers(&self) -> usize;kv_view returns live references into the cache's own storage, valid for as long as the borrow lasts, zero bytes copied. It is an Option because not every cache can answer: the paged cache from Act 3 stores a sequence as scattered fixed-size blocks, and there is no contiguous &Tensor to lend. Its default of None means "I can't lend a view, materialize instead," so the paged cache keeps working, untouched, on its existing path. BasicKvCache can always answer:
fn push_row(&mut self, layer: usize, k_row: &Tensor, v_row: &Tensor) {
self.layers[layer].push_row(k_row, v_row);
}
fn materialize(&self, layer: usize) -> (Tensor, Tensor) {
self.layers[layer].materialize()
}
fn kv_view(&self, layer: usize) -> Option<(&Tensor, &Tensor)> {
let l = &self.layers[layer];
Some((&l.k, &l.v))
}materialize stays in the trait and in the impl: snapshots and the paged path still use it, and it is the fallback below.
The decode paths choose the free option
src/model/qwen3/forward.rs has two decode paths, and both get the same change: try kv_view, fall back to materialize. The single-token path from II.2:
cache.push_row(li, &k_rope, &v);
let attn_out = match cache.kv_view(li) {
Some((k_full, v_full)) => gqa_attention_decode_one_query(
ops,
&q,
k_full,
v_full,
&layer.o_proj,
cfg.num_attention_heads,
cfg.num_key_value_heads,
cfg.head_dim,
),
None => {
let (k_full, v_full) = cache.materialize(li);
gqa_attention_decode_one_query(
ops,
&q,
&k_full,
&v_full,
&layer.o_proj,
cfg.num_attention_heads,
cfg.num_key_value_heads,
cfg.head_dim,
)
}
};Note the ordering: push_row first, then borrow. The borrow checker enforces what correctness requires anyway, that the append (a mutable borrow) finishes before the view (a shared borrow) begins, so attention always sees a cache that already includes the current token. The Some arm passes k_full and v_full straight through as references; the None arm is the old clone-based code, unchanged, for caches that can't lend.
The batched decode path from III.7 is the identical transformation inside its per-sequence loop:
cache.push_row(li, &k_rope, &v_row);
let attn_out = match cache.kv_view(li) {
Some((k_full, v_full)) => gqa_attention_decode_one_query(
ops,
&q_rope,
k_full,
v_full,
&layer.o_proj,
cfg.num_attention_heads,
cfg.num_key_value_heads,
cfg.head_dim,
),
None => {
let (k_full, v_full) = cache.materialize(li);
gqa_attention_decode_one_query(
ops,
&q_rope,
&k_full,
&v_full,
&layer.o_proj,
cfg.num_attention_heads,
cfg.num_key_value_heads,
cfg.head_dim,
)
}
};
let dst = &mut attn_data[bi * hidden..(bi + 1) * hidden];
ops.copy_contiguous_into(&attn_out, 0, dst);Server decode (which runs batched, usually on the paged cache) and CLI decode (basic cache) both compile against the same code; each cache picks its lane at runtime.
That is the whole chapter: three Tensor methods, a rewritten LayerKvState, one trait method, and a match in two decode loops. Per decode step, the basic cache now costs one 4 KB append per tensor per layer and zero copies to read.
Running it
Same two prompts as IV.1. Short prompt (17 tokens, 64 generated), IV.1 baseline:
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)With the zero-copy cache:
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 263.060
decode_tokens_per_second: 46.071
per_forward_ms: min 20.569 max 263.060 mean 25.477 (n=64)44.1 to 46.1 tokens/sec, +4%. At a 30-ish token cache the copies were small, so the win is modest. The 511-token prompt is where the copies were big:
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)becomes
backend: simd
kv cache: basic
metrics:
time_to_first_token_ms: 10937.147
decode_tokens_per_second: 24.304
per_forward_ms: min 38.799 max 10937.147 mean 211.395 (n=64)20.6 to 24.3 tokens/sec, +18%. The per-forward dumps tell the same story in a different way: the IV.1 run's decode times sawtooth between about 42 and 56 ms as the allocator churns through cache-sized allocations every step, while this run sits flat at 39 to 41 ms. And the profiler's histogram closes the loop on where this chapter started: every decode step used to make 56 concat_dim0 calls, 392 across just the first seven steps of a generation; rerun the traced histogram now and concat_dim0 reads zero, because nothing on the decode path calls it anymore.
The output, as always, is token-for-token identical: the same floats end up in the same rows, we just stopped photocopying them first.
Where this leaves us
Two chapters into the act, decode is at 46 tokens/sec short-context and 24.3 at 511 tokens of context, from 41.3 and 15.9, and both wins came from deleting work that produced nothing: op-dispatch overhead in IV.1, redundant copies here. Both wins also grow with context length, because that is how the waste grew.
But everything so far still runs on one CPU core, while nine others sit idle. II.4 already tried threads once and decode barely moved; the next chapter figures out why, and builds a threadpool that actually earns its keep on decode.