minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Scheduler] Count complete prefill bursts and their tokens (#40006)

Browse Source 
Co-authored-by: pranjalssh <pranjalssh@fb.com>
Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Co-authored-by: Jialin Ouyang <jialino@meta.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
metamergebot
2026-09-19 12:56:15 -07:00
committed by GitHub
co-authored by pranjalssh Jialin Ouyang Jialin Ouyang Lianmin Zheng
parent 7a6c652c77
commit 8139a1740e
2 changed files with 36 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+3
View File
@@ -2345,6 +2345,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
split_prefill_finished: bool = False
split_forward_count: int = 1
split_forward_batch: ForwardBatch = None
# A full prefill has one result but can span several run_batch calls.
split_prefill_start: Optional[Tuple[int, float]] = None
# CPU mirror of req_pool_indices; schedule-path only (used in overlap_utils,
# not read by ForwardBatch), stale in spec draft window
@@ -3790,6 +3792,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
forward_iter=self.forward_iter,
launch_ts=self.launch_ts,
after_idle_gap=self.after_idle_gap,
split_prefill_start=self.split_prefill_start,
extend_num_tokens=self.extend_num_tokens,
)
python/sglang/srt/managers/scheduler.py
+33 -7
View File
@@ -2449,6 +2449,7 @@ class Scheduler(
self.total_prefill_busy_us = 0
self.decode_moment_totals: list[float] = [0.0] * 6
self._prev_step: Optional[Tuple[int, float, bool]] = None
self._prev_prefill_end_ts: Optional[float] = None
self._sched_idled = False
self.load_inquirer = SchedulerLoadInquirer(
disaggregation_mode=self.disaggregation_mode,
@@ -4333,7 +4334,11 @@ class Scheduler(
self.forward_ct += 1
batch.forward_iter = self.forward_ct
batch.launch_ts = time.monotonic()
batch.after_idle_gap = self._sched_idled
is_split_prefill = batch.forward_mode.is_split_prefill()
if not is_split_prefill or batch.split_index == 0:
batch.after_idle_gap = self._sched_idled
if is_split_prefill:
batch.split_prefill_start = (batch.forward_iter, batch.launch_ts)
self._sched_idled = False
# Accumulate the prefill-token counter used by the HRRN scheduling policy. Decode / prebuilt batches contribute 0.
@@ -4798,15 +4803,36 @@ class Scheduler(
return
if all(is_health_check_generate_req(req) for req in batch.reqs):
return
if is_prefill and mode.is_split_prefill():
start_iter, start_ts = batch.split_prefill_start
else:
start_iter, start_ts = batch.forward_iter, batch.launch_ts
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)
contiguous = False
if prev is not None:
prev_iter, prev_launch_ts, prev_is_prefill = prev
contiguous = (
not batch.after_idle_gap
and prev_iter + 1 == start_iter
and prev_is_prefill == is_prefill
)
if is_prefill:
# Completion boundaries include the burst tail and scheduling overhead
# without double-counting overlapping launch-to-result spans.
end_ts = time.monotonic()
prev_end_ts = self._prev_prefill_end_ts
self._prev_prefill_end_ts = end_ts
if prev_end_ts is not None:
# A new burst or mode can still overlap an earlier prefill.
start_ts = prev_end_ts if contiguous else max(start_ts, prev_end_ts)
step_us = int((end_ts - start_ts) * 1e6)
else:
# Decode regression uses launch cadence and skips boundaries.
if not contiguous:
return
step_us = int((batch.launch_ts - prev_launch_ts) * 1e6)
if not 0 < step_us < STEP_MAX_US:
return
if is_prefill: