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

2026-01-05 16:28:54 +08:00
parent 9d6029fb92
commit 4d737db857
2 changed files with 111 additions and 61 deletions
@@ -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."
@@ -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):
-    def test_1_online_serving_has_prefill_delayer(self):
-        self._run_throughput_test_online_serving(prefill_delayer=True)
+class TestPrefillDelayerThroughputOnlineServing(CustomTestCase):
+    def test_1_has_prefill_delayer(self):
+        self._run(prefill_delayer=True)

-    def test_2_online_serving_no_prefill_delayer(self):
-        self._run_throughput_test_online_serving(prefill_delayer=False)
+    def test_2_no_prefill_delayer(self):
+        self._run(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(
+    def _run(self, prefill_delayer: bool):
+        _run_throughput_test(
            debug_name=f"online_serving ({prefill_delayer=})",
            prefill_delayer=prefill_delayer,
            other_launch_args=[
-                "--mem-fraction-static",
-                "0.6",
+                # Not really needed, only to test support non-FCFS algorithms
+                "--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,57 +60,105 @@ 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"
+def _run_throughput_test(
+    debug_name: str,
+    prefill_delayer: bool,
+    other_launch_args,
+    other_benchmark_args,
+):
+    model = "Qwen/Qwen3-0.6B"
+    base_url = DEFAULT_URL_FOR_TEST
+
+    process = _launch_server(
+        prefill_delayer=prefill_delayer,
+        model=model,
+        base_url=base_url,
+        other_args=other_launch_args,
+    )
+
+    try:
+        args = get_benchmark_args(
+            base_url=base_url,
+            dataset_name="random",
+            tokenizer=model,
+            **other_benchmark_args,
+        )
+        res = run_benchmark(args)
+    finally:
+        kill_process_tree(process.pid)
+
+    print(f"=== {debug_name} ===")
+    print(f"Input throughput: {res['input_throughput']:.2f} token/s")
+    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
-
-        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",
-                    "8",
-                    "--enable-dp-attention",
-                    "--dp",
-                    "8",
-                    "--chunked-prefill-size",
-                    "131072",
-                    *other_launch_args,
-                ],
-            )
-
+        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 = get_benchmark_args(
+            args = SimpleNamespace(
                base_url=base_url,
-                dataset_name="random",
-                tokenizer=model,
-                **other_benchmark_args,
+                model=model,
+                eval_name="mgsm_en",
+                num_examples=None,
+                num_threads=1024,
            )
-            res = run_benchmark(args)
+            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)

-        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 _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__":