| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import errno |
| 5 | import mmap |
| 6 | import os |
| 7 | import time |
| 8 | from collections.abc import Callable |
| 9 | |
| 10 | import numpy as np |
| 11 | import torch |
| 12 | |
| 13 | from vllm.distributed.device_communicators.shm_broadcast import ( |
| 14 | check_shm_free_space, |
| 15 | ) |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | logger = init_logger(__name__) |
| 20 | |
| 21 | # MADV_POPULATE_WRITE was added in Linux 5.14 (value 23). |
| 22 | _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) |
| 23 | |
| 24 | |
| 25 | def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None: |
| 26 | """Spin-wait until the file reaches expected_size (creator truncated it).""" |
| 27 | deadline = time.monotonic() + timeout |
| 28 | while True: |
| 29 | if os.fstat(fd).st_size >= expected_size: |
| 30 | return |
| 31 | if os.fstat(fd).st_nlink == 0: |
| 32 | raise RuntimeError( |
| 33 | "Shared offload region creator failed during initialization." |
| 34 | ) |
| 35 | if time.monotonic() > deadline: |
| 36 | raise TimeoutError( |
| 37 | f"Timed out waiting for mmap file to reach {expected_size} bytes" |
| 38 | ) |
| 39 | time.sleep(0.005) |
| 40 | |
| 41 | |
| 42 | def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: |
| 43 | mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length) |
| 44 | |
| 45 | |
| 46 | def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: |
| 47 | # Touch one byte per page via a read-modify-write so existing bytes are |
| 48 | # preserved — a peer worker may have already written KV data into this |
| 49 | # shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE. |
| 50 | arr = np.frombuffer(mmap_obj, dtype=np.uint8) |
| 51 | arr[offset : offset + length : mmap.PAGESIZE] |= 0 |
| 52 | |
| 53 | |
| 54 | def _get_populate_write_fn( |
| 55 | mmap_obj: mmap.mmap, |
| 56 | ) -> Callable[[mmap.mmap, int, int], None]: |
| 57 | """Select the pre-faulting method once for this mmap.""" |
| 58 | try: |
| 59 | _madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE) |
| 60 | except OSError as e: |
| 61 | if e.errno != errno.EINVAL: |
| 62 | raise |
| 63 | logger.warning( |
| 64 | "MADV_POPULATE_WRITE is not supported; falling back to per-page " |
| 65 | "writes for mmap pre-population. Startup may be slower." |
| 66 | ) |
| 67 | return _fallback_populate_write |
| 68 | return _madvise_populate_write |
| 69 | |
| 70 | |
| 71 | class SharedOffloadRegion: |
| 72 | """Single mmap-backed memory region shared across all workers for a |
| 73 | vLLM instance. Workers coordinate via the filesystem: the first worker |
| 74 | to open the file with O_EXCL becomes the creator and calls ftruncate; |
| 75 | the rest open the existing file and wait until it reaches the expected |
| 76 | size. Each worker then mmap()s the full file. |
| 77 | |
| 78 | File path: /dev/shm/vllm_offload_{engine_id}.mmap. When a barrier is |
| 79 | given, the path is unlinked once every worker has mapped the file, so |
| 80 | the kernel reclaims the memory when the last worker exits, no matter |
| 81 | how it exits; mappings taken before the unlink stay valid. |
| 82 | |
| 83 | Creator-only population pre-faults the entire region before the barrier |
| 84 | and requires that barrier to keep joiners from using unpopulated pages. |
| 85 | """ |
| 86 | |
| 87 | BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE |
| 88 | |
| 89 | def __init__( |
| 90 | self, |
| 91 | engine_id: str, |
| 92 | num_chunks: int, |
| 93 | rank: int | None, |
| 94 | kv_bytes_per_chunk: int, |
| 95 | cpu_page_size: int, |
| 96 | barrier: Callable[[], None] | None = None, |
| 97 | *, |
| 98 | creator_memory_check: Callable[[int], None] | None = None, |
| 99 | populate_only_on_creator: bool = False, |
| 100 | ) -> None: |
| 101 | if populate_only_on_creator and barrier is None: |
| 102 | raise ValueError("Creator-only population requires a barrier.") |
| 103 | self.page_size = mmap.PAGESIZE |
| 104 | assert kv_bytes_per_chunk % self.page_size == 0 |
| 105 | |
| 106 | self.num_chunks = num_chunks |
| 107 | self._row_stride = kv_bytes_per_chunk |
| 108 | self.total_size_bytes = self.num_chunks * self._row_stride |
| 109 | |
| 110 | self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap" |
| 111 | self._creator = False # set True only if this worker creates the file |
| 112 | self.rank = rank |
| 113 | if rank is not None: |
| 114 | # byte offset to this worker's first slot within each chunk row |
| 115 | self._worker_offset = rank * cpu_page_size |
| 116 | # exclusive upper bound for this worker's area within each row |
| 117 | self._worker_area_end = (rank + 1) * cpu_page_size |
| 118 | try: |
| 119 | try: |
| 120 | self.fd: int | None = os.open( |
| 121 | self.mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600 |
| 122 | ) |
| 123 | except FileExistsError: |
| 124 | # Joiner path — another worker won O_EXCL. Reopen and wait |
| 125 | # for the file to reach expected size. |
| 126 | self.fd = os.open(self.mmap_path, os.O_RDWR) |
| 127 | _wait_for_file_size(self.fd, self.total_size_bytes) |
| 128 | logger.info("Opened existing mmap file %s", self.mmap_path) |
| 129 | else: |
| 130 | # Creator path. We won O_EXCL, so we own the file: any |
| 131 | # failure here must clean up so concurrent joiners don't |
| 132 | # land on a 0-byte stub and spin in _wait_for_file_size |
| 133 | # for the full 30 s timeout. |
| 134 | self._creator = True |
| 135 | if creator_memory_check is not None: |
| 136 | creator_memory_check(self.total_size_bytes) |
| 137 | check_shm_free_space( |
| 138 | self.total_size_bytes, |
| 139 | allocation_name="CPU KV offload shared region in /dev/shm", |
| 140 | ) |
| 141 | os.ftruncate(self.fd, self.total_size_bytes) |
| 142 | logger.info( |
| 143 | "Created mmap file %s (%.2f GB)", |
| 144 | self.mmap_path, |
| 145 | self.total_size_bytes / 1e9, |
| 146 | ) |
| 147 | |
| 148 | self.mmap_obj: mmap.mmap | None = mmap.mmap( |
| 149 | self.fd, |
| 150 | self.total_size_bytes, |
| 151 | flags=mmap.MAP_SHARED, |
| 152 | prot=mmap.PROT_READ | mmap.PROT_WRITE, |
| 153 | ) |
| 154 | |
| 155 | if populate_only_on_creator and self._creator: |
| 156 | populate_write_fn = _get_populate_write_fn(self.mmap_obj) |
| 157 | populate_write_fn(self.mmap_obj, 0, self.total_size_bytes) |
| 158 | except Exception: |
| 159 | if self._creator: |
| 160 | with contextlib.suppress(FileNotFoundError): |
| 161 | os.unlink(self.mmap_path) |
| 162 | self._creator = False |
| 163 | if hasattr(self, "mmap_obj") and self.mmap_obj is not None: |
| 164 | self.mmap_obj.close() |
| 165 | self.mmap_obj = None |
| 166 | if hasattr(self, "fd") and self.fd is not None: |
| 167 | os.close(self.fd) |
| 168 | self.fd = None |
| 169 | # Peers block inside the barrier until the collective times out if |
| 170 | # we die before reaching it. Arrive anyway so every worker calls |
| 171 | # barrier() exactly once and they fail on their own errors instead |
| 172 | # of hanging; a failure here must not replace ours. |
| 173 | if barrier is not None: |
| 174 | try: |
| 175 | barrier() |
| 176 | except Exception: |
| 177 | logger.warning( |
| 178 | "Failed to release peers waiting at the mmap barrier", |
| 179 | exc_info=True, |
| 180 | ) |
| 181 | raise |
| 182 | |
| 183 | if barrier is not None: |
| 184 | # Every worker has mapped the file once the barrier releases, so |
| 185 | # its name is no longer needed and dropping it here means no exit |
| 186 | # path — including SIGKILL — can leak the file. |
| 187 | try: |
| 188 | barrier() |
| 189 | except Exception: |
| 190 | if self._creator: |
| 191 | with contextlib.suppress(FileNotFoundError): |
| 192 | os.unlink(self.mmap_path) |
| 193 | self._creator = False |
| 194 | self.mmap_obj.close() |
| 195 | os.close(self.fd) |
| 196 | self.mmap_obj = None |
| 197 | self.fd = None |
| 198 | raise |
| 199 | if self._creator: |
| 200 | os.unlink(self.mmap_path) |
| 201 | self._creator = False |
| 202 | logger.info("Unlinked mmap file %s", self.mmap_path) |
| 203 | |
| 204 | self._base = torch.frombuffer(memoryview(self.mmap_obj), dtype=torch.int8) |
| 205 | self._views: list[torch.Tensor] = [] |
| 206 | self._canonical_offset = 0 |
| 207 | self.is_pinned: bool = False |
| 208 | self.pinned_addresses: list[int] = [] |
| 209 | |
| 210 | if populate_only_on_creator: |
| 211 | return |
| 212 | |
| 213 | populate_write_fn = _get_populate_write_fn(self.mmap_obj) |
| 214 | |
| 215 | if rank is not None: |
| 216 | # Populate only this worker's pages (one slot per chunk row). |
| 217 | worker_offset = rank * cpu_page_size |
| 218 | _t0 = time.perf_counter() |
| 219 | page_size = self.page_size |
| 220 | for chunk in range(num_chunks): |
| 221 | raw_offset = chunk * self._row_stride + worker_offset |
| 222 | aligned_offset = (raw_offset // page_size) * page_size |
| 223 | end = raw_offset + cpu_page_size |
| 224 | aligned_length = end - aligned_offset |
| 225 | populate_write_fn(self.mmap_obj, aligned_offset, aligned_length) |
| 226 | logger.debug( |
| 227 | "MADV_POPULATE_WRITE loop: %d chunks in %.3f s", |
| 228 | num_chunks, |
| 229 | time.perf_counter() - _t0, |
| 230 | ) |
| 231 | else: |
| 232 | # No rank — populate the entire shared region in one call. |
| 233 | _t0 = time.perf_counter() |
| 234 | populate_write_fn(self.mmap_obj, 0, self.total_size_bytes) |
| 235 | logger.debug( |
| 236 | "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0 |
| 237 | ) |
| 238 | |
| 239 | @property |
| 240 | def base_tensor(self) -> torch.Tensor: |
| 241 | if self._base is None: |
| 242 | raise RuntimeError("Shared offload region has been released.") |
| 243 | return self._base |
| 244 | |
| 245 | def create_next_worker_view(self, tensor_page_size: int) -> torch.Tensor: |
| 246 | """Allocate a strided int8 view for this worker, one canonical tensor. |
| 247 | |
| 248 | Must be called once per canonical tensor. The full mmap layout is: |
| 249 | |
| 250 | worker0_chunk0 | worker1_chunk0 | ... | worker{M-1}_chunk0 |
| 251 | worker0_chunk1 | worker1_chunk1 | ... | worker{M-1}_chunk1 |
| 252 | ... |
| 253 | |
| 254 | Each worker_chunk cell is cpu_page_size bytes and holds all canonical |
| 255 | tensors for that worker and chunk concatenated: |
| 256 | [ tensor0_data | tensor1_data | ... | tensor{L-1}_data ] |
| 257 | |
| 258 | Consecutive rows are separated by row_stride = cpu_page_size * M. |
| 259 | |
| 260 | Returns an int8 tensor of shape (num_chunks, tensor_page_size) with stride |
| 261 | (row_stride, 1). Using int8 keeps stride == bytes, so swap_blocks |
| 262 | address arithmetic works without any dtype conversion. |
| 263 | |
| 264 | Args: |
| 265 | tensor_page_size: Bytes per chunk for this tensor. |
| 266 | |
| 267 | """ |
| 268 | assert self.rank is not None |
| 269 | new_offset = self._worker_offset + tensor_page_size |
| 270 | assert new_offset <= self._worker_area_end, ( |
| 271 | f"Worker offset {new_offset} exceeds worker area end " |
| 272 | f"{self._worker_area_end} (overflowed by " |
| 273 | f"{new_offset - self._worker_area_end} bytes)" |
| 274 | ) |
| 275 | worker_layer_view = torch.as_strided( |
| 276 | self._base, |
| 277 | size=(self.num_chunks, tensor_page_size), |
| 278 | stride=(self._row_stride, 1), |
| 279 | storage_offset=self._worker_offset, |
| 280 | ) |
| 281 | self._worker_offset = new_offset |
| 282 | self._views.append(worker_layer_view) |
| 283 | return worker_layer_view |
| 284 | |
| 285 | def create_next_canonical_view(self, tensor_page_size: int) -> torch.Tensor: |
| 286 | """Allocate a strided int8 view shared by all workers for one |
| 287 | canonical tensor (canonical layout). |
| 288 | |
| 289 | Must be called once per canonical tensor, instead of |
| 290 | create_next_worker_view. The full mmap layout is: |
| 291 | |
| 292 | |<-------- canonical area ------->|<-------- unused ------->| |
| 293 | | all workers share this area | | |
| 294 | | | | |
| 295 | | [ canonical_t0 | canonical_t1 ] | | |
| 296 | | [ canonical_t0 | canonical_t1 ] | | |
| 297 | | [ canonical_t0 | canonical_t1 ] | | |
| 298 | ^ ^ |
| 299 | _canonical_offset=0, then advances by each tensor's size |
| 300 | |
| 301 | Each canonical_t{i} cell is that tensor's canonical page for the |
| 302 | chunk. Canonical areas are carved consecutively from the start of |
| 303 | each chunk row; consecutive rows are separated by row_stride. Every |
| 304 | worker gets the identical byte ranges and writes only its disjoint |
| 305 | bytes within them, as described by its canonical mappings — unlike |
| 306 | create_next_worker_view, which gives each worker a private |
| 307 | cpu_page_size slot per row. |
| 308 | |
| 309 | The trailing unused bytes exist only when the canonical pages sum to |
| 310 | less than row_stride: page-alignment padding of the row, or |
| 311 | deduplication of KV replicated across workers (e.g. the MLA latent), |
| 312 | where one canonical copy replaces world_size worker copies. |
| 313 | |
| 314 | Args: |
| 315 | tensor_page_size: Canonical bytes per chunk for this tensor. |
| 316 | |
| 317 | """ |
| 318 | new_offset = self._canonical_offset + tensor_page_size |
| 319 | assert new_offset <= self._row_stride |
| 320 | view = torch.as_strided( |
| 321 | self._base, |
| 322 | size=(self.num_chunks, tensor_page_size), |
| 323 | stride=(self._row_stride, 1), |
| 324 | storage_offset=self._canonical_offset, |
| 325 | ) |
| 326 | self._canonical_offset = new_offset |
| 327 | self._views.append(view) |
| 328 | return view |
| 329 | |
| 330 | def create_kv_memoryview(self) -> memoryview: |
| 331 | """Return a zero-copy memoryview over the entire KV buffer. |
| 332 | |
| 333 | Shape: (num_chunks, row_stride_bytes). Secondary tiers address |
| 334 | chunk *b* as ``view[b]``. |
| 335 | """ |
| 336 | kv_tensor = self._base.view(self.num_chunks, self._row_stride) |
| 337 | np_arr = kv_tensor.numpy() |
| 338 | assert np_arr.ctypes.data == self._base.data_ptr(), ( |
| 339 | "view()/numpy() created a copy instead of sharing the mmap buffer; " |
| 340 | "secondary tiers require zero-copy access to primary KV data" |
| 341 | ) |
| 342 | return memoryview(np_arr) |
| 343 | |
| 344 | def cleanup(self) -> None: |
| 345 | if self.is_pinned and self._base is not None: |
| 346 | if current_platform.is_cuda_alike(): |
| 347 | base_ptr = self._base.data_ptr() |
| 348 | addresses = self.pinned_addresses or [base_ptr] |
| 349 | for address in reversed(addresses): |
| 350 | result = torch.cuda.cudart().cudaHostUnregister(address) |
| 351 | if result.value != 0: |
| 352 | logger.warning( |
| 353 | "cudaHostUnregister failed for rank=%d, " |
| 354 | "address=%#x (code=%d)", |
| 355 | self.rank, |
| 356 | address, |
| 357 | result.value, |
| 358 | ) |
| 359 | self.pinned_addresses.clear() |
| 360 | self.is_pinned = False |
| 361 | # Release views before _base: each view holds a _base reference and a |
| 362 | # direct StorageImpl reference. Freeing views first lets both refcounts |
| 363 | # drop so the storage (which holds the mmap_obj buffer export) is freed |
| 364 | # before mmap_obj.close() is called below. |
| 365 | if self._views is not None: |
| 366 | self._views.clear() |
| 367 | self._base = None |
| 368 | if self.mmap_obj: |
| 369 | try: |
| 370 | self.mmap_obj.close() |
| 371 | except Exception: |
| 372 | logger.warning("Failed to close mmap_obj", exc_info=True) |
| 373 | self.mmap_obj = None |
| 374 | if self.fd is not None: |
| 375 | try: |
| 376 | os.close(self.fd) |
| 377 | except Exception: |
| 378 | logger.warning("Failed to close fd %s", self.fd, exc_info=True) |
| 379 | self.fd = None |
| 380 | if self._creator and getattr(self, "mmap_path", None): |
| 381 | try: |
| 382 | os.unlink(self.mmap_path) |
| 383 | logger.info("Removed mmap file %s", self.mmap_path) |
| 384 | except Exception: |
| 385 | logger.warning( |
| 386 | "Failed to unlink path %s", self.mmap_path, exc_info=True |
| 387 | ) |
| 388 | self._creator = False |