From d4c4ceb31e3d35222a15f01ac2349d4a39141738 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Tue, 25 Aug 2026 10:13:27 +0800 Subject: [PATCH] [Bugfix][NIXL] Fix Mamba prefill truncation ordering (#53523) Signed-off-by: Ziming Huang Co-authored-by: Roger Wang --- .../unit/test_nixl_connector_hma.py | 19 ++++++++++-------- .../unit/test_nixl_push_connector.py | 20 ++++++++++++++++++- .../kv_connector/v1/nixl/base_scheduler.py | 3 +++ .../kv_connector/v1/nixl/pull_scheduler.py | 3 --- .../kv_connector/v1/nixl/push_scheduler.py | 3 --- 5 files changed, 33 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 66a90d64e5..11d5328ba1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index d7b6edf07c..a7ea078d6c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index bd93c23233..9b623b53c7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/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"): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index d0a5383c5c..5393989c9a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/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") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index 6efe491be8..095d90e7b1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/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(