A Rust inference server for Qwen2 on Apple Silicon

A from-scratch LLM inference engine — paged KV cache, continuous batching, KV quantization, YaRN — built on candle-core's Metal backend, targeting DeepSeek-R1-Distill-Qwen-1.5B (Qwen2 architecture, 1.5B).

Built as a 5-day learning project.

Lessons

• Day 0 — workspace, Metal smoke test, checkpoint verified
• Day 1 — forward pass, greedy decode, golden-logit gate
• Day 2 — KV cache, sampling, SSE streaming server
• Day 3 — paged KV cache, continuous batching
• Day 4 — KV quantization, prefix caching, benchmarks
• Day 5 — YaRN, long context, OpenAI compatibility

Layout

Commands

cargo run --release --bin smoke      # Day 0 gate: Metal + checkpoint
cargo clippy --workspace --all-targets -- -D warnings

# Reference logits for the Day 1 correctness gate (pin 3.12; torch wheels lag)
uv run --python 3.12 --with torch --with transformers --with numpy \
    python scripts/golden.py

Target architecture

Qwen2ForCausalLM   28 layers   hidden 1536   intermediate 8960 (SwiGLU)
12 Q heads / 2 KV heads (GQA 6:1)   head_dim 128   vocab 151936
rms_norm_eps 1e-6   rope_theta 10000   tie_word_embeddings false
q/k/v_proj have bias; o_proj and all MLP projections do not

KV per token = 2 × 28 × 2 × 128 × 2 bytes = 28,672 bytes. A 32k context is 0.94 GB of KV for a single sequence — which is why paging and KV quantization exist, and the thread running through Days 3–4.

Day 0 measurements

M5 Pro, 24 GB unified memory, 16-core GPU, macOS 26.5.2, rustc 1.96.1.

bf16 being only ~10% ahead of f32 says this size is launch/bandwidth bound rather than compute bound in candle's Metal GEMM — worth re-measuring at 4096² and at real prefill shapes before drawing conclusions about dtype choice.

Checkpoint: 339 tensors (28 × 12 + embed + norm + lm_head), 3.55 GB bf16, q/k/v biases present and o_proj bias absent as expected.

Day 2 measurements

F32 weights on Metal, single stream, measured over HTTP/SSE against a warm server. Median of three runs.

Measure warm. The first request after startup shows a TTFT of ~858 ms — Metal kernel compilation and paging in the mmap'd weights, not prefill.

What the KV cache bought

Same prompt, 20 tokens, byte-identical output at every stage:

Day 1 was O(N²); extrapolating it to 512 tokens gives roughly 80 minutes against 13 seconds measured.

Prefill and decode are different machines

Taking 1.78e9 parameters and 2·P·N FLOPs, against the 5.96 TFLOP/s ceiling measured on Day 0:

A 27× gap on identical weights. Prefill does T tokens' work per pass over the weights and is compute-bound; decode does one token's work per pass over the same 3.5 GB and is bandwidth-bound at ~2 FLOP/byte. That ratio is the argument for batching — 16 sequences read those weights once — and it is what Day 3 cashes in.

Decode falls from 40.8 to 35.3 tok/s as context grows to 480 tokens: attention is O(context) per step, and the cache readback copies the whole span with .contiguous() on every layer.

Sampling costs ~10 ms/token when enabled (37.3 → 27.4 tok/s), almost entirely the 608 KB device→host to_vec1 of the logits row. The fix is on-device top-k so only survivors cross the bus; not yet done.

Day 3 measurements

Paged KV cache (16-token blocks) plus continuous batching. Same machine, F32 weights on Metal, 40 tokens per request over HTTP/SSE.

Throughput against concurrency

Batching is worth 4.8x at concurrency 32, and the curve bends immediately — efficiency is already 58% at two concurrent requests. Three costs scale with batch size instead of amortising:

1. The KV gather. Every decode step copies each sequence's blocks into a contiguous tensor, pads to the longest, and concatenates. That is O(B x L) per layer per step. vLLM avoids it entirely by walking the block table inside the attention kernel; candle has no Metal equivalent (see gaps below).
2. repeat_kv in the batched path. GQA expansion is materialised 6x because the batched decode cannot use sdpa (see gaps).
3. Sampling. Sampler::sample copies a 151,936-float logits row to the host per sequence per step — 608 KB each, so 19 MB per step at concurrency 32, plus a CPU pass per sequence. On-device top-k would move only the survivors.

Weight reads do amortise, which is where the 4.8x comes from: one pass over 3.5 GB now serves the whole batch instead of one sequence.

TTFT grows roughly linearly with concurrency because prefill and decode alternate: a step is either one or the other, so an arriving request stalls every running sequence for one long step. Chunked prefill is the fix and is deliberately not implemented (see below).

Against Day 2

Known gaps

Gather, not a kernel. KVPool::gather uses index_select to make scattered blocks contiguous before sdpa. The memory-management win is intact — no fragmentation, no over-reservation, and prefix sharing becomes possible — but part of the bandwidth win is given up, because KV is copied every step rather than read in place. A hand-written MSL paged-attention kernel is the fix.

sdpa ignores masks during decode. On Metal, q_seq == 1 selects the vectorised kernel and call_sdpa_vector takes no mask argument. Batched decode pads ragged sequence lengths and must mask the padding — a zero key is not neutral, exp(q.0) = 1 gives it real attention weight — so the batched path uses hand-rolled attention instead.

KV pool capped at 2 GB. Above roughly 4 GB, candle's Metal backend hands out buffers that alias ones still in use: logits stay bit-exact while argmax returns the contents of the gather index tensor. Not a paging bug — verified max|diff| = 0 against the contiguous cache at every pool size.

Alternate prefill/decode, not chunked. Chosen for simplicity; the TTFT cost is visible in the table above.

Correctness

Paged KV is bit-identical to the contiguous cache: max|diff| = 0.000000 across prefill and 20 decode steps, with deliberately non-adjacent block tables. Deliberate over-subscription forces preemption (95 scheduler steps versus 61 with a roomy pool) and the output is byte-identical, so recompute is invisible to the client. The block invariant — free + held == total, two independently maintained numbers — is asserted every scheduler iteration.

Day 4 measurements — int8 KV

Symmetric int8, stored as u8 with a +128 offset because candle has no I8. Per-channel scales for K, per-token for V, following KIVI. Runtime-selectable:

KV_DTYPE=int8 cargo run --release -p server

The scorecard

Bytes per token: 57,344 (F32) against 14,784 (int8 = 14,336 codes + 448 scales; scales are 3% overhead).

Grouping is the whole question

Error is roughly scale/2 per element, and the scale is set by the largest magnitude in the group — so one outlier makes every other value in that group noise. Measured on a synthetic group, values sharing a scale with a 400x outlier came back 73x worse than the same values in a clean group.

Key and value tensors need different grouping. Keys carry persistent per-channel outliers, and K errors are amplified because they pass through q.k^T and then softmax; V errors pass through a linear weighted sum. Measured on real keys at layer granularity, K's absolute reconstruction error is 32x V's, at identical relative error — K values are simply larger.

Switching K from per-token to per-channel, on the real checkpoint:

Perplexity says the remaining 5.70% logit error costs 0.78% of predictive quality. Greedy decoding still diverges from the F32 trajectory after ~10 tokens, because argmax amplifies small logit differences at genuinely uncertain choices — the text differs, the model is not meaningfully worse.

The throughput result, which contradicts the premise

The usual argument is that halving KV bytes buys throughput directly, because decode is memory-bandwidth-bound. Measured here it does the opposite, and the gap widens with concurrency:

Bytes saved on the pool read are paid for twice in compute:

1. gather dequantizes on every read — to_dtype, affine, broadcast_mul, plus a second index_select for K's block-indexed scales. Three extra full-size tensor ops per layer per step.
2. write requantizes a whole 16-slot block per token. Per-channel K scales span a block, so a decode write dequantizes 16 slots, overwrites one row, recomputes the scale, and writes all 16 back — ~16x the quantize work, on every token, on every one of 28 layers.

Cost (2) scales with batch size, which is why the penalty grows from 7% to 28%. The premise holds only when dequantization is fused into the attention kernel; with dequant-on-gather it is the same trade as Day 3's index_select — the capacity win is kept in full and the bandwidth win is lost rather than merely forgone.

So int8 here buys 4x the concurrent sequences for 0.8% quality and a 28% throughput penalty. Worth it if capacity-bound, not if throughput-bound.

Partial blocks

A per-channel K scale spans all 16 tokens of a block, so it is not final until the block fills — and decode writes one token at a time. Rather than staging partial blocks in F32 (which would need per-sequence state that KVPool, being shared, does not have), a write requantizes the whole block: dequantize its 16 slots, overwrite the incoming rows, recompute the scale, write back. Unwritten slots are zeros and cannot raise a max, so the scale is correct without tracking fill level, and error does not compound — once a value sits on the quantization grid, requantizing at an unchanged scale returns it exactly.