[Bugfix][HiSparse][PD] Align region-mapped pulls across logical block sizes (#57077)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Sep 16, 2026 at 18:39 UTC
804e5377aa0b0cea9544061beb2ebff162543c72
4 files changed
+267
-11
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+169
@@ -3,6 +3,8 @@
3
"""Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""
4
5
import gc
6
+import queue
7
+from types import SimpleNamespace
8
from unittest.mock import MagicMock, patch
9
10
import msgspec
@@ -39,6 +41,173 @@ from .utils import (
41
)
42
43
44
+@pytest.fixture
45
+def region_pull_worker():
46
+ """Real page mapping and descriptors, with only the transport mocked."""
47
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping
48
+
49
+ worker = object.__new__(NixlConnectorWorker)
50
+ worker.engine_id = "D"
51
+ worker.tp_rank = 0
52
+ worker.block_size = 64
53
+ worker._physical_blocks_per_logical_kv_block = 1
54
+ worker.enable_permute_local_kv = False
55
+ worker.enable_heterogeneous_attn_post_process = False
56
+ worker._engine_last_active = {}
57
+ worker._bidirectional_kv_xfer_enabled = False
58
+ worker._recving_transfers = {}
59
+ worker.use_mla, worker._has_mamba = True, False
60
+ worker.dcp_size = 1
61
+ spec = MLAAttentionSpec(
62
+ block_size=64, num_kv_heads=1, head_size=8, dtype=torch.bfloat16
63
+ )
64
+ worker.kv_cache_config = KVCacheConfig(
65
+ num_blocks=100,
66
+ kv_cache_tensors=[],
67
+ kv_cache_groups=[
68
+ KVCacheGroupSpec(["resident"], spec, enable_kv_transfer=False),
69
+ KVCacheGroupSpec(["host"], spec),
70
+ KVCacheGroupSpec(["indexer"], spec),
71
+ ],
72
+ )
73
+ worker._group_spec_types = (MLAAttentionSpec, MLAAttentionSpec)
74
+ worker.region_group_ids = [0, 1]
75
+ worker.dst_region_group_ids = {"P": [0, 0]}
76
+ worker.num_regions = 2
77
+ worker.block_len_per_layer = [1024, 1024]
78
+ worker.transfer_topo = MagicMock()
79
+ worker.transfer_topo.tp_ratio.return_value = 1
80
+ worker.transfer_topo.block_size_ratio.return_value = 1
81
+ worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
82
+ remote_tp_size=4,
83
+ remote_dcp_size=1,
84
+ remote_block_size=64,
85
+ remote_physical_blocks_per_logical=2,
86
+ )
87
+ worker.tp_mappings = {"P": TPMapping(((0,), (0,)), (0,), {0: 0}, 0)}
88
+ worker._mixed_mem_types = True
89
+ worker.src_xfer_handles_by_block_size = {64: 1}
90
+ worker._dram_src_handles_by_block_size = {64: 2}
91
+ worker.dst_xfer_side_handles = {"P": {0: 3}}
92
+ worker.dst_num_blocks = {"P": 100, "D": 100}
93
+ worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]}
94
+ worker._remote_agents = {"P": {(0, 0): "P-rank0"}}
95
+ worker._read_blocks_mixed = MagicMock()
96
+ worker.nixl_wrapper = MagicMock()
97
+ return worker
98
+
99
+
100
+@pytest.mark.cpu_test
101
+@pytest.mark.parametrize("local_ratio,remote_ratio", [(1, 2), (2, 1), (2, 2)])
102
+@pytest.mark.parametrize("num_tokens", [160, 224])
103
+@pytest.mark.parametrize("cached", [0, 1])
104
+def test_region_pull_ignores_allocation_padding(
105
+ region_pull_worker, local_ratio, remote_ratio, num_tokens, cached
106
+):
107
+ """Padded capacity must not shift either the host or indexer token range."""
108
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
109
+ NixlConnectorMetadata,
110
+ )
111
+ from vllm.utils.math_utils import cdiv
112
+
113
+ worker = region_pull_worker
114
+ worker._physical_blocks_per_logical_kv_block = local_ratio
115
+ worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(
116
+ remote_tp_size=4,
117
+ remote_dcp_size=1,
118
+ remote_block_size=64,
119
+ remote_physical_blocks_per_logical=remote_ratio,
120
+ )
121
+ local_count = cdiv(num_tokens, 64 * local_ratio)
122
+ remote_count = cdiv(num_tokens, 64 * remote_ratio)
123
+ # The host has a cached prefix, while the indexer still needs the full range.
124
+ local = [
125
+ list(range(30 + cached, 30 + local_count)),
126
+ list(range(40, 40 + local_count)),
127
+ ]
128
+ metadata = NixlConnectorMetadata()
129
+ metadata.add_new_req_to_recv(
130
+ request_id="request",
131
+ local_block_ids=local,
132
+ local_num_computed_blocks=(0, cached, 0),
133
+ awaiting_kvs=True,
134
+ kv_transfer_params={
135
+ "remote_engine_id": "P",
136
+ "remote_request_id": "request-P",
137
+ "remote_host": "localhost",
138
+ "remote_port": 1,
139
+ "remote_num_tokens": num_tokens,
140
+ "remote_block_ids": [list(range(10, 10 + remote_count))],
141
+ },
142
+ )
143
+ meta = metadata.reqs_to_recv["request"]
144
+ meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
145
+ local, local_ratio
146
+ )
147
+ worker._read_blocks_for_req("request", meta)
148
+ read = worker._read_blocks_mixed.call_args.kwargs
149
+ valid_pages = cdiv(num_tokens, 64)
150
+ source_start = 10 * remote_ratio
151
+ assert read["remote_block_descs_ids"].tolist() == (
152
+ list(range(source_start + cached * local_ratio, source_start + valid_pages))
153
+ + list(range(100 + source_start, 100 + source_start + valid_pages))
154
+ )
155
+ assert read["local_block_descs_ids"].tolist() == (
156
+ list(range((30 + cached) * local_ratio, 30 * local_ratio + valid_pages))
157
+ + list(range(100 + 40 * local_ratio, 100 + 40 * local_ratio + valid_pages))
158
+ )
159
+ assert meta.region_blocks_to_zero == [
160
+ list(
161
+ range(base * local_ratio + valid_pages, (base + local_count) * local_ratio)
162
+ )
163
+ for base in (30, 40)
164
+ ]
165
+
166
+
167
+@pytest.mark.cpu_test
168
+def test_region_pull_completion_zeros_only_own_padding(region_pull_worker):
169
+ """Completion must preserve neighboring layers and already imported pages."""
170
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
171
+ RemoteMeta,
172
+ ReqMeta,
173
+ )
174
+
175
+ worker = region_pull_worker
176
+ # Two strided layer regions in one backing, with nonzero layer offsets.
177
+ backing = torch.full((4, 3, 8), 7, dtype=torch.uint8)
178
+ caches = {"host": backing[:, 1], "indexer": backing[:, 2]}
179
+ worker.device_kv_caches = caches
180
+ worker.region_names = list(caches)
181
+ worker.kv_caches_base_addr = {"D": {0: [c.data_ptr() for c in caches.values()]}}
182
+ worker.region_num_blocks = [4, 4]
183
+ worker.block_len_per_layer = [8, 8]
184
+ worker.block_stride_per_layer = [24, 24]
185
+ worker._recving_metadata = {
186
+ "request": ReqMeta(
187
+ local_block_ids=[[1, 2], [0, 1]],
188
+ local_physical_block_ids=[[1, 2], [0, 1]],
189
+ tp_size=1,
190
+ remote=RemoteMeta([[3]], "localhost", 1, "P", "request-P"),
191
+ region_blocks_to_zero=[[2], [1]],
192
+ )
193
+ }
194
+ worker._get_new_notifs = MagicMock(return_value=set())
195
+ worker._pop_done_transfers = MagicMock(return_value=({"request"}, set()))
196
+ worker._replicated_pcp_done_sending = set()
197
+ worker._failed_recv_reqs = queue.Queue()
198
+ worker._recv_failures = set()
199
+ worker._send_pending_recv_notifs = MagicMock()
200
+ worker._sync_device_after_direct_recv = MagicMock()
201
+ worker.use_host_buffer = False
202
+ worker._reqs_to_send = {}
203
+ result = worker.get_transfer_results()
204
+ expected = torch.full_like(backing, 7)
205
+ expected[2, 1] = 0
206
+ expected[1, 2] = 0
207
+ assert torch.equal(backing, expected)
208
+ assert result.finished_recving == {"request"}
209
+
210
+
211
@pytest.mark.cpu_test
212
@pytest.mark.parametrize(
213
"swa_enabled,expected_sw_sizes",
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+45
-1
@@ -2785,6 +2785,25 @@ class NixlBaseConnectorWorker:
2785
indices=indices,
2786
)
2787
2788
+ def _zero_region_blocks(self, block_ids: BlockIds) -> None:
2789
+ """Clear clipped physical pages in their owning region only."""
2790
+ if not any(block_ids):
2791
+ return
2792
+ bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
2793
+ for region, blocks in enumerate(block_ids):
2794
+ if not blocks:
2795
+ continue
2796
+ cache = self.device_kv_caches[self.region_names[region]]
2797
+ storage = cache.untyped_storage()
2798
+ pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_(
2799
+ storage,
2800
+ bases[region] - storage.data_ptr(),
2801
+ (self.region_num_blocks[region], self.block_len_per_layer[region]),
2802
+ (self.block_stride_per_layer[region], 1),
2803
+ )
2804
+ for block in blocks:
2805
+ pages[block].zero_()
2806
+
2807
def get_transfer_results(self) -> KVConnectorTransferResults:
2808
"""
2809
Get transfers that completed on this specific worker.
@@ -2857,6 +2876,12 @@ class NixlBaseConnectorWorker:
2876
2877
direct_device_recving.add(req_id)
2878
2879
+ if meta.region_blocks_to_zero is not None:
2880
+ # P/D group positions differ. Use the actual region read plan,
2881
+ # including any local allocation padding the read did not cover.
2882
+ self._zero_region_blocks(meta.region_blocks_to_zero)
2883
+ continue
2884
+
2885
# Post processing for heteroblocksize/layout, and for blocks the
2886
# transfer clipped. The latter happens either at remote-block
2887
# granularity (block_size_ratio > 1) or at kernel-block
@@ -3254,13 +3279,32 @@ class NixlBaseConnectorWorker:
3279
3280
@staticmethod
3281
def _apply_prefix_caching_by_region(
3257
- decode_block_ids: BlockIds, prefill_block_ids: BlockIds
3282
+ decode_block_ids: BlockIds,
3283
+ prefill_block_ids: BlockIds,
3284
+ *,
3285
+ num_computed_blocks: list[int] | None = None,
3286
+ num_remote_blocks: int | None = None,
3287
) -> tuple[BlockIds, BlockIds]:
3288
"""Pair an uncached decode suffix with the same prefill regions."""
3289
assert len(decode_block_ids) == len(prefill_block_ids)
3290
if not any(decode_block_ids):
3291
return [], prefill_block_ids
3292
3293
+ if num_computed_blocks is not None:
3294
+ assert num_remote_blocks is not None
3295
+ matched_decode, matched_prefill = [], []
3296
+ for decode_region, prefill_region, start in zip(
3297
+ decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True
3298
+ ):
3299
+ count = min(len(decode_region), max(num_remote_blocks - start, 0))
3300
+ if start + count > len(prefill_region):
3301
+ raise ValueError("Remote KV pages do not cover the requested range")
3302
+ # An extra producer page can be allocation padding, not a
3303
+ # cached prefix. Select by token position instead of list length.
3304
+ matched_decode.append(list(decode_region[:count]))
3305
+ matched_prefill.append(list(prefill_region[start : start + count]))
3306
+ return matched_decode, matched_prefill
3307
+
3308
trimmed_prefill: list[list[int]] = []
3309
for decode_region, prefill_region in zip(
3310
decode_block_ids, prefill_block_ids, strict=True
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+5
@@ -230,6 +230,7 @@ class RemoteMeta:
230
engine_id: str
231
request_id: str
232
blocks_expiry_time: float | None = None
233
+ num_tokens: int | None = None
234
235
236
@dataclass
@@ -252,6 +253,9 @@ class ReqMeta:
253
# True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS
254
# and expects it in finished_recving; notify-only recvs must not be reported.
255
awaiting_kvs: bool = False
256
+ # Worker-only, per-region physical pages to zero after a successful pull.
257
+ # None selects group-based completion; empty lists mean no zeroing.
258
+ region_blocks_to_zero: BlockIds | None = None
259
260
261
class NixlConnectorMetadata(KVConnectorMetadata):
@@ -326,5 +330,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):
330
host=kv_transfer_params["remote_host"],
331
port=kv_transfer_params["remote_port"],
332
blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"),
333
+ num_tokens=kv_transfer_params.get("remote_num_tokens"),
334
)
335
self.reqs_to_recv[request_id] = req
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+48
-10
@@ -20,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
20
_is_attention_spec,
21
)
22
from vllm.logger import init_logger
23
+from vllm.utils.math_utils import cdiv
24
25
if TYPE_CHECKING:
26
from vllm.config import VllmConfig
@@ -176,6 +177,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
177
"pure MLA models"
178
)
179
assert len(plan.all_source_ranks) == 1
180
+ if self.block_size != remote_info.remote_block_size:
181
+ raise NotImplementedError(
182
+ "Region-mapped NIXL transfers require matching physical block sizes"
183
+ )
184
remote_physical_block_ids = self._logical_to_kernel_block_ids(
185
meta.remote.block_ids,
186
remote_info.remote_physical_blocks_per_logical,
@@ -183,13 +188,51 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
188
remote_by_region = self._block_ids_by_region(
189
remote_physical_block_ids, remote_region_groups
190
)
191
+ local_by_region = self._block_ids_by_region(
192
+ local_block_ids, local_region_groups
193
+ )
194
+ num_computed_blocks = None
195
+ num_remote_blocks = None
196
+ if (
197
+ meta.remote.num_tokens is not None
198
+ and meta.local_num_computed_blocks
199
+ and all(group >= 0 for group in local_region_groups)
200
+ and all(group >= 0 for group in remote_region_groups)
201
+ and not dcp_active
202
+ ):
203
+ transfer_groups = self.kv_cache_config.transfer_group_ids
204
+ num_computed_blocks = [
205
+ meta.local_num_computed_blocks[transfer_groups[group]]
206
+ * self._physical_blocks_per_logical_kv_block
207
+ for group in local_region_groups
208
+ ]
209
+ num_remote_blocks = cdiv(
210
+ meta.remote.num_tokens, remote_info.remote_block_size
211
+ )
212
+ elif (
213
+ remote_info.remote_physical_blocks_per_logical
214
+ != self._physical_blocks_per_logical_kv_block
215
+ ):
216
+ raise NotImplementedError(
217
+ "Region-mapped pulls with different logical block sizes require "
218
+ "remote_num_tokens, per-group prefix counts, unshared regions "
219
+ "and DCP=1"
220
+ )
221
+ matched_local, matched_remote = self._apply_prefix_caching_by_region(
222
+ local_by_region,
223
+ remote_by_region,
224
+ num_computed_blocks=num_computed_blocks,
225
+ num_remote_blocks=num_remote_blocks,
226
+ )
227
+ meta.region_blocks_to_zero = [
228
+ list(blocks[len(matched) :])
229
+ for blocks, matched in zip(local_by_region, matched_local, strict=True)
230
+ ]
231
read_specs = [
232
ReadSpec(
233
remote_rank=plan.all_source_ranks[0],
189
- local_block_ids=self._block_ids_by_region(
190
- local_block_ids, local_region_groups
191
- ),
192
- remote_block_ids=remote_by_region,
234
+ local_block_ids=matched_local,
235
+ remote_block_ids=matched_remote,
236
block_ids_by_region=True,
237
)
238
]
@@ -396,12 +439,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
439
self._recving_transfers.setdefault(request_id, [])
440
return True
441
399
- if read_spec.block_ids_by_region:
400
- local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region(
401
- decode_block_ids=local_block_ids,
402
- prefill_block_ids=remote_block_ids,
403
- )
404
- else:
442
+ if not read_spec.block_ids_by_region:
443
assert (
444
len(remote_block_ids)
445
== len(local_block_ids)