[Bugfix][KV Connector] Only enforce disk block alignment for O_DIRECT (#55424)
Signed-off-by: Vincent <vincexxchan@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
mevince committed
Sep 11, 2026 at 13:12 UTC
c1b69aa0d49b88016b5b2156ae61649c2c2df282
2 files changed
+63
-3
tests/v1/simple_kv_offload/test_worker.py
+57
@@ -43,6 +43,7 @@ from vllm.v1.simple_kv_offload.cuda_mem_ops import (
43
build_params,
44
pin_tensor,
45
)
46
+from vllm.v1.simple_kv_offload.disk_backend import DiskBackend
47
from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
48
from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
49
from vllm.v1.worker.utils import allocate_kv_cache
@@ -519,3 +520,59 @@ def test_mixed_page_byte_placement_is_dcp_invariant():
520
)
521
522
assert placements[0] == placements[1]
523
+
524
+
525
+@pytest.mark.parametrize(
526
+ ("block_bytes", "use_page_cache", "expect_ok"),
527
+ [
528
+ (4096, True, True), # aligned control
529
+ (3 * 512, True, True), # 512-aligned only: fine for buffered I/O
530
+ (3 * 512, False, False), # O_DIRECT keeps enforcing 4096 alignment
531
+ ],
532
+)
533
+def test_disk_backend_alignment_only_required_for_direct_io(
534
+ tmp_path, monkeypatch, block_bytes: int, use_page_cache: bool, expect_ok: bool
535
+):
536
+ """The 4096 stride assert belongs to O_DIRECT; page-cache I/O must not hit it.
537
+
538
+ Hybrid models (e.g. 101 per-rank cache tensors) sum to a block stride that
539
+ is 512- but not 4096-aligned, which previously made disk offload unbootable
540
+ even with use_page_cache=True.
541
+ """
542
+ # Host pinning is a separate lifecycle concern with its own tests; the
543
+ # alignment gate must not depend on registration state, so stub both ends.
544
+ monkeypatch.setattr(
545
+ "vllm.v1.simple_kv_offload.disk_backend.pin_tensor", lambda t: None
546
+ )
547
+ monkeypatch.setattr(
548
+ "vllm.v1.simple_kv_offload.disk_backend.unpin_tensor",
549
+ lambda t: None,
550
+ raising=False, # harmless on branches without teardown-side unpins
551
+ )
552
+ num_blocks = 4
553
+ gpu = {"k": torch.zeros((num_blocks, block_bytes), dtype=torch.int8, device="cuda")}
554
+ backend = DiskBackend()
555
+ disk_path = tmp_path / "kv-offload.bin"
556
+
557
+ def _init() -> None:
558
+ backend.init(
559
+ gpu,
560
+ gpu["k"].device,
561
+ torch.cuda.Stream(),
562
+ torch.cuda.Stream(),
563
+ str(disk_path),
564
+ num_blocks,
565
+ block_bytes,
566
+ use_page_cache=use_page_cache,
567
+ )
568
+
569
+ if not expect_ok:
570
+ with pytest.raises(AssertionError, match="not aligned"):
571
+ _init()
572
+ return
573
+
574
+ _init()
575
+ try:
576
+ assert disk_path.stat().st_size == num_blocks * block_bytes
577
+ finally:
578
+ backend.shutdown()
vllm/v1/simple_kv_offload/disk_backend.py
+6
-3
@@ -99,9 +99,12 @@ class DiskBackend:
99
t.stride(0) * t.element_size() for t in gpu_caches.values()
100
]
101
102
- assert total_block_bytes % _ALIGNMENT == 0, (
103
- f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
104
- )
102
+ # The stride alignment is an O_DIRECT kernel-DMA requirement; buffered
103
+ # I/O through the page cache has no alignment constraint.
104
+ if not use_page_cache:
105
+ assert total_block_bytes % _ALIGNMENT == 0, (
106
+ f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
107
+ )
108
109
# Separate buffer pools for store and load threads
110
self._store_buffer_caches = {}