Mixed style of chunked prefill (#1013)

test/srt/test_chunked_prefill.py

@@ -11,11 +11,14 @@ from sglang.test.test_utils import (
 
 
 class TestChunkedPrefill(unittest.TestCase):
-    def run_mmlu(self, disable_radix_cache):
+    def run_mmlu(self, disable_radix_cache, enable_mixed_chunk):
         other_args = ["--chunked-prefill-size", "32"]
         if disable_radix_cache:
             other_args += ["--disable-radix-cache"]
 
+        if enable_mixed_chunk:
+            other_args += ["--enable-mixed-chunk"]
+
         model = DEFAULT_MODEL_NAME_FOR_TEST
         base_url = DEFAULT_URL_FOR_UNIT_TEST
         process = popen_launch_server(
@@ -40,10 +43,16 @@ class TestChunkedPrefill(unittest.TestCase):
             kill_child_process(process.pid)
 
     def test_chunked_prefill(self):
-        self.run_mmlu(disable_radix_cache=False)
+        self.run_mmlu(disable_radix_cache=False, enable_mixed_chunk=False)
+
+    def test_mixed_chunked_prefill(self):
+        self.run_mmlu(disable_radix_cache=False, enable_mixed_chunk=True)
 
     def test_chunked_prefill_without_radix_cache(self):
-        self.run_mmlu(disable_radix_cache=True)
+        self.run_mmlu(disable_radix_cache=True, enable_mixed_chunk=False)
+
+    def test_mixed_chunked_prefill_without_radix_cache(self):
+        self.run_mmlu(disable_radix_cache=True, enable_mixed_chunk=True)
 
 
 if __name__ == "__main__":

test/srt/test_eval_accuracy_large_chunked_prefill.py

@@ -6,7 +6,6 @@ from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MODEL_NAME_FOR_TEST,
     DEFAULT_URL_FOR_ACCURACY_TEST,
-    DEFAULT_URL_FOR_UNIT_TEST,
     popen_launch_server,
 )
 

test/srt/test_eval_accuracy_large_mixed_chunked_prefill.py

@@ -0,0 +1,73 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_child_process
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_URL_FOR_ACCURACY_TEST,
+    popen_launch_server,
+)
+
+
+class TestEvalAccuracyLargeChunkedPrefill(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_ACCURACY_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=300,
+            other_args=[
+                "--log-level-http",
+                "warning",
+                "--chunked-prefill-size",
+                "256",
+                "--enable-mixed-chunk",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_child_process(cls.process.pid)
+
+    def test_mmlu(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=3000,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        assert metrics["score"] >= 0.71, f"{metrics}"
+
+    def test_human_eval(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="humaneval",
+            num_examples=None,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        assert metrics["score"] >= 0.64, f"{metrics}"
+
+    def test_mgsm_en(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mgsm_en",
+            num_examples=None,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        assert metrics["score"] >= 0.84, f"{metrics}"
+
+
+if __name__ == "__main__":
+    unittest.main()