Support offline generation scenario for prefill delayer (#16363)

python/sglang/srt/managers/prefill_delayer.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional
 
 import torch
 
@@ -35,21 +36,7 @@ class PrefillDelayer:
             not server_args.disable_overlap_schedule
         ), "To use PrefillDelayer, disable_overlap_schedule must be False."
 
-    def _gather_info(self, local_prefillable: int):
+    def _negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
-        local_info = torch.tensor(
-            [local_prefillable],
-            device="cpu",
-            dtype=torch.int64,
-        )
-        torch.distributed.all_gather_into_tensor(
-            self.global_info.flatten(),
-            local_info,
-            group=self.cpu_group,
-        )
-        tp0_info = self.global_info[:, 0, :]
-        return tp0_info
-
-    def should_allow_prefill(self, local_prefillable: int) -> bool:
         tp0_info = self._gather_info(local_prefillable=local_prefillable)
         global_prefillable = tp0_info[:, 0]
         global_exists_not_prefillable = global_prefillable.min().item() == 0
@@ -70,3 +57,42 @@ class PrefillDelayer:
 
         self.curr_delayed_count = 0
         return True
+
+    def _gather_info(self, local_prefillable: bool):
+        local_info = torch.tensor(
+            [int(local_prefillable)],
+            device="cpu",
+            dtype=torch.int64,
+        )
+        torch.distributed.all_gather_into_tensor(
+            self.global_info.flatten(),
+            local_info,
+            group=self.cpu_group,
+        )
+        tp0_info = self.global_info[:, 0, :]
+        return tp0_info
+
+
+class PrefillDelayerSinglePassExecutor:
+    def __init__(self, prefill_delayer: PrefillDelayer):
+        self._prefill_delayer = prefill_delayer
+        self._result: Optional[bool] = None
+
+    @property
+    def _called(self) -> bool:
+        return self._result is not None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if not self._called:
+            self.negotiate_should_allow_prefill(local_prefillable=False)
+        return False
+
+    def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool:
+        if not self._called:
+            self._result = self._prefill_delayer._negotiate_should_allow_prefill(
+                local_prefillable=local_prefillable
+            )
+        return self._result

python/sglang/srt/managers/schedule_policy.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
+
 # Copyright 2023-2024 SGLang Team
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -327,6 +329,7 @@ class PrefillAdder:
         mixed_with_decode_tokens: int = 0,
         priority_scheduling_preemption_threshold: int = 0,
         prefill_max_requests: Optional[int] = None,
+        prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
     ):
         self.page_size = page_size
         self.tree_cache = tree_cache
@@ -366,6 +369,7 @@ class PrefillAdder:
         )
         self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split()
         self.prefill_max_requests = prefill_max_requests
+        self.prefill_delayer_single_pass = prefill_delayer_single_pass
 
     def _get_running_request_total_token_offset(self, req: Req) -> int:
         return (
@@ -614,6 +618,13 @@ class PrefillAdder:
             if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0:
                 return AddReqResult.OTHER
 
+            if (self.prefill_delayer_single_pass is not None) and (
+                not self.prefill_delayer_single_pass.negotiate_should_allow_prefill(
+                    local_prefillable=True
+                )
+            ):
+                return AddReqResult.OTHER
+
             if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:
                 # Non-chunked prefill
                 self.can_run_list.append(req)

python/sglang/srt/managers/scheduler.py

@@ -22,6 +22,7 @@ import sys
 import time
 from collections import deque
 from concurrent import futures
+from contextlib import nullcontext
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, Deque, Dict, List, Optional, Tuple, Union
@@ -120,7 +121,10 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.mm_utils import init_mm_embedding_cache
 from sglang.srt.managers.overlap_utils import FutureMap
-from sglang.srt.managers.prefill_delayer import PrefillDelayer
+from sglang.srt.managers.prefill_delayer import (
+    PrefillDelayer,
+    PrefillDelayerSinglePassExecutor,
+)
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     ModelWorkerBatch,
@@ -1911,6 +1915,18 @@ class Scheduler(
         return res
 
     def get_new_batch_prefill(self) -> Optional[ScheduleBatch]:
+        with (
+            PrefillDelayerSinglePassExecutor(self.prefill_delayer)
+            if self.prefill_delayer
+            else nullcontext()
+        ) as prefill_delayer_single_pass:
+            return self._get_new_batch_prefill_raw(
+                prefill_delayer_single_pass=prefill_delayer_single_pass
+            )
+
+    def _get_new_batch_prefill_raw(
+        self, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor]
+    ) -> Optional[ScheduleBatch]:
         # Check if the grammar is ready in the grammar queue
         if self.grammar_queue:
             self.move_ready_grammar_requests()
@@ -1919,19 +1935,8 @@ class Scheduler(
             # Reset batch_is_full to try preemption with a prefill adder.
             self.running_batch.batch_is_full = False
 
-        # The `should_allow_prefill` needs to be called on all ranks since contains communication
-        delayer_allow_prefill = (
-            self.prefill_delayer.should_allow_prefill(
-                # TODO: consider offline generation cases when there are a lot of waiting requests
-                local_prefillable=len(self.waiting_queue)
-                > 0,
-            )
-            if self.prefill_delayer
-            else True
-        )
         if (
-            (not delayer_allow_prefill)
+            self.running_batch.batch_is_full or len(self.waiting_queue) == 0
-            or ((self.running_batch.batch_is_full or len(self.waiting_queue) == 0))
         ) and self.chunked_req is None:
             return None
 
@@ -1981,6 +1986,7 @@ class Scheduler(
             running_bs if self.is_mixed_chunk else 0,
             self.priority_scheduling_preemption_threshold,
             prefill_max_requests=self.server_args.prefill_max_requests,
+            prefill_delayer_single_pass=prefill_delayer_single_pass,
         )
 
         if self.chunked_req is not None:

test/srt/test_prefill_delayer.py

@@ -14,11 +14,71 @@ from sglang.test.test_utils import (
 
 
 class TestPrefillDelayerThroughput(CustomTestCase):
-    def _run_throughput_test(self, with_prefill_delayer: bool):
+    def test_1_online_serving_has_prefill_delayer(self):
+        self._run_throughput_test_online_serving(prefill_delayer=True)
+
+    def test_2_online_serving_no_prefill_delayer(self):
+        self._run_throughput_test_online_serving(prefill_delayer=False)
+
+    def test_3_offline_gen_has_prefill_delayer(self):
+        self._run_throughput_test_offline_gen(prefill_delayer=True)
+
+    def test_4_offline_gen_no_prefill_delayer(self):
+        self._run_throughput_test_offline_gen(prefill_delayer=False)
+
+    def _run_throughput_test_online_serving(self, prefill_delayer: bool):
+        self._run_throughput_test(
+            debug_name=f"online_serving ({prefill_delayer=})",
+            prefill_delayer=prefill_delayer,
+            other_launch_args=[
+                "--mem-fraction-static",
+                "0.6",
+            ],
+            other_benchmark_args=dict(
+                num_prompts=500,
+                # trigger chunked prefill
+                random_input_len=30000,
+                random_output_len=256,
+                request_rate=32,
+            ),
+        )
+
+    def _run_throughput_test_offline_gen(self, prefill_delayer: bool):
+        self._run_throughput_test(
+            debug_name=f"offline_gen ({prefill_delayer=})",
+            prefill_delayer=prefill_delayer,
+            other_benchmark_args=dict(
+                num_prompts=800,
+                random_input_len=30000,
+                random_output_len=500,
+            ),
+            other_launch_args=[
+                "--max-total-tokens",
+                "200000",
+                "--mem-fraction-static",
+                "0.6",
+            ],
+        )
+
+    def _run_throughput_test(
+        self,
+        debug_name: str,
+        prefill_delayer: bool,
+        other_launch_args,
+        other_benchmark_args,
+    ):
         os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
 
         model = "Qwen/Qwen3-0.6B"
         base_url = DEFAULT_URL_FOR_TEST
+
+        with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
+            prefill_delayer
+        ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(100):
+            process = popen_launch_server(
+                model,
+                base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
                 other_args=[
                     "--trust-remote-code",
                     "--tp",
@@ -28,45 +88,25 @@ class TestPrefillDelayerThroughput(CustomTestCase):
                     "8",
                     "--chunked-prefill-size",
                     "131072",
-            "--mem-fraction-static",
+                    *other_launch_args,
-            "0.6",
+                ],
-        ]
-
-        with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
-            with_prefill_delayer
-        ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(100):
-            process = popen_launch_server(
-                model,
-                base_url,
-                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-                other_args=other_args,
             )
 
         try:
             args = get_benchmark_args(
                 base_url=base_url,
                 dataset_name="random",
-                num_prompts=500,
-                # trigger chunked prefill
-                random_input_len=30000,
-                random_output_len=256,
-                request_rate=32,
                 tokenizer=model,
+                **other_benchmark_args,
             )
             res = run_benchmark(args)
         finally:
             kill_process_tree(process.pid)
 
-        print(f"=== {with_prefill_delayer=} ===")
+        print(f"=== {debug_name} ===")
         print(f"Input throughput: {res['input_throughput']:.2f} token/s")
         print(f"Output throughput: {res['output_throughput']:.2f} token/s")
 
-    def test_1_dp_attention_throughput_with_prefill_delayer(self):
-        self._run_throughput_test(with_prefill_delayer=True)
-
-    def test_2_dp_attention_throughput_without_prefill_delayer(self):
-        self._run_throughput_test(with_prefill_delayer=False)
-
 
 if __name__ == "__main__":
     unittest.main()