diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a45d122f0..e20854b33 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1979,6 +1979,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): prefill_stats: Optional[PrefillStats] = None forward_iter: Optional[int] = None launch_ts: Optional[float] = None + after_idle_gap: bool = False # === GPU tensors crossing to ForwardBatch (clone targets for stream isolation) === # Batched arguments to model runner @@ -3201,6 +3202,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): fpm_start_time=self.fpm_start_time, forward_iter=self.forward_iter, launch_ts=self.launch_ts, + after_idle_gap=self.after_idle_gap, extend_num_tokens=self.extend_num_tokens, ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a9615a14c..5480ddacd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -335,7 +335,7 @@ TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get() TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get() -DECODE_STEP_MAX_US = 2_000_000 +STEP_MAX_US = 2_000_000 def _accumulate_decode_moment( @@ -1705,6 +1705,7 @@ class Scheduler( self.process_batch_result(batch, result) else: # When the server is idle, do self-check and re-init some states. + self._sched_idled = True self.on_idle() # Update last_batch @@ -1766,6 +1767,7 @@ class Scheduler( self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None + self._sched_idled = True # Process the last batch if self.last_batch: @@ -2044,7 +2046,8 @@ class Scheduler( self.total_prefill_uncached_tokens = 0 self.total_prefill_busy_us = 0 self.decode_moment_totals: list[float] = [0.0] * 6 - self._prev_decode_launch_ts: Optional[float] = None + self._prev_step: Optional[Tuple[int, float, bool]] = None + self._sched_idled = False self.load_inquirer = SchedulerLoadInquirer( disaggregation_mode=self.disaggregation_mode, ps=self.ps, @@ -3534,6 +3537,8 @@ class Scheduler( self.forward_ct += 1 batch.forward_iter = self.forward_ct batch.launch_ts = time.monotonic() + batch.after_idle_gap = self._sched_idled + self._sched_idled = False if self.scripted_scheduler_hook is not None: self.scripted_scheduler_hook.on_run_batch(batch) @@ -3853,23 +3858,28 @@ class Scheduler( return if all(is_health_check_generate_req(req) for req in batch.reqs): return + prev = self._prev_step + self._prev_step = (batch.forward_iter, batch.launch_ts, is_prefill) + # An idle pass keeps forward_iter contiguous (forward_ct advances in run_batch). + if prev is None or batch.after_idle_gap: + return + prev_iter, prev_ts, prev_is_prefill = prev + if prev_iter + 1 != batch.forward_iter or prev_is_prefill != is_prefill: + return + step_us = int((batch.launch_ts - prev_ts) * 1e6) + if not 0 < step_us < STEP_MAX_US: + return if is_prefill: - # Busy span = run_batch entry -> result processed. - span_us = int((time.monotonic() - batch.launch_ts) * 1e6) - self.total_prefill_busy_us += span_us + self.total_prefill_busy_us += step_us self.total_prefill_uncached_tokens += batch.extend_num_tokens else: batch_size = len(batch.reqs) - if self._prev_decode_launch_ts is not None: - step_us = int((batch.launch_ts - self._prev_decode_launch_ts) * 1e6) - if 0 < step_us < DECODE_STEP_MAX_US: - _accumulate_decode_moment( - self.decode_moment_totals, - batch_size, - step_us, - batch_size + result.num_correct_drafts, - ) - self._prev_decode_launch_ts = batch.launch_ts + _accumulate_decode_moment( + self.decode_moment_totals, + batch_size, + step_us, + batch_size + result.num_correct_drafts, + ) def maybe_send_health_check_signal(self): if self.return_health_check_ipcs: