[ROCm][P/D] Support MiniMax-M3 mixed KV layouts in MoRIIO READ mode (#46039)
Signed-off-by: Jun Kang Chow <junkangchow@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Tan Pin Siang <tanpinsiang@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Chun Fang <chun.fang@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
junkang1991 committed
Jun 21, 2026 at 20:55 UTC
b91b7726e068d5eed2aa3bd084cdeee9456da424
4 files changed
+566
-72
tests/v1/kv_connector/unit/test_moriio_connector.py
+33
-4
@@ -36,13 +36,33 @@ from vllm.utils.network_utils import (
36
get_ip,
37
make_zmq_path,
38
)
39
-from vllm.v1.kv_cache_interface import KVCacheConfig
39
+from vllm.v1.kv_cache_interface import (
40
+ FullAttentionSpec,
41
+ KVCacheConfig,
42
+ KVCacheGroupSpec,
43
+ KVCacheTensor,
44
+)
45
46
from .utils import create_request, create_scheduler
47
48
49
def _make_test_kv_cache_config() -> KVCacheConfig:
45
- return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
50
+ layer_names = ["layer0", "layer1", "layer2"]
51
+ return KVCacheConfig(
52
+ num_blocks=2,
53
+ kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
54
+ kv_cache_groups=[
55
+ KVCacheGroupSpec(
56
+ layer_names=layer_names,
57
+ kv_cache_spec=FullAttentionSpec(
58
+ block_size=16,
59
+ num_kv_heads=4,
60
+ head_size=64,
61
+ dtype=torch.float16,
62
+ ),
63
+ )
64
+ ],
65
+ )
66
67
68
aiter_available = importlib.util.find_spec("aiter") is not None
@@ -175,9 +195,18 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
195
REMOTE_ENGINE_ID = "remote_engine"
196
197
def __init__(
178
- self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
198
+ self,
199
+ vllm_config,
200
+ engine_id,
201
+ *args,
202
+ hand_shake_latency: float = 1.8,
203
+ kv_cache_layout="HND",
204
+ kv_cache_config=None,
205
+ **kwargs,
206
):
180
- super().__init__(*args, **kwargs)
207
+ super().__init__(
208
+ vllm_config, engine_id, kv_cache_config or _make_test_kv_cache_config()
209
+ )
210
211
212
def create_vllm_config(
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
new
+228
@@ -0,0 +1,228 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib.util
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.platforms import current_platform
11
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
12
+
13
+aiter_available = importlib.util.find_spec("aiter") is not None
14
+mori_available = importlib.util.find_spec("mori") is not None
15
+
16
+if not (current_platform.is_rocm() and mori_available):
17
+ pytest.skip(
18
+ "MoRIIOs are only available on ROCm with mori package installed",
19
+ allow_module_level=True,
20
+ )
21
+
22
+moriio_layout = importlib.import_module(
23
+ "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
24
+)
25
+
26
+
27
+def _full_spec(block_size: int = 4) -> FullAttentionSpec:
28
+ return FullAttentionSpec(
29
+ block_size=block_size,
30
+ num_kv_heads=2,
31
+ head_size=3,
32
+ dtype=torch.bfloat16,
33
+ )
34
+
35
+
36
+def _mla_spec(block_size: int = 4) -> MLAAttentionSpec:
37
+ return MLAAttentionSpec(
38
+ block_size=block_size,
39
+ num_kv_heads=1,
40
+ head_size=3,
41
+ dtype=torch.bfloat16,
42
+ )
43
+
44
+
45
+def _worker(
46
+ kv_caches: dict[str, torch.Tensor],
47
+ layer_to_spec: dict[str, object],
48
+ num_blocks: int = 8,
49
+) -> SimpleNamespace:
50
+ return SimpleNamespace(
51
+ kv_caches=kv_caches,
52
+ layer_to_spec=layer_to_spec,
53
+ num_blocks=num_blocks,
54
+ block_size=4,
55
+ )
56
+
57
+
58
+def _remote_meta(num_blocks: int = 16) -> SimpleNamespace:
59
+ return SimpleNamespace(num_blocks=num_blocks)
60
+
61
+
62
+def test_separated_kv_layout_uses_kv_axis_zero_and_block_axis_one():
63
+ cache = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
64
+ worker = _worker({"layer": cache}, {"layer": _full_spec()})
65
+
66
+ geometry = moriio_layout.get_layer_transfer_geometry(
67
+ "layer", cache, worker.layer_to_spec, remote_num_blocks=16
68
+ )
69
+ assert geometry.block_stride == 24
70
+ assert geometry.local_kv_stride == 192
71
+ assert geometry.remote_kv_stride == 384
72
+ assert geometry.split_kv_regions
73
+
74
+ assert moriio_layout.compute_block_transfer_offsets(
75
+ "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
76
+ ) == ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48])
77
+
78
+
79
+def test_interleaved_kv_layout_uses_block_axis_zero_and_kv_axis_one():
80
+ cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
81
+ worker = _worker({"layer": cache}, {"layer": _full_spec()})
82
+
83
+ geometry = moriio_layout.get_layer_transfer_geometry(
84
+ "layer", cache, worker.layer_to_spec, remote_num_blocks=16
85
+ )
86
+ assert geometry.block_stride == 48
87
+ assert geometry.local_kv_stride == 24
88
+ assert geometry.remote_kv_stride == 24
89
+ assert not geometry.split_kv_regions
90
+
91
+ assert moriio_layout.compute_block_transfer_offsets(
92
+ "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
93
+ ) == ([96, 288], [384, 480], [96, 96])
94
+
95
+
96
+def test_mla_key_only_layout_transfers_one_slab_per_block():
97
+ cache = torch.empty((8, 4, 3), dtype=torch.bfloat16)
98
+ worker = _worker({"layer": cache}, {"layer": _mla_spec()})
99
+
100
+ geometry = moriio_layout.get_layer_transfer_geometry(
101
+ "layer", cache, worker.layer_to_spec, remote_num_blocks=16
102
+ )
103
+ assert geometry.block_stride == 12
104
+ assert geometry.local_kv_stride is None
105
+ assert geometry.remote_kv_stride is None
106
+ assert geometry.transfers_per_block == 1
107
+
108
+ assert moriio_layout.compute_block_transfer_offsets(
109
+ "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
110
+ ) == ([24, 72], [96, 120], [24, 24])
111
+
112
+
113
+def test_mixed_layers_compute_distinct_offsets_per_layer():
114
+ kv_caches = {
115
+ "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16),
116
+ "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
117
+ "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
118
+ }
119
+ worker = _worker(
120
+ kv_caches,
121
+ {
122
+ "separated": _full_spec(),
123
+ "interleaved": _full_spec(),
124
+ "indexer": _mla_spec(),
125
+ },
126
+ )
127
+
128
+ separated = moriio_layout.compute_block_transfer_offsets(
129
+ "separated",
130
+ kv_caches["separated"],
131
+ worker.layer_to_spec,
132
+ [1, 3],
133
+ [4, 5],
134
+ _remote_meta().num_blocks,
135
+ )
136
+ interleaved = moriio_layout.compute_block_transfer_offsets(
137
+ "interleaved",
138
+ kv_caches["interleaved"],
139
+ worker.layer_to_spec,
140
+ [1, 3],
141
+ [4, 5],
142
+ _remote_meta().num_blocks,
143
+ )
144
+ indexer = moriio_layout.compute_block_transfer_offsets(
145
+ "indexer",
146
+ kv_caches["indexer"],
147
+ worker.layer_to_spec,
148
+ [1, 3],
149
+ [4, 5],
150
+ _remote_meta().num_blocks,
151
+ )
152
+
153
+ assert separated != interleaved
154
+ assert separated != indexer
155
+ assert interleaved != indexer
156
+
157
+
158
+def test_block_id_length_mismatch_raises_value_error():
159
+ cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
160
+ worker = _worker({"layer": cache}, {"layer": _full_spec()})
161
+
162
+ with pytest.raises(ValueError, match="must have the same length"):
163
+ moriio_layout.compute_block_transfer_offsets(
164
+ "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks
165
+ )
166
+
167
+
168
+def test_registration_regions_do_not_split_interleaved_or_mla_cache():
169
+ separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
170
+ interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
171
+ indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16)
172
+ worker = _worker(
173
+ {
174
+ "separated": separated,
175
+ "interleaved": interleaved,
176
+ "indexer": indexer,
177
+ },
178
+ {
179
+ "separated": _full_spec(),
180
+ "interleaved": _full_spec(),
181
+ "indexer": _mla_spec(),
182
+ },
183
+ )
184
+
185
+ separated_regions = moriio_layout.iter_layer_registration_regions(
186
+ "separated", separated, worker.layer_to_spec
187
+ )
188
+ interleaved_regions = moriio_layout.iter_layer_registration_regions(
189
+ "interleaved", interleaved, worker.layer_to_spec
190
+ )
191
+ indexer_regions = moriio_layout.iter_layer_registration_regions(
192
+ "indexer", indexer, worker.layer_to_spec
193
+ )
194
+
195
+ assert [region[0].data_ptr() for region in separated_regions] == [
196
+ separated[0].data_ptr(),
197
+ separated[1].data_ptr(),
198
+ ]
199
+ assert separated_regions[0][1] == 8 * 48
200
+ assert separated_regions[1][1] == 8 * 48
201
+
202
+ assert len(interleaved_regions) == 1
203
+ assert interleaved_regions[0][0].data_ptr() == interleaved.data_ptr()
204
+ assert interleaved_regions[0][1] == 8 * 2 * 48
205
+
206
+ assert len(indexer_regions) == 1
207
+ assert indexer_regions[0][0].data_ptr() == indexer.data_ptr()
208
+ assert indexer_regions[0][1] == 8 * 24
209
+
210
+
211
+def test_registration_regions_use_layer_num_blocks():
212
+ cache = torch.empty((4, 2, 4, 2, 3), dtype=torch.bfloat16)
213
+ worker = _worker({"layer": cache}, {"layer": _full_spec()}, num_blocks=8)
214
+
215
+ regions = moriio_layout.iter_layer_registration_regions(
216
+ "layer", cache, worker.layer_to_spec
217
+ )
218
+
219
+ assert len(regions) == 1
220
+ assert regions[0][1] == 4 * 2 * 48
221
+
222
+
223
+def test_unsupported_shape_raises_value_error():
224
+ cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16)
225
+ worker = _worker({"layer": cache}, {"layer": _full_spec()})
226
+
227
+ with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
228
+ moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+92
-68
@@ -47,6 +47,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
47
MoRIIOWrapper,
48
MoRIIOWriter,
49
)
50
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
51
+ LayerTransferGeometry,
52
+ build_layer_to_spec,
53
+ compute_block_transfer_offsets,
54
+ get_layer_transfer_geometry,
55
+ is_mla_cache_layer,
56
+ iter_layer_registration_regions,
57
+)
58
from vllm.distributed.parallel_state import (
59
get_tensor_model_parallel_world_size,
60
get_tp_group,
@@ -71,6 +79,7 @@ if TYPE_CHECKING:
79
80
logger = init_logger(__name__)
81
82
+
83
try:
84
from mori.io import (
85
BackendType,
@@ -117,7 +126,9 @@ class MoRIIOConnector(KVConnectorBase_V1):
126
self.connector_worker: MoRIIOConnectorWorker | None = None
127
elif role == KVConnectorRole.WORKER:
128
self.connector_scheduler = None
120
- self.connector_worker = MoRIIOConnectorWorker(vllm_config, self.engine_id)
129
+ self.connector_worker = MoRIIOConnectorWorker(
130
+ vllm_config, self.engine_id, kv_cache_config
131
+ )
132
logger.info(
133
"Initialized MoRIIO Connector,engine_id:%s,role: %s",
134
self.engine_id,
@@ -683,7 +694,12 @@ class MoRIIOConnectorScheduler:
694
class MoRIIOConnectorWorker:
695
"""Implementation of Worker side methods"""
696
686
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
697
+ def __init__(
698
+ self,
699
+ vllm_config: VllmConfig,
700
+ engine_id: str,
701
+ kv_cache_config: "KVCacheConfig",
702
+ ):
703
if not is_moriio_available():
704
raise RuntimeError(
705
"MoRIIO is not available. Please ensure the 'mori' package "
@@ -707,6 +723,7 @@ class MoRIIOConnectorWorker:
723
)
724
self.kv_transfer_config = vllm_config.kv_transfer_config
725
self.is_producer = self.kv_transfer_config.is_kv_producer
726
+ self.layer_to_spec = build_layer_to_spec(kv_cache_config)
727
728
if self.is_producer:
729
set_role(ROLE.PRODUCER)
@@ -809,6 +826,8 @@ class MoRIIOConnectorWorker:
826
self.kv_cache_shape = None
827
self.block_shape = None
828
self.kv_element_size = 0
829
+ self.kv_cache_shapes: dict[str, torch.Size] = {}
830
+ self.block_lens: dict[str, int] = {}
831
832
# Map of engine_id -> {agent_name0, agent_name1..}.
833
self._remote_agents: dict[EngineId, set[str]] = {}
@@ -1218,51 +1237,86 @@ class MoRIIOConnectorWorker:
1237
all_done_future = self._handshake_initiation_executor.submit(wait_all_dp)
1238
all_done_future.add_done_callback(request_ready)
1239
1240
+ def _is_mla_cache_layer(self, layer_name: str) -> bool:
1241
+ return is_mla_cache_layer(self.layer_to_spec, layer_name)
1242
+
1243
+ def _get_layer_transfer_geometry(
1244
+ self, layer_name: str, remote_num_blocks: int | None = None
1245
+ ) -> LayerTransferGeometry:
1246
+ return get_layer_transfer_geometry(
1247
+ layer_name,
1248
+ self.kv_caches[layer_name],
1249
+ self.layer_to_spec,
1250
+ remote_num_blocks,
1251
+ )
1252
+
1253
+ def _iter_layer_registration_regions(
1254
+ self, layer_name: str
1255
+ ) -> list[tuple[torch.Tensor, int]]:
1256
+ return iter_layer_registration_regions(
1257
+ layer_name,
1258
+ self.kv_caches[layer_name],
1259
+ self.layer_to_spec,
1260
+ )
1261
+
1262
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
1263
"""Register the KV Cache data in moriio."""
1264
1224
- _, first_kv_cache = next(iter(kv_caches.items()))
1265
+ self.kv_caches = kv_caches # layer name to kv cache
1266
+ self.kv_cache_shapes = {
1267
+ layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items()
1268
+ }
1269
+
1270
+ first_layer_name, first_kv_cache = next(
1271
+ (
1272
+ (layer_name, kv_cache)
1273
+ for layer_name, kv_cache in kv_caches.items()
1274
+ if (
1275
+ not self._is_mla_cache_layer(layer_name)
1276
+ and len(kv_cache.shape) == 5
1277
+ and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2)
1278
+ )
1279
+ ),
1280
+ next(iter(kv_caches.items())),
1281
+ )
1282
kv_elem_size = first_kv_cache.element_size()
1283
1227
- use_mla = len(first_kv_cache.shape) == 3
1228
- assert use_mla == self.use_mla
1284
+ use_mla = self._is_mla_cache_layer(first_layer_name)
1285
+ first_geometry = self._get_layer_transfer_geometry(first_layer_name)
1286
1287
if use_mla:
1288
# MLA case.
1232
- self.num_blocks = first_kv_cache.shape[0]
1289
block_rank = 2 # [block_size, latent_dim]
1290
block_shape = first_kv_cache.shape[-block_rank:]
1235
- block_size, kv_latent_dim = block_shape
1236
- self.slot_size_bytes = kv_elem_size * kv_latent_dim
1291
else:
1238
- # [2 (k and v), num_blocks, ...]
1239
- self.num_blocks = first_kv_cache.shape[1]
1292
+ # [2, num_blocks, ...] or [num_blocks, 2, ...]
1293
block_rank = 3 # [block_size, kv_heads, head_dim]
1294
block_shape = first_kv_cache.shape[-block_rank:]
1242
- block_size, n_kv_heads, head_dim = block_shape[-3:]
1243
- # head size in bytes.
1244
- self.slot_size_bytes = (
1245
- kv_elem_size * n_kv_heads * head_dim
1246
- ) # 1 token 1 layer size , slot size
1247
- assert block_size == self.block_size
1295
+ self.num_blocks = first_geometry.num_blocks
1296
+ self.slot_size_bytes = first_geometry.slot_size_bytes
1297
+ assert first_geometry.block_size == self.block_size
1298
# TODO(tms): self.block_len needs to be per-layer for sliding window,
1299
# hybrid attn, etc
1300
# block size in bytes
1251
- self.block_len = kv_elem_size * math.prod(block_shape)
1301
+ self.block_len = first_geometry.block_len
1302
self.kv_cache_shape = first_kv_cache.shape
1303
self.block_shape = block_shape
1304
self.kv_element_size = kv_elem_size
1305
1306
self.dst_num_blocks[self.engine_id] = self.num_blocks
1257
- self.kv_caches = kv_caches # layer name to kv cache
1307
kv_caches_base_addr = []
1308
caches_data = []
1309
1261
- for cache_or_caches in kv_caches.values():
1262
- cache_list = [cache_or_caches] if use_mla else cache_or_caches
1263
- for cache in cache_list:
1310
+ for layer_name in kv_caches:
1311
+ geometry = self._get_layer_transfer_geometry(layer_name)
1312
+ if geometry.block_size != self.block_size:
1313
+ raise ValueError(
1314
+ "MoRIIO KV cache block size mismatch for layer "
1315
+ f"{layer_name}: {geometry.block_size} != {self.block_size}"
1316
+ )
1317
+ self.block_lens[layer_name] = geometry.block_len
1318
+ for cache, region_len in self._iter_layer_registration_regions(layer_name):
1319
base_addr = cache.data_ptr()
1265
- region_len = self.num_blocks * self.block_len
1320
caches_data.append((base_addr, region_len, cache.device.index, ""))
1321
kv_caches_base_addr.append(base_addr)
1322
@@ -1275,7 +1329,9 @@ class MoRIIOConnectorWorker:
1329
moriio_mem_metadata
1330
)
1331
1278
- self.local_kv_cache_size.append(cache.nelement() * cache.element_size())
1332
+ self.local_kv_cache_size.append(
1333
+ kv_cache.nelement() * kv_cache.element_size()
1334
+ )
1335
1336
self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
1337
self.num_regions = len(caches_data)
@@ -1666,47 +1722,17 @@ class MoRIIOConnectorWorker:
1722
Returns:
1723
Tuple of (local_offsets, remote_offsets, transfer_sizes)
1724
"""
1669
- assert self.kv_cache_shape is not None, "KV caches shape not initialized"
1670
- is_mla = len(self.kv_cache_shape) == 3
1671
- stride = self.kv_caches[layer_name].stride()
1672
- sz = self.kv_caches[layer_name].element_size()
1673
- if is_mla:
1674
- blknum, blksize, hs = self.kv_cache_shape
1675
- hn = 1
1676
- block_stride = stride[0]
1677
- else:
1678
- _, blknum, blksize, hn, hs = self.kv_cache_shape
1679
- local_ktov_stride = stride[0]
1680
- block_stride = stride[1]
1681
- remote_ktov_stride = block_stride * remote_moriio_meta.num_blocks
1682
-
1683
- transfer_size_byte = blksize * hn * hs * sz
1684
- per_block = 1 if is_mla else 2
1685
- total = len(local_block_ids) * per_block
1686
- offset_local = [0] * total
1687
- offset_remote = [0] * total
1688
- sizes = [transfer_size_byte] * total
1689
-
1690
- w = 0
1691
- for i, lb in enumerate(local_block_ids):
1692
- rb = remote_block_ids[i]
1693
- # K
1694
- offset_local[w] = sz * (lb * block_stride)
1695
- offset_remote[w] = sz * (rb * block_stride)
1696
- w += 1
1697
- if not is_mla:
1698
- # V
1699
- # Handle num_block variations originating from PD (different kv strides)
1700
- # TODO: address block_sz differences in heterogeneous TP scenarios
1701
- # In MLA, we don't need to consider these two cases.
1702
- offset_local[w] = sz * (1 * local_ktov_stride + lb * block_stride)
1703
- offset_remote[w] = sz * (1 * remote_ktov_stride + rb * block_stride)
1704
- w += 1
1705
-
1706
- merged_l, merged_r, merged_s = self.merge_contiguous_blocks(
1707
- offset_local, offset_remote, sizes, assume_sorted=False
1725
+ return compute_block_transfer_offsets(
1726
+ layer_name=layer_name,
1727
+ kv_cache=self.kv_caches[layer_name],
1728
+ layer_to_spec=self.layer_to_spec,
1729
+ local_block_ids=local_block_ids,
1730
+ remote_block_ids=remote_block_ids,
1731
+ remote_num_blocks=remote_moriio_meta.num_blocks,
1732
+ merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks(
1733
+ local, remote, sizes, assume_sorted=False
1734
+ ),
1735
)
1709
- return merged_l, merged_r, merged_s
1736
1737
def _read_blocks(
1738
self,
@@ -1724,15 +1750,13 @@ class MoRIIOConnectorWorker:
1750
dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
1751
sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)
1752
1727
- first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0]
1728
- offs = self._compute_block_transfer_offsets(
1729
- first_layer, local_block_ids, remote_block_ids, remote_moriio_meta
1730
- )
1731
-
1753
for layer_name in self.layer_name_to_local_kv_cache_metadata:
1754
sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
1755
layer_name
1756
)
1757
+ offs = self._compute_block_transfer_offsets(
1758
+ layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
1759
+ )
1760
# TODO : apply multi-session batch-read when moriio support it
1761
transfer_status = self.moriio_wrapper.read_remote_data(
1762
offs[2], offs[0], offs[1], sessions[sess_idx]
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
new
+213
@@ -0,0 +1,213 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable, Mapping
5
+from typing import NamedTuple
6
+
7
+import torch
8
+
9
+from vllm.v1.kv_cache_interface import (
10
+ KVCacheConfig,
11
+ KVCacheSpec,
12
+ MLAAttentionSpec,
13
+ SlidingWindowMLASpec,
14
+ UniformTypeKVCacheSpecs,
15
+)
16
+
17
+
18
+class LayerTransferGeometry(NamedTuple):
19
+ num_blocks: int
20
+ block_size: int
21
+ block_len: int
22
+ slot_size_bytes: int
23
+ block_stride: int
24
+ local_kv_stride: int | None
25
+ remote_kv_stride: int | None
26
+ transfers_per_block: int
27
+ regions_per_block: int
28
+ split_kv_regions: bool
29
+
30
+
31
+def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
32
+ layer_to_spec: dict[str, KVCacheSpec] = {}
33
+ for group in kv_cache_config.kv_cache_groups:
34
+ group_spec = group.kv_cache_spec
35
+ if isinstance(group_spec, UniformTypeKVCacheSpecs):
36
+ layer_to_spec.update(
37
+ {
38
+ layer_name: group_spec.kv_cache_specs[layer_name]
39
+ for layer_name in group.layer_names
40
+ }
41
+ )
42
+ else:
43
+ layer_to_spec.update(
44
+ {layer_name: group_spec for layer_name in group.layer_names}
45
+ )
46
+ return layer_to_spec
47
+
48
+
49
+def is_mla_cache_layer(
50
+ layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str
51
+) -> bool:
52
+ try:
53
+ spec = layer_to_spec[layer_name]
54
+ except KeyError as e:
55
+ raise ValueError(f"Missing KV cache spec for layer {layer_name}") from e
56
+ return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
57
+
58
+
59
+def get_layer_transfer_geometry(
60
+ layer_name: str,
61
+ kv_cache: torch.Tensor,
62
+ layer_to_spec: Mapping[str, KVCacheSpec],
63
+ remote_num_blocks: int | None = None,
64
+) -> LayerTransferGeometry:
65
+ shape = kv_cache.shape
66
+ stride = kv_cache.stride()
67
+ element_size = kv_cache.element_size()
68
+ is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
69
+
70
+ if is_mla_cache and len(shape) == 3:
71
+ num_blocks, block_size, latent_dim = shape
72
+ slot_size_bytes = latent_dim * element_size
73
+ block_len = block_size * slot_size_bytes
74
+ return LayerTransferGeometry(
75
+ num_blocks=num_blocks,
76
+ block_size=block_size,
77
+ block_len=block_len,
78
+ slot_size_bytes=slot_size_bytes,
79
+ block_stride=stride[0],
80
+ local_kv_stride=None,
81
+ remote_kv_stride=None,
82
+ transfers_per_block=1,
83
+ regions_per_block=1,
84
+ split_kv_regions=False,
85
+ )
86
+
87
+ if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
88
+ _, num_blocks, block_size, num_kv_heads, head_dim = shape
89
+ slot_size_bytes = num_kv_heads * head_dim * element_size
90
+ block_len = block_size * slot_size_bytes
91
+ remote_kv_stride = stride[1] * (remote_num_blocks or num_blocks)
92
+ return LayerTransferGeometry(
93
+ num_blocks=num_blocks,
94
+ block_size=block_size,
95
+ block_len=block_len,
96
+ slot_size_bytes=slot_size_bytes,
97
+ block_stride=stride[1],
98
+ local_kv_stride=stride[0],
99
+ remote_kv_stride=remote_kv_stride,
100
+ transfers_per_block=2,
101
+ regions_per_block=1,
102
+ split_kv_regions=True,
103
+ )
104
+
105
+ if not is_mla_cache and len(shape) == 5 and shape[1] == 2:
106
+ num_blocks, _, block_size, num_kv_heads, head_dim = shape
107
+ slot_size_bytes = num_kv_heads * head_dim * element_size
108
+ block_len = block_size * slot_size_bytes
109
+ return LayerTransferGeometry(
110
+ num_blocks=num_blocks,
111
+ block_size=block_size,
112
+ block_len=block_len,
113
+ slot_size_bytes=slot_size_bytes,
114
+ block_stride=stride[0],
115
+ local_kv_stride=stride[1],
116
+ remote_kv_stride=stride[1],
117
+ transfers_per_block=2,
118
+ regions_per_block=2,
119
+ split_kv_regions=False,
120
+ )
121
+
122
+ cache_kind = "MLA" if is_mla_cache else "K/V"
123
+ raise ValueError(
124
+ f"Unsupported MoRIIO {cache_kind} cache shape for layer "
125
+ f"{layer_name}: {tuple(shape)}"
126
+ )
127
+
128
+
129
+def iter_layer_registration_regions(
130
+ layer_name: str,
131
+ kv_cache: torch.Tensor,
132
+ layer_to_spec: Mapping[str, KVCacheSpec],
133
+) -> list[tuple[torch.Tensor, int]]:
134
+ geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
135
+ region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
136
+ if geometry.split_kv_regions:
137
+ return [(cache, region_len) for cache in kv_cache]
138
+ return [(kv_cache, region_len)]
139
+
140
+
141
+def merge_contiguous_offsets(
142
+ offsets_local: list[int],
143
+ offsets_remote: list[int],
144
+ sizes: list[int],
145
+) -> tuple[list[int], list[int], list[int]]:
146
+ if not offsets_local:
147
+ return [], [], []
148
+ if not (len(offsets_local) == len(offsets_remote) == len(sizes)):
149
+ raise ValueError("Input list lengths mismatch")
150
+
151
+ rows = sorted(zip(offsets_local, offsets_remote, sizes), key=lambda row: row[0])
152
+ merged: list[list[int]] = []
153
+ for local, remote, size in rows:
154
+ if (
155
+ merged
156
+ and local == merged[-1][0] + merged[-1][2]
157
+ and remote == merged[-1][1] + merged[-1][2]
158
+ ):
159
+ merged[-1][2] += size
160
+ else:
161
+ merged.append([local, remote, size])
162
+
163
+ return (
164
+ [row[0] for row in merged],
165
+ [row[1] for row in merged],
166
+ [row[2] for row in merged],
167
+ )
168
+
169
+
170
+def compute_block_transfer_offsets(
171
+ layer_name: str,
172
+ kv_cache: torch.Tensor,
173
+ layer_to_spec: Mapping[str, KVCacheSpec],
174
+ local_block_ids: list[int],
175
+ remote_block_ids: list[int],
176
+ remote_num_blocks: int,
177
+ merge_fn: Callable[
178
+ [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]
179
+ ] = merge_contiguous_offsets,
180
+) -> tuple[list[int], list[int], list[int]]:
181
+ if len(local_block_ids) != len(remote_block_ids):
182
+ raise ValueError(
183
+ "local_block_ids and remote_block_ids must have the same length: "
184
+ f"{len(local_block_ids)} != {len(remote_block_ids)}"
185
+ )
186
+ geometry = get_layer_transfer_geometry(
187
+ layer_name, kv_cache, layer_to_spec, remote_num_blocks
188
+ )
189
+ element_size = kv_cache.element_size()
190
+ transfer_size_byte = geometry.block_len
191
+ per_block = geometry.transfers_per_block
192
+ total = len(local_block_ids) * per_block
193
+ offset_local = [0] * total
194
+ offset_remote = [0] * total
195
+ sizes = [transfer_size_byte] * total
196
+
197
+ w = 0
198
+ for lb, rb in zip(local_block_ids, remote_block_ids):
199
+ offset_local[w] = element_size * (lb * geometry.block_stride)
200
+ offset_remote[w] = element_size * (rb * geometry.block_stride)
201
+ w += 1
202
+ if per_block == 2:
203
+ assert geometry.local_kv_stride is not None
204
+ assert geometry.remote_kv_stride is not None
205
+ offset_local[w] = element_size * (
206
+ geometry.local_kv_stride + lb * geometry.block_stride
207
+ )
208
+ offset_remote[w] = element_size * (
209
+ geometry.remote_kv_stride + rb * geometry.block_stride
210
+ )
211
+ w += 1
212
+
213
+ return merge_fn(offset_local, offset_remote, sizes)