Demystifying KV Cache Math: Why Llama 3.1 8B Needs 16 GiB for 128K Tokens

The model fit. That was the part that confused me initially.

Llama 3.1 8B in bf16 takes up about 16 GB of weights. My consumer GPU had 24 GB of VRAM. I happily set max_model_len=131072 because the model card claimed support for 128K context, hit run, and vLLM immediately refused to start. The error message pointed to the maximum sequence length being larger than the number of tokens that could be stored in the KV cache.

I had 8 GB of headroom. How was that not enough for a single request?

It turns out the KV cache math for Llama 3.1 8B states that a single 128K-token sequence needs 16 GiB on its own. Not the weights. Just the cache. This post covers the exact arithmetic required before picking context limits, which applies universally to every transformer you self-host or manage via scalable infrastructure providers like n1n.ai.

The Core Formula for KV Cache Size

The memory consumption of the Key-Value cache is deterministic and follows a straightforward formula:

KV cache size per token = 2 * layers * KV heads * head_dim * bytes per value

For Llama 3.1 8B in bf16 format, the numbers look like this:

• Layers: 32
• KV Heads: 8 (thanks to Grouped-Query Attention)
• Head Dimension: 128
• Bytes per value: 2 (bf16)

Plugging these into the equation: 2 * 32 * 8 * 128 * 2 = 131,072 bytes, which is exactly 128 KiB per token.

When multiplying 128K tokens by 128 KiB, you get 16 GiB for one full-length sequence. Combined with the ~16 GB of weights, a 24 GB card cannot hold both. vLLM checks this at startup: if a sequence of max_model_len cannot fit into the remaining KV memory after weights are loaded, it safely refuses to run rather than crashing mid-inference.

Why the KV Cache Exists and Its Cost

The KV cache stores the key and value vectors for every past token across every layer, preventing the model from recomputing them during autoregressive generation. Without it, generating token 10,001 would require re-running attention projections for the previous 10,000 tokens.

Picture it as 32 shelves (layers), each containing two rows (Key and Value). Every new token drops a tensor box onto every shelf in both rows. Nothing is removed until the request finishes.

Here is a quick Python snippet to compute this programmatically in your infrastructure pipeline:

import json

def kv_bytes_per_token(config_path, bytes_per_value=2):
    with open(config_path) as f:
        cfg = json.load(f)
    layers = cfg["num_hidden_layers"]
    kv_heads = cfg.get("num_key_value_heads", cfg["num_attention_heads"])
    head_dim = cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    return 2 * layers * kv_heads * head_dim * bytes_per_value

per_tok = kv_bytes_per_token("Llama-3.1-8B/config.json")
print(f"Bytes per token: {per_tok}")                      # 131072
print(f"GiB for 128K: {per_tok * 131072 / 2**30}")     # 16.0 GiB

Context Length vs. Memory Scale

The scaling is strictly linear. Double the context length, and you double the cache size:

When deploying models at scale, managing these allocations manually across disparate GPU clusters can become a bottleneck. Many development teams bypass infrastructure overhead entirely by integrating high-performance APIs through n1n.ai, ensuring optimal resource pooling and low-latency response times without managing raw VRAM allocation.

Grouped-Query Attention (GQA) Saves the Day

Llama 3.1 8B utilizes Grouped-Query Attention. It has 32 query heads but only 8 KV heads, meaning groups of 4 query heads share one set of K and V vectors.

If we compare this to older architectures like Llama 2 7B which used classic Multi-Head Attention (MHA) with 32 KV heads:

• Llama 2 7B (MHA): 2 * 32 * 32 * 128 * 2 = 512 KiB per token
• Llama 3.1 8B (GQA): 2 * 32 * 8 * 128 * 2 = 128 KiB per token

Without GQA, a 128K-token sequence on an 8B model would consume an astronomical 64 GiB of cache. GQA is the primary engineering breakthrough that makes long-context windows viable on commodity hardware.

Four Practical Levers for Optimization

When you hit memory limits during self-hosting, you have four distinct levers to pull:

1. Realistic Context Caps: Set max_model_len to what your traffic actually demands. If your application prompts rarely exceed 20K tokens, advertising 128K is purely vanity.
2. FP8 KV Cache: Store the cache in fp8 format. This cuts the memory footprint per token in half (from 128 KiB to 64 KiB), doubling your token pool capacity.
3. Weight Quantization: Use 4-bit AWQ or GPTQ quantization for the model weights. Every gigabyte saved from the model weights goes directly into expanding the KV cache pool.
4. Tune GPU Memory Utilization: Safely adjust parameters like gpu_memory_utilization=0.92 to capture residual headroom, provided background processes aren't starved.

Here is a production-ready vLLM configuration example:

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    max_model_len=32768,
    kv_cache_dtype="fp8",
    gpu_memory_utilization=0.92,
)

Alternatively, if you prefer avoiding the operational burdens of cluster sizing, weight quantization, and cache tuning, you can leverage fully managed, high-speed model endpoints directly via n1n.ai.

Get a free API key at n1n.ai