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

[KVConnector][MoRIIO] Transfer hybrid mamba/KDA recurrent state in READ mode (#51052)

Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

YukioZzz committed Sep 21, 2026 at 05:30 UTC 96791737887224c8bfec0a791517759f62e4d93b
11 files changed +2746 -184
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+2
@@ -367,6 +367,8 @@ async def handle_request(api: str, request: Request):
367 "remote_block_ids"
368 ]
369 req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
370 + if prefill_kv.get("tp_size") is not None:
371 + req_data["kv_transfer_params"]["tp_size"] = prefill_kv["tp_size"]
372
373 req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
374 "dp_size"
tests/v1/kv_connector/unit/test_moriio_connector.py
+11
@@ -775,6 +775,17 @@ def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
775 assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
776
777
778 +def test_transfer_disabled_hybrid_group_is_ignored():
779 + config = _make_hybrid_kv_cache_config()
780 + config.kv_cache_groups[1].enable_kv_transfer = False
781 + scheduler = _read_scheduler(config)
782 +
783 + assert scheduler._is_hma_required is False
784 + assert scheduler.get_exchange_clipped_blocks([[1, 2, 3], [20, 21, 22]]) == [
785 + [1, 2, 3]
786 + ]
787 +
788 +
789 def test_non_sliding_window_hybrid_is_rejected():
790 """A hybrid group that is not sliding-window (e.g. chunked-local
791 attention) must fail closed rather than be silently mistransferred."""
tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py new
+809
@@ -0,0 +1,809 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Pure-CPU unit tests for the MoRIIO hybrid (attention + mamba/KDA) scheduler
4 +and the KDA offset-template cache.
5 +
6 +These exercise the highest-risk, otherwise-untested scheduler logic without a
7 +GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE
8 +``N-1`` token accounting, P-side prompt truncation, and offset-template cache
9 +wiring.
10 +
11 +Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is
12 +running on ROCm with ``mori`` installed (importing the connector pulls in
13 +``mori``). The authoritative run happens on the MIA recipe image.
14 +"""
15 +
16 +import importlib
17 +import importlib.util
18 +from types import SimpleNamespace
19 +
20 +import pytest
21 +import torch
22 +
23 +from vllm.platforms import current_platform
24 +
25 +mori_available = importlib.util.find_spec("mori") is not None
26 +
27 +if not (current_platform.is_rocm() and mori_available):
28 + pytest.skip(
29 + "MoRIIOs are only available on ROCm with mori package installed",
30 + allow_module_level=True,
31 + )
32 +
33 +moriio_connector = importlib.import_module(
34 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector"
35 +)
36 +moriio_layout = importlib.import_module(
37 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
38 +)
39 +ssm_conv_transfer_utils = importlib.import_module(
40 + "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils"
41 +)
42 +MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo
43 +MoRIIOMode = moriio_connector.MoRIIOMode
44 +moriio_common = importlib.import_module(
45 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
46 +)
47 +
48 +
49 +class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore[name-defined]
50 + """Constructs a scheduler without the heavy real ``__init__``; only the
51 + attributes read by the method under test are set."""
52 +
53 + def __init__(self, **attrs):
54 + self._mamba_group_ids: list[int] = []
55 + self._attn_group_ids: list[int] = [0]
56 + self._ssm_state_slots_are_positional = False
57 + self._is_hma_required = False
58 + self.kv_cache_config = SimpleNamespace(
59 + kv_cache_groups=[None, None],
60 + select_transfer_block_ids=lambda block_ids: tuple(block_ids),
61 + )
62 + for k, v in attrs.items():
63 + setattr(self, k, v)
64 +
65 +
66 +class _FakeConnector(moriio_connector.MoRIIOConnector): # type: ignore[name-defined]
67 + def __init__(self, connector_scheduler):
68 + self.connector_scheduler = connector_scheduler
69 +
70 +
71 +class _FakeWorker(moriio_connector.MoRIIOConnectorWorker): # type: ignore[name-defined]
72 + def __init__(self, **attrs):
73 + for k, v in attrs.items():
74 + setattr(self, k, v)
75 +
76 +
77 +class _FakeBlocks:
78 + def __init__(self, all_groups):
79 + self._all_groups = all_groups
80 +
81 + def get_block_ids(self):
82 + return self._all_groups
83 +
84 +
85 +def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16):
86 + inner = 1
87 + for d in per_slot_shape:
88 + inner *= d
89 + assert slot_stride >= inner
90 + backing = torch.zeros(num_slots * slot_stride, dtype=dtype)
91 + inner_strides: list[int] = []
92 + acc = 1
93 + for d in reversed(per_slot_shape):
94 + inner_strides.insert(0, acc)
95 + acc *= d
96 + return backing.as_strided(
97 + (num_slots, *per_slot_shape), (slot_stride, *inner_strides)
98 + )
99 +
100 +
101 +def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2):
102 + conv_dim = 2 * key_dim + value_dim
103 + conv_state_bytes = conv_dim * conv_rows * dtype_size
104 + ssm_state_bytes = 64
105 + return MambaConvSplitInfo(
106 + conv_rows=conv_rows,
107 + local_proj_dims=(key_dim, key_dim, value_dim),
108 + conv_dtype_size=dtype_size,
109 + ssm_sizes=(conv_state_bytes, ssm_state_bytes),
110 + )
111 +
112 +
113 +def _mamba_spec(num_states: int = 2):
114 + shapes = tuple((1, 1) for _ in range(num_states))
115 + return moriio_connector.MambaSpec(
116 + block_size=16,
117 + shapes=shapes,
118 + dtypes=(torch.float32,) * num_states,
119 + mamba_cache_mode="all",
120 + )
121 +
122 +
123 +def _gate_vllm_config(*, read_mode: bool = True, speculative_config=None):
124 + return SimpleNamespace(
125 + kv_transfer_config=SimpleNamespace(
126 + kv_connector_extra_config={"read_mode": read_mode}
127 + ),
128 + speculative_config=speculative_config,
129 + )
130 +
131 +
132 +# --------------------------------------------------------------------------
133 +# split_block_groups
134 +# --------------------------------------------------------------------------
135 +def test_split_block_groups_separates_attention_and_mamba():
136 + sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1])
137 + block_ids = ([1, 2], [99])
138 + attn, mamba = sched.split_block_groups(block_ids)
139 + assert attn == [1, 2]
140 + assert mamba == [[99]]
141 +
142 +
143 +def test_split_block_groups_ignores_transfer_disabled_group():
144 + mamba_spec = moriio_connector.MambaSpec(
145 + block_size=16,
146 + shapes=((1, 1),),
147 + dtypes=(torch.float32,),
148 + mamba_cache_mode="all",
149 + )
150 + config = SimpleNamespace(
151 + kv_cache_groups=[
152 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
153 + SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object()),
154 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=mamba_spec),
155 + ],
156 + )
157 + config.transfer_groups = tuple(
158 + group for group in config.kv_cache_groups if group.enable_kv_transfer
159 + )
160 + config.select_transfer_block_ids = lambda block_ids: tuple(
161 + block_ids[group_id]
162 + for group_id, group in enumerate(config.kv_cache_groups)
163 + if group.enable_kv_transfer
164 + )
165 + attn_groups, mamba_groups = moriio_connector._split_kv_cache_group_kinds(config)
166 + assert attn_groups == [0]
167 + assert mamba_groups == [1]
168 +
169 + sched = _FakeScheduler(
170 + _has_mamba=True,
171 + _attn_group_ids=attn_groups,
172 + _mamba_group_ids=mamba_groups,
173 + kv_cache_config=config,
174 + )
175 + assert sched.split_block_groups(([1, 2], [70, 71], [99])) == (
176 + [1, 2],
177 + [[99]],
178 + )
179 +
180 +
181 +def test_exchange_blocks_ignore_transfer_disabled_group():
182 + config = SimpleNamespace(
183 + kv_cache_groups=[None, None, None],
184 + select_transfer_block_ids=lambda block_ids: (block_ids[0], block_ids[2]),
185 + )
186 + sched = _FakeScheduler(
187 + kv_cache_config=config,
188 + _is_hma_required=True,
189 + blocks_per_sw=[0, 2],
190 + )
191 +
192 + assert sched.get_exchange_clipped_blocks(([1, 2], [70, 71], [90, 91, 92])) == [
193 + [1, 2],
194 + [91, 92],
195 + ]
196 +
197 +
198 +def test_scheduler_rejects_multiple_attention_groups_with_mamba():
199 + config = SimpleNamespace(
200 + kv_cache_groups=[
201 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
202 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
203 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
204 + ],
205 + )
206 + config.transfer_groups = tuple(config.kv_cache_groups)
207 +
208 + with pytest.raises(moriio_common.MoRIIOError, match="exactly one transferable"):
209 + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
210 +
211 +
212 +def test_split_block_groups_preserves_multiple_mamba_groups():
213 + sched = _FakeScheduler(
214 + _has_mamba=True,
215 + _attn_group_ids=[0],
216 + _mamba_group_ids=[1, 2],
217 + )
218 +
219 + assert sched.split_block_groups(([1, 2], [40], [70])) == (
220 + [1, 2],
221 + [[40], [70]],
222 + )
223 +
224 +
225 +def test_scheduler_rejects_incompatible_mamba_groups():
226 + incompatible = _mamba_spec()
227 + object.__setattr__(incompatible, "page_size_padded", 4096)
228 + config = SimpleNamespace(
229 + kv_cache_groups=[
230 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
231 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
232 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=incompatible),
233 + ]
234 + )
235 + config.transfer_groups = tuple(config.kv_cache_groups)
236 +
237 + with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"):
238 + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
239 +
240 +
241 +def test_scheduler_rejects_mamba_group_without_two_states():
242 + config = SimpleNamespace(
243 + kv_cache_groups=[
244 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
245 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec(1)),
246 + ],
247 + )
248 + config.transfer_groups = tuple(config.kv_cache_groups)
249 +
250 + with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"):
251 + moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)
252 +
253 +
254 +def test_scheduler_rejects_speculative_hybrid_read():
255 + config = SimpleNamespace(
256 + kv_cache_groups=[
257 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
258 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
259 + ],
260 + )
261 + config.transfer_groups = tuple(config.kv_cache_groups)
262 +
263 + with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"):
264 + moriio_connector.MoRIIOConnectorScheduler(
265 + _gate_vllm_config(speculative_config=object()), "engine", config
266 + )
267 +
268 +
269 +def test_scheduler_rejects_hybrid_write():
270 + config = SimpleNamespace(
271 + kv_cache_groups=[
272 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),
273 + SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()),
274 + ],
275 + )
276 + config.transfer_groups = tuple(config.kv_cache_groups)
277 +
278 + with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
279 + moriio_connector.MoRIIOConnectorScheduler(
280 + _gate_vllm_config(read_mode=False), "engine", config
281 + )
282 +
283 +
284 +def test_split_block_groups_rejects_non_hybrid_use():
285 + sched = _FakeScheduler(_has_mamba=False)
286 +
287 + with pytest.raises(moriio_common.MoRIIOError, match="non-hybrid"):
288 + sched.split_block_groups(([5, 6, 7],))
289 +
290 +
291 +def test_split_block_groups_accepts_empty_abort_payload():
292 + sched = _FakeScheduler(_has_mamba=True)
293 +
294 + assert sched.split_block_groups(()) == ([], [])
295 +
296 +
297 +def test_split_block_groups_keeps_positional_slots_in_all_mode():
298 + # mamba_cache_mode="all" keeps a state per block position.
299 + sched = _FakeScheduler(
300 + _has_mamba=True,
301 + _attn_group_ids=[0],
302 + _mamba_group_ids=[1],
303 + _ssm_state_slots_are_positional=True,
304 + )
305 + attn, mamba = sched.split_block_groups(([1], [40, 41, 42]))
306 + assert attn == [1]
307 + assert mamba == [[40, 41, 42]]
308 +
309 +
310 +def test_split_block_groups_keeps_only_running_state_outside_all_mode():
311 + sched = _FakeScheduler(
312 + _has_mamba=True,
313 + _attn_group_ids=[0],
314 + _mamba_group_ids=[1],
315 + )
316 + assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]])
317 +
318 +
319 +# --------------------------------------------------------------------------
320 +# request_finished_all_groups
321 +# --------------------------------------------------------------------------
322 +def test_request_finished_all_groups_carries_attn_and_mamba_in_one_field():
323 + seen = {}
324 +
325 + def _fake_request_finished(request, attn_block_ids, mamba_block_groups):
326 + seen["attn"] = list(attn_block_ids)
327 + seen["mamba"] = [list(group) for group in mamba_block_groups]
328 + return True, {
329 + "do_remote_prefill": True,
330 + "remote_block_ids": [attn_block_ids, *mamba_block_groups],
331 + }
332 +
333 + sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1])
334 + sched.request_finished = _fake_request_finished
335 + conn = _FakeConnector(sched)
336 +
337 + request = SimpleNamespace(request_id="r0")
338 + block_ids = ([10, 11], [42])
339 + delay_free, params = conn.request_finished_all_groups(request, block_ids)
340 +
341 + assert delay_free is True
342 + # Attention ids drive request_finished; the mamba slot rides the SAME
343 + # remote_block_ids channel as [attn, *mamba_groups], so the proxy/router
344 + # needs no KDA-specific field.
345 + assert seen["attn"] == [10, 11]
346 + assert seen["mamba"] == [[42]]
347 + assert params["remote_block_ids"] == [[10, 11], [42]]
348 + assert "remote_mamba_block_ids" not in params
349 +
350 +
351 +def test_request_finished_all_groups_pure_attention_stays_grouped():
352 + def _fake_request_finished(request, block_ids):
353 + return True, {"remote_block_ids": block_ids}
354 +
355 + sched = _FakeScheduler(_has_mamba=False, _mamba_group_ids=[])
356 + sched.request_finished = _fake_request_finished
357 + conn = _FakeConnector(sched)
358 +
359 + delay_free, params = conn.request_finished_all_groups(
360 + SimpleNamespace(request_id="r1"), ([7, 8],)
361 + )
362 + # Main's HMA contract keeps pure-attention block ids grouped.
363 + assert params["remote_block_ids"] == ([7, 8],)
364 + assert "remote_mamba_block_ids" not in params
365 +
366 +
367 +def test_connector_supports_divergent_local_hybrid_hits():
368 + connector = _FakeConnector(None)
369 + connector._transfers_mamba_state = True
370 + connector.mode = MoRIIOMode.READ
371 + assert connector.supports_divergent_local_hybrid_hits is True
372 + connector.mode = MoRIIOMode.WRITE
373 + assert connector.supports_divergent_local_hybrid_hits is False
374 +
375 + connector.mode = MoRIIOMode.READ
376 + connector._transfers_mamba_state = False
377 + assert connector.supports_divergent_local_hybrid_hits is False
378 +
379 +
380 +def _make_read_scheduler():
381 + return _FakeScheduler(
382 + is_producer=False,
383 + mode=MoRIIOMode.READ,
384 + _has_mamba=True,
385 + _attn_group_ids=[0],
386 + _mamba_group_ids=[1],
387 + _ssm_state_slots_are_positional=True,
388 + request_id_to_transfer_id={},
389 + transfer_id_to_request_id={},
390 + _reqs_need_recv={},
391 + _req_kv_params={},
392 + _max_decode_tail_blocks=1,
393 + )
394 +
395 +
396 +def _make_read_request(remote_block_ids):
397 + return SimpleNamespace(
398 + request_id="req",
399 + kv_transfer_params={
400 + "do_remote_prefill": True,
401 + "transfer_id": "tx",
402 + "remote_engine_id": "prefill",
403 + "remote_block_ids": remote_block_ids,
404 + },
405 + )
406 +
407 +
408 +def test_update_state_drops_decode_recompute_tail_block():
409 + sched = _make_read_scheduler()
410 + request = _make_read_request([[10, 11], [90, 91]])
411 + blocks = _FakeBlocks(
412 + all_groups=([100, 101, 102], [200, 201, 202]),
413 + )
414 +
415 + sched.update_state_after_alloc(request, blocks, num_external_tokens=256)
416 +
417 + assert sched._reqs_need_recv["req"][1] == [[100, 101], [200, 201]]
418 + assert sched._req_kv_params["req"]["remote_block_ids"] == [
419 + [10, 11],
420 + [90, 91],
421 + ]
422 +
423 +
424 +def test_update_state_pairs_shorter_local_blocks_with_remote_suffix():
425 + sched = _make_read_scheduler()
426 + request = _make_read_request([[10, 11, 12], [92]])
427 + blocks = _FakeBlocks(
428 + all_groups=([102], [202]),
429 + )
430 +
431 + sched.update_state_after_alloc(request, blocks, num_external_tokens=64)
432 +
433 + assert sched._reqs_need_recv["req"][1] == [[102], [202]]
434 + assert sched._req_kv_params["req"]["remote_block_ids"] == [[12], [92]]
435 +
436 +
437 +def test_update_state_full_attention_hit_still_carries_mamba_state():
438 + sched = _make_read_scheduler()
439 + request = _make_read_request([[10], [90]])
440 +
441 + sched.update_state_after_alloc(
442 + request,
443 + _FakeBlocks(all_groups=([100], [200])),
444 + num_external_tokens=0,
445 + )
446 +
447 + assert sched._reqs_need_recv["req"][1] == [[], [200]]
448 +
449 +
450 +def test_update_state_preserves_each_mamba_group():
451 + sched = _make_read_scheduler()
452 + sched._mamba_group_ids = [1, 2]
453 + request = _make_read_request([[10], [90], [80]])
454 +
455 + sched.update_state_after_alloc(
456 + request,
457 + _FakeBlocks(all_groups=([100], [200], [300])),
458 + num_external_tokens=64,
459 + )
460 +
461 + assert sched._reqs_need_recv["req"][1] == [[100], [200], [300]]
462 + assert sched._req_kv_params["req"]["remote_block_ids"] == [
463 + [10],
464 + [90],
465 + [80],
466 + ]
467 +
468 +
469 +def test_align_read_blocks_allows_configured_lookahead_tail():
470 + assert moriio_connector.MoRIIOConnectorScheduler._align_read_blocks(
471 + [100, 101, 102], [10], max_decode_tail_blocks=2
472 + ) == ([100], [10])
473 + with pytest.raises(ValueError, match="allowed local tail"):
474 + moriio_connector.MoRIIOConnectorScheduler._align_read_blocks(
475 + [100, 101, 102], [10], max_decode_tail_blocks=1
476 + )
477 +
478 +
479 +def test_update_state_aligns_attention_only_decode_tail():
480 + sched = _make_read_scheduler()
481 + sched._has_mamba = False
482 + sched._mamba_group_ids = []
483 + request = _make_read_request([[10, 11]])
484 +
485 + sched.update_state_after_alloc(
486 + request,
487 + _FakeBlocks(all_groups=([100, 101, 102],)),
488 + num_external_tokens=256,
489 + )
490 +
491 + assert sched._reqs_need_recv["req"][1] == [[100, 101]]
492 + assert sched._req_kv_params["req"]["remote_block_ids"] == [[10, 11]]
493 +
494 +
495 +def test_read_producer_does_not_queue_write_only_save_state():
496 + sched = _FakeScheduler(
497 + is_producer=True,
498 + mode=MoRIIOMode.READ,
499 + _has_mamba=True,
500 + _attn_group_ids=[0],
501 + _mamba_group_ids=[1],
502 + request_id_to_transfer_id={},
503 + transfer_id_to_request_id={},
504 + _reqs_need_save={},
505 + _req_kv_params={},
506 + )
507 + request = SimpleNamespace(
508 + request_id="req",
509 + kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"},
510 + )
511 +
512 + sched.update_state_after_alloc(
513 + request,
514 + _FakeBlocks(all_groups=([100, 101], [200])),
515 + num_external_tokens=0,
516 + )
517 +
518 + assert sched._reqs_need_save == {}
519 + assert sched._req_kv_params == {}
520 +
521 +
522 +def test_write_producer_queues_grouped_attention_blocks():
523 + sched = _FakeScheduler(
524 + is_producer=True,
525 + mode=MoRIIOMode.WRITE,
526 + _has_mamba=False,
527 + _attn_group_ids=[0],
528 + _mamba_group_ids=[],
529 + request_id_to_transfer_id={},
530 + transfer_id_to_request_id={},
531 + _reqs_need_save={},
532 + _req_kv_params={},
533 + )
534 + request = SimpleNamespace(
535 + request_id="req",
536 + kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"},
537 + )
538 +
539 + sched.update_state_after_alloc(
540 + request,
541 + _FakeBlocks(all_groups=([100, 101],)),
542 + num_external_tokens=0,
543 + )
544 +
545 + assert sched._reqs_need_save["req"][1] == ([100, 101],)
546 + assert sched._req_kv_params["req"] == request.kv_transfer_params
547 +
548 +
549 +def test_session_build_rejects_per_layer_region_count_mismatch():
550 + worker = _FakeWorker(
551 + built_write_session={},
552 + layer_name_to_local_kv_cache_metadata={"kda.0": ["conv", "ssm"]},
553 + layer_name_to_remote_kv_cache_metadata={"prefill": {"kda.0": ["conv"]}},
554 + )
555 +
556 + with pytest.raises(moriio_common.MoRIIOError, match="registered 1 region"):
557 + worker._get_built_session("prefill")
558 +
559 +
560 +def test_register_kv_caches_rejects_hybrid_write_before_registration():
561 + worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"})
562 + worker._is_mamba_layer = lambda _layer_name: True
563 +
564 + with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"):
565 + worker.register_kv_caches({"kda.0": object()})
566 +
567 +
568 +def test_register_kv_caches_rejects_no_transfer_enabled_layers():
569 + worker = _FakeWorker(mode=MoRIIOMode.READ, _transfer_layer_names=set())
570 +
571 + with pytest.raises(moriio_common.MoRIIOError, match="no transfer-enabled"):
572 + worker.register_kv_caches({"disabled.0": object()})
573 +
574 +
575 +def test_worker_rejects_incompatible_mamba_specs_before_registration():
576 + incompatible = _mamba_spec()
577 + object.__setattr__(incompatible, "page_size_padded", 4096)
578 + worker = _FakeWorker(
579 + mode=MoRIIOMode.READ,
580 + world_size=1,
581 + _transfer_layer_names={"kda.0", "kda.1"},
582 + layer_to_spec={"kda.0": _mamba_spec(), "kda.1": incompatible},
583 + )
584 +
585 + with pytest.raises(moriio_common.MoRIIOError, match="same cache spec"):
586 + worker.register_kv_caches({"kda.0": object(), "kda.1": object()})
587 +
588 +
589 +# --------------------------------------------------------------------------
590 +# _truncate_mamba_request_for_prefill
591 +# --------------------------------------------------------------------------
592 +def _mk_request(prompt, max_tokens=64, params=None):
593 + return SimpleNamespace(
594 + kv_transfer_params={} if params is None else params,
595 + num_prompt_tokens=len(prompt),
596 + prompt_token_ids=list(prompt),
597 + prompt_embeds=None,
598 + _all_token_ids=list(prompt),
599 + max_tokens=max_tokens,
600 + )
601 +
602 +
603 +def test_truncate_mamba_request_pops_last_token_and_caps_tokens():
604 + sched = _FakeScheduler()
605 + req = _mk_request([10, 11, 12, 13, 14])
606 + sched._truncate_mamba_request_for_prefill(req)
607 +
608 + assert req.prompt_token_ids == [10, 11, 12, 13] # last token popped
609 + assert req._all_token_ids == [10, 11, 12, 13]
610 + assert req.num_prompt_tokens == 4 # N-1 accounting
611 + assert req.max_tokens == 1
612 + assert req.kv_transfer_params["_p_side_truncated"] is True
613 +
614 +
615 +def test_truncate_mamba_request_is_idempotent_across_reschedule():
616 + sched = _FakeScheduler()
617 + req = _mk_request([10, 11, 12])
618 + sched._truncate_mamba_request_for_prefill(req)
619 + assert req.num_prompt_tokens == 2
620 + # A second call (e.g. after a preemption) must not truncate again.
621 + sched._truncate_mamba_request_for_prefill(req)
622 + assert req.num_prompt_tokens == 2
623 + assert req.prompt_token_ids == [10, 11]
624 +
625 +
626 +def test_truncate_mamba_request_noop_for_single_token_prompt():
627 + sched = _FakeScheduler()
628 + req = _mk_request([10])
629 + sched._truncate_mamba_request_for_prefill(req)
630 + assert req.num_prompt_tokens == 1
631 + assert req.prompt_token_ids == [10]
632 + assert "_p_side_truncated" not in req.kv_transfer_params
633 +
634 +
635 +# --------------------------------------------------------------------------
636 +# get_num_new_matched_tokens (hybrid N-1 accounting)
637 +# --------------------------------------------------------------------------
638 +def test_get_num_new_matched_tokens_read_recomputes_last_token():
639 + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.READ, _has_mamba=True)
640 + req = SimpleNamespace(
641 + num_prompt_tokens=10,
642 + prompt_token_ids=list(range(10)),
643 + kv_transfer_params=None,
644 + )
645 + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
646 + # READ always recomputes the final token locally: N-1 - computed.
647 + assert n == 9
648 + assert is_async is False
649 +
650 +
651 +def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens():
652 + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False)
653 + req = SimpleNamespace(
654 + num_prompt_tokens=10,
655 + prompt_token_ids=list(range(10)),
656 + kv_transfer_params=None,
657 + )
658 + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=2)
659 + # Pure-attention WRITE: no N-1 drop; full length minus already-computed.
660 + assert n == 8
661 + assert is_async is True
662 +
663 +
664 +@pytest.mark.parametrize(
665 + ("mode", "num_computed_tokens", "expected", "is_async"),
666 + [
667 + (MoRIIOMode.READ, 0, 9, False),
668 + (MoRIIOMode.READ, 10, 0, False),
669 + (MoRIIOMode.WRITE, 2, 8, True),
670 + ],
671 +)
672 +def test_get_num_new_matched_tokens_supports_embeds_only_prompts(
673 + mode, num_computed_tokens, expected, is_async
674 +):
675 + sched = _FakeScheduler(is_producer=False, mode=mode, _has_mamba=True)
676 + req = SimpleNamespace(
677 + num_prompt_tokens=10,
678 + prompt_token_ids=None,
679 + prompt_embeds=object(),
680 + kv_transfer_params=None,
681 + )
682 +
683 + assert sched.get_num_new_matched_tokens(req, num_computed_tokens) == (
684 + expected,
685 + is_async,
686 + )
687 +
688 +
689 +def test_get_num_new_matched_tokens_producer_returns_zero():
690 + sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True)
691 + req = _mk_request(list(range(5)), params={"do_remote_decode": True})
692 + n, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
693 + assert (n, is_async) == (0, False)
694 + # The producer never truncates here: the scheduler has already measured the
695 + # prefix-cache hit against the full prompt by this point.
696 + assert req.num_prompt_tokens == 5
697 +
698 +
699 +def test_on_new_request_truncates_producer_prompt():
700 + sched = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=True)
701 + req = _mk_request(list(range(5)), params={"do_remote_decode": True})
702 + sched.on_new_request(req)
703 + # Producer stops at h(N-1): the last prompt token was dropped.
704 + assert req.num_prompt_tokens == 4
705 + assert req.max_tokens == 1
706 + assert req.kv_transfer_params["_p_side_truncated"] is True
707 +
708 +
709 +def test_on_new_request_is_noop_without_mamba_or_on_decode():
710 + plain = _FakeScheduler(is_producer=True, mode=MoRIIOMode.WRITE, _has_mamba=False)
711 + req = _mk_request(list(range(5)), params={"do_remote_decode": True})
712 + plain.on_new_request(req)
713 + assert req.num_prompt_tokens == 5
714 +
715 + decode_side = _FakeScheduler(
716 + is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True
717 + )
718 + req2 = _mk_request(list(range(5)), params={"do_remote_prefill": True})
719 + decode_side.on_new_request(req2)
720 + assert req2.num_prompt_tokens == 5
721 +
722 +
723 +# --------------------------------------------------------------------------
724 +# Offset-template cache
725 +# --------------------------------------------------------------------------
726 +_SAMPLE_SLOT_SETS = [
727 + ([1, 2], [1, 2]),
728 + ([0], [3]),
729 + ([3, 0, 2], [3, 0, 2]),
730 + ([2, 2], [5, 5]),
731 + ([], []),
732 +]
733 +
734 +
735 +def test_worker_compute_mamba_offsets_caches_and_matches_template():
736 + split = _gdn_split_info()
737 + conv_dim = sum(split.local_proj_dims)
738 + conv_shape = (conv_dim, split.conv_rows)
739 + ssm_shape = (2, 4, 4)
740 + page_bytes = 256
741 + pages = torch.zeros((8, 1, 1, page_bytes), dtype=torch.uint8)
742 + spec = moriio_connector.MambaSpec(
743 + block_size=1,
744 + shapes=(conv_shape, ssm_shape),
745 + dtypes=(torch.bfloat16, torch.bfloat16),
746 + page_size_padded=page_bytes,
747 + )
748 + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
749 +
750 + worker = _FakeWorker(
751 + kv_caches={"kda.0": pages},
752 + layer_to_spec={"kda.0": spec},
753 + _conv_decomp=split,
754 + _mamba_offset_templates={},
755 + tp_rank=0,
756 + world_size=1,
757 + )
758 +
759 + cached_template = None
760 + for local_slots, remote_slots in _SAMPLE_SLOT_SETS:
761 + lo, ro, sz, n_conv = worker._compute_mamba_transfer_offsets(
762 + "kda.0", local_slots, remote_slots
763 + )
764 + template = moriio_layout.build_mamba_offset_template(
765 + conv, ssm, split, tp_ratio=1
766 + )
767 + expected = moriio_layout.apply_mamba_offset_template(
768 + template, local_slots, remote_slots
769 + )
770 + assert (lo, ro, sz) == expected
771 + assert n_conv == moriio_layout.compute_mamba_conv_split_count(
772 + local_slots, split
773 + )
774 +
775 + current_template = worker._mamba_offset_templates["kda.0"]
776 + if cached_template is None:
777 + cached_template = current_template
778 + else:
779 + assert current_template is cached_template
780 +
781 + assert "kda.0" in worker._mamba_offset_templates
782 +
783 +
784 +@pytest.mark.parametrize("remote_tp_size", [1, None])
785 +def test_mamba_tp_ratio_accepts_only_equal_tp(remote_tp_size):
786 + worker = _FakeWorker(world_size=1)
787 + assert worker._mamba_tp_ratio(remote_tp_size) == 1
788 +
789 +
790 +@pytest.mark.parametrize("remote_tp_size", [2, 3])
791 +def test_mamba_tp_ratio_rejects_heterogeneous_tp(remote_tp_size):
792 + worker = _FakeWorker(world_size=1)
793 + with pytest.raises(NotImplementedError, match="heterogeneous-TP"):
794 + worker._mamba_tp_ratio(remote_tp_size)
795 +
796 +
797 +def test_worker_selects_each_mamba_layers_own_block_group():
798 + worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 1, "kda.1": 2})
799 + groups = ([10], [40], [70])
800 +
801 + assert worker._mamba_blocks_for_layer("kda.0", groups) == [40]
802 + assert worker._mamba_blocks_for_layer("kda.1", groups) == [70]
803 +
804 +
805 +def test_worker_rejects_missing_mamba_block_group():
806 + worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 2})
807 +
808 + with pytest.raises(moriio_common.MoRIIOError, match="no block group"):
809 + worker._mamba_blocks_for_layer("kda.0", ([10], [40]))
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+275
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib.util
5 +import math
6 import threading
7 from collections import OrderedDict, defaultdict
8 from queue import Queue
@@ -14,6 +15,7 @@ import torch
15 from vllm.platforms import current_platform
16 from vllm.v1.kv_cache_interface import (
17 FullAttentionSpec,
18 + MambaSpec,
19 MLAAttentionSpec,
20 )
21
@@ -36,6 +38,10 @@ moriio_layout = importlib.import_module(
38 "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
39 )
40 msgpack = importlib.import_module("msgpack")
41 +ssm_conv_transfer_utils = importlib.import_module(
42 + "vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils"
43 +)
44 +MambaConvSplitInfo = ssm_conv_transfer_utils.MambaConvSplitInfo
45
46 ROLE = moriio_common.ROLE
47 MoRIIOError = moriio_common.MoRIIOError
@@ -353,6 +359,9 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry():
359 kv_caches: dict[str, torch.Tensor]
360 layer_name_to_local_kv_cache_metadata: dict[str, list[Any]]
361
362 + def _region_session_indices(self, layer_name):
363 + return [list(self.kv_caches).index(layer_name)]
364 +
365 def _compute_block_transfer_offsets(
366 self,
367 layer_name,
@@ -673,6 +682,30 @@ def test_local_block_ids_longer_than_remote_raises_value_error():
682 )
683
684
685 +@pytest.mark.parametrize(
686 + ("local_block_ids", "remote_block_ids", "match"),
687 + [
688 + ([-1], [0], "local block id -1"),
689 + ([0], [8], "remote block id 8"),
690 + ],
691 +)
692 +def test_block_ids_must_fit_registered_regions(
693 + local_block_ids, remote_block_ids, match
694 +):
695 + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
696 + worker = _worker({"layer": cache}, {"layer": _full_spec()})
697 +
698 + with pytest.raises(ValueError, match=match):
699 + moriio_layout.compute_block_transfer_offsets(
700 + "layer",
701 + cache,
702 + worker.layer_to_spec,
703 + local_block_ids,
704 + remote_block_ids,
705 + _remote_meta(num_blocks=8).num_blocks,
706 + )
707 +
708 +
709 def test_empty_local_block_ids_is_free_only_noop():
710 cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
711 worker = _worker({"layer": cache}, {"layer": _full_spec()})
@@ -745,3 +778,245 @@ def test_unsupported_shape_raises_value_error():
778
779 with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
780 moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec)
781 +
782 +
783 +# ---------------------------------------------------------------------------
784 +# Mamba/KDA (conv + ssm) transfer geometry
785 +# ---------------------------------------------------------------------------
786 +
787 +
788 +def _slot_strided(num_slots, per_slot_shape, slot_stride, dtype=torch.bfloat16):
789 + """Build a slot-strided view: shape (num_slots, *per_slot_shape) whose slot
790 + stride (dim 0) exceeds the per-slot element count, mimicking the K3 KDA
791 + conv/ssm cache tensors."""
792 + inner = 1
793 + for d in per_slot_shape:
794 + inner *= d
795 + assert slot_stride >= inner
796 + backing = torch.zeros(num_slots * slot_stride, dtype=dtype)
797 + inner_strides: list[int] = []
798 + acc = 1
799 + for d in reversed(per_slot_shape):
800 + inner_strides.insert(0, acc)
801 + acc *= d
802 + return backing.as_strided(
803 + (num_slots, *per_slot_shape), (slot_stride, *inner_strides)
804 + )
805 +
806 +
807 +def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2):
808 + """A GDN-style conv split: sub-projections (key, key, value)."""
809 + conv_dim = 2 * key_dim + value_dim
810 + conv_state_bytes = conv_dim * conv_rows * dtype_size
811 + ssm_state_bytes = 64 # opaque here; ssm size is taken from the tensor
812 + return MambaConvSplitInfo(
813 + conv_rows=conv_rows,
814 + local_proj_dims=(key_dim, key_dim, value_dim),
815 + conv_dtype_size=dtype_size,
816 + ssm_sizes=(conv_state_bytes, ssm_state_bytes),
817 + )
818 +
819 +
820 +def test_mamba_transfer_geometry_is_slot_strided():
821 + split = _gdn_split_info()
822 + conv_dim = sum(split.local_proj_dims)
823 + conv = _slot_strided(5, (conv_dim, split.conv_rows), slot_stride=200)
824 + ssm = _slot_strided(5, (2, 4, 4), slot_stride=64)
825 +
826 + geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm)
827 +
828 + assert geom.conv_slot_stride == 200
829 + assert geom.ssm_slot_stride == 64
830 + assert geom.ssm_slot_bytes == (2 * 4 * 4) * ssm.element_size()
831 + # Region spans the full slot-strided extent (shape[0] * stride(0)), not just
832 + # numel(), so the highest-slot transfers stay in-bounds.
833 + assert geom.conv_region_len == conv.shape[0] * conv.stride(0) * conv.element_size()
834 + assert geom.ssm_region_len == ssm.shape[0] * ssm.stride(0) * ssm.element_size()
835 +
836 +
837 +def test_kda_conv_ssm_unpacks_padded_pages_as_aliasing_strided_views():
838 + num_blocks = 4
839 + page_bytes = 128
840 + conv_shape = (3, 4)
841 + ssm_shape = (2, 2, 4)
842 + conv_bytes = math.prod(conv_shape) * torch.bfloat16.itemsize
843 + ssm_bytes = math.prod(ssm_shape) * torch.bfloat16.itemsize
844 + assert conv_bytes + ssm_bytes < page_bytes
845 +
846 + pages = torch.zeros((num_blocks, 1, 1, page_bytes), dtype=torch.uint8)
847 + spec = MambaSpec(
848 + block_size=1,
849 + shapes=(conv_shape, ssm_shape),
850 + dtypes=(torch.bfloat16, torch.bfloat16),
851 + page_size_padded=page_bytes,
852 + )
853 +
854 + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
855 +
856 + assert conv.shape == (num_blocks, *conv_shape)
857 + assert ssm.shape == (num_blocks, *ssm_shape)
858 + assert conv.stride() == (page_bytes // 2, 4, 1)
859 + assert ssm.stride() == (page_bytes // 2, 8, 4, 1)
860 + assert conv.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr()
861 + assert ssm.untyped_storage().data_ptr() == pages.untyped_storage().data_ptr()
862 + assert conv.storage_offset() == 0
863 + assert ssm.storage_offset() * ssm.element_size() == conv_bytes
864 +
865 + conv[1].fill_(1)
866 + ssm[2].fill_(2)
867 + assert torch.all(pages[1, 0, 0, :conv_bytes].view(torch.bfloat16) == 1)
868 + assert torch.all(
869 + pages[2, 0, 0, conv_bytes : conv_bytes + ssm_bytes].view(torch.bfloat16) == 2
870 + )
871 +
872 +
873 +def test_kda_conv_ssm_accepts_strided_blocks_without_copying():
874 + backing = torch.zeros(4 * 256, dtype=torch.uint8)
875 + pages = backing.as_strided((4, 1, 1, 128), (256, 128, 128, 1))
876 + spec = MambaSpec(
877 + block_size=1,
878 + shapes=((3, 4), (2, 2, 4)),
879 + dtypes=(torch.bfloat16, torch.bfloat16),
880 + page_size_padded=128,
881 + )
882 +
883 + conv, ssm = moriio_layout.kda_conv_ssm(pages, spec)
884 +
885 + assert not pages.is_contiguous()
886 + assert conv.stride(0) == 128
887 + assert ssm.stride(0) == 128
888 + assert conv.untyped_storage().data_ptr() == backing.data_ptr()
889 + assert ssm.untyped_storage().data_ptr() == backing.data_ptr()
890 +
891 +
892 +def test_kda_conv_ssm_rejects_noncontiguous_page_bytes():
893 + pages = torch.zeros((4, 1, 1, 256), dtype=torch.uint8)[..., ::2]
894 + spec = MambaSpec(
895 + block_size=1,
896 + shapes=((3, 4), (2, 2, 4)),
897 + dtypes=(torch.bfloat16, torch.bfloat16),
898 + page_size_padded=128,
899 + )
900 +
901 + with pytest.raises(ValueError, match="contiguous byte dimension"):
902 + moriio_layout.kda_conv_ssm(pages, spec)
903 +
904 +
905 +def test_kda_conv_ssm_rejects_non_byte_packed_cache():
906 + pages = torch.zeros((4, 1, 1, 128), dtype=torch.bfloat16)
907 + spec = MambaSpec(
908 + block_size=1,
909 + shapes=((3, 4), (2, 2, 4)),
910 + dtypes=(torch.bfloat16, torch.bfloat16),
911 + page_size_padded=128,
912 + )
913 +
914 + with pytest.raises(ValueError, match="byte-sized dtype"):
915 + moriio_layout.kda_conv_ssm(pages, spec)
916 +
917 +
918 +def test_mamba_transfer_geometry_requires_matching_slot_counts():
919 + conv = torch.empty((5, 4, 3), dtype=torch.bfloat16)
920 + ssm = torch.empty((4, 2, 4, 4), dtype=torch.bfloat16)
921 +
922 + with pytest.raises(ValueError, match="conv/ssm slot count mismatch"):
923 + moriio_layout.get_mamba_transfer_geometry(conv, ssm)
924 +
925 +
926 +def test_mamba_transfer_geometry_rejects_empty_caches():
927 + conv = torch.empty((0, 4, 3), dtype=torch.bfloat16)
928 + ssm = torch.empty((0, 2, 4, 4), dtype=torch.bfloat16)
929 +
930 + with pytest.raises(ValueError, match="at least one slot"):
931 + moriio_layout.get_mamba_transfer_geometry(conv, ssm)
932 +
933 +
934 +def test_mamba_conv_ssm_offsets_homogeneous_tp():
935 + split = _gdn_split_info()
936 + conv_dim = sum(split.local_proj_dims)
937 + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
938 + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
939 + geom = moriio_layout.get_mamba_transfer_geometry(conv, ssm)
940 +
941 + local_slots = [1, 2]
942 + remote_slots = [1, 2]
943 + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
944 + local_offs, remote_offs, sizes = moriio_layout.apply_mamba_offset_template(
945 + template, local_slots, remote_slots
946 + )
947 +
948 + n_conv = moriio_layout.compute_mamba_conv_split_count(local_slots, split)
949 + assert n_conv == len(local_slots) * len(split.local_proj_dims)
950 + assert len(local_offs) == n_conv + len(local_slots)
951 +
952 + # Homogeneous TP: remote conv slices equal local slices; since local==remote
953 + # slots, conv local and remote offsets coincide.
954 + loc_co = split.local_conv_offsets
955 + rem_co = split.remote_conv_offsets(0, 1)
956 + k = 0
957 + esz = conv.element_size()
958 + for ls, rs in zip(local_slots, remote_slots):
959 + lbase = ls * geom.conv_slot_stride * esz
960 + rbase = rs * geom.conv_slot_stride * esz
961 + for (loff, _), (roff, rsz) in zip(loc_co, rem_co):
962 + assert local_offs[k] == lbase + loff
963 + assert remote_offs[k] == rbase + roff
964 + assert sizes[k] == rsz
965 + k += 1
966 +
967 + # SSM: whole per-slot block, no head sharding at tp_ratio == 1.
968 + essz = ssm.element_size()
969 + for j, (ls, rs) in enumerate(zip(local_slots, remote_slots)):
970 + assert local_offs[n_conv + j] == ls * geom.ssm_slot_stride * essz
971 + assert remote_offs[n_conv + j] == rs * geom.ssm_slot_stride * essz
972 + assert sizes[n_conv + j] == geom.ssm_slot_bytes
973 +
974 +
975 +def test_mamba_conv_ssm_offsets_heterogeneous_tp_is_gated():
976 + # Heterogeneous-TP mamba transfer is a documented follow-up (remote slot
977 + # stride + multi-rank fan-in gather). Until then the offset helper must fail
978 + # loudly rather than silently transfer wrong bytes, since the mamba path
979 + # bypasses the attention KV-head validator.
980 + split = _gdn_split_info()
981 + conv_dim = sum(split.local_proj_dims)
982 + conv = _slot_strided(4, (conv_dim, split.conv_rows), slot_stride=200)
983 + ssm = _slot_strided(4, (2, 4, 4), slot_stride=64)
984 + with pytest.raises(NotImplementedError):
985 + moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=2)
986 +
987 +
988 +def test_local_mamba_slots_must_fit_registered_regions():
989 + split = _gdn_split_info()
990 + conv_dim = sum(split.local_proj_dims)
991 + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
992 + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
993 + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
994 +
995 + with pytest.raises(ValueError, match="local mamba slot 6"):
996 + moriio_layout.apply_mamba_offset_template(template, [6], [0])
997 +
998 +
999 +@pytest.mark.parametrize(
1000 + ("local_slots", "remote_slots"),
1001 + [([1, 2], [3]), ([1], [2, 3])],
1002 +)
1003 +def test_mamba_slot_counts_must_match(local_slots, remote_slots):
1004 + split = _gdn_split_info()
1005 + conv_dim = sum(split.local_proj_dims)
1006 + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
1007 + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
1008 + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
1009 +
1010 + with pytest.raises(ValueError, match="slot counts must match"):
1011 + moriio_layout.apply_mamba_offset_template(template, local_slots, remote_slots)
1012 +
1013 +
1014 +def test_remote_mamba_slots_must_be_non_negative():
1015 + split = _gdn_split_info()
1016 + conv_dim = sum(split.local_proj_dims)
1017 + conv = _slot_strided(6, (conv_dim, split.conv_rows), slot_stride=200)
1018 + ssm = _slot_strided(6, (2, 4, 4), slot_stride=64)
1019 + template = moriio_layout.build_mamba_offset_template(conv, ssm, split, tp_ratio=1)
1020 +
1021 + with pytest.raises(ValueError, match="must be non-negative"):
1022 + moriio_layout.apply_mamba_offset_template(template, [1], [-1])
tests/v1/kv_connector/unit/test_moriio_routing.py
+1
@@ -28,6 +28,7 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
28 s = SimpleNamespace(
29 is_producer=True,
30 _is_hma_required=False,
31 + _has_mamba=False,
32 blocks_per_sw=[0],
33 _global_dp_rank=global_dp_rank,
34 engine_id="engine-abc",
tests/v1/kv_connector/unit/test_moriio_tp_ack.py
+141 -1
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import threading
5 +from types import SimpleNamespace
6
7 import pytest
8
9 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
10 MoRIIOMode,
11 MoRIIOTransferAck,
12 + TransferError,
13 )
14 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
15 MoRIIOConnector,
@@ -17,6 +19,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import
19 resolve_moriio_transfer_ack,
20 validate_moriio_heterogeneous_tp_kv_heads,
21 )
22 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
23 + MoRIIOWrapper,
24 +)
25
26
27 def test_remote_tp_rank_same_tp_maps_to_self():
@@ -277,6 +282,9 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
282 self.lock = threading.Lock()
283 self.sent = []
284
285 + # Exercise the real batch verdict rather than restating it here.
286 + poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch
287 +
288 def send_notify(
289 self,
290 transfer_id,
@@ -292,13 +300,18 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
300
301 worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
302 worker.world_size = 8
303 + worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0)
304 worker.moriio_wrapper = FakeWrapper()
296 - worker._recving_transfers = {"req": {"layer0": DoneStatus()}}
305 + # A layer maps to the list of reads posted for it (a KDA layer posts two).
306 + worker._recving_transfers = {"req": {"layer0": [DoneStatus(), DoneStatus()]}}
307 worker._recving_transfers_callback_addr = {
308 "req": ("127.0.0.1", "7000", "tx-release")
309 }
310 # Transfer-timeout reaping state consulted by _pop_done_transfers.
311 worker._recving_transfers_start = {}
312 + # Load-error bookkeeping cleared alongside a completed transfer.
313 + worker._recving_local_blocks = {}
314 + worker._invalid_block_ids = set()
315
316 assert worker._pop_done_transfers() == {"tx-release"}
317 assert worker.moriio_wrapper.sent == [
@@ -314,6 +327,133 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
327 assert worker._recving_transfers_callback_addr == {}
328
329
330 +def test_aborted_read_without_local_blocks_only_releases_remote_blocks():
331 + class FakeWrapper:
332 + def __init__(self):
333 + self.sent = []
334 +
335 + def send_notify(
336 + self,
337 + transfer_id,
338 + host,
339 + port,
340 + message_type=None,
341 + message_fields=None,
342 + ):
343 + self.sent.append((transfer_id, host, port, message_type, message_fields))
344 +
345 + def shutdown(self):
346 + pass
347 +
348 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
349 + worker.mode = MoRIIOMode.READ
350 + worker.world_size = 8
351 + worker.moriio_wrapper = FakeWrapper()
352 +
353 + worker._read_blocks(
354 + local_block_ids=[],
355 + remote_block_ids=[[10, 11], [90]],
356 + dst_engine_id="prefill",
357 + request_id="req-aborted",
358 + transfer_id="tx-aborted",
359 + remote_host="127.0.0.1",
360 + remote_notify_port=7000,
361 + remote_tp_size=8,
362 + remote_dp_rank=0,
363 + chosen_tp=2,
364 + )
365 +
366 + assert worker.moriio_wrapper.sent == [
367 + (
368 + "tx-aborted",
369 + "127.0.0.1",
370 + "7002",
371 + "release",
372 + {"consumer_tp_size": 8},
373 + )
374 + ]
375 +
376 +
377 +def test_read_completion_waits_for_every_posted_read():
378 + class Status:
379 + def __init__(self, done: bool):
380 + self.done = done
381 +
382 + def Succeeded(self):
383 + return self.done
384 +
385 + def Failed(self):
386 + return False
387 +
388 + class FakeWrapper:
389 + def __init__(self):
390 + self.lock = threading.Lock()
391 + self.sent = []
392 +
393 + poll_transfer_batch = MoRIIOWrapper.poll_transfer_batch
394 +
395 + def send_notify(self, *args, **kwargs):
396 + self.sent.append((args, kwargs))
397 +
398 + def shutdown(self):
399 + pass
400 +
401 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
402 + worker.world_size = 8
403 + worker.moriio_config = SimpleNamespace(recv_abort_timeout=600.0)
404 + worker.moriio_wrapper = FakeWrapper()
405 + worker._recving_transfers = {
406 + "req": {"kda_layer": [Status(done=True), Status(done=False)]}
407 + }
408 + worker._recving_transfers_callback_addr = {
409 + "req": ("127.0.0.1", "7000", "tx-pending")
410 + }
411 + worker._recving_transfers_start = {"req": float("inf")}
412 + worker._recving_local_blocks = {"req": [1, 2]}
413 + worker._invalid_block_ids = set()
414 +
415 + assert worker._pop_done_transfers() == set()
416 + assert "req" in worker._recving_transfers
417 + assert worker.moriio_wrapper.sent == []
418 +
419 +
420 +@pytest.mark.parametrize(
421 + ("has_mamba", "expected_invalid"),
422 + [(False, {1, 2}), (True, set())],
423 +)
424 +def test_failed_read_reports_blocks_only_without_hma(has_mamba, expected_invalid):
425 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
426 + worker._has_mamba = has_mamba
427 + worker._recving_local_blocks = {"req": [1, 2]}
428 + worker._invalid_block_ids = set()
429 +
430 + worker._record_failed_recv("req")
431 +
432 + assert worker.get_block_ids_with_load_errors() == expected_invalid
433 +
434 +
435 +def test_hybrid_step_barrier_fails_closed(monkeypatch):
436 + class FailingWrapper:
437 + def waiting_for_transfer_complete(self, _statuses):
438 + raise TransferError("failed")
439 +
440 + def shutdown(self):
441 + pass
442 +
443 + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
444 + worker._has_mamba = True
445 + worker._reads_issued_this_step = [object()]
446 + worker._mamba_reads_this_step = [object()]
447 + worker.moriio_wrapper = FailingWrapper()
448 + monkeypatch.setattr(
449 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.get_forward_context",
450 + lambda: SimpleNamespace(cudagraph_runtime_mode=None),
451 + )
452 +
453 + with pytest.raises(TransferError, match="failed"):
454 + worker._await_reads_issued_this_step()
455 +
456 +
457 def test_requested_cudagraph_mode_is_never_overridden():
458 # The configured cudagraph mode is always honored: the barrier fires when
459 # the operator sets cudagraph_mode=PIECEWISE, and READ mode with full
tests/v1/kv_connector/unit/test_moriio_transfer_completion.py new
+129
@@ -0,0 +1,129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Transfer-completion semantics of the MoRIIO wrapper.
4 +
5 +Covers the two behaviours the READ barrier depends on:
6 + * poll_transfer_batch considers EVERY status, not just the newest one
7 + * waiting_for_transfer_complete blocks until terminal and raises on failure.
8 +"""
9 +
10 +import importlib
11 +import importlib.util
12 +import threading
13 +import time
14 +
15 +import pytest
16 +
17 +from vllm.platforms import current_platform
18 +
19 +mori_available = importlib.util.find_spec("mori") is not None
20 +
21 +if not (current_platform.is_rocm() and mori_available):
22 + pytest.skip(
23 + "MoRIIOs are only available on ROCm with mori package installed",
24 + allow_module_level=True,
25 + )
26 +
27 +moriio_common = importlib.import_module(
28 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
29 +)
30 +moriio_engine = importlib.import_module(
31 + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine"
32 +)
33 +
34 +MoRIIOWrapper = moriio_engine.MoRIIOWrapper
35 +TransferBatchState = moriio_common.TransferBatchState
36 +TransferError = moriio_common.TransferError
37 +
38 +
39 +@pytest.mark.parametrize("value", [0, -1, float("inf"), float("nan")])
40 +def test_timeout_must_be_positive_and_finite(value):
41 + with pytest.raises(ValueError, match="must be finite and greater than zero"):
42 + moriio_common._positive_finite_timeout("test_timeout", value)
43 +
44 +
45 +def test_positive_finite_timeout_is_preserved():
46 + assert moriio_common._positive_finite_timeout("test_timeout", "1.5") == 1.5
47 +
48 +
49 +class FakeStatus:
50 + """Mirrors the mori TransferStatus surface the wrapper uses."""
51 +
52 + def __init__(self, state: str = "in_progress", message: str = ""):
53 + self.state = state
54 + self.message = message
55 +
56 + def Succeeded(self) -> bool:
57 + return self.state == "success"
58 +
59 + def Failed(self) -> bool:
60 + return self.state == "failed"
61 +
62 + def Message(self) -> str:
63 + return self.message
64 +
65 + def Code(self) -> str:
66 + return self.state
67 +
68 +
69 +def _ok():
70 + return FakeStatus("success")
71 +
72 +
73 +def _pending():
74 + return FakeStatus("in_progress")
75 +
76 +
77 +def _bad():
78 + return FakeStatus("failed", "SQ full")
79 +
80 +
81 +@pytest.mark.parametrize(
82 + "statuses,expected",
83 + [
84 + ([], TransferBatchState.DONE),
85 + ([_ok(), _ok()], TransferBatchState.DONE),
86 + # Each case below puts the interesting status BEFORE a success, which a
87 + # statuses[-1] check would have missed entirely.
88 + ([_bad(), _ok()], TransferBatchState.FAILED),
89 + ([_pending(), _ok()], TransferBatchState.PENDING),
90 + ([_pending(), _bad(), _ok()], TransferBatchState.FAILED),
91 + ],
92 +)
93 +def test_poll_transfer_batch_considers_every_status(statuses, expected):
94 + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
95 + assert wrapper.poll_transfer_batch(statuses) is expected
96 +
97 +
98 +@pytest.mark.parametrize(
99 + "statuses_factory,expected_message",
100 + [
101 + (lambda: [_ok(), _bad()], "SQ full"),
102 + (lambda: [_ok(), _pending()], "timed out"),
103 + ],
104 + ids=["failed", "timeout"],
105 +)
106 +def test_wait_raises_with_detail(statuses_factory, expected_message):
107 + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
108 + with pytest.raises(TransferError, match=expected_message):
109 + wrapper.waiting_for_transfer_complete(statuses_factory())
110 +
111 +
112 +def test_wait_returns_immediately_when_nothing_pending():
113 + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=0.05)
114 + wrapper.waiting_for_transfer_complete([])
115 + wrapper.waiting_for_transfer_complete([_ok(), _ok()])
116 +
117 +
118 +def test_wait_blocks_until_terminal():
119 + """The barrier is only worth anything if it actually waits."""
120 + wrapper = MoRIIOWrapper(moriio_engine=object(), transfer_timeout=5.0)
121 + late = FakeStatus("in_progress")
122 + timer = threading.Timer(0.2, lambda: setattr(late, "state", "success"))
123 + timer.start()
124 + try:
125 + start = time.monotonic()
126 + wrapper.waiting_for_transfer_complete([late])
127 + assert time.monotonic() - start >= 0.15
128 + finally:
129 + timer.cancel()
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+56 -7
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import contextlib
4 +import math
5 import os
6 import threading
7 import time
@@ -45,6 +46,13 @@ TransferId = str
46 TransferOffsetsKey = tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]
47
48
49 +def _positive_finite_timeout(name: str, value: Any) -> float:
50 + timeout = float(value)
51 + if not math.isfinite(timeout) or timeout <= 0:
52 + raise ValueError(f"{name} must be finite and greater than zero")
53 + return timeout
54 +
55 +
56 class MoRIIOTransferAck(NamedTuple):
57 transfer_id: TransferId
58 consumer_tp_size: int = 1
@@ -160,6 +168,20 @@ class MoRIIOMode(Enum):
168 WRITE = "write"
169
170
171 +class TransferBatchState(Enum):
172 + """Verdict for a group of transfer statuses that belong to one request.
173 +
174 + A request's KV transfer is spread over one status per layer (two for a KDA
175 + layer), so a single status never decides the request: PENDING means at
176 + least one is still in flight and none has failed, FAILED means at least
177 + one failed, DONE means all succeeded.
178 + """
179 +
180 + DONE = "done"
181 + FAILED = "failed"
182 + PENDING = "pending"
183 +
184 +
185 class MoRIIOError(Exception):
186 """Base exception for MoRIIO operations."""
187
@@ -260,6 +282,7 @@ class MoRIIOConfig:
282 tp_size: int
283 transfer_timeout: float
284 defer_timeout: float
285 + recv_abort_timeout: float
286 read_mode: bool = False
287 qp_per_transfer: int = 1
288 post_batch_size: int = -1
@@ -284,6 +307,8 @@ class MoRIIOConfig:
307 # raising TransferError (sec).
308 # defer_timeout -> Timeout before a deferred send with no finished_sending
309 # notification is reaped and its blocks force-freed (sec).
310 + # recv_abort_timeout -> Timeout before an in-flight recv whose RDMA
311 + # completion never arrived is aborted (sec).
312
313 # Knobs for RDMA transfers, ignored if on xgmi backend
314 # qp_per_transfer -> Number of RDMA Queue Pairs per KV transfer.
@@ -315,13 +340,29 @@ class MoRIIOConfig:
340 "must be one of 'rdma' or 'xgmi'."
341 )
342
318 - transfer_timeout = float(
343 + transfer_timeout = _positive_finite_timeout(
344 + "kv_connector_extra_config.transfer_timeout",
345 extra_config.get(
346 "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT
321 - )
347 + ),
348 )
323 - defer_timeout = float(
324 - extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT)
349 + defer_timeout = _positive_finite_timeout(
350 + "kv_connector_extra_config.defer_timeout",
351 + extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT),
352 + )
353 + legacy_recv_abort_timeout = os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S")
354 + if legacy_recv_abort_timeout is not None:
355 + logger.warning_once(
356 + "The environment variable VLLM_MORIIO_TRANSFER_TIMEOUT_S is "
357 + "deprecated. Set 'recv_abort_timeout' inside "
358 + "kv_transfer_config.kv_connector_extra_config instead."
359 + )
360 + recv_abort_timeout = _positive_finite_timeout(
361 + "kv_connector_extra_config.recv_abort_timeout",
362 + extra_config.get(
363 + "recv_abort_timeout",
364 + legacy_recv_abort_timeout or MoRIIOConstants.DEFAULT_RECV_ABORT_TIMEOUT,
365 + ),
366 )
367
368 return cls(
@@ -347,6 +388,7 @@ class MoRIIOConfig:
388 backend=backend,
389 transfer_timeout=transfer_timeout,
390 defer_timeout=defer_timeout,
391 + recv_abort_timeout=recv_abort_timeout,
392 )
393
394
@@ -374,6 +416,10 @@ class MoRIIOConstants:
416 # notification is reaped and its blocks force-freed.
417 # Overridable via kv_connector_extra_config["defer_timeout"].
418 DEFAULT_DEFER_TIMEOUT = 60.0
419 + # Timeout (seconds) before an in-flight recv whose RDMA completion was lost
420 + # is aborted, so the decode worker does not hang on it forever.
421 + # Overridable via kv_connector_extra_config["recv_abort_timeout"].
422 + DEFAULT_RECV_ABORT_TIMEOUT = 120.0
423
424
425 # The router embeds both zmq_addresses in the request_id:
@@ -567,9 +613,12 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
613 remote_handshake_port=int(remote_handshake_port),
614 remote_notify_port=int(remote_notify_port),
615 # Remote peer TP degree (used as remote_tp_size downstream). The
570 - # proxy advertises it under "remote_tp_size"; #46332 read "tp_size"
571 - # which is absent on WRITE producer requests -> defaulted to 1 ->
572 - # rank collapse. Read the right key; 0 == unknown (== homogeneous).
616 + # decode side must know the prefiller's TP degree to map producer
617 + # ranks onto the correct decode rank (and vice versa). The old plain
618 + # "tp_size" key (#46332) is absent on WRITE producer requests and
619 + # defaulted to 1, collapsing every producer rank onto decode rank 0
620 + # and corrupting the transfer; prefer "remote_tp_size". 0 == unknown
621 + # (treated as homogeneous downstream).
622 tp_size=int(
623 kv_transfer_params.get("remote_tp_size")
624 or kv_transfer_params.get("tp_size")
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+1003 -167
@@ -2,14 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import logging
4 import math
5 -import os
5 import queue
6 import threading
7 import time
8 from collections import defaultdict
9 from collections.abc import Collection
10 from concurrent.futures import Future, ThreadPoolExecutor
12 -from typing import TYPE_CHECKING, Any
11 +from typing import TYPE_CHECKING, Any, cast
12
13 import msgpack
14 import msgspec
@@ -33,10 +32,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
32 MoRIIOConfig,
33 MoRIIOConnectorMetadata,
34 MoRIIOConstants,
35 + MoRIIOError,
36 MoRIIOMode,
37 MoRIIOTransferAck,
38 ReqId,
39 ReqMeta,
40 + TransferBatchState,
41 + TransferError,
42 TransferId,
43 WriteTask,
44 fold_local_rank,
@@ -56,11 +58,20 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
58 )
59 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
60 LayerTransferGeometry,
61 + MambaOffsetTemplate,
62 + apply_mamba_offset_template,
63 build_layer_to_spec,
64 + build_mamba_offset_template,
65 compute_block_transfer_offsets,
66 + compute_mamba_conv_split_count,
67 get_layer_transfer_geometry,
68 is_mla_cache_layer,
69 iter_layer_registration_regions,
70 + kda_conv_ssm,
71 +)
72 +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
73 + MambaConvSplitInfo,
74 + derive_mamba_conv_split,
75 )
76 from vllm.distributed.parallel_state import (
77 get_tensor_model_parallel_world_size,
@@ -79,6 +90,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
90 from vllm.v1.kv_cache_interface import (
91 FullAttentionSpec,
92 KVCacheConfig,
93 + MambaSpec,
94 SlidingWindowSpec,
95 )
96 from vllm.v1.outputs import KVConnectorOutput
@@ -91,6 +103,13 @@ if TYPE_CHECKING:
103
104 logger = init_logger(__name__)
105
106 +# Backoff bounds for re-posting an RDMA READ that the NIC rejected because the
107 +# send queue was full. A separate CQ-poll thread drains completions, so the
108 +# wait only needs to be long enough for it to make progress.
109 +_SQ_FULL_BACKOFF_INITIAL_S = 0.001
110 +_SQ_FULL_BACKOFF_MAX_S = 0.05
111 +_MAX_LOCAL_DECODE_TAIL_BLOCKS = 1
112 +
113
114 try:
115 from mori.io import (
@@ -197,6 +216,12 @@ def resolve_moriio_transfer_ack(
216
217
218 class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
219 + @property
220 + def supports_divergent_local_hybrid_hits(self) -> bool:
221 + # The READ path always transfers the recurrent-state slot, including
222 + # when the attention prefix is already present locally.
223 + return self.mode == MoRIIOMode.READ and self._transfers_mamba_state
224 +
225 def __init__(
226 self,
227 vllm_config: VllmConfig,
@@ -217,19 +242,10 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
242 + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
243 )
244 self.mode = get_moriio_mode(self.kv_transfer_config)
220 - if (
221 - self.mode == MoRIIOMode.READ
222 - and self.kv_transfer_config.is_kv_consumer
223 - and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
224 - ):
225 - # warn only; kv-read barrier requires PIECEWISE cudagraph mode
226 - logger.warning_once(
227 - "MoRIIO READ mode is running with %s CUDA graphs: per-layer "
228 - "KV-read barrier can't fire inside full graph; accuracy may "
229 - "degrade at high concurrency. Set cudagraph_mode=PIECEWISE "
230 - "in --compilation-config.",
231 - vllm_config.compilation_config.cudagraph_mode.name,
232 - )
245 + self._transfers_mamba_state = any(
246 + group.enable_kv_transfer and isinstance(group.kv_cache_spec, MambaSpec)
247 + for group in kv_cache_config.kv_cache_groups
248 + )
249 if role == KVConnectorRole.SCHEDULER:
250 self.connector_scheduler: MoRIIOConnectorScheduler | None = (
251 MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
@@ -279,6 +295,11 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
295 request, blocks, num_external_tokens, self.connector_worker
296 )
297
298 + def on_new_request(self, request: "Request") -> None:
299 + """Fired when the scheduler admits a request, before any allocation."""
300 + assert self.connector_scheduler is not None
301 + self.connector_scheduler.on_new_request(request)
302 +
303 def build_connector_meta(
304 self,
305 scheduler_output: SchedulerOutput,
@@ -300,6 +321,14 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
321 block_ids: BlockIds,
322 ) -> tuple[bool, dict[str, Any] | None]:
323 assert self.connector_scheduler is not None
324 + if self.connector_scheduler._has_mamba:
325 + attn_block_ids, mamba_block_groups = (
326 + self.connector_scheduler.split_block_groups(block_ids)
327 + )
328 + delay_free, params = self.connector_scheduler.request_finished(
329 + request, attn_block_ids, mamba_block_groups
330 + )
331 + return delay_free, params
332 return self.connector_scheduler.request_finished(request, block_ids)
333
334 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
@@ -318,6 +347,12 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
347 assert self.connector_worker is not None
348 return self.connector_worker.get_finished()
349
350 + def get_block_ids_with_load_errors(self) -> set[int]:
351 + """Blocks whose MoRIIO read failed, for the scheduler to recompute."""
352 + if self.connector_worker is None:
353 + return set()
354 + return self.connector_worker.get_block_ids_with_load_errors()
355 +
356 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
357 assert self.connector_worker is not None
358 if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.CONSUMER:
@@ -381,6 +416,47 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
416 return False
417
418
419 +def _split_kv_cache_group_kinds(
420 + kv_cache_config: "KVCacheConfig",
421 +) -> tuple[list[int], list[int]]:
422 + """Classify kv_cache_groups into (attention_group_indices,
423 + mamba_group_indices) by whether each group's spec is a MambaSpec.
424 +
425 + Groups that opt out of external transfer are absent from
426 + ``KVCacheConfig.transfer_groups``. Returned indices address that filtered
427 + sequence and therefore also the connector's transfer block payload.
428 + """
429 + attn: list[int] = []
430 + mamba: list[int] = []
431 + for gi, group in enumerate(kv_cache_config.transfer_groups):
432 + if isinstance(group.kv_cache_spec, MambaSpec):
433 + mamba.append(gi)
434 + else:
435 + attn.append(gi)
436 + return attn, mamba
437 +
438 +
439 +def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None:
440 + """Validate the common packed-state contract and return its representative."""
441 + specs = tuple(specs)
442 + if not specs:
443 + return None
444 + for spec in specs:
445 + if len(spec.shapes) != 2 or len(spec.dtypes) != 2:
446 + raise MoRIIOError(
447 + "MoRIIO hybrid READ supports exactly two Mamba states "
448 + "(conv and SSM), got "
449 + f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes"
450 + )
451 + reference = specs[0]
452 + if any(spec != reference for spec in specs[1:]):
453 + raise MoRIIOError(
454 + "MoRIIO hybrid READ requires all transferable Mamba groups "
455 + "to use the same cache spec"
456 + )
457 + return reference
458 +
459 +
460 class MoRIIOConnectorScheduler:
461 """Implementation of Scheduler side methods."""
462
@@ -391,35 +467,73 @@ class MoRIIOConnectorScheduler:
467 kv_cache_config: "KVCacheConfig",
468 ):
469 self.vllm_config = vllm_config
394 -
470 + self.kv_cache_config = kv_cache_config
471 assert vllm_config.kv_transfer_config is not None, (
472 "kv_transfer_config must be set for MoRIIOConnector"
473 )
474 self.kv_transfer_config = vllm_config.kv_transfer_config
475 + self.mode = get_moriio_mode(self.kv_transfer_config)
476 + self._attn_group_ids, self._mamba_group_ids = _split_kv_cache_group_kinds(
477 + kv_cache_config
478 + )
479 + self._has_mamba = bool(self._mamba_group_ids)
480 + if self._has_mamba:
481 + if len(self._attn_group_ids) != 1:
482 + raise MoRIIOError(
483 + "MoRIIO hybrid READ requires exactly one transferable "
484 + "attention group, got "
485 + f"{len(self._attn_group_ids)} attention groups"
486 + )
487 + mamba_specs = [
488 + cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec)
489 + for group_id in self._mamba_group_ids
490 + ]
491 + _validate_mamba_specs(mamba_specs)
492 + if vllm_config.speculative_config is not None:
493 + raise MoRIIOError(
494 + "MoRIIO hybrid READ does not support speculative decoding"
495 + )
496 + if self.mode != MoRIIOMode.READ:
497 + raise MoRIIOError(
498 + "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
499 + "mode only; set kv_connector_extra_config.read_mode=true"
500 + )
501 + # Only "all" mode keeps a state per block position; the other modes keep
502 + # a single running state in the last slot.
503 + self._ssm_state_slots_are_positional = (
504 + vllm_config.cache_config.mamba_cache_mode == "all"
505 + )
506 self.block_size = vllm_config.cache_config.block_size
507 + self._max_decode_tail_blocks = (
508 + _MAX_LOCAL_DECODE_TAIL_BLOCKS
509 + if self._has_mamba
510 + else math.ceil((vllm_config.num_lookahead_tokens + 1) / self.block_size)
511 + )
512 self.engine_id: EngineId = engine_id
401 - self.mode = get_moriio_mode(self.kv_transfer_config)
513
514 self._is_hma_required = (
515 not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
516 and any(
517 not isinstance(g.kv_cache_spec, FullAttentionSpec)
407 - for g in kv_cache_config.kv_cache_groups
518 + for g in kv_cache_config.transfer_groups
519 )
520 )
521 if self._is_hma_required:
411 - # TODO(simondanielsson): support non-sliding window hybrids
522 + # MoRIIO handles sliding-window attention and mamba/KDA groups in
523 + # READ mode. Other hybrid cache specs still fail closed.
524 unsupported = [
525 type(g.kv_cache_spec).__name__
414 - for g in kv_cache_config.kv_cache_groups
526 + for g in kv_cache_config.transfer_groups
527 if not isinstance(
416 - g.kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec)
528 + g.kv_cache_spec,
529 + (FullAttentionSpec, SlidingWindowSpec, MambaSpec),
530 )
531 ]
532 if unsupported:
533 raise NotImplementedError(
534 "MoRIIO only supports sliding-window hybrid models (e.g. "
422 - "Gemma) under the hybrid KV cache manager. Unsupported KV "
535 + "Gemma) and mamba/KDA groups under the hybrid KV cache "
536 + "manager. Unsupported KV "
537 f"cache group spec(s): {sorted(set(unsupported))}. Pass "
538 "--disable-hybrid-kv-cache-manager to run without HMA."
539 )
@@ -435,7 +549,7 @@ class MoRIIOConnectorScheduler:
549 (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
550 if isinstance(g.kv_cache_spec, SlidingWindowSpec)
551 else (0, self.block_size)
438 - for g in kv_cache_config.kv_cache_groups
552 + for g in kv_cache_config.transfer_groups
553 ]
554 # add 1 to conservatively account for boundary overlap eg window isn't fully
555 # aligned with blocks.
@@ -451,8 +565,7 @@ class MoRIIOConnectorScheduler:
565 self._full_attn_group_idx = 0
566 self._full_attn_block_size = self.block_size
567 for gi, group in enumerate(kv_cache_config.kv_cache_groups):
454 - is_full_attn = self.blocks_per_sw[gi] == 0
455 - if is_full_attn:
568 + if isinstance(group.kv_cache_spec, FullAttentionSpec):
569 self._full_attn_group_idx = gi
570 self._full_attn_block_size = getattr(
571 group.kv_cache_spec, "block_size", self.block_size
@@ -533,7 +646,11 @@ class MoRIIOConnectorScheduler:
646 - sliding window: keep only the last window size blocks
647 - full attention: keep all blocks
648 """
536 - if len(block_ids) == 0 or not self._is_hma_required:
649 + if len(block_ids) == 0:
650 + return []
651 + if len(block_ids) == len(self.kv_cache_config.kv_cache_groups):
652 + block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)
653 + if not self._is_hma_required:
654 return [list(blocks) for blocks in block_ids]
655 return [
656 blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
@@ -595,21 +712,75 @@ class MoRIIOConnectorScheduler:
712
713 """
714 if self.is_producer:
715 + # The P-side prompt truncation for hybrid models happens in
716 + # on_new_request, before the scheduler looks the request up in the
717 + # prefix cache.
718 return 0, False
719
600 - token_ids = request.prompt_token_ids or []
720 + num_prompt_tokens = request.num_prompt_tokens
721 + num_external_tokens = max(num_prompt_tokens - num_computed_tokens, 0)
722 if self.mode == MoRIIOMode.WRITE:
602 - # MoriiO in write mode, no remote prefill
723 + # MoriiO in write mode, no remote prefill. Hybrid models never get
724 + # here: register_kv_caches refuses WRITE mode for them, so there is
725 + # no recurrent-state accounting to do.
726 + return num_external_tokens, True
727
604 - return len(token_ids) - num_computed_tokens, True
728 + # READ mode always recomputes the last token locally on the decoder.
729 + #
730 + # The second element declares who waits for the KV, not whether it has
731 + # arrived. False keeps the wait on this side: the forward runs in the
732 + # same step, and the connector blocks in wait_for_layer_load (or in
733 + # start_load_kv under full-graph capture) until the KV is really there.
734 + return max(num_external_tokens - 1, 0), False
735 +
736 + def on_new_request(self, request: "Request") -> None:
737 + """Scheduler hook fired when a request is admitted.
738 +
739 + The hybrid P-side truncation must land here rather than in
740 + ``get_num_new_matched_tokens``: the scheduler runs the local prefix
741 + cache lookup first and only then asks the connector, so truncating
742 + there shortens a prompt whose cache hit was already measured against
743 + the full length. On a full hit that leaves ``num_computed_tokens ==
744 + num_tokens`` and trips the scheduler's ``num_new_tokens > 0`` assert.
745 + """
746 + if not (self.is_producer and self._has_mamba):
747 + return
748 + params = request.kv_transfer_params
749 + if params is not None and params.get("do_remote_decode"):
750 + self._truncate_mamba_request_for_prefill(request)
751
606 - return len(token_ids) - 1 - num_computed_tokens, False
752 + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
753 + """P-side only: drop the last prompt token so the prefiller computes
754 + h(N-1) instead of h(N).
755 +
756 + The decoder recomputes the last token locally to derive h(N) correctly
757 + (see the READ len-1 accounting in get_num_new_matched_tokens). Guarded
758 + by ``_p_side_truncated`` to avoid repeated truncation across a
759 + preemption/reschedule. Mirrors the NIXL scheduler.
760 + """
761 + params = request.kv_transfer_params
762 + if (
763 + params is not None
764 + and not params.get("_p_side_truncated")
765 + and request.num_prompt_tokens > 1
766 + ):
767 + if request.prompt_token_ids is not None:
768 + request.prompt_token_ids.pop()
769 + elif request.prompt_embeds is not None:
770 + request.prompt_embeds = request.prompt_embeds[:-1]
771 + else:
772 + return
773 +
774 + request._all_token_ids.pop()
775 + request.num_prompt_tokens -= 1
776 + request.max_tokens = 1
777 + params["_p_side_truncated"] = True
778
779 def send_notify_block(
780 self,
781 req_id: ReqId,
782 transfer_id: TransferId,
612 - block_notify_list: list[int],
783 + block_notify_list: list,
784 host=None,
785 port=None,
786 ):
@@ -739,7 +910,7 @@ class MoRIIOConnectorScheduler:
910 params.setdefault("transfer_id", transfer_id)
911 request_id = request.request_id
912 self.map_request_id(request_id, transfer_id)
742 - if params.get("do_remote_decode"):
913 + if params.get("do_remote_decode") and self.mode == MoRIIOMode.WRITE:
914 local_block_ids: BlockIds = blocks.get_block_ids()
915 self._reqs_need_save[request.request_id] = (request, local_block_ids)
916 # Snapshot params now so chunked-prefill build_connector_meta
@@ -753,24 +924,88 @@ class MoRIIOConnectorScheduler:
924 # remote_engine_id is returned by the prefill's request_finished.
925 # host/ports come from the request_id (parsed in add_new_req).
926 if "remote_engine_id" in params:
927 + adjusted_remote_block_ids = remote_block_ids
928 if num_external_tokens > 0:
757 - # Get unhashed blocks to pull from remote.
758 - local_block_ids = self.get_exchange_clipped_blocks(
759 - blocks.get_block_ids()
760 - )
761 - # Partial cache hits can lead to fewer local blocks vs
762 - # remote blocks, but never more
763 - assert all(
764 - len(local) <= len(remote)
765 - for local, remote in zip(
766 - local_block_ids, remote_block_ids, strict=True
929 + if self._has_mamba:
930 + # READ is synchronous, so allocate_slots may hash
931 + # destination blocks before this hook runs. Pair
932 + # the attention and recurrent-state suffixes
933 + # independently.
934 + remote_attn = list(remote_block_ids[0])
935 + remote_mamba_groups = [
936 + list(group) for group in remote_block_ids[1:]
937 + ]
938 + attn_block_ids, mamba_block_groups = (
939 + self.split_block_groups(blocks.get_block_ids())
940 )
768 - )
941 + if len(mamba_block_groups) != len(remote_mamba_groups):
942 + raise MoRIIOError(
943 + "MoRIIO hybrid READ Mamba group "
944 + "count mismatch: "
945 + f"local={len(mamba_block_groups)}, "
946 + f"remote={len(remote_mamba_groups)}"
947 + )
948 + local_attn = attn_block_ids
949 + local_attn, remote_attn = self._align_read_blocks(
950 + local_attn,
951 + remote_attn,
952 + self._max_decode_tail_blocks,
953 + )
954 + aligned_local_mamba: list[list[int]] = []
955 + aligned_remote_mamba: list[list[int]] = []
956 + for local_group, remote_group in zip(
957 + mamba_block_groups,
958 + remote_mamba_groups,
959 + strict=True,
960 + ):
961 + aligned_local, aligned_remote = (
962 + self._align_read_blocks(
963 + local_group, remote_group, 1
964 + )
965 + )
966 + aligned_local_mamba.append(aligned_local)
967 + aligned_remote_mamba.append(aligned_remote)
968 + local_block_ids = [local_attn, *aligned_local_mamba]
969 + adjusted_remote_block_ids = [
970 + remote_attn,
971 + *aligned_remote_mamba,
972 + ]
973 + else:
974 + # Preserve main's per-group HMA representation
975 + # for attention-only hybrid models while
976 + # accounting for the locally recomputed tail.
977 + local_block_ids = self.get_exchange_clipped_blocks(
978 + blocks.get_block_ids()
979 + )
980 + adjusted_remote_block_ids = []
981 + aligned_local_block_ids: list[list[int]] = []
982 + for local_group, remote_group in zip(
983 + local_block_ids,
984 + remote_block_ids,
985 + strict=True,
986 + ):
987 + aligned_local, aligned_remote = (
988 + self._align_read_blocks(
989 + local_group,
990 + remote_group,
991 + self._max_decode_tail_blocks,
992 + )
993 + )
994 + aligned_local_block_ids.append(aligned_local)
995 + adjusted_remote_block_ids.append(aligned_remote)
996 + local_block_ids = aligned_local_block_ids
997 else:
770 - # If remote_blocks and num_external_tokens = 0, we have
771 - # a full prefix cache hit on the D worker. We need to call
772 - # send_notify in _read_blocks to free the memory on the P.
773 - local_block_ids = [[] for _ in remote_block_ids]
998 + if self._has_mamba:
999 + # Attention can be a complete local hit, but the
1000 + # recurrent state is never prefix-cacheable.
1001 + _, mamba_block_groups = self.split_block_groups(
1002 + blocks.get_block_ids()
1003 + )
1004 + local_block_ids = [[], *mamba_block_groups]
1005 + else:
1006 + # Nothing to pull, but retain one empty list per
1007 + # cache group so the worker can notify P.
1008 + local_block_ids = [[] for _ in remote_block_ids]
1009
1010 self._reqs_need_recv[request.request_id] = (
1011 request,
@@ -778,7 +1013,10 @@ class MoRIIOConnectorScheduler:
1013 )
1014 # Snapshot params for chunked prefill consumption
1015 # in build_connector_meta (see comment above).
781 - self._req_kv_params[request.request_id] = dict(params)
1016 + saved_params = dict(params)
1017 + if num_external_tokens > 0:
1018 + saved_params["remote_block_ids"] = adjusted_remote_block_ids
1019 + self._req_kv_params[request.request_id] = saved_params
1020 else:
1021 logger.warning(
1022 "Got invalid KVTransferParams: %s. This "
@@ -1003,13 +1241,78 @@ class MoRIIOConnectorScheduler:
1241 logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
1242 self.paths.clear()
1243
1244 + def split_block_groups(
1245 + self, block_ids: list[list[int]] | tuple[list[int], ...]
1246 + ) -> tuple[list[int], list[list[int]]]:
1247 + """Select transferable attention and Mamba block groups.
1248 +
1249 + The wire payload stores the attention group first, followed by every
1250 + Mamba group in transfer-group order. Outside ``mamba_cache_mode="all"``,
1251 + only each group's running-state slot is transferred.
1252 + """
1253 + if not block_ids:
1254 + return [], []
1255 + if not self._has_mamba:
1256 + raise MoRIIOError("Mamba block groups requested for a non-hybrid model")
1257 + transfer_block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)
1258 + attn = list(transfer_block_ids[self._attn_group_ids[0]])
1259 + mamba_groups = [
1260 + self._clip_mamba_group(list(transfer_block_ids[group_id]))
1261 + for group_id in self._mamba_group_ids
1262 + ]
1263 + return attn, mamba_groups
1264 +
1265 + def _clip_mamba_group(self, blocks: list[int]) -> list[int]:
1266 + """Keep only the state-bearing slots of one mamba kv cache group."""
1267 + if not blocks:
1268 + return blocks
1269 + if not self._ssm_state_slots_are_positional:
1270 + # Single running state: everything before it is a null placeholder
1271 + # or the previous step's superseded state.
1272 + blocks = blocks[-1:]
1273 + return blocks
1274 +
1275 + @staticmethod
1276 + def _align_read_blocks(
1277 + local_block_ids: list[int],
1278 + remote_block_ids: list[int],
1279 + max_decode_tail_blocks: int,
1280 + ) -> tuple[list[int], list[int]]:
1281 + """Pair a decode suffix with the corresponding prefill blocks.
1282 +
1283 + Prefix hits can make the decode's unhashed suffix shorter than the
1284 + prefill's full list, so read from the matching remote suffix. The
1285 + decoder can instead have one extra tail block for the locally
1286 + recomputed final token; that block has no remote source and is not
1287 + transferred.
1288 + """
1289 + if len(local_block_ids) > len(remote_block_ids):
1290 + extra_decode_blocks = len(local_block_ids) - len(remote_block_ids)
1291 + if extra_decode_blocks > max_decode_tail_blocks:
1292 + raise ValueError(
1293 + "decode block ids exceed prefill block ids beyond the "
1294 + f"allowed local tail ({max_decode_tail_blocks} blocks): "
1295 + f"{len(local_block_ids)} > {len(remote_block_ids)}"
1296 + )
1297 + local_block_ids = local_block_ids[: len(remote_block_ids)]
1298 + elif len(local_block_ids) < len(remote_block_ids):
1299 + remote_block_ids = (
1300 + remote_block_ids[-len(local_block_ids) :] if local_block_ids else []
1301 + )
1302 + return local_block_ids, remote_block_ids
1303 +
1304 def request_finished(
1305 self,
1306 request: "Request",
1009 - block_ids: BlockIds,
1307 + block_ids: BlockIds | list[int],
1308 + mamba_block_groups: list[list[int]] | None = None,
1309 ) -> tuple[bool, dict[str, Any] | None]:
1310 """Once a request is finished, determine whether request blocks
1311 should be freed now or will be sent asynchronously and freed later.
1312 +
1313 + ``block_ids`` normally contains all cache groups. Hybrid callers from
1314 + ``request_finished_all_groups`` pass a flat attention list and the
1315 + recurrent-state groups separately in ``mamba_block_groups``.
1316 """
1317 request_id = request.request_id
1318 params = request.kv_transfer_params
@@ -1042,8 +1345,25 @@ class MoRIIOConnectorScheduler:
1345 if self.mode == MoRIIOMode.WRITE:
1346 self._release_write_prefill_blocks(request.request_id, params)
1347 else:
1045 - empty_groups: BlockIds = [[] for _ in self.blocks_per_sw]
1046 - self._reqs_need_recv[request.request_id] = (request, empty_groups)
1348 + # Carry the actually-allocated blocks (incl. the per-request KDA
1349 + # mamba slot) instead of []: a do_remote_prefill request finishing
1350 + # before update_state_after_alloc still has blocks, and dropping the
1351 + # mamba slot here would re-trip the KDA guard in _read_blocks.
1352 + if mamba_block_groups:
1353 + recv_blocks: BlockIds = [
1354 + list(cast(list[int], block_ids)),
1355 + *[list(group) for group in mamba_block_groups],
1356 + ]
1357 + elif block_ids and isinstance(block_ids[0], int):
1358 + recv_blocks = [list(cast(list[int], block_ids))]
1359 + else:
1360 + grouped_block_ids = cast(BlockIds, block_ids)
1361 + recv_blocks = (
1362 + [list(group) for group in grouped_block_ids]
1363 + if self._has_mamba
1364 + else [[] for _ in self.blocks_per_sw]
1365 + )
1366 + self._reqs_need_recv[request.request_id] = (request, recv_blocks)
1367 params["do_remote_prefill"] = False
1368 return False, None
1369
@@ -1054,7 +1374,17 @@ class MoRIIOConnectorScheduler:
1374 return False, None
1375
1376 # On producer
1057 - computed_block_ids = self.get_exchange_clipped_blocks(block_ids)
1377 + if mamba_block_groups is not None:
1378 + computed_block_ids: BlockIds = [
1379 + list(cast(list[int], block_ids)),
1380 + *[list(group) for group in mamba_block_groups],
1381 + ]
1382 + elif self._has_mamba:
1383 + raise MoRIIOError("Hybrid request completion requires all KV cache groups")
1384 + else:
1385 + computed_block_ids = self.get_exchange_clipped_blocks(
1386 + cast(BlockIds, block_ids)
1387 + )
1388 delay_free_blocks = any(len(group) > 0 for group in computed_block_ids)
1389
1390 if delay_free_blocks:
@@ -1209,10 +1539,14 @@ class MoRIIOConnectorWorker:
1539 self.is_producer = self.kv_transfer_config.is_kv_producer
1540 self.layer_to_spec = build_layer_to_spec(kv_cache_config)
1541 # we transfer by layer, but store by kv cache group
1212 - self.layer_to_group: dict[str, int] = {
1213 - layer_name: group_idx
1214 - for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
1215 - for layer_name in group.layer_names
1542 + self.layer_to_group = dict(kv_cache_config.transfer_group_index_by_layer)
1543 + self._transfer_layer_names = set(self.layer_to_group)
1544 + _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config)
1545 + self._num_mamba_transfer_groups = len(mamba_transfer_group_ids)
1546 + self._mamba_payload_index_by_layer = {
1547 + layer_name: payload_index
1548 + for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1)
1549 + for layer_name in kv_cache_config.transfer_groups[group_id].layer_names
1550 }
1551
1552 if self.is_producer:
@@ -1353,7 +1687,8 @@ class MoRIIOConnectorWorker:
1687 # Map of engine_id -> num_blocks. All ranks in the same deployment will
1688 # have the same number of blocks.
1689 self.dst_num_blocks: dict[EngineId, int] = {}
1356 - # In-progress READ transfers: req_id -> {layer_name: status}.
1690 + # In-progress READ transfers: req_id -> {layer_name: [statuses]}.
1691 + # A KDA layer posts two reads (conv + ssm), so a layer maps to a list.
1692 self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict)
1693 # Values are (remote_host, remote_notify_port, transfer_id).
1694 self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {}
@@ -1361,6 +1696,18 @@ class MoRIIOConnectorWorker:
1696 # Used by _pop_done_transfers to abort transfers whose RDMA
1697 # completion is lost, instead of hanging forever.
1698 self._recving_transfers_start: dict[str, float] = {}
1699 + # Destination attention blocks per in-flight recv. Failed non-HMA reads
1700 + # report these to the scheduler's single-group recovery path.
1701 + self._recving_local_blocks: dict[ReqId, list[int]] = {}
1702 + # Drained by get_block_ids_with_load_errors.
1703 + self._invalid_block_ids: set[int] = set()
1704 + # This step's posted reads, awaited before the forward when the
1705 + # per-layer barrier cannot run.
1706 + self._reads_issued_this_step: list = []
1707 + # This step's posted recurrent-state reads. Mamba layers are not
1708 + # wrapped by the attention KV-transfer decorator, so wait_for_layer_load
1709 + # never fires for them and these must be awaited in start_load_kv.
1710 + self._mamba_reads_this_step: list = []
1711
1712 # Track the expiration time of requests that are waiting to be sent.
1713 self._reqs_to_send: dict[ReqId, float] = {}
@@ -1386,6 +1733,14 @@ class MoRIIOConnectorWorker:
1733
1734 self.block_window_per_layer: list[int | None] = []
1735 self.use_mla = self.model_config.use_mla
1736 + # Hybrid (mamba/KDA) recurrent-state transfer. Populated in
1737 + # register_kv_caches when the model has a MambaSpec kv_cache group.
1738 + self._has_mamba = False
1739 + self._conv_decomp: MambaConvSplitInfo | None = None
1740 + self._mamba_offset_templates: dict[str, MambaOffsetTemplate] = {}
1741 + # layer_name -> flat session indices (one per registered region:
1742 + # 1 for attention, 2 (conv, ssm) for a KDA layer). Built once lazily.
1743 + self._region_session_index: dict[str, list[int]] | None = None
1744 self.built_session = False
1745 self.built_write_session: defaultdict[str, list] = defaultdict(list)
1746 backend = get_attn_backend(
@@ -1460,23 +1815,39 @@ class MoRIIOConnectorWorker:
1815
1816 def _get_built_session(self, remote_engine_id):
1817 if remote_engine_id not in self.built_write_session:
1818 + remote_layer_map = self.layer_name_to_remote_kv_cache_metadata[
1819 + remote_engine_id
1820 + ]
1821 cur_remote_engine_sessions = []
1464 - for ln, local_meta in self.layer_name_to_local_kv_cache_metadata.items():
1465 - unpacked_local_memory_meta = (
1466 - self.moriio_wrapper.get_unpack_memory_metadata(local_meta[0])
1467 - )
1468 - unpacked_remote_memory_meta = (
1469 - self.moriio_wrapper.get_unpack_memory_metadata(
1470 - self.layer_name_to_remote_kv_cache_metadata[remote_engine_id][
1471 - ln
1472 - ][0]
1822 + for ln, local_metas in self.layer_name_to_local_kv_cache_metadata.items():
1823 + remote_metas = remote_layer_map[ln]
1824 + if len(remote_metas) != len(local_metas):
1825 + # Sessions are addressed by region index, so a peer that
1826 + # registered a different number of regions for this layer
1827 + # would not just lose a region to zip(): every later index
1828 + # would point into the wrong layer's session.
1829 + raise MoRIIOError(
1830 + f"peer {remote_engine_id} registered "
1831 + f"{len(remote_metas)} region(s) for layer {ln}, this "
1832 + f"engine has {len(local_metas)}; the two sides must run "
1833 + "the same model and KV cache layout"
1834 )
1474 - )
1475 - cur_remote_engine_sessions.append(
1476 - self.moriio_wrapper.build_session(
1477 - unpacked_local_memory_meta, unpacked_remote_memory_meta
1835 + # One session per registered region, in registration order:
1836 + # attention layers have a single region; KDA layers have two
1837 + # (conv then ssm). This flat layout is indexed via
1838 + # _region_session_indices(layer_name).
1839 + for local_meta, remote_meta in zip(local_metas, remote_metas):
1840 + unpacked_local_memory_meta = (
1841 + self.moriio_wrapper.get_unpack_memory_metadata(local_meta)
1842 + )
1843 + unpacked_remote_memory_meta = (
1844 + self.moriio_wrapper.get_unpack_memory_metadata(remote_meta)
1845 + )
1846 + cur_remote_engine_sessions.append(
1847 + self.moriio_wrapper.build_session(
1848 + unpacked_local_memory_meta, unpacked_remote_memory_meta
1849 + )
1850 )
1479 - )
1851 self.built_write_session[remote_engine_id] = cur_remote_engine_sessions
1852 return self.built_write_session[remote_engine_id], self.remote_moriio_metadata[
1853 remote_engine_id
@@ -1788,15 +2159,60 @@ class MoRIIOConnectorWorker:
2159 def _is_mla_cache_layer(self, layer_name: str) -> bool:
2160 return is_mla_cache_layer(self.layer_to_spec, layer_name)
2161
2162 + def _is_mamba_layer(self, layer_name: str) -> bool:
2163 + """True for a hybrid/KDA layer backed by a packed state page."""
2164 + return isinstance(self.layer_to_spec.get(layer_name), MambaSpec)
2165 +
2166 + @staticmethod
2167 + def _contiguous_byte_alias(tensor: torch.Tensor) -> torch.Tensor:
2168 + """Return a contiguous uint8 view over ``tensor``'s full byte extent.
2169 +
2170 + KDA conv/ssm caches are slot-strided views (``stride(0)`` may exceed the
2171 + per-slot element count), so they are non-contiguous and cannot be passed
2172 + to ``register_torch_tensor``. This aliases the same storage (zero-copy,
2173 + identical ``data_ptr``) as a flat uint8 tensor spanning
2174 + ``shape[0] * stride(0)`` elements -- the full slot-strided extent that
2175 + transfer offsets (``slot * stride(0) * elem``) address.
2176 + """
2177 + esz = tensor.element_size()
2178 + storage_nbytes = tensor.untyped_storage().nbytes()
2179 + offset_bytes = tensor.storage_offset() * esz
2180 + extent_bytes = tensor.shape[0] * tensor.stride(0) * esz
2181 + # Clamp to the bytes remaining from the view's storage offset. In a
2182 + # packed page, the ssm view starts after conv, so its slot-strided extent
2183 + # would otherwise run past the backing allocation.
2184 + extent_bytes = min(extent_bytes, storage_nbytes - offset_bytes)
2185 + assert extent_bytes > 0, (
2186 + "KDA byte alias empty: storage_offset="
2187 + f"{tensor.storage_offset()} * esz={esz} >= "
2188 + f"storage_nbytes={storage_nbytes}"
2189 + )
2190 + alias = torch.empty(0, dtype=torch.uint8, device=tensor.device)
2191 + alias.set_(
2192 + tensor.untyped_storage(),
2193 + tensor.storage_offset() * esz,
2194 + (extent_bytes,),
2195 + (1,),
2196 + )
2197 + return alias
2198 +
2199 def _get_layer_transfer_geometry(
2200 self, layer_name: str, remote_num_blocks: int | None = None
2201 ) -> LayerTransferGeometry:
1794 - return get_layer_transfer_geometry(
2202 + geometry = get_layer_transfer_geometry(
2203 layer_name,
2204 self.kv_caches[layer_name],
2205 self.layer_to_spec,
2206 remote_num_blocks,
2207 )
2208 + # get_layer_transfer_geometry also serves mamba layers, whose geometry
2209 + # is a different NamedTuple with no block fields. Every caller here is
2210 + # on the attention path (mamba layers return before reaching it), so
2211 + # narrow the type instead of leaking the union.
2212 + assert isinstance(geometry, LayerTransferGeometry), (
2213 + f"layer {layer_name} is a mamba layer; use get_mamba_transfer_geometry"
2214 + )
2215 + return geometry
2216
2217 def _iter_layer_registration_regions(
2218 self, layer_name: str
@@ -1853,22 +2269,71 @@ class MoRIIOConnectorWorker:
2269
2270 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
2271 """Register the KV Cache data in moriio."""
2272 + kv_caches = {
2273 + layer_name: kv_cache
2274 + for layer_name, kv_cache in kv_caches.items()
2275 + if layer_name in self._transfer_layer_names
2276 + }
2277 + if not kv_caches:
2278 + raise MoRIIOError("MoRIIO has no transfer-enabled KV cache layers")
2279 self.kv_caches = kv_caches # layer name to kv cache
2280 + self._has_mamba = any(
2281 + self._is_mamba_layer(layer_name) for layer_name in kv_caches
2282 + )
2283 + if self._has_mamba and self.mode == MoRIIOMode.WRITE:
2284 + # WRITE needs a distinct conv/ssm transfer plan. Reject before
2285 + # allocating or registering any regions rather than silently
2286 + # serving a zero recurrent state.
2287 + raise MoRIIOError(
2288 + "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
2289 + "mode only; set kv_connector_extra_config.read_mode=true "
2290 + "for hybrid models"
2291 + )
2292 + if self._has_mamba:
2293 + mamba_specs = [
2294 + cast(MambaSpec, self.layer_to_spec[layer_name])
2295 + for layer_name in kv_caches
2296 + if self._is_mamba_layer(layer_name)
2297 + ]
2298 + mamba_spec = _validate_mamba_specs(mamba_specs)
2299 + assert mamba_spec is not None
2300 +
2301 + from vllm.model_executor.layers.mamba.mamba_utils import (
2302 + is_conv_state_dim_first,
2303 + )
2304 +
2305 + if not is_conv_state_dim_first():
2306 + raise MoRIIOError(
2307 + "MoRIIO KDA conv transfer requires the DS (dim-first) "
2308 + "state layout; set VLLM_SSM_CONV_STATE_LAYOUT=DS"
2309 + )
2310 + self._conv_decomp = derive_mamba_conv_split(mamba_spec, self.world_size)
2311 + # KDA layers use packed pages with different geometry; only record
2312 + # attention tensor shapes here.
2313 self.kv_cache_shapes = {
1858 - layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items()
2314 + layer_name: kv_cache.shape
2315 + for layer_name, kv_cache in kv_caches.items()
2316 + if not self._is_mamba_layer(layer_name)
2317 }
2318
2319 + # Geometry selection must pick an attention layer (never a KDA tuple):
2320 + # prefer a standard 5-D K/V layer, else the first attention layer.
2321 + attn_items = [
2322 + (layer_name, kv_cache)
2323 + for layer_name, kv_cache in kv_caches.items()
2324 + if not self._is_mamba_layer(layer_name)
2325 + ]
2326 first_layer_name, first_kv_cache = next(
2327 (
2328 (layer_name, kv_cache)
1864 - for layer_name, kv_cache in kv_caches.items()
2329 + for layer_name, kv_cache in attn_items
2330 if (
2331 not self._is_mla_cache_layer(layer_name)
2332 and len(kv_cache.shape) == 5
2333 and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2)
2334 )
2335 ),
1871 - next(iter(kv_caches.items())),
2336 + attn_items[0] if attn_items else next(iter(kv_caches.items())),
2337 )
2338 kv_elem_size = first_kv_cache.element_size()
2339
@@ -1915,6 +2380,17 @@ class MoRIIOConnectorWorker:
2380 )
2381 for layer_name in kv_caches:
2382 self.layer_base_addr_index[layer_name] = base_addr_idx
2383 + if self._is_mamba_layer(layer_name):
2384 + # KDA layer: two whole-tensor regions (conv, ssm); no block-size
2385 + # geometry (recurrent state is per-slot, not paged).
2386 + for cache, region_len in self._iter_layer_registration_regions(
2387 + layer_name
2388 + ):
2389 + base_addr = cache.data_ptr()
2390 + caches_data.append((base_addr, region_len, cache.device.index, ""))
2391 + kv_caches_base_addr.append(base_addr)
2392 + base_addr_idx += 1
2393 + continue
2394 geometry = self._get_layer_transfer_geometry(layer_name)
2395 if hma_enabled:
2396 if geometry.block_len != self.block_len:
@@ -1941,31 +2417,87 @@ class MoRIIOConnectorWorker:
2417 kv_caches_base_addr.append(base_addr)
2418 base_addr_idx += 1
2419
1944 - shared_backing = len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1
2420 + shared_backing = (
2421 + all(isinstance(t, torch.Tensor) for t in kv_caches.values())
2422 + and len({id(t.untyped_storage()) for t in kv_caches.values()}) == 1
2423 + )
2424 shared_mr = None
2425 if shared_backing:
2426 reg_tensor, self.kv_layer_mr_offset = self._build_shared_kv_mr(kv_caches)
2427 shared_mr = self.moriio_wrapper.register_local_tensor(reg_tensor)
2428
2429 + # Say it out loud when the groups disagree. This is expected on a hybrid
2430 + # model and used to be fatal, so a run that silently had one page size
2431 + # and a run that has three should not look identical in the log.
2432 + distinct_block_lens = sorted(set(self.block_lens.values()))
2433 + if len(distinct_block_lens) > 1:
2434 + logger.info(
2435 + "MoRIIO registered %d attention layers with %d distinct block "
2436 + "lengths (bytes): %s; transfer sizes and strides are per layer, "
2437 + "and %d is the page advertised to the peer.",
2438 + len(self.block_lens),
2439 + len(distinct_block_lens),
2440 + distinct_block_lens,
2441 + self.block_len,
2442 + )
2443 +
2444 for layer_name, kv_cache in kv_caches.items():
2445 if layer_name not in self.layer_name_to_local_kv_cache_metadata:
2446 self.layer_name_to_local_kv_cache_metadata[layer_name] = []
2447
1954 - moriio_mem_metadata = shared_mr or (
1955 - self.moriio_wrapper.register_local_tensor(kv_cache)
1956 - )
1957 - self.layer_name_to_local_kv_cache_metadata[layer_name].append(
1958 - moriio_mem_metadata
1959 - )
1960 -
1961 - self.local_kv_cache_size.append(
1962 - kv_cache.nelement() * kv_cache.element_size()
1963 - )
2448 + if self._is_mamba_layer(layer_name):
2449 + # Register conv and ssm as two whole-tensor regions, appending
2450 + # both metadata blobs (order: conv then ssm) so the session
2451 + # layout matches iter_layer_registration_regions. conv/ssm are
2452 + # slot-strided views and are not contiguous, which
2453 + # register_torch_tensor rejects; register a contiguous uint8
2454 + # alias over each tensor's full byte extent instead (zero-copy:
2455 + # same storage and data_ptr).
2456 + conv, ssm = kda_conv_ssm(
2457 + kv_cache, cast(MambaSpec, self.layer_to_spec[layer_name])
2458 + )
2459 + logger.debug(
2460 + "MoRIIO KDA register %s: conv shape=%s stride=%s contig=%s "
2461 + "storage_nbytes=%s ; ssm shape=%s stride=%s contig=%s "
2462 + "storage_nbytes=%s",
2463 + layer_name,
2464 + tuple(conv.shape),
2465 + tuple(conv.stride()),
2466 + conv.is_contiguous(),
2467 + conv.untyped_storage().nbytes(),
2468 + tuple(ssm.shape),
2469 + tuple(ssm.stride()),
2470 + ssm.is_contiguous(),
2471 + ssm.untyped_storage().nbytes(),
2472 + )
2473 + for tensor in (conv, ssm):
2474 + alias = self._contiguous_byte_alias(tensor)
2475 + meta = self.moriio_wrapper.register_local_tensor(alias)
2476 + self.layer_name_to_local_kv_cache_metadata[layer_name].append(meta)
2477 + self.local_kv_cache_size.append(alias.numel())
2478 + else:
2479 + moriio_mem_metadata = shared_mr or (
2480 + self.moriio_wrapper.register_local_tensor(kv_cache)
2481 + )
2482 + self.layer_name_to_local_kv_cache_metadata[layer_name].append(
2483 + moriio_mem_metadata
2484 + )
2485 + self.local_kv_cache_size.append(
2486 + kv_cache.nelement() * kv_cache.element_size()
2487 + )
2488
2489 self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
2490 self.num_regions = len(caches_data)
2491 self.num_layers = len(self.kv_caches.keys())
2492
2493 + # Report the prevalidated KDA decomposition after registration.
2494 + if self._has_mamba:
2495 + assert self._conv_decomp is not None
2496 + logger.info(
2497 + "MoRIIO registered %d KDA (conv+ssm) layer(s); ssm sizes=%s",
2498 + sum(1 for ln in kv_caches if self._is_mamba_layer(ln)),
2499 + self._conv_decomp.ssm_sizes,
2500 + )
2501 # Optimization for models with local attention (Llama 4)
2502 if self.vllm_config.model_config.hf_config.model_type == "llama4":
2503 from transformers import Llama4TextConfig
@@ -2101,49 +2633,49 @@ class MoRIIOConnectorWorker:
2633 if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
2634 return
2635
2104 - deadline = time.monotonic() + self.moriio_config.transfer_timeout
2105 - while True:
2106 - with self.moriio_wrapper.lock:
2107 - pending = [
2108 - status_by_layer[layer_name]
2109 - for status_by_layer in self._recving_transfers.values()
2110 - if layer_name in status_by_layer
2111 - ]
2112 -
2113 - if not pending:
2114 - return
2115 -
2116 - still_running = False
2117 - for status in pending:
2118 - # A failed read is dropped in _pop_done_transfers.
2119 - if status.Succeeded() or status.Failed():
2120 - continue
2121 - still_running = True
2636 + with self.moriio_wrapper.lock:
2637 + pending = [
2638 + status
2639 + for status_by_layer in self._recving_transfers.values()
2640 + for status in status_by_layer.get(layer_name, ())
2641 + ]
2642
2123 - if not still_running:
2124 - return
2643 + if not pending:
2644 + return
2645
2126 - if time.monotonic() > deadline:
2127 - logger.warning(
2128 - "MoRIIO READ barrier timed out for layer %s; proceeding "
2129 - "(request dropped via get_finished).",
2646 + try:
2647 + self.moriio_wrapper.waiting_for_transfer_complete(pending)
2648 + except TransferError:
2649 + if self._has_mamba:
2650 + logger.exception(
2651 + "MoRIIO hybrid READ failed before layer %s; aborting the "
2652 + "forward because HMA load-error recovery is unsupported.",
2653 layer_name,
2654 )
2132 - return
2133 -
2134 - time.sleep(0.001)
2655 + raise
2656 + logger.warning(
2657 + "MoRIIO READ barrier did not complete for layer %s; proceeding "
2658 + "to transfer cleanup.",
2659 + layer_name,
2660 + exc_info=True,
2661 + )
2662
2663 def _pop_done_transfers(self) -> set[str]:
2664 done_req_ids: set[str] = set()
2138 - _xfer_timeout = int(os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S", "120"))
2665 + _xfer_timeout = self.moriio_config.recv_abort_timeout
2666 with self.moriio_wrapper.lock:
2667 to_remove = []
2668 for req_id, status_by_layer in self._recving_transfers.items():
2142 - statuses = list(status_by_layer.values())
2143 - failed_status = next(
2144 - (status for status in statuses if status.Failed()), None
2145 - )
2146 - if statuses and all(status.Succeeded() for status in statuses):
2669 + # Every layer's read must land before the request is complete,
2670 + # so the whole batch is polled rather than just the newest
2671 + # status.
2672 + statuses = [
2673 + status
2674 + for layer_statuses in status_by_layer.values()
2675 + for status in layer_statuses
2676 + ]
2677 + state = self.moriio_wrapper.poll_transfer_batch(statuses)
2678 + if statuses and state is TransferBatchState.DONE:
2679 host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
2680 done_req_ids.add(xfer_id)
2681 self.moriio_wrapper.send_notify(
@@ -2154,15 +2686,19 @@ class MoRIIOConnectorWorker:
2686 message_fields={"consumer_tp_size": self.world_size},
2687 )
2688 to_remove.append(req_id)
2157 - elif failed_status is not None:
2689 + elif state is TransferBatchState.FAILED:
2690 + failed_status = next(
2691 + (status for status in statuses if status.Failed()), None
2692 + )
2693 logger.error(
2694 "RDMA transfer failed for request %s: %s (code=%s). "
2160 - "Notifying prefill to free blocks; request will be "
2161 - "aborted by timeout.",
2695 + "Notifying prefill to free blocks and recording the "
2696 + "load error where scheduler recovery is supported.",
2697 req_id,
2163 - failed_status.Message(),
2164 - failed_status.Code(),
2698 + failed_status.Message() if failed_status else "unknown",
2699 + failed_status.Code() if failed_status else "unknown",
2700 )
2701 + self._record_failed_recv(req_id)
2702 host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
2703 try:
2704 self.moriio_wrapper.send_notify(
@@ -2178,8 +2714,9 @@ class MoRIIOConnectorWorker:
2714 req_id,
2715 )
2716 to_remove.append(req_id)
2181 - # Do NOT add to done_req_ids: decode KV cache is incomplete.
2182 - # The request will expire via the normal request timeout.
2717 + # Deliberately not in done_req_ids: the decode KV is
2718 + # incomplete, so this is a load error rather than a
2719 + # completion.
2720 elif req_id in self._recving_transfers_start:
2721 # Abort still-in-flight transfers that exceed the
2722 # configured deadline. Otherwise a lost RDMA
@@ -2189,19 +2726,84 @@ class MoRIIOConnectorWorker:
2726 if _age > _xfer_timeout:
2727 logger.error(
2728 "RDMA read TIMED OUT for req %s after %.1fs "
2192 - "(VLLM_MORIIO_TRANSFER_TIMEOUT_S=%d)",
2729 + "(kv_connector_extra_config.recv_abort_timeout=%.0f)",
2730 req_id,
2731 _age,
2732 _xfer_timeout,
2733 )
2734 + self._record_failed_recv(req_id)
2735 to_remove.append(req_id)
2736 for req_id in to_remove:
2737 del self._recving_transfers[req_id]
2738 del self._recving_transfers_callback_addr[req_id]
2739 self._recving_transfers_start.pop(req_id, None)
2740 + self._recving_local_blocks.pop(req_id, None)
2741
2742 return done_req_ids
2743
2744 + def _await_reads_issued_this_step(self) -> None:
2745 + """Step-level READ barrier, for the reads no per-layer barrier covers.
2746 +
2747 + load_kv_async=False promises the KV is in place before the forward runs.
2748 + wait_for_layer_load keeps that promise per layer, letting early layers
2749 + overlap later transfers, but it only fires for layers that go through
2750 + the attention KV-transfer decorator. Two cases fall through it:
2751 +
2752 + * mamba/KDA layers have no such hook at all, so their conv+ssm reads
2753 + would otherwise be consumed while still in flight;
2754 + * under CUDAGraphMode.FULL a host-side blocking wait is illegal
2755 + anywhere in the graph, so every read is awaited here instead.
2756 +
2757 + start_load_kv runs outside the graph, so blocking here is safe.
2758 + Attention reads keep their per-layer overlap in the non-FULL case.
2759 +
2760 + Attention-only failures are reported via
2761 + get_block_ids_with_load_errors. Hybrid failures are re-raised because
2762 + the scheduler's recovery path does not support multiple cache groups;
2763 + this fails the step closed instead of running forward with incomplete
2764 + recurrent state.
2765 + """
2766 + all_statuses = self._reads_issued_this_step
2767 + self._reads_issued_this_step = []
2768 + mamba_statuses = self._mamba_reads_this_step
2769 + self._mamba_reads_this_step = []
2770 + if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
2771 + statuses = all_statuses
2772 + else:
2773 + statuses = mamba_statuses
2774 + if not statuses:
2775 + return
2776 + try:
2777 + self.moriio_wrapper.waiting_for_transfer_complete(statuses)
2778 + except TransferError:
2779 + if self._has_mamba:
2780 + logger.exception(
2781 + "MoRIIO hybrid READ failed before the forward; aborting "
2782 + "because HMA load-error recovery is unsupported."
2783 + )
2784 + raise
2785 + logger.exception(
2786 + "MoRIIO reads did not complete before the forward; proceeding "
2787 + "to transfer cleanup."
2788 + )
2789 +
2790 + def _record_failed_recv(self, req_id: ReqId) -> None:
2791 + """Hand a failed attention-only read to scheduler recovery.
2792 +
2793 + The scheduler's invalid-block recovery currently assumes exactly one
2794 + KV cache group. Reporting IDs for a hybrid request would crash it while
2795 + unpacking the request's groups, so match NIXL and keep HMA failures out
2796 + of this API until scheduler recovery carries group identity.
2797 + """
2798 + if self._has_mamba:
2799 + return
2800 + self._invalid_block_ids.update(self._recving_local_blocks.get(req_id, ()))
2801 +
2802 + def get_block_ids_with_load_errors(self) -> set[int]:
2803 + """Drain the blocks whose read failed, for the scheduler to recompute."""
2804 + invalid, self._invalid_block_ids = self._invalid_block_ids, set()
2805 + return invalid
2806 +
2807 def save_kv_layer(
2808 self,
2809 metadata: MoRIIOConnectorMetadata,
@@ -2462,11 +3064,12 @@ class MoRIIOConnectorWorker:
3064
3065 continue
3066
2465 - # Handshake already completed, start async read xfer.
3067 + # Handshake already completed, start the read xfer.
3068 self._read_blocks_for_req(req_id, meta)
3069 # Start transfers for requests whose handshakes have now finished.
3070
3071 if remote_engine_id is None and not wait_handshake_readd_req:
3072 + self._await_reads_issued_this_step()
3073 return
3074 _deadline = time.monotonic() + self.moriio_config.transfer_timeout
3075 while True:
@@ -2493,6 +3096,7 @@ class MoRIIOConnectorWorker:
3096 else:
3097 break
3098
3099 + self._await_reads_issued_this_step()
3100 self._reqs_to_send.update(metadata.reqs_to_send)
3101
3102 def wait_for_save(self, metadata: MoRIIOConnectorMetadata):
@@ -2717,6 +3321,163 @@ class MoRIIOConnectorWorker:
3321 remote = [x + remote_off for x in remote]
3322 return local, remote, sizes
3323
3324 + def _region_session_indices(self, layer_name: str) -> list[int]:
3325 + """Flat session indices for a layer's registered regions.
3326 +
3327 + Sessions are built one per region in registration order (see
3328 + _get_built_session), so an attention layer maps to a single index and a
3329 + KDA layer to two indices ([conv, ssm]). Built once and cached.
3330 + """
3331 + if self._region_session_index is None:
3332 + mapping: dict[str, list[int]] = {}
3333 + idx = 0
3334 + for ln, metas in self.layer_name_to_local_kv_cache_metadata.items():
3335 + mapping[ln] = list(range(idx, idx + len(metas)))
3336 + idx += len(metas)
3337 + self._region_session_index = mapping
3338 + return self._region_session_index[layer_name]
3339 +
3340 + def _mamba_tp_ratio(self, remote_tp_size: int | None) -> int:
3341 + """Validate the homogeneous-TP KDA contract."""
3342 + remote = self.world_size if remote_tp_size is None else remote_tp_size
3343 + if remote != self.world_size:
3344 + raise NotImplementedError(
3345 + "heterogeneous-TP mamba/KDA transfer is not supported yet "
3346 + f"(local_tp={self.world_size}, remote_tp={remote}); use equal "
3347 + "prefill/decode TP for hybrid models"
3348 + )
3349 + return 1
3350 +
3351 + def _mamba_blocks_for_layer(
3352 + self,
3353 + layer_name: str,
3354 + block_groups: BlockIds,
3355 + ) -> list[int]:
3356 + payload_index = self._mamba_payload_index_by_layer.get(layer_name)
3357 + if payload_index is None or payload_index >= len(block_groups):
3358 + raise MoRIIOError(
3359 + f"KDA layer {layer_name} has no block group in the transfer payload"
3360 + )
3361 + return list(block_groups[payload_index])
3362 +
3363 + def _compute_mamba_transfer_offsets(
3364 + self,
3365 + layer_name: str,
3366 + local_slots: list[int],
3367 + remote_slots: list[int],
3368 + remote_tp_size: int | None = None,
3369 + ) -> tuple[list[int], list[int], list[int], int]:
3370 + """Compute (local, remote, sizes, n_conv) for a KDA layer.
3371 +
3372 + Entries [:n_conv] address the conv region/session; entries [n_conv:]
3373 + address the ssm region/session.
3374 + """
3375 + assert self._conv_decomp is not None, "KDA decomposition not derived"
3376 + conv, ssm = kda_conv_ssm(
3377 + self.kv_caches[layer_name],
3378 + cast(MambaSpec, self.layer_to_spec[layer_name]),
3379 + )
3380 + tp_ratio = self._mamba_tp_ratio(remote_tp_size)
3381 + # Cache the slot-independent geometry per layer; only the per-request
3382 + # slot bases vary, and heterogeneous TP is rejected above.
3383 + template = self._mamba_offset_templates.get(layer_name)
3384 + if template is None:
3385 + template = build_mamba_offset_template(
3386 + conv,
3387 + ssm,
3388 + self._conv_decomp,
3389 + tp_ratio,
3390 + )
3391 + self._mamba_offset_templates[layer_name] = template
3392 + local_offs, remote_offs, sizes = apply_mamba_offset_template(
3393 + template, local_slots, remote_slots
3394 + )
3395 + n_conv = compute_mamba_conv_split_count(local_slots, self._conv_decomp)
3396 + return local_offs, remote_offs, sizes, n_conv
3397 +
3398 + def _post_mamba_reads(
3399 + self,
3400 + layer_name: str,
3401 + sessions: list,
3402 + region_sessions: list[int],
3403 + local_slots: list[int],
3404 + remote_slots: list[int],
3405 + remote_tp_size: int,
3406 + request_id: str,
3407 + deadline: float,
3408 + ) -> list:
3409 + """Post one KDA layer's conv and ssm reads; return their statuses.
3410 +
3411 + The layer owns two registered regions, so the offset list is split at
3412 + the conv/ssm boundary and each half goes to its own session. A half
3413 + with nothing to move is skipped rather than posted empty.
3414 + """
3415 + local, remote, sizes, n_conv = self._compute_mamba_transfer_offsets(
3416 + layer_name,
3417 + local_slots,
3418 + remote_slots,
3419 + remote_tp_size=remote_tp_size,
3420 + )
3421 + statuses: list = []
3422 + for sess_idx, sl in (
3423 + (region_sessions[0], slice(0, n_conv)),
3424 + (region_sessions[1], slice(n_conv, None)),
3425 + ):
3426 + region_sizes = sizes[sl]
3427 + if not region_sizes:
3428 + continue
3429 + statuses.append(
3430 + self._post_read_with_backoff(
3431 + sessions[sess_idx],
3432 + region_sizes,
3433 + local[sl],
3434 + remote[sl],
3435 + request_id,
3436 + layer_name,
3437 + deadline,
3438 + )
3439 + )
3440 + return statuses
3441 +
3442 + def _post_read_with_backoff(
3443 + self,
3444 + session,
3445 + sizes: list[int],
3446 + local_offsets: list[int],
3447 + remote_offsets: list[int],
3448 + request_id: str,
3449 + layer_name: str,
3450 + deadline: float,
3451 + ):
3452 + """Post one RDMA READ, backing off on transient SQ-full rejection.
3453 +
3454 + read_remote_data posts synchronously, so a send-queue-full rejection is
3455 + a Failed() status on return; a separate CQ-poll thread drains
3456 + completions and frees SQ depth, so we back off and re-post until
3457 + transfer_timeout, then store the failed status (get_finished notifies
3458 + prefill and drops the request non-fatally).
3459 + """
3460 + _backoff = _SQ_FULL_BACKOFF_INITIAL_S
3461 + while True:
3462 + transfer_status = self.moriio_wrapper.read_remote_data(
3463 + sizes, local_offsets, remote_offsets, session
3464 + )
3465 + if not self._is_sq_full_status(transfer_status):
3466 + break
3467 + if time.monotonic() > deadline:
3468 + logger.warning(
3469 + "MoRIIO READ send queue stayed full past "
3470 + "transfer_timeout for req %s layer %s; storing failed "
3471 + "status (get_finished notifies prefill and drops the "
3472 + "request). Raise qp_per_transfer if frequent.",
3473 + request_id,
3474 + layer_name,
3475 + )
3476 + break
3477 + time.sleep(_backoff)
3478 + _backoff = min(_backoff * 2, _SQ_FULL_BACKOFF_MAX_S)
3479 + return transfer_status
3480 +
3481 @staticmethod
3482 def _is_sq_full_status(status) -> bool:
3483 """True if a MoRIIO transfer status is a transient RDMA send-queue-full
@@ -2766,6 +3527,44 @@ class MoRIIOConnectorWorker:
3527 if chosen_tp is not None
3528 else self._remote_tp_rank(remote_tp_size)
3529 )
3530 + if not any(local_block_ids):
3531 + # A request aborted before decode allocation still carries the
3532 + # producer's block table. There is no local destination or forward,
3533 + # so release the producer blocks without building sessions or
3534 + # posting RDMA reads.
3535 + self.moriio_wrapper.send_notify(
3536 + transfer_id,
3537 + remote_host,
3538 + str(
3539 + remote_notify_port
3540 + + get_port_offset(
3541 + int(remote_dp_rank),
3542 + eff_tp,
3543 + remote_tp_size,
3544 + )
3545 + ),
3546 + message_type="release",
3547 + message_fields={"consumer_tp_size": self.world_size},
3548 + )
3549 + return
3550 +
3551 + if self._has_mamba:
3552 + expected_groups = 1 + self._num_mamba_transfer_groups
3553 + if (
3554 + len(local_block_ids) != expected_groups
3555 + or len(remote_block_ids) != expected_groups
3556 + ):
3557 + raise MoRIIOError(
3558 + "MoRIIO hybrid READ block group count mismatch: "
3559 + f"expected={expected_groups}, local={len(local_block_ids)}, "
3560 + f"remote={len(remote_block_ids)}"
3561 + )
3562 + local_attn = list(local_block_ids[0])
3563 + remote_attn = list(remote_block_ids[0])
3564 + self._mamba_tp_ratio(remote_tp_size)
3565 + else:
3566 + local_attn = remote_attn = []
3567 +
3568 if flexible:
3569 remote_dp_engine_id = self.get_engine_name_with_dp_tp(
3570 dst_engine_id, int(remote_dp_rank), eff_tp
@@ -2778,52 +3577,89 @@ class MoRIIOConnectorWorker:
3577
3578 # SQ-full backpressure deadline, shared across this request's layers.
3579 _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout
3580 + # TODO: apply multi-session batch-read when MoRIIO supports it.
3581 for layer_name in self.layer_name_to_local_kv_cache_metadata:
2782 - sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
2783 - layer_name
2784 - )
2785 - group_idx = self.layer_to_group[layer_name]
2786 - offs = self._compute_block_transfer_offsets(
2787 - layer_name,
2788 - local_block_ids[group_idx],
2789 - remote_block_ids[group_idx],
2790 - remote_moriio_meta,
2791 - remote_tp_size=remote_tp_size,
2792 - remote_engine_id=remote_dp_engine_id,
2793 - )
2794 - # TODO : apply multi-session batch-read when moriio support it
2795 - #
2796 - # SQ-full backpressure: read_remote_data posts the RDMA READ
2797 - # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes
2798 - # back as a Failed() status right here. A SEPARATE CQ-poll thread
2799 - # drains completions and frees SQ depth, so back off and RE-POST
2800 - # rather than let a transient rejection abort the request. No
2801 - # self-deadlock (the drain is off-thread); the reserve is
2802 - # all-or-nothing (nothing posted on a rejected attempt). Bounded by
2803 - # transfer_timeout; on sustained overload store the failed status and
2804 - # let get_finished handle it non-fatally (notify prefill + drop).
2805 - _backoff = 0.001
2806 - while True:
2807 - transfer_status = self.moriio_wrapper.read_remote_data(
2808 - offs[2], offs[0], offs[1], sessions[sess_idx]
3582 + region_sessions = self._region_session_indices(layer_name)
3583 + statuses = []
3584 + if self._is_mamba_layer(layer_name):
3585 + local_mamba = self._mamba_blocks_for_layer(layer_name, local_block_ids)
3586 + remote_mamba = self._mamba_blocks_for_layer(
3587 + layer_name, remote_block_ids
3588 )
2810 - if not self._is_sq_full_status(transfer_status):
2811 - break
2812 - if time.monotonic() > _sq_deadline:
2813 - logger.warning(
2814 - "MoRIIO READ send queue stayed full past "
2815 - "transfer_timeout for req %s layer %s; storing failed "
2816 - "status (get_finished notifies prefill and drops the "
2817 - "request). Raise qp_per_transfer if frequent.",
3589 + # KDA layer: pull conv (sub-projection offsets) and ssm at the
3590 + # request's recurrent-state slot. Two regions -> two sessions.
3591 + if not local_mamba or not remote_mamba:
3592 + if not local_attn and not remote_attn:
3593 + # Whole-request no-op (e.g. full prefix cache hit):
3594 + # nothing to transfer for any layer.
3595 + continue
3596 + # A KDA layer must have its recurrent-state slot in the
3597 + # block-id tuple; missing it (while attention blocks are
3598 + # present) means the mamba KV-cache group is absent -> bug.
3599 + raise MoRIIOError(
3600 + f"KDA layer {layer_name}: missing mamba recurrent-state "
3601 + f"block ids (local={local_mamba}, remote={remote_mamba}) "
3602 + f"with attention blocks present for request {request_id}; "
3603 + "the mamba KV-cache group is absent from the transfer "
3604 + "block-id tuple"
3605 + )
3606 + statuses = self._post_mamba_reads(
3607 + layer_name,
3608 + sessions,
3609 + region_sessions,
3610 + local_mamba,
3611 + remote_mamba,
3612 + remote_tp_size,
3613 + request_id,
3614 + _sq_deadline,
3615 + )
3616 + else:
3617 + if self._has_mamba:
3618 + local_layer_blocks = local_attn
3619 + remote_layer_blocks = remote_attn
3620 + else:
3621 + group_idx = self.layer_to_group[layer_name]
3622 + local_layer_blocks = local_block_ids[group_idx]
3623 + remote_layer_blocks = remote_block_ids[group_idx]
3624 + offs = self._compute_block_transfer_offsets(
3625 + layer_name,
3626 + local_layer_blocks,
3627 + remote_layer_blocks,
3628 + remote_moriio_meta,
3629 + remote_tp_size=remote_tp_size,
3630 + remote_engine_id=remote_dp_engine_id,
3631 + )
3632 + statuses.append(
3633 + self._post_read_with_backoff(
3634 + sessions[region_sessions[0]],
3635 + offs[2],
3636 + offs[0],
3637 + offs[1],
3638 request_id,
3639 layer_name,
3640 + _sq_deadline,
3641 )
2821 - break
2822 - time.sleep(_backoff)
2823 - _backoff = min(_backoff * 2, 0.05)
3642 + )
3643 with self.moriio_wrapper.lock:
2825 - self._recving_transfers[request_id][layer_name] = transfer_status
3644 + self._recving_transfers[request_id][layer_name] = statuses
3645 self._recving_transfers_start.setdefault(request_id, time.monotonic())
3646 + # Destination blocks, kept for _record_failed_recv.
3647 + if self._has_mamba:
3648 + failed_recv_blocks = local_attn
3649 + else:
3650 + failed_recv_blocks = [
3651 + block_id for group in local_block_ids for block_id in group
3652 + ]
3653 + self._recving_local_blocks.setdefault(
3654 + request_id, list(failed_recv_blocks)
3655 + )
3656 + # Awaited before the forward when the per-layer barrier cannot
3657 + # run (see _await_reads_issued_this_step).
3658 + self._reads_issued_this_step.extend(statuses)
3659 + if self._is_mamba_layer(layer_name):
3660 + # No per-layer barrier exists for mamba layers, so these are
3661 + # always awaited in start_load_kv.
3662 + self._mamba_reads_this_step.extend(statuses)
3663 self._recving_transfers_callback_addr[request_id] = (
3664 remote_host,
3665 str(
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+34 -4
@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
29 MoRIIOError,
30 MoRIIOTransferAck,
31 RemoteAllocInfo,
32 + TransferBatchState,
33 TransferError,
34 TransferId,
35 WriteTask,
@@ -58,7 +59,6 @@ try:
59 except ImportError:
60 logger.error("MoRIIO is not available")
61
61 -
62 """Write task execution logic for MoRIIO connector."""
63
64
@@ -395,9 +395,9 @@ class MoRIIOWriter:
395 )
396 request_info.transfer_offsets[key] = offsets
397
398 - # Get session index
399 - layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
400 - sess_idx = layer_names.index(task.layer_name)
398 + # One session per registered region; attention layers map to a single
399 + # index (see MoRIIOConnectorWorker._region_session_indices).
400 + sess_idx = self.worker._region_session_indices(task.layer_name)[0]
401
402 local_off, remote_off, sizes = offsets
403
@@ -662,6 +662,25 @@ class MoRIIOWrapper:
662 )
663 return transfer_status
664
665 + def poll_transfer_batch(self, transfer_statuses: list[Any]) -> TransferBatchState:
666 + """Non-blocking verdict over every status of one request.
667 +
668 + Judging a request by its newest status alone calls it done while an
669 + earlier layer is unfinished or already failed; failures do land on
670 + layers other than the last, because _post_read_with_backoff returns a
671 + failed status when the send queue never drains.
672 +
673 + This stays a Python scan so it does not drive the backend's progress
674 + callback from the engine thread alongside MoRIIO's own poller.
675 + """
676 + state = TransferBatchState.DONE
677 + for status in transfer_statuses:
678 + if status.Failed():
679 + return TransferBatchState.FAILED
680 + if not status.Succeeded():
681 + state = TransferBatchState.PENDING
682 + return state
683 +
684 def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None):
685 if transfer_statuses is None:
686 with self.lock:
@@ -673,6 +692,15 @@ class MoRIIOWrapper:
692 if not transfers_to_wait:
693 return
694
695 + self._poll_transfers_until_done(transfers_to_wait)
696 +
697 + def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None:
698 + """Fallback for mori builds without the batched wait.
699 +
700 + Gives failure precedence over pending: once any status has failed the
701 + request is lost, so waiting out the remaining ones would only delay the
702 + barrier by up to transfer_timeout.
703 + """
704 timeout = self._transfer_timeout
705 deadline = time.monotonic() + timeout
706 remaining = list(transfers_to_wait)
@@ -698,6 +726,8 @@ class MoRIIOWrapper:
726 )
727 else:
728 still_waiting.append(status)
729 + if errors:
730 + break
731 remaining = still_waiting
732 if remaining:
733 time.sleep(0.001)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py
+285 -5
@@ -2,14 +2,18 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable, Mapping
5 -from typing import NamedTuple
5 +from typing import NamedTuple, cast
6
7 import torch
8
9 +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
10 + MambaConvSplitInfo,
11 +)
12 from vllm.v1.kv_cache_interface import (
13 AttentionSpec,
14 KVCacheConfig,
15 KVCacheSpec,
16 + MambaSpec,
17 MLAAttentionSpec,
18 SlidingWindowMLASpec,
19 UniformTypeKVCacheSpecs,
@@ -29,6 +33,85 @@ class LayerTransferGeometry(NamedTuple):
33 split_kv_regions: bool
34
35
36 +class MambaTransferGeometry(NamedTuple):
37 + """Geometry of a hybrid (mamba/KDA) layer's recurrent state.
38 +
39 + A KDA layer's packed page is exposed as ``conv_state`` and ``ssm_state``
40 + views. Each is slot-strided, so a slot's bytes live at
41 + ``slot * slot_stride * element_size`` and span ``slot_bytes``. The two
42 + tensors are registered as two separate MoRIIO regions per layer.
43 + """
44 +
45 + conv_slot_stride: int
46 + ssm_slot_stride: int
47 + ssm_slot_bytes: int
48 + conv_element_size: int
49 + ssm_element_size: int
50 + conv_region_len: int
51 + ssm_region_len: int
52 +
53 +
54 +def get_mamba_transfer_geometry(
55 + conv: torch.Tensor,
56 + ssm: torch.Tensor,
57 +) -> MambaTransferGeometry:
58 + """Derive the slot-strided geometry of a KDA layer's conv + ssm state.
59 +
60 + ``conv``/``ssm`` are zero-copy views over the layer's packed page. Slot byte
61 + offset is ``slot * slot_stride * element_size`` and each
62 + slot spans ``subtensor[0].numel() * element_size`` bytes; the whole tensor
63 + is registered as one region (``*_region_len`` bytes).
64 + """
65 + if conv.ndim == 0 or ssm.ndim == 0 or conv.shape[0] == 0 or ssm.shape[0] == 0:
66 + raise ValueError("KDA conv/ssm caches must contain at least one slot")
67 + if conv.shape[0] != ssm.shape[0]:
68 + raise ValueError(
69 + f"KDA conv/ssm slot count mismatch: {conv.shape[0]} != {ssm.shape[0]}"
70 + )
71 + conv_slot_bytes = conv[0].numel() * conv.element_size()
72 + ssm_slot_bytes = ssm[0].numel() * ssm.element_size()
73 + conv_available_bytes = (
74 + conv.untyped_storage().nbytes() - conv.storage_offset() * conv.element_size()
75 + )
76 + ssm_available_bytes = (
77 + ssm.untyped_storage().nbytes() - ssm.storage_offset() * ssm.element_size()
78 + )
79 + if conv.shape[0]:
80 + conv_required_bytes = (conv.shape[0] - 1) * conv.stride(
81 + 0
82 + ) * conv.element_size() + conv_slot_bytes
83 + ssm_required_bytes = (ssm.shape[0] - 1) * ssm.stride(
84 + 0
85 + ) * ssm.element_size() + ssm_slot_bytes
86 + if conv_required_bytes > conv_available_bytes:
87 + raise ValueError("KDA conv slots exceed the registered storage extent")
88 + if ssm_required_bytes > ssm_available_bytes:
89 + raise ValueError("KDA ssm slots exceed the registered storage extent")
90 +
91 + return MambaTransferGeometry(
92 + conv_slot_stride=conv.stride(0),
93 + ssm_slot_stride=ssm.stride(0),
94 + ssm_slot_bytes=ssm_slot_bytes,
95 + conv_element_size=conv.element_size(),
96 + ssm_element_size=ssm.element_size(),
97 + # Registered region must span the full slot-strided extent: stride(0)
98 + # may exceed the per-slot element count, so numel() would under-size the
99 + # region and truncate the highest-slot transfers. shape[0]*stride(0)
100 + # covers every byte a ``slot * slot_stride * elem`` offset can address.
101 + # Cap at the bytes remaining from each view's storage offset. A packed
102 + # conv+ssm page places ssm at a non-zero intra-buffer offset,
103 + # so do not register past the end of the shared backing allocation.
104 + conv_region_len=min(
105 + conv.shape[0] * conv.stride(0) * conv.element_size(),
106 + conv_available_bytes,
107 + ),
108 + ssm_region_len=min(
109 + ssm.shape[0] * ssm.stride(0) * ssm.element_size(),
110 + ssm_available_bytes,
111 + ),
112 + )
113 +
114 +
115 def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
116 layer_to_spec: dict[str, KVCacheSpec] = {}
117 for group in kv_cache_config.kv_cache_groups:
@@ -99,10 +182,10 @@ def get_layer_transfer_geometry(
182 layer_to_spec: Mapping[str, KVCacheSpec],
183 remote_num_blocks: int | None = None,
184 ) -> LayerTransferGeometry:
185 + spec = layer_to_spec[layer_name]
186 shape = kv_cache.shape
187 stride = kv_cache.stride()
188 element_size = kv_cache.element_size()
105 - spec = layer_to_spec[layer_name]
189 is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
190
191 if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
@@ -238,12 +321,60 @@ def get_layer_transfer_geometry(
321 )
322
323
324 +def kda_conv_ssm(
325 + kv_cache: torch.Tensor,
326 + spec: MambaSpec,
327 +) -> tuple[torch.Tensor, torch.Tensor]:
328 + """Return zero-copy conv and SSM views over a packed two-state page."""
329 + from math import prod
330 +
331 + from vllm.utils.torch_utils import get_dtype_size
332 +
333 + if len(spec.shapes) != 2 or len(spec.dtypes) != 2:
334 + raise ValueError(
335 + "MoRIIO KDA transfer requires exactly two states (conv and SSM)"
336 + )
337 + if kv_cache.element_size() != 1:
338 + raise ValueError("MoRIIO KDA packed cache must use a byte-sized dtype")
339 + if kv_cache.ndim != 4 or kv_cache.shape[1:3] != (1, 1) or kv_cache.stride(3) != 1:
340 + raise ValueError(
341 + "MoRIIO KDA packed cache must have shape [B, 1, 1, C] with "
342 + "a contiguous byte dimension"
343 + )
344 + # Mirror MambaBase.bind_kv_cache: blocks may be strided within a shared
345 + # multi-layer allocation, while every individual byte page is contiguous.
346 + pages = kv_cache.squeeze(dim=(1, 2))
347 + states: list[torch.Tensor] = []
348 + offset = 0
349 + for shp, dt in zip(spec.shapes, spec.dtypes, strict=True):
350 + nbytes = prod(shp) * get_dtype_size(dt)
351 + if offset + nbytes > pages.shape[1]:
352 + raise ValueError(
353 + "MoRIIO KDA state layout exceeds the packed cache page size"
354 + )
355 + state = pages[:, offset : offset + nbytes].view(dt)
356 + states.append(state.view(-1, *shp))
357 + offset += nbytes
358 + return states[0], states[1]
359 +
360 +
361 def iter_layer_registration_regions(
362 layer_name: str,
363 kv_cache: torch.Tensor,
364 layer_to_spec: Mapping[str, KVCacheSpec],
365 ) -> list[tuple[torch.Tensor, int]]:
246 - geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
366 + spec = layer_to_spec[layer_name]
367 + if isinstance(spec, MambaSpec):
368 + # KDA layer: register the conv and ssm tensors as two whole-tensor
369 + # regions (the conv 3-subprojection split is expressed as transfer
370 + # offset triples, not extra registrations).
371 + conv, ssm = kda_conv_ssm(kv_cache, spec)
372 + geom = get_mamba_transfer_geometry(conv, ssm)
373 + return [(conv, geom.conv_region_len), (ssm, geom.ssm_region_len)]
374 + geometry = cast(
375 + LayerTransferGeometry,
376 + get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec),
377 + )
378 region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
379 if geometry.regions_per_block == 1:
380 # With padded or interleaved pages the block stride exceeds the
@@ -311,9 +442,36 @@ def compute_block_transfer_offsets(
442 "local_block_ids longer than remote_block_ids: "
443 f"{len(local_block_ids)} > {len(remote_block_ids)}"
444 )
314 - geometry = get_layer_transfer_geometry(
315 - layer_name, kv_cache, layer_to_spec, remote_num_blocks
445 + geometry = cast(
446 + LayerTransferGeometry,
447 + get_layer_transfer_geometry(
448 + layer_name, kv_cache, layer_to_spec, remote_num_blocks
449 + ),
450 + )
451 + invalid_local = next(
452 + (
453 + block_id
454 + for block_id in local_block_ids
455 + if not 0 <= block_id < geometry.num_blocks
456 + ),
457 + None,
458 + )
459 + if invalid_local is not None:
460 + raise ValueError(
461 + f"local block id {invalid_local} is outside [0, {geometry.num_blocks})"
462 + )
463 + invalid_remote = next(
464 + (
465 + block_id
466 + for block_id in remote_block_ids
467 + if not 0 <= block_id < remote_num_blocks
468 + ),
469 + None,
470 )
471 + if invalid_remote is not None:
472 + raise ValueError(
473 + f"remote block id {invalid_remote} is outside [0, {remote_num_blocks})"
474 + )
475 element_size = kv_cache.element_size()
476 transfer_size_byte = geometry.block_len
477 per_block = geometry.transfers_per_block
@@ -339,3 +497,125 @@ def compute_block_transfer_offsets(
497 w += 1
498
499 return merge_fn(offset_local, offset_remote, sizes)
500 +
501 +
502 +class MambaOffsetTemplate(NamedTuple):
503 + """Slot-independent conv+ssm offset decomposition for a KDA layer.
504 +
505 + Homogeneous-TP geometry (conv/ssm slot strides, conv sub-projection
506 + offsets, ssm per-slot size) is identical across requests and across the
507 + homogeneous GDN layers, so it is computed once and reused; only the
508 + per-request slot bases vary (see ``apply_mamba_offset_template``).
509 + """
510 +
511 + num_slots: int
512 + conv_slot_stride_bytes: int
513 + ssm_slot_stride_bytes: int
514 + ssm_read_bytes: int
515 + conv_subprojs: tuple[tuple[int, int, int], ...]
516 +
517 +
518 +def build_mamba_offset_template(
519 + conv: torch.Tensor,
520 + ssm: torch.Tensor,
521 + split_info: MambaConvSplitInfo,
522 + tp_ratio: int,
523 +) -> MambaOffsetTemplate:
524 + """Precompute the slot-independent conv+ssm offset decomposition.
525 +
526 + Only homogeneous TP (``tp_ratio == 1``) is supported; heterogeneous TP is
527 + gated with ``NotImplementedError`` (see the design doc). The returned
528 + template is a pure function of the layer geometry and ``split_info``, so it
529 + can be cached and reused across requests without recomputation.
530 + """
531 + if tp_ratio != 1:
532 + # Heterogeneous-TP recurrent-state transfer needs the remote page's
533 + # slot stride (which differs from the local page) and, for
534 + # P_TP > D_TP, a multi-rank gather of conv/ssm slices. Both are
535 + # follow-ups. The mamba path bypasses the attention KV-head validator,
536 + # so fail loudly here instead of silently transferring wrong bytes.
537 + raise NotImplementedError(
538 + "heterogeneous-TP mamba/KDA transfer is not supported yet "
539 + f"(tp_ratio={tp_ratio}); use equal prefill/decode TP for "
540 + "hybrid models"
541 + )
542 + geom = get_mamba_transfer_geometry(conv, ssm)
543 + conv_local_offsets = split_info.local_conv_offsets
544 + conv_remote_offsets = split_info.remote_conv_offsets(0, 1)
545 + ssm_read_bytes = geom.ssm_slot_bytes
546 + conv_subprojs = tuple(
547 + (loff, roff, rsz)
548 + for (loff, _lsz), (roff, rsz) in zip(conv_local_offsets, conv_remote_offsets)
549 + )
550 + return MambaOffsetTemplate(
551 + num_slots=conv.shape[0],
552 + conv_slot_stride_bytes=geom.conv_slot_stride * geom.conv_element_size,
553 + ssm_slot_stride_bytes=geom.ssm_slot_stride * geom.ssm_element_size,
554 + ssm_read_bytes=ssm_read_bytes,
555 + conv_subprojs=conv_subprojs,
556 + )
557 +
558 +
559 +def apply_mamba_offset_template(
560 + template: MambaOffsetTemplate,
561 + local_slots: list[int],
562 + remote_slots: list[int],
563 +) -> tuple[list[int], list[int], list[int]]:
564 + """Apply per-request slot bases to a cached ``MambaOffsetTemplate``.
565 +
566 + Conv sub-projection entries (all slots) come first, followed by one ssm
567 + entry per slot.
568 + """
569 + if len(local_slots) != len(remote_slots):
570 + raise ValueError(
571 + "local and remote mamba slot counts must match: "
572 + f"{len(local_slots)} != {len(remote_slots)}"
573 + )
574 + invalid_local = next(
575 + (slot for slot in local_slots if not 0 <= slot < template.num_slots), None
576 + )
577 + if invalid_local is not None:
578 + raise ValueError(
579 + f"local mamba slot {invalid_local} is outside [0, {template.num_slots})"
580 + )
581 + invalid_remote = next((slot for slot in remote_slots if slot < 0), None)
582 + if invalid_remote is not None:
583 + raise ValueError(f"remote mamba slot {invalid_remote} must be non-negative")
584 + # The peer's KDA slot count is not part of MoRIIOAgentMetadata yet. Do not
585 + # compare remote slot IDs against the local pool size: two otherwise
586 + # compatible engines may reserve different numbers of slots.
587 + local_offs: list[int] = []
588 + remote_offs: list[int] = []
589 + sizes: list[int] = []
590 +
591 + # Conv sub-projections (region 0), all slots first.
592 + conv_slot_stride_bytes = template.conv_slot_stride_bytes
593 + for ls, rs in zip(local_slots, remote_slots):
594 + lbase = ls * conv_slot_stride_bytes
595 + rbase = rs * conv_slot_stride_bytes
596 + for loff, roff, rsz in template.conv_subprojs:
597 + local_offs.append(lbase + loff)
598 + remote_offs.append(rbase + roff)
599 + sizes.append(rsz)
600 +
601 + # SSM temporal state (region 1), one entry per slot.
602 + ssm_slot_stride_bytes = template.ssm_slot_stride_bytes
603 + ssm_read_bytes = template.ssm_read_bytes
604 + for ls, rs in zip(local_slots, remote_slots):
605 + local_offs.append(ls * ssm_slot_stride_bytes)
606 + remote_offs.append(rs * ssm_slot_stride_bytes)
607 + sizes.append(ssm_read_bytes)
608 +
609 + return local_offs, remote_offs, sizes
610 +
611 +
612 +def compute_mamba_conv_split_count(
613 + local_slots: list[int],
614 + split_info: MambaConvSplitInfo,
615 +) -> int:
616 + """Number of leading conv entries in a KDA layer transfer plan.
617 +
618 + Entries ``[:count]`` are conv sub-projections (conv region/session);
619 + entries ``[count:]`` are the per-slot ssm state (ssm region/session).
620 + """
621 + return len(local_slots) * len(split_info.local_conv_offsets)