[ROCm][Connector] SWA+HMA-support in MoRI-IO connector (Gemma4) (#53721)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 14, 2026 at 14:52 UTC
1d0d1081c409176a0e7c24742f7c60b9943657aa
4 files changed
+408
-46
tests/v1/kv_connector/unit/test_moriio_connector.py
+228
-4
@@ -39,10 +39,12 @@ from vllm.utils.network_utils import (
39
make_zmq_path,
40
)
41
from vllm.v1.kv_cache_interface import (
42
+ ChunkedLocalAttentionSpec,
43
FullAttentionSpec,
44
KVCacheConfig,
45
KVCacheGroupSpec,
46
KVCacheTensor,
47
+ SlidingWindowSpec,
48
compute_layer_kv_cache_shape_bytes,
49
)
50
@@ -338,7 +340,7 @@ def test_write_mode_saves_local_block_ids():
340
req_meta = kv_connector_metadata.reqs_to_save[request_id]
341
342
for block_id, block in zip(
341
- req_meta.local_block_ids,
343
+ req_meta.local_block_ids[0],
344
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
345
request_id
346
],
@@ -393,7 +395,7 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
395
req_meta = kv_connector_metadata.reqs_to_save[request_id]
396
397
for block_id, block in zip(
396
- req_meta.local_block_ids,
398
+ req_meta.local_block_ids[0],
399
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
400
request_id
401
],
@@ -432,7 +434,7 @@ def test_read_mode_loads_remote_block_ids():
434
].req_to_blocks[request_id]
435
436
# Set remote block ids to be fetched.
435
- request.kv_transfer_params["remote_block_ids"] = block_list
437
+ request.kv_transfer_params["remote_block_ids"] = [block_list]
438
439
# Remote Prefill, triggers MoRIIOConnectorMetadata.
440
@@ -457,7 +459,7 @@ def test_read_mode_loads_remote_block_ids():
459
req_meta = kv_connector_metadata.reqs_to_recv[request_id]
460
461
for block_id, block in zip(
460
- req_meta.local_block_ids,
462
+ req_meta.local_block_ids[0],
463
scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
464
request_id
465
],
@@ -699,3 +701,225 @@ def test_resolve_host_ip_prefers_extra_config():
701
fallback = get_ip()
702
assert resolve_host_ip({}) == fallback
703
assert resolve_host_ip({"host_ip": ""}) == fallback
704
+
705
+
706
+def _make_hybrid_kv_cache_config() -> KVCacheConfig:
707
+ """One full-attention group + one sliding-window group (Gemma-like)."""
708
+ full_spec = FullAttentionSpec(
709
+ block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
710
+ )
711
+ sw_spec = SlidingWindowSpec(
712
+ block_size=16,
713
+ num_kv_heads=4,
714
+ head_size=64,
715
+ dtype=torch.float16,
716
+ sliding_window=32,
717
+ )
718
+ num_blocks = 2
719
+ page = full_spec.page_size_bytes
720
+ return KVCacheConfig(
721
+ num_blocks=num_blocks,
722
+ kv_cache_tensors=[
723
+ KVCacheTensor(
724
+ size=2 * num_blocks * page,
725
+ layers=["full0", "sw0"],
726
+ layer_stride=num_blocks * page,
727
+ block_stride=page,
728
+ )
729
+ ],
730
+ kv_cache_groups=[
731
+ KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec),
732
+ KVCacheGroupSpec(layer_names=["sw0"], kv_cache_spec=sw_spec),
733
+ ],
734
+ )
735
+
736
+
737
+def _read_scheduler(
738
+ kv_cache_config: KVCacheConfig, disable_hma: bool = False
739
+) -> MoRIIOConnectorScheduler:
740
+ vllm_config = create_vllm_config(role="kv_producer", read_mode=True)
741
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = disable_hma
742
+ with set_current_vllm_config(vllm_config):
743
+ connector = MoRIIOConnector(
744
+ vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
745
+ )
746
+ assert connector.connector_scheduler is not None
747
+ return connector.connector_scheduler
748
+
749
+
750
+def test_hma_blocks_per_sw_two_groups():
751
+ """A Full + SlidingWindow config runs with HMA and computes block budgets
752
+ correctly"""
753
+ scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
754
+ assert scheduler._is_hma_required is True
755
+ # cdiv(32, 16) + 1 == 3 for the sliding-window group, 0 for full attention.
756
+ assert scheduler.blocks_per_sw == [0, 3]
757
+
758
+
759
+@pytest.mark.parametrize(
760
+ "swa_enabled, disable_hma, expected_is_hma",
761
+ [
762
+ (True, False, True), # sliding-window group present, HMA enabled
763
+ (True, True, False), # sliding-window group present but HMA disabled
764
+ (False, False, False), # full-attention only, HMA not needed
765
+ ],
766
+)
767
+def test_is_hma_required(swa_enabled, disable_hma, expected_is_hma):
768
+ """_is_hma_required tracks both the KV cache groups and the
769
+ --disable-hybrid-kv-cache-manager flag. When HMA is off,
770
+ get_exchange_clipped_blocks must be a no-op."""
771
+ config = (
772
+ _make_hybrid_kv_cache_config() if swa_enabled else _make_test_kv_cache_config()
773
+ )
774
+ scheduler = _read_scheduler(config, disable_hma=disable_hma)
775
+ assert scheduler._is_hma_required is expected_is_hma
776
+ if not expected_is_hma:
777
+ blocks = [[1, 2, 3, 4, 5]]
778
+ assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
779
+
780
+
781
+def test_non_sliding_window_hybrid_is_rejected():
782
+ """A hybrid group that is not sliding-window (e.g. chunked-local
783
+ attention) must fail closed rather than be silently mistransferred."""
784
+ full_spec = FullAttentionSpec(
785
+ block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
786
+ )
787
+ local_spec = ChunkedLocalAttentionSpec(
788
+ block_size=16,
789
+ num_kv_heads=4,
790
+ head_size=64,
791
+ dtype=torch.float16,
792
+ attention_chunk_size=32,
793
+ )
794
+ num_blocks = 2
795
+ page = full_spec.page_size_bytes
796
+ config = KVCacheConfig(
797
+ num_blocks=num_blocks,
798
+ kv_cache_tensors=[
799
+ KVCacheTensor(
800
+ size=2 * num_blocks * page,
801
+ layers=["full0", "local0"],
802
+ layer_stride=num_blocks * page,
803
+ block_stride=page,
804
+ )
805
+ ],
806
+ kv_cache_groups=[
807
+ KVCacheGroupSpec(layer_names=["full0"], kv_cache_spec=full_spec),
808
+ KVCacheGroupSpec(layer_names=["local0"], kv_cache_spec=local_spec),
809
+ ],
810
+ )
811
+ with pytest.raises(NotImplementedError, match="sliding-window hybrid"):
812
+ _read_scheduler(config)
813
+
814
+
815
+def test_token_count_basis_uses_full_attention_group():
816
+ """Chunked-prefill token counting must use an unclipped full-attention
817
+ group (blocks_per_sw == 0), not a clipped sliding-window group."""
818
+ scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
819
+ # Group 0 is the full-attention group
820
+ assert scheduler._full_attn_group_idx == 0
821
+ assert scheduler._full_attn_block_size == 16
822
+
823
+
824
+def test_get_exchange_clipped_blocks_clips_only_sw_group():
825
+ """get_exchange_clipped_blocks keeps the full attn group intact and clips the
826
+ sliding-window group to its window tail."""
827
+ scheduler = _read_scheduler(_make_hybrid_kv_cache_config())
828
+ full = [10, 11, 12, 13, 14]
829
+ sw = [20, 21, 22, 23, 24]
830
+ clipped = scheduler.get_exchange_clipped_blocks([full, sw])
831
+ assert clipped[0] == full
832
+ assert clipped[1] == [22, 23, 24]
833
+
834
+
835
+def test_metadata_hma_block_ids_preserved_per_group():
836
+ """add_new_req stores per-group (BlockIds) block lists unchanged for both
837
+ read and write, so the hybrid group structure is retained."""
838
+ metadata = MoRIIOConnectorMetadata()
839
+
840
+ # Assume:
841
+ # - Full-attention group (6 blocks) +
842
+ # - sliding-window group already clipped to its window tail (3 blocks).
843
+ fa_blocks = [0, 1, 2, 3, 4, 5]
844
+ sw_blocks = [10, 11, 12]
845
+ local_block_ids = [fa_blocks, sw_blocks]
846
+ remote_block_ids = [[100, 101, 102, 103, 104, 105], [200, 201, 202]]
847
+
848
+ base_params = {
849
+ "remote_engine_id": "remote-engine",
850
+ "remote_host": "127.0.0.1",
851
+ "remote_handshake_port": 6301,
852
+ "remote_notify_port": 61005,
853
+ "remote_block_ids": remote_block_ids,
854
+ }
855
+
856
+ # Read mode: both local and remote ids stay per-group.
857
+ metadata.add_new_req(
858
+ request_id="recv-req",
859
+ local_block_ids=local_block_ids,
860
+ kv_transfer_params={**base_params, "transfer_id": "recv-req"},
861
+ )
862
+ recv_meta = metadata.reqs_to_recv["recv-req"]
863
+ assert recv_meta.local_block_ids == [fa_blocks, sw_blocks]
864
+ assert recv_meta.remote_block_ids == remote_block_ids
865
+
866
+ # Write mode: the decode peer allocates its own blocks, so #remote_block_ids may
867
+ # be empty, Local group structure is kept.
868
+ metadata.add_new_req(
869
+ request_id="save-req",
870
+ local_block_ids=local_block_ids,
871
+ kv_transfer_params={
872
+ **base_params,
873
+ "transfer_id": "save-req",
874
+ "remote_block_ids": [],
875
+ },
876
+ write_mode=True,
877
+ )
878
+ save_meta = metadata.reqs_to_save["save-req"]
879
+ assert save_meta.local_block_ids == [fa_blocks, sw_blocks]
880
+ assert save_meta.remote_block_ids == []
881
+
882
+
883
+def test_single_group_path_unchanged():
884
+ """A full attn config does not enable HMA and never clips blocks."""
885
+ scheduler = _read_scheduler(_make_test_kv_cache_config())
886
+ assert scheduler._is_hma_required is False
887
+ assert scheduler.blocks_per_sw == [0]
888
+ blocks = [[1, 2, 3, 4, 5]]
889
+ assert scheduler.get_exchange_clipped_blocks(blocks) == blocks
890
+
891
+
892
+def test_hybrid_write_mode_rejected():
893
+ """Hybrid KV cache groups are unsupported in WRITE mode and fail closed."""
894
+ vllm_config = create_vllm_config(role="kv_producer", read_mode=False)
895
+ with (
896
+ set_current_vllm_config(vllm_config),
897
+ pytest.raises(NotImplementedError),
898
+ ):
899
+ MoRIIOConnector(
900
+ vllm_config,
901
+ KVConnectorRole.SCHEDULER,
902
+ _make_hybrid_kv_cache_config(),
903
+ )
904
+
905
+
906
+def test_worker_layer_to_group_routing(mock_parallel_groups):
907
+ """The worker maps every layer to its KV cache group correctly."""
908
+ vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
909
+ # Building the worker directly bypasses MoRIIOConnector._set_port_defaults,
910
+ # so provide the ports manually.
911
+ vllm_config.kv_transfer_config.kv_connector_extra_config.update(
912
+ {
913
+ "http_port": 12346,
914
+ "handshake_port": 12347,
915
+ "notify_port": 12348,
916
+ }
917
+ )
918
+ with set_current_vllm_config(vllm_config):
919
+ worker = FakeMoRIIOConnectorWorker(
920
+ vllm_config,
921
+ "engine0",
922
+ hand_shake_latency=0,
923
+ kv_cache_config=_make_hybrid_kv_cache_config(),
924
+ )
925
+ assert worker.layer_to_group == {"full0": 0, "sw0": 1}
tests/v1/kv_connector/unit/test_moriio_routing.py
+6
-1
@@ -20,12 +20,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import
20
from vllm.v1.request import RequestStatus
21
22
_request_finished = MoRIIOConnectorScheduler.request_finished
23
+_clip_blocks = MoRIIOConnectorScheduler.get_exchange_clipped_blocks
24
25
26
def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
27
"""Minimal prefill-side (producer) scheduler stand-in for request_finished."""
27
- return SimpleNamespace(
28
+ s = SimpleNamespace(
29
is_producer=True,
30
+ _is_hma_required=False,
31
+ blocks_per_sw=[0],
32
_global_dp_rank=global_dp_rank,
33
engine_id="engine-abc",
34
host_ip="10.0.0.1",
@@ -39,6 +42,8 @@ def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
42
)
43
),
44
)
45
+ s.get_exchange_clipped_blocks = _clip_blocks.__get__(s)
46
+ return s
47
48
49
def _req(
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+4
-3
@@ -14,6 +14,7 @@ import torch
14
import zmq
15
16
from vllm.config import KVTransferConfig, VllmConfig
17
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
18
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19
KVConnectorMetadata,
20
)
@@ -444,8 +445,8 @@ class ReqMeta:
445
"""Metadata for a single request."""
446
447
transfer_id: TransferId
447
- local_block_ids: list[int]
448
- remote_block_ids: list[int]
448
+ local_block_ids: BlockIds
449
+ remote_block_ids: BlockIds
450
remote_host: str
451
remote_port: int
452
remote_handshake_port: int
@@ -481,7 +482,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
482
def add_new_req(
483
self,
484
request_id: ReqId,
484
- local_block_ids: list[int],
485
+ local_block_ids: BlockIds,
486
kv_transfer_params: dict[str, Any],
487
write_mode=False,
488
):
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+170
-38
@@ -18,10 +18,12 @@ import torch
18
import zmq
19
20
from vllm.config import CUDAGraphMode, VllmConfig
21
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
22
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
23
KVConnectorBase_V1,
24
KVConnectorMetadata,
25
KVConnectorRole,
26
+ SupportsHMA,
27
)
28
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
29
ROLE,
@@ -67,19 +69,24 @@ from vllm.distributed.parallel_state import (
69
)
70
from vllm.forward_context import ForwardContext, get_forward_context
71
from vllm.logger import init_logger
72
+from vllm.utils.math_utils import cdiv
73
from vllm.utils.network_utils import (
74
make_zmq_path,
75
make_zmq_socket,
76
)
77
from vllm.v1.attention.selector import get_attn_backend
78
from vllm.v1.core.sched.output import SchedulerOutput
79
+from vllm.v1.kv_cache_interface import (
80
+ FullAttentionSpec,
81
+ KVCacheConfig,
82
+ SlidingWindowSpec,
83
+)
84
from vllm.v1.outputs import KVConnectorOutput
85
from vllm.v1.request import RequestStatus
86
87
if TYPE_CHECKING:
88
from vllm.v1.attention.backend import AttentionMetadata
89
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
82
- from vllm.v1.kv_cache_interface import KVCacheConfig
90
from vllm.v1.request import Request
91
92
logger = init_logger(__name__)
@@ -189,7 +196,7 @@ def resolve_moriio_transfer_ack(
196
return transfer_id
197
198
192
-class MoRIIOConnector(KVConnectorBase_V1):
199
+class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA):
200
def __init__(
201
self,
202
vllm_config: VllmConfig,
@@ -225,7 +232,7 @@ class MoRIIOConnector(KVConnectorBase_V1):
232
)
233
if role == KVConnectorRole.SCHEDULER:
234
self.connector_scheduler: MoRIIOConnectorScheduler | None = (
228
- MoRIIOConnectorScheduler(vllm_config, self.engine_id)
235
+ MoRIIOConnectorScheduler(vllm_config, self.engine_id, kv_cache_config)
236
)
237
self.connector_worker: MoRIIOConnectorWorker | None = None
238
elif role == KVConnectorRole.WORKER:
@@ -283,6 +290,14 @@ class MoRIIOConnector(KVConnectorBase_V1):
290
self,
291
request: "Request",
292
block_ids: list[int],
293
+ ) -> tuple[bool, dict[str, Any] | None]:
294
+ assert self.connector_scheduler is not None
295
+ return self.connector_scheduler.request_finished(request, (block_ids,))
296
+
297
+ def request_finished_all_groups(
298
+ self,
299
+ request: "Request",
300
+ block_ids: BlockIds,
301
) -> tuple[bool, dict[str, Any] | None]:
302
assert self.connector_scheduler is not None
303
return self.connector_scheduler.request_finished(request, block_ids)
@@ -368,7 +383,12 @@ class MoRIIOConnector(KVConnectorBase_V1):
383
class MoRIIOConnectorScheduler:
384
"""Implementation of Scheduler side methods"""
385
371
- def __init__(self, vllm_config: VllmConfig, engine_id: str):
386
+ def __init__(
387
+ self,
388
+ vllm_config: VllmConfig,
389
+ engine_id: str,
390
+ kv_cache_config: "KVCacheConfig",
391
+ ):
392
self.vllm_config = vllm_config
393
394
assert vllm_config.kv_transfer_config is not None, (
@@ -378,6 +398,66 @@ class MoRIIOConnectorScheduler:
398
self.block_size = vllm_config.cache_config.block_size
399
self.engine_id: EngineId = engine_id
400
self.mode = get_moriio_mode(self.kv_transfer_config)
401
+
402
+ self._is_hma_required = (
403
+ not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
404
+ and any(
405
+ not isinstance(g.kv_cache_spec, FullAttentionSpec)
406
+ for g in kv_cache_config.kv_cache_groups
407
+ )
408
+ )
409
+ if self._is_hma_required:
410
+ # TODO(simondanielsson): support non-sliding window hybrids
411
+ unsupported = [
412
+ type(g.kv_cache_spec).__name__
413
+ for g in kv_cache_config.kv_cache_groups
414
+ if not isinstance(
415
+ g.kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec)
416
+ )
417
+ ]
418
+ if unsupported:
419
+ raise NotImplementedError(
420
+ "MoRIIO only supports sliding-window hybrid models (e.g. "
421
+ "Gemma) under the hybrid KV cache manager. Unsupported KV "
422
+ f"cache group spec(s): {sorted(set(unsupported))}. Pass "
423
+ "--disable-hybrid-kv-cache-manager to run without HMA."
424
+ )
425
+ if self.mode == MoRIIOMode.WRITE:
426
+ # TODO(simondanielsson): support HMA in WRITE mode
427
+ raise NotImplementedError(
428
+ "MoRIIO WRITE mode does not support hybrid KV cache groups "
429
+ "(sliding-window attention). Use READ mode, or pass "
430
+ "--disable-hybrid-kv-cache-manager."
431
+ )
432
+
433
+ sw_sizes_tokens: list[tuple[int, int]] = [
434
+ (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
435
+ if isinstance(g.kv_cache_spec, SlidingWindowSpec)
436
+ else (0, self.block_size)
437
+ for g in kv_cache_config.kv_cache_groups
438
+ ]
439
+ # add 1 to conservatively account for boundary overlap eg window isn't fully
440
+ # aligned with blocks.
441
+ self.blocks_per_sw = [
442
+ cdiv(n_tokens, block_size) + 1 if n_tokens else 0
443
+ for n_tokens, block_size in sw_sizes_tokens
444
+ ]
445
+ # In WRITE mode and chunked prefill, we perform the write after the last chunk.
446
+ # Hence we need to track once we are in the last chunk. We do this by computing
447
+ # len(block_ids[g]) * block_size[g] for a full attn group g, as hybrid groups
448
+ # may be clipped and hence not reflect the full sequence length. Here we store
449
+ # g and block_size[g] for later.
450
+ self._full_attn_group_idx = 0
451
+ self._full_attn_block_size = self.block_size
452
+ for gi, group in enumerate(kv_cache_config.kv_cache_groups):
453
+ is_full_attn = self.blocks_per_sw[gi] == 0
454
+ if is_full_attn:
455
+ self._full_attn_group_idx = gi
456
+ self._full_attn_block_size = getattr(
457
+ group.kv_cache_spec, "block_size", self.block_size
458
+ )
459
+ break
460
+
461
self.host_ip = resolve_host_ip(
462
self.kv_transfer_config.kv_connector_extra_config
463
)
@@ -409,14 +489,14 @@ class MoRIIOConnectorScheduler:
489
# Requests that need to start recv/send.
490
# New requests are added by update_state_after_alloc in
491
# the scheduler. Used to make metadata passed to Worker.
412
- self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
413
- self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
492
+ self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
493
+ self._reqs_need_save: dict[ReqId, tuple[Request, BlockIds]] = {}
494
# Snapshot of kv_transfer_params for chunked prefill recovery.
495
self._req_kv_params: dict[ReqId, dict] = {}
496
497
# For chunked prefill, we perform layer-wise access within the final chunk.
498
# TODO: Perform transfer at end chunk.
419
- self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {}
499
+ self._reqs_need_pending_save: dict[ReqId, tuple[Request, BlockIds]] = {}
500
501
if self.is_producer:
502
set_role(ROLE.PRODUCER)
@@ -442,6 +522,23 @@ class MoRIIOConnectorScheduler:
522
self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
523
self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}
524
525
+ def get_exchange_clipped_blocks(
526
+ self,
527
+ block_ids: BlockIds,
528
+ ) -> BlockIds:
529
+ """Clip the block list based on the layer type.
530
+
531
+ Behavior:
532
+ - sliding window: keep only the last window size blocks
533
+ - full attention: keep all blocks
534
+ """
535
+ if len(block_ids) == 0 or not self._is_hma_required:
536
+ return [list(blocks) for blocks in block_ids]
537
+ return [
538
+ blocks[-self.blocks_per_sw[i] :] if self.blocks_per_sw[i] > 0 else blocks
539
+ for i, blocks in enumerate(block_ids)
540
+ ]
541
+
542
def map_request_id(self, request_id: ReqId, transfer_id: TransferId):
543
self.transfer_id_to_request_id[transfer_id] = request_id
544
self.request_id_to_transfer_id[request_id] = transfer_id
@@ -642,7 +739,7 @@ class MoRIIOConnectorScheduler:
739
request_id = request.request_id
740
self.map_request_id(request_id, transfer_id)
741
if params.get("do_remote_decode"):
645
- local_block_ids = blocks.get_block_ids()[0]
742
+ local_block_ids: BlockIds = blocks.get_block_ids()
743
self._reqs_need_save[request.request_id] = (request, local_block_ids)
744
# Snapshot params now so chunked-prefill build_connector_meta
745
# can recover them on the final chunk even if the live
@@ -657,17 +754,22 @@ class MoRIIOConnectorScheduler:
754
if "remote_engine_id" in params:
755
if num_external_tokens > 0:
756
# Get unhashed blocks to pull from remote.
660
- local_block_ids = blocks.get_block_ids()[0]
661
- assert len(local_block_ids) <= len(remote_block_ids)
662
- if len(local_block_ids) != len(remote_block_ids):
663
- local_block_ids = remote_block_ids[
664
- -len(local_block_ids) :
665
- ]
757
+ local_block_ids = self.get_exchange_clipped_blocks(
758
+ blocks.get_block_ids()
759
+ )
760
+ # Partial cache hits can lead to fewer local blocks vs
761
+ # remote blocks, but never more
762
+ assert all(
763
+ len(local) <= len(remote)
764
+ for local, remote in zip(
765
+ local_block_ids, remote_block_ids, strict=True
766
+ )
767
+ )
768
else:
769
# If remote_blocks and num_external_tokens = 0, we have
770
# a full prefix cache hit on the D worker. We need to call
771
# send_notify in _read_blocks to free the memory on the P.
670
- local_block_ids = []
772
+ local_block_ids = [[] for _ in remote_block_ids]
773
774
self._reqs_need_recv[request.request_id] = (
775
request,
@@ -818,8 +920,6 @@ class MoRIIOConnectorScheduler:
920
new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]
921
922
if new_block_ids is not None:
821
- block_ids = new_block_ids[0]
822
- # TODO : hybrid attn, etc
923
# A non-disagg request (no kv_transfer_params, e.g. smoke
924
# test) is never registered in _reqs_need_pending_save;
925
# indexing it unconditionally would KeyError and crash the
@@ -827,12 +927,16 @@ class MoRIIOConnectorScheduler:
927
if req_id not in self._reqs_need_pending_save:
928
continue
929
req, existing_blocks = self._reqs_need_pending_save[req_id]
830
- updated_blocks = list(existing_blocks) + (block_ids)
930
+ updated_blocks = [
931
+ existing_blocks[g] + new_block_ids[g]
932
+ for g in range(len(new_block_ids))
933
+ ]
934
self._reqs_need_pending_save[req_id] = (req, updated_blocks)
832
- if (
833
- len(self._reqs_need_pending_save[req_id][1]) * self.block_size
834
- >= req.num_prompt_tokens
835
- ):
935
+ saved_tokens = (
936
+ len(updated_blocks[self._full_attn_group_idx])
937
+ * self._full_attn_block_size
938
+ )
939
+ if saved_tokens >= req.num_prompt_tokens:
940
# Final chunk: live kv_transfer_params may be cleared,
941
# so prefer the snapshot from update_state_after_alloc.
942
kv_params = self._req_kv_params.pop(
@@ -857,7 +961,10 @@ class MoRIIOConnectorScheduler:
961
962
for req_id, (req, block_ids) in self._reqs_need_save.items():
963
kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
860
- if req.num_prompt_tokens > len(block_ids) * self.block_size:
964
+ tokens_covered = (
965
+ len(block_ids[self._full_attn_group_idx]) * self._full_attn_block_size
966
+ )
967
+ if req.num_prompt_tokens > tokens_covered:
968
# not last chunk prefill
969
self._reqs_need_pending_save[req_id] = (req, block_ids)
970
continue
@@ -898,7 +1005,7 @@ class MoRIIOConnectorScheduler:
1005
def request_finished(
1006
self,
1007
request: "Request",
901
- block_ids: list[int],
1008
+ block_ids: BlockIds,
1009
) -> tuple[bool, dict[str, Any] | None]:
1010
"""
1011
Once a request is finished, determine whether request blocks
@@ -936,7 +1043,8 @@ class MoRIIOConnectorScheduler:
1043
if self.mode == MoRIIOMode.WRITE:
1044
self._release_write_prefill_blocks(request.request_id, params)
1045
else:
939
- self._reqs_need_recv[request.request_id] = (request, [])
1046
+ empty_groups: BlockIds = [[] for _ in self.blocks_per_sw]
1047
+ self._reqs_need_recv[request.request_id] = (request, empty_groups)
1048
params["do_remote_prefill"] = False
1049
return False, None
1050
@@ -946,10 +1054,9 @@ class MoRIIOConnectorScheduler:
1054
):
1055
return False, None
1056
949
- # computed_block_ids = block_ids if all_full else block_ids[:-1]
950
- computed_block_ids = block_ids
951
- # If prompt < block_size, no xfer so free blocks immediately.
952
- delay_free_blocks = len(computed_block_ids) > 0
1057
+ # On producer
1058
+ computed_block_ids = self.get_exchange_clipped_blocks(block_ids)
1059
+ delay_free_blocks = any(len(group) > 0 for group in computed_block_ids)
1060
1061
if delay_free_blocks:
1062
# Prefill request on remote. It will be read from D upon completion
@@ -1102,6 +1209,12 @@ class MoRIIOConnectorWorker:
1209
self.kv_transfer_config = vllm_config.kv_transfer_config
1210
self.is_producer = self.kv_transfer_config.is_kv_producer
1211
self.layer_to_spec = build_layer_to_spec(kv_cache_config)
1212
+ # we transfer by layer, but store by kv cache group
1213
+ self.layer_to_group: dict[str, int] = {
1214
+ layer_name: group_idx
1215
+ for group_idx, group in enumerate(kv_cache_config.kv_cache_groups)
1216
+ for layer_name in group.layer_names
1217
+ }
1218
1219
if self.is_producer:
1220
set_role(ROLE.PRODUCER)
@@ -1789,8 +1902,6 @@ class MoRIIOConnectorWorker:
1902
self.block_size,
1903
)
1904
self.block_size = first_geometry.block_size
1792
- # TODO(tms): self.block_len needs to be per-layer for sliding window,
1793
- # hybrid attn, etc
1905
# block size in bytes
1906
self.block_len = first_geometry.block_len
1907
self.kv_cache_shape = first_kv_cache.shape
@@ -1803,14 +1914,30 @@ class MoRIIOConnectorWorker:
1914
self.layer_base_addr_index = {}
1915
base_addr_idx = 0
1916
1917
+ hma_enabled = (
1918
+ not self.vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
1919
+ )
1920
for layer_name in kv_caches:
1921
self.layer_base_addr_index[layer_name] = base_addr_idx
1922
geometry = self._get_layer_transfer_geometry(layer_name)
1809
- if geometry.block_size != self.block_size:
1923
+ if hma_enabled:
1924
+ if geometry.block_len != self.block_len:
1925
+ raise ValueError(
1926
+ "MoRIIO KV cache block length mismatch for layer "
1927
+ f"{layer_name}: {geometry.block_len} != {self.block_len}"
1928
+ )
1929
+ elif geometry.block_size != self.block_size:
1930
raise ValueError(
1931
"MoRIIO KV cache block size mismatch for layer "
1932
f"{layer_name}: {geometry.block_size} != {self.block_size}"
1933
)
1934
+ # num_blocks is advertised as a single scalar to the peer, so must it
1935
+ # be uniform
1936
+ if geometry.num_blocks != self.num_blocks:
1937
+ raise ValueError(
1938
+ "MoRIIO KV cache num_blocks mismatch for layer "
1939
+ f"{layer_name}: {geometry.num_blocks} != {self.num_blocks}"
1940
+ )
1941
self.block_lens[layer_name] = geometry.block_len
1942
for cache, region_len in self._iter_layer_registration_regions(layer_name):
1943
base_addr = cache.data_ptr()
@@ -2455,12 +2582,16 @@ class MoRIIOConnectorWorker:
2582
self.remote_dp_size_local = int(meta.remote_dp_size_local)
2583
else:
2584
self.remote_dp_size_local = int(meta.remote_dp_size)
2585
+ # WRITE does not support HMA, so unwrap blocks into flat lists.
2586
+ # remote_block_ids can itself be empty so need to be careful when unwrapping.
2587
+ local_block_ids = meta.local_block_ids[0]
2588
+ remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else []
2589
self.schedule_write_blocks(
2590
request_id=req_id,
2591
transfer_id=meta.transfer_id,
2592
dst_engine_id=meta.remote_engine_id,
2462
- local_block_ids=meta.local_block_ids,
2463
- remote_block_ids=meta.remote_block_ids,
2593
+ local_block_ids=local_block_ids,
2594
+ remote_block_ids=remote_block_ids,
2595
layer_name=layer_name,
2596
kv_layer=kv_layer,
2597
remote_notify_port=meta.remote_notify_port,
@@ -2607,8 +2738,8 @@ class MoRIIOConnectorWorker:
2738
2739
def _read_blocks(
2740
self,
2610
- local_block_ids: list[int],
2611
- remote_block_ids: list[int],
2741
+ local_block_ids: BlockIds,
2742
+ remote_block_ids: BlockIds,
2743
dst_engine_id: str,
2744
request_id: str,
2745
transfer_id: str,
@@ -2653,10 +2784,11 @@ class MoRIIOConnectorWorker:
2784
sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
2785
layer_name
2786
)
2787
+ group_idx = self.layer_to_group[layer_name]
2788
offs = self._compute_block_transfer_offsets(
2789
layer_name,
2658
- local_block_ids,
2659
- remote_block_ids,
2790
+ local_block_ids[group_idx],
2791
+ remote_block_ids[group_idx],
2792
remote_moriio_meta,
2793
remote_tp_size=remote_tp_size,
2794
remote_engine_id=remote_dp_engine_id,