vLLM/vllm/v1/kv_offload/base.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Core abstractions for KV cache offloading in vLLM v1.
"""

from abc import ABC, abstractmethod
from collections.abc import Collection, Iterable, Sequence
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar

import numpy as np
import torch

if TYPE_CHECKING:
    from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
        OffloadingConnectorStats,
    )

from vllm.v1.kv_offload.config import OffloadingConfig

# `OffloadKey` identifies an offloaded block. It combines a block hash with
# its KV cache group index, encoded as raw bytes to avoid tuple GC overhead.
# Use the helper functions below to construct / decompose keys.
OffloadKey = NewType("OffloadKey", bytes)


def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey:
    """Pack a block hash and group index into an `OffloadKey`."""
    return OffloadKey(block_hash + group_idx.to_bytes(4, "big", signed=False))


def get_offload_block_hash(key: OffloadKey) -> bytes:
    """Extract the block hash from an `OffloadKey`."""
    return key[:-4]


def get_offload_group_idx(key: OffloadKey) -> int:
    """Extract the group index from an `OffloadKey`."""
    return int.from_bytes(key[-4:], "big", signed=False)


_T = TypeVar("_T")


class Medium(Enum):
    """Storage medium of an offloading tier."""

    CPU = "CPU"
    STORAGE = "STORAGE"


class Locality(Enum):
    """Locality of a tier's storage relative to the publishing instance."""

    LOCAL = "LOCAL"
    REMOTE = "REMOTE"


class TierMatcher(NamedTuple):
    medium: Medium | None = None
    locality: Locality | None = None

    def matches(self, medium: Medium | None, locality: Locality | None) -> bool:
        medium_matches = self.medium is None or medium is None or self.medium == medium
        locality_matches = (
            self.locality is None or locality is None or self.locality == locality
        )
        return medium_matches and locality_matches


@dataclass(frozen=True)
class TierFilter:
    """Per-request filter controlling which tiers participate."""

    matchers: tuple[TierMatcher, ...] = ()

    ALL: ClassVar["TierFilter"]

    def allows(self, medium: Medium | None, locality: Locality | None) -> bool:
        if self is TierFilter.ALL:
            return True
        return any(m.matches(medium, locality) for m in self.matchers)


TierFilter.ALL = TierFilter(matchers=(TierMatcher(),))


@dataclass
class ReqContext:
    req_id: str
    kv_transfer_params: dict[str, Any] | None = None
    load_tier_filter: TierFilter = TierFilter.ALL
    # Per-request scratch space keyed by value type, so a tier can parse
    # kv_transfer_params once (in on_new_request) and read the result back
    # on later calls for the same request.
    _state: dict[type, Any] = field(default_factory=dict, repr=False, init=False)

    def set_state(self, val: Any) -> None:
        self._state[type(val)] = val

    def get_state(self, cls: type[_T]) -> _T | None:
        return self._state.get(cls)


class LookupResult(Enum):
    """Result of OffloadingManager.lookup()."""

    MISS = auto()
    HIT = auto()
    HIT_PENDING = auto()
    RETRY = auto()


class OffloadPolicy(Enum):
    # Offload only newly-computed blocks as they arrive; prefix-hit
    # blocks (already offloaded by a prior request) are skipped.
    BLOCK_LEVEL = "block_level"
    # Offload all blocks for the request, including prefix hits.
    # Used by tiers that need the complete KV context for a request.
    REQUEST_LEVEL = "request_level"


@dataclass
class RequestOffloadingContext:
    policy: OffloadPolicy = OffloadPolicy.BLOCK_LEVEL


class ScheduleEndContext(NamedTuple):
    """Per-step scheduling info passed to on_schedule_end()."""

    # Request IDs scheduled for the first time this step.
    new_req_ids: Collection[str]
    # Request IDs preempted this step.
    preempted_req_ids: Collection[str]


class LoadStoreSpec:
    """
    Metadata that encapsulates information allowing a worker
    to load, and optionally also to store, blocks of KV data.
    """


@dataclass
class PrepareStoreOutput:
    keys_to_store: list[OffloadKey]
    store_spec: LoadStoreSpec
    evicted_keys: list[OffloadKey]


@dataclass
class OffloadingEvent:
    keys: list[OffloadKey]
    medium: Medium
    # True if blocks are removed, False if stored
    removed: bool
    locality: Locality | None = None
    # Secondary tier identifier that generated the event, or None for primary.
    ownership: str | None = None
    removal_expected: bool = False


"""
OffloadingManager class for managing KV data offloading in vLLM v1

This class runs in the scheduler, tracks which blocks are offloaded
and their address.

The class provides the following primitives:
    lookup() - check whether a single block is offloaded and ready.
    prepare_load() - prepare given blocks to be read.
        The given blocks will be protected from eviction.
        This function returns a LoadSpec which encapsulates
        information required for performing the load.
    touch() - marks the give blocks as recently used. Can be used
        to track block's LRU. This function is separated from the
        prepare_load function to allow setting block recency even
        for blocks which do not need reading from the cache, such as
        blocks that are cached by the GPU prefix cache.
    complete_load() - mark blocks which were previously prepared to be
        loaded as done loading. This is to re-allow their eviction.
    prepare_store() - prepare the given blocks to be written.
        Returns a StoreSpec encapsulating offloading information,
        as well as a list of blocks that were evicted as a result.
    complete_store() - marks a previous store as completed.
        Following this call, the given blocks will become loadable.
"""


@dataclass(frozen=True)
class OffloadingMetricMetadata:
    documentation: str
    labelnames: tuple[str, ...] = ()


@dataclass(frozen=True)
class OffloadingCounterMetadata(OffloadingMetricMetadata):
    pass


@dataclass(frozen=True)
class OffloadingGaugeMetadata(OffloadingMetricMetadata):
    pass


@dataclass(frozen=True)
class OffloadingHistogramMetadata(OffloadingMetricMetadata):
    buckets: tuple[float, ...] | None = None


@dataclass(frozen=True)
class OffloadingKVEventsConfig:
    # Global vLLM KV event publishing flag. When false, connector-specific
    # event capture must stay inert because take_events() is not drained.
    enable_kv_cache_events: bool
    # OffloadingConnector opt-in for self-describing BlockStored payloads.
    # Effective only when enable_kv_cache_events is true.
    self_describing_kv_events: bool


class OffloadingManager(ABC):
    @abstractmethod
    def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
        """
        Checks whether a single block is offloaded and ready to be read.

        Args:
            key: the key identifying the block to lookup.
            req_context: per-request context (e.g. kv_transfer_params).

        Returns:
            HIT if the block is offloaded and ready, MISS if not found,
            HIT_PENDING if found but not yet readable, or RETRY if the
            lookup should be retried later.
        """
        pass

    @abstractmethod
    def prepare_load(
        self,
        keys: Collection[OffloadKey],
        req_context: ReqContext,
    ) -> LoadStoreSpec:
        """
        Prepare the given blocks to be read.
        The given blocks will be protected from eviction until
        complete_load is called.
        It assumes all given blocks are offloaded.

        Args:
            keys: the keys identifying the blocks.
            req_context: per-request context (e.g. kv_transfer_params).

        Returns:
            A LoadStoreSpec that can be used by a worker to locate and load
            the actual offloaded KV data.
        """
        pass

    def touch(self, keys: Collection[OffloadKey], req_context: ReqContext):
        """
        Mark the given blocks as recently used.
        This could in practice mean moving them to the end of an LRU list.

        Args:
            keys: the keys identifying the blocks.
            req_context: per-request context (e.g. kv_transfer_params).
        """
        return

    def complete_load(self, keys: Collection[OffloadKey], req_context: ReqContext):
        """
        Marks previous blocks that were prepared to load as done loading.

        Args:
            keys: the keys identifying the blocks.
            req_context: per-request context (e.g. kv_transfer_params).
        """
        return

    @abstractmethod
    def prepare_store(
        self,
        keys: Collection[OffloadKey],
        req_context: ReqContext,
    ) -> PrepareStoreOutput | None:
        """
        Prepare the given blocks to be offloaded.
        The given blocks will be protected from eviction until
        complete_store is called.

        Args:
            keys: the keys identifying the blocks.
            req_context: per-request context (e.g. kv_transfer_params).

        Returns:
            A PrepareStoreOutput indicating which blocks need storing,
            where to store them (LoadStoreSpec), and list of blocks that
            were evicted as a result.
            None is returned if the blocks cannot be stored.
        """
        pass

    def complete_store(
        self,
        keys: Collection[OffloadKey],
        req_context: ReqContext,
        success: bool = True,
    ):
        """
        Marks blocks which were previously prepared to be stored, as stored.
        Following this call, the blocks become loadable.
        If success is False, blocks that were not marked as stored will be
        removed.

        Args:
            keys: the keys identifying the blocks.
            req_context: per-request context (e.g. kv_transfer_params).
            success: whether the blocks were stored successfully.
        """
        return

    @abstractmethod
    def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
        """
        Called when a new request is first seen by the scheduler.

        Returns a RequestOffloadingContext indicating how this request's
        blocks should be offloaded.

        Args:
            req_context: per-request context.
        """
        pass

    def on_request_finished(self, req_context: ReqContext) -> None:
        """
        Called when a request has finished.

        By the time this is called, the scheduler will issue no more
        submit-side calls for this request, such as prepare_store() and
        prepare_load(). Completion callbacks for already-submitted transfers
        (complete_store() and complete_load()) may still arrive afterward.

        This hook does NOT imply the data has been persisted. Asynchronous
        transfers already submitted for this request may still be in flight.
        Managers that cascade to lower tiers should delay those tiers'
        on_request_finished() calls until no more lower-tier submit calls can
        be issued for this request.

        Args:
            req_context: per-request context.
        """
        return

    def take_events(self) -> Iterable[OffloadingEvent]:
        """
        Take the offloading events from the manager.

        A tier manager emits only events for storage state it owns. A
        composing manager may aggregate child event streams, but should not
        synthesize events on behalf of a child tier.

        Yields:
            New OffloadingEvents collected since the last call.
        """
        return ()

    def on_schedule_end(self, context: ScheduleEndContext) -> None:
        """Called once at the end of each scheduler step.

        Managers may override this to flush deferred work accumulated
        during the step (e.g., batched promotions).
        """
        return

    def has_pending_work(self) -> bool:
        """Whether this manager needs the engine to keep stepping.

        While True, on_schedule_end() and get_finished_jobs() continue
        to be called even when no requests are scheduled.
        """
        return False

    def reset_cache(self) -> None:
        """Evict all tracked blocks and reset internal state."""
        return

    def get_stats(self) -> "OffloadingConnectorStats | None":
        """Return collected metrics since last call, or None if disabled."""
        return None

    def shutdown(self) -> None:
        """Shutdown the manager and release any resources."""
        return


class BlockIDsLoadStoreSpec(LoadStoreSpec, ABC):
    """
    Spec for loading/storing KV blocks from given block numbers.
    """

    def __init__(self, block_ids: list[int]):
        self.block_ids = np.array(block_ids, dtype=np.int32)

    def __repr__(self) -> str:
        return repr(self.block_ids)


class GPULoadStoreSpec(BlockIDsLoadStoreSpec):
    """
    Spec for loading/storing a KV block to GPU memory.

    If there are multiple KV groups, the blocks are expected to be
    ordered by the group index.
    In that case, group_sizes[i] determines the number of blocks
    per the i-th KV group, and thus sum(group_sizes) == len(block_ids).
    group_sizes=None indicates a single KV group.

    If block_indices is given, each group (determined by group_sizes) of block IDs
    will correspond to logically contiguous blocks, e.g. blocks 5-10 of a some request.
    block_indices[i] will represent the block index of the first block in group #i.
    Thus, len(block_indices) == len(group_sizes) = number of KV cache groups.
    This information is required in order to support off/loading from offloaded blocks
    which are larger than GPU blocks.
    In such cases, the first GPU block per each group may be unaligned to the offloaded
    block size, and so knowing block_indices[i] allows the worker to correctly
    skip part of the first matching offloaded block.
    """

    def __init__(
        self,
        block_ids: list[int],
        group_sizes: Sequence[int],
        block_indices: Sequence[int],
    ):
        super().__init__(block_ids)
        assert sum(group_sizes) == len(block_ids)
        assert len(block_indices) == len(group_sizes)
        self.group_sizes: Sequence[int] = group_sizes
        self.block_indices: Sequence[int] = block_indices


@dataclass
class CanonicalKVCacheTensor:
    """
    A canonicalized KV cache tensor whose first dimension is num_blocks.

    With standardized layouts (RFC #42082) num_blocks is always the leading
    logical dimension.
    """

    # The KV cache tensor with shape (num_blocks, ...)
    tensor: torch.Tensor
    # The (possibly padded) page size per block in bytes
    page_size_bytes: int


@dataclass(frozen=True)
class CopyRun:
    """A strided byte correspondence between this worker's physical page and
    a canonical page: for i in range(num_fragments), fragment i spans
    [local_offset + i * local_stride, +fragment_size) in the worker's page and
    [canonical_offset + i * canonical_stride, +fragment_size) canonically."""

    local_offset: int
    canonical_offset: int
    fragment_size: int
    num_fragments: int
    local_stride: int
    canonical_stride: int


@dataclass(frozen=True)
class CanonicalPageMapping:
    """How this worker's page maps into a canonical (parallelism-free) page.
    In-process only, never serialized. Runs cover the full local page in both
    directions; ranks holding identical bytes take turns writing them.
    """

    # Size of the canonical page in bytes
    canonical_page_size_bytes: int
    # Size of this worker's (un-padded) page in bytes
    local_page_size_bytes: int
    # Byte correspondences between this worker's page and a canonical page
    runs: tuple[CopyRun, ...]
    # Number of ranks holding these exact bytes
    num_writers: int
    # This worker's index among those ranks
    writer_index: int
    # Canonical bytes identical under any parallel config with this block span
    parallelism_agnostic: bool

    def is_writer(self, block_id: int) -> bool:
        """Whether this worker stores the canonical page of the given block.
        Rotating by block spreads writes across ranks holding identical bytes.
        """
        return block_id % self.num_writers == self.writer_index


@dataclass
class CanonicalKVCacheRef:
    """
    Per-layer (or group of layers) reference to a specific (by index)
    CanonicalKVCacheTensor and records the un-padded page size used by that layer.
    """

    # Index into the list of CanonicalKVCacheTensor objects
    tensor_idx: int
    # The un-padded page size per block in bytes
    page_size_bytes: int
    # How this worker's page maps into a canonical page; None = uncertified
    mapping: CanonicalPageMapping | None = None


@dataclass
class CanonicalKVCaches:
    """
    Canonicalized block-level representation of the KV caches.

    Composed of:
        - Unique list of KV cache data tensors,
          each with shape (num_blocks, page_size_in_bytes) and int8 dtype.
        - Per-group data references of the tensors.
          i.e. how each KV cache group maps to the tensors.
    """

    # Ordered list of unique block tensors, each with shape
    # (num_blocks, ...).
    tensors: list[CanonicalKVCacheTensor]
    # Per-KV-cache-group list of data references that map each layer
    # in the group to the appropriate entry in the tensors list.
    group_data_refs: list[list[CanonicalKVCacheRef]]


@dataclass
class TransferResult:
    job_id: int
    success: bool
    transfer_size: int | None = None
    transfer_time: float | None = None


class OffloadingWorker(ABC):
    """Runs in the worker process. Performs async KV transfers for ONE
    offloaded medium (e.g. CPU). Direction is explicit via submit_store /
    submit_load, so there is no (src_medium, dst_medium) routing."""

    @abstractmethod
    def submit_store(
        self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec
    ) -> bool:
        """Async GPU -> offloaded medium."""

    @abstractmethod
    def submit_load(
        self, job_id: int, src_spec: LoadStoreSpec, dst_spec: GPULoadStoreSpec
    ) -> bool:
        """Async offloaded medium -> GPU."""

    @abstractmethod
    def get_finished(self) -> list[TransferResult]: ...

    @abstractmethod
    def wait(self, job_ids: set[int]) -> None: ...

    def shutdown(self) -> None:
        return


class OffloadingSpec(ABC):
    """Spec for an offloading connector"""

    @classmethod
    def build_metric_definitions(
        cls, extra_config: dict[str, Any]
    ) -> dict[str, "OffloadingMetricMetadata"]:
        """Return Prometheus metric definitions emitted by this spec."""
        return {}

    def __init__(self, config: OffloadingConfig):
        self.config = config
        self.extra_config = config.extra_config
        self.replicated_layout: bool = False
        self.kv_events_config = OffloadingKVEventsConfig(
            enable_kv_cache_events=config.enable_kv_cache_events,
            self_describing_kv_events=bool(
                self.extra_config.get("self_describing_kv_events", False)
            ),
        )

        # When True, only prompt (prefill) blocks are offloaded; decode-phase
        # blocks (KV generated after the prompt) are skipped. Useful when prior
        # turns' generated tokens are dropped before the next turn (e.g.
        # reasoning models that strip thinking).
        self.offload_prompt_only: bool = bool(
            self.extra_config.get("offload_prompt_only", True)
        )

        self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups)
        self.tokens_per_hash = config.cache.tokens_per_hash
        self.blocks_per_chunk = config.cache.blocks_per_chunk

    @abstractmethod
    def get_manager(self) -> OffloadingManager:
        """
        Get an OffloadingManager that will be used
        by the scheduler-side offloading connector to track
        offloaded blocks and manage evictions.
        """
        pass

    @abstractmethod
    def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
        """
        Get an OffloadingWorker that handles async KV transfers for this spec.

        Args:
            kv_caches: Canonicalized KV caches.

        Returns:
            An OffloadingWorker instance for this medium.
        """
        pass