From f926c30c5706147c5fe5569fb04753f86c6a4d1f Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:03:10 -0700 Subject: [PATCH] Revert "Fix mamba track-boundary seqlen under overlap scheduler (#31369)" (#31622) --- .../batch_result_processor.py | 26 +++++++------------ 1 file changed, 9 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 33fd7391c..aca94e616 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -956,20 +956,13 @@ class SchedulerBatchResultProcessor: def _mamba_check_track_boundary(self, req, batch, result, i): """Check if this decode step crosses a mamba track interval boundary. - Returns (at_boundary, track_seqlen). ``track_seqlen`` must equal the - seq_len the forward's tracking mask used, so the tracked state and its - recorded ``mamba_last_track_seqlen`` describe the same token position. - That seq_len is a pure function of the tokens the request has produced: - ``len(origin_input_ids) + len(output_ids) - 1`` (the just-decoded token - is already appended to ``output_ids`` before this runs). - - ``kv_committed_len`` must NOT be used here: under the overlap scheduler, - ``prepare_for_decode`` for the *next* batch increments it before this - result is processed, so it leads seq_len by a jittering lookahead - (0 or 1 depending on prefill interleaving). Using it fires the boundary - one decode step early on most steps, mislabeling the tracked mamba - state; a later request that reuses/donates that tracked prefix then - extends from a state a cold prefill recompute would not produce. + Returns (at_boundary, track_seqlen). The boundary condition + matches what the forward's tracking mask used: + ``prepare_for_decode`` increments both ``seq_lens_cpu`` and + ``kv_committed_len`` by 1, then checks + ``seq_lens_cpu % interval == 0``. Using ``kv_committed_len`` + here reproduces that check exactly, and the value is always a + multiple of ``interval`` (hence page-aligned). For spec decode, the boundary is detected by comparing the accepted seq_len range against interval boundaries. @@ -977,9 +970,8 @@ class SchedulerBatchResultProcessor: interval = get_server_args().mamba_track_interval if batch.spec_algorithm.is_none(): - seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1 - if seq_len % interval == 0: - return True, seq_len + if req.kv_committed_len % interval == 0: + return True, req.kv_committed_len elif result.num_correct_drafts_per_req_cpu is not None: cur = req.seqlen - 1 prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1