LLM Inference: KV Caching, Batching and Speculative Decoding

Training an LLM happens once; inference happens billions of times. The cost, speed and energy of serving determine whether an LLM application is practical. LLM inference has unusual characteristics — autoregressive, memory-bound, with a growing cache — that have inspired a rich set of optimisations. This lecture explains the essentials.

Two phases: prefill and decode#

1. Prefill: process the whole prompt in parallel, computing keys and values for every prompt token and the first output token. Compute-bound (large matrix multiplications), similar to training's forward pass.
2. Decode: generate tokens one at a time; each step processes a single new token per sequence but must read all model weights (and the KV cache) from memory. Memory-bandwidth-bound: the GPU's arithmetic units are mostly idle.

A rough decode-speed bound for a single sequence: tokens/second ≈ memory bandwidth ÷ bytes read per token (≈ model size in bytes). A 7B model in 16-bit (14 GB) on a GPU with 1 TB/s bandwidth: at most about 70 tokens/second — which is why quantisation and batching matter so much.

Metrics#

• Time to first token (TTFT) — dominated by prefill; matters for interactive feel.
• Time per output token (TPOT) / inter-token latency — dominated by decode.
• Throughput — total tokens per second across all users.
• Cost per million tokens and energy per token.

There is a trade-off: larger batches improve throughput but can increase per-user latency.

The KV cache#

Without caching, each decode step would recompute attention keys and values for the entire prefix — quadratic work. The KV cache stores them per layer; each step computes $\mathbf{q}, \mathbf{k}, \mathbf{v}$ only for the new token and appends to the cache. The cost is memory:

For long contexts and many concurrent users, the KV cache can exceed the model weights in size. Reductions: grouped-query / multi-query attention, KV quantisation, sliding-window attention, and evicting less important tokens.

PagedAttention and memory management#

Pre-allocating a contiguous KV buffer for each request's maximum length wastes memory (fragmentation). PagedAttention (Kwon et al., 2023, the basis of vLLM) stores the KV cache in fixed-size blocks, like virtual-memory pages, allocated on demand. This nearly eliminates waste, allows many more concurrent sequences, and enables prefix sharing — requests with the same system prompt reuse the same cached blocks (prefix caching).

Continuous batching#

Static batching waits for a whole batch to finish before starting new requests; short responses idle while long ones complete. Continuous (in-flight) batching (Orca, Yu et al., 2022) schedules at the iteration level: after every decode step, finished sequences leave and new requests join. This dramatically improves GPU utilisation and throughput. Advanced schedulers also split long prefills into chunks and interleave them with decodes ("chunked prefill") to keep latency stable.

Speculative decoding#

Decode steps are memory-bound: verifying several tokens costs about the same as generating one. Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) exploits this:

1. A small, fast draft model proposes $\gamma$ tokens.
2. The large target model evaluates all of them in one parallel forward pass.
3. Each proposed token is accepted with probability $\min\left(1, \frac{p(x)}{q(x)}\right)$ (target probability over draft probability); at the first rejection, a replacement is sampled from the normalised residual distribution $\max(0, p - q)$.

This acceptance scheme guarantees the output distribution is exactly that of the target model — a lossless speed-up, often 2–3× when the draft agrees with the target frequently. Variants avoid a separate draft model: Medusa adds extra decoding heads, EAGLE drafts at the feature level, and n-gram / prompt-lookup decoding copies spans from the context (great for editing and RAG).

import torch

def speculative_step(target_probs, draft_probs, draft_tokens):
    """target_probs/draft_probs: (gamma, V) distributions at each drafted position.
    Returns the accepted tokens (lossless w.r.t. the target distribution)."""
    accepted = []
    for i, tok in enumerate(draft_tokens):
        p, q = target_probs[i, tok], draft_probs[i, tok]
        if torch.rand(()) < torch.clamp(p / q, max=1.0):
            accepted.append(int(tok))                          # accept the draft token
        else:
            residual = torch.clamp(target_probs[i] - draft_probs[i], min=0)
            accepted.append(int(torch.multinomial(residual / residual.sum(), 1)))
            break                                              # stop at first rejection
    return accepted
# (If all gamma tokens are accepted, one extra token is sampled from the target's next distribution.)

Other key optimisations#

• Quantisation (weights, activations, KV cache) — fewer bytes to move.
• Fused kernels and FlashAttention / FlashDecoding for attention.
• Tensor parallelism across GPUs for large models; disaggregated serving runs prefill and decode on separate GPU pools.
• CUDA graphs and compilation to reduce kernel-launch overhead.
• Caching at the application level: reuse answers for repeated questions; semantic caches (with care for correctness and privacy).
• Model routing: send easy requests to small models and hard ones to large models.

Serving engines#

Popular open-source engines — vLLM, SGLang, TensorRT-LLM, text-generation-inference, llama.cpp (CPU/edge), Ollama (local convenience) — implement many of these techniques. Choose based on hardware, model support, quantisation formats, and latency vs throughput goals.