@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6d3ea3c2c9

[KVConnector][NIXL] Support packed MLA KV layouts in pipeline-parallel push prefill (#50499)

Signed-off-by: zixi-qi <zixi@inferact.ai>

qizixi committed Sep 28, 2026 at 09:50 UTC 6d3ea3c2c9037deadb9fb70d40bc69b7947af463

6 files changed +357 -28

docs/design/nixl_kv_push_connector.md

+29 -11

@@ -195,35 +195,53 @@ combined with a hybrid (HMA) KV layout:

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

198 +**by layer name** 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

201 +needs layer routing derives its transfer layer order 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

208 +layer order. Descriptor offsets use each layer'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.

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

211 +

212 +Packed MLA caches (block-outermost layouts such as DeepSeek-V4's) store the

213 +pages of every layer for one block in a single row. A PP stage holds only its

214 +own layers, so its rows are shorter and each layer sits at a different offset:

215 +

216 +```text

217 +D (PP=1), block b: | L0 | L1 | L2 | L3 | L4 | L5 |

218 +P (stage 2), block b: | L3 | L4 | L5 |

219 +```

220 +

221 +`PP=1` peers transfer whole rows. They also advertise where each layer's page

222 +sits in the row (`NixlAgentMetadata.packed_member_layouts`), so a PP producer

223 +writes each of its pages to the matching slot in D's row. Page sizes must

224 +match; offsets, row strides, and block counts may differ. Because PP stages can

225 +list the same attention backends in a different order, packed push compares the

226 +sorted set of backend names during the handshake. Pull mode is unchanged.

227

228 Invariants enforced when the remote regions are aligned:

229

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

230 +* every locally owned layer must be advertised exactly once by the

231 + remote; a missing layer fails the handshake rather than silently

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

233 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

234 +* a remote that omits layer metadata while the local layout requires

235 + layer routing fails the handshake instead of falling back to

236 region-index routing;

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

237 +* layer order is a property of the local layout alone, so the same local

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

239 remote descriptor list is rebuilt per rank.

240

241 Decode-side PP is unsupported because completions are counted per

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

242 +consumer rank. Mamba/SSM hybrids are unsupported under PP. Packed PP push

243 +requires MLA caches and equal P/D block sizes. MLA pages are replicated

244 +across TP ranks, so producer TP1 to consumer TP2 is supported.

245

246 ## Scheduler-side responsibilities

247

docs/features/nixl_connector_compatibility.md

+2 -1

@@ -120,7 +120,8 @@ Current push PP + HMA limitations:

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.

123 +- Non-MLA attention-HMA layer routing requires decode TP to be no greater than prefill TP.

124 +- Packed layouts under PP require MLA caches. Local and remote packed block strides may differ; layer pages must have equal sizes. MLA supports decode TP greater than prefill TP because its KV is replicated.

125

126 ### Quantized KV cache

127

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

+203 -2

@@ -16,10 +16,27 @@ from threading import Event, Lock

16 from types import SimpleNamespace

17 from unittest.mock import MagicMock, patch

19 +import msgspec

20 import numpy as np

21 import pytest

22 import torch

24 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw

25 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

26 + NixlAgentMetadata,

27 +)

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

29 + NixlPushConnectorWorker,

30 +)

31 +from vllm.v1.kv_cache_interface import (

32 + KVCacheConfig,

33 + KVCacheGroupSpec,

34 + KVCacheLayout,

35 + KVCacheTensor,

36 + MLAAttentionSpec,

37 + create_kv_cache_views,

38 +)

39 +

40 from .utils import create_vllm_config

@@ -102,6 +119,189 @@ class _RecordingNixl:

119 pass

120

121

122 +def _make_packed_mla_view_worker(

123 + layouts,

124 + block_stride,

125 + *,

126 + num_blocks=4,

127 + pp_size=1,

128 + backend_names=("FLASHMLA",),

129 +):

130 + """Register real strided views; only NIXL and distributed runtime are fake."""

131 + block_size = 16

132 + raw = torch.zeros(num_blocks * block_stride, dtype=torch.int8)

133 + tensors, groups, caches = [], [], {}

134 + for name, (offset, page_size) in layouts.items():

135 + spec = MLAAttentionSpec(

136 + block_size=block_size,

137 + num_kv_heads=1,

138 + head_size=page_size // block_size,

139 + dtype=torch.uint8,

140 + )

141 + tensor = KVCacheTensor(

142 + size=raw.nbytes,

143 + layers=[name],

144 + layer_stride=page_size,

145 + block_stride=block_stride,

146 + offset=offset,

147 + )

148 + tensors.append(tensor)

149 + groups.append(KVCacheGroupSpec([name], spec))

150 + (caches[name],) = create_kv_cache_views(

151 + raw, spec, num_blocks, KVCacheLayout.BLHNC, tensor

152 + )

153 +

154 + config = create_vllm_config(block_size=block_size)

155 + config.parallel_config.pipeline_parallel_size = pp_size

156 + config.kv_transfer_config.kv_role = "kv_producer"

157 + config.cache_config.kv_cache_layout = "BLHNC"

158 + backends = []

159 + for name in backend_names:

160 + backend = MagicMock()

161 + backend.get_supported_kernel_block_sizes.return_value = [block_size]

162 + backend.get_name.return_value = name

163 + backend.full_cls_name.return_value = f"fake.{name}"

164 + backends.append(backend)

165 + with (

166 + patch.object(bw, "NixlWrapper", _RecordingNixl),

167 + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),

168 + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),

169 + patch.object(bw, "get_current_attn_backends", return_value=backends),

170 + patch("threading.Thread"),

171 + ):

172 + worker = NixlPushConnectorWorker(

173 + config,

174 + "producer" if pp_size > 1 else "consumer",

175 + KVCacheConfig(num_blocks, tensors, groups),

176 + )

177 + worker.use_mla = True

178 + worker.register_kv_caches(caches)

179 + return worker, raw

180 +

181 +

182 +@pytest.mark.cpu_test

183 +@pytest.mark.parametrize("num_layers", [1, 2])

184 +def test_packed_mla_pp1_push_peer_transfers_whole_rows(num_layers):

185 + layouts = {f"L{i}": (i * 128, 128) for i in range(num_layers)}

186 + stride = num_layers * 128

187 + worker, raw = _make_packed_mla_view_worker(layouts, stride)

188 + assert worker._registered_descs == [[(raw.data_ptr(), raw.nbytes, 0, "")]]

189 + assert worker.src_blocks_data.tolist() == [

190 + [raw.data_ptr() + block * stride, stride, 0] for block in range(4)

191 + ]

192 + assert worker._transfer_layer_group_ids == ()

193 + metadata = msgspec.msgpack.decode(

194 + worker.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata

195 + )

196 + assert metadata.packed_member_layouts == layouts

197 +

198 +

199 +@pytest.mark.cpu_test

200 +@pytest.mark.parametrize(

201 + "remote_backends, compatible",

202 + [

203 + (("INDEXER", "FLASHMLA"), True),

204 + (("FLASHMLA", "OTHER"), False),

205 + (("FLASHMLA",), False),

206 + ],

207 +)

208 +def test_packed_push_compatibility_hash_uses_all_backends_in_stable_order(

209 + remote_backends, compatible

210 +):

211 + producer, _ = _make_packed_mla_view_worker(

212 + {"L0": (0, 128)},

213 + 128,

214 + pp_size=2,

215 + backend_names=("FLASHMLA", "INDEXER"),

216 + )

217 + consumer, _ = _make_packed_mla_view_worker(

218 + {"L0": (0, 128)}, 128, backend_names=remote_backends

219 + )

220 + assert (producer.compat_hash == consumer.compat_hash) is compatible

221 +

222 +

223 +@pytest.mark.cpu_test

224 +@pytest.mark.parametrize("p_num_blocks, d_num_blocks", [(4, 4), (4, 6), (6, 4)])

225 +def test_packed_mla_pp_pairs_asymmetric_strides_and_overlapping_layers(

226 + p_num_blocks, d_num_blocks

227 +):

228 + # Different cache groups overlay pages of different sizes at the same address.

229 + # PP also changes both the placement and the stride of the matching D pages.

230 + producer, p_raw = _make_packed_mla_view_worker(

231 + {"L2": (0, 128), "L2.swa": (0, 64), "L3": (128, 64)},

232 + 192,

233 + num_blocks=p_num_blocks,

234 + pp_size=2,

235 + )

236 + consumer, d_raw = _make_packed_mla_view_worker(

237 + {"L0": (0, 128), "L2": (128, 128), "L2.swa": (0, 64), "L3": (64, 64)},

238 + 256,

239 + num_blocks=d_num_blocks,

240 + )

241 + assert producer.block_len_per_layer == [128, 64, 64]

242 + assert producer._transfer_layer_group_ids == (0, 1, 2)

243 + assert producer.num_descs == 3 * p_num_blocks

244 + local_ids = producer._compute_desc_ids(

245 + [[1], [2], [3]],

246 + producer.num_blocks,

247 + None,

248 + 1,

249 + region_num_blocks=producer.dst_region_num_blocks[producer.engine_id],

250 + )

251 + assert producer.src_blocks_data[local_ids].tolist() == [

252 + [p_raw.data_ptr() + 192, 128, 0],

253 + [p_raw.data_ptr() + 2 * 192, 64, 0],

254 + [p_raw.data_ptr() + 128 + 3 * 192, 64, 0],

255 + ]

256 + # Each decoder TP rank receives the complete MLA page, without head slicing.

257 + for rank in (0, 1):

258 + metadata = msgspec.msgpack.decode(

259 + consumer.xfer_handshake_metadata.agent_metadata_bytes,

260 + type=NixlAgentMetadata,

261 + )

262 + producer.add_remote_agent(metadata, remote_tp_rank=rank, remote_tp_size=2)

263 + assert metadata.region_num_blocks == [d_num_blocks] * 3

264 + assert metadata.region_names == ["L2", "L2.swa", "L3"]

265 + remote_ids = producer._compute_desc_ids(

266 + [[3], [1], [2]],

267 + metadata.num_blocks,

268 + None,

269 + 1,

270 + region_num_blocks=producer.dst_region_num_blocks[consumer.engine_id],

271 + )

272 + handle = producer.dst_xfer_side_handles[consumer.engine_id][rank]

273 + remote_descs = producer.nixl_wrapper.dlists[handle]

274 + assert len(remote_descs) == 3 * d_num_blocks

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

276 + [d_raw.data_ptr() + 128 + 3 * 256, 128, 0],

277 + [d_raw.data_ptr() + 256, 64, 0],

278 + [d_raw.data_ptr() + 64 + 2 * 256, 64, 0],

279 + ]

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

281 + producer._align_remote_regions_by_layer(metadata)

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

283 +

284 +

285 +@pytest.mark.cpu_test

286 +@pytest.mark.parametrize("remote_block_size", [8, 32])

287 +def test_packed_mla_rejects_unequal_block_sizes_before_peer_registration(

288 + remote_block_size,

289 +):

290 + producer, _ = _make_packed_mla_view_worker({"L0": (0, 128)}, 128, pp_size=2)

291 + metadata = msgspec.msgpack.decode(

292 + producer.xfer_handshake_metadata.agent_metadata_bytes, type=NixlAgentMetadata

293 + )

294 + metadata.engine_id = "remote"

295 + metadata.block_size = remote_block_size

296 + with pytest.raises(NotImplementedError, match="identical P/D block sizes"):

297 + producer.add_remote_agent(metadata)

298 + assert len(producer.nixl_wrapper.dlists) == 1 # Only the local list exists.

299 + assert producer.tp_mappings == {}

300 + assert "remote" not in producer.dst_num_blocks

301 + with pytest.raises(KeyError):

302 + producer.transfer_topo.get_engine_info("remote")

303 +

304 +

305 @pytest.mark.cpu_test

306 def test_local_descriptors_follow_each_region_pool_capacity():

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

@@ -200,8 +400,8 @@ def _register_overlaid_mla_worker(

400 worker.src_xfer_handles_by_block_size = {}

401 worker.kv_caches_base_addr = defaultdict(dict)

402 worker._mamba_ssm_size = (0, 0)

203 - worker.kv_cache_layout = "NHD"

204 - worker.host_buffer_kv_cache_layout = "NHD"

403 + worker.kv_cache_layout = "LBNHC"

404 + worker.host_buffer_kv_cache_layout = "LBNHC"

405 worker._physical_blocks_per_logical_kv_block = 1

406 worker._logical_num_blocks = num_blocks

407 worker.region_mem_types = []

@@ -221,6 +421,7 @@ def _register_overlaid_mla_worker(

421 worker.device_kv_caches = {}

422 worker.pp_size = 2 if push_pp else 1

423 worker._is_hma_required = True

424 + worker._has_packed_cache = False

425 worker.dcp_size = 1

426 worker.pcp_size = 1

427 worker.kv_buffer_device = "cuda"

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+36

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

368 w._recving_transfers = defaultdict(list)

369 w._is_hma_required = False

370 w.xfer_stats = NixlKVConnectorStats()

371 + w._has_packed_cache = False

372 w._reqs_to_process = set()

373 w._reqs_to_send = {}

374 w.consumer_notification_counts_by_req = defaultdict(int)

@@ -1843,6 +1844,40 @@ def test_layer_metadata_round_trip():

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

1845

1846

1847 +@pytest.mark.parametrize(

1848 + "layouts, error",

1849 + [

1850 + ({"L0": (-1, 64)}, "escapes its block"),

1851 + ({"L0": (0, 0)}, "escapes its block"),

1852 + ({"L0": (200, 64)}, "escapes its block"),

1853 + ({"L1": (0, 64)}, "has no layout"),

1854 + ({"L0": (0, 32)}, "page sizes must match"),

1855 + ],

1856 +)

1857 +def test_packed_layer_alignment_rejects_invalid_page_layout(layouts, error):

1858 + worker = _layer_routing_worker([["L0"]], {"L0": 0})

1859 + worker.block_len_per_layer = [64]

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

1861 + metadata.packed_member_layouts = layouts

1862 + with pytest.raises(AssertionError, match=error):

1863 + worker._align_remote_regions_by_layer(metadata)

1864 +

1865 +

1866 +def test_layer_alignment_preserves_nonpacked_regions_alongside_packed_layers():

1867 + worker = _layer_routing_worker([["a"], ["b"]], {"a": 0, "b": 1})

1868 + worker.block_len_per_layer = [64, 128]

1869 + metadata = _agent_metadata(

1870 + [["a", "other"], ["b"]], [0x1000, 0x2000], [192, 128], [256, 128]

1871 + )

1872 + metadata.packed_member_layouts = {"a": (64, 64), "other": (128, 64)}

1873 +

1874 + worker._align_remote_regions_by_layer(metadata)

1875 +

1876 + assert metadata.kv_caches_base_addr == [0x1040, 0x2000]

1877 + assert metadata.block_lens == [64, 128]

1878 + assert metadata.block_strides == [256, 128]

1879 +

1880 +

1881 def test_layer_identity_gate_preserves_the_non_hma_path():

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

1883 0,

@@ -1850,6 +1885,7 @@ def test_layer_identity_gate_preserves_the_non_hma_path():

1885

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

1887 pull = object.__new__(NixlBaseConnectorWorker)

1888 + pull.pp_size = 2

1889 pull._transfer_layer_names = ()

1890 pull._transfer_layer_region_indices = ()

1891 pull._transfer_layer_group_ids = ()

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

+83 -13

@@ -165,8 +165,8 @@ class NixlBaseConnectorWorker:

165 # Overridden by NixlPushConnectorWorker.

166 _TRANSFER_MODE: str = "pull"

167

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

169 - # caches under PP.

168 + # Layer-name routing is supported only by NixlPushConnector for HMA or

169 + # packed KV caches under PP.

170 _supports_pp_hma = False

171

172 def _compute_desc_ids(

@@ -432,15 +432,22 @@ class NixlBaseConnectorWorker:

432 group_by_layer[name] for name in layer_names

433 )

434

435 - def _requires_layer_name_routing(self) -> bool:

436 - """Whether PP push must match HMA layers by name, not region index."""

435 + def _tracks_region_layers(self) -> bool:

436 + """Whether regions advertise their member layer names.

437 +

438 + Push workers do so for HMA and packed layouts, so a PP producer can

439 + match layers by name instead of by region index.

440 + """

441 return (

442 self._supports_pp_hma

439 - and self.pp_size > 1

440 - and self._is_hma_required

443 + and (self._is_hma_required or self._has_packed_cache)

444 and not self._has_mamba

445 )

446

447 + def _requires_layer_name_routing(self) -> bool:

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

449 + return self.pp_size > 1 and self._tracks_region_layers()

450 +

451 def _align_remote_regions_by_layer(

452 self, nixl_agent_meta: NixlAgentMetadata

453 ) -> None:

@@ -479,6 +486,10 @@ class NixlBaseConnectorWorker:

486 ]

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

488

489 + packed_regions = {

490 + remote_region_by_layer[name]

491 + for name in nixl_agent_meta.packed_member_layouts

492 + }

493 remote_regions = [

494 remote_region_by_layer[name] for name in self._transfer_layer_names

495 ]

@@ -505,6 +516,24 @@ class NixlBaseConnectorWorker:

516 ]

517 if nixl_agent_meta.region_names is not None:

518 nixl_agent_meta.region_names = list(self._transfer_layer_names)

519 + if nixl_agent_meta.packed_member_layouts:

520 + for i, layer_name in enumerate(self._transfer_layer_names):

521 + if remote_regions[i] not in packed_regions:

522 + continue

523 + assert layer_name in nixl_agent_meta.packed_member_layouts, (

524 + f"Remote packed layer {layer_name!r} has no layout"

525 + )

526 + offset, page_size = nixl_agent_meta.packed_member_layouts[layer_name]

527 + assert (

528 + 0 <= offset < offset + page_size <= nixl_agent_meta.block_lens[i]

529 + ), f"Remote packed layer {layer_name!r} escapes its block"

530 + local_region = self._transfer_layer_region_indices[i]

531 + assert page_size == self.block_len_per_layer[local_region], (

532 + f"Packed MLA page sizes must match for layer {layer_name!r}"

533 + )

534 + nixl_agent_meta.kv_caches_base_addr[i] += offset

535 + nixl_agent_meta.block_lens[i] = page_size

536 + nixl_agent_meta.packed_member_layouts = {}

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

538 nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names]

539

@@ -899,7 +928,10 @@ class NixlBaseConnectorWorker:

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

929 self.block_stride_per_layer = list[int]()

930

902 - # Region layer membership is populated for HMA layouts.

931 + # Block-outermost layouts (BLHNC/BHLNC) pack every layer's page into one

932 + # block row of a shared allocation.

933 + self._has_packed_cache = KVCacheLayout[self.kv_cache_layout].is_block_outermost

934 + # Region layer membership is populated for HMA and packed layouts.

935 self.region_members = []

936 # Local layer order, set by ``_set_region_layers`` only when

937 # routing by layer name. Empty means region-index routing.

@@ -1352,6 +1384,8 @@ class NixlBaseConnectorWorker:

1384

1385 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

1386 """Register the KV Cache data in nixl."""

1387 + use_layer_name_routing = self._requires_layer_name_routing()

1388 + route_packed_layers = self._has_packed_cache and use_layer_name_routing

1389 self.transfer_topo = TransferTopology(

1390 tp_rank=self.transfer_tp_rank,

1391 tp_size=self.transfer_tp_size,

@@ -1369,9 +1403,15 @@ class NixlBaseConnectorWorker:

1403 else None,

1404 is_mamba=self._has_mamba,

1405 )

1406 + backend_name = self.backend_name

1407 + if self._supports_pp_hma and self._has_packed_cache:

1408 + # PP can reorder the same backends across stages of a packed model.

1409 + backend_name = ",".join(

1410 + sorted(backend.get_name() for backend in self.attn_backends)

1411 + )

1412 self.compat_hash = compute_nixl_compatibility_hash(

1413 self.vllm_config,

1374 - self.backend_name,

1414 + backend_name,

1415 transfer_mode=self._TRANSFER_MODE,

1416 )

1417

@@ -1413,6 +1453,11 @@ class NixlBaseConnectorWorker:

1453 # CSA-linear needs separate logical regions for attention and state

1454 # aliases even though every view shares one block-major allocation.

1455 packed_storage = packed_storage and not self._is_csa_linear

1456 + if route_packed_layers and any(

1457 + not isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))

1458 + for spec in self._layer_specs.values()

1459 + ):

1460 + raise NotImplementedError("PP push with packed KV caches requires MLA")

1461

1462 layer_specs: dict[str, KVCacheSpec] = {}

1463 compressed_region_owners: dict[int, torch.Tensor] = {}

@@ -1439,10 +1484,9 @@ class NixlBaseConnectorWorker:

1484 ):

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

1486

1442 - track_region_layers = (

1443 - self._supports_pp_hma and self._is_hma_required and not self._has_mamba

1444 - )

1487 + track_region_layers = self._tracks_region_layers()

1488 region_layers: list[list[str]] = []

1489 + packed_member_layouts: dict[str, tuple[int, int]] = {}

1490

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

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

@@ -1587,11 +1631,34 @@ class NixlBaseConnectorWorker:

1631 virtual_transfer_pages = _uses_dense_virtual_transfer_pages(

1632 layer_spec, cache, physical_page_size, num_blocks

1633 )

1634 + # Push workers address the pages of packed MLA rows by layer.

1635 + packed_mla_push = (

1636 + track_region_layers

1637 + and self._has_packed_cache

1638 + and packed_storage

1639 + and storage_is_block_major

1640 + and is_mla_region

1641 + )

1642 if virtual_transfer_pages:

1643 # A compressed kernel row can contain multiple NIXL transfer pages.

1644 region_specs = [

1645 (cache.data_ptr(), physical_page_size, physical_page_size)

1646 ]

1647 + elif packed_mla_push and use_layer_name_routing:

1648 + # PP>1 prefill stage: register only this layer's page of each row.

1649 + region_specs = [

1650 + (cache.data_ptr(), physical_page_size, block_stride)

1651 + ]

1652 + elif packed_mla_push:

1653 + # PP=1 decode or prefill: register whole rows, and advertise where

1654 + # this layer's page sits so a PP>1 prefill can write into it.

1655 + storage_block_len = registration_len // num_blocks

1656 + region_specs = [

1657 + (registration_base, storage_block_len, storage_block_len)

1658 + ]

1659 + offset = cache.data_ptr() - storage_addr

1660 + assert offset >= 0 and offset + physical_page_size <= block_stride

1661 + packed_member_layouts[layer_name] = (offset, physical_page_size)

1662 elif storage_is_block_major and (

1663 not page_contiguous

1664 and ((packed_storage and is_mla_region) or not self._is_csa_linear)

@@ -1633,7 +1700,9 @@ class NixlBaseConnectorWorker:

1700 ]

1701

1702 for base_addr, block_len, block_stride in region_specs:

1636 - if base_addr in seen_base_addresses:

1703 + # A packed PP producer keeps one region per layer, so aliases at

1704 + # one address (a layer and its SWA view) stay distinct.

1705 + if base_addr in seen_base_addresses and not route_packed_layers:

1706 region_index = seen_base_addresses.index(base_addr)

1707 assert region_mem_types[region_index] == mem_type

1708 self._region_is_mla[region_index] |= is_mla_region

@@ -1717,7 +1786,7 @@ class NixlBaseConnectorWorker:

1786

1787 self._set_region_layers(region_layers)

1788

1720 - if self.pp_size > 1 and not self._is_hma_required:

1789 + if self.pp_size > 1 and not use_layer_name_routing:

1790 start_layer, end_layer = self.model_config.get_layers_start_end_indices(

1791 self.vllm_config.parallel_config

1792 )

@@ -1807,6 +1876,7 @@ class NixlBaseConnectorWorker:

1876 dcp_size=self.dcp_size,

1877 pcp_size=self.pcp_size,

1878 region_members=self.region_members,

1879 + packed_member_layouts=packed_member_layouts,

1880 )

1881 # Wrap metadata in payload with hash for defensive decoding

1882 assert self.compat_hash is not None

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

+4 -1

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

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 +# 13: Add packed-member layouts and order-independent packed-push backend hashes

52 #

52 -NIXL_CONNECTOR_VERSION: int = 12

53 +NIXL_CONNECTOR_VERSION: int = 13

56 @dataclass

@@ -74,6 +75,8 @@ class NixlAgentMetadata:

75 pcp_size: int = 1

76 # Layer names sharing each advertised region, in region order.

77 region_members: list[list[str]] = field(default_factory=list)

78 + # Packed member -> (byte offset in its region's block, bytes per page).

79 + packed_member_layouts: dict[str, tuple[int, int]] = field(default_factory=dict)

82 @dataclass