[Bugfix][NIXL] Fix DCP pulls across MLA cache regions (#57389)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Lucas Wilkinson committed
Sep 21, 2026 at 11:57 UTC
8c0825090df8b7939a92ff22b02bc7d03cb3605a
3 files changed
+133
-28
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+77
@@ -58,6 +58,7 @@ def region_pull_worker():
58
worker._recving_transfers = {}
59
worker.use_mla, worker._has_mamba = True, False
60
worker.dcp_size = 1
61
+ worker.dcp_rank = 0
62
spec = MLAAttentionSpec(
63
block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16
64
)
@@ -165,6 +166,82 @@ def test_region_pull_ignores_allocation_padding(
166
]
167
168
169
+@pytest.mark.cpu_test
170
+@pytest.mark.parametrize("num_pages", [0, 3, 19])
171
+@pytest.mark.parametrize("region_groups", [(0, 1), (1, 0, 1)])
172
+def test_dcp_region_pull(region_pull_worker, num_pages, region_groups):
173
+ """Read each uncached page from its DCP rank; notify ranks with no pages."""
174
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
175
+ RemoteMeta,
176
+ ReqMeta,
177
+ )
178
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
179
+
180
+ worker = region_pull_worker
181
+ worker.region_group_ids = list(region_groups)
182
+ worker.num_regions = len(region_groups)
183
+ worker.dst_region_group_ids = {"P": [0] * worker.num_regions}
184
+ worker.dst_region_num_blocks = {
185
+ engine: [100] * worker.num_regions for engine in ("P", "D")
186
+ }
187
+ worker.block_len_per_layer = [1024] * worker.num_regions
188
+ worker.dcp_rank = 0
189
+ remote = worker.transfer_topo.get_engine_info.return_value
190
+ remote.remote_tp_size = remote.remote_dcp_size = 8
191
+ remote.remote_physical_blocks_per_logical = 1
192
+ ranks = tuple(range(8))
193
+ worker.tp_mappings["P"] = TPMapping((ranks, ranks), ranks, {r: r for r in ranks}, 8)
194
+ worker.dst_xfer_side_handles = {"P": {r: 1000 + r for r in ranks}}
195
+ worker._remote_agents = {"P": {(0, r): f"P-rank{r}" for r in ranks}}
196
+ # Two regions with different prefix hits, plus one padding page each.
197
+ local = (
198
+ [list(range(31, 31 + num_pages)), list(range(42, 41 + num_pages))]
199
+ if num_pages
200
+ else []
201
+ )
202
+ meta = ReqMeta(
203
+ local,
204
+ local,
205
+ tp_size=8,
206
+ local_num_computed_blocks=(0, 1, 2),
207
+ remote=RemoteMeta(
208
+ [[10, 11, 12]],
209
+ "localhost",
210
+ 1,
211
+ "P",
212
+ "request-P",
213
+ num_tokens=max(0, num_pages * 64 - 7),
214
+ ),
215
+ )
216
+ worker._read_blocks_for_req("request", meta)
217
+ reads = {
218
+ call.kwargs["remote_xfer_side_handle"] - 1000: call.kwargs
219
+ for call in worker._read_blocks_mixed.call_args_list
220
+ }
221
+ notified = {call.args[0] for call in worker.nixl_wrapper.send_notif.call_args_list}
222
+ for rank in ranks:
223
+ expected = [
224
+ (region * 100 + 30 + group * 10 + page, region * 100 + 10 + page // 8)
225
+ for region, group in enumerate(region_groups)
226
+ for page in range(group + 1, num_pages)
227
+ if page % 8 == rank
228
+ ]
229
+ if expected:
230
+ read = reads[rank]
231
+ assert (
232
+ list(zip(read["local_block_descs_ids"], read["remote_block_descs_ids"]))
233
+ == expected
234
+ )
235
+ assert f"P-rank{rank}" not in notified
236
+ else:
237
+ assert rank not in reads and f"P-rank{rank}" in notified
238
+ assert meta.region_blocks_to_zero == (
239
+ [[30 + group * 10 + num_pages] for group in region_groups]
240
+ if num_pages
241
+ else None
242
+ )
243
+
244
+
245
@pytest.mark.cpu_test
246
def test_region_pull_completion_zeros_only_own_padding(region_pull_worker):
247
"""Completion must preserve neighboring layers and already imported pages."""
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+36
-8
@@ -75,6 +75,7 @@ from vllm.distributed.parallel_state import (
75
from vllm.logger import init_logger
76
from vllm.platforms import current_platform
77
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
78
+from vllm.utils.math_utils import cdiv
79
from vllm.utils.network_utils import make_zmq_path
80
from vllm.utils.torch_utils import async_tensor_h2d
81
from vllm.v1.kv_cache_interface import (
@@ -2430,6 +2431,18 @@ class NixlBaseConnectorWorker:
2431
remote_physical_per_logical = (
2432
nixl_agent_meta.physical_blocks_per_logical_kv_block
2433
)
2434
+ if (
2435
+ (self.dcp_size > 1 or remote_dcp_size > 1)
2436
+ and self.region_group_ids
2437
+ and self.region_group_ids != self.dst_region_group_ids[remote_engine_id]
2438
+ and (
2439
+ self._physical_blocks_per_logical_kv_block != 1
2440
+ or remote_physical_per_logical != 1
2441
+ )
2442
+ ):
2443
+ raise NotImplementedError(
2444
+ "DCP region pulls require matching logical and physical block sizes"
2445
+ )
2446
if (
2447
self._has_mamba
2448
and remote_physical_per_logical
@@ -3234,11 +3247,14 @@ class NixlBaseConnectorWorker:
3247
local_dcp_rank: int,
3248
remote_dcp_size: int,
3249
local_num_computed_blocks: int,
3250
+ num_remote_blocks: int | None = None,
3251
) -> tuple[list[int], list[int]]:
3252
"""Match local and remote block IDs by global DCP position.
3253
3254
``local_ids`` excludes blocks already satisfied by the local prefix
3255
cache, while ``remote_ids`` contains the full transferable remote list.
3256
+ ``num_remote_blocks``, when provided, is the global block count before
3257
+ DCP partitioning and excludes allocation padding.
3258
3259
Example:
3260
local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6]
@@ -3254,6 +3270,9 @@ class NixlBaseConnectorWorker:
3270
3271
"""
3272
local_size, remote_size = local_dcp_size, remote_dcp_size
3273
+ if num_remote_blocks is not None:
3274
+ num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size)
3275
+ local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)]
3276
3277
if local_size == remote_size:
3278
local_slice = local_ids
@@ -3274,6 +3293,8 @@ class NixlBaseConnectorWorker:
3293
remote_slice = remote_ids[start_remote::k]
3294
3295
matched_blocks = min(len(local_slice), len(remote_slice))
3296
+ if num_remote_blocks is not None and matched_blocks < len(local_slice):
3297
+ raise ValueError("Remote KV pages do not cover the requested range")
3298
return local_slice[:matched_blocks], remote_slice[:matched_blocks]
3299
3300
@staticmethod
@@ -3290,13 +3311,15 @@ class NixlBaseConnectorWorker:
3311
block_ids_by_region.append(list(block_ids[group_id]))
3312
return block_ids_by_region
3313
3293
- @staticmethod
3314
def _apply_prefix_caching_by_region(
3315
+ self,
3316
decode_block_ids: BlockIds,
3317
prefill_block_ids: BlockIds,
3318
*,
3319
num_computed_blocks: list[int] | None = None,
3320
num_remote_blocks: int | None = None,
3321
+ remote_rank: int = 0,
3322
+ remote_dcp_size: int = 1,
3323
) -> tuple[BlockIds, BlockIds]:
3324
"""Pair an uncached decode suffix with the same prefill regions."""
3325
assert len(decode_block_ids) == len(prefill_block_ids)
@@ -3310,13 +3333,18 @@ class NixlBaseConnectorWorker:
3333
for decode_region, prefill_region, start in zip(
3334
decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True
3335
):
3313
- count = min(len(decode_region), max(num_remote_blocks - start, 0))
3314
- if start + count > len(prefill_region):
3315
- raise ValueError("Remote KV pages do not cover the requested range")
3316
- # An extra producer page can be allocation padding, not a
3317
- # cached prefix. Select by token position instead of list length.
3318
- matched_decode.append(list(decode_region[:count]))
3319
- matched_prefill.append(list(prefill_region[start : start + count]))
3336
+ local, remote = self._apply_dcp_prefix_caching(
3337
+ decode_region,
3338
+ prefill_region,
3339
+ remote_rank,
3340
+ self.dcp_size,
3341
+ self.dcp_rank,
3342
+ remote_dcp_size,
3343
+ start,
3344
+ num_remote_blocks=num_remote_blocks,
3345
+ )
3346
+ matched_decode.append(local)
3347
+ matched_prefill.append(remote)
3348
return matched_decode, matched_prefill
3349
3350
trimmed_prefill: list[list[int]] = []
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+20
-20
@@ -191,7 +191,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
191
"Different NIXL cache-group layouts are only supported for "
192
"pure MLA models"
193
)
194
- assert len(plan.all_source_ranks) == 1
194
if self.block_size != remote_info.remote_block_size:
195
raise NotImplementedError(
196
"Region-mapped NIXL transfers require matching physical block sizes"
@@ -213,7 +212,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
212
and meta.local_num_computed_blocks
213
and all(group >= 0 for group in local_region_groups)
214
and all(group >= 0 for group in remote_region_groups)
216
- and not dcp_active
215
):
216
transfer_groups = self.kv_cache_config.transfer_group_ids
217
num_computed_blocks = [
@@ -224,32 +222,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
222
num_remote_blocks = cdiv(
223
meta.remote.num_tokens, remote_info.remote_block_size
224
)
227
- elif (
228
- remote_info.remote_physical_blocks_per_logical
225
+ elif any(local_by_region) and (
226
+ dcp_active
227
+ or remote_info.remote_physical_blocks_per_logical
228
!= self._physical_blocks_per_logical_kv_block
229
):
230
raise NotImplementedError(
232
- "Region-mapped pulls with different logical block sizes require "
233
- "remote_num_tokens, per-group prefix counts, unshared regions "
234
- "and DCP=1"
231
+ "Region-mapped pulls with DCP or different logical block sizes "
232
+ "require remote_num_tokens, per-group prefix counts "
233
+ "and unshared regions"
234
)
236
- matched_local, matched_remote = self._apply_prefix_caching_by_region(
237
- local_by_region,
238
- remote_by_region,
239
- num_computed_blocks=num_computed_blocks,
240
- num_remote_blocks=num_remote_blocks,
241
- )
242
- meta.region_blocks_to_zero = [
243
- list(blocks[len(matched) :])
244
- for blocks, matched in zip(local_by_region, matched_local, strict=True)
245
- ]
235
read_specs = [
236
ReadSpec(
248
- remote_rank=plan.all_source_ranks[0],
249
- local_block_ids=matched_local,
250
- remote_block_ids=matched_remote,
237
+ rank,
238
+ *self._apply_prefix_caching_by_region(
239
+ local_by_region,
240
+ remote_by_region,
241
+ num_computed_blocks=num_computed_blocks,
242
+ num_remote_blocks=num_remote_blocks,
243
+ remote_rank=rank,
244
+ remote_dcp_size=remote_info.remote_dcp_size,
245
+ ),
246
block_ids_by_region=True,
247
)
248
+ for rank in plan.all_source_ranks
249
+ ]
250
+ meta.region_blocks_to_zero = [
251
+ list(blocks[sum(len(spec.local_block_ids[r]) for spec in read_specs) :])
252
+ for r, blocks in enumerate(local_by_region)
253
]
254
else:
255
remote_logical_block_ids = meta.remote.block_ids