"""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,
    n_full_attn_layers: int = 0,
    head_dim_override: int = 0,
) -> float:
    """Total KV cache size in bytes for the whole context window.

    Hybrid-attention models interleave linear/recurrent layers (no O(n_ctx) KV)
    with full-attention layers; pass ``n_full_attn_layers`` to charge KV only to
    the full-attention layers (+ MTP heads). 0 ⇒ all layers (old behavior).
    ``head_dim_override`` handles archs where head_dim ≠ n_embd/n_head (Qwen3.6
    sets head_dim=256 while n_embd/n_head = 213). 0 ⇒ derive n_embd//n_head.
    """
    head_dim = head_dim_override or _head_dim(n_embd, n_head)
    per_layer = n_ctx * 2 * n_head_kv * head_dim * cache_dtype_bytes(cache_dtype)
    layers = (n_full_attn_layers or 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,
    n_ubatch: int | None = None,
    n_vocab: int = 0,
    n_full_attn_layers: int = 0,
    head_dim_override: int = 0,
) -> float:
    """Working/compute scratch memory in bytes.

    Two distinct buffers (llama.cpp discussion #6328):

    - **Physical compute tile** (``-ub`` / ``--ubatch-size``): the activations that
      ride the GPU during one forward pass. Scales with the *physical* batch, not the
      logical one. When FA is off with a quantized cache, also includes the per-tile
      f32 dequantization scratch for K and V across all layers.
    - **Logits/embeddings buffer** (``-b`` / ``--batch-size``): the *logical* max
      batch — a separate buffer of logits sized ``n_batch * vocab``. Without a known
      vocab (``n_vocab``) this term is omitted (conservative; the activation term
      dominates and the logits buffer is small relative to weights + KV at scale).

    Constraint enforced by the caller: ``batch >= ubatch``. Here ``n_ubatch``
    defaults to ``n_batch`` for backward compatibility with older callers that
    passed only ``n_batch``.
    """
    if compute_dtype not in COMPUTE_DTYPE_BYTES:
        raise ValueError(f"Unknown compute dtype: {compute_dtype!r}")
    cb = COMPUTE_DTYPE_BYTES[compute_dtype]

    ub = n_batch if n_ubatch is None else n_ubatch

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

    # activations: physical tile * n_embd * compute_dtype (a couple of buffers)
    activations = ub * n_embd * cb * 2

    scratch = 0.0
    if not flash_attn and cache_dtype not in ("f16", "bf16", "f32"):
        # f32 dequant buffer for one physical tile across all layers, K and V.
        # Linear/recurrent layers have no KV so contribute no dequant scratch —
        # the same `layers` (full-attn + MTP) count applies.
        scratch = layers * ub * 2 * n_head_kv * head_dim * 4.0

    # logits buffer: logical batch * vocab * compute_dtype. Only when vocab known.
    logits = 0.0
    if n_vocab > 0:
        logits = n_batch * n_vocab * cb
    return activations + scratch + logits