@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e0aaef85f3

[2/N][KV Connector][NIXL] Support per-region transfer geometry (#53780)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Matthew Bonanni committed Sep 8, 2026 at 16:01 UTC e0aaef85f3de7a2376aee3708a3cba96286b41b1
9 files changed +892 -110
tests/v1/kv_connector/unit/test_nixl_connector.py
+125 -4
@@ -518,8 +518,15 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
518 slot_size_bytes = 4096
519 self.slot_size_per_layer = [slot_size_bytes]
520 self.block_len_per_layer = [slot_size_bytes * self.block_size]
521 - self.num_blocks = 1
521 + self.num_regions = 1
522 + self.block_stride_per_layer = list(self.block_len_per_layer)
523 + self.region_num_blocks = [self.num_blocks]
524 + self.region_group_ids = [0]
525 + self.region_mem_types = [self.nixl_memory_type]
526 self.dst_num_blocks[self.engine_id] = self.num_blocks
527 + self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
528 + self.dst_region_group_ids[self.engine_id] = self.region_group_ids
529 + self.dst_region_mem_types[self.engine_id] = self.region_mem_types
530
531 assert expected_engine_id == self.REMOTE_ENGINE_ID
532
@@ -548,7 +555,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
555 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
556 kv_caches_base_addr=[0],
557 device_id=remote_tp_rank,
551 - num_blocks=1,
558 + num_blocks=self.num_blocks,
559 block_lens=remote_block_lens,
560 block_strides=remote_block_lens,
561 kv_cache_layout="LBHNC",
@@ -556,6 +563,9 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
563 ssm_sizes=(0, 0),
564 attn_backend_name=self.backend_name,
565 physical_blocks_per_logical_kv_block=1,
566 + region_num_blocks=self.region_num_blocks,
567 + region_group_ids=self.region_group_ids,
568 + region_mem_types=self.region_mem_types,
569 ),
570 remote_tp_rank=remote_tp_rank,
571 remote_tp_size=remote_tp_size,
@@ -640,7 +650,7 @@ class TestNixlHandshake:
650 request_id = "req_id"
651
652 # Test worker role in decode server.
643 - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
653 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10)
654 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
655 connector.connector_worker = FakeNixlConnectorWorker(
656 vllm_config,
@@ -1167,12 +1177,19 @@ class TestNixlHandshake:
1177 source_ranks_per_group=((0,), (0,)),
1178 rank_offset_factor=1,
1179 )
1170 - meta = MagicMock(
1180 + meta = NixlAgentMetadata(
1181 + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
1182 + agent_metadata=FakeNixlWrapper.AGENT_METADATA,
1183 kv_caches_base_addr=[0x1000],
1184 device_id=0,
1185 num_blocks=1,
1186 block_lens=[remote_block_len],
1187 block_strides=[remote_block_len],
1188 + kv_cache_layout="HND",
1189 + block_size=worker.block_size,
1190 + ssm_sizes=(0, 0),
1191 + attn_backend_name=worker.backend_name,
1192 + physical_blocks_per_logical_kv_block=1,
1193 )
1194
1195 assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1846,6 +1863,110 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1863 llm.llm_engine.engine_core.shutdown()
1864
1865
1866 +def test_mixed_memory_local_descriptors_split_by_memory_type():
1867 + worker = object.__new__(NixlConnectorWorker)
1868 + worker.transfer_topo = MagicMock()
1869 + worker.block_size = 16
1870 + worker.engine_id = "local"
1871 + worker.tp_rank = 0
1872 + worker.device_id = 3
1873 + worker.kv_caches_base_addr = {"local": {0: [100, 200]}}
1874 + worker._has_mamba = False
1875 + worker._mixed_mem_types = True
1876 + worker.region_mem_types = ["DRAM", "VRAM"]
1877 + worker.region_num_blocks = [2, 2]
1878 + worker._desc_is_dram_by_block_size = {}
1879 + worker._desc_pos_by_block_size = {}
1880 + worker._dram_src_handles_by_block_size = {}
1881 + worker.nixl_memory_type = "VRAM"
1882 + worker._build_fa_local = MagicMock( # type: ignore[method-assign]
1883 + return_value=np.array(
1884 + [
1885 + [100, 10, 3],
1886 + [110, 10, 3],
1887 + [200, 10, 3],
1888 + [210, 10, 3],
1889 + ]
1890 + )
1891 + )
1892 + worker.nixl_wrapper = MagicMock()
1893 + worker.nixl_wrapper.get_xfer_descs.side_effect = lambda blocks, memory_type: (
1894 + memory_type,
1895 + blocks,
1896 + )
1897 + worker.nixl_wrapper.prep_xfer_dlist.side_effect = [11, 22]
1898 +
1899 + handle, blocks = worker.register_local_xfer_handler(worker.block_size)
1900 +
1901 + assert handle == 22
1902 + assert worker._dram_src_handles_by_block_size[worker.block_size] == 11
1903 + assert [block[2] for block in blocks] == [0, 0, 3, 3]
1904 + assert [
1905 + call.args[1] for call in worker.nixl_wrapper.get_xfer_descs.call_args_list
1906 + ] == ["DRAM", "VRAM"]
1907 +
1908 +
1909 +def test_mixed_memory_read_notifies_after_both_transfers_finish():
1910 + worker = object.__new__(NixlConnectorWorker)
1911 + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}
1912 + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}
1913 + worker._dram_src_handles_by_block_size = {16: 10}
1914 + worker._recving_metadata = {"request": MagicMock()}
1915 + worker._recving_transfers = defaultdict(list)
1916 + worker._pending_recv_notifs = {}
1917 + worker.xfer_stats = MagicMock()
1918 + worker.nixl_wrapper = MagicMock()
1919 + worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]
1920 + worker.nixl_wrapper.check_xfer_state.return_value = "DONE"
1921 +
1922 + worker._read_blocks_mixed(
1923 + request_id="request",
1924 + local_block_size_key=16,
1925 + local_device_handle=20,
1926 + local_dram_handle=10,
1927 + remote_xfer_side_handle=30,
1928 + local_block_descs_ids=np.array([0, 2]),
1929 + remote_block_descs_ids=np.array([5, 7]),
1930 + notif_agent="prefill",
1931 + notif_id=b"request:1",
1932 + )
1933 +
1934 + dram_read, device_read = worker.nixl_wrapper.make_prepped_xfer.call_args_list
1935 + assert dram_read.args[:2] == ("READ", 10)
1936 + assert dram_read.args[3] == 30
1937 + np.testing.assert_array_equal(dram_read.args[2], [0])
1938 + np.testing.assert_array_equal(dram_read.args[4], [5])
1939 + assert device_read.args[:2] == ("READ", 20)
1940 + assert device_read.args[3] == 30
1941 + np.testing.assert_array_equal(device_read.args[2], [0])
1942 + np.testing.assert_array_equal(device_read.args[4], [7])
1943 + worker.nixl_wrapper.send_notif.assert_not_called()
1944 +
1945 + assert worker._pop_done_transfers(worker._recving_transfers) == {"request"}
1946 + worker.nixl_wrapper.send_notif.assert_called_once_with(
1947 + "prefill", notif_msg=b"request:1"
1948 + )
1949 +
1950 +
1951 +def test_mixed_memory_read_failure_does_not_notify_producer():
1952 + """A failed half of a split READ must suppress its success notification."""
1953 + worker = object.__new__(NixlConnectorWorker)
1954 + worker._recving_metadata = {"request": MagicMock()}
1955 + worker._recving_transfers = {"request": [101, 102]}
1956 + worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}
1957 + worker._is_hma_required = True
1958 + worker._failed_recv_reqs = MagicMock()
1959 + worker._log_failure = MagicMock() # type: ignore[method-assign]
1960 + worker.xfer_stats = MagicMock()
1961 + worker.nixl_wrapper = MagicMock()
1962 + worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]
1963 +
1964 + worker._pop_done_transfers(worker._recving_transfers)
1965 +
1966 + assert "request" not in worker._pending_recv_notifs
1967 + worker.nixl_wrapper.send_notif.assert_not_called()
1968 +
1969 +
1970 def element_byte_addrs(view: torch.Tensor) -> list[int]:
1971 """Absolute byte addresses of every element of a (possibly strided) view."""
1972 offsets = torch.zeros(view.shape, dtype=torch.int64)
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+59
@@ -861,6 +861,7 @@ def _make_mock_worker_for_desc_ids(
861
862 worker = MagicMock(spec=NixlConnectorWorker)
863 worker.num_regions = num_regions
864 + worker._uses_region_group_mapping = False
865 worker._has_mamba = has_mamba
866 worker._group_spec_types = group_spec_types
867 worker.block_len_per_layer = block_len_per_layer or [100]
@@ -913,6 +914,64 @@ def test_get_block_descs_ids_hybrid_ssm():
914 assert list(result) == expected, f"Expected {expected}, got {list(result)}"
915
916
917 +@pytest.mark.cpu_test
918 +def test_get_block_descs_ids_selects_attention_regions_by_group():
919 + """Each attention group's blocks address only that group's regions."""
920 + from vllm.v1.kv_cache_interface import FullAttentionSpec
921 +
922 + worker = _make_mock_worker_for_desc_ids(
923 + num_regions=3,
924 + has_mamba=False,
925 + group_spec_types=(FullAttentionSpec, FullAttentionSpec),
926 + block_len_per_layer=[100, 100, 100],
927 + )
928 + worker.region_group_ids = [0, 0, 1]
929 + worker._uses_region_group_mapping = True
930 +
931 + result = worker._compute_desc_ids(
932 + block_ids=([1, 2], [7]),
933 + dst_num_blocks=10,
934 + block_size_ratio=None,
935 + physical_blocks_per_logical=1,
936 + )
937 +
938 + assert result.tolist() == [1, 2, 11, 12, 27]
939 +
940 + remapped = worker._compute_desc_ids(
941 + block_ids=([1, 2], [7]),
942 + dst_num_blocks=10,
943 + block_size_ratio=None,
944 + physical_blocks_per_logical=1,
945 + region_group_ids=[1, 1, 0],
946 + )
947 + assert remapped.tolist() == [21, 22, 7, 17]
948 +
949 +
950 +@pytest.mark.cpu_test
951 +def test_get_block_descs_ids_uses_per_region_pool_capacity():
952 + """Independent host and device pools use cumulative descriptor offsets."""
953 + from vllm.v1.kv_cache_interface import FullAttentionSpec
954 +
955 + worker = _make_mock_worker_for_desc_ids(
956 + num_regions=2,
957 + has_mamba=False,
958 + group_spec_types=(FullAttentionSpec, FullAttentionSpec),
959 + block_len_per_layer=[100, 100],
960 + )
961 + worker.region_group_ids = [0, 1]
962 + worker._uses_region_group_mapping = True
963 +
964 + result = worker._compute_desc_ids(
965 + block_ids=([4], [8]),
966 + dst_num_blocks=10,
967 + block_size_ratio=None,
968 + physical_blocks_per_logical=1,
969 + region_num_blocks=[5, 10],
970 + )
971 +
972 + assert result.tolist() == [4, 13]
973 +
974 +
975 @pytest.mark.cpu_test
976 def test_get_block_descs_ids_kernel_block_mismatch():
977 """Test _compute_desc_ids uses different strides for FA
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+148 -2
@@ -11,7 +11,8 @@ 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
14 +from collections import defaultdict
15 +from unittest.mock import MagicMock, patch
16
17 import numpy as np
18 import pytest
@@ -27,13 +28,14 @@ class _RecordingNixl:
28 def __init__(self, *args, **kwargs):
29 self.dlists: dict[int, np.ndarray] = {}
30 self.xfers: list[tuple] = []
31 + self.registered: list[tuple[list[tuple], object]] = []
32 self._next_handle = 1
33
34 def get_reg_descs(self, caches_data, mem_type):
35 return caches_data
36
37 def register_memory(self, descs, backends=None):
36 - pass
38 + self.registered.append((descs, backends))
39
40 def deregister_memory(self, descs):
41 pass
@@ -98,6 +100,150 @@ class _RecordingNixl:
100 pass
101
102
103 +@pytest.mark.cpu_test
104 +def test_local_descriptors_follow_each_region_pool_capacity():
105 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
106 + NixlConnectorWorker,
107 + )
108 +
109 + worker = object.__new__(NixlConnectorWorker)
110 + worker.transfer_topo = MagicMock()
111 + worker.device_id = 0
112 + worker.block_len_per_layer = [16, 16]
113 + worker.block_stride_per_layer = [16, 16]
114 + worker.region_num_blocks = [2, 3]
115 +
116 + descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1)
117 +
118 + assert descriptors[:, 0].tolist() == [100, 116, 1000, 1016, 1032]
119 +
120 +
121 +@pytest.mark.cpu_test
122 +def test_overlaid_transfer_groups_share_region_geometry():
123 + """Groups overlaid on one allocation share its transfer region."""
124 + import msgspec
125 +
126 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
127 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
128 + NixlAgentMetadata,
129 + )
130 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
131 + NixlConnectorWorker,
132 + )
133 + from vllm.v1.kv_cache_interface import (
134 + KVCacheConfig,
135 + KVCacheGroupSpec,
136 + KVCacheTensor,
137 + MLAAttentionSpec,
138 + )
139 +
140 + num_blocks = 4
141 + spec = MLAAttentionSpec(
142 + block_size=4,
143 + num_kv_heads=1,
144 + head_size=8,
145 + dtype=torch.uint8,
146 + )
147 + page_size = spec.page_size_bytes
148 + block_stride = 2 * page_size
149 + backing = torch.zeros(num_blocks, block_stride, dtype=torch.uint8)
150 + caches = {
151 + "layer.0": backing[:, :page_size],
152 + "layer.1": backing[:, :page_size],
153 + }
154 + groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches]
155 +
156 + worker = object.__new__(NixlConnectorWorker)
157 + worker.tp_rank = 0
158 + worker.world_size = 1
159 + worker.block_size = 4
160 + worker.engine_id = "local-engine"
161 + worker.use_mla = True
162 + worker.model_config = MagicMock()
163 + worker.model_config.get_total_num_kv_heads.return_value = 1
164 + worker.attn_backends = []
165 + worker._has_mamba = False
166 + worker._is_csa_linear = False
167 + worker.vllm_config = MagicMock()
168 + worker.backend_name = "FLASHMLA"
169 + worker.num_blocks = num_blocks
170 + worker.nixl_memory_type = "VRAM"
171 + worker.nixl_backends = None
172 + worker.nixl_wrapper = _RecordingNixl()
173 + worker._registered_descs = []
174 + worker.dst_num_blocks = {}
175 + worker.dst_region_num_blocks = {}
176 + worker.dst_region_group_ids = {}
177 + worker.dst_uses_region_group_mapping = {}
178 + worker.dst_region_mem_types = {}
179 + worker._desc_is_dram_by_block_size = {}
180 + worker._desc_pos_by_block_size = {}
181 + worker._dram_src_handles_by_block_size = {}
182 + worker.src_xfer_handles_by_block_size = {}
183 + worker.kv_caches_base_addr = defaultdict(dict)
184 + worker._mamba_ssm_size = (0, 0)
185 + worker.kv_cache_layout = "NHD"
186 + worker.host_buffer_kv_cache_layout = "NHD"
187 + worker._physical_blocks_per_logical_kv_block = 1
188 + worker._logical_num_blocks = num_blocks
189 + worker.region_group_ids = []
190 + worker.region_mem_types = []
191 + worker._mixed_mem_types = False
192 + worker.region_names = []
193 + worker.region_num_blocks = []
194 + worker._region_is_mla = []
195 + worker.block_len_per_layer = []
196 + worker.block_stride_per_layer = []
197 + worker.device_id = 0
198 + worker.use_host_buffer = False
199 + worker.host_xfer_buffers = {}
200 + worker.device_kv_caches = {}
201 + worker.pp_size = 1
202 + worker.dcp_size = 1
203 + worker.pcp_size = 1
204 + worker.kv_buffer_device = "cuda"
205 + worker._layer_specs = {name: spec for name in caches}
206 + worker.kv_cache_config = KVCacheConfig(
207 + num_blocks=num_blocks,
208 + kv_cache_tensors=[
209 + KVCacheTensor(
210 + size=backing.nbytes,
211 + layers=[name],
212 + layer_stride=page_size,
213 + block_stride=block_stride,
214 + )
215 + for name in caches
216 + ],
217 + kv_cache_groups=groups,
218 + )
219 +
220 + transfer_topology = MagicMock()
221 +
222 + with (
223 + patch.object(bw, "TransferTopology", return_value=transfer_topology),
224 + patch.object(bw, "compute_nixl_compatibility_hash", return_value="hash"),
225 + ):
226 + worker.register_kv_caches(caches)
227 +
228 + assert worker.region_group_ids == [-1]
229 + assert worker.block_stride_per_layer == [block_stride]
230 + assert worker.nixl_wrapper.registered[0][0] == [
231 + (backing.data_ptr(), backing.nbytes, 0, "")
232 + ]
233 + expected_addrs = [
234 + backing.data_ptr() + block * block_stride for block in range(num_blocks)
235 + ]
236 + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs
237 +
238 + metadata = msgspec.msgpack.decode(
239 + worker.xfer_handshake_metadata.agent_metadata_bytes,
240 + type=NixlAgentMetadata,
241 + )
242 + assert metadata.region_group_ids == [-1]
243 + assert metadata.region_num_blocks == [num_blocks]
244 + assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]
245 +
246 +
247 def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
248 """Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
249 from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+15
@@ -357,6 +357,12 @@ class _StubWriterWorker(NixlPushConnectorWorker):
357 w.engine_id = "test-decode-engine"
358 w._remote_agents = {}
359 w._physical_blocks_per_logical_kv_block = 1
360 + w._uses_region_group_mapping = False
361 + w.region_group_ids = [0]
362 + w.dst_region_num_blocks = {}
363 + w.dst_region_group_ids = {}
364 + w.dst_uses_region_group_mapping = {}
365 + w.dst_region_mem_types = {}
366 # Single non-hybrid attention group, matching the stub block id lists.
367 w._has_mamba = False
368 w._is_csa_linear = False
@@ -1314,6 +1320,15 @@ class TestPushPrefixCaching:
1320 )
1321 }
1322 w.dst_num_blocks = {engine_id: 10_000, w.engine_id: 10_000}
1323 + w.dst_region_num_blocks = {
1324 + engine_id: [10_000],
1325 + w.engine_id: [10_000],
1326 + }
1327 + w.dst_region_group_ids = {engine_id: [0], w.engine_id: [0]}
1328 + w.dst_uses_region_group_mapping = {
1329 + engine_id: False,
1330 + w.engine_id: False,
1331 + }
1332 w.dst_xfer_side_handles = {engine_id: {0: 5000}}
1333 w.src_xfer_handles_by_block_size = {16: 2000}
1334
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+321 -36
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, cast
17
18 import msgspec
19 import numpy as np
20 +import regex as re
21 import torch
22 import zmq
23
@@ -92,6 +93,16 @@ if TYPE_CHECKING:
93
94 logger = init_logger(__name__)
95
96 +_SHARED_REGION_GROUP_ID = -1
97 +
98 +
99 +def _region_sort_key(layer_name: str) -> tuple[tuple[int, int | str], ...]:
100 + """Sort transfer regions in model-layer order, then by cache name."""
101 + return tuple(
102 + (0, int(part)) if part.isdigit() else (1, part)
103 + for part in re.split(r"(\d+)", layer_name)
104 + )
105 +
106
107 def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
108 """Return whether all views share storage and a block stride."""
@@ -154,6 +165,9 @@ class NixlBaseConnectorWorker:
165 dst_num_blocks: int,
166 block_size_ratio: float | None,
167 physical_blocks_per_logical: int,
168 + region_num_blocks: list[int] | None = None,
169 + region_group_ids: list[int] | None = None,
170 + uses_region_group_mapping: bool | None = None,
171 ) -> np.ndarray:
172 """Compute NIXL descriptor IDs for given block IDs."""
173 num_ssm_regions = 0
@@ -172,21 +186,51 @@ class NixlBaseConnectorWorker:
186 num_blocks = dst_num_blocks
187 if block_size_ratio is not None:
188 num_blocks = int(num_blocks * block_size_ratio)
175 - num_fa_descs = self.num_regions * num_blocks
189 + if region_num_blocks is None:
190 + region_num_blocks = [num_blocks] * self.num_regions
191 + elif block_size_ratio is not None:
192 + region_num_blocks = [
193 + int(count * block_size_ratio) for count in region_num_blocks
194 + ]
195 + assert len(region_num_blocks) == self.num_regions
196 + region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
197 + num_fa_descs = sum(region_num_blocks)
198
199 # All-attention fast path: single vectorized broadcast.
200 if num_ssm_regions == 0:
201 + if region_group_ids is None:
202 + region_group_ids = self.region_group_ids
203 + if uses_region_group_mapping is None:
204 + uses_region_group_mapping = self._uses_region_group_mapping
205 + assert len(region_group_ids) == self.num_regions
206 + if uses_region_group_mapping is None:
207 + uses_region_group_mapping = len(set(region_group_ids)) > 1
208 + region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
209 # NOTE (NickLucche) With HMA, every kv group has the same number of layers
210 # and layers from different groups share the same kv tensor.
211 # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
212 # read across all regions, same for [3], but group0-group1 blocks will
213 # always differ (different areas). Therefore we can just flatten the
214 # block_ids and compute the descs ids for all groups at once.
185 - block_arr = np.concatenate(
186 - [np.asarray(g, dtype=np.int32) for g in block_ids]
187 - )[None, :]
188 - region_ids = np.arange(self.num_regions, dtype=np.int32)[:, None]
189 - return (region_ids * num_blocks + block_arr).ravel()
215 + if not uses_region_group_mapping:
216 + block_arr = np.concatenate(
217 + [np.asarray(group, dtype=np.int32) for group in block_ids]
218 + )[None, :]
219 + return (region_offsets[:, None] + block_arr).ravel()
220 + desc_ids = []
221 + for group_id, group in enumerate(block_ids):
222 + if not group:
223 + continue
224 + region_ids = np.flatnonzero(
225 + (region_group_ids_array == group_id)
226 + | (region_group_ids_array == _SHARED_REGION_GROUP_ID)
227 + )[:, None]
228 + assert region_ids.size > 0
229 + group_blocks = np.asarray(group, dtype=np.int32)[None, :]
230 + desc_ids.append((region_offsets[region_ids] + group_blocks).ravel())
231 + if not desc_ids:
232 + return np.array([], dtype=np.int64)
233 + return np.concatenate(desc_ids)
234
235 # Compute desc ids per group using the right stride: FA descs have
236 # num_blocks entries per region (kernel granularity, expanded by
@@ -556,6 +600,12 @@ class NixlBaseConnectorWorker:
600 # Number of NIXL regions. Currently one region per cache
601 # (so 1 per layer for MLA, otherwise 2 per layer)
602 self.num_regions = 0
603 + self.region_group_ids: list[int] = []
604 + self._uses_region_group_mapping = False
605 + self.region_names: list[str] = []
606 + self.region_num_blocks: list[int] = []
607 + self.region_mem_types: list[str] = []
608 + self._mixed_mem_types = False
609
610 # PP>1 (push mode): this worker holds a contiguous layer slice and
611 # transfers into the matching sub-range of a PP=1 remote's regions.
@@ -584,12 +634,20 @@ class NixlBaseConnectorWorker:
634 # Populated dynamically during handshake based on remote configuration.
635 # Per-source split handles, keyed by (tp_ratio, remote_block_size).
636 self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
637 + self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
638 # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
639 self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
640
641 # Map of engine_id -> num_blocks. All ranks in the same deployment will
642 # have the same number of blocks.
643 self.dst_num_blocks: dict[EngineId, int] = {}
644 + self.dst_region_num_blocks: dict[EngineId, list[int]] = {}
645 + self.dst_region_group_ids: dict[EngineId, list[int]] = {}
646 + self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
647 + self.dst_region_mem_types: dict[EngineId, list[str]] = {}
648 + self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
649 + self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
650 + self._dram_src_handles_by_block_size: dict[int, int] = {}
651 self._registered_descs: list[Any] = []
652
653 # In progress transfers.
@@ -607,6 +665,7 @@ class NixlBaseConnectorWorker:
665 # Uses Queue for thread-safe cross-thread coordination with the
666 # background handshake thread, matching the _ready_requests pattern.
667 self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
668 + self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
669
670 # Handshake metadata of this worker for NIXL transfers.
671 self.xfer_handshake_metadata: NixlHandshakePayload | None = None
@@ -1179,8 +1238,8 @@ class NixlBaseConnectorWorker:
1238 self.use_host_buffer,
1239 )
1240
1182 - caches_data = []
1183 - seen_storage_addresses: set[int] = set()
1241 + registration_ranges: dict[tuple[int, str], tuple[int, int, int, str]] = {}
1242 + region_mem_types: list[str] = []
1243 seen_base_addresses: list[int] = []
1244 self._ssm_region_indices = []
1245 self._scratch_region_indices = []
@@ -1219,7 +1278,13 @@ class NixlBaseConnectorWorker:
1278 # K and V are packed into the content dim, so each attention layer is a
1279 # single NIXL region whose block transfers as one unit. Mamba layers instead
1280 # register separate conv/ssm sub-regions (see `_build_mamba_local`).
1222 - for layer_name, cache in xfer_buffers.items():
1281 + layer_names = (
1282 + xfer_buffers
1283 + if self._is_csa_linear
1284 + else sorted(xfer_buffers, key=_region_sort_key)
1285 + )
1286 + for layer_name in layer_names:
1287 + cache = xfer_buffers[layer_name]
1288 # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
1289 # kernel requires a specific block size. This leads to SSM and FA layers
1290 # having different num_blocks.
@@ -1240,6 +1305,7 @@ class NixlBaseConnectorWorker:
1305 else layer_spec.page_size_bytes
1306 // self._physical_blocks_per_logical_kv_block
1307 )
1308 + group_id = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1309 num_blocks = (
1310 self._logical_num_blocks
1311 if isinstance(layer_spec, MambaSpec)
@@ -1277,13 +1343,17 @@ class NixlBaseConnectorWorker:
1343 )
1344 continue
1345
1280 - # Memory registration follows allocations, while transfer regions follow
1281 - # logical layers (or contiguous head segments). This keeps strided
1282 - # cross-layer views inside their registered allocation.
1283 - if storage_addr not in seen_storage_addresses:
1284 - seen_storage_addresses.add(storage_addr)
1285 - self.device_id = max(cache.get_device(), 0)
1286 - caches_data.append((storage_addr, storage.nbytes(), self.device_id, ""))
1346 + if cache.device.type == "cpu":
1347 + mem_type = "DRAM"
1348 + region_device_id = 0
1349 + else:
1350 + mem_type = self.nixl_memory_type
1351 + region_device_id = max(cache.get_device(), 0)
1352 + self.device_id = region_device_id
1353 + registration_ranges.setdefault(
1354 + (storage_addr, mem_type),
1355 + (storage_addr, storage.nbytes(), region_device_id, ""),
1356 + )
1357
1358 is_mla_region = isinstance(
1359 layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
@@ -1366,12 +1436,19 @@ class NixlBaseConnectorWorker:
1436 for base_addr, block_len, block_stride in region_specs:
1437 if base_addr in seen_base_addresses:
1438 region_index = seen_base_addresses.index(base_addr)
1439 + assert region_mem_types[region_index] == mem_type
1440 self._region_is_mla[region_index] |= is_mla_region
1441 + if self.region_group_ids[region_index] != group_id:
1442 + self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
1443 else:
1444 region_index = len(seen_base_addresses)
1445 seen_base_addresses.append(base_addr)
1446 self.block_len_per_layer.append(block_len)
1447 self.block_stride_per_layer.append(block_stride)
1448 + self.region_group_ids.append(group_id)
1449 + self.region_names.append(layer_name)
1450 + self.region_num_blocks.append(num_blocks)
1451 + region_mem_types.append(mem_type)
1452 self._region_is_mla.append(is_mla_region)
1453
1454 if isinstance(layer_spec, MambaSpec):
@@ -1391,10 +1468,10 @@ class NixlBaseConnectorWorker:
1468 # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
1469 if (
1470 self._physical_blocks_per_logical_kv_block == 1
1394 - and cache.shape[0] != num_blocks
1471 + and cache.shape[0] < num_blocks
1472 ):
1473 raise AssertionError(
1397 - "All kv cache tensors must have the same number of "
1474 + "KV cache tensor has too few "
1475 f"blocks; layer={layer_name}, "
1476 f"expected_num_blocks={num_blocks}, "
1477 f"cache_shape={tuple(cache.shape)}, "
@@ -1414,12 +1491,21 @@ class NixlBaseConnectorWorker:
1491 == len(seen_base_addresses)
1492 == len(self._region_is_mla)
1493 == len(self.block_stride_per_layer)
1494 + == len(self.region_group_ids)
1495 + == len(self.region_names)
1496 + == len(self.region_num_blocks)
1497 + == len(region_mem_types)
1498 )
1499 # Descriptor ids must be region-ordered, matching the remote side.
1500 self._scratch_region_indices.sort()
1501
1502 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
1503 self.num_regions = len(seen_base_addresses)
1504 + self.region_mem_types = region_mem_types
1505 + self._uses_region_group_mapping = len(set(self.region_group_ids)) > 1
1506 + self._mixed_mem_types = len(set(region_mem_types)) > 1
1507 + if self._has_mamba:
1508 + self.region_num_blocks = [self.num_blocks] * self.num_regions
1509
1510 if self.pp_size > 1:
1511 start_layer, end_layer = self.model_config.get_layers_start_end_indices(
@@ -1431,16 +1517,29 @@ class NixlBaseConnectorWorker:
1517 self._remote_region_offset = regions_per_layer * start_layer
1518
1519 # Total local FA descriptors (boundary between FA and mamba descs).
1434 - self.num_descs = self.num_regions * self.num_blocks
1435 -
1436 - descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
1437 - logger.debug("Registering descs: %s", caches_data)
1438 - self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1439 - logger.debug("Done registering descs")
1440 - self._registered_descs.append(descs)
1520 + self.num_descs = sum(self.region_num_blocks)
1521 +
1522 + for mem_type in sorted(set(region_mem_types)):
1523 + ranges = [
1524 + registration
1525 + for (
1526 + _,
1527 + registration_mem_type,
1528 + ), registration in registration_ranges.items()
1529 + if registration_mem_type == mem_type
1530 + ]
1531 + descs = self.nixl_wrapper.get_reg_descs(ranges, mem_type)
1532 + self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
1533 + self._registered_descs.append(descs)
1534
1535 self.device_kv_caches = kv_caches
1536 self.dst_num_blocks[self.engine_id] = self.num_blocks
1537 + self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks
1538 + self.dst_region_group_ids[self.engine_id] = self.region_group_ids
1539 + self.dst_uses_region_group_mapping[self.engine_id] = (
1540 + self._uses_region_group_mapping
1541 + )
1542 + self.dst_region_mem_types[self.engine_id] = self.region_mem_types
1543
1544 if self._has_mamba:
1545 logger.info(
@@ -1479,6 +1578,10 @@ class NixlBaseConnectorWorker:
1578 physical_blocks_per_logical_kv_block=(
1579 self._physical_blocks_per_logical_kv_block
1580 ),
1581 + region_num_blocks=self.region_num_blocks,
1582 + region_group_ids=self.region_group_ids,
1583 + region_names=self.region_names,
1584 + region_mem_types=self.region_mem_types,
1585 dcp_size=self.dcp_size,
1586 pcp_size=self.pcp_size,
1587 )
@@ -1642,14 +1745,23 @@ class NixlBaseConnectorWorker:
1745 """Build local FA descriptors for all layers as an Nx3 uint64 array."""
1746 assert self.transfer_topo is not None
1747 assert base_addresses, "Local KV cache base addresses must not be empty."
1645 - num_blocks = self.num_blocks * block_size_ratio
1748 device_id = self.device_id
1647 - block_arange = np.arange(num_blocks, dtype=np.uint64)
1749 parts: list[np.ndarray] = []
1750 for i, base_addr in enumerate(base_addresses):
1751 block_len = self.block_len_per_layer[i] // block_size_ratio
1651 - block_stride = self.block_stride_per_layer[i] // block_size_ratio
1652 - addrs = base_addr + block_arange * block_stride
1752 + logical_blocks = np.repeat(
1753 + np.arange(self.region_num_blocks[i], dtype=np.uint64),
1754 + block_size_ratio,
1755 + )
1756 + split_offsets = np.tile(
1757 + np.arange(block_size_ratio, dtype=np.uint64),
1758 + self.region_num_blocks[i],
1759 + )
1760 + addrs = (
1761 + base_addr
1762 + + logical_blocks * self.block_stride_per_layer[i]
1763 + + split_offsets * block_len
1764 + )
1765 parts.append(self._stack_descs(addrs, block_len, device_id))
1766 return np.concatenate(parts)
1767
@@ -1670,9 +1782,11 @@ class NixlBaseConnectorWorker:
1782 # SPLIT regions read their head slice from this many remote ranks at a
1783 # per-rank offset; REPLICATE regions read the whole block once.
1784 split_reads = len(plan.source_ranks_per_group[fa_group_idx])
1673 - num_blocks = nixl_agent_meta.num_blocks
1785 + region_num_blocks = nixl_agent_meta.region_num_blocks or [
1786 + nixl_agent_meta.num_blocks
1787 + ] * len(nixl_agent_meta.kv_caches_base_addr)
1788 device_id = nixl_agent_meta.device_id
1675 - block_arange = np.arange(num_blocks, dtype=np.uint64)
1789 + block_strides = nixl_agent_meta.block_strides
1790 parts: list[np.ndarray] = []
1791 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1792 replicated = self._is_region_replicated(i)
@@ -1691,8 +1805,8 @@ class NixlBaseConnectorWorker:
1805 )
1806 local_block_len = local_block_len // num_reads
1807
1694 - page_size = nixl_agent_meta.block_strides[i]
1695 - addrs = base_addr + rank_offset + block_arange * page_size
1808 + block_arange = np.arange(region_num_blocks[i], dtype=np.uint64)
1809 + addrs = base_addr + rank_offset + block_arange * block_strides[i]
1810 parts.append(self._stack_descs(addrs, local_block_len, device_id))
1811 return np.concatenate(parts)
1812
@@ -1735,7 +1849,42 @@ class NixlBaseConnectorWorker:
1849 mamba = self._build_mamba_local(local_base_addresses)
1850 blocks_data = np.concatenate([blocks_data, mamba])
1851
1738 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
1852 + if self._mixed_mem_types:
1853 + desc_is_dram = np.concatenate(
1854 + [
1855 + np.full(count * block_size_ratio, mem_type == "DRAM", dtype=bool)
1856 + for mem_type, count in zip(
1857 + self.region_mem_types,
1858 + self.region_num_blocks,
1859 + strict=True,
1860 + )
1861 + ]
1862 + )
1863 + assert len(desc_is_dram) == len(blocks_data)
1864 + desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
1865 + dram_indices = np.flatnonzero(desc_is_dram)
1866 + device_indices = np.flatnonzero(~desc_is_dram)
1867 + desc_pos[dram_indices] = np.arange(len(dram_indices), dtype=np.int64)
1868 + desc_pos[device_indices] = np.arange(len(device_indices), dtype=np.int64)
1869 + self._desc_is_dram_by_block_size[block_size] = desc_is_dram
1870 + self._desc_pos_by_block_size[block_size] = desc_pos
1871 + blocks_data[desc_is_dram, 2] = 0
1872 + dram_descs = self.nixl_wrapper.get_xfer_descs(
1873 + blocks_data[desc_is_dram], "DRAM"
1874 + )
1875 + self._dram_src_handles_by_block_size[block_size] = (
1876 + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", dram_descs)
1877 + )
1878 + device_descs = self.nixl_wrapper.get_xfer_descs(
1879 + blocks_data[~desc_is_dram], self.nixl_memory_type
1880 + )
1881 + return (
1882 + self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", device_descs),
1883 + blocks_data,
1884 + )
1885 +
1886 + mem_type = self.region_mem_types[0]
1887 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, mem_type)
1888 # NIXL_INIT_AGENT to be used for preparations of local descs.
1889 return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data
1890
@@ -1817,6 +1966,20 @@ class NixlBaseConnectorWorker:
1966 ]
1967 nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
1968 nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
1969 + if nixl_agent_meta.region_num_blocks is not None:
1970 + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
1971 + start:end
1972 + ]
1973 + if nixl_agent_meta.region_group_ids is not None:
1974 + nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
1975 + start:end
1976 + ]
1977 + if nixl_agent_meta.region_names is not None:
1978 + nixl_agent_meta.region_names = nixl_agent_meta.region_names[start:end]
1979 + if nixl_agent_meta.region_mem_types is not None:
1980 + nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
1981 + start:end
1982 + ]
1983
1984 ### Register remote engine in TransferTopology (idempotent).
1985 assert self.transfer_topo is not None
@@ -1855,7 +2018,24 @@ class NixlBaseConnectorWorker:
2018 block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
2019
2020 if engine_id not in self.dst_num_blocks:
2021 + num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
2022 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
2023 + self.dst_region_num_blocks[engine_id] = (
2024 + nixl_agent_meta.region_num_blocks
2025 + or [nixl_agent_meta.num_blocks] * num_remote_regions
2026 + )
2027 + self.dst_region_group_ids[engine_id] = nixl_agent_meta.region_group_ids or (
2028 + self.region_group_ids
2029 + if len(self.region_group_ids) == num_remote_regions
2030 + else [0] * num_remote_regions
2031 + )
2032 + self.dst_uses_region_group_mapping[engine_id] = (
2033 + len(set(self.dst_region_group_ids[engine_id])) > 1
2034 + )
2035 + self.dst_region_mem_types[engine_id] = (
2036 + nixl_agent_meta.region_mem_types
2037 + or [self.nixl_memory_type] * num_remote_regions
2038 + )
2039
2040 # Keep track of remote agent kv caches base addresses.
2041 self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
@@ -1902,6 +2082,8 @@ class NixlBaseConnectorWorker:
2082 # while the SSM state is sharded across every remote TP rank.
2083 # We only do this once per remote (tp_size, block_size).
2084 self.src_xfer_handles_by_tp_ratio[split_key] = []
2085 + if self._mixed_mem_types:
2086 + self._dram_src_handles_by_tp_ratio[split_key] = []
2087
2088 for handle_data in self._build_local_splits_from_plan(
2089 plan,
@@ -1909,6 +2091,17 @@ class NixlBaseConnectorWorker:
2091 self.num_descs * block_size_ratio,
2092 block_size_ratio,
2093 ):
2094 + if self._mixed_mem_types:
2095 + handle_data = np.asarray(handle_data, dtype=np.uint64)
2096 + desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size]
2097 + dram_descs = self.nixl_wrapper.get_xfer_descs(
2098 + handle_data[desc_is_dram], "DRAM"
2099 + )
2100 + dram_handle = self.nixl_wrapper.prep_xfer_dlist(
2101 + "NIXL_INIT_AGENT", dram_descs
2102 + )
2103 + self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle)
2104 + handle_data = handle_data[~desc_is_dram]
2105 descs = self.nixl_wrapper.get_xfer_descs(
2106 handle_data, self.nixl_memory_type
2107 )
@@ -1943,8 +2136,14 @@ class NixlBaseConnectorWorker:
2136 mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
2137 blocks_data = np.concatenate([blocks_data, mamba])
2138
1946 - # Register with NIXL.
1947 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
2139 + # A mixed local destination is split into multiple READs below. The
2140 + # producer remains a single descriptor list in this prerequisite.
2141 + remote_mem_types = set(self.dst_region_mem_types[engine_id])
2142 + if len(remote_mem_types) != 1:
2143 + raise NotImplementedError(
2144 + "NIXL pull does not yet support a mixed-memory producer"
2145 + )
2146 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, remote_mem_types.pop())
2147 self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
2148 self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)
2149 )
@@ -2133,6 +2332,27 @@ class NixlBaseConnectorWorker:
2332 assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
2333 # Same number of regions/~layers.
2334 assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)
2335 + num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
2336 + if nixl_agent_meta.region_num_blocks is not None:
2337 + assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions
2338 + assert len(nixl_agent_meta.block_strides) == num_remote_regions
2339 + assert all(
2340 + stride >= block_len
2341 + for stride, block_len in zip(
2342 + nixl_agent_meta.block_strides,
2343 + nixl_agent_meta.block_lens,
2344 + strict=True,
2345 + )
2346 + )
2347 + if nixl_agent_meta.region_group_ids is not None:
2348 + assert len(nixl_agent_meta.region_group_ids) == num_remote_regions
2349 + if nixl_agent_meta.region_names is not None:
2350 + assert nixl_agent_meta.region_names == self.region_names, (
2351 + "NIXL transfer regions must name the same model caches in the "
2352 + "same order"
2353 + )
2354 + if nixl_agent_meta.region_mem_types is not None:
2355 + assert len(nixl_agent_meta.region_mem_types) == num_remote_regions
2356
2357 def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
2358 """copy recved kv from host buffer to device."""
@@ -2526,6 +2746,7 @@ class NixlBaseConnectorWorker:
2746 # don't report it again, just drop the remaining handles.
2747 if req_id in self._recving_metadata:
2748 done_req_ids.add(req_id)
2749 + self._send_pending_recv_notifs(req_id)
2750 del transfers[req_id]
2751 else:
2752 transfers[req_id] = in_progress
@@ -2545,6 +2766,8 @@ class NixlBaseConnectorWorker:
2766 # reported failed on the first one, which pops its metadata in
2767 # get_finished(); on the later failures only the handle cleanup is left.
2768 # TODO (NickLucche) handle failed transfer for HMA.
2769 + # A split READ's notification is sent only after every handle succeeds.
2770 + self._pending_recv_notifs.pop(req_id, None)
2771 if (meta := self._recving_metadata.get(req_id)) is not None:
2772 if not self._is_hma_required:
2773 self._invalid_block_ids.put(set(meta.local_block_ids[0]))
@@ -2553,6 +2776,21 @@ class NixlBaseConnectorWorker:
2776 self.nixl_wrapper.release_xfer_handle(handle)
2777 self.xfer_stats.record_failed_transfer()
2778
2779 + def _send_pending_recv_notifs(self, req_id: str) -> None:
2780 + for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
2781 + try:
2782 + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
2783 + except Exception as e:
2784 + self._log_failure(
2785 + failure_type="notification_failed",
2786 + msg="P worker blocks will be freed after timeout. "
2787 + "This may indicate network issues.",
2788 + req_id=req_id,
2789 + error=e,
2790 + remote_agent_name=agent_name,
2791 + )
2792 + self.xfer_stats.record_failed_notification()
2793 +
2794 def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
2795 """
2796 Send heartbeat notifications to remote engines, extending lease on KV blocks.
@@ -2726,6 +2964,42 @@ class NixlBaseConnectorWorker:
2964 matched_blocks = min(len(local_slice), len(remote_slice))
2965 return local_slice[:matched_blocks], remote_slice[:matched_blocks]
2966
2967 + @staticmethod
2968 + def _block_ids_by_region(
2969 + block_ids: BlockIds, region_group_ids: list[int]
2970 + ) -> BlockIds:
2971 + """Expand group block IDs into the corresponding region order."""
2972 + shared_block_ids = list(itertools.chain.from_iterable(block_ids))
2973 + block_ids_by_region = []
2974 + for group_id in region_group_ids:
2975 + if group_id == _SHARED_REGION_GROUP_ID:
2976 + block_ids_by_region.append(shared_block_ids.copy())
2977 + continue
2978 + block_ids_by_region.append(list(block_ids[group_id]))
2979 + return block_ids_by_region
2980 +
2981 + @staticmethod
2982 + def _apply_prefix_caching_by_region(
2983 + decode_block_ids: BlockIds, prefill_block_ids: BlockIds
2984 + ) -> tuple[BlockIds, BlockIds]:
2985 + """Pair an uncached decode suffix with the same prefill regions."""
2986 + assert len(decode_block_ids) == len(prefill_block_ids)
2987 + if not any(decode_block_ids):
2988 + return [], prefill_block_ids
2989 +
2990 + trimmed_prefill: list[list[int]] = []
2991 + for decode_region, prefill_region in zip(
2992 + decode_block_ids, prefill_block_ids, strict=True
2993 + ):
2994 + if len(decode_region) > len(prefill_region):
2995 + raise ValueError(
2996 + "Decode allocated more KV pages than the prefill worker supplied"
2997 + )
2998 + trimmed_prefill.append(
2999 + list(prefill_region[-len(decode_region) :]) if decode_region else []
3000 + )
3001 + return decode_block_ids, trimmed_prefill
3002 +
3003 def _apply_prefix_caching(
3004 self,
3005 decode_block_ids: BlockIds,
@@ -2893,6 +3167,10 @@ class NixlBaseConnectorWorker:
3167
3168 self.kv_caches_base_addr.pop(engine_id, None)
3169 self.dst_num_blocks.pop(engine_id, None)
3170 + self.dst_region_num_blocks.pop(engine_id, None)
3171 + self.dst_region_group_ids.pop(engine_id, None)
3172 + self.dst_uses_region_group_mapping.pop(engine_id, None)
3173 + self.dst_region_mem_types.pop(engine_id, None)
3174 self.tp_mappings.pop(engine_id, None)
3175 if self.transfer_topo is not None:
3176 self.transfer_topo.unregister_remote_engine(engine_id)
@@ -2929,6 +3207,13 @@ class NixlBaseConnectorWorker:
3207 for handle in handles:
3208 self.nixl_wrapper.release_dlist_handle(handle)
3209 self.src_xfer_handles_by_tp_ratio.clear()
3210 + for handles in self._dram_src_handles_by_tp_ratio.values():
3211 + for handle in handles:
3212 + self.nixl_wrapper.release_dlist_handle(handle)
3213 + self._dram_src_handles_by_tp_ratio.clear()
3214 + for handle in self._dram_src_handles_by_block_size.values():
3215 + self.nixl_wrapper.release_dlist_handle(handle)
3216 + self._dram_src_handles_by_block_size.clear()
3217 for engine_id in list(self._remote_agents):
3218 self._cleanup_remote_engine(engine_id, log_eviction=False)
3219 for desc in self._registered_descs:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+6 -1
@@ -46,8 +46,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
46 # 8: Add dcp_size and pcp_size to NixlAgentMetadata
47 # 9: Add block_strides
48 # 10: Add dense virtual transfer pages for compressed MLA caches
49 +# 11: Add per-region transfer geometry and memory types to NixlAgentMetadata
50 #
50 -NIXL_CONNECTOR_VERSION: int = 10
51 +NIXL_CONNECTOR_VERSION: int = 11
52
53
54 @dataclass
@@ -64,6 +65,10 @@ class NixlAgentMetadata:
65 ssm_sizes: tuple[int, int]
66 attn_backend_name: str
67 physical_blocks_per_logical_kv_block: int
68 + region_num_blocks: list[int] | None = None
69 + region_group_ids: list[int] | None = None
70 + region_names: list[str] | None = None
71 + region_mem_types: list[str] | None = None
72 dcp_size: int = 1
73 pcp_size: int = 1
74
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+207 -67
@@ -5,6 +5,8 @@
5 import time
6 from typing import TYPE_CHECKING
7
8 +import numpy as np
9 +
10 from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
11 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
12 NixlBaseConnectorWorker,
@@ -156,60 +158,86 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
158 remote_info = self.transfer_topo.get_engine_info(engine_id)
159 tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
160
159 - remote_logical_block_ids = meta.remote.block_ids
160 - meta.remote.block_ids = self._logical_to_kernel_block_ids(
161 - remote_logical_block_ids,
162 - remote_info.remote_physical_blocks_per_logical,
163 - )
164 - num_groups = len(meta.local_block_ids)
161 dcp_active = self.dcp_size > 1 or remote_info.remote_dcp_size > 1
166 -
167 - def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
168 - return [
169 - list(block_ids[g]) if rank in plan.source_ranks_per_group[g] else []
170 - for g in range(num_groups)
162 + local_block_ids = meta.local_physical_block_ids
163 + local_region_groups = getattr(self, "region_group_ids", [])
164 + remote_region_groups = getattr(self, "dst_region_group_ids", {}).get(
165 + engine_id, local_region_groups
166 + )
167 + if local_region_groups != remote_region_groups:
168 + if not self.use_mla or self._has_mamba:
169 + raise NotImplementedError(
170 + "Different NIXL cache-group layouts are only supported for "
171 + "pure MLA models"
172 + )
173 + assert len(plan.all_source_ranks) == 1
174 + remote_physical_block_ids = self._logical_to_kernel_block_ids(
175 + meta.remote.block_ids,
176 + remote_info.remote_physical_blocks_per_logical,
177 + )
178 + remote_by_region = self._block_ids_by_region(
179 + remote_physical_block_ids, remote_region_groups
180 + )
181 + read_specs = [
182 + ReadSpec(
183 + remote_rank=plan.all_source_ranks[0],
184 + local_block_ids=self._block_ids_by_region(
185 + local_block_ids, local_region_groups
186 + ),
187 + remote_block_ids=remote_by_region,
188 + block_ids_by_region=True,
189 + )
190 ]
191 + else:
192 + remote_logical_block_ids = meta.remote.block_ids
193 + meta.remote.block_ids = self._logical_to_kernel_block_ids(
194 + remote_logical_block_ids,
195 + remote_info.remote_physical_blocks_per_logical,
196 + )
197 + num_groups = len(meta.local_block_ids)
198
173 - read_specs = []
174 - for rank in plan.all_source_ranks:
175 - if dcp_active:
176 - # DCP interleaves at block granularity, so slicing happens
177 - # here on logical blocks, before kernel-block expansion.
178 - local_ids = group_ids(meta.local_block_ids, rank)
179 - remote_ids = group_ids(remote_logical_block_ids, rank)
180 - for g in range(num_groups):
181 - if not local_ids[g]:
182 - continue
183 - # Prefix cache hit may lead to skip some of the remote reads
184 - # TODO (NickLucche) consider unifying prefix cache handling on
185 - # logical blocks here for both dcp and non-dcp
186 - local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
187 - local_ids[g],
188 - remote_ids[g],
199 + def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]:
200 + return [
201 + list(block_ids[g]) if rank in plan.source_ranks_per_group[g] else []
202 + for g in range(num_groups)
203 + ]
204 +
205 + read_specs = []
206 + for rank in plan.all_source_ranks:
207 + if dcp_active:
208 + local_ids = group_ids(meta.local_block_ids, rank)
209 + remote_ids = group_ids(remote_logical_block_ids, rank)
210 + for g in range(num_groups):
211 + if not local_ids[g]:
212 + continue
213 + local_ids[g], remote_ids[g] = self._apply_dcp_prefix_caching(
214 + local_ids[g],
215 + remote_ids[g],
216 + remote_rank=rank,
217 + local_dcp_size=self.dcp_size,
218 + local_dcp_rank=self.dcp_rank,
219 + remote_dcp_size=remote_info.remote_dcp_size,
220 + local_num_computed_blocks=(
221 + meta.local_num_computed_blocks[g]
222 + ),
223 + )
224 + local_physical_ids = self._logical_to_kernel_block_ids(
225 + local_ids, self._physical_blocks_per_logical_kv_block
226 + )
227 + remote_physical_ids = self._logical_to_kernel_block_ids(
228 + remote_ids,
229 + remote_info.remote_physical_blocks_per_logical,
230 + )
231 + else:
232 + local_physical_ids = group_ids(meta.local_physical_block_ids, rank)
233 + remote_physical_ids = group_ids(meta.remote.block_ids, rank)
234 + read_specs.append(
235 + ReadSpec(
236 remote_rank=rank,
190 - local_dcp_size=self.dcp_size,
191 - local_dcp_rank=self.dcp_rank,
192 - remote_dcp_size=remote_info.remote_dcp_size,
193 - local_num_computed_blocks=meta.local_num_computed_blocks[g],
237 + local_block_ids=local_physical_ids,
238 + remote_block_ids=remote_physical_ids,
239 )
195 - local_physical_ids = self._logical_to_kernel_block_ids(
196 - local_ids, self._physical_blocks_per_logical_kv_block
197 - )
198 - remote_physical_ids = self._logical_to_kernel_block_ids(
199 - remote_ids, remote_info.remote_physical_blocks_per_logical
240 )
201 - else:
202 - # No DCP realignment needed: reuse the already-expanded full
203 - # physical lists instead of re-deriving them from logical ids.
204 - local_physical_ids = group_ids(meta.local_physical_block_ids, rank)
205 - remote_physical_ids = group_ids(meta.remote.block_ids, rank)
206 - read_specs.append(
207 - ReadSpec(
208 - remote_rank=rank,
209 - local_block_ids=local_physical_ids,
210 - remote_block_ids=remote_physical_ids,
211 - )
212 - )
241
242 # D may have to perform multiple reads from different remote ranks.
243 # Pure MLA reads once because its cache is replicated. Hybrid
@@ -234,27 +262,39 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
262 # reads. Get the memory chunk onto which we will write to.
263 split_key = (tp_ratio, remote_block_size)
264 local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[split_key][i]
265 + local_dram_handle = (
266 + self._dram_src_handles_by_tp_ratio[split_key][i]
267 + if self._mixed_mem_types
268 + else None
269 + )
270 else:
271 # Single read from remote, we write to the whole memory region.
272 # Also handle remote block size different from local block size.
273 local_xfer_side_handle = self.src_xfer_handles_by_block_size[
274 remote_block_size
275 ]
276 + local_dram_handle = (
277 + self._dram_src_handles_by_block_size[remote_block_size]
278 + if self._mixed_mem_types
279 + else None
280 + )
281
282 # Destination handle: remote_engine_id -> remote_rank -> handle.
283 remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
284 spec.remote_rank
285 ]
286
249 - self._read_blocks(
287 + if not self._read_blocks(
288 read_spec=spec,
289 request_id=req_id,
290 dst_engine_id=meta.remote.engine_id,
291 remote_request_id=meta.remote.request_id,
292 local_xfer_side_handle=local_xfer_side_handle,
293 + local_dram_handle=local_dram_handle,
294 remote_xfer_side_handle=remote_xfer_side_handle,
295 expected_consumers=plan.local_consumers,
257 - )
296 + ):
297 + return
298
299 if self.use_mla and tp_ratio < 0 and len(read_specs) == 1:
300 # ..but we still need to notify the other remote ranks that we
@@ -274,9 +314,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
314 request_id: str,
315 remote_request_id: str,
316 local_xfer_side_handle: int,
317 + local_dram_handle: int | None,
318 remote_xfer_side_handle: int,
319 expected_consumers: int,
279 - ):
320 + ) -> bool:
321 """
322 Post a READ point-to-point xfer request from a single local worker to
323 a single remote worker.
@@ -291,6 +332,10 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
332 remote_info.remote_block_size
333 )
334 if block_size_ratio > 1:
335 + if read_spec.block_ids_by_region:
336 + raise NotImplementedError(
337 + "Region-mapped NIXL transfers require matching physical block sizes"
338 + )
339 local_block_ids, remote_block_ids = (
340 self._map_block_ids_for_block_size_ratio(
341 local_block_ids, remote_block_ids, block_size_ratio
@@ -328,26 +373,30 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
373 remote_agent_name=agent_name,
374 )
375 self.xfer_stats.record_failed_notification()
331 - return
376 + return True
377
333 - assert (
334 - len(remote_block_ids)
335 - == len(local_block_ids)
336 - == len(self.kv_cache_config.transfer_groups)
337 - )
338 - if not (self.dcp_size > 1 or remote_info.remote_dcp_size > 1):
339 - # DCP-active reads were already trimmed (and DCP-realigned, when
340 - # sizes mismatch) in _read_blocks_for_req, at logical granularity.
341 - local_block_ids, remote_block_ids = self._apply_prefix_caching(
378 + if read_spec.block_ids_by_region:
379 + local_block_ids, remote_block_ids = self._apply_prefix_caching_by_region(
380 decode_block_ids=local_block_ids,
381 prefill_block_ids=remote_block_ids,
344 - decode_physical_per_logical=(
345 - self._physical_blocks_per_logical_kv_block
346 - ),
347 - prefill_physical_per_logical=(
348 - remote_info.remote_physical_blocks_per_logical
349 - ),
382 )
383 + else:
384 + assert (
385 + len(remote_block_ids)
386 + == len(local_block_ids)
387 + == len(self.kv_cache_config.transfer_groups)
388 + )
389 + if not (self.dcp_size > 1 or remote_info.remote_dcp_size > 1):
390 + local_block_ids, remote_block_ids = self._apply_prefix_caching(
391 + decode_block_ids=local_block_ids,
392 + prefill_block_ids=remote_block_ids,
393 + decode_physical_per_logical=(
394 + self._physical_blocks_per_logical_kv_block
395 + ),
396 + prefill_physical_per_logical=(
397 + remote_info.remote_physical_blocks_per_logical
398 + ),
399 + )
400
401 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
402 # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
@@ -359,12 +408,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
408 dst_num_blocks=self.dst_num_blocks[dst_engine_id],
409 block_size_ratio=None,
410 physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
411 + region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
412 + region_group_ids=(
413 + list(range(self.num_regions))
414 + if read_spec.block_ids_by_region
415 + else self.dst_region_group_ids[dst_engine_id]
416 + ),
417 + uses_region_group_mapping=(
418 + self.num_regions > 1
419 + if read_spec.block_ids_by_region
420 + else self.dst_uses_region_group_mapping[dst_engine_id]
421 + ),
422 )
423 local_block_descs_ids = self._compute_desc_ids(
424 block_ids=local_block_ids,
425 dst_num_blocks=self.dst_num_blocks[self.engine_id],
426 block_size_ratio=block_size_ratio,
427 physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
428 + region_num_blocks=self.dst_region_num_blocks[self.engine_id],
429 + region_group_ids=(
430 + list(range(self.num_regions))
431 + if read_spec.block_ids_by_region
432 + else self.region_group_ids
433 + ),
434 + uses_region_group_mapping=(
435 + self.num_regions > 1
436 + if read_spec.block_ids_by_region
437 + else self._uses_region_group_mapping
438 + ),
439 )
440
441 assert len(local_block_descs_ids) == len(remote_block_descs_ids)
@@ -372,6 +443,19 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
443 # Prepare transfer with Nixl.
444 handle = None
445 try:
446 + if self._mixed_mem_types:
447 + self._read_blocks_mixed(
448 + request_id=request_id,
449 + local_block_size_key=remote_info.remote_block_size,
450 + local_device_handle=local_xfer_side_handle,
451 + local_dram_handle=local_dram_handle,
452 + remote_xfer_side_handle=remote_xfer_side_handle,
453 + local_block_descs_ids=local_block_descs_ids,
454 + remote_block_descs_ids=remote_block_descs_ids,
455 + notif_agent=self._remote_agents[dst_engine_id][(0, remote_rank)],
456 + notif_id=notif_id,
457 + )
458 + return True
459 handle = self.nixl_wrapper.make_prepped_xfer(
460 "READ",
461 local_xfer_side_handle,
@@ -386,6 +470,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
470
471 # Use handle to check completion in future step().
472 self._recving_transfers[request_id].append(handle)
473 + return True
474 except Exception as e:
475 # mark all (logical) blocks for this request as invalid
476 self._log_failure(
@@ -397,6 +482,61 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
482 remote_rank=remote_rank,
483 )
484 self._handle_failed_transfer(request_id, handle)
485 + return False
486 +
487 + def _read_blocks_mixed(
488 + self,
489 + request_id: str,
490 + local_block_size_key: int,
491 + local_device_handle: int,
492 + local_dram_handle: int | None,
493 + remote_xfer_side_handle: int,
494 + local_block_descs_ids: np.ndarray,
495 + remote_block_descs_ids: np.ndarray,
496 + notif_agent: str,
497 + notif_id: bytes,
498 + ) -> None:
499 + """Split a READ across the local DRAM and device descriptor lists."""
500 + desc_is_dram = self._desc_is_dram_by_block_size[local_block_size_key]
501 + desc_pos = self._desc_pos_by_block_size[local_block_size_key]
502 + local_ids = np.asarray(local_block_descs_ids)
503 + remote_ids = np.asarray(remote_block_descs_ids)
504 + is_dram = desc_is_dram[local_ids]
505 +
506 + assert local_dram_handle is not None
507 + reads = (
508 + (is_dram, local_dram_handle),
509 + (~is_dram, local_device_handle),
510 + )
511 + handles: list[int] = []
512 + try:
513 + for mask, local_handle in reads:
514 + if mask.any():
515 + handles.append(
516 + self.nixl_wrapper.make_prepped_xfer(
517 + "READ",
518 + local_handle,
519 + desc_pos[local_ids[mask]],
520 + remote_xfer_side_handle,
521 + remote_ids[mask],
522 + )
523 + )
524 + except Exception:
525 + for handle in handles:
526 + self.nixl_wrapper.release_xfer_handle(handle)
527 + raise
528 +
529 + self._pending_recv_notifs.setdefault(request_id, []).append(
530 + (notif_agent, notif_id)
531 + )
532 + for index, handle in enumerate(handles):
533 + try:
534 + self.nixl_wrapper.transfer(handle)
535 + except Exception:
536 + for unstarted in handles[index:]:
537 + self.nixl_wrapper.release_xfer_handle(unstarted)
538 + raise
539 + self._recving_transfers[request_id].append(handle)
540
541 def _get_new_notifs(self) -> set[str]:
542 """
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+10
@@ -132,6 +132,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
132
133 def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
134 super().register_kv_caches(kv_caches)
135 + if self._mixed_mem_types:
136 + raise NotImplementedError(
137 + "NixlPushConnector does not support mixed-memory KV caches"
138 + )
139 if self._push_writer_thread is None:
140 self._push_writer_thread = threading.Thread(
141 target=self._push_writer_loop,
@@ -661,12 +665,18 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
665 dst_num_blocks=self.dst_num_blocks[dst_engine_id],
666 block_size_ratio=None,
667 physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
668 + region_num_blocks=self.dst_region_num_blocks[dst_engine_id],
669 + region_group_ids=self.dst_region_group_ids[dst_engine_id],
670 + uses_region_group_mapping=self.dst_uses_region_group_mapping[dst_engine_id],
671 )
672 local_block_descs_ids = self._compute_desc_ids(
673 block_ids=local_block_ids,
674 dst_num_blocks=self.dst_num_blocks[self.engine_id],
675 block_size_ratio=block_size_ratio,
676 physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
677 + region_num_blocks=self.dst_region_num_blocks[self.engine_id],
678 + region_group_ids=self.region_group_ids,
679 + uses_region_group_mapping=self._uses_region_group_mapping,
680 )
681
682 assert len(local_block_descs_ids) == len(remote_block_descs_ids)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py
+1
@@ -26,6 +26,7 @@ class ReadSpec:
26 remote_rank: int
27 local_block_ids: BlockIds
28 remote_block_ids: BlockIds
29 + block_ids_by_region: bool = False
30
31
32 def _is_attention_spec(spec_type: type[KVCacheSpec]) -> bool: