kv_cache.py

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class KVCache:
    keys: torch.Tensor | None = None
    values: torch.Tensor | None = None
    max_seq_len: int | None = None
    seq_lens: torch.Tensor | None = None

    @classmethod
    def allocate(
        cls,
        *,
        batch_size: int,
        num_heads: int,
        max_seq_len: int,
        head_dim: int,
        device: torch.device | str,
        dtype: torch.dtype,
    ) -> KVCache:
        device = torch.device(device)
        keys = torch.empty((batch_size, num_heads, max_seq_len, head_dim), device=device, dtype=dtype)
        values = torch.empty((batch_size, num_heads, max_seq_len, head_dim), device=device, dtype=dtype)
        seq_lens = torch.zeros(batch_size, dtype=torch.int32, device=device)
        return cls(keys=keys, values=values, max_seq_len=max_seq_len, seq_lens=seq_lens)

    def update(self, new_keys: torch.Tensor, new_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.keys is None or self.values is None:
            self.keys = new_keys
            self.values = new_values
            return self.keys, self.values

        self._validate_compatible(new_keys, new_values)
        if self.is_preallocated:
            start = self._current_uniform_seq_len()
            end = start + new_keys.size(-2)
            if self.max_seq_len is not None and end > self.max_seq_len:
                raise ValueError(
                    f"KV cache capacity {self.max_seq_len} is too small for requested sequence length {end}."
                )
            self.keys[:, :, start:end, :] = new_keys
            self.values[:, :, start:end, :] = new_values
            self.seq_lens = self.seq_lens + new_keys.size(-2)
            return self.current()

        self.keys = torch.cat((self.keys, new_keys), dim=-2)
        self.values = torch.cat((self.values, new_values), dim=-2)
        return self.keys, self.values

    def append(self, new_keys: torch.Tensor, new_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.keys is None or self.values is None:
            return new_keys, new_values
        self._validate_compatible(new_keys, new_values)
        if self.is_preallocated:
            current_keys, current_values = self.current()
            return (
                torch.cat((current_keys, new_keys), dim=-2),
                torch.cat((current_values, new_values), dim=-2),
            )
        return torch.cat((self.keys, new_keys), dim=-2), torch.cat((self.values, new_values), dim=-2)

    def current(self) -> tuple[torch.Tensor, torch.Tensor]:
        if self.keys is None or self.values is None:
            raise ValueError("Cannot get current KV tensors from an empty cache.")
        if not self.is_preallocated:
            return self.keys, self.values
        seq_len = self._current_uniform_seq_len()
        return self.keys[:, :, :seq_len, :], self.values[:, :, :seq_len, :]

    def branch(self, batch_size: int | None = None) -> KVCache:
        if self.keys is None or self.values is None:
            return KVCache(max_seq_len=self.max_seq_len, seq_lens=None if self.seq_lens is None else self.seq_lens.clone())
        if not self.is_preallocated:
            keys = self.keys.clone()
            values = self.values.clone()
            if batch_size is not None and keys.size(0) == 1 and batch_size > 1:
                keys = keys.expand(batch_size, -1, -1, -1).clone()
                values = values.expand(batch_size, -1, -1, -1).clone()
            return KVCache(keys=keys, values=values, max_seq_len=self.max_seq_len)

        seq_len = self._current_uniform_seq_len()
        target_batch = batch_size or self.keys.size(0)
        branched = KVCache.allocate(
            batch_size=target_batch,
            num_heads=self.keys.size(1),
            max_seq_len=self.max_seq_len or self.keys.size(2),
            head_dim=self.keys.size(-1),
            device=self.keys.device,
            dtype=self.keys.dtype,
        )
        src_keys = self.keys[:, :, :seq_len, :]
        src_values = self.values[:, :, :seq_len, :]
        if src_keys.size(0) == 1 and target_batch > 1:
            src_keys = src_keys.expand(target_batch, -1, -1, -1)
            src_values = src_values.expand(target_batch, -1, -1, -1)
        branched.keys[:, :, :seq_len, :] = src_keys
        branched.values[:, :, :seq_len, :] = src_values
        branched.seq_lens.fill_(seq_len)
        return branched

    def reset(self) -> None:
        if self.is_preallocated and self.seq_lens is not None:
            self.seq_lens.zero_()
            return
        self.keys = None
        self.values = None

    @property
    def is_preallocated(self) -> bool:
        return self.seq_lens is not None and self.max_seq_len is not None

    def get_seq_len(self) -> int:
        if self.keys is None or self.values is None:
            return 0
        if self.is_preallocated:
            return self._current_uniform_seq_len()
        return self.keys.size(-2)

    def is_empty(self) -> bool:
        return self.get_seq_len() == 0

    def _validate_compatible(self, new_keys: torch.Tensor, new_values: torch.Tensor) -> None:
        if self.keys is None or self.values is None:
            return
        if self.keys.shape[:-2] != new_keys.shape[:-2] or self.keys.shape[-1] != new_keys.shape[-1]:
            raise ValueError(
                f"New keys shape {tuple(new_keys.shape)} is incompatible with cached keys shape {tuple(self.keys.shape)}."
            )
        if self.values.shape[:-2] != new_values.shape[:-2] or self.values.shape[-1] != new_values.shape[-1]:
            raise ValueError(
                f"New values shape {tuple(new_values.shape)} is incompatible with cached values shape {tuple(self.values.shape)}."
            )

    def _current_uniform_seq_len(self) -> int:
        if self.seq_lens is None:
            return self.get_seq_len()
        if self.seq_lens.numel() == 0:
            return 0
        if not torch.equal(self.seq_lens, self.seq_lens[:1].expand_as(self.seq_lens)):
            raise ValueError("KV cache currently assumes uniform sequence lengths across the batch.")
        return int(self.seq_lens[0].item())


def initialize_kv_caches(
    num_layers: int,
    *,
    batch_size: int | None = None,
    num_heads: int | None = None,
    max_seq_len: int | None = None,
    head_dim: int | None = None,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> list[KVCache]:
    if batch_size is None:
        return [KVCache() for _ in range(num_layers)]
    if num_heads is None or max_seq_len is None or head_dim is None or device is None or dtype is None:
        raise ValueError(
            "num_heads, max_seq_len, head_dim, device, and dtype must be provided when preallocating KV caches."
        )
    return [
        KVCache.allocate(
            batch_size=batch_size,
            num_heads=num_heads,
            max_seq_len=max_seq_len,
            head_dim=head_dim,
            device=device,
            dtype=dtype,
        )
        for _ in range(num_layers)
    ]