@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 601fa9a74e

[KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models (#49612)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Nick Hill committed Jul 28, 2026 at 06:45 UTC 601fa9a74e8cd2ae0a4fa127d68b91bbf363a24e
7 files changed +1014 -111
tests/v1/kv_connector/unit/test_nixl_connector.py
+6 -4
@@ -479,8 +479,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
479 super().__init__(*args, kv_cache_config=kv_cache_config, **kwargs)
480 self._hand_shake_latency = hand_shake_latency
481 self.kv_cache_layout = kv_cache_layout
482 - # Mock register_kv_caches attribute needed for tests that do not call it.
482 + # Mock register_kv_caches attributes needed for tests that do not call it.
483 self.src_xfer_handles_by_block_size = {self.block_size: 1}
484 + self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
485 test_shape = self.attn_backends[0].get_kv_cache_shape(
486 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
487 )
@@ -765,8 +766,9 @@ class TestNixlHandshake:
766 assert remote_info.remote_tp_size == remote_tp_size
767 assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)
768 # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks
768 - assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio
769 - assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio
769 + split_key = (-tp_ratio, worker.block_size)
770 + assert split_key in worker.src_xfer_handles_by_tp_ratio
771 + assert len(worker.src_xfer_handles_by_tp_ratio[split_key]) == tp_ratio
772 assert remote_engine_id in worker.dst_xfer_side_handles
773 assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set(
774 range(tp_ratio)
@@ -2091,7 +2093,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
2093 # Mock register_kv_cache which registers local handle
2094 worker.src_xfer_handles_by_block_size = {worker.block_size: 455}
2095 # P TP = 2 * D TP case, we should register 2 local handles
2094 - worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]}
2096 + worker.src_xfer_handles_by_tp_ratio = {(-2, 16): [456, 457]}
2097 worker.dst_xfer_side_handles = {"engine1": {0: 789}}
2098 worker._remote_agents = {"engine1": {(0, 0): "agent1"}}
2099 # _cleanup_remote_engine (called by shutdown) also clears these:
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+114
@@ -708,6 +708,120 @@ def test_get_block_descs_ids_kernel_block_mismatch():
708 assert list(result) == expected, f"Expected {expected}, got {list(result)}"
709
710
711 +@pytest.mark.cpu_test
712 +def test_get_block_descs_ids_hetero_block_size_hybrid():
713 + """With a block-size ratio, FA desc ids are ratio-expanded while SSM
714 + desc ids keep the unexpanded logical stride (state blocks are never
715 + sub-split)."""
716 + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
717 +
718 + worker = _make_mock_worker_for_desc_ids(
719 + num_regions=2,
720 + has_mamba=True,
721 + group_spec_types=(FullAttentionSpec, MambaSpec),
722 + block_len_per_layer=[100],
723 + )
724 +
725 + ratio = 4
726 + # FA ids are already remote-granularity (expanded) sub-block ids.
727 + fa_sub_blocks = [3, 5]
728 + ssm_blocks = [1]
729 + result = worker._compute_desc_ids(
730 + block_ids=(fa_sub_blocks, ssm_blocks),
731 + dst_num_blocks=100,
732 + block_size_ratio=ratio,
733 + physical_blocks_per_logical=1,
734 + )
735 +
736 + # FA regions have 100*4 entries each; SSM regions (4 per layer) start at
737 + # 2*400 and stride by the unexpanded 100 logical blocks.
738 + expected = [3, 5, 403, 405, 801, 901, 1001, 1101]
739 + assert list(result) == expected, f"Expected {expected}, got {list(result)}"
740 +
741 +
742 +def _bind_worker_method(worker, name):
743 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
744 + NixlConnectorWorker,
745 + )
746 +
747 + method = getattr(NixlConnectorWorker, name)
748 + setattr(worker, name, method.__get__(worker, NixlConnectorWorker))
749 +
750 +
751 +@pytest.mark.cpu_test
752 +def test_map_block_ids_for_block_size_ratio_hybrid():
753 + """Attention groups expand to remote granularity and clip to the remote
754 + coverage; mamba state blocks pass through 1:1."""
755 + from unittest.mock import MagicMock
756 +
757 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
758 + NixlConnectorWorker,
759 + )
760 + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
761 +
762 + worker = MagicMock(spec=NixlConnectorWorker)
763 + worker._group_spec_types = (FullAttentionSpec, MambaSpec)
764 + _bind_worker_method(worker, "get_mapped_blocks")
765 + _bind_worker_method(worker, "_map_block_ids_for_block_size_ratio")
766 +
767 + local, remote = worker._map_block_ids_for_block_size_ratio(
768 + [[1, 2, 3], [7]],
769 + [list(range(30, 40)), [42]],
770 + 4,
771 + )
772 + # [1, 2, 3] expand to sub-blocks [4..15], clipped to the 10 remote blocks.
773 + assert local == [list(range(4, 14)), [7]]
774 + assert remote == [list(range(30, 40)), [42]]
775 +
776 + # Attention-only full prefix hit: empty local list is preserved.
777 + worker._group_spec_types = (FullAttentionSpec,)
778 + local, remote = worker._map_block_ids_for_block_size_ratio([[]], [[30, 31]], 4)
779 + assert local == []
780 +
781 +
782 +@pytest.mark.cpu_test
783 +def test_post_process_zeroes_untransferred_tail():
784 + """The untransferred sub-blocks of the last local block are zeroed on
785 + receive; mamba state caches are untouched by the attention permute."""
786 + from unittest.mock import MagicMock
787 +
788 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
789 + NixlConnectorWorker,
790 + )
791 + from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
792 +
793 + ratio = 4
794 + block_tokens = 8 # 2 tokens per remote sub-block
795 +
796 + worker = MagicMock(spec=NixlConnectorWorker)
797 + worker._group_spec_types = (FullAttentionSpec, MambaSpec)
798 + worker.transfer_topo = MagicMock()
799 + worker.device_type = "cpu"
800 + worker.enable_permute_local_kv = False
801 + attn_cache = torch.ones(6, block_tokens, 2, 4)
802 + mamba_cache = torch.ones(6, 16)
803 + worker.device_kv_caches = {"attn.0": attn_cache, "mamba.0": mamba_cache}
804 + fa_group = MagicMock(layer_names=["attn.0"])
805 + ssm_group = MagicMock(layer_names=["mamba.0"])
806 + worker.kv_cache_config = MagicMock(kv_cache_groups=[fa_group, ssm_group])
807 + # The cached property filters mamba layers out of the permuted caches.
808 + attn_caches = NixlConnectorWorker._attention_kv_caches.func(worker)
809 + assert len(attn_caches) == 1 and attn_caches[0] is attn_cache
810 + worker._attention_kv_caches = attn_caches
811 + _bind_worker_method(worker, "post_process_device_kv_on_receive")
812 +
813 + # Request occupies blocks [2, 3]; only 6 of 8 sub-blocks were received.
814 + worker.post_process_device_kv_on_receive(ratio, [([2, 3], 6)])
815 +
816 + # Block 2 fully covered; block 3 covered for 2 sub-blocks (4 tokens).
817 + assert torch.all(attn_cache[2] == 1)
818 + assert torch.all(attn_cache[3, :4] == 1)
819 + assert torch.all(attn_cache[3, 4:] == 0)
820 + # Untouched blocks and the mamba cache keep their content.
821 + assert torch.all(attn_cache[4] == 1)
822 + assert torch.all(mamba_cache == 1)
823 +
824 +
825 @pytest.mark.cpu_test
826 def test_nixl_metadata_hybrid_ssm_block_ids():
827 """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py new
+619
@@ -0,0 +1,619 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end NIXL descriptor geometry invariants for hybrid MLA+SSM models
4 +under heterogeneous P/D block geometry (TP-sharded KDA-style state, so
5 +the mamba-aligned logical block size differs between P and D while the
6 +kernel-granularity pages stay equal).
7 +
8 +The invariant under test: every LOCAL byte range a request's READ transfers
9 +into must lie within that request's own blocks. A violation means an
10 +incoming transfer can overwrite a co-resident request's KV or mamba state
11 +mid-decode (silent corruption of an unrelated request).
12 +"""
13 +
14 +from unittest.mock import patch
15 +
16 +import numpy as np
17 +import pytest
18 +import torch
19 +
20 +from .utils import create_vllm_config
21 +
22 +
23 +class _RecordingNixl:
24 + """Minimal NIXL wrapper stand-in that records descriptor lists and
25 + prepared transfers so tests can resolve desc ids to byte ranges."""
26 +
27 + def __init__(self, *args, **kwargs):
28 + self.dlists: dict[int, np.ndarray] = {}
29 + self.xfers: list[tuple] = []
30 + self._next_handle = 1
31 +
32 + def get_reg_descs(self, caches_data, mem_type):
33 + return caches_data
34 +
35 + def register_memory(self, descs, backends=None):
36 + pass
37 +
38 + def deregister_memory(self, descs):
39 + pass
40 +
41 + def get_agent_metadata(self):
42 + return b"agent-meta"
43 +
44 + def get_xfer_descs(self, blocks_data, mem_type):
45 + return blocks_data
46 +
47 + def prep_xfer_dlist(self, agent, descs):
48 + handle = self._next_handle
49 + self._next_handle += 1
50 + self.dlists[handle] = np.asarray(descs, dtype=np.uint64).reshape(-1, 3)
51 + return handle
52 +
53 + def add_remote_agent(self, metadata):
54 + return "remote-agent"
55 +
56 + def make_prepped_xfer(
57 + self, op, local_handle, local_ids, remote_handle, remote_ids, notif_msg=None
58 + ):
59 + handle = self._next_handle
60 + self._next_handle += 1
61 + self.xfers.append(
62 + (
63 + op,
64 + local_handle,
65 + np.asarray(local_ids),
66 + remote_handle,
67 + np.asarray(remote_ids),
68 + )
69 + )
70 + return handle
71 +
72 + def transfer(self, handle):
73 + pass
74 +
75 + def check_xfer_state(self, handle):
76 + return "DONE"
77 +
78 + def get_xfer_telemetry(self, handle):
79 + from types import SimpleNamespace
80 +
81 + return SimpleNamespace(
82 + xferDuration=1.0, postDuration=1.0, totalBytes=1, descCount=1
83 + )
84 +
85 + def release_xfer_handle(self, handle):
86 + pass
87 +
88 + def release_dlist_handle(self, handle):
89 + pass
90 +
91 + def send_notif(self, agent, notif_msg=None):
92 + pass
93 +
94 + def get_new_notifs(self):
95 + return {}
96 +
97 + def remove_remote_agent(self, agent):
98 + pass
99 +
100 +
101 +def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
102 + """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
103 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
104 + base_worker as bw,
105 + )
106 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
107 + NixlConnectorWorker,
108 + )
109 + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
110 + from vllm.v1.kv_cache_interface import (
111 + KVCacheConfig,
112 + KVCacheGroupSpec,
113 + KVCacheTensor,
114 + MambaSpec,
115 + MLAAttentionSpec,
116 + )
117 +
118 + mla_spec = MLAAttentionSpec(
119 + block_size=local_block_size,
120 + num_kv_heads=1,
121 + head_size=6,
122 + dtype=torch.float16,
123 + )
124 + unified_page = mla_spec.page_size_bytes
125 + kda_spec = MambaSpec(
126 + block_size=local_block_size,
127 + shapes=((8, 3), (1, 4, 4)),
128 + dtypes=(torch.float16, torch.float32),
129 + page_size_padded=unified_page,
130 + mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
131 + )
132 + kv_cache_config = KVCacheConfig(
133 + num_blocks=num_logical_blocks,
134 + kv_cache_tensors=[
135 + KVCacheTensor(
136 + size=num_logical_blocks * unified_page,
137 + shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
138 + )
139 + for i in range(2)
140 + ],
141 + kv_cache_groups=[
142 + KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
143 + KVCacheGroupSpec(["kda_a.0", "kda_a.1"], kda_spec),
144 + KVCacheGroupSpec(["kda_b.0", "kda_b.1"], kda_spec),
145 + ],
146 + )
147 +
148 + vllm_config = create_vllm_config(block_size=local_block_size)
149 + vllm_config.cache_config.enable_prefix_caching = False
150 + # kv_buffer_device defaults to the *real* platform's device type, which on
151 + # a CPU-only test host would make this a host-buffer worker: host xfer
152 + # buffers are per-layer, so the HMA shared-tensor regions this test builds
153 + # would not be deduplicated. Pin it to the faked device type.
154 + vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
155 +
156 + from unittest.mock import MagicMock
157 +
158 + fake_backend = MagicMock()
159 + fake_backend.get_supported_kernel_block_sizes.return_value = [kernel_block_size]
160 + fake_backend.get_name.return_value = "FLASHMLA"
161 + fake_backend.full_cls_name.return_value = "fake.FLASHMLA"
162 + fake_platform = MagicMock()
163 + fake_platform.device_type = "cuda"
164 + fake_platform.get_nixl_memory_type.return_value = "VRAM"
165 +
166 + from vllm.config import set_current_vllm_config
167 +
168 + with (
169 + patch.object(bw, "NixlWrapper", _RecordingNixl),
170 + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
171 + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
172 + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
173 + patch.object(bw, "current_platform", fake_platform),
174 + patch(
175 + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
176 + return_value="DS",
177 + ),
178 + set_current_vllm_config(vllm_config),
179 + ):
180 + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
181 + worker.use_mla = True
182 +
183 + # Attention caches are kernel-block granular on dim 0, as the
184 + # receive post-process assumes.
185 + ppl = local_block_size // kernel_block_size
186 + tensors = [
187 + torch.zeros(
188 + num_logical_blocks * ppl, unified_page // ppl, dtype=torch.uint8
189 + )
190 + for _ in range(2)
191 + ]
192 + worker.register_kv_caches(
193 + {
194 + "kda_a.0": tensors[0],
195 + "mla.0": tensors[0],
196 + "kda_b.0": tensors[0],
197 + "kda_a.1": tensors[1],
198 + "mla.1": tensors[1],
199 + "kda_b.1": tensors[1],
200 + }
201 + )
202 + # Keep tensors alive alongside the worker; flat views for byte checks.
203 + worker._test_tensors = [t.view(-1) for t in tensors]
204 + worker._test_tensors_2d = tensors
205 + worker._test_unified_page = unified_page
206 + return worker
207 +
208 +
209 +def _make_remote_meta(
210 + worker,
211 + remote_block_size,
212 + remote_kernel_block_size,
213 + remote_num_logical,
214 + remote_ssm_sizes,
215 +):
216 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
217 + NixlAgentMetadata,
218 + )
219 +
220 + remote_ppl = remote_block_size // remote_kernel_block_size
221 + # Kernel-granularity pages are TP-independent for MLA hybrids and must
222 + # match the local ones for the handshake to pass, scaled down by the
223 + # block-size ratio when the remote's kernel block is smaller.
224 + block_size_ratio = worker.block_size // remote_kernel_block_size
225 + kernel_page = worker.block_len_per_layer[0] // block_size_ratio
226 + return NixlAgentMetadata(
227 + engine_id="remote-engine",
228 + agent_metadata=b"remote-agent-meta",
229 + device_id=0,
230 + kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
231 + num_blocks=remote_num_logical * remote_ppl,
232 + block_lens=[kernel_page, kernel_page],
233 + kv_cache_layout=worker.kv_cache_layout,
234 + block_size=remote_kernel_block_size,
235 + ssm_sizes=remote_ssm_sizes,
236 + attn_backend_name=worker.backend_name,
237 + physical_blocks_per_logical_kv_block=remote_ppl,
238 + )
239 +
240 +
241 +def _owned_byte_ranges(worker, group_logical_ids):
242 + """Byte ranges owned by a request: for each HMA region tensor, every
243 + logical block id of every group maps to one unified page."""
244 + unified_page = worker._test_unified_page
245 + bases = [t.data_ptr() for t in worker._test_tensors]
246 + owned = []
247 + for base in bases:
248 + for ids in group_logical_ids:
249 + for b in ids:
250 + owned.append((base + b * unified_page, base + (b + 1) * unified_page))
251 + return owned
252 +
253 +
254 +def _assert_local_writes_within(worker, owned_ranges):
255 + nixl = worker.nixl_wrapper
256 + assert nixl.xfers, "no transfers were posted"
257 + violations = []
258 + total_descs = 0
259 + for op, local_handle, local_ids, _, remote_ids in nixl.xfers:
260 + assert len(local_ids) == len(remote_ids)
261 + desc_arr = nixl.dlists[local_handle]
262 + for i in local_ids:
263 + addr, length, _dev = desc_arr[int(i)]
264 + addr, length = int(addr), int(length)
265 + total_descs += 1
266 + if not any(lo <= addr and addr + length <= hi for lo, hi in owned_ranges):
267 + violations.append((int(i), hex(addr), length))
268 + assert not violations, (
269 + f"{len(violations)}/{total_descs} local descriptors write outside "
270 + f"the request's own blocks: {violations[:10]}"
271 + )
272 + return total_descs
273 +
274 +
275 +@pytest.mark.cpu_test
276 +def test_hetero_ppl_multi_read_writes_stay_within_request_blocks():
277 + """MLA-hybrid hetero geometry: local (D, TP1) logical blocks of 12 tokens
278 + (kernel 4, ppl=3) vs remote (P, TP2) logical blocks of 8 tokens (ppl=2),
279 + equal kernel pages, tp_ratio=-2 multi-read with replicated MLA and
280 + TP-sharded KDA state. Every local descriptor of the request's reads must
281 + stay within its own blocks."""
282 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
283 + NixlConnectorMetadata,
284 + )
285 +
286 + worker = _make_mla_hybrid_worker(
287 + local_block_size=12, kernel_block_size=4, num_logical_blocks=8
288 + )
289 + assert worker._physical_blocks_per_logical_kv_block == 3
290 +
291 + meta_r = _make_remote_meta(
292 + worker,
293 + remote_block_size=8,
294 + remote_kernel_block_size=4,
295 + remote_num_logical=12,
296 + remote_ssm_sizes=(24, 32),
297 + )
298 + for rank in (0, 1):
299 + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=2)
300 +
301 + # Request B: 17 matched tokens. Local: 2 logical blocks (24 tok
302 + # capacity); remote: 16 prefilled tokens -> 2 remote logical blocks.
303 + # Sparse, non-contiguous ids so neighbor blocks exist on all sides.
304 + local_ids = ([2, 5], [1], [7])
305 + remote_ids = [[1, 4], [5], [2]]
306 +
307 + metadata = NixlConnectorMetadata()
308 + metadata.add_new_req_to_recv(
309 + request_id="req-b",
310 + local_block_ids=local_ids,
311 + kv_transfer_params={
312 + "remote_block_ids": remote_ids,
313 + "remote_engine_id": "remote-engine",
314 + "remote_request_id": "prefill-req-b",
315 + "remote_host": "localhost",
316 + "remote_port": 1234,
317 + "tp_size": 2,
318 + },
319 + )
320 + meta = metadata.reqs_to_recv["req-b"]
321 + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
322 + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
323 + )
324 + worker._recving_metadata["req-b"] = meta
325 +
326 + worker._read_blocks_for_req("req-b", meta)
327 +
328 + owned = _owned_byte_ranges(worker, local_ids)
329 + total = _assert_local_writes_within(worker, owned)
330 + # Multi-read: rank 0 carries the replicated MLA + its SSM shard,
331 + # rank 1 carries only its SSM shard.
332 + assert len(worker.nixl_wrapper.xfers) == 2
333 + assert total > 0
334 +
335 +
336 +def _resolve(
337 + desc_arr,
338 + idx,
339 + bases,
340 + region_size,
341 + unified_page,
342 + desc_page,
343 + logical_ids_attn,
344 + block_tokens,
345 +):
346 + """Resolve a desc id to (region, kind, token_start) where kind is 'attn'
347 + (desc-page sized, sub-block-aligned, in the request's attention blocks)
348 + or 'mamba'. token_start is the request-relative token offset, so local
349 + and remote are comparable even when their kernel blocks differ in size."""
350 + addr, length, _ = (int(x) for x in desc_arr[int(idx)])
351 + for region, base in enumerate(bases):
352 + off = addr - base
353 + if 0 <= off < region_size:
354 + b = off // unified_page
355 + rem = off % unified_page
356 + if length == desc_page and rem % desc_page == 0 and b in logical_ids_attn:
357 + pos = logical_ids_attn.index(b)
358 + tokens_per_desc = block_tokens * desc_page // unified_page
359 + sub = rem // desc_page
360 + return (region, "attn", pos * block_tokens + sub * tokens_per_desc)
361 + return (region, "mamba", None)
362 + raise AssertionError(f"desc {idx} addr {addr:#x} not in any region")
363 +
364 +
365 +def _run_hetero_case(
366 + local_block, kernel, remote_block, num_tokens, tp_size=2, remote_kernel=None
367 +):
368 + """Full pull-path run for one geometry; returns pairing records.
369 +
370 + ``remote_kernel`` defaults to the local kernel block size; a smaller
371 + value additionally exercises block_size_ratio > 1.
372 + """
373 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
374 + NixlConnectorMetadata,
375 + )
376 +
377 + remote_kernel = remote_kernel or kernel
378 + block_size_ratio = kernel // remote_kernel
379 + remote_ppl = remote_block // remote_kernel
380 + matched = num_tokens - 1 # mamba N-1 rule
381 + n_local = -(-num_tokens // local_block)
382 + n_remote = -(-matched // remote_block)
383 +
384 + worker = _make_mla_hybrid_worker(
385 + local_block_size=local_block,
386 + kernel_block_size=kernel,
387 + num_logical_blocks=max(2 * n_local + 4, 8),
388 + )
389 + # Local KDA state pages are (48, 64) bytes; the remote holds 1/tp_size
390 + # shards of each.
391 + meta_r = _make_remote_meta(
392 + worker,
393 + remote_block_size=remote_block,
394 + remote_kernel_block_size=remote_kernel,
395 + remote_num_logical=max(2 * n_remote + 4, 8),
396 + remote_ssm_sizes=(48 // tp_size, 64 // tp_size),
397 + )
398 + for rank in range(tp_size):
399 + worker.add_remote_agent(meta_r, remote_tp_rank=rank, remote_tp_size=tp_size)
400 +
401 + # Sparse ids so neighbors exist between the request's blocks.
402 + local_attn = [2 * i + 1 for i in range(n_local)]
403 + remote_attn = [2 * i + 2 for i in range(n_remote)]
404 + local_ids = (local_attn, [0], [2 * n_local + 2])
405 + remote_ids = [remote_attn, [1], [0]]
406 +
407 + metadata = NixlConnectorMetadata()
408 + metadata.add_new_req_to_recv(
409 + request_id="req-b",
410 + local_block_ids=local_ids,
411 + kv_transfer_params={
412 + "remote_block_ids": remote_ids,
413 + "remote_engine_id": "remote-engine",
414 + "remote_request_id": "prefill-req-b",
415 + "remote_host": "localhost",
416 + "remote_port": 1234,
417 + "tp_size": tp_size,
418 + },
419 + )
420 + meta = metadata.reqs_to_recv["req-b"]
421 + meta.local_physical_block_ids = worker._logical_to_kernel_block_ids(
422 + meta.local_block_ids, worker._physical_blocks_per_logical_kv_block
423 + )
424 + worker._recving_metadata["req-b"] = meta
425 +
426 + # Sentinel-fill the local KV so untouched bytes are detectable.
427 + for t in worker._test_tensors:
428 + t.fill_(0xAA)
429 +
430 + worker._read_blocks_for_req("req-b", meta)
431 +
432 + # Invariant 1: all local writes within the request's own blocks.
433 + owned = _owned_byte_ranges(worker, local_ids)
434 + _assert_local_writes_within(worker, owned)
435 +
436 + # Invariant 2: local<->remote attention pairs are token-aligned.
437 + nixl = worker.nixl_wrapper
438 + local_bases = [t.data_ptr() for t in worker._test_tensors]
439 + remote_bases = [0x10_000_000, 0x20_000_000]
440 + local_unified = worker._test_unified_page
441 + remote_unified = (local_unified // local_block) * remote_block
442 + # With block_size_ratio > 1 the local page is split into ratio sub-descs,
443 + # each the size of a whole remote kernel page.
444 + desc_page = worker.block_len_per_layer[0] // block_size_ratio
445 + meta_r_num_blocks_bytes = (meta_r.num_blocks // remote_ppl) * remote_unified
446 + covered_tokens = set()
447 + for op, lh, lids, rh, rids in nixl.xfers:
448 + larr, rarr = nixl.dlists[lh], nixl.dlists[rh]
449 + for li, ri in zip(lids, rids):
450 + lreg, lkind, ltok = _resolve(
451 + larr,
452 + li,
453 + local_bases,
454 + len(worker._test_tensors[0]),
455 + local_unified,
456 + desc_page,
457 + local_attn,
458 + local_block,
459 + )
460 + rreg, rkind, rtok = _resolve(
461 + rarr,
462 + ri,
463 + remote_bases,
464 + meta_r_num_blocks_bytes,
465 + remote_unified,
466 + desc_page,
467 + remote_attn,
468 + remote_block,
469 + )
470 + assert lkind == rkind, (
471 + f"pair kind mismatch: local {lkind} vs remote {rkind} "
472 + f"(local desc {li}, remote desc {ri})"
473 + )
474 + assert lreg == rreg, (
475 + f"region mismatch: local {lreg} vs remote {rreg} for "
476 + f"tokens {ltok} vs {rtok}"
477 + )
478 + if lkind == "attn":
479 + assert ltok == rtok, (
480 + f"TOKEN MISALIGNMENT: local sub-block holds tokens "
481 + f"[{ltok}..) but receives remote tokens [{rtok}..) "
482 + f"(geometry local_block={local_block}, "
483 + f"remote_block={remote_block}, N={num_tokens})"
484 + )
485 + covered_tokens.add(ltok)
486 +
487 + # Invariant 3: full coverage of the matched tokens, at the finest
488 + # transfer granularity (the remote kernel block).
489 + needed = {t for t in range(0, matched - matched % remote_kernel, remote_kernel)}
490 + missing = needed - covered_tokens
491 + assert not missing, (
492 + f"tokens never transferred: {sorted(missing)[:8]} "
493 + f"(geometry local_block={local_block}, remote_block={remote_block}, "
494 + f"N={num_tokens}, matched={matched})"
495 + )
496 +
497 + # Invariant 4: no stale bytes after receive completion. The scheduler
498 + # excludes the blocks covering the matched tokens from alloc-time KV
499 + # zeroing (the zeroing would race the RDMA write), so every byte of
500 + # those blocks must be either written by the transfer or zeroed by the
501 + # receive post-process. Stale bytes surface as mid-response garbage
502 + # once decode grows into the untransferred tail.
503 + for op, lh, lids, rh, rids in nixl.xfers:
504 + larr = nixl.dlists[lh]
505 + for li in lids:
506 + addr, length, _ = (int(x) for x in larr[int(li)])
507 + for t in worker._test_tensors:
508 + off = addr - t.data_ptr()
509 + if 0 <= off < t.numel():
510 + t[off : off + length] = 0 # simulate the RDMA write
511 + break
512 + done_sending, done_recving = worker.get_finished()
513 + assert "req-b" in done_recving
514 + n_excluded = -(-matched // local_block)
515 + stale = []
516 + for b in local_attn[:n_excluded]:
517 + for region, t in enumerate(worker._test_tensors):
518 + page = t[b * local_unified : (b + 1) * local_unified]
519 + n_stale = int((page == 0xAA).sum())
520 + if n_stale:
521 + stale.append((region, b, n_stale))
522 + assert not stale, (
523 + f"stale (unzeroed, untransferred) bytes in matched-range attention "
524 + f"blocks (region, block, bytes): {stale} "
525 + f"(geometry local_block={local_block}, remote_block={remote_block}, "
526 + f"N={num_tokens}, matched={matched})"
527 + )
528 +
529 +
530 +@pytest.mark.cpu_test
531 +@pytest.mark.parametrize(
532 + "local_block,remote_block",
533 + [
534 + (12, 8), # ppl 3 vs 2
535 + (36, 8), # ppl 9 vs 2 (large ppl asymmetry, scaled)
536 + (24, 4), # ppl 6 vs 1
537 + (16, 24), # remote larger than local (D_TP > P_TP direction)
538 + ],
539 +)
540 +@pytest.mark.parametrize("num_tokens", list(range(2, 40)))
541 +def test_hetero_ppl_token_alignment_sweep(local_block, remote_block, num_tokens):
542 + """Sweep prompt lengths across block-boundary residues for several
543 + hetero-ppl geometries; assert neighbor-safety, token alignment, and
544 + coverage of every transferred kernel block."""
545 + _run_hetero_case(
546 + local_block, kernel=4, remote_block=remote_block, num_tokens=num_tokens
547 + )
548 +
549 +
550 +@pytest.mark.cpu_test
551 +@pytest.mark.parametrize(
552 + "num_tokens",
553 + # Residues around the remote kernel block (4), the local kernel block
554 + # (8), the remote logical block (8) and the local logical block (24).
555 + [2, 5, 8, 9, 13, 16, 17, 21, 24, 25, 29, 32, 33, 41, 48, 49],
556 +)
557 +def test_hetero_ppl_with_block_size_ratio(num_tokens):
558 + """Both hetero regimes at once: kernel blocks differ (local 8 / remote
559 + 4, block_size_ratio=2) *and* physical_blocks_per_logical differs (3 vs
560 + 2). The transfer is clipped at remote sub-block granularity by the
561 + pairing and front-trimmed by _apply_prefix_caching, so the
562 + untransferred tail can span both a partial block and whole blocks —
563 + the case each of the two former zeroing paths handled only half of."""
564 + _run_hetero_case(
565 + local_block=24,
566 + kernel=8,
567 + remote_block=8,
568 + remote_kernel=4,
569 + num_tokens=num_tokens,
570 + )
571 +
572 +
573 +@pytest.mark.cpu_test
574 +@pytest.mark.parametrize(
575 + "num_tokens",
576 + # Residues around every geometric boundary: kernel block (64), remote
577 + # logical block (768), local logical block (5760), plus odd offsets.
578 + [
579 + 2,
580 + 63,
581 + 64,
582 + 65,
583 + 127,
584 + 128,
585 + 300,
586 + 640,
587 + 767,
588 + 768,
589 + 769,
590 + 831,
591 + 832,
592 + 1000,
593 + 1535,
594 + 1536,
595 + 1537,
596 + 2303,
597 + 2304,
598 + 2305,
599 + 3001,
600 + 5759,
601 + 5760,
602 + 5761,
603 + 5824,
604 + 6528,
605 + 6529,
606 + ],
607 +)
608 +def test_mla_hybrid_large_ppl_geometry(num_tokens):
609 + """KimiLinear-scale MLA-hybrid geometry (TP8 prefill -> TP1 decode):
610 + decode (local) logical block 5760 / kernel 64 (ppl=90), prefill
611 + (remote) logical block 768 (ppl=12), tp_ratio=-8 multi-read with
612 + replicated MLA and 8-way TP-sharded KDA state."""
613 + _run_hetero_case(
614 + local_block=5760,
615 + kernel=64,
616 + remote_block=768,
617 + num_tokens=num_tokens,
618 + tp_size=8,
619 + )
tests/v1/kv_connector/unit/test_tp_mapping.py
+54
@@ -161,3 +161,57 @@ class TestMambaPlanSplitHandles:
161 # FA: chunk=200//1=200, slot=0 (skip_fa) → (1000, 200, 0), (2000, 200, 0)
162 # SSM: chunk=400//2=200, idx=1 → (3200, 200, 0)
163 assert splits[1] == [(1000, 200, 0), (2000, 200, 0), (3200, 200, 0)]
164 +
165 + def test_hetero_block_size_splits(self):
166 + """With a block-size ratio, single-source FA sub-block descs pass
167 + through whole; SSM descs are unexpanded and split per source."""
168 + plan = TPMapping(
169 + source_ranks_per_group=((0,), (0, 1)),
170 + all_source_ranks=(0, 1),
171 + rank_to_attention_slot={0: 0, 1: 0},
172 + rank_offset_factor=0,
173 + )
174 +
175 + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
176 + # 2 FA blocks x ratio 2 sub-blocks + 1 SSM desc (never expanded).
177 + src_blocks_data = np.array(
178 + [
179 + (1000, 100, 0),
180 + (1100, 100, 0),
181 + (2000, 100, 0),
182 + (2100, 100, 0),
183 + (3000, 400, 0),
184 + ],
185 + dtype=np.uint64,
186 + )
187 +
188 + splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 4, 2))
189 +
190 + assert len(splits) == 2
191 + fa_passthrough = [
192 + (1000, 100, 0),
193 + (1100, 100, 0),
194 + (2000, 100, 0),
195 + (2100, 100, 0),
196 + ]
197 + assert splits[0] == fa_passthrough + [(3000, 200, 0)]
198 + assert splits[1] == fa_passthrough + [(3200, 200, 0)]
199 +
200 + def test_hetero_block_size_head_sharded_asserts(self):
201 + """Head-sharded FA reads (multiple FA sources) are incompatible with
202 + a block-size mismatch and must fail loudly."""
203 + plan = TPMapping(
204 + source_ranks_per_group=((0, 1), (0, 1)),
205 + all_source_ranks=(0, 1),
206 + rank_to_attention_slot={0: 0, 1: 1},
207 + rank_offset_factor=0,
208 + )
209 +
210 + worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec))
211 + src_blocks_data = np.array(
212 + [(1000, 100, 0), (1100, 100, 0), (3000, 400, 0)],
213 + dtype=np.uint64,
214 + )
215 +
216 + with pytest.raises(AssertionError, match="Head-sharded"):
217 + list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2))
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+207 -64
@@ -12,6 +12,7 @@ import uuid
12 from collections import defaultdict
13 from collections.abc import Iterator
14 from concurrent.futures import Future, ThreadPoolExecutor
15 +from functools import cached_property
16 from typing import TYPE_CHECKING, Any, cast
17
18 import msgspec
@@ -67,6 +68,7 @@ from vllm.distributed.parallel_state import (
68 from vllm.logger import init_logger
69 from vllm.platforms import current_platform
70 from vllm.utils.network_utils import make_zmq_path
71 +from vllm.utils.torch_utils import async_tensor_h2d
72 from vllm.v1.attention.backends.utils import get_kv_cache_layout
73 from vllm.v1.kv_cache_interface import (
74 FullAttentionSpec,
@@ -122,9 +124,11 @@ class NixlBaseConnectorWorker:
124 return (region_ids * num_blocks + block_arr).flatten()
125
126 # Compute desc ids per group using the right stride: FA descs have
125 - # num_blocks entries per region (kernel granularity), SSM descs have
126 - # logical_blocks entries per region (no kernel splitting).
127 - logical_blocks = num_blocks // physical_blocks_per_logical
127 + # num_blocks entries per region (kernel granularity, expanded by
128 + # block_size_ratio for heterogeneous block sizes), SSM descs have
129 + # logical_blocks entries per region (no kernel splitting, and never
130 + # ratio-expanded since state blocks are indivisible).
131 + logical_blocks = dst_num_blocks // physical_blocks_per_logical
132 all_descs: list[np.ndarray] = []
133 for i, group in enumerate(block_ids):
134 group_arr = np.asarray(group)
@@ -161,6 +165,7 @@ class NixlBaseConnectorWorker:
165 plan: TPMapping,
166 src_blocks_data: np.ndarray,
167 num_fa_descs: int,
168 + block_size_ratio: int = 1,
169 ) -> Iterator[list[tuple[int, int, int]]]:
170 """Build split handle data for P_TP > D_TP scenario.
171
@@ -187,6 +192,11 @@ class NixlBaseConnectorWorker:
192
193 # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
194 fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)
195 +
196 + assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), (
197 + "Head-sharded attention reads with P_TP > D_TP and heterogeneous "
198 + "block sizes are not supported"
199 + )
200 src_blocks_list = src_blocks_data.tolist()
201
202 for p_idx, p_rank in enumerate(plan.all_source_ranks):
@@ -442,9 +452,12 @@ class NixlBaseConnectorWorker:
452
453 # nixl_prepped_dlist_handle.
454 self.src_xfer_handles_by_block_size: dict[int, int] = {}
455 + # Local descriptor arrays per remote block size (block_size_ratio>1),
456 + # kept for building per-tp-ratio splits at the same granularity.
457 + self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {}
458 # Populated dynamically during handshake based on remote configuration.
446 - # Keep track of regions at different tp_ratio values. tp_ratio->handles
447 - self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {}
459 + # Per-source split handles, keyed by (tp_ratio, remote_block_size).
460 + self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
461 # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
462 self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
463
@@ -1258,11 +1271,7 @@ class NixlBaseConnectorWorker:
1271 agent_metadata_bytes=encoder.encode(agent_metadata),
1272 )
1273
1261 - def _build_mamba_local(
1262 - self,
1263 - base_addresses: list[int],
1264 - block_size_ratio: int,
1265 - ) -> np.ndarray:
1274 + def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray:
1275 """Build desc regions (conv sub-projections + ssm) per layer for
1276 local mamba blocks with DS conv layout, as an Nx3 uint64 array.
1277
@@ -1289,16 +1298,17 @@ class NixlBaseConnectorWorker:
1298 | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
1299 +-------------------+ +--------------------+
1300 |1st_split-2nd_split| |1st_split-2nd_split |
1301 +
1302 + Mamba state blocks are indivisible (not token-extent data), so the
1303 + descriptors always use the local page geometry regardless of any
1304 + attention block-size ratio; their desc ids are likewise never
1305 + ratio-expanded (see _compute_desc_ids).
1306 """
1293 - assert block_size_ratio == 1, (
1294 - "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1295 - f"Got block_size_ratio={block_size_ratio}."
1296 - )
1307 assert base_addresses, "Local KV cache base addresses must not be empty."
1308 assert self._conv_decomp is not None
1309 conv_offsets = self._conv_decomp.local_conv_offsets
1310 conv_size, ssm_size = self._mamba_ssm_size
1301 - num_blocks = self._logical_num_blocks * block_size_ratio
1311 + num_blocks = self._logical_num_blocks
1312 physical_per_logical = self._physical_blocks_per_logical_kv_block
1313 device_id = self.device_id
1314 block_arange = np.arange(num_blocks, dtype=np.uint64)
@@ -1307,9 +1317,7 @@ class NixlBaseConnectorWorker:
1317 for i, base_addr in enumerate(base_addresses):
1318 # Jump one page_size, but ssm page_size may be bigger when kernel
1319 # locks block size to a specific value (physical_per_logical scale).
1310 - page_stride = (
1311 - self.block_len_per_layer[i] // block_size_ratio * physical_per_logical
1312 - )
1320 + page_stride = self.block_len_per_layer[i] * physical_per_logical
1321 blk_addrs = base_addr + block_arange * page_stride
1322 for off, sz in conv_offsets:
1323 parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
@@ -1459,7 +1467,7 @@ class NixlBaseConnectorWorker:
1467 self.device_id,
1468 )
1469 if self._has_mamba:
1462 - assert self.num_descs == len(blocks_data)
1470 + assert self.num_descs * block_size_ratio == len(blocks_data)
1471 # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
1472 # is unnecessary — a single conv desc per block suffices. Consider
1473 # adding a fast path that falls back to the standard 2-region
@@ -1467,7 +1475,7 @@ class NixlBaseConnectorWorker:
1475 # remote has been seen. Currently we always register 4 regions
1476 # because local descs are created before knowing the remote TP.
1477 logger.debug("Registering local Mamba descriptors (4 regions/layer)")
1470 - mamba = self._build_mamba_local(local_base_addresses, block_size_ratio)
1478 + mamba = self._build_mamba_local(local_base_addresses)
1479 blocks_data = np.concatenate([blocks_data, mamba])
1480
1481 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -1607,29 +1615,44 @@ class NixlBaseConnectorWorker:
1615
1616 plan = self.tp_mappings[engine_id]
1617
1618 + ### (Optional) Register a local handler at the remote engine's block
1619 + ### granularity (remote/prefill blocks smaller than local).
1620 + remote_block_size = nixl_agent_meta.block_size
1621 + src_blocks_data = self.src_blocks_data
1622 + if block_size_ratio > 1:
1623 + if remote_block_size not in self.src_xfer_handles_by_block_size:
1624 + handle, blocks_data = self.register_local_xfer_handler(
1625 + remote_block_size
1626 + )
1627 + self.src_xfer_handles_by_block_size[remote_block_size] = handle
1628 + self.src_blocks_data_by_block_size[remote_block_size] = blocks_data
1629 + src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size]
1630 +
1631 ### (Optional) Register local agent memory regions. MLA is not split.
1632 + split_key = (tp_ratio, remote_block_size)
1633 if (
1634 tp_ratio < 0
1635 and (not self.use_mla or len(plan.all_source_ranks) > 1)
1614 - and tp_ratio not in self.src_xfer_handles_by_tp_ratio
1636 + and split_key not in self.src_xfer_handles_by_tp_ratio
1637 ):
1638 # Remote tp_size > local tp_size: read from multiple remote ranks.
1639 # Logically "split" own regions into per-source chunks. Hybrid
1640 # MLA+SSM also needs this path: MLA is replicated and read once,
1641 # while the SSM state is sharded across every remote TP rank.
1620 - # We only do this once per remote tp_size (replica-friendly).
1621 - self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1642 + # We only do this once per remote (tp_size, block_size).
1643 + self.src_xfer_handles_by_tp_ratio[split_key] = []
1644
1645 for handle_data in self._build_local_splits_from_plan(
1646 plan,
1625 - self.src_blocks_data,
1626 - self.num_descs,
1647 + src_blocks_data,
1648 + self.num_descs * block_size_ratio,
1649 + block_size_ratio,
1650 ):
1651 descs = self.nixl_wrapper.get_xfer_descs(
1652 handle_data, self.nixl_memory_type
1653 )
1654 handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1632 - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
1655 + self.src_xfer_handles_by_tp_ratio[split_key].append(handle)
1656
1657 ### Register remote agent memory regions
1658 # With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
@@ -1665,13 +1688,6 @@ class NixlBaseConnectorWorker:
1688 self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
1689 )
1690
1668 - if block_size_ratio > 1:
1669 - # when prefill with smaller block_size, we need to init a
1670 - # new handler with same block_len to match
1671 - self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = (
1672 - self.register_local_xfer_handler(nixl_agent_meta.block_size)[0]
1673 - )
1674 -
1691 return remote_agent_name
1692
1693 def _validate_remote_agent_handshake(
@@ -1716,9 +1732,13 @@ class NixlBaseConnectorWorker:
1732 "Disable prefix caching with --no-enable-prefix-caching."
1733 )
1734
1719 - if self._is_hma_required:
1720 - assert block_size_ratio == 1, (
1721 - "HMA does not support different remote block size yet"
1735 + if block_size_ratio != 1:
1736 + # Heterogeneous block sizes transfer at remote-block granularity;
1737 + # the untransferred tail of the last local attention block is
1738 + # zeroed in the receive post-process, and mamba state pages
1739 + # transfer 1:1 (never sub-split).
1740 + assert not self.use_host_buffer, (
1741 + "Heterogeneous block sizes are not supported with host buffer"
1742 )
1743 kv_cache_layout = (
1744 self.kv_cache_layout
@@ -1875,27 +1895,57 @@ class NixlBaseConnectorWorker:
1895 "d2h",
1896 )
1897
1898 + @cached_property
1899 + def _attention_kv_caches(self) -> list[torch.Tensor]:
1900 + """Device KV caches of attention layers (mamba states excluded),
1901 + as consumed by the receive post-process."""
1902 + assert self.device_kv_caches, (
1903 + "_attention_kv_caches accessed before register_kv_caches"
1904 + )
1905 + mamba_layers = {
1906 + name
1907 + for g, group in enumerate(self.kv_cache_config.kv_cache_groups)
1908 + if _is_ssm_spec(self._group_spec_types[g])
1909 + for name in group.layer_names
1910 + }
1911 + kv_caches = self.device_kv_caches
1912 + return [cache for name, cache in kv_caches.items() if name not in mamba_layers]
1913 +
1914 def post_process_device_kv_on_receive(
1915 self,
1916 block_size_ratio: int,
1881 - block_ids_list: list[list[int]],
1917 + block_ids_list: list[tuple[list[int], int]],
1918 + convert: bool = True,
1919 ):
1920 """
1921 Post process device kv cache after receiving from remote.
1922
1886 - 3 types of post processing supported:
1923 + 3 types of conversion supported (``convert``):
1924 * kv_cache_postprocess_layout => convert from HND to NHD
1925 * kv_cache_postprocess_blksize => convert from small block size
1926 to large block size
1927 * kv_cache_postprocess_blksize_and_layout => convert from small
1928 block size to large block size and convert from HND to NHD
1929
1930 + The transfer only covers ``covered_sub_blocks`` remote-sized
1931 + sub-blocks of each request's local attention blocks; the rest was
1932 + clipped, either by remote-block pairing (block-size ratio) or by the
1933 + hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were
1934 + excluded from the scheduler's alloc-time KV zeroing (which would race
1935 + the RDMA write), so everything past the covered range is zeroed here.
1936 + Stale bytes would otherwise surface as garbage or NaNs once decode
1937 + grows into the untransferred tail.
1938 """
1939 if len(self.device_kv_caches) == 0:
1940 return
1941 assert block_size_ratio >= 1, "Only nP < nD supported currently."
1942 assert self.transfer_topo is not None
1898 - if self.enable_permute_local_kv and block_size_ratio > 1:
1943 + if not convert:
1944 + logger.debug(
1945 + "Post-processing device kv cache on receive by zeroing "
1946 + "untransferred blocks."
1947 + )
1948 + elif self.enable_permute_local_kv and block_size_ratio > 1:
1949 logger.debug(
1950 "Post-processing device kv cache on receive by converting "
1951 "block_size with %sx bigger and permuting layout from HND"
@@ -1914,18 +1964,45 @@ class NixlBaseConnectorWorker:
1964 block_size_ratio,
1965 )
1966
1917 - for block_ids in block_ids_list:
1918 - indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1967 + attn_caches = self._attention_kv_caches
1968 + device = attn_caches[0].device
1969 + for block_ids, covered_sub_blocks in block_ids_list:
1970 + # Blocks the transfer didn't write: the token tail of the last
1971 + # partially covered block, then everything beyond it.
1972 + covered_blocks, sub_blocks_in_last = divmod(
1973 + covered_sub_blocks, block_size_ratio
1974 + )
1975 + first_stale = covered_blocks + (1 if sub_blocks_in_last else 0)
1976 + has_stale = first_stale < len(block_ids)
1977 + indices = None
1978 + if convert or has_stale:
1979 + indices = async_tensor_h2d(block_ids, device, torch.long)
1980 +
1981 + if convert:
1982 + for cache in attn_caches:
1983 + if self.enable_permute_local_kv and block_size_ratio > 1:
1984 + kv_postprocess_blksize_and_layout_on_receive(
1985 + cache, indices, block_size_ratio
1986 + )
1987 + elif self.enable_permute_local_kv:
1988 + kv_postprocess_layout_on_receive(cache, indices)
1989 + else:
1990 + kv_postprocess_blksize_on_receive(
1991 + cache, indices, block_size_ratio
1992 + )
1993
1920 - for cache in self.device_kv_caches.values():
1921 - if self.enable_permute_local_kv and block_size_ratio > 1:
1922 - kv_postprocess_blksize_and_layout_on_receive(
1923 - cache, indices, block_size_ratio
1924 - )
1925 - elif self.enable_permute_local_kv:
1926 - kv_postprocess_layout_on_receive(cache, indices)
1927 - else:
1928 - kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio)
1994 + if sub_blocks_in_last:
1995 + last_block_id = block_ids[covered_blocks]
1996 + for cache in attn_caches:
1997 + # Both post-processed layouts leave tokens on dim 1.
1998 + sub_block_tokens = cache.shape[1] // block_size_ratio
1999 + zero_from = sub_blocks_in_last * sub_block_tokens
2000 + cache[last_block_id, zero_from:].zero_()
2001 + if has_stale:
2002 + assert indices is not None
2003 + stale_ids = indices[first_stale:]
2004 + for cache in attn_caches:
2005 + cache.index_fill_(0, stale_ids, 0)
2006
2007 def post_process_device_kv_on_receive_heterogeneous_attn(
2008 self, block_ids: list[int]
@@ -1995,18 +2072,33 @@ class NixlBaseConnectorWorker:
2072 if self.use_host_buffer:
2073 self.sync_recved_kv_to_device(req_id, meta)
2074
1998 - # post processing for heteroblocksize
2075 + # Post processing for heteroblocksize/layout, and for blocks the
2076 + # transfer clipped. The latter happens either at remote-block
2077 + # granularity (block_size_ratio > 1) or at kernel-block
2078 + # granularity, when equal kernel pages meet differing logical
2079 + # block sizes and _apply_prefix_caching front-trims to the
2080 + # minimum count (hybrid heterogeneous TP).
2081 remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
2082 block_size_ratio = self.transfer_topo.block_size_ratio(
2083 remote_info.remote_block_size
2084 )
2003 - if not self.use_mla and (
2004 - block_size_ratio > 1 or self.enable_permute_local_kv
2005 - ):
2006 - assert not self._is_hma_required
2007 - block_ids_for_blocksize_post_process[block_size_ratio].append(
2008 - meta.local_physical_block_ids[0]
2009 - )
2085 + hetero_ppl = (
2086 + remote_info.remote_physical_blocks_per_logical
2087 + != self._physical_blocks_per_logical_kv_block
2088 + )
2089 + if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl:
2090 + for g, local_group in enumerate(meta.local_physical_block_ids):
2091 + if not local_group or _is_ssm_spec(self._group_spec_types[g]):
2092 + continue
2093 + # Number of remote-sized sub-blocks the transfer covered;
2094 + # everything past this was clipped and must be zeroed.
2095 + covered_sub_blocks = min(
2096 + len(local_group) * block_size_ratio,
2097 + len(meta.remote.block_ids[g]),
2098 + )
2099 + block_ids_for_blocksize_post_process[block_size_ratio].append(
2100 + (local_group, covered_sub_blocks)
2101 + )
2102 # post processing for heterogeneous attention
2103 if self.enable_heterogeneous_attn_post_process:
2104 block_ids_for_heterogeneous_attn_post_process.append(
@@ -2016,7 +2108,14 @@ class NixlBaseConnectorWorker:
2108 block_size_ratio,
2109 block_ids_list,
2110 ) in block_ids_for_blocksize_post_process.items():
2019 - self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)
2111 + # MLA never needs the block-size/layout conversion, but its
2112 + # clipped blocks still need zeroing.
2113 + convert = not self.use_mla and (
2114 + block_size_ratio > 1 or self.enable_permute_local_kv
2115 + )
2116 + self.post_process_device_kv_on_receive(
2117 + block_size_ratio, block_ids_list, convert
2118 + )
2119
2120 for block_ids in block_ids_for_heterogeneous_attn_post_process:
2121 self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
@@ -2206,6 +2305,45 @@ class NixlBaseConnectorWorker:
2305
2306 return mapped_2d.flatten().astype(np.int64)
2307
2308 + def _map_block_ids_for_block_size_ratio(
2309 + self,
2310 + local_block_ids: BlockIds,
2311 + remote_block_ids: BlockIds,
2312 + block_size_ratio: int,
2313 + ) -> tuple[BlockIds, BlockIds]:
2314 + """Map attention-group block ids to remote-block granularity.
2315 +
2316 + Each local attention block is split into ``block_size_ratio``
2317 + sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the
2318 + remote list — the untransferred tail of the last local block — are
2319 + clipped here and zeroed in the receive post-process. Mamba state
2320 + blocks are indivisible and transfer 1:1, unexpanded.
2321 +
2322 + ex: remote (prefill) block ids with block_size 4:
2323 + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2324 + Local (decode) block ids with block_size 16: [1, 2, 3] expand to
2325 + [4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote.
2326 + """
2327 + mapped_local: list[list[int]] = []
2328 + mapped_remote: list[list[int]] = []
2329 + for i, remote_group in enumerate(remote_block_ids):
2330 + local_group = local_block_ids[i] if local_block_ids else []
2331 + if _is_ssm_spec(self._group_spec_types[i]):
2332 + mapped_local.append(list(local_group))
2333 + mapped_remote.append(list(remote_group))
2334 + continue
2335 + mapped = self.get_mapped_blocks(
2336 + np.asarray(local_group), block_size_ratio
2337 + ).tolist()
2338 + if len(mapped) > len(remote_group):
2339 + mapped = mapped[: len(remote_group)]
2340 + mapped_local.append(mapped)
2341 + mapped_remote.append(list(remote_group))
2342 + if not any(mapped_local):
2343 + # Full prefix cache hit is indicated with an empty list.
2344 + return [], mapped_remote
2345 + return mapped_local, mapped_remote
2346 +
2347 def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds:
2348 """
2349 Convert block ids to kernel physical block ids.
@@ -2300,13 +2438,18 @@ class NixlBaseConnectorWorker:
2438 remote_block_ids[i] = remote_group[-num_local_blocks:]
2439 else:
2440 # TODO Handle prefix caching with different block_sizes
2303 - max_padding = max(
2304 - self._physical_blocks_per_logical_kv_block,
2305 - remote_physical_per_logical,
2441 + # Allocation rounding legitimately leaves up to
2442 + # ppl - 1 trailing dead kernel blocks per side (plus one
2443 + # extra local block for the recomputed final token), so
2444 + # the counts may differ by up to the sum of the two
2445 + # ratios; anything larger indicates mismatched lists.
2446 + max_padding = (
2447 + self._physical_blocks_per_logical_kv_block
2448 + + remote_physical_per_logical
2449 )
2307 - assert abs(num_local_blocks - num_remote_blocks) < max_padding, (
2450 + assert abs(num_local_blocks - num_remote_blocks) <= max_padding, (
2451 f"Group {i}: |{num_local_blocks} - "
2309 - f"{num_remote_blocks}| >= {max_padding}"
2452 + f"{num_remote_blocks}| > {max_padding}"
2453 )
2454 num_blocks = min(num_local_blocks, num_remote_blocks)
2455 local_block_ids[i] = local_block_ids[i][:num_blocks]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+7 -28
@@ -5,8 +5,6 @@
5 import time
6 from typing import TYPE_CHECKING
7
8 -import numpy as np
9 -
8 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
9 NixlBaseConnectorWorker,
10 )
@@ -178,10 +176,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
176 )
177 # Get side handles.
178 if tp_ratio < 0 and (not self.use_mla or len(read_specs) > 1):
181 - assert remote_block_size == self.block_size
179 # Remote tp_size > local tp_size: we must perform multiple
180 # reads. Get the memory chunk onto which we will write to.
184 - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
181 + split_key = (tp_ratio, remote_block_size)
182 + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
183 else:
184 # Single read from remote, we write to the whole memory region.
185 # Also handle remote block size different from local block size.
@@ -235,30 +233,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
233 remote_info.remote_block_size
234 )
235 if block_size_ratio > 1:
238 - # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
239 - assert not self._is_hma_required
240 - local_block_ids0 = local_block_ids[0] if local_block_ids else []
241 - remote_block_ids0 = remote_block_ids[0]
242 - local_block_ids_mapped = self.get_mapped_blocks(
243 - np.asarray(local_block_ids0), block_size_ratio
244 - ).tolist()
245 - if len(local_block_ids_mapped) > len(remote_block_ids0):
246 - # NOTE:
247 - # get_mapped_blocks will always expand block_ids for n times.
248 - # ex:
249 - # prefill block_ids with block_size as 4:
250 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
251 - # Local decode block_ids with block_size as 16: [1, 2, 3]
252 - # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
253 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
254 - # Then we clip local to align with prefill
255 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
256 - # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
257 - local_block_ids_mapped = local_block_ids_mapped[
258 - : len(remote_block_ids0)
259 - ]
260 - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
261 - remote_block_ids = [remote_block_ids0]
236 + local_block_ids, remote_block_ids = (
237 + self._map_block_ids_for_block_size_ratio(
238 + local_block_ids, remote_block_ids, block_size_ratio
239 + )
240 + )
241 # NOTE(rob): having the staging blocks be on the READER side is
242 # not going to work well (since we will have to call rearrange tensors).
243 # after we detect the txn is complete (which means we cannot make the
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+7 -15
@@ -39,7 +39,6 @@ from concurrent.futures import Future
39 from typing import TYPE_CHECKING, Any
40
41 import msgspec
42 -import numpy as np
42
43 from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
44 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
@@ -553,8 +552,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
552 req_id,
553 )
554 if tp_ratio < 0 and not self.use_mla:
556 - assert remote_block_size == self.block_size
557 - local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
555 + split_key = (tp_ratio, remote_block_size)
556 + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
557 else:
558 local_xfer_side_handle = self.src_xfer_handles_by_block_size[
559 remote_block_size
@@ -606,18 +605,11 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
605 remote_info.remote_block_size
606 )
607 if block_size_ratio > 1:
609 - assert not self._is_hma_required
610 - local_block_ids0 = local_block_ids[0] if local_block_ids else []
611 - remote_block_ids0 = remote_block_ids[0]
612 - local_block_ids_mapped = self.get_mapped_blocks(
613 - np.asarray(local_block_ids0), block_size_ratio
614 - ).tolist()
615 - if len(local_block_ids_mapped) > len(remote_block_ids0):
616 - local_block_ids_mapped = local_block_ids_mapped[
617 - : len(remote_block_ids0)
618 - ]
619 - local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
620 - remote_block_ids = [remote_block_ids0]
608 + local_block_ids, remote_block_ids = (
609 + self._map_block_ids_for_block_size_ratio(
610 + local_block_ids, remote_block_ids, block_size_ratio
611 + )
612 + )
613
614 notif_id = f"{remote_request_id}:{self.world_size}".encode()
615