diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 0eb65e803..c09d1b354 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -65,7 +65,6 @@ from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, set_time_batch, ) -from sglang.srt.utils import get_int_env_var from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) @@ -75,7 +74,7 @@ if TYPE_CHECKING: from sglang.srt.managers.scheduler import Scheduler from sglang.srt.server_args import ServerArgs -CLIP_MAX_NEW_TOKEN = get_int_env_var("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", 4096) +CLIP_MAX_NEW_TOKEN = envs.SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION.get() def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool: @@ -1036,39 +1035,27 @@ class SchedulerDisaggregationDecodeMixin: def get_next_disagg_decode_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: - """Create fake completed prefill if possible and merge with running batch""" - # Merge the prefill batch into the running batch - last_batch = self.last_batch - if last_batch and last_batch.forward_mode.is_prebuilt(): - # chunked prefill doesn't happen in decode instance. - assert self.chunked_req is None - # Filter finished batches. - last_batch.filter_batch() - if not last_batch.is_empty(): - if self.running_batch.is_empty(): - self.running_batch = last_batch - else: - # merge running_batch with prefill batch - self.running_batch.merge_batch(last_batch) - + """Process prebuilt batch and schedule the next decode batch.""" + # Process pending prebuilt batch: output processing + filter + merge new_prebuilt_batch = self.get_new_prebuilt_batch() - - ret: Optional[ScheduleBatch] = None if new_prebuilt_batch: - ret = new_prebuilt_batch + assert self.chunked_req is None + self.process_batch_result_prebuilt(new_prebuilt_batch) + new_prebuilt_batch.filter_batch() + if not new_prebuilt_batch.is_empty(): + if self.running_batch.is_empty(): + self.running_batch = new_prebuilt_batch + else: + self.running_batch.merge_batch(new_prebuilt_batch) + + # Schedule decode batch + if self.running_batch.is_empty(): + ret = None else: - if self.running_batch.is_empty(): - ret = None - else: - self.running_batch = self.update_running_batch(self.running_batch) - ret = self.running_batch if not self.running_batch.is_empty() else None + self.running_batch = self.update_running_batch(self.running_batch) + ret = self.running_batch if not self.running_batch.is_empty() else None - # 1. decode + None -> decode + idle - # 2. decode + prebuilt -> decode + idle (idle forward, prebuilt returns) - # 3. prebuilt + None -> None (None forward, prebuilt returns) + None - # 4. prebuilt + decode + None -> idle (idle forward, prebuilt returns) + decode + idle ret = self.maybe_prepare_mlp_sync_batch(ret) - if ret: set_schedule_time_batch(ret) return ret diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 03f86c615..9de2fa09a 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2088,9 +2088,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) def merge_batch(self, other: "ScheduleBatch"): - # NOTE: in spec v2 mode, we do not need wait verify here because - # 1) current batch is always prefill, whose seq_lens is not a future - # 2) other batch is always decode, which is finished in previous step + # In the regular scheduler path: + # 1) self is always prefill, whose seq_lens is not a future + # 2) other is always decode, which is finished in previous step + # so verify_done is already synced and this is a no-op. + # In disagg decode + overlap, merge_batch can be called before + # filter_batch, so running_batch.seq_lens may still be a forward_stream + # future. Synchronize here to avoid a cross-stream data race. + self.maybe_wait_verify_done() # Penalizer orchestrator must be merged before Batch.reqs is merged. This is because # orchestrator.merge() depends on Batch.reqs during preparation of each penalizers, so it @@ -2213,9 +2218,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ) def copy(self): - # Only contain fields that will be used by process_batch_result + # Only contain fields that will be used by process_batch_result. + # Shallow-copy the reqs list so that in-place mutations (filter_batch, + # merge_batch) on the original don't corrupt this snapshot. return ScheduleBatch( - reqs=self.reqs, + reqs=self.reqs[:], req_to_token_pool=self.req_to_token_pool, req_pool_indices=self.req_pool_indices, model_config=self.model_config, diff --git a/test/registered/distributed/test_disaggregation_dp_attention.py b/test/registered/distributed/test_disaggregation_dp_attention.py index fed913b88..a09a9c661 100644 --- a/test/registered/distributed/test_disaggregation_dp_attention.py +++ b/test/registered/distributed/test_disaggregation_dp_attention.py @@ -125,6 +125,23 @@ class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention): self.assertLess(result["mean_tpot_ms"], 20) self.assertEqual(result["completed"], 1000) + def test_bench_serving_itl(self): + num_prompts = 512 + args = get_benchmark_args( + base_url=f"http://{self.base_host}:{self.lb_port}", + dataset_name="random", + tokenizer=self.model, + num_prompts=num_prompts, + random_input_len=512, + random_output_len=64, + request_rate=float("inf"), + max_concurrency=64, + ) + result = run_benchmark(args) + + self.assertEqual(result["completed"], num_prompts) + self.assertLess(result["p99_itl_ms"], 20) + @unittest.skip( "Skip this test until new testing logic in mini-lb has been updated in docker image."