Add more testing for chunked prefill (#27506)

test/manual/chunked_prefill/test_scripted_piecewise_cuda_graph.py

@@ -0,0 +1,67 @@
+import unittest
+
+from sglang.test.scripted_runtime.context import ScriptedContext
+from sglang.test.scripted_runtime.test_case import ScriptedTestCase
+from sglang.test.scripted_runtime_chunked_helpers import (
+    DEFAULT_CHUNK_SIZE,
+    VERY_LONG_PROMPT_LEN,
+    base_engine_kwargs,
+    run_until,
+    run_until_finished,
+)
+
+
+class TestPiecewiseBasic(ScriptedTestCase):
+    ENGINE_KWARGS = base_engine_kwargs(
+        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
+        disable_cuda_graph=False,
+    )
+
+    def test_naive_piecewise_cg_chunked(self):
+        self.server.execute_script(self._script_naive_piecewise_cg_chunked)
+
+    @staticmethod
+    def _script_naive_piecewise_cg_chunked(t: ScriptedContext):
+        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=8)
+        yield from run_until_finished(r)
+        assert r.finished
+        assert r.chunks_done >= 2
+
+    def test_piecewise_cg_tail_chunk_tiny(self):
+        self.server.execute_script(self._script_piecewise_cg_tail_chunk_tiny)
+
+    @staticmethod
+    def _script_piecewise_cg_tail_chunk_tiny(t: ScriptedContext):
+        r = t.start_req(
+            prompt_len=4 * DEFAULT_CHUNK_SIZE + 1,
+            max_new_tokens=2,
+        )
+        yield from run_until_finished(r, max_steps=800)
+        assert r.finished
+        assert r.chunks_done == 5
+
+
+class TestPiecewiseRetractResume(ScriptedTestCase):
+    ENGINE_KWARGS = base_engine_kwargs(
+        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
+        disable_cuda_graph=False,
+    )
+
+    def test_piecewise_cg_retract_resume(self):
+        self.server.execute_script(self._script_piecewise_cg_retract_resume)
+
+    @staticmethod
+    def _script_piecewise_cg_retract_resume(t: ScriptedContext):
+        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4)
+        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
+        t.pause_generation(mode="retract")
+        yield
+        t.continue_generation()
+        yield from run_until_finished(r, max_steps=800)
+        assert r.finished
+        assert r.kv_pages == 0
+        assert r.lock_refs == 0
+
+
+if __name__ == "__main__":
+    unittest.main()