KV cache, PagedAttention and FlashAttention

The KV cache decides how many users fit on a GPU. PagedAttention decides how much of that memory is actually used. FlashAttention decides how fast attention runs once the data is there. This page sizes the cache with real numbers, then explains the two ideas that made modern serving possible.

Source: Vishal Singh's post of 25 inference-engineering interview problems · saved in sources/09-vishal-inference-25.md

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

Problem 3 · Size the KV cache for a Llama-style 70B at 32K context and 100 users, with GQA and precision

What the interviewer is testing: can you write the formula from memory, plug in the right shape (KV heads, not query heads), convert units without slipping, and then turn the answer into a hardware decision with headroom?

The concepts

Every layer stores one key vector and one value vector per KV head per token. Query heads don't count: under grouped-query attention (GQA) several query heads share one KV head. Module 08 derives this; here's the formula to use:

KV bytes = 2 × L × n_kv × d_head × bytes_per_value × T × users
           (K and V) (layers) (KV heads) (head dim)  (precision) (tokens) (sequences)

per-token cost = 2 × L × n_kv × d_head × bytes_per_value

The rest of GPU memory goes to the weights (parameters × bytes per weight), activations and temporary buffers, the CUDA context, and the engine's own overhead. Engines like vLLM reserve a fraction of memory up front and carve the KV pool out of what's left, so plan with headroom. This page assumes 90% of HBM is usable for weights plus KV; that's an assumption to tune.

A worked answer with real numbers

Model shape. The Llama-2/3 70B published configs: 80 layers, 64 query heads, 8 KV heads, head dim 128 (hidden size 8,192 = 64 × 128). About 70 billion parameters.

GPUs (NVIDIA product pages): H100 SXM 80 GB at 3.35 TB/s; H200 141 GB at 4.8 TB/s. Note GB here is 10⁹ bytes while the cache maths below lands in GiB (2³⁰); mixing them is a classic 7% error.

Values per token = 2 × 80 × 8 × 128 = 163,840

Per token:  BF16 (2 B)  = 327,680 B = 320 KiB
            FP8  (1 B)  = 160 KiB
            4-bit       =  80 KiB  (+ a little for scales)

Per user at 32K (32,768 tokens):
            BF16 = 320 KiB × 32,768 = 10 GiB
            FP8  = 5 GiB
            4-bit = 2.5 GiB

100 users, all at full 32K:
            BF16 = 1,000 GiB = 1,073.7 GB
            FP8  =   500 GiB =   536.9 GB
            4-bit =  250 GiB =   268.4 GB

Without GQA (64 KV heads, BF16): 2.5 MiB/token → 8,000 GiB ≈ 8,590 GB. GQA saves 8×.

Now add the weights and divide by usable memory per GPU. Weights at FP8 are about 70 GB (at BF16, about 140 GB).

Turning that into a deployment. Tensor-parallel groups usually come in powers of two, and 8 GPUs is one node. FP8 weights plus FP8 KV need 9 H100s at 90% usable: just over one node. Options: one 8×H200 node (lots of room), two H100 nodes, or trim the requirement. With TP = 8, each GPU holds exactly one of the 8 KV heads, so per-GPU KV cost is 20 KiB per token at FP8. Beyond TP = 8, KV heads must be replicated, and the per-GPU saving stops scaling.

The question behind the question. "100 users at 32K" is the worst case. Real conversations are mostly short. If the average live context is 8K, the FP8 cache for 100 users is 125 GiB, and the whole thing fits on one 8×H100 node with room to spare. Whether you can bank on the average instead of the maximum depends on how the cache is allocated, which is problem 4.

Trade-offs

Common mistakes

• Using query heads (64) instead of KV heads (8): an 8× overestimate.
• Forgetting the factor 2 for keys and values.
• Mixing GiB and GB, or forgetting the weights and headroom.
• Assuming tensor parallelism always divides the cache: with more GPUs than KV heads, heads are replicated.

Problem 4 · Explain PagedAttention and the problem it solves

What the interviewer is testing: whether you know the problem is memory waste limiting batch size, not attention speed. Bonus points for block tables, copy-on-write, and what it costs.

The concepts

The problem. A sequence's KV cache grows by one token per step, and you don't know in advance how long the answer will be. Early serving systems stored each sequence's cache in one contiguous tensor sized for the maximum length. That wastes memory three ways:

• Reservation: space for tokens not yet generated sits idle for the whole request.
• Internal fragmentation: most requests finish far short of the maximum, so most of the slab is never used.
• External fragmentation: slabs of different sizes leave gaps between them that are too small to reuse.

Wasted KV memory means fewer sequences in the batch, and in the bandwidth-bound decode phase a bigger batch is almost free throughput. So waste directly caps throughput.

The idea (Kwon et al., 2023), borrowed from operating-system virtual memory:

• Split the KV pool into fixed-size blocks, each holding the keys and values for a fixed number of tokens (commonly 16) across all layers' heads for that block.
• Each sequence has a block table mapping its logical block 0, 1, 2… to physical blocks anywhere in the pool. Blocks are allocated only when the previous one fills.
• The attention kernel reads keys and values through the block table, gathering non-contiguous blocks.
• Because blocks are indirect, sequences can share them. Parallel samples or beam-search branches point to the same prompt blocks, with a reference count. On write to a shared block, copy it first (copy-on-write). The same mechanism enables prefix caching across requests (page 04).
• When memory runs out, the scheduler can preempt a sequence: free its blocks and later recompute its prefix, or swap the blocks to CPU memory.

logical token position p  →  block = p // B,  offset = p % B
physical address          =  block_table[seq][block] × block_bytes + offset × token_bytes

waste per sequence        ≤  B − 1 tokens (only the last block is partly empty)
average waste             ≈  B / 2 tokens

The paper reports near-zero KV waste and 2–4× higher throughput than the systems it compared against (FasterTransformer, Orca) at the same latency, with larger gains for longer sequences, bigger models and complex decoding.

A worked answer with real numbers

Assumptions. The 70B shape from problem 3, FP8 weights (70 GB) and FP8 KV (160 KiB per token), one node of 8 H100 SXMs, 90% of HBM usable, max context 32K, real requests averaging 2,048 tokens (an assumed traffic mix).

KV budget  = 8 × 80 GB × 0.9 − 70 GB = 506 GB ≈ 471 GiB

Contiguous (reserve 32K per request):
  per request = 32,768 × 160 KiB = 5 GiB
  fits        = 471 / 5 ≈ 94 requests
  utilisation = 2,048 / 32,768 ≈ 6% of reserved KV memory holds real tokens

Paged (16-token blocks):
  per request ≈ (2,048 + 8 wasted) × 160 KiB ≈ 321 MiB
  fits        ≈ 471 × 1,024 / 321 ≈ 1,500 requests
  utilisation ≈ 99.6%

One block = 16 tokens × 160 KiB = 2.5 MiB (the whole model; 320 KiB per GPU at TP = 8)

About 16× more concurrent sequences from the same memory. In practice you won't run 1,500 at once on one node, because TPOT climbs with batch size and you hit compute limits or latency targets first. The point is that memory stops being the binding constraint, and the scheduler can choose the batch size that meets latency targets.

Trade-offs

Common mistakes

• Describing PagedAttention as a faster attention algorithm. Its gain is batch size via memory efficiency.
• Confusing it with offloading to disk or CPU. Paging here is within GPU memory; swapping is an optional extra.
• Forgetting sharing and copy-on-write, which make parallel sampling and prefix caching cheap.

Problem 5 · Explain why FlashAttention is faster when the maths is the same

What the interviewer is testing: whether you think in memory traffic, not FLOPs. The key sentence: standard attention is limited by reading and writing the N×N matrix to HBM; FlashAttention never writes it.

The concepts

Standard attention for one head, sequence length N, head dim d:

S = Q Kᵀ / √d        N×N   written to HBM
P = softmax(S)       N×N   read from HBM, written back
O = P V              N×d   P read again from HBM

FLOPs ≈ 4 N² d (two matmuls)
HBM traffic ≈ Θ(N d + N²): the N² term dominates

A GPU has a large, relatively slow HBM and a tiny, very fast on-chip SRAM. Softmax and masking do few FLOPs per byte, so the kernels that read and write S and P are memory-bound. The matmuls are fast; the round trips through HBM are what cost time.

FlashAttention (Dao et al., 2022) changes the order of work, not the result:

• Tiling. Load a block of Q rows and a block of K and V rows into SRAM. Compute that tile of S, use it, discard it. S and P never go to HBM.
• Online softmax. Softmax needs a row's max and sum over all columns, which you don't have when working tile by tile. Keep a running max m and running sum ℓ per row, and rescale the partial output when the max changes:

for each K,V tile j:
    S_j   = Q_i K_jᵀ / √d
    m_new = max(m, rowmax(S_j))
    ℓ     = e^(m − m_new) · ℓ + rowsum(e^(S_j − m_new))
    O     = e^(m − m_new) · O + e^(S_j − m_new) · V_j
    m     = m_new
O = O / ℓ                      exact same result as softmax(QKᵀ/√d) V

• Kernel fusion. Matmul, scale, mask, softmax and the second matmul run in one kernel, so there's one launch and no intermediate tensors.
• Recompute in the backward pass. For training, it stores only m and ℓ and recomputes S tile by tile during backprop. More FLOPs, much less memory traffic, still faster.

The paper shows HBM accesses fall to Θ(N² d² / M), where M is the SRAM size, and that this is optimal for a range of SRAM sizes. Memory use goes from quadratic to linear in N, which is what made long contexts practical.

Versions. FlashAttention-2 (Dao, 2023) reorganised the work: fewer non-matmul operations, parallelism across the sequence dimension, and better splitting of work between warps. FlashAttention-3 (Shah et al., 2024) targets Hopper GPUs: it overlaps data movement and computation using the Tensor Cores' and TMA's asynchrony, and adds FP8. Its abstract reports 1.5–2.0× over FlashAttention-2 on H100, up to 740 TFLOPS in FP16 (75% utilisation) and close to 1.2 PFLOPS in FP8.

Serving note. Decode is different: one query token against a long cache. There's no N×N matrix to avoid, and the cost is reading K and V once. Decode kernels split the long KV sequence across thread blocks so the GPU stays busy at batch 1, and they read through PagedAttention's block table. Modern engines combine both ideas.

A worked answer with real numbers

Assumptions. One sequence of N = 8,192 tokens, head dim 128, FP16, 32 query heads per layer (an 8B-class shape), one H100 SXM: 3.35 TB/s HBM, about 989 TFLOPS dense FP16 (half the page's 1,979 "with sparsity" figure). Count only the big N×N traffic for standard attention: write S, read S, write P, read P.

Standard attention, per head:
  S size         = 8,192² × 2 B = 128 MiB
  HBM traffic    ≈ 4 × 128 MiB  = 512 MiB ≈ 0.54 GB
  time at 3.35 TB/s ≈ 160 µs
  × 32 heads     ≈ 5.1 ms per layer   (memory-bound)

Compute, per layer:
  FLOPs = 32 × 4 × 8,192² × 128 ≈ 1.1e12
  time at 989 TFLOPS ≈ 1.1 ms

FlashAttention, per head:
  HBM traffic ≈ read Q, K, V + write O = 4 × 8,192 × 128 × 2 B = 8 MiB
  time ≈ 2.5 µs                        → now compute-bound (~1.1 ms per layer)

Memory for S if materialised: 128 MiB × 32 heads = 4 GiB per layer, per sequence

Same 1.1e12 FLOPs either way. Standard attention spends roughly 5 ms per layer just moving S and P; FlashAttention moves about 64× less data per head and becomes limited by the matmuls. These are idealised floors, but they show the shape of the win. Real speedups depend on N (bigger N, bigger win) and the kernel: the original paper reports 3× on GPT-2 training at 1K tokens.

Trade-offs

• Exact vs approximate. FlashAttention is exact (up to floating-point ordering). Approximate attention methods cut FLOPs but often didn't speed up wall-clock time, a point the paper makes.
• Hardware-specific code. Peak performance needs kernels tuned per GPU generation (FA3 is Hopper-specific). Portable fallbacks are slower.
• Short sequences gain less. At small N the N² traffic is small and other overheads dominate.
• Recompute costs FLOPs in training's backward pass; it pays off because memory, not compute, was the limit.

Common mistakes

• Saying it "reduces FLOPs" or "approximates attention". It does neither.
• Not mentioning online softmax, the trick that makes tiling possible.
• Assuming the same speedup at decode time. Decode has a different bottleneck: reading the KV cache.

Practise it

1. Check the formula against an engine. On a GPU with 16–24 GB, start vLLM with a small GQA model (for example a 1B–3B Llama-family instruct model). At startup it logs how many KV-cache blocks or tokens it allocated. Compute the expected number from the formula, the model's config.json (layers, KV heads, head dim) and the memory left after weights. They should agree to within the engine's overheads. Repeat with --kv-cache-dtype fp8 and confirm the token capacity roughly doubles.
2. Simulate fragmentation. In 30 lines of Python, draw 10,000 request lengths from a log-normal distribution, then compute utilisation for contiguous max-length allocation and for block sizes 1, 16, 64 and 256. Plot waste against block size.
3. Feel FlashAttention on Colab. On a free T4 GPU, time a naive PyTorch attention (explicit softmax(Q @ K.T) @ V) against torch.nn.functional.scaled_dot_product_attention for N = 512 to 8,192, recording time and peak memory with torch.cuda.max_memory_allocated. Watch the naive version's memory grow with N² and eventually run out. Note which backend PyTorch picks on your GPU; FlashAttention kernels need recent GPUs, so a T4 may use the memory-efficient backend instead.
4. Write online softmax. Implement the tiled loop above in NumPy for one head, check it matches the naive result to 1e-5, and count bytes "moved" with a tile size you choose.

In this course

• Module 08 · Inference-aware design: MHA vs GQA vs MQA, KV quantisation (KIVI, KVQuant), and a KV-memory demo. Read it before this page if the formula is new.
• Module 05 · Serving: two cache types in one paged pool, prefix caching with checkpoints, and a production-scale serving stack.
• Module 08 · Tensor maths and autograd: Triton and FlashAttention-style tiling from the kernel side.
• Module 08 · Distributed training: tensor parallelism, and why KV heads get replicated beyond TP = n_kv.
• Page 01 · Request lifecycle: where prefill and decode sit in a request.

Why it matters for robotics

Self-check

Resources

• Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)
• vLLM docs · Paged attention kernel design and vLLM documentation
• SGLang documentation and TensorRT-LLM documentation
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023)
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (Shah et al., 2024)
• flash-attention repository (kernels and supported GPUs)
• Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron et al., 2023): the 70B model with GQA
• NVIDIA Nsight Compute: measure a kernel's memory traffic yourself
• NVIDIA product pages: H100 and H200 (memory and bandwidth used above)