KV Cache Memory › The Memory Bottleneck

Michael BrenndoerferJanuary 7, 20266 min read
Language AI HandbookAll chapters
Part 42 of 67Inference Optimization
Chapter 314KV Cache Memory › The Memory Bottleneck
Reading section 8 of 14 · The Memory Bottleneck · 6 minSection 8 of 14 · 6 min

Book contents

Section 8 of 14 · 6 min read

KV Cache Memory › The Memory BottleneckOne section · Sections in this chapter8 / 14 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

When KV cache memory dominates, we encounter a fundamental constraint that differs from the compute bottleneck that traditionally limits deep learning. Understanding this distinction is essential because the remedies for each bottleneck are completely different. Buying a faster GPU solves a compute bottleneck; it does nothing for a memory bottleneck. Reducing model depth helps with weight memory but not cache memory. Knowing which bottleneck you face determines which optimizations are worth pursuing.

In a compute-bound scenario, performance is limited by how fast the GPU can perform arithmetic operations. The GPU's thousands of CUDA cores are running at full utilization, and the only way to go faster is to reduce the number of operations or use hardware with more compute units. In a memory-bound scenario, performance is limited by how fast data can move between memory and compute units. The arithmetic units sit idle while waiting for data to arrive, regardless of how many of them exist. An A100 GPU has 312 teraFLOPS of FP16 compute capacity but only 2 terabytes per second of memory bandwidth. The ratio between these two numbers, approximately 156 FLOPs per byte, is the threshold that determines which resource dominates.

During token generation, the model performs two qualitatively different kinds of work. Attention computation reads the entire KV cache to compute attention scores, then reads it again to compute the weighted value sum. This requires moving massive amounts of data from GPU high-bandwidth memory to the compute units, with relatively modest arithmetic. Feed-forward layers perform dense matrix multiplications where the same weight matrix applies to the input, achieving high arithmetic intensity because the weights are reused across the batch. The attention part is almost always memory-bound for realistic sequence lengths; the feed-forward part can be compute-bound when batches are large enough.

We can formalize this relationship using the roofline model:

I=FLOPsBytes\mathcal{I} = \frac{\text{FLOPs}}{\text{Bytes}}

where:

  • I\mathcal{I}: the arithmetic intensity (measured in floating-point operations per byte transferred)
  • FLOPs\text{FLOPs}: the total floating-point operations performed for one generation step
  • Bytes\text{Bytes}: the total data transferred from GPU memory, dominated by KV cache reads

If I\mathcal{I} is lower than the GPU's ratio of peak compute to peak bandwidth, the process is memory-bound. For an A100, this ridge point is approximately 156 FLOPs/byte (312 TFLOPS divided by 2 TB/s). Any attention computation with arithmetic intensity below this threshold cannot fully utilize GPU compute because the data arrives too slowly.

The practical implication is that for single-sequence generation at typical context lengths, the attention computation is memory-bound. Each token generation requires reading the entire KV cache from memory. This memory transfer takes longer than the actual arithmetic, which is why techniques like FlashAttention (which reduces memory traffic through tiling) and KV cache quantization (which reduces the size of each transfer) provide significant speedups without reducing the number of arithmetic operations.

In[20]:
Code
# Compute vs Memory bound analysis
def analyze_bottleneck(
    batch_size: int,
    seq_len: int,
    d_model: int,
    num_layers: int,
    gpu_tflops: float = 312,  # A100 FP16
    gpu_bandwidth_tb: float = 2.0,  # A100 HBM
):
    """
    Analyze whether generation is compute-bound or memory-bound.

    For generation, we process one token at a time. Attention: O(seq_len) memory reads, O(seq_len x d_model) compute. FFN: O(d_model^2) memory reads, O(d_model^2) compute.
    """
    # FLOPs for generating one token (simplified)
    # Attention: 4 x d_model x seq_len (Q x K, softmax, x V, output proj)
    attention_flops = 4 * d_model * seq_len * batch_size * num_layers

    # FFN: 8 x d_model x d_ffn approx 8 x 4 x d_model^2 (for d_ffn = 4 x d_model)
    ffn_flops = 8 * 4 * d_model * d_model * batch_size * num_layers

    total_flops = attention_flops + ffn_flops

    # Memory reads (bytes) - dominated by KV cache for attention
    # KV cache read: 2 x seq_len x d_model x bytes per layer
    kv_read_bytes = 2 * seq_len * d_model * 2 * num_layers * batch_size

    # Time estimates (seconds)
    compute_time = total_flops / (gpu_tflops * 1e12)
    memory_time = kv_read_bytes / (gpu_bandwidth_tb * 1e12)

    arithmetic_intensity = total_flops / kv_read_bytes

    return {
        "total_flops": total_flops,
        "kv_read_bytes": kv_read_bytes,
        "compute_time_us": compute_time * 1e6,
        "memory_time_us": memory_time * 1e6,
        "arithmetic_intensity": arithmetic_intensity,
        "bottleneck": "memory" if memory_time > compute_time else "compute",
    }
In[21]:
Code
# Analyze for different sequence lengths
seq_lengths_to_analyze = [1024, 4096, 16384, 65536]
bottleneck_results = []

for seq_len in seq_lengths_to_analyze:
    analysis = analyze_bottleneck(
        batch_size=1, seq_len=seq_len, d_model=4096, num_layers=32
    )
    bottleneck_results.append((seq_len, analysis))
Out[22]:
Console
Bottleneck Analysis: LLaMA 7B Token Generation (A100 GPU)
=================================================================

Context: 1,024 tokens
  Compute time: 56.8 us
  Memory time:  268.4 us
  Arithmetic intensity: 33.0 FLOPs/byte
  Bottleneck: MEMORY

Context: 4,096 tokens
  Compute time: 61.9 us
  Memory time:  1073.7 us
  Arithmetic intensity: 9.0 FLOPs/byte
  Bottleneck: MEMORY

Context: 16,384 tokens
  Compute time: 82.6 us
  Memory time:  4295.0 us
  Arithmetic intensity: 3.0 FLOPs/byte
  Bottleneck: MEMORY

Context: 65,536 tokens
  Compute time: 165.2 us
  Memory time:  17179.9 us
  Arithmetic intensity: 1.5 FLOPs/byte
  Bottleneck: MEMORY
Out[23]:
Visualization
Grouped bar chart comparing compute time versus memory time across different context lengths.
Simplified roofline estimate of compute time versus KV-cache transfer time per token for LLaMA 7B on an A100 GPU across context lengths from 1K to 64K tokens. KV-cache transfer time dominates at every length and grows much faster than the estimated compute time, making the modeled operation increasingly memory-bound as context expands.

These results reveal a critical pattern in transformer inference performance. Memory transfer time dominates compute time at every context length in this simplified estimate. As context grows, the modeled KV-cache transfer time increases more steeply than compute time, so the memory bottleneck widens rather than narrows. For typical production scenarios at 4K-8K contexts, memory bandwidth is therefore the primary constraint. This is why KV cache quantization and memory-efficient attention implementations can provide significant speedups by reducing memory traffic. Buying a GPU with more FLOPS alone would not resolve this memory-bound regime; bandwidth and memory-access efficiency matter directly.

This bottleneck analysis demonstrates that at reasonable context lengths, single-sequence generation is memory-bound. The GPU spends most of its time waiting for KV cache data to transfer from high-bandwidth memory to the compute units. Increasing batch size helps amortize memory transfer costs, but it in turn increases total memory requirements, creating a constrained optimization problem. Every solution to the memory bottleneck creates a new version of the same problem at a different scale.