"""kv-cache-arithmetic.py -- KV-cache bytes per token, per request, and batch capacity from a config.json.

What it does: takes the attention geometry of two open models exactly as published in their
Hugging Face config.json files (read 2026-09-05):
  Qwen/Qwen3-8B:  36 layers, 32 Q heads, 8 KV heads, head_dim 128, bf16, 8.2B parameters
  Qwen/Qwen3-32B: 64 layers, 64 Q heads, 8 KV heads, head_dim 128, bf16, 32.8B parameters
(model cards: 32,768 tokens native context, 131,072 with YaRN)
and computes: bytes of KV cache per token (2 * layers * kv_heads * head_dim * bytes_per_elem),
the cache for one request at several context lengths, the weight memory at bf16, and how many
concurrent requests of a given length fit in the memory left on an 80 GB device after weights.
Also shows what happens without grouped-query attention (kv_heads = attention heads) and with
an 8-bit KV cache, since those two levers are what change the answer.

All numbers are arithmetic on published shapes; no GPU is touched. Real servers reserve
extra memory for activations, fragmentation, and CUDA context, so treat the capacities as
upper bounds.

Inputs:  none (edit MODELS / DEVICE_GB)
Run:     python code/kv-cache-arithmetic.py
Output:  datasets/kv-cache-arithmetic.csv
Requires: Python 3.11+ standard library only.
"""
from __future__ import annotations

import csv

MODELS = {
    # from config.json: hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, head_dim, params (bn)
    "Qwen3-8B": {"layers": 36, "heads": 32, "kv_heads": 8, "head_dim": 128, "params_b": 8.2},
    "Qwen3-32B": {"layers": 64, "heads": 64, "kv_heads": 8, "head_dim": 128, "params_b": 32.8},
}
DEVICE_GB = 80.0
RESERVE_GB = 4.0          # activations, CUDA context, fragmentation; a placeholder you should measure
CONTEXTS = (4096, 8192, 32768, 131072)   # 32,768 native and 131,072 with YaRN per the model cards


def kv_bytes_per_token(m: dict, bytes_per_elem: int = 2, kv_heads: int | None = None) -> int:
    kvh = m["kv_heads"] if kv_heads is None else kv_heads
    return 2 * m["layers"] * kvh * m["head_dim"] * bytes_per_elem   # 2 = K and V


def main() -> None:
    rows = []
    for name, m in MODELS.items():
        weights_gb = m["params_b"] * 2 / 1.0  # bf16: 2 bytes per parameter, in GB (1e9 bytes)
        free_gb = DEVICE_GB - weights_gb - RESERVE_GB
        print(f"\n=== {name}: {m['layers']} layers, {m['heads']} heads, {m['kv_heads']} KV heads, head_dim {m['head_dim']} ===")
        print(f"weights bf16 ~{weights_gb:.1f} GB; free for KV on {DEVICE_GB:.0f} GB after {RESERVE_GB:.0f} GB reserve: {free_gb:.1f} GB")
        variants = {
            "bf16 KV, GQA (as shipped)": (2, None),
            "bf16 KV, no GQA (kv_heads = heads)": (2, m["heads"]),
            "int8 KV, GQA": (1, None),
        }
        for label, (bpe, kvh) in variants.items():
            per_tok = kv_bytes_per_token(m, bpe, kvh)
            print(f"  {label:36s} {per_tok / 1024:8.1f} KiB/token")
            for ctx in CONTEXTS:
                per_req_gb = per_tok * ctx / 1e9
                fits = int(free_gb // per_req_gb) if per_req_gb > 0 else 0
                print(f"      ctx {ctx:>6}: {per_req_gb:6.2f} GB per request -> {fits:>4} concurrent")
                rows.append({"model": name, "variant": label, "kv_bytes_per_token": per_tok, "context": ctx,
                             "kv_gb_per_request": round(per_req_gb, 3), "concurrent_requests_80gb": fits})
    with open("datasets/kv-cache-arithmetic.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()