kv-cache-size 0.1.1

Exact KV-cache arithmetic for transformer inference: bytes per token, total cache bytes, and the context length that fits a memory budget.
Documentation
  • Coverage
  • 100%
    24 out of 24 items documented1 out of 12 items with examples
  • Size
  • Source code size: 8.24 MB This is the summed size of all the files inside the crates.io package for this release.
  • Documentation size: 316.43 kB This is the summed size of all files generated by rustdoc for all configured targets
  • Ø build duration
  • this release: 2s Average build duration of successful builds.
  • all releases: 2s Average build duration of successful builds in releases after 2024-10-23.
  • Links
  • Homepage
  • Documentation
  • theluckystrike/kv-cache-size
    0 0 0
  • crates.io
  • Dependencies
    • —
  • Versions
  • Owners
  • theluckystrike

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.