diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 3ed34f169d..e198192217 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -41,6 +41,15 @@ logger = init_logger(__name__) # Default CPU capacity: 8 GB DEFAULT_CPU_CAPACITY_BYTES = 8 * (1024**3) +VALID_KV_OFFLOAD_BACKENDS = ("cpu", "disk") +# Keys that only apply to the disk backend, warned about under "cpu". +_DISK_ONLY_KEYS = ( + "disk_path", + "disk_capacity_bytes", + "disk_buffer_slots", + "use_page_cache", +) + class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): """CPU KV cache offloading with custom kernel transfers and BlockPool LRU.""" @@ -82,6 +91,36 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): lazy_offload = bool(extra_config.get("lazy_offload", False)) + kv_offload_backend = str(extra_config.get("kv_offload_backend", "cpu")) + if kv_offload_backend not in VALID_KV_OFFLOAD_BACKENDS: + raise ValueError( + f"Unknown kv_offload_backend {kv_offload_backend!r}; " + f"expected one of {VALID_KV_OFFLOAD_BACKENDS}" + ) + disk_mode = kv_offload_backend == "disk" + + disk_path = extra_config.get("disk_path", None) or None + disk_capacity_bytes = int( + extra_config.get("disk_capacity_bytes", 100 * (1024**3)) + ) + disk_buffer_slots = max(1, int(extra_config.get("disk_buffer_slots", 2))) + use_page_cache = bool(extra_config.get("use_page_cache", False)) + + if disk_mode: + if disk_path is None: + raise ValueError( + 'kv_offload_backend="disk" requires disk_path to be set.' + ) + else: + ignored = [k for k in _DISK_ONLY_KEYS if k in extra_config] + if ignored: + logger.warning( + 'kv_offload_backend is "cpu", ignoring disk-only config: %s. ' + 'Set kv_offload_backend="disk" to enable the disk backend.', + ", ".join(ignored), + ) + disk_path = None + self.scheduler_manager: SimpleCPUOffloadScheduler | None = None self.worker_handler: SimpleCPUOffloadWorker | None = None @@ -94,11 +133,13 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): logger.info( "SimpleCPUOffloadConnector: role=%s, " - "per_rank=%.2f GB, world_size=%d, mode=%s", + "per_rank=%.2f GB, world_size=%d, mode=%s, backend=%s, disk=%s", role.name, cpu_capacity_per_rank / (1024**3), world_size, "lazy" if lazy_offload else "eager", + kv_offload_backend, + disk_path or "none", ) if role == KVConnectorRole.SCHEDULER: @@ -115,10 +156,18 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, lazy_offload=lazy_offload, + disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0, ) elif role == KVConnectorRole.WORKER: self.worker_handler = SimpleCPUOffloadWorker( - vllm_config, kv_cache_config, cpu_capacity_per_rank + vllm_config, + kv_cache_config, + cpu_capacity_per_rank, + kv_offload_backend=kv_offload_backend, + disk_path=disk_path, + disk_capacity_bytes=disk_capacity_bytes, + disk_buffer_slots=disk_buffer_slots, + use_page_cache=use_page_cache, ) # --- Worker-side methods --- diff --git a/vllm/v1/simple_kv_offload/disk_backend.py b/vllm/v1/simple_kv_offload/disk_backend.py new file mode 100644 index 0000000000..f48dfd46bb --- /dev/null +++ b/vllm/v1/simple_kv_offload/disk_backend.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Disk I/O backend for GPU<->NVMe block transfers via pinned staging buffers. + +Uses separate IO threads for store and load so that loads (latency-critical) +never block behind stores (background work). Each thread owns its own pinned +staging buffers to avoid contention. +""" + +from __future__ import annotations + +import contextlib +import os +import queue +import threading + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.v1.simple_kv_offload.cuda_mem_ops import ( + CU_MEMCPY_SRC_ACCESS_ORDER_ANY, + CU_MEMCPY_SRC_ACCESS_ORDER_STREAM, + BatchMemcpyParams, + build_params, + copy_blocks, + pin_tensor, +) + +logger = init_logger(__name__) + +O_DIRECT = getattr(os, "O_DIRECT", 0) +_ALIGNMENT = 4096 + + +def _alloc_aligned(num_slots: int, bpb: int) -> torch.Tensor: + """Allocate a staging buffer whose base address is O_DIRECT aligned. + + The CPU allocator only guarantees 64-byte alignment, so over-allocate by + one alignment unit and return an aligned view. The view keeps the backing + storage alive. + """ + nbytes = num_slots * bpb + raw = torch.zeros(nbytes + _ALIGNMENT, dtype=torch.int8, device="cpu") + offset = -raw.data_ptr() % _ALIGNMENT + return raw[offset : offset + nbytes].view(num_slots, bpb) + + +class DiskBackend: + """Async disk offload backend with pipelined GPU DMA and interleaved IO. + + Architecture: + - Separate coordinator threads for store and load (never block each other) + - Interleaved pipeline: DMA slot N while preadv/pwritev slot N-1 + - O_DIRECT by default; page cache is opt-in via use_page_cache + + Same launch_copy interface as DmaCopyBackend so the worker can swap + backends without changing calling code. + """ + + def __init__(self) -> None: + self._store_params: BatchMemcpyParams | None = None + self._load_params: BatchMemcpyParams | None = None + self._load_stream: torch.cuda.Stream | None = None + self._store_stream: torch.cuda.Stream | None = None + self._store_queue: queue.SimpleQueue = queue.SimpleQueue() + self._load_queue: queue.SimpleQueue = queue.SimpleQueue() + self._store_thread: threading.Thread | None = None + self._load_thread: threading.Thread | None = None + self._shutdown: bool = False + self._fd: int = -1 + self._disk_path: str = "" + self._total_block_bytes: int = 0 + self._store_buffer_caches: dict[str, torch.Tensor] = {} + self._load_buffer_caches: dict[str, torch.Tensor] = {} + self._store_slot_views: list[list[memoryview]] = [] + self._load_slot_views: list[list[memoryview]] = [] + self._per_tensor_bpb: list[int] = [] + self._tensor_names: list[str] = [] + + def init( + self, + gpu_caches: dict[str, torch.Tensor], + device: torch.device, + load_stream: torch.cuda.Stream, + store_stream: torch.cuda.Stream, + disk_path: str, + num_disk_slots: int, + total_block_bytes: int, + num_buffer_slots: int = 2, + use_page_cache: bool = False, + ) -> None: + self._load_stream = load_stream + self._store_stream = store_stream + self._total_block_bytes = total_block_bytes + self._num_buffer_slots = num_buffer_slots + self._tensor_names = list(gpu_caches.keys()) + self._per_tensor_bpb = [ + t.stride(0) * t.element_size() for t in gpu_caches.values() + ] + + assert total_block_bytes % _ALIGNMENT == 0, ( + f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}" + ) + + # Separate buffer pools for store and load threads + self._store_buffer_caches = {} + self._load_buffer_caches = {} + for name, gpu_t in gpu_caches.items(): + bpb = gpu_t.stride(0) * gpu_t.element_size() + store_buf = _alloc_aligned(num_buffer_slots, bpb) + pin_tensor(store_buf) + self._store_buffer_caches[name] = store_buf + load_buf = _alloc_aligned(num_buffer_slots, bpb) + pin_tensor(load_buf) + self._load_buffer_caches[name] = load_buf + + # Pre-built iovec views per slot (avoid per-transfer .numpy() calls) + self._store_slot_views = [ + [ + memoryview(self._store_buffer_caches[name][slot].numpy()) + for name in self._tensor_names + ] + for slot in range(num_buffer_slots) + ] + self._load_slot_views = [ + [ + memoryview(self._load_buffer_caches[name][slot].numpy()) + for name in self._tensor_names + ] + for slot in range(num_buffer_slots) + ] + + self._store_params = build_params( + gpu_caches, + self._store_buffer_caches, + store_stream, + src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM, + ) + self._load_params = build_params( + self._load_buffer_caches, + gpu_caches, + load_stream, + src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY, + ) + + os.makedirs(os.path.dirname(disk_path) or ".", exist_ok=True) + # Slot contents never outlive the process, so unlink then O_EXCL rather + # than reopening: a pre-existing file would otherwise keep its own + # (possibly world-readable) mode, and blocks may encode user prompts. + with contextlib.suppress(FileNotFoundError): + os.unlink(disk_path) + # O_DIRECT by default: page cache would consume the very host DRAM this + # backend exists to conserve, and doubles the copy on the store path. + flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW + if not use_page_cache: + flags |= O_DIRECT + self._fd = os.open(disk_path, flags, 0o600) + self._disk_path = disk_path + os.ftruncate(self._fd, num_disk_slots * total_block_bytes) + + logger.info( + "DiskBackend: path=%s, slots=%d, total=%.2f GB, buf=%dx%d bytes" + " (page_cache=%s)", + disk_path, + num_disk_slots, + (num_disk_slots * total_block_bytes) / (1024**3), + num_buffer_slots, + total_block_bytes, + use_page_cache, + ) + + self._store_thread = threading.Thread( + target=self._store_loop, + args=(device, store_stream), + daemon=True, + ) + self._load_thread = threading.Thread( + target=self._load_loop, + args=(device, load_stream), + daemon=True, + ) + self._store_thread.start() + self._load_thread.start() + + def launch_copy( + self, + src_blocks: list[int], + dst_blocks: list[int], + is_store: bool, + event_idx: int, + events_list: list[tuple[int, torch.Event]], + wait_event: torch.Event | None = None, + ) -> None: + q = self._store_queue if is_store else self._load_queue + q.put((src_blocks, dst_blocks, event_idx, events_list, wait_event)) + + def shutdown(self) -> None: + if self._shutdown: + return + self._shutdown = True + self._store_queue.put(None) + self._load_queue.put(None) + if self._store_thread is not None: + self._store_thread.join(timeout=10.0) + if self._load_thread is not None: + self._load_thread.join(timeout=10.0) + if self._fd < 0: + return + # Slot contents can encode user prompts, so drop the name now rather + # than leaving them readable until the next run overwrites the file. + # Unlinking only removes the directory entry: any thread still holding + # the fd keeps writing to the (now anonymous) inode, which the kernel + # frees once the last fd goes away. + with contextlib.suppress(OSError): + os.unlink(self._disk_path) + # Closing under a still-running IO thread would let the fd number be + # reused by an unrelated open(), turning its next pwritev into a write + # into that file. Leaking one fd for the remaining process lifetime is + # the cheaper failure. + if any( + t is not None and t.is_alive() + for t in (self._store_thread, self._load_thread) + ): + logger.warning( + "IO thread still running after shutdown timeout; leaking fd %d", + self._fd, + ) + return + os.close(self._fd) + self._fd = -1 + + def _store_loop( + self, + device: torch.device, + stream: torch.cuda.Stream, + ) -> None: + current_platform.set_device(device) + while True: + item = self._store_queue.get() + if item is None: + return + (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item + if wait_event is not None: + stream.wait_event(wait_event) + self._do_store(src_blocks, dst_blocks, stream) + event = torch.Event() + event.record(stream) + events_list.append((event_idx, event)) + + def _writev_slot(self, buf_slot: int, file_offset: int) -> None: + written = os.pwritev(self._fd, self._store_slot_views[buf_slot], file_offset) + if written < self._total_block_bytes: + raise OSError( + f"Short write: expected {self._total_block_bytes} bytes, " + f"wrote {written}" + ) + + def _readv_slot(self, buf_slot: int, file_offset: int) -> None: + bytes_read = os.preadv(self._fd, self._load_slot_views[buf_slot], file_offset) + if bytes_read < self._total_block_bytes: + raise OSError( + f"Short read: expected {self._total_block_bytes} bytes, " + f"read {bytes_read}" + ) + + def _load_loop( + self, + device: torch.device, + stream: torch.cuda.Stream, + ) -> None: + current_platform.set_device(device) + while True: + item = self._load_queue.get() + if item is None: + return + (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item + if wait_event is not None: + stream.wait_event(wait_event) + self._do_load(src_blocks, dst_blocks, stream) + event = torch.Event() + event.record(stream) + events_list.append((event_idx, event)) + + def _do_store( + self, + gpu_blocks: list[int], + disk_slots: list[int], + stream: torch.cuda.Stream, + ) -> None: + """GPU -> buffer (DMA) -> disk (pwritev), interleaved double-buffer.""" + assert self._store_params is not None + n = self._num_buffer_slots + # (DMA event, file offset) of the block already staged in each slot. + pending: list[tuple[torch.Event, int] | None] = [None] * n + + for i, (gpu_blk, disk_slot) in enumerate(zip(gpu_blocks, disk_slots)): + buf_slot = i % n + prev = pending[buf_slot] + if prev is not None: + prev[0].synchronize() + self._writev_slot(buf_slot, prev[1]) + + copy_blocks([gpu_blk], [buf_slot], self._store_params) + ev = torch.Event() + ev.record(stream) + pending[buf_slot] = (ev, disk_slot * self._total_block_bytes) + + for slot, last in enumerate(pending): + if last is not None: + last[0].synchronize() + self._writev_slot(slot, last[1]) + + def _do_load( + self, + disk_slots: list[int], + gpu_blocks: list[int], + stream: torch.cuda.Stream, + ) -> None: + """Disk (preadv) -> buffer -> GPU (DMA), interleaved double-buffer.""" + assert self._load_params is not None + n = self._num_buffer_slots + prev_dma_events: list[torch.Event | None] = [None] * n + + for i, (disk_slot, gpu_blk) in enumerate(zip(disk_slots, gpu_blocks)): + buf_slot = i % n + prev = prev_dma_events[buf_slot] + if prev is not None: + prev.synchronize() + + self._readv_slot(buf_slot, disk_slot * self._total_block_bytes) + + copy_blocks([buf_slot], [gpu_blk], self._load_params) + ev = torch.Event() + ev.record(stream) + prev_dma_events[buf_slot] = ev diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 67c6f66463..2e6839fef5 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -75,9 +75,14 @@ class SimpleCPUOffloadScheduler: scheduler_block_size: int, hash_block_size: int, lazy_offload: bool = False, + disk_capacity_bytes: int = 0, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config + # When disk mode is active, the offload pool size is disk-based. + offload_capacity = ( + disk_capacity_bytes if disk_capacity_bytes > 0 else cpu_capacity_bytes + ) self.enable_kv_cache_events = ( vllm_config.kv_events_config is not None and vllm_config.kv_events_config.enable_kv_cache_events @@ -90,7 +95,7 @@ class SimpleCPUOffloadScheduler: # Derive a CPU KVCacheConfig from the GPU config and build a coordinator assert kv_cache_config is not None self.cpu_kv_cache_config = self._derive_cpu_config( - kv_cache_config, cpu_capacity_bytes + kv_cache_config, offload_capacity ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks # Find the full attention kv group for prefix cache matching. @@ -111,10 +116,12 @@ class SimpleCPUOffloadScheduler: assert self.block_size % self.fa_block_size == 0 logger.info( - "SimpleCPUOffloadScheduler: Allocating %d CPU blocks (%.2f GB, mode=%s)", + "SimpleCPUOffloadScheduler: Allocating %d offload blocks " + "(%.2f GB, mode=%s, backend=%s)", self.num_cpu_blocks, - cpu_capacity_bytes / (1024**3), + offload_capacity / (1024**3), "lazy" if lazy_offload else "eager", + "disk" if disk_capacity_bytes > 0 else "cpu", ) spec_config = vllm_config.speculative_config diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index 9cb9c02ed7..e519bca261 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -11,6 +11,7 @@ from vllm.logger import init_logger from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor +from vllm.v1.simple_kv_offload.disk_backend import DiskBackend from vllm.v1.simple_kv_offload.metadata import ( SimpleCPUOffloadMetadata, SimpleCPUOffloadWorkerMetadata, @@ -30,10 +31,20 @@ class SimpleCPUOffloadWorker: vllm_config: VllmConfig, kv_cache_config: "KVCacheConfig | None", cpu_capacity_bytes: int, + kv_offload_backend: str = "cpu", + disk_path: str | None = None, + disk_capacity_bytes: int = 0, + disk_buffer_slots: int = 2, + use_page_cache: bool = False, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config self.cpu_capacity_bytes = cpu_capacity_bytes + self.disk_path = disk_path + self.disk_capacity_bytes = disk_capacity_bytes + self.disk_buffer_slots = disk_buffer_slots + self.use_page_cache = use_page_cache + self.disk_mode = kv_offload_backend == "disk" self.gpu_kv_caches: dict[str, torch.Tensor] | None = None self.cpu_kv_caches: dict[str, torch.Tensor] | None = None @@ -44,7 +55,7 @@ class SimpleCPUOffloadWorker: self.load_stream: torch.cuda.Stream | None = None self.store_stream: torch.cuda.Stream | None = None - self._backend = DmaCopyBackend() + self._backend: DmaCopyBackend | DiskBackend | None = None # Ordered (event_idx, Event). Events pre-allocated on main thread. self._load_events: list[tuple[int, torch.Event]] = [] @@ -145,9 +156,57 @@ class SimpleCPUOffloadWorker: self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block) + # Use lowest priority so KV cache I/O yields to compute streams. + low_pri, _ = torch.cuda.Stream.priority_range() + self.load_stream = torch.cuda.Stream(priority=low_pri) + self.store_stream = torch.cuda.Stream(priority=low_pri) + + self.gpu_kv_caches = unique_gpu_caches + + if self.disk_mode: + self._init_disk_mode(unique_gpu_caches, total_bytes_per_block, self.device) + else: + self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device) + + def _init_disk_mode( + self, + unique_gpu_caches: dict[str, torch.Tensor], + total_bytes_per_block: int, + device: torch.device, + ) -> None: + num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block) + self.num_cpu_blocks = num_disk_slots + logger.info( - "SimpleCPUOffloadWorker: %d unique GPU KV tensors, " - "allocating %d CPU blocks (%.2f GB)", + "SimpleCPUOffloadWorker [DISK]: %d tensors, %d disk slots (%.2f GB)", + len(unique_gpu_caches), + num_disk_slots, + (num_disk_slots * total_bytes_per_block) / (1024**3), + ) + + assert self.disk_path is not None + rank_path = f"{self.disk_path}.rank_{device.index or 0}" + self._backend = DiskBackend() + self._backend.init( + unique_gpu_caches, + device, + self.load_stream, + self.store_stream, + rank_path, + num_disk_slots, + total_bytes_per_block, + self.disk_buffer_slots, + self.use_page_cache, + ) + + def _init_cpu_mode( + self, + unique_gpu_caches: dict[str, torch.Tensor], + total_bytes_per_block: int, + device: torch.device, + ) -> None: + logger.info( + "SimpleCPUOffloadWorker [CPU]: %d tensors, %d CPU blocks (%.2f GB)", len(unique_gpu_caches), self.num_cpu_blocks, (self.num_cpu_blocks * total_bytes_per_block) / (1024**3), @@ -159,7 +218,6 @@ class SimpleCPUOffloadWorker: "Pinned memory not available. CPU offload performance may be degraded." ) - self.gpu_kv_caches = unique_gpu_caches self.cpu_kv_caches = {} for name, gpu_tensor in unique_gpu_caches.items(): cpu_shape = (self.num_cpu_blocks,) + gpu_tensor.shape[1:] @@ -171,16 +229,11 @@ class SimpleCPUOffloadWorker: pin_tensor(tensor) self.cpu_kv_caches[name] = tensor - # Use lowest priority so KV cache I/O yields to compute streams. - low_pri, _ = torch.cuda.Stream.priority_range() - self.load_stream = torch.cuda.Stream(priority=low_pri) - self.store_stream = torch.cuda.Stream(priority=low_pri) - - # Initialize copy backend with caches and streams. + self._backend = DmaCopyBackend() self._backend.init( - self.gpu_kv_caches, + unique_gpu_caches, self.cpu_kv_caches, - self.device, + device, self.load_stream, self.store_stream, ) @@ -224,8 +277,10 @@ class SimpleCPUOffloadWorker: # (1) Submit transfers metadata = self._connector_metadata if metadata is not None: + backend = self._backend + assert backend is not None if metadata.load_cpu_blocks: - self._backend.launch_copy( + backend.launch_copy( metadata.load_cpu_blocks, metadata.load_gpu_blocks, is_store=False, @@ -236,7 +291,7 @@ class SimpleCPUOffloadWorker: if self._store_compute_done is None: self._store_compute_done = torch.Event() self._store_compute_done.record(torch.cuda.current_stream()) - self._backend.launch_copy( + backend.launch_copy( metadata.store_gpu_blocks, metadata.store_cpu_blocks, is_store=True,