KV Cache Memory Calculation

Recommended resources

The freely available resources this lesson will direct you to.

Welcome back. In the previous lesson, you learned to use nvidia-smi to observe the real-time state of a GPU, including its memory usage. This is a crucial reactive skill for debugging and health monitoring. Today, we shift from observing to predicting.

Your goal for this lesson is to learn how to precisely calculate the GPU memory consumed by the KV cache. This is a fundamental capacity planning skill for any LLM Ops Engineer. Instead of discovering you're out of memory when a user sends a long prompt, you'll be able to calculate your memory limits beforehand. This allows you to provision resources effectively, set appropriate usage limits, and understand the cost-performance trade-offs inherent in your system's configuration.

1. Recapping the KV Cache

As a brief refresher, the KV cache is the mechanism that makes autoregressive generation efficient. Instead of re-calculating the Key (K) and Value (V) vectors for all previous tokens at every step, we store them in GPU memory. This turns a compute-intensive operation into a memory-bound one.

This diagram illustrates the prefill and decode phases. During prefill (left), K and V vectors are computed for the input prompt and stored in the cache. During the decode phase (right), only the K and V vectors for the newest token are computed; the rest are loaded directly from the cache, saving significant computation.

The trade-off is memory. As the sequence length and number of concurrent requests (batch size) grow, so does the cache. Your job is to quantify exactly how much.

2. The KV Cache Memory Formula

The memory required for the KV cache isn't a mystery; it can be calculated with a straightforward formula based on the model's architecture and the inference configuration. Understanding each component of this formula is key.

KV Cache Memory: Calculating GPU Requirements for LLM ...

This article by Michael Brenndoerfer provides a detailed derivation of the KV cache memory formula and explains each component. It's an excellent resource for understanding the underlying mechanics.

Please read the following sections: "Anatomy of KV Cache Storage": Focus on understanding the shape of the cached tensors: [batch_size, num_kv_heads, seq_len, head_dim]. "Cache Size Calculation": This is the core of the lesson. Study the full formula and the simplified version for Multi-Head Attention (MHA). Pay close attention to how Grouped-Query Attention (GQA) reduces the num_kv_heads and thus the cache size. "Key Parameters": This provides a quick summary of all the variables you need for the calculation.

After completing the reading, the full formula for calculating the KV cache size in bytes should be clear:

KV Cache (bytes) = 2 \times B \times S \times L \times H_{k v} \times D_{h} \times P

Where:

2: Accounts for storing two separate matrices, Keys (K) and Values (V).
B: The batch_size, or the number of concurrent sequences being processed.
S: The sequence_length (or context window), which is the total number of tokens (prompt + generated) in the sequence.
L: The number of transformer layers in the model.
$H_{k v}$ : The number of Key/Value heads. For models with Grouped-Query Attention (GQA) or Multi-Query Attention (MQA), this is less than the total number of attention heads. This is a critical factor in modern model efficiency.
$D_{h}$ : The head_dimension. This is the size of each attention head.
P: The precision in bytes per parameter (e.g., 2 for FP16/BF16, 4 for FP32, 1 for INT8).

Your experience with infrastructure capacity planning for applications is directly applicable here. Just as you'd calculate memory requirements based on heap size, thread count, and cache sizes for a traditional service, here we are doing the same for a model based on its specific architectural parameters.

3. A Practical Calculation

Let's apply this formula to a real-world model. We'll use Llama-3-8B. Finding the architectural specifications is the first step. You can usually find these in the model card, the official paper, or technical blogs.

This table from a VMware blog post is a great example of the kind of data you'll need.

LLM Inference Sizing and Performance Guidance - VMware Blogs

This resource provides a handy table with specifications for popular LLMs.

Focus on the table "LLM_Specification" and locate the row for Llama-3 8B. Note the values for n_layers, n_heads, d_model, and n_kv_heads.

From the table, we get the following for Llama-3-8B:

n_layers (L): 32
n_heads: 32
d_model: 4096
n_kv_heads ( $H_{k v}$ ): 8 (Note: This is an example of GQA, as $H_{k v} < n_{h e a d s}$ )
head_dim ( $D_{h}$ ): This is not always listed directly. It's calculated as d_model / n_heads = 4096 / 32 = 128.

Now, let's calculate the KV cache size for a single sequence (batch_size = 1) with a sequence length of 8192 tokens in FP16 precision (bytes_per_element = 2).

KV Cache = 2 \times 1 \times 8192 \times 32 \times 8 \times 128 \times 2 bytes

KV Cache = 1, 073, 741, 824 bytes

Converting this to more familiar units:

1,073,741,824 bytes / 1024 = 1,048,576 KB
1,048,576 KB / 1024 = 1024 MB
1024 MB / 1024 = 1 GB

So, for every request that uses an 8192-token context window, you must reserve 1 GB of VRAM just for its KV cache.

A Useful Heuristic: Memory per Token

For quick estimations, it's often useful to calculate the memory cost per token. You can do this by removing the sequence length from the main formula.

Memory per Token = 2 \times L \times H_{k v} \times D_{h} \times P

For Llama-3-8B:

Memory per Token = 2 \times 32 \times 8 \times 128 \times 2 = 131, 072 bytes/token = 128 KB/token

With this heuristic, the calculation becomes simple: 128 KB/token * 8192 tokens = 1024 MB = 1 GB.

4. Impact on Capacity Planning

This calculation is the foundation of LLM capacity planning. The total memory footprint on a GPU is not just the model's weights. A more complete formula is:

Total VRAM \approx Model Weights + KV Cache + Activation Memory + Overhead

Let's put this into practice.

Scenario: You have an NVIDIA A100 GPU with 80 GB of VRAM. You are serving the Llama-3-8B model (FP16) and want to know the maximum number of concurrent requests (batch_size) you can support if the average sequence_length is 4096 tokens.

Model Weight Memory:
- Llama-3-8B has ~8 billion parameters.
- In FP16 (2 bytes/parameter), this is 8 * 10^9 * 2 = 16 * 10^9 bytes ≈ 16 GB.
KV Cache Memory per Request:
- Using our per-token heuristic: 128 KB/token * 4096 tokens = 512 MB = 0.5 GB.
Activation & Overhead Memory:
- This is memory for intermediate calculations. It's harder to predict exactly, but a common rule of thumb is to reserve 10-20% of the model weight size, plus a small fixed amount for the CUDA context. Let's be conservative and estimate ~4 GB.
Putting it Together:
- Available memory for KV caches = Total VRAM - Weight Memory - Overhead
- Available memory = 80 GB - 16 GB - 4 GB = 60 GB.
- Max batch_size = Available Memory / Cache per Request
- Max batch_size = 60 GB / 0.5 GB/request = 120 requests.

This tells you that, under these conditions, your server can theoretically handle up to 120 concurrent requests. If you exceed this, you risk an out-of-memory (OOM) error, which would crash the inference process. You can now see how increasing sequence length or using a model with a less efficient architecture (e.g., standard MHA instead of GQA) would drastically reduce the number of users you can serve concurrently.

KV Cache Memory: Calculating GPU Requirements for LLM ...

The same article provides excellent code examples and visualizations for these concepts.

Please review the following sections: Skim the Python code under "Cache Size Calculation" to see how this logic translates to code. Read "Sequence Length Effects" and "Batch Size Effects". The plots in these sections clearly visualize the linear scaling and the critical "crossover point" where the cache can consume more memory than the model weights themselves.

Conclusion

In this lesson, you have moved from simply observing GPU memory usage to quantitatively predicting it. This is a critical step towards mastering the systems layer of LLM operations.

Here are the key takeaways:

The KV cache memory consumption can be calculated with the formula: $2 \times B \times S \times L \times H_{k v} \times D_{h} \times P$ .
Finding the model's architectural parameters (L, $H_{k v}$ , $D_{h}$ ) is the first step in any capacity planning exercise.
KV cache size scales linearly with both batch_size and sequence_length.
For long-context or high-throughput scenarios, the KV cache often becomes the dominant consumer of VRAM, surpassing even the model's own weight.
By estimating the total memory footprint, you can determine the maximum concurrent load a given GPU can handle, which is essential for setting admission control policies and provisioning infrastructure.

In the previous lesson, you learned how to monitor GPU state. In this lesson, you learned how to calculate the memory impact of batching. In our next lesson, we will complete the picture by empirically measuring the impact of batching on performance, analyzing the trade-offs between GPU throughput and latency with a hands-on benchmark script.