104 lines
4.1 KiB
Python
104 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
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.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.worker.worker import OffloadingHandler
|
|
|
|
|
|
class CPUOffloadingSpec(OffloadingSpec):
|
|
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
|
|
super().__init__(vllm_config, kv_cache_config)
|
|
|
|
cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use")
|
|
if not cpu_bytes_to_use:
|
|
raise Exception(
|
|
"cpu_bytes_to_use must be specified in kv_connector_extra_config"
|
|
)
|
|
|
|
# calculate kv_bytes_per_offloaded_block
|
|
assert kv_cache_config is not None
|
|
if kv_cache_config.num_blocks > 0:
|
|
total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
|
|
kv_bytes_per_block = (
|
|
total_gpu_kv_bytes // kv_cache_config.num_blocks
|
|
) * vllm_config.parallel_config.world_size
|
|
else:
|
|
kv_bytes_per_block = 0
|
|
|
|
kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
|
|
self.num_blocks = (
|
|
int(cpu_bytes_to_use) // kv_bytes_per_offloaded_block
|
|
if kv_bytes_per_offloaded_block > 0
|
|
else 0
|
|
)
|
|
world_size = vllm_config.parallel_config.world_size
|
|
self.cpu_page_size_per_worker: int = (
|
|
kv_bytes_per_offloaded_block // world_size if world_size > 0 else 0
|
|
)
|
|
|
|
# scheduler-side
|
|
self._manager: OffloadingManager | None = None
|
|
|
|
# worker-side
|
|
self._handlers: CpuGpuOffloadingHandlers | None = None
|
|
|
|
self.eviction_policy: str = self.extra_config.get("eviction_policy", "lru")
|
|
|
|
def get_manager(self) -> OffloadingManager:
|
|
if not self._manager:
|
|
kv_events_config = self.vllm_config.kv_events_config
|
|
enable_events = (
|
|
kv_events_config is not None and kv_events_config.enable_kv_cache_events
|
|
)
|
|
|
|
# store_threshold: how many times a block must appear in lookup()
|
|
# before it is eligible for CPU offloading. Values < 2 disable
|
|
# filtering (a threshold of 1 equals no filter; 0 is the default).
|
|
store_threshold = int(self.extra_config.get("store_threshold", 0))
|
|
|
|
# Maximum entries in the internal tracker's LRU table.
|
|
max_tracker_size = int(self.extra_config.get("max_tracker_size", 64_000))
|
|
|
|
self._manager = CPUOffloadingManager(
|
|
num_blocks=self.num_blocks,
|
|
cache_policy=self.eviction_policy, # type: ignore[arg-type]
|
|
enable_events=enable_events,
|
|
store_threshold=store_threshold,
|
|
max_tracker_size=max_tracker_size,
|
|
)
|
|
return self._manager
|
|
|
|
def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
|
|
return CpuGpuOffloadingHandlers(
|
|
kv_caches=kv_caches,
|
|
block_size_factor=self.block_size_factor,
|
|
num_cpu_blocks=self.num_blocks,
|
|
)
|
|
|
|
def get_handlers(
|
|
self, kv_caches: CanonicalKVCaches
|
|
) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
|
|
if not self._handlers:
|
|
if not current_platform.is_cuda_alike():
|
|
raise Exception(
|
|
"CPU Offloading is currently only supported on CUDA-alike GPUs"
|
|
)
|
|
self._handlers = self.create_handlers(kv_caches)
|
|
|
|
assert self._handlers is not None
|
|
yield GPULoadStoreSpec, CPULoadStoreSpec, self._handlers.gpu_to_cpu_handler
|
|
yield CPULoadStoreSpec, GPULoadStoreSpec, self._handlers.cpu_to_gpu_handler
|