[KV Offload] Make compact secondary identity TP-independent (#49858)
Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: GPT-5.6 Sol <noreply@openai.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Chang Guo committed
Jul 27, 2026 at 21:10 UTC
a8f296083f43e236f908d4d898ccfb0b108db1f6
4 files changed
+226
-11
tests/v1/kv_offload/test_file_mapper.py
+105
@@ -51,6 +51,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
51
data_parallel_index=0,
52
is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False),
53
),
54
+ replicated_layout=kwargs.get("replicated_layout", False),
55
)
56
spec = MagicMock(spec=OffloadingSpec)
57
spec.config = config
@@ -205,3 +206,107 @@ def test_parallel_agnostic_separates_persistent_layouts():
206
assert agnostic.base_path != specific.base_path
207
assert "parallel_agnostic" not in agnostic.fields
208
assert specific.fields["parallel_agnostic"] is False
209
+
210
+
211
+# ---------------------------------------------------------------------------
212
+# replicated_layout: OR'd into parallel-agnostic identity for compact rows
213
+# ---------------------------------------------------------------------------
214
+
215
+
216
+def test_replicated_layout_collapses_parallel_identity():
217
+ shared = dict(
218
+ model_name="mla-model",
219
+ groups=((16, "mla_layer"),),
220
+ replicated_layout=True,
221
+ parallel_agnostic=True,
222
+ )
223
+ tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=1, **shared)
224
+ tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=3, **shared)
225
+
226
+ assert tp2.base_path == tp4.base_path
227
+ for fm in (tp2, tp4):
228
+ assert fm.fields["tp_size"] == 1
229
+ assert fm.fields["pp_size"] == 1
230
+ assert fm.fields["pcp_size"] == 1
231
+ assert fm.fields["dcp_size"] == 1
232
+ assert fm.rank == 0
233
+ assert "parallel_agnostic" not in fm.fields
234
+ assert fm.fields["replicated_layout"] is True
235
+ assert fm.get_file_name(make_offload_key(b"\x01" * 8, 0)).startswith(
236
+ f"{fm.base_path}_r0/"
237
+ )
238
+
239
+
240
+def test_replicated_layout_requires_caller_opt_in():
241
+ fm = make_mapper_from_offloading_spec(
242
+ tp_size=2,
243
+ world_size=2,
244
+ rank=1,
245
+ replicated_layout=True,
246
+ parallel_agnostic=False,
247
+ )
248
+ assert fm.fields["tp_size"] == 2
249
+ assert fm.rank == 1
250
+ assert fm.fields["parallel_agnostic"] is False
251
+ assert "replicated_layout" not in fm.fields
252
+ baseline = make_mapper_from_offloading_spec(
253
+ tp_size=2,
254
+ world_size=2,
255
+ rank=1,
256
+ replicated_layout=False,
257
+ parallel_agnostic=False,
258
+ )
259
+ assert fm.base_path == baseline.base_path
260
+
261
+
262
+def test_non_replicated_keeps_parallel_identity():
263
+ fm = make_mapper_from_offloading_spec(
264
+ tp_size=4,
265
+ world_size=4,
266
+ rank=2,
267
+ replicated_layout=False,
268
+ is_parallelism_agnostic=False,
269
+ parallel_agnostic=True,
270
+ )
271
+ assert fm.fields["tp_size"] == 4
272
+ assert fm.rank == 2
273
+ assert fm.fields["parallel_agnostic"] is False
274
+ assert fm.get_file_name(make_offload_key(b"\x02" * 8, 0)).startswith(
275
+ f"{fm.base_path}_r2/"
276
+ )
277
+
278
+
279
+def test_replicated_and_parallelism_agnostic_separate_layouts():
280
+ shared = dict(
281
+ model_name="shared-model",
282
+ groups=((16, "layer0"),),
283
+ tp_size=2,
284
+ world_size=2,
285
+ rank=1,
286
+ parallel_agnostic=True,
287
+ )
288
+ via_agnostic = make_mapper_from_offloading_spec(
289
+ is_parallelism_agnostic=True,
290
+ replicated_layout=False,
291
+ **shared,
292
+ )
293
+ via_replicated = make_mapper_from_offloading_spec(
294
+ is_parallelism_agnostic=False,
295
+ replicated_layout=True,
296
+ **shared,
297
+ )
298
+ assert via_agnostic.base_path != via_replicated.base_path
299
+ assert "replicated_layout" not in via_agnostic.fields
300
+ assert via_replicated.fields["replicated_layout"] is True
301
+
302
+
303
+def test_replicated_layout_run_config_tp_invariant():
304
+ shared = dict(
305
+ model_name="mla-model",
306
+ groups=((16, "mla_layer"),),
307
+ replicated_layout=True,
308
+ parallel_agnostic=True,
309
+ )
310
+ tp2 = make_mapper_from_offloading_spec(tp_size=2, world_size=2, rank=0, **shared)
311
+ tp4 = make_mapper_from_offloading_spec(tp_size=4, world_size=4, rank=2, **shared)
312
+ assert tp2.get_run_config() == tp4.get_run_config()
tests/v1/kv_offload/tiering/test_fs_tier.py
+61
-5
@@ -52,8 +52,18 @@ _DTYPE: torch.dtype = torch.float32
52
_CTX = ReqContext(req_id="test")
53
54
55
-def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock:
55
+def _make_offloading_spec(
56
+ enable_kv_cache_events: bool = False,
57
+ *,
58
+ tp_size: int = 1,
59
+ rank: int = 0,
60
+ world_size: int | None = None,
61
+ replicated_layout: bool = False,
62
+ is_parallelism_agnostic: bool = False,
63
+) -> MagicMock:
64
"""Mock spec with an explicit global KV events flag."""
65
+ if world_size is None:
66
+ world_size = tp_size
67
spec = MagicMock()
68
spec.config = OffloadingConfig(
69
groups=(),
@@ -64,15 +74,16 @@ def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock:
74
model=OffloadingModelConfig(name="test-model", dtype="float32"),
75
cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
76
parallel=OffloadingParallelConfig(
67
- rank=0,
68
- world_size=1,
69
- tp_size=1,
77
+ rank=rank,
78
+ world_size=world_size,
79
+ tp_size=tp_size,
80
pp_size=1,
81
pcp_size=1,
82
dcp_size=1,
83
data_parallel_index=0,
74
- is_parallelism_agnostic=False,
84
+ is_parallelism_agnostic=is_parallelism_agnostic,
85
),
86
+ replicated_layout=replicated_layout,
87
)
88
spec.blocks_per_chunk = 1
89
spec.kv_events_config = OffloadingKVEventsConfig(
@@ -725,3 +736,48 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path):
736
assert not fs_events[0].removed
737
finally:
738
tier.shutdown()
739
+
740
+
741
+def test_fs_tier_cross_tp_round_trip(tmp_path):
742
+ """TP=2 replicated writer and TP=4 reader share namespace and bytes."""
743
+ root = str(tmp_path)
744
+ writer_tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS)
745
+ expected = writer_tensor[0].clone()
746
+ writer = FileSystemTierManager(
747
+ offloading_spec=_make_offloading_spec(
748
+ tp_size=2, world_size=2, rank=0, replicated_layout=True
749
+ ),
750
+ primary_kv_view=memoryview(writer_tensor.numpy()),
751
+ tier_type="fs",
752
+ root_dir=root,
753
+ n_read_threads=2,
754
+ n_write_threads=2,
755
+ )
756
+ try:
757
+ writer.submit_store(make_job(1, [key(7)], [0]))
758
+ assert all(r.success for r in drain(writer))
759
+ writer_base = writer.file_mapper.base_path
760
+ writer_path = writer.file_mapper.get_file_name(key(7))
761
+ finally:
762
+ writer.shutdown()
763
+
764
+ reader_tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS)
765
+ reader = FileSystemTierManager(
766
+ offloading_spec=_make_offloading_spec(
767
+ tp_size=4, world_size=4, rank=3, replicated_layout=True
768
+ ),
769
+ primary_kv_view=memoryview(reader_tensor.numpy()),
770
+ tier_type="fs",
771
+ root_dir=root,
772
+ n_read_threads=2,
773
+ n_write_threads=2,
774
+ )
775
+ try:
776
+ assert reader.file_mapper.base_path == writer_base
777
+ assert reader.file_mapper.get_file_name(key(7)) == writer_path
778
+ assert lookup_and_wait(reader, [key(7)]) == [LookupResult.HIT]
779
+ reader.submit_load(make_job(2, [key(7)], [1], is_promotion=True))
780
+ assert all(r.success for r in drain(reader))
781
+ assert torch.allclose(reader_tensor[1], expected)
782
+ finally:
783
+ reader.shutdown()
tests/v1/kv_offload/tiering/test_obj_tier.py
+50
-5
@@ -43,7 +43,17 @@ from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManag
43
# ---------------------------------------------------------------------------
44
45
46
-def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig:
46
+def _make_offloading_config(
47
+ enable_kv_cache_events: bool,
48
+ *,
49
+ tp_size: int = 1,
50
+ rank: int = 0,
51
+ world_size: int | None = None,
52
+ replicated_layout: bool = False,
53
+ is_parallelism_agnostic: bool = False,
54
+) -> OffloadingConfig:
55
+ if world_size is None:
56
+ world_size = tp_size
57
return OffloadingConfig(
58
groups=(),
59
worker_kv_bytes_per_block=0,
@@ -53,15 +63,16 @@ def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig:
63
model=OffloadingModelConfig(name="test/model", dtype="float16"),
64
cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
65
parallel=OffloadingParallelConfig(
56
- rank=0,
57
- world_size=1,
58
- tp_size=1,
66
+ rank=rank,
67
+ world_size=world_size,
68
+ tp_size=tp_size,
69
pp_size=1,
70
pcp_size=1,
71
dcp_size=1,
72
data_parallel_index=0,
63
- is_parallelism_agnostic=False,
73
+ is_parallelism_agnostic=is_parallelism_agnostic,
74
),
75
+ replicated_layout=replicated_layout,
76
)
77
78
@@ -617,3 +628,37 @@ class TestObjStoreConfig:
628
params = cfg.to_nixl_params()
629
assert params["ca_bundle"] == "/path/to/ca.pem"
630
assert "access_key" not in params
631
+
632
+
633
+def test_obj_tier_replicated_layout_collapses_mapper_identity():
634
+ """TP=2 and TP=4 replicated configs share the obj FileMapper namespace."""
635
+ tp2_spec = SimpleNamespace(
636
+ config=_make_offloading_config(
637
+ False, tp_size=2, world_size=2, rank=1, replicated_layout=True
638
+ ),
639
+ kv_events_config=OffloadingKVEventsConfig(
640
+ enable_kv_cache_events=False,
641
+ self_describing_kv_events=False,
642
+ ),
643
+ )
644
+ tp4_spec = SimpleNamespace(
645
+ config=_make_offloading_config(
646
+ False, tp_size=4, world_size=4, rank=3, replicated_layout=True
647
+ ),
648
+ kv_events_config=OffloadingKVEventsConfig(
649
+ enable_kv_cache_events=False,
650
+ self_describing_kv_events=False,
651
+ ),
652
+ )
653
+ tp2_tier, _ = _make_tier(offloading_spec=tp2_spec)
654
+ tp4_tier, _ = _make_tier(offloading_spec=tp4_spec)
655
+ try:
656
+ assert tp2_tier._file_mapper.base_path == tp4_tier._file_mapper.base_path
657
+ assert tp2_tier._file_mapper.rank == 0
658
+ assert tp4_tier._file_mapper.rank == 0
659
+ assert tp2_tier._file_mapper.get_run_config() == (
660
+ tp4_tier._file_mapper.get_run_config()
661
+ )
662
+ finally:
663
+ tp2_tier.shutdown()
664
+ tp4_tier.shutdown()
vllm/v1/kv_offload/file_mapper.py
+10
-1
@@ -35,6 +35,7 @@ class FileMapper:
35
kv_cache_groups: list[dict] | None = None,
36
inference_engine: str = "vllm",
37
parallel_agnostic: bool = False,
38
+ replicated_layout: bool = False,
39
):
40
"""
41
Initialize the file mapper. Each worker constructs its own, but
@@ -60,6 +61,10 @@ class FileMapper:
61
}
62
if not parallel_agnostic:
63
self.fields["parallel_agnostic"] = False
64
+ # Only written when True so existing deployments' hashed fields are
65
+ # unchanged (False is the historical default and must not appear).
66
+ if replicated_layout:
67
+ self.fields["replicated_layout"] = True
68
self.base_path: str = self._compute_base_path(root_dir, self.fields)
69
70
@classmethod
@@ -92,7 +97,11 @@ class FileMapper:
97
rank=parallel.rank,
98
dtype=config.model.dtype,
99
kv_cache_groups=kv_cache_groups,
95
- parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic),
100
+ parallel_agnostic=(
101
+ parallel_agnostic
102
+ and (parallel.is_parallelism_agnostic or config.replicated_layout)
103
+ ),
104
+ replicated_layout=(parallel_agnostic and config.replicated_layout),
105
)
106
107
def get_file_name(self, key: OffloadKey) -> str: