KV Cache Size Calculation

Introduction

Welcome back. In our last lesson, we dissected the forward pass of a transformer block, revealing a crucial dichotomy: the prefill stage is often compute-bound, while the autoregressive decode stage is almost always memory-bound. The reason for the latter is the need to read the entire, ever-growing Key-Value (KV) cache from GPU memory at every single token generation step.

This lesson focuses squarely on that cache. You'll learn how to answer the question: "Exactly how much VRAM is my KV cache consuming?" This is the third and final pillar of our VRAM arithmetic, following model weights and activations.

Our learning outcome is to calculate the total KV cache size for a given model architecture, batch size, and sequence length. Mastering this calculation is essential for system design, capacity planning, and diagnosing memory-related performance issues in any LLM serving environment.

Why the KV Cache Exists

Before we calculate its size, it's vital to understand why the KV cache is a non-negotiable component of efficient LLM inference. During autoregressive generation, the model predicts one token at a time, appends it to the input sequence, and repeats the process. A naive implementation would recompute the attention states for all previous tokens in every step. This is incredibly wasteful.

The KV cache solves this by storing the Key (K) and Value (V) vectors for each token once they are computed. For subsequent steps, the model only needs to compute the K and V vectors for the new token and append them to the cache, avoiding redundant calculations.

To build a strong intuition for this process, let's watch a detailed walkthrough.

As the video explains, this caching mechanism is what transforms the computational complexity of generating a new token from being quadratic with respect to the sequence length to being linear. This is a massive performance gain, but it comes at the cost of memory.

The KV Cache Size Formula

Now, let's quantify that memory cost. The size of the KV cache is a direct function of the model's architecture and the inference job's parameters.

The formula to calculate the total KV cache size in bytes is:

Let's break down each component:

• 2: Represents the two matrices we are caching: Keys (K) and Values (V).
• B (Batch Size): The number of sequences being processed in parallel.
• S (Sequence Length): The maximum number of tokens in the context window (prompt + generated tokens).
• L (Number of Layers): The number of decoder blocks in the model.
• H_kv (Number of KV Heads): The number of attention heads for the Key and Value projections. For standard Multi-Head Attention (MHA), this equals the number of query heads. For optimizations like Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), this number is smaller, significantly reducing the cache size.
• D_head (Head Dimension): The dimension of each attention head's vector.
• P (Precision): The number of bytes per element. Common values are 4 for fp32, 2 for fp16/bf16, and 1 for INT8.

Note that the product H_kv × D_head gives the total size of the KV projection per layer. For a standard model, this equals the hidden dimension d_model.

For a clear presentation of this formula and its components, the following resource is very helpful.

From Formula to Implementation

Given your practical focus, let's translate this formula directly into code. This allows you to programmatically estimate VRAM requirements for different models and scenarios.

The following resource provides a Python function that does exactly this. It's a perfect example of turning theory into a practical engineering tool.

Here is the core logic from that resource, isolated for clarity:

def estimate_kv_cache_gb(
    batch_size: int,
    sequence_length: int,
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    bytes_per_element: int = 2,  # Default to 2 for fp16/bf16
) -> float:
    """Estimates the KV cache size in Gigabytes (GB)."""
    



    # The '2' is for storing both Keys and Values
    kv_cache_bytes = (
        2 * 
        batch_size * 
        sequence_length * 
        num_layers * 
        num_kv_heads * 
        head_dim * 
        bytes_per_element
    )
    



    # Convert bytes to Gigabytes (using 1024^3 for GiB)
    kv_cache_gb = kv_cache_bytes / (1024**3)
    
    return round(kv_cache_gb, 2)

Worked Example: Llama 3 8B

Let's apply this to a real model, Llama 3 8B, which uses Grouped-Query Attention (GQA).

• L (Layers): 32
• H_kv (KV Heads): 8
• D_head (Head Dimension): 128
• P (Precision): 2 bytes (for bfloat16)

Scenario: A single user request (B=1) with a maximum context length of S=8192.

# Llama 3 8B Parameters
L = 32
H_kv = 8
D_head = 128
P = 2 # bf16




# Scenario Parameters
B = 1
S = 8192




# Calculation
kv_cache_gb = estimate_kv_cache_gb(
    batch_size=B,
    sequence_length=S,
    num_layers=L,
    num_kv_heads=H_kv,
    head_dim=D_head,
    bytes_per_element=P
)

print(f"KV Cache for Llama 3 8B with {S} context: {kv_cache_gb} GB")



# Expected Output: KV Cache for Llama 3 8B with 8192 context: 1.0 GB

As you can see, even for an 8B model, the KV cache for a single 8K sequence consumes a full Gigabyte of VRAM. If you were serving a batch of 16 such users, this would grow to 16 GB—a figure that already exceeds the capacity of many consumer GPUs, just for the KV cache alone. This linear scaling with both batch size and sequence length is why the KV cache is a dominant factor in memory planning.

Self-Check Exercise

Now it's your turn. Let's calculate the KV cache size for a larger model under a small batch load.

Model: Llama 3 70B

• L (Layers): 80
• H_kv (KV Heads): 8
• D_head (Head Dimension): 128

Scenario:

• B (Batch Size): 4
• S (Sequence Length): 4096
• P (Precision): 2 bytes (bfloat16)

Use the formula or the Python function to calculate the required VRAM for the KV cache in GB.

...

kv_cache_70b = estimate_kv_cache_gb(
    batch_size=4,
    sequence_length=4096,
    num_layers=80,
    num_kv_heads=8,
    head_dim=128,
    bytes_per_element=2
)
print(f"KV Cache for 70B model: {kv_cache_70b} GB")

Conclusion

In this lesson, we demystified the KV cache, the single largest consumer of dynamic VRAM during LLM inference.

Key Takeaways:

• The KV cache is a performance optimization that avoids recomputing attention states for the entire context at every step, making autoregressive decoding feasible.
• This performance gain comes at a significant memory cost, which can often exceed the memory required for the model weights themselves, especially with long contexts or large batches.
• The size of the KV cache can be calculated precisely with the formula: 2 * B * S * L * H_kv * D_head * P.
• The cache size scales linearly with batch size (B) and sequence length (S), making it the primary bottleneck for achieving high throughput and supporting long contexts.

You can now calculate all three major components of VRAM usage during inference: model weights, activations, and the KV cache. This comprehensive understanding is the foundation of LLM systems engineering.

Preview of the Next Lesson:

We have now completed our analysis of the fundamental memory requirements of a transformer model. However, just knowing the size isn't enough; we need to explain why a 70B model can feel so large and unwieldy in practice. In the next lesson, we will synthesize everything we've learned by applying VRAM and KV cache calculations to a 70B model to explain its large memory footprint. This will solidify your mental model before we move on to optimization techniques.