@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7f4793eaa3

[Nixl][PD] DCP support for MLA models (#50611)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Aug 29, 2026 at 18:03 UTC 7f4793eaa335a3667927a0191868d01f36b170af
23 files changed +818 -104
tests/test_config.py
+30
@@ -17,8 +17,11 @@ import vllm.config.vllm as vllm_config_module
17 import vllm.envs as envs
18 from vllm.compilation.backends import VllmBackend
19 from vllm.config import (
20 + CacheConfig,
21 CompilationConfig,
22 + DeviceConfig,
23 KernelConfig,
24 + KVTransferConfig,
25 ModelConfig,
26 ObservabilityConfig,
27 ParallelConfig,
@@ -100,6 +103,32 @@ def test_per_request_spec_decode_metrics_requires_spec_decode():
103 )
104
105
106 +def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog):
107 + config = VllmConfig(
108 + cache_config=CacheConfig(block_size=16),
109 + device_config=DeviceConfig(device="cpu"),
110 + parallel_config=ParallelConfig(
111 + tensor_parallel_size=2,
112 + decode_context_parallel_size=2,
113 + cp_kv_cache_interleave_size=3,
114 + distributed_executor_backend="mp",
115 + ),
116 + kv_transfer_config=KVTransferConfig(
117 + kv_connector="NixlConnector",
118 + kv_role="kv_both",
119 + ),
120 + )
121 +
122 + kv_cache_config = SimpleNamespace(
123 + kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))]
124 + )
125 + with caplog.at_level(logging.INFO):
126 + config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
127 +
128 + assert config.parallel_config.cp_kv_cache_interleave_size == 16
129 + assert "automatically adjusted from 3 to block_size 16" in caplog.text
130 +
131 +
132 def test_compile_config_repr_succeeds():
133 # setup: VllmBackend mutates the config object
134 config = VllmConfig()
@@ -1730,6 +1759,7 @@ def test_validate_mamba_align_subblock_prefill():
1759 long_prefill_token_threshold=4096,
1760 disable_chunked_mm_input=False,
1761 ),
1762 + kv_transfer_config=None,
1763 )
1764
1765 VllmConfig.validate_block_size(config)
tests/v1/kv_connector/unit/test_nixl_connector.py
+1
@@ -507,6 +507,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
507 port: int,
508 remote_tp_size: int,
509 expected_engine_id: str,
510 + remote_dcp_size: int = 1,
511 remote_pp_size: int = 1,
512 notif_agents_only: bool = False,
513 ) -> tuple[dict[tuple[int, int], str], float]:
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+181 -1
@@ -3,7 +3,7 @@
3 """Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""
4
5 import gc
6 -from unittest.mock import patch
6 +from unittest.mock import MagicMock, patch
7
8 import pytest
9 import torch
@@ -63,6 +63,180 @@ def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes):
63 )
64
65
66 +@pytest.mark.cpu_test
67 +@pytest.mark.parametrize(
68 + "use_mla,source_ranks,tp_ratio,expected",
69 + [
70 + pytest.param(True, (0,), -2, False, id="pure_mla_with_remote_dcp"),
71 + pytest.param(True, (0, 1), -2, True, id="hybrid_mla_ssm"),
72 + pytest.param(False, (0,), -2, True, id="full_attention"),
73 + pytest.param(False, (0,), 2, False, id="local_tp_greater"),
74 + ],
75 +)
76 +def test_needs_split_local_xfer_handles(use_mla, source_ranks, tp_ratio, expected):
77 + """Handle creation and selection must use the same TP-mapping predicate.
78 +
79 + In pure MLA, DCP can target several physical remote workers even though
80 + the TP mapping has one replicated attention source. Those workers own
81 + disjoint blocks, so every read uses the whole local-region handle.
82 + """
83 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
84 + NixlConnectorWorker,
85 + )
86 +
87 + worker = object.__new__(NixlConnectorWorker)
88 + worker.use_mla = use_mla
89 + plan = MagicMock()
90 + plan.all_source_ranks = source_ranks
91 +
92 + assert worker._needs_split_local_xfer_handles(tp_ratio, plan) is expected
93 +
94 +
95 +@pytest.mark.cpu_test
96 +def test_update_state_after_alloc_tracks_cached_blocks_per_group():
97 + """Hybrid SWA+FA requests can have different prefix-cache-hit counts per
98 + KV cache group. Each group's count must be tracked independently rather
99 + than collapsed to a single scalar, or DCP read-phase alignment
100 + (local_num_computed_blocks) would misalign one of the groups."""
101 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
102 + NixlPullConnectorScheduler,
103 + )
104 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
105 + from vllm.v1.core.kv_cache_utils import KVCacheBlock
106 +
107 + scheduler = object.__new__(NixlPullConnectorScheduler)
108 + scheduler._reqs_in_batch = set()
109 + scheduler._reqs_need_save = {}
110 + scheduler._reqs_need_recv = {}
111 + scheduler.use_host_buffer = False
112 + scheduler.is_bidirectional_kv_xfer_enabled = False
113 + scheduler._is_hma_required = False
114 + scheduler.kv_cache_config = MagicMock(
115 + select_transfer_block_ids=lambda block_ids: block_ids
116 + )
117 +
118 + def cached(block_id):
119 + return KVCacheBlock(block_id=block_id, _block_hash=object())
120 +
121 + def uncached(block_id):
122 + return KVCacheBlock(block_id=block_id)
123 +
124 + # Group 0 (full-attention): 2 of 3 blocks cached.
125 + # Group 1 (sliding-window): 1 of 3 blocks cached.
126 + blocks = KVCacheBlocks(
127 + blocks=(
128 + [cached(0), cached(1), uncached(2)],
129 + [cached(10), uncached(11), uncached(12)],
130 + )
131 + )
132 +
133 + request = create_request(do_remote_prefill=True)
134 +
135 + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2)
136 +
137 + _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id]
138 + assert local_num_computed_blocks == (2, 1)
139 +
140 +
141 +@pytest.mark.cpu_test
142 +@pytest.mark.parametrize(
143 + "local_ids,remote_ids,remote_rank,local_dcp_size,local_dcp_rank,"
144 + "remote_dcp_size,local_num_computed_blocks,expected_local,expected_remote",
145 + [
146 + # local more finely sharded than remote: remote's raw slice has an
147 + # extra tail block beyond what local's phase-aligned slice covers.
148 + pytest.param(
149 + [100, 101, 102],
150 + [200, 201],
151 + 1,
152 + 1,
153 + 0,
154 + 2,
155 + 0,
156 + [101],
157 + [200],
158 + id="remote_tail_padding",
159 + ),
160 + # remote more finely sharded than local: local's raw slice has a
161 + # block beyond what remote's phase-aligned slice covers.
162 + pytest.param(
163 + [100],
164 + [200],
165 + 0,
166 + 2,
167 + 1,
168 + 1,
169 + 0,
170 + [],
171 + [],
172 + id="local_tail_padding",
173 + ),
174 + # Asymmetric DCP with a partial prefix cache hit: the cached count
175 + # shifts both the local phase and the remote start position.
176 + pytest.param(
177 + [100],
178 + [200, 201],
179 + 2,
180 + 2,
181 + 0,
182 + 4,
183 + 3,
184 + [100],
185 + [201],
186 + id="prefix_phase",
187 + ),
188 + # Equal (but active) DCP sizes still need to skip the locally cached
189 + # prefix from the front of remote's list.
190 + pytest.param(
191 + [100, 101],
192 + [200, 201, 202, 203],
193 + 1,
194 + 2,
195 + 1,
196 + 2,
197 + 2,
198 + [100, 101],
199 + [202, 203],
200 + id="equal_dcp_with_prefix",
201 + ),
202 + ],
203 +)
204 +def test_apply_dcp_prefix_caching_matches_global_logical_positions(
205 + local_ids,
206 + remote_ids,
207 + remote_rank,
208 + local_dcp_size,
209 + local_dcp_rank,
210 + remote_dcp_size,
211 + local_num_computed_blocks,
212 + expected_local,
213 + expected_remote,
214 +):
215 + """_apply_dcp_prefix_caching must operate on logical block IDs (DCP
216 + interleaves at logical-block granularity) and always return equal-length
217 + local/remote slices, so a later, unconditional _apply_prefix_caching
218 + trim is a guaranteed no-op rather than double-processing the read."""
219 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
220 + NixlConnectorWorker,
221 + )
222 +
223 + worker = object.__new__(NixlConnectorWorker)
224 +
225 + actual_local, actual_remote = worker._apply_dcp_prefix_caching(
226 + local_ids,
227 + remote_ids,
228 + remote_rank=remote_rank,
229 + local_dcp_size=local_dcp_size,
230 + local_dcp_rank=local_dcp_rank,
231 + remote_dcp_size=remote_dcp_size,
232 + local_num_computed_blocks=local_num_computed_blocks,
233 + )
234 +
235 + assert actual_local == expected_local
236 + assert actual_remote == expected_remote
237 + assert len(actual_local) == len(actual_remote)
238 +
239 +
240 @pytest.mark.cpu_test
241 def test_logical_to_kernel_block_ids_with_hma():
242 """Test _logical_to_kernel_block_ids expands blocks when HMA is enabled.
@@ -211,7 +385,12 @@ def test_read_blocks_for_req_expands_remote_ids(
385 worker = object.__new__(NixlConnectorWorker)
386 worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
387 worker._engine_last_active = {}
388 + worker._recving_transfers = {}
389 worker._bidirectional_kv_xfer_enabled = False
390 + worker.dcp_size = 1
391 + worker.dcp_rank = 0
392 + worker._group_spec_types = resolved_types
393 + worker._has_mamba = any(t is MambaSpec for t in resolved_types)
394
395 has_mamba = any(t is MambaSpec for t in resolved_types)
396 has_swa = any(t is SlidingWindowSpec for t in resolved_types)
@@ -227,6 +406,7 @@ def test_read_blocks_for_req_expands_remote_ids(
406 worker.transfer_topo.tp_ratio.return_value = tp_ratio
407 remote_info = MagicMock()
408 remote_info.remote_physical_blocks_per_logical = remote_physical_per_logical
409 + remote_info.remote_dcp_size = 1
410 worker.transfer_topo.get_engine_info.return_value = remote_info
411 worker.use_mla = False
412
tests/v1/kv_connector/unit/test_tp_mapping.py
+73 -6
@@ -14,6 +14,7 @@ from types import SimpleNamespace
14 import numpy as np
15 import pytest
16
17 +from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology
18 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
19 TPMapping,
20 compute_tp_mapping,
@@ -35,17 +36,20 @@ def _compute_mapping(
36 is_mla: bool = False,
37 num_kv_heads: int = 8,
38 group_spec_types: tuple[type, ...] = (FullAttentionSpec,),
39 + dcp_size: int = 1,
40 + remote_dcp_size: int = 1,
41 ) -> TPMapping:
39 - transfer_topology = SimpleNamespace(
40 - tp_rank=tp_rank,
41 - tp_size=tp_size,
42 - is_mla=is_mla,
43 - total_num_kv_heads=num_kv_heads,
44 - )
42 + transfer_topology = object.__new__(TransferTopology)
43 + transfer_topology.tp_rank = tp_rank
44 + transfer_topology.tp_size = tp_size
45 + transfer_topology.is_mla = is_mla
46 + transfer_topology.total_num_kv_heads = num_kv_heads
47 + transfer_topology.dcp_size = dcp_size
48 return compute_tp_mapping(
49 transfer_topology=transfer_topology,
50 remote_tp_size=remote_tp_size,
51 group_spec_types=group_spec_types,
52 + remote_dcp_size=remote_dcp_size,
53 )
54
55
@@ -68,6 +72,69 @@ class TestTPMappingStructure:
72 assert m.all_source_ranks == (0, 1)
73
74
75 +@pytest.mark.parametrize(
76 + "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks",
77 + [
78 + (0, 4, 4, 1, 4, (0, 1, 2, 3)),
79 + (2, 4, 4, 4, 4, (2,)),
80 + (0, 2, 4, 2, 4, (0, 2)),
81 + (3, 4, 2, 4, 2, (1,)),
82 + ],
83 +)
84 +def test_mla_dcp_source_ranks(
85 + tp_rank,
86 + tp_size,
87 + remote_tp_size,
88 + dcp_size,
89 + remote_dcp_size,
90 + expected_ranks,
91 +):
92 + mapping = _compute_mapping(
93 + tp_rank=tp_rank,
94 + tp_size=tp_size,
95 + remote_tp_size=remote_tp_size,
96 + is_mla=True,
97 + num_kv_heads=1,
98 + dcp_size=dcp_size,
99 + remote_dcp_size=remote_dcp_size,
100 + )
101 +
102 + assert mapping.all_source_ranks == expected_ranks
103 +
104 +
105 +@pytest.mark.parametrize(
106 + "tp_size,remote_tp_size,dcp_size,remote_dcp_size",
107 + [
108 + (4, 2, 4, 2),
109 + (4, 4, 4, 4),
110 + (4, 1, 4, 1),
111 + (4, 4, 1, 4),
112 + (2, 4, 2, 4),
113 + ],
114 +)
115 +def test_dcp_consumer_count_matches_readers(
116 + tp_size, remote_tp_size, dcp_size, remote_dcp_size
117 +):
118 + """Each producer waits for exactly the local ranks that read from it."""
119 + mappings = [
120 + _compute_mapping(
121 + tp_rank=tp_rank,
122 + tp_size=tp_size,
123 + remote_tp_size=remote_tp_size,
124 + is_mla=True,
125 + num_kv_heads=1,
126 + dcp_size=dcp_size,
127 + remote_dcp_size=remote_dcp_size,
128 + )
129 + for tp_rank in range(tp_size)
130 + ]
131 +
132 + for remote_rank in range(remote_tp_size):
133 + readers = [m for m in mappings if remote_rank in m.all_source_ranks]
134 + for mapping in readers:
135 + assert mapping.local_consumers == len(readers)
136 +
137 +
138 # ======================================================================
139 # Split handle tests
140 # ======================================================================
tests/v1/worker/test_gpu_model_runner.py
+71
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import gc
5 +from contextlib import nullcontext
6 from types import SimpleNamespace
7 from unittest.mock import Mock
8
@@ -1380,6 +1381,76 @@ def test_hybrid_attention_mamba_tensor_shapes():
1381 assert torch.equal(actual_ssm, expected_ssm)
1382
1383
1384 +def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch):
1385 + runner = object.__new__(GPUModelRunner)
1386 + runner.vllm_config = SimpleNamespace(reasoning_config=None)
1387 + runner.parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=16)
1388 + runner.cache_config = SimpleNamespace(use_replayssm=False)
1389 + runner.model_config = SimpleNamespace(get_vocab_size=lambda: 32)
1390 + runner.max_model_len = 64
1391 + runner.max_encoder_len = 0
1392 + runner.max_num_reqs = 1
1393 + runner.max_num_tokens = 64
1394 + runner.num_spec_tokens = 0
1395 + runner.device = torch.device("cpu")
1396 + runner.is_pooling_model = False
1397 + runner._init_block_sizes = [16]
1398 + runner._init_kernel_block_sizes = [16]
1399 + runner._init_max_num_blocks = [4]
1400 + runner._init_slot_mapping_modes = [
1401 + gpu_model_runner_module.SlotMappingMode.TOKEN_TO_KV_SLOT
1402 + ]
1403 + runner.cp_kv_cache_interleave_size = 1
1404 + runner.input_batch = SimpleNamespace(
1405 + logitsprocs=None,
1406 + logitsprocs_need_output_token_ids=False,
1407 + )
1408 + runner.jit_warmup_registry = Mock()
1409 + runner.jit_warmup_registry.activate.return_value = nullcontext()
1410 +
1411 + spec = SimpleNamespace(
1412 + block_size=16,
1413 + max_num_blocks_per_req=lambda *_: 4,
1414 + )
1415 + kv_cache_config = SimpleNamespace(
1416 + kv_cache_groups=[SimpleNamespace(kv_cache_spec=spec)]
1417 + )
1418 + input_batch_cls = Mock(return_value=SimpleNamespace())
1419 + monkeypatch.setattr(gpu_model_runner_module, "InputBatch", input_batch_cls)
1420 + monkeypatch.setattr(
1421 + gpu_model_runner_module,
1422 + "get_kv_cache_spec_kind",
1423 + lambda _: gpu_model_runner_module.KVCacheSpecKind.FULL_ATTENTION,
1424 + )
1425 +
1426 + runner.may_reinitialize_input_batch(kv_cache_config, [16])
1427 +
1428 + assert input_batch_cls.call_count == 1
1429 + assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16
1430 +
1431 +
1432 +def test_v2_runner_snapshots_late_interleave_adjustment(monkeypatch):
1433 + from vllm.v1.worker.gpu import model_runner as v2_model_runner_module
1434 +
1435 + runner = object.__new__(v2_model_runner_module.GPUModelRunner)
1436 + runner.parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=16)
1437 + runner.cp_interleave = 1
1438 +
1439 + class StopInitialization(Exception):
1440 + pass
1441 +
1442 + monkeypatch.setattr(
1443 + v2_model_runner_module,
1444 + "deepcopy",
1445 + Mock(side_effect=StopInitialization),
1446 + )
1447 +
1448 + with pytest.raises(StopInitialization):
1449 + runner.initialize_kv_cache(SimpleNamespace())
1450 +
1451 + assert runner.cp_interleave == 16
1452 +
1453 +
1454 def test_hybrid_block_table_initialization():
1455 """Test hybrid block table with different kernel and kvcache_manager block
1456 sizes."""
tests/v1/worker/test_gpu_model_runner_v2.py
+6 -1
@@ -57,14 +57,19 @@ def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table(
57 KVCacheGroupSpec(["kda"], mamba_spec),
58 ],
59 )
60 + parallel_config = SimpleNamespace(
61 + decode_context_parallel_size=dcp_size,
62 + cp_kv_cache_interleave_size=1,
63 + )
64 vllm_config = SimpleNamespace(
61 - parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size),
65 + parallel_config=parallel_config,
66 cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode),
67 )
68 runner = SimpleNamespace(
69 max_model_len=max_model_len,
70 is_encoder_decoder=False,
71 vllm_config=vllm_config,
72 + parallel_config=parallel_config,
73 )
74
75 class _CapturedWidths(Exception):
vllm/config/kv_transfer.py
+9
@@ -119,3 +119,12 @@ class KVTransferConfig:
119
120 def get_from_extra_config(self, key, default) -> Any:
121 return self.kv_connector_extra_config.get(key, default)
122 +
123 + def has_connector(self, connector_name: str) -> bool:
124 + """Whether ``connector_name`` is configured, directly or in MultiConnector."""
125 + if self.kv_connector == connector_name:
126 + return True
127 + return self.kv_connector == "MultiConnector" and any(
128 + child.get("kv_connector") == connector_name
129 + for child in self.kv_connector_extra_config.get("connectors", [])
130 + )
vllm/config/vllm.py
+78 -15
@@ -1141,6 +1141,30 @@ class VllmConfig:
1141 self._verify_sampling_replay_config()
1142 self._verify_trace_replay_config()
1143
1144 + # A NIXL side is either fully replicated or fully DCP-sharded; MLA only.
1145 + if (
1146 + self.kv_transfer_config is not None
1147 + and self.kv_transfer_config.has_connector("NixlConnector")
1148 + ):
1149 + assert self.parallel_config.prefill_context_parallel_size == 1, (
1150 + "NIXL does not support prefill context parallelism."
1151 + )
1152 + dcp_size = self.parallel_config.decode_context_parallel_size
1153 + tp_size = self.parallel_config.tensor_parallel_size
1154 + assert dcp_size in (1, tp_size), (
1155 + f"decode_context_parallel_size={dcp_size} must be 1 or equal "
1156 + f"to tensor_parallel_size={tp_size} when using NixlConnector."
1157 + )
1158 + if self.model_config is not None:
1159 + assert self.model_config.use_mla or dcp_size == 1, (
1160 + "PD with decode_context_parallel_size > 1 is only "
1161 + "supported for MLA models."
1162 + )
1163 + assert not (self.model_config.is_hybrid and dcp_size > 1), (
1164 + "PD with decode_context_parallel_size > 1 is not "
1165 + "supported for hybrid Mamba/SSM models."
1166 + )
1167 +
1168 if self.lora_config is not None:
1169 self.lora_config.verify_with_model_config(self.model_config)
1170
@@ -2714,6 +2738,52 @@ class VllmConfig:
2738 f"Model Runner V1 does not support: {', '.join(unsupported)}"
2739 )
2740
2741 + def adjust_dcp_kv_cache_interleave_size(
2742 + self, kv_cache_config: "KVCacheConfig"
2743 + ) -> None:
2744 + """Normalize DCP interleave size against the resolved block_size for PD.
2745 +
2746 + Called by each worker (via ensure_kv_transfer_initialized), once it knows its
2747 + own final block_size via kv_cache_config.
2748 + """
2749 + dcp_size = self.parallel_config.decode_context_parallel_size
2750 + if dcp_size <= 1:
2751 + return
2752 + # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid
2753 + # scaling by dcp_size (we need the local block_size here).
2754 + local_block_size = min(
2755 + g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups
2756 + )
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
2760 + ):
2761 + self.parallel_config.cp_kv_cache_interleave_size = (
2762 + self.parallel_config.dcp_kv_cache_interleave_size
2763 + )
2764 + logger.warning_once(
2765 + "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
2766 + "_interleave_size. And dcp-kv-cache-interleave-size will be "
2767 + "deprecated when PCP is fully supported."
2768 + )
2769 +
2770 + if (
2771 + self.kv_transfer_config is not None
2772 + and self.kv_transfer_config.kv_connector is not None
2773 + and self.parallel_config.cp_kv_cache_interleave_size != local_block_size
2774 + ):
2775 + interleave = self.parallel_config.cp_kv_cache_interleave_size
2776 + self.parallel_config.cp_kv_cache_interleave_size = local_block_size
2777 + logger.info_once(
2778 + "When using PD disaggregation with DCP "
2779 + "(decode_context_parallel_size=%d), "
2780 + "cp_kv_cache_interleave_size is automatically adjusted "
2781 + "from %d to block_size %d for block-level alignment.",
2782 + dcp_size,
2783 + interleave,
2784 + local_block_size,
2785 + )
2786 +
2787 def validate_block_size(self) -> None:
2788 """Validate block_size against DCP and mamba constraints.
2789
@@ -2722,20 +2792,14 @@ class VllmConfig:
2792 """
2793 block_size = self.cache_config.block_size
2794
2725 - # DCP interleave-size compatibility
2726 - if self.parallel_config.decode_context_parallel_size > 1:
2727 - if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
2728 - self.parallel_config.cp_kv_cache_interleave_size
2729 - != self.parallel_config.dcp_kv_cache_interleave_size
2730 - ):
2731 - self.parallel_config.cp_kv_cache_interleave_size = (
2732 - self.parallel_config.dcp_kv_cache_interleave_size
2733 - )
2734 - logger.warning_once(
2735 - "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
2736 - "_interleave_size. And dcp-kv-cache-interleave-size will be "
2737 - "deprecated when PCP is fully supported."
2738 - )
2795 + # Skip DCP interleave-size compatibility when a KV connector is configured:
2796 + # cp_kv_cache_interleave_size is pinned to block_size for PD by each worker
2797 + pd_active = (
2798 + self.kv_transfer_config is not None
2799 + and self.kv_transfer_config.kv_connector is not None
2800 + and self.kv_transfer_config.is_kv_transfer_instance
2801 + )
2802 + if self.parallel_config.decode_context_parallel_size > 1 and not pd_active:
2803 assert (
2804 self.parallel_config.cp_kv_cache_interleave_size <= block_size
2805 and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0
@@ -2744,7 +2808,6 @@ class VllmConfig:
2808 "than or equal to and divisible by cp_kv_cache_interleave_size "
2809 f"({self.parallel_config.cp_kv_cache_interleave_size})."
2810 )
2747 -
2811 # Mamba cache align-mode constraints
2812 if self.cache_config.mamba_cache_mode == "align":
2813 assert not self.scheduler_config.disable_chunked_mm_input, (
vllm/distributed/kv_transfer/kv_connector/utils.py
+70 -3
@@ -404,6 +404,9 @@ class EngineTransferInfo:
404 end_layer: int = 0
405 """Exclusive global index after the last layer owned by this PP rank."""
406
407 + remote_dcp_size: int = 1
408 + """Remote decode context parallel size."""
409 +
410
411 # ---- Transfer topology ----
412
@@ -420,6 +423,7 @@ class TransferTopology:
423 is_mamba: bool
424 total_num_kv_heads: int
425 attn_backends: list[type[AttentionBackend]]
426 + dcp_size: int = 1
427 tensor_shape: torch.Size | None = None
428
429 def __post_init__(self):
@@ -461,6 +465,15 @@ class TransferTopology:
465 for key in [k for k in self._engines if k[0] == remote_engine_id]:
466 del self._engines[key]
467
468 + @property
469 + def dcp_rank(self) -> int:
470 + """This rank's DCP coverage rank.
471 +
472 + with ``dcp_size in (1, tp_size)`` enforced at the connector boundary, a
473 + rank's DCP identity is always exactly ``tp_rank % dcp_size``.
474 + """
475 + return self.tp_rank % self.dcp_size
476 +
477 # ============================================================
478 # Common methods
479 # ============================================================
@@ -514,18 +527,70 @@ class TransferTopology:
527 """Whether the local engine's KV cache is replicated."""
528 return self.is_mla or self.tp_size > self.total_num_kv_heads
529
517 - def handshake_target_ranks(self, remote_tp_size: int) -> list[int]:
530 + def dcp_source_ranks(self, remote_tp_size: int, remote_dcp_size: int) -> list[int]:
531 + """Remote ranks whose DCP slice overlaps mine (MLA, ``remote_dcp_size > 1``).
532 +
533 + Shared by ``handshake_target_ranks`` (who to query metadata from) and
534 + ``compute_tp_mapping`` (who to actually read from) — for MLA the two
535 + questions have the identical answer, since DCP sharding is the only
536 + thing keeping a remote rank from being interchangeable with any other.
537 + """
538 + local_dcp_size = self.dcp_size
539 + local_dcp_rank = self.dcp_rank
540 + if local_dcp_size <= remote_dcp_size:
541 + # Keep every remote rank whose slice sits inside mine. When
542 + # local_dcp_size == 1 (replicated locally), local_dcp_rank == 0 reduces to
543 + # every remote rank, since no single one holds the whole sequence.
544 + return [
545 + r for r in range(remote_tp_size) if r % local_dcp_size == local_dcp_rank
546 + ]
547 + # Local finer-grained: exactly one remote rank covers my whole slice
548 + return [local_dcp_rank % remote_dcp_size]
549 +
550 + def handshake_target_ranks(
551 + self, remote_tp_size: int, remote_dcp_size: int = 1
552 + ) -> list[int]:
553 """Pre-registration: compute which remote TP ranks to handshake with.
554
520 - Pure math based on local/remote TP sizes — does not require
521 - the remote engine to be registered yet.
555 + Pure math based on local/remote TP (and DCP, when the remote shards
556 + its KV cache) sizes — does not require the remote engine to be
557 + registered yet.
558 +
559 + DCP support is scoped to ``dcp_size in (1, tp_size)`` on each side
560 + and DCP sizes that divide one another: neither side ever has a
561 + partially-duplicated, partially-sharded KV cache. When the
562 + remote is not sharded (``remote_dcp_size == 1``) this reduces
563 + exactly to the DTP>=PTP case, since a sharded local side
564 + already has ``tp_size == dcp_size``.
565 """
566 + if remote_dcp_size > 1:
567 + return self.dcp_source_ranks(remote_tp_size, remote_dcp_size)
568 +
569 tp_ratio = self.tp_ratio(remote_tp_size)
570 if tp_ratio > 0:
571 return [self.tp_rank // tp_ratio]
572 abs_ratio = -tp_ratio
573 return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
574
575 + def dcp_consumer_count(self, remote_tp_size: int, remote_dcp_size: int) -> int:
576 + """How many local ranks (in aggregate) read from a given remote rank.
577 +
578 + Used by the producer side to know how many reader notifications to
579 + wait for before freeing a request's blocks. Reuses ``tp_ratio``
580 + whenever the remote isn't sharded — a sharded local side already
581 + has ``tp_size == dcp_size``, so the existing TP-ratio formula is
582 + already correct there unmodified.
583 + """
584 + if remote_dcp_size > 1:
585 + if self.dcp_size == 1:
586 + # Replicated locally: every local rank reads every shard.
587 + return self.tp_size
588 + # Both sharded, different degrees.
589 + return max(1, self.dcp_size // remote_dcp_size)
590 + # Remote replicated: `tp_ratio` local ranks share each remote rank
591 + # when local_tp >= remote_tp, else each remote rank has one reader.
592 + return max(1, self.tp_ratio(remote_tp_size))
593 +
594 def target_remote_ranks(
595 self, remote_engine_id: EngineId, remote_pp_rank: int = 0
596 ) -> list[int]:
@@ -551,6 +616,8 @@ class TransferTopology:
616 f"local_tp={self.tp_size}, "
617 f"remote_tp={info.remote_tp_size}, "
618 f"remote_pp={remote_pp_rank}, "
619 + f"local_dcp={self.dcp_size}, "
620 + f"remote_dcp={info.remote_dcp_size}, "
621 f"local_rank={self.tp_rank}, "
622 f"remote_block_len={info.remote_block_len})"
623 )
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+7 -2
@@ -109,7 +109,9 @@ class NixlBaseConnectorScheduler:
109 # Requests that need to start recv/send.
110 # New requests are added by update_state_after_alloc in
111 # the scheduler. Used to make metadata passed to Worker.
112 - self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
112 + self._reqs_need_recv: dict[
113 + ReqId, tuple[Request, BlockIds, tuple[int, ...]]
114 + ] = {}
115 self._reqs_need_save: dict[ReqId, Request] = {}
116 # Reqs to send and their expiration time
117 self._reqs_need_send: dict[ReqId, float] = {}
@@ -207,6 +209,7 @@ class NixlBaseConnectorScheduler:
209 host = params.get("remote_host")
210 port = params.get("remote_port")
211 tp_size = params.get("tp_size")
212 + dcp_size = params.get("dcp_size", 1)
213 pp_size = params.get("pp_size", 1)
214 if (
215 remote_engine_id is None
@@ -222,6 +225,7 @@ class NixlBaseConnectorScheduler:
225 host=host,
226 port=port,
227 tp_size=tp_size,
228 + dcp_size=dcp_size,
229 pp_size=pp_size,
230 )
231 self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)
@@ -452,12 +456,13 @@ class NixlBaseConnectorScheduler:
456 meta = NixlConnectorMetadata()
457
458 # Loop through scheduled reqs and convert to ReqMeta.
455 - for req_id, (req, block_ids) in self._reqs_need_recv.items():
459 + for req_id, (req, block_ids, cached) in self._reqs_need_recv.items():
460 assert req.kv_transfer_params is not None
461 meta.add_new_req_to_recv(
462 request_id=req_id,
463 local_block_ids=block_ids,
464 kv_transfer_params=req.kv_transfer_params,
465 + local_num_computed_blocks=cached,
466 )
467
468 if self.use_host_buffer:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+111 -15
@@ -231,6 +231,17 @@ class NixlBaseConnectorWorker:
231 handle.append((addr + p_idx * chunk, chunk, dev))
232 yield handle
233
234 + def _needs_split_local_xfer_handles(self, tp_ratio: int, plan: TPMapping) -> bool:
235 + """Whether reads need per-source slices of the local KV region.
236 +
237 + Pure MLA attention is replicated across TP ranks and writes the whole
238 + local region. Multiple physical remote workers may still participate
239 + because DCP assigns them disjoint blocks, but that does not require
240 + splitting the local region. Hybrid MLA+SSM is different: its mapping
241 + contains multiple source ranks for the sharded SSM state.
242 + """
243 + return tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1)
244 +
245 def _fa_desc_replicated(self, num_fa_descs: int) -> list[bool]:
246 """Per-FA-descriptor replicate flag, in _build_fa_local emission order
247 (region-major; one desc per block, with K/V packed). Length ``num_fa_descs``.
@@ -380,6 +391,13 @@ class NixlBaseConnectorWorker:
391 self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
392 self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
393
394 + # DCP support is scoped to MLA, with dcp_size in (1, tp_size): either fully
395 + # replicated or fully sharded. A DCP rank is always derivable this way.
396 + self.dcp_rank = self.tp_rank % self.dcp_size
397 + if self._has_mamba and self.dcp_size > 1:
398 + # Prefix-cache-aware DCP slicing isn't implemented for the Mamba group.
399 + raise ValueError("DCP is not supported for hybrid MLA+Mamba models.")
400 +
401 self.num_blocks = kv_cache_config.num_blocks
402 self.enable_permute_local_kv = False
403 self.enable_heterogeneous_attn_post_process = False
@@ -543,9 +561,10 @@ class NixlBaseConnectorWorker:
561 self.compat_hash: str | None = None
562 self.transfer_topo: TransferTopology | None = None
563
546 - # With heterogeneous TP, P must wait for all assigned D TP workers to
547 - # finish reading before safely freeing the blocks.
564 + # With heterogeneous TP (or DCP), P must wait for all assigned D
565 + # workers to finish reading before safely freeing the blocks.
566 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
567 + self.expected_consumer_notifications_by_req: dict[ReqId, int] = {}
568 self.xfer_stats = NixlKVConnectorStats()
569
570 self._physical_blocks_per_logical_kv_block = 1
@@ -556,7 +575,6 @@ class NixlBaseConnectorWorker:
575 get_representative_spec_type(g.kv_cache_spec)
576 for g in self.kv_cache_config.transfer_groups
577 )
559 -
578 # Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE
579 # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
580 # combining both (e.g. GQA main + MLA Eagle-3 draft).
@@ -620,6 +638,7 @@ class NixlBaseConnectorWorker:
638 port: int,
639 remote_tp_size: int,
640 expected_engine_id: str,
641 + remote_dcp_size: int = 1,
642 remote_pp_size: int = 1,
643 notif_agents_only: bool = False,
644 ) -> tuple[dict[tuple[int, int], str], float]:
@@ -643,7 +662,9 @@ class NixlBaseConnectorWorker:
662 # local rank will read from. Note that With homogeneous TP,
663 # this happens to be the same single rank_i.
664 assert self.transfer_topo is not None
646 - p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)
665 + p_remote_ranks = self.transfer_topo.handshake_target_ranks(
666 + remote_tp_size, remote_dcp_size
667 + )
668 remote_rank_to_agent_name: dict[tuple[int, int], str] = {}
669 path = make_zmq_path("tcp", host, port)
670 # Clock offset to the peer, estimated from the handshake round-trip.
@@ -747,11 +768,11 @@ class NixlBaseConnectorWorker:
768 # Register Remote agent.
769 if notif_agents_only:
770 remote_agent_name = self._add_notif_only_remote_agent(
750 - metadata, remote_tp_size
771 + metadata, remote_tp_size, metadata.dcp_size
772 )
773 else:
774 remote_agent_name = self.add_remote_agent(
754 - metadata, remote_rank, remote_tp_size
775 + metadata, remote_rank, remote_tp_size, metadata.dcp_size
776 )
777 setup_agent_time = time.perf_counter()
778 logger.debug(
@@ -766,7 +787,7 @@ class NixlBaseConnectorWorker:
787 return remote_rank_to_agent_name, best_offset
788
789 def _add_notif_only_remote_agent(
769 - self, metadata: NixlAgentMetadata, remote_tp_size: int
790 + self, metadata: NixlAgentMetadata, remote_tp_size: int, remote_dcp_size: int = 1
791 ) -> str:
792 """Load a remote agent for notifs only on the push-mode decode side.
793
@@ -782,6 +803,7 @@ class NixlBaseConnectorWorker:
803 remote_physical_blocks_per_logical=(
804 metadata.physical_blocks_per_logical_kv_block
805 ),
806 + remote_dcp_size=remote_dcp_size,
807 ),
808 )
809 return self.nixl_wrapper.add_remote_agent(metadata.agent_metadata)
@@ -898,6 +920,7 @@ class NixlBaseConnectorWorker:
920 host: str,
921 port: int,
922 tp_size: int,
923 + dcp_size: int = 1,
924 pp_size: int = 1,
925 notif_agents_only: bool = False,
926 ) -> Future[tuple[dict[tuple[int, int], str], float]] | None:
@@ -923,6 +946,7 @@ class NixlBaseConnectorWorker:
946 port,
947 tp_size,
948 engine_id,
949 + dcp_size,
950 pp_size,
951 notif_agents_only,
952 )
@@ -960,6 +984,8 @@ class NixlBaseConnectorWorker:
984 meta.remote.host,
985 meta.remote.port,
986 meta.tp_size,
987 + meta.dcp_size,
988 + meta.pp_size,
989 )
990 if fut is None:
991 # Already handshaked — only happens if caller does not pre-check.
@@ -991,8 +1017,11 @@ class NixlBaseConnectorWorker:
1017 block_size=self.block_size,
1018 engine_id=self.engine_id,
1019 is_mla=self.use_mla,
994 - total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
1020 + total_num_kv_heads=1
1021 + if self.use_mla
1022 + else self.model_config.get_total_num_kv_heads(),
1023 attn_backends=self.attn_backends,
1024 + dcp_size=self.dcp_size,
1025 # SSM States come in tuples (ssm, conv)
1026 tensor_shape=next(iter(kv_caches.values())).shape
1027 if not self._has_mamba
@@ -1469,6 +1498,7 @@ class NixlBaseConnectorWorker:
1498 nixl_agent_meta: NixlAgentMetadata,
1499 remote_tp_rank: int = 0,
1500 remote_tp_size: int = 1,
1501 + remote_dcp_size: int = 1,
1502 ) -> str:
1503 """
1504 Add the remote NIXL agent and prepare the descriptors for reading cache
@@ -1553,6 +1583,7 @@ class NixlBaseConnectorWorker:
1583 remote_block_size=nixl_agent_meta.block_size,
1584 remote_block_len=nixl_agent_meta.block_lens[0],
1585 remote_physical_blocks_per_logical=physical_blocks_per_logical,
1586 + remote_dcp_size=remote_dcp_size,
1587 )
1588 transfer_topo.register_remote_engine(engine_id, transfer_info)
1589 logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
@@ -1561,6 +1592,7 @@ class NixlBaseConnectorWorker:
1592 transfer_topology=transfer_topo,
1593 remote_tp_size=remote_tp_size,
1594 group_spec_types=self._group_spec_types,
1595 + remote_dcp_size=remote_dcp_size,
1596 )
1597
1598 remote_agent_name = self.nixl_wrapper.add_remote_agent(
@@ -1583,7 +1615,9 @@ class NixlBaseConnectorWorker:
1615 self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
1616 nixl_agent_meta.kv_caches_base_addr
1617 )
1586 - self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size)
1618 + self._validate_remote_agent_handshake(
1619 + nixl_agent_meta, remote_tp_size, remote_dcp_size
1620 + )
1621
1622 # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
1623 # this is the ratio between the two sizes.
@@ -1613,10 +1647,8 @@ class NixlBaseConnectorWorker:
1647
1648 ### (Optional) Register local agent memory regions. MLA is not split.
1649 split_key = (tp_ratio, remote_block_size)
1616 - if (
1617 - tp_ratio < 0
1618 - and (not self.use_mla or len(plan.all_source_ranks) > 1)
1619 - and split_key not in self.src_xfer_handles_by_tp_ratio
1650 + if self._needs_split_local_xfer_handles(tp_ratio, plan) and (
1651 + split_key not in self.src_xfer_handles_by_tp_ratio
1652 ):
1653 # Remote tp_size > local tp_size: read from multiple remote ranks.
1654 # Logically "split" own regions into per-source chunks. Hybrid
@@ -1674,7 +1706,10 @@ class NixlBaseConnectorWorker:
1706 return remote_agent_name
1707
1708 def _validate_remote_agent_handshake(
1677 - self, nixl_agent_meta: NixlAgentMetadata, remote_tp_size: int
1709 + self,
1710 + nixl_agent_meta: NixlAgentMetadata,
1711 + remote_tp_size: int,
1712 + remote_dcp_size: int = 1,
1713 ):
1714 """
1715 Validate the remote agent handshake metadata ensuring the
@@ -1685,6 +1720,15 @@ class NixlBaseConnectorWorker:
1720 assert self.transfer_topo is not None
1721 remote_info = self.transfer_topo.get_engine_info(remote_engine_id)
1722 assert remote_info.remote_tp_size == remote_tp_size
1723 + assert remote_info.remote_dcp_size == remote_dcp_size
1724 + # DCP sizes must divide one another; this is what keeps the
1725 + # read-slicing math in pull_worker a closed form.
1726 + assert (
1727 + self.dcp_size % remote_dcp_size == 0 or remote_dcp_size % self.dcp_size == 0
1728 + ), (
1729 + f"DCP sizes must divide one another: local={self.dcp_size}, "
1730 + f"remote={remote_dcp_size} (engine {remote_engine_id})."
1731 + )
1732
1733 tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
1734 block_size_ratio = self.transfer_topo.block_size_ratio(
@@ -2132,6 +2176,7 @@ class NixlBaseConnectorWorker:
2176 if now < expires:
2177 break
2178 count = self.consumer_notification_counts_by_req.pop(req_id, 0)
2179 + self.expected_consumer_notifications_by_req.pop(req_id, None)
2180 self.xfer_stats.record_kv_expired_req()
2181 logger.warning(
2182 "Releasing expired KV blocks for request %s which were "
@@ -2267,6 +2312,7 @@ class NixlBaseConnectorWorker:
2312 hb_info.host,
2313 hb_info.port,
2314 hb_info.tp_size,
2315 + hb_info.dcp_size,
2316 hb_info.pp_size,
2317 self._hb_handshake_notif_only and hb_info.pp_size > 1,
2318 )
@@ -2352,7 +2398,7 @@ class NixlBaseConnectorWorker:
2398 This is required when the logical block size (the one set by the user)
2399 does not match the one required by the attn backend.
2400 `ratio` is the number of physical blocks per logical block.
2355 - We always receive logical blocks from the engine, so we expand them here eg:
2401 + We always receive logical blocks from the engine, so we expand them here eg:
2402 logical block ids: [(SW-clipped) [1], (FA) [2, 3]], ratio=2
2403 physical block ids: [(SW-clipped) [2, 3], (FA) [4, 5, 6, 7]]
2404 """
@@ -2377,6 +2423,56 @@ class NixlBaseConnectorWorker:
2423 )
2424 return physical_block_ids
2425
2426 + def _apply_dcp_prefix_caching(
2427 + self,
2428 + local_ids: list[int],
2429 + remote_ids: list[int],
2430 + remote_rank: int,
2431 + local_dcp_size: int,
2432 + local_dcp_rank: int,
2433 + remote_dcp_size: int,
2434 + local_num_computed_blocks: int,
2435 + ) -> tuple[list[int], list[int]]:
2436 + """Match local and remote block IDs by global DCP position.
2437 +
2438 + ``local_ids`` excludes blocks already satisfied by the local prefix
2439 + cache, while ``remote_ids`` contains the full transferable remote list.
2440 +
2441 + Example:
2442 + local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6]
2443 + cached_blocks=[0,2,4]
2444 + to transfer: [6] (the rest is cached)
2445 +
2446 + remote DCP size 4
2447 + remote rank 0 owns [0, 4] -> no more match, skip the read
2448 + remote rank 2 owns [2, 6] -> match position 6 and read
2449 +
2450 + An empty result (rank 0 above) takes the notification-only path, allowing the
2451 + remote to release its blocks without performing a transfer.
2452 + """
2453 + local_size, remote_size = local_dcp_size, remote_dcp_size
2454 +
2455 + if local_size == remote_size:
2456 + local_slice = local_ids
2457 + remote_slice = remote_ids[local_num_computed_blocks:]
2458 + elif local_size < remote_size:
2459 + k = remote_size // local_size
2460 + p = (remote_rank - local_dcp_rank) // local_size
2461 + start_local = (p - local_num_computed_blocks) % k
2462 + start_remote = (local_num_computed_blocks + start_local - p) // k
2463 + local_slice = local_ids[start_local::k]
2464 + remote_slice = remote_ids[start_remote:]
2465 + else:
2466 + k = local_size // remote_size
2467 + remote_dcp_rank = remote_rank % remote_size
2468 + c = (local_dcp_rank - remote_dcp_rank) // remote_size
2469 + start_remote = c + local_num_computed_blocks * k
2470 + local_slice = local_ids
2471 + remote_slice = remote_ids[start_remote::k]
2472 +
2473 + matched_blocks = min(len(local_slice), len(remote_slice))
2474 + return local_slice[:matched_blocks], remote_slice[:matched_blocks]
2475 +
2476 def _apply_prefix_caching(
2477 self,
2478 decode_block_ids: BlockIds,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+4
@@ -369,6 +369,10 @@ class NixlPushConnector(NixlBaseConnector):
369 kv_cache_config: "KVCacheConfig",
370 ):
371 super().__init__(vllm_config, role, kv_cache_config)
372 + if vllm_config.parallel_config.decode_context_parallel_size > 1:
373 + raise ValueError(
374 + "NixlPushConnector does not support decode_context_parallel_size > 1."
375 + )
376 self.connector_scheduler: NixlPushConnectorScheduler | None = None
377 self.connector_worker: NixlPushConnectorWorker | None = None
378 if role == KVConnectorRole.SCHEDULER:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+15 -2
@@ -209,6 +209,7 @@ class HeartbeatInfo:
209 host: str
210 port: int
211 tp_size: int
212 + dcp_size: int = 1
213 pp_size: int = 1
214
215
@@ -228,6 +229,12 @@ class ReqMeta:
229 # To be used when logical block size does not match the kernel block size
230 local_physical_block_ids: BlockIds
231 tp_size: int
232 + dcp_size: int = 1
233 + # Per-KV-cache-group logical blocks this rank already holds, i.e. its
234 + # prefix-cache hit. Fixes where this rank's DCP slice starts relative to
235 + # the remote's; kept per-group since hybrid models (e.g. SWA+FA) can have
236 + # different cache-hit counts per group.
237 + local_num_computed_blocks: tuple[int, ...] = ()
238 remote: RemoteMeta | None = None
239 # Remote block size, discovered during NIXL handshake (push mode).
240 remote_block_size: int | None = None
@@ -261,14 +268,17 @@ class NixlConnectorMetadata(KVConnectorMetadata):
268 self,
269 local_block_ids: BlockIds,
270 kv_transfer_params: dict[str, Any],
271 + local_num_computed_blocks: tuple[int, ...] = (),
272 ) -> ReqMeta:
273 return ReqMeta(
274 local_block_ids=local_block_ids,
275 local_physical_block_ids=local_block_ids,
268 - # P workers don't need to receive tp_size from proxy here.
276 + # P workers don't need to receive these from proxy here.
277 tp_size=kv_transfer_params.get("tp_size", 1),
278 + dcp_size=kv_transfer_params.get("dcp_size", 1),
279 remote_block_size=kv_transfer_params.get("remote_block_size"),
280 pp_size=kv_transfer_params.get("pp_size", 1),
281 + local_num_computed_blocks=local_num_computed_blocks,
282 )
283
284 def add_new_req_to_save(
@@ -286,8 +296,11 @@ class NixlConnectorMetadata(KVConnectorMetadata):
296 request_id: ReqId,
297 local_block_ids: BlockIds,
298 kv_transfer_params: dict[str, Any],
299 + local_num_computed_blocks: tuple[int, ...] = (),
300 ):
290 - req = self._add_new_req(local_block_ids, kv_transfer_params)
301 + req = self._add_new_req(
302 + local_block_ids, kv_transfer_params, local_num_computed_blocks
303 + )
304 req.remote = RemoteMeta(
305 block_ids=kv_transfer_params["remote_block_ids"],
306 engine_id=kv_transfer_params["remote_engine_id"],
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
+14 -1
@@ -155,12 +155,23 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
155 local_block_ids = self.get_exchange_clipped_blocks(
156 unhashed_local_block_ids
157 )
158 + # Blocks covered by the local prefix cache, per KV cache group.
159 + # Each count fixes where that group's DCP slice starts, which the
160 + # worker needs to line up with the remote's slice.
161 + local_num_computed_blocks = tuple(
162 + sum(
163 + block.block_hash is not None and not block.is_null
164 + for block in group
165 + )
166 + for group in blocks.blocks
167 + )
168
169 # Get unhashed blocks to pull from remote. Mind that a full prefix
170 # cache hit is indicated with an empty list.
171 self._reqs_need_recv[request.request_id] = (
172 request,
173 local_block_ids,
174 + local_num_computed_blocks,
175 )
176
177 else:
@@ -213,7 +224,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
224 # To avoid stranding the prefill blocks in the prefill instance,
225 # we must add empty block_ids to _reqs_need_recv so that our
226 # worker side will notify and free blocks in the prefill instance.
216 - self._reqs_need_recv[request.request_id] = (request, [])
227 + self._reqs_need_recv[request.request_id] = (request, [], ())
228 params["do_remote_prefill"] = False
229 return False, None
230
@@ -272,6 +283,8 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
283 remote_host=self.side_channel_host,
284 remote_port=self.side_channel_port,
285 tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
286 + dcp_size=self.vllm_config.parallel_config.decode_context_parallel_size,
287 + pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,
288 remote_num_tokens=remote_num_tokens,
289 remote_blocks_expiry_time=blocks_expiry_time,
290 transfer_mode=self._TRANSFER_MODE,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+94 -46
@@ -5,6 +5,7 @@
5 import time
6 from typing import TYPE_CHECKING
7
8 +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
9 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
10 NixlBaseConnectorWorker,
11 )
@@ -143,40 +144,78 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
144 self._handle_failed_transfer(req_id, None)
145 return
146
147 + if any(len(group) > 0 for group in meta.local_block_ids):
148 + # The scheduler waits for finished_recving from *every* worker.
149 + # Under DCP a rank's slice can legitimately come out empty when its
150 + # interleaved positions fall past the end of the sequence. _read_blocks
151 + # then takes the notify-only path without registering a transfer.
152 + # Seed the entry so this rank still reports completion.
153 + self._recving_transfers.setdefault(req_id, [])
154 +
155 plan = self.tp_mappings[engine_id]
156 remote_info = self.transfer_topo.get_engine_info(engine_id)
157 tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
158
159 + remote_logical_block_ids = meta.remote.block_ids
160 meta.remote.block_ids = self._logical_to_kernel_block_ids(
151 - meta.remote.block_ids,
161 + remote_logical_block_ids,
162 remote_info.remote_physical_blocks_per_logical,
163 )
154 - remote_block_ids = meta.remote.block_ids
155 - local_block_ids = meta.local_physical_block_ids
156 - num_groups = len(local_block_ids)
157 - read_specs = [
158 - ReadSpec(
159 - remote_rank=rank,
160 - local_block_ids=[
161 - list(local_block_ids[g])
162 - if rank in plan.source_ranks_per_group[g]
163 - else []
164 - for g in range(num_groups)
165 - ],
166 - remote_block_ids=[
167 - list(remote_block_ids[g])
168 - if rank in plan.source_ranks_per_group[g]
169 - else []
170 - for g in range(num_groups)
171 - ],
164 + num_groups = len(meta.local_block_ids)
165 + dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
166 +
167 + def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
168 + return [
169 + list(block_ids[g]) if rank in plan.source_ranks_per_group[g] else []
170 + for g in range(num_groups)
171 + ]
172 +
173 + read_specs = []
174 + for rank in plan.all_source_ranks:
175 + if dcp_active:
176 + # DCP interleaves at block granularity, so slicing happens
177 + # here on logical blocks, before kernel-block expansion.
178 + local_ids = group_ids(meta.local_block_ids, rank)
179 + remote_ids = group_ids(remote_logical_block_ids, rank)
180 + for g in range(num_groups):
181 + if not local_ids[g]:
182 + continue
183 + # Prefix cache hit may lead to skip some of the remote reads
184 + # TODO (NickLucche) consider unifying prefix cache handling on
185 + # logical blocks here for both dcp and non-dcp
186 + local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
187 + local_ids[g],
188 + remote_ids[g],
189 + remote_rank=rank,
190 + local_dcp_size=self.dcp_size,
191 + local_dcp_rank=self.dcp_rank,
192 + remote_dcp_size=remote_info.remote_dcp_size,
193 + local_num_computed_blocks=meta.local_num_computed_blocks[g],
194 + )
195 + local_physical_ids = self._logical_to_kernel_block_ids(
196 + local_ids, self._physical_blocks_per_logical_kv_block
197 + )
198 + remote_physical_ids = self._logical_to_kernel_block_ids(
199 + remote_ids, remote_info.remote_physical_blocks_per_logical
200 + )
201 + else:
202 + # No DCP realignment needed: reuse the already-expanded full
203 + # physical lists instead of re-deriving them from logical ids.
204 + local_physical_ids = group_ids(meta.local_physical_block_ids, rank)
205 + remote_physical_ids = group_ids(meta.remote.block_ids, rank)
206 + read_specs.append(
207 + ReadSpec(
208 + remote_rank=rank,
209 + local_block_ids=local_physical_ids,
210 + remote_block_ids=remote_physical_ids,
211 + )
212 )
173 - for rank in plan.all_source_ranks
174 - ]
213
214 # D may have to perform multiple reads from different remote ranks.
215 # Pure MLA reads once because its cache is replicated. Hybrid
178 - # MLA+SSM still needs one read per SSM source rank.
179 - if self.use_mla and tp_ratio < 0 and not self._has_mamba:
216 + # MLA+SSM still needs one read per SSM source rank. With DCP, pure
217 + # MLA may also read from multiple ranks (disjoint token slices).
218 + if self.use_mla and tp_ratio < 0 and not self._has_mamba and not dcp_active:
219 assert len(read_specs) == 1
220
221 for i, spec in enumerate(read_specs):
@@ -214,12 +253,15 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
253 remote_request_id=meta.remote.request_id,
254 local_xfer_side_handle=local_xfer_side_handle,
255 remote_xfer_side_handle=remote_xfer_side_handle,
256 + expected_consumers=plan.local_consumers,
257 )
258
259 if self.use_mla and tp_ratio < 0 and len(read_specs) == 1:
260 # ..but we still need to notify the other remote ranks that we
261 # have the blocks we need so they can update the request state.
222 - notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
262 + # Same thing for DCP (tp_size == dcp_size), so the raw tp_ratio already
263 + # reflects whether any remote replica is left unchosen.
264 + notif_id = f"{meta.remote.request_id}:{plan.local_consumers}".encode()
265 remote_agents = self._remote_agents[meta.remote.engine_id]
266 for rank_to_notify, agent in remote_agents.items():
267 if rank_to_notify != (0, read_specs[0].remote_rank):
@@ -233,6 +275,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
275 remote_request_id: str,
276 local_xfer_side_handle: int,
277 remote_xfer_side_handle: int,
278 + expected_consumers: int,
279 ):
280 """
281 Post a READ point-to-point xfer request from a single local worker to
@@ -262,13 +305,13 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
305 # NOTE(rob): according to nvidia the staging blocks are used to
306 # saturate IB with heterogeneous TP sizes.
307
265 - # Number of D TP workers that will read from dst P. Propagate info
266 - # on notification so that dst worker can wait before freeing blocks.
267 - notif_id = f"{remote_request_id}:{self.world_size}".encode()
308 + # Number of local workers that will notify this producer worker.
309 + # Propagate on notification so dst worker can wait before freeing.
310 + notif_id = f"{remote_request_id}:{expected_consumers}".encode()
311
312 # Full prefix cache hit: do not need to read remote blocks,
313 # just notify P worker that we have the blocks we need.
271 - if len(local_block_ids) == 0:
314 + if not any(len(group) > 0 for group in local_block_ids):
315 # A full prefix cache hit is indicated with an empty list.
316 agent_name = self._remote_agents[dst_engine_id][(0, remote_rank)]
317 try:
@@ -292,12 +335,19 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
335 == len(local_block_ids)
336 == len(self.kv_cache_config.transfer_groups)
337 )
295 - local_block_ids, remote_block_ids = self._apply_prefix_caching(
296 - decode_block_ids=local_block_ids,
297 - prefill_block_ids=remote_block_ids,
298 - decode_physical_per_logical=self._physical_blocks_per_logical_kv_block,
299 - prefill_physical_per_logical=remote_info.remote_physical_blocks_per_logical,
300 - )
338 + if not (self.dcp_size > 1 or remote_info.remote_dcp_size > 1):
339 + # DCP-active reads were already trimmed (and DCP-realigned, when
340 + # sizes mismatch) in _read_blocks_for_req, at logical granularity.
341 + local_block_ids, remote_block_ids = self._apply_prefix_caching(
342 + decode_block_ids=local_block_ids,
343 + prefill_block_ids=remote_block_ids,
344 + decode_physical_per_logical=(
345 + self._physical_blocks_per_logical_kv_block
346 + ),
347 + prefill_physical_per_logical=(
348 + remote_info.remote_physical_blocks_per_logical
349 + ),
350 + )
351
352 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
353 # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
@@ -351,8 +401,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
401 def _get_new_notifs(self) -> set[str]:
402 """
403 Get req_ids which got a remote xfer message. When multiple consumers
354 - are reading from the same producer (heterogeneous TP scenario), wait
355 - for all consumers to be done pulling.
404 + are reading from the same producer (heterogeneous TP or DCP
405 + scenario), wait for all consumers to be done pulling.
406
407 Also handles heartbeat notifications ("HB:req1,req2,...") by
408 extending the lease on the referenced requests.
@@ -368,7 +418,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
418 self._handle_heartbeat(msg[3:])
419 continue
420
371 - req_id, tp_size = msg.rsplit(":", 1)
421 + req_id, expected_consumers = msg.rsplit(":", 1)
422 if (
423 req_id not in self._reqs_to_send
424 and req_id not in self._reqs_to_process
@@ -381,24 +431,22 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
431 )
432 continue
433
384 - # NOTE: `tp_ratio` is the opposite when swapping local<>remote
385 - n_consumers = int(tp_size)
386 - tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
387 -
388 - # Number of reads *per producer* to wait for.
389 - # When remote D TP > local P TP we expect `tp_ratio` reads.
390 - consumers_per_producer = (
391 - -tp_ratio if n_consumers > self.world_size else 1
434 + # Every reader of this req_id reports the same count (it's
435 + # derived from aggregate topology, not the specific rank),
436 + # so repeated notifications never disagree on it.
437 + self.expected_consumer_notifications_by_req[req_id] = int(
438 + expected_consumers
439 )
440
441 self.consumer_notification_counts_by_req[req_id] += 1
442 # Wait all consumers (D) to be done reading before freeing.
443 if (
444 self.consumer_notification_counts_by_req[req_id]
398 - == consumers_per_producer
445 + == self.expected_consumer_notifications_by_req[req_id]
446 ):
447 notified_req_ids.add(req_id)
448 del self.consumer_notification_counts_by_req[req_id]
449 + del self.expected_consumer_notifications_by_req[req_id]
450 self._reqs_to_process.remove(req_id)
451 self._reqs_to_send.pop(req_id, None)
452 return notified_req_ids
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
+2 -2
@@ -198,7 +198,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
198 # ReqMeta without a KeyError — the actual remote block IDs are
199 # learned by P over the NIXL handshake at WRITE time.
200 params["remote_block_ids"] = ()
201 - self._reqs_need_recv[request.request_id] = (request, local_block_ids)
201 + self._reqs_need_recv[request.request_id] = (request, local_block_ids, ())
202
203 # Mark as processed so a re-entry (e.g. preemption + reschedule)
204 # doesn't re-stage the registration.
@@ -240,7 +240,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
240 # recv so the worker emits a notif that lets P free them.
241 # Seed remote_block_ids so add_new_req_to_recv won't KeyError.
242 params["remote_block_ids"] = ()
243 - self._reqs_need_recv[request.request_id] = (request, [])
243 + self._reqs_need_recv[request.request_id] = (request, [], ())
244 params["do_remote_prefill"] = False
245 return False, None
246
vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+23 -5
@@ -56,6 +56,10 @@ class TPMapping:
56 # FA head offset factor for hetero-TP (D_TP > P_TP).
57 rank_offset_factor: int
58
59 + # Local ranks (in aggregate) that read from a given source rank. The producer frees
60 + # a request's blocks only once that many notifications have come in.
61 + local_consumers: int = 1
62 +
63
64 # ======================================================================
65 # TP mapping computation
@@ -66,22 +70,31 @@ def compute_tp_mapping(
70 transfer_topology: TransferTopology,
71 remote_tp_size: int,
72 group_spec_types: tuple[type[KVCacheSpec], ...],
73 + remote_dcp_size: int = 1,
74 ) -> TPMapping:
75 """Build the complete local-to-remote TP mapping.
76
77 Computes source ranks, head slot assignments, and the rank offset
78 factor in a single pass.
79 +
80 + DCP support is scoped to MLA only, with a side is either fully replicated or fully
81 + sharded. DCP-branch reuses the same rank set used at handshake selection.
82 """
83 tp_rank = transfer_topology.tp_rank
84 tp_size = transfer_topology.tp_size
85 total_num_kv_heads = transfer_topology.total_num_kv_heads
86 # --- Attention source ranks ---
87 if transfer_topology.is_mla or tp_size >= remote_tp_size:
80 - # D (local TP) > P (remote TP): multiple local ranks read different chunks from
81 - # *one* remote rank, corresponding to different kv heads.
82 - # For MLA, we only need one remote since cache is duplicated. When P TP=k*TP k,
83 - # this will spread mla ranks to read from remote k*tp_rank.
84 - attn_ranks = [tp_rank * remote_tp_size // tp_size]
88 + if transfer_topology.is_mla and remote_dcp_size > 1:
89 + attn_ranks = transfer_topology.dcp_source_ranks(
90 + remote_tp_size, remote_dcp_size
91 + )
92 + else:
93 + # D (local TP) > P (remote TP): multiple local ranks read different chunks
94 + # from *one* remote rank, corresponding to different kv heads.
95 + # For MLA, we only need one remote since cache is duplicated. When
96 + # P TP=k*TP k, this will spread mla ranks to read from remote k*tp_rank.
97 + attn_ranks = [tp_rank * remote_tp_size // tp_size]
98 else:
99 # P (remote TP) > D (local TP): one local rank
100 # reads from multiple remote ranks.
@@ -134,9 +147,14 @@ def compute_tp_mapping(
147 # D TP > P TP: we index into remote to read different heads depending on rank.
148 rank_offset_factor = tp_rank % (tp_size // remote_tp_size)
149
150 + local_consumers = transfer_topology.dcp_consumer_count(
151 + remote_tp_size, remote_dcp_size
152 + )
153 +
154 return TPMapping(
155 source_ranks_per_group=source_ranks_per_group,
156 all_source_ranks=tuple(all_ranks),
157 rank_to_attention_slot=rank_to_attention_slot,
158 rank_offset_factor=rank_offset_factor,
159 + local_consumers=local_consumers,
160 )
vllm/distributed/kv_transfer/kv_transfer_state.py
+2
@@ -85,6 +85,8 @@ 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.
89 + vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config)
90 _sync_engine_id_across_tp(vllm_config)
91
92 _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector(
vllm/model_executor/layers/attention/mla_attention.py
-3
@@ -3157,9 +3157,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
3157 parallel_config = get_current_vllm_config().parallel_config
3158 # Avoid requiring an initialized DCP group in tests.
3159 self.dcp_world_size: int = parallel_config.decode_context_parallel_size
3160 - self.cp_kv_cache_interleave_size: int = (
3161 - parallel_config.cp_kv_cache_interleave_size
3162 - )
3160
3161 @abstractmethod
3162 def forward_mqa(
vllm/model_executor/layers/sparse_attn_indexer.py
+17 -1
@@ -771,14 +771,30 @@ class SparseAttnIndexer(CustomOp):
771 parallel_config = get_current_vllm_config().parallel_config
772 self.dcp_world_size = parallel_config.decode_context_parallel_size
773 self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
774 - self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
774 self.use_pcp = parallel_config.prefill_context_parallel_size > 1
775 + self._cp_kv_cache_interleave_size: int | None = None
776 if current_platform.is_cuda() and not has_deep_gemm():
777 raise RuntimeError(
778 "Sparse Attention Indexer CUDA op requires DeepGEMM support in "
779 "the current vLLM environment."
780 )
781
782 + @property
783 + def cp_kv_cache_interleave_size(self) -> int:
784 + """With PD+DCP, the real value isn't known until block_size is finalized,
785 + which happens after this layer is built. Safe to cache after the first access,
786 + as long as the adjustment always runs before any forward pass
787 + (it's set up in Worker.initialize_from_config, ahead of warmup/serving).
788 + """
789 + if self._cp_kv_cache_interleave_size is None:
790 + value = (
791 + get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
792 + )
793 + if isinstance(get_forward_context().attn_metadata, dict):
794 + self._cp_kv_cache_interleave_size = value
795 + return value
796 + return self._cp_kv_cache_interleave_size
797 +
798 def forward_native(
799 self,
800 hidden_states: torch.Tensor,
vllm/v1/attention/backends/mla/flashinfer_mla.py
-1
@@ -324,7 +324,6 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
324 extra_kwargs: dict[str, Any] = {}
325 decode_backend: str | None
326 if self.dcp_world_size > 1:
327 - assert self.cp_kv_cache_interleave_size == 1
327 causal_seqlens_kv_global = attn_metadata.decode.dcp_tot_seq_lens
328 assert causal_seqlens_kv_global is not None
329 if not attn_metadata.causal and query_len > 1:
vllm/v1/worker/gpu/model_runner.py
+2
@@ -531,6 +531,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
531 is_profiling: bool = False,
532 kv_cache_allocation_context: AbstractContextManager | None = None,
533 ) -> None:
534 + # GPUWorker finalizes the PD interleave before KV cache initialization.
535 + self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
536 kv_cache_config = deepcopy(kv_cache_config)
537 self.kv_cache_config = kv_cache_config
538
vllm/v1/worker/gpu_model_runner.py
+8
@@ -749,6 +749,9 @@ class GPUModelRunner(
749 self._init_kernel_block_sizes = [placeholder_block_size]
750 self._init_max_num_blocks = [placeholder_max_num_blocks]
751 self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT]
752 + self.cp_kv_cache_interleave_size = (
753 + self.parallel_config.cp_kv_cache_interleave_size
754 + )
755 # Capture warmup providers registered by the initial placeholder InputBatch
756 with self.jit_warmup_registry.activate():
757 self.input_batch = InputBatch(
@@ -7350,11 +7353,16 @@ class GPUModelRunner(
7353 or kernel_block_sizes != self._init_kernel_block_sizes
7354 or max_num_blocks != self._init_max_num_blocks
7355 or slot_mapping_modes != self._init_slot_mapping_modes
7356 + or self.cp_kv_cache_interleave_size
7357 + != self.parallel_config.cp_kv_cache_interleave_size
7358 ):
7359 self._init_block_sizes = block_sizes
7360 self._init_kernel_block_sizes = kernel_block_sizes
7361 self._init_max_num_blocks = max_num_blocks
7362 self._init_slot_mapping_modes = slot_mapping_modes
7363 + self.cp_kv_cache_interleave_size = (
7364 + self.parallel_config.cp_kv_cache_interleave_size
7365 + )
7366 # Capture warmup providers registered after final KV-cache geometry is known
7367 with self.jit_warmup_registry.activate():
7368 self.input_batch = InputBatch(