[Bugfix] `adjust_dcp_kv_cache_interleave_size` for NixlConnector only (#54803)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: cjackal <44624812+cjackal@users.noreply.github.com>
Nicolò Lucchesi committed
Sep 2, 2026 at 12:31 UTC
f81eb4193431675918f37483d1541f6b1b96a69f
3 files changed
+91
-22
tests/test_config.py
+74
-5
@@ -114,7 +114,34 @@ def test_per_request_spec_decode_metrics_requires_spec_decode():
114
)
115
116
117
-def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog):
117
+@pytest.mark.parametrize(
118
+ "kv_transfer_config",
119
+ [
120
+ KVTransferConfig(
121
+ kv_connector="NixlConnector",
122
+ kv_role="kv_both",
123
+ ),
124
+ KVTransferConfig(
125
+ kv_connector="MultiConnector",
126
+ kv_role="kv_both",
127
+ kv_connector_extra_config={
128
+ "connectors": [
129
+ {
130
+ "kv_connector": "NixlConnector",
131
+ "kv_role": "kv_both",
132
+ },
133
+ {
134
+ "kv_connector": "OffloadingConnector",
135
+ "kv_role": "kv_both",
136
+ },
137
+ ]
138
+ },
139
+ ),
140
+ ],
141
+)
142
+def test_pd_dcp_interleave_size_is_adjusted_to_block_size(
143
+ caplog, disable_log_dedup, kv_transfer_config
144
+):
145
config = VllmConfig(
146
cache_config=CacheConfig(block_size=16),
147
device_config=DeviceConfig(device="cpu"),
@@ -124,10 +151,7 @@ def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog):
151
cp_kv_cache_interleave_size=3,
152
distributed_executor_backend="mp",
153
),
127
- kv_transfer_config=KVTransferConfig(
128
- kv_connector="NixlConnector",
129
- kv_role="kv_both",
130
- ),
154
+ kv_transfer_config=kv_transfer_config,
155
)
156
157
kv_cache_config = SimpleNamespace(
@@ -140,6 +164,51 @@ def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog):
164
assert "automatically adjusted from 3 to block_size 16" in caplog.text
165
166
167
+def test_kv_offloading_does_not_adjust_dcp_interleave_size():
168
+ config = VllmConfig(
169
+ cache_config=CacheConfig(block_size=16),
170
+ device_config=DeviceConfig(device="cpu"),
171
+ parallel_config=ParallelConfig(
172
+ tensor_parallel_size=2,
173
+ decode_context_parallel_size=2,
174
+ cp_kv_cache_interleave_size=1,
175
+ distributed_executor_backend="mp",
176
+ ),
177
+ kv_transfer_config=KVTransferConfig(
178
+ kv_connector="OffloadingConnector",
179
+ kv_role="kv_both",
180
+ ),
181
+ )
182
+
183
+ kv_cache_config = SimpleNamespace(
184
+ kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))]
185
+ )
186
+ config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
187
+
188
+ assert config.parallel_config.cp_kv_cache_interleave_size == 1
189
+
190
+
191
+def test_kv_offloading_does_not_skip_dcp_interleave_validation():
192
+ config = SimpleNamespace(
193
+ cache_config=SimpleNamespace(
194
+ block_size=16,
195
+ mamba_cache_mode="none",
196
+ ),
197
+ parallel_config=SimpleNamespace(
198
+ decode_context_parallel_size=2,
199
+ cp_kv_cache_interleave_size=3,
200
+ ),
201
+ scheduler_config=SimpleNamespace(disable_chunked_mm_input=False),
202
+ kv_transfer_config=KVTransferConfig(
203
+ kv_connector="OffloadingConnector",
204
+ kv_role="kv_both",
205
+ ),
206
+ )
207
+
208
+ with pytest.raises(AssertionError, match="divisible by"):
209
+ VllmConfig.validate_block_size(config)
210
+
211
+
212
def test_compile_config_repr_succeeds():
213
# setup: VllmBackend mutates the config object
214
config = VllmConfig()
vllm/config/vllm.py
+16
-16
@@ -2746,7 +2746,7 @@ class VllmConfig:
2746
def adjust_dcp_kv_cache_interleave_size(
2747
self, kv_cache_config: "KVCacheConfig"
2748
) -> None:
2749
- """Normalize DCP interleave size against the resolved block_size for PD.
2749
+ """Normalize DCP interleave size against block_size for NIXL P/D.
2750
2751
Called by each worker (via ensure_kv_transfer_initialized), once it knows its
2752
own final block_size via kv_cache_config.
@@ -2754,11 +2754,6 @@ class VllmConfig:
2754
dcp_size = self.parallel_config.decode_context_parallel_size
2755
if dcp_size <= 1:
2756
return
2757
- # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid
2758
- # scaling by dcp_size (we need the local block_size here).
2759
- local_block_size = min(
2760
- g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups
2761
- )
2757
if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
2758
self.parallel_config.cp_kv_cache_interleave_size
2759
!= self.parallel_config.dcp_kv_cache_interleave_size
@@ -2772,11 +2767,17 @@ class VllmConfig:
2767
"deprecated when PCP is fully supported."
2768
)
2769
2775
- if (
2776
- self.kv_transfer_config is not None
2777
- and self.kv_transfer_config.kv_connector is not None
2778
- and self.parallel_config.cp_kv_cache_interleave_size != local_block_size
2770
+ if self.kv_transfer_config is None or not self.kv_transfer_config.has_connector(
2771
+ "NixlConnector"
2772
):
2773
+ return
2774
+
2775
+ # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid
2776
+ # scaling by dcp_size (we need the local block_size here).
2777
+ local_block_size = min(
2778
+ g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups
2779
+ )
2780
+ if self.parallel_config.cp_kv_cache_interleave_size != local_block_size:
2781
interleave = self.parallel_config.cp_kv_cache_interleave_size
2782
self.parallel_config.cp_kv_cache_interleave_size = local_block_size
2783
logger.info_once(
@@ -2797,14 +2798,13 @@ class VllmConfig:
2798
"""
2799
block_size = self.cache_config.block_size
2800
2800
- # Skip DCP interleave-size compatibility when a KV connector is configured:
2801
- # cp_kv_cache_interleave_size is pinned to block_size for PD by each worker
2802
- pd_active = (
2801
+ # Skip DCP interleave-size compatibility for NIXL P/D: the interleave
2802
+ # size is pinned to block_size by each worker.
2803
+ nixl_pd_active = (
2804
self.kv_transfer_config is not None
2804
- and self.kv_transfer_config.kv_connector is not None
2805
- and self.kv_transfer_config.is_kv_transfer_instance
2805
+ and self.kv_transfer_config.has_connector("NixlConnector")
2806
)
2807
- if self.parallel_config.decode_context_parallel_size > 1 and not pd_active:
2807
+ if self.parallel_config.decode_context_parallel_size > 1 and not nixl_pd_active:
2808
assert (
2809
self.parallel_config.cp_kv_cache_interleave_size <= block_size
2810
and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0
vllm/distributed/kv_transfer/kv_transfer_state.py
+1
-1
@@ -85,7 +85,7 @@ def ensure_kv_transfer_initialized(
85
vllm_config.kv_transfer_config.is_kv_transfer_instance
86
and _KV_CONNECTOR_AGENT is None
87
):
88
- # PD only supports an interleave_size equal to block_size.
88
+ # NIXL P/D requires an interleave_size equal to block_size.
89
vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
90
_sync_engine_id_across_tp(vllm_config)
91