From f14369c7f9fcd5f29c33c3ff6753a59664db59e5 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Wed, 26 Aug 2026 16:08:37 +0800 Subject: [PATCH] [Bugfix][Mooncake] Fix Mamba prefill truncation ordering (#53663) Signed-off-by: Ziming Huang Co-authored-by: OpenAI Codex --- .../test_mooncake_connector_hybrid_mamba.py | 38 ++++++++++++++----- .../v1/mooncake/mooncake_connector.py | 12 ++++-- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py index fa6482781a..948f95a461 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py @@ -101,26 +101,46 @@ def test_hybrid_gdn_remote_prefill_uses_mamba_n_minus_one(): @pytest.mark.cpu_test -def test_hybrid_gdn_remote_decode_truncates_prefill_once(): - scheduler = make_hybrid_gdn_scheduler(kv_role="kv_producer") +def test_hybrid_gdn_remote_decode_truncates_prefill_before_cache_lookup(): + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", + kv_role="kv_producer", + ) + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + connector = MooncakeConnector( + vllm_config, + KVConnectorRole.SCHEDULER, + make_hybrid_gdn_kv_cache_config(vllm_config.cache_config.block_size), + ) + scheduler = connector.connector_scheduler + assert scheduler is not None request = create_request(num_tokens=10, do_remote_decode=True) original_tokens = list(request.prompt_token_ids) - num_new_tokens, is_async = scheduler.get_num_new_matched_tokens( - request, num_computed_tokens=0 - ) - - assert num_new_tokens == 0 - assert is_async is False + connector.on_new_request(request) assert request.prompt_token_ids == original_tokens[:-1] assert request._all_token_ids == original_tokens[:-1] assert request.num_prompt_tokens == len(original_tokens) - 1 assert request.max_tokens == 1 assert request.kv_transfer_params["_p_side_truncated"] is True - scheduler.get_num_new_matched_tokens(request, num_computed_tokens=0) + # Re-adding or rescheduling the request must not truncate another token. + scheduler.on_new_request(request) assert request.prompt_token_ids == original_tokens[:-1] + # Prefix-cache matching must not mutate the request after its local lookup. + with patch.object( + scheduler, + "_truncate_mamba_request_for_prefill", + side_effect=AssertionError("late Mamba truncation"), + ): + num_new_tokens, is_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + + assert num_new_tokens == 0 + assert is_async is False + def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 0d4afbd56d..b7e8dbde36 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -516,6 +516,10 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) + def on_new_request(self, request: "Request") -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + def request_finished( self, request: "Request", @@ -687,6 +691,11 @@ class MooncakeConnectorScheduler: request.max_tokens = 1 params["_p_side_truncated"] = True + def on_new_request(self, request: "Request") -> None: + params = request.kv_transfer_params + if params is not None and params.get("do_remote_decode") and self._has_mamba: + self._truncate_mamba_request_for_prefill(request) + def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int, bool]: @@ -726,9 +735,6 @@ class MooncakeConnectorScheduler: if count > 0: return count, True - if params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) - # No remote prefill for this request. return 0, False