[Refactor][kv_offload] KV Offloading maintainability improvements (#40538)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Martin Hickey
co-authored by
Or Ozeri
Or Ozeri
parent
0ff1bf9bb1
commit
c42981d034
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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,
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
+1
-1
@@ -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):
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user