Chapter 97 — LLM Inference & Serving Optimisation

📖 Definition

Inference is running a trained model to generate output. For LLMs it has two phases:

Most serving optimisations either make decode use the GPU better (batching), shrink bytes moved (quantisation, GQA), or skip work (caching, speculative decoding).

📊 Metrics You Must Know

The core trade-off: bigger batches → higher throughput (cheaper) but slower per-user TPOT and TTFT.

🔍 GPU Memory Math (Whiteboard Favourite)

Weights   = parameters × bytes per parameter
            7B in bf16 (2 bytes)  = 14 GB
            70B in bf16           = 140 GB  → needs ≥ 2 × 80 GB GPUs
            70B in INT4 (0.5 byte) ≈ 35 GB  (+ some overhead)

KV cache per token = 2 (K and V) × layers × kv_heads × head_dim × bytes
            Llama-2-7B (32 layers, 32 KV heads, head_dim 128, fp16):
            2 × 32 × 32 × 128 × 2 bytes = 524,288 bytes ≈ 0.5 MB per token
            4,096-token sequence ≈ 2 GB   → 30 concurrent users ≈ 60 GB of KV cache!

            Llama-3-8B uses GQA with 8 KV heads → 4× smaller: ≈ 0.125 MB per token

This is why KV-cache memory decides how many users you can batch, and why GQA, KV-cache quantisation and paging matter.

🔍 KV Cache

Without caching, generating token t would recompute keys and values for all previous tokens — O(T²) work per sequence. The KV cache stores each layer's K and V for past tokens, so each decode step only computes K/V for the new token and attends to the cache.

💻 Code Example — Decode Loop with KV Cache and Top-p Sampling

import torch

@torch.no_grad()
def generate(model, tokenizer, prompt, max_new_tokens=200, temperature=0.7, top_p=0.9):
    ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    out = model(ids, use_cache=True)                       # PREFILL: whole prompt at once
    past = out.past_key_values
    generated = []

    for _ in range(max_new_tokens):                        # DECODE: one token per step
        logits = out.logits[:, -1, :] / temperature
        probs = torch.softmax(logits, dim=-1)

        # top-p (nucleus): keep the smallest set of tokens whose probs sum to top_p
        sorted_p, sorted_idx = probs.sort(descending=True)
        keep = sorted_p.cumsum(-1) - sorted_p < top_p
        sorted_p = sorted_p * keep
        next_id = sorted_idx.gather(-1, torch.multinomial(sorted_p / sorted_p.sum(), 1))

        if next_id.item() == tokenizer.eos_token_id:
            break
        generated.append(next_id.item())
        out = model(next_id, past_key_values=past, use_cache=True)   # only the new token
        past = out.past_key_values

    return tokenizer.decode(generated)

🔍 Batching Strategies

🔍 PagedAttention (vLLM)

Naively each request reserves KV-cache memory for its maximum length → 60–80% wasted by fragmentation. PagedAttention stores the KV cache in fixed-size blocks (like OS virtual memory pages), allocated on demand and shareable between requests with the same prefix. Result: near-zero waste → far larger batches → multiple-times higher throughput.

🔍 Prefix (Prompt) Caching

If many requests share a prefix (system prompt, tool definitions, a long document), keep its KV cache and reuse it → much lower TTFT and cost. Design prompts with static content first, variable content last. Hosted APIs expose this as prompt caching (see AI / LLM Engineering: Cost & Latency chapter).

🔍 FlashAttention

Standard attention writes the full T × T score matrix to slow GPU memory (HBM). FlashAttention computes attention in tiles inside fast on-chip SRAM, using an online softmax, never materialising the full matrix.

• Exact (not an approximation).
• Memory O(T) instead of O(T²); 2–4× faster for long sequences.
• Used for training and prefill; FlashDecoding variants speed up decode with long contexts.

🔍 Quantisation

Store weights (and sometimes activations / KV cache) in fewer bits.

Why weight-only 4-bit speeds up decode: decode is memory-bandwidth-bound, so moving 4× fewer bytes ≈ faster tokens, even though maths is done after dequantising.

Always evaluate after quantising — reasoning, code and long-context tasks degrade first.

🔍 Speculative Decoding

1. A small, fast draft model proposes k tokens (e.g. 4).
2. The big target model checks all k in one forward pass (like prefill).
3. Accept the longest prefix that matches the target's distribution (rejection sampling); the target supplies the next token.
4. Output distribution is identical to the target model alone — it's a pure speed-up (often 2–3×) when the draft is usually right.

Variants: Medusa (extra prediction heads instead of a draft model), EAGLE, lookahead decoding, n-gram/prompt-lookup drafting (great for editing/code where output copies input).

🔍 Parallelism for Serving

📊 Serving Engines

🔍 Other Levers

• Smaller model / distillation / routing: send easy requests to a small model.
• Shorter outputs: ask for concise answers, set max_tokens, stop sequences — output tokens are the slowest part.
• Structured / constrained decoding prevents retries from invalid JSON.
• Streaming improves perceived latency (first token fast).
• Autoscaling on queue depth / KV-cache usage, not CPU.

🧠 What a Senior Answer Adds

• Does the memory maths (weights + KV cache) before choosing GPUs.
• Separates prefill vs decode bottlenecks and picks techniques accordingly.
• Frames everything as latency SLO vs cost per token, and measures goodput.

🎯 Likely Interview Questions

1. Why is decode slower per token than prefill? — Decode is memory-bandwidth-bound: each step reads all weights + KV cache for one token per sequence.
2. What is the KV cache and how big is it? — Stored K/V for past tokens; 2 × layers × kv_heads × head_dim × bytes per token.
3. Explain continuous batching. — Scheduling at every decode step so finished requests leave and new ones join immediately.
4. What problem does PagedAttention solve? — KV-cache fragmentation/over-reservation; enables bigger batches and prefix sharing.
5. How does FlashAttention speed things up? — IO-aware tiled exact attention in SRAM; avoids writing the T×T matrix to HBM.
6. Explain speculative decoding and why output quality is unchanged. — Draft proposes, target verifies in parallel with rejection sampling preserving the target distribution.
7. GPTQ vs AWQ vs INT8? — 4-bit weight-only via second-order error minimisation vs activation-aware salient weight protection vs 8-bit weights+activations.
8. Why does 4-bit weight quantisation speed up decoding? — Fewer bytes to read per step in a bandwidth-bound phase.
9. TTFT vs TPOT — what affects each? — TTFT: queueing, prompt length, prefix cache; TPOT: model size, batch size, bandwidth, quantisation.
10. How many GPUs to serve a 70B model? — bf16 weights 140 GB → at least 2 × 80 GB plus KV cache headroom; or INT4/FP8 on fewer GPUs.
11. How would you cut serving cost by 50%? — Quantisation, better batching, prefix caching, routing to smaller models, shorter outputs, distillation, autoscaling.
12. What is disaggregated serving? — Separate prefill and decode onto different GPU pools to optimise each and meet SLOs.

← Fine-Tuning & LoRA | Index | Next: Embeddings, Multimodal & Diffusion →