[Bugfix][KV Offload] Fall back when MADV_POPULATE_WRITE is unsupported (#51116)
Signed-off-by: Guy Girmonsky <guygir@gmail.com> Signed-off-by: Alex <jihui.huang@daocloud.io> Co-authored-by: Guy Girmonsky <guygir@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
AlexHuang committed
Aug 6, 2026 at 18:48 UTC
ef2615c2e011d9e3b064f483574f47ad45fa8c38
2 files changed
+197
-6
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+158
@@ -3,6 +3,8 @@
3
"""Unit tests for SharedOffloadRegion."""
4
5
import contextlib
6
+import ctypes
7
+import errno
8
import mmap
9
import os
10
import threading
@@ -57,6 +59,29 @@ def _cleanup_file(path: str) -> None:
59
os.unlink(path)
60
61
62
+def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]:
63
+ """Return Linux page-residency bits for a writable mmap."""
64
+ try:
65
+ libc = ctypes.CDLL(None, use_errno=True)
66
+ mincore = libc.mincore
67
+ except AttributeError:
68
+ pytest.skip("mincore is unavailable")
69
+
70
+ page_count = (length + PAGE_SIZE - 1) // PAGE_SIZE
71
+ mincore.argtypes = [
72
+ ctypes.c_void_p,
73
+ ctypes.c_size_t,
74
+ ctypes.POINTER(ctypes.c_ubyte),
75
+ ]
76
+ mincore.restype = ctypes.c_int
77
+ vector = (ctypes.c_ubyte * page_count)()
78
+ address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj))
79
+ result = mincore(ctypes.c_void_p(address), length, vector)
80
+ if result != 0:
81
+ raise OSError(ctypes.get_errno())
82
+ return [bool(value & 1) for value in vector]
83
+
84
+
85
@contextlib.contextmanager
86
def _region(engine_id: str, **kwargs):
87
"""Context manager: create one region, clean up on exit."""
@@ -405,6 +430,139 @@ def test_file_has_correct_size(iid):
430
assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE
431
432
433
+def test_madvise_success_selects_madvise_population(iid, monkeypatch):
434
+ """A successful probe must keep using MADV_POPULATE_WRITE — the fallback
435
+ helper must never be invoked on a kernel that accepts the advice.
436
+
437
+ Spies on both helpers so we can verify call counts: 1 probe call +
438
+ N populate calls, fallback 0 times.
439
+ """
440
+ from vllm.v1.kv_offload.cpu import shared_offload_region as sor
441
+
442
+ madvise_calls: list[tuple[int, int, int]] = []
443
+ fallback_calls: list[tuple[int, int]] = []
444
+
445
+ def _spy_madvise(mm, off, ln):
446
+ madvise_calls.append((off, ln, id(mm)))
447
+
448
+ def _spy_fallback(mm, off, ln):
449
+ fallback_calls.append((off, ln))
450
+
451
+ monkeypatch.setattr(sor, "_madvise_populate_write", _spy_madvise)
452
+ monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
453
+
454
+ num_blocks = 3
455
+ num_workers = 2
456
+ with _region(iid, num_blocks=num_blocks, num_workers=num_workers, rank=1):
457
+ # 1 probe (PAGESIZE) + N populate calls (one per block per worker column).
458
+ expected_populate = num_blocks # ranked path: one call per block
459
+ mmap_id = madvise_calls[0][2]
460
+ assert madvise_calls == [
461
+ (0, mmap.PAGESIZE, mmap_id),
462
+ (mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
463
+ (3 * mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
464
+ (5 * mmap.PAGESIZE, mmap.PAGESIZE, mmap_id),
465
+ ]
466
+ assert len(madvise_calls) == 1 + expected_populate
467
+ assert fallback_calls == [], (
468
+ "native-path constructor must not invoke the fallback helper"
469
+ )
470
+
471
+
472
+def test_madvise_einval_selects_fallback_for_ranked_region(iid, monkeypatch):
473
+ """An EINVAL probe must select fallback for every ranked block."""
474
+ from vllm.v1.kv_offload.cpu import shared_offload_region as sor
475
+
476
+ fallback_calls: list[tuple[int, int]] = []
477
+ real_fallback = sor._fallback_populate_write
478
+
479
+ def _raise_einval(mm, off, ln):
480
+ raise OSError(errno.EINVAL, "simulated unsupported kernel")
481
+
482
+ def _spy_fallback(mm, off, ln):
483
+ fallback_calls.append((off, ln))
484
+ return real_fallback(mm, off, ln)
485
+
486
+ monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
487
+ monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
488
+
489
+ with _region(iid, num_blocks=3, num_workers=2, rank=1):
490
+ assert fallback_calls == [
491
+ (mmap.PAGESIZE, mmap.PAGESIZE),
492
+ (3 * mmap.PAGESIZE, mmap.PAGESIZE),
493
+ (5 * mmap.PAGESIZE, mmap.PAGESIZE),
494
+ ]
495
+
496
+
497
+def test_madvise_einval_selects_fallback_for_unranked_region(iid, monkeypatch):
498
+ """An EINVAL probe must select fallback for the whole unranked region."""
499
+ from vllm.v1.kv_offload.cpu import shared_offload_region as sor
500
+
501
+ fallback_calls: list[tuple[int, int]] = []
502
+ real_fallback = sor._fallback_populate_write
503
+
504
+ def _raise_einval(mm, off, ln):
505
+ raise OSError(errno.EINVAL, "simulated unsupported kernel")
506
+
507
+ def _spy_fallback(mm, off, ln):
508
+ fallback_calls.append((off, ln))
509
+ return real_fallback(mm, off, ln)
510
+
511
+ monkeypatch.setattr(sor, "_madvise_populate_write", _raise_einval)
512
+ monkeypatch.setattr(sor, "_fallback_populate_write", _spy_fallback)
513
+
514
+ with _region(iid, num_blocks=3, num_workers=2, rank=None):
515
+ assert fallback_calls == [(0, 6 * mmap.PAGESIZE)]
516
+
517
+
518
+def test_fallback_populate_write_preserves_bytes_and_faults_pages():
519
+ """Fallback preserves existing bytes and touches every target page."""
520
+ from vllm.v1.kv_offload.cpu import shared_offload_region as sor
521
+
522
+ size = 3 * mmap.PAGESIZE
523
+ if not hasattr(mmap, "MADV_DONTNEED"):
524
+ pytest.skip("MADV_DONTNEED is unavailable")
525
+
526
+ mmap_obj = mmap.mmap(
527
+ -1,
528
+ size,
529
+ flags=mmap.MAP_SHARED,
530
+ prot=mmap.PROT_READ | mmap.PROT_WRITE,
531
+ )
532
+ try:
533
+ mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size)
534
+ if any(_page_residency(mmap_obj, size)):
535
+ pytest.skip("kernel did not discard anonymous mmap pages")
536
+
537
+ mmap_obj[0] = 0xAB
538
+ assert _page_residency(mmap_obj, size) == [True, False, False]
539
+
540
+ sor._fallback_populate_write(mmap_obj, 0, size)
541
+
542
+ assert _page_residency(mmap_obj, size) == [True, True, True]
543
+ assert mmap_obj[0] == 0xAB
544
+ finally:
545
+ mmap_obj.close()
546
+
547
+
548
+def test_madvise_unexpected_oserror_propagates(iid, monkeypatch):
549
+ """Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must
550
+ propagate out of __init__, not be silently masked by the fallback branch.
551
+ """
552
+ from vllm.v1.kv_offload.cpu import shared_offload_region as sor
553
+
554
+ monkeypatch.setattr(
555
+ sor,
556
+ "_madvise_populate_write",
557
+ lambda mm, off, ln: (_ for _ in ()).throw(
558
+ OSError(errno.EIO, "simulated I/O failure")
559
+ ),
560
+ )
561
+ with pytest.raises(OSError) as exc_info:
562
+ _make_region(iid)
563
+ assert exc_info.value.errno == errno.EIO
564
+
565
+
566
# ---------------------------------------------------------------------------
567
# Multi-worker race — concurrent construction
568
# ---------------------------------------------------------------------------
vllm/v1/kv_offload/cpu/shared_offload_region.py
+39
-6
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import errno
4
import mmap
5
import os
6
import time
7
+from collections.abc import Callable
8
9
+import numpy as np
10
import torch
11
12
from vllm.distributed.device_communicators.shm_broadcast import (
@@ -14,6 +17,9 @@ from vllm.platforms import current_platform
17
18
logger = init_logger(__name__)
19
20
+# MADV_POPULATE_WRITE was added in Linux 5.14 (value 23).
21
+_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
22
+
23
24
def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> None:
25
"""Spin-wait until the file reaches expected_size (creator truncated it)."""
@@ -28,6 +34,35 @@ def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0) -> N
34
time.sleep(0.005)
35
36
37
+def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
38
+ mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length)
39
+
40
+
41
+def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
42
+ # Touch one byte per page via a read-modify-write so existing bytes are
43
+ # preserved — a peer worker may have already written KV data into this
44
+ # shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE.
45
+ arr = np.frombuffer(mmap_obj, dtype=np.uint8)
46
+ arr[offset : offset + length : mmap.PAGESIZE] |= 0
47
+
48
+
49
+def _get_populate_write_fn(
50
+ mmap_obj: mmap.mmap,
51
+) -> Callable[[mmap.mmap, int, int], None]:
52
+ """Select the pre-faulting method once for this mmap."""
53
+ try:
54
+ _madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE)
55
+ except OSError as e:
56
+ if e.errno != errno.EINVAL:
57
+ raise
58
+ logger.warning(
59
+ "MADV_POPULATE_WRITE is not supported; falling back to per-page "
60
+ "writes for mmap pre-population. Startup may be slower."
61
+ )
62
+ return _fallback_populate_write
63
+ return _madvise_populate_write
64
+
65
+
66
class SharedOffloadRegion:
67
"""
68
Single mmap-backed memory region shared across all workers for a
@@ -104,8 +139,8 @@ class SharedOffloadRegion:
139
prot=mmap.PROT_READ | mmap.PROT_WRITE,
140
)
141
107
- # MADV_POPULATE_WRITE was added in Linux 5.14 (value 23).
108
- _MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
142
+ populate_write_fn = _get_populate_write_fn(self.mmap_obj)
143
+
144
if rank is not None:
145
# Populate only this worker's pages (one slot per block row).
146
worker_offset = rank * cpu_page_size
@@ -116,9 +151,7 @@ class SharedOffloadRegion:
151
aligned_offset = (raw_offset // page_size) * page_size
152
end = raw_offset + cpu_page_size
153
aligned_length = end - aligned_offset
119
- self.mmap_obj.madvise(
120
- _MADV_POPULATE_WRITE, aligned_offset, aligned_length
121
- )
154
+ populate_write_fn(self.mmap_obj, aligned_offset, aligned_length)
155
logger.debug(
156
"MADV_POPULATE_WRITE loop: %d blocks in %.3f s",
157
num_blocks,
@@ -127,7 +160,7 @@ class SharedOffloadRegion:
160
else:
161
# No rank — populate the entire shared region in one call.
162
_t0 = time.perf_counter()
130
- self.mmap_obj.madvise(_MADV_POPULATE_WRITE, 0, self.total_size_bytes)
163
+ populate_write_fn(self.mmap_obj, 0, self.total_size_bytes)
164
logger.debug(
165
"MADV_POPULATE_WRITE entire region: %.3f s", time.perf_counter() - _t0
166
)