[Bugfix][Mooncake] Fix Mamba prefill truncation ordering (#53663)
Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Ziming Huang committed
Aug 26, 2026 at 16:08 UTC
f14369c7f9fcd5f29c33c3ff6753a59664db59e5
2 files changed
+38
-12
tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py
+29
-9
@@ -101,26 +101,46 @@ def test_hybrid_gdn_remote_prefill_uses_mamba_n_minus_one():
101
102
103
@pytest.mark.cpu_test
104
-def test_hybrid_gdn_remote_decode_truncates_prefill_once():
105
- scheduler = make_hybrid_gdn_scheduler(kv_role="kv_producer")
104
+def test_hybrid_gdn_remote_decode_truncates_prefill_before_cache_lookup():
105
+ vllm_config = create_vllm_config(
106
+ kv_connector="MooncakeConnector",
107
+ kv_role="kv_producer",
108
+ )
109
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
110
+ connector = MooncakeConnector(
111
+ vllm_config,
112
+ KVConnectorRole.SCHEDULER,
113
+ make_hybrid_gdn_kv_cache_config(vllm_config.cache_config.block_size),
114
+ )
115
+ scheduler = connector.connector_scheduler
116
+ assert scheduler is not None
117
request = create_request(num_tokens=10, do_remote_decode=True)
118
original_tokens = list(request.prompt_token_ids)
119
109
- num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
110
- request, num_computed_tokens=0
111
- )
112
-
113
- assert num_new_tokens == 0
114
- assert is_async is False
120
+ connector.on_new_request(request)
121
assert request.prompt_token_ids == original_tokens[:-1]
122
assert request._all_token_ids == original_tokens[:-1]
123
assert request.num_prompt_tokens == len(original_tokens) - 1
124
assert request.max_tokens == 1
125
assert request.kv_transfer_params["_p_side_truncated"] is True
126
121
- scheduler.get_num_new_matched_tokens(request, num_computed_tokens=0)
127
+ # Re-adding or rescheduling the request must not truncate another token.
128
+ scheduler.on_new_request(request)
129
assert request.prompt_token_ids == original_tokens[:-1]
130
131
+ # Prefix-cache matching must not mutate the request after its local lookup.
132
+ with patch.object(
133
+ scheduler,
134
+ "_truncate_mamba_request_for_prefill",
135
+ side_effect=AssertionError("late Mamba truncation"),
136
+ ):
137
+ num_new_tokens, is_async = scheduler.get_num_new_matched_tokens(
138
+ request, num_computed_tokens=0
139
+ )
140
+
141
+ assert num_new_tokens == 0
142
+ assert is_async is False
143
+
144
145
def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
146
monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+9
-3
@@ -516,6 +516,10 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
516
assert self.connector_scheduler is not None
517
return self.connector_scheduler.build_connector_meta(scheduler_output)
518
519
+ def on_new_request(self, request: "Request") -> None:
520
+ assert self.connector_scheduler is not None
521
+ self.connector_scheduler.on_new_request(request)
522
+
523
def request_finished(
524
self,
525
request: "Request",
@@ -687,6 +691,11 @@ class MooncakeConnectorScheduler:
691
request.max_tokens = 1
692
params["_p_side_truncated"] = True
693
694
+ def on_new_request(self, request: "Request") -> None:
695
+ params = request.kv_transfer_params
696
+ if params is not None and params.get("do_remote_decode") and self._has_mamba:
697
+ self._truncate_mamba_request_for_prefill(request)
698
+
699
def get_num_new_matched_tokens(
700
self, request: "Request", num_computed_tokens: int
701
) -> tuple[int, bool]:
@@ -726,9 +735,6 @@ class MooncakeConnectorScheduler:
735
if count > 0:
736
return count, True
737
729
- if params.get("do_remote_decode") and self._has_mamba:
730
- self._truncate_mamba_request_for_prefill(request)
731
-
738
# No remote prefill for this request.
739
return 0, False
740