[Bugfix][ROCm][Disagg] Fix MoRIIO shared KV memory region registration (#53698)
Signed-off-by: avininjamay8 <avpaul@amd.com> Co-authored-by: avininjamay8 <avpaul@amd.com>
Avinash Paul committed
Aug 26, 2026 at 21:40 UTC
8d301f075b970427ae2486194f3694cdc04fde71
5 files changed
+97
-31
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+14
-24
@@ -341,7 +341,7 @@ def test_mixed_layers_compute_distinct_offsets_per_layer():
341
assert interleaved != indexer
342
343
344
-def test_write_transfer_plan_caches_offsets_per_geometry():
344
+def test_write_transfer_plan_caches_offsets_per_layer_geometry():
345
kv_caches = {
346
"dense0": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
347
"dense1": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
@@ -354,7 +354,12 @@ def test_write_transfer_plan_caches_offsets_per_geometry():
354
layer_name_to_local_kv_cache_metadata: dict[str, list[Any]]
355
356
def _compute_block_transfer_offsets(
357
- self, layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
357
+ self,
358
+ layer_name,
359
+ local_block_ids,
360
+ remote_block_ids,
361
+ remote_moriio_meta,
362
+ remote_engine_id=None,
363
):
364
calls.append(layer_name)
365
call_id = len(calls)
@@ -369,41 +374,26 @@ def test_write_transfer_plan_caches_offsets_per_geometry():
374
remote_meta = _remote_meta()
375
376
dense0_plan = writer._prepare_transfer_plan(
372
- SimpleNamespace(
373
- layer_name="dense0",
374
- local_block_ids=[1, 3],
375
- request_id="req",
376
- transfer_id="xfer",
377
- ),
377
+ _write_task("dense0"),
378
request_info,
379
remote_meta,
380
)
381
dense1_plan = writer._prepare_transfer_plan(
382
- SimpleNamespace(
383
- layer_name="dense1",
384
- local_block_ids=[1, 3],
385
- request_id="req",
386
- transfer_id="xfer",
387
- ),
382
+ _write_task("dense1"),
383
request_info,
384
remote_meta,
385
)
386
indexer_plan = writer._prepare_transfer_plan(
392
- SimpleNamespace(
393
- layer_name="indexer",
394
- local_block_ids=[1, 3],
395
- request_id="req",
396
- transfer_id="xfer",
397
- ),
387
+ _write_task("indexer"),
388
request_info,
389
remote_meta,
390
)
391
402
- assert calls == ["dense0", "indexer"]
392
+ assert calls == ["dense0", "dense1", "indexer"]
393
assert dense0_plan.transfer_local_offsets == [1]
404
- assert dense1_plan.transfer_local_offsets == [1]
405
- assert indexer_plan.transfer_local_offsets == [2]
406
- assert len(request_info.transfer_offsets) == 2
394
+ assert dense1_plan.transfer_local_offsets == [2]
395
+ assert indexer_plan.transfer_local_offsets == [3]
396
+ assert len(request_info.transfer_offsets) == 3
397
398
399
def test_write_scheduler_deduplicates_layers_and_seals_expected_count():
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+2
-1
@@ -41,6 +41,7 @@ Transfer = tuple[int, float]
41
EngineId = str
42
ReqId = str
43
TransferId = str
44
+TransferOffsetsKey = tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]
45
46
47
class MoRIIOTransferAck(NamedTuple):
@@ -91,7 +92,7 @@ class RemoteAllocInfo:
92
completion_notified: bool = False
93
transfer_statuses: list[Any] = field(default_factory=list)
94
transfer_offsets: dict[
94
- tuple[tuple[int, ...], tuple[int, ...], torch.dtype],
95
+ TransferOffsetsKey,
96
tuple[list[int], list[int], list[int]],
97
] = field(default_factory=dict)
98
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+70
-2
@@ -1226,6 +1226,8 @@ class MoRIIOConnectorWorker:
1226
1227
# KV Caches and moriio tracking data.
1228
self.kv_caches: dict[str, torch.Tensor] = {}
1229
+ self.kv_layer_mr_offset: dict[str, int] = {}
1230
+ self.layer_base_addr_index: dict[str, int] = {}
1231
1232
# Map of engine_id -> kv_caches_base_addr. For TP case, each local
1233
# rank will still only pull from a single remote TP worker.
@@ -1695,6 +1697,50 @@ class MoRIIOConnectorWorker:
1697
self.layer_to_spec,
1698
)
1699
1700
+ def _build_shared_kv_mr(
1701
+ self, kv_caches: dict[str, torch.Tensor]
1702
+ ) -> tuple[torch.Tensor, dict[str, int]]:
1703
+ page_size = 4096
1704
+ backing = next(iter(kv_caches.values())).view(torch.uint8)
1705
+ nbytes = backing.untyped_storage().nbytes()
1706
+ base = torch.as_strided(backing, (nbytes,), (1,), 0)
1707
+ ptr = base.data_ptr()
1708
+ slack = ptr % page_size
1709
+ end = 0
1710
+ for layer_name in kv_caches:
1711
+ _, region_len = next(
1712
+ iter(self._iter_layer_registration_regions(layer_name))
1713
+ )
1714
+ end = max(end, kv_caches[layer_name].data_ptr() - ptr + region_len)
1715
+ reg_nbytes = ((slack + end + page_size - 1) // page_size) * page_size
1716
+ if reg_nbytes > nbytes:
1717
+ raise ValueError(
1718
+ f"shared KV backing too small for page-aligned MR: "
1719
+ f"need {reg_nbytes}, have {nbytes}"
1720
+ )
1721
+ reg = torch.as_strided(base, (reg_nbytes,), (1,), 0 if slack == 0 else -slack)
1722
+ mr_ptr = reg.data_ptr()
1723
+ return reg, {
1724
+ name: cache.data_ptr() - mr_ptr for name, cache in kv_caches.items()
1725
+ }
1726
+
1727
+ def _remote_layer_mr_offset(
1728
+ self,
1729
+ layer_name: str,
1730
+ remote_moriio_meta: MoRIIOAgentMetadata,
1731
+ remote_engine_id: EngineId | None,
1732
+ ) -> int:
1733
+ if not self.kv_layer_mr_offset or not remote_engine_id:
1734
+ return 0
1735
+ idx = self.layer_base_addr_index[layer_name]
1736
+ metas = self.layer_name_to_remote_kv_cache_metadata.get(
1737
+ remote_engine_id, {}
1738
+ ).get(layer_name)
1739
+ if not metas or idx >= len(remote_moriio_meta.kv_caches_base_addr):
1740
+ return 0
1741
+ mr_base = self.moriio_wrapper.get_unpack_memory_metadata(metas[0]).data
1742
+ return remote_moriio_meta.kv_caches_base_addr[idx] - mr_base
1743
+
1744
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1745
"""Register the KV Cache data in moriio."""
1746
@@ -1754,8 +1800,11 @@ class MoRIIOConnectorWorker:
1800
self.dst_num_blocks[self.engine_id] = self.num_blocks
1801
kv_caches_base_addr = []
1802
caches_data = []
1803
+ self.layer_base_addr_index = {}
1804
+ base_addr_idx = 0
1805
1806
for layer_name in kv_caches:
1807
+ self.layer_base_addr_index[layer_name] = base_addr_idx
1808
geometry = self._get_layer_transfer_geometry(layer_name)
1809
if geometry.block_size != self.block_size:
1810
raise ValueError(
@@ -1767,12 +1816,21 @@ class MoRIIOConnectorWorker:
1816
base_addr = cache.data_ptr()
1817
caches_data.append((base_addr, region_len, cache.device.index, ""))
1818
kv_caches_base_addr.append(base_addr)
1819
+ base_addr_idx += 1
1820
+
1821
+ shared_backing = len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1
1822
+ shared_mr = None
1823
+ if shared_backing:
1824
+ reg_tensor, self.kv_layer_mr_offset = self._build_shared_kv_mr(kv_caches)
1825
+ shared_mr = self.moriio_wrapper.register_local_tensor(reg_tensor)
1826
1827
for layer_name, kv_cache in kv_caches.items():
1828
if layer_name not in self.layer_name_to_local_kv_cache_metadata:
1829
self.layer_name_to_local_kv_cache_metadata[layer_name] = []
1830
1775
- moriio_mem_metadata = self.moriio_wrapper.register_local_tensor(kv_cache)
1831
+ moriio_mem_metadata = shared_mr or (
1832
+ self.moriio_wrapper.register_local_tensor(kv_cache)
1833
+ )
1834
self.layer_name_to_local_kv_cache_metadata[layer_name].append(
1835
moriio_mem_metadata
1836
)
@@ -2488,6 +2546,7 @@ class MoRIIOConnectorWorker:
2546
remote_block_ids: list[int],
2547
remote_moriio_meta: MoRIIOAgentMetadata,
2548
remote_tp_size: int | None = None,
2549
+ remote_engine_id: EngineId | None = None,
2550
) -> tuple[list[int], list[int], list[int]]:
2551
"""Compute transfer offsets for block data.
2552
@@ -2509,7 +2568,7 @@ class MoRIIOConnectorWorker:
2568
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
2569
is_mla=self._is_mla_cache_layer(layer_name),
2570
)
2512
- return compute_block_transfer_offsets(
2571
+ local, remote, sizes = compute_block_transfer_offsets(
2572
layer_name=layer_name,
2573
kv_cache=self.kv_caches[layer_name],
2574
layer_to_spec=self.layer_to_spec,
@@ -2520,6 +2579,14 @@ class MoRIIOConnectorWorker:
2579
local, remote, sizes, assume_sorted=False
2580
),
2581
)
2582
+ if self.kv_layer_mr_offset:
2583
+ local_off = self.kv_layer_mr_offset[layer_name]
2584
+ remote_off = self._remote_layer_mr_offset(
2585
+ layer_name, remote_moriio_meta, remote_engine_id
2586
+ )
2587
+ local = [x + local_off for x in local]
2588
+ remote = [x + remote_off for x in remote]
2589
+ return local, remote, sizes
2590
2591
@staticmethod
2592
def _is_sq_full_status(status) -> bool:
@@ -2592,6 +2659,7 @@ class MoRIIOConnectorWorker:
2659
remote_block_ids,
2660
remote_moriio_meta,
2661
remote_tp_size=remote_tp_size,
2662
+ remote_engine_id=remote_dp_engine_id,
2663
)
2664
# TODO : apply multi-session batch-read when moriio support it
2665
#
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+4
-3
@@ -377,16 +377,17 @@ class MoRIIOWriter:
377
The transfer plan
378
"""
379
layer_cache = self.worker.kv_caches[task.layer_name]
380
- geometry_key = _get_write_geometry_key(layer_cache)
381
- offsets = request_info.transfer_offsets.get(geometry_key)
380
+ key = (task.layer_name, *_get_write_geometry_key(layer_cache))
381
+ offsets = request_info.transfer_offsets.get(key)
382
if offsets is None:
383
offsets = self.worker._compute_block_transfer_offsets(
384
task.layer_name,
385
task.local_block_ids,
386
request_info.block_ids,
387
remote_moriio_meta,
388
+ remote_engine_id=task.dst_engine_id,
389
)
389
- request_info.transfer_offsets[geometry_key] = offsets
390
+ request_info.transfer_offsets[key] = offsets
391
392
# Get session index
393
layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
vllm/v1/worker/utils.py
+7
-1
@@ -392,7 +392,13 @@ def allocate_kv_cache(
392
393
sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors}
394
assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
395
- buf = torch.zeros(sizes.pop(), dtype=torch.int8, device=device)
395
+ raw_size = sizes.pop()
396
+ page_size = 4096
397
+ buf = torch.zeros(
398
+ ((raw_size + page_size - 1) // page_size) * page_size,
399
+ dtype=torch.int8,
400
+ device=device,
401
+ )
402
403
kv_caches: dict[str, torch.Tensor] = {}
404
for tensor in kv_cache_config.kv_cache_tensors: