[Bugfix][NIXL] Fix Mamba prefill truncation ordering (#53523)

Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1054,7 +1054,7 @@ def test_mamba_n1_d_side_builds_decode_metadata():
 
 @pytest.mark.cpu_test
 def test_mamba_n1_p_side_truncation():
-    """P-side: Mamba truncates prompt to N-1, sets max_tokens=1.
+    """P-side truncates to N-1 before the scheduler's prefix-cache lookup.
 
     Also verifies idempotency (calling again is a no-op) which is
     needed for preemption safety via the _p_side_truncated guard,
@@ -1065,17 +1065,20 @@ def test_mamba_n1_p_side_truncation():
     req.max_tokens = 128
     original_len = len(req.prompt_token_ids)
 
-    count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
-
-    assert count == 0
-    assert is_async is False
+    sched.on_new_request(req)
     assert len(req.prompt_token_ids) == original_len - 1
     assert req.num_prompt_tokens == original_len - 1
     assert req.max_tokens == 1
     assert req.kv_transfer_params["_p_side_truncated"] is True
 
-    # Idempotency: second call must not truncate further
-    sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
+    count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
+
+    assert count == 0
+    assert is_async is False
+    assert len(req.prompt_token_ids) == original_len - 1
+
+    # Idempotency: re-adding a preempted request must not truncate further.
+    sched.on_new_request(req)
     assert len(req.prompt_token_ids) == original_len - 1
 
     # Non-Mamba: truncation is skipped
@@ -1083,7 +1086,7 @@ def test_mamba_n1_p_side_truncation():
     fa_req = create_request(num_tokens=10, do_remote_decode=True)
     fa_original = len(fa_req.prompt_token_ids)
 
-    fa_sched.get_num_new_matched_tokens(fa_req, num_computed_tokens=0)
+    fa_sched.on_new_request(fa_req)
     assert len(fa_req.prompt_token_ids) == fa_original
 
 

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -46,7 +46,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 from vllm.v1.outputs import KVConnectorOutput
 
-from .utils import make_nixl_push_scheduler
+from .utils import create_request, make_nixl_push_scheduler
 
 # ----------------------------------------------------------------- #
 #  Helpers / fakes                                                   #
@@ -117,6 +117,24 @@ def _stub_sw_clipping(scheduler) -> None:
 
 
 class TestPushScheduler:
+    def test_p_side_mamba_truncates_before_cache_lookup(self):
+        """Push mode normalizes P-side Mamba requests before cache lookup."""
+        sched = make_nixl_push_scheduler(has_mamba=True)
+        request = create_request(num_tokens=10, do_remote_decode=True)
+        original_len = len(request.prompt_token_ids)
+
+        sched.on_new_request(request)
+
+        assert len(request.prompt_token_ids) == original_len - 1
+        assert request.kv_transfer_params["_p_side_truncated"] is True
+
+        with patch.object(
+            sched,
+            "_truncate_mamba_request_for_prefill",
+            side_effect=AssertionError("must not truncate after cache lookup"),
+        ):
+            assert sched.get_num_new_matched_tokens(request, 0) == (0, False)
+
     def test_d_side_update_state_after_alloc_stages_registration(self):
         """D scheduler stashes registration data + arms watchdog deadline."""
         sched = make_nixl_push_scheduler()

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -194,6 +194,9 @@ class NixlBaseConnectorScheduler:
     def on_new_request(self, request: "Request") -> None:
         """Track a request that may need heartbeats."""
         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)
+
         # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
         # effectively disabled for Bidirectional KV transfer.
         if params is None or not params.get("do_remote_prefill"):

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -65,9 +65,6 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
             if count > 0:
                 return count, True
 
-        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")

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py

@@ -123,9 +123,6 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
             if count > 0:
                 return count, True
 
-        if params is not None and params.get("do_remote_decode") and self._has_mamba:
-            self._truncate_mamba_request_for_prefill(request)
-
         return 0, False
 
     def update_state_after_alloc(