
Architflix Original
mini-vLLM
A small vLLM: Qwen2 written from scratch in about 700 lines of PyTorch, a paged KV cache, continuous batching with preemption, and an OpenAI-compatible streaming server. Checked token-for-token against Hugging Face, and up to 1.55× Hugging Face's generate throughput on a T4.
Overview
mini-vLLM is a small LLM inference engine by Archit Jain with the two ideas that make vLLM fast, a paged KV cache and continuous batching, serving a from-scratch PyTorch implementation of Qwen2 behind an OpenAI-compatible API, checked against Hugging Face's implementation and benchmarked on an NVIDIA T4.
Season 1 · The Build
5 episodesA small vLLM: Qwen2 written from scratch in about 700 lines of PyTorch, a paged KV cache, continuous batching with preemption, and an OpenAI-compatible streaming server. Checked token-for-token against Hugging Face, and up to 1.55× Hugging Face's generate throughput on a T4.
Built with Python, PyTorch, FastAPI, Hugging Face, CUDA.
mini-vLLM loads Qwen2.5 safetensors into a hand-written PyTorch Qwen2 (RMSNorm, rotary embeddings, grouped-query attention with qkv bias, SwiGLU, tied embeddings) whose forward pass takes a flat batch of tokens: a prefill is several prompts concatenated, each attending causally to itself, and a decode is one token per sequence attending to that sequence's cached keys and values. The KV cache is a preallocated pool of 16-token blocks per layer with an O(1) free-list allocator and a block table per sequence, so memory is claimed as sequences grow instead of being reserved for the maximum length. The scheduler does continuous batching: every step is either a prefill that admits waiting sequences under a token budget or a decode step for every running sequence, and when a decode needs a block and the pool is empty, the most recently admitted sequence is preempted, its blocks freed, and it is recomputed later from its prompt plus the tokens it already produced. A single engine thread owns the model while FastAPI handlers feed it through a thread-safe inbox, so concurrent HTTP requests land in the same batch, and a client that disconnects frees its blocks. Tests check logits and greedy output against Hugging Face transformers, batched output against running alone, and that recomputation after preemption reproduces the same tokens and returns every block. On a Kaggle T4 the engine reaches 1.07-1.55× Hugging Face generate throughput on mixed output lengths at 1-32 concurrent requests (243 vs 157 tokens/s at 8), but 0.95× at 64, because decode still copies blocks into a padded tensor with no fused paged-attention kernel or CUDA graphs. Greedy fp16 output matches Hugging Face exactly on 57 of 64 prompts, with the rest diverging after 13-59 tokens on fp16 near-ties, and the chat template matches on 128 of 128 prompts. Built by Archit with Claude as the AI coding assistant.
- Qwen2 written from scratch (RMSNorm, rotary embeddings, grouped-query attention with qkv bias, SwiGLU, tied embeddings), loaded from Hugging Face safetensors
- Paged KV cache: per-layer pools of 16-token blocks, an O(1) free-list allocator, a block table per sequence
- Continuous batching with prefill priority, a token budget and one block of headroom per running sequence
- Preemption by recomputation when the pool runs dry, tested to reproduce identical output and return every block
- OpenAI-compatible /v1/chat/completions and /v1/completions with streaming, temperature, top-p, seed and stop strings, plus /v1/models and /health
- Multi-byte-safe streaming; a disconnected client's blocks are freed immediately
- Runs on CUDA in fp16 or on a laptop CPU in fp32
Every project hits walls. Here's what nearly broke mini-vLLM — and how it shipped anyway.
- Token-for-token parity with Hugging Face's Qwen2 across KV block sizes of 1, 4 and 16
- Proving preemption is invisible: recomputed sequences must produce the same tokens and hand back every block
- Fair benchmarking: turned off Hugging Face's default repetition_penalty of 1.05 for true greedy decoding, added a mixed-length workload, and counted only useful tokens
- Explaining fp16 divergences (near-ties in the logits) against exact fp32 parity on CPU, instead of hiding them
- Tracking down an intermittent OpenMP deadlock in the engine thread on CPU, fixed by serving with a single torch thread
Every build leaves residue. The takeaways from mini-vLLM that carry forward.
- Most of vLLM's speed comes from memory management, not model code: paging the KV cache is what lets batches grow
- Benchmarks mislead by default; matching sampling settings and workloads took more care than the engine's first version
- Where it loses (64 concurrent requests) is as informative as where it wins: the next step is a fused paged-attention kernel and CUDA graphs
- Parity tests against a reference implementation turn "it seems to work" into a measurable claim