From 21f4b6dc84d73d53818af8cdb6e4dfde4577d6cd Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:17:40 +0300 Subject: [PATCH] [KV Offload] Split the CPU offload region across shm files (#59546) Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Codex Co-authored-by: Kevin H. Luu --- .../cpu/test_shared_offload_region.py | 36 ++++++++ .../kv_offload/cpu/shared_offload_region.py | 84 +++++++++++++++++-- vllm/v1/kv_offload/cpu/spec.py | 5 ++ 3 files changed, 116 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_shared_offload_region.py b/tests/v1/kv_offload/cpu/test_shared_offload_region.py index 21ef6d2761..6fc557c3a8 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -653,6 +653,41 @@ def test_madvise_unexpected_oserror_propagates(iid, monkeypatch): assert exc_info.value.errno == errno.EIO +def test_segmented_region_maps_files_back_to_back(iid, monkeypatch): + """Rows split over three shm files mapped back to back: each worker sees + the other's writes in every file, and the creator removes every file.""" + monkeypatch.setattr(region_module, "_MIN_SEGMENT_BYTES", PAGE_SIZE) + regions = [ + SharedOffloadRegion( + engine_id=iid, + num_chunks=5, + rank=rank, + kv_bytes_per_chunk=2 * PAGE_SIZE, + cpu_page_size=PAGE_SIZE, + num_segments=3, + ) + for rank in range(2) + ] + segments = regions[0]._segments + try: + assert [length // PAGE_SIZE for _, _, length in segments] == [4, 4, 2] + for rank, region in enumerate(regions): + ranges = list(region.iter_worker_slot_ranges()) + first_row = rank * 5 // 2 + assert ranges[0] == ((first_row * 2 + rank) * PAGE_SIZE, PAGE_SIZE) + assert sorted(ranges) == [ + ((row * 2 + rank) * PAGE_SIZE, PAGE_SIZE) for row in range(5) + ] + region.create_next_worker_view(PAGE_SIZE).fill_(rank + 1) + rows = regions[0].base_tensor.view(5, 2, PAGE_SIZE) + assert (rows[:, 0] == 1).all() and (rows[:, 1] == 2).all() + del rows + finally: + for region in regions: + region.cleanup() + assert not any(os.path.exists(path) for path, _, _ in segments) + + # --------------------------------------------------------------------------- # Multi-worker race — concurrent construction # --------------------------------------------------------------------------- @@ -899,6 +934,7 @@ def _unmapped_registration_region(num_chunks, slot_size, num_workers): region.cpu_page_size = slot_size region._row_stride = num_workers * slot_size region.total_size_bytes = region.num_chunks * region._row_stride + region._segments = [("unused", 0, region.total_size_bytes)] region._base = SimpleNamespace(data_ptr=lambda: 0x100000) region.pinned_addresses = [] region.is_pinned = False diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 34faf28b94..3457a34ff0 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +import ctypes import errno +import functools import mmap import os import time from collections.abc import Callable, Iterator +from ctypes import c_int, c_long, c_size_t, c_void_p import numpy as np import torch @@ -21,6 +24,10 @@ logger = init_logger(__name__) # MADV_POPULATE_WRITE was added in Linux 5.14 (value 23). _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) +_MAP_FIXED = 0x10 +# Page-cache inserts are serialized per shm file, capping pre-faulting of one +# file near 10 GiB/s; smaller segments are not worth a file. +_MIN_SEGMENT_BYTES = 1 << 30 def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None: @@ -40,6 +47,14 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N time.sleep(0.005) +@functools.cache +def _libc_mmap() -> Callable[..., int | None]: + mmap_fn = ctypes.CDLL(None, use_errno=True).mmap + mmap_fn.argtypes = [c_void_p, c_size_t, c_int, c_int, c_int, c_long] + mmap_fn.restype = c_void_p + return mmap_fn + + def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None: mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length) @@ -84,6 +99,7 @@ class SharedOffloadRegion: ``populate_shard=(index, count)`` pre-faults only this worker's contiguous share, for layouts where every worker maps the same bytes. The barrier also waits for all shards to finish pre-faulting before returning. + ``num_segments`` splits backing files beneath the contiguous mapping. """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE @@ -99,6 +115,7 @@ class SharedOffloadRegion: *, creator_memory_check: Callable[[int], None] | None = None, populate_shard: tuple[int, int] | None = None, + num_segments: int = 1, ) -> None: self.page_size = mmap.PAGESIZE assert kv_bytes_per_chunk % self.page_size == 0 @@ -109,6 +126,17 @@ class SharedOffloadRegion: self.cpu_page_size = cpu_page_size self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap" + num_segments = min(num_segments, self.total_size_bytes // _MIN_SEGMENT_BYTES) + rows = max(1, cdiv(num_chunks, max(1, num_segments))) + # (path, offset, length) per file. + self._segments = [ + ( + f"{self.mmap_path}.{i}" if i else self.mmap_path, + row * self._row_stride, + min(rows, num_chunks - row) * self._row_stride, + ) + for i, row in enumerate(range(0, max(num_chunks, 1), rows)) + ] self._creator = False # set True only if this worker creates the file self.rank = rank if rank is not None: @@ -126,7 +154,7 @@ class SharedOffloadRegion: # Joiner path — another worker won O_EXCL. Reopen and wait # for the file to reach expected size. self.fd = os.open(self.mmap_path, os.O_RDWR) - _wait_for_file_size(self.fd, self.total_size_bytes) + _wait_for_file_size(self.fd, self._segments[0][2]) logger.info("Opened existing mmap file %s", self.mmap_path) else: # Creator path. We won O_EXCL, so we own the file: any @@ -140,24 +168,32 @@ class SharedOffloadRegion: self.total_size_bytes, allocation_name="CPU KV offload shared region in /dev/shm", ) - os.ftruncate(self.fd, self.total_size_bytes) + for path, _, length in self._segments[1:]: + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600) + os.ftruncate(fd, length) + os.close(fd) + # Joiners wait for the first file, so size it last. + os.ftruncate(self.fd, self._segments[0][2]) logger.info( "Created mmap file %s (%.2f GB)", self.mmap_path, self.total_size_bytes / 1e9, ) + single = len(self._segments) == 1 self.mmap_obj: mmap.mmap | None = mmap.mmap( - self.fd, + self.fd if single else -1, self.total_size_bytes, - flags=mmap.MAP_SHARED, + flags=mmap.MAP_SHARED if single else mmap.MAP_PRIVATE, prot=mmap.PROT_READ | mmap.PROT_WRITE, ) + if not single: + self._map_segments() if barrier is not None: barriers_remaining -= 1 barrier() if self._creator: - os.unlink(self.mmap_path) + self._unlink() self._creator = False logger.info("Unlinked mmap file %s", self.mmap_path) if populate_shard is not None: @@ -173,8 +209,7 @@ class SharedOffloadRegion: barrier() except Exception: if self._creator: - with contextlib.suppress(FileNotFoundError): - os.unlink(self.mmap_path) + self._unlink() self._creator = False if hasattr(self, "mmap_obj") and self.mmap_obj is not None: # The prefault traceback may retain an exported buffer. @@ -225,11 +260,42 @@ class SharedOffloadRegion: "MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0 ) + def _map_segments(self) -> None: + """Replace the anonymous mapping with the segment files, back to back.""" + assert self.mmap_obj is not None and self.fd is not None + base = ctypes.addressof(ctypes.c_char.from_buffer(self.mmap_obj)) + for path, offset, length in self._segments: + fd = self.fd if offset == 0 else os.open(path, os.O_RDWR) + address = _libc_mmap()( + base + offset, + length, + mmap.PROT_READ | mmap.PROT_WRITE, + mmap.MAP_SHARED | _MAP_FIXED, + fd, + 0, + ) + if fd != self.fd: + os.close(fd) + if address != base + offset: + raise OSError(ctypes.get_errno(), f"Failed to map {path}") + + def _unlink(self) -> None: + for path, _, _ in self._segments: + with contextlib.suppress(FileNotFoundError): + os.unlink(path) + def iter_worker_slot_ranges(self) -> Iterator[tuple[int, int]]: """Yield this worker's slots as page-aligned (offset, length) ranges.""" assert self.rank is not None worker_offset = self.rank * self.cpu_page_size - for row in range(self.num_chunks): + first = 0 + if len(self._segments) > 1: + # Start each worker on a different file. + first = ( + self.rank * self.num_chunks // (self._row_stride // self.cpu_page_size) + ) + for i in range(self.num_chunks): + row = (first + i) % self.num_chunks start = row * self._row_stride + worker_offset aligned = start - start % self.page_size end = round_up(start + self.cpu_page_size, self.page_size) @@ -380,7 +446,7 @@ class SharedOffloadRegion: self.fd = None if self._creator and getattr(self, "mmap_path", None): try: - os.unlink(self.mmap_path) + self._unlink() logger.info("Removed mmap file %s", self.mmap_path) except Exception: logger.warning( diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 4f6275d028..807f1b6e8f 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -22,6 +22,10 @@ from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion +# Split the shared region into this many shm files so workers pre-fault them +# in parallel (page-cache inserts are serialized per file). +_REGION_SEGMENTS = 32 + def _all_workers_barrier() -> None: """Block until every worker rank has reached this point (gloo cpu group). @@ -188,6 +192,7 @@ class CPUOffloadingSpec(OffloadingSpec): kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, barrier=_all_workers_barrier, + num_segments=_REGION_SEGMENTS, populate_shard=( (worker_index, world_size) if self.replicated_layout else None ),