@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0fefffc934

[KVConnector][NIXL] Support attention-HMA layouts in pipeline-parallel push prefill (#50494)

Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yiliu Dong <91178480+qianlihuang@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com>

qizixi committed Sep 15, 2026 at 15:43 UTC 0fefffc9346644a01dc132e544b881ad646ea81b

8 files changed +713 -49

docs/design/nixl_kv_push_connector.md

+44

@@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing

181 is D's own request id, taken from the registration), so the D-side

182 accounting code is unchanged.

183

184 +## Pipeline parallelism and hybrid KV caches

185 +

186 +The pull connector requires the local and remote workers to expose a

187 +congruent list of KV regions — region *i* here corresponds to region

188 +*i* there. That assumption breaks under pipeline parallelism (PP)

189 +combined with a hybrid (HMA) KV layout:

190 +

191 +* a PP-sharded prefiller (**P**) holds only a slice of the model's

192 + layers while the `PP=1` decoder (**D**) holds them all, so region

193 + counts differ;

194 +* with HMA, several layer names are pooled into one region and the layer

195 + that represents a pooled region can differ between P and D.

196 +

197 +`NixlPushConnector` handles this for PP-sharded producers by routing

198 +**by layer-name (member) identity** instead of by region index. Each

199 +worker advertises which layer names back each of its NIXL regions in the

200 +handshake metadata (`NixlAgentMetadata.region_members`). A producer that

201 +needs member routing derives its member-major layout once, when it

202 +registers its KV caches, and every transfer it issues uses that order.

203 +`add_remote_agent` then selects exactly the remote regions this stage

204 +owns and reorders them to match, so both sides stay paired regardless of

205 +how each remote rank happens to order its metadata.

206 +

207 +Addresses, block lengths, strides, and per-region capacities follow the same

208 +member order. Descriptor offsets use each member's region capacity, so P and

209 +D need not allocate the same number of blocks. Physical allocations are still

210 +registered once, even when multiple members share them.

211 +

212 +Invariants enforced when the remote regions are aligned:

213 +

214 +* every locally owned member must be advertised exactly once by the

215 + remote; a missing member fails the handshake rather than silently

216 + leaving that layer's KV stale, and remote-only members (owned by other

217 + PP stages) are ignored;

218 +* a remote that omits member metadata while the local layout requires

219 + member routing fails the handshake instead of falling back to

220 + region-index routing;

221 +* member order is a property of the local layout alone, so the same local

222 + source descriptors serve every remote engine and TP rank; only the

223 + remote descriptor list is rebuilt per rank.

224 +

225 +Decode-side PP is unsupported because completions are counted per

226 +consumer rank. Mamba/SSM hybrids are unsupported under PP.

227 +

228 ## Scheduler-side responsibilities

229

230 `NixlPushConnectorScheduler` extends the base scheduler with:

docs/features/nixl_connector_compatibility.md

+18

@@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan

104 - `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting.

105 - Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.

106

107 +### Pipeline parallelism

108 +

109 +The default (pull) NixlConnector does **not** support

110 +`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA)

111 +layout: region indices are not a stable identity across the

112 +prefill/decode layer split, so the connector raises at startup. Use the

113 +push connector (`NixlPushConnector`) for pipeline parallelism with

114 +hybrid KV caches — it routes transfers by layer-name (member) identity,

115 +so a PP-sharded prefiller can write into a `PP=1` decoder. See

116 +[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md).

117 +

118 +Current push PP + HMA limitations:

119 +

120 +- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported.

121 +- Hybrid SSM/Mamba layouts are not supported under PP.

122 +- HMA requires the same block size on P and D.

123 +- Attention-HMA member routing requires decode TP to be no greater than prefill TP.

124 +

125 ### Quantized KV cache

126

127 [Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake.

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -918,6 +918,7 @@ def _make_mock_worker_for_desc_ids(

918 worker = MagicMock(spec=NixlConnectorWorker)

919 worker.num_regions = num_regions

920 worker._uses_region_group_mapping = False

921 + worker._transfer_layer_group_ids = ()

922 worker._has_mamba = has_mamba

923 worker._group_spec_types = group_spec_types

924 worker.block_len_per_layer = block_len_per_layer or [100]

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

+28 -4

@@ -12,6 +12,7 @@ mid-decode (silent corruption of an unrelated request).

12 """

14 from collections import defaultdict

15 +from threading import Event, Lock

16 from unittest.mock import MagicMock, patch

18 import numpy as np

@@ -112,6 +113,7 @@ def test_local_descriptors_follow_each_region_pool_capacity():

113 worker.block_len_per_layer = [16, 16]

114 worker.block_stride_per_layer = [16, 16]

115 worker.region_num_blocks = [2, 3]

116 + worker._transfer_layer_region_indices = ()

117

118 descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1)

119

@@ -119,7 +121,8 @@ def test_local_descriptors_follow_each_region_pool_capacity():

121

122

123 @pytest.mark.cpu_test

122 -def test_overlaid_transfer_groups_share_region_geometry():

124 +@pytest.mark.parametrize("push_pp", [False, True])

125 +def test_overlaid_transfer_groups_share_region_geometry(push_pp):

126 """Groups overlaid on one allocation share its transfer region."""

127 import msgspec

128

@@ -127,6 +130,9 @@ def test_overlaid_transfer_groups_share_region_geometry():

130 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

131 NixlAgentMetadata,

132 )

133 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (

134 + NixlPushConnectorWorker,

135 + )

136 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (

137 NixlConnectorWorker,

138 )

@@ -153,7 +159,14 @@ def test_overlaid_transfer_groups_share_region_geometry():

159 }

160 groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches]

161

156 - worker = object.__new__(NixlConnectorWorker)

162 + worker_cls = NixlPushConnectorWorker if push_pp else NixlConnectorWorker

163 + worker = object.__new__(worker_cls)

164 + if push_pp:

165 + worker._push_writer_stop = Event()

166 + worker._push_writer_wake = Event()

167 + worker._push_writer_thread = MagicMock()

168 + worker._sending_transfers_lock = Lock()

169 + worker._sending_transfers = defaultdict(list)

170 worker.tp_rank = 0

171 worker.world_size = 1

172 worker.transfer_tp_rank = 0

@@ -198,6 +211,9 @@ def test_overlaid_transfer_groups_share_region_geometry():

211 worker._desc_is_dram_by_block_size = {}

212 worker._desc_pos_by_block_size = {}

213 worker._dram_src_handles_by_block_size = {}

214 + worker._transfer_layer_names = ()

215 + worker._transfer_layer_region_indices = ()

216 + worker._transfer_layer_group_ids = ()

217 worker._region_is_mla = []

218 worker.block_len_per_layer = []

219 worker.block_stride_per_layer = []

@@ -205,7 +221,8 @@ def test_overlaid_transfer_groups_share_region_geometry():

221 worker.use_host_buffer = False

222 worker.host_xfer_buffers = {}

223 worker.device_kv_caches = {}

208 - worker.pp_size = 1

224 + worker.pp_size = 2 if push_pp else 1

225 + worker._is_hma_required = True

226 worker.dcp_size = 1

227 worker.pcp_size = 1

228 worker.kv_buffer_device = "cuda"

@@ -240,7 +257,13 @@ def test_overlaid_transfer_groups_share_region_geometry():

257 expected_addrs = [

258 backing.data_ptr() + block * block_stride for block in range(num_blocks)

259 ]

243 - assert worker.src_blocks_data[:, 0].tolist() == expected_addrs

260 + num_desc_regions = 2 if push_pp else 1

261 + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions

262 + assert worker.num_descs == num_blocks * num_desc_regions

263 + assert (

264 + worker.dst_region_num_blocks[worker.engine_id]

265 + == [num_blocks] * num_desc_regions

266 + )

267

268 metadata = msgspec.msgpack.decode(

269 worker.xfer_handshake_metadata.agent_metadata_bytes,

@@ -248,6 +271,7 @@ def test_overlaid_transfer_groups_share_region_geometry():

271 )

272 assert metadata.region_group_ids == [-1]

273 assert metadata.region_num_blocks == [num_blocks]

274 + assert metadata.region_members == ([["layer.0", "layer.1"]] if push_pp else [])

275 assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]

276

277

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+398 -2

@@ -32,9 +32,13 @@ from unittest.mock import MagicMock, patch

32 import msgspec

33 import pytest

35 +from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology

36 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

37 KVConnectorTransferResults,

38 )

39 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (

40 + NixlBaseConnectorWorker,

41 +)

42 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

43 PUSH_REG_NOTIF_PREFIX,

44 NixlAgentMetadata,

@@ -43,6 +47,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

47 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (

48 NixlPushConnectorWorker,

49 )

50 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping

51 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (

52 get_base_request_id,

53 )

@@ -354,6 +359,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

359 w._failed_recv_reqs = queue.Queue()

360 w._recv_failures = set()

361 w._recving_transfers = defaultdict(list)

362 + w._is_hma_required = False

363 w._reqs_to_process = set()

364 w._reqs_to_send = {}

365 w.consumer_notification_counts_by_req = defaultdict(int)

@@ -361,12 +367,16 @@ class _StubWriterWorker(NixlPushConnectorWorker):

367 w.pcp_rank = 0

368 w.pcp_dcp_sharded = False

369 w.world_size = 1

370 + w.pp_size = 1

371 w.engine_id = "test-decode-engine"

372 w._remote_agents = {}

373 w._handshake_lock = threading.RLock()

374 w._physical_blocks_per_logical_kv_block = 1

375 w._uses_region_group_mapping = False

376 w.region_group_ids = [0]

377 + w._transfer_layer_names = ()

378 + w._transfer_layer_region_indices = ()

379 + w._transfer_layer_group_ids = ()

380 w.dst_region_num_blocks = {}

381 w.dst_region_group_ids = {}

382 w.dst_uses_region_group_mapping = {}

@@ -922,12 +932,20 @@ class TestPushWriterNotifs:

932 # Both handles were still cleaned up.

933 assert request_id not in w._sending_transfers

934

925 - def test_completed_push_send_waits_for_consumer_notif(self):

935 + @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)])

936 + def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma):

937 """P-side sends are completed by consumer notifs / lease expiry,

938 not by local WRITE-handle completion."""

939 w = self._pollable_worker()

940 + w.pp_size = pp_size

941 + w._is_hma_required = is_hma

942 request_id = self._make_sending_req(w)

930 - w.nixl_wrapper.check_xfer_state.return_value = "DONE"

943 + w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"]

944 +

945 + results = w.get_transfer_results()

946 + assert request_id not in results.finished_sending

947 + assert w._sending_transfers[request_id] == [102]

948 + assert request_id in w._reqs_to_send

949

950 results = w.get_transfer_results()

951

@@ -940,6 +958,12 @@ class TestPushWriterNotifs:

958 results = w.get_transfer_results()

959 assert request_id in results.finished_sending

960 assert request_id not in w._reqs_to_send

961 + assert request_id not in w._reqs_to_process

962 + assert w._evict_finished_inbox.get_nowait() == request_id

963 + assert w.get_transfer_results().finished_sending == set()

964 + assert w._evict_finished_inbox.empty()

965 + assert w.nixl_wrapper.release_xfer_handle.call_count == 2

966 + w.xfer_stats.record_kv_expired_req.assert_not_called()

967

968

969 # ----------------------------------------------------------------- #

@@ -1576,3 +1600,375 @@ class TestPushPrefixCaching:

1600 local, remote = self._written_block_ids(w)

1601 assert local == [10, 11, 12]

1602 assert remote == [500, 501, 502]

1603 +

1604 + @pytest.mark.parametrize("layer_name_routing", [False, True])

1605 + def test_different_region_groups_require_layer_routing(self, layer_name_routing):

1606 + """A layer mapped to a pooled remote region can still receive its blocks."""

1607 + w, engine_id = self._worker_driving_xfer()

1608 + w.dst_region_group_ids[engine_id] = [-1]

1609 + if layer_name_routing:

1610 + w._transfer_layer_group_ids = (0,)

1611 + reg = _registration_data("req-groups", local_block_ids=([500, 501],))

1612 +

1613 + if not layer_name_routing:

1614 + with pytest.raises(NotImplementedError, match="cache-group layouts"):

1615 + NixlPushConnectorWorker._do_start_push_kv(

1616 + w, "req-groups", ([10, 11],), reg

1617 + )

1618 + w.nixl_wrapper.make_prepped_xfer.assert_not_called()

1619 + return

1620 +

1621 + NixlPushConnectorWorker._do_start_push_kv(w, "req-groups", ([10, 11],), reg)

1622 +

1623 + local, remote = self._written_block_ids(w)

1624 + assert local == [10, 11]

1625 + assert remote == [500, 501]

1626 +

1627 +

1628 +def _agent_metadata(

1629 + region_layers: list[list[str]],

1630 + base_addresses: list[int],

1631 + block_lens: list[int],

1632 + block_strides: list[int] | None = None,

1633 +) -> NixlAgentMetadata:

1634 + return NixlAgentMetadata(

1635 + engine_id="remote-engine",

1636 + agent_metadata=b"agent",

1637 + kv_caches_base_addr=base_addresses,

1638 + device_id=7,

1639 + num_blocks=2,

1640 + block_lens=block_lens,

1641 + block_strides=list(block_lens if block_strides is None else block_strides),

1642 + kv_cache_layout="LBHNC",

1643 + block_size=16,

1644 + ssm_sizes=(0, 0),

1645 + attn_backend_name="FLASH_ATTN",

1646 + physical_blocks_per_logical_kv_block=1,

1647 + region_members=region_layers,

1648 + )

1649 +

1650 +

1651 +def _layer_routing_worker(

1652 + region_layers: list[list[str]],

1653 + group_by_layer: dict[str, int],

1654 + pp_size: int = 2,

1655 + is_hma: bool = True,

1656 +) -> _StubWriterWorker:

1657 + worker = _StubWriterWorker.fresh()

1658 + worker.pp_size = pp_size

1659 + worker.dcp_size = 1

1660 + worker.block_size = 16

1661 + worker._has_mamba = False

1662 + worker._is_hma_required = is_hma

1663 + worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer)

1664 + worker.transfer_topo = MagicMock()

1665 + worker._set_region_layers(region_layers)

1666 + return worker

1667 +

1668 +

1669 +@pytest.mark.parametrize(

1670 + ("region_num_blocks", "expected"),

1671 + [(None, [1, 2, 15, 21, 22]), ([4, 7, 5], [1, 2, 9, 12, 13])],

1672 +)

1673 +def test_layer_group_ids_route_descriptor_blocks(region_num_blocks, expected):

1674 + worker = _layer_routing_worker(

1675 + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}

1676 + )

1677 + desc_ids = worker._compute_desc_ids(

1678 + block_ids=[[1, 2], [5]],

1679 + dst_num_blocks=10,

1680 + block_size_ratio=None,

1681 + physical_blocks_per_logical=1,

1682 + region_num_blocks=region_num_blocks,

1683 + )

1684 + assert desc_ids.tolist() == expected

1685 +

1686 +

1687 +def test_layer_metadata_round_trip():

1688 + metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256])

1689 + metadata.region_num_blocks = [4]

1690 + metadata.region_group_ids = [-1]

1691 + metadata.region_names = ["L0"]

1692 + metadata.region_mem_types = ["VRAM"]

1693 +

1694 + encoded = msgspec.msgpack.encode(metadata)

1695 + assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata

1696 +

1697 +

1698 +def test_layer_identity_gate_preserves_the_non_hma_path():

1699 + assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == (

1700 + 0,

1701 + )

1702 +

1703 + # A base worker (pull) never routes by layer name.

1704 + pull = object.__new__(NixlBaseConnectorWorker)

1705 + pull._transfer_layer_names = ()

1706 + pull._transfer_layer_region_indices = ()

1707 + pull._transfer_layer_group_ids = ()

1708 + pull._set_region_layers([["a"]])

1709 + assert pull._transfer_layer_region_indices == ()

1710 +

1711 + # A non-HMA local layout does not require layer-name routing.

1712 + assert (

1713 + _layer_routing_worker(

1714 + [["a"]], {"a": 0}, is_hma=False

1715 + )._transfer_layer_region_indices

1716 + == ()

1717 + )

1718 +

1719 + # PP=1 has congruent local/remote regions and keeps the region route even

1720 + # when its allocator is hybrid.

1721 + assert (

1722 + _layer_routing_worker(

1723 + [["a"]], {"a": 0}, pp_size=1

1724 + )._transfer_layer_region_indices

1725 + == ()

1726 + )

1727 +

1728 +

1729 +def test_layer_alignment_fails_loud_when_remote_omits_layers():

1730 + # Falling back to region-index routing here would silently transfer stale

1731 + # KV, so an unannounced peer must fail the handshake instead.

1732 + worker = _layer_routing_worker([["a"]], {"a": 0})

1733 + with pytest.raises(AssertionError, match="no region_members"):

1734 + worker._align_remote_regions_by_layer(_agent_metadata([], [0xA000], [128]))

1735 +

1736 +

1737 +def test_layer_alignment_expands_pooled_regions():

1738 + worker = _layer_routing_worker(

1739 + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}

1740 + )

1741 + assert worker._transfer_layer_names == ("a", "a.swa", "b")

1742 + assert worker._transfer_layer_region_indices == (0, 0, 1)

1743 + assert worker._transfer_layer_group_ids == (0, 1, 0)

1744 +

1745 + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128])

1746 + metadata.region_num_blocks = [4, 6]

1747 + metadata.region_group_ids = [-1, 0]

1748 + metadata.region_names = ["a", "b"]

1749 + metadata.region_mem_types = ["VRAM", "VRAM"]

1750 + worker._align_remote_regions_by_layer(metadata)

1751 +

1752 + assert metadata.kv_caches_base_addr == [0xA000, 0xA000, 0xB000]

1753 + assert metadata.block_lens == [128, 128, 128]

1754 + assert metadata.block_strides == [128, 128, 128]

1755 + assert metadata.region_members == [["a"], ["a.swa"], ["b"]]

1756 + assert metadata.region_num_blocks == [4, 4, 6]

1757 + assert metadata.region_group_ids == [-1, -1, 0]

1758 + assert metadata.region_names == ["a", "a.swa", "b"]

1759 + assert metadata.region_mem_types == ["VRAM", "VRAM", "VRAM"]

1760 +

1761 +

1762 +def test_layer_alignment_filters_and_reorders_a_pp_stage():

1763 + worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1})

1764 + metadata = _agent_metadata(

1765 + [["l2"], ["l0"], ["l3"], ["l1"]],

1766 + [0xC000, 0xA000, 0xD000, 0xB000],

1767 + [65536, 65536, 32768, 32768],

1768 + [131072, 131072, 65536, 65536],

1769 + )

1770 + metadata.region_num_blocks = [4, 6, 8, 10]

1771 + metadata.region_group_ids = [0, 0, 1, 1]

1772 + metadata.region_names = ["l2", "l0", "l3", "l1"]

1773 + metadata.region_mem_types = ["VRAM"] * 4

1774 +

1775 + worker._align_remote_regions_by_layer(metadata)

1776 +

1777 + assert worker._transfer_layer_region_indices == (0, 1)

1778 + assert worker._transfer_layer_group_ids == (0, 1)

1779 + assert metadata.kv_caches_base_addr == [0xC000, 0xD000]

1780 + assert metadata.block_lens == [65536, 32768]

1781 + assert metadata.block_strides == [131072, 65536]

1782 + assert metadata.region_num_blocks == [4, 8]

1783 + assert metadata.region_group_ids == [0, 1]

1784 + assert metadata.region_names == ["l2", "l3"]

1785 + assert metadata.region_mem_types == ["VRAM", "VRAM"]

1786 +

1787 +

1788 +@pytest.mark.parametrize(

1789 + ("local_num_blocks", "remote_num_blocks"),

1790 + [([4, 4], [4, 4, 4]), ([4, 6], [3, 5, 7])],

1791 +)

1792 +def test_layer_descriptors_pair_layers_across_asymmetric_pp_split(

1793 + local_num_blocks, remote_num_blocks

1794 +):

1795 + """A PP split can leave each stage a different mix of attention types.

1796 +

1797 + Stage 1 owns L3 (full) pooled with L4 (sliding) in region 0, and L5 (full)

1798 + alone in region 1. The consumer holds every layer and pools them

1799 + differently, so region indices cannot be paired positionally.

1800 + """

1801 + worker = _layer_routing_worker([["L3", "L4"], ["L5"]], {"L3": 0, "L4": 1, "L5": 0})

1802 + worker.block_len_per_layer = [128, 128]

1803 + worker.block_stride_per_layer = [256, 512]

1804 + worker._region_is_mla = [False, False]

1805 + worker.num_blocks = 4

1806 + worker.region_num_blocks = local_num_blocks

1807 + worker.device_id = 0

1808 +

1809 + consumer = _agent_metadata(

1810 + [["L0", "L1"], ["L3", "L2"], ["L5", "L4"]],

1811 + [0xA000, 0xB000, 0xC000],

1812 + [128, 128, 128],

1813 + [256, 512, 1024],

1814 + )

1815 + consumer.num_blocks = 4

1816 + consumer.region_num_blocks = remote_num_blocks

1817 + worker._align_remote_regions_by_layer(consumer)

1818 +

1819 + # L3 -> remote region 1, L4 -> 2, L5 -> 2. Pairing by index would have sent

1820 + # L5 to region 1 and L4 to region 2's sibling.

1821 + assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000]

1822 +

1823 + plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0)

1824 + local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1)

1825 + remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1)

1826 + local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices]

1827 + local_ids = worker._compute_desc_ids(

1828 + [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts

1829 + )

1830 + remote_ids = worker._compute_desc_ids(

1831 + [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks

1832 + )

1833 + assert len(local_descs) == sum(local_counts)

1834 + assert len(remote_descs) == sum(consumer.region_num_blocks)

1835 +

1836 + # Each pair addresses the same layer's blocks despite different pooling

1837 + # and strides. L3 and L5 use group 0; L4 uses group 1.

1838 + assert local_descs[local_ids].tolist() == [

1839 + [0x1100, 128, 0],

1840 + [0x1200, 128, 0],

1841 + [0x1300, 128, 0],

1842 + [0x2200, 128, 0],

1843 + [0x2400, 128, 0],

1844 + ]

1845 + assert remote_descs[remote_ids].tolist() == [

1846 + [0xB000, 128, 7],

1847 + [0xB200, 128, 7],

1848 + [0xC800, 128, 7],

1849 + [0xC000, 128, 7],

1850 + [0xC400, 128, 7],

1851 + ]

1852 +

1853 +

1854 +def test_layer_alignment_is_canonical_across_remote_orderings():

1855 + local, groups = [["x"], ["y"]], {"x": 0, "y": 1}

1856 + rank0 = _agent_metadata([["x"], ["y"]], [0x1000, 0x2000], [64, 128], [256, 512])

1857 + rank1 = _agent_metadata([["y"], ["x"]], [0x2000, 0x1000], [128, 64], [512, 256])

1858 +

1859 + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank0)

1860 + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank1)

1861 +

1862 + assert rank0.kv_caches_base_addr == rank1.kv_caches_base_addr == [0x1000, 0x2000]

1863 + assert rank0.block_lens == rank1.block_lens == [64, 128]

1864 + assert rank0.block_strides == rank1.block_strides == [256, 512]

1865 +

1866 +

1867 +def test_layer_alignment_is_idempotent():

1868 + worker = _layer_routing_worker(

1869 + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0}

1870 + )

1871 + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128])

1872 + metadata.region_num_blocks = [4, 6]

1873 + metadata.region_group_ids = [-1, 0]

1874 + metadata.region_names = ["a", "b"]

1875 + metadata.region_mem_types = ["VRAM", "VRAM"]

1876 +

1877 + worker._align_remote_regions_by_layer(metadata)

1878 + aligned = msgspec.msgpack.encode(metadata)

1879 + worker._align_remote_regions_by_layer(metadata)

1880 +

1881 + assert msgspec.msgpack.encode(metadata) == aligned

1882 +

1883 +

1884 +def test_layer_alignment_rejects_missing_local_layer():

1885 + worker = _layer_routing_worker([["l0"], ["l1"]], {"l0": 0, "l1": 1})

1886 + metadata = _agent_metadata([["l0"]], [0xA000], [128])

1887 +

1888 + with pytest.raises(AssertionError, match="missing locally owned layers"):

1889 + worker._align_remote_regions_by_layer(metadata)

1890 +

1891 +

1892 +def test_layer_alignment_rejects_duplicate_remote_layer():

1893 + worker = _layer_routing_worker([["a"]], {"a": 0})

1894 + metadata = _agent_metadata([["a"], ["a"]], [0xA000, 0xB000], [128, 128])

1895 +

1896 + with pytest.raises(AssertionError, match="in multiple regions"):

1897 + worker._align_remote_regions_by_layer(metadata)

1898 +

1899 +

1900 +def test_layer_alignment_rejects_inconsistent_remote_metadata():

1901 + worker = _layer_routing_worker([["a"]], {"a": 0})

1902 + metadata = _agent_metadata([["a"], ["b"]], [0xA000], [128])

1903 +

1904 + with pytest.raises(AssertionError, match="lengths disagree"):

1905 + worker._align_remote_regions_by_layer(metadata)

1906 +

1907 +

1908 +@pytest.mark.parametrize(

1909 + "field",

1910 + ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"],

1911 +)

1912 +def test_layer_alignment_rejects_inconsistent_region_geometry(field):

1913 + worker = _layer_routing_worker([["a"]], {"a": 0})

1914 + metadata = _agent_metadata([["a"], ["b"]], [0xA000, 0xB000], [128, 128])

1915 + setattr(metadata, field, [])

1916 + original = msgspec.msgpack.encode(metadata)

1917 +

1918 + with pytest.raises(AssertionError, match="lengths disagree"):

1919 + worker._align_remote_regions_by_layer(metadata)

1920 + assert msgspec.msgpack.encode(metadata) == original

1921 +

1922 +

1923 +def test_set_region_layers_rejects_duplicate_local_layer():

1924 + with pytest.raises(AssertionError, match="spans multiple NIXL regions"):

1925 + _layer_routing_worker([["a"], ["a"]], {"a": 0})

1926 +

1927 +

1928 +def test_set_region_layers_rejects_layer_outside_any_kv_group():

1929 + with pytest.raises(AssertionError, match="outside any local group"):

1930 + _layer_routing_worker([["a"], ["b"]], {"a": 0})

1931 +

1932 +

1933 +@pytest.mark.parametrize(

1934 + ("local_block_size", "remote_block_size", "remote_tp_size", "error"),

1935 + [

1936 + (32, 16, 1, "identical P/D block sizes"),

1937 + (16, 32, 1, "identical P/D block sizes"),

1938 + (16, 16, 2, "decode TP greater"),

1939 + ],

1940 +)

1941 +def test_layer_handshake_rejects_unsupported_geometry(

1942 + local_block_size: int, remote_block_size: int, remote_tp_size: int, error: str

1943 +):

1944 + """Reject unsupported peers without registering agents or transfer state."""

1945 + metadata = _agent_metadata([["a"]], [0xA000], [128])

1946 + metadata.block_size = remote_block_size

1947 + worker = _layer_routing_worker([["a"]], {"a": 0})

1948 + worker.block_size = local_block_size

1949 + worker.block_len_per_layer = [128]

1950 + worker.use_mla = False

1951 + worker.nixl_wrapper = MagicMock()

1952 + worker.kv_caches_base_addr = defaultdict(dict)

1953 + worker.dst_num_blocks = {}

1954 + worker.tp_mappings = {}

1955 + worker.transfer_topo = TransferTopology(

1956 + tp_rank=0,

1957 + tp_size=1,

1958 + block_size=local_block_size,

1959 + engine_id=worker.engine_id,

1960 + is_mla=False,

1961 + is_mamba=False,

1962 + total_num_kv_heads=8,

1963 + attn_backends=[],

1964 + )

1965 +

1966 + with pytest.raises(NotImplementedError, match=error):

1967 + worker.add_remote_agent(metadata, remote_tp_size=remote_tp_size)

1968 + worker.nixl_wrapper.add_remote_agent.assert_not_called()

1969 + worker.nixl_wrapper.prep_xfer_dlist.assert_not_called()

1970 + assert not worker.tp_mappings

1971 + assert not worker.dst_num_blocks

1972 + assert not worker.kv_caches_base_addr

1973 + with pytest.raises(KeyError):

1974 + worker.transfer_topo.get_engine_info(metadata.engine_id)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+216 -40

@@ -164,6 +164,10 @@ class NixlBaseConnectorWorker:

164 # Overridden by NixlPushConnectorWorker.

165 _TRANSFER_MODE: str = "pull"

166

167 + # Layer-name routing is supported only by NixlPushConnector for HMA KV

168 + # caches under PP.

169 + _supports_pp_hma = False

170 +

171 def _compute_desc_ids(

172 self,

173 block_ids: BlockIds,

@@ -191,28 +195,42 @@ class NixlBaseConnectorWorker:

195 num_blocks = dst_num_blocks

196 if block_size_ratio is not None:

197 num_blocks = int(num_blocks * block_size_ratio)

194 - num_regions = self.num_regions or len(self.block_len_per_layer)

198 + use_layer_name_routing = bool(self._transfer_layer_group_ids)

199 + num_desc_regions = (

200 + len(self._transfer_layer_group_ids)

201 + if use_layer_name_routing

202 + else self.num_regions or len(self.block_len_per_layer)

203 + )

204 if region_num_blocks is None:

196 - region_num_blocks = [num_blocks] * num_regions

205 + region_num_blocks = [num_blocks] * num_desc_regions

206 elif block_size_ratio is not None:

207 region_num_blocks = [

208 int(count * block_size_ratio) for count in region_num_blocks

209 ]

201 - if region_num_blocks:

202 - num_regions = len(region_num_blocks)

203 - assert len(region_num_blocks) == num_regions

210 + if region_num_blocks and not use_layer_name_routing:

211 + num_desc_regions = len(region_num_blocks)

212 + assert len(region_num_blocks) == num_desc_regions

213 region_offsets = np.cumsum([0, *region_num_blocks[:-1]])

214 num_fa_descs = sum(region_num_blocks)

215

216 + if use_layer_name_routing:

217 + # Descriptors in layer order, using each layer's KV-cache group.

218 + return np.concatenate(

219 + [

220 + np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k]

221 + for k, g in enumerate(self._transfer_layer_group_ids)

222 + ]

223 + )

224 +

225 # All-attention fast path.

226 if num_ssm_regions == 0:

227 if region_group_ids is None:

228 region_group_ids = self.region_group_ids

229 if uses_region_group_mapping is None:

230 uses_region_group_mapping = self._uses_region_group_mapping

213 - if not region_group_ids and num_regions == 1:

231 + if not region_group_ids and num_desc_regions == 1:

232 region_group_ids = [0]

215 - assert len(region_group_ids) == num_regions

233 + assert len(region_group_ids) == num_desc_regions

234 if uses_region_group_mapping is None:

235 uses_region_group_mapping = len(set(region_group_ids)) > 1

236 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)

@@ -392,6 +410,103 @@ class NixlBaseConnectorWorker:

410 """

411 return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]

412

413 + def _set_region_layers(self, region_layers: list[list[str]]) -> None:

414 + layer_names = [layer_name for region in region_layers for layer_name in region]

415 + assert len(layer_names) == len(set(layer_names)), (

416 + "A KV cache layer spans multiple NIXL regions"

417 + )

418 + self.region_members: list[list[str]] = region_layers

419 + if not self._requires_layer_name_routing():

420 + return

421 + group_by_layer = self.kv_cache_config.transfer_group_index_by_layer

422 + ungrouped = [name for name in layer_names if name not in group_by_layer]

423 + assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}"

424 + self._transfer_layer_names = tuple(layer_names)

425 + self._transfer_layer_region_indices = tuple(

426 + region_idx

427 + for region_idx, region in enumerate(region_layers)

428 + for _ in region

429 + )

430 + self._transfer_layer_group_ids = tuple(

431 + group_by_layer[name] for name in layer_names

432 + )

433 +

434 + def _requires_layer_name_routing(self) -> bool:

435 + """Whether PP push must match HMA layers by name, not region index."""

436 + return (

437 + self._supports_pp_hma

438 + and self.pp_size > 1

439 + and self._is_hma_required

440 + and not self._has_mamba

441 + )

442 +

443 + def _align_remote_regions_by_layer(

444 + self, nixl_agent_meta: NixlAgentMetadata

445 + ) -> None:

446 + """Select remote regions for this PP stage in local layer order."""

447 + remote_region_layers = nixl_agent_meta.region_members

448 + # Region-index routing here would silently transfer stale KV.

449 + assert remote_region_layers, "Remote advertised no region_members"

450 + assert (

451 + len(nixl_agent_meta.kv_caches_base_addr)

452 + == len(nixl_agent_meta.block_lens)

453 + == len(nixl_agent_meta.block_strides)

454 + == len(remote_region_layers)

455 + ), "Remote region metadata lengths disagree"

456 + assert all(

457 + values is None or len(values) == len(remote_region_layers)

458 + for values in (

459 + nixl_agent_meta.region_num_blocks,

460 + nixl_agent_meta.region_group_ids,

461 + nixl_agent_meta.region_names,

462 + nixl_agent_meta.region_mem_types,

463 + )

464 + ), "Remote region metadata lengths disagree"

465 +

466 + remote_region_by_layer: dict[str, int] = {}

467 + for region_idx, layer_names in enumerate(remote_region_layers):

468 + for layer_name in layer_names:

469 + assert layer_name not in remote_region_by_layer, (

470 + f"Remote advertised layer {layer_name!r} in multiple regions"

471 + )

472 + remote_region_by_layer[layer_name] = region_idx

473 +

474 + missing = [

475 + name

476 + for name in self._transfer_layer_names

477 + if name not in remote_region_by_layer

478 + ]

479 + assert not missing, f"Remote is missing locally owned layers: {missing}"

480 +

481 + remote_regions = [

482 + remote_region_by_layer[name] for name in self._transfer_layer_names

483 + ]

484 + nixl_agent_meta.kv_caches_base_addr = [

485 + nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions

486 + ]

487 + nixl_agent_meta.block_lens = [

488 + nixl_agent_meta.block_lens[i] for i in remote_regions

489 + ]

490 + nixl_agent_meta.block_strides = [

491 + nixl_agent_meta.block_strides[i] for i in remote_regions

492 + ]

493 + if nixl_agent_meta.region_num_blocks is not None:

494 + nixl_agent_meta.region_num_blocks = [

495 + nixl_agent_meta.region_num_blocks[i] for i in remote_regions

496 + ]

497 + if nixl_agent_meta.region_group_ids is not None:

498 + nixl_agent_meta.region_group_ids = [

499 + nixl_agent_meta.region_group_ids[i] for i in remote_regions

500 + ]

501 + if nixl_agent_meta.region_mem_types is not None:

502 + nixl_agent_meta.region_mem_types = [

503 + nixl_agent_meta.region_mem_types[i] for i in remote_regions

504 + ]

505 + if nixl_agent_meta.region_names is not None:

506 + nixl_agent_meta.region_names = list(self._transfer_layer_names)

507 + # One layer per region keeps a second alignment pass a no-op.

508 + nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names]

509 +

510 def __init__(

511 self,

512 vllm_config: "VllmConfig",

@@ -638,18 +753,26 @@ class NixlBaseConnectorWorker:

753 # transfers into the matching sub-range of a PP=1 remote's regions.

754 self.pp_size = vllm_config.parallel_config.pipeline_parallel_size

755 self._remote_region_offset = 0

641 - # PP push slices regions per layer (uniform count); HMA breaks that.

642 - if self.pp_size > 1 and self._is_hma_required:

643 - raise NotImplementedError(

644 - "NixlPushConnector does not support pipeline_parallel_size > 1 "

645 - "with hybrid KV cache layouts (HMA) yet."

646 - )

647 - # Decode-side PP is unsupported (completions counted per consumer rank).

648 - if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1:

649 - raise NotImplementedError(

650 - "NixlPushConnector consumer (decode) does not support "

651 - "pipeline_parallel_size > 1."

652 - )

756 + if self.pp_size > 1:

757 + if self._is_hma_required and not self._supports_pp_hma:

758 + raise NotImplementedError(

759 + "NixlConnector (pull) does not support pipeline_parallel_size "

760 + "> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector "

761 + "for PP + HMA."

762 + )

763 + # PP push routes HMA (hybrid) attention layouts by layer name;

764 + # Mamba/SSM hybrids are not yet supported under PP.

765 + if self._has_mamba:

766 + raise NotImplementedError(

767 + "NixlPushConnector does not support pipeline_parallel_size > 1 "

768 + "with Mamba/SSM hybrid KV cache layouts yet."

769 + )

770 + # Decode-side PP is unsupported (completions counted per consumer rank).

771 + if vllm_config.kv_transfer_config.kv_role == "kv_consumer":

772 + raise NotImplementedError(

773 + "NixlPushConnector consumer (decode) does not support "

774 + "pipeline_parallel_size > 1."

775 + )

776 # Keep heartbeat handshakes to a PP-sharded producer notif-only.

777 self._hb_handshake_notif_only = False

778

@@ -773,6 +896,13 @@ class NixlBaseConnectorWorker:

896 # within a block (BLHNC/BHLNC), where stride > block_len.

897 self.block_stride_per_layer = list[int]()

898

899 + # Region layer membership is populated for HMA layouts.

900 + self.region_members = []

901 + # Local layer order, set by ``_set_region_layers`` only when

902 + # routing by layer name. Empty means region-index routing.

903 + self._transfer_layer_names = tuple[str, ...]()

904 + self._transfer_layer_region_indices = tuple[int, ...]()

905 + self._transfer_layer_group_ids = tuple[int, ...]()

906 # Per-engine TP mappings. Generated during handshake.

907 self.tp_mappings: dict[EngineId, TPMapping] = {}

908

@@ -1306,6 +1436,11 @@ class NixlBaseConnectorWorker:

1436 ):

1437 compressed_region_owners.setdefault(cache.data_ptr(), cache)

1438

1439 + track_region_layers = (

1440 + self._supports_pp_hma and self._is_hma_required and not self._has_mamba

1441 + )

1442 + region_layers: list[list[str]] = []

1443 +

1444 # K and V are packed into the content dim, so each attention layer is a

1445 # single NIXL region whose block transfers as one unit. Mamba layers instead

1446 # register separate conv/ssm sub-regions (see `_build_mamba_local`).

@@ -1508,6 +1643,12 @@ class NixlBaseConnectorWorker:

1643 self._region_is_mla.append(is_mla_region)

1644 region_mem_types.append(mem_type)

1645

1646 + if track_region_layers:

1647 + if region_index == len(region_layers):

1648 + region_layers.append([layer_name])

1649 + else:

1650 + region_layers[region_index].append(layer_name)

1651 +

1652 if isinstance(layer_spec, MambaSpec):

1653 if layer_spec.tp_replicated:

1654 assert self._is_csa_linear

@@ -1563,7 +1704,9 @@ class NixlBaseConnectorWorker:

1704 if self._has_mamba:

1705 self.region_num_blocks = [self.num_blocks] * self.num_regions

1706

1566 - if self.pp_size > 1:

1707 + self._set_region_layers(region_layers)

1708 +

1709 + if self.pp_size > 1 and not self._is_hma_required:

1710 start_layer, end_layer = self.model_config.get_layers_start_end_indices(

1711 self.vllm_config.parallel_config

1712 )

@@ -1573,7 +1716,12 @@ class NixlBaseConnectorWorker:

1716 self._remote_region_offset = regions_per_layer * start_layer

1717

1718 # Total local FA descriptors (boundary between FA and mamba descs).

1576 - self.num_descs = sum(self.region_num_blocks)

1719 + xfer_region_num_blocks = (

1720 + [self.region_num_blocks[i] for i in self._transfer_layer_region_indices]

1721 + if self._transfer_layer_region_indices

1722 + else self.region_num_blocks

1723 + )

1724 + self.num_descs = sum(xfer_region_num_blocks)

1725

1726 self._mixed_mem_types = len(set(region_mem_types)) > 1

1727 if self._mixed_mem_types:

@@ -1597,7 +1745,7 @@ class NixlBaseConnectorWorker:

1745

1746 self.device_kv_caches = kv_caches

1747 self.dst_num_blocks[self.engine_id] = self.num_blocks

1600 - self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks

1748 + self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks

1749 self.dst_region_group_ids[self.engine_id] = self.region_group_ids

1750 self.dst_uses_region_group_mapping[self.engine_id] = (

1751 self._uses_region_group_mapping

@@ -1647,6 +1795,7 @@ class NixlBaseConnectorWorker:

1795 region_mem_types=self.region_mem_types,

1796 dcp_size=self.dcp_size,

1797 pcp_size=self.pcp_size,

1798 + region_members=self.region_members,

1799 )

1800 # Wrap metadata in payload with hash for defensive decoding

1801 assert self.compat_hash is not None

@@ -1805,12 +1954,16 @@ class NixlBaseConnectorWorker:

1954 base_addresses: list[int],

1955 block_size_ratio: int,

1956 ) -> np.ndarray:

1808 - """Build local FA descriptors for all layers as an Nx3 uint64 array."""

1957 + """Build local FA descriptors as an Nx3 uint64 array."""

1958 assert self.transfer_topo is not None

1959 assert base_addresses, "Local KV cache base addresses must not be empty."

1960 device_id = self.device_id

1961 + region_indices = self._transfer_layer_region_indices or range(

1962 + len(base_addresses)

1963 + )

1964 parts: list[np.ndarray] = []

1813 - for i, base_addr in enumerate(base_addresses):

1965 + for i in region_indices:

1966 + base_addr = base_addresses[i]

1967 block_len = self.block_len_per_layer[i] // block_size_ratio

1968 logical_blocks = np.repeat(

1969 np.arange(self.region_num_blocks[i], dtype=np.uint64),

@@ -1834,7 +1987,7 @@ class NixlBaseConnectorWorker:

1987 nixl_agent_meta: NixlAgentMetadata,

1988 block_size_ratio: int,

1989 ) -> np.ndarray:

1837 - """Build remote FA descriptors for all layers as an Nx3 uint64 array."""

1990 + """Build remote FA descriptors as an Nx3 uint64 array."""

1991 assert self.transfer_topo is not None

1992 assert nixl_agent_meta.kv_caches_base_addr, (

1993 "Remote KV cache base addresses must not be empty."

@@ -1851,10 +2004,14 @@ class NixlBaseConnectorWorker:

2004 device_id = nixl_agent_meta.device_id

2005 block_strides = nixl_agent_meta.block_strides

2006 parts: list[np.ndarray] = []

2007 + use_layer_name_routing = bool(self._transfer_layer_region_indices)

2008 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1855 - replicated = self._is_region_replicated(i)

2009 + local_region = (

2010 + self._transfer_layer_region_indices[i] if use_layer_name_routing else i

2011 + )

2012 + replicated = self._is_region_replicated(local_region)

2013 # Read our whole local region size from remote..

1857 - local_block_len = self.block_len_per_layer[i]

2014 + local_block_len = self.block_len_per_layer[local_region]

2015 remote_kv_block_len = local_block_len // block_size_ratio

2016 if block_size_ratio > 1:

2017 # ..using remote kv_block_len as transfer unit

@@ -2021,9 +2178,22 @@ class NixlBaseConnectorWorker:

2178 )

2179 return self._remote_agents[engine_id][(0, remote_tp_rank)]

2180

2181 + assert self.transfer_topo is not None

2182 + transfer_topo = self.transfer_topo

2183 # Number of physical regions registered locally (one per layer/tensor).

2184 num_local_regions = len(self.block_len_per_layer)

2026 - if (

2185 + if self._transfer_layer_region_indices:

2186 + if self.block_size != nixl_agent_meta.block_size:

2187 + raise NotImplementedError(

2188 + "Attention-HMA push requires identical P/D block sizes."

2189 + )

2190 + if remote_tp_size > transfer_topo.tp_size and not self.use_mla:

2191 + raise NotImplementedError(

2192 + "Attention-HMA push does not support decode TP greater than "

2193 + "prefill TP yet"

2194 + )

2195 + self._align_remote_regions_by_layer(nixl_agent_meta)

2196 + elif (

2197 self.pp_size > 1

2198 and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions

2199 ):

@@ -2038,6 +2208,10 @@ class NixlBaseConnectorWorker:

2208 ]

2209 nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]

2210 nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]

2211 + if nixl_agent_meta.region_num_blocks is not None:

2212 + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[

2213 + start:end

2214 + ]

2215 if nixl_agent_meta.region_group_ids is not None:

2216 nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[

2217 start:end

@@ -2048,14 +2222,8 @@ class NixlBaseConnectorWorker:

2222 nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[

2223 start:end

2224 ]

2051 - if nixl_agent_meta.region_num_blocks is not None:

2052 - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[

2053 - start:end

2054 - ]

2225

2226 ### Register remote engine in TransferTopology (idempotent).

2057 - assert self.transfer_topo is not None

2058 - transfer_topo = self.transfer_topo

2227 physical_blocks_per_logical = (

2228 nixl_agent_meta.physical_blocks_per_logical_kv_block

2229 )

@@ -2349,6 +2517,9 @@ class NixlBaseConnectorWorker:

2517 # the per-rank KV head ratio rather than the raw tp_ratio, because GQA

2518 # replication caps per-rank heads at 1 when tp > total_kv_heads

2519 # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.

2520 + local_regions = self._transfer_layer_region_indices or range(

2521 + len(self.block_len_per_layer)

2522 + )

2523 if self._has_mamba and self.use_mla:

2524 # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are

2525 # kernel-granularity views of the mamba-unified page. The MLA

@@ -2365,7 +2536,7 @@ class NixlBaseConnectorWorker:

2536 f"remote={nixl_agent_meta.block_lens}."

2537 )

2538 elif not self._has_mamba:

2368 - assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), (

2539 + assert len(local_regions) == len(nixl_agent_meta.block_lens), (

2540 "Number of KV layers must match between prefill and decode"

2541 )

2542 model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(

@@ -2374,8 +2545,9 @@ class NixlBaseConnectorWorker:

2545 total_kv_heads = self.transfer_topo.total_num_kv_heads

2546 local_heads = self.transfer_topo.local_physical_heads

2547 remote_heads = max(1, total_kv_heads // remote_tp_size)

2377 - for i, local_len in enumerate(self.block_len_per_layer):

2378 - replicated = model_replicated or self._is_region_replicated(i)

2548 + for i, region_idx in enumerate(local_regions):

2549 + local_len = self.block_len_per_layer[region_idx]

2550 + replicated = model_replicated or self._is_region_replicated(region_idx)

2551 remote_len = nixl_agent_meta.block_lens[i]

2552 if replicated:

2553 assert local_len // block_size_ratio == remote_len, (

@@ -2406,8 +2578,7 @@ class NixlBaseConnectorWorker:

2578

2579 # TP workers that handhshake with same remote have same #blocks.

2580 assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks

2409 - # Same number of regions/~layers.

2410 - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)

2581 + assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions)

2582 num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)

2583 if nixl_agent_meta.region_num_blocks is not None:

2584 assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions

@@ -2423,7 +2594,12 @@ class NixlBaseConnectorWorker:

2594 if nixl_agent_meta.region_group_ids is not None:

2595 assert len(nixl_agent_meta.region_group_ids) == num_remote_regions

2596 if nixl_agent_meta.region_names is not None:

2426 - assert nixl_agent_meta.region_names == self.region_names, (

2597 + expected_names = (

2598 + list(self._transfer_layer_names)

2599 + if self._transfer_layer_names

2600 + else self.region_names

2601 + )

2602 + assert nixl_agent_meta.region_names == expected_names, (

2603 "NIXL transfer regions must name the same model caches in the "

2604 "same order"

2605 )

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

+5 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Metadata dataclasses and helpers for the NIXL connector."""

5 -from dataclasses import dataclass

5 +from dataclasses import dataclass, field

6 from typing import Any

8 from vllm.config import VllmConfig

@@ -47,8 +47,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"

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 +# 12: Add per-region member names for PP push

51 #

51 -NIXL_CONNECTOR_VERSION: int = 11

52 +NIXL_CONNECTOR_VERSION: int = 12

55 @dataclass

@@ -71,6 +72,8 @@ class NixlAgentMetadata:

72 region_mem_types: list[str] | None = None

73 dcp_size: int = 1

74 pcp_size: int = 1

75 + # Layer names sharing each advertised region, in region order.

76 + region_members: list[list[str]] = field(default_factory=list)

79 @dataclass

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

+3 -1

@@ -83,6 +83,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

83 # Distinguishes push from pull in the NIXL compatibility hash.

84 _TRANSFER_MODE: str = "push"

86 + _supports_pp_hma = True

87 +

88 def __init__(

89 self,

90 vllm_config: "VllmConfig",

@@ -530,7 +532,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

532 local_region_groups = self.region_group_ids

533 remote_region_groups = self.dst_region_group_ids[engine_id]

534 groups_differ = local_region_groups != remote_region_groups

533 - if groups_differ:

535 + if groups_differ and not self._transfer_layer_group_ids:

536 raise NotImplementedError(

537 "NixlPushConnector does not support different producer and "

538 "consumer cache-group layouts"