A from-scratch model of the systems core of a modern LLM serving engine — paged KV-cache, continuous batching, preemption, and prefix caching — that runs on a laptop with no GPU, and quantifies what each mechanism buys you with reproducible benchmarks.
The expensive, well-trodden part of inference is the GPU kernels. The interesting part — the part that decides whether your cluster serves 1× or several× the traffic on the same hardware — is memory management and scheduling. This project implements that part for real and isolates it from the matmuls: the compute is a deterministic latency model (so it runs anywhere and benchmarks reproduce exactly), while the block manager and scheduler are real, unit-tested implementations of the algorithms behind vLLM.
$ python -m mini_vllm.demo
mini-vLLM demo — continuous batching on a deliberately small KV pool
step 1 running= 7 prefill_tok=1024 decode= 0 KV= 38.3% preempt=0 done=0
step 2 running=11 prefill_tok=1024 decode= 6 KV= 75.0% preempt=0 done=0
step 3 running=13 prefill_tok= 637 decode=10 KV= 97.8% preempt=0 done=0
step 4 running=13 prefill_tok= 0 decode=13 KV= 97.8% preempt=0 done=0
...
final: requests=64/64 throughput=670.1 tok/s TTFT p50/p99=4147/10047 ms peakKV=100.0%
PagedAttention is, explicitly, virtual memory for the KV cache. Every mechanism here has an OS twin — which is what makes it a good way to learn the systems behind inference:
| OS concept | mini-vLLM | where |
|---|---|---|
| physical frame | physical block (fixed block_size tokens) |
block_manager.py |
| virtual page | logical block (index into a block table) | request.py |
| page table | Sequence.block_table (logical → physical) |
block_manager.py |
| demand paging | blocks allocated as tokens are computed | BlockManager.grow |
| internal fragmentation | the partially-filled last block | analysis.py |
| no external fragmentation | any free block backs any sequence | (the point) |
| page cache / dedup | prefix caching (hashed, shared full blocks) | BlockManager.admit_prefix |
| copy-on-write | fork a shared partial block before a write | BlockManager.grow |
| swapping to disk | swap KV out to a CPU block pool | swap_out / swap_in |
| CPU scheduling (preempt) | evict the lowest-priority running sequence | Scheduler._preempt |
| admission control | watermark + token budget per step | Scheduler._admit_waiting |
┌──────────────────────────────────────────────────┐
add_request │ LLMEngine │
───────────▶ │ simulated clock · metrics · run loop │
└───────┬───────────────────────────────┬──────────┘
│ schedule() │ step_latency_ms()
▼ ▼
┌────────────────────┐ ┌────────────────────┐
│ Scheduler │ │ ModelRunner │
│ waiting/running/ │ │ simulated cost: │
│ swapped queues │ │ prefill linear, │
│ continuous|static │ │ decode sublinear │
│ preempt:recompute │ │ in batch size │
│ or swap │ └────────────────────┘
└─────────┬──────────┘
│ can_grow / grow / admit_prefix / swap
▼
┌────────────────────────────────────────────────┐
│ BlockManager │
│ GPU block pool · CPU swap pool │
│ ref-counted blocks · prefix-cache hash table │
│ copy-on-write · per-sequence block tables │
└────────────────────────────────────────────────┘
All numbers are deterministic (seeded) and reproduced by
python benchmarks/run_all.py. Raw values: docs/results.json.
A contiguous allocator must reserve max_seq_len per sequence up front; a paged
allocator reserves ceil(len / block_size) blocks on demand, so the only waste
is one partial block per sequence. On a 32,768-slot KV budget with realistic
(log-normal, mean ≈ 240) lengths and 2048 max context:
| sequences that fit | KV utilization | wasted | |
|---|---|---|---|
| contiguous (reserve max) | 16 | 14% | 86% |
| paged (on demand) | 112 | 98% | 2% |
The right panel is the kicker: as you support longer context, contiguous capacity halves with every doubling (64 → 32 → 16 → 8 → 4) while paged stays flat.
Static batching drains a whole batch before admitting more, leaving the accelerator idle as sequences finish at different times. Continuous batching refills the decode batch every step. Sweeping offered load to saturation:
- static throughput ceiling: 746 tok/s
- continuous throughput ceiling: 1139 tok/s (1.53×)
Same Poisson sweep, restricted to the sustainable load range so both policies serve the same throughput — the difference is pure latency. Static batching makes an arriving request wait for the current batch to drain (head-of-line blocking), so its tail explodes; continuous batching admits immediately. Across the whole range continuous keeps p99 TTFT under 0.16 s while static climbs past 10 s. At 6 req/s: static p99 = 7,257 ms vs continuous 134 ms (≈54×).
For a chat/RAG/agent workload where requests share an identical system prompt, the shared blocks are hashed and computed once. With a 512-token shared prompt across 64 requests: 98% block hit-rate, 97% fewer prefill tokens computed, 50% lower median TTFT.
git clone https://github.com/WeishuZ/mini-vllm.git && cd mini-vllm
make setup # create .venv and install dev dependencies
make test # 30 tests: cache, scheduler, engine, trace, benchmarks
make demo # watch continuous batching + preemption live
make trace # write docs/trace.html
make bench # regenerate every plot + docs/results.jsonWithout make, run the same commands with .venv/bin/python -m ...; the core
package has zero runtime dependencies, while [dev] adds pytest and plotting.
Minimal API:
from mini_vllm import LLMEngine, CacheConfig, SchedulerConfig, ModelConfig, workloads
engine = LLMEngine(
CacheConfig(block_size=16, num_gpu_blocks=600, enable_prefix_caching=True,
prefix_cache_max_blocks=300),
SchedulerConfig(policy="continuous", max_num_seqs=128,
max_num_batched_tokens=2048, preemption_mode="recompute"),
ModelConfig(),
)
engine.add_requests(workloads.poisson(n=300, rate_rps=20, prompt_mean=256, gen_mean=128))
print(engine.run().summary())This repository is meant to be used as a teaching instrument, not just a code sample. Start with the staged guide, then work through the labs:
docs/learning-path.md— stage-by-stage curriculum.docs/labs/— hands-on exercises for token accounting, paged KV, continuous batching, preemption, prefix caching, and benchmarking.docs/vllm-mapping.md— how the mini-vLLM classes map to real vLLM concepts.docs/real-vllm-runbook.md— final bridge to a real OpenAI-compatible vLLM server.examples/— optional real-vLLM client commands and KV sizing.
For a portfolio-friendly walkthrough of the serving loop, generate the standalone trace page:
make traceThe page embeds one deterministic run and lets you scrub through each scheduler
step: waiting/running/swapped/finished queues, token-budget usage, GPU KV-block
occupancy, recompute preemptions, swap events, prefix-cache block hits, cache
evictions, and saved prefill tokens. It has no runtime dependencies and can be
opened directly from
docs/trace.html.
- Block manager. Fixed-size blocks in a ref-counted pool. Sequences grow by
demand-paging blocks as tokens are computed; the only waste is the partial
tail block. Prefix caching publishes completed prompt blocks to a content-
hash table (a block's hash includes its whole prefix, vLLM-style). Cached
blocks are ref-counted/pinned while active, then stay resident after the last
sequence frees them; an LRU evictor reclaims idle cached blocks when the cache
exceeds
prefix_cache_max_blocksor when allocation pressure needs space. A write to a shared partial block triggers copy-on-write. - Scheduler. Continuous batching advances in-flight decodes first, then admits new prefills with the leftover token budget — but only above a memory watermark and only when nothing was preempted this step. Those two guards are what keep it out of a preemption thrashing loop under oversubscription: during development, removing them turned a ~110-preemption run into a 6,600-preemption collapse. Chunked prefill keeps long prompts from blocking decodes. Preemption is recompute (default) or swap-to-CPU.
- Cost model. Prefill is linear in tokens; decode is sub-linear per token in batch size — that sub-linearity is exactly why batching pays off and is the property the throughput/latency benchmarks probe.
This is a learning instrument, not a serving framework. The compute is simulated — there are no real weights, kernels, or generated-text correctness; the contribution is the memory/scheduling control plane and the measurement of its behavior. The prefix cache now keeps an LRU-managed resident set for cross-request reuse, but it is still a compact model: there is no multi-tenant cache partitioning, CPU/off-GPU prefix spill, or sophisticated admission policy for deciding which prompts deserve cache space. Swap-based preemption is enabled only without prefix caching (so every block is privately owned and swap is lossless). These are the natural next steps, and each maps to a real vLLM subsystem.
mini_vllm/
block_manager.py paged KV cache: paging, prefix cache, COW, swap
scheduler.py continuous + static batching, admission, preemption
engine.py the serving loop, clock, and metrics
model_runner.py simulated forward-pass cost model
request.py Request / Sequence token accounting
analysis.py contiguous-vs-paged capacity math
trace_viewer.py standalone HTML trace generator
workloads.py deterministic burst / Poisson / shared-prefix workloads
benchmarks/ one script per result above + run_all.py
examples/ optional bridge scripts for real vLLM usage
tests/ 30 pytest unit + end-to-end tests
docs/ generated plots, trace, labs, and vLLM mapping notes
MIT © Weishu Zhang