Zain Ali

← All work

spindle

A faster engine for running language models.

2×

faster on realistic traffic, with output identical to the standard Hugging Face library.

The technical write-up

An inference engine built from the memory manager up: block-paged KV cache with copy-on-write sharing, iteration-level scheduling with preemption — over a from-scratch GPT-2 forward pass, because the interesting parts are exactly the parts a generate() API owns and hides. Correctness is pinned against HuggingFace: greedy output token-identical, prefill logits within 2e-3.

2.50×
less peak memory, paged vs contiguous
1.16×
the batching speed-up once generation lengths are uniform
GPT-2 124M · ~120-line forward pass · Apple M5, fp32, CPU · PyTorch
2026 · Inference engineering
  1. requests → scheduler (iteration-level, preemption) → block table → block pool → GPT-2 forward → tokensshared block, refcount
  2. scheduler (iteration-level, preemption) → block table → block pool

Why write the forward pass

The subject is memory management and scheduling. Both are owned and hidden by any generate() API — the library allocates the KV cache, decides its layout, and runs the loop — so handing it a block table is not an option. The GPT-2 forward pass is about 120 lines here; transformers is used only to fetch checkpoints and tokenize.

That carries an obligation to prove correctness rather than assume it. Prefill logits match HuggingFace to 2e-3 in float32 and greedy generation is token-identical. The test suite also checks that batched decoding produces exactly the same tokens as decoding each sequence alone, which is what catches attention leaking across sequence boundaries — the most likely way a hand-written paged attention goes subtly wrong.

Paging

paged vs contiguous · 64 requests, 8192 token slots, block size 16
paged vs contiguous · 64 requests, 8192 token slots, block size 16
peak slots heldcontiguous 8161 → paged 3264 (2.50×)
mean slot utilisationcontiguous 0.574 → paged 0.929 (1.62×)
contiguouspaged
peak slots held816132642.50× less
mean slot utilisation0.5740.9291.62×
steps to drain236192

64 requests replayed through both allocators, 8192 token slots, block size 16.

A contiguous per-sequence buffer must be sized at admission to prompt + max_new_tokens, because that is the only bound available then. A request that stops at 20 of a possible 192 tokens has still held all 192 slots for its entire lifetime. Blocks instead: a pool of fixed-size blocks and a per-sequence block table, so waste is bounded by the final partial block regardless of how long the sequence might have become.

Reference counts give copy-on-write forking, which contiguous allocation cannot do at all. Four sequences forked from a 16-token parent hold 64 logical tokens in 16 physical ones.

Continuous batching, and the control

continuous vs static batching · 24 requests, long-tailed lengths
continuous vs static batching · 24 requests, long-tailed lengths
throughputstatic 158.9 tok/s → continuous 325.6 tok/s (2.05×)
TTFT p50static 2356 ms → continuous 1124 ms (2.1×)
TTFT p99static 5262 ms → continuous 2447 ms (2.15×)
static (batch 8)continuous
throughput158.9 tok/s325.6 tok/s2.05×
forward steps383246
TTFT p502356 ms1124 ms2.1×
TTFT p995262 ms2447 ms2.15×

24 requests, GPT-2 124M, long-tailed generation lengths. Apple M5, fp32, CPU.

The control is more informative than the headline. With uniform generation lengths both policies take exactly 48 forward steps and the speedup collapses to 1.16×. Continuous batching is not faster per step — a static batch finishes when its longest member finishes, and every slot freed before then is idle. Remove the length variance and there is nothing to reclaim.

An earlier version of this benchmark reported static batching having better TTFT, 32 ms against 586 ms. It constructed each group’s requests just before running that group, so a request in the third group had its clock started after the first two finished — not measuring the queueing delay static batching imposes. With a common arrival time the numbers invert. A latency benchmark that does not model arrival is measuring service time.

Two bugs worth recording

The cache inferred positions
Sequence length advanced on write, and to stop twelve layers advancing it twelve times it advanced only on the last — so prefill at layer 0 saw a context length of zero and every sequence attended over nothing. Output was fluent English and 16.4 off HuggingFace in max absolute logit. The cache no longer infers: append takes an explicit start position, gather takes an explicit length, the runner advances once per step.
A stale distribution in speculation
When every proposal was accepted, the cache was already at prompt + output, so the resync had nothing to process and the next-token distribution silently kept its value from the start of the round. The bonus token was drawn from it. The symptom was not a crash but fluent text with tokens spliced in — “the capital, of France, is Paris., and the capital” — which reads as a mediocre model rather than a broken cache.