Inference optimization decoded: KV cache, Flash Attention, and why your LLM is slower than it should be

Every team running LLMs in production hits the same wall eventually: inference is slow, and adding GPUs doesn't help nearly as much as expected.

The reason is almost always architectural, not computational. Modern transformer inference is a chain of optimizations stacked on top of each other, each with its own cost model. Most engineers know the names — KV cache, Flash Attention, quantization, speculative decoding — but far fewer know exactly what each one fixes, when it pays off, and when it gives you nothing.

I spent four years working on real-time AI systems for AAA gaming, where inference latency wasn't a metric on a dashboard — it was the difference between a believable in-game agent and a janky one. After that, building agentic systems in production, the same chain kept appearing in different costumes. The names of the tricks change. The order in which they pay off doesn't.

This is a walk-through of that chain, in the order the optimizations actually matter, with the math that explains why.

1. The sequential problem

LLMs generate text one token at a time. Each new token depends on every previous token through self-attention. This is what makes language modelling work; it is also what makes it slow.

Naively implemented, generating token n+1 means:

• Running the full input plus all previously generated tokens through every layer
• Computing Q (query), K (key), and V (value) projections for every position
• Computing the (n+1)×(n+1) attention matrix
• Performing softmax, the attention-weighted sum, and the feed-forward block
• Sampling one token, appending it, and repeating

For a sequence of length n, attention alone is O(n²). Generating 1000 tokens means n² work growing from 1 to 1,000,000 — a million attention operations for the 1000th token, summing to roughly 333M operations across the whole sequence.

This is the cost model every subsequent optimization attacks.

2. KV cache: from O(n²) to O(n)

The key observation: for every previous token, the K and V values are fully determined by its input embedding and the model weights. They don't change when a new token is appended.

So cache them.

At generation time, for a new token at position n+1:

• Compute Q only for position n+1
• Compute K, V only for position n+1
• Retrieve cached K, V from positions 0 through n
• Run attention using one Q against n+1 K,V pairs
• Append the new K, V to the cache

Attention compute per new token drops from O((n+1)²) to O(n+1). Linear, not quadratic.

What this looks like in numbers

For a 7B-parameter model at batch=1, fp16, generating at 2k context length on a single H100, the difference is roughly:

• Without KV cache: 3-5 tokens/sec, compute-bound
• With KV cache: 50-80 tokens/sec, memory-bandwidth-bound

Notice what happened. The bottleneck shifted. KV cache makes attention linear, which means compute is no longer the dominant cost. Now every new token requires reading the entire growing KV cache from memory. A 32k context with 32 layers, 32 heads, and 128 head dimensions produces a KV cache of roughly 500 MB per sample. Reading 500 MB per token at 1 TB/s memory bandwidth costs half a millisecond — not much, until you multiply by hundreds of generation steps and dozens of concurrent requests.

This is where Flash Attention enters.

3. Flash Attention: it's about memory, not compute

The most common misconception about Flash Attention is that it is a faster algorithm. It is not. It computes the same mathematical result as standard attention.

What it changes is memory access patterns.

The cost model of standard attention

Standard attention allocates and writes an N×N attention matrix to HBM (high-bandwidth GPU memory), then reads it back for softmax, then writes again, then reads it for the attention-weighted sum. That matrix materialization is the killer.

For a 4096-token sequence at head dimension 128, the attention matrix is 4096 × 4096 floats — 64 MB for fp16, per head, per layer. Multi-head with 32 heads: 2 GB of transient attention data per layer. For a 32-layer model: 64 GB of ephemeral memory touched per forward pass. This is why long-context models used to OOM on modest GPUs.

What Flash Attention does

Fuse attention into a single tiled kernel. Each tile:

• Loads a small block of Q, K, V into SRAM (on-chip memory)
• Computes the partial attention for that block
• Accumulates the running softmax and output without ever materializing the full matrix
• Writes only the final output to HBM

The N×N matrix never exists in HBM. The number of HBM reads drops from O(N²) to O(N).

The numbers

On an H100, Flash Attention typically delivers:

• 2-4× end-to-end inference speedup for transformer workloads with reasonable context lengths (1k-16k)
• Memory footprint reductions enabling 2-4× longer contexts without OOM
• Larger gains for training than for inference (but still material for inference)

It gives you nothing when:

• Batch is small and context is short (attention was already cheap)
• The model uses non-standard attention patterns (some MoE, sparse, ALiBi variants)
• The stack is compute-bound for other reasons (massive FFN, bad quantization)

Why "20×" is rarely achievable

You'll see marketing numbers claiming 10-20× speedups. These are for attention-in-isolation microbenchmarks, not end-to-end inference. In a real stack, you still pay for:

• Token embedding lookup
• Feed-forward network
• Layer norms
• Output projection
• Sampling

Flash Attention accelerates attention only. The rest of the stack is your floor.

4. Quantization: cheaper, with a trap

Once attention is handled, memory bandwidth becomes the dominant cost for inference. Quantization attacks this directly by reducing the memory footprint of model weights and activations.

INT8 instead of fp16 halves the weight memory. INT4 quarters it. The naïve win: 2× or 4× less bandwidth pressure, often 1.5-3× faster inference, with a modest accuracy hit.

The reality is more nuanced.

Weight-only vs. activation quantization

Weight quantization (weights in INT8, activations remaining in fp16/bf16) is the safe version. It's what most "INT8 inference" offerings actually do. Accuracy drop is minimal, often negligible, on standard benchmarks.

Activation quantization is where it gets hard. Activations have much wider dynamic range and contain outlier channels that carry disproportionate signal. Quantizing them naively crashes accuracy.

The long-context trap

Here's something rarely discussed: quantized models tend to degrade more on long-context tasks than on short-context ones. A model that passes benchmarks at 2k context may produce noticeably worse outputs at 32k.

Why: quantization error accumulates. Each attention step compounds the approximation. At long contexts, the compounding overwhelms the accuracy margin.

Mitigation:

• Benchmark at your actual target context length, not whatever the open benchmarks use
• Consider mixed precision (INT8 for early layers, higher precision for later layers)
• For critical applications, use fp16/bf16 and accept the cost

Rough practical guidance

• Weight-only INT8: ~1.5-2× speedup, less than 1% accuracy drop on standard benchmarks
• Weight + activation INT8 with care: ~2-3× speedup, 1-3% accuracy drop
• INT4 (GPTQ, AWQ family): ~3-4× speedup, 3-10% accuracy drop, sensitive to calibration

5. Speculative decoding: running two models to run one faster

At first glance, this sounds contradictory. Run a small draft model, have it generate 5-10 tokens ahead, then have the big model verify them all in a single forward pass. The acceptance rate determines the speedup.

Why it works

The big model's forward pass is dominated by weight loading, not by sequential dependency. Processing 1 token and processing 8 tokens with already-computed Q, K, V has nearly the same wall time — the weights still must be loaded layer by layer.

If the draft model gets 4 tokens right, the big model verifies them in one pass that would otherwise have done one token's worth of work. Effective speedup: 4× for that batch.

The math

Let:

• α = acceptance rate (fraction of draft tokens accepted)
• c = cost ratio (draft forward / target forward)
• k = draft lookahead (how many tokens drafted per round)

Expected tokens per target-model step: (1 − α^(k+1)) / (1 − α).

With good draft/target alignment (α = 0.7, k = 5), you get roughly 2.5 tokens per target step, minus draft overhead. Net speedup typically 2-3× in production.

Where it breaks

• Domains where draft and target disagree frequently (low α → speedup vanishes)
• Small batches or latency-critical single-request workloads where draft overhead dominates
• When you don't have a well-aligned draft model available

6. When none of this helps

Sometimes you deploy KV cache, Flash Attention, quantization, speculative decoding — and inference is still slow.

Likely suspects, in the order I check them:

• Serialized tool calls in an agent pipeline. Optimizing the LLM won't help if you await five tools in a row. Fix the orchestration first. This is usually the single biggest win in production agent systems.
• Poor batching. Continuous batching (vLLM, TGI) matters more than any attention optimization when you have multiple concurrent requests. Static batching is leaving 30-60% throughput on the floor.
• Unnecessary re-prompting. Many systems re-run the full prompt every turn when a cached conversation state would do. The most expensive bug is doing the same work twice.
• Over-large models. A 70B model solving a task a 7B can handle is the most common premature optimization failure. Right-size before you tune.

The order in which to apply these

For a typical production inference system, apply optimizations in this order:

1. KV cache (non-negotiable; nobody ships without it)
2. Continuous batching if you serve multiple requests
3. Flash Attention (1-click in most frameworks)
4. Quantization (weight-only INT8 first; benchmark at your real context length)
5. Speculative decoding (requires a well-aligned draft model)
6. Right-size the model (often worth more than any of the above)

Skip step 1 and step 3 won't be visible. Over-invest in step 5 with a poor draft model and you'll lose what step 4 bought you.

Closing thought

Inference optimization is not one trick. It is a sequence of independent optimizations, each with its own cost model, each paying off only when the previous one has been applied. The engineering discipline is to measure, understand which regime you are in, and pick the next optimization accordingly. The rest is noise.

If you want to see this stack in motion — when each optimization helps, when it doesn't, and how the bottleneck shifts as you toggle each one — there's an interactive Inference Optimization Explorer in the AP Lab Academy module that visualizes exactly this:

→ aplab.academy/en/courses/nlp-advanced/lessons/inference-optimization

If this was useful, the next deep dive will cover production RAG architecture — specifically why naive embedding-similarity retrieval falls apart in real systems, and what a four-stage retrieval pipeline (rewrite → hybrid retrieve → rerank → answer) actually looks like in code. Subscribe if that's the kind of thing you'd read.

— Aleksei