"""KV cache and compute-scratch memory math.

Based on llama.cpp's KV cache layout:

    KV = n_layer * n_ctx * 2 * n_head_kv * head_dim * sizeof(cache_dtype)

where head_dim = n_embd / n_head. The factor 2 covers both K and V. With MTP
heads (e.g. DeepSeek-V3 n_mtp=1) each head adds ~one extra layer's worth of
KV, so the layer count becomes (n_layer + n_mtp).
"""

from __future__ import annotations


# cache dtype -> bytes per element (effective, including scales for quantized)
CACHE_DTYPE_BYTES: dict[str, float] = {
    "f16": 2.0,
    "bf16": 2.0,
    "f32": 4.0,
    "q8_0": 1.0,
    "q8_1": 1.0625,
    "q4_0": 0.5,
    "q4_1": 0.5625,
    "q5_0": 0.625,
    "q5_1": 0.6875,
}

# compute dtype -> bytes per element
COMPUTE_DTYPE_BYTES: dict[str, float] = {
    "f16": 2.0,
    "bf16": 2.0,
    "f32": 4.0,
}


def cache_dtype_bytes(dtype: str) -> float:
    if dtype not in CACHE_DTYPE_BYTES:
        raise ValueError(f"Unknown cache dtype: {dtype!r}")
    return CACHE_DTYPE_BYTES[dtype]


def _head_dim(n_embd: int, n_head: int) -> int:
    if n_head <= 0:
        return 0
    return n_embd // n_head


def kv_cache_bytes(
    *,
    n_layer: int,
    n_embd: int,
    n_head: int,
    n_head_kv: int,
    n_ctx: int,
    cache_dtype: str,
    n_mtp: int = 0,
    flash_attn: bool = True,
) -> float:
    """Total KV cache size in bytes for the whole context window."""
    head_dim = _head_dim(n_embd, n_head)
    per_layer = n_ctx * 2 * n_head_kv * head_dim * cache_dtype_bytes(cache_dtype)
    layers = n_layer + max(0, n_mtp)
    base = per_layer * layers

    # When flash attention is off AND the cache is quantized, llama.cpp keeps a
    # dequantized f32 scratch for the active batch slice. We approximate the
    # worst-case scratch as one batch-sized slice in f32, per layer (K+V).
    if not flash_attn and cache_dtype not in ("f16", "bf16", "f32"):
        # scratch is per-batch, not per-ctx; caller passes batch separately via
        # compute_scratch_bytes. Here we leave it out to avoid double counting.
        pass
    return base


def compute_scratch_bytes(
    *,
    n_layer: int,
    n_embd: int,
    n_head: int,
    n_head_kv: int,
    n_batch: int,
    compute_dtype: str = "f16",
    cache_dtype: str = "f16",
    flash_attn: bool = True,
    n_mtp: int = 0,
) -> float:
    """Working/compute scratch memory in bytes.

    Includes the batch activations and, when FA is off with a quantized cache,
    the per-batch f32 dequantization scratch for K and V across all layers.
    """
    if compute_dtype not in COMPUTE_DTYPE_BYTES:
        raise ValueError(f"Unknown compute dtype: {compute_dtype!r}")
    cb = COMPUTE_DTYPE_BYTES[compute_dtype]

    head_dim = _head_dim(n_embd, n_head)
    layers = n_layer + max(0, n_mtp)

    # activations: roughly batch * n_embd * compute_dtype (a couple of buffers)
    activations = n_batch * n_embd * cb * 2

    scratch = 0.0
    if not flash_attn and cache_dtype not in ("f16", "bf16", "f32"):
        # f32 dequant buffer for one batch across all layers, K and V
        scratch = layers * n_batch * 2 * n_head_kv * head_dim * 4.0
    return activations + scratch