Crate kv_cache_size 

Exact KV-cache arithmetic for transformer inference.

The cache a decoder keeps while generating is one key tensor and one value tensor per layer, per key-value head:

kv_bytes = 2 x bytes_per_element x num_hidden_layers x num_key_value_heads
           x head_dim x context_length x batch_size

Two terms in that product are the ones that go wrong in practice.

The head count is KvCacheConfig::num_key_value_heads, not the number of attention (query) heads. Grouped-query attention keeps an intermediate number of key-value heads — more than one, fewer than the query heads (Ainslie et al., GQA: Training Generalized Multi-Query Transformer Checkpoints, arXiv:2305.13245) — so substituting the query-head count overstates the cache by the whole GQA group size.

The other is head_dim. Read it from the model config when the config publishes it. head_dim_from_hidden exists for the configs that genuinely omit it, and is kept off the main path so that the fallback is visible where it is used.

This crate sizes the KV cache and nothing else — weights, activations, runtime context and allocator fragmentation are excluded, so treat the result as a floor. An interactive version with per-model configs is at https://ml0x.com/calculators/kv-cache-size-calculator.html. For the whole budget — weights, this cache, optimizer state and an activation estimate, sized from each model’s published head count — see https://ml0x.com/calculators/llm-memory-calculator.html.

Example

use kv_cache_size::{KvCacheConfig, KvPrecision};

// Llama 3.1 8B: 32 layers, 8 key-value heads, head_dim 128.
let cfg = KvCacheConfig::new(32, 8, 128).unwrap();
assert_eq!(cfg.bytes_per_token(KvPrecision::Bf16), 131_072.0);

Structs

KvCacheConfig

The three model-config fields the KV-cache formula actually needs.

Enums

KvError

Why a configuration or a query could not be evaluated.

KvPrecision

The element type the cache is stored in.

Constants

GIB

Bytes in one gibibyte, for turning a VRAM figure into a budget.

MIB

Bytes in one mebibyte.

Functions

head_dim_from_hidden

Derives head_dim from hidden_size / num_attention_heads.