LLM Serving & Inference

Background

Serving latency and throughput are distinct from training concerns - a model that trains in hours must respond in milliseconds
GPU batch inference, model quantization, and distillation emerged as the three primary production optimization strategies
TensorRT, ONNX Runtime, and vLLM are the dominant high-performance inference engines

The Concept

ML inference serving takes a trained model and makes it available for production prediction requests at scale. Key concerns: latency (how fast per request), throughput (requests per second), cost (GPU utilization), and reliability (uptime, graceful degradation). Optimization techniques include quantization, distillation, batching, and caching.

How it Links

Quantization (INT8, FP16, 4-bit) connects to numerical methods and the precision-performance trade-off in deep learning
Model distillation (compressing a large model into a smaller one) connects to knowledge transfer and ensemble methods
vLLM paged attention connects to operating systems memory management (virtual memory, paging)

Concrete Example

Serving GPT-4 at ChatGPT scale: a naive implementation would process one request at a time. vLLM uses continuous batching to process hundreds of requests in parallel, paged KV cache to manage GPU memory like a virtual memory system, and quantization to reduce each parameter from 32-bit to 4-bit float - achieving 10x the throughput at 1/4 the memory cost.

Prerequisites

The Transformer Architecture, Fine-tuning and PEFT, Attention Mechanisms

Learning outcomes

Derive the KV-cache memory formula and compute the exact memory cost for a given LLM at a specific batch size and sequence length.
Explain PagedAttention and how vLLM uses OS-style virtual memory paging to eliminate KV-cache fragmentation.
Derive the expected speedup from speculative decoding as a function of the draft acceptance rate.
Compare GPTQ and AWQ as post-training quantization methods and state which is better for weight-only vs activation quantization.
Identify the throughput-latency trade-off and explain when to prefer continuous batching over static batching.

TL;DR

LLM inference is memory-bandwidth-bound, not compute-bound. The dominant cost at generation time is loading the KV-cache: for each new token, every layer must read its keys and values for all previous tokens. PagedAttention [4] treats the KV-cache like virtual memory, eliminating fragmentation and allowing higher batch utilisation. Speculative decoding [2] uses a small draft model to propose $\gamma$ tokens at once, which the target model verifies in a single forward pass, yielding 2-3x throughput gains when the draft acceptance rate is high. Post-training quantization (GPTQ [5], AWQ [6]) reduces weight size from 16-bit to 4-bit with under 1% accuracy loss on most benchmarks.

Intuition

A standard transformer forward pass at training time processes an entire sequence of $n$ tokens in parallel. At inference time, generation is autoregressive: tokens are produced one at a time, and generating token $t$ requires attending over all previous tokens $1, \dots, t-1$ .

Recomputing the keys and values for all previous tokens at every step would be $O(t^2)$ per sequence. The KV-cache avoids this by storing the keys and values from every previous step. Each new token only needs to compute its own keys and values and append them; the attention over all previous tokens reads from cache. This reduces the per-token cost from $O(t^2)$ to $O(t)$ , but the cache itself consumes memory proportional to $t$ , batch size, and model depth.

The two regimes of LLM inference

LLM serving has two distinct phases with different bottlenecks.

Prefill: the prompt tokens are processed in parallel (a standard forward pass). This phase is compute-bound: a long prompt can use the full GPU FLOP capacity.

Decode: tokens are generated one at a time. This phase is memory-bandwidth-bound: each decode step reads the full KV-cache from HBM. A 13 B parameter model generating a 2048-token sequence has its decode step dominated by the ~2 GB KV-cache read, not by the ~26 GB weight read. Optimising throughput means keeping the GPU’s HBM bandwidth fully saturated with large batches.

KV-cache memory math

Deriving the KV-cache memory formula

For a transformer with $L$ layers, $H$ attention heads of dimension $d_h$ each (so head dimension $d_h = d_{\text{model}} / H$ ), processing a sequence of $s$ tokens at batch size $b$ , the KV-cache stores:

$K$ tensor per layer: shape $(b, H, s, d_h)$
$V$ tensor per layer: shape $(b, H, s, d_h)$

Two tensors (K and V) per layer, $L$ layers total. In bytes (assuming BF16 = 2 bytes per element):

\text{KV bytes} = 2 \times L \times 2 \times H \times d_h \times s \times b \times 2 = 8 \times L \times d_{\text{model}} \times s \times b \text{ bytes},

where the first factor of 2 is for K and V, the second 2 is for the BF16 element size, and $H \times d_h = d_{\text{model}}$ .

Worked example: Llama-2-7B. $L = 32$ layers, $d_{\text{model}} = 4096$ , $b = 1$ , $s = 4096$ :

\text{KV bytes} = 8 \times 32 \times 4096 \times 4096 \times 1 = 4{,}294{,}967{,}296 \approx 4 \text{ GB}.

Compare to the model weights: 7 B parameters $\times$ 2 bytes/param $\approx 14$ GB. The KV-cache is roughly 30% of weight size at this sequence length and doubles with every doubling of sequence length. At $s = 8192$ the KV-cache reaches 8 GB, at $s = 32768$ it reaches 32 GB---larger than the model itself.

For a batch of 8 concurrent requests each at sequence length 2048, the KV-cache is 8 GB, and the total GPU memory requirement with weights is 22 GB, saturating a 24 GB A10 with no headroom for the decode logits and activations. This is why KV-cache management is the central engineering challenge in LLM serving.

Grouped-query attention (GQA). In GQA (used by Llama-2 70B, Mistral, and others), $H_K = H_V < H$ keys and values are shared across groups of query heads. If $G$ is the number of query heads per KV head, the KV-cache shrinks by $G\times$ :

\text{KV bytes}^{\text{GQA}} = \frac{8 \times L \times d_{\text{model}} \times s \times b}{G}.

Llama-2-70B uses $H = 64$ query heads and $H_K = 8$ KV heads ( $G = 8$ ), reducing its KV-cache by $8\times$ relative to multi-head attention at the same $d_{\text{model}}$ .

PagedAttention and vLLM

In a naive serving system the KV-cache is allocated as a contiguous block for each request at the maximum possible sequence length. This causes two problems: (1) internal fragmentation - most slots in the pre-allocated block are unused for short-context requests; (2) no sharing - if two requests share a common prefix (e.g., a system prompt), the KV-cache is duplicated for each.

Kwon et al. [4] introduced PagedAttention, which manages KV-cache memory using the same virtual-memory paging abstraction that operating systems use for RAM.

The KV-cache is divided into fixed-size blocks (pages) of $B$ tokens each. Each request maintains a block table mapping logical block indices to physical block slots in GPU HBM. The attention kernel is modified to follow these indirections at runtime. Physical blocks are allocated on demand and freed when the request completes. Blocks can be shared across requests (copy-on-write for beam search) and between requests with a common prefix (prompt caching).

The result: vLLM achieves 2-4x higher throughput than Hugging Face generate() on identical hardware by packing more concurrent requests into the same KV-cache memory budget.

Continuous batching adds to this: rather than processing fixed-size batches of $B$ requests from start to finish, vLLM uses an iteration-level scheduler that inserts new requests into the batch as others finish. Each decode step the scheduler adds new requests up to the available KV-cache budget. This eliminates the “last request finishes 10x faster than the first but the batch is held back” problem of static batching.

Speculative decoding

Speculative decoding speedup derivation

Speculative decoding [2] uses a small draft model $M_q$ (e.g., a distilled 7B version of a 70B target) to produce $\gamma$ candidate tokens $\tilde{x}_{t+1}, \dots, \tilde{x}_{t+\gamma}$ in $\gamma$ serial steps. The large target model $M_p$ then verifies all $\gamma$ tokens in a single batched forward pass by computing $p(x | x_{1:t}), p(x | x_{1:t+1}), \dots, p(x | x_{1:t+\gamma})$ in parallel.

The acceptance criterion for token $i$ is:

\text{accept } \tilde{x}_{t+i} \quad \text{if} \quad u_i \le \frac{p(\tilde{x}_{t+i} | x_{1:t+i-1})}{q(\tilde{x}_{t+i} | x_{1:t+i-1})}, \quad u_i \sim \text{Uniform}(0, 1).

If token $i$ is rejected, the algorithm samples a correction token from $\max(0, p - q)$ (renormalised) and discards tokens $i+1, \dots, \gamma$ . This rejection-sampling scheme guarantees that the output distribution of speculative decoding is identical to the target model’s distribution.

Expected accepted tokens per verification round. Let $\alpha$ be the per-token acceptance probability averaged over positions. In the simplest analysis (i.i.d. accepts), the probability of accepting all $k$ tokens is $\alpha^k$ , so the expected number of accepted tokens is

E[\text{accepted}] = \sum_{k=1}^{\gamma} \alpha^k = \frac{\alpha(1 - \alpha^{\gamma})}{1 - \alpha}.

Including the one guaranteed new token from the correction step, the expected tokens per round is $E[\text{accepted}] + 1$ .

Expected wallclock speedup. Let $T_q$ = time for one draft step, $T_p$ = time for one target step. A standard round takes $T_p$ for one token. A speculative round takes $\gamma T_q + T_p$ but produces $E[\text{accepted}] + 1$ tokens:

\text{speedup} = \frac{(E[\text{accepted}] + 1) \times T_p}{\gamma T_q + T_p}.

For $\alpha = 0.8$ , $\gamma = 5$ , $T_p / T_q = 10$ (a 10x larger target model), $E[\text{accepted}] \approx 3.3$ :

\text{speedup} = \frac{4.3 \times T_p}{0.5 T_p + T_p} = \frac{4.3}{1.5} \approx 2.9\times.

Empirically, speedups of 2-3x on chat-model benchmarks are typical. Speedup collapses when the draft and target models diverge (low $\alpha$ ) or when the draft model is slow relative to the target.

Decoding strategies

Strategy	Deterministic?	Diversity	Quality	Best use case	Key parameter
Greedy	Yes	None (always picks argmax)	Good for factual tasks	Code generation, translation	None
Beam search	Yes (given beam width)	Low	Best BLEU for seq2seq	Machine translation, summarisation	Beam width k
Top-k sampling	No	Medium	Good; avoids very rare tokens	Creative text, chat	k (typical 50)
Top-p (nucleus)	No	Adaptive	Best for open-ended generation	Chat, creative writing	p (typical 0.9)
Temperature scaling	No	Controlled	Trades coherence vs diversity	Always combined with sampling	temp (typical 0.7-1.0)
Speculative decoding	Matches target distribution	Same as target sampling	Identical to target model	High-throughput serving	Draft model, gamma

Nucleus (top-p) sampling [1] selects from the smallest set of tokens whose cumulative probability exceeds $p$ . At each step the set size varies: on easy-to-predict tokens the nucleus may have just 1-3 tokens; on ambiguous tokens it may have hundreds. This adaptive vocabulary is what makes nucleus sampling superior to top-k for open-ended generation.

Quantization for inference

Post-training quantization (PTQ) reduces model weights from FP16/BF16 to INT8 or INT4, shrinking both memory footprint and HBM bandwidth requirements.

GPTQ [5] applies layer-wise quantization with a second-order correction. For each weight matrix $W$ , GPTQ minimises:

\| WX - \hat{W}X \|_F^2,

where $X$ is a small calibration set of activations and $\hat{W}$ is the quantized weight matrix. The correction uses the inverse Hessian $H = XX^\top$ (via the Optimal Brain Surgeon framework) to update the remaining unquantized weights after each column is quantized, absorbing the quantization error. GPTQ at INT4 achieves under 1% perplexity loss on LLaMA and OPT families.

AWQ [6] takes a different approach: instead of correcting weights post-hoc, it identifies salient channels (those with large activations that are hardest to quantize accurately) and applies a per-channel scaling factor to them before quantization, then absorbs the scale into the next layer. This activation-aware scaling preserves accuracy for the 1% of channels that matter most, while the other 99% are quantized aggressively.

Method	Precision	Approach	Speed (vs FP16)	Accuracy loss	Best for
FP16 / BF16	16-bit	Baseline	1x	None	Training, small batch inference
INT8 (W8A8)	8-bit weights + activations	Round-to-nearest per tensor	1.3-1.6x	Small for most models	Balanced throughput
GPTQ (W4A16)	4-bit weights, 16-bit activations	Second-order weight correction	1.8-2.2x	Under 1% perplexity	Weight-only quant, CPU offload
AWQ (W4A16)	4-bit weights, activation-aware	Per-channel scale before quant	1.8-2.2x	Slightly better than GPTQ	Activation-sensitive models
GGUF (llama.cpp)	Mixed 2-8 bit per tensor	k-quants with importance scoring	CPU-optimised	Varies by quant level	Consumer hardware, edge

Serving runtimes

Runtime	PagedAttention	Continuous batching	Multi-LoRA	Speculative decoding	License
vLLM	Yes (native)	Yes	Yes (S-LoRA)	Yes	Apache 2.0
TGI (HuggingFace)	Yes	Yes	Limited	Yes	Apache 2.0 / HFOIL
TensorRT-LLM (NVIDIA)	Yes	Yes	Yes	Yes	Apache 2.0
SGLang	Yes	Yes	Yes	Yes (RadixAttention)	Apache 2.0
llama.cpp	No (KV slots)	Partial	No	No	MIT

Throughput vs latency

The fundamental trade-off in LLM serving: throughput (tokens generated per second across all users) and latency (time to first token, time between tokens for one user) pull in opposite directions. Larger batches amortise the weight load across more tokens per step (throughput up) but increase the queue time for each individual request (latency up).

Tail latency (P99 time to first token) is the metric that determines whether interactive applications feel responsive. For a chat interface, users tolerate about 1 second TTFT and 20-40 tokens/second per-user streaming throughput. For batch inference (document processing, embedding generation), throughput dominates and latency is irrelevant.

The operating point is chosen by the serving system’s scheduler. vLLM’s scheduler uses a priority queue with preemption: long-running requests can be paused (their KV-cache swapped to CPU) to let shorter requests complete, reducing P99 TTFT at the cost of some throughput.

Pitfalls

KV-cache invalidation on prompt changes

If the prompt is modified between requests (e.g., the system prompt changes per user, or chain-of-thought steps are appended), the KV-cache prefix is invalidated and the prefill must be recomputed from the divergence point. For systems with long shared system prompts, prompt caching (storing the KV-cache for the fixed prefix) is a major latency win. Always keep the variable part of the prompt at the end.

Beam search hurts open-ended generation

Beam search maximises sequence likelihood, which biases toward short, high-confidence tokens. On open-ended tasks (creative writing, chat) the result is repetitive, generic text. Nucleus sampling is strongly preferred for open-ended generation, and beam search should be used only for structured prediction tasks (translation, summarisation, constrained code generation).

Profile memory before choosing quantization

Start by profiling: how much of your decode latency is weight loading vs KV-cache loading? On long sequences with large batches, the KV-cache dominates and weight quantization alone gives diminishing returns. On short-context high-concurrency workloads, weight quantization is the right lever. Use nvtx or a vendor profiler to distinguish the two before committing to a quantization strategy.

Pop quiz

Check yourself

Score: 0 / 5(0 answered)

1. For Llama-2-7B (32 layers, 4096 hidden, BF16) at batch size 1 and sequence length 2048, the KV-cache size is approximately:
2. PagedAttention reduces KV-cache memory waste primarily by:
3. The key correctness guarantee of speculative decoding is:
4. Top-p (nucleus) sampling is preferred over top-k for open-ended generation because:
5. AWQ (Activation-aware Weight Quantization) differs from GPTQ mainly by:

Interview Q&A

Derive the KV-cache memory formula and compute it for Llama-2-13B at batch 4, sequence 4096.onsite

For $L$ layers, $d_{\text{model}}$ hidden dimension, batch $b$ , sequence $s$ , BF16 (2 bytes/element):

\text{KV bytes} = 2 \times L \times 2 \times d_{\text{model}} \times s \times b = 8 L d_{\text{model}} s b.

For Llama-2-13B: $L = 40$ , $d_{\text{model}} = 5120$ , $b = 4$ , $s = 4096$ :

8 \times 40 \times 5120 \times 4096 \times 4 = 26{,}843{,}545{,}600 \approx 25 \text{ GB}.

Llama-2-13B weights are $13 \times 10^9 \times 2 \approx 26$ GB. So at batch 4 and 4K sequence length the KV-cache equals the weight size. This is why serving 13B at scale requires either GQA, quantized KV-cache, or an A100 with 80 GB HBM.

Note that Llama-2-13B uses grouped-query attention with $G = 5$ (40 query heads, 8 KV heads), so the actual KV-cache is $5\times$ smaller: $\approx 5$ GB.

Derive the expected speedup from speculative decoding. Under what conditions does it fail to give a speedup?staff

Let $\alpha$ = per-token acceptance probability, $\gamma$ = number of draft tokens, $T_p$ = target model latency per step, $T_q$ = draft model latency per step.

Expected accepted tokens:

E[k] = \sum_{i=1}^{\gamma} \alpha^i = \frac{\alpha(1-\alpha^\gamma)}{1-\alpha}.

Expected tokens per speculative round = $E[k] + 1$ (the correction token is always emitted).

Expected wallclock per token (normal): $T_p$ .

Expected wallclock per token (speculative): $\frac{\gamma T_q + T_p}{E[k] + 1}$ .

Speedup: $\frac{(E[k]+1) T_p}{\gamma T_q + T_p}$ .

Failure conditions: (1) Low $\alpha$ : if the draft and target disagree often, few tokens are accepted and the overhead of running the draft model outweighs the savings. Typical minimum useful $\alpha \approx 0.7$ . (2) Slow draft model ( $T_q \approx T_p$ ): if the draft model is nearly as expensive as the target, the $\gamma T_q$ term dominates. (3) Long sequences with large batches: the target model becomes more throughput-efficient at high batch sizes, reducing the relative cost of its decode step. Speculative decoding primarily helps at small batch sizes (latency-optimised serving).

References

Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. (2019). The Curious Case of Neural Text Degeneration. ICLR 2020. arXiv:1904.09751
Leviathan, Y., Kalman, M., and Matias, Y. (2022). Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192
Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint. arXiv:2302.01318
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. arXiv:2309.06180
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ICLR 2023. arXiv:2210.17323
Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. (2023). AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration. MLSys 2024. arXiv:2306.00978
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, U., and Dean, J. (2022). Efficiently Scaling Transformer Inference. MLSys 2023. arXiv:2211.05102
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762

Capstone: RAG System MLOps