diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 4bd32f786a28..1c7b4092dc61 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -10,6 +10,8 @@ """ import importlib +import threading +from collections import defaultdict from dataclasses import asdict from types import SimpleNamespace @@ -17,9 +19,18 @@ import torch from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_utils import get_kv_cache_groups from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager -from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KpoolTailSpec, + KVCacheConfig, + KVCacheSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.request import RequestStatus moriio_connector = importlib.import_module( "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector" @@ -54,6 +65,7 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore def __init__(self, **attrs): self._mamba_group_ids: list[int] = [] self._attn_group_ids: list[int] = [0] + self._attn_group_is_state: list[bool] = [False] self._num_ssm_scratch_blocks = 0 self._is_hma_required = False self.kv_cache_config = SimpleNamespace( @@ -171,7 +183,7 @@ def test_split_block_groups_separates_attention_and_mamba(): sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1]) block_ids = ([1, 2], [99]) attn, mamba = sched.split_block_groups(block_ids) - assert attn == [1, 2] + assert attn == [[1, 2]] assert mamba == [[99]] @@ -208,7 +220,7 @@ def test_split_block_groups_ignores_transfer_disabled_group(): kv_cache_config=config, ) assert sched.split_block_groups(([1, 2], [70, 71], [99])) == ( - [1, 2], + [[1, 2]], [[99]], ) @@ -230,21 +242,220 @@ def test_exchange_blocks_ignore_transfer_disabled_group(): ] -def test_scheduler_rejects_multiple_attention_groups_with_mamba(): - """A drafter owning its own transfer group is refused, with a hint.""" +def test_exchange_clipped_blocks_match_select_for_single_full_attention(): + # Back-compat invariant (K3/Qwen3.5 single full-attention group): when every + # group has blocks_per_sw == 0, get_exchange_clipped_blocks must reduce to + # select_transfer_block_ids, so the wire payload is byte-identical to the + # pre-multigroup path. Locks the single-group claim against future edits to + # get_exchange_clipped_blocks. + def select(block_ids): + # One full-attention attention group (0) + one mamba group (2); the + # middle group is not transfer-enabled. + return (block_ids[0], block_ids[2]) + + config = SimpleNamespace( + kv_cache_groups=[None, None, None], + select_transfer_block_ids=select, + ) + raw = ([1, 2, 3], [70, 71], [90, 91, 92]) + + hma = _FakeScheduler( + kv_cache_config=config, _is_hma_required=True, blocks_per_sw=[0, 0] + ) + plain = _FakeScheduler(kv_cache_config=config, _is_hma_required=False) + + clipped = hma.get_exchange_clipped_blocks(raw) + + # Equal to select, and identical to the non-HMA path (no clipping at all). + assert clipped == [list(group) for group in select(raw)] + assert clipped == plain.get_exchange_clipped_blocks(raw) + + +def test_scheduler_rejects_hybrid_without_attention_group(): config = SimpleNamespace( kv_cache_groups=[ - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), ], ) config.transfer_groups = tuple(config.kv_cache_groups) - with pytest.raises(moriio_common.MoRIIOError, match="separate attention group"): + with pytest.raises(moriio_common.MoRIIOError, match="attention group"): moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) +def _glm5_next_mtp_kv_cache_config() -> KVCacheConfig: + """GLM-5.3-Flash-shaped groups from the real grouping, with an MTP layer. + + KDA, KDA, MLA + kpool indexer + kpool tail, repeated twice, then the MTP + layer 6 (an MLA layer registered last). Grouping yields + [MLA + indexer, kpool tail, KDA, KDA]. + """ + kpool, num_spec = 4, 1 + mamba = _mamba_spec(num_speculative_blocks=num_spec) + specs: dict[str, KVCacheSpec] = {} + for i in range(7): + if i % 3 != 2 and i != 6: + specs[f"layers.{i}.linear_attn"] = mamba + continue + specs[f"layers.{i}.attn"] = MLAAttentionSpec( + block_size=1024, num_kv_heads=1, head_size=576, dtype=torch.bfloat16 + ) + specs[f"layers.{i}.indexer"] = MLAAttentionSpec( + block_size=1024, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=kpool, + ) + specs[f"layers.{i}.tail"] = KpoolTailSpec( + block_size=2 * kpool, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=2 * kpool, + ) + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + speculative_config=None, + ) + groups = get_kv_cache_groups(vllm_config, specs) + return KVCacheConfig(num_blocks=64, kv_cache_tensors=[], kv_cache_groups=groups) + + +def _glm5_next_mtp_scheduler(kv_role: str = "kv_consumer"): + vllm_config = _gate_vllm_config( + speculative_config=_spec_config("mtp"), + block_size=1024, + num_lookahead_tokens=1, + num_speculative_tokens=1, + ) + vllm_config.kv_transfer_config.kv_role = kv_role + config = _glm5_next_mtp_kv_cache_config() + return moriio_connector.MoRIIOConnectorScheduler(vllm_config, "engine", config) + + +@pytest.mark.usefixtures("restore_moriio_role") +def test_scheduler_accepts_glm5_next_groups_with_mtp(): + """Several attention groups (MLA + indexer, kpool tail) plus KDA, with MTP.""" + scheduler = _glm5_next_mtp_scheduler() + groups = scheduler.kv_cache_config.transfer_groups + + # The MTP layer shares the target's MLA and tail groups. + assert {"layers.6.attn", "layers.6.indexer"} <= set(groups[0].layer_names) + assert "layers.6.tail" in groups[1].layer_names + assert scheduler._attn_group_ids == [0, 1] + assert scheduler._mamba_group_ids == [2, 3] + assert scheduler._attn_group_is_state == [False, True] + assert scheduler.blocks_per_sw == [0, 2, 0, 0] + assert scheduler._num_ssm_scratch_blocks == 1 + # 1 recomputed token + 1 padded spec row + 1 drafter lookahead slot. + assert scheduler._max_decode_tail_blocks == 1 + + +@pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize("external_tokens", [0, 1500]) +def test_glm5_next_mtp_read_pairs_every_group(external_tokens): + """P payload -> D alignment -> zeroing filter across all transfer groups.""" + producer = _glm5_next_mtp_scheduler("kv_producer") + p_request = SimpleNamespace( + request_id="req", + kv_transfer_params={"do_remote_decode": True, "transfer_id": "tx"}, + status=RequestStatus.FINISHED_LENGTH_CAPPED, + ) + # Mamba groups: [running state, MTP scratch slot]. + p_blocks = ([10, 11], [50], [90, 91], [80, 81]) + delay_free, p_params = _FakeConnector(producer).request_finished_all_groups( + p_request, p_blocks + ) + assert delay_free is True + assert p_params["remote_block_ids"] == [[10, 11], [50], [90], [80]] + + consumer = _glm5_next_mtp_scheduler() + d_request = _make_read_request(p_params["remote_block_ids"]) + # The decode's MLA group holds one extra local block for the spec tail. + d_blocks = _FakeBlocks(([100, 101, 102], [300], [200, 201], [210, 211])) + consumer.update_state_after_alloc( + d_request, d_blocks, num_external_tokens=external_tokens + ) + + local = consumer._reqs_need_recv["req"][1] + if external_tokens: + assert local == [[100, 101], [300], [200], [210]] + assert consumer._req_kv_params["req"]["remote_block_ids"] == [ + [10, 11], + [50], + [90], + [80], + ] + else: + # A full local MLA hit still reads the kpool tail ring and KDA state. + assert local == [[], [300], [200], [210]] + + output = SchedulerOutput.make_empty() + output.new_block_ids_to_zero = [100, 101, 102, 300] + consumer.build_connector_meta(output) + expected_zero = [102] if external_tokens else [100, 101, 102] + assert output.new_block_ids_to_zero == expected_zero + + +def test_read_blocks_routes_each_hybrid_layer_to_its_group(): + offsets: dict[str, tuple[list[int], list[int]]] = {} + worker = _FakeWorker( + mode=MoRIIOMode.READ, + world_size=1, + _has_mamba=True, + _num_attn_transfer_groups=2, + _num_mamba_transfer_groups=1, + _hybrid_payload_index_by_layer={"mla.0": 0, "tail.0": 1, "kda.0": 2}, + layer_name_to_local_kv_cache_metadata={"mla.0": 0, "tail.0": 0, "kda.0": 0}, + moriio_wrapper=SimpleNamespace(lock=threading.Lock(), shutdown=lambda: None), + moriio_config=SimpleNamespace(transfer_timeout=1.0), + _recving_transfers=defaultdict(dict), + _recving_transfers_start={}, + _recving_local_blocks={}, + _reads_issued_this_step=[], + _mamba_reads_this_step=[], + _recving_transfers_callback_addr={}, + ) + worker._remote_tp_rank = lambda remote_tp_size: 0 + worker.get_engine_name_with_dp = lambda engine_id, dp_rank: engine_id + worker._get_built_session = lambda engine_id: ([None, None], None) + worker._region_session_indices = lambda layer_name: [0, 1] + worker._is_mamba_layer = lambda layer_name: layer_name.startswith("kda") + + def _offsets(layer_name, local, remote, meta, **kwargs): + offsets[layer_name] = (local, remote) + return [], [], [] + + def _mamba_reads(layer_name, sessions, regions, local, remote, *args): + offsets[layer_name] = (local, remote) + return [] + + worker._compute_block_transfer_offsets = _offsets + worker._post_read_with_backoff = lambda *args: None + worker._post_mamba_reads = _mamba_reads + + worker._read_blocks( + [[100, 101], [300], [200]], + [[10, 11], [50], [90]], + dst_engine_id="prefill", + request_id="req", + transfer_id="tx", + remote_host="127.0.0.1", + remote_notify_port=6001, + remote_tp_size=1, + ) + + assert offsets == { + "mla.0": ([100, 101], [10, 11]), + "tail.0": ([300], [50]), + "kda.0": ([200], [90]), + } + assert worker._recving_local_blocks["req"] == [100, 101, 300] + + def test_split_block_groups_preserves_multiple_mamba_groups(): sched = _FakeScheduler( _has_mamba=True, @@ -253,7 +464,7 @@ def test_split_block_groups_preserves_multiple_mamba_groups(): ) assert sched.split_block_groups(([1, 2], [40], [70])) == ( - [1, 2], + [[1, 2]], [[40], [70]], ) @@ -287,23 +498,6 @@ def test_scheduler_rejects_mamba_group_without_two_states(): moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) -def test_scheduler_rejects_non_dspark_speculative_hybrid_read(): - config = SimpleNamespace( - kv_cache_groups=[ - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), - ], - ) - config.transfer_groups = tuple(config.kv_cache_groups) - - with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"): - moriio_connector.MoRIIOConnectorScheduler( - _gate_vllm_config(speculative_config=_spec_config("ngram")), - "engine", - config, - ) - - def _dspark_wrapped_attention_config(*, dcp_sharded: bool = True): """Mamba first, so a wrapped full-attention group at index 1 has to be found. @@ -418,7 +612,7 @@ def test_split_block_groups_keeps_only_running_state(): _attn_group_ids=[0], _mamba_group_ids=[1], ) - assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]]) + assert sched.split_block_groups(([1], [40, 41])) == ([[1]], [[41]]) @pytest.mark.parametrize( @@ -439,7 +633,7 @@ def test_split_block_groups_strips_dspark_scratch_slots(blocks, expected): _num_ssm_scratch_blocks=2, ) - assert sched.split_block_groups(([1], blocks)) == ([1], [expected]) + assert sched.split_block_groups(([1], blocks)) == ([[1]], [expected]) # -------------------------------------------------------------------------- @@ -448,12 +642,12 @@ def test_split_block_groups_strips_dspark_scratch_slots(blocks, expected): def test_request_finished_all_groups_carries_attn_and_mamba_in_one_field(): seen = {} - def _fake_request_finished(request, attn_block_ids, mamba_block_groups): - seen["attn"] = list(attn_block_ids) + def _fake_request_finished(request, attn_block_groups, mamba_block_groups): + seen["attn"] = [list(group) for group in attn_block_groups] seen["mamba"] = [list(group) for group in mamba_block_groups] return True, { "do_remote_prefill": True, - "remote_block_ids": [attn_block_ids, *mamba_block_groups], + "remote_block_ids": [*attn_block_groups, *mamba_block_groups], } sched = _FakeScheduler(_has_mamba=True, _attn_group_ids=[0], _mamba_group_ids=[1]) @@ -468,7 +662,7 @@ def _fake_request_finished(request, attn_block_ids, mamba_block_groups): # Attention ids drive request_finished; the mamba slot rides the SAME # remote_block_ids channel as [attn, *mamba_groups], so the proxy/router # needs no KDA-specific field. - assert seen["attn"] == [10, 11] + assert seen["attn"] == [[10, 11]] assert seen["mamba"] == [[42]] assert params["remote_block_ids"] == [[10, 11], [42]] assert "remote_mamba_block_ids" not in params @@ -966,7 +1160,7 @@ def test_mamba_tp_ratio_rejects_heterogeneous_tp(remote_tp_size): def test_worker_selects_each_mamba_layers_own_block_group(): - worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 1, "kda.1": 2}) + worker = _FakeWorker(_hybrid_payload_index_by_layer={"kda.0": 1, "kda.1": 2}) groups = ([10], [40], [70]) assert worker._mamba_blocks_for_layer("kda.0", groups) == [40] @@ -974,7 +1168,7 @@ def test_worker_selects_each_mamba_layers_own_block_group(): def test_worker_rejects_missing_mamba_block_group(): - worker = _FakeWorker(_mamba_payload_index_by_layer={"kda.0": 2}) + worker = _FakeWorker(_hybrid_payload_index_by_layer={"kda.0": 2}) with pytest.raises(moriio_common.MoRIIOError, match="no block group"): worker._mamba_blocks_for_layer("kda.0", ([10], [40])) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 082643778e7d..6a3d8fa92d6b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -94,10 +94,13 @@ from vllm.v1.attention.selector import get_attn_backend from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, KVCacheConfig, MambaSpec, SlidingWindowSpec, + get_kv_cache_spec_sliding_window, is_full_attention_spec, + iter_layer_specs, ) from vllm.v1.metrics.cache_hit_source import CacheHitSource from vllm.v1.outputs import KVConnectorOutput @@ -445,19 +448,6 @@ def _split_kv_cache_group_kinds( return attn, mamba -def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None: - speculative_config = vllm_config.speculative_config - if speculative_config is None: - return - # DSpark is allowlisted because it is the only hybrid speculative method - # validated end to end with MoRIIO READ. - if not speculative_config.use_dspark(): - raise MoRIIOError( - "MoRIIO hybrid READ supports DSpark speculative decoding only, got " - f"method={speculative_config.method!r}" - ) - - def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None: """Validate the common packed-state contract and return its representative.""" specs = tuple(specs) @@ -501,20 +491,22 @@ def __init__( self._has_mamba = bool(self._mamba_group_ids) mamba_spec = None if self._has_mamba: - if len(self._attn_group_ids) != 1: + # Scope: hybrid READ carries every target attention group and + # supports any speculative method whose drafter *shares* the + # target's attention group(s) (e.g. MTP, DSpark). A drafter that + # owns a *separate* transferable attention group is out of scope: + # output stays correct because the target re-verifies every draft, + # but that layout is untested here, so we only require that at + # least one attention group is present. + if not self._attn_group_ids: raise MoRIIOError( - "MoRIIO hybrid READ requires exactly one transferable " - "attention group, got " - f"{len(self._attn_group_ids)}; a drafter that owns a " - "separate attention group is not supported, give it the " - "target's group or disable KV transfer for it" + "MoRIIO hybrid READ requires a transferable attention group" ) mamba_specs = [ cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec) for group_id in self._mamba_group_ids ] mamba_spec = _validate_mamba_specs(mamba_specs) - _validate_hybrid_speculation(vllm_config) if self.mode != MoRIIOMode.READ: raise MoRIIOError( "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " @@ -526,15 +518,23 @@ def __init__( self.block_size = vllm_config.cache_config.block_size attn_block_size = self.block_size max_decode_tail_tokens = 1 + vllm_config.num_lookahead_tokens + # Per attention group: whether its blocks hold per-request state (e.g. a + # kpool tail ring) that must be read even on a full local prefix hit. + self._attn_group_is_state: list[bool] = [] if self._has_mamba: - attn_spec = kv_cache_config.transfer_groups[ - self._attn_group_ids[0] - ].kv_cache_spec - attn_block_size = attn_spec.block_size - if attn_spec.dcp_sharded: - attn_block_size *= ( - vllm_config.parallel_config.decode_context_parallel_size - ) + attn_specs = [ + kv_cache_config.transfer_groups[group_id].kv_cache_spec + for group_id in self._attn_group_ids + ] + self._attn_group_is_state = [ + not spec.prefix_cacheable for spec in attn_specs + ] + dcp_size = vllm_config.parallel_config.decode_context_parallel_size + # The smallest attention block bounds every group's local tail. + attn_block_size = min( + spec.block_size * (dcp_size if spec.dcp_sharded else 1) + for spec in attn_specs + ) # Admission can pad the final prompt token to 1 + num_spec rows; # allocate_slots reserves drafter lookahead in addition to those. max_decode_tail_tokens += vllm_config.num_speculative_tokens @@ -554,9 +554,9 @@ def __init__( unsupported = [ type(g.kv_cache_spec).__name__ for g in kv_cache_config.transfer_groups - if not ( - is_full_attention_spec(g.kv_cache_spec) - or isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec)) + if not all( + isinstance(spec, (FullAttentionSpec, SlidingWindowSpec, MambaSpec)) + for spec in iter_layer_specs(g.kv_cache_spec) ) ] if unsupported: @@ -576,8 +576,8 @@ def __init__( ) sw_sizes_tokens: list[tuple[int, int]] = [ - (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) - if isinstance(g.kv_cache_spec, SlidingWindowSpec) + (sliding_window, g.kv_cache_spec.block_size) + if (sliding_window := get_kv_cache_spec_sliding_window(g.kv_cache_spec)) else (0, self.block_size) for g in kv_cache_config.transfer_groups ] @@ -983,26 +983,38 @@ def update_state_after_alloc( # destination blocks before this hook runs. Pair # the attention and recurrent-state suffixes # independently. - remote_attn = list(remote_block_ids[0]) - remote_mamba_groups = [ - list(group) for group in remote_block_ids[1:] - ] - attn_block_ids, mamba_block_groups = ( + attn_block_groups, mamba_block_groups = ( self.split_block_groups(blocks.get_block_ids()) ) - if len(mamba_block_groups) != len(remote_mamba_groups): + num_attn = len(attn_block_groups) + if len(remote_block_ids) != num_attn + len( + mamba_block_groups + ): raise MoRIIOError( - "MoRIIO hybrid READ Mamba group " - "count mismatch: " - f"local={len(mamba_block_groups)}, " - f"remote={len(remote_mamba_groups)}" + "MoRIIO hybrid READ block group " + "count mismatch: local=" + f"{num_attn + len(mamba_block_groups)}, " + f"remote={len(remote_block_ids)}" ) - local_attn = attn_block_ids - local_attn, remote_attn = self._align_read_blocks( - local_attn, - remote_attn, - self._max_decode_tail_blocks, - ) + remote_mamba_groups = [ + list(group) for group in remote_block_ids[num_attn:] + ] + aligned_local_attn: list[list[int]] = [] + aligned_remote_attn: list[list[int]] = [] + for local_group, remote_group in zip( + attn_block_groups, + remote_block_ids[:num_attn], + strict=True, + ): + aligned_local, aligned_remote = ( + self._align_read_blocks( + local_group, + list(remote_group), + self._max_decode_tail_blocks, + ) + ) + aligned_local_attn.append(aligned_local) + aligned_remote_attn.append(aligned_remote) aligned_local_mamba: list[list[int]] = [] aligned_remote_mamba: list[list[int]] = [] # Scratch slots were removed by @@ -1022,9 +1034,12 @@ def update_state_after_alloc( ) aligned_local_mamba.append(aligned_local) aligned_remote_mamba.append(aligned_remote) - local_block_ids = [local_attn, *aligned_local_mamba] + local_block_ids = [ + *aligned_local_attn, + *aligned_local_mamba, + ] adjusted_remote_block_ids = [ - remote_attn, + *aligned_remote_attn, *aligned_remote_mamba, ] else: @@ -1055,10 +1070,18 @@ def update_state_after_alloc( if self._has_mamba: # Attention can be a complete local hit, but the # recurrent state is never prefix-cacheable. - _, mamba_block_groups = self.split_block_groups( - blocks.get_block_ids() + attn_block_groups, mamba_block_groups = ( + self.split_block_groups(blocks.get_block_ids()) ) - local_block_ids = [[], *mamba_block_groups] + local_block_ids = [ + group if is_state else [] + for group, is_state in zip( + attn_block_groups, + self._attn_group_is_state, + strict=True, + ) + ] + local_block_ids.extend(mamba_block_groups) else: # Nothing to pull, but retain one empty list per # cache group so the worker can notify P. @@ -1298,11 +1321,12 @@ def build_connector_meta( ): # Hybrid models zero recycled attention pages that held Mamba state. # Host-submitted READs overwrite these pages and can race zeroing. - # Hybrid READ metadata puts the aligned attention pages first. + # Hybrid READ metadata puts the aligned attention groups first. read_dst_block_ids = { b for _, block_ids in self._reqs_need_recv.values() - for b in block_ids[0] + for group in block_ids[: len(self._attn_group_ids)] + for b in group } scheduler_output.new_block_ids_to_zero = [ b @@ -1352,19 +1376,24 @@ def shutdown(self): def split_block_groups( self, block_ids: list[list[int]] | tuple[list[int], ...] - ) -> tuple[list[int], list[list[int]]]: + ) -> tuple[list[list[int]], list[list[int]]]: """Select transferable attention and Mamba block groups. - The wire payload stores the attention group first, followed by every - Mamba group in transfer-group order. Only each group's running-state - slot is transferred. + The wire payload stores every attention group first, followed by every + Mamba group, each in transfer-group order. Sliding-window groups keep + only their window; Mamba groups only their running-state slot. """ if not block_ids: return [], [] if not self._has_mamba: raise MoRIIOError("Mamba block groups requested for a non-hybrid model") - transfer_block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids) - attn = list(transfer_block_ids[self._attn_group_ids[0]]) + # NOTE: get_exchange_clipped_blocks also window-clips a *sliding-window + # attention* group (blocks_per_sw > 0). The single-attention-group + # wire back-compat guarantee is therefore scoped to full-attention + # groups (K3/Qwen3.5 MLA), for which blocks_per_sw == 0 and this + # reduces to select_transfer_block_ids (see the back-compat test). + transfer_block_ids = self.get_exchange_clipped_blocks(block_ids) + attn = [list(transfer_block_ids[group_id]) for group_id in self._attn_group_ids] mamba_groups = [ clip_ssm_state_blocks( list(transfer_block_ids[group_id]), @@ -1413,7 +1442,7 @@ def request_finished( should be freed now or will be sent asynchronously and freed later. ``block_ids`` normally contains all cache groups. Hybrid callers from - ``request_finished_all_groups`` pass a flat attention list and the + ``request_finished_all_groups`` pass the attention groups and the recurrent-state groups separately in ``mamba_block_groups``. """ request_id = request.request_id @@ -1453,7 +1482,7 @@ def request_finished( # mamba slot here would re-trip the KDA guard in _read_blocks. if mamba_block_groups: recv_blocks: BlockIds = [ - list(cast(list[int], block_ids)), + *[list(group) for group in cast(BlockIds, block_ids)], *[list(group) for group in mamba_block_groups], ] elif block_ids and isinstance(block_ids[0], int): @@ -1478,7 +1507,7 @@ def request_finished( # On producer if mamba_block_groups is not None: computed_block_ids: BlockIds = [ - list(cast(list[int], block_ids)), + *[list(group) for group in cast(BlockIds, block_ids)], *[list(group) for group in mamba_block_groups], ] elif self._has_mamba: @@ -1643,11 +1672,17 @@ def __init__( # we transfer by layer, but store by kv cache group self.layer_to_group = dict(kv_cache_config.transfer_group_index_by_layer) self._transfer_layer_names = set(self.layer_to_group) - _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config) + attn_transfer_group_ids, mamba_transfer_group_ids = _split_kv_cache_group_kinds( + kv_cache_config + ) + self._num_attn_transfer_groups = len(attn_transfer_group_ids) self._num_mamba_transfer_groups = len(mamba_transfer_group_ids) - self._mamba_payload_index_by_layer = { + # Hybrid payload order: attention groups, then Mamba groups. + self._hybrid_payload_index_by_layer = { layer_name: payload_index - for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1) + for payload_index, group_id in enumerate( + [*attn_transfer_group_ids, *mamba_transfer_group_ids] + ) for layer_name in kv_cache_config.transfer_groups[group_id].layer_names } @@ -3418,7 +3453,7 @@ def _mamba_blocks_for_layer( layer_name: str, block_groups: BlockIds, ) -> list[int]: - payload_index = self._mamba_payload_index_by_layer.get(layer_name) + payload_index = self._hybrid_payload_index_by_layer.get(layer_name) if payload_index is None or payload_index >= len(block_groups): raise MoRIIOError( f"KDA layer {layer_name} has no block group in the transfer payload" @@ -3614,7 +3649,8 @@ def _read_blocks( return if self._has_mamba: - expected_groups = 1 + self._num_mamba_transfer_groups + num_attn = self._num_attn_transfer_groups + expected_groups = num_attn + self._num_mamba_transfer_groups if ( len(local_block_ids) != expected_groups or len(remote_block_ids) != expected_groups @@ -3624,8 +3660,8 @@ def _read_blocks( f"expected={expected_groups}, local={len(local_block_ids)}, " f"remote={len(remote_block_ids)}" ) - local_attn = list(local_block_ids[0]) - remote_attn = list(remote_block_ids[0]) + local_attn = [b for group in local_block_ids[:num_attn] for b in group] + remote_attn = [b for group in remote_block_ids[:num_attn] for b in group] self._mamba_tp_ratio(remote_tp_size) else: local_attn = remote_attn = [] @@ -3679,13 +3715,13 @@ def _read_blocks( _sq_deadline, ) else: - if self._has_mamba: - local_layer_blocks = local_attn - remote_layer_blocks = remote_attn - else: - group_idx = self.layer_to_group[layer_name] - local_layer_blocks = local_block_ids[group_idx] - remote_layer_blocks = remote_block_ids[group_idx] + group_idx = ( + self._hybrid_payload_index_by_layer + if self._has_mamba + else self.layer_to_group + )[layer_name] + local_layer_blocks = local_block_ids[group_idx] + remote_layer_blocks = remote_block_ids[group_idx] offs = self._compute_block_transfer_offsets( layer_name, local_layer_blocks,