code/models/common/llm_runtime/paged_kv_cache.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Exclusive physical ownership for one paged KV-cache pool."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch

import ttnn
from models.common.llm_runtime.config import PagedKVCacheConfig


@dataclass(frozen=True)
class PagedKVCacheContext:
    """Read-only compile context containing borrowed tensor references."""

    config: PagedKVCacheConfig
    tensors: tuple[tuple[Any, Any], ...]
    cache_shapes: tuple[tuple[int, int, int, int], ...]
    per_layer_dtypes: tuple[ttnn.DataType, ...]


def torch_dtype_for_ttnn(dtype: ttnn.DataType) -> torch.dtype:
    """Return the documented torch storage/request surrogate for a TT dtype."""

    mapping = {
        ttnn.bfloat16: torch.bfloat16,
        ttnn.bfloat8_b: torch.bfloat16,
        ttnn.bfloat4_b: torch.bfloat16,
        ttnn.float32: torch.float32,
        ttnn.int32: torch.int32,
        ttnn.uint32: torch.uint32,
        ttnn.uint16: torch.uint16,
        ttnn.uint8: torch.uint8,
    }
    try:
        return mapping[dtype]
    except KeyError as error:
        raise ValueError(f"No torch compatibility mapping exists for TT dtype {dtype!r}") from error


class PagedKVCacheManager:
    """Own one model-bound paged KV-cache allocation from configure to release.

    ``Llama3Executor`` constructs the manager from model-owned layer metadata.
    vLLM may resolve the final block geometry once through `configure`, then
    `allocate` binds the physical tensors to the model. Program
    compilation borrows `bound_context`; request execution may present
    only the exact handle returned by `allocate`.

    The returned cache is a borrowed compatibility handle. The manager remains
    the only owner allowed to replace or deallocate its physical tensors.
    """

    def __init__(self, model: Any, config: PagedKVCacheConfig):
        self._model = model
        self._mesh_device, self._layer_specs, model_paged_configs = _model_contract(model)
        self._validate_static_model_contract(config, model_paged_configs)

        self._config = config
        self._state = "configured" if config.is_resolved() else "unresolved"
        self._configuration_replaced = False
        self._bound_cache: list[list[Any]] | None = None
        self._bound_context: PagedKVCacheContext | None = None
        self._owned_tensors: tuple[Any, ...] = ()
        self._release_in_progress = False

    # Public API

    @property
    def config(self) -> PagedKVCacheConfig:
        return self._config

    @property
    def bound_context(self) -> PagedKVCacheContext | None:
        """Return immutable metadata and borrowed tensor references for compile."""

        return self._bound_context

    @property
    def per_layer_dtypes(self) -> tuple[ttnn.DataType, ...]:
        return tuple(spec.dtype for spec in self._layer_specs)

    @property
    def cache_shapes(self) -> tuple[tuple[int, int, int, int], ...]:
        if self._config.num_blocks is None:
            return ()
        return tuple(
            (
                self._config.num_blocks,
                spec.local_kv_heads,
                self._config.block_size,
                spec.head_dim,
            )
            for spec in self._layer_specs
        )

    def configure(self, config: PagedKVCacheConfig) -> None:
        """Install one immutable resolved replacement before allocation."""

        if self._state in ("bound", "released"):
            raise RuntimeError(f"Cannot configure paged KV cache while manager is {self._state}")
        if self._config.is_resolved() or self._configuration_replaced:
            raise RuntimeError("Paged KV cache configuration can be resolved only once")
        if not config.is_resolved():
            raise ValueError("Replacement PagedKVCacheConfig must contain num_blocks")

        for field in ("dtype", "memory_config"):
            if getattr(config, field) != getattr(self._config, field):
                raise ValueError(f"Resolved PagedKVCacheConfig cannot replace {field}")

        _, _, model_paged_configs = _model_contract(self._model)
        self._validate_static_model_contract(config, model_paged_configs)
        self._config = config
        self._configuration_replaced = True
        self._state = "configured"

    def validate_vllm_cache_spec(
        self,
        *,
        block_size: int,
        dtype: torch.dtype,
        num_blocks: int | None = None,
    ) -> None:
        """Validate vLLM's torch-facing cache request against model-owned TT policy."""

        if block_size != self._config.block_size:
            raise ValueError(
                f"vLLM block_size {block_size} does not match configured block_size {self._config.block_size}"
            )
        if not isinstance(dtype, torch.dtype):
            raise TypeError(f"vLLM cache dtype must be torch.dtype, got {type(dtype).__name__}")

        mismatches = [
            (layer, spec.dtype, torch_dtype_for_ttnn(spec.dtype))
            for layer, spec in enumerate(self._layer_specs)
            if dtype != torch_dtype_for_ttnn(spec.dtype)
        ]
        if mismatches:
            details = ", ".join(
                f"layer {layer}: {device_dtype!r} expects {torch_dtype!r}"
                for layer, device_dtype, torch_dtype in mismatches
            )
            raise ValueError(f"vLLM cache dtype {dtype!r} is incompatible with model KV dtype ({details})")

        if num_blocks is not None:
            if not isinstance(num_blocks, int) or isinstance(num_blocks, bool) or num_blocks <= 0:
                raise ValueError("vLLM num_blocks must be a positive integer")
            if num_blocks > self._config.max_num_blocks:
                raise ValueError(
                    f"vLLM num_blocks ({num_blocks}) exceeds configured maximum ({self._config.max_num_blocks})"
                )

    def allocate(self) -> list[list[Any]]:
        """Allocate, bind, and return one borrowed physical cache handle."""

        if self._state == "unresolved":
            raise RuntimeError("Paged KV cache capacity must be resolved before allocation")
        if self._state == "bound":
            raise RuntimeError("Paged KV cache has already been allocated")
        if self._state == "released":
            raise RuntimeError("Paged KV cache manager is released and terminal")
        if self._owned_tensors:
            failures = self._deallocate_owned_tensors()
            if failures:
                raise RuntimeError(
                    f"Failed to finish cleaning {len(failures)} tensor(s) from a previous KV allocation"
                ) from failures[0]

        cache: list[list[Any]] = []
        allocated: list[Any] = []
        host_staging: dict[tuple[tuple[int, int, int, int], torch.dtype], torch.Tensor] = {}
        model_args = getattr(self._model, "model_args", None)
        cache_path_value = getattr(model_args, "model_cache_path", None)
        cache_path = Path(cache_path_value) if cache_path_value else None
        dtypes_by_shape: dict[tuple[int, int, int, int], set[ttnn.DataType]] = {}
        for shape, spec in zip(self.cache_shapes, self._layer_specs):
            dtypes_by_shape.setdefault(shape, set()).add(spec.dtype)
        try:
            for shape, spec in zip(self.cache_shapes, self._layer_specs):
                host_key = (shape, torch_dtype_for_ttnn(spec.dtype))
                host_tensor = host_staging.get(host_key)
                if host_tensor is None:
                    host_tensor = torch.zeros(shape, dtype=host_key[1])
                    host_staging[host_key] = host_tensor
                pair = []
                for kv in ("k", "v"):
                    cache_file_name = None
                    if cache_path is not None and len(dtypes_by_shape[shape]) == 1:
                        cache_file_name = cache_path / f"empty_{kv}cache_paged_attention{shape}"
                    tensor = ttnn.as_tensor(
                        host_tensor,
                        device=self._mesh_device,
                        mesh_mapper=ttnn.ReplicateTensorToMesh(self._mesh_device),
                        layout=ttnn.TILE_LAYOUT,
                        memory_config=self._config.memory_config,
                        dtype=spec.dtype,
                        cache_file_name=cache_file_name,
                    )
                    allocated.append(tensor)
                    pair.append(tensor)
                cache.append(pair)
        except BaseException as primary:
            self._owned_tensors = tuple(allocated)
            cleanup_failures = self._deallocate_owned_tensors(reverse=True)
            _attach_cleanup_failures(primary, cleanup_failures)
            raise

        self._owned_tensors = tuple(allocated)
        context = PagedKVCacheContext(
            config=self._config,
            tensors=tuple(tuple(pair) for pair in cache),
            cache_shapes=self.cache_shapes,
            per_layer_dtypes=self.per_layer_dtypes,
        )
        try:
            self._model.set_kv_cache(cache)
        except BaseException as primary:
            try:
                self._model.set_kv_cache(None)
            except BaseException as cleanup_error:
                # The model may retain a partial binding. Keep manager ownership
                # so a later release can retry unbinding before deallocation.
                self._bound_cache = cache
                self._bound_context = context
                self._state = "bound"
                _attach_cleanup_failures(primary, [cleanup_error])
            else:
                cleanup_failures = self._deallocate_owned_tensors(reverse=True)
                _attach_cleanup_failures(primary, cleanup_failures)
            raise

        self._bound_cache = cache
        self._bound_context = context
        self._state = "bound"
        return cache

    def validate_borrowed_handle(self, cache: Any) -> None:
        """Require the exact borrowed handle and unchanged tensor identities."""

        if self._state != "bound" or self._bound_cache is None:
            raise RuntimeError("Paged KV cache is not allocated and bound")
        if cache is not self._bound_cache:
            raise ValueError("Request KV cache is not the exact manager-owned borrowed handle")
        try:
            supplied_tensors = tuple(tensor for pair in cache for tensor in pair)
        except TypeError as error:
            raise ValueError("Request KV cache no longer contains the manager-owned K/V tensor pairs") from error
        if len(supplied_tensors) != len(self._owned_tensors) or any(
            supplied is not owned for supplied, owned in zip(supplied_tensors, self._owned_tensors)
        ):
            raise ValueError("Request KV cache no longer contains the exact manager-owned K/V tensors")

    def release(self) -> None:
        """Unbind then deallocate every owned tensor exactly once."""

        if self._state == "released":
            return
        if self._state == "bound" and not self._release_in_progress:
            # Never deallocate while the model still retains the installed handles.
            self._model.set_kv_cache(None)
            self._bound_cache = None
            self._bound_context = None
            self._release_in_progress = True

        failures = self._deallocate_owned_tensors()
        if failures:
            raise RuntimeError(f"Failed to deallocate {len(failures)} paged KV cache tensor(s)") from failures[0]
        self._bound_cache = None
        self._bound_context = None
        self._release_in_progress = False
        self._state = "released"

    # Private implementation

    def _deallocate_owned_tensors(self, *, reverse: bool = False) -> list[BaseException]:
        failures = []
        remaining = []
        tensors = reversed(self._owned_tensors) if reverse else self._owned_tensors
        for tensor in tensors:
            try:
                ttnn.deallocate(tensor)
            except BaseException as error:
                failures.append(error)
                remaining.append(tensor)
        if reverse:
            remaining.reverse()
        self._owned_tensors = tuple(remaining)
        return failures

    def _validate_static_model_contract(self, config, model_paged_configs) -> None:
        for layer, paged in enumerate(model_paged_configs):
            if paged is None:
                raise ValueError(f"Model layer {layer} is not configured for paged attention")
            if paged.block_size != config.block_size:
                raise ValueError(
                    f"Model layer {layer} block_size {paged.block_size} does not match {config.block_size}"
                )
            if paged.max_num_blocks != config.max_num_blocks:
                raise ValueError(
                    f"Model layer {layer} max_num_blocks {paged.max_num_blocks} does not match "
                    f"{config.max_num_blocks}"
                )

        model_dtypes = self.per_layer_dtypes
        if len(set(model_dtypes)) == 1 and model_dtypes[0] != config.dtype:
            raise ValueError(
                f"Configured KV dtype {config.dtype!r} does not match model-owned dtype {model_dtypes[0]!r}"
            )


@dataclass(frozen=True)
class _LayerKVSpec:
    local_kv_heads: int
    head_dim: int
    dtype: ttnn.DataType


def _model_contract(model: Any):
    model_config = getattr(model, "config", None)
    mesh_device = getattr(model_config, "mesh_device", None) or getattr(model, "mesh_device", None)
    if mesh_device is None:
        raise ValueError("Model config must provide mesh_device")
    num_devices = getattr(model_config, "num_devices", None) or getattr(model, "num_devices", None)
    if num_devices is None and hasattr(mesh_device, "get_num_devices"):
        num_devices = mesh_device.get_num_devices()
    if not isinstance(num_devices, int) or isinstance(num_devices, bool) or num_devices <= 0:
        raise ValueError("Model config must provide a positive num_devices")
    block_configs = getattr(model_config, "block_configs", None)
    if block_configs is not None:
        attention_configs = [getattr(block, "attention_config", None) for block in block_configs]
    else:
        layers = getattr(model, "layers", None)
        if layers is None:
            raise ValueError("Model must expose config.block_configs or layers")
        attention_configs = [getattr(getattr(layer, "attention", None), "config", None) for layer in layers]

    expected_layers = getattr(model_config, "n_layers", None) or getattr(model, "n_layers", None)
    if expected_layers is not None and len(attention_configs) != expected_layers:
        raise ValueError(
            f"Model exposes {len(attention_configs)} attention configs but declares {expected_layers} layers"
        )
    if not attention_configs or any(config is None for config in attention_configs):
        raise ValueError("Every model layer must expose an attention config")

    specs = []
    paged_configs = []
    for layer, attention in enumerate(attention_configs):
        local_kv_heads = getattr(attention, "n_local_kv_heads", None)
        if local_kv_heads is None:
            n_kv_heads = getattr(attention, "n_kv_heads", None)
            if not isinstance(n_kv_heads, int) or n_kv_heads <= 0 or n_kv_heads % num_devices:
                raise ValueError(f"Model layer {layer} n_kv_heads must be positive and divisible by num_devices")
            local_kv_heads = n_kv_heads // num_devices
        head_dim = getattr(attention, "head_dim", None)
        if not isinstance(head_dim, int) or head_dim <= 0:
            raise ValueError(f"Model layer {layer} must provide a positive head_dim")
        dtype = getattr(attention, "kv_cache_dtype", None)
        if dtype is None:
            raise ValueError(f"Model layer {layer} must own kv_cache_dtype")
        specs.append(_LayerKVSpec(local_kv_heads=local_kv_heads, head_dim=head_dim, dtype=dtype))
        paged_configs.append(getattr(attention, "paged_attention_config", None))
    return mesh_device, tuple(specs), tuple(paged_configs)


def _attach_cleanup_failures(primary: BaseException, failures: list[BaseException]) -> None:
    if not failures:
        return
    try:
        primary.cleanup_failures = tuple(failures)
    except BaseException:
        pass
    add_note = getattr(primary, "add_note", None)
    if add_note is not None:
        for failure in failures:
            add_note(f"cleanup failure: {type(failure).__name__}: {failure}")