Estimate KV-cache memory before choosing a context length

Weights are only one part of the budget

A model can have the same weights while its generation-memory requirements change with the request. A KV cache retains attention keys and values for reuse during autoregressive inference. As more tokens are retained, a dense growing cache needs more storage. The model checkpoint size alone therefore cannot answer a context-length or concurrency question. [1]

This guide estimates the payload of a conventional full-attention cache. It does not load a model or measure GPU usage. The distinction matters: an exact multiplication can still be a poor estimate of total runtime memory if its assumptions do not match the serving implementation.

References: [1]

Count the values you intend to store

Assume every layer retains both keys and values, with equal head dimensions, a uniform cache dtype, and the same retained token count in each batch element. For L layers, Hkv stored KV heads, D values per head, T retained tokens, B sequences and S bytes per value, the payload is 2 × L × Hkv × D × T × B × S.

The factor of two counts keys and values. The remaining dimensions follow the documented per-layer cache tensor shape. Treat the formula as a storage-counting derivation under these stated assumptions. Do not substitute a parameter count or a model-file size for any of its inputs. [1]

References: [1]

Use KV heads, not query heads

Some attention configurations share keys and values across groups of query heads. Transformers exposes a num_key_value_heads setting for this distinction; the Llama configuration documentation describes its relation to multi-head, multi-query and grouped-query attention. Inspect the actual model configuration and cache implementation before choosing the stored head count. [3]

Our illustrative configuration uses 32 layers, eight KV heads, head dimension 128, 4,096 retained tokens, one sequence and two bytes per value. These are explicit example inputs, not the specification of a named checkpoint. That avoids implying that all similarly sized models share a memory profile.

References: [3]

Run the arithmetic without a model download

Run examples/kv_memory.py with Python 3.12.14. It uses the standard library and needs neither a GPU nor a model account. Assertions check the base result, three scaling cases, and rejection of an invalid token count.

The base payload is 536,870,912 bytes, or 0.5 GiB. GiB here means 2^30 bytes. Doubling retained context gives 1 GiB; batching four equal-length sequences gives 2 GiB; replacing eight stored KV heads with 32 gives 2 GiB. These are calculated payloads, not observed allocations.

Executed on CPU · Python 3.12.14 · 21 October 2026

"""Dense full-attention KV payload arithmetic, not a GPU measurement."""
def kv_bytes(layers, kv_heads, head_dim, tokens, batch=1, bytes_per_value=2):
    values = (layers, kv_heads, head_dim, tokens, batch, bytes_per_value)
    if any(type(value) is not int or value <= 0 for value in values):
        raise ValueError('All inputs must be positive integers')
    return 2 * layers * kv_heads * head_dim * tokens * batch * bytes_per_value

base = kv_bytes(32, 8, 128, 4096)
assert base == 536_870_912
assert kv_bytes(32, 8, 128, 8192) == 2 * base
assert kv_bytes(32, 8, 128, 4096, batch=4) == 4 * base
assert kv_bytes(32, 32, 128, 4096) == 4 * base
print('base GiB:', base / 2**30)
print('double context GiB:', kv_bytes(32, 8, 128, 8192) / 2**30)
print('batch four GiB:', kv_bytes(32, 8, 128, 4096, batch=4) / 2**30)
print('32 KV heads GiB:', kv_bytes(32, 32, 128, 4096) / 2**30)
try:
    kv_bytes(32, 8, 128, 0)
except ValueError:
    print('invalid input rejected')
else:
    raise AssertionError('Expected invalid input rejection')

Recorded output

base GiB: 0.5
double context GiB: 1.0
batch four GiB: 2.0
32 KV heads GiB: 2.0
invalid input rejected

Turn the estimate into a question for the server

The table holds all unmentioned inputs at the example baseline. It is useful for checking which configuration change caused a budget to move. It should not be presented as a comparison of inference engines or as a benchmark.

Where the simple estimate stops

Actual cache strategies can reserve storage differently. Static caches preallocate capacity; sliding-window or chunked layers can limit retained lengths; quantized caches change the storage representation and can carry additional bookkeeping. A single full-attention formula should not silently stand in for those cases. [2]

The estimate also excludes weights, temporary activations, allocation overhead and other runtime storage. Padded batches, cache sharing and device placement need implementation-specific inspection. Measure the intended workload on the intended stack before calling a configuration safe to fit. No GPU-fit or latency claim follows from this calculator.

References: [2]

Sources

1. Hugging Face Transformers: how caching works
2. Hugging Face Transformers: cache strategies
3. Hugging Face Transformers: Llama configuration