Make tiering offload region DP-replica aware (#47987)
Signed-off-by: Liran Schour <lirans@il.ibm.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
liranschour committed
Jul 12, 2026 at 13:10 UTC
5c0c987c03593a1c05b8728fb265c94deb3d0a4c
4 files changed
+28
-18
tests/v1/kv_offload/cpu/test_gpu_worker.py
+1
-1
@@ -96,7 +96,7 @@ def test_transfer(
96
SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
97
)
98
mmap_region = SharedOffloadRegion(
99
- instance_id=str(uuid.uuid4()),
99
+ engine_id=str(uuid.uuid4()),
100
num_blocks=num_cpu_blocks,
101
rank=0,
102
kv_bytes_per_block=cpu_page_size,
tests/v1/kv_offload/cpu/test_shared_offload_region.py
+11
-11
@@ -34,7 +34,7 @@ def _set_spawn_method(monkeypatch):
34
35
36
def _make_region(
37
- instance_id: str,
37
+ engine_id: str,
38
num_blocks: int = 4,
39
cpu_page_size: int = PAGE_SIZE,
40
num_workers: int = 1,
@@ -42,7 +42,7 @@ def _make_region(
42
) -> SharedOffloadRegion:
43
assert cpu_page_size % PAGE_SIZE == 0
44
return SharedOffloadRegion(
45
- instance_id=instance_id,
45
+ engine_id=engine_id,
46
num_blocks=num_blocks,
47
rank=rank,
48
kv_bytes_per_block=num_workers * cpu_page_size,
@@ -57,9 +57,9 @@ def _cleanup_file(path: str) -> None:
57
58
59
@contextlib.contextmanager
60
-def _region(instance_id: str, **kwargs):
60
+def _region(engine_id: str, **kwargs):
61
"""Context manager: create one region, clean up on exit."""
62
- r = _make_region(instance_id, **kwargs)
62
+ r = _make_region(engine_id, **kwargs)
63
try:
64
yield r
65
finally:
@@ -69,7 +69,7 @@ def _region(instance_id: str, **kwargs):
69
70
@contextlib.contextmanager
71
def _multi_region(
72
- instance_id: str,
72
+ engine_id: str,
73
num_workers: int,
74
num_blocks: int = 4,
75
cpu_page_size: int = PAGE_SIZE,
@@ -77,7 +77,7 @@ def _multi_region(
77
"""Context manager: create one SharedOffloadRegion per rank, clean up on exit."""
78
regions = [
79
SharedOffloadRegion(
80
- instance_id=instance_id,
80
+ engine_id=engine_id,
81
num_blocks=num_blocks,
82
rank=rank,
83
kv_bytes_per_block=num_workers * cpu_page_size,
@@ -94,7 +94,7 @@ def _multi_region(
94
95
96
def _race_construct(
97
- instance_id: str,
97
+ engine_id: str,
98
num_workers: int,
99
num_blocks: int = 4,
100
cpu_page_size: int = PAGE_SIZE,
@@ -108,7 +108,7 @@ def _race_construct(
108
barrier.wait() # all threads start at the same instant
109
try:
110
regions[rank] = SharedOffloadRegion(
111
- instance_id=instance_id,
111
+ engine_id=engine_id,
112
num_blocks=num_blocks,
113
rank=rank,
114
kv_bytes_per_block=num_workers * cpu_page_size,
@@ -127,7 +127,7 @@ def _race_construct(
127
128
129
def _mp_race_construct_and_write(
130
- instance_id: str,
130
+ engine_id: str,
131
num_blocks: int,
132
rank: int,
133
num_workers: int,
@@ -141,7 +141,7 @@ def _mp_race_construct_and_write(
141
parent a window to read the raw mmap before the creator removes the file."""
142
try:
143
region = SharedOffloadRegion(
144
- instance_id=instance_id,
144
+ engine_id=engine_id,
145
num_blocks=num_blocks,
146
rank=rank,
147
kv_bytes_per_block=num_workers * cpu_page_size,
@@ -308,7 +308,7 @@ def test_create_next_view_multiprocess_slots(iid):
308
309
# Parent is rank 0 (creator); child is rank 1 (joiner).
310
region = SharedOffloadRegion(
311
- instance_id=iid,
311
+ engine_id=iid,
312
num_blocks=num_blocks,
313
rank=0,
314
kv_bytes_per_block=num_workers * PAGE_SIZE,
vllm/v1/kv_offload/cpu/shared_offload_region.py
+3
-3
@@ -33,14 +33,14 @@ class SharedOffloadRegion:
33
the rest open the existing file and wait until it reaches the expected
34
size. Each worker then mmap()s the full file.
35
36
- File path: /dev/shm/vllm_offload_{instance_id}.mmap
36
+ File path: /dev/shm/vllm_offload_{engine_id}.mmap
37
"""
38
39
BLOCK_SIZE_ALIGNMENT: int = mmap.PAGESIZE
40
41
def __init__(
42
self,
43
- instance_id: str,
43
+ engine_id: str,
44
num_blocks: int,
45
rank: int | None,
46
kv_bytes_per_block: int,
@@ -53,7 +53,7 @@ class SharedOffloadRegion:
53
self._row_stride = kv_bytes_per_block
54
self.total_size_bytes = self.num_blocks * self._row_stride
55
56
- self.mmap_path = f"/dev/shm/vllm_offload_{instance_id}.mmap"
56
+ self.mmap_path = f"/dev/shm/vllm_offload_{engine_id}.mmap"
57
self._creator = False # set True only if this worker creates the file
58
self.rank = rank
59
if rank is not None:
vllm/v1/kv_offload/tiering/spec.py
+13
-3
@@ -108,6 +108,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
108
# Scheduler-side mmap (rank=None); kept for cleanup
109
self._scheduler_mmap: SharedOffloadRegion | None = None
110
111
+ # engine_id is unique per DP replica (suffixed with _dp{rank} in both
112
+ # the Ray and multiprocessing paths), so it names a per-replica offload
113
+ # region. Non-None is guaranteed by OffloadingSpec.__init__.
114
+ assert vllm_config.kv_transfer_config is not None
115
+ assert vllm_config.kv_transfer_config.engine_id is not None
116
+ self._engine_id: str = vllm_config.kv_transfer_config.engine_id
117
+
118
@override
119
def get_manager(self) -> OffloadingManager:
120
"""
@@ -124,7 +131,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
131
# Create scheduler-side SharedOffloadRegion (rank=None) so the
132
# primary tier can eagerly create a memoryview over _base.
133
scheduler_mmap = SharedOffloadRegion(
127
- instance_id=self.vllm_config.instance_id,
134
+ engine_id=self._engine_id,
135
num_blocks=self.num_blocks,
136
rank=None,
137
kv_bytes_per_block=self.kv_bytes_per_offloaded_block,
@@ -187,9 +194,12 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
194
195
@override
196
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
190
- rank = torch.accelerator.current_device_index()
197
+ # Fold the global physical device index into the replica-local
198
+ # [0, world_size) slot range.
199
+ world_size = self.vllm_config.parallel_config.world_size
200
+ rank = torch.accelerator.current_device_index() % world_size
201
worker_mmap = SharedOffloadRegion(
192
- instance_id=self.vllm_config.instance_id,
202
+ engine_id=self._engine_id,
203
num_blocks=self.num_blocks,
204
rank=rank,
205
kv_bytes_per_block=self.kv_bytes_per_offloaded_block,