diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py index 57d0f65d5820..bc31e62396d5 100644 --- a/python/sglang/srt/constrained/reasoner_grammar_backend.py +++ b/python/sglang/srt/constrained/reasoner_grammar_backend.py @@ -78,6 +78,7 @@ def __init__( def maybe_init_reasoning(self, reasoning: bool): if reasoning: self.tokens_in_think = 0 + self.tokens_after_end = -1 # reset: cache hit may have stale GENERATION state else: self.tokens_in_think = -1 self.tokens_after_end = 0 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 022446bf4122..2529710e49ad 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1682,6 +1682,15 @@ def event_loop_normal_disagg_decode(self: Scheduler): # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) + + # Match the normal overlap scheduler: for spec_v2 + grammar, commit + # accepted tokens to req.output_ids / req.grammar before launching the + # next decode batch. Otherwise disagg overlap can prepare the next MTP + # verify from advanced spec state while the grammar matcher is still + # one result behind. + if disable_overlap_for_batch and self.last_batch: + pop_and_process() # Launch the current batch if batch: @@ -1699,6 +1708,10 @@ def event_loop_overlap_disagg_decode(self: Scheduler): self.result_queue = deque() self.last_batch: Optional[ScheduleBatch] = None + def pop_and_process(): + tmp_batch, tmp_result = self.result_queue.popleft() + self.process_batch_result(tmp_batch, tmp_result) + while True: # Receive requests recv_reqs = self.request_receiver.recv_requests() @@ -1714,6 +1727,15 @@ def event_loop_overlap_disagg_decode(self: Scheduler): # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) + + # Match the normal overlap scheduler: for spec_v2 + grammar, commit + # accepted tokens to req.output_ids / req.grammar before launching the + # next decode batch. Otherwise disagg overlap can prepare the next MTP + # verify from advanced spec state while the grammar matcher is still + # one result behind. + if disable_overlap_for_batch and self.last_batch: + pop_and_process() # Launch the current batch if batch: @@ -1724,8 +1746,8 @@ def event_loop_overlap_disagg_decode(self: Scheduler): # Process the last batch if self.last_batch: - tmp_batch, tmp_result = self.result_queue.popleft() - self.process_batch_result(tmp_batch, tmp_result) + if not disable_overlap_for_batch: + pop_and_process() elif batch is None: self.on_idle()