[Bugfix][KV Connector][Mooncake] Apply multi-module MTP prefill backoff in P/D

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(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -717,41 +717,53 @@ 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)."""
-        if self._has_mamba 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 needs the last prompt token to derive its final 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 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
         ):
             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:
-            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