KV Cache Memory › Practical Memory Planning

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

Book contents

Section 10 of 14 · 6 min read

KV Cache Memory › Practical Memory PlanningOne section · Sections in this chapter10 / 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

Before deploying a model, estimate memory requirements for your target configuration. This should be a routine step in any model deployment process, not an afterthought. The calculations are inexpensive and the alternative, discovering out-of-memory errors in production, is far more costly. Here is a practical workflow that combines all the components we have discussed.

The planning process involves four inputs you must know before you can make hardware decisions: the model architecture specifications (available in any model card or configuration file), the maximum context length your application requires, the batch size you need to achieve your throughput targets, and the data type you plan to use for inference. Once you have these, the calculation takes seconds.

In[25]:
Code
def can_fit_on_gpu(
    gpu_memory_gb: float,
    model_config: dict,
    batch_size: int,
    max_seq_len: int,
    dtype_bytes: int = 2,
    safety_margin: float = 0.9,  # Leave 10 percent buffer
) -> dict:
    """
    Check if a configuration fits in GPU memory.

    Returns feasibility and utilization details.
    """
    available = gpu_memory_gb * safety_margin

    mem = estimate_inference_memory(
        model_params=model_config["total_params"],
        num_layers=model_config["num_layers"],
        num_kv_heads=model_config["num_kv_heads"],
        head_dim=model_config["head_dim"],
        batch_size=batch_size,
        seq_len=max_seq_len,
        dtype_bytes=dtype_bytes,
    )

    fits = mem["total_gb"] <= available
    utilization = (mem["total_gb"] / available) * 100

    return {
        "fits": fits,
        "required_gb": mem["total_gb"],
        "available_gb": available,
        "utilization_pct": utilization,
    }
In[26]:
Code
# Test configurations on common GPUs
gpus = {"RTX 4090": 24, "A100 40GB": 40, "A100 80GB": 80, "H100 80GB": 80}

test_configs = [
    {"model": "LLaMA 7B", "batch": 1, "seq": 4096},
    {"model": "LLaMA 7B", "batch": 8, "seq": 4096},
    {"model": "LLaMA 7B", "batch": 26, "seq": 4096},
]

feasibility_results = []
for config in test_configs:
    model = models[config["model"]]
    config_results = []
    for gpu_name, gpu_mem in gpus.items():
        result = can_fit_on_gpu(gpu_mem, model, config["batch"], config["seq"])
        config_results.append((gpu_name, result))
    feasibility_results.append((config, config_results))
Out[27]:
Console
Deployment Feasibility Check
======================================================================

LLaMA 7B | batch=1 | seq=4096
--------------------------------------------------
  RTX 4090: OK (16.8/21.6 GB)
  A100 40GB: OK (16.8/36.0 GB)
  A100 80GB: OK (16.8/72.0 GB)
  H100 80GB: OK (16.8/72.0 GB)

LLaMA 7B | batch=8 | seq=4096
--------------------------------------------------
  RTX 4090: OOM (30.8/21.6 GB)
  A100 40GB: OK (30.8/36.0 GB)
  A100 80GB: OK (30.8/72.0 GB)
  H100 80GB: OK (30.8/72.0 GB)

LLaMA 7B | batch=26 | seq=4096
--------------------------------------------------
  RTX 4090: OOM (66.8/21.6 GB)
  A100 40GB: OOM (66.8/36.0 GB)
  A100 80GB: OK (66.8/72.0 GB)
  H100 80GB: OK (66.8/72.0 GB)
Out[28]:
Visualization
Heatmap showing memory utilization percentages for different model configurations across GPU types.
Heatmap showing GPU memory utilization across three LLaMA 7B configurations and four GPU types (RTX 4090, A100 40GB, A100 80GB, H100 80GB). Green cells indicate safe utilization under 90% of available memory, yellow indicates tight packing between 90-100%, and red indicates out-of-memory conditions where the configuration exceeds hardware capacity. Single-sequence inference at 4K context fits on all hardware, while larger batches require datacenter-grade GPUs.

These results illustrate the hardware requirements for different deployment scenarios. LLaMA 7B fits comfortably on all tested GPUs for single-sequence inference at 4K context, with even the 24 GB RTX 4090 providing adequate headroom. Scaling to batch size 8 pushes requirements beyond consumer hardware, requiring datacenter GPUs like the A100. These constraints directly inform deployment decisions: smaller models offer flexibility across hardware tiers, while larger batches demand specialized infrastructure.

The feasibility analysis demonstrates the practical constraints of different GPU configurations. For production deployment, understanding these memory limits is essential for selecting appropriate hardware. A team that does not work through these calculations before launch will discover the limits the hard way, during peak load when requests begin failing with out-of-memory errors.

Choosing a Data Type

The data type used for cache storage is an often-overlooked lever for reducing memory requirements. Every parameter in our formula is fixed by the model architecture or the deployment scenario, except for the bytes-per-element term, which you can control directly. Switching from FP32 to FP16 halves cache memory at no algorithmic cost. Switching from FP16 to INT8 halves it again, using quantization to represent each key and value with 8 bits instead of 16. INT4 quantization reduces it by another factor of 2.

The tradeoff is accuracy. FP32 provides full floating-point precision. FP16 and BF16 maintain good numerical range while halving memory. INT8 quantization introduces small rounding errors that research has shown are generally tolerable for KV cache values, since the attention mechanism is somewhat robust to small perturbations in the attended values. INT4 quantization is more aggressive and may affect output quality for some tasks, particularly those requiring precise recall of specific numerical values from long contexts.

The key insight is that the KV cache stores intermediate representations, not the raw training weights of the model. These representations are produced and consumed by the attention mechanism, which involves softmax normalization that provides some robustness to quantization noise. This makes KV cache quantization practically easier than weight quantization, where precision matters more directly for the model's learned computation.

The memory requirements change as follows:

  • bytes=4\text{bytes} = 4: FP32, full precision
  • bytes=2\text{bytes} = 2: FP16 or BF16, standard inference precision, recommended default
  • bytes=1\text{bytes} = 1: INT8, halves cache size with minimal quality impact in most scenarios
  • bytes=0.5\text{bytes} = 0.5: INT4, quarters cache size with some quality tradeoff

Practitioners serving large models at scale should treat INT8 KV cache as the sensible default and reserve FP16 for applications requiring maximum precision. The cache savings of INT8 over FP16 are large enough to support twice the batch size or twice the context length on the same hardware.