LLM Inference Server
Build a high-throughput, quantized LLM inference engine and streaming HTTP server with KV-cache management, INT8/INT4 quantization, and real-time latency telemetry
Production Inference Security & Reliability Requirements
Your inference server must implement production security guardrails: strict request schema validation (Pydantic), maximum prompt character/token length enforcement, maximum output token caps (`max_new_tokens <= 1024`), per-client rate limiting, request timeout handling, and structured JSON error responses that never leak internal stack traces.
What You Will Build
Serving a Large Language Model in production is fundamentally different from training it: prefill is compute-bound, while autoregressive token decoding is strictly memory-bandwidth-bound by loading model weights and KV caches from GPU HBM at every step. In this project, you will build a high-performance LLM Inference Engine and streaming HTTP server. You will separate prefill and decode phases with pre-allocated KV-cache buffers and system-prompt prefix caching, implement absmax INT8 and group-wise INT4 weight-only quantization from scratch, implement speculative decoding with a draft model for lossless 2× speedups, expose a streaming Server-Sent Events (SSE) HTTP API with strict input validation, rate limiting, and timeout handling, and benchmark Time-To-First-Token (TTFT), Time-Per-Output-Token (TPOT), and concurrent throughput.
Why This Project Matters in the AI/LLM Stack
Understanding why naive `model.generate()` recalculates $O(T^2)$ attention on every token—and how KV caching, INT8/INT4 quantization, dynamic batching, and speculative decoding cut latency and memory by 3×–8×—is essential for AI systems and inference engineers.
Learning Objectives
- Implement a stateful KV-Cache manager supporting prefill, single-token decode, and shared system-prompt prefix reuse.
- Implement symmetric INT8 per-channel and group-wise INT4 weight quantization and measure VRAM reduction vs. perplexity delta.
- Implement speculative decoding using a small draft model and verify mathematically lossless target distribution sampling.
- Build a streaming FastAPI Server-Sent Events (SSE) endpoint hardened with input validation, rate limiting, token caps, and timeouts.
- Benchmark TTFT, TPOT, peak VRAM, and tokens/sec across batch sizes and quantization modes.
Prerequisites
- Phase 5 Project: Mini GPT Architecture (`milestone-10-mini-gpt-architecture`)
- Phase 6 Project: Pretrained Mini LLM (`milestone-11-pretrained-mini-llm`)
- Phase 8 Modules 8.1 & 8.4: GPU Memory Hierarchy and Inference & Deployment
Required Mathematical Concepts
- KV-Cache memory footprint: 2 * n_layers * n_kv_heads * d_head * seq_len * bytes_per_elem
- Symmetric absmax INT8 quantization: s = max(|W|) / 127, W_q = round(W / s), W_hat = s * W_q
- Group-wise INT4 asymmetric quantization with group size g = 64 or 128
- Speculative decoding acceptance probability: min(1, p_target(x) / p_draft(x))
- Latency percentiles (P50, P95, P99) for TTFT and TPOT (ITL)
Step-by-Step Implementation Tasks (15 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
KV-Cache Engine, Prefill/Decode Split & Prefix Caching
Eliminate $O(T^2)$ redundant recomputation during decoding by pre-allocating KV cache tensors and caching shared system prompts.
Calling `torch.cat` on every generated token allocates new GPU tensors repeatedly; pre-allocating `(B, n_kv_heads, max_seq_len, d_head)` buffers writes new K/V slices in-place.
Prefill processes all `L_prompt` tokens in parallel (`q_len = L_prompt`), while each subsequent decode step passes only the newly sampled token `(B, 1)` with `start_pos = current_len`.
KV cache grows linearly with batch size and sequence length; pre-computing exact byte footprint prevents out-of-memory (OOM) crashes.
When thousands of requests share the same system prompt, pre-computing its KV cache once and cloning/reusing that prefix cuts Time-To-First-Token (TTFT).
INT8 & Group-Wise INT4 Weight Quantization
Quantize linear layer weights to 8-bit and 4-bit integers and measure memory footprint savings and perplexity preservation.
Storing linear weights as `int8` with a per-output-channel FP16/FP32 scale $s_i = \max_j |W_{i,j}| / 127$ cuts model weight memory by 2×–4×.
Packing two 4-bit integers (`[-8, 7]` or `[0, 15]`) into a single `uint8` byte with per-group scales and zero-points compresses weights by 4× over FP16 while isolating outlier channels.
Quantization is only useful if model quality holds up; measuring validation perplexity across all three precisions quantifies the accuracy-memory Pareto frontier.
Batched Scheduling & Speculative Decoding
Support continuous/dynamic multi-request batching and speculative draft-then-verify decoding.
A tiny draft model proposes $K=4$ candidate tokens cheaply; the target model verifies all $K$ tokens in a single parallel forward pass using rejection sampling.
Batching multiple concurrent user prompts together amortizes weight loading across requests.
Streaming tokens as an async generator allows HTTP clients to render words immediately upon first-token completion.
Hardened Streaming HTTP Server & Latency Benchmarks
Expose a streaming SSE API with full security guardrails and run an automated load benchmark suite.
Exposing `POST /v1/generate` (and `/v1/generate/stream`) via SSE (`text/event-stream`) matches industry-standard LLM serving APIs.
Unbounded prompts or `max_new_tokens=1000000` can trigger instant GPU OOM or denial of service.
Production inference servers must protect against request floods, cancel stalled generations, and return safe JSON error envelopes.
Measuring P50/P95/P99 Time-To-First-Token (TTFT), Time-Per-Output-Token (TPOT), and aggregate tokens/sec across concurrency levels (`c = 1, 4, 8, 16`) reveals the server's saturation point.
Synthesize your measurements comparing uncached vs. KV-cached decoding, FP32/BF16 vs. INT8 vs. INT4 memory/perplexity, speculative decoding acceptance rate, and concurrent server latency.
Starter & Verification Code Examples
Symmetric Per-Channel INT8 Linear Quantization & KV-Cache Buffer
Core implementation of weight-only INT8 quantization and a pre-allocated KV-cache tensor buffer.
import torch
import torch.nn as nn
import torch.nn.functional as F
class Int8Linear(nn.Module):
def __init__(self, linear: nn.Linear):
super().__init__()
w = linear.weight.data.float()
scale = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-8) / 127.0
self.register_buffer("weight_int8", torch.round(w / scale).to(torch.int8))
self.register_buffer("scale", scale.to(linear.weight.dtype))
self.bias = linear.bias
def forward(self, x: torch.Tensor) -> torch.Tensor:
w_dequant = self.weight_int8.to(x.dtype) * self.scale
return F.linear(x, w_dequant, self.bias)
class LayerKVCache:
def __init__(self, batch_size: int, n_heads: int, max_seq_len: int, head_dim: int, dtype: torch.dtype, device: torch.device):
self.k = torch.zeros((batch_size, n_heads, max_seq_len, head_dim), dtype=dtype, device=device)
self.v = torch.zeros((batch_size, n_heads, max_seq_len, head_dim), dtype=dtype, device=device)
self.seq_len = 0
def update(self, k_new: torch.Tensor, v_new: torch.Tensor):
t = k_new.size(2)
self.k[:, :, self.seq_len : self.seq_len + t, :] = k_new
self.v[:, :, self.seq_len : self.seq_len + t, :] = v_new
self.seq_len += t
return self.k[:, :, :self.seq_len, :], self.v[:, :, :self.seq_len, :]How to Test Your Implementation
Run unit and integration tests covering KV-cache equivalence, INT8/INT4 quantization accuracy, speculative decoding correctness, and HTTP security guardrails.
KV-Cache Logit Equivalence Test
Confirm prefill + step-by-step KV-cached decoding matches full-sequence forward passes.
Steps: Generate 32 tokens greedily with `use_kv_cache=True` and `use_kv_cache=False` from the same prompt.
Pass Criteria: Generated token ID sequences and step logits match.
Speculative Decoding Greedy Exactness Test
Verify that speculative decoding never alters greedy target model outputs.
Steps: Run `speculative_generate(..., temperature=0.0, K=4)` and compare against standard target greedy decoding.
Pass Criteria: Exact 100% token sequence match with `< N` target forward calls for `N` tokens.
Server Security Guardrails Test (`TestClient`)
Verify input bounds, rate limiting, and structured error envelopes.
Steps: Send requests with `max_new_tokens=99999` (expect 422), empty prompt (expect 422), and burst above rate limit (expect 429).
Pass Criteria: All security status codes and structured JSON error schemas pass.
Overall Project Acceptance Criteria
- KV-cached decoding matches uncached logits while achieving multi-fold speedup on long sequences.
- From-scratch INT8 and group-wise INT4 weight quantizers compress linear weights by ~2× and ~4× (vs FP16) with documented perplexity impact.
- Speculative decoding produces identical greedy outputs to standard decoding with fewer target model forward passes.
- FastAPI SSE server enforces input length caps, max output tokens, rate limiting (HTTP 429), and timeouts (HTTP 504) with structured error responses.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
Paged KV-Cache Block Allocator
Replace contiguous `(B, H, max_seq_len, D)` KV buffers with fixed-size 16-token physical memory blocks mapped via a logical block table (inspired by vLLM PagedAttention).
What You Learn: How production serving engines eliminate external and internal KV memory fragmentation.
INT8 KV-Cache Quantization
Quantize the cached $K$ and $V$ tensors dynamically to `int8` per token head to cut long-context KV memory by 2×.
What You Learn: Combined weight + activation/KV-cache quantization for long-context serving.
What Your Completed Repository Should Contain
A high-throughput LLM Inference Engine and hardened FastAPI SSE streaming server supporting KV caching, INT8/INT4 quantization, speculative decoding, and TTFT/TPOT load benchmarking.
Deliverable Artifacts
- `server/kv_cache.py` & `server/engine.py` with prefill/decode split and prefix caching
- `server/quantization.py` with INT8 per-channel and group-wise INT4 weight quantization
- `server/speculative.py` implementing draft-and-verify speculative decoding
- `server/app.py` FastAPI SSE streaming server with validation, rate limiting, and timeouts
- `bench_server.py` & `INFERENCE_BENCHMARK_REPORT.md` with P50/P95 TTFT, TPOT, and quantization metrics
After Completing This Project, You Should Be Able To...
- Implement pre-allocated KV-cache management and system-prompt prefix caching.
- Quantize LLM weights to INT8 and group-wise INT4 from scratch and measure perplexity vs. memory.
- Implement lossless speculative decoding to accelerate memory-bandwidth-bound token generation.
- Build, secure, and load-test a streaming HTTP/SSE LLM inference server.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: llm-inference-server-from-scratch
Recommended Directory Layout:
server/__init__.py server/kv_cache.py server/quantization.py server/speculative.py server/engine.py server/app.py bench_server.py INFERENCE_BENCHMARK_REPORT.md tests/test_inference_server.py README.md
Recommended README Documentation Checklist
- Include the KV-cached vs. uncached latency speedup chart across sequence lengths.
- Show the FP32/BF16 vs. INT8 vs. INT4 model memory and validation perplexity table.
- Include P50/P95 TTFT and TPOT results from `bench_server.py`.