From c42981d034713da814913ebd2e53269346f3ecea Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Thu, 30 Apr 2026 03:55:31 +0100 Subject: [PATCH] [Refactor][kv_offload] KV Offloading maintainability improvements (#40538) Signed-off-by: Martin Hickey Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 2 +- .../unit/offloading_connector/test_worker.py | 2 +- .../unit/offloading_connector/utils.py | 6 +- .../unit/test_offloading_connector.py} | 0 .../test_gpu_worker.py} | 9 +- .../test_manager.py} | 4 +- .../{ => cpu}/test_shared_offload_region.py | 0 tests/v1/kv_offload/test_worker.py | 2 +- .../kv_connector/v1/offloading/scheduler.py | 6 +- .../kv_connector/v1/offloading/worker.py | 2 +- vllm/v1/kv_offload/abstract.py | 197 --------- vllm/v1/kv_offload/base.py | 390 ++++++++++++++++++ vllm/v1/kv_offload/cpu/common.py | 13 + .../{worker/cpu_gpu.py => cpu/gpu_worker.py} | 8 +- vllm/v1/kv_offload/cpu/manager.py | 6 +- vllm/v1/kv_offload/cpu/policies/arc.py | 4 +- .../cpu/policies/{abstract.py => base.py} | 2 +- vllm/v1/kv_offload/cpu/policies/lru.py | 4 +- vllm/v1/kv_offload/cpu/spec.py | 13 +- vllm/v1/kv_offload/factory.py | 2 +- vllm/v1/kv_offload/mediums.py | 68 --- vllm/v1/kv_offload/reuse_manager.py | 2 +- vllm/v1/kv_offload/spec.py | 142 ------- vllm/v1/kv_offload/worker/worker.py | 2 +- 24 files changed, 446 insertions(+), 440 deletions(-) rename tests/v1/{kv_offload/test_cpu_offloading.py => kv_connector/unit/test_offloading_connector.py} (100%) rename tests/v1/kv_offload/{test_cpu_gpu.py => cpu/test_gpu_worker.py} (98%) rename tests/v1/kv_offload/{test_cpu_manager.py => cpu/test_manager.py} (99%) rename tests/v1/kv_offload/{ => cpu}/test_shared_offload_region.py (100%) delete mode 100644 vllm/v1/kv_offload/abstract.py create mode 100644 vllm/v1/kv_offload/base.py create mode 100644 vllm/v1/kv_offload/cpu/common.py rename vllm/v1/kv_offload/{worker/cpu_gpu.py => cpu/gpu_worker.py} (99%) rename vllm/v1/kv_offload/cpu/policies/{abstract.py => base.py} (97%) delete mode 100644 vllm/v1/kv_offload/mediums.py delete mode 100644 vllm/v1/kv_offload/spec.py diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 3ea935b090..1d76d6bdc0 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -15,7 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, ) from vllm.v1.core.kv_cache_utils import BlockHash -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( OffloadingEvent, OffloadingManager, ReqContext, diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 3cfb32b306..b2fb0846d9 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -20,7 +20,7 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, UniformTypeKVCacheSpecs, ) -from vllm.v1.kv_offload.spec import ( +from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, CanonicalKVCaches, OffloadingSpec, diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index d5adcd3f77..0268766eb7 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -37,15 +37,15 @@ from vllm.v1.kv_cache_interface import ( KVCacheConfig, KVCacheGroupSpec, ) -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( + GPULoadStoreSpec, LoadStoreSpec, OffloadingManager, + OffloadingSpec, OffloadKey, PrepareStoreOutput, make_offload_key, ) -from vllm.v1.kv_offload.mediums import GPULoadStoreSpec -from vllm.v1.kv_offload.spec import OffloadingSpec from vllm.v1.kv_offload.worker.worker import ( OffloadingHandler, TransferResult, diff --git a/tests/v1/kv_offload/test_cpu_offloading.py b/tests/v1/kv_connector/unit/test_offloading_connector.py similarity index 100% rename from tests/v1/kv_offload/test_cpu_offloading.py rename to tests/v1/kv_connector/unit/test_offloading_connector.py diff --git a/tests/v1/kv_offload/test_cpu_gpu.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py similarity index 98% rename from tests/v1/kv_offload/test_cpu_gpu.py rename to tests/v1/kv_offload/cpu/test_gpu_worker.py index db851edbcc..e4ed635b9b 100644 --- a/tests/v1/kv_offload/test_cpu_gpu.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -9,14 +9,15 @@ import torch from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec -from vllm.v1.kv_offload.spec import ( +from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, CanonicalKVCaches, CanonicalKVCacheTensor, + GPULoadStoreSpec, ) -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion NUM_GPU_BLOCKS = [64] NUM_CPU_BLOCKS = [256] diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/cpu/test_manager.py similarity index 99% rename from tests/v1/kv_offload/test_cpu_manager.py rename to tests/v1/kv_offload/cpu/test_manager.py index 733f9bf519..ef5d61e7b3 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -6,7 +6,7 @@ from dataclasses import dataclass import numpy as np import pytest -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( LoadStoreSpec, OffloadingEvent, OffloadKey, @@ -14,9 +14,9 @@ from vllm.v1.kv_offload.abstract import ( ReqContext, make_offload_key, ) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager diff --git a/tests/v1/kv_offload/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py similarity index 100% rename from tests/v1/kv_offload/test_shared_offload_region.py rename to tests/v1/kv_offload/cpu/test_shared_offload_region.py diff --git a/tests/v1/kv_offload/test_worker.py b/tests/v1/kv_offload/test_worker.py index fbdac5f9dc..b291fcf1b8 100644 --- a/tests/v1/kv_offload/test_worker.py +++ b/tests/v1/kv_offload/test_worker.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.v1.kv_offload.abstract import LoadStoreSpec +from vllm.v1.kv_offload.base import LoadStoreSpec from vllm.v1.kv_offload.worker.worker import ( OffloadingHandler, OffloadingWorker, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index cb8af41cdd..f437782070 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -18,15 +18,15 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( + GPULoadStoreSpec, OffloadingManager, + OffloadingSpec, OffloadKey, ReqContext, get_offload_block_hash, make_offload_key, ) -from vllm.v1.kv_offload.mediums import GPULoadStoreSpec -from vllm.v1.kv_offload.spec import OffloadingSpec from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 78547d569d..9acfd9c66c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import ( MambaSpec, UniformTypeKVCacheSpecs, ) -from vllm.v1.kv_offload.spec import ( +from vllm.v1.kv_offload.base import ( CanonicalKVCacheRef, CanonicalKVCaches, CanonicalKVCacheTensor, diff --git a/vllm/v1/kv_offload/abstract.py b/vllm/v1/kv_offload/abstract.py deleted file mode 100644 index be54611356..0000000000 --- a/vllm/v1/kv_offload/abstract.py +++ /dev/null @@ -1,197 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -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. -""" - -from abc import ABC, abstractmethod -from collections.abc import Iterable, Sequence -from dataclasses import dataclass -from typing import Any, NewType - -# `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) - - -@dataclass -class ReqContext: - kv_transfer_params: dict[str, Any] | None = None - - -class LoadStoreSpec(ABC): - """ - Abstract metadata that encapsulates information allowing a worker - to load, and optionally also to store, blocks of KV data. - """ - - @staticmethod - @abstractmethod - def medium() -> str: - """ - Returns a string representation of the medium type - this store/load targets. - """ - pass - - -@dataclass -class PrepareStoreOutput: - keys_to_store: list[OffloadKey] - store_spec: LoadStoreSpec - evicted_keys: list[OffloadKey] - - -@dataclass -class OffloadingEvent: - keys: list[OffloadKey] - medium: str - # True if blocks are removed, False if stored - removed: bool - - -class OffloadingManager(ABC): - @abstractmethod - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - """ - 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: - True if the block is offloaded and ready, False if not, - or None if the lookup should be retried later. - Returning None will delay the request handling by the vLLM - scheduler. - """ - pass - - @abstractmethod - def prepare_load( - self, - keys: Sequence[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: Sequence[OffloadKey]): - """ - 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. - """ - return - - def complete_load(self, keys: Iterable[OffloadKey]): - """ - Marks previous blocks that were prepared to load as done loading. - - Args: - keys: the keys identifying the blocks. - """ - return - - @abstractmethod - def prepare_store( - self, - keys: Sequence[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: Iterable[OffloadKey], success: bool = True): - """ - Marks blocks which were previously prepared to be stored, as stored. - Following this call, the blocks become loadable. - If if_success is False, blocks that were not marked as stored will be - removed. - - Args: - keys: the keys identifying the blocks. - success: whether the blocks were stored successfully. - """ - return - - def take_events(self) -> Iterable[OffloadingEvent]: - """ - Take the offloading events from the manager. - - Yields: - New OffloadingEvents collected since the last call. - """ - return () - - def shutdown(self) -> None: - """Shutdown the manager and release any resources.""" - return diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py new file mode 100644 index 0000000000..b30c5d066a --- /dev/null +++ b/vllm/v1/kv_offload/base.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Core abstractions for KV cache offloading in vLLM v1. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterable, Iterator, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, NewType + +import numpy as np +import torch + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.kv_offload.worker.worker import OffloadingHandler + +# `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) + +logger = init_logger(__name__) + + +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) + + +@dataclass +class ReqContext: + kv_transfer_params: dict[str, Any] | None = None + + +class LoadStoreSpec(ABC): + """ + Abstract metadata that encapsulates information allowing a worker + to load, and optionally also to store, blocks of KV data. + """ + + @staticmethod + @abstractmethod + def medium() -> str: + """ + Returns a string representation of the medium type + this store/load targets. + """ + pass + + +@dataclass +class PrepareStoreOutput: + keys_to_store: list[OffloadKey] + store_spec: LoadStoreSpec + evicted_keys: list[OffloadKey] + + +@dataclass +class OffloadingEvent: + keys: list[OffloadKey] + medium: str + # True if blocks are removed, False if stored + removed: bool + + +""" +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. +""" + + +class OffloadingManager(ABC): + @abstractmethod + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + """ + 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: + True if the block is offloaded and ready, False if not, + or None if the lookup should be retried later. + Returning None will delay the request handling by the vLLM + scheduler. + """ + pass + + @abstractmethod + def prepare_load( + self, + keys: Sequence[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: Sequence[OffloadKey]): + """ + 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. + """ + return + + def complete_load(self, keys: Iterable[OffloadKey]): + """ + Marks previous blocks that were prepared to load as done loading. + + Args: + keys: the keys identifying the blocks. + """ + return + + @abstractmethod + def prepare_store( + self, + keys: Sequence[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: Iterable[OffloadKey], success: bool = True): + """ + Marks blocks which were previously prepared to be stored, as stored. + Following this call, the blocks become loadable. + If if_success is False, blocks that were not marked as stored will be + removed. + + Args: + keys: the keys identifying the blocks. + success: whether the blocks were stored successfully. + """ + return + + def take_events(self) -> Iterable[OffloadingEvent]: + """ + Take the offloading events from the manager. + + Yields: + New OffloadingEvents collected since the last call. + """ + return () + + 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.int64) + + 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 + + @staticmethod + def medium() -> str: + return "GPU" + + +@dataclass +class CanonicalKVCacheTensor: + """ + A canonicalized KV cache tensor whose first dimension is num_blocks. + + For attention backends where the raw tensor has num_blocks at a + non-leading physical dimension (e.g. FlashAttention's + (2, num_blocks, ...) layout), the tensor is split so that each + resulting CanonicalKVCacheTensor starts with (num_blocks, ...). + """ + + # 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 +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 + + +@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]] + + +class OffloadingSpec(ABC): + """Spec for an offloading connector""" + + def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): + logger.warning( + "Initializing OffloadingSpec. This API is experimental and " + "subject to change in the future as we iterate the design." + ) + self.vllm_config = vllm_config + self.kv_cache_config = kv_cache_config + + kv_transfer_config = vllm_config.kv_transfer_config + assert kv_transfer_config is not None + self.extra_config = kv_transfer_config.kv_connector_extra_config + + # block size used by vLLM for hashing request tokens for the sake + # of enabling prefix caching + self.hash_block_size = vllm_config.cache_config.block_size + # gpu block size per group + self.gpu_block_size: tuple[int, ...] = tuple( + kv_cache_group.kv_cache_spec.block_size + for kv_cache_group in kv_cache_config.kv_cache_groups + ) + + for block_size in self.gpu_block_size: + assert block_size % self.hash_block_size == 0, ( + f"gpu_block_size={block_size} not divisible by " + f"hash_block_size={self.hash_block_size}. " + f"Hybrid models (e.g. Mamba+Attention) need " + f"--enable-prefix-caching to align block sizes." + ) + + # offloaded_block_size / gpu_block_size + self.block_size_factor: int = 1 + + offloaded_block_size = self.extra_config.get("block_size") + if offloaded_block_size is not None: + offloaded_block_size_int = int(offloaded_block_size) + gpu_block_sizes = set(self.gpu_block_size) + assert len(gpu_block_sizes) == 1, ( + "If 'block_size' is specified in kv_connector_extra_config, " + "there must be at least one KV cache group, " + "and all groups must have the same block size." + ) + gpu_block_size = gpu_block_sizes.pop() + + assert offloaded_block_size_int % gpu_block_size == 0 + self.block_size_factor = offloaded_block_size_int // gpu_block_size + + @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_handlers( + self, kv_caches: CanonicalKVCaches + ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]: + """ + Get offloading handlers along with their respective src and dst types. + + Args: + kv_caches: Canonicalized KV caches. + + Yields: + Tuples of (src_type, dst_type, offloading_handler). + """ + pass diff --git a/vllm/v1/kv_offload/cpu/common.py b/vllm/v1/kv_offload/cpu/common.py new file mode 100644 index 0000000000..cf5b2b39dd --- /dev/null +++ b/vllm/v1/kv_offload/cpu/common.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec + + +class CPULoadStoreSpec(BlockIDsLoadStoreSpec): + """ + Spec for loading/storing a KV block to CPU memory. + """ + + @staticmethod + def medium() -> str: + return "CPU" diff --git a/vllm/v1/kv_offload/worker/cpu_gpu.py b/vllm/v1/kv_offload/cpu/gpu_worker.py similarity index 99% rename from vllm/v1/kv_offload/worker/cpu_gpu.py rename to vllm/v1/kv_offload/cpu/gpu_worker.py index aab57ef2be..c8cd0ad3b4 100644 --- a/vllm/v1/kv_offload/worker/cpu_gpu.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -11,9 +11,13 @@ from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import is_pin_memory_available +from vllm.v1.kv_offload.base import ( + BlockIDsLoadStoreSpec, + CanonicalKVCacheRef, + CanonicalKVCaches, + GPULoadStoreSpec, +) from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion -from vllm.v1.kv_offload.mediums import BlockIDsLoadStoreSpec, GPULoadStoreSpec -from vllm.v1.kv_offload.spec import CanonicalKVCacheRef, CanonicalKVCaches from vllm.v1.kv_offload.worker.worker import ( OffloadingHandler, TransferResult, diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index fb03846105..3527773a20 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -3,7 +3,7 @@ from collections.abc import Iterable, Sequence from typing import Literal -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( LoadStoreSpec, OffloadingEvent, OffloadingManager, @@ -11,10 +11,10 @@ from vllm.v1.kv_offload.abstract import ( PrepareStoreOutput, ReqContext, ) -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec _CACHE_POLICIES: dict[str, type[CachePolicy]] = { "lru": LRUCachePolicy, diff --git a/vllm/v1/kv_offload/cpu/policies/arc.py b/vllm/v1/kv_offload/cpu/policies/arc.py index 07dfa4af4b..e2af991c0c 100644 --- a/vllm/v1/kv_offload/cpu/policies/arc.py +++ b/vllm/v1/kv_offload/cpu/policies/arc.py @@ -3,8 +3,8 @@ from collections import OrderedDict from collections.abc import Iterable -from vllm.v1.kv_offload.abstract import OffloadKey -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy +from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy class ARCCachePolicy(CachePolicy): diff --git a/vllm/v1/kv_offload/cpu/policies/abstract.py b/vllm/v1/kv_offload/cpu/policies/base.py similarity index 97% rename from vllm/v1/kv_offload/cpu/policies/abstract.py rename to vllm/v1/kv_offload/cpu/policies/base.py index 181df36baa..ee4916956e 100644 --- a/vllm/v1/kv_offload/cpu/policies/abstract.py +++ b/vllm/v1/kv_offload/cpu/policies/base.py @@ -4,7 +4,7 @@ import ctypes from abc import ABC, abstractmethod from collections.abc import Iterable -from vllm.v1.kv_offload.abstract import OffloadKey +from vllm.v1.kv_offload.base import OffloadKey class BlockStatus(ctypes.Structure): diff --git a/vllm/v1/kv_offload/cpu/policies/lru.py b/vllm/v1/kv_offload/cpu/policies/lru.py index be04831cf5..bf9a8b66e6 100644 --- a/vllm/v1/kv_offload/cpu/policies/lru.py +++ b/vllm/v1/kv_offload/cpu/policies/lru.py @@ -3,8 +3,8 @@ from collections import OrderedDict from collections.abc import Iterable -from vllm.v1.kv_offload.abstract import OffloadKey -from vllm.v1.kv_offload.cpu.policies.abstract import BlockStatus, CachePolicy +from vllm.v1.kv_offload.base import OffloadKey +from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy class LRUCachePolicy(CachePolicy): diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 8d4b744a0b..54046d98f4 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -5,12 +5,17 @@ from collections.abc import Iterator from vllm.config import VllmConfig from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager +from vllm.v1.kv_offload.base import ( + CanonicalKVCaches, + GPULoadStoreSpec, + LoadStoreSpec, + OffloadingManager, + OffloadingSpec, +) +from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec +from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager -from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager -from vllm.v1.kv_offload.spec import CanonicalKVCaches, OffloadingSpec -from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers from vllm.v1.kv_offload.worker.worker import OffloadingHandler diff --git a/vllm/v1/kv_offload/factory.py b/vllm/v1/kv_offload/factory.py index ecbaebb0d9..ecfc270b64 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -5,7 +5,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING from vllm.logger import init_logger -from vllm.v1.kv_offload.spec import OffloadingSpec +from vllm.v1.kv_offload.base import OffloadingSpec if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/vllm/v1/kv_offload/mediums.py b/vllm/v1/kv_offload/mediums.py deleted file mode 100644 index 02e36a80a8..0000000000 --- a/vllm/v1/kv_offload/mediums.py +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from abc import ABC -from collections.abc import Sequence - -import numpy as np - -from vllm.v1.kv_offload.abstract import LoadStoreSpec - - -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.int64) - - 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 - - @staticmethod - def medium() -> str: - return "GPU" - - -class CPULoadStoreSpec(BlockIDsLoadStoreSpec): - """ - Spec for loading/storing a KV block to CPU memory. - """ - - @staticmethod - def medium() -> str: - return "CPU" diff --git a/vllm/v1/kv_offload/reuse_manager.py b/vllm/v1/kv_offload/reuse_manager.py index 751f3a93df..2e556ca8d0 100644 --- a/vllm/v1/kv_offload/reuse_manager.py +++ b/vllm/v1/kv_offload/reuse_manager.py @@ -10,7 +10,7 @@ FilterReusedOffloadingManager — OffloadingManager decorator that skips from collections import OrderedDict from collections.abc import Iterable, Sequence -from vllm.v1.kv_offload.abstract import ( +from vllm.v1.kv_offload.base import ( LoadStoreSpec, OffloadingEvent, OffloadingManager, diff --git a/vllm/v1/kv_offload/spec.py b/vllm/v1/kv_offload/spec.py deleted file mode 100644 index b66b04ffbd..0000000000 --- a/vllm/v1/kv_offload/spec.py +++ /dev/null @@ -1,142 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from abc import ABC, abstractmethod -from collections.abc import Iterator -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import torch - -from vllm.logger import init_logger -from vllm.v1.kv_offload.abstract import LoadStoreSpec, OffloadingManager -from vllm.v1.kv_offload.worker.worker import OffloadingHandler - -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig - -logger = init_logger(__name__) - - -@dataclass -class CanonicalKVCacheTensor: - """ - A canonicalized KV cache tensor whose first dimension is num_blocks. - - For attention backends where the raw tensor has num_blocks at a - non-leading physical dimension (e.g. FlashAttention's - (2, num_blocks, ...) layout), the tensor is split so that each - resulting CanonicalKVCacheTensor starts with (num_blocks, ...). - """ - - # 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 -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 - - -@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]] - - -class OffloadingSpec(ABC): - """Spec for an offloading connector""" - - def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"): - logger.warning( - "Initializing OffloadingSpec. This API is experimental and " - "subject to change in the future as we iterate the design." - ) - self.vllm_config = vllm_config - self.kv_cache_config = kv_cache_config - - kv_transfer_config = vllm_config.kv_transfer_config - assert kv_transfer_config is not None - self.extra_config = kv_transfer_config.kv_connector_extra_config - - # block size used by vLLM for hashing request tokens for the sake - # of enabling prefix caching - self.hash_block_size = vllm_config.cache_config.block_size - # gpu block size per group - self.gpu_block_size: tuple[int, ...] = tuple( - kv_cache_group.kv_cache_spec.block_size - for kv_cache_group in kv_cache_config.kv_cache_groups - ) - - for block_size in self.gpu_block_size: - assert block_size % self.hash_block_size == 0, ( - f"gpu_block_size={block_size} not divisible by " - f"hash_block_size={self.hash_block_size}. " - f"Hybrid models (e.g. Mamba+Attention) need " - f"--enable-prefix-caching to align block sizes." - ) - - # offloaded_block_size / gpu_block_size - self.block_size_factor: int = 1 - - offloaded_block_size = self.extra_config.get("block_size") - if offloaded_block_size is not None: - offloaded_block_size_int = int(offloaded_block_size) - gpu_block_sizes = set(self.gpu_block_size) - assert len(gpu_block_sizes) == 1, ( - "If 'block_size' is specified in kv_connector_extra_config, " - "there must be at least one KV cache group, " - "and all groups must have the same block size." - ) - gpu_block_size = gpu_block_sizes.pop() - - assert offloaded_block_size_int % gpu_block_size == 0 - self.block_size_factor = offloaded_block_size_int // gpu_block_size - - @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_handlers( - self, kv_caches: CanonicalKVCaches - ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]: - """ - Get offloading handlers along with their respective src and dst types. - - Args: - kv_caches: Canonicalized KV caches. - - Yields: - Tuples of (src_type, dst_type, offloading_handler). - """ - pass diff --git a/vllm/v1/kv_offload/worker/worker.py b/vllm/v1/kv_offload/worker/worker.py index 2a47cc8e9d..2f0dd24716 100644 --- a/vllm/v1/kv_offload/worker/worker.py +++ b/vllm/v1/kv_offload/worker/worker.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from vllm.logger import init_logger -from vllm.v1.kv_offload.abstract import LoadStoreSpec +from vllm.v1.kv_offload.base import LoadStoreSpec # a single transfer spec (src_blocks_spec, dst_blocks_spec) TransferSpec = tuple[LoadStoreSpec, LoadStoreSpec]