kv_cache_manager.py

"""
Gestione KV cache sulla P40.

La KV cache risiede sulla P40 (24 GB) perché:
- È troppo grande per la RTX 3050 (~7.4 GB disponibili)
- Non va trasferita ad ogni step (overhead PCIe)
- Viene aggiornata incrementalmente

Funzionalità:
- Allocazione della KV cache sulla P40
- Aggiornamento incrementale (append)
- Lettura per il decoding
- Sharding per evitare OOM
- Gestione della sequenza dinamica
"""

from __future__ import annotations

import logging
from typing import Optional, Tuple

import torch

from config import RouterConfig

logger = logging.getLogger("cuda_router.kv_cache")


class KvCacheBlock:
    """Singolo blocco KV cache per un layer."""

    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        max_seq_len: int,
        device: torch.device,
        dtype: torch.dtype = torch.float16,
    ):
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.device = device
        self.dtype = dtype

        # Key e Value cache — shape: [num_heads, head_dim, max_seq_len]
        self.key_cache = torch.zeros(
            (num_heads, head_dim, max_seq_len),
            dtype=dtype,
            device=device,
        )
        self.value_cache = torch.zeros(
            (num_heads, head_dim, max_seq_len),
            dtype=dtype,
            device=device,
        )

        # Posizione corrente nella cache
        self.current_len = 0

    def update(
        self,
        new_keys: torch.Tensor,
        new_values: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Aggiorna la cache con nuovi key/value e restituisce la cache completa.

        Args:
            new_keys: Shape [num_heads, head_dim, new_seq_len]
            new_values: Shape [num_heads, head_dim, new_seq_len]

        Returns:
            Tuple (key_cache, value_cache) completi
        """
        start = self.current_len
        end = start + new_keys.shape[-1]

        if end > self.max_seq_len:
            raise RuntimeError(
                f"KV cache overflow: layer richiede {end} positions, "
                f"max è {self.max_seq_len}"
            )

        self.key_cache[:, :, start:end] = new_keys
        self.value_cache[:, :, start:end] = new_values
        self.current_len = end

        return self.key_cache[:, :, :end], self.value_cache[:, :, :end]

    def get(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """Restituisce la cache corrente."""
        return self.key_cache[:, :, :self.current_len], self.value_cache[:, :, :self.current_len]

    def reset(self) -> None:
        """Resetta la cache."""
        self.current_len = 0
        self.key_cache.zero_()
        self.value_cache.zero_()

    @property
    def is_full(self) -> bool:
        return self.current_len >= self.max_seq_len

    @property
    def remaining(self) -> int:
        return self.max_seq_len - self.current_len


class KvCacheManager:
    """
    Gestisce la KV cache per tutti i layer del modello.

    La cache risiede sulla P40 (weights_device) per massimizzare la VRAM
    disponibile e minimizzare il trasferimento PCIe.
    """

    def __init__(
        self,
        num_layers: int,
        num_heads: int,
        head_dim: int,
        max_seq_len: int,
        config: RouterConfig,
        dtype: torch.dtype = torch.float16,
    ):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.config = config
        self.dtype = dtype
        self.cache_device = torch.device(config.kv_cache.cache_device)

        # Blocco KV cache per ogni layer
        self.blocks: list[KvCacheBlock] = []

        # Alloca la cache
        self._allocate()

        logger.info(
            "KV cache allocata: %d layer × %d heads × %d dim × %d seq len su %s",
            num_layers,
            num_heads,
            head_dim,
            max_seq_len,
            self.cache_device,
        )

        # Calcola memoria stimata
        bytes_per_block = num_heads * head_dim * max_seq_len * 2  # float16 = 2 bytes
        total_bytes = bytes_per_block * num_layers
        total_gb = total_bytes / 1024**3
        logger.info("Memoria KV cache stimata: %.2f GB", total_gb)

    def _allocate(self) -> None:
        """Alloca tutti i blocchi KV cache sulla P40."""
        with torch.cuda.device(self.cache_device):
            for i in range(self.num_layers):
                block = KvCacheBlock(
                    num_heads=self.num_heads,
                    head_dim=self.head_dim,
                    max_seq_len=self.max_seq_len,
                    device=self.cache_device,
                    dtype=self.dtype,
                )
                self.blocks.append(block)
                logger.debug("Layer %d KV cache allocata: %.2f GB", i, bytes_per_block / 1024**3)

    def update(
        self,
        layer_idx: int,
        new_keys: torch.Tensor,
        new_values: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Aggiorna la KV cache per un layer specifico.

        Args:
            layer_idx: Indice del layer.
            new_keys: Shape [batch, num_heads, seq_len, head_dim] o [num_heads, head_dim, seq_len]
            new_values: Shape [batch, num_heads, seq_len, head_dim] o [num_heads, head_dim, seq_len]

        Returns:
            Tuple (keys, values) completi per il layer
        """
        if layer_idx < 0 or layer_idx >= self.num_layers:
            raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})")

        # Trasforma shape se necessario: [batch, heads, seq, dim] → [heads, dim, seq]
        if new_keys.dim() == 4:
            # [batch, num_heads, seq_len, head_dim] → [num_heads, head_dim, seq_len]
            new_keys = new_keys.squeeze(0).transpose(1, 2)
            new_values = new_values.squeeze(0).transpose(1, 2)

        return self.blocks[layer_idx].update(new_keys, new_values)

    def get(
        self,
        layer_idx: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Recupera la KV cache per un layer.

        Args:
            layer_idx: Indice del layer.

        Returns:
            Tuple (keys, values) per il layer
        """
        if layer_idx < 0 or layer_idx >= self.num_layers:
            raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})")

        return self.blocks[layer_idx].get()

    def get_all(self) -> list[Tuple[torch.Tensor, torch.Tensor]]:
        """Recupera la KV cache per tutti i layer."""
        return [block.get() for block in self.blocks]

    def reset(self, layer_idx: Optional[int] = None) -> None:
        """
        Resetta la KV cache.

        Args:
            layer_idx: Se None, resetta tutti i layer. Altrimenti solo quel layer.
        """
        if layer_idx is None:
            for block in self.blocks:
                block.reset()
            logger.debug("Tutta la KV cache resettata")
        else:
            self.blocks[layer_idx].reset()
            logger.debug("KV cache layer %d resettata", layer_idx)

    def get_sequence_length(self, layer_idx: Optional[int] = None) -> int:
        """Restituisce la lunghezza corrente della sequenza."""
        if layer_idx is None:
            # Usa il primo layer come riferimento (tutti dovrebbero avere la stessa length)
            return self.blocks[0].current_len if self.blocks else 0
        return self.blocks[layer_idx].current_len

    def get_remaining_slots(self, layer_idx: Optional[int] = None) -> int:
        """Restituisce gli slot rimanenti nella cache."""
        if layer_idx is None:
            return min(block.remaining for block in self.blocks) if self.blocks else 0
        return self.blocks[layer_idx].remaining

    def is_full(self, layer_idx: Optional[int] = None) -> bool:
        """Verifica se la cache è piena."""
        if layer_idx is None:
            return any(block.is_full for block in self.blocks)
        return self.blocks[layer_idx].is_full

    def get_memory_usage(self) -> dict:
        """Restituisce le statistiche di memoria della KV cache."""
        total_used = 0
        total_max = 0
        for block in self.blocks:
            total_used += block.current_len
            total_max += block.max_seq_len

        return {
            "total_layers": self.num_layers,
            "current_seq_len": self.get_sequence_length(),
            "max_seq_len": self.max_seq_len,
            "total_slots_used": total_used,
            "total_slots_max": total_max * self.num_layers,
            "cache_device": str(self.cache_device),
            "dtype": str(self.dtype),
        }

    def __repr__(self) -> str:
        stats = self.get_memory_usage()
        return (
            f"KvCacheManager(layers={self.num_layers}, "
            f"seq={stats['current_seq_len']}/{stats['max_seq_len']}, "
            f"device={stats['cache_device']})"
        )