diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 1ecd04ae1..25506ca33 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1754,6 +1754,10 @@ class SchedulerDisaggregationDecodeMixin: 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() @@ -1769,6 +1773,11 @@ class SchedulerDisaggregationDecodeMixin: # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() self.cur_batch = batch + # overlap + spec + grammar is unsupported (would desync DP ranks). + disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) + + if disable_overlap_for_batch and self.last_batch: + pop_and_process() # Launch the current batch if batch: @@ -1779,8 +1788,8 @@ class SchedulerDisaggregationDecodeMixin: # 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() diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 50708530f..048b3e014 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -10,6 +10,7 @@ import requests from transformers import AutoTokenizer from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin from sglang.test.run_eval import run_eval from sglang.test.server_fixtures.disaggregation_fixture import ( @@ -21,7 +22,7 @@ from sglang.test.test_utils import ( DEFAULT_TARGET_MODEL_EAGLE3, ) -register_cuda_ci(est_time=509, stage="base-b", runner_config="2-gpu-large") +register_cuda_ci(est_time=560, stage="base-b", runner_config="2-gpu-large") class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase): @@ -215,7 +216,7 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase): raise e from health_check_error -class TestDisaggregationMooncakeSpec(PDDisaggregationServerBase): +class TestDisaggregationMooncakeSpec(JSONConstrainedMixin, PDDisaggregationServerBase): @classmethod def setUpClass(cls): super().setUpClass()