From e81abe4dd7f7d5eb689f918d8ad18bb212db0093 Mon Sep 17 00:00:00 2001 From: Taoyu Zhu Date: Sun, 18 Oct 2026 14:23:06 +0000 Subject: [PATCH] [Bugfix][KV Connector][ROCm] MoRIIO READ: skip zeroing READ destinations for mixed-precision caches Co-authored-by: Claude Opus 5.5 Signed-off-by: Taoyu Zhu --- .../unit/test_moriio_hma_scheduler.py | 17 +++++++++++++++-- .../kv_connector/v1/moriio/moriio_connector.py | 13 +++++-------- 2 files changed, 20 insertions(+), 10 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 1120062e5410..36c739c7ed3d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -763,7 +763,7 @@ def _make_read_request(remote_block_ids): 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.""" + """Full READ destinations skip zeroing; local tails still zero.""" sched = _make_read_scheduler() sched._has_mamba = has_mamba remote_blocks = [[10, 11], [91]] if has_mamba else [[10, 11]] @@ -787,7 +787,7 @@ def test_update_state_drops_decode_recompute_tail_block( 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: + if zero_ids and external_tokens: before["new_block_ids_to_zero"] = zero_ids[1:] assert asdict(output) == before @@ -798,6 +798,19 @@ def test_update_state_drops_decode_recompute_tail_block( assert output.new_block_ids_to_zero == [100, 101, 102] +def test_attention_only_read_skips_zeroing_in_every_group(): + """Mixed-precision attention-only hybrids zero pages in every READ group.""" + sched = _make_read_scheduler() + sched._has_mamba = False + request = _make_read_request([[10, 11], [12]]) + sched._reqs_need_recv["req"] = (request, [[100, 101], [300]]) + + output = SchedulerOutput.make_empty() + output.new_block_ids_to_zero = [100, 101, 102, 300, 301] + sched.build_connector_meta(output) + assert output.new_block_ids_to_zero == [102, 301] + + def test_update_state_pairs_shorter_local_blocks_with_remote_suffix(): sched = _make_read_scheduler() request = _make_read_request([[10, 11, 12], [92]]) 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 9760d8bef1a0..808e73192018 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 @@ -1361,18 +1361,15 @@ def build_connector_meta( 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. + if self.mode == MoRIIOMode.READ and scheduler_output.new_block_ids_to_zero: + # Mamba and mixed-precision caches zero new pages on the compute + # stream. Host-submitted READs overwrite these pages and can race it. # 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] + for group in (block_ids[:1] if self._has_mamba else block_ids) + for b in group } scheduler_output.new_block_ids_to_zero = [ b