From 8f01a12d43fede93963f4606f1d9d9481b5cf518 Mon Sep 17 00:00:00 2001 From: vipwangerxiao Date: Wed, 12 Nov 2025 11:24:04 +0800 Subject: [PATCH] Improve overlap scheduling for better TTFT (#11856) Co-authored-by: Peng Wang --- docs/references/environment_variables.md | 1 + python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/scheduler.py | 24 +++++++++++++++++++++--- 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 616a02ee4..4814748c1 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -29,6 +29,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_SKIP_P2P_CHECK` | Skip P2P (peer-to-peer) access check | `false` | | `SGL_CHUNKED_PREFIX_CACHE_THRESHOLD` | Sets the threshold for enabling chunked prefix caching | `8192` | | `SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION` | Enable RoPE fusion in Fused Multi-Layer Attention | `1` | +| `SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP` | Disable overlap schedule for consecutive prefill batches | `false` | | `SGLANG_DISABLE_FA4_WARMUP` | Disable Flash Attention 4 warmup passes (set to `1`, `true`, `yes`, or `on` to disable) | `false` | ## DeepGEMM Configuration (Advanced Optimization) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 25da56d87..c822f6820 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -162,6 +162,7 @@ class Envs: # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. + SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False) SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False) # Test: pd-disaggregation diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4e8612ef5..e5117f3aa 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -988,6 +988,14 @@ class Scheduler( def event_loop_overlap(self): """A scheduler loop that overlaps the CPU processing and GPU computation.""" self.result_queue: Deque[Tuple[ScheduleBatch, GenerationBatchResult]] = deque() + disable_consecutive_prefill_overlap = ( + envs.SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP.get() + ) + + def pop_and_process(): + # Process the results of the last batch + tmp_batch, tmp_result = self.result_queue.popleft() + self.process_batch_result(tmp_batch, tmp_result) while True: recv_reqs = self.recv_requests() @@ -996,15 +1004,25 @@ class Scheduler( batch = self.get_next_batch_to_run() self.cur_batch = batch + disable_overlap_for_batch = ( + disable_consecutive_prefill_overlap + and batch + and batch.forward_mode.is_extend() + and self.last_batch + and self.last_batch.forward_mode.is_extend() + ) + + if disable_overlap_for_batch: + pop_and_process() + batch_result = None if batch: batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) if self.last_batch: - # Process the results of the 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: # When the server is idle, do self-check and re-init some states self.self_check_during_idle()