[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 <noreply@openai.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
This commit is contained in:
Itay Etelis
co-authored by
Robert Shaw
Codex
Kevin H. Luu
Robert Shaw
Codex
Kevin H. Luu
parent
16532876e7
commit
21f4b6dc84
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user