LLM Inference Optimization: KV Cache, Quantization, Batching and Speculative Decoding

Serving a large language model is a strange workload. A GPU capable of hundreds of trillions of operations per second spends much of its time idle, waiting for memory. A model that fits comfortably in GPU memory can still run out of it after a few dozen users. And the same hardware can produce 30 or 3,000 tokens per second depending entirely on how requests are scheduled. Understanding why is the key to every inference optimization, so this guide starts with the mechanics and the memory math, then covers each major technique with the numbers that justify it.

Model architectures quoted below (layer counts, head counts) are those of publicly documented open-weight models such as Llama 3; hardware figures are approximate published specifications. Use them as worked examples, and plug in your own model's config.

How LLM Inference Works: Prefill and Decode

A decoder-only transformer generates text autoregressively: given tokens 1..n, it predicts token n+1, appends it, and repeats. If you need a refresher on tokens and transformers, read how ChatGPT works first. A request goes through two very different phases:

A chatbot request with a 2,000-token prompt (system prompt, retrieved documents, history) and a 300-token answer spends one prefill pass on the 2,000 tokens and 300 sequential decode passes. The decode steps usually dominate total time, which is why most optimizations target them.

Why Decoding Is Memory-Bandwidth Bound

In a decode step at batch size 1, the GPU must read every weight of the model from high-bandwidth memory (HBM) to produce a single token, and it does about 2 floating-point operations per weight (a multiply and an add). In 16-bit precision each weight is 2 bytes, so the arithmetic intensity is about 1 FLOP per byte.

Compare that with the hardware. An NVIDIA H100 SXM is specified at roughly 989 TFLOPS of dense BF16 compute and about 3.35 TB/s of memory bandwidth. To keep the compute units busy you need about 989 / 3.35 ≈ 295 FLOPs per byte. Decode at batch 1 delivers about 1. The GPU is more than 99% idle on compute, waiting for memory.

That gives a hard upper bound on single-stream speed:

Two consequences drive everything else:

1. Batching is nearly free throughput. If 32 sequences decode together, each weight read serves 32 tokens. Intensity rises about 32×, and total tokens per second rises almost proportionally until you approach the compute limit or run out of memory for the KV cache.
2. Fewer bytes means faster tokens. Halving the bytes per weight (quantization) roughly halves the time per decode step in the bandwidth-bound regime. That is why quantization speeds up inference, not just shrinks it.

The KV Cache: What It Is and the Memory Math

In self-attention, each token produces a query, a key and a value vector in every layer. A new token's query must attend to the keys and values of all previous tokens. Without caching, generating token n would recompute keys and values for all n−1 earlier tokens at every step. The KV cache stores them once, so each decode step computes K and V only for the new token.

The formula

Worked examples (16-bit cache)

The last two rows show why grouped-query attention (GQA) became standard: sharing each key/value head across several query heads cut the 70B model's KV cache by 8×.

Sizing a server

Serve the 8B model in BF16 on one 80 GB GPU. Weights take about 16 GB. Reserve a few GB for activations, CUDA graphs and framework overhead, leaving roughly 58 GiB for the KV cache. At 128 KiB per token, that is about 475,000 tokens of cache, for example around 58 concurrent sequences at 8K context, or about 230 at 2K. The KV cache, not the weights, is what limits concurrency, and therefore throughput.

def kv_bytes_per_token(layers, kv_heads, head_dim, bytes_per_value=2):
    return 2 * layers * kv_heads * head_dim * bytes_per_value

def max_concurrent_seqs(gpu_gb, weight_gb, overhead_gb, seq_len, **cfg):
    free_bytes = (gpu_gb - weight_gb - overhead_gb) * 2**30
    return int(free_bytes // (kv_bytes_per_token(**cfg) * seq_len))

llama3_8b = dict(layers=32, kv_heads=8, head_dim=128)
print(kv_bytes_per_token(**llama3_8b) // 1024, "KiB/token")              # 128
print(max_concurrent_seqs(80, 16, 6, 8192, **llama3_8b), "seqs @ 8K")    # 58
print(max_concurrent_seqs(80, 16, 6, 8192, bytes_per_value=1,
                          **llama3_8b), "seqs @ 8K with FP8 KV cache")   # 116

PagedAttention and Prefix Caching

Early serving systems allocated each request a contiguous KV region sized for the maximum possible length. Most requests finish far shorter, so much of the reserved memory was wasted, and fragmentation wasted more. The vLLM paper (Kwon et al., SOSP 2023) reported that existing systems wasted a large majority of KV cache memory this way.

PagedAttention applies the operating-system idea of virtual memory paging. The KV cache is split into fixed-size blocks (for example 16 tokens each). Each sequence has a block table mapping logical blocks to physical blocks anywhere in GPU memory, and blocks are allocated on demand as the sequence grows. Waste is limited to the last, partially filled block per sequence.

Paging also makes sharing cheap. If many requests start with the same system prompt or few-shot examples, their first blocks can point to the same physical blocks (copy-on-write when they diverge). This is prefix caching: the shared prefix is prefilled once and reused, cutting both TTFT and memory. SGLang's RadixAttention generalises this with a radix tree of cached prefixes across requests, which helps multi-turn chat and agent workloads. Hosted LLM APIs expose the same idea as prompt caching, usually with discounted pricing for cached input tokens, which is why you should put stable content (system prompt, tool definitions, documents) at the start of prompts and variable content at the end.

Continuous Batching, Chunked Prefill and Disaggregation

Static vs continuous batching

With static batching, the server groups N requests, runs them together until the longest one finishes, then starts the next batch. A request generating 20 tokens waits for a neighbour generating 800, and the GPU runs with shrinking batches as sequences finish.

Continuous batching (iteration-level scheduling, introduced in the Orca paper at OSDI 2022) re-forms the batch at every decode step: finished sequences leave, waiting requests join. Combined with paged KV memory, it keeps the batch as large as memory allows at all times. This is the single biggest throughput win in modern serving engines.

Chunked prefill

A long prompt's prefill can take hundreds of milliseconds and stall all decoding sequences sharing the GPU, causing latency spikes. Chunked prefill splits long prefills into pieces and interleaves them with decode steps, smoothing inter-token latency for everyone.

Prefill-decode disaggregation

Because prefill is compute-bound and decode is bandwidth-bound, some large deployments run them on separate GPU pools, transferring the KV cache from prefill workers to decode workers (as explored in research systems such as DistServe and Splitwise). Each pool can then be sized and tuned for its own bottleneck and latency target. It adds complexity and network transfer, so it pays off mainly at scale.

Quantization: Fewer Bits, Faster Tokens

Quantization stores numbers in fewer bits. Since decode speed is governed by bytes moved, it improves both memory capacity and speed.

Worked example: a 70B-parameter model is about 140 GB in BF16, which needs at least two 80 GB GPUs. At 4 bits per weight it is about 35 GB plus quantization scales, which fits on a single 80 GB GPU with room for KV cache, or even a 48 GB card with a smaller cache.

How weight-only methods preserve quality: GPTQ quantizes layer by layer while compensating for the error using second-order information; AWQ identifies the small fraction of weight channels that matter most (judged by activation magnitudes) and scales them to protect them. Outlier handling is the central challenge in LLM quantization, because a few activation channels have very large values.

Speculative Decoding: Guess Cheaply, Verify in Bulk

Decode is slow because each step is a full pass of the big model for a single token, and that pass is memory-bound, so verifying several tokens in one pass costs about the same as generating one. Speculative decoding (Leviathan et al. and Chen et al., 2023) exploits this:

1. A small, fast draft model proposes the next k tokens.
2. The large target model scores all k positions in one forward pass.
3. Tokens are accepted left to right using a rejection-sampling rule; at the first rejection, the target's own corrected token is used and the rest discarded.

The acceptance rule is constructed so the output distribution is identical to sampling from the target model alone. You get the same quality, faster. If each draft token is accepted with probability α, independently, the expected number of tokens produced per target pass with k drafted tokens is (1 − α^k+1) / (1 − α). With α = 0.8 and k = 4 that is about 3.4 tokens per expensive pass instead of 1. Real speedups are lower because the draft has its own cost and acceptance varies by content: code and structured text, which are predictable, tend to accept more than creative writing.

def expected_tokens_per_pass(alpha, k):
    return (1 - alpha ** (k + 1)) / (1 - alpha)

for alpha in (0.6, 0.7, 0.8, 0.9):
    print(alpha, [round(expected_tokens_per_pass(alpha, k), 2) for k in (2, 4, 8)])
# 0.8 -> [2.44, 3.36, 4.33]: diminishing returns as k grows

Variants avoid a separate draft model: Medusa adds extra decoding heads to the target model, EAGLE drafts at the feature level for higher acceptance, some models are trained with multi-token prediction heads usable for drafting, and n-gram or prompt-lookup drafting copies likely continuations from the prompt itself, which works well for editing and RAG tasks that quote context. Speculative decoding helps most at low batch sizes, where the GPU has spare compute; at very high batch sizes, the system is already closer to compute-bound and gains shrink.

Attention Kernels, Architecture Choices and Parallelism

• FlashAttention computes exact attention in tiles that stay in fast on-chip SRAM, never materialising the full n×n attention matrix in HBM. It dramatically speeds prefill and long-context attention and reduces memory. Newer versions exploit newer GPU features.
• GQA and MQA reduce the number of KV heads (shown in the table above). Multi-head latent attention (MLA), used in DeepSeek's models, compresses keys and values into a smaller latent vector to shrink the cache further.
• Sliding-window or hybrid attention limits some layers to recent tokens, capping KV growth for long contexts.
• Mixture-of-experts (MoE) models activate only a few experts per token, so compute per token is far below total parameters, but all experts must still be in memory; serving them involves expert parallelism.
• Tensor parallelism splits each layer's matrices across GPUs (needs fast interconnect such as NVLink); pipeline parallelism splits layers across GPUs. Tensor parallelism also multiplies available memory bandwidth, reducing per-token latency for large models.
• CUDA graphs and kernel fusion remove per-step launch overheads that matter when each decode step is only a few milliseconds.

Metrics That Matter for LLM Serving

The central trade-off: larger batches raise throughput but also raise TPOT for each user. A chat product might target TTFT under one second and 20–50 ms per token (faster than people read), while an offline batch summarisation job ignores latency and maximises throughput. Always state your latency SLO before optimizing, and look at p95 and p99, not just averages.

Cost per Token: The Math

Illustrative example (plug in your own numbers): suppose a GPU costs $2.50 per hour and, with continuous batching at your latency target, it sustains 2,500 output tokens per second aggregated across users. That is 9 million tokens per hour, or about $0.28 per million output tokens. If you only achieve 250 tokens per second because requests are served one at a time, the same GPU costs about $2.78 per million tokens: ten times more for the same hardware. Utilisation, not list price, dominates cost.

Further levers on the same equation:

• Quantization raises tokens per second and frees KV memory for larger batches.
• Prefix caching removes repeated prefill work, which dominates cost in RAG and agent workloads with long, repeated prompts.
• Right-sizing the model: a fine-tuned or distilled small model may serve a narrow task at a fraction of a large model's cost.
• Routing: send easy requests to a small model and hard ones to a large one.
• Autoscaling and batching windows so GPUs are not idle at night; offline jobs can fill troughs.

Compare self-hosting against API pricing honestly: include engineering time, idle capacity, redundancy and the effect of provider-side prompt caching discounts. For many Indian startups at modest volume, APIs are cheaper; self-hosting wins at high, steady volume, for data residency, or for customised models.

An Optimization Playbook

1. Measure first. Record prompt and output length distributions, TTFT, TPOT, throughput and GPU memory under realistic load.
2. Use a modern engine. vLLM, SGLang or TensorRT-LLM give paged KV cache, continuous batching and optimized kernels out of the box. This step alone is often the largest gain.
3. Turn on prefix caching and restructure prompts so shared content comes first.
4. Quantize: try FP8 weights (if supported) and an FP8 KV cache, then 4-bit weights if memory-bound; validate with evals.
5. Tune batching limits (max concurrent sequences, max batched tokens, chunked prefill) against your latency SLO.
6. Add speculative decoding for latency-sensitive, low-concurrency traffic.
7. Revisit the model: smaller, distilled or task-tuned models, and request routing.

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    max_model_len=8192,             # caps per-sequence KV usage
    gpu_memory_utilization=0.90,    # fraction of GPU memory vLLM may use
    enable_prefix_caching=True,     # reuse KV blocks for shared prefixes
    kv_cache_dtype="fp8",           # halve KV memory (check model/GPU support)
)
params = SamplingParams(temperature=0.2, max_tokens=256)
system = "You are a support assistant for an Indian e-commerce site. Be concise.\n\n"
prompts = [system + q for q in ["Where is my order?", "How do refunds work?"]]
for out in llm.generate(prompts, params):
    print(out.outputs[0].text)

The same options are available as flags on vLLM's OpenAI-compatible server. Exact flag names and defaults change between versions, so check the current documentation.

How Inference Optimization Comes Up in Interviews

AI infrastructure and ML engineering interviews, and increasingly backend system design rounds for AI products, probe this material. Common questions:

• "How many concurrent users can one GPU serve for this model?" Do the KV cache math out loud: weights, overhead, bytes per token, context length, then concurrency. This is the LLM version of back-of-the-envelope estimation.
• "Why is generation slow even though the GPU shows low compute utilisation?" Memory-bandwidth-bound decode; fix with batching and quantization.
• "Explain the KV cache and its trade-off." Compute saved versus memory consumed; GQA, paging, FP8 cache.
• "How does speculative decoding preserve quality?" Verification with rejection sampling keeps the target distribution.
• "Design an LLM serving platform for an internal chatbot." Engine choice, autoscaling, prefix caching, routing, SLOs, monitoring; see designing a ChatGPT-like application.

Candidates who can connect a hardware number (bandwidth) to a product number (tokens per second, cost per million tokens) stand out immediately. Practise one full calculation for a model you know until you can do it on a whiteboard in five minutes.

Frequently Asked Questions

Pranjal Jain

Ex-Microsoft SDE · IIT Kanpur · Founder of Prepflix. Helps engineers crack startup and product company interviews across India.