Inference Optimization

Decode is memory-bound: continuous batching is the first lever

Background & principle

Inference has two phases with completely different bottlenecks. Prefill (processing the prompt) computes all tokens in parallel—compute-bound. Decode (generating one token at a time) computes just 1 token per step yet must re-stream the entire model weights from HBM—arithmetic intensity is tiny, the Tensor Cores sit mostly idle, and the real bottleneck is memory bandwidth. That's why at batch=1 even the priciest card stays underutilized.

The fix is to pack the same decode step from many requests into one batch: weights are streamed once, a token is computed for N requests simultaneously, and throughput scales near-linearly. But traditional static batching has a fatal flaw—it can't release the batch until the longest sequence in it finishes, so short requests are dragged down by long ones and GPU utilization collapses as generation proceeds. Orca (OSDI'22) introduced iteration-level scheduling (i.e. continuous batching): after every decode step, it evicts finished requests and admits queued ones, so the GPU never idles. Anyscale measured 23× throughput versus a naive HF implementation. This is the single highest-payoff change.

Hands-on

# vLLM does continuous batching by default — serve in one line
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --max-num-seqs 256          # concurrent request cap = how big batch can grow

# Diagnose whether you're memory-bound / under-batched:
# while load-testing, watch nvidia-smi — if GPU-Util is high but SM
# compute is low, mem bandwidth is maxed, and throughput rises
# near-linearly with batch → you're memory-bound, batching pays off
nvidia-smi dmon -s u        # compare sm% vs mem%

Remember the division of labor: throughput comes from bigger batches, latency from smaller batches—a zero-sum pair. Decide which one you're optimizing first.

KV cache is the real memory hog: PagedAttention + prefix caching

Background & principle

Continuous batching wants larger batches, but the bigger the batch and the longer the sequence, the more memory the KV cache eats—at long context it can exceed the weights themselves. A size estimate (standard formula, not a cited figure):

KV_bytes ≈ 2 × n_layers × n_kv_heads × head_dim
           × seq_len × batch × dtype_bytes
# the "2" = K and V; one 8K-context request can reach hundreds of MB

The traditional approach pre-reserves contiguous memory for each request at "max length," but actual lengths vary, so 60-80% of memory is wasted on fragmentation and over-reservation—squeezing your achievable batch. PagedAttention (vLLM's core, SOSP'23) borrows the OS's virtual memory: it splits the KV cache into fixed-size blocks (pages), stored non-contiguously and allocated on demand, cutting waste below 4%—effectively doubling the batch and yielding 2-4× throughput.

The second trick is prefix caching: if several requests share an identical prefix (system prompt, few-shot examples, a long document), that prefix's KV is computed once and reused across requests. This is exactly how Anthropic prompt caching works server-side—put stable content first and volatile content last, and a cache hit skips the redundant prefill.

Hands-on

# vLLM: enable prefix caching + quantize KV cache (save half the memory)
vllm serve  \
  --enable-prefix-caching \    # KV of identical prefixes is reused
  --kv-cache-dtype fp8         # fp8 KV halves memory, lifts throughput

# Prompt layout: stable prefix first, volatile last — required for a hit
[system + tool defs + few-shot]  ← byte-identical across reqs → cache hit
[this turn's user input]         ← differs each time → prefill only this bit

Speculative decoding: a small model runs ahead, doubling throughput with identical output

Background & principle

The key insight follows from §1: during decode the GPU's compute sits idle, so why not compute the probabilities of several tokens at once. Speculative decoding (Leviathan et al., ICML'23) lets a cheap draft model quickly guess k candidate tokens, then has the big model verify all k in a single parallel forward—accepting the longest correct prefix and resampling at the first disagreement. When the draft guesses well, one step advances several tokens; when it's wrong, you simply fall back to normal speed.

The most counterintuitive part: via rejection sampling, the final output distribution is provably identical to the original model—this is lossless acceleration, not an approximation. The speedup depends on the acceptance rate (how often the draft is right) and the draft/target cost ratio. Medusa (arXiv 2401.10774) goes further: instead of a separate draft model, it adds a few lightweight decoding heads + tree attention so the model speculates on itself, removing the burden of maintaining two models—measured at 2.3-3.6×.

Hands-on

# vLLM: use a small model as draft (must share tokenizer, ~10-20× smaller)
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --speculative-model meta-llama/Llama-3.2-1B-Instruct \
  --num-speculative-tokens 5     # guess 5; too many backfires

# No draft model handy? Use n-gram speculation (zero extra model)
  --speculative-model "[ngram]" --ngram-prompt-lookup-max 4

Choosing an engine & quantization: the vLLM / TGI / TensorRT-LLM trade-offs

Background & principle

All three major open engines implement the core optimizations of §1-§3; they differ in the triangle of peak performance vs deployment flexibility vs onboarding cost:

• vLLM: native PagedAttention, fastest iteration, Python-friendly, lowest onboarding. The default starting point for most self-hosting.
• TGI (HuggingFace): production-ready, seamless with the HF ecosystem/Endpoints, multi-backend (incl. TRT-LLM). Pick it for stable ops.
• TensorRT-LLM (NVIDIA): peak performance on NVIDIA cards, but requires pre-compiling an engine per GPU architecture—heaviest ops, least flexible. Worth it only when you chase max throughput-per-card and your model/hardware are stable.

The second-biggest lever is quantization: compress weights from fp16 to int4/fp8. Because decode is memory-bound (§1), smaller weights = faster streaming = directly faster decode, while freeing memory to enlarge the batch. weight-only int4 (AWQ / GPTQ) halves memory and adds speed; fp8 (weights + activations) needs Hopper or newer and loses less precision.

Hands-on

# Quick decision sheet
fast validation / self-host / rapid iteration  → vLLM
HF ecosystem / stable production ops            → TGI
squeeze NVIDIA cards / stable model             → TensorRT-LLM

# Quantization picks
memory-tight, must fit a big model              → AWQ int4 (weight-only)
Hopper+, want precision and speed               → fp8 (weights + KV cache)
vllm serve  --quantization awq   # one flag