vllm-omni/vllm_omni/core/prefix_cache/interface.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Interface types for the omni prefix cache.

Naming aligns with vLLM's v1/core KV-cache design.
"""

from collections.abc import Iterable
from dataclasses import dataclass
from enum import Enum
from typing import Any, NamedTuple, TypeAlias

import torch

# Four identities (still str/int at runtime, not interchangeable):
#   TensorName  which cached tensor (hidden / mm), not a prefix hash
#   ReqId       vLLM request id
#   Tid         WriteTask handle
#   StepId      save_outputs id; materialize or discard exactly once
TensorName: TypeAlias = str
ReqId: TypeAlias = str
Tid: TypeAlias = int
StepId: TypeAlias = int

# Reserved pool key for hidden states (mm keys are flat dotted names).
# Identity only — whether a model caches it lives on ModelCachePolicy.
HIDDEN_KEY: TensorName = "__hidden_states__"


def is_hidden_key(key: TensorName) -> bool:
    return key == HIDDEN_KEY


def without_hidden(keys: Iterable[TensorName]) -> set[TensorName]:
    """Drop the reserved hidden identity; the rest are mm names."""
    return {k for k in keys if k != HIDDEN_KEY}


class WriteSchedule(Enum):
    """Write scheduling policy for one WriteTask."""

    # Immediately-cached keys: device→host launched at save into the
    # staging pool; committer waits that event and writes the CPU pool.
    # The next save waits done (a reused slot must be pool-complete).
    JOIN_NEXT_STEP = "join_next_step"
    # Deferred mm: stays on the GPU clone until finish/abort (GPU-byte
    # budget may force a copy earlier). One WriteTask per request.
    JOIN_ON_FINISH = "join_on_finish"


@dataclass(frozen=True)
class PrefixCacheConfig:
    """Sizing and flow-control knobs (mirrors KVCacheConfig)."""

    num_blocks: int
    block_size: int
    # GPU-clone byte budget for JOIN_ON_FINISH; exceeding it forces a copy.
    gpu_staging_bytes: int = 512 * 1024 * 1024
    # Device→host staging: circular slots, one whole step each (not per request).
    # Host memory per key ≈ staging_depth * staging_capacity_tokens * width * dtype.
    # Prefer from_vllm_config so staging_capacity_tokens tracks max_num_batched_tokens.
    staging_depth: int = 4
    staging_capacity_tokens: int = 1024
    # How long save waits for a free staging slot, and join/join_host_ready
    # for the committer; exceeding it raises.
    staging_claim_timeout_s: float = 30.0
    # Device→host chunk size for JOIN_ON_FINISH (copied a piece at a time).
    copy_chunk_bytes: int = 16 * 1024 * 1024

    @classmethod
    def from_vllm_config(
        cls,
        *,
        num_blocks: int,
        block_size: int,
        scheduler_config: Any = None,
        model_config: Any = None,
    ) -> "PrefixCacheConfig":
        """Size device→host staging from the running scheduler.

        A slot holds one *step* (the whole batch), not one request:
        ``staging_capacity_tokens`` is ``max_num_batched_tokens`` (falls back
        to ``max_model_len``); ``staging_depth`` is how many unconsumed
        step ids may exist at once, not ``max_num_seqs``.

        Pinned staging is allocated lazily per key at
        ``depth * capacity_tokens * width * dtype``. There is no clamp: a
        step larger than capacity raises. A 16k-token thinking batch at
        hidden=2048 bf16 is ~256 MiB for hidden alone; each mm key adds
        another pool tensor.

        ``staging_depth`` is the dataclass default (4). There is no CLI or
        deploy YAML knob — changing it is a code change. Every save that
        issues a step id claims one slot, including saves with only leftover
        mm. A full pool waits for materialize/discard;
        ``staging_claim_timeout_s`` then errors.
        """
        batched = getattr(scheduler_config, "max_num_batched_tokens", None)
        model_len = getattr(scheduler_config, "max_model_len", None)
        if not model_len and model_config is not None:
            model_len = getattr(model_config, "max_model_len", None)
        try:
            capacity = int(batched or model_len or 1024)
        except (TypeError, ValueError):
            capacity = 1024
        return cls(
            num_blocks=num_blocks,
            block_size=block_size,
            staging_capacity_tokens=max(1, capacity),
        )


class StageCacheOutputs(NamedTuple):
    """Plain value object: per-request merged stage outputs."""

    # req -> full-prompt hidden states (None when policy skips them)
    hidden_states: dict[ReqId, torch.Tensor] | None
    # tensor name -> req -> payload element
    mm_outputs: dict[TensorName, dict[ReqId, Any]]


class OmniPrefixCacheUnmatchError(RuntimeError):
    """Contract, config, or KV-occupancy error that must raise.

    Includes hit spans that resolve to absent slots (omni cache diverged
    from vLLM KV), a step id consumed twice or never saved, a step larger
    than the staging page, and a failed write. Do not pretend these
    were a miss.
    """


class OmniPrefixCacheStagingTimeoutError(OmniPrefixCacheUnmatchError):
    """Save waited for a free staging slot and timed out."""


@dataclass(frozen=True)
class ModelCachePolicy:
    """Replaces getattr probing on models for cache behavior decisions.

    Hidden's *name* is ``HIDDEN_KEY`` (shared identity). This object
    answers whether this model caches it, and which mm keys are deferred.
    """

    needs_full_hidden_states: bool = True
    # Mm whose first dim is this step's token count; stays on the GPU
    # clone until finish/abort (JOIN_ON_FINISH). Also skipped by the
    # immediate on-device clone / device→host path.
    deferred_keys: frozenset[TensorName] = frozenset()

    @property
    def hidden_key(self) -> TensorName | None:
        """Pool key for hidden, or None when this model opts out."""
        return HIDDEN_KEY if self.needs_full_hidden_states else None

    def get_hit_keys(self, keys: Iterable[TensorName]) -> list[TensorName]:
        """Keys to plan/prefetch for a hit: hidden first (if cached), then mm."""
        mm = sorted(without_hidden(keys))
        return [HIDDEN_KEY, *mm] if self.needs_full_hidden_states else mm

    def skip_immediate_mm(self, key: TensorName) -> bool:
        """Immediate on-device clone must not take hidden or deferred keys from mm."""
        return is_hidden_key(key) or key in self.deferred_keys

    @classmethod
    def from_model(cls, model: Any) -> "ModelCachePolicy":
        """Shim over legacy per-model attributes (deprecation window)."""
        deferred = frozenset(getattr(model, "deferred_prefix_cache_mm_keys", ()) or ())
        return cls(
            needs_full_hidden_states=bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)),
            deferred_keys=deferred,
        )