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 65abd910bbcf..1ad43e246e72 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 @@ -131,7 +131,7 @@ def test_hybrid_gdn_remote_decode_truncates_prefill_before_cache_lookup(): # Prefix-cache matching must not mutate the request after its local lookup. with patch.object( scheduler, - "_truncate_mamba_request_for_prefill", + "_truncate_request_for_prefill", side_effect=AssertionError("late Mamba truncation"), ): num_new_tokens, is_async = scheduler.get_num_new_matched_tokens( @@ -142,6 +142,52 @@ def test_hybrid_gdn_remote_decode_truncates_prefill_before_cache_lookup(): assert is_async is False +@pytest.mark.cpu_test +def test_multi_module_mtp_remote_prefill_recomputes_lookahead_window(): + """Mooncake P/D must match the multi-module MTP prefill backoff.""" + scheduler = object.__new__(MooncakeConnectorScheduler) + scheduler._has_mamba = False + scheduler.is_kv_producer = False + scheduler.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=3) + + request = create_request(num_tokens=10, do_remote_prefill=True) + num_new_tokens, is_async = scheduler.get_num_new_matched_tokens( + request, num_computed_tokens=0 + ) + assert num_new_tokens == 8 + assert is_async is True + + request = create_request(num_tokens=10, do_remote_decode=True) + original_tokens = list(request.prompt_token_ids) + scheduler.on_new_request(request) + + assert request.prompt_token_ids == original_tokens[:-2] + assert request._all_token_ids == original_tokens[:-2] + assert request.num_prompt_tokens == 8 + assert request.max_tokens == 1 + assert request.kv_transfer_params["_p_side_truncated"] is True + scheduler.on_new_request(request) + assert request.prompt_token_ids == original_tokens[:-2] + assert request.num_prompt_tokens == 8 + + +@pytest.mark.cpu_test +def test_dense_remote_prefill_keeps_full_prompt_without_mtp(): + """The generalized backoff must preserve ordinary dense Mooncake P/D.""" + scheduler = object.__new__(MooncakeConnectorScheduler) + scheduler._has_mamba = False + scheduler.is_kv_producer = False + scheduler.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0) + + request = create_request(num_tokens=10, do_remote_prefill=True) + assert scheduler.get_num_new_matched_tokens(request, 0) == (10, True) + + request = create_request(num_tokens=10, do_remote_decode=True) + scheduler.on_new_request(request) + assert request.num_prompt_tokens == 10 + assert "_p_side_truncated" not in request.kv_transfer_params + + def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch): monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5") vllm_config = create_vllm_config( 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 41c1337ecd85..1fa2f5d76fb3 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 @@ -944,50 +944,56 @@ def get_sw_clipped_blocks( ] def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1), and for - DSV41 SWA bounded replay, whose window is not replayed after a P/D load - and must be laid out for that token.""" - if (self._has_mamba or self._bounded_replay) and num_prompt_tokens > 1: - return num_prompt_tokens - 1 + """D-side only. Exclude prompt tokens recomputed by the decoder.""" + backoff = self._prefill_backoff() + if backoff and num_prompt_tokens > backoff: + return num_prompt_tokens - backoff return num_prompt_tokens - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. + def _prefill_backoff(self) -> int: + """Number of prompt tokens the decoder recomputes locally. + + Mamba and bounded replay need the last prompt token to derive the + decoder-side state. Multi-module MTP needs the trailing lookahead + window to avoid transferring unverified draft KV from the prefiller. + """ + return max( + 1 if (self._has_mamba or self._bounded_replay) else 0, + self.vllm_config.num_prefill_lookahead_tokens - 1, + ) + + def _truncate_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop tokens the decoder recomputes locally. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the request is preempted and rescheduled.""" params = request.kv_transfer_params + backoff = self._prefill_backoff() if ( - params is not None + backoff + and params is not None and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 + and request.num_prompt_tokens > backoff # A request that skips reading the prefix cache (prompt logprobs) # loads nothing on the decoder, which recomputes its whole prompt. and not request.get_skip_reading_prefix_cache() ): if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() + del request.prompt_token_ids[-backoff:] elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] + request.prompt_embeds = request.prompt_embeds[:-backoff] else: return - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 + del request._all_token_ids[-backoff:] + request.num_prompt_tokens -= backoff 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 or self._bounded_replay) - ): - self._truncate_mamba_request_for_prefill(request) + if params is not None and params.get("do_remote_decode"): + self._truncate_request_for_prefill(request) def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int