inference-optimization

A serving-layer study for LLM inference: a production-shaped control plane in front of vLLM, plus a from-scratch mini-engine that rebuilds paged KV caching, prefix caching and continuous batching so the mechanisms are readable instead of vendored — with a measured ablation, not just an argument.

Headline result, measured on 2× RTX 2080 Ti (Llama-3.2-3B, fp16):

	prefix cache ON	OFF	speedup
Mini-engine, cacheable requests (mean TTFT)	133 ms	736 ms	5.5×
Full stack via gateway, steady state (median TTFT)	64 ms	328 ms	5.1×
Control: uncacheable prompt	93 ms	91 ms	1.0×

The flat control line is the point: the effect appears exactly where prefixes repeat and nowhere else. Outputs were byte-identical across both runs (greedy decoding), so the win costs no quality. Every request is traced to Weights & Biases Weave with cached_tokens, hit rate and per-phase timings.

Run it yourself: python -m miniengine.experiment_cache_ablation.

---

Why this exists

Serving an LLM economically is not one problem, it's several with opposite shapes, and the interesting engineering lives in the tension between them:

• Batch size buys cost (it amortizes the weight read across sequences) — its latency curve is gentle.
• Utilization costs tail latency — its curve is hyperbolic; queue wait scales as ρ/(1-ρ) and worsens with service-time variance (M/G/1).
• Prefill is compute-bound and sets TTFT; decode is memory-bandwidth-bound and sets TPOT. They want opposite batch sizes.
• KV cache memory, not FLOPs, is usually the binding resource, so fragmentation directly caps the batch size that pays for everything else.

Most of these are settled in the literature. What isn't handed to you is the layer above the engine — admission control, class-aware scheduling, cache policy, and the instrumentation that proves any of it still works next month. That layer is what this repo is actually about.

Three layers, three relationships to the code

serving/ — the layer you write. Session-affinity routing over a consistent-hash ring, an EDF scheduler with a starvation floor, an admission ladder anchored to a measured goodput knee, tier-aware cache eviction, a preemptible batch lane, adaptive speculative-decode steering, and a verification suite. None of this lives inside an inference engine, because it encodes business rules an engine cannot know for you.

serving/workers/ — the layer you configure. Continuous batching and paged KV are vLLM defaults; prefix caching, chunked prefill and P/D disaggregation are flags. Re-implementing them in production would be a mistake. Knowing which flag encodes which mechanism — and what it costs on Turing-class hardware — is the skill.

miniengine/ — the layer you learn. ~600 lines of PyTorch reproducing the block manager, the hash-chain prefix cache, and a continuous-batching scheduler with chunked prefill, running against a real transformers model. It exists so the mechanism has somewhere to be read and stepped through. The gather it performs on every decode step is deliberately visible: that copy is precisely the work PagedAttention's CUDA kernel eliminates.

Start here

The five trace documents in docs/ walk one request through the system, each with a sequence diagram and real numbers. Read them in this order — each supplies vocabulary the next assumes:

1. Prefix cache: insert and match — how blocks enter the cache and how the next request finds them (hash chains, refcounts)
2. One prefill chunk — how a 475-token prompt costs 27 tokens of work, and where TTFT is actually spent
3. One decode step — the block gather, the 6.4 GB weight read for a single token, and why batching is the answer
4. EDF + floor admission — why pure deadline scheduling starves the class you're paid to protect, with a runnable demo
5. Multi-user cache lifecycle — what's shared between tenants, what isn't, and the state machine of one KV block

Two things are runnable without a GPU:

python scenarios/starvation_trap.py      # pure EDF vs EDF+floor, simulated
python -m loadgen.sweep --help           # goodput-knee load sweep

starvation_trap.py prints the result that motivates the whole scheduler: under a flood of low-priority traffic, pure EDF pushes the protected class to a 7.11 s p95 wait; an 8% admission floor brings it to 0.01 s while the free tier moves 10.46 s → 10.42 s. The guarantee is nearly free — when the protected class is small.

Running the full stack

pip install -e .[miniengine,tracing]
cp deploy/.env.example deploy/.env        # set HF_TOKEN, MODEL_ID
cd deploy && docker compose --profile replicated up

Gateway on :8080 (OpenAI-compatible, SSE), Prometheus :9090, Grafana :3000. A disagg profile runs prefill and decode on separate GPUs with KV transfer between them. To reproduce the ablation end to end, drive identical traffic with loadgen/stack_ablation_driver.py once normally and once with the deploy/docker-compose.nocache.yml overlay, then compare the two experiment tags in Weave.

What it measures about itself

Every audited production codebase I've looked at logs prompt and completion tokens and never cached_tokens — which makes prompt caching invisible on every dashboard the team owns. This repo treats that as the first-class metric: usage.prompt_tokens_details.cached_tokens from the worker flows into Prometheus (serving_cached_tokens_total per tier) and into Weave, tagged by experiment so cache-on and cache-off runs compare directly.

Three continuous controls ship with it, each aimed at a failure mode that is otherwise silent:

• Cache correctness — cached-prefix output is asserted byte-identical to cold-prefill output on a sampled basis. A prefix-cache bug corrupts generations rather than crashing.
• Starvation probe — synthetic protected-class requests injected continuously, asserting the scheduler floor still holds under load.
• Burn-rate alerting — error-budget consumption rate, not absolute thresholds, so slow degradation pages someone before the budget is gone.

Measured: the eviction pressure sweep

The 5x cache ablation above answers "does caching help." The follow-up experiment answers the harder questions: how high does multi-turn hit rate actually climb, and what happens when concurrent load pushes the KV cache past capacity and LRU eviction starts. A frozen 64-persona / 8-turn workload (loadgen/datasets/eviction_sweep_v1.json) replayed at four auto-calibrated pressure levels, cache on/off, 960 requests total:

• The hit-rate ceiling is structural, ~88% here — not the ~98% folklore number. Each turn's prompt necessarily contains its own uncached newest turn: ceiling ≈ 1 − newest_turn_tokens / prompt_tokens.
• Eviction has a sharp two-phase knee. At exactly 1× capacity, eviction begins but costs nothing (294 evictions, zero re-missed tokens). At 2×, it costs plenty: hit rate 88% → 54%, 32k re-missed tokens. At 4×: 31%.
• The stress test found a real scheduler bug — the admission gate over-committed the KV pool because it never accounted for blocks already promised to running sequences (fixed with a virtual reservation ledger; the war story and CPU repro are in the write-up).
• The pattern was confirmed off the toy engine: the same workload replayed through the full gateway → vLLM stack at 2× shows cache-on beating cache-off at every turn depth (1.82× mean TTFT).

Full write-up with per-turn tables and war stories: docs/eviction-sweep.md. The analysis itself was run through W&B's ARIA research agent — six natural-language questions, every number cross-checked against local ground truth (docs/aria-session-log.md); the published report is here.

Hardware notes (Turing / RTX 2080 Ti)

Compute capability 7.5 means: fp16 only (no bf16, no FP8), xFormers attention backend (no FlashAttention-2), and vLLM falls back to its V0 engine — which reports prompt_tokens_details: null, so cached_tokens reads 0 on this hardware even at a 95% hit rate. The evidence in the table above therefore comes from A/B TTFT and the worker's own hit-rate metric; on Ampere and later the field populates and the gateway reports it directly. GPU 0 also drives the desktop on this machine, leaving too little VRAM for a second worker — so the stack runs single-worker and the router's dead-node path gets exercised for real.

Pin vLLM to 0.8.x for Turing; later versions drop V0.

What this is not

A production system. There's no auth, no TLS, no multi-region, no autoscaling, and no test suite — verification here is compile checks, runnable scenarios, and the measured experiment. It's a study rig built to understand the tradeoffs and to have something concrete to argue from. deploy/k8s/ sketches what changes at fleet scale (HPA on queue depth rather than CPU, PDBs sized so evictions can't break the scheduling floor, priority classes for batch preemption) without pretending to be deployable.

How this was built

With an agentic workflow: a written spec, an implementation plan decomposed into tasks, one agent implementing each task, and an independent review agent gating every diff before it landed. The commit history is the artifact — fix-round commits like "edf_floor empty-queue crash, full-queue EDF scan, cancellation + tick supervision" are review findings caught before merge, not bugs found later. Roughly ten defects were caught this way, including a scheduler deadlock that only triggered on fully-cached prompts.

The measurement work was the most valuable part, and not for the reason expected: instrumenting the cache ablation surfaced a gateway bug that had nothing to do with caching. A fresh HTTP client per request was paying 0.3–1.2 s of DNS and connect time, burying a 5× effect under connection churn until per-phase timing exposed it. Measurement paid for itself before it measured the thing it was built for.

References

• PagedAttention / vLLM — Kwon et al., SOSP 2023
• Continuous batching — Orca, Yu et al., OSDI 2022
• Chunked prefill — Sarathi-Serve, Agrawal et al., OSDI 2024
• Prefill/decode disaggregation and goodput — DistServe, Zhong et al., OSDI 2024; Splitwise, Patel et al., ISCA 2024
• Radix-tree prefix reuse — SGLang, Zheng et al.
• EDF optimality and overload behavior — Liu & Layland, 1973
• Speculative decoding — Leviathan et al., 2023
• Error budgets and burn-rate alerting — Google SRE Workbook

License

MIT — see LICENSE.