Fix mamba track-boundary seqlen under overlap scheduler (#31369)

This commit is contained in:
Sam Shleifer
2026-07-17 10:22:07 +08:00
committed by GitHub
parent 37f94cb7a0
commit 0675d3033f
@@ -956,13 +956,20 @@ class SchedulerBatchResultProcessor:
def _mamba_check_track_boundary(self, req, batch, result, i): def _mamba_check_track_boundary(self, req, batch, result, i):
"""Check if this decode step crosses a mamba track interval boundary. """Check if this decode step crosses a mamba track interval boundary.
Returns (at_boundary, track_seqlen). The boundary condition Returns (at_boundary, track_seqlen). ``track_seqlen`` must equal the
matches what the forward's tracking mask used: seq_len the forward's tracking mask used, so the tracked state and its
``prepare_for_decode`` increments both ``seq_lens_cpu`` and recorded ``mamba_last_track_seqlen`` describe the same token position.
``kv_committed_len`` by 1, then checks That seq_len is a pure function of the tokens the request has produced:
``seq_lens_cpu % interval == 0``. Using ``kv_committed_len`` ``len(origin_input_ids) + len(output_ids) - 1`` (the just-decoded token
here reproduces that check exactly, and the value is always a is already appended to ``output_ids`` before this runs).
multiple of ``interval`` (hence page-aligned).
``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.
For spec decode, the boundary is detected by comparing the For spec decode, the boundary is detected by comparing the
accepted seq_len range against interval boundaries. accepted seq_len range against interval boundaries.
@@ -970,8 +977,9 @@ class SchedulerBatchResultProcessor:
interval = get_server_args().mamba_track_interval interval = get_server_args().mamba_track_interval
if batch.spec_algorithm.is_none(): if batch.spec_algorithm.is_none():
if req.kv_committed_len % interval == 0: seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
return True, req.kv_committed_len if seq_len % interval == 0:
return True, seq_len
elif result.num_correct_drafts_per_req_cpu is not None: elif result.num_correct_drafts_per_req_cpu is not None:
cur = req.seqlen - 1 cur = req.seqlen - 1
prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1 prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1