kv-cache-size

Exact KV-cache arithmetic for transformer inference. No dependencies, no allocation, no model downloads, no network.

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

The leading 2 is one K tensor plus one V tensor.

The two terms people get wrong

1. num_key_value_heads, not num_attention_heads. Grouped-query attention keeps fewer key/value heads than query heads (Ainslie et al., GQA, arXiv:2305.13245), so using the query-head count overstates the cache by the GQA group size — 4x on Llama 3.1 8B, 8x on Qwen2.5-7B.
2. head_dim as published, not derived. hidden_size / num_attention_heads is a habit that breaks on configs where the published head_dim disagrees with it. head_dim_from_hidden exists for configs that genuinely omit the field, and is deliberately separate from the main path so the fallback is visible at the call site.

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); // 128 KiB per token
assert_eq!(cfg.total_bytes(KvPrecision::Bf16, 8192, 1).unwrap(), 1_073_741_824.0);

// How much context fits in 16 GiB of spare VRAM at fp8, batch 4?
let budget = 16.0 * 1024.0 * 1024.0 * 1024.0;
assert_eq!(cfg.max_context(KvPrecision::Fp8, budget, 4).unwrap(), 65_536);

What this crate does not do

It sizes the KV cache only. Weights, activations, CUDA context and allocator fragmentation are not included, so the number is a floor for planning, not a capacity guarantee. Interactive version and per-model configs: https://ml0x.com/calculators/kv-cache-size-calculator.html.

Licensed under MIT OR Apache-2.0.