"""How much GPU memory a language model needs to answer chats: weights, attention cache, working memory.

Python standard library only. The model shapes come from public files saved in this folder:
llama32_1b_config.json, llama31_8b_config.json and llama31_70b_config.json are copies of the
config.json of each model on Hugging Face (read 2026-10-19); llama31_405b_shape.json holds the
numbers of Table 3 in the Llama 3 paper. Every number is a size in bytes unless its name says GiB.
"""
import json
from math import ceil
from pathlib import Path

HERE = Path(__file__).parent
GIB = 1024**3

# key -> (name on the page, saved file)
MODELS = {
    "1b": ("Llama 3.2 1B", "llama32_1b_config.json"),
    "8b": ("Llama 3.1 8B", "llama31_8b_config.json"),
    "70b": ("Llama 3.1 70B", "llama31_70b_config.json"),
    "405b": ("Llama 3.1 405B", "llama31_405b_shape.json"),
}
BITS = (16, 8, 4)                                   # bits stored per weight
CONTEXTS = (512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072)   # tokens per chat the page offers
CHATS = (1, 2, 4, 8, 16, 32, 64)                    # chats served at once the page offers
# Graphics-card sizes from the vendor pages (read 2026-10-19). The pages write GB; this code reads the label as GiB.
# The L40S is the exception: its 48 GB includes error-correcting memory (ECC), which takes 6.25% of it, and
# nvidia-smi shows 46,068 MiB with ECC on (public forum outputs, read 2026-10-19), so it is 45 GiB here.
GPUS = ((24, "RTX 4090"), (32, "RTX 5090"), (45, "L40S"), (80, "H100"), (141, "H200"))

FULL_BITS = 16            # the token table and output matrix stay at 16 bits in published 8-bit and 4-bit files
KV_VALUE_BYTES = 2        # the attention cache keeps each number in 16 bits, whatever the weights use
ACTIVATION_BYTES = 2      # the working tensors of a layer are 16-bit as well
ACTIVATION_TOKENS = 2048  # example: tokens the card works on in one pass
OVERHEAD_GIB = 1          # example allowance for the CUDA context, graph capture and fragmentation
USABLE = 0.92             # share of a card a server claims by default (vLLM gpu_memory_utilization)


def load_config(key: str) -> dict:
    return json.loads((HERE / MODELS[key][1]).read_text())


def head_dim(cfg: dict) -> int:
    return cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]


def param_count(cfg: dict) -> int:
    """Count the weights of a Llama-style model from its shape.

    Per layer: attention (query and output h x h, key and value h x kv_heads x head_dim), the
    gated feed-forward block (three h x ffn matrices) and two norm vectors of length h. Around
    the layers: the token table (vocab x h), the final norm, and an output matrix of the same
    size as the table unless the model reuses the table ("tied")."""
    h, ffn = cfg["hidden_size"], cfg["intermediate_size"]
    kv_width = cfg["num_key_value_heads"] * head_dim(cfg)
    attention = 2 * h * h + 2 * h * kv_width
    per_layer = attention + 3 * h * ffn + 2 * h
    return cfg["num_hidden_layers"] * per_layer + embedding_params(cfg) + h


def embedding_params(cfg: dict) -> int:
    """The token table (vocab x h) and, unless the model reuses the table ("tied"), an output matrix of the same size."""
    return cfg["vocab_size"] * cfg["hidden_size"] * (1 if cfg.get("tie_word_embeddings") else 2)


def weight_bytes(cfg: dict, bits: int) -> int:
    """The embeddings stay at 16 bits (counted once when tied); every other weight uses `bits`.
    Scale values that real 4-bit files add are not counted."""
    if bits not in BITS:
        raise ValueError(f"bits must be one of {BITS}")
    embeddings = embedding_params(cfg)
    return embeddings * FULL_BITS // 8 + (param_count(cfg) - embeddings) * bits // 8


def kv_bytes_per_token(cfg: dict) -> int:
    """One key and one value (the 2) for every layer, key/value head and head slot."""
    return 2 * cfg["num_hidden_layers"] * cfg["num_key_value_heads"] * head_dim(cfg) * KV_VALUE_BYTES


def kv_cache_bytes(cfg: dict, tokens: int, chats: int) -> int:
    return kv_bytes_per_token(cfg) * tokens * chats


def activation_bytes(cfg: dict, tokens: int = ACTIVATION_TOKENS) -> int:
    """Rough estimate, not a measurement: the biggest tensors one layer holds while it works on
    `tokens` tokens are the feed-forward block's two wide outputs and their product (3 x ffn per
    token) and a few hidden-size copies (4 x h per token). Real servers measure this at start-up."""
    return tokens * (3 * cfg["intermediate_size"] + 4 * cfg["hidden_size"]) * ACTIVATION_BYTES


def overhead_bytes() -> int:
    return OVERHEAD_GIB * GIB


def memory_need(key: str, bits: int, tokens: int, chats: int) -> dict:
    cfg = load_config(key)
    parts = {
        "weights": weight_bytes(cfg, bits),
        "kv_cache": kv_cache_bytes(cfg, tokens, chats),
        "activations": activation_bytes(cfg),
        "overhead": overhead_bytes(),
    }
    parts["total"] = sum(parts.values())
    return parts


def usable_bytes(gpu_gib: int) -> int:
    return int(gpu_gib * GIB * USABLE)


def fits(total: int, gpu_gib: int) -> bool:
    return total <= usable_bytes(gpu_gib)


def cards_needed(total: int, gpu_gib: int) -> int:
    """Cards of this size if the load were split evenly; splitting adds costs this code ignores."""
    return ceil(total / usable_bytes(gpu_gib))


def chats_that_fit(key: str, bits: int, tokens: int, gpu_gib: int) -> int:
    """How many chats of `tokens` tokens each fit beside the weights and fixed parts on one card."""
    base = memory_need(key, bits, tokens, 0)["total"]
    return max(0, (usable_bytes(gpu_gib) - base) // kv_cache_bytes(load_config(key), tokens, 1))