[KV Offload] Expose data-parallel topology to offloading backends (#51879)
Signed-off-by: Ziqi Fan <ziqif@nvidia.com>
Ziqi Fan committed
Aug 12, 2026 at 20:27 UTC
f9a0f629bdb6ffa0202ed79f566e052e3cef42f4
8 files changed
+23
-1
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+5
-1
@@ -326,13 +326,17 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks():
326
] == [1, 3]
327
328
329
-def test_preserves_data_parallel_index():
329
+def test_preserves_data_parallel_config():
330
config = _make_vllm_config()
331
config.parallel_config.data_parallel_index = 2
332
+ config.parallel_config.data_parallel_size = 4
333
+ config.parallel_config.data_parallel_rank_local = 1
334
335
offloading_config = build_offloading_config(config, _make_kv_cache_config())
336
337
assert offloading_config.parallel.data_parallel_index == 2
338
+ assert offloading_config.parallel.data_parallel_size == 4
339
+ assert offloading_config.parallel.data_parallel_rank_local == 1
340
341
342
def test_resolves_heterogeneous_hybrid_block_sizes():
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+2
@@ -191,6 +191,8 @@ def _offloading_config(rank: int = 0) -> OffloadingConfig:
191
pcp_size=1,
192
dcp_size=1,
193
data_parallel_index=0,
194
+ data_parallel_size=1,
195
+ data_parallel_rank_local=None,
196
is_parallelism_agnostic=False,
197
),
198
)
tests/v1/kv_offload/test_factory.py
+4
@@ -49,6 +49,8 @@ def _make_offloading_config(
49
pcp_size: int = 1,
50
dcp_size: int = 1,
51
data_parallel_index: int = 0,
52
+ data_parallel_size: int = 1,
53
+ data_parallel_rank_local: int | None = None,
54
is_parallelism_agnostic: bool = False,
55
replicated_layout: bool = False,
56
extra_config: dict[str, Any] | None = None,
@@ -81,6 +83,8 @@ def _make_offloading_config(
83
pcp_size=pcp_size,
84
dcp_size=dcp_size,
85
data_parallel_index=data_parallel_index,
86
+ data_parallel_size=data_parallel_size,
87
+ data_parallel_rank_local=data_parallel_rank_local,
88
is_parallelism_agnostic=is_parallelism_agnostic,
89
),
90
replicated_layout=replicated_layout,
tests/v1/kv_offload/test_file_mapper.py
+2
@@ -49,6 +49,8 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
49
pcp_size=kwargs.get("pcp_size", 1),
50
dcp_size=kwargs.get("dcp_size", 1),
51
data_parallel_index=0,
52
+ data_parallel_size=1,
53
+ data_parallel_rank_local=None,
54
is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False),
55
),
56
replicated_layout=kwargs.get("replicated_layout", False),
tests/v1/kv_offload/tiering/test_fs_tier.py
+2
@@ -81,6 +81,8 @@ def _make_offloading_spec(
81
pcp_size=1,
82
dcp_size=1,
83
data_parallel_index=0,
84
+ data_parallel_size=1,
85
+ data_parallel_rank_local=None,
86
is_parallelism_agnostic=is_parallelism_agnostic,
87
),
88
replicated_layout=replicated_layout,
tests/v1/kv_offload/tiering/test_obj_tier.py
+2
@@ -74,6 +74,8 @@ def _make_offloading_config(
74
pcp_size=1,
75
dcp_size=1,
76
data_parallel_index=0,
77
+ data_parallel_size=1,
78
+ data_parallel_rank_local=None,
79
is_parallelism_agnostic=is_parallelism_agnostic,
80
),
81
replicated_layout=replicated_layout,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+2
@@ -213,6 +213,8 @@ def build_offloading_config(
213
pcp_size=parallel_config.prefill_context_parallel_size,
214
dcp_size=parallel_config.decode_context_parallel_size,
215
data_parallel_index=parallel_config.data_parallel_index,
216
+ data_parallel_size=parallel_config.data_parallel_size,
217
+ data_parallel_rank_local=parallel_config.data_parallel_rank_local,
218
is_parallelism_agnostic=is_parallelism_agnostic,
219
),
220
replicated_layout=replicated_layout,
vllm/v1/kv_offload/config.py
+4
@@ -48,6 +48,10 @@ class OffloadingParallelConfig:
48
dcp_size: int
49
# Data parallel replica index of this engine.
50
data_parallel_index: int
51
+ # Number of data parallel replicas.
52
+ data_parallel_size: int
53
+ # Local rank of the data parallel group, set only in SPMD mode.
54
+ data_parallel_rank_local: int | None
55
# True when the bytes that will be persisted for a block are portable
56
# across parallelism configurations: for the direct layout, concatenating
57
# the block's data across all workers in rank order yields the same bytes