[KV Connector] Support symmetric DCP disagg for hybrid mamba models (#55531)
Signed-off-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster> Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Wei Zhao <weizha@oci-aga-slurm-1-vscode-02.cm.cluster> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wei Zhao committed
Sep 10, 2026 at 14:12 UTC
7cdd9304ae2e46572f220741bf86e0b3c2da569c
5 files changed
+31
-12
vllm/config/parallel.py
+7
@@ -393,8 +393,15 @@ class ParallelConfig:
393
in (rank i+1, block j) only after (rank i, block j) is fully occupied.
394
Block_size should be greater than or equal to cp_kv_cache_interleave_size.
395
Block_size should be divisible by cp_kv_cache_interleave_size.
396
+
397
+ When --cp-kv-cache-interleave-size is omitted (None), the interleave size
398
+ is resolved automatically based on NIXL transfer requirements.
399
+ Explicit settings take priority.
400
"""
401
402
+ _allow_auto_resolve_cp_interleave_size: bool = True
403
+ """Whether NIXL may select the interleave size automatically."""
404
+
405
data_parallel_index: int = Field(init=False)
406
"""Equal to the data parallel rank but not used for torch process groups
407
and not overridden for dense models."""
vllm/config/vllm.py
+2
-5
@@ -1237,11 +1237,6 @@ class VllmConfig:
1237
"PD with decode_context_parallel_size > 1 is only "
1238
"supported for MLA models."
1239
)
1240
- assert not (self.model_config.is_hybrid and dcp_size > 1), (
1241
- "PD with decode_context_parallel_size > 1 is not "
1242
- "supported for hybrid Mamba/SSM models."
1243
- )
1244
-
1240
if self.lora_config is not None:
1241
self.lora_config.verify_with_model_config(self.model_config)
1242
@@ -2883,6 +2878,8 @@ class VllmConfig:
2878
"NixlConnector"
2879
):
2880
return
2881
+ if not self.parallel_config._allow_auto_resolve_cp_interleave_size:
2882
+ return
2883
2884
# Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid
2885
# scaling by dcp_size (we need the local block_size here).
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+5
-3
@@ -530,9 +530,6 @@ class NixlBaseConnectorWorker:
530
# DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully
531
# replicated or fully sharded. A DCP rank is always derivable this way.
532
self.dcp_rank = self.tp_rank % self.dcp_size
533
- if self._has_mamba and self.dcp_size > 1:
534
- # Prefix-cache-aware DCP slicing isn't implemented for the Mamba group.
535
- raise ValueError("DCP is not supported for hybrid MLA+Mamba models.")
533
534
self.num_blocks = kv_cache_config.num_blocks
535
self.enable_permute_local_kv = False
@@ -2177,6 +2174,11 @@ class NixlBaseConnectorWorker:
2174
f"DCP sizes must divide one another: local={self.dcp_size}, "
2175
f"remote={remote_dcp_size} (engine {remote_engine_id})."
2176
)
2177
+ if self._has_mamba and self.dcp_size != remote_dcp_size:
2178
+ raise RuntimeError(
2179
+ "Hybrid MLA+Mamba NIXL transfers require matching DCP sizes, "
2180
+ f"got local={self.dcp_size}, remote={remote_dcp_size}."
2181
+ )
2182
2183
tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
2184
block_size_ratio = self.transfer_topo.block_size_ratio(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+4
-1
@@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
17
)
18
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
19
ReadSpec,
20
+ _is_attention_spec,
21
)
22
from vllm.logger import init_logger
23
@@ -213,7 +214,9 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
214
local_ids = group_ids(meta.local_block_ids, rank)
215
remote_ids = group_ids(remote_logical_block_ids, rank)
216
for g in range(num_groups):
216
- if not local_ids[g]:
217
+ if not local_ids[g] or not _is_attention_spec(
218
+ self._group_spec_types[g]
219
+ ):
220
continue
221
local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
222
local_ids[g],
vllm/engine/arg_utils.py
+13
-3
@@ -487,7 +487,7 @@ class EngineArgs:
487
dcp_comm_backend: DCPCommBackend | None = ParallelConfig.dcp_comm_backend
488
dcp_q_replicate: bool | None = ParallelConfig.dcp_q_replicate
489
dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
490
- cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
490
+ cp_kv_cache_interleave_size: int | None = None
491
data_parallel_size: int = ParallelConfig.data_parallel_size
492
data_parallel_rank: int | None = None
493
data_parallel_start_rank: int | None = None
@@ -1104,7 +1104,10 @@ class EngineArgs:
1104
)
1105
parallel_group.add_argument(
1106
"--cp-kv-cache-interleave-size",
1107
- **parallel_kwargs["cp_kv_cache_interleave_size"],
1107
+ **{
1108
+ **parallel_kwargs["cp_kv_cache_interleave_size"],
1109
+ "default": None,
1110
+ },
1111
)
1112
parallel_group.add_argument(
1113
"--prefill-context-parallel-size",
@@ -2363,7 +2366,14 @@ class EngineArgs:
2366
dcp_comm_backend=self.dcp_comm_backend,
2367
dcp_q_replicate=self.dcp_q_replicate,
2368
dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
2366
- cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
2369
+ cp_kv_cache_interleave_size=(
2370
+ self.cp_kv_cache_interleave_size
2371
+ if self.cp_kv_cache_interleave_size is not None
2372
+ else ParallelConfig.cp_kv_cache_interleave_size
2373
+ ),
2374
+ _allow_auto_resolve_cp_interleave_size=(
2375
+ self.cp_kv_cache_interleave_size is None
2376
+ ),
2377
_api_process_count=self._api_process_count,
2378
_api_process_rank=self._api_process_rank,
2379
assigned_physical_gpu_ids=self._resolve_device_ids(),