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

[Perf][KV Connector][Mooncake] Pack hybrid/MLA KV into coalesced transfer regions (#57952)

Signed-off-by: staryxchen <staryxchen@tencent.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Stary committed Sep 29, 2026 at 22:29 UTC faacc13565312d29e9596d182fc808b452a4508e
3 files changed +1386 -209
tests/v1/kv_connector/unit/test_mooncake_connector.py
+867 -167
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
23 mooncake_utils,
24 )
25 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
26 + _SHARED_REGION_GROUP_ID,
27 KVConnectorRole,
28 MooncakeConnector,
29 MooncakeConnectorMetadata,
@@ -34,7 +35,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im
35 SendBlockMeta,
36 TransferRegion,
37 _align_transfer_regions,
38 + _block_ids_for_region,
39 + _coalesce_contiguous_transfer_regions,
40 _compute_sender_transfer_plan,
41 + _has_opaque_packed_row,
42 + _pp_mismatch_hides_packed_layers,
43 _validate_asymmetric_region_lengths,
44 get_mooncake_bootstrap_addr,
45 should_launch_bootstrap_server,
@@ -48,6 +53,7 @@ from vllm.v1.kv_cache_interface import (
53 KVCacheConfig,
54 KVCacheGroupSpec,
55 KVCacheLayout,
56 + MLAAttentionSpec,
57 )
58 from vllm.v1.request import RequestStatus
59
@@ -164,6 +170,34 @@ def _make_test_kv_cache_config() -> KVCacheConfig:
170 )
171
172
173 +def _make_packed_mla_kv_cache_config(
174 + num_blocks: int, *, num_groups: int = 1
175 +) -> KVCacheConfig:
176 + spec = MLAAttentionSpec(
177 + block_size=16,
178 + num_kv_heads=1,
179 + head_size=64,
180 + dtype=torch.uint8,
181 + )
182 + if num_groups == 1:
183 + groups = [
184 + KVCacheGroupSpec(
185 + ["model.layers.0.mla_attn", "model.layers.1.indexer"],
186 + spec,
187 + )
188 + ]
189 + else:
190 + groups = [
191 + KVCacheGroupSpec(["model.layers.0.mla_attn"], spec),
192 + KVCacheGroupSpec(["model.layers.1.indexer"], spec),
193 + ]
194 + return KVCacheConfig(
195 + num_blocks=num_blocks,
196 + kv_cache_tensors=[],
197 + kv_cache_groups=groups,
198 + )
199 +
200 +
201 class FakeMooncakeWrapper:
202 """Mock Mooncake TransferEngine for unit testing environments."""
203
@@ -235,6 +269,37 @@ def test_align_transfer_regions_uses_layer_name_occurrences():
269 assert [r.base_addr for r in aligned_local] == [0x1000, 0x1100]
270 assert [r.base_addr for r in aligned_remote] == [0xB000, 0xB100]
271
272 + # Hetero PP with no shared layers is an empty success. Same PP fails closed.
273 + p_regions = [
274 + _region(
275 + 0x1000,
276 + block_len=2048,
277 + kv_block_len=1024,
278 + row_offset=0,
279 + layer_name="model.layers.2.mla_attn",
280 + layer_index=2,
281 + )
282 + ]
283 + early_d = [
284 + _region(
285 + 0xA000,
286 + block_len=2048,
287 + kv_block_len=1024,
288 + row_offset=0,
289 + layer_name="model.layers.0.mla_attn",
290 + layer_index=0,
291 + )
292 + ]
293 + none_p, none_d, none_err = _align_transfer_regions(
294 + p_regions, early_d, allow_partial_layers=True
295 + )
296 + assert none_err is None
297 + assert none_p == [] and none_d == []
298 + _, _, mismatch_err = _align_transfer_regions(
299 + p_regions, early_d, allow_partial_layers=False
300 + )
301 + assert mismatch_err is not None
302 +
303
304 @pytest.mark.asyncio
305 async def test_build_transfer_params_separates_prefill_pp_layers():
@@ -323,18 +388,7 @@ async def test_build_transfer_params_separates_prefill_pp_layers():
388 local_block_ids=[[10, 11]],
389 ready=asyncio.Event(),
390 )
326 - xfer_meta = MooncakeXferMetadata(
327 - remote_hostname="consumer-host",
328 - remote_port=54321,
329 - remote_tp_size=1,
330 - remote_tp_rank=0,
331 - req_blocks={"d-req-pp": (transfer_id, [[20, 21]])},
332 - kv_caches_base_addr=[region.base_addr for region in remote_regions],
333 - block_lens=[region.block_len for region in remote_regions],
334 - kv_block_lens=[region.kv_block_len for region in remote_regions],
335 - registered_layer_names=[region.layer_name for region in remote_regions],
336 - registered_layer_indices=[region.layer_index for region in remote_regions],
337 - )
391 + xfer_meta = _xfer_meta(remote_regions, {"d-req-pp": (transfer_id, [[20, 21]])})
392
393 for pp_rank, local_regions in producer_pp_regions.items():
394 aligned_local, aligned_remote, err = _align_transfer_regions(
@@ -411,20 +465,24 @@ async def test_send_kv_to_decode_aligns_consumer_regions_by_layer_metadata(
465 prefill_worker.reqs_need_send[transfer_id] = send_meta
466 send_meta.ready.set()
467
414 - xfer_meta = MooncakeXferMetadata(
415 - remote_hostname="consumer-host",
416 - remote_port=54321,
417 - remote_tp_size=1,
418 - remote_tp_rank=0,
419 - req_blocks={"d-req-layer-align": (transfer_id, [[20]])},
420 - kv_caches_base_addr=[0xA000, 0xB000],
421 - block_lens=[block_len, block_len],
422 - kv_block_lens=[block_len, block_len],
423 - registered_layer_names=[
424 - "model.layers.0.self_attn",
425 - "model.layers.1.self_attn",
468 + xfer_meta = _xfer_meta(
469 + [
470 + _region(
471 + 0xA000,
472 + block_len=block_len,
473 + row_offset=-1,
474 + layer_name="model.layers.0.self_attn",
475 + layer_index=0,
476 + ),
477 + _region(
478 + 0xB000,
479 + block_len=block_len,
480 + row_offset=-1,
481 + layer_name="model.layers.1.self_attn",
482 + layer_index=1,
483 + ),
484 ],
427 - registered_layer_indices=[0, 1],
485 + {"d-req-layer-align": (transfer_id, [[20]])},
486 )
487 mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
488 mock_socket.send_multipart = AsyncMock()
@@ -1093,6 +1151,81 @@ def patch_worker_dependencies():
1151 }
1152
1153
1154 +def _region(
1155 + base: int,
1156 + *,
1157 + block_len: int = 256,
1158 + kv_block_len: int | None = None,
1159 + row_offset: int | None = None,
1160 + layer_name: str = "layer",
1161 + layer_index: int = 0,
1162 + group_index: int = 0,
1163 + shared_group_ids: tuple[int, ...] = (),
1164 +) -> TransferRegion:
1165 + """One transfer slice. `row_offset` defaults to `base` for in-row addresses."""
1166 + return TransferRegion(
1167 + layer_name=layer_name,
1168 + layer_index=layer_index,
1169 + base_addr=base,
1170 + block_len=block_len,
1171 + kv_block_len=block_len if kv_block_len is None else kv_block_len,
1172 + group_index=group_index,
1173 + shared_group_ids=shared_group_ids,
1174 + row_offset=base if row_offset is None else row_offset,
1175 + )
1176 +
1177 +
1178 +def _xfer_meta(
1179 + regions: list[TransferRegion],
1180 + req_blocks: dict,
1181 + **overrides,
1182 +) -> MooncakeXferMetadata:
1183 + """Handshake metadata derived from the regions under test."""
1184 + fields = dict(
1185 + remote_hostname="consumer-host",
1186 + remote_port=54321,
1187 + remote_tp_size=1,
1188 + remote_tp_rank=0,
1189 + req_blocks=req_blocks,
1190 + kv_caches_base_addr=[region.base_addr for region in regions],
1191 + block_lens=[region.block_len for region in regions],
1192 + kv_block_lens=[region.kv_block_len for region in regions],
1193 + registered_layer_names=[region.layer_name for region in regions],
1194 + registered_layer_indices=[region.layer_index for region in regions],
1195 + registered_group_indices=[region.group_index for region in regions],
1196 + registered_shared_group_ids=[
1197 + list(region.shared_group_ids) for region in regions
1198 + ],
1199 + registered_row_offsets=[region.row_offset for region in regions],
1200 + )
1201 + fields.update(overrides)
1202 + return MooncakeXferMetadata(**fields)
1203 +
1204 +
1205 +@contextlib.contextmanager
1206 +def mooncake_register_worker(kv_cache_config, *, kv_role: str = "kv_consumer"):
1207 + """Build a worker whose engine registration and sender thread are patched."""
1208 + vllm_config = create_vllm_config(kv_connector="MooncakeConnector", kv_role=kv_role)
1209 + connector_mod = (
1210 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector"
1211 + )
1212 + with (
1213 + set_current_vllm_config(vllm_config),
1214 + patch_worker_dependencies(),
1215 + patch(f"{connector_mod}.threading.Event"),
1216 + patch(f"{connector_mod}.threading.Thread") as mock_thread,
1217 + ):
1218 + connector = MooncakeConnector(
1219 + vllm_config, KVConnectorRole.WORKER, kv_cache_config
1220 + )
1221 + worker = connector.connector_worker
1222 + mock_thread.return_value.is_alive.return_value = False
1223 + with patch.object(
1224 + worker.engine, "batch_register_memory", return_value=0
1225 + ) as reg:
1226 + yield worker, reg
1227 +
1228 +
1229 @pytest.mark.asyncio
1230 @pytest.mark.parametrize(
1231 ("local_pp_size", "local_pp_rank", "expected_addrs"),
@@ -1211,18 +1344,19 @@ async def test_heterogeneous_pp_waits_for_consumer_without_shared_layers():
1344 patch.object(worker, "_send_blocks", return_value=0) as send,
1345 ):
1346 for pp_rank in range(2):
1214 - metadata = MooncakeXferMetadata(
1347 + metadata = _xfer_meta(
1348 + [
1349 + _region(
1350 + 0x2000,
1351 + row_offset=-1,
1352 + layer_name=f"model.layers.{pp_rank}.self_attn",
1353 + layer_index=pp_rank,
1354 + )
1355 + ],
1356 + {"d-req": ("transfer", [[2]])},
1357 remote_hostname="consumer",
1358 remote_port=1234 + pp_rank,
1217 - remote_tp_size=1,
1218 - remote_tp_rank=0,
1359 remote_pp_size=2,
1220 - req_blocks={"d-req": ("transfer", [[2]])},
1221 - kv_caches_base_addr=[0x2000],
1222 - block_lens=[256],
1223 - kv_block_lens=[256],
1224 - registered_layer_names=[f"model.layers.{pp_rank}.self_attn"],
1225 - registered_layer_indices=[pp_rank],
1360 )
1361 await worker.send_kv_to_decode(b"consumer", sock, metadata)
1362 response = worker._xfer_resp_decoder.decode(
@@ -1286,17 +1420,17 @@ async def test_kv_producer(monkeypatch):
1420 send_meta.ready.set()
1421
1422 # Remote consumer request metadata
1289 - xfer_meta = MooncakeXferMetadata(
1290 - remote_hostname="consumer-host",
1291 - remote_port=54321,
1292 - remote_tp_size=1,
1293 - remote_tp_rank=0,
1294 - req_blocks={"d-req-1": (transfer_id, [[20, 21]])},
1295 - kv_caches_base_addr=[0x2000],
1296 - block_lens=[block_len],
1297 - kv_block_lens=[block_len],
1298 - registered_layer_names=["model.layers.0.self_attn"],
1299 - registered_layer_indices=[0],
1423 + xfer_meta = _xfer_meta(
1424 + [
1425 + _region(
1426 + 0x2000,
1427 + block_len=block_len,
1428 + row_offset=-1,
1429 + layer_name="model.layers.0.self_attn",
1430 + layer_index=0,
1431 + )
1432 + ],
1433 + {"d-req-1": (transfer_id, [[20, 21]])},
1434 )
1435
1436 mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
@@ -1544,139 +1678,703 @@ async def test_worker_get_finished_timeout(monkeypatch):
1678 )
1679 def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
1680 """Tests the memory registration logic with the underlying Mooncake engine."""
1547 - vllm_config = create_vllm_config(
1548 - kv_connector="MooncakeConnector", kv_role="kv_consumer"
1681 + spec = FullAttentionSpec(
1682 + block_size=16,
1683 + num_kv_heads=4,
1684 + head_size=64,
1685 + dtype=torch.float16,
1686 + num_head_slots=2 if separate_kv_head_groups else None,
1687 + state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None,
1688 )
1550 -
1551 - with (
1552 - set_current_vllm_config(vllm_config),
1553 - patch_worker_dependencies(),
1554 - patch(
1555 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
1556 - ),
1557 - patch(
1558 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
1559 - ) as mock_thread,
1560 - ):
1561 - connector = MooncakeConnector(
1562 - vllm_config,
1563 - KVConnectorRole.WORKER,
1564 - _make_test_kv_cache_config(),
1565 - )
1566 - worker = connector.connector_worker
1567 - mock_thread.return_value.is_alive.return_value = False
1568 -
1569 - spec = FullAttentionSpec(
1570 - block_size=16,
1571 - num_kv_heads=4,
1572 - head_size=64,
1573 - dtype=torch.float16,
1574 - num_head_slots=2 if separate_kv_head_groups else None,
1575 - state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None,
1576 - )
1577 - layer_names = [
1578 - "model.layers.0.self_attn",
1579 - "model.layers.1.self_attn",
1580 - ]
1689 + layer_names = [
1690 + "model.layers.0.self_attn",
1691 + "model.layers.1.self_attn",
1692 + ]
1693 + with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg):
1694 for layer_name in layer_names:
1695 worker._layer_specs[layer_name] = spec
1696 raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
1697 tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, layout)
1698 kv_caches = dict(zip(layer_names, (tensor1, tensor2)))
1699 + worker.register_kv_caches(kv_caches)
1700 +
1701 + reg.assert_called_once()
1702 + registered_ptrs, registered_lens = reg.call_args[0]
1703 + assert registered_ptrs == [raw.data_ptr()]
1704 + assert registered_lens == [raw.nbytes]
1705 +
1706 + n_blocks = tensor1.shape[0] if tensor1.ndim > 1 else 1
1707 + storage_end = raw.data_ptr() + raw.nbytes
1708 + for base, stride, kv_len in zip(
1709 + worker.kv_caches_base_addr,
1710 + worker.block_len_per_layer,
1711 + worker.kv_block_len_per_layer,
1712 + ):
1713 + assert base + (n_blocks - 1) * stride + kv_len <= storage_end
1714 +
1715 + # Expected shape is the layout, not a recomputation of the pack predicate.
1716 + kind = {
1717 + KVCacheLayout.LBHNC: "page",
1718 + KVCacheLayout.BLHNC: "page",
1719 + KVCacheLayout.BHLNC: "opaque",
1720 + KVCacheLayout.LHBNC: "per_head",
1721 + }[layout]
1722 + if kind == "opaque":
1723 + packed_row = raw.nbytes // n_blocks
1724 + assert worker.kv_caches_base_addr == [raw.data_ptr()]
1725 + assert worker.block_len_per_layer == [packed_row]
1726 + assert worker.kv_block_len_per_layer == [packed_row]
1727 + assert worker.registered_layer_names == [layer_names[0]]
1728 + elif kind == "per_head":
1729 + expected_addrs = [
1730 + cache[:, head_idx].data_ptr()
1731 + for cache in (tensor1, tensor2)
1732 + for head_idx in range(cache.shape[1])
1733 + ]
1734 + head_block_bytes = tensor1.stride(0) * tensor1.element_size()
1735 + assert worker.kv_caches_base_addr == expected_addrs
1736 + assert worker.block_len_per_layer == [head_block_bytes] * len(
1737 + expected_addrs
1738 + )
1739 + assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
1740 + expected_addrs
1741 + )
1742 + assert worker.registered_layer_names == [
1743 + layer_name
1744 + for layer_name in layer_names
1745 + for _ in range(tensor1.shape[1])
1746 + ]
1747 + else:
1748 + page_stride = tensor1.stride(0) * tensor1.element_size()
1749 + assert len(worker.block_len_per_layer) == len(kv_caches)
1750 + assert worker.block_len_per_layer == [page_stride, page_stride]
1751 + assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
1752 + assert worker.registered_layer_names == list(kv_caches)
1753 + assert worker.registered_layer_indices == [0, 1]
1754
1587 - with patch.object(
1588 - worker.engine, "batch_register_memory", return_value=0
1589 - ) as mock_batch_register:
1590 - connector.register_kv_caches(kv_caches)
1591 -
1592 - mock_batch_register.assert_called_once()
1593 - registered_ptrs, registered_lens = mock_batch_register.call_args[0]
1594 - assert registered_ptrs == [raw.data_ptr()]
1595 - assert registered_lens == [raw.nbytes]
1596 -
1597 - if not layout.is_block_compact:
1598 - expected_addrs = [
1599 - cache[:, head_idx].data_ptr()
1600 - for cache in (tensor1, tensor2)
1601 - for head_idx in range(cache.shape[1])
1602 - ]
1603 - head_block_bytes = tensor1.stride(0) * tensor1.element_size()
1604 - assert worker.kv_caches_base_addr == expected_addrs
1605 - assert worker.block_len_per_layer == [head_block_bytes] * len(
1606 - expected_addrs
1607 - )
1608 - assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
1609 - expected_addrs
1610 - )
1611 - assert worker.registered_layer_names == [
1612 - layer_name
1613 - for layer_name in layer_names
1614 - for _ in range(tensor1.shape[1])
1615 - ]
1616 - else:
1617 - assert len(worker.block_len_per_layer) == len(kv_caches)
1618 - for bl in worker.block_len_per_layer:
1619 - assert bl == tensor1.stride(0) * tensor1.element_size()
1620 - assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
1621 - assert worker.registered_layer_names == list(kv_caches)
1622 - assert worker.registered_layer_indices == [0, 1]
1755
1756 +def test_register_noncontiguous_packed_allows_matching_pp():
1757 + """Same PP on both sides can still register a non-contiguous packed row.
1758
1625 -def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1626 - """Mixed MLA+Eagle caches should register by byte length, not shape."""
1627 - vllm_config = create_vllm_config(
1628 - kv_connector="MooncakeConnector", kv_role="kv_consumer"
1759 + A remote PP mismatch is rejected at send time, once the peer PP size is
1760 + known. Rejecting every local PP>1 layout at startup is too strict.
1761 + """
1762 + spec = FullAttentionSpec(
1763 + block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
1764 )
1765 + layer_names = ["model.layers.0.self_attn", "model.layers.1.self_attn"]
1766 + with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, _reg):
1767 + worker.pp_size = 2
1768 + for layer_name in layer_names:
1769 + worker._layer_specs[layer_name] = spec
1770 + raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
1771 + tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, KVCacheLayout.BHLNC)
1772 + assert not tensor1[0].is_contiguous()
1773 + worker.register_kv_caches(dict(zip(layer_names, (tensor1, tensor2))))
1774 + assert worker.kv_caches_base_addr == [raw.data_ptr()]
1775 + assert worker.region_row_offsets == [0]
1776
1631 - with (
1632 - set_current_vllm_config(vllm_config),
1633 - patch_worker_dependencies(),
1634 - patch(
1635 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
1636 - ),
1637 - patch(
1638 - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
1639 - ) as mock_thread,
1640 - ):
1641 - connector = MooncakeConnector(
1642 - vllm_config,
1643 - KVConnectorRole.WORKER,
1644 - _make_test_kv_cache_config(),
1645 - )
1646 - worker = connector.connector_worker
1647 - mock_thread.return_value.is_alive.return_value = False
1777
1778 +def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
1779 + """Mixed MLA+Eagle caches should register by byte length, not shape."""
1780 + with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, reg):
1781 worker.use_mla = True
1782 worker.transfer_topo.is_mla = True
1651 -
1652 - # MLA cache tensor: shape[-2] is the block size and each block's
1653 - # byte stride matches the cache spec page size.
1783 mla_cache = torch.zeros((2, 16, 512), dtype=torch.float16)
1655 - # Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size.
1784 eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
1785 kv_caches = {
1786 "model.layers.0.mla_attn": mla_cache,
1787 "model.layers.1.eagle_attn": eagle_cache,
1788 }
1789 + worker.register_kv_caches(kv_caches)
1790 +
1791 + reg.assert_called_once()
1792 + registered_ptrs, registered_lens = reg.call_args[0]
1793 + assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()]
1794 + assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes]
1795 + assert worker.block_len_per_layer == [
1796 + mla_cache.nbytes // mla_cache.shape[0],
1797 + eagle_cache.nbytes // eagle_cache.shape[0],
1798 + ]
1799 + assert worker.registered_layer_names == [
1800 + "model.layers.0.mla_attn",
1801 + "model.layers.1.eagle_attn",
1802 + ]
1803 + assert worker.registered_layer_indices == [0, 1]
1804 + assert len(worker.kv_caches_base_addr) == 2
1805 +
1806 +
1807 +def test_packed_and_unpacked_region_lengths_fail_homogeneous_tp_handshake():
1808 + packed = TransferRegion(
1809 + layer_name="model.layers.0.mla_attn",
1810 + layer_index=0,
1811 + base_addr=0x1000,
1812 + block_len=2048,
1813 + kv_block_len=2048,
1814 + group_index=0,
1815 + )
1816 + unpacked = TransferRegion(
1817 + layer_name="model.layers.0.mla_attn",
1818 + layer_index=0,
1819 + base_addr=0x1000,
1820 + block_len=64,
1821 + kv_block_len=64,
1822 + group_index=0,
1823 + )
1824 + err = _validate_asymmetric_region_lengths(
1825 + local_regions=[packed],
1826 + remote_regions=[unpacked],
1827 + local_tp_size=1,
1828 + remote_tp_size=1,
1829 + producer_cache_replicated=False,
1830 + )
1831 + assert err is not None
1832 + assert "length mismatch" in err
1833
1662 - with patch.object(
1663 - worker.engine, "batch_register_memory", return_value=0
1664 - ) as mock_batch_register:
1665 - connector.register_kv_caches(kv_caches)
1666 -
1667 - mock_batch_register.assert_called_once()
1668 - registered_ptrs, registered_lens = mock_batch_register.call_args[0]
1669 - assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()]
1670 - assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes]
1671 - assert worker.block_len_per_layer == [
1672 - mla_cache.nbytes // mla_cache.shape[0],
1673 - eagle_cache.nbytes // eagle_cache.shape[0],
1674 - ]
1675 - assert worker.registered_layer_names == [
1676 - "model.layers.0.mla_attn",
1677 - "model.layers.1.eagle_attn",
1834 + count_err = _validate_asymmetric_region_lengths(
1835 + local_regions=[packed],
1836 + remote_regions=[unpacked, unpacked],
1837 + local_tp_size=1,
1838 + remote_tp_size=1,
1839 + producer_cache_replicated=False,
1840 + )
1841 + assert count_err is not None
1842 + assert "region counts" in count_err
1843 +
1844 +
1845 +@pytest.mark.parametrize(
1846 + ("num_groups", "expected_group_index"),
1847 + [(1, 0), (2, _SHARED_REGION_GROUP_ID)],
1848 +)
1849 +def test_register_kv_caches_collapses_shared_mla_storage(
1850 + num_groups, expected_group_index
1851 +):
1852 + """Shared backing collapses to one packed region; two groups mark it -1."""
1853 + num_blocks = 4
1854 + packed_row = 256
1855 + with mooncake_register_worker(
1856 + _make_packed_mla_kv_cache_config(num_blocks, num_groups=num_groups)
1857 + ) as (worker, reg):
1858 + backing = torch.zeros((num_blocks, packed_row), dtype=torch.uint8)
1859 + worker.register_kv_caches(
1860 + {
1861 + "model.layers.0.mla_attn": backing,
1862 + "model.layers.1.indexer": backing,
1863 + }
1864 + )
1865 +
1866 + reg.assert_called_once()
1867 + registered_ptrs, registered_lens = reg.call_args[0]
1868 + assert registered_ptrs == [backing.data_ptr()]
1869 + assert registered_lens == [backing.untyped_storage().nbytes()]
1870 + assert worker.kv_caches_base_addr == [backing.data_ptr()]
1871 + assert worker.block_len_per_layer == [packed_row]
1872 + assert worker.kv_block_len_per_layer == [packed_row]
1873 + assert worker.registered_layer_names == ["model.layers.0.mla_attn"]
1874 + assert worker.registered_group_indices == [expected_group_index]
1875 + if expected_group_index == _SHARED_REGION_GROUP_ID:
1876 + assert worker.region_shared_groups == [(0, 1)]
1877 +
1878 +
1879 +def test_block_ids_for_region_flattens_shared_groups():
1880 + ids = [[10, 11], [12]]
1881 + assert _block_ids_for_region(ids, 0) == [10, 11]
1882 + assert _block_ids_for_region(ids, 1) == [12]
1883 + assert _block_ids_for_region(ids, _SHARED_REGION_GROUP_ID, (0, 1)) == [
1884 + 10,
1885 + 11,
1886 + 12,
1887 + ]
1888 + # A group that does not share the allocation stays out.
1889 + assert _block_ids_for_region(
1890 + [[10, 11], [12], [99]], _SHARED_REGION_GROUP_ID, (0, 1)
1891 + ) == [10, 11, 12]
1892 +
1893 +
1894 +@pytest.mark.asyncio
1895 +@pytest.mark.parametrize(
1896 + ("group_index", "local_ids", "remote_ids", "n_blocks"),
1897 + [
1898 + (0, [[10, 11]], [[20, 21]], 2),
1899 + (_SHARED_REGION_GROUP_ID, [[10, 11], [12]], [[20, 21], [22]], 3),
1900 + ],
1901 +)
1902 +async def test_build_transfer_params_sends_packed_region_once(
1903 + group_index, local_ids, remote_ids, n_blocks
1904 +):
1905 + """Packed region coalesces contiguous blocks, including shared-group flatten."""
1906 + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker)
1907 + worker.async_zmq_ctx = MagicMock()
1908 + worker.is_kv_consumer = True
1909 + worker.is_kv_producer = True
1910 + worker.tp_rank = 0
1911 + worker.tp_size = 1
1912 + worker.use_mla = True
1913 + worker.kv_cache_config = _make_packed_mla_kv_cache_config(
1914 + num_blocks=4, num_groups=len(local_ids)
1915 + )
1916 + worker._physical_blocks_per_logical_kv_block = 1
1917 + worker.transfer_topo = SimpleNamespace(
1918 + local_replicates_kv_cache=False,
1919 + total_num_kv_heads=1,
1920 + )
1921 +
1922 + block_len = 256
1923 + shared = (0, 1) if group_index == _SHARED_REGION_GROUP_ID else ()
1924 + local_region = _region(
1925 + 0x1000,
1926 + block_len=block_len,
1927 + row_offset=0,
1928 + layer_name="model.layers.0.mla_attn",
1929 + group_index=group_index,
1930 + shared_group_ids=shared,
1931 + )
1932 + remote_region = _region(
1933 + 0xA000,
1934 + block_len=block_len,
1935 + row_offset=0,
1936 + layer_name="model.layers.0.mla_attn",
1937 + group_index=group_index,
1938 + shared_group_ids=shared,
1939 + )
1940 + transfer_id = "xfer-packed"
1941 + send_meta = SendBlockMeta(
1942 + p_req_id="p-packed",
1943 + transfer_id=transfer_id,
1944 + local_block_ids=local_ids,
1945 + ready=asyncio.Event(),
1946 + )
1947 + xfer_meta = _xfer_meta([remote_region], {"d-packed": (transfer_id, remote_ids)})
1948 +
1949 + (
1950 + src_ptrs,
1951 + dst_ptrs,
1952 + lengths,
1953 + err_reqs,
1954 + err_msg,
1955 + ) = await worker._build_transfer_params(
1956 + ready_reqs=[("d-packed", send_meta)],
1957 + agent_meta=xfer_meta,
1958 + local_regions=[local_region],
1959 + remote_regions=[remote_region],
1960 + )
1961 +
1962 + assert err_reqs == []
1963 + assert err_msg is None
1964 + assert src_ptrs == [0x1000 + 10 * block_len]
1965 + assert dst_ptrs == [0xA000 + 20 * block_len]
1966 + assert lengths == [n_blocks * block_len]
1967 +
1968 +
1969 +def test_coalesce_promotes_padding_only_for_a_full_row():
1970 + """A full-row merge includes the padding tail; a partial row does not."""
1971 + page, row = 1024, 2560
1972 + local = [
1973 + _region(0, block_len=row, kv_block_len=page, row_offset=0),
1974 + _region(page, block_len=row, kv_block_len=page, row_offset=page),
1975 + ]
1976 + remote = [
1977 + _region(100, block_len=row, kv_block_len=page, row_offset=0),
1978 + _region(100 + page, block_len=row, kv_block_len=page, row_offset=page),
1979 + ]
1980 + full_local, full_remote = _coalesce_contiguous_transfer_regions(
1981 + local, remote, promote_full_row=True
1982 + )
1983 + assert len(full_local) == 1
1984 + assert full_local[0].kv_block_len == full_remote[0].kv_block_len == row
1985 +
1986 + partial_local, _ = _coalesce_contiguous_transfer_regions(local, remote)
1987 + assert partial_local[0].kv_block_len == 2 * page
1988 +
1989 +
1990 +def test_coalesce_spans_padding_between_pages():
1991 + """Unpadded pages merge across the padding gap without a full-row lift."""
1992 + row = 2560
1993 + merged, _ = _coalesce_contiguous_transfer_regions(
1994 + [
1995 + _region(0, block_len=row, kv_block_len=1000),
1996 + _region(1024, block_len=row, kv_block_len=1000),
1997 + ],
1998 + [
1999 + _region(5000, block_len=row, kv_block_len=1000, row_offset=0),
2000 + _region(6024, block_len=row, kv_block_len=1000, row_offset=1024),
2001 + ],
2002 + )
2003 + assert len(merged) == 1
2004 + assert merged[0].kv_block_len == 2024
2005 + assert merged[0].row_offset == 0
2006 +
2007 +
2008 +def test_coalesce_promotes_each_group_that_starts_at_row_zero():
2009 + """Groups overlay one row, so each row-start run is its own full-row copy."""
2010 + page, gap, row = 1000, 1024, 4096
2011 +
2012 + def side(origin: int) -> list[TransferRegion]:
2013 + return [
2014 + _region(
2015 + origin + base,
2016 + block_len=row,
2017 + kv_block_len=page,
2018 + row_offset=offset,
2019 + layer_name=f"group-{group}",
2020 + layer_index=group,
2021 + group_index=group,
2022 + )
2023 + for group, base, offset in (
2024 + (0, 0, 0),
2025 + (0, gap, gap),
2026 + (1, 0, 0),
2027 + (1, gap, gap),
2028 + )
2029 ]
1679 - assert worker.registered_layer_indices == [0, 1]
2030 +
2031 + merged_local, merged_remote = _coalesce_contiguous_transfer_regions(
2032 + side(0), side(8000), promote_full_row=True
2033 + )
2034 + assert [region.group_index for region in merged_local] == [0, 1]
2035 + assert [region.kv_block_len for region in merged_local] == [row, row]
2036 + assert [region.kv_block_len for region in merged_remote] == [row, row]
2037 + assert merged_local[1].base_addr == 0
2038 + assert merged_remote[1].base_addr == 8000
2039 +
2040 +
2041 +def test_coalesce_does_not_promote_a_mid_row_run():
2042 + """A matched PP slice that does not start at row offset 0 stays partial."""
2043 + page, row = 1000, 4096
2044 + merged, _ = _coalesce_contiguous_transfer_regions(
2045 + [
2046 + _region(
2047 + 1024, block_len=row, kv_block_len=page, row_offset=1024, layer_index=2
2048 + ),
2049 + _region(
2050 + 2048, block_len=row, kv_block_len=page, row_offset=2048, layer_index=2
2051 + ),
2052 + ],
2053 + [
2054 + _region(
2055 + 9000 + 1024,
2056 + block_len=row,
2057 + kv_block_len=page,
2058 + row_offset=1024,
2059 + layer_index=2,
2060 + ),
2061 + _region(
2062 + 9000 + 2048,
2063 + block_len=row,
2064 + kv_block_len=page,
2065 + row_offset=2048,
2066 + layer_index=2,
2067 + ),
2068 + ],
2069 + promote_full_row=True,
2070 + )
2071 + assert len(merged) == 1
2072 + assert merged[0].row_offset == 1024
2073 + assert merged[0].kv_block_len == 1024 + page
2074 +
2075 +
2076 +def test_opaque_packed_row_is_full_stride_at_offset_zero():
2077 + """Branch B is row-sized at offset 0. A page view at offset 0 is not."""
2078 + assert _has_opaque_packed_row([0], [4096], [4096])
2079 + assert not _has_opaque_packed_row([0, 1024], [1000, 1000], [4096, 4096])
2080 + assert not _has_opaque_packed_row([], [4096], [4096])
2081 +
2082 +
2083 +def test_pp_mismatch_rejects_opaque_row_in_both_directions():
2084 + """The smaller PP side can be either the producer or the consumer."""
2085 + page_view = dict(
2086 + remote_row_offsets=[0, 1024],
2087 + remote_kv_block_lens=[1000, 1000],
2088 + remote_block_lens=[4096, 4096],
2089 + )
2090 + opaque_row = dict(
2091 + remote_row_offsets=[0],
2092 + remote_kv_block_lens=[4096],
2093 + remote_block_lens=[4096],
2094 + )
2095 + # P pp=2, D pp=1, local row is opaque.
2096 + assert _pp_mismatch_hides_packed_layers(
2097 + 2, 1, local_has_opaque_row=True, **page_view
2098 + )
2099 + # P pp=1, D pp=2, remote row is opaque.
2100 + assert _pp_mismatch_hides_packed_layers(
2101 + 1, 2, local_has_opaque_row=False, **opaque_row
2102 + )
2103 + # Matching PP still aligns by the shared first layer name.
2104 + assert not _pp_mismatch_hides_packed_layers(
2105 + 2, 2, local_has_opaque_row=True, **opaque_row
2106 + )
2107 + # Hetero PP with page-contiguous views is not this rejection.
2108 + assert not _pp_mismatch_hides_packed_layers(
2109 + 2, 1, local_has_opaque_row=False, **page_view
2110 + )
2111 +
2112 +
2113 +def test_prepare_transfer_regions_reuses_success_and_error():
2114 + """A peer layout is built once. Re-registering drops the cache."""
2115 + worker, _, _, _ = _register_sliced_packed_mla([0, 1])
2116 + regions = worker._get_transfer_regions(
2117 + worker.kv_caches_base_addr,
2118 + worker.block_len_per_layer,
2119 + worker.kv_block_len_per_layer,
2120 + worker.registered_layer_names,
2121 + worker.registered_layer_indices,
2122 + worker.registered_group_indices,
2123 + [tuple(groups) for groups in worker.region_shared_groups],
2124 + worker.region_row_offsets,
2125 + )
2126 + meta = _xfer_meta(
2127 + regions,
2128 + {},
2129 + remote_hostname="peer",
2130 + remote_port=1,
2131 + remote_tp_size=worker.tp_size,
2132 + remote_pp_size=worker.pp_size,
2133 + )
2134 + calls = 0
2135 + original = worker._get_transfer_regions
2136 +
2137 + def counting(*args, **kwargs):
2138 + nonlocal calls
2139 + calls += 1
2140 + return original(*args, **kwargs)
2141 +
2142 + worker._get_transfer_regions = counting
2143 + prepared = worker._prepare_transfer_regions(meta)
2144 + assert prepared[2] is None
2145 + assert calls == 2
2146 + assert worker._prepare_transfer_regions(meta) == prepared
2147 + assert calls == 2
2148 +
2149 + mismatch = _xfer_meta(
2150 + [
2151 + _region(
2152 + 0x2000,
2153 + block_len=1024,
2154 + row_offset=0,
2155 + layer_name="model.layers.9.mla_attn",
2156 + layer_index=9,
2157 + )
2158 + ],
2159 + {},
2160 + remote_hostname="peer",
2161 + remote_port=2,
2162 + remote_tp_size=worker.tp_size,
2163 + remote_pp_size=worker.pp_size,
2164 + )
2165 + failed = worker._prepare_transfer_regions(mismatch)
2166 + assert failed[2] is not None
2167 + assert calls == 4
2168 + assert worker._prepare_transfer_regions(mismatch)[2] == failed[2]
2169 + assert calls == 4
2170 +
2171 + page = 1024
2172 + backing = torch.zeros((4, page * 2), dtype=torch.uint8)
2173 + kv_caches = {
2174 + "model.layers.0.mla_attn": backing[:, :page],
2175 + "model.layers.1.mla_attn": backing[:, page:],
2176 + }
2177 + worker.is_kv_consumer = True
2178 + # The first register built a producer. Its thread patch is already gone,
2179 + # so stay on the consumer path and return before the sender thread starts.
2180 + with patch.object(worker.engine, "batch_register_memory", return_value=0):
2181 + worker.register_kv_caches(kv_caches)
2182 + assert worker._prepared_transfer_regions == {}
2183 + worker._prepare_transfer_regions(meta)
2184 + assert calls == 6
2185 +
2186 +
2187 +def test_coalesce_keeps_a_run_inside_its_starting_row():
2188 + """The row bound is the run start, not the previous slice."""
2189 + row = 2500
2190 + slices = [
2191 + _region(base, block_len=row, kv_block_len=1000) for base in (0, 1000, 2000)
2192 + ]
2193 + merged, _ = _coalesce_contiguous_transfer_regions(slices, slices)
2194 + assert len(merged) == 2
2195 + assert merged[0].kv_block_len == 2000
2196 + assert merged[1].base_addr == 2000
2197 + assert merged[1].kv_block_len == 1000
2198 +
2199 +
2200 +def _layer_name(idx: int) -> str:
2201 + return f"model.layers.{idx}.mla_attn"
2202 +
2203 +
2204 +def _register_sliced_packed_mla(
2205 + layer_idxs: list[int], *, page: int = 1024, num_blocks: int = 4
2206 +):
2207 + """Register column slices of one MLA packed row. Returns worker fields."""
2208 + spec = MLAAttentionSpec(
2209 + block_size=16,
2210 + num_kv_heads=1,
2211 + head_size=64,
2212 + dtype=torch.uint8,
2213 + )
2214 + kv_cache_config = KVCacheConfig(
2215 + num_blocks=num_blocks,
2216 + kv_cache_tensors=[],
2217 + kv_cache_groups=[
2218 + KVCacheGroupSpec([_layer_name(i) for i in layer_idxs], spec),
2219 + ],
2220 + )
2221 + backing = torch.zeros((num_blocks, page * len(layer_idxs)), dtype=torch.uint8)
2222 + kv_caches = {
2223 + _layer_name(idx): backing[:, i * page : (i + 1) * page]
2224 + for i, idx in enumerate(layer_idxs)
2225 + }
2226 + with mooncake_register_worker(kv_cache_config, kv_role="kv_producer") as (
2227 + worker,
2228 + _,
2229 + ):
2230 + worker.register_kv_caches(kv_caches)
2231 + regions = worker._get_transfer_regions(
2232 + worker.kv_caches_base_addr,
2233 + worker.block_len_per_layer,
2234 + worker.kv_block_len_per_layer,
2235 + worker.registered_layer_names,
2236 + worker.registered_layer_indices,
2237 + worker.registered_group_indices,
2238 + row_offsets=worker.region_row_offsets,
2239 + )
2240 + return worker, kv_cache_config, regions, backing.data_ptr()
2241 +
2242 +
2243 +def test_sliced_packed_mla_registers_each_layer_view():
2244 + """Each column of a packed MLA row stays its own view, in row order."""
2245 + page = 1024
2246 + layer_idxs = [0, 1]
2247 + with mooncake_register_worker(
2248 + KVCacheConfig(
2249 + num_blocks=4,
2250 + kv_cache_tensors=[],
2251 + kv_cache_groups=[
2252 + KVCacheGroupSpec(
2253 + [_layer_name(i) for i in layer_idxs],
2254 + MLAAttentionSpec(
2255 + block_size=16,
2256 + num_kv_heads=1,
2257 + head_size=64,
2258 + dtype=torch.uint8,
2259 + ),
2260 + )
2261 + ],
2262 + ),
2263 + kv_role="kv_producer",
2264 + ) as (worker, reg):
2265 + backing = torch.zeros((4, page * len(layer_idxs)), dtype=torch.uint8)
2266 + worker.register_kv_caches(
2267 + {
2268 + _layer_name(idx): backing[:, i * page : (i + 1) * page]
2269 + for i, idx in enumerate(layer_idxs)
2270 + }
2271 + )
2272 + reg.assert_called_once()
2273 + assert reg.call_args[0][0] == [backing.data_ptr()]
2274 + assert worker.region_row_offsets == [0, page]
2275 +
2276 +
2277 +async def _copy_packed_slices(
2278 + p_layers: list[int],
2279 + d_layers: list[int],
2280 + *,
2281 + allow_partial: bool,
2282 + local_ids: list[list[int]],
2283 + remote_ids: list[list[int]],
2284 + page: int = 1024,
2285 +):
2286 + p_worker, p_config, p_regions, p_base = _register_sliced_packed_mla(
2287 + p_layers, page=page
2288 + )
2289 + _, _, d_regions, d_base = _register_sliced_packed_mla(d_layers, page=page)
2290 + aligned_p, aligned_d, err = _align_transfer_regions(
2291 + p_regions, d_regions, allow_partial_layers=allow_partial
2292 + )
2293 + assert err is None
2294 + coalesced_p, coalesced_d = _coalesce_contiguous_transfer_regions(
2295 + aligned_p, aligned_d
2296 + )
2297 + transfer_id = "xfer-packed-slices"
2298 + send_meta = SendBlockMeta(
2299 + p_req_id="p-slices",
2300 + transfer_id=transfer_id,
2301 + local_block_ids=local_ids,
2302 + ready=asyncio.Event(),
2303 + )
2304 + p_worker.kv_cache_config = p_config
2305 + src, dst, lengths, err_reqs, err_msg = await p_worker._build_transfer_params(
2306 + ready_reqs=[("d-slices", send_meta)],
2307 + agent_meta=_xfer_meta(coalesced_d, {"d-slices": (transfer_id, remote_ids)}),
2308 + local_regions=coalesced_p,
2309 + remote_regions=coalesced_d,
2310 + )
2311 + return SimpleNamespace(
2312 + p_regions=p_regions,
2313 + d_regions=d_regions,
2314 + coalesced_p=coalesced_p,
2315 + coalesced_d=coalesced_d,
2316 + p_base=p_base,
2317 + d_base=d_base,
2318 + src=src,
2319 + dst=dst,
2320 + lengths=lengths,
2321 + err_reqs=err_reqs,
2322 + err_msg=err_msg,
2323 + )
2324 +
2325 +
2326 +@pytest.mark.asyncio
2327 +async def test_packed_hetero_pp_keeps_shared_layers():
2328 + """P-PP4/D-PP2 copies the shared MLA span, not either side's full row."""
2329 + page = 1024
2330 + copied = await _copy_packed_slices(
2331 + [2, 3],
2332 + [0, 1, 2, 3],
2333 + allow_partial=True,
2334 + local_ids=[[10, 11]],
2335 + remote_ids=[[10, 11]],
2336 + page=page,
2337 + )
2338 + assert [region.layer_index for region in copied.p_regions] == [2, 3]
2339 + assert [region.kv_block_len for region in copied.p_regions] == [page, page]
2340 + assert copied.p_regions[0].block_len == 2 * page
2341 + assert copied.d_regions[0].block_len == 4 * page
2342 + assert len(copied.coalesced_p) == 1
2343 + assert copied.coalesced_p[0].kv_block_len == 2 * page
2344 + assert copied.coalesced_d[0].base_addr == copied.d_base + 2 * page
2345 + assert copied.err_reqs == []
2346 + assert copied.err_msg is None
2347 + assert copied.lengths == [2 * page, 2 * page]
2348 + assert copied.src == [
2349 + copied.p_base + 10 * (2 * page),
2350 + copied.p_base + 11 * (2 * page),
2351 + ]
2352 + assert copied.dst == [
2353 + copied.d_base + 2 * page + 10 * (4 * page),
2354 + copied.d_base + 2 * page + 11 * (4 * page),
2355 + ]
2356 +
2357 +
2358 +@pytest.mark.asyncio
2359 +async def test_packed_matching_layers_coalesce_to_one_copy():
2360 + """Identical packed layer sets still emit one copy for contiguous blocks."""
2361 + page = 1024
2362 + copied = await _copy_packed_slices(
2363 + [0, 1],
2364 + [0, 1],
2365 + allow_partial=False,
2366 + local_ids=[[10, 11]],
2367 + remote_ids=[[20, 21]],
2368 + page=page,
2369 + )
2370 + assert copied.err_reqs == []
2371 + assert copied.err_msg is None
2372 + assert len(copied.coalesced_p) == 1
2373 + assert copied.coalesced_p[0].kv_block_len == copied.coalesced_p[0].block_len
2374 + block_len = 2 * page
2375 + assert copied.src == [copied.p_base + 10 * block_len]
2376 + assert copied.dst == [copied.d_base + 20 * block_len]
2377 + assert copied.lengths == [2 * block_len]
2378
2379
2380 @pytest.mark.asyncio
@@ -1760,22 +2458,24 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
2458 ) as mock_send_blocks:
2459 for d_rank in target_d_ranks:
2460 remote_block_ids = d_rank_remote_blocks[d_rank]
1763 - xfer_meta = MooncakeXferMetadata(
1764 - remote_hostname="consumer-host",
1765 - remote_port=54321,
1766 - remote_tp_size=d_tp_size,
1767 - remote_tp_rank=d_rank,
1768 - req_blocks={
2461 + xfer_meta = _xfer_meta(
2462 + [
2463 + _region(
2464 + 0x2000,
2465 + block_len=remote_block_len,
2466 + row_offset=-1,
2467 + layer_name="model.layers.0.self_attn",
2468 + layer_index=0,
2469 + )
2470 + ],
2471 + {
2472 f"d-req-h1-r{d_rank}": (
2473 transfer_id,
2474 remote_block_ids,
2475 )
2476 },
1774 - kv_caches_base_addr=[0x2000],
1775 - block_lens=[remote_block_len],
1776 - kv_block_lens=[remote_block_len],
1777 - registered_layer_names=["model.layers.0.self_attn"],
1778 - registered_layer_indices=[0],
2477 + remote_tp_size=d_tp_size,
2478 + remote_tp_rank=d_rank,
2479 )
2480
2481 mock_send_blocks.reset_mock()
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+2
@@ -242,6 +242,8 @@ async def test_build_transfer_params_uses_non_overlapping_physical_pages():
242 kv_cache_config = make_hybrid_gdn_kv_cache_config(block_size=16)
243 worker = object.__new__(MooncakeConnectorWorker)
244 worker.shutdown = noop_shutdown
245 + # __new__ skips __init__, where this cache is created. register_kv_caches clears it.
246 + worker._prepared_transfer_regions = {}
247 worker.use_mla = False
248 worker.engine = SimpleNamespace(batch_register_memory=lambda *_: 0)
249 worker.is_kv_consumer = True
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+517 -42
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 +import itertools
5 import logging
6 import queue
7 import threading
@@ -62,6 +63,8 @@ from vllm.v1.kv_cache_interface import (
63 KpoolTailSpec,
64 KVCacheSpec,
65 MambaSpec,
66 + MLAAttentionSpec,
67 + SlidingWindowMLASpec,
68 SlidingWindowSpec,
69 )
70 from vllm.v1.request import RequestStatus
@@ -90,6 +93,12 @@ if TYPE_CHECKING:
93 ReqId = str # Internal scheduler request ID
94 TransferId = str # KV transfer coordination ID (shared by P/D)
95
96 +# Follow NIXL (`nixl/base_worker.py`): HMA packed pages can put more than one
97 +# KV group on a single allocation. group_index=-1 flattens every group's block
98 +# ids onto that region. Hybrid KV manager assigns distinct ids per group, so
99 +# flatten does not copy the same block twice.
100 +_SHARED_REGION_GROUP_ID = -1
101 +
102
103 @dataclass(frozen=True)
104 class TransferRegion:
@@ -99,6 +108,34 @@ class TransferRegion:
108 block_len: int
109 kv_block_len: int
110 group_index: int = 0
111 + # Groups that share this allocation. Empty unless group_index is -1.
112 + shared_group_ids: tuple[int, ...] = ()
113 + # Byte offset of this view from the start of its block row.
114 + row_offset: int = 0
115 +
116 +
117 +def _block_ids_for_region(
118 + block_ids_by_group: list[list[int]],
119 + group_index: int,
120 + shared_group_ids: tuple[int, ...] = (),
121 +) -> list[int]:
122 + """Follow NIXL `_block_ids_by_region`.
123 +
124 + -1 flattens only the groups that share this allocation. NIXL can flatten
125 + every group because its packed flag means they all share one storage;
126 + Mooncake can still have a separate Mamba group in the same request.
127 + """
128 + if group_index == _SHARED_REGION_GROUP_ID:
129 + assert shared_group_ids, "Shared packed region is missing its group ids."
130 + return list(
131 + itertools.chain.from_iterable(
132 + block_ids_by_group[group] for group in shared_group_ids
133 + )
134 + )
135 + assert 0 <= group_index < len(block_ids_by_group), (
136 + "Transfer region references a missing KV group."
137 + )
138 + return list(block_ids_by_group[group_index])
139
140
141 def _get_tp_ratio(local_tp_size: int, remote_tp_size: int) -> int:
@@ -129,6 +166,8 @@ def _expand_transfer_regions(
166 layer_names: list[str],
167 layer_indices: list[int],
168 group_indices: list[int] | None = None,
169 + shared_group_ids: list[tuple[int, ...]] | None = None,
170 + row_offsets: list[int] | None = None,
171 ) -> list[TransferRegion]:
172 """Expand registered KV tensors into the regions transferred by Mooncake."""
173 assert (
@@ -150,6 +189,14 @@ def _expand_transfer_regions(
189 "Mooncake transfer regions require matching group metadata lengths, "
190 f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}."
191 )
192 + if shared_group_ids is None:
193 + shared_group_ids = [() for _ in layer_names]
194 + assert len(shared_group_ids) == len(layer_names)
195 + # Missing offsets must not look like row starts, or every run would be
196 + # promoted to the full row.
197 + if row_offsets is None:
198 + row_offsets = [-1] * len(layer_names)
199 + assert len(row_offsets) == len(layer_names)
200 regions: list[TransferRegion] = []
201 for (
202 base_addr,
@@ -158,6 +205,8 @@ def _expand_transfer_regions(
205 layer_name,
206 layer_index,
207 group_index,
208 + region_groups,
209 + row_offset,
210 ) in zip(
211 base_addrs,
212 block_lens,
@@ -165,6 +214,8 @@ def _expand_transfer_regions(
214 layer_names,
215 layer_indices,
216 group_indices,
217 + shared_group_ids,
218 + row_offsets,
219 ):
220 regions.append(
221 TransferRegion(
@@ -174,6 +225,8 @@ def _expand_transfer_regions(
225 block_len=block_len,
226 kv_block_len=kv_block_len,
227 group_index=group_index,
228 + shared_group_ids=region_groups,
229 + row_offset=row_offset,
230 )
231 )
232 return regions
@@ -377,12 +430,168 @@ def _align_transfer_regions(
430 f"{remote_region.group_index}."
431 ),
432 )
433 + if (
434 + local_region.shared_group_ids
435 + and remote_region.shared_group_ids
436 + and local_region.shared_group_ids != remote_region.shared_group_ids
437 + ):
438 + return (
439 + [],
440 + [],
441 + (
442 + "Mooncake shared-group set mismatch for "
443 + f"{local_region.layer_name}: producer="
444 + f"{local_region.shared_group_ids}, consumer="
445 + f"{remote_region.shared_group_ids}."
446 + ),
447 + )
448 aligned_local.append(local_region)
449 aligned_remote.append(remote_region)
450
451 return aligned_local, aligned_remote, None
452
453
454 +def _coalesce_contiguous_transfer_regions(
455 + local_regions: list[TransferRegion],
456 + remote_regions: list[TransferRegion],
457 + *,
458 + promote_full_row: bool = False,
459 +) -> tuple[list[TransferRegion], list[TransferRegion]]:
460 + """Merge aligned packed slices that are adjacent on both sides.
461 +
462 + Page-contiguous views stay addressable by layer name for hetero PP
463 + (#56033). KV groups overlay one row and must stay separate: each group
464 + owns a different block-id space. A run that starts at row offset 0 is
465 + lifted to the full row so consecutive blocks can coalesce.
466 + """
467 + if len(local_regions) != len(remote_regions):
468 + return local_regions, remote_regions
469 +
470 + def adjacent(
471 + run_start: TransferRegion, prev: TransferRegion, nxt: TransferRegion
472 + ) -> bool:
473 + # Page padding sits between unpadded payloads. Allow that gap, but
474 + # keep the whole run inside the row that `run_start` began. Anchoring
475 + # on `prev` would slide the bound forward as the run grows.
476 + return (
477 + prev.group_index == nxt.group_index
478 + and prev.shared_group_ids == nxt.shared_group_ids
479 + and prev.block_len == nxt.block_len
480 + and prev.base_addr + prev.kv_block_len <= nxt.base_addr
481 + and nxt.base_addr + nxt.kv_block_len
482 + <= run_start.base_addr + run_start.block_len
483 + )
484 +
485 + def span(region: TransferRegion, kv_block_len: int) -> TransferRegion:
486 + return TransferRegion(
487 + layer_name=region.layer_name,
488 + layer_index=region.layer_index,
489 + base_addr=region.base_addr,
490 + block_len=region.block_len,
491 + kv_block_len=kv_block_len,
492 + group_index=region.group_index,
493 + shared_group_ids=region.shared_group_ids,
494 + row_offset=region.row_offset,
495 + )
496 +
497 + def promote(
498 + merged_local: list[TransferRegion], merged_remote: list[TransferRegion]
499 + ) -> tuple[list[TransferRegion], list[TransferRegion]]:
500 + if not promote_full_row:
501 + return merged_local, merged_remote
502 + # Each KV group overlays the row from offset 0. Promoting that run to
503 + # the full row copies bytes the group's block ids exclusively own.
504 + # Equal block_len is the row-structure check: both peers packed the
505 + # same stride, so row_offset does not need a second compare in align.
506 + promoted_local: list[TransferRegion] = []
507 + promoted_remote: list[TransferRegion] = []
508 + for local_region, remote_region in zip(merged_local, merged_remote):
509 + if (
510 + local_region.row_offset == 0
511 + and remote_region.row_offset == 0
512 + and local_region.block_len == remote_region.block_len
513 + and local_region.kv_block_len <= local_region.block_len
514 + and remote_region.kv_block_len <= remote_region.block_len
515 + ):
516 + promoted_local.append(span(local_region, local_region.block_len))
517 + promoted_remote.append(span(remote_region, remote_region.block_len))
518 + else:
519 + promoted_local.append(local_region)
520 + promoted_remote.append(remote_region)
521 + return promoted_local, promoted_remote
522 +
523 + if len(local_regions) <= 1:
524 + return promote(list(local_regions), list(remote_regions))
525 +
526 + merged_local: list[TransferRegion] = []
527 + merged_remote: list[TransferRegion] = []
528 + run_local = [local_regions[0]]
529 + run_remote = [remote_regions[0]]
530 +
531 + def flush() -> None:
532 + if len(run_local) == 1:
533 + merged_local.append(run_local[0])
534 + merged_remote.append(run_remote[0])
535 + return
536 + for merged, run in ((merged_local, run_local), (merged_remote, run_remote)):
537 + span_len = run[-1].base_addr + run[-1].kv_block_len - run[0].base_addr
538 + merged.append(span(run[0], span_len))
539 +
540 + for local_region, remote_region in zip(local_regions[1:], remote_regions[1:]):
541 + if adjacent(run_local[0], run_local[-1], local_region) and adjacent(
542 + run_remote[0], run_remote[-1], remote_region
543 + ):
544 + run_local.append(local_region)
545 + run_remote.append(remote_region)
546 + continue
547 + flush()
548 + run_local = [local_region]
549 + run_remote = [remote_region]
550 + flush()
551 + return promote(merged_local, merged_remote)
552 +
553 +
554 +def _has_opaque_packed_row(
555 + row_offsets: list[int], kv_block_lens: list[int], block_lens: list[int]
556 +) -> bool:
557 + """Branch B is one row-sized region at offset 0, so later layer names are gone.
558 +
559 + Page-contiguous views keep kv_block_len below the stride until promotion,
560 + and promotion happens after this handshake check.
561 + """
562 + if (
563 + not row_offsets
564 + or len(row_offsets) != len(kv_block_lens)
565 + or len(row_offsets) != len(block_lens)
566 + ):
567 + return False
568 + return any(
569 + offset == 0 and length == stride
570 + for offset, length, stride in zip(row_offsets, kv_block_lens, block_lens)
571 + )
572 +
573 +
574 +def _pp_mismatch_hides_packed_layers(
575 + local_pp_size: int,
576 + remote_pp_size: int,
577 + *,
578 + local_has_opaque_row: bool,
579 + remote_row_offsets: list[int],
580 + remote_kv_block_lens: list[int],
581 + remote_block_lens: list[int],
582 +) -> bool:
583 + """A row that dropped per-layer names cannot be aligned across PP sizes.
584 +
585 + Either side may be the one that collapsed the row. Matching PP still
586 + lines up, because both stages keep the same first layer name.
587 + """
588 + if local_pp_size == remote_pp_size:
589 + return False
590 + return local_has_opaque_row or _has_opaque_packed_row(
591 + remote_row_offsets, remote_kv_block_lens, remote_block_lens
592 + )
593 +
594 +
595 def _get_tensor_dense_flag(tensor: torch.Tensor) -> bool | None:
596 is_dense = getattr(tensor, "is_non_overlapping_and_dense", None)
597 if callable(is_dense):
@@ -405,6 +614,12 @@ class MooncakeXferMetadata(
614 registered_layer_names: list[str] = msgspec.field(default_factory=list)
615 registered_layer_indices: list[int] = msgspec.field(default_factory=list)
616 registered_group_indices: list[int] = msgspec.field(default_factory=list)
617 + # Parallel to registered regions. Each entry is the groups that share
618 + # that region. Empty means the peer did not send the field.
619 + registered_shared_group_ids: list[list[int]] = msgspec.field(default_factory=list)
620 + # Byte offset of each registered view from its block row. -1 means the
621 + # peer did not send the field, so the run must not be promoted.
622 + registered_row_offsets: list[int] = msgspec.field(default_factory=list)
623 remote_pp_size: int = 1
624
625
@@ -999,7 +1214,18 @@ class MooncakeConnectorWorker:
1214 self.registered_layer_names: list[str] = []
1215 self.registered_layer_indices: list[int] = []
1216 self.registered_group_indices: list[int] = []
1217 + # Parallel to registered regions. -1 flattens only these groups.
1218 + self.region_shared_groups: list[tuple[int, ...]] = []
1219 + # Parallel to registered regions. -1 is not a packed row view.
1220 + self.region_row_offsets: list[int] = []
1221 + self.opaque_packed_storages: set[int] = set()
1222 self.seen_base_addresses: list[int] = []
1223 + # Aligned regions depend only on the peer's registered layout.
1224 + # The third item is an error string when alignment cannot proceed.
1225 + self._prepared_transfer_regions: dict[
1226 + tuple,
1227 + tuple[list[TransferRegion], list[TransferRegion], str | None],
1228 + ] = {}
1229
1230 assert (parallel_config := vllm_config.parallel_config)
1231 dp_rank = parallel_config.data_parallel_index
@@ -1266,39 +1492,35 @@ class MooncakeConnectorWorker:
1492 )
1493 await sock.send_multipart((identity, self._encoder.encode(response)))
1494 return
1269 - local_regions = self._get_transfer_regions(
1270 - self.kv_caches_base_addr,
1271 - self.block_len_per_layer,
1272 - self.kv_block_len_per_layer,
1273 - self.registered_layer_names,
1274 - self.registered_layer_indices,
1275 - self.registered_group_indices,
1276 - )
1277 - remote_regions = self._get_transfer_regions(
1278 - meta.kv_caches_base_addr,
1279 - meta.block_lens,
1280 - meta.kv_block_lens,
1281 - meta.registered_layer_names,
1282 - meta.registered_layer_indices,
1283 - meta.registered_group_indices,
1284 - )
1285 - local_regions, remote_regions, align_err = _align_transfer_regions(
1286 - local_regions,
1287 - remote_regions,
1288 - allow_partial_layers=(
1289 - meta.remote_pp_size > 1 and meta.remote_pp_size != self.pp_size
1290 - ),
1291 - )
1292 - if align_err is not None:
1495 + if _pp_mismatch_hides_packed_layers(
1496 + self.pp_size,
1497 + meta.remote_pp_size,
1498 + local_has_opaque_row=bool(self.opaque_packed_storages),
1499 + remote_row_offsets=meta.registered_row_offsets,
1500 + remote_kv_block_lens=meta.kv_block_lens,
1501 + remote_block_lens=meta.block_lens,
1502 + ):
1503 + msg = (
1504 + "Mooncake non-page-contiguous packed regions keep only the "
1505 + "first layer name, so they cannot be aligned across PP sizes "
1506 + f"{self.pp_size} and {meta.remote_pp_size}."
1507 + )
1508 + logger.error(msg)
1509 response = MooncakeXferResponse(
1510 status=MooncakeXferResponseStatus.ERROR,
1295 - err_msg=align_err,
1511 + err_msg=msg,
1512 )
1513 await sock.send_multipart((identity, self._encoder.encode(response)))
1514 return
1299 - validation_err = self._validate_head_resharding_layout(
1300 - meta.remote_tp_size, local_regions
1301 - ) or _validate_asymmetric_region_lengths(
1515 + local_regions, remote_regions, prep_err = self._prepare_transfer_regions(meta)
1516 + if prep_err is not None:
1517 + response = MooncakeXferResponse(
1518 + status=MooncakeXferResponseStatus.ERROR,
1519 + err_msg=prep_err,
1520 + )
1521 + await sock.send_multipart((identity, self._encoder.encode(response)))
1522 + return
1523 + validation_err = _validate_asymmetric_region_lengths(
1524 local_regions=local_regions,
1525 remote_regions=remote_regions,
1526 local_tp_size=self.tp_size,
@@ -1600,12 +1822,16 @@ class MooncakeConnectorWorker:
1822 "Aligned Mooncake transfer regions must belong to the same "
1823 "KV group."
1824 )
1603 - group_index = local_region.group_index
1604 - assert group_index < len(local_block_ids_by_group), (
1605 - "Transfer region references a missing KV group."
1825 + local_block_ids = _block_ids_for_region(
1826 + local_block_ids_by_group,
1827 + local_region.group_index,
1828 + local_region.shared_group_ids,
1829 + )
1830 + remote_block_ids = _block_ids_for_region(
1831 + remote_block_ids_by_group,
1832 + remote_region.group_index,
1833 + local_region.shared_group_ids,
1834 )
1607 - local_block_ids = local_block_ids_by_group[group_index]
1608 - remote_block_ids = remote_block_ids_by_group[group_index]
1835 if not local_block_ids:
1836 continue
1837
@@ -1738,6 +1964,17 @@ class MooncakeConnectorWorker:
1964 self.registered_layer_names = []
1965 self.registered_layer_indices = []
1966 self.registered_group_indices = []
1967 + self.region_shared_groups = []
1968 + self.region_row_offsets = []
1969 + self.opaque_packed_storages = set()
1970 + self._prepared_transfer_regions.clear()
1971 +
1972 + packed_storage_to_region: dict[int, int] = {}
1973 + packed_view_to_region: dict[tuple[int, int, int], int] = {}
1974 + collapsed_views = 0
1975 + num_blocks = (
1976 + self.kv_cache_config.num_blocks * self._physical_blocks_per_logical_kv_block
1977 + )
1978
1979 for layer_name, cache in kv_caches.items():
1980 layer_index = extract_layer_index(layer_name)
@@ -1751,6 +1988,104 @@ class MooncakeConnectorWorker:
1988 # One raw page tensor per layer; for Mamba that page holds all the
1989 # recurrent states, unpacked only when binding the cache for execution.
1990 self._log_debug_cache_registration(layer_name, cache)
1991 + group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1992 + storage = cache.untyped_storage()
1993 + storage_addr = storage.data_ptr()
1994 + if storage_addr not in seen_storage_ptrs:
1995 + seen_storage_ptrs.add(storage_addr)
1996 + kv_data_ptrs.append(storage_addr)
1997 + kv_data_lens.append(storage.nbytes())
1998 +
1999 + is_mla_region = isinstance(
2000 + layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
2001 + )
2002 + tensor_blocks = cache.shape[0] if cache.ndim > 1 else num_blocks
2003 + if cache.ndim == 1:
2004 + # A flat buffer is block-major only when it is an integral run
2005 + # of this layer's pages. NIXL uses the same remainder check.
2006 + hnc_contiguous = False
2007 + page = (
2008 + layer_spec.page_size_bytes
2009 + if isinstance(layer_spec, AttentionSpec)
2010 + else 0
2011 + )
2012 + remainder = (
2013 + storage.nbytes() - tensor_blocks * page
2014 + if tensor_blocks > 0 and page > 0
2015 + else -1
2016 + )
2017 + if 0 <= remainder < page:
2018 + block_stride = page
2019 + storage_is_block_major = True
2020 + else:
2021 + block_stride = 0
2022 + storage_is_block_major = False
2023 + else:
2024 + block_stride = cache.stride(0) * cache.element_size()
2025 + storage_is_block_major = (
2026 + tensor_blocks > 0
2027 + and tensor_blocks * block_stride == storage.nbytes()
2028 + )
2029 + hnc_contiguous = (
2030 + cache.ndim == 4
2031 + and cache.stride(2) == cache.shape[3]
2032 + and cache.stride(1) == cache.shape[2] * cache.shape[3]
2033 + )
2034 + # Follow NIXL packed pages. Page-contiguous slices stay per layer so
2035 + # hetero PP can align by name (#56033). A non-contiguous row cannot
2036 + # name a sub-span, so it is one region anchored at the allocation.
2037 + use_packed_row = (
2038 + not isinstance(layer_spec, (MambaSpec, KpoolTailSpec))
2039 + and storage_is_block_major
2040 + and block_stride > 0
2041 + and (is_mla_region or not hnc_contiguous)
2042 + )
2043 + if use_packed_row:
2044 + page_contiguous = (
2045 + cache.ndim > 1
2046 + and cache[0].is_contiguous()
2047 + and 0 < cache[0].nbytes <= block_stride
2048 + )
2049 + if page_contiguous:
2050 + view_addr = cache.data_ptr()
2051 + kv_block_len = cache[0].nbytes
2052 + else:
2053 + # NIXL branch B: length is the row, base is the allocation.
2054 + view_addr = storage_addr
2055 + kv_block_len = block_stride
2056 + self.opaque_packed_storages.add(storage_addr)
2057 + # Same PP on both sides still aligns by the first layer name.
2058 + # A remote PP mismatch is rejected at send time, once the
2059 + # peer's PP size is known.
2060 + dedupe_key = (view_addr, block_stride, kv_block_len)
2061 + region_idx = packed_view_to_region.get(dedupe_key)
2062 + if region_idx is not None:
2063 + collapsed_views += 1
2064 + # -1 means these groups share the view. Flatten only them.
2065 + groups = list(self.region_shared_groups[region_idx])
2066 + if group_index not in groups:
2067 + groups.append(group_index)
2068 + self.region_shared_groups[region_idx] = tuple(groups)
2069 + if len(self.region_shared_groups[region_idx]) > 1:
2070 + self.registered_group_indices[region_idx] = (
2071 + _SHARED_REGION_GROUP_ID
2072 + )
2073 + continue
2074 + if not page_contiguous:
2075 + self.opaque_packed_storages.add(storage_addr)
2076 + if storage_addr not in packed_storage_to_region:
2077 + packed_storage_to_region[storage_addr] = len(region_base_addresses)
2078 + packed_view_to_region[dedupe_key] = len(region_base_addresses)
2079 + region_base_addresses.append(view_addr)
2080 + self.block_len_per_layer.append(block_stride)
2081 + self.kv_block_len_per_layer.append(kv_block_len)
2082 + self.registered_layer_names.append(layer_name)
2083 + self.registered_layer_indices.append(layer_index)
2084 + self.registered_group_indices.append(group_index)
2085 + self.region_shared_groups.append((group_index,))
2086 + self.region_row_offsets.append(view_addr - storage_addr)
2087 + continue
2088 +
2089 block_is_contiguous = is_non_overlapping_and_dense(cache[0])
2090 if not block_is_contiguous:
2091 # Non-block-compact layouts scatter a block across per-head
@@ -1791,15 +2126,10 @@ class MooncakeConnectorWorker:
2126 self.kv_block_len_per_layer.append(kv_block_len)
2127 self.registered_layer_names.append(layer_name)
2128 self.registered_layer_indices.append(layer_index)
1794 - self.registered_group_indices.append(
1795 - self.kv_cache_config.transfer_group_index_by_layer[layer_name]
1796 - )
1797 - storage = cache.untyped_storage()
1798 - storage_addr = storage.data_ptr()
1799 - if storage_addr not in seen_storage_ptrs:
1800 - seen_storage_ptrs.add(storage_addr)
1801 - kv_data_ptrs.append(storage_addr)
1802 - kv_data_lens.append(storage.nbytes())
2129 + self.registered_group_indices.append(group_index)
2130 + self.region_shared_groups.append((group_index,))
2131 + # Not a view into a shared packed row, so do not promote it.
2132 + self.region_row_offsets.append(-1)
2133
2134 self.kv_caches_base_addr = region_base_addresses
2135 self.seen_base_addresses = kv_data_ptrs
@@ -1807,6 +2137,18 @@ class MooncakeConnectorWorker:
2137 if not kv_data_ptrs:
2138 raise RuntimeError("No KV cache tensors were registered with Mooncake.")
2139
2140 + if packed_storage_to_region and (
2141 + collapsed_views
2142 + or len(packed_view_to_region) != len(packed_storage_to_region)
2143 + ):
2144 + logger.info(
2145 + "Mooncake packed KV xfer: %s layer regions over %s storage "
2146 + "regions (%s duplicate views collapsed).",
2147 + len(packed_view_to_region),
2148 + len(packed_storage_to_region),
2149 + collapsed_views,
2150 + )
2151 +
2152 ret_value = self.engine.batch_register_memory(kv_data_ptrs, kv_data_lens)
2153 if ret_value != 0:
2154 raise RuntimeError("Mooncake batch memory registration failed.")
@@ -1962,6 +2304,10 @@ class MooncakeConnectorWorker:
2304 registered_layer_names=self.registered_layer_names,
2305 registered_layer_indices=self.registered_layer_indices,
2306 registered_group_indices=self.registered_group_indices,
2307 + registered_shared_group_ids=[
2308 + list(groups) for groups in self.region_shared_groups
2309 + ],
2310 + registered_row_offsets=self.region_row_offsets,
2311 )
2312
2313 encoded_data = self._encoder.encode(metadata)
@@ -2209,6 +2555,131 @@ class MooncakeConnectorWorker:
2555 def _producer_cache_is_replicated(self) -> bool:
2556 return self.transfer_topo.local_replicates_kv_cache
2557
2558 + def _prepare_transfer_regions(
2559 + self, meta: MooncakeXferMetadata
2560 + ) -> tuple[list[TransferRegion], list[TransferRegion], str | None]:
2561 + """Align and coalesce once per peer layout.
2562 +
2563 + ``send_kv_to_decode`` runs on every scheduler batch. The registered
2564 + layout does not change until the next ``register_kv_caches``.
2565 + """
2566 + cache_key = (
2567 + meta.remote_hostname,
2568 + meta.remote_port,
2569 + meta.remote_tp_size,
2570 + meta.remote_tp_rank,
2571 + meta.remote_pp_size,
2572 + tuple(meta.kv_caches_base_addr),
2573 + tuple(meta.block_lens),
2574 + tuple(meta.kv_block_lens),
2575 + tuple(meta.registered_layer_names),
2576 + tuple(meta.registered_layer_indices),
2577 + tuple(meta.registered_group_indices),
2578 + tuple(tuple(groups) for groups in meta.registered_shared_group_ids),
2579 + tuple(meta.registered_row_offsets),
2580 + )
2581 + cached = self._prepared_transfer_regions.get(cache_key)
2582 + if cached is not None:
2583 + return cached
2584 +
2585 + def finish(
2586 + local: list[TransferRegion],
2587 + remote: list[TransferRegion],
2588 + err: str | None,
2589 + ) -> tuple[list[TransferRegion], list[TransferRegion], str | None]:
2590 + prepared = (local, remote, err)
2591 + self._prepared_transfer_regions[cache_key] = prepared
2592 + return prepared
2593 +
2594 + shared_groups: list[tuple[int, ...]] | None = self.region_shared_groups
2595 + if len(self.region_shared_groups) != len(self.registered_layer_names):
2596 + if _SHARED_REGION_GROUP_ID in self.registered_group_indices:
2597 + return finish(
2598 + [],
2599 + [],
2600 + (
2601 + "Mooncake shared-group metadata does not match "
2602 + "registered layers, so a shared packed region cannot "
2603 + "be transferred."
2604 + ),
2605 + )
2606 + shared_groups = None
2607 + local_rows: list[int] | None = self.region_row_offsets
2608 + if len(self.region_row_offsets) != len(self.registered_layer_names):
2609 + local_rows = None
2610 + local_regions = self._get_transfer_regions(
2611 + self.kv_caches_base_addr,
2612 + self.block_len_per_layer,
2613 + self.kv_block_len_per_layer,
2614 + self.registered_layer_names,
2615 + self.registered_layer_indices,
2616 + self.registered_group_indices,
2617 + shared_groups,
2618 + local_rows,
2619 + )
2620 + remote_shared_groups = [
2621 + tuple(groups) for groups in meta.registered_shared_group_ids
2622 + ]
2623 + remote_shared: list[tuple[int, ...]] | None = remote_shared_groups
2624 + if len(remote_shared_groups) != len(meta.registered_layer_names):
2625 + remote_shared = None
2626 + remote_rows = meta.registered_row_offsets
2627 + remote_row_offsets: list[int] | None = remote_rows
2628 + if len(remote_rows) != len(meta.registered_layer_names):
2629 + remote_row_offsets = None
2630 + remote_regions = self._get_transfer_regions(
2631 + meta.kv_caches_base_addr,
2632 + meta.block_lens,
2633 + meta.kv_block_lens,
2634 + meta.registered_layer_names,
2635 + meta.registered_layer_indices,
2636 + meta.registered_group_indices,
2637 + remote_shared,
2638 + remote_row_offsets,
2639 + )
2640 + pre_align_local = local_regions
2641 + pre_align_remote = remote_regions
2642 + local_regions, remote_regions, align_err = _align_transfer_regions(
2643 + local_regions,
2644 + remote_regions,
2645 + allow_partial_layers=meta.remote_pp_size != self.pp_size,
2646 + )
2647 + if align_err is not None:
2648 + return finish([], [], align_err)
2649 + # Head checks must see every layer name. Coalesce keeps only the first.
2650 + head_err = self._validate_head_resharding_layout(
2651 + meta.remote_tp_size, local_regions
2652 + )
2653 + if head_err is not None:
2654 + return finish([], [], head_err)
2655 + # Hetero TP indexes kv_block_len as one head-sharded payload. A merged
2656 + # multi-layer span would land that offset on the wrong layer.
2657 + # Replicated KV keeps that offset at 0. MLA is replicated in practice;
2658 + # gating on use_mla alone would promote and then fail the length check
2659 + # for a non-replicated MLA rank with tp_ratio != 1.
2660 + tp_ratio = _get_tp_ratio(self.tp_size, meta.remote_tp_size)
2661 + if tp_ratio == 1 or self._producer_cache_is_replicated():
2662 + covered_both_sides = len(local_regions) == len(pre_align_local) and len(
2663 + remote_regions
2664 + ) == len(pre_align_remote)
2665 + n_aligned = len(local_regions)
2666 + local_regions, remote_regions = _coalesce_contiguous_transfer_regions(
2667 + local_regions,
2668 + remote_regions,
2669 + promote_full_row=covered_both_sides,
2670 + )
2671 + logger.info(
2672 + "Mooncake packed coalesce: %s aligned regions -> %s transfer regions.",
2673 + n_aligned,
2674 + len(local_regions),
2675 + )
2676 + self._prepared_transfer_regions[cache_key] = (
2677 + local_regions,
2678 + remote_regions,
2679 + None,
2680 + )
2681 + return local_regions, remote_regions, None
2682 +
2683 def _validate_head_resharding_layout(
2684 self, remote_tp_size: int, local_regions: list[TransferRegion]
2685 ) -> str | None:
@@ -2242,6 +2713,8 @@ class MooncakeConnectorWorker:
2713 layer_names: list[str],
2714 layer_indices: list[int],
2715 group_indices: list[int] | None = None,
2716 + shared_group_ids: list[tuple[int, ...]] | None = None,
2717 + row_offsets: list[int] | None = None,
2718 ) -> list[TransferRegion]:
2719 if not group_indices:
2720 group_indices = [
@@ -2255,6 +2728,8 @@ class MooncakeConnectorWorker:
2728 layer_names=layer_names,
2729 layer_indices=layer_indices,
2730 group_indices=group_indices,
2731 + shared_group_ids=shared_group_ids,
2732 + row_offsets=row_offsets,
2733 )
2734
2735 def _get_sender_transfer_plan(