⚡ Inference Optimization

Real-world scenario: GPT-4 serves millions of requests per day. A single GPT-4 forward pass on a modern GPU takes ~100ms. Without optimisation, that's 10 requests/second per GPU, costing millions of dollars per day. With KV caching, continuous batching, and quantisation, the same hardware serves 10-100× more requests. Inference optimisation is the difference between a research model and a production service.

🌟 Why Inference Optimization Matters

Training a model happens once. Inference happens billions of times. Every millisecond of latency and every watt of compute directly impacts:

• Cost: cloud GPU time is expensive; 10× efficiency = 10× fewer GPUs
• User experience: users expect responses in <2 seconds
• Scale: a 100ms model can serve 10 req/s; a 10ms model can serve 100 req/s

The four main levers are: KV Cache (avoid redundant computation), Batching (share computation), Quantisation (smaller numbers = faster math), and Memory Optimization (fit bigger batches).

KV Cache

In transformer attention, for every token in the sequence, the model computes Key and Value matrices from the token's embedding. During autoregressive generation, when you generate the Nth token, you've already computed K and V for tokens 1 through N-1. Why recompute them?

KV cache saves the Key and Value tensors from previous steps and reuses them. Only the new token's K and V are computed at each step.

Without KV cache:
  Step 1: compute K,V for [prompt]          → 1 forward pass on N tokens
  Step 2: compute K,V for [prompt + tok_1]  → 1 forward pass on N+1 tokens
  Step 3: compute K,V for [prompt + tok_1 + tok_2] → N+2 tokens
  ...
  Total compute: O(T²) where T = total tokens

With KV cache:
  Step 1: compute K,V for [prompt]          → full forward pass, cache K,V
  Step 2: compute K,V for [tok_1] only      → 1-token forward pass
  Step 3: compute K,V for [tok_2] only      → 1-token forward pass
  ...
  Total compute: O(T) after prefill — massive speedup!

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model.eval()

prompt = "The transformer architecture was introduced in"
inputs = tokenizer(prompt, return_tensors="pt")

with torch.no_grad():
    # use_cache=True (default): KV cache enabled — fast generation
    output_cached = model.generate(
        **inputs,
        max_new_tokens=20,
        use_cache=True,     # saves ~4x compute on longer sequences
    )
    # use_cache=False: recompute everything at each step — slow
    output_nocache = model.generate(
        **inputs,
        max_new_tokens=20,
        use_cache=False,
    )

# Both produce the same output — the cache is just a performance optimization
assert tokenizer.decode(output_cached[0]) == tokenizer.decode(output_nocache[0])
print("KV cache produces identical results, just faster ✅")

Trade-off: The KV cache uses memory. A long context window with large batch sizes can exhaust GPU memory — this is one of the key constraints in serving large models.

Batching

Batching processes multiple requests simultaneously, sharing the fixed cost of loading weights into GPU registers.

Without batching, 100 users waiting for responses each trigger a separate forward pass — the GPU processes them sequentially, one at a time. With a batch size of 32, the GPU can process 32 requests in roughly the same time as 1.

from transformers import pipeline

generator = pipeline("text-generation", model="gpt2", device=0)

# Sequential (slow): each prompt triggers a separate forward pass
import time
prompts = ["Hello world"] * 8

t0 = time.time()
results_seq = [generator(p, max_new_tokens=20)[0] for p in prompts]
t1 = time.time()
print(f"Sequential: {t1-t0:.2f}s")

# Batched (fast): all prompts processed together
t2 = time.time()
results_batch = generator(prompts, max_new_tokens=20, batch_size=8)
t3 = time.time()
print(f"Batched:    {t3-t2:.2f}s")
# Batched is typically 3-6x faster

Static batching (traditional): wait until you have a full batch, then process. Wastes time if requests arrive slowly.

Continuous batching (modern, used in vLLM): add new requests to in-flight batches as soon as a slot opens up — no waiting, much higher throughput.

Memory Optimization

GPU memory is the primary constraint in LLM serving. A 7B parameter model in FP16 requires ~14GB just for weights. Add the KV cache for a batch of requests and you can easily exceed 80GB (an A100's limit).

Key techniques:

1. Flash Attention: recomputes attention scores on the fly rather than materialising the full N×N attention matrix. Reduces memory from O(N²) to O(N). Standard in all modern inference frameworks.

# Enabling Flash Attention with transformers:
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    attn_implementation="flash_attention_2",   # requires flash-attn package
    torch_dtype=torch.bfloat16,
)

2. Paged Attention (vLLM): manages KV cache memory like virtual memory in an OS — allocates and frees cache blocks dynamically, allowing more concurrent sequences.

3. Tensor Parallelism: splits model weights across multiple GPUs, each GPU holds a fraction of each layer.

# With vLLM (production inference engine):
# pip install vllm
from vllm import LLM, SamplingParams

llm = LLM(
    model="mistralai/Mistral-7B-v0.1",
    tensor_parallel_size=4,    # split across 4 GPUs
    dtype="bfloat16",
)
params = SamplingParams(temperature=0.8, max_tokens=100)
outputs = llm.generate(["Tell me about attention mechanisms"], params)
print(outputs[0].outputs[0].text)

Latency vs Throughput

These are the two key metrics, and they're often in tension:

Increasing batch size improves throughput (more tokens/sec total) but increases latency (each user waits longer for their request to be processed).

Latency-optimized:   small batches, low waiting time per user
Throughput-optimized: large batches, high total output rate

Production systems run dynamic batching that tries to balance both:
- Prefer larger batches when load is high (throughput mode)
- Shrink batches when load is low (latency mode)

Most production LLM APIs target:

• Time to first token (TTFT): < 500ms
• Tokens per second (TPS): > 20 tok/s per user
• P99 latency: < 5 seconds for 200-token responses

GPU Utilization

A GPU is only fast when it's fully utilised — all CUDA cores busy, memory bandwidth maxed out. LLM inference is surprisingly hard to keep the GPU busy:

• Memory-bound, not compute-bound: reading weights from HBM (high-bandwidth memory) is the bottleneck, not the matrix multiplications themselves. Larger batches help by amortising the memory reads.
• Attention is sequential per user: each user's KV cache is different; you can't share computation across users for the attention step.
• KV cache limits batch size: every concurrent request occupies KV cache memory; once memory is full, you can't add more requests.

Key utilization metrics to monitor:

GPU Utilization (SM %)  — should be > 80% during inference
GPU Memory Usage        — should be > 75% (leaving headroom for KV cache growth)
Memory Bandwidth Usage  — the real bottleneck for LLM inference