| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Disk I/O backend for GPU<->NVMe block transfers via pinned staging buffers. |
| 4 | |
| 5 | Uses separate IO threads for store and load so that loads (latency-critical) |
| 6 | never block behind stores (background work). Each thread owns its own pinned |
| 7 | staging buffers to avoid contention. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import contextlib |
| 13 | import os |
| 14 | import queue |
| 15 | import threading |
| 16 | |
| 17 | import torch |
| 18 | |
| 19 | from vllm.logger import init_logger |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.v1.simple_kv_offload.cuda_mem_ops import ( |
| 22 | CU_MEMCPY_SRC_ACCESS_ORDER_ANY, |
| 23 | CU_MEMCPY_SRC_ACCESS_ORDER_STREAM, |
| 24 | BatchMemcpyParams, |
| 25 | build_params, |
| 26 | copy_blocks, |
| 27 | pin_tensor, |
| 28 | ) |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | O_DIRECT = getattr(os, "O_DIRECT", 0) |
| 33 | _ALIGNMENT = 4096 |
| 34 | |
| 35 | |
| 36 | def _alloc_aligned(num_slots: int, bpb: int) -> torch.Tensor: |
| 37 | """Allocate a staging buffer whose base address is O_DIRECT aligned. |
| 38 | |
| 39 | The CPU allocator only guarantees 64-byte alignment, so over-allocate by |
| 40 | one alignment unit and return an aligned view. The view keeps the backing |
| 41 | storage alive. |
| 42 | """ |
| 43 | nbytes = num_slots * bpb |
| 44 | raw = torch.zeros(nbytes + _ALIGNMENT, dtype=torch.int8, device="cpu") |
| 45 | offset = -raw.data_ptr() % _ALIGNMENT |
| 46 | return raw[offset : offset + nbytes].view(num_slots, bpb) |
| 47 | |
| 48 | |
| 49 | class DiskBackend: |
| 50 | """Async disk offload backend with pipelined GPU DMA and interleaved IO. |
| 51 | |
| 52 | Architecture: |
| 53 | - Separate coordinator threads for store and load (never block each other) |
| 54 | - Interleaved pipeline: DMA slot N while preadv/pwritev slot N-1 |
| 55 | - O_DIRECT by default; page cache is opt-in via use_page_cache |
| 56 | |
| 57 | Same launch_copy interface as DmaCopyBackend so the worker can swap |
| 58 | backends without changing calling code. |
| 59 | """ |
| 60 | |
| 61 | def __init__(self) -> None: |
| 62 | self._store_params: BatchMemcpyParams | None = None |
| 63 | self._load_params: BatchMemcpyParams | None = None |
| 64 | self._load_stream: torch.cuda.Stream | None = None |
| 65 | self._store_stream: torch.cuda.Stream | None = None |
| 66 | self._store_queue: queue.SimpleQueue = queue.SimpleQueue() |
| 67 | self._load_queue: queue.SimpleQueue = queue.SimpleQueue() |
| 68 | self._store_thread: threading.Thread | None = None |
| 69 | self._load_thread: threading.Thread | None = None |
| 70 | self._shutdown: bool = False |
| 71 | self._fd: int = -1 |
| 72 | self._disk_path: str = "" |
| 73 | self._total_block_bytes: int = 0 |
| 74 | self._store_buffer_caches: dict[str, torch.Tensor] = {} |
| 75 | self._load_buffer_caches: dict[str, torch.Tensor] = {} |
| 76 | self._store_slot_views: list[list[memoryview]] = [] |
| 77 | self._load_slot_views: list[list[memoryview]] = [] |
| 78 | self._per_tensor_bpb: list[int] = [] |
| 79 | self._tensor_names: list[str] = [] |
| 80 | |
| 81 | def init( |
| 82 | self, |
| 83 | gpu_caches: dict[str, torch.Tensor], |
| 84 | device: torch.device, |
| 85 | load_stream: torch.cuda.Stream, |
| 86 | store_stream: torch.cuda.Stream, |
| 87 | disk_path: str, |
| 88 | num_disk_slots: int, |
| 89 | total_block_bytes: int, |
| 90 | num_buffer_slots: int = 2, |
| 91 | use_page_cache: bool = False, |
| 92 | ) -> None: |
| 93 | self._load_stream = load_stream |
| 94 | self._store_stream = store_stream |
| 95 | self._total_block_bytes = total_block_bytes |
| 96 | self._num_buffer_slots = num_buffer_slots |
| 97 | self._tensor_names = list(gpu_caches.keys()) |
| 98 | self._per_tensor_bpb = [ |
| 99 | t.stride(0) * t.element_size() for t in gpu_caches.values() |
| 100 | ] |
| 101 | |
| 102 | # The stride alignment is an O_DIRECT kernel-DMA requirement; buffered |
| 103 | # I/O through the page cache has no alignment constraint. |
| 104 | if not use_page_cache: |
| 105 | assert total_block_bytes % _ALIGNMENT == 0, ( |
| 106 | f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}" |
| 107 | ) |
| 108 | |
| 109 | # Separate buffer pools for store and load threads |
| 110 | self._store_buffer_caches = {} |
| 111 | self._load_buffer_caches = {} |
| 112 | for name, gpu_t in gpu_caches.items(): |
| 113 | bpb = gpu_t.stride(0) * gpu_t.element_size() |
| 114 | store_buf = _alloc_aligned(num_buffer_slots, bpb) |
| 115 | pin_tensor(store_buf) |
| 116 | self._store_buffer_caches[name] = store_buf |
| 117 | load_buf = _alloc_aligned(num_buffer_slots, bpb) |
| 118 | pin_tensor(load_buf) |
| 119 | self._load_buffer_caches[name] = load_buf |
| 120 | |
| 121 | # Pre-built iovec views per slot (avoid per-transfer .numpy() calls) |
| 122 | self._store_slot_views = [ |
| 123 | [ |
| 124 | memoryview(self._store_buffer_caches[name][slot].numpy()) |
| 125 | for name in self._tensor_names |
| 126 | ] |
| 127 | for slot in range(num_buffer_slots) |
| 128 | ] |
| 129 | self._load_slot_views = [ |
| 130 | [ |
| 131 | memoryview(self._load_buffer_caches[name][slot].numpy()) |
| 132 | for name in self._tensor_names |
| 133 | ] |
| 134 | for slot in range(num_buffer_slots) |
| 135 | ] |
| 136 | |
| 137 | self._store_params = build_params( |
| 138 | gpu_caches, |
| 139 | self._store_buffer_caches, |
| 140 | store_stream, |
| 141 | src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM, |
| 142 | ) |
| 143 | self._load_params = build_params( |
| 144 | self._load_buffer_caches, |
| 145 | gpu_caches, |
| 146 | load_stream, |
| 147 | src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY, |
| 148 | ) |
| 149 | |
| 150 | os.makedirs(os.path.dirname(disk_path) or ".", exist_ok=True) |
| 151 | # Slot contents never outlive the process, so unlink then O_EXCL rather |
| 152 | # than reopening: a pre-existing file would otherwise keep its own |
| 153 | # (possibly world-readable) mode, and blocks may encode user prompts. |
| 154 | with contextlib.suppress(FileNotFoundError): |
| 155 | os.unlink(disk_path) |
| 156 | # O_DIRECT by default: page cache would consume the very host DRAM this |
| 157 | # backend exists to conserve, and doubles the copy on the store path. |
| 158 | flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW |
| 159 | if not use_page_cache: |
| 160 | flags |= O_DIRECT |
| 161 | self._fd = os.open(disk_path, flags, 0o600) |
| 162 | self._disk_path = disk_path |
| 163 | os.ftruncate(self._fd, num_disk_slots * total_block_bytes) |
| 164 | |
| 165 | logger.info( |
| 166 | "DiskBackend: path=%s, slots=%d, total=%.2f GB, buf=%dx%d bytes" |
| 167 | " (page_cache=%s)", |
| 168 | disk_path, |
| 169 | num_disk_slots, |
| 170 | (num_disk_slots * total_block_bytes) / (1024**3), |
| 171 | num_buffer_slots, |
| 172 | total_block_bytes, |
| 173 | use_page_cache, |
| 174 | ) |
| 175 | |
| 176 | self._store_thread = threading.Thread( |
| 177 | target=self._store_loop, |
| 178 | args=(device, store_stream), |
| 179 | daemon=True, |
| 180 | ) |
| 181 | self._load_thread = threading.Thread( |
| 182 | target=self._load_loop, |
| 183 | args=(device, load_stream), |
| 184 | daemon=True, |
| 185 | ) |
| 186 | self._store_thread.start() |
| 187 | self._load_thread.start() |
| 188 | |
| 189 | def launch_copy( |
| 190 | self, |
| 191 | src_blocks: list[int], |
| 192 | dst_blocks: list[int], |
| 193 | is_store: bool, |
| 194 | event_idx: int, |
| 195 | events_list: list[tuple[int, torch.Event]], |
| 196 | wait_event: torch.Event | None = None, |
| 197 | ) -> None: |
| 198 | q = self._store_queue if is_store else self._load_queue |
| 199 | q.put((src_blocks, dst_blocks, event_idx, events_list, wait_event)) |
| 200 | |
| 201 | def shutdown(self) -> None: |
| 202 | if self._shutdown: |
| 203 | return |
| 204 | self._shutdown = True |
| 205 | self._store_queue.put(None) |
| 206 | self._load_queue.put(None) |
| 207 | if self._store_thread is not None: |
| 208 | self._store_thread.join(timeout=10.0) |
| 209 | if self._load_thread is not None: |
| 210 | self._load_thread.join(timeout=10.0) |
| 211 | if self._fd < 0: |
| 212 | return |
| 213 | # Slot contents can encode user prompts, so drop the name now rather |
| 214 | # than leaving them readable until the next run overwrites the file. |
| 215 | # Unlinking only removes the directory entry: any thread still holding |
| 216 | # the fd keeps writing to the (now anonymous) inode, which the kernel |
| 217 | # frees once the last fd goes away. |
| 218 | with contextlib.suppress(OSError): |
| 219 | os.unlink(self._disk_path) |
| 220 | # Closing under a still-running IO thread would let the fd number be |
| 221 | # reused by an unrelated open(), turning its next pwritev into a write |
| 222 | # into that file. Leaking one fd for the remaining process lifetime is |
| 223 | # the cheaper failure. |
| 224 | if any( |
| 225 | t is not None and t.is_alive() |
| 226 | for t in (self._store_thread, self._load_thread) |
| 227 | ): |
| 228 | logger.warning( |
| 229 | "IO thread still running after shutdown timeout; leaking fd %d", |
| 230 | self._fd, |
| 231 | ) |
| 232 | return |
| 233 | os.close(self._fd) |
| 234 | self._fd = -1 |
| 235 | |
| 236 | def _store_loop( |
| 237 | self, |
| 238 | device: torch.device, |
| 239 | stream: torch.cuda.Stream, |
| 240 | ) -> None: |
| 241 | current_platform.set_device(device) |
| 242 | while True: |
| 243 | item = self._store_queue.get() |
| 244 | if item is None: |
| 245 | return |
| 246 | (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item |
| 247 | if wait_event is not None: |
| 248 | stream.wait_event(wait_event) |
| 249 | self._do_store(src_blocks, dst_blocks, stream) |
| 250 | event = torch.Event() |
| 251 | event.record(stream) |
| 252 | events_list.append((event_idx, event)) |
| 253 | |
| 254 | def _writev_slot(self, buf_slot: int, file_offset: int) -> None: |
| 255 | written = os.pwritev(self._fd, self._store_slot_views[buf_slot], file_offset) |
| 256 | if written < self._total_block_bytes: |
| 257 | raise OSError( |
| 258 | f"Short write: expected {self._total_block_bytes} bytes, " |
| 259 | f"wrote {written}" |
| 260 | ) |
| 261 | |
| 262 | def _readv_slot(self, buf_slot: int, file_offset: int) -> None: |
| 263 | bytes_read = os.preadv(self._fd, self._load_slot_views[buf_slot], file_offset) |
| 264 | if bytes_read < self._total_block_bytes: |
| 265 | raise OSError( |
| 266 | f"Short read: expected {self._total_block_bytes} bytes, " |
| 267 | f"read {bytes_read}" |
| 268 | ) |
| 269 | |
| 270 | def _load_loop( |
| 271 | self, |
| 272 | device: torch.device, |
| 273 | stream: torch.cuda.Stream, |
| 274 | ) -> None: |
| 275 | current_platform.set_device(device) |
| 276 | while True: |
| 277 | item = self._load_queue.get() |
| 278 | if item is None: |
| 279 | return |
| 280 | (src_blocks, dst_blocks, event_idx, events_list, wait_event) = item |
| 281 | if wait_event is not None: |
| 282 | stream.wait_event(wait_event) |
| 283 | self._do_load(src_blocks, dst_blocks, stream) |
| 284 | event = torch.Event() |
| 285 | event.record(stream) |
| 286 | events_list.append((event_idx, event)) |
| 287 | |
| 288 | def _do_store( |
| 289 | self, |
| 290 | gpu_blocks: list[int], |
| 291 | disk_slots: list[int], |
| 292 | stream: torch.cuda.Stream, |
| 293 | ) -> None: |
| 294 | """GPU -> buffer (DMA) -> disk (pwritev), interleaved double-buffer.""" |
| 295 | assert self._store_params is not None |
| 296 | n = self._num_buffer_slots |
| 297 | # (DMA event, file offset) of the block already staged in each slot. |
| 298 | pending: list[tuple[torch.Event, int] | None] = [None] * n |
| 299 | |
| 300 | for i, (gpu_blk, disk_slot) in enumerate(zip(gpu_blocks, disk_slots)): |
| 301 | buf_slot = i % n |
| 302 | prev = pending[buf_slot] |
| 303 | if prev is not None: |
| 304 | prev[0].synchronize() |
| 305 | self._writev_slot(buf_slot, prev[1]) |
| 306 | |
| 307 | copy_blocks([gpu_blk], [buf_slot], self._store_params) |
| 308 | ev = torch.Event() |
| 309 | ev.record(stream) |
| 310 | pending[buf_slot] = (ev, disk_slot * self._total_block_bytes) |
| 311 | |
| 312 | for slot, last in enumerate(pending): |
| 313 | if last is not None: |
| 314 | last[0].synchronize() |
| 315 | self._writev_slot(slot, last[1]) |
| 316 | |
| 317 | def _do_load( |
| 318 | self, |
| 319 | disk_slots: list[int], |
| 320 | gpu_blocks: list[int], |
| 321 | stream: torch.cuda.Stream, |
| 322 | ) -> None: |
| 323 | """Disk (preadv) -> buffer -> GPU (DMA), interleaved double-buffer.""" |
| 324 | assert self._load_params is not None |
| 325 | n = self._num_buffer_slots |
| 326 | prev_dma_events: list[torch.Event | None] = [None] * n |
| 327 | |
| 328 | for i, (disk_slot, gpu_blk) in enumerate(zip(disk_slots, gpu_blocks)): |
| 329 | buf_slot = i % n |
| 330 | prev = prev_dma_events[buf_slot] |
| 331 | if prev is not None: |
| 332 | prev.synchronize() |
| 333 | |
| 334 | self._readv_slot(buf_slot, disk_slot * self._total_block_bytes) |
| 335 | |
| 336 | copy_blocks([buf_slot], [gpu_blk], self._load_params) |
| 337 | ev = torch.Event() |
| 338 | ev.record(stream) |
| 339 | prev_dma_events[buf_slot] = ev |
| 340 | |
| 341 | for ev in prev_dma_events: |
| 342 | if ev is not None: |
| 343 | ev.synchronize() |