nano-vLLM

A small LLM inference engine, built to understand how engines like vLLM work by building the core ideas and measuring each one:

• an OpenAI-compatible POST /v1/completions endpoint that streams tokens over SSE;
• continuous batching: users join and leave between decode steps;
• our own SmolLM2 forward pass (RMSNorm, RoPE, GQA, SwiGLU), checked against Hugging Face;
• a paged KV cache: one shared pool of 16-token blocks, a block table per user, nothing re-padded or copied between steps;
• prefix caching: users whose prompts start the same share those blocks and skip prefilling them.

Every change is measured against a baseline server (baseline_server.py, kept unchanged so it can be re-run) that keeps a separate KV cache per user and pads all of them to the same length every decode step.

Headline results

Apple M4, 16 GB, Mac GPU (MPS), SmolLM2-135M at fp32, each user asks for 100 tokens, at most 30 users decoding at once.

How it works

 client ──POST /v1/completions──▶ FastAPI handler (event loop)
                                    │ tokenize, check the length
                                    ▼
                               waiting line ──▶ BatchEngine thread, one round at a time:
                                                 1. free the blocks of finished or disconnected users
                                                 2. seat users from the front of the line while a seat
                                                    and enough blocks for prompt + max_tokens are free;
                                                    reuse cached blocks for a known prompt prefix
                                                 3. prefill each new user (only the uncached part)
                                                 4. decode one token for every seated user
 client ◀──SSE, one token each── handler ◀──── tokens through an asyncio.Queue

 inside one decode step, for each of the 30 layers:
   new token ─▶ RMSNorm ─▶ Q, K, V ─▶ RoPE at position n
                                  ├─▶ write K, V into block table[n // 16], slot n % 16
                                  └─▶ gather every user's blocks (one gather per layer, short
                                      tables filled out with dummy block 0) ─▶ masked attention
                                      (query heads grouped onto their shared K/V head) ─▶ o_proj
                ─▶ RMSNorm ─▶ SwiGLU MLP

 shared pool (one K and one V tensor per layer: [blocks, 16, 3 KV heads, 64])
 ┌───┬───┬───┬───┬───┬───┬───┬───┬───┐
 │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │ 6 │ 7 │...│   block 0: dummy, never handed out
 └───┴───┴───┴───┴───┴───┴───┴───┴───┘
 user A table: [1, 2, 5]    user B table: [1, 2, 6]    ← both start with the same
                                                          32-token system prompt,
                                                          so blocks 1 and 2 are shared

How to run it

python3 server.py

In a second terminal:

curl -N -X POST localhost:8080/v1/completions -H "Content-Type: application/json" \
     -d '{"prompt": "The cat sat on the", "max_tokens": 20}'

python3 baseline_server.py runs the baseline instead, on the same port, so run only one at a time. NANO_PREFIX_CACHE=0 python3 server.py turns prefix caching off.

Load test (the server must be running), then redraw the charts:

nice -n 10 python3 bench.py --label my-run --users 1 10 30 40
nice -n 10 python3 bench.py --label shared --users 1 10 30 --system-repeats 35
python3 plot.py

Tests, each printing what it compares:

python3 test_blocks.py      # allocator: toy pool, sharing, reviving, eviction
python3 test_kv_pool.py     # write path: values read back exactly through tables
python3 test_forward.py     # our forward pass against Hugging Face, piece by piece
python3 test_paged.py       # paged prefill/decode/prefix against Hugging Face
python3 test_engine.py      # the server's engine: greedy answers equal HF generate

Where the time went

Before: half of every step was rearranging the cache

The baseline's stopwatches (waiting for the GPU before each reading) showed where one decode step goes:

At 30 users, every step padded each user's K and V separately in every layer: 1,800 pads and 60 concatenates, all to build a rectangle that was thrown away after one step. Slicing back was almost free because slicing makes a view, not a copy.

After: no cache rearranging, and faster attention

1. Paged KV cache (58 ms per step). K and V are written once into their block and never moved. A step does one gather per layer for K and one for V, whatever the number of users.
2. Attention by reshape instead of SDPA (48 ms). On MPS, PyTorch's SDPA with a mask and grouped heads took 3.3–3.5× longer than a plain matmul-softmax-matmul where each group of 3 query heads is reshaped onto its shared K/V head (no copy). Same results to 6e-7.

What is left: attention grows with users × context

profile_decode.py with every part timed separately:

(The parts add up to more than the step because of the extra GPU waits.)

• With 1 user, the step is mostly fixed overhead. Reading all 540 MB of weights at about 120 GB/s takes about 4.5 ms, but the step takes 19 ms: about 30 layers × ~20 small GPU operations, each with a launch cost. Paging cannot fix that; fusing operations (or CUDA graphs on NVIDIA) would.
• With long contexts, the gather and attention dominate. A step costs about 13–14 ms per 1,000 user-tokens, the same for 30 × 540, 16 × 1,000 and 8 × 2,000. The gather is still a copy: real vLLM's PagedAttention kernel reads blocks in place, which needs a custom GPU kernel.
• Shared prompt blocks are gathered once per user. With a 500-token shared system prompt, 30 users gather the same blocks 30 times every step, so prefix caching helps prefill (TTFT) but not decode.

Prefix caching

Every user sends the same ~500-token system prompt plus their own question. With caching on, the first user prefills everything; everyone after reuses the cached full blocks and prefills only the last block and their question.

Correctness

Speed means nothing if the answers change, so every step has a test that compares against Hugging Face with sampling turned off:

Two lessons from testing:

• Comparing only the chosen token hides bugs. RoPE broken on purpose (theta 10,000 instead of 100,000) moved the logits by 1.49 but kept the same top token. The tests compare logits.
• The garbage pool proves the mask works. Masked slots get weight 0, so huge garbage values there change nothing (0 × garbage = 0). NaN would not work as garbage: 0 × NaN is still NaN.

Decision log

Limits

• A tiny model (135M) on a Mac GPU, fp32, one machine, no tensor parallelism, no quantization.
• The paged attention is a gather plus plain PyTorch operations, not a fused GPU kernel, so the gather is still a copy and long contexts cost about 13–14 ms per 1,000 user-tokens per step.
• New users are prefilled one at a time, so TTFT grows with the number of users arriving together (575 ms p50 for 30 at once).
• Reserving prompt + max_tokens at the door wastes blocks when answers stop early; there is no preemption.
• Measurements on MPS vary between runs (the same decode step measured 98 ms and 190 ms in different sessions); compare runs from the same session.
• The load test client runs on the same Mac as the server.
• No live metrics endpoint: the numbers come from bench.py and the server's stopwatch reports.

Next steps

1. Batched prefill: prefill everyone who arrives together in one forward pass (with a token budget) to cut TTFT.
2. Preemption instead of reserving max_tokens up front.
3. A real paged attention kernel (Metal or Triton on NVIDIA) that reads blocks in place, so shared prefix blocks are read once, not gathered per user.
4. A /metrics endpoint with TTFT, ITL, batch size and free blocks while the server runs.