A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 02:15 UTC
[SharedOffloadRegion] Align blocks to page-size (#43689)
Signed-off-by: varun sundar rabindranath <vsundarr@redhat.com> Co-authored-by: varun sundar rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Jun 3, 2026 at 07:25 UTC
3d76f395e3e25ec3cb9c63fb3c7662fd734222f1
6 files changed
+83
-59
tests/v1/kv_offload/cpu/test_gpu_worker.py
+6
-3
@@ -8,6 +8,7 @@ import pytest
8
import torch
9
10
from vllm.platforms import current_platform
11
+from vllm.utils.math_utils import round_up
12
from vllm.utils.torch_utils import set_random_seed
13
from vllm.v1.kv_offload.base import (
14
CanonicalKVCacheRef,
@@ -90,13 +91,15 @@ def test_transfer(
91
92
mmap_region: SharedOffloadRegion | None = None
93
if use_shared_memory:
93
- cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor
94
+ cpu_page_size = round_up(
95
+ gpu_page_size_bytes * num_tensors * block_size_factor,
96
+ SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
97
+ )
98
mmap_region = SharedOffloadRegion(
99
instance_id=str(uuid.uuid4()),
96
- total_size_bytes=num_cpu_blocks * cpu_page_size,
100
num_blocks=num_cpu_blocks,
101
rank=0,
99
- num_workers=1,
102
+ kv_bytes_per_block=cpu_page_size,
103
cpu_page_size=cpu_page_size,
104
)
105
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+6
-19
@@ -40,14 +40,12 @@ def _make_region(
40
num_workers: int = 1,
41
rank: int = 0,
42
) -> SharedOffloadRegion:
43
- total_size_bytes = num_blocks * num_workers * cpu_page_size
44
- assert total_size_bytes % PAGE_SIZE == 0
43
+ assert cpu_page_size % PAGE_SIZE == 0
44
return SharedOffloadRegion(
45
instance_id=instance_id,
47
- total_size_bytes=total_size_bytes,
46
num_blocks=num_blocks,
47
rank=rank,
50
- num_workers=num_workers,
48
+ kv_bytes_per_block=num_workers * cpu_page_size,
49
cpu_page_size=cpu_page_size,
50
)
51
@@ -77,14 +75,12 @@ def _multi_region(
75
cpu_page_size: int = PAGE_SIZE,
76
):
77
"""Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
80
- total = num_blocks * num_workers * cpu_page_size
78
regions = [
79
SharedOffloadRegion(
80
instance_id=instance_id,
84
- total_size_bytes=total,
81
num_blocks=num_blocks,
82
rank=rank,
87
- num_workers=num_workers,
83
+ kv_bytes_per_block=num_workers * cpu_page_size,
84
cpu_page_size=cpu_page_size,
85
)
86
for rank in range(num_workers)
@@ -104,7 +100,6 @@ def _race_construct(
100
cpu_page_size: int = PAGE_SIZE,
101
) -> tuple[list[SharedOffloadRegion], list[Exception]]:
102
"""Spawn num_workers threads that all race to construct SharedOffloadRegion."""
107
- total = num_blocks * num_workers * cpu_page_size
103
regions: list[SharedOffloadRegion | None] = [None] * num_workers
104
errors: list[Exception] = []
105
barrier = threading.Barrier(num_workers)
@@ -114,10 +109,9 @@ def _race_construct(
109
try:
110
regions[rank] = SharedOffloadRegion(
111
instance_id=instance_id,
117
- total_size_bytes=total,
112
num_blocks=num_blocks,
113
rank=rank,
120
- num_workers=num_workers,
114
+ kv_bytes_per_block=num_workers * cpu_page_size,
115
cpu_page_size=cpu_page_size,
116
)
117
except Exception as e:
@@ -134,7 +128,6 @@ def _race_construct(
128
129
def _mp_race_construct_and_write(
130
instance_id: str,
137
- total_bytes: int,
131
num_blocks: int,
132
rank: int,
133
num_workers: int,
@@ -149,10 +142,9 @@ def _mp_race_construct_and_write(
142
try:
143
region = SharedOffloadRegion(
144
instance_id=instance_id,
152
- total_size_bytes=total_bytes,
145
num_blocks=num_blocks,
146
rank=rank,
155
- num_workers=num_workers,
147
+ kv_bytes_per_block=num_workers * cpu_page_size,
148
cpu_page_size=cpu_page_size,
149
)
150
t = region.create_next_view(cpu_page_size)
@@ -309,7 +301,6 @@ def test_create_next_view_multiprocess_slots(iid):
301
the parent verifies each slot lands at the correct interleaved offset."""
302
num_workers = 2
303
num_blocks = 4
312
- total_bytes = num_blocks * num_workers * PAGE_SIZE
304
305
ctx = get_mp_context()
306
done_queue = ctx.Queue()
@@ -318,10 +309,9 @@ def test_create_next_view_multiprocess_slots(iid):
309
# Parent is rank 0 (creator); child is rank 1 (joiner).
310
region = SharedOffloadRegion(
311
instance_id=iid,
321
- total_size_bytes=total_bytes,
312
num_blocks=num_blocks,
313
rank=0,
324
- num_workers=num_workers,
314
+ kv_bytes_per_block=num_workers * PAGE_SIZE,
315
cpu_page_size=PAGE_SIZE,
316
)
317
try:
@@ -329,7 +319,6 @@ def test_create_next_view_multiprocess_slots(iid):
319
target=_mp_race_construct_and_write,
320
args=(
321
iid,
332
- total_bytes,
322
num_blocks,
323
1,
324
num_workers,
@@ -464,7 +453,6 @@ def test_multiprocess_race_construct_and_write(iid):
453
fill_value = rank+1 into their slot; parent verifies interleaved layout."""
454
num_workers = 4
455
num_blocks = 3
467
- total_bytes = num_blocks * num_workers * PAGE_SIZE
456
457
ctx = get_mp_context()
458
done_queue = ctx.Queue()
@@ -475,7 +463,6 @@ def test_multiprocess_race_construct_and_write(iid):
463
target=_mp_race_construct_and_write,
464
args=(
465
iid,
478
- total_bytes,
466
num_blocks,
467
rank,
468
num_workers,
tests/v1/kv_offload/tiering/test_fs_tier.py
+32
-5
@@ -8,6 +8,7 @@ The tier manager writes KV cache blocks to disk and reads them back, verifying
8
data integrity throughout the process.
9
"""
10
11
+import mmap
12
import os
13
import time
14
from unittest.mock import MagicMock
@@ -26,8 +27,8 @@ from vllm.v1.kv_offload.tiering.fs.manager import (
27
# Helpers
28
# ---------------------------------------------------------------------------
29
29
-_BLOCK_ELEMENTS = 512 * 1024 # 2 MB per block (float32 × 512K = 2MB)
30
-_DTYPE = torch.float32
30
+_BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096.
31
+_DTYPE: torch.dtype = torch.float32
32
_CTX = ReqContext(req_id="test")
33
34
_MOCK_VLLM_CONFIG = MagicMock()
@@ -90,6 +91,32 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list:
91
return results
92
93
94
+def _page_aligned_zero_tensor(
95
+ num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE
96
+) -> torch.Tensor:
97
+ page_size = mmap.PAGESIZE
98
+ dtype_num_bytes = torch.tensor([], dtype=dtype).element_size()
99
+
100
+ num_bytes = num_blocks * block_elements * dtype_num_bytes
101
+ num_bytes_aligned = num_bytes + page_size
102
+ t = torch.zeros(num_bytes_aligned, dtype=torch.uint8)
103
+
104
+ ptr = t.data_ptr()
105
+ alignment_offset = ptr % page_size
106
+ # Move tensor to next page regardless.
107
+ shift = page_size - alignment_offset
108
+ t = t[shift : shift + num_bytes]
109
+ return t.view(dtype).view(num_blocks, block_elements)
110
+
111
+
112
+def _page_aligned_rand_tensor(
113
+ num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE
114
+) -> torch.Tensor:
115
+ rand_tensor = _page_aligned_zero_tensor(num_blocks, block_elements)
116
+ rand_tensor[:] = torch.rand(num_blocks, block_elements, dtype=dtype)
117
+ return rand_tensor
118
+
119
+
120
# ---------------------------------------------------------------------------
121
# Fixtures
122
# ---------------------------------------------------------------------------
@@ -97,7 +124,7 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list:
124
125
@pytest.fixture
126
def fs_tier(tmp_path):
100
- tensor = torch.zeros((4, _BLOCK_ELEMENTS), dtype=_DTYPE)
127
+ tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
128
mock_view = memoryview(tensor.numpy())
129
tier = FileSystemTierManager(
130
offloading_spec=_MOCK_OFFLOADING_SPEC,
@@ -155,7 +182,7 @@ def test_store_then_load_roundtrip(fs_tier):
182
183
def test_invalid_path_raises_at_construction():
184
"""Construction must fail immediately when the config file cannot be written."""
158
- tensor = torch.zeros((32, _BLOCK_ELEMENTS), dtype=_DTYPE)
185
+ tensor = _page_aligned_zero_tensor(32, _BLOCK_ELEMENTS)
186
mock_view = memoryview(tensor.numpy())
187
188
with pytest.raises(OSError):
@@ -228,7 +255,7 @@ def test_store_load_data_integrity(fs_tier):
255
"""Data written by store must be exactly recovered by load."""
256
tier, tensor = fs_tier
257
# Populate tensor with random data
231
- tensor[:] = torch.rand((4, _BLOCK_ELEMENTS), dtype=_DTYPE)
258
+ tensor[:] = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS)
259
260
# Store first 2 blocks
261
num_store = 2
vllm/v1/kv_offload/cpu/shared_offload_region.py
+8
-6
@@ -35,24 +35,26 @@ class SharedOffloadRegion:
35
File path: /dev/shm/vllm_offload_{instance_id}.mmap
36
"""
37
38
+ BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE
39
+
40
def __init__(
41
self,
42
instance_id: str,
41
- total_size_bytes: int,
43
num_blocks: int,
44
rank: int | None,
44
- num_workers: int,
45
+ kv_bytes_per_block: int,
46
cpu_page_size: int,
47
) -> None:
48
self.page_size = mmap.PAGESIZE
49
+ assert kv_bytes_per_block % self.page_size == 0
50
+
51
+ self.num_blocks = num_blocks
52
+ self._row_stride = kv_bytes_per_block
53
+ self.total_size_bytes = self.num_blocks * self._row_stride
54
49
- self.total_size_bytes = total_size_bytes
55
self.mmap_path = f"/dev/shm/vllm_offload_{instance_id}.mmap"
56
self._creator = False # set True only if this worker creates the file
52
- self.num_blocks = num_blocks
57
self.rank = rank
54
- # interleaved-layout stride: one row = all workers' data for one block
55
- self._row_stride = cpu_page_size * num_workers
58
if rank is not None:
59
# byte offset to this worker's first slot within each block row
60
self._worker_offset = rank * cpu_page_size
vllm/v1/kv_offload/cpu/spec.py
+27
-16
@@ -6,6 +6,7 @@ from typing_extensions import override
6
7
from vllm.config import VllmConfig
8
from vllm.platforms import current_platform
9
+from vllm.utils.math_utils import round_up
10
from vllm.v1.kv_cache_interface import KVCacheConfig
11
from vllm.v1.kv_offload.base import (
12
CanonicalKVCaches,
@@ -21,6 +22,8 @@ from vllm.v1.kv_offload.worker.worker import OffloadingHandler
22
23
24
class CPUOffloadingSpec(OffloadingSpec):
25
+ BLOCK_SIZE_ALIGNMENT = 1
26
+
27
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
28
super().__init__(vllm_config, kv_cache_config)
29
@@ -30,26 +33,34 @@ class CPUOffloadingSpec(OffloadingSpec):
33
"cpu_bytes_to_use must be specified in kv_connector_extra_config"
34
)
35
33
- # calculate kv_bytes_per_offloaded_block
36
+ world_size = vllm_config.parallel_config.world_size
37
+ self.num_blocks = 0
38
+ self.kv_bytes_per_offloaded_block = 0
39
+ self.cpu_page_size_per_worker = 0
40
assert kv_cache_config is not None
35
- if kv_cache_config.num_blocks > 0:
41
+ if kv_cache_config.num_blocks > 0 and world_size > 0:
42
total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
43
kv_bytes_per_block = (
44
total_gpu_kv_bytes // kv_cache_config.num_blocks
39
- ) * vllm_config.parallel_config.world_size
40
- else:
41
- kv_bytes_per_block = 0
42
-
43
- kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
44
- self.num_blocks = (
45
- int(cpu_bytes_to_use) // kv_bytes_per_offloaded_block
46
- if kv_bytes_per_offloaded_block > 0
47
- else 0
48
- )
49
- world_size = vllm_config.parallel_config.world_size
50
- self.cpu_page_size_per_worker: int = (
51
- kv_bytes_per_offloaded_block // world_size if world_size > 0 else 0
52
- )
45
+ ) * world_size
46
+ kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
47
+
48
+ # calculate cpu_page_size_per_worker
49
+ self.cpu_page_size_per_worker = kv_bytes_per_offloaded_block // world_size
50
+
51
+ # calculate num_blocks
52
+ aligned_kv_bytes_per_offloaded_block = round_up(
53
+ kv_bytes_per_offloaded_block, self.BLOCK_SIZE_ALIGNMENT
54
+ )
55
+ self.num_blocks = (
56
+ int(cpu_bytes_to_use) // aligned_kv_bytes_per_offloaded_block
57
+ )
58
+
59
+ # Expose aligned_kv_bytes_per_offloaded_block as
60
+ # kv_bytes_per_offloaded_block. Note that this might contain
61
+ # some padding. i.e. each offloaded block is of the form,
62
+ # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
63
+ self.kv_bytes_per_offloaded_block = aligned_kv_bytes_per_offloaded_block
64
65
# scheduler-side
66
self._manager: OffloadingManager | None = None
vllm/v1/kv_offload/tiering/spec.py
+4
-10
@@ -63,6 +63,8 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
63
memory and must transfer data through the primary tier.
64
"""
65
66
+ BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
67
+
68
def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
69
super().__init__(vllm_config, kv_cache_config)
70
# Redeclare for mypy: parent sets this but `--follow-imports skip` hides it
@@ -96,15 +98,11 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
98
99
# Create scheduler-side SharedOffloadRegion (rank=None) so the
100
# primary tier can eagerly create a memoryview over _base.
99
- world_size = self.vllm_config.parallel_config.world_size
101
scheduler_mmap = SharedOffloadRegion(
102
instance_id=self.vllm_config.instance_id,
102
- total_size_bytes=self.cpu_page_size_per_worker
103
- * world_size
104
- * self.num_blocks,
103
num_blocks=self.num_blocks,
104
rank=None,
107
- num_workers=world_size,
105
+ kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
106
cpu_page_size=self.cpu_page_size_per_worker,
107
)
108
self._scheduler_mmap = scheduler_mmap
@@ -165,16 +163,12 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
163
164
@override
165
def create_handlers(self, kv_caches: CanonicalKVCaches) -> CpuGpuOffloadingHandlers:
168
- world_size = self.vllm_config.parallel_config.world_size
166
rank = torch.accelerator.current_device_index()
167
worker_mmap = SharedOffloadRegion(
168
instance_id=self.vllm_config.instance_id,
172
- total_size_bytes=self.cpu_page_size_per_worker
173
- * world_size
174
- * self.num_blocks,
169
num_blocks=self.num_blocks,
170
rank=rank,
177
- num_workers=world_size,
171
+ kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
172
cpu_page_size=self.cpu_page_size_per_worker,
173
)
174
return CpuGpuOffloadingHandlers(