iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][NIXL] Fix DCP pulls across MLA cache regions (#57389)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Lucas Wilkinson
2026-09-21 11:57:49 -04:00
committed by GitHub
co-authored by OpenAI Codex
parent 6b858751f6
commit 8c0825090d
3 changed files with 133 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+77
View File
@@ -58,6 +58,7 @@ def region_pull_worker():
worker._recving_transfers = {}
worker.use_mla, worker._has_mamba = True, False
worker.dcp_size = 1
worker.dcp_rank = 0
spec = MLAAttentionSpec(
block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16
)
@@ -165,6 +166,82 @@ def test_region_pull_ignores_allocation_padding(
]
@pytest.mark.cpu_test
@pytest.mark.parametrize("num_pages", [0, 3, 19])
@pytest.mark.parametrize("region_groups", [(0, 1), (1, 0, 1)])
def test_dcp_region_pull(region_pull_worker, num_pages, region_groups):
"""Read each uncached page from its DCP rank; notify ranks with no pages."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
RemoteMeta,
ReqMeta,
)
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
worker = region_pull_worker
worker.region_group_ids = list(region_groups)
worker.num_regions = len(region_groups)
worker.dst_region_group_ids = {"P": [0] * worker.num_regions}
worker.dst_region_num_blocks = {
engine: [100] * worker.num_regions for engine in ("P", "D")
}
worker.block_len_per_layer = [1024] * worker.num_regions
worker.dcp_rank = 0
remote = worker.transfer_topo.get_engine_info.return_value
remote.remote_tp_size = remote.remote_dcp_size = 8
remote.remote_physical_blocks_per_logical = 1
ranks = tuple(range(8))
worker.tp_mappings["P"] = TPMapping((ranks, ranks), ranks, {r: r for r in ranks}, 8)
worker.dst_xfer_side_handles = {"P": {r: 1000 + r for r in ranks}}
worker._remote_agents = {"P": {(0, r): f"P-rank{r}" for r in ranks}}
# Two regions with different prefix hits, plus one padding page each.
local = (
[list(range(31, 31 + num_pages)), list(range(42, 41 + num_pages))]
if num_pages
else []
)
meta = ReqMeta(
local,
local,
tp_size=8,
local_num_computed_blocks=(0, 1, 2),
remote=RemoteMeta(
[[10, 11, 12]],
"localhost",
1,
"P",
"request-P",
num_tokens=max(0, num_pages * 64 - 7),
),
)
worker._read_blocks_for_req("request", meta)
reads = {
call.kwargs["remote_xfer_side_handle"] - 1000: call.kwargs
for call in worker._read_blocks_mixed.call_args_list
}
notified = {call.args[0] for call in worker.nixl_wrapper.send_notif.call_args_list}
for rank in ranks:
expected = [
(region * 100 + 30 + group * 10 + page, region * 100 + 10 + page // 8)
for region, group in enumerate(region_groups)
for page in range(group + 1, num_pages)
if page % 8 == rank
]
if expected:
read = reads[rank]
assert (
list(zip(read["local_block_descs_ids"], read["remote_block_descs_ids"]))
== expected
)
assert f"P-rank{rank}" not in notified
else:
assert rank not in reads and f"P-rank{rank}" in notified
assert meta.region_blocks_to_zero == (
[[30 + group * 10 + num_pages] for group in region_groups]
if num_pages
else None
)
@pytest.mark.cpu_test
def test_region_pull_completion_zeros_only_own_padding(region_pull_worker):
"""Completion must preserve neighboring layers and already imported pages."""
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+36 -8
View File
@@ -75,6 +75,7 @@ from vllm.distributed.parallel_state import (
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.math_utils import cdiv
from vllm.utils.network_utils import make_zmq_path
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.kv_cache_interface import (
@@ -2430,6 +2431,18 @@ class NixlBaseConnectorWorker:
remote_physical_per_logical = (
nixl_agent_meta.physical_blocks_per_logical_kv_block
)
if (
(self.dcp_size > 1 or remote_dcp_size > 1)
and self.region_group_ids
and self.region_group_ids != self.dst_region_group_ids[remote_engine_id]
and (
self._physical_blocks_per_logical_kv_block != 1
or remote_physical_per_logical != 1
)
):
raise NotImplementedError(
"DCP region pulls require matching logical and physical block sizes"
)
if (
self._has_mamba
and remote_physical_per_logical
@@ -3234,11 +3247,14 @@ class NixlBaseConnectorWorker:
local_dcp_rank: int,
remote_dcp_size: int,
local_num_computed_blocks: int,
num_remote_blocks: int | None = None,
) -> tuple[list[int], list[int]]:
"""Match local and remote block IDs by global DCP position.
``local_ids`` excludes blocks already satisfied by the local prefix
cache, while ``remote_ids`` contains the full transferable remote list.
``num_remote_blocks``, when provided, is the global block count before
DCP partitioning and excludes allocation padding.
Example:
local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6]
@@ -3254,6 +3270,9 @@ class NixlBaseConnectorWorker:
"""
local_size, remote_size = local_dcp_size, remote_dcp_size
if num_remote_blocks is not None:
num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size)
local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)]
if local_size == remote_size:
local_slice = local_ids
@@ -3274,6 +3293,8 @@ class NixlBaseConnectorWorker:
remote_slice = remote_ids[start_remote::k]
matched_blocks = min(len(local_slice), len(remote_slice))
if num_remote_blocks is not None and matched_blocks < len(local_slice):
raise ValueError("Remote KV pages do not cover the requested range")
return local_slice[:matched_blocks], remote_slice[:matched_blocks]
@staticmethod
@@ -3290,13 +3311,15 @@ class NixlBaseConnectorWorker:
block_ids_by_region.append(list(block_ids[group_id]))
return block_ids_by_region
@staticmethod
def _apply_prefix_caching_by_region(
self,
decode_block_ids: BlockIds,
prefill_block_ids: BlockIds,
*,
num_computed_blocks: list[int] | None = None,
num_remote_blocks: int | None = None,
remote_rank: int = 0,
remote_dcp_size: int = 1,
) -> tuple[BlockIds, BlockIds]:
"""Pair an uncached decode suffix with the same prefill regions."""
assert len(decode_block_ids) == len(prefill_block_ids)
@@ -3310,13 +3333,18 @@ class NixlBaseConnectorWorker:
for decode_region, prefill_region, start in zip(
decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True
):
count = min(len(decode_region), max(num_remote_blocks - start, 0))
if start + count > len(prefill_region):
raise ValueError("Remote KV pages do not cover the requested range")
# An extra producer page can be allocation padding, not a
# cached prefix. Select by token position instead of list length.
matched_decode.append(list(decode_region[:count]))
matched_prefill.append(list(prefill_region[start : start + count]))
local, remote = self._apply_dcp_prefix_caching(
decode_region,
prefill_region,
remote_rank,
self.dcp_size,
self.dcp_rank,
remote_dcp_size,
start,
num_remote_blocks=num_remote_blocks,
)
matched_decode.append(local)
matched_prefill.append(remote)
return matched_decode, matched_prefill
trimmed_prefill: list[list[int]] = []
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+20 -20
View File
@@ -191,7 +191,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
"Different NIXL cache-group layouts are only supported for "
"pure MLA models"
)
assert len(plan.all_source_ranks) == 1
if self.block_size != remote_info.remote_block_size:
raise NotImplementedError(
"Region-mapped NIXL transfers require matching physical block sizes"
@@ -213,7 +212,6 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
and meta.local_num_computed_blocks
and all(group >= 0 for group in local_region_groups)
and all(group >= 0 for group in remote_region_groups)
and not dcp_active
):
transfer_groups = self.kv_cache_config.transfer_group_ids
num_computed_blocks = [
@@ -224,32 +222,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
num_remote_blocks = cdiv(
meta.remote.num_tokens, remote_info.remote_block_size
)
elif (
remote_info.remote_physical_blocks_per_logical
elif any(local_by_region) and (
dcp_active
or remote_info.remote_physical_blocks_per_logical
!= self._physical_blocks_per_logical_kv_block
):
raise NotImplementedError(
"Region-mapped pulls with different logical block sizes require "
"remote_num_tokens, per-group prefix counts, unshared regions "
"and DCP=1"
"Region-mapped pulls with DCP or different logical block sizes "
"require remote_num_tokens, per-group prefix counts "
"and unshared regions"
)
matched_local, matched_remote = self._apply_prefix_caching_by_region(
local_by_region,
remote_by_region,
num_computed_blocks=num_computed_blocks,
num_remote_blocks=num_remote_blocks,
)
meta.region_blocks_to_zero = [
list(blocks[len(matched) :])
for blocks, matched in zip(local_by_region, matched_local, strict=True)
]
read_specs = [
ReadSpec(
remote_rank=plan.all_source_ranks[0],
local_block_ids=matched_local,
remote_block_ids=matched_remote,
rank,
*self._apply_prefix_caching_by_region(
local_by_region,
remote_by_region,
num_computed_blocks=num_computed_blocks,
num_remote_blocks=num_remote_blocks,
remote_rank=rank,
remote_dcp_size=remote_info.remote_dcp_size,
),
block_ids_by_region=True,
)
for rank in plan.all_source_ranks
]
meta.region_blocks_to_zero = [
list(blocks[sum(len(spec.local_block_ids[r]) for spec in read_specs) :])
for r, blocks in enumerate(local_by_region)
]
else:
remote_logical_block_ids = meta.remote.block_ids