[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>
Ziming Huang committed
Aug 25, 2026 at 10:13 UTC
d4c4ceb31e3d35222a15f01ac2349d4a39141738
5 files changed
+32
-14
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+10
-7
@@ -1054,7 +1054,7 @@ def test_mamba_n1_d_side_builds_decode_metadata():
1054
1055
@pytest.mark.cpu_test
1056
def test_mamba_n1_p_side_truncation():
1057
- """P-side: Mamba truncates prompt to N-1, sets max_tokens=1.
1057
+ """P-side truncates to N-1 before the scheduler's prefix-cache lookup.
1058
1059
Also verifies idempotency (calling again is a no-op) which is
1060
needed for preemption safety via the _p_side_truncated guard,
@@ -1065,17 +1065,20 @@ def test_mamba_n1_p_side_truncation():
1065
req.max_tokens = 128
1066
original_len = len(req.prompt_token_ids)
1067
1068
+ sched.on_new_request(req)
1069
+ assert len(req.prompt_token_ids) == original_len - 1
1070
+ assert req.num_prompt_tokens == original_len - 1
1071
+ assert req.max_tokens == 1
1072
+ assert req.kv_transfer_params["_p_side_truncated"] is True
1073
+
1074
count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
1075
1076
assert count == 0
1077
assert is_async is False
1078
assert len(req.prompt_token_ids) == original_len - 1
1073
- assert req.num_prompt_tokens == original_len - 1
1074
- assert req.max_tokens == 1
1075
- assert req.kv_transfer_params["_p_side_truncated"] is True
1079
1077
- # Idempotency: second call must not truncate further
1078
- sched.get_num_new_matched_tokens(req, num_computed_tokens=0)
1080
+ # Idempotency: re-adding a preempted request must not truncate further.
1081
+ sched.on_new_request(req)
1082
assert len(req.prompt_token_ids) == original_len - 1
1083
1084
# Non-Mamba: truncation is skipped
@@ -1083,7 +1086,7 @@ def test_mamba_n1_p_side_truncation():
1086
fa_req = create_request(num_tokens=10, do_remote_decode=True)
1087
fa_original = len(fa_req.prompt_token_ids)
1088
1086
- fa_sched.get_num_new_matched_tokens(fa_req, num_computed_tokens=0)
1089
+ fa_sched.on_new_request(fa_req)
1090
assert len(fa_req.prompt_token_ids) == fa_original
1091
1092
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+19
-1
@@ -46,7 +46,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
46
from vllm.v1.kv_cache_interface import FullAttentionSpec
47
from vllm.v1.outputs import KVConnectorOutput
48
49
-from .utils import make_nixl_push_scheduler
49
+from .utils import create_request, make_nixl_push_scheduler
50
51
# ----------------------------------------------------------------- #
52
# Helpers / fakes #
@@ -117,6 +117,24 @@ def _stub_sw_clipping(scheduler) -> None:
117
118
119
class TestPushScheduler:
120
+ def test_p_side_mamba_truncates_before_cache_lookup(self):
121
+ """Push mode normalizes P-side Mamba requests before cache lookup."""
122
+ sched = make_nixl_push_scheduler(has_mamba=True)
123
+ request = create_request(num_tokens=10, do_remote_decode=True)
124
+ original_len = len(request.prompt_token_ids)
125
+
126
+ sched.on_new_request(request)
127
+
128
+ assert len(request.prompt_token_ids) == original_len - 1
129
+ assert request.kv_transfer_params["_p_side_truncated"] is True
130
+
131
+ with patch.object(
132
+ sched,
133
+ "_truncate_mamba_request_for_prefill",
134
+ side_effect=AssertionError("must not truncate after cache lookup"),
135
+ ):
136
+ assert sched.get_num_new_matched_tokens(request, 0) == (0, False)
137
+
138
def test_d_side_update_state_after_alloc_stages_registration(self):
139
"""D scheduler stashes registration data + arms watchdog deadline."""
140
sched = make_nixl_push_scheduler()
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+3
@@ -194,6 +194,9 @@ class NixlBaseConnectorScheduler:
194
def on_new_request(self, request: "Request") -> None:
195
"""Track a request that may need heartbeats."""
196
params = request.kv_transfer_params
197
+ if params is not None and params.get("do_remote_decode") and self._has_mamba:
198
+ self._truncate_mamba_request_for_prefill(request)
199
+
200
# NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
201
# effectively disabled for Bidirectional KV transfer.
202
if params is None or not params.get("do_remote_prefill"):
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py
-3
@@ -65,9 +65,6 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
65
if count > 0:
66
return count, True
67
68
- if params is not None and params.get("do_remote_decode") and self._has_mamba:
69
- self._truncate_mamba_request_for_prefill(request)
70
-
68
if (
69
params is not None
70
and params.get("do_remote_decode")
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py
-3
@@ -123,9 +123,6 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
123
if count > 0:
124
return count, True
125
126
- if params is not None and params.get("do_remote_decode") and self._has_mamba:
127
- self._truncate_mamba_request_for_prefill(request)
128
-
126
return 0, False
127
128
def update_state_after_alloc(