Tiny fix prefill delayer not support non-fcfs schedule policy (#16471)

python/sglang/srt/managers/prefill_delayer.py

@@ -23,9 +23,6 @@ class PrefillDelayer:
         self.curr_delayed_count = 0
         self.max_delay_passes = envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get()
 
-        assert (
-            server_args.schedule_policy == "fcfs"
-        ), f"To use PrefillDelayer, schedule_policy must be 'fcfs'. '{server_args.schedule_policy}' is not supported."
         assert (
             server_args.enable_dp_attention
         ), "To use PrefillDelayer, enable_dp_attention must be enabled."

test/srt/test_prefill_delayer.py

@@ -1,10 +1,13 @@
 import os
 import unittest
+from types import SimpleNamespace
 
 from sglang.bench_serving import run_benchmark
 from sglang.srt.environ import envs
 from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
@@ -13,38 +16,40 @@ from sglang.test.test_utils import (
 )
 
 
-class TestPrefillDelayerThroughput(CustomTestCase):
+class TestPrefillDelayerThroughputOnlineServing(CustomTestCase):
-    def test_1_online_serving_has_prefill_delayer(self):
+    def test_1_has_prefill_delayer(self):
-        self._run_throughput_test_online_serving(prefill_delayer=True)
+        self._run(prefill_delayer=True)
 
-    def test_2_online_serving_no_prefill_delayer(self):
+    def test_2_no_prefill_delayer(self):
-        self._run_throughput_test_online_serving(prefill_delayer=False)
+        self._run(prefill_delayer=False)
 
-    def test_3_offline_gen_has_prefill_delayer(self):
+    def _run(self, prefill_delayer: bool):
-        self._run_throughput_test_offline_gen(prefill_delayer=True)
+        _run_throughput_test(
-
-    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",
+                # Not really needed, only to test support non-FCFS algorithms
-                "0.6",
+                "--schedule-policy",
+                "lpm",
             ],
             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(
+class TestPrefillDelayerThroughputOfflineGen(CustomTestCase):
+    def test_1_has_prefill_delayer(self):
+        self._run(prefill_delayer=True)
+
+    def test_2_no_prefill_delayer(self):
+        self._run(prefill_delayer=False)
+
+    def _run(self, prefill_delayer: bool):
+        _run_throughput_test(
             debug_name=f"offline_gen ({prefill_delayer=})",
             prefill_delayer=prefill_delayer,
             other_benchmark_args=dict(
@@ -55,41 +60,24 @@ class TestPrefillDelayerThroughput(CustomTestCase):
             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(
+    process = _launch_server(
-            prefill_delayer
+        prefill_delayer=prefill_delayer,
-        ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(100):
+        model=model,
-            process = popen_launch_server(
+        base_url=base_url,
-                model,
+        other_args=other_launch_args,
-                base_url,
-                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-                other_args=[
-                    "--trust-remote-code",
-                    "--tp",
-                    "8",
-                    "--enable-dp-attention",
-                    "--dp",
-                    "8",
-                    "--chunked-prefill-size",
-                    "131072",
-                    *other_launch_args,
-                ],
     )
 
     try:
@@ -108,5 +96,70 @@ class TestPrefillDelayerThroughput(CustomTestCase):
     print(f"Output throughput: {res['output_throughput']:.2f} token/s")
 
 
+class TestPrefillDelayerAccuracy(CustomTestCase):
+    def test_1_mgsm_en_has_prefill_delayer(self):
+        self._run_accuracy_test(prefill_delayer=True)
+
+    def test_2_mgsm_en_no_prefill_delayer(self):
+        self._run_accuracy_test(prefill_delayer=False)
+
+    def _run_accuracy_test(self, prefill_delayer: bool):
+        model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        base_url = DEFAULT_URL_FOR_TEST
+        process = _launch_server(
+            prefill_delayer=prefill_delayer,
+            model=model,
+            base_url=base_url,
+            other_args=[
+                # Not really needed, only to test support non-FCFS algorithms
+                "--schedule-policy",
+                "lpm",
+                # Use this to ensure prefill delayer will be run
+                "--max-total-tokens",
+                "4096",
+            ],
+        )
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                model=model,
+                eval_name="mgsm_en",
+                num_examples=None,
+                num_threads=1024,
+            )
+            metrics = run_eval(args)
+            print(f"=== mgsm_en ({prefill_delayer=}) ===")
+            print(f"{metrics=}")
+            self.assertGreater(metrics["score"], 0.87)
+        finally:
+            kill_process_tree(process.pid)
+
+
+def _launch_server(*, model, base_url, prefill_delayer: bool, other_args):
+    os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
+
+    with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
+        prefill_delayer
+    ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(100):
+        return popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "8",
+                "--enable-dp-attention",
+                "--dp",
+                "8",
+                "--chunked-prefill-size",
+                "131072",
+                "--mem-fraction-static",
+                "0.6",
+                *(other_args or []),
+            ],
+        )
+
+
 if __name__ == "__main__":
     unittest.main()