# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/inference/kv_cache.py
# ─────────────────────────────────────────────────────────

"""
Per-loop-step KV cache for AXIOM-1 inference.

AXIOM-1 uses a SEPARATE KV cache per loop step, not one global cache.
This is because each loop iteration processes the same sequence positions
with different hidden states, so KV projections differ at each step.

This is fundamentally different from standard transformer KV caching
where each layer has its own cache.
"""

import torch


class PerStepKVCache:
    """
    Per-loop-step KV cache for efficient autoregressive generation.

    Maintains separate key and value caches for each loop step.
    During generation, only the new token's KV projections are computed
    and appended to the cache for each step.

    Usage:
        cache = PerStepKVCache(num_steps=12, n_kv_heads=2, head_dim=64)
        # During generation:
        for step in range(num_loops):
            k, v = cache.update(step, new_k, new_v)
            # k, v now contain full sequence history for this step
    """

    def __init__(
        self,
        num_steps: int = 12,
        n_kv_heads: int = 2,
        head_dim: int = 64,
        max_seq_len: int = 2048,
        device: str = 'cuda',
        dtype: torch.dtype = torch.bfloat16
    ) -> None:
        """
        Args:
            num_steps: Number of loop steps (= num_loops)
            n_kv_heads: Number of KV heads in GQA
            head_dim: Dimension per head
            max_seq_len: Maximum sequence length
            device: Device for cache tensors
            dtype: Data type for cache tensors (bfloat16 — never float16)
        """
        self.num_steps = num_steps
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.device = device
        self.dtype = dtype

        # Cache storage: one (k, v) pair per loop step
        # Each is initially None and grows as tokens are generated
        self.k_cache: list[torch.Tensor | None] = [None] * num_steps
        self.v_cache: list[torch.Tensor | None] = [None] * num_steps

    def update(
        self,
        step_idx: int,
        k: torch.Tensor,
        v: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Append new k, v to the cache for a given loop step.

        Args:
            step_idx: Loop iteration index (0 to num_steps-1)
            k: New key tensor (batch, n_kv_heads, new_len, head_dim)
            v: New value tensor (batch, n_kv_heads, new_len, head_dim)

        Returns:
            Tuple of (full_k, full_v) with complete sequence history
        """
        if self.k_cache[step_idx] is None:
            self.k_cache[step_idx] = k
            self.v_cache[step_idx] = v
        else:
            self.k_cache[step_idx] = torch.cat(
                [self.k_cache[step_idx], k], dim=2
            )
            self.v_cache[step_idx] = torch.cat(
                [self.v_cache[step_idx], v], dim=2
            )

        return self.k_cache[step_idx], self.v_cache[step_idx]

    def reset(self) -> None:
        """Clear all cached keys and values."""
        self.k_cache = [None] * self.num_steps
        self.v_cache = [None] * self.num_steps

    @property
    def seq_len(self) -> int:
        """Current cached sequence length (from first step's cache)."""
        if self.k_cache[0] is None:
            return 0
        return self.k_cache[0].shape[2]

    def memory_usage_bytes(self) -> int:
        """
        Report total memory consumed by all cached KV tensors.

        Returns:
            Total bytes across all loop steps' K and V caches.
        """
        total = 0
        for k, v in zip(self.k_cache, self.v_cache):
            if k is not None:
                total += k.nelement() * k.element_size()
            if v is not None:
                total += v.nelement() * v.element_size()
        return total

    def trim(self, max_seq_len: int) -> None:
        """
        Discard oldest tokens if cache exceeds max_seq_len.

        Keeps only the most recent max_seq_len tokens in each
        step's cache. Useful for long generation runs where
        memory must be bounded.

        Args:
            max_seq_len: Maximum number of tokens to keep in cache.
        """
        for i in range(self.num_steps):
            if self.k_cache[i] is not None and self.k_cache[i].shape[2] > max_seq_len:
                self.k_cache[i] = self.k_cache[i][:, :, -max_seq_len:, :]
                self.v_cache[i] = self.v_cache[i][:, :, -max_seq_len:, :]
