LLM serving: KV cache, GQA and memory budgets
Estimate what fits on an accelerator before promising a concurrency target. Separate model weights, KV cache, activations and runtime reserve.
- •Calculate conventional KV-cache memory from a model configuration.
- •Distinguish GQA, PagedAttention and FlashAttention.
- •Set token-aware admission limits from measured capacity.
Sources & further reading
GQA ↗PagedAttention ↗FlashAttention ↗In-Depth Architecture & Mechanics
Interview brief
You have a hypothetical dense 70-billion-parameter decoder, 80 layers, 8 KV heads and head dimension 128. Requests can reach 32,768 cached tokens. Explain why a model that fits in GPU memory may still fail under load.
Work the memory calculation
For a conventional attention KV cache, bytes per sequence = 2 × layers × KV heads × head dimension × cached tokens × bytes per value. With BF16, the example is 2 × 80 × 8 × 128 × 32,768 × 2 = 10,737,418,240 bytes = 10 GiB per sequence across all shards. This formula does not describe compressed latent caches or every hybrid architecture.
Weights are only one part of the budget
Seventy billion BF16 parameters require approximately 140 GB, or 130.39 GiB, before overhead. A raw 4-bit representation would be about 35 GB, but quantization scales, unquantized layers, kernel workspaces and layout matter. Quantized weights do not automatically quantize KV cache. Budget separately for activations, communication buffers, fragmentation and a safety reserve.
Understand the mechanisms
Grouped-query attention shares KV heads across multiple query heads, reducing cache relative to otherwise comparable multi-head attention. PagedAttention manages KV storage in blocks so differently sized requests need not reserve contiguous maximum-length buffers. FlashAttention reduces attention memory traffic through tiled computation; it does not make exact dense attention linear in sequence length.
Translate memory into admission control
If measured free cache budget is 120 GiB, twelve fully allocated 10-GiB sequences are an idealized upper bound, not a safe operating concurrency. Leave reserve, account for tensor-parallel placement and use measured allocator behavior. Mixed request lengths require token-aware admission rather than a fixed request count. Define preemption and rejection behavior explicitly.
Follow-up exercise
Repeat the estimate for 32 KV heads, an 8k context and an 8-bit KV representation. Then explain which quality and performance checks you need before changing precision. Measure time to first token, per-token latency, queue age, cache occupancy and out-of-memory failures under a realistic length distribution.
Memory budget = weights + KV cache + activations + workspaces + reserve Token-aware queue → admission → prefill → decode → release KV blocks
def kv_bytes(layers, kv_heads, head_dim, tokens, bytes_per_value):
return 2 * layers * kv_heads * head_dim * tokens * bytes_per_value
size = kv_bytes(80, 8, 128, 32768, 2)
print(f"Aggregate KV cache per sequence: {size / 2**30:.2f} GiB")
assert size == 10 * 2**30
# Excludes weights, activations, allocator overhead and reserve.
# Latent attention and other cache architectures need different formulas.