A from-scratch LLM inference server with paged KV-cache, continuous batching, and speculative decoding.
vLLM-lite: the serving stack — forward pass, memory manager, scheduler, and API — implemented by hand, not called from a library.
tinyserve implements the machinery an inference engine like vLLM is built from — a manual transformer forward pass over a paged KV-cache, an iteration-level (continuous batching) scheduler with preemption, an OpenAI-compatible HTTP server, and speculative decoding with exact rejection-sampling verification. transformers is used for nothing but loading checkpoint weights and the tokenizer; every attention kernel, cache block, scheduling decision, and sampling step is original code.
The whole thing runs CPU-only (CUDA optional), so you can read it, test it, and benchmark the policies end-to-end on a laptop with no GPU.
The headline result: under Poisson load at 8 req/s, continuous batching cuts p99 time-to-first-token from 3.10 s → 0.46 s (6.7×) while delivering 1.58× the throughput of static batching — and it wins on both axes of the latency/throughput frontier simultaneously. Jump to benchmarks ↓
┌──────────────────────────────────────────────────┐
HTTP (asyncio) │ /v1/completions /health /metrics │
│ │ bounded FairRequestQueue (429 ⇠ full) │
└────────┼─────────────────────────────────────────┘
│ admit when working set has room
Engine thread ┌────────▼─────────────────────────────────────────┐
│ EngineCore.step() ── scheduler ── KV manager │
│ │ one fused forward pass / step │
│ ModelRunner → Qwen2ForCausalLM ⇄ Paged KVPool │
│ ▲ drafts (ngram / draft model) │
└────────┴─────────────────────────────────────────┘
- Highlights
- Quickstart
- Benchmarks
- How it works
- The HTTP API
- Memory accounting
- Correctness & testing
- Configuration
- Project layout
- Roadmap
| Subsystem | What's implemented from scratch |
|---|---|
| Forward pass | Qwen2 decoder — RMSNorm, rotary embeddings, grouped-query attention, SwiGLU MLP — operating on a flattened ragged batch so sequences at any length/phase fuse into one pass |
| Paged KV-cache | Fixed-size blocks, per-sequence block tables, a reference-counted allocator, copy-on-write on shared blocks, and chained-hash prefix caching with LRU eviction |
| Scheduler | Iteration-level admission/eviction, chunked prefill, per-step token budget, and preemption by swapping KV to a host arena (or recompute when the arena is full) |
| Serving | OpenAI-compatible /v1/completions with SSE streaming, a priority + client-fairness queue, 429 + Retry-After backpressure, and Prometheus /metrics |
| Speculative decoding | N-gram (prompt-lookup) and draft-model proposers, verified by rejection sampling that provably preserves the target distribution |
| Rigor | Golden-output tests vs the HuggingFace reference (logit parity within 1e-4), plus a statistical test that verified sampling matches the target law |
Design invariant that runs through the whole codebase: outputs are correctness-preserving under every optimization. Greedy generation is bit-identical whether prefix caching, preemption, or speculation is on or off — the KV either moves losslessly or is recomputed, and speculative verification is mathematically exact.
pip install -e ".[dev]"
# Run the full suite — tiny random-weight models, no downloads (~40s on CPU)
pytest -m "not slow"
# Serve a real checkpoint (downloads Qwen2.5-0.5B, ~1 GB, on first run)
python -m tinyserve.server.app --config configs/default.yaml# Stream a completion (OpenAI-compatible)
curl -N http://127.0.0.1:8000/v1/completions \
-H 'Content-Type: application/json' \
-d '{"prompt": "The three laws of robotics are", "max_tokens": 64, "stream": true}'Point any OpenAI client at http://127.0.0.1:8000/v1:
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="not-needed")
for chunk in client.completions.create(
model="tinyserve", prompt="def fibonacci(n):", max_tokens=128, stream=True
):
print(chunk.choices[0].text, end="", flush=True)The offline (batch) API skips HTTP entirely:
from tinyserve.engine.llm import LLM
from tinyserve.engine.sampling import SamplingParams
llm = LLM(model="Qwen/Qwen2.5-0.5B")
out = llm.generate(["Once upon a time"], SamplingParams(temperature=0.8, max_tokens=64))
print(out[0].text)Static vs. continuous batching under Poisson arrivals, 30 requests, prompts 16–96 tokens, outputs 16–64 tokens. Reproduce with:
python -m tinyserve.bench.sweep_pareto --size small --rates 1,2,4,8 --num-requests 30| rate (req/s) | policy | throughput (tok/s) | p99 TTFT (s) | p99 e2e (s) |
|---|---|---|---|---|
| 4 | static | 154.6 | 1.168 | 2.304 |
| 4 | continuous | 167.2 | 0.074 | 1.592 |
| 4 | continuous + spec | 178.4 | 0.121 | 1.198 |
| 8 | static | 181.3 | 3.102 | 4.127 |
| 8 | continuous | 287.0 | 0.462 | 1.732 |
| 8 | continuous + spec | 329.5 | 0.142 | 1.027 |
What the numbers say:
- p99 time-to-first-token is the story. Static batching makes a new arrival wait for the entire current batch to finish decoding before it gets a slot, so TTFT scales with batch duration. At 8 req/s that's 3.10 s vs 0.46 s for continuous — a 6.7× reduction, and 22× against continuous+spec.
- Throughput diverges under load. Static wastes slots on finished-but-not-yet-drained sequences; continuous frees a sequence's blocks the step it finishes and admits the next arrival immediately. At 8 req/s that's 1.58× the throughput (287 vs 181 tok/s).
- Continuous batching is strictly Pareto-dominant here — not a latency/throughput trade. See
benchmarks/pareto.pngand the full methodology + memory accounting in benchmarks/REPORT.md.
Numbers are from a CPU run on a synthetic model — policy comparisons depend on step-time structure (how prefill/decode interleave, when blocks free), not weight values, which is why a small model gives fast, download-free, meaningful ratios. Every bench script takes
--model Qwen/Qwen2.5-0.5Bto run against the real checkpoint.
Full diagrams live in ARCHITECTURE.md; the essentials:
K/V for every layer live in two contiguous slabs shaped [num_layers, num_blocks × block_size, num_kv_heads, head_dim]. A sequence never owns a contiguous cache region — it owns a block table mapping logical positions to physical slots:
seq A (19 tokens) block table → [7, 2] block 7 (rc=3) ← shared 16-tok prefix
seq B (shares pfx) block table → [7, 5] block 2 (rc=1) ← A's tail
block 5 (rc=1) ← B's tail
prefix cache: hash(chain) → block 7
- A reference-counted allocator underpins sharing. Forking a sequence or matching a cached prefix just bumps refcounts; a write into a block with
refcount > 1triggers copy-on-write so writers never corrupt shared readers. - Prefix caching keys full blocks by a chained hash
hᵢ = H(hᵢ₋₁, tokensᵢ), so a block's identity encodes its entire prefix, not just its own tokens. A repeated system prompt resolves to the same hash chain and its KV is reused without recompute (the last matched token is always recomputed so prefill still emits logits). - Allocation is all-or-nothing per step: growth + COW blocks are counted first, LRU cache entries evicted if the free list is short, and only then are tables mutated. A failure surfaces cleanly to the scheduler as a preemption decision instead of corrupting state mid-forward.
One step() = one scheduler decision + one fused forward pass over a flattened ragged batch (prefill chunks and single-token decodes mixed freely). Each step the scheduler rebuilds the working set:
- Decodes first (
(priority, arrival)order) — protects running requests' inter-token latency from arrival bursts. - Swap-ins — preempted sequences return oldest-first as blocks free.
- Chunked-prefill admission — new prompts fill whatever token budget remains, so one huge prompt can't stall everyone's decode.
Under memory pressure the lowest-priority running sequence is preempted: its KV is swapped out to a host arena (or, if that's full, dropped and recomputed later). Either path is output-preserving — swapped KV is lossless, recompute replays deterministically. A guard detects the genuinely-stuck case (no sequence can ever fit) and raises instead of spinning.
A decode step is expanded from 1 input token to 1 + k (last accepted token + k drafts) inside the same fused forward pass; the stepped causal mask scores every draft position at once. Drafts come from either proposer:
- N-gram / prompt-lookup — finds the longest recent n-gram that occurred earlier and proposes what followed. Free, no second model, strong on code/structured/repetitive output.
- Draft model — a smaller model runs
ksteps with its own paged shadow KV-cache, rolled back to the longest common prefix with the target after each rejection.
Verification is rejection sampling (Leviathan et al. 2023): draft dᵢ is accepted w.p. min(1, pᵢ(dᵢ)/qᵢ(dᵢ)); on rejection the replacement is drawn from the residual normalize(max(p−q, 0)); if all k survive, a bonus token is drawn from the final distribution. This reproduces sampling from the target distribution exactly — verified both by bit-identical greedy output and by a statistical test of the sampler against the target law.
Honest about when it hurts. Every round pays for k+1 positions to emit ≥1 token; with low acceptance that's a net slowdown (a mismatched draft model measured at 0.16×). The classic win comes from decode being memory-bandwidth-bound on GPU, where verifying k+1 tokens costs barely more than 1 — a structural advantage that mostly isn't there on a compute-bound CPU. spec_bench.py reports acceptance and speedup across temperatures and workloads, including the losses.
| Endpoint | Purpose |
|---|---|
POST /v1/completions |
OpenAI-compatible. String prompts (tokenized server-side) or raw token-id arrays; stream=true → SSE with incremental detokenization and a [DONE] terminator. Extensions: top_k, priority, ignore_eos. |
GET /health |
Queue depth and engine occupancy (waiting / running / swapped). |
GET /metrics |
Prometheus text: request counters, TTFT & e2e latency histograms, queue depth, and engine gauges (free blocks, preemptions, prefix-cache hits, COW copies, speculative acceptance). |
Backpressure is real, not cosmetic. The engine runs on its own thread behind an asyncio bridge; admission goes through a bounded FairRequestQueue (priority level → round-robin across clients → FIFO within a client). A full queue returns 429 with Retry-After before the engine ever sees the request — load is shed, not silently buffered into blown SLOs. Client disconnects and per-request timeouts abort the sequence and free its blocks between steps.
An open-loop HTTP load generator ships in-tree:
python -m tinyserve.bench.loadgen --url http://127.0.0.1:8000 --rate 4 --num-requests 50EngineCore.memory_report() itemizes every resident byte. Per-block KV is
2 (K+V) × num_layers × block_size × num_kv_heads × head_dim × dtype.
For Qwen2.5-0.5B (24 layers, 2 KV heads, head_dim 64) that's 384 KiB per 16-token block; the default 2048-block pool is 768 MiB of KV (32,768 token-slots) on top of ~1 GiB of fp32 weights. Paging is what makes that budget stretch: a contiguous per-request cache sized for max_model_len would reserve ~48 MiB per request whether used or not; blocks are handed out 384 KiB at a time on demand. CacheConfig.from_memory_budget(model, kv_cache_gb=...) sizes the pool to a target, and live utilization is exported at /metrics.
95 tests, all green in CI (ruff + pytest matrix on Python 3.10–3.12), every fast test on random-weight tiny models so CI needs no downloads.
| Layer | What's proven |
|---|---|
| Model parity | Logits match the HF reference within 1e-4 through ragged batches, chunked prefill, and incremental decode; greedy generation is token-identical. TINYSERVE_REAL_MODEL=Qwen/Qwen2.5-0.5B pytest -m slow runs the same parity against the real checkpoint. |
| Paged cache | Refcount lifecycle, COW on shared blocks, swap round-trip losslessness, LRU prefix eviction, OOM leaves state untouched. |
| Scheduler | Per-step token budget & seq cap, priority ordering, abort lifecycle, and output-parity under both swap and recompute preemption. |
| Serving | OpenAI response shape, SSE framing, 429 + Retry-After, Prometheus format, fairness-queue ordering. |
| Speculative | Rejection sampler preserves the target distribution (statistical test), greedy bit-parity spec-on vs spec-off, correct stop-token handling mid-round. |
Everything is config-driven via configs/default.yaml (or TinyServeConfig.from_yaml):
cache:
block_size: 16
num_blocks: 2048 # device KV pool
num_swap_blocks: 1024 # host arena for preemption
enable_prefix_caching: true
scheduler:
max_num_seqs: 16
max_num_batched_tokens: 512 # per-step token budget
enable_chunked_prefill: true
max_model_len: 4096
spec:
enabled: false
method: ngram # "ngram" | "draft_model"
num_speculative_tokens: 4src/tinyserve/
├── models/ manual Qwen2 forward pass + HF weight loader
├── kv/ paged pool, refcounted allocator, cache manager, prefix cache
├── engine/ sequence state, sampling, model runner, scheduler, engine loop
├── spec/ speculative proposers + rejection-sampling verifier
├── server/ FastAPI app, async engine bridge, fairness queue, metrics
└── bench/ Poisson workloads, static/continuous runners, Pareto & spec sweeps, load gen
tests/ 95 tests — golden parity, cache mechanics, scheduler, serving, spec
- CUDA paged-attention kernel (the current gather+SDPA path is CPU-first and correctness-oriented)
- Tensor parallelism across devices
- Guided / constrained decoding (JSON schema, regex)
- Multi-LoRA serving with per-adapter block tables
- Disaggregated prefill/decode
MIT — see LICENSE.