From df417f780a19b68cebc4379942de6ea0e6f6aceb Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:35:56 +0800 Subject: [PATCH] [Bugfix][MoRIIO] Exclude synchronous READ destinations from KV zeroing (#59164) Signed-off-by: Yichao Zhu Signed-off-by: whx-sjtu Signed-off-by: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Co-authored-by: Yichao Zhu Co-authored-by: Codex Co-authored-by: Yifan Qiao --- .../unit/test_moriio_hma_scheduler.py | 64 ++++++++++++------- .../kv_transfer/kv_connector/v1/base.py | 6 +- .../v1/moriio/moriio_connector.py | 18 ++++++ 3 files changed, 64 insertions(+), 24 deletions(-) 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 28bc8bebe2..1de729416e 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -7,32 +7,20 @@ These exercise the highest-risk, otherwise-untested scheduler logic without a GPU or the ``mori`` runtime: per-group block-id splitting, the READ/WRITE ``N-1`` token accounting, P-side prompt truncation, and offset-template cache wiring. - -Like ``test_moriio_kv_layout.py`` the whole module is skipped unless it is -running on ROCm with ``mori`` installed (importing the connector pulls in -``mori``). The authoritative run happens on the MIA recipe image. """ import importlib -import importlib.util +from dataclasses import asdict from types import SimpleNamespace import pytest import torch -from vllm.platforms import current_platform from vllm.v1.core.block_pool import BlockPool +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 -mori_available = importlib.util.find_spec("mori") is not None - -if not (current_platform.is_rocm() and mori_available): - pytest.skip( - "MoRIIOs are only available on ROCm with mori package installed", - allow_module_level=True, - ) - moriio_connector = importlib.import_module( "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector" ) @@ -525,6 +513,8 @@ def _make_read_scheduler(): request_id_to_transfer_id={}, transfer_id_to_request_id={}, _reqs_need_recv={}, + _reqs_need_save={}, + _reqs_need_send={}, _req_kv_params={}, _max_decode_tail_blocks=1, ) @@ -538,24 +528,52 @@ def _make_read_request(remote_block_ids): "transfer_id": "tx", "remote_engine_id": "prefill", "remote_block_ids": remote_block_ids, + "remote_host": "127.0.0.1", + "remote_handshake_port": 6000, + "remote_notify_port": 6001, }, ) -def test_update_state_drops_decode_recompute_tail_block(): +@pytest.mark.parametrize("external_tokens", [0, 256]) +@pytest.mark.parametrize("has_mamba", [False, True]) +@pytest.mark.parametrize("zero_ids", [None, [], [101], [101, 102, 202, 999]]) +def test_update_state_drops_decode_recompute_tail_block( + external_tokens, has_mamba, zero_ids +): + """Only full hybrid READ destinations skip zeroing; local tails still zero.""" sched = _make_read_scheduler() - request = _make_read_request([[10, 11], [91]]) + sched._has_mamba = has_mamba + remote_blocks = [[10, 11], [91]] if has_mamba else [[10, 11]] + request = _make_read_request(remote_blocks) blocks = _FakeBlocks( - all_groups=([100, 101, 102], [200, 201, 202]), + all_groups=([100, 101, 102], [200, 201, 202]) + if has_mamba + else ([100, 101, 102],), ) - sched.update_state_after_alloc(request, blocks, num_external_tokens=256) + sched.update_state_after_alloc(request, blocks, num_external_tokens=external_tokens) - assert sched._reqs_need_recv["req"][1] == [[100, 101], [202]] - assert sched._req_kv_params["req"]["remote_block_ids"] == [ - [10, 11], - [91], - ] + expected_blocks = [[100, 101] if external_tokens else []] + if has_mamba: + expected_blocks.append([202]) + assert sched._reqs_need_recv["req"][1] == expected_blocks + assert sched._req_kv_params["req"]["remote_block_ids"] == remote_blocks + + output = SchedulerOutput.make_empty() + output.new_block_ids_to_zero = None if zero_ids is None else list(zero_ids) + before = asdict(output) + meta = sched.build_connector_meta(output) + assert meta.reqs_to_recv["req"].local_block_ids == expected_blocks + if zero_ids and has_mamba and external_tokens: + before["new_block_ids_to_zero"] = zero_ids[1:] + assert asdict(output) == before + + # A later local allocation of the same pages must still be zeroed. + output.new_block_ids_to_zero = [100, 101, 102] + meta = sched.build_connector_meta(output) + assert not meta.reqs_to_recv + assert output.new_block_ids_to_zero == [100, 101, 102] def test_update_state_pairs_shorter_local_blocks_with_remote_suffix(): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index bb9d4597ce..71593f59e5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -583,7 +583,11 @@ class KVConnectorBase_V1(ABC): """Build the connector metadata for this step. This function should NOT modify fields in the scheduler_output. - Also, calling this function will reset the state of the connector. + FIXME: one exception: + synchronous READ connectors may remove attention blocks they fully + overwrite this step from new_block_ids_to_zero. They must complete + those loads before the blocks are used and fail the step if a load fails. + Calling this function will reset the state of the connector. Args: scheduler_output (SchedulerOutput): the scheduler output object. 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 167c7ef3c2..5a8a917664 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 @@ -1232,6 +1232,24 @@ class MoRIIOConnectorScheduler: local_block_ids=block_ids, kv_transfer_params=kv_params, ) + if ( + self.mode == MoRIIOMode.READ + and self._has_mamba + and scheduler_output.new_block_ids_to_zero + ): + # 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. + read_dst_block_ids = { + b + for _, block_ids in self._reqs_need_recv.values() + for b in block_ids[0] + } + scheduler_output.new_block_ids_to_zero = [ + b + for b in scheduler_output.new_block_ids_to_zero + if b not in read_dst_block_ids + ] for req_id, (req, block_ids) in self._reqs_need_save.items(): kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})