diff --git a/vllm/utils/extensible_tensor.py b/vllm/utils/extensible_tensor.py index 2bef3525e11..d1b746bc810 100644 --- a/vllm/utils/extensible_tensor.py +++ b/vllm/utils/extensible_tensor.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Growable CUDA byte buffers backed by CUDA virtual memory management.""" +"""Growable GPU byte buffers backed by driver virtual memory management.""" from __future__ import annotations @@ -9,167 +9,7 @@ from contextlib import suppress import torch -_CUDA_SUCCESS = 0 -_CU_MEM_ALLOCATION_TYPE_PINNED = 1 -_CU_MEM_LOCATION_TYPE_DEVICE = 1 -_CU_MEM_ALLOC_GRANULARITY_MINIMUM = 0 -_CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 -_CU_MEM_ALLOCATION_COMP_NONE = 0 - - -class _CUmemLocation(ctypes.Structure): - _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] - - -class _CUmemAllocFlags(ctypes.Structure): - _fields_ = [ - ("compressionType", ctypes.c_ubyte), - ("gpuDirectRDMACapable", ctypes.c_ubyte), - ("usage", ctypes.c_ushort), - ("reserved", ctypes.c_ubyte * 4), - ] - - -class _CUmemAllocationProp(ctypes.Structure): - _fields_ = [ - ("type", ctypes.c_int), - ("requestedHandleTypes", ctypes.c_int), - ("location", _CUmemLocation), - ("win32HandleMetaData", ctypes.c_void_p), - ("allocFlags", _CUmemAllocFlags), - ] - - -class _CUmemAccessDesc(ctypes.Structure): - _fields_ = [("location", _CUmemLocation), ("flags", ctypes.c_int)] - - -_CUdeviceptr = ctypes.c_ulonglong -_CUmemHandle = ctypes.c_ulonglong -_CUcontext = ctypes.c_void_p - -_libcuda: ctypes.CDLL | None = None - - -def _find_loaded_library(lib_name: str) -> str | None: - try: - with open("/proc/self/maps") as f: - for line in f: - if lib_name not in line: - continue - start = line.index("/") - return line[start:].strip() - except (OSError, ValueError): - return None - return None - - -def _load_libcuda() -> ctypes.CDLL: - for name in ("libcuda.so.1", "libcuda.so"): - try: - return ctypes.CDLL(name) - except OSError: - continue - if path := _find_loaded_library("libcuda"): - return ctypes.CDLL(path) - raise RuntimeError( - "Could not load libcuda. The CUDA driver library is required for " - "ExtensibleTensor." - ) - - -def _configure_signatures(lib: ctypes.CDLL) -> None: - pointer = ctypes.POINTER - lib.cuGetErrorString.argtypes = [ctypes.c_int, pointer(ctypes.c_char_p)] - lib.cuCtxGetCurrent.argtypes = [pointer(_CUcontext)] - lib.cuDevicePrimaryCtxRetain.argtypes = [pointer(_CUcontext), ctypes.c_int] - lib.cuCtxSetCurrent.argtypes = [_CUcontext] - lib.cuMemGetAllocationGranularity.argtypes = [ - pointer(ctypes.c_size_t), - pointer(_CUmemAllocationProp), - ctypes.c_int, - ] - lib.cuMemAddressReserve.argtypes = [ - pointer(_CUdeviceptr), - ctypes.c_size_t, - ctypes.c_size_t, - _CUdeviceptr, - ctypes.c_ulonglong, - ] - lib.cuMemCreate.argtypes = [ - pointer(_CUmemHandle), - ctypes.c_size_t, - pointer(_CUmemAllocationProp), - ctypes.c_ulonglong, - ] - lib.cuMemMap.argtypes = [ - _CUdeviceptr, - ctypes.c_size_t, - ctypes.c_size_t, - _CUmemHandle, - ctypes.c_ulonglong, - ] - lib.cuMemSetAccess.argtypes = [ - _CUdeviceptr, - ctypes.c_size_t, - pointer(_CUmemAccessDesc), - ctypes.c_size_t, - ] - lib.cuMemUnmap.argtypes = [_CUdeviceptr, ctypes.c_size_t] - lib.cuMemRelease.argtypes = [_CUmemHandle] - lib.cuMemAddressFree.argtypes = [_CUdeviceptr, ctypes.c_size_t] - - for fn in ( - lib.cuGetErrorString, - lib.cuCtxGetCurrent, - lib.cuDevicePrimaryCtxRetain, - lib.cuCtxSetCurrent, - lib.cuMemGetAllocationGranularity, - lib.cuMemAddressReserve, - lib.cuMemCreate, - lib.cuMemMap, - lib.cuMemSetAccess, - lib.cuMemUnmap, - lib.cuMemRelease, - lib.cuMemAddressFree, - ): - fn.restype = ctypes.c_int - - -def _cuda() -> ctypes.CDLL: - global _libcuda - if _libcuda is None: - lib = _load_libcuda() - _configure_signatures(lib) - _libcuda = lib - return _libcuda - - -def _check(result: int) -> None: - if result == _CUDA_SUCCESS: - return - msg = ctypes.c_char_p() - _cuda().cuGetErrorString(result, ctypes.byref(msg)) - detail = msg.value.decode() if msg.value else "unknown error" - raise RuntimeError(f"CUDA driver error {result}: {detail}") - - -def _ensure_context(device_index: int) -> None: - pctx = _CUcontext() - _check(_cuda().cuCtxGetCurrent(ctypes.byref(pctx))) - if pctx.value: - return - _check(_cuda().cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index)) - _check(_cuda().cuCtxSetCurrent(pctx)) - - -def _make_alloc_prop(device_index: int) -> _CUmemAllocationProp: - prop = _CUmemAllocationProp() - prop.type = _CU_MEM_ALLOCATION_TYPE_PINNED - prop.location.type = _CU_MEM_LOCATION_TYPE_DEVICE - prop.location.id = device_index - prop.allocFlags.compressionType = _CU_MEM_ALLOCATION_COMP_NONE - return prop +from vllm.utils.vmm_driver import get_vmm_driver def _round_up(value: int, multiple: int) -> int: @@ -186,26 +26,13 @@ class _VirtualBuffer: """ def __init__(self, max_bytes: int, device_index: int) -> None: - _ensure_context(device_index) + self._driver = get_vmm_driver() + self._driver.ensure_context(device_index) self.device_index = device_index - prop = _make_alloc_prop(device_index) - granularity = ctypes.c_size_t() - _check( - _cuda().cuMemGetAllocationGranularity( - ctypes.byref(granularity), - ctypes.byref(prop), - _CU_MEM_ALLOC_GRANULARITY_MINIMUM, - ) - ) - self.granularity: int = granularity.value + self.granularity: int = self._driver.granularity(device_index) self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity) - - dptr = _CUdeviceptr() - _check( - _cuda().cuMemAddressReserve(ctypes.byref(dptr), self.reserved_size, 0, 0, 0) - ) - self.base_ptr: int = dptr.value + self.base_ptr: int = self._driver.reserve(self.reserved_size) # Granule indices (VA offset // granularity) that have physical # memory mapped. @@ -255,39 +82,33 @@ class _VirtualBuffer: def _map_chunk_at(self, offset: int, size: int) -> None: """Create one physical chunk of `size` bytes and map it at `offset`.""" - _ensure_context(self.device_index) - prop = _make_alloc_prop(self.device_index) - - handle = _CUmemHandle() - _check(_cuda().cuMemCreate(ctypes.byref(handle), size, ctypes.byref(prop), 0)) + driver = self._driver + driver.ensure_context(self.device_index) + handle = driver.create(size, self.device_index) addr = self.base_ptr + offset try: - _check(_cuda().cuMemMap(addr, size, 0, handle, 0)) + driver.map(addr, size, handle) except RuntimeError: - _cuda().cuMemRelease(handle) + driver.release(handle) raise + driver.set_access(addr, size, self.device_index) - desc = _CUmemAccessDesc() - desc.location.type = _CU_MEM_LOCATION_TYPE_DEVICE - desc.location.id = self.device_index - desc.flags = _CU_MEM_ACCESS_FLAGS_PROT_READWRITE - _check(_cuda().cuMemSetAccess(addr, size, ctypes.byref(desc), 1)) - - self._handles.append((handle.value, offset, size)) + self._handles.append((handle, offset, size)) def free(self) -> None: if self._freed: return self._freed = True - _ensure_context(self.device_index) + driver = self._driver + driver.ensure_context(self.device_index) if self._handles: torch.cuda.synchronize(self.device_index) for handle, offset, size in self._handles: - _check(_cuda().cuMemUnmap(self.base_ptr + offset, size)) - _check(_cuda().cuMemRelease(handle)) + driver.unmap(self.base_ptr + offset, size) + driver.release(handle) if self.base_ptr: - _check(_cuda().cuMemAddressFree(self.base_ptr, self.reserved_size)) + driver.free_reserved(self.base_ptr, self.reserved_size) self._handles = [] self._mapped_granules = set() self.base_ptr = 0 @@ -297,7 +118,6 @@ class _VirtualBuffer: self.free() -_K_DL_CUDA = 2 _K_DL_UINT = 1 _UINT8_BITS = 8 @@ -348,7 +168,8 @@ def _uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch managed = _DLManagedTensor() managed.dl_tensor.data = ctypes.c_void_p(ptr) - managed.dl_tensor.device = _DLDevice(_K_DL_CUDA, device_index) + device_type = get_vmm_driver().dlpack_device_type + managed.dl_tensor.device = _DLDevice(device_type, device_index) managed.dl_tensor.ndim = 1 managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1) managed.dl_tensor.shape = ctypes.cast(shape_arr, ctypes.POINTER(ctypes.c_int64)) diff --git a/vllm/utils/vmm_driver.py b/vllm/utils/vmm_driver.py new file mode 100644 index 00000000000..55d27327db4 --- /dev/null +++ b/vllm/utils/vmm_driver.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ctypes bindings for GPU virtual-memory-management (VMM) driver APIs. + +Exposes a uniform driver interface over the CUDA driver's ``cuMem*`` entry +points (and, structurally, HIP's mirrored ``hipMem*`` entry points) used by +:class:`vllm.utils.extensible_tensor.ExtensibleTensor`: reserve a virtual +address range, create physical memory handles, map/unmap them into the +reservation, and set access permissions. +""" + +from __future__ import annotations + +import ctypes +from functools import cache + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_SUCCESS = 0 +_MEM_ALLOCATION_TYPE_PINNED = 1 +_MEM_LOCATION_TYPE_DEVICE = 1 +_MEM_ALLOC_GRANULARITY_MINIMUM = 0 +_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 +_MEM_ALLOCATION_COMP_NONE = 0 + +DevicePtr = ctypes.c_ulonglong +MemHandle = ctypes.c_ulonglong +_Context = ctypes.c_void_p + + +class _MemLocation(ctypes.Structure): + _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] + + +class _MemAllocFlags(ctypes.Structure): + _fields_ = [ + ("compressionType", ctypes.c_ubyte), + ("gpuDirectRDMACapable", ctypes.c_ubyte), + ("usage", ctypes.c_ushort), + ("reserved", ctypes.c_ubyte * 4), + ] + + +class _MemAllocationProp(ctypes.Structure): + # Layout shared by CUmemAllocationProp and hipMemAllocationProp. + _fields_ = [ + ("type", ctypes.c_int), + ("requestedHandleTypes", ctypes.c_int), + ("location", _MemLocation), + ("win32HandleMetaData", ctypes.c_void_p), + ("allocFlags", _MemAllocFlags), + ] + + +class _MemAccessDesc(ctypes.Structure): + _fields_ = [("location", _MemLocation), ("flags", ctypes.c_int)] + + +def _find_loaded_library(lib_name: str) -> str | None: + try: + with open("/proc/self/maps") as f: + for line in f: + if lib_name not in line: + continue + start = line.index("/") + return line[start:].strip() + except (OSError, ValueError): + return None + return None + + +class VmmDriver: + """Uniform interface over a GPU driver's virtual memory management API. + + Subclasses supply the driver library candidates and symbol names; the + call signatures and struct layouts are shared between CUDA and HIP. + """ + + # DLPack device type for tensors viewing driver-mapped memory. + dlpack_device_type: int + _lib_candidates: tuple[str, ...] + _lib_search_name: str + # Logical name -> library symbol. + _symbols: dict[str, str] + + def __init__(self) -> None: + self._lib = self._load_library() + self._fns: dict[str, ctypes._CFuncPtr] = {} + for logical, symbol in self._symbols.items(): + self._fns[logical] = getattr(self._lib, symbol) + self._configure_signatures() + + def _load_library(self) -> ctypes.CDLL: + for name in self._lib_candidates: + try: + return ctypes.CDLL(name) + except OSError: + continue + if path := _find_loaded_library(self._lib_search_name): + return ctypes.CDLL(path) + raise RuntimeError( + f"Could not load {self._lib_candidates[0]}. The GPU driver " + "library is required for VMM-backed tensors." + ) + + def _configure_signatures(self) -> None: + pointer = ctypes.POINTER + fns = self._fns + fns["get_granularity"].argtypes = [ + pointer(ctypes.c_size_t), + pointer(_MemAllocationProp), + ctypes.c_int, + ] + fns["address_reserve"].argtypes = [ + pointer(DevicePtr), + ctypes.c_size_t, + ctypes.c_size_t, + DevicePtr, + ctypes.c_ulonglong, + ] + fns["create"].argtypes = [ + pointer(MemHandle), + ctypes.c_size_t, + pointer(_MemAllocationProp), + ctypes.c_ulonglong, + ] + fns["map"].argtypes = [ + DevicePtr, + ctypes.c_size_t, + ctypes.c_size_t, + MemHandle, + ctypes.c_ulonglong, + ] + fns["set_access"].argtypes = [ + DevicePtr, + ctypes.c_size_t, + pointer(_MemAccessDesc), + ctypes.c_size_t, + ] + fns["unmap"].argtypes = [DevicePtr, ctypes.c_size_t] + fns["release"].argtypes = [MemHandle] + fns["address_free"].argtypes = [DevicePtr, ctypes.c_size_t] + for fn in fns.values(): + fn.restype = ctypes.c_int + + def error_string(self, code: int) -> str: + raise NotImplementedError + + def ensure_context(self, device_index: int) -> None: + """Make sure a driver context for `device_index` is current.""" + raise NotImplementedError + + def _check(self, result: int) -> None: + if result == _SUCCESS: + return + raise RuntimeError(f"GPU driver error {result}: {self.error_string(result)}") + + def _make_alloc_prop( + self, device_index: int, rdma_capable: bool = False + ) -> _MemAllocationProp: + prop = _MemAllocationProp() + prop.type = _MEM_ALLOCATION_TYPE_PINNED + prop.location.type = _MEM_LOCATION_TYPE_DEVICE + prop.location.id = device_index + prop.allocFlags.compressionType = _MEM_ALLOCATION_COMP_NONE + prop.allocFlags.gpuDirectRDMACapable = 1 if rdma_capable else 0 + return prop + + def granularity(self, device_index: int) -> int: + prop = self._make_alloc_prop(device_index) + granularity = ctypes.c_size_t() + self._check( + self._fns["get_granularity"]( + ctypes.byref(granularity), + ctypes.byref(prop), + _MEM_ALLOC_GRANULARITY_MINIMUM, + ) + ) + return granularity.value + + def reserve(self, size: int) -> int: + """Reserve a virtual address range and return its base pointer.""" + dptr = DevicePtr() + self._check(self._fns["address_reserve"](ctypes.byref(dptr), size, 0, 0, 0)) + return dptr.value + + def free_reserved(self, ptr: int, size: int) -> None: + self._check(self._fns["address_free"](ptr, size)) + + def create(self, size: int, device_index: int, rdma_capable: bool = False) -> int: + """Create a physical memory handle of `size` bytes.""" + prop = self._make_alloc_prop(device_index, rdma_capable) + handle = MemHandle() + self._check( + self._fns["create"](ctypes.byref(handle), size, ctypes.byref(prop), 0) + ) + return handle.value + + def map(self, ptr: int, size: int, handle: int) -> None: + self._check(self._fns["map"](ptr, size, 0, handle, 0)) + + def set_access(self, ptr: int, size: int, device_index: int) -> None: + desc = _MemAccessDesc() + desc.location.type = _MEM_LOCATION_TYPE_DEVICE + desc.location.id = device_index + desc.flags = _MEM_ACCESS_FLAGS_PROT_READWRITE + self._check(self._fns["set_access"](ptr, size, ctypes.byref(desc), 1)) + + def unmap(self, ptr: int, size: int) -> None: + self._check(self._fns["unmap"](ptr, size)) + + def release(self, handle: int) -> None: + self._check(self._fns["release"](handle)) + + +class CudaVmmDriver(VmmDriver): + dlpack_device_type = 2 # kDLCUDA + _lib_candidates = ("libcuda.so.1", "libcuda.so") + _lib_search_name = "libcuda" + _symbols = { + "get_granularity": "cuMemGetAllocationGranularity", + "address_reserve": "cuMemAddressReserve", + "create": "cuMemCreate", + "map": "cuMemMap", + "set_access": "cuMemSetAccess", + "unmap": "cuMemUnmap", + "release": "cuMemRelease", + "address_free": "cuMemAddressFree", + } + + def __init__(self) -> None: + super().__init__() + lib = self._lib + lib.cuGetErrorString.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_char_p), + ] + lib.cuGetErrorString.restype = ctypes.c_int + lib.cuCtxGetCurrent.argtypes = [ctypes.POINTER(_Context)] + lib.cuCtxGetCurrent.restype = ctypes.c_int + lib.cuDevicePrimaryCtxRetain.argtypes = [ + ctypes.POINTER(_Context), + ctypes.c_int, + ] + lib.cuDevicePrimaryCtxRetain.restype = ctypes.c_int + lib.cuCtxSetCurrent.argtypes = [_Context] + lib.cuCtxSetCurrent.restype = ctypes.c_int + + def error_string(self, code: int) -> str: + msg = ctypes.c_char_p() + self._lib.cuGetErrorString(code, ctypes.byref(msg)) + return msg.value.decode() if msg.value else "unknown error" + + def ensure_context(self, device_index: int) -> None: + pctx = _Context() + self._check(self._lib.cuCtxGetCurrent(ctypes.byref(pctx))) + if pctx.value: + return + self._check( + self._lib.cuDevicePrimaryCtxRetain(ctypes.byref(pctx), device_index) + ) + self._check(self._lib.cuCtxSetCurrent(pctx)) + + +@cache +def get_vmm_driver() -> VmmDriver: + import torch + + if torch.version.hip is not None: + raise RuntimeError("VMM-backed tensors are not yet supported on ROCm.") + return CudaVmmDriver() + + +@cache +def vmm_unavailable_reason() -> str | None: + """Probe VMM support; returns None if usable, else a reason string. + + Checks that the driver library loads, exposes the VMM entry points, and + can actually reserve (and release) a virtual address range on the current + device. Notably returns a reason on platforms whose driver lacks VMM + support (e.g. WSL2) and on non-CUDA/ROCm builds. + """ + try: + import torch + + if not torch.cuda.is_available(): + return "no CUDA/ROCm device is available" + torch.cuda.init() + driver = get_vmm_driver() + device_index = torch.cuda.current_device() + driver.ensure_context(device_index) + granularity = driver.granularity(device_index) + ptr = driver.reserve(granularity) + driver.free_reserved(ptr, granularity) + except Exception as e: + return str(e) + return None diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b7be19791b7..211fde56588 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -296,18 +296,6 @@ class EngineCore: has_kv_cache and vllm_config.cache_config.enable_extensible_kv_cache ) if use_extensible_kv_cache: - from vllm.platforms import current_platform - - if vllm_config.cache_config.kv_cache_memory_bytes is not None: - raise ValueError( - "enable_extensible_kv_cache=True is not supported with " - "kv_cache_memory_bytes. The extensible path requires " - "automatic KV cache sizing." - ) - if not current_platform.is_cuda(): - raise ValueError( - "enable_extensible_kv_cache=True is only supported on CUDA." - ) if vllm_config.kv_transfer_config is not None: raise ValueError( "enable_extensible_kv_cache=True is not supported with " @@ -319,6 +307,16 @@ class EngineCore: "enable_extensible_kv_cache=True is not supported with " "sleep mode: the KV cache bypasses the CuMem allocator." ) + # The workers' drivers must support virtual memory management + # (e.g. WSL2 and non-GPU platforms do not); fall back gracefully. + reasons = self.collective_rpc("extensible_kv_cache_unsupported_reason") + if reason := next((r for r in reasons if r), None): + logger.warning( + "Disabling extensible KV cache; falling back to standard " + "KV cache allocation: %s", + reason, + ) + use_extensible_kv_cache = False # Track max_model_len before KV cache config to detect auto-fit changes # made by get_kv_cache_configs(). @@ -341,35 +339,41 @@ class EngineCore: extensible=use_extensible_kv_cache, ) if use_extensible_kv_cache: - if len(compilation_times) != len(available_gpu_memory): - raise RuntimeError( - "Expected one CompilationTimes result per worker when " - "initializing extensible KV cache, but got " - f"{len(compilation_times)} results for " - f"{len(available_gpu_memory)} workers." + if vllm_config.cache_config.kv_cache_memory_bytes is None: + # Automatic sizing: re-derive the KV cache size from the + # memory actually consumed by warmup and CUDA graph capture. + # With an explicit kv_cache_memory_bytes, the requested size + # is committed as-is (the extensible path still defers the + # commit until after warmup). + if len(compilation_times) != len(available_gpu_memory): + raise RuntimeError( + "Expected one CompilationTimes result per worker when " + "initializing extensible KV cache, but got " + f"{len(compilation_times)} results for " + f"{len(available_gpu_memory)} workers." + ) + final_available_gpu_memory = [ + max( + available_memory + - times.warmup_memory + - _WARMUP_MEMORY_BUFFER_BYTES, + 0, + ) + for available_memory, times in zip( + available_gpu_memory, compilation_times, strict=True + ) + ] + max_model_len_before = vllm_config.model_config.max_model_len + kv_cache_configs = get_kv_cache_configs( + vllm_config, + kv_cache_specs, + final_available_gpu_memory, ) - final_available_gpu_memory = [ - max( - available_memory - - times.warmup_memory - - _WARMUP_MEMORY_BUFFER_BYTES, - 0, + scheduler_kv_cache_config = self._apply_kv_cache_config( + vllm_config, + kv_cache_configs, + max_model_len_before, ) - for available_memory, times in zip( - available_gpu_memory, compilation_times, strict=True - ) - ] - max_model_len_before = vllm_config.model_config.max_model_len - kv_cache_configs = get_kv_cache_configs( - vllm_config, - kv_cache_specs, - final_available_gpu_memory, - ) - scheduler_kv_cache_config = self._apply_kv_cache_config( - vllm_config, - kv_cache_configs, - max_model_len_before, - ) self.model_executor.extend_kv_cache(scheduler_kv_cache_config.num_blocks) elapsed = time.time() - start diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 700ab4f2416..8c2fa558edc 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -754,6 +754,11 @@ class Worker(WorkerBase): self.cache_config.num_gpu_blocks = num_blocks self.model_runner.extend_kv_cache(num_blocks) + def extensible_kv_cache_unsupported_reason(self) -> str | None: + from vllm.utils.vmm_driver import vmm_unavailable_reason + + return vmm_unavailable_reason() + @instrument(span_name="Warmup (GPU)") def compile_or_warm_up_model(self) -> CompilationTimes: warmup_sizes: list[int] = [] diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 22c362dd9b6..0ed8494a23a 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -108,6 +108,10 @@ class WorkerBase: f"{self.__class__.__name__} does not support extensible KV cache." ) + def extensible_kv_cache_unsupported_reason(self) -> str | None: + """Return why this worker cannot use the extensible KV cache, or None.""" + return f"not supported by {self.__class__.__name__}" + def compile_or_warm_up_model(self) -> CompilationTimes: """Prepare model for execution through compilation/warmup.