forked from Karylab-cklius/vllm
[Core] Extensible KV cache: VMM driver probe/fallback, manual size support
- Extract the driver ctypes bindings into vllm/utils/vmm_driver.py behind a small VmmDriver interface (CUDA implementation; struct layouts and call signatures are shared with HIP for a future ROCm backend). - Probe VMM support on the workers (driver loads, VA reservation works) and fall back to standard KV cache allocation with a warning instead of failing on platforms without VMM (e.g. WSL2, non-GPU workers). - Support kv_cache_memory_bytes: the requested size is committed as-is after warmup (single sizing pass), still avoiding warmup-time OOM. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D Signed-off-by: Nick Hill <nickhill@us.ibm.com>
This commit is contained in:
Nick Hill
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
a7fd4c7482
commit
fca040885b
+20
-199
@@ -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))
|
||||
|
||||
@@ -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
|
||||
+43
-39
@@ -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
|
||||
|
||||
@@ -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] = []
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user