Accelerator Architecture intermediate 8 min read 7 flashcards

KV-Cache Memory and Bandwidth

The key-value cache trades GPU memory capacity for inference speed, and understanding how that trade interacts with memory bandwidth is what separates fast serving systems from slow ones.

A 70B-parameter LLaMA model in fp16 weighs about 140 GB. Run it with a batch of 32 requests at sequence length 4096 and the KV cache alone adds another 80+ GB. On most accelerators, the cache exceeds the model weights before you even start worrying about activations. That is the memory wall for transformer inference, and it is not a software bug you can patch away.

What the cache stores and why it must grow

Every transformer decoder layer computes attention over the full context. For a given token at position t, the attention mechanism needs the key and value projections for every previous position 0 … t-1. Without caching, you would recompute those projections from scratch on every new token - O(t) work per step, O(t²) total. The cache trades memory for a flat O(1) projection cost per step.

The size of that trade is exact and predictable. For a single request:

bytes_per_token = 2 * num_layers * num_kv_heads * head_dim * dtype_bytes

For Llama-3-70B in fp16 (80 layers, 8 KV heads, 128-dim heads, 2 bytes):

2 * 80 * 8 * 128 * 2 = 327,680 bytes ≈ 0.32 MB per token

At 4096 tokens that is 1.3 GB per request. With 32 concurrent requests: 41 GB before any model weights are loaded. This arithmetic is why serving systems spend so much effort on memory management.

The bandwidth bottleneck during autoregressive generation

Prefill (processing the prompt) is compute-bound: you have a full matrix of queries attending to all prompt positions simultaneously, which saturates tensor cores. Generation (producing tokens one at a time) is a completely different regime.

During generation, each forward pass processes a single new token against a growing cache. The compute load is tiny - one query vector per head per layer. But before any floating-point arithmetic happens, the GPU must stream the entire KV cache for that request from HBM (high-bandwidth memory) into SRAM. That transfer is proportional to the context length and happens on every single step.

Phase Bottleneck Typical utilisation
Prefill Compute (tensor cores) 40-80% MFU
Generation (short context) Memory bandwidth < 10% MFU
Generation (long context) Memory bandwidth (dominant) < 5% MFU

An A100 SXM5 has 2 TB/s of HBM bandwidth. A 70B model in fp16 is 140 GB of weights alone - reading them once takes 70 ms at peak bandwidth. Add a 4096-token KV cache for 32 requests (41 GB) and each generation step requires streaming roughly 180 GB, giving a theoretical lower bound around 90 ms per token just from bandwidth. In practice you are well below peak, so observed latency is higher.

This is why generation throughput scales badly with batch size past a certain point: each new request adds cache pressure, and the memory bus saturates before the compute does.

IO-aware algorithms and architectural mitigations

The core insight of FlashAttention (Dao et al., 2022) is that naive attention incurs redundant HBM reads and writes of the QK^T and softmax intermediate matrices. By fusing the softmax and the value-weighted sum into a single tiled kernel that stays in SRAM, FlashAttention reduces HBM traffic by a factor proportional to the block size. This matters enormously during prefill but has a more limited effect on generation-phase cache reads, because the bottleneck there is loading K and V themselves, not intermediate matrices.

Two architectural changes in the model itself directly shrink the cache:

Multi-Query Attention (MQA) collapses all KV heads to one. For Llama-3-70B this would shrink the per-token cache from 0.32 MB to 0.04 MB - an 8x reduction - at some quality cost.

Grouped-Query Attention (GQA), standardised by Ainslie et al. (2023), uses G KV heads shared across H/G query heads. Llama-3 uses GQA with G=8, halving cache size relative to MHA while recovering most quality loss compared to MQA. The per-token cost becomes:

bytes = 2 * layers * G * head_dim * dtype_bytes

GQA is now the default in most production models precisely because the cache arithmetic changes so dramatically.

Memory management: from contiguous blocks to pages

Early serving frameworks allocated a contiguous memory region for each request's KV cache up front. The problem is that you rarely know the final sequence length at the start: allocate too little and you must copy; allocate too much and you waste HBM. With 80+ GB of cache competing for space, fragmentation compounds until 30-50% of HBM is wasted.

PagedAttention (Kwon et al., 2023), the core innovation in vLLM, applies the virtual-memory metaphor to KV cache management. Cache entries are stored in fixed-size pages (typically 16 tokens), and pages need not be contiguous. A logical block table maps (request, layer, position) to physical page locations. The attention kernel is rewritten to follow this indirection.

The result is near-zero fragmentation. The serving system can also share pages across requests - for example, multiple requests with the same system prompt can share the prefill pages read-only, a technique called prefix caching or prompt caching.

# Simplified block table structure
request_id -> [block_0, block_7, block_23, ...]
                  |          |          |
               HBM pages (each 16 tokens * bytes_per_token)

The bookkeeping cost is modest: a block table lookup per attention call. The throughput gains are 2-4x over contiguous allocation at realistic batch sizes.

When it falls down

Quantisation mismatch with activations. Quantising KV cache to int8 or fp8 saves 2x-4x memory and bandwidth, but the precision loss is not uniform. Certain attention patterns (spiky distributions, long-range dependencies) degrade noticeably. Some frameworks apply per-head or per-token scaling to compensate, but this adds implementation complexity and can reintroduce bandwidth overhead.

Very long contexts. At 128K tokens, even a quantised GQA cache can exceed 2 GB per request. Batch size collapses to single digits on an 80 GB A100. Streaming the cache per generation step becomes the sole bottleneck. Sparse attention variants (sliding-window, retrieval-augmented) mitigate this but require architectural changes and careful evaluation of which tokens you can safely evict.

Prefill-decode disaggregation. Because prefill and generation have different bottlenecks (compute vs. bandwidth), some systems route them to separate machines. Transferring the populated KV cache between hosts over NVLink or Infiniband adds latency and bandwidth pressure at the network layer, creating new failure modes if the interconnect is saturated.

Cache eviction under pressure. When HBM fills up, a serving system must either stall new requests or evict existing cache blocks to secondary memory (PCIe-attached CPU DRAM, NVMe). PCIe bandwidth (60-128 GB/s) is 15-30x slower than HBM. A cold cache hit on NVMe is three orders of magnitude worse. Eviction policies that get this wrong destroy latency for the affected requests.

Speculative decoding and verification. Speculative decoding runs a small draft model to propose multiple tokens, then verifies them in parallel with the target model. The verification step requires extending the KV cache for all draft tokens simultaneously, then rolling back rejected ones. The cache management logic here interacts poorly with paged allocators that assume monotone growth.

Further reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track