KV Cache, Flash Attention & Inference Optimization

A KV cache stores the key and value vectors of every prefix token so a decoder can reuse them instead of recomputing them. Without one, a naive autoregressive decoder does O(N²) work to generate N tokens, because at each step it recomputes attention over the full prefix. For a 4K-token response that is 16M attention operations, most of them redundant. A prefix token's hidden state is deterministic once computed, so the new token needs to run its query against the cached keys and values of everything before it.

Attention also moves a lot of data. Standard attention materializes an N×N score matrix, an N×d softmax output, and an N×d final output, which means many reads and writes to HBM. For N≥2K, attention becomes memory-bound before it becomes FLOP-bound. Classic attention kernels underuse modern GPUs by 4–10×.

Two optimizations, both from Dao et al., took frontier inference from slow to fast:

1. KV cache. Store the K and V vectors of each prefix token. A new token's attention is one query against the cached keys, which takes inference from O(N²) to O(N) per generation step.
2. Flash Attention. Tile the attention computation so the full N×N matrix does not have to be written to HBM. The softmax and the matmuls run in SRAM. That gives a 2–4× wall-clock speedup on A100 and 5–10× on H100 with FP8.

By 2026 both are standard. vLLM, TensorRT-LLM, SGLang, and llama.cpp all assume them, and frontier models ship with Flash Attention enabled.

The KV cache is the stored K and V of each prefix token, and this section covers what it costs in memory, how Flash Attention keeps attention on-chip, and three serving techniques that cut latency or raise throughput.

KV cache math

Per decoder layer, per token, the cache holds one K vector and one V vector for every KV head:

bytes_per_token_per_layer = 2 * n_kv_heads * d_head * dtype_size
                            ^
                            K and V

One head alone costs 2 * 128 * 2 = 512 bytes per token per layer (d_head=128, fp16). The head count multiplies that. For a 7B-class model with 32 layers and 32 KV heads (full multi-head attention):

per token per layer = 2 * 32 * 128 * 2 = 16,384 bytes (16 KB)
per token (32 layers) = 512 KB
per 32K context = 16 GB

For Llama 3 70B (80 layers, d_head=128, GQA with 8 KV heads):

per token per layer = 2 * 8 * 128 * 2 = 4096 bytes (4 KB)
per token (80 layers) = 320 KB
per 32K context = 10.7 GB

At 128K context that is about 43 GB, more than an entire 40 GB A100, for KV cache alone at batch size 1 and before a single weight is loaded.

GQA is where the KV-cache savings come from. The cache scales with KV heads, not query heads. Llama 3 70B has 64 query heads, so MHA with 64 KV heads would be 8× larger, about 86 GB at 32K. MLA compresses further still.

Drag the dimensions and watch the cache size move. Push the sequence length or batch size up and see how quickly it outgrows a single GPU:

Flash Attention: the tiling trick

Flash Attention computes attention in tiles so the N×N score matrix is not materialized in HBM. For contrast, here is standard attention:

S = Q @ K^T          (HBM read, N×N, HBM write)
P = softmax(S)       (HBM read, HBM write)
O = P @ V            (HBM read, HBM write)

That is three HBM round trips. On H100, HBM bandwidth is 3 TB/s and SRAM is 30 TB/s, so an HBM trip costs roughly 10× what staying on-chip does.

Flash Attention:

for each block of Q (tile size ~128 × 128):
    load Q_tile into SRAM
    for each block of K, V:
        load K_tile, V_tile into SRAM
        compute S_tile = Q_tile @ K_tile^T     (SRAM)
        running softmax aggregation             (SRAM)
        accumulate into O_tile                  (SRAM)
    write O_tile to HBM

Each tile takes one HBM trip, and the total memory footprint drops from O(N²) to O(N). The backward pass recomputes some values from the forward pass instead of storing them, which saves more memory.

Numerical trick. The running softmax keeps (max, sum) across tiles, so the final normalization is exact. Flash Attention is not an approximation. Its output matches standard attention up to fp16 non-associativity.

Version evolution:

Flash 4 is forward-pass only at launch. Training still uses Flash 3. GQA and varlen support for Flash 4 is pending (mid-2026).

Speculative decoding: the other latency win

Speculative decoding has a cheap model propose N tokens and a big model verify all N in parallel. If verification accepts k tokens, you paid one big-model forward pass for k generations. Typical k is 3–5 on code and prose.

2026 defaults:

• EAGLE 2 / Medusa. Draft heads built into the model that share the verifier's hidden states. 2–3× speedup with no quality loss.
• Speculative decoding with a draft model. 2–4× speedup on consumer hardware.
• Lookahead decoding. Uses Jacobi iteration, so no draft model is needed. Niche but free.

Continuous batching

Continuous batching swaps new requests into the batch as soon as old ones finish. It first shipped in Orca and now lives in vLLM, TensorRT-LLM, and SGLang. Typical chat workloads see a 5–10× throughput gain.

Classic batched inference waits for the slowest sequence to finish before starting a new batch, which wastes GPU time when short responses finish early.

PagedAttention: KV cache as virtual memory

PagedAttention is vLLM's headline feature. The KV cache is allocated in 16-token blocks, and a page table maps logical positions to physical blocks. That lets you share KV across parallel samples (beam search, parallel sampling), hot-swap prefixes for prompt caching, and defragment memory. Throughput improves 4× over naive contiguous allocation.

code/main.py implements:

1. A naive incremental decoder that recomputes attention over the full prefix at every step.
2. A KV-cached decoder that appends one K and V per step and queries against the cache.
3. A tiled softmax that simulates Flash Attention's running-max algorithm.
4. kv_cache_bytes, which turns the KV cache math above into a size table.

Step 1: KV cache

class KVCache:
    def __init__(self):
        self.K = []
        self.V = []
 
    def append(self, k, v):
        self.K.append(k)
        self.V.append(v)
 
    def __len__(self):
        return len(self.K)

The cache appends one K and one V vector per generated token. This demo models a single head, so it is two flat lists. A real cache keeps one such pair per layer and per KV head, which is exactly the multiplier kv_cache_bytes applies.

Step 2: tiled softmax

def tiled_softmax_dot(q, Ks, Vs, tile=4):
    """Flash-attention-style incremental softmax(qK^T)V with tile size `tile`."""
    d_head = len(Vs[0])
    scale = 1.0 / math.sqrt(len(q))
    m = float("-inf")
    s = 0.0
    out = [0.0] * d_head
    for start in range(0, len(Ks), tile):
        k_block = Ks[start:start + tile]
        v_block = Vs[start:start + tile]
        scores = [dot(q, k) * scale for k in k_block]
        new_m = max(m, *scores)
        if m == float("-inf"):
            exp_old = 0.0
        else:
            exp_old = math.exp(m - new_m)
        exp_new = [math.exp(sc - new_m) for sc in scores]
        s = s * exp_old + sum(exp_new)
        for j in range(d_head):
            out[j] = out[j] * exp_old + sum(e * v[j] for e, v in zip(exp_new, v_block))
        m = new_m
    return [o / s for o in out]

The output matches softmax(qK^T / sqrt(d)) V computed in one shot (attention_full in main.py), but the working set at any moment is a tile × d_head block instead of the full N × d_head.

Step 3: compare naive vs cached decoding on a 10-token generation

The demo runs N=10 and counts one attention op per key attended to. Both decoders report 55 (N(N+1)/2), because each step still attends to the whole prefix. The cache's saving is not recomputing the earlier tokens' K and V (hidden states) at every step, and this pure-Python demo does not count that work. What it does verify is that the two decoders produce the same outputs (max abs diff 0).

Step 4: print the KV cache size table

kv_cache_bytes(N, n_layers, n_heads_kv, d_head, dtype) is the formula from the KV cache math section. main.py prints it for six configs at 2K, 32K, and 128K context: Llama-3.2-3B, Llama-3-8B, Llama-3-70B, Llama-3.1-405B, Qwen2.5-72B, and DeepSeek-V3. The DeepSeek-V3 row is a rough stand-in for MLA, modeled as one head with a 512-dim latent. Llama-3-70B comes out at 10.74 GB for 32K and 42.95 GB for 128K, matching the math above.

# HuggingFace transformers auto-enables KV cache on decoder-only generate().
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B",
    attn_implementation="flash_attention_2",  # use FA3 if Hopper
    torch_dtype="bfloat16",
)
# generate() uses KV cache automatically

vLLM production:

pip install vllm
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --max-model-len 32768 \
    --enable-prefix-caching \
    --kv-cache-dtype fp8

Prefix caching across requests is a big 2026 win. The same system prompt, few-shot examples, or long context document reuses its KV across calls. For agent workloads with repeated tool prompts, prefix caching often gives a 5× throughput gain.

outputs/skill-inference-optimizer.md holds the skill. It picks the attention implementation, KV cache strategy, quantization, and speculative decoding for a new inference deployment.

1. Easy. Run code/main.py. Confirm the naive and cached decoders produce the same output; note that the op counts are equal and why (the saving is the K/V recomputation the demo does not count).
2. Medium. Implement prefix caching: given a prompt P and several completions, run one forward pass over P to fill the KV cache, then branch per-completion. Measure speedup vs re-encoding P for each.
3. Hard. Implement a toy PagedAttention: KV cache in fixed 16-token blocks with a free-list. When a sequence finishes, return its blocks to the pool. Simulate 1,000 chat completions with varying lengths. Compare memory fragmentation vs contiguous allocation.

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Flash 1)
• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Flash 2)
• Shah et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (Flash 3)
• FlashAttention-4 release notes (Dao-AILab, 2026) (Blackwell 5-stage pipeline and the software-exp2 trick). Read the repo README for the forward-only launch caveats.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention (the vLLM paper)
• Leviathan et al. (2023). Fast Inference from Transformers via Speculative Decoding (speculative decoding)
• Li et al. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty (EAGLE-1/2, the integrated-draft approach)
• Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (the Medusa approach)
• vLLM docs: PagedAttention (the canonical deep dive on the 16-token block and page-table design)