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

[Bugfix][HiSparse][PD] Align region-mapped pulls across logical block sizes (#57077)

Browse Source 
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
This commit is contained in:
Nicolò Lucchesi
2026-09-16 18:39:07 +02:00
committed by GitHub
parent 42919b49c5
commit 804e5377aa
4 changed files with 267 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+169
View File
@@ -3,6 +3,8 @@
"""Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""
import gc
import queue
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import msgspec
@@ -39,6 +41,173 @@ from .utils import (
)
@pytest.fixture
def region_pull_worker():
"""Real page mapping and descriptors, with only the transport mocked."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
worker = object.__new__(NixlConnectorWorker)
worker.engine_id = "D"
worker.tp_rank = 0
worker.block_size = 64
worker._physical_blocks_per_logical_kv_block = 1
worker.enable_permute_local_kv = False
worker.enable_heterogeneous_attn_post_process = False
worker._engine_last_active = {}
worker._bidirectional_kv_xfer_enabled = False
worker._recving_transfers = {}
worker.use_mla, worker._has_mamba = True, False
worker.dcp_size = 1
spec = MLAAttentionSpec(
block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16
)
worker.kv_cache_config = KVCacheConfig(
num_blocks=100,
kv_cache_tensors=[],
kv_cache_groups=[
KVCacheGroupSpec(["resident"], spec, enable_kv_transfer=False),
KVCacheGroupSpec(["host"], spec),
KVCacheGroupSpec(["indexer"], spec),
],
)
worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec)
worker.region_group_ids = [0, 1]
worker.dst_region_group_ids = {"P": [0, 0]}
worker.num_regions = 2
worker.block_len_per_layer = [1024, 1024]
worker.transfer_topo = MagicMock()
worker.transfer_topo.tp_ratio.return_value = 1
worker.transfer_topo.block_size_ratio.return_value = 1
worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
remote_tp_size=4,
remote_dcp_size=1,
remote_block_size=64,
remote_physical_blocks_per_logical=2,
)
worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)}
worker._mixed_mem_types = True
worker.src_xfer_handles_by_block_size = {64: 1}
worker._dram_src_handles_by_block_size = {64: 2}
worker.dst_xfer_side_handles = {"P": {0: 3}}
worker.dst_num_blocks = {"P": 100, "D": 100}
worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]}
worker._remote_agents = {"P": {(0, 0): "P-rank0"}}
worker._read_blocks_mixed = MagicMock()
worker.nixl_wrapper = MagicMock()
return worker
@pytest.mark.cpu_test
@pytest.mark.parametrize("local_ratio,remote_ratio", [(1, 2), (2, 1), (2, 2)])
@pytest.mark.parametrize("num_tokens", [160, 224])
@pytest.mark.parametrize("cached", [0, 1])
def test_region_pull_ignores_allocation_padding(
region_pull_worker, local_ratio, remote_ratio, num_tokens, cached
):
"""Padded capacity must not shift either the host or indexer token range."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
NixlConnectorMetadata,
)
from vllm.utils.math_utils import cdiv
worker = region_pull_worker
worker._physical_blocks_per_logical_kv_block = local_ratio
worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
remote_tp_size=4,
remote_dcp_size=1,
remote_block_size=64,
remote_physical_blocks_per_logical=remote_ratio,
)
local_count = cdiv(num_tokens, 64 * local_ratio)
remote_count = cdiv(num_tokens, 64 * remote_ratio)
# The host has a cached prefix, while the indexer still needs the full range.
local = [
list(range(30 + cached, 30 + local_count)),
list(range(40, 40 + local_count)),
]
metadata = NixlConnectorMetadata()
metadata.add_new_req_to_recv(
request_id="request",
local_block_ids=local,
local_num_computed_blocks=(0, cached, 0),
awaiting_kvs=True,
kv_transfer_params={
"remote_engine_id": "P",
"remote_request_id": "request-P",
"remote_host": "localhost",
"remote_port": 1,
"remote_num_tokens": num_tokens,
"remote_block_ids": [list(range(10, 10 + remote_count))],
},
)
meta = metadata.reqs_to_recv["request"]
meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
local, local_ratio
)
worker._read_blocks_for_req("request", meta)
read = worker._read_blocks_mixed.call_args.kwargs
valid_pages = cdiv(num_tokens, 64)
source_start = 10 * remote_ratio
assert read["remote_block_descs_ids"].tolist() == (
list(range(source_start + cached * local_ratio, source_start + valid_pages))
+ list(range(100 + source_start, 100 + source_start + valid_pages))
)
assert read["local_block_descs_ids"].tolist() == (
list(range((30 + cached) * local_ratio, 30 * local_ratio + valid_pages))
+ list(range(100 + 40 * local_ratio, 100 + 40 * local_ratio + valid_pages))
)
assert meta.region_blocks_to_zero == [
list(
range(base * local_ratio + valid_pages, (base + local_count) * local_ratio)
)
for base in (30, 40)
]
@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."""
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
RemoteMeta,
ReqMeta,
)
worker = region_pull_worker
# Two strided layer regions in one backing, with nonzero layer offsets.
backing = torch.full((4, 3, 8), 7, dtype=torch.uint8)
caches = {"host": backing[:, 1], "indexer": backing[:, 2]}
worker.device_kv_caches = caches
worker.region_names = list(caches)
worker.kv_caches_base_addr = {"D": {0: [c.data_ptr() for c in caches.values()]}}
worker.region_num_blocks = [4, 4]
worker.block_len_per_layer = [8, 8]
worker.block_stride_per_layer = [24, 24]
worker._recving_metadata = {
"request": ReqMeta(
local_block_ids=[[1, 2], [0, 1]],
local_physical_block_ids=[[1, 2], [0, 1]],
tp_size=1,
remote=RemoteMeta([[3]], "localhost", 1, "P", "request-P"),
region_blocks_to_zero=[[2], [1]],
)
}
worker._get_new_notifs = MagicMock(return_value=set())
worker._pop_done_transfers = MagicMock(return_value=({"request"}, set()))
worker._replicated_pcp_done_sending = set()
worker._failed_recv_reqs = queue.Queue()
worker._recv_failures = set()
worker._send_pending_recv_notifs = MagicMock()
worker._sync_device_after_direct_recv = MagicMock()
worker.use_host_buffer = False
worker._reqs_to_send = {}
result = worker.get_transfer_results()
expected = torch.full_like(backing, 7)
expected[2, 1] = 0
expected[1, 2] = 0
assert torch.equal(backing, expected)
assert result.finished_recving == {"request"}
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"swa_enabled,expected_sw_sizes",
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+45 -1
View File
@@ -2785,6 +2785,25 @@ class NixlBaseConnectorWorker:
indices=indices,
)
def _zero_region_blocks(self, block_ids: BlockIds) -> None:
"""Clear clipped physical pages in their owning region only."""
if not any(block_ids):
return
bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
for region, blocks in enumerate(block_ids):
if not blocks:
continue
cache = self.device_kv_caches[self.region_names[region]]
storage = cache.untyped_storage()
pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_(
storage,
bases[region] - storage.data_ptr(),
(self.region_num_blocks[region], self.block_len_per_layer[region]),
(self.block_stride_per_layer[region], 1),
)
for block in blocks:
pages[block].zero_()
def get_transfer_results(self) -> KVConnectorTransferResults:
"""
Get transfers that completed on this specific worker.
@@ -2857,6 +2876,12 @@ class NixlBaseConnectorWorker:
direct_device_recving.add(req_id)
if meta.region_blocks_to_zero is not None:
# P/D group positions differ. Use the actual region read plan,
# including any local allocation padding the read did not cover.
self._zero_region_blocks(meta.region_blocks_to_zero)
continue
# Post processing for heteroblocksize/layout, and for blocks the
# transfer clipped. The latter happens either at remote-block
# granularity (block_size_ratio > 1) or at kernel-block
@@ -3254,13 +3279,32 @@ class NixlBaseConnectorWorker:
@staticmethod
def _apply_prefix_caching_by_region(
decode_block_ids: BlockIds, prefill_block_ids: BlockIds
decode_block_ids: BlockIds,
prefill_block_ids: BlockIds,
*,
num_computed_blocks: list[int] | None = None,
num_remote_blocks: int | None = None,
) -> tuple[BlockIds, BlockIds]:
"""Pair an uncached decode suffix with the same prefill regions."""
assert len(decode_block_ids) == len(prefill_block_ids)
if not any(decode_block_ids):
return [], prefill_block_ids
if num_computed_blocks is not None:
assert num_remote_blocks is not None
matched_decode, matched_prefill = [], []
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]))
return matched_decode, matched_prefill
trimmed_prefill: list[list[int]] = []
for decode_region, prefill_region in zip(
decode_block_ids, prefill_block_ids, strict=True
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+5
View File
@@ -230,6 +230,7 @@ class RemoteMeta:
engine_id: str
request_id: str
blocks_expiry_time: float | None = None
num_tokens: int | None = None
@dataclass
@@ -252,6 +253,9 @@ class ReqMeta:
# True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS
# and expects it in finished_recving; notify-only recvs must not be reported.
awaiting_kvs: bool = False
# Worker-only, per-region physical pages to zero after a successful pull.
# None selects group-based completion; empty lists mean no zeroing.
region_blocks_to_zero: BlockIds | None = None
class NixlConnectorMetadata(KVConnectorMetadata):
@@ -326,5 +330,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):
host=kv_transfer_params["remote_host"],
port=kv_transfer_params["remote_port"],
blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"),
num_tokens=kv_transfer_params.get("remote_num_tokens"),
)
self.reqs_to_recv[request_id] = req
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+48 -10
View File
@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
_is_attention_spec,
)
from vllm.logger import init_logger
from vllm.utils.math_utils import cdiv
if TYPE_CHECKING:
from vllm.config import VllmConfig
@@ -176,6 +177,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
"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"
)
remote_physical_block_ids = self._logical_to_kernel_block_ids(
meta.remote.block_ids,
remote_info.remote_physical_blocks_per_logical,
@@ -183,13 +188,51 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
remote_by_region = self._block_ids_by_region(
remote_physical_block_ids, remote_region_groups
)
local_by_region = self._block_ids_by_region(
local_block_ids, local_region_groups
)
num_computed_blocks = None
num_remote_blocks = None
if (
meta.remote.num_tokens is not None
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 = [
meta.local_num_computed_blocks[transfer_groups[group]]
* self._physical_blocks_per_logical_kv_block
for group in local_region_groups
]
num_remote_blocks = cdiv(
meta.remote.num_tokens, remote_info.remote_block_size
)
elif (
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"
)
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=self._block_ids_by_region(
local_block_ids, local_region_groups
),
remote_block_ids=remote_by_region,
local_block_ids=matched_local,
remote_block_ids=matched_remote,
block_ids_by_region=True,
)
]
@@ -396,12 +439,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
self._recving_transfers.setdefault(request_id, [])
return True
if read_spec.block_ids_by_region:
local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region(
decode_block_ids=local_block_ids,
prefill_block_ids=remote_block_ids,
)
else:
if not read_spec.block_ids_by_region:
assert (
len(remote_block_ids)
== len(local_block_ids)