KV Cache › Profiling Memory Usage

Michael BrenndoerferJanuary 6, 20265 min read
Language AI HandbookAll chapters
Part 42 of 67Inference Optimization
Chapter 313KV Cache › Profiling Memory Usage
Reading section 10 of 14 · Profiling Memory Usage · 5 minSection 10 of 14 · 5 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 313 of 513

Section 10 of 14 · 5 min read

KV Cache › Profiling Memory UsageOne section · Sections in this chapter10 / 14 ↓
Sections in this chapter14
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

Let's measure actual memory consumption for different configurations:

In[17]:
Code
def calculate_cache_memory(
    num_layers: int,
    d_model: int,
    max_seq_len: int,
    batch_size: int,
    dtype_bytes: int = 2,  # FP16
) -> dict:
    """
    Calculate KV cache memory requirements.
    """
    # Memory per layer: 2 (K and V) x batch x seq_len x d_model x bytes
    memory_per_layer = 2 * batch_size * max_seq_len * d_model * dtype_bytes
    total_memory = num_layers * memory_per_layer

    return {
        "memory_per_layer_mb": memory_per_layer / (1024**2),
        "total_memory_mb": total_memory / (1024**2),
        "total_memory_gb": total_memory / (1024**3),
    }


# Model configurations
models = {
    "LLaMA-7B": {"layers": 32, "d_model": 4096},
    "LLaMA-13B": {"layers": 40, "d_model": 5120},
    "LLaMA-70B": {"layers": 80, "d_model": 8192},
}

# Calculate for different context lengths
context_lengths = [2048, 4096, 8192, 16384]
batch_size_1 = 1
memory_table_1 = []

for name, config in models.items():
    row = []
    for ctx in context_lengths:
        mem = calculate_cache_memory(
            config["layers"], config["d_model"], ctx, batch_size_1
        )
        row.append(mem["total_memory_gb"])
    memory_table_1.append((name, row))

# Calculate for different batch sizes (Context 4096)
target_ctx = 4096
batch_sizes = [1, 8, 32]
memory_table_2 = []

for name, config in models.items():
    row = []
    for bs in batch_sizes:
        mem = calculate_cache_memory(
            config["layers"], config["d_model"], target_ctx, bs
        )
        row.append(mem["total_memory_gb"])
    memory_table_2.append((name, row))
Out[18]:
Console
KV Cache Memory (GB) - Batch Size 1, FP16

Model        |   2048 |   4096 |   8192 |  16384
-------------------------------------------------------
LLaMA-7B     |   1.0G |   2.0G |   4.0G |   8.0G |
LLaMA-13B    |   1.6G |   3.1G |   6.2G |  12.5G |
LLaMA-70B    |   5.0G |  10.0G |  20.0G |  40.0G |

As context length increases, memory usage grows linearly. For the LLaMA-70B model, a single sequence at 8192 context requires 20 GB of cache memory, which is significant even for high-end hardware.

Out[19]:
Console


KV Cache Memory (GB) - Context 4096, FP16

Model        | Batch 1 | Batch 8 | Batch 32
--------------------------------------------------
LLaMA-7B     |    2.0G |   16.0G |   64.0G |
LLaMA-13B    |    3.1G |   25.0G |  100.0G |
LLaMA-70B    |   10.0G |   80.0G |  320.0G |

These numbers reveal why KV cache memory management is necessary for production systems. A LLaMA-70B model serving 32 concurrent requests at 4K context requires 320 GB just for the KV cache, exceeding what most single GPUs can give.

Out[20]:
Visualization
Line chart of KV cache memory in GB versus batch size (1-32) at 4096 context length for LLaMA 7B, 13B, and 70B. All three lines increase linearly.
KV cache memory scaling with batch size at fixed 4096 context length in FP16 for LLaMA 7B, 13B, and 70B models. All three curves increase linearly with batch size. The dashed horizontal reference line at 80 GB marks the capacity limit of an A100/H100 GPU. LLaMA 70B reaches 80 GB at batch size 8 and exceeds it beyond that point, meaning a single 80 GB GPU cannot serve more than eight concurrent 4K-context requests with this model.
Out[21]:
Visualization
Grouped bar chart comparing KV cache memory in GB at batch sizes 1, 8, and 16 for LLaMA 7B, 13B, and 70B at 4096 context length.
Comparison of KV cache memory versus model weights for batch sizes 1, 8, and 16 at 4096 context length. Bars show cache memory; dashed horizontal lines mark the model weight sizes for reference. For LLaMA 70B at batch size 16, the 160 GB cache exceeds the 140 GB model-weight footprint, and at batch size 32 (not shown) it grows to 320 GB. This comparison illustrates that cache memory, not model weight memory, often becomes the bottleneck in production inference.