AskAiml
Sign in Get started

Mini inference engine (KV cache + sampling + quantization)

Serve a model in a compact engine: implement the KV cache, sampling, and INT4 quantization end to end.

🛠 Capstone

Goal

Build a small but real LLM inference server. It must support:

  • Continuous batching.
  • A KV cache (PagedAttention-style block allocator is the stretch goal).
  • Standard sampling (greedy, top-k, top-p, min-p, temperature, frequency penalty).
  • Streaming SSE responses to multiple concurrent clients.
  • Optional INT4 / FP8 weight loading.
  • An OpenAI-compatible /v1/chat/completions endpoint.

This is the project that proves you can do systems work, not just notebook code. Inference is where most of the production cost lives.

Time: 3-5 weeks part-time.

Prerequisites

  • 02_transformer/04_build_gpt_from_scratch.md
  • 05_inference/02_kv_cache.md
  • 05_inference/01_sampling_decoding.md

Tech stack

  • PyTorch 2.x (F.scaled_dot_product_attention)
  • FastAPI + uvicorn
  • A model loader for HF safetensors weights (no need to use the HF model class; you'll load weights into your own).
  • bitsandbytes or awq-py for quantized weights (optional)

Architecture

mermaid
flowchart LR
    CLIENT -->|/v1/chat/completions| API[FastAPI server]
    API --> SCHED[Scheduler queue]
    SCHED -->|micro-batch| ENG[Engine: forward + sample]
    ENG --> KV[KV cache pool]
    ENG --> WEIGHTS[(model weights)]
    ENG -->|tokens| SCHED
    SCHED --> STREAM[SSE streamer]
    STREAM --> CLIENT

Step-by-step

1. Reload a real model into your own architecture

Pick a small, well-known model (gpt2, Qwen2.5-0.5B, Llama-3.2-1B). Load its safetensors weights into the modern decoder you built in Project 2 (or write a fresh one for this model). Verify outputs match HF's model.generate(do_sample=False) exactly on a fixed prompt — bit-exact agreement on the first 32 tokens is the gold-standard sanity check.

2. KV cache

Implement a KVCache per request:

  • Pre-allocate (L, h_kv, max_len, d_h) per request.
  • Append on each step.
  • Provide a clear() method.

(Stretch: a paged allocator with 16-token blocks shared across requests.)

3. Sampling module

Implement, configurable per request:

  • temperature (with T=0 → argmax)
  • top_k
  • top_p
  • min_p
  • frequency_penalty
  • presence_penalty
  • seed (deterministic generation)

Match OpenAI parameter names where possible.

4. Continuous batching scheduler

Maintain three queues:

  • pending — requests waiting for prefill.
  • running — requests in decode.
  • done — finished; awaiting return.

Each iteration:

  1. Move new requests into running, prefilling their prompts.
  2. Run one forward step for the batched running requests.
  3. Sample one token per request; append.
  4. Stream tokens to clients via asyncio.Queues.
  5. Move EOS/maxlen requests to done.

Cap concurrency on KV memory budget.

5. FastAPI server

Endpoints:

  • POST /v1/chat/completions — OpenAI-compatible chat. Supports stream: true (SSE).
  • POST /v1/completions — legacy completion endpoint.
  • GET /v1/models — list loaded models.
  • GET /metrics — Prometheus metrics (request count, latency, tokens/s).

Test with the official openai Python SDK pointing at your server's URL.

6. Quantization (stretch)

Add a path to load AWQ-quantized weights (or just bitsandbytes 4-bit linear layers). Verify quality with a small perplexity check.

7. Speculative decoding (stretch)

Plug in a small draft model (Lesson 5.3). Show ~2× decode speedup at batch=1.

8. Benchmark

Compare to vLLM and TGI on:

  • Tokens/s at batch=1.
  • Tokens/s at batch=8.
  • Time-to-first-token under load.
  • Memory footprint at T=8k.

Be honest: vLLM will likely beat your engine. Document by how much and why (PagedAttention, fused kernels, optimised CUDA).

Acceptance criteria

  • Loads at least one real HF model and matches its outputs bit-exact at temperature 0 for the first 32 tokens.
  • Implements KV cache; per-token cost roughly constant for a single request as T grows (no recompute).
  • Implements all sampling controls listed above.
  • Continuous batching of ≥ 8 concurrent requests.
  • OpenAI-compatible JSON over HTTP, with SSE streaming.
  • Prometheus /metrics with tokens_generated_total, request_latency_seconds, kv_cache_used_bytes.
  • README with benchmark numbers vs vLLM (small model is fine).

Stretch goals

  • PagedAttention-style block KV pool with prefix cache sharing.
  • AWQ INT4 weight loading.
  • Speculative decoding with a draft model.
  • Tensor parallel across 2 GPUs.
  • Tool-call streaming (tool_calls field deltas).
  • Support JSON-grammar sampling via outlines integration.
  • Add a Triton or CUDA fused kernel for one operation (e.g., RMSNorm).

Common pitfalls

  • Forgetting to detach KV from the autograd graph → memory blow-up.
  • Mismatched sampling RNG across requests in a batch → bad determinism.
  • Holding the GIL in CPU-bound Python sampling → throughput cap. Move to a separate process or use Triton/Numba for sampling.
  • Streaming protocol bugs (SSE headers, Content-Type: text/event-stream, data: prefix, [DONE] terminator).
  • Off-by-one between prompt length and KV-cache write index.

Story / portfolio

  • Title: "Building vLLM (a worse one) to learn how vLLM works."
  • Diagram: scheduler + KV manager.
  • Benchmarks: tokens/s vs batch size; latency CDF; KV memory vs context length.
  • Honest comparison: "vLLM is 1.7× faster because they fuse qkv_proj + RoPE + attention and use PagedAttention; here is what I would build next to close that gap."
  • Demo: curl a streaming chat completion + a tiny web UI.

This is the systems project. With Projects 2 and 5 in your portfolio, you can credibly apply for "ML systems" or "inference platform" roles — among the highest-paying tracks in the industry.