kv-cache-size 0.1.1

//! 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:
//!
//! ```text
//! 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);
//! ```

#![forbid(unsafe_code)]
#![deny(missing_docs)]
#![no_std]

use core::fmt;

/// Why a configuration or a query could not be evaluated.
///
/// Every variant is a refusal, never a silently approximated answer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KvError {
    /// A structural field of the model config was zero.
    ///
    /// A transformer with no layers, no key-value heads or a zero head dimension is
    /// not a model this formula describes.
    ZeroDimension(&'static str),
    /// `context_length` or `batch_size` was zero.
    ZeroWorkload(&'static str),
    /// The requested cache does not fit in a `u64` byte count.
    Overflow,
    /// A memory budget was negative or not finite.
    InvalidBudget,
}

impl fmt::Display for KvError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            KvError::ZeroDimension(field) => write!(f, "model config field `{field}` is zero"),
            KvError::ZeroWorkload(field) => write!(f, "workload field `{field}` is zero"),
            KvError::Overflow => f.write_str("KV cache size does not fit in u64 bytes"),
            KvError::InvalidBudget => f.write_str("memory budget must be finite and non-negative"),
        }
    }
}

/// The element type the cache is stored in.
///
/// `Int4` is half a byte per element, which is why element size is a `f64` rather
/// than an integer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KvPrecision {
    /// `bf16` or `fp16`: two bytes per element. The default for most runtimes.
    Bf16,
    /// `fp8` KV cache: one byte per element.
    Fp8,
    /// `int4` KV cache: half a byte per element.
    Int4,
}

impl KvPrecision {
    /// Bytes occupied by one cached element.
    #[must_use]
    pub fn bytes_per_element(self) -> f64 {
        match self {
            KvPrecision::Bf16 => 2.0,
            KvPrecision::Fp8 => 1.0,
            KvPrecision::Int4 => 0.5,
        }
    }
}

/// The three model-config fields the KV-cache formula actually needs.
///
/// These map onto `num_hidden_layers`, `num_key_value_heads` and `head_dim` in a
/// Hugging Face `config.json`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct KvCacheConfig {
    /// `num_hidden_layers` — every layer keeps its own K and V tensors.
    pub num_hidden_layers: u32,
    /// `num_key_value_heads` — the GQA key-value head count, not the query head count.
    pub num_key_value_heads: u32,
    /// `head_dim` — the per-head dimension of the key and value vectors.
    pub head_dim: u32,
}

impl KvCacheConfig {
    /// Builds a config, refusing any zero dimension.
    ///
    /// # Errors
    ///
    /// [`KvError::ZeroDimension`] naming the offending field.
    pub fn new(
        num_hidden_layers: u32,
        num_key_value_heads: u32,
        head_dim: u32,
    ) -> Result<Self, KvError> {
        if num_hidden_layers == 0 {
            return Err(KvError::ZeroDimension("num_hidden_layers"));
        }
        if num_key_value_heads == 0 {
            return Err(KvError::ZeroDimension("num_key_value_heads"));
        }
        if head_dim == 0 {
            return Err(KvError::ZeroDimension("head_dim"));
        }
        Ok(Self {
            num_hidden_layers,
            num_key_value_heads,
            head_dim,
        })
    }

    /// Bytes the cache grows by for one additional token in one sequence.
    ///
    /// This is the number that decides whether cache size or weight size dominates,
    /// and it is independent of context length and batch size.
    #[must_use]
    pub fn bytes_per_token(self, precision: KvPrecision) -> f64 {
        2.0 * precision.bytes_per_element()
            * f64::from(self.num_hidden_layers)
            * f64::from(self.num_key_value_heads)
            * f64::from(self.head_dim)
    }

    /// Total cache bytes for `context_length` tokens across `batch_size` sequences.
    ///
    /// # Errors
    ///
    /// [`KvError::ZeroWorkload`] if either argument is zero — a zero-token or
    /// zero-sequence workload is a caller mistake, not a zero-byte answer.
    pub fn total_bytes(
        self,
        precision: KvPrecision,
        context_length: u64,
        batch_size: u32,
    ) -> Result<f64, KvError> {
        if context_length == 0 {
            return Err(KvError::ZeroWorkload("context_length"));
        }
        if batch_size == 0 {
            return Err(KvError::ZeroWorkload("batch_size"));
        }
        Ok(self.bytes_per_token(precision) * context_length as f64 * f64::from(batch_size))
    }

    /// The same total, rounded down to whole bytes.
    ///
    /// # Errors
    ///
    /// As [`Self::total_bytes`], plus [`KvError::Overflow`] when the result exceeds
    /// `u64::MAX` bytes.
    pub fn total_bytes_u64(
        self,
        precision: KvPrecision,
        context_length: u64,
        batch_size: u32,
    ) -> Result<u64, KvError> {
        let bytes = self.total_bytes(precision, context_length, batch_size)?;
        if !bytes.is_finite() || bytes >= u64::MAX as f64 {
            return Err(KvError::Overflow);
        }
        Ok(bytes as u64)
    }

    /// The longest context that fits `budget_bytes`, at this precision and batch size.
    ///
    /// Rounds down: the returned length is guaranteed to fit. Returns `0` when not
    /// even one token fits, which is a real answer and not an error.
    ///
    /// # Errors
    ///
    /// [`KvError::InvalidBudget`] for a negative or non-finite budget, and
    /// [`KvError::ZeroWorkload`] for a zero batch size.
    pub fn max_context(
        self,
        precision: KvPrecision,
        budget_bytes: f64,
        batch_size: u32,
    ) -> Result<u64, KvError> {
        if !budget_bytes.is_finite() || budget_bytes < 0.0 {
            return Err(KvError::InvalidBudget);
        }
        if batch_size == 0 {
            return Err(KvError::ZeroWorkload("batch_size"));
        }
        let per_token = self.bytes_per_token(precision) * f64::from(batch_size);
        let tokens = budget_bytes / per_token;
        if tokens >= u64::MAX as f64 {
            return Err(KvError::Overflow);
        }
        // `floor` on a non-negative finite quotient: the result always fits.
        Ok(tokens as u64)
    }

    /// The factor by which using `num_attention_heads` instead of
    /// `num_key_value_heads` would overstate the cache.
    ///
    /// This is the GQA group size. It is `1.0` for a model without GQA, and the
    /// full head count for multi-query attention. Returns `None` if
    /// `num_attention_heads` is zero.
    #[must_use]
    pub fn gqa_overstatement(self, num_attention_heads: u32) -> Option<f64> {
        if num_attention_heads == 0 {
            return None;
        }
        Some(f64::from(num_attention_heads) / f64::from(self.num_key_value_heads))
    }
}

/// Derives `head_dim` from `hidden_size / num_attention_heads`.
///
/// Use this **only** when the model config does not publish `head_dim`. When the
/// config publishes it, the published value wins: the two disagree on real configs,
/// and this function has no way to know that.
///
/// Returns `None` when `num_attention_heads` is zero or does not divide
/// `hidden_size` exactly — a non-integer head dimension means the fallback does not
/// describe this architecture, and guessing is worse than refusing.
#[must_use]
pub fn head_dim_from_hidden(hidden_size: u32, num_attention_heads: u32) -> Option<u32> {
    if num_attention_heads == 0 || hidden_size == 0 {
        return None;
    }
    if hidden_size % num_attention_heads != 0 {
        return None;
    }
    Some(hidden_size / num_attention_heads)
}

/// Bytes in one gibibyte, for turning a VRAM figure into a budget.
pub const GIB: f64 = 1024.0 * 1024.0 * 1024.0;

/// Bytes in one mebibyte.
pub const MIB: f64 = 1024.0 * 1024.0;

#[cfg(test)]
mod tests {
    use super::*;

    /// Llama 3.1 8B: 32 layers, 8 KV heads, head_dim 128.
    fn llama31_8b() -> KvCacheConfig {
        KvCacheConfig::new(32, 8, 128).unwrap()
    }

    #[test]
    fn llama31_8b_is_128_kib_per_token() {
        // 2 * 2 * 32 * 8 * 128 = 131072 bytes = 128 KiB.
        assert_eq!(llama31_8b().bytes_per_token(KvPrecision::Bf16), 131_072.0);
    }

    #[test]
    fn qwen25_7b_is_56_kib_per_token() {
        // 28 layers, 4 KV heads, head_dim 128 -> 2 * 2 * 28 * 4 * 128 = 57344.
        let cfg = KvCacheConfig::new(28, 4, 128).unwrap();
        assert_eq!(cfg.bytes_per_token(KvPrecision::Bf16), 57_344.0);
    }

    #[test]
    fn precision_scales_linearly() {
        let cfg = llama31_8b();
        let bf16 = cfg.bytes_per_token(KvPrecision::Bf16);
        assert_eq!(cfg.bytes_per_token(KvPrecision::Fp8), bf16 / 2.0);
        assert_eq!(cfg.bytes_per_token(KvPrecision::Int4), bf16 / 4.0);
    }

    #[test]
    fn eight_k_context_is_one_gib() {
        let cfg = llama31_8b();
        assert_eq!(cfg.total_bytes(KvPrecision::Bf16, 8192, 1).unwrap(), GIB);
        assert_eq!(
            cfg.total_bytes_u64(KvPrecision::Bf16, 8192, 1).unwrap(),
            1_073_741_824
        );
    }

    #[test]
    fn batch_multiplies_the_total() {
        let cfg = llama31_8b();
        let one = cfg.total_bytes(KvPrecision::Bf16, 4096, 1).unwrap();
        let four = cfg.total_bytes(KvPrecision::Bf16, 4096, 4).unwrap();
        assert_eq!(four, one * 4.0);
    }

    #[test]
    fn max_context_round_trips_against_total_bytes() {
        let cfg = llama31_8b();
        // 16 GiB at fp8, batch 4: 16 GiB / (65_536 bytes/token * 4) = 65_536 tokens.
        let tokens = cfg.max_context(KvPrecision::Fp8, 16.0 * GIB, 4).unwrap();
        assert_eq!(tokens, 65_536);
        // The answer fits, and one more token does not.
        assert!(cfg.total_bytes(KvPrecision::Fp8, tokens, 4).unwrap() <= 16.0 * GIB);
        assert!(cfg.total_bytes(KvPrecision::Fp8, tokens + 1, 4).unwrap() > 16.0 * GIB);
    }

    #[test]
    fn max_context_rounds_down_and_can_be_zero() {
        let cfg = llama31_8b();
        // Half a token's worth of budget fits no tokens at all.
        let half = cfg.bytes_per_token(KvPrecision::Bf16) / 2.0;
        assert_eq!(cfg.max_context(KvPrecision::Bf16, half, 1).unwrap(), 0);
        // One and a half tokens' worth fits exactly one.
        assert_eq!(cfg.max_context(KvPrecision::Bf16, half * 3.0, 1).unwrap(), 1);
    }

    #[test]
    fn gqa_overstatement_is_the_group_size() {
        // Llama 3.1 8B has 32 query heads against 8 KV heads.
        assert_eq!(llama31_8b().gqa_overstatement(32), Some(4.0));
        // No GQA: query heads equal KV heads.
        let mha = KvCacheConfig::new(32, 32, 128).unwrap();
        assert_eq!(mha.gqa_overstatement(32), Some(1.0));
        assert_eq!(mha.gqa_overstatement(0), None);
    }

    #[test]
    fn zero_dimensions_are_refused_by_name() {
        assert_eq!(
            KvCacheConfig::new(0, 8, 128),
            Err(KvError::ZeroDimension("num_hidden_layers"))
        );
        assert_eq!(
            KvCacheConfig::new(32, 0, 128),
            Err(KvError::ZeroDimension("num_key_value_heads"))
        );
        assert_eq!(
            KvCacheConfig::new(32, 8, 0),
            Err(KvError::ZeroDimension("head_dim"))
        );
    }

    #[test]
    fn zero_workload_is_refused_not_zeroed() {
        let cfg = llama31_8b();
        assert_eq!(
            cfg.total_bytes(KvPrecision::Bf16, 0, 1),
            Err(KvError::ZeroWorkload("context_length"))
        );
        assert_eq!(
            cfg.total_bytes(KvPrecision::Bf16, 1024, 0),
            Err(KvError::ZeroWorkload("batch_size"))
        );
    }

    #[test]
    fn invalid_budgets_are_refused() {
        let cfg = llama31_8b();
        assert_eq!(
            cfg.max_context(KvPrecision::Bf16, -1.0, 1),
            Err(KvError::InvalidBudget)
        );
        assert_eq!(
            cfg.max_context(KvPrecision::Bf16, f64::NAN, 1),
            Err(KvError::InvalidBudget)
        );
        assert_eq!(
            cfg.max_context(KvPrecision::Bf16, f64::INFINITY, 1),
            Err(KvError::InvalidBudget)
        );
    }

    #[test]
    fn head_dim_fallback_refuses_non_integer_results() {
        assert_eq!(head_dim_from_hidden(4096, 32), Some(128));
        assert_eq!(head_dim_from_hidden(4096, 0), None);
        assert_eq!(head_dim_from_hidden(0, 32), None);
        // 4096 / 33 is not an integer: refuse rather than round.
        assert_eq!(head_dim_from_hidden(4096, 33), None);
    }

    #[test]
    fn overflow_is_reported_not_wrapped() {
        let cfg = KvCacheConfig::new(u32::MAX, u32::MAX, u32::MAX).unwrap();
        assert_eq!(
            cfg.total_bytes_u64(KvPrecision::Bf16, u64::MAX, u32::MAX),
            Err(KvError::Overflow)
        );
    }
}