VRAM Calculator — Python source

"""VRAM Calculator — estimate the VRAM an LLM needs (weights + KV cache).

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the VRAM Calculator tool,
          ported from src/lib/vramCalculator.ts (the canonical TypeScript
          implementation).
Tool page: https://dev.cosmolabs.org/tools/vram-calculator
License:  display source — part of CosmoDev's polyglot tool pages.

Sizes use decimal gigabytes (GB = 10**9 bytes). Two components are
estimated: weights (params x bytes-per-param) and the KV cache
(2 x layers x context x kv_heads x head_dim x batch x bytes per KV
element). Activations and CUDA-context overhead are not modeled — treat
the total as a floor and leave headroom on the card.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple

#: Every supported quantization id, in the order the UI lists them.
QUANTS: Tuple[str, ...] = ("fp32", "fp16", "bf16", "int8", "int4", "q4_K_M")

#: Quantization table: bytes stored per weight for each format
#: (q4_K_M = 4.85 bits/weight, the llama.cpp mix).
BYTES_PER_PARAM: Dict[str, float] = {
    "fp32": 4.0,
    "fp16": 2.0,
    "bf16": 2.0,
    "int8": 1.0,
    "int4": 0.5,
    "q4_K_M": 4.85 / 8.0,
}

#: Default attention architecture: a modern GQA-style layout (32 layers,
#: 8 KV heads, 128-dim heads). Override per model.
DEFAULT_LAYERS = 32
DEFAULT_KV_HEADS = 8
DEFAULT_HEAD_DIM = 128

#: Bytes per KV-cache element (fp16 K and V tensors) unless overridden.
DEFAULT_KV_BYTES = 2

#: A decimal gigabyte.
GB = 1e9


def _require_finite_min(name: str, value: float, minimum: float) -> float:
    """Return ``value``; raise ValueError unless it is finite and >= minimum."""
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{name} must be a finite number ≥ {minimum} (got {value!r})")
    if not math.isfinite(value) or value < minimum:
        raise ValueError(f"{name} must be a finite number ≥ {minimum} (got {value})")
    return float(value)


@dataclass(frozen=True)
class VramBreakdown:
    """The estimated VRAM footprint of a model."""

    quant: str
    bytes_per_param: float
    weights_gb: float
    kv_cache_gb: float
    total_gb: float


def vram(
    params_b: float,
    quant: str,
    context: float,
    *,
    layers: Optional[int] = None,
    kv_heads: Optional[int] = None,
    head_dim: Optional[int] = None,
    batch: Optional[int] = None,
    kv_bytes: Optional[int] = None,
) -> VramBreakdown:
    """Estimate the VRAM footprint of a model: weights plus KV cache.

    ::

        weights_gb  = params_b * bytes_per_param
        kv_cache_gb = 2 * layers * context * kv_heads * head_dim * kv_bytes * batch / 1e9

    Raises ValueError on params_b <= 0, an unknown quantization, negative
    context, or any option below 1 (context 0 is allowed — no context, no
    cache).
    """
    if isinstance(params_b, bool) or not isinstance(params_b, (int, float)):
        raise ValueError(f"paramsB must be a finite number > 0 (got {params_b!r})")
    if not math.isfinite(params_b) or params_b <= 0:
        raise ValueError(f"paramsB must be a finite number > 0 (got {params_b})")
    try:
        bytes_per_param = BYTES_PER_PARAM[quant]
    except KeyError:
        expected = ", ".join(QUANTS)
        raise ValueError(f"unknown quantization {quant!r} — expected one of {expected}") from None

    _require_finite_min("context", context, 0)
    n_layers = _require_finite_min("layers", DEFAULT_LAYERS if layers is None else layers, 1)
    n_kv_heads = _require_finite_min("kvHeads", DEFAULT_KV_HEADS if kv_heads is None else kv_heads, 1)
    n_head_dim = _require_finite_min("headDim", DEFAULT_HEAD_DIM if head_dim is None else head_dim, 1)
    n_batch = _require_finite_min("batch", 1 if batch is None else batch, 1)
    n_kv_bytes = _require_finite_min("kvBytes", DEFAULT_KV_BYTES if kv_bytes is None else kv_bytes, 1)

    weights_gb = params_b * 1e9 * bytes_per_param / GB
    kv_cache_gb = 2 * n_layers * context * n_kv_heads * n_head_dim * n_kv_bytes * n_batch / GB
    return VramBreakdown(
        quant=quant,
        bytes_per_param=bytes_per_param,
        weights_gb=weights_gb,
        kv_cache_gb=kv_cache_gb,
        total_gb=weights_gb + kv_cache_gb,
    )


@dataclass(frozen=True)
class GpuCard:
    """One GPU memory tier."""

    name: str
    size_gb: float


#: Common GPU memory tiers, from consumer boards to datacenter cards.
GPU_CARDS: Tuple[GpuCard, ...] = (
    GpuCard("RTX 3060 Ti / RTX 4060 / RX 7600", 8),
    GpuCard("RTX 3060 12 GB / RTX 4070", 12),
    GpuCard("RTX 4060 Ti 16 GB / RTX 5080", 16),
    GpuCard("RTX 3090 / RTX 4090", 24),
    GpuCard("RTX A6000 / L40S", 48),
    GpuCard("A100 80 GB / H100 / H200", 80),
)


@dataclass(frozen=True)
class GpuFit:
    """One card scored against a total footprint."""

    name: str
    size_gb: float
    fits: bool
    headroom_gb: float


def gpu_fits(
    total_gb: float,
    cards: Optional[Sequence[GpuCard]] = None,
) -> List[GpuFit]:
    """Score every card against a total footprint.

    ``fits`` is inclusive: a total exactly equal to the card size fits
    (headroom 0). Pass a custom ``cards`` list to score other tiers.
    """
    _require_finite_min("totalGB", total_gb, 0)
    tiers = GPU_CARDS if cards is None else tuple(cards)
    return [
        GpuFit(
            name=c.name,
            size_gb=c.size_gb,
            fits=(c.size_gb - total_gb) >= 0,
            headroom_gb=c.size_gb - total_gb,
        )
        for c in tiers
    ]