From 483e0c24c757094df710fc3da3170a5f2ebf0b8d Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:18:56 +0300 Subject: [PATCH] [KV Offload] Split replicated-layout pre-faulting across workers (#59545) 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: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../cpu/test_shared_offload_region.py | 21 +++++++++++++++++++ .../kv_offload/cpu/shared_offload_region.py | 13 +++++++++++- vllm/v1/kv_offload/cpu/spec.py | 13 ++++++++---- 3 files changed, 42 insertions(+), 5 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 b7259e6fd6..d74165a965 100644 --- a/tests/v1/kv_offload/cpu/test_shared_offload_region.py +++ b/tests/v1/kv_offload/cpu/test_shared_offload_region.py @@ -45,6 +45,7 @@ def _make_region( num_workers: int = 1, rank: int = 0, barrier=None, + populate_shard: tuple[int, int] | None = None, ) -> SharedOffloadRegion: assert cpu_page_size % PAGE_SIZE == 0 return SharedOffloadRegion( @@ -54,6 +55,7 @@ def _make_region( kv_bytes_per_chunk=num_workers * cpu_page_size, cpu_page_size=cpu_page_size, barrier=barrier, + populate_shard=populate_shard, ) @@ -526,6 +528,25 @@ def test_madvise_success_selects_madvise_population(iid, monkeypatch): ) +@pytest.mark.parametrize( + "shard, expected", + [((1, 2), [(3 * PAGE_SIZE, 3 * PAGE_SIZE)]), ((3, 4), [])], +) +def test_populate_shard_prefaults_only_its_share(iid, monkeypatch, shard, expected): + """Workers that map the same bytes each pre-fault one contiguous share + (page-rounded; a trailing share may be empty) instead of the whole region.""" + from vllm.v1.kv_offload.cpu import shared_offload_region as sor + + madvise_calls: list[tuple[int, int]] = [] + monkeypatch.setattr( + sor, + "_madvise_populate_write", + lambda mm, off, ln: madvise_calls.append((off, ln)), + ) + with _region(iid, num_chunks=3, num_workers=2, populate_shard=shard): + assert madvise_calls == [(0, PAGE_SIZE), *expected] + + def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch): """An EINVAL probe must select fallback for every ranked chunk.""" from vllm.v1.kv_offload.cpu import shared_offload_region as sor diff --git a/vllm/v1/kv_offload/cpu/shared_offload_region.py b/vllm/v1/kv_offload/cpu/shared_offload_region.py index 7e0e97aa12..aa3ff8838b 100644 --- a/vllm/v1/kv_offload/cpu/shared_offload_region.py +++ b/vllm/v1/kv_offload/cpu/shared_offload_region.py @@ -14,6 +14,7 @@ from vllm.distributed.device_communicators.shm_broadcast import ( check_shm_free_space, ) from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv, round_up from vllm.v1.kv_offload.cpu.host_register import host_unregister logger = init_logger(__name__) @@ -82,6 +83,8 @@ class SharedOffloadRegion: Creator-only population pre-faults the entire region before the barrier and requires that barrier to keep joiners from using unpopulated pages. + ``populate_shard=(index, count)`` pre-faults only this worker's contiguous + share, for layouts where every worker maps the same bytes. """ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE @@ -97,6 +100,7 @@ class SharedOffloadRegion: *, creator_memory_check: Callable[[int], None] | None = None, populate_only_on_creator: bool = False, + populate_shard: tuple[int, int] | None = None, ) -> None: if populate_only_on_creator and barrier is None: raise ValueError("Creator-only population requires a barrier.") @@ -212,7 +216,14 @@ class SharedOffloadRegion: populate_write_fn = _get_populate_write_fn(self.mmap_obj) - if rank is not None: + if populate_shard is not None: + index, count = populate_shard + share = round_up(cdiv(self.total_size_bytes, count), self.page_size) + start = min(index * share, self.total_size_bytes) + if start < self.total_size_bytes: + length = min(share, self.total_size_bytes - start) + populate_write_fn(self.mmap_obj, start, length) + elif rank is not None: # Populate only this worker's pages (one slot per chunk row). worker_offset = rank * cpu_page_size _t0 = time.perf_counter() diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index c26cd32f97..3cb1f51722 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -170,16 +170,21 @@ class CPUOffloadingSpec(OffloadingSpec): # num_chunks == 0 would size the region to zero bytes, which cannot be # mmap'd; fall back to the tensor path (empty tensors) as before. if self._uses_shared_region() and self.num_chunks > 0: - # Replicated layout puts all ranks on slot 0 (single MLA copy); - # otherwise each worker uses its own rank-indexed slot. - rank = 0 if self.replicated_layout else self.config.parallel.rank + # Use the logical worker rank for each slot, and distribute the + # replicated region's pre-faulting work across the worker group. + world_size = self.config.parallel.world_size + worker_index = self.config.parallel.rank + mmap_region = SharedOffloadRegion( engine_id=self.config.engine_id, num_chunks=self.num_chunks, - rank=rank, + rank=0 if self.replicated_layout else worker_index, kv_bytes_per_chunk=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, barrier=_all_workers_barrier, + populate_shard=( + (worker_index, world_size) if self.replicated_layout else None + ), ) try: return CPUOffloadingWorker(