diff --git a/python/sglang/test/chunked_prefill_test_utils.py b/python/sglang/test/chunked_prefill_test_utils.py new file mode 100644 index 000000000..ddcaf1333 --- /dev/null +++ b/python/sglang/test/chunked_prefill_test_utils.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import time +from types import SimpleNamespace +from typing import ClassVar, List, Optional + +from sglang.srt.utils import kill_process_tree +from sglang.test.run_eval import run_eval +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, + try_cached_model, +) + +DEFAULT_MODEL: str = "Qwen/Qwen3-0.6B" + +DEFAULT_CHUNKED_PREFILL_SIZE: int = 256 +DEFAULT_NUM_EXAMPLES: int = 100 +DEFAULT_NUM_SHOTS: int = 10 +LONG_PROMPT_NUM_SHOTS: int = 24 +DEFAULT_NUM_THREADS: int = 128 +DEFAULT_MAX_TOKENS: int = 512 +DEFAULT_SEED: int = 42 + +KV_CANARY_ARGS: List[str] = [ + "--kv-canary", + "raise", + "--kv-canary-real-data", + "partial", + "--kv-canary-sweep-interval", + "100", + "--disable-piecewise-cuda-graph", +] + + +class ChunkedGsm8kMixin: + __test__ = False + use_kv_canary: ClassVar[bool] = True + model: ClassVar[str] = DEFAULT_MODEL + feature_args: ClassVar[List[str]] = [] + + chunked_prefill_size: ClassVar[int] = DEFAULT_CHUNKED_PREFILL_SIZE + num_shots: ClassVar[int] = DEFAULT_NUM_SHOTS + num_examples: ClassVar[int] = DEFAULT_NUM_EXAMPLES + num_threads: ClassVar[int] = DEFAULT_NUM_THREADS + max_tokens: ClassVar[int] = DEFAULT_MAX_TOKENS + gsm8k_threshold: ClassVar[float] + + def build_prefill_side_args(self) -> List[str]: + canary = list(KV_CANARY_ARGS) if self.use_kv_canary else [] + return ( + ["--chunked-prefill-size", str(self.chunked_prefill_size)] + + list(self.feature_args) + + canary + ) + + def test_mixed_prefix_gsm8k_chunked(self): + fixture_name = type(self).__name__ + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mixed_prefix_gsm8k", + api="chat_completion", + max_tokens=self.max_tokens, + num_examples=self.num_examples, + num_threads=self.num_threads, + num_shots=self.num_shots, + mixed_prefix_gsm8k_secondary_pool_size=15, + mixed_prefix_gsm8k_seed=DEFAULT_SEED, + gsm8k_data_path=None, + temperature=0.0, + ) + tic = time.perf_counter() + metrics = run_eval(args) + metrics["elapsed_sec"] = time.perf_counter() - tic + print(f"[{fixture_name}] {metrics} threshold={self.gsm8k_threshold:.4f}") + + score = metrics.get("score") + self.assertIsNotNone(score, "run_eval returned no score") + self.assertGreaterEqual(score, self.gsm8k_threshold) + + +class ChunkedTestBase(ChunkedGsm8kMixin, CustomTestCase): + __test__ = False + + base_url: ClassVar[str] = DEFAULT_URL_FOR_TEST + launch_timeout: ClassVar[int] = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + + process: ClassVar[Optional[object]] = None + + @classmethod + def setUpClass(cls): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=cls.launch_timeout, + other_args=cls("test_mixed_prefix_gsm8k_chunked").build_prefill_side_args(), + ) + + @classmethod + def tearDownClass(cls): + if cls.process is not None: + kill_process_tree(cls.process.pid) + + +class ChunkedTestPDBase(ChunkedGsm8kMixin, PDDisaggregationServerBase): + __test__ = False + decode_feature_args: ClassVar[List[str]] = [] + + @classmethod + def setUpClass(cls): + cls.extra_prefill_args = cls( + "test_mixed_prefix_gsm8k_chunked" + ).build_prefill_side_args() + canary = list(KV_CANARY_ARGS) if cls.use_kv_canary else [] + cls.extra_decode_args = canary + list(cls.decode_feature_args) + PDDisaggregationServerBase.setUpClass() + cls.model = try_cached_model(cls.model) + cls.launch_all() + + @classmethod + def tearDownClass(cls): + PDDisaggregationServerBase.tearDownClass() diff --git a/python/sglang/test/scripted_runtime/context/api.py b/python/sglang/test/scripted_runtime/context/api.py index 26090db1a..8f924748a 100644 --- a/python/sglang/test/scripted_runtime/context/api.py +++ b/python/sglang/test/scripted_runtime/context/api.py @@ -64,6 +64,8 @@ class ScriptedContext: return_logprob: bool = False, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, lora_path: Optional[str] = None, ) -> "ScriptedReqHandle": return self._req_starter.start_req( @@ -77,6 +79,8 @@ class ScriptedContext: return_logprob=return_logprob, logprob_start_len=logprob_start_len, top_logprobs_num=top_logprobs_num, + stop_token_ids=stop_token_ids, + temperature=temperature, lora_path=lora_path, ) @@ -89,8 +93,8 @@ class ScriptedContext: def abort_all(self) -> None: return lifecycle.abort_all(self) - def abort(self, handle: "ScriptedReqHandle") -> None: - return lifecycle.abort(self, rid=handle.rid) + def abort(self, handle: "ScriptedReqHandle", *, await_arrival: bool = True) -> None: + return lifecycle.abort(self, rid=handle.rid, await_arrival=await_arrival) def flush_cache(self) -> None: return lifecycle.flush_cache(self) diff --git a/python/sglang/test/scripted_runtime/context/http_post.py b/python/sglang/test/scripted_runtime/context/http_post.py index 447759589..54e26cdb4 100644 --- a/python/sglang/test/scripted_runtime/context/http_post.py +++ b/python/sglang/test/scripted_runtime/context/http_post.py @@ -20,18 +20,33 @@ def _http_post_and_await_recv_msg( description: str, timeout_s: float = RECV_MSG_ARRIVAL_TIMEOUT_S, ) -> None: - server_args = ctx.scheduler.server_args - url = f"http://{server_args.host}:{server_args.port}{path}" - - async def _post() -> None: - try: - await ctx._http_poster.post(url, json) - except Exception: # noqa: BLE001 — fire-and-forget background POST - logger.exception("scripted_runtime: POST %s failed", path) - - ctx._http_poster.submit_coro(_post()) + _submit_post(ctx, path=path, json=json) ctx._tokenizer_recv_proxy.wait_until_arrived( predicate, timeout_s=timeout_s, description=description, ) + + +def _http_post_fire_and_forget( + ctx: "ScriptedContext", + *, + path: str, + json: Optional[Dict[str, Any]], +) -> None: + _submit_post(ctx, path=path, json=json) + + +def _submit_post( + ctx: "ScriptedContext", + *, + path: str, + json: Optional[Dict[str, Any]], +) -> None: + server_args = ctx.scheduler.server_args + url = f"http://{server_args.host}:{server_args.port}{path}" + + async def _post() -> None: + await ctx._http_poster.post(url, json) + + ctx._http_poster.submit_coro(_post()) diff --git a/python/sglang/test/scripted_runtime/context/lifecycle.py b/python/sglang/test/scripted_runtime/context/lifecycle.py index d66c838a9..4d966b470 100644 --- a/python/sglang/test/scripted_runtime/context/lifecycle.py +++ b/python/sglang/test/scripted_runtime/context/lifecycle.py @@ -10,6 +10,7 @@ from sglang.srt.managers.io_struct import ( ) from sglang.test.scripted_runtime.context.http_post import ( _http_post_and_await_recv_msg, + _http_post_fire_and_forget, ) if TYPE_CHECKING: @@ -17,8 +18,16 @@ if TYPE_CHECKING: def _await_control( - ctx: "ScriptedContext", *, path: str, json, expect_type: type + ctx: "ScriptedContext", + *, + path: str, + json, + expect_type: type, + await_arrival: bool = True, ) -> None: + if not await_arrival: + _http_post_fire_and_forget(ctx, path=path, json=json) + return _http_post_and_await_recv_msg( ctx, path=path, @@ -57,12 +66,13 @@ def abort_all(ctx: "ScriptedContext") -> None: ) -def abort(ctx: "ScriptedContext", *, rid: str) -> None: +def abort(ctx: "ScriptedContext", *, rid: str, await_arrival: bool = True) -> None: _await_control( ctx, path="/abort_request", json={"rid": rid, "abort_all": False}, expect_type=AbortReq, + await_arrival=await_arrival, ) diff --git a/python/sglang/test/scripted_runtime/context/queries.py b/python/sglang/test/scripted_runtime/context/queries.py index 45632a731..ce0639532 100644 --- a/python/sglang/test/scripted_runtime/context/queries.py +++ b/python/sglang/test/scripted_runtime/context/queries.py @@ -80,9 +80,20 @@ def find_req_by_rid(ctx: "ScriptedContext", rid: str) -> Optional["Req"]: def is_finished(ctx: "ScriptedContext", rid: str) -> bool: req = find_req_by_rid(ctx, rid) - if req is None: - return rid in ctx._seen_rids - return req.finished() + if req is not None: + return req.finished() + if rid in ctx._seen_rids: + return True + # Fallback: if the req ran in a forward batch (recorded in _batch_log) but + # is now absent from all active scheduler sets, it must have finished. + # This catches requests that completed without ever being observed via + # find_req_by_rid (e.g. when Python short-circuit evaluation prevents the + # query while another request is still running). + log = ctx._scheduler_hook._batch_log + if any(rid in record.rids for record in log): + ctx._seen_rids.add(rid) + return True + return False def is_chunking(ctx: "ScriptedContext", rid: str) -> bool: diff --git a/python/sglang/test/scripted_runtime/context/req_starter.py b/python/sglang/test/scripted_runtime/context/req_starter.py index ff9ee8cab..99c188ad4 100644 --- a/python/sglang/test/scripted_runtime/context/req_starter.py +++ b/python/sglang/test/scripted_runtime/context/req_starter.py @@ -1,7 +1,7 @@ from __future__ import annotations import uuid -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, List, Optional from sglang.test.scripted_runtime.context.http_post import ( _http_post_and_await_recv_msg, @@ -30,6 +30,8 @@ class ScriptedContextReqStarter: return_logprob: bool = False, logprob_start_len: Optional[int] = None, top_logprobs_num: Optional[int] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, lora_path: Optional[str] = None, ) -> ScriptedReqHandle: ctx = self._ctx @@ -39,6 +41,10 @@ class ScriptedContextReqStarter: self._req_counter += 1 sampling_params = {"max_new_tokens": max_new_tokens, "ignore_eos": ignore_eos} + if stop_token_ids is not None: + sampling_params["stop_token_ids"] = stop_token_ids + if temperature is not None: + sampling_params["temperature"] = temperature payload = { "input_ids": [prompt_token] * prompt_len, "sampling_params": sampling_params, diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index 945d763cf..d5d1301a3 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -3,6 +3,8 @@ from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING, Optional +from sglang.test.scripted_runtime.context.radix import _node_lock_ref + if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.test.scripted_runtime.context.api import ScriptedContext @@ -47,5 +49,10 @@ class ScriptedReqHandle: @property def lock_refs(self) -> int: - node = self.req.last_node - return node.lock_ref if node is not None else 0 + req = self.req + if req is None: + return 0 + node = req.last_node + if node is None: + return 0 + return _node_lock_ref(node) diff --git a/python/sglang/test/scripted_runtime/scheduler_hook.py b/python/sglang/test/scripted_runtime/scheduler_hook.py index 494638f30..4411945b3 100644 --- a/python/sglang/test/scripted_runtime/scheduler_hook.py +++ b/python/sglang/test/scripted_runtime/scheduler_hook.py @@ -95,6 +95,9 @@ def _drive_engine_through_warmup(ctx: ScriptedContext) -> Generator: def _reset_engine_state(ctx: ScriptedContext) -> Generator: scheduler = ctx.scheduler + if scheduler._engine_paused: + ctx.continue_generation() + ctx._release_exhausted_pools() ctx.abort_all() for _ in range(RESET_DRAIN_MAX_STEPS): diff --git a/python/sglang/test/scripted_runtime_chunked_helpers.py b/python/sglang/test/scripted_runtime_chunked_helpers.py index a5b0fac15..393d064ff 100644 --- a/python/sglang/test/scripted_runtime_chunked_helpers.py +++ b/python/sglang/test/scripted_runtime_chunked_helpers.py @@ -41,13 +41,16 @@ def run_until_finished(handle, *, max_steps: int = DEFAULT_MAX_STEPS): def run_until_all_finished(handles: List[Any], *, max_steps: int = DEFAULT_MAX_STEPS): + done = [False] * len(handles) for _ in range(max_steps): - if all(h.finished for h in handles): + for i, h in enumerate(handles): + done[i] = done[i] or h.finished + if all(done): return yield raise AssertionError( f"run_until_all_finished: not all reqs finished after {max_steps} " - f"steps (finished={[h.finished for h in handles]})" + f"steps (finished={done})" ) @@ -65,6 +68,12 @@ def warmup_radix(t, prompt_tokens: List[int], *, max_steps: int = DEFAULT_MAX_ST BALLAST_MAX_NEW_TOKENS: int = 30000 +SMALL_KV_POOL_MAX_TOTAL_TOKENS: int = 4096 + +SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS: int = 512 + +SMALL_KV_POOL_BALLAST_PROMPT_LEN: int = 1536 + def exhaust_row_pool(t, *, leave_rows: int, max_steps: int = DEFAULT_MAX_STEPS): target: int = t.scheduler.req_to_token_pool.available_size() - leave_rows diff --git a/test/manual/chunked_prefill/__init__.py b/test/manual/chunked_prefill/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/manual/chunked_prefill/test_e2e_disagg.py b/test/manual/chunked_prefill/test_e2e_disagg.py new file mode 100644 index 000000000..0c596cd31 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_disagg.py @@ -0,0 +1,13 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestPDBase + + +class TestChunkedFeatureDisagg(ChunkedTestPDBase): + __test__ = True + use_kv_canary = False + gsm8k_threshold = 0.50 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_dp_attention.py b/test/manual/chunked_prefill/test_e2e_dp_attention.py new file mode 100644 index 000000000..e5c018d27 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_dp_attention.py @@ -0,0 +1,23 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase +from sglang.test.test_utils import DEFAULT_MLA_MODEL_NAME_FOR_TEST + + +class TestChunkedFeatureDPAttention(ChunkedTestBase): + __test__ = True + use_kv_canary = False + model = DEFAULT_MLA_MODEL_NAME_FOR_TEST + gsm8k_threshold = 0.50 + feature_args = [ + "--trust-remote-code", + "--tp", + "2", + "--enable-dp-attention", + "--dp", + "2", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_hybrid_swa.py b/test/manual/chunked_prefill/test_e2e_hybrid_swa.py new file mode 100644 index 000000000..6d2e5c72a --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_hybrid_swa.py @@ -0,0 +1,22 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ( + LONG_PROMPT_NUM_SHOTS, + ChunkedTestBase, +) + + +class TestChunkedFeatureHybridSWA(ChunkedTestBase): + __test__ = True + model = "openai/gpt-oss-20b" + num_shots = LONG_PROMPT_NUM_SHOTS + gsm8k_threshold = 0.50 + feature_args = [ + "--mem-fraction-static", + "0.70", + "--disable-piecewise-cuda-graph", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_lora.py b/test/manual/chunked_prefill/test_e2e_lora.py new file mode 100644 index 000000000..2dc41f465 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_lora.py @@ -0,0 +1,25 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase + + +class TestChunkedFeatureLoRA(ChunkedTestBase): + __test__ = True + model = "meta-llama/Llama-3.2-1B-Instruct" + gsm8k_threshold = 0.20 + feature_args = [ + "--enable-lora", + "--lora-paths", + "nicoboss/Llama-3.2-1B-Instruct-Uncensored-Lora", + "codelion/Llama-3.2-1B-Instruct-tool-calling-lora", + "--max-loras-per-batch", + "2", + "--max-loaded-loras", + "4", + "--mem-fraction-static", + "0.75", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_lora_overlap.py b/test/manual/chunked_prefill/test_e2e_lora_overlap.py new file mode 100644 index 000000000..b7cad5bc5 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_lora_overlap.py @@ -0,0 +1,26 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase + + +class TestChunkedFeatureLoRAOverlap(ChunkedTestBase): + __test__ = True + model = "meta-llama/Llama-3.2-1B-Instruct" + gsm8k_threshold = 0.20 + feature_args = [ + "--enable-lora", + "--enable-lora-overlap-loading", + "--lora-paths", + "nicoboss/Llama-3.2-1B-Instruct-Uncensored-Lora", + "codelion/Llama-3.2-1B-Instruct-tool-calling-lora", + "--max-loras-per-batch", + "2", + "--max-loaded-loras", + "4", + "--mem-fraction-static", + "0.75", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_page_size.py b/test/manual/chunked_prefill/test_e2e_page_size.py new file mode 100644 index 000000000..e31bdcb92 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_page_size.py @@ -0,0 +1,18 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + + +class TestChunkedFeaturePageSize(ChunkedTestBase): + __test__ = True + model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + gsm8k_threshold = 0.30 + feature_args = [ + "--page-size", + "16", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_pd_pp.py b/test/manual/chunked_prefill/test_e2e_pd_pp.py new file mode 100644 index 000000000..eabbd2a40 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_pd_pp.py @@ -0,0 +1,25 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestPDBase + + +class TestChunkedFeaturePDPP(ChunkedTestPDBase): + __test__ = True + gsm8k_threshold = 0.50 + feature_args = [ + "--tp-size", + "2", + "--pp-size", + "2", + "--disable-overlap-schedule", + ] + decode_feature_args = [ + "--tp-size", + "2", + "--base-gpu-id", + "4", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_piecewise_cuda_graph.py b/test/manual/chunked_prefill/test_e2e_piecewise_cuda_graph.py new file mode 100644 index 000000000..4f1b7fc9a --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_piecewise_cuda_graph.py @@ -0,0 +1,16 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + + +class TestChunkedFeaturePiecewiseCudaGraph(ChunkedTestBase): + __test__ = True + use_kv_canary = False + model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + gsm8k_threshold = 0.30 + feature_args = [] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_pp.py b/test/manual/chunked_prefill/test_e2e_pp.py new file mode 100644 index 000000000..10e005eed --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_pp.py @@ -0,0 +1,19 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase + + +class TestChunkedFeaturePP(ChunkedTestBase): + __test__ = True + gsm8k_threshold = 0.50 + feature_args = [ + "--tp-size", + "2", + "--pp-size", + "2", + "--enable-dynamic-chunking", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_priority.py b/test/manual/chunked_prefill/test_e2e_priority.py new file mode 100644 index 000000000..d1344abac --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_priority.py @@ -0,0 +1,21 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + + +class TestChunkedFeaturePriority(ChunkedTestBase): + __test__ = True + model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + gsm8k_threshold = 0.30 + feature_args = [ + "--enable-priority-scheduling", + "--max-running-requests", + "8", + "--max-queued-requests", + "128", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_radix.py b/test/manual/chunked_prefill/test_e2e_radix.py new file mode 100644 index 000000000..60191a1e1 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_radix.py @@ -0,0 +1,20 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + + +class TestChunkedFeatureRadix(ChunkedTestBase): + __test__ = True + model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST + gsm8k_threshold = 0.30 + feature_args = [ + "--max-total-tokens", + "20000", + "--schedule-policy", + "fcfs", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_e2e_spec.py b/test/manual/chunked_prefill/test_e2e_spec.py new file mode 100644 index 000000000..140dad806 --- /dev/null +++ b/test/manual/chunked_prefill/test_e2e_spec.py @@ -0,0 +1,28 @@ +import unittest + +from sglang.test.chunked_prefill_test_utils import ChunkedTestBase + + +class TestChunkedFeatureSpec(ChunkedTestBase): + __test__ = True + use_kv_canary = False + model = "Qwen/Qwen3-8B" + gsm8k_threshold = 0.50 + feature_args = [ + "--speculative-algorithm", + "EAGLE3", + "--speculative-draft-model-path", + "Tengyunw/qwen3_8b_eagle3", + "--speculative-num-steps", + "5", + "--speculative-eagle-topk", + "8", + "--speculative-num-draft-tokens", + "64", + "--mem-fraction-static", + "0.7", + ] + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_abort.py b/test/manual/chunked_prefill/test_scripted_abort.py new file mode 100644 index 000000000..fe399a222 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_abort.py @@ -0,0 +1,643 @@ +import unittest + +from sglang.srt.environ import envs +from sglang.test.scripted_runtime.context import ScriptedContext +from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle +from sglang.test.scripted_runtime.test_case import ScriptedTestCase +from sglang.test.scripted_runtime_chunked_helpers import ( + DEFAULT_CHUNK_SIZE, + DEFAULT_MAX_STEPS, + SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + SMALL_KV_POOL_BALLAST_PROMPT_LEN, + SMALL_KV_POOL_MAX_TOTAL_TOKENS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle): + for _ in range(12): + if all( + h.kv_pages == 0 + and h.lock_refs == 0 + and (h.req is None or h.req.req_pool_idx is None) + for h in handles + ): + return + yield + + +class TestAbortBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_abort_waiting_chunked_resume(self): + self.server.execute_script(self._script_abort_waiting_chunked_resume) + + @staticmethod + def _script_abort_waiting_chunked_resume(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + ) + yield from run_until(r, lambda h: h.is_chunking) + + pages_before = r.kv_pages + assert pages_before > 0, "chunked req should own KV pages mid-chunk" + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.status in ( + "finished", + "unknown", + ), f"after abort r should be finished/unknown, got {r.status}" + assert ( + r.kv_pages == 0 + ), f"abort must release KV; r.kv_pages={r.kv_pages} after abort" + assert ( + r.req is None or r.req.req_pool_idx is None + ), f"abort must release row; r.req={r.req} after abort" + assert ( + r.lock_refs == 0 + ), f"abort must release lock_refs; r.lock_refs={r.lock_refs}" + + def test_abort_at_chunk_0(self): + self.server.execute_script(self._script_abort_at_chunk_0) + + @staticmethod + def _script_abort_at_chunk_0(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=110 + ) + yield + yield from run_until(r, lambda h: h.is_chunking) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + + def test_abort_at_chunk_mid(self): + self.server.execute_script(self._script_abort_at_chunk_mid) + + @staticmethod + def _script_abort_at_chunk_mid(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=120 + ) + yield from run_until(r, lambda h: h.chunks_done >= 2 and h.is_chunking) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + + def test_abort_one_does_not_disturb_other(self): + self.server.execute_script(self._script_abort_one_does_not_disturb_other) + + @staticmethod + def _script_abort_one_does_not_disturb_other(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=130 + ) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=131 + ) + yield from run_until(r1, lambda h: h.is_chunking) + + t.abort(r1) + yield from _drain_until_released(t, r1) + + assert r1.kv_pages == 0 + yield from run_until_finished(r2) + assert r2.finished, "r2 should still complete after r1 is aborted" + + def test_abort_with_zero_yield(self): + self.server.execute_script(self._script_abort_with_zero_yield) + + @staticmethod + def _script_abort_with_zero_yield(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + assert r.lock_refs == 0 + + def test_abort_at_admission_step(self): + self.server.execute_script(self._script_abort_at_admission_step) + + @staticmethod + def _script_abort_at_admission_step(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + + def test_abort_then_start_same_step_new_rid(self): + self.server.execute_script(self._script_abort_then_start_same_step_new_rid) + + @staticmethod + def _script_abort_then_start_same_step_new_rid(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=140 + ) + yield from run_until(r1, lambda h: h.is_chunking) + t.abort(r1) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=141 + ) + yield from run_until_finished(r2) + assert r2.finished + assert r1.kv_pages == 0 + + def test_abort_then_start_same_step_same_rid(self): + self.server.execute_script(self._script_abort_then_start_same_step_same_rid) + + @staticmethod + def _script_abort_then_start_same_step_same_rid(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + rid="abort-reuse", + prompt_token=150, + ) + yield from run_until(r1, lambda h: h.is_chunking) + t.abort(r1) + yield + r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="abort-reuse") + yield from run_until_finished(r2) + assert r2.finished + + def test_abort_five_chunked_in_a_row(self): + self.server.execute_script(self._script_abort_five_chunked_in_a_row) + + @staticmethod + def _script_abort_five_chunked_in_a_row(t: ScriptedContext): + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=160 + i + ) + for i in range(5) + ] + yield from run_until(reqs[0], lambda h: h.is_chunking) + for r in reqs: + t.abort(r) + yield from _drain_until_released(t, *reqs) + for r in reqs: + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + + def test_abort_unknown_rid_noop(self): + self.server.execute_script(self._script_abort_unknown_rid_noop) + + @staticmethod + def _script_abort_unknown_rid_noop(t: ScriptedContext): + bogus = ScriptedReqHandle(rid="never-submitted-rid", context=t) + t.abort(bogus, await_arrival=False) + yield + r = t.start_req(prompt_len=16, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + for _ in range(12): + if r.kv_pages == 0 and r.lock_refs == 0: + break + yield + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_abort_after_finish_noop(self): + self.server.execute_script(self._script_abort_after_finish_noop) + + @staticmethod + def _script_abort_after_finish_noop(t: ScriptedContext): + r = t.start_req(prompt_len=16, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + + for _ in range(12): + if t.is_fully_idle: + break + yield + assert r.kv_pages == 0 + assert r.lock_refs == 0 + kv_pool_free_before = t.engine_stats()["kv_pool_free"] + + t.abort(r, await_arrival=False) + yield + assert r.kv_pages == 0 + assert r.lock_refs == 0 + kv_pool_free_after = t.engine_stats()["kv_pool_free"] + assert kv_pool_free_after == kv_pool_free_before, ( + f"abort-after-finish must not move KV pool; " + f"before={kv_pool_free_before} after={kv_pool_free_after}" + ) + + def test_abort_chunk_last(self): + self.server.execute_script(self._script_abort_chunk_last) + + @staticmethod + def _script_abort_chunk_last(t: ScriptedContext): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, prompt_token=170 + ) + yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking) + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.req is None or r.req.inflight_middle_chunks == 0 + + def test_abort_penultimate_chunk(self): + self.server.execute_script(self._script_abort_penultimate_chunk) + + @staticmethod + def _script_abort_penultimate_chunk(t: ScriptedContext): + r = t.start_req( + prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=180 + ) + yield from run_until(r, lambda h: h.chunks_done >= 2 and h.is_chunking) + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + assert r.lock_refs == 0 + + def test_double_abort_idempotent(self): + self.server.execute_script(self._script_double_abort_idempotent) + + @staticmethod + def _script_double_abort_idempotent(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=190 + ) + yield from run_until(r, lambda h: h.is_chunking) + t.abort(r) + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_abort_during_decode(self): + self.server.execute_script(self._script_abort_during_decode) + + @staticmethod + def _script_abort_during_decode(t: ScriptedContext): + r = t.start_req(prompt_len=16, max_new_tokens=64) + yield from run_until(r, lambda h: h.status == "running") + assert r.kv_pages > 0, "decode req must own KV before abort" + t.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_abort_one_of_three_others_finish(self): + self.server.execute_script(self._script_abort_one_of_three_others_finish) + + @staticmethod + def _script_abort_one_of_three_others_finish(t: ScriptedContext): + r1 = t.start_req(prompt_len=16, max_new_tokens=4) + r2 = t.start_req(prompt_len=16, max_new_tokens=4) + r3 = t.start_req(prompt_len=16, max_new_tokens=4) + yield from run_until(r2, lambda h: h.status == "running") + t.abort(r2) + yield from run_until_all_finished([r1, r3]) + assert r2.kv_pages == 0 + + def test_abort_in_separate_yields(self): + self.server.execute_script(self._script_abort_in_separate_yields) + + @staticmethod + def _script_abort_in_separate_yields(t: ScriptedContext): + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=200 + i + ) + for i in range(3) + ] + yield from run_until(reqs[0], lambda h: h.is_chunking) + for r in reqs: + t.abort(r) + yield + yield from _drain_until_released(t, *reqs) + for r in reqs: + assert r.kv_pages == 0 + + def test_abort_at_chunk_boundary_race(self): + self.server.execute_script(self._script_abort_at_chunk_boundary_race) + + @staticmethod + def _script_abort_at_chunk_boundary_race(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=210 + ) + yield from run_until(r, lambda h: h.is_chunking) + yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + + chunks_after_abort = r.chunks_done + yield + yield + assert not r.is_chunking, "aborted req must not resume chunking" + assert r.chunks_done == chunks_after_abort, ( + f"aborted req revived and ran another chunk; " + f"chunks_done went {chunks_after_abort} -> {r.chunks_done}" + ) + assert r.req is None or r.req.req_pool_idx is None + + def test_abort_mid_chunk_no_extra_radix_node(self): + self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node) + + @staticmethod + def _script_abort_mid_chunk_no_extra_radix_node( + t: ScriptedContext, + ): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=220 + ) + yield from run_until(r, lambda h: h.is_chunking) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + chunks_after_release = r.chunks_done + for _ in range(4): + yield + assert r.chunks_done == chunks_after_release, ( + f"aborted mid-chunk req revived after release; chunks_done went " + f"{chunks_after_release} -> {r.chunks_done}" + ) + + lock_refs_after = t.get_all_node_lock_refs() + assert all(ref == 0 for ref in lock_refs_after.values()), ( + f"abort mid-chunk left a locked radix node behind; " + f"node lock_refs={lock_refs_after!r}" + ) + + def test_abort_then_resubmit_same_rid_same_step(self): + self.server.execute_script(self._script_abort_then_resubmit_same_rid_same_step) + + @staticmethod + def _script_abort_then_resubmit_same_rid_same_step(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + rid="abort-resubmit-same-step", + prompt_token=230, + ) + yield from run_until(r1, lambda h: h.is_chunking) + t.abort(r1) + r2 = t.start_req( + prompt_len=16, + max_new_tokens=2, + rid="abort-resubmit-same-step", + ) + yield + yield from run_until_finished(r2) + assert r2.finished, "resubmit under same rid must complete independently" + assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit" + assert r1.req is None or r1.req.req_pool_idx is None + assert r1.lock_refs == 0 + + def test_abort_during_gap_inflight_middle_chunks_positive(self): + self.server.execute_script( + self._script_abort_during_gap_inflight_middle_chunks_positive + ) + + @staticmethod + def _script_abort_during_gap_inflight_middle_chunks_positive(t: ScriptedContext): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=240 + ) + yield from run_until( + r, + lambda h: h.is_chunking and h.chunks_done >= 1, + ) + assert r.req.inflight_middle_chunks > 0 + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + + assert not r.is_chunking, "aborted gap req must not re-enter chunking" + yield + assert not r.is_chunking, "aborted gap req must stay out of chunking" + + if r.req is not None: + assert ( + r.req.inflight_middle_chunks == 0 + ), f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}" + + def test_abort_when_chunked_only_then_idle(self): + self.server.execute_script(self._script_abort_when_chunked_only_then_idle) + + @staticmethod + def _script_abort_when_chunked_only_then_idle(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=250 + ) + yield from run_until(r, lambda h: h.is_chunking) + assert (1 if t.scheduler.chunked_req is not None else 0) == 1 + + t.abort(r) + yield from _drain_until_released(t, r) + + for _ in range(12): + if t.scheduler.chunked_req is None and t.is_idle: + break + yield + + assert r.kv_pages == 0 + assert (1 if t.scheduler.chunked_req is not None else 0) == 0 + assert t.is_idle, "engine must be idle after the only chunked req is aborted" + + def test_chunked_req_then_abort_then_new_short_in_one_yield(self): + self.server.execute_script( + self._script_chunked_req_then_abort_then_new_short_in_one_yield + ) + + @staticmethod + def _script_chunked_req_then_abort_then_new_short_in_one_yield(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=260 + ) + yield from run_until(r1, lambda h: h.is_chunking) + assert ( + t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None + ) == r1.rid, ( + f"r1 should hold the chunked slot before abort; got " + f"{(t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None)!r}" + ) + + t.abort(r1) + r2 = t.start_req(prompt_len=16, max_new_tokens=2) + yield from _drain_until_released(t, r1) + + cur = ( + t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None + ) + assert cur != r1.rid, f"chunked slot still points to aborted r1; got {cur!r}" + assert r1.kv_pages == 0 + yield from run_until_finished(r2) + assert r2.finished, "fresh r2 must admit and complete after combo step" + + def test_force_retract_then_abort_same_yield(self): + self.server.execute_script(self._script_force_retract_then_abort_same_yield) + + @staticmethod + def _script_force_retract_then_abort_same_yield(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=270 + ) + yield from run_until(r1, lambda h: h.is_chunking) + assert r1.kv_pages > 0 + + t.pause_generation(mode="retract") + t.abort(r1) + yield from _drain_until_released(t, r1) + + assert r1.kv_pages == 0, ( + f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}" + ) + assert r1.req is None or r1.req.req_pool_idx is None, ( + f"force_retract + abort same yield must release row; got " f"{r1.req}" + ) + assert r1.lock_refs == 0, ( + f"force_retract + abort same yield must release lock_refs; " + f"got {r1.lock_refs}" + ) + yield + t.continue_generation() + + def test_abort_chunked_with_baton_handoff(self): + self.server.execute_script(self._script_abort_chunked_with_baton_handoff) + + @staticmethod + def _script_abort_chunked_with_baton_handoff(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=280 + ) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=281 + ) + yield from run_until(r1, lambda h: h.is_chunking) + assert (1 if t.scheduler.chunked_req is not None else 0) == 1 + + t.abort(r1) + yield from _drain_until_released(t, r1) + + yield from run_until(r2, lambda h: h.is_chunking) + assert r1.kv_pages == 0 + assert r1.req is None or r1.req.req_pool_idx is None + assert r1.lock_refs == 0 + yield from run_until_finished(r2) + assert r2.finished, "baton handoff must let r2 complete" + assert r2.lock_refs == 0 + + +class TestAbortPP(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + pp_size=2, + ) + + def test_abort_at_last_chunk_in_flight_pp(self): + self.server.execute_script(self._script_abort_at_last_chunk_in_flight_pp) + + @staticmethod + def _script_abort_at_last_chunk_in_flight_pp(t: ScriptedContext): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, prompt_token=290 + ) + yield from run_until( + r, + lambda h: h.chunks_done >= 1 and h.is_chunking, + ) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + assert r.lock_refs == 0 + assert r.finished + + +class TestAbortSmallPool(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + max_total_tokens=SMALL_KV_POOL_MAX_TOTAL_TOKENS, + ) + + def test_waiting_timeout_sweep_aborts_pressured_waiting_req(self): + self.server.execute_script( + self._script_waiting_timeout_sweep_aborts_pressured_waiting_req + ) + + @staticmethod + def _script_waiting_timeout_sweep_aborts_pressured_waiting_req(t: ScriptedContext): + + b1 = t.start_req( + prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN, + max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + prompt_token=300, + ) + b2 = t.start_req( + prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN, + max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + prompt_token=301, + ) + yield from run_until(b1, lambda h: h.status == "running") + yield from run_until(b2, lambda h: h.status == "running") + + r = t.start_req(prompt_len=16, max_new_tokens=2) + + def waiting_rids(): + return {req.rid for req in t.scheduler.waiting_queue} + + yield from run_until(r, lambda h: r.rid in waiting_rids()) + assert r.kv_pages == 0, "pressured waiting req must not own KV before admission" + + with envs.SGLANG_REQ_WAITING_TIMEOUT.override(1e-6): + for _ in range(DEFAULT_MAX_STEPS): + if r.rid not in waiting_rids(): + break + yield + else: + raise AssertionError( + f"waiting-timeout sweep never removed the req from " + f"waiting_queue after {DEFAULT_MAX_STEPS} steps; " + f"waiting_rids={waiting_rids()!r}" + ) + + assert r.rid not in waiting_rids(), ( + f"the loop's waiting-timeout sweep must drop the timed-out waiting " + f"req from waiting_queue; got {waiting_rids()!r}" + ) + assert r.status in ("finished", "unknown"), ( + f"swept-out req must be aborted (gone from every live scheduler " + f"structure); got status={r.status!r}" + ) + assert r.kv_pages == 0, "timeout-abort of an unadmitted req owns no KV" + + t.abort(b1) + t.abort(b2) + yield from _drain_until_released(t, b1, b2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_chunk_size.py b/test/manual/chunked_prefill/test_scripted_chunk_size.py new file mode 100644 index 000000000..975461f67 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_chunk_size.py @@ -0,0 +1,361 @@ +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, + base_engine_kwargs, + run_until, + run_until_finished, +) + + +def _expected_chunks(prompt_len: int, chunk_size: int) -> int: + if prompt_len <= chunk_size: + return 0 + return (prompt_len + chunk_size - 1) // chunk_size + + +class TestChunkSizeDefault(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_exact_chunk_size(self): + self.server.execute_script(self._script_exact_chunk_size) + + @staticmethod + def _script_exact_chunk_size(t: ScriptedContext): + r = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0, ( + f"prompt_len == chunk_size completes in one non-chunked shot, " + f"got chunks_done={r.chunks_done}" + ) + + def test_one_token_over(self): + self.server.execute_script(self._script_one_token_over) + + @staticmethod + def _script_one_token_over(t: ScriptedContext): + r = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE + 1, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2, ( + f"prompt_len == chunk_size + 1 expected exactly 2 chunks, " + f"got chunks_done={r.chunks_done}" + ) + + def test_two_chunks_exact(self): + self.server.execute_script(self._script_two_chunks_exact) + + @staticmethod + def _script_two_chunks_exact(t: ScriptedContext): + r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2 + + def test_three_chunks_one_decode(self): + self.server.execute_script(self._script_three_chunks_one_decode) + + @staticmethod + def _script_three_chunks_one_decode(t: ScriptedContext): + r = t.start_req(prompt_len=3 * DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 3 + + def test_tiny_prompt(self): + self.server.execute_script(self._script_tiny_prompt) + + @staticmethod + def _script_tiny_prompt(t: ScriptedContext): + r = t.start_req(prompt_len=1, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert ( + r.chunks_done == 0 + ), f"single-token prompt should not chunk, got chunks_done={r.chunks_done}" + + def test_chunk_size_256_prompt_100x(self): + self.server.execute_script(self._script_chunk_size_256_prompt_100x) + + @staticmethod + def _script_chunk_size_256_prompt_100x(t: ScriptedContext): + r = t.start_req(prompt_len=100 * 256, max_new_tokens=2) + yield from run_until(r, lambda h: h.finished, max_steps=4000) + assert r.finished + assert r.chunks_done == 100 + + def test_prompt_n_chunks_plus_minus_1(self): + self.server.execute_script(self._script_prompt_n_chunks_plus_minus_1) + + @staticmethod + def _script_prompt_n_chunks_plus_minus_1(t: ScriptedContext): + for n in range(1, 6): + prompt_minus = n * DEFAULT_CHUNK_SIZE - 1 + r_minus = t.start_req( + prompt_len=prompt_minus, max_new_tokens=1, prompt_token=2 * n + ) + yield from run_until_finished(r_minus, max_steps=800) + assert r_minus.finished + expected_minus = _expected_chunks(prompt_minus, DEFAULT_CHUNK_SIZE) + assert r_minus.chunks_done == expected_minus, ( + f"N={n} prompt_len={prompt_minus}: " + f"expected chunks_done={expected_minus}, got {r_minus.chunks_done}" + ) + + prompt_plus = n * DEFAULT_CHUNK_SIZE + 1 + r_plus = t.start_req( + prompt_len=prompt_plus, max_new_tokens=1, prompt_token=2 * n + 1 + ) + yield from run_until_finished(r_plus, max_steps=800) + assert r_plus.finished + expected_plus = _expected_chunks(prompt_plus, DEFAULT_CHUNK_SIZE) + assert r_plus.chunks_done == expected_plus, ( + f"N={n} prompt_len={prompt_plus}: " + f"expected chunks_done={expected_plus}, got {r_plus.chunks_done}" + ) + + +class TestChunkSizeDefaultRadixExplicit(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + disable_radix_cache=False, + ) + + def test_radix_hit_minus_one(self): + self.server.execute_script(self._script_radix_hit_minus_one) + + @staticmethod + def _script_radix_hit_minus_one(t: ScriptedContext): + r_warm = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE - 1, max_new_tokens=1) + yield from run_until_finished(r_warm) + assert r_warm.finished + + yield + + r = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0, ( + f"radix hit should leave only 1 fresh token; chunked path should " + f"not engage. Got chunks_done={r.chunks_done}" + ) + + +class TestChunkSize1(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=1) + + def test_chunk_size_one_token(self): + self.server.execute_script(self._script_chunk_size_one_token) + + @staticmethod + def _script_chunk_size_one_token(t: ScriptedContext): + r = t.start_req(prompt_len=8, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 8, f"expected 8 chunks, got {r.chunks_done}" + + def test_chunks_done_monotone_under_chunk_size_1(self): + self.server.execute_script(self._script_chunks_done_monotone_under_chunk_size_1) + + @staticmethod + def _script_chunks_done_monotone_under_chunk_size_1(t: ScriptedContext): + r = t.start_req(prompt_len=16, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 16, ( + f"chunk_size=1 with prompt_len=16 chunks one token at a time " + f"(ceil(16/1)=16), got chunks_done={r.chunks_done}" + ) + + def test_chunk_size_one_long_prompt(self): + self.server.execute_script(self._script_chunk_size_one_long_prompt) + + @staticmethod + def _script_chunk_size_one_long_prompt(t: ScriptedContext): + r = t.start_req(prompt_len=64, max_new_tokens=2) + yield from run_until(r, lambda h: h.finished, max_steps=500) + assert r.finished + assert r.chunks_done == 64 + + +class TestChunkSize2(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=2) + + def test_chunk_size_two_prompt_two(self): + self.server.execute_script(self._script_chunk_size_two_prompt_two) + + @staticmethod + def _script_chunk_size_two_prompt_two(t: ScriptedContext): + r = t.start_req(prompt_len=2, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + + def test_chunk_size_two_prompt_five(self): + self.server.execute_script(self._script_chunk_size_two_prompt_five) + + @staticmethod + def _script_chunk_size_two_prompt_five(t: ScriptedContext): + r = t.start_req(prompt_len=5, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 3, f"expected 3 chunks, got {r.chunks_done}" + + +class TestChunkSize4(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=4) + + def test_chunk_size_4_prompt_4(self): + self.server.execute_script(self._script_chunk_size_4_prompt_4) + + @staticmethod + def _script_chunk_size_4_prompt_4(t: ScriptedContext): + r = t.start_req(prompt_len=4, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + + def test_chunk_size_4_prompt_5(self): + self.server.execute_script(self._script_chunk_size_4_prompt_5) + + @staticmethod + def _script_chunk_size_4_prompt_5(t: ScriptedContext): + r = t.start_req(prompt_len=5, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2 + + +class TestChunkSize16(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=16) + + def test_chunk_size_16_prompt_1024(self): + self.server.execute_script(self._script_chunk_size_16_prompt_1024) + + @staticmethod + def _script_chunk_size_16_prompt_1024(t: ScriptedContext): + r = t.start_req(prompt_len=1024, max_new_tokens=2) + yield from run_until(r, lambda h: h.finished, max_steps=1000) + assert r.finished + assert r.chunks_done == 64 + + +class TestChunkSize32(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=32) + + def test_chunk_size_32_prompt_33(self): + self.server.execute_script(self._script_chunk_size_32_prompt_33) + + @staticmethod + def _script_chunk_size_32_prompt_33(t: ScriptedContext): + r = t.start_req(prompt_len=33, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2 + + +class TestChunkSize64(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=64) + + def test_chunk_size_64_prompt_129(self): + self.server.execute_script(self._script_chunk_size_64_prompt_129) + + @staticmethod + def _script_chunk_size_64_prompt_129(t: ScriptedContext): + r = t.start_req(prompt_len=129, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 3 + + +class TestChunkSize128(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=128) + + def test_chunk_size_128_prompt_512(self): + self.server.execute_script(self._script_chunk_size_128_prompt_512) + + @staticmethod + def _script_chunk_size_128_prompt_512(t: ScriptedContext): + r = t.start_req(prompt_len=512, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 4 + + +class TestChunkSize1024(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=1024) + + def test_chunk_size_1024_prompt_4096(self): + self.server.execute_script(self._script_chunk_size_1024_prompt_4096) + + @staticmethod + def _script_chunk_size_1024_prompt_4096(t: ScriptedContext): + r = t.start_req(prompt_len=4096, max_new_tokens=2) + yield from run_until(r, lambda h: h.finished, max_steps=1000) + assert r.finished + assert r.chunks_done == 4 + + +class TestChunkSize4096(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=4096) + + def test_chunk_size_4096_prompt_4097(self): + self.server.execute_script(self._script_chunk_size_4096_prompt_4097) + + @staticmethod + def _script_chunk_size_4096_prompt_4097(t: ScriptedContext): + r = t.start_req(prompt_len=4097, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2 + + +class TestChunkSize1024MaxPrefill1024(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=1024, max_prefill_tokens=1024 + ) + + def test_chunk_size_equals_max_prefill_tokens(self): + self.server.execute_script(self._script_chunk_size_equals_max_prefill_tokens) + + @staticmethod + def _script_chunk_size_equals_max_prefill_tokens(t: ScriptedContext): + r = t.start_req(prompt_len=1024, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + + +class TestChunkSize2048MaxPrefill1024(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=2048, max_prefill_tokens=1024 + ) + + def test_chunk_size_exceeds_max_prefill_tokens(self): + self.server.execute_script(self._script_chunk_size_exceeds_max_prefill_tokens) + + @staticmethod + def _script_chunk_size_exceeds_max_prefill_tokens(t: ScriptedContext): + r = t.start_req(prompt_len=512, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + + def test_prompt_between_max_prefill_and_chunk_size(self): + self.server.execute_script( + self._script_prompt_between_max_prefill_and_chunk_size + ) + + @staticmethod + def _script_prompt_between_max_prefill_and_chunk_size(t: ScriptedContext): + r = t.start_req(prompt_len=1536, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_http_smoke.py b/test/manual/chunked_prefill/test_scripted_http_smoke.py new file mode 100644 index 000000000..74943a699 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_http_smoke.py @@ -0,0 +1,49 @@ +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, + base_engine_kwargs, + run_until_finished, +) + + +class TestScriptedHttpSmoke(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_chunked_req_is_chunking_then_finishes(self): + self.server.execute_script(self._script_chunked_req_is_chunking_then_finishes) + + @staticmethod + def _script_chunked_req_is_chunking_then_finishes(t: ScriptedContext): + r = t.start_req(prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + saw_chunking = False + for _ in range(800): + if r.is_chunking: + saw_chunking = True + if r.finished: + break + yield + assert r.finished + assert ( + saw_chunking + ), "expected the req to hold the chunked_req slot at least once" + + def test_two_reqs_finish(self): + self.server.execute_script(self._script_two_reqs_finish) + + @staticmethod + def _script_two_reqs_finish(t: ScriptedContext): + r1 = t.start_req(prompt_len=8, max_new_tokens=4) + yield from run_until_finished(r1) + r2 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until_finished(r2) + assert r1.finished + assert r2.finished + assert r1.chunks_done == 0 + assert r2.chunks_done == 2 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_hybrid_swa.py b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py new file mode 100644 index 000000000..3dc808ae6 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_hybrid_swa.py @@ -0,0 +1,180 @@ +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, +) + +_SWA_MODEL = "openai/gpt-oss-20b" +_SWA_WINDOW = 4096 + + +class TestSWABasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_SWA_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + mem_fraction_static=0.70, + disable_piecewise_cuda_graph=True, + ) + + def test_naive_swa_chunked(self): + self.server.execute_script(self._script_naive_swa_chunked) + + @staticmethod + def _script_naive_swa_chunked(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 4096, max_new_tokens=4) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 24 + assert len(r.req.output_ids) == 4 + + def test_swa_prompt_equals_window(self): + self.server.execute_script(self._script_swa_prompt_equals_window) + + @staticmethod + def _script_swa_prompt_equals_window(t: ScriptedContext): + r = t.start_req(prompt_len=_SWA_WINDOW, max_new_tokens=4) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_swa_budget_for_chunked_req_math(self): + self.server.execute_script(self._script_swa_budget_for_chunked_req_math) + + @staticmethod + def _script_swa_budget_for_chunked_req_math(t: ScriptedContext): + baseline_free = t.engine_stats()["kv_pool_free"] + r = t.start_req(prompt_len=_SWA_WINDOW + 13, max_new_tokens=2) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + assert t.engine_stats()["kv_pool_free"] >= baseline_free, ( + "SWA pool failed to recover after a window-straddling chunked req: " + f"baseline={baseline_free}, " + f"final={t.engine_stats()['kv_pool_free']}" + ) + + def test_swa_chunked_resume_under_swa_pressure(self): + self.server.execute_script(self._script_swa_chunked_resume_under_swa_pressure) + + @staticmethod + def _script_swa_chunked_resume_under_swa_pressure(t: ScriptedContext): + r = t.start_req(prompt_len=_SWA_WINDOW + VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + chunks_at_pressure = r.chunks_done + t.exhaust_kv(leave_pages=1000) + yield from run_until_finished(r, max_steps=2000) + assert r.finished + assert r.chunks_done > chunks_at_pressure, ( + f"chunked prefill stalled under SWA pressure: chunks_done=" + f"{r.chunks_done} did not advance past chunks_at_pressure=" + f"{chunks_at_pressure}" + ) + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_swa_chunked_resume_kv_committed_bound(self): + self.server.execute_script(self._script_swa_chunked_resume_kv_committed_bound) + + @staticmethod + def _script_swa_chunked_resume_kv_committed_bound(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(400): + if r.is_chunking: + assert len(r.req.prefix_indices) <= r.req.kv_committed_len, ( + f"prefix_indices must be bounded by kv_committed_len, " + f"got prefix_indices_len={len(r.req.prefix_indices)}, " + f"kv_committed_len={r.req.kv_committed_len}" + ) + if r.finished: + break + yield + assert r.finished + + +class TestSWAHalfWindowChunk(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_SWA_MODEL, + chunked_prefill_size=_SWA_WINDOW // 2, + mem_fraction_static=0.70, + disable_piecewise_cuda_graph=True, + ) + + def test_swa_prompt_2x_window_half_chunks(self): + self.server.execute_script(self._script_swa_prompt_2x_window_half_chunks) + + @staticmethod + def _script_swa_prompt_2x_window_half_chunks(t: ScriptedContext): + r = t.start_req(prompt_len=2 * _SWA_WINDOW, max_new_tokens=4) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert ( + r.chunks_done >= 4 + ), f"expected >=4 chunks for 2*window / (window/2), got {r.chunks_done}" + assert len(r.req.output_ids) == 4 + + +class TestSWAChunkSizeExceedsWindow(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_SWA_MODEL, + chunked_prefill_size=_SWA_WINDOW * 2, + mem_fraction_static=0.70, + disable_piecewise_cuda_graph=True, + ) + + def test_swa_chunk_size_exceeds_window(self): + self.server.execute_script(self._script_swa_chunk_size_exceeds_window) + + @staticmethod + def _script_swa_chunk_size_exceeds_window(t: ScriptedContext): + r = t.start_req(prompt_len=3 * _SWA_WINDOW, max_new_tokens=2) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 2 + + +class TestSWARadix(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_SWA_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + mem_fraction_static=0.70, + disable_radix_cache=False, + disable_piecewise_cuda_graph=True, + ) + + def test_swa_radix_partial_hit_straddles_window(self): + self.server.execute_script(self._script_swa_radix_partial_hit_straddles_window) + + @staticmethod + def _script_swa_radix_partial_hit_straddles_window(t: ScriptedContext): + r1 = t.start_req(prompt_len=_SWA_WINDOW + DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until_finished(r1, max_steps=800) + assert r1.finished + r2 = t.start_req( + prompt_len=_SWA_WINDOW + DEFAULT_CHUNK_SIZE * 2, max_new_tokens=2 + ) + yield from run_until_finished(r2, max_steps=800) + assert r2.finished + assert ( + r2.req.cached_tokens > 0 + ), f"r2 must hit the radix prefix, got cached_tokens={r2.req.cached_tokens}" + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_invariants.py b/test/manual/chunked_prefill/test_scripted_invariants.py new file mode 100644 index 000000000..451efd9eb --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_invariants.py @@ -0,0 +1,512 @@ +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, + DEFAULT_MAX_STEPS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +class TestInvariantsBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_kv_pages_zero_after_finish(self): + self.server.execute_script(self._script_kv_pages_zero_after_finish) + + @staticmethod + def _script_kv_pages_zero_after_finish(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r) + assert r.kv_pages == 0 + + def test_kv_pages_positive_continuously_mid_chunk(self): + self.server.execute_script( + self._script_kv_pages_positive_continuously_mid_chunk + ) + + @staticmethod + def _script_kv_pages_positive_continuously_mid_chunk(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + observed_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + observed_chunking = True + assert ( + r.kv_pages > 0 + ), f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}" + if r.finished: + break + yield + assert observed_chunking, "test must observe at least one mid-chunk iter" + assert r.finished + + def test_batch_composition_consistent_with_status(self): + self.server.execute_script( + self._script_batch_composition_consistent_with_status + ) + + @staticmethod + def _script_batch_composition_consistent_with_status(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(DEFAULT_MAX_STEPS): + if r.status == "running": + comp = t.batch_composition() + all_rids = ( + comp.get("prefill", []) + + comp.get("decode", []) + + comp.get("chunked", []) + ) + assert ( + r.rid in all_rids + ), f"running but not in batch_composition: {comp}" + if r.finished: + return + yield + raise AssertionError("req never finished") + + def test_active_reqs_listing(self): + self.server.execute_script(self._script_active_reqs_listing) + + @staticmethod + def _script_active_reqs_listing(t: ScriptedContext): + r1 = t.start_req(prompt_len=16, max_new_tokens=4) + r2 = t.start_req(prompt_len=16, max_new_tokens=4) + yield + actives = t.list_active_reqs() + rids = {h.rid for h in actives} + assert r1.rid in rids or r2.rid in rids + yield from run_until_all_finished([r1, r2]) + for _ in range(12): + actives_after = t.list_active_reqs() + if all(h.rid not in (r1.rid, r2.rid) for h in actives_after): + break + yield + actives_after = t.list_active_reqs() + assert all(h.rid not in (r1.rid, r2.rid) for h in actives_after) + + def test_finished_means_chunks_done_stable(self): + self.server.execute_script(self._script_finished_means_chunks_done_stable) + + @staticmethod + def _script_finished_means_chunks_done_stable(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r) + snap = r.chunks_done + for _ in range(10): + yield + assert r.chunks_done == snap + + def test_finished_means_kv_pages_stays_zero(self): + self.server.execute_script(self._script_finished_means_kv_pages_stays_zero) + + @staticmethod + def _script_finished_means_kv_pages_stays_zero(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r) + for _ in range(10): + yield + assert r.kv_pages == 0 + + def test_engine_stats_keys_present(self): + self.server.execute_script(self._script_engine_stats_keys_present) + + @staticmethod + def _script_engine_stats_keys_present(t: ScriptedContext): + stats = t.engine_stats() + assert isinstance(stats, dict) + assert "kv_pool_free" in stats + assert "req_pool_free" in stats + yield + + def test_kv_pool_recovers_to_baseline(self): + self.server.execute_script(self._script_kv_pool_recovers_to_baseline) + + @staticmethod + def _script_kv_pool_recovers_to_baseline(t: ScriptedContext): + before = t.engine_stats()["kv_pool_free"] + reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(8)] + yield from run_until_all_finished(reqs) + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + after = t.engine_stats()["kv_pool_free"] + assert after >= before + + def test_hundred_reqs_no_leak(self): + self.server.execute_script(self._script_hundred_reqs_no_leak) + + @staticmethod + def _script_hundred_reqs_no_leak(t: ScriptedContext): + baseline = t.engine_stats() + reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(100)] + yield from run_until_all_finished(reqs, max_steps=4000) + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats() + assert ( + final["kv_pool_free"] >= baseline["kv_pool_free"] + ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + assert final["req_pool_free"] >= baseline["req_pool_free"] + + def test_long_lived_engine_reps_chunked(self): + self.server.execute_script(self._script_long_lived_engine_reps_chunked) + + @staticmethod + def _script_long_lived_engine_reps_chunked(t: ScriptedContext): + baseline = t.engine_stats() + for _ in range(20): + reqs = [ + t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(5) + ] + yield from run_until_all_finished(reqs, max_steps=2000) + for r in reqs: + assert r.finished + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats() + assert final["kv_pool_free"] >= baseline["kv_pool_free"] + + def test_sustained_long_chunked_load(self): + self.server.execute_script(self._script_sustained_long_chunked_load) + + @staticmethod + def _script_sustained_long_chunked_load(t: ScriptedContext): + expected_chunks_done = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + baseline_kv = t.engine_stats()["kv_pool_free"] + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10 + i + ) + for i in range(30) + ] + yield from run_until_all_finished(reqs, max_steps=DEFAULT_MAX_STEPS * 20) + for r in reqs: + assert r.finished + assert r.chunks_done == expected_chunks_done, ( + f"VERY_LONG_PROMPT_LEN must take exactly {expected_chunks_done} " + f"chunks; got chunks_done={r.chunks_done}" + ) + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final_kv = t.engine_stats()["kv_pool_free"] + assert ( + final_kv >= baseline_kv + ), f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}" + + def test_round_robin_short_and_chunked(self): + self.server.execute_script(self._script_round_robin_short_and_chunked) + + @staticmethod + def _script_round_robin_short_and_chunked(t: ScriptedContext): + baseline = t.engine_stats() + for _ in range(5): + shorts = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(10)] + chunked = [ + t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(1) + ] + yield from run_until_all_finished(shorts + chunked, max_steps=2000) + for r in shorts + chunked: + assert r.finished + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats() + assert final["kv_pool_free"] >= baseline["kv_pool_free"] + + def test_long_decode_then_many_short(self): + self.server.execute_script(self._script_long_decode_then_many_short) + + @staticmethod + def _script_long_decode_then_many_short(t: ScriptedContext): + expected_chunks_done = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + long_max_new_tokens = 256 + long_decode = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=long_max_new_tokens, + ignore_eos=True, + ) + shorts = [t.start_req(prompt_len=8, max_new_tokens=2) for _ in range(50)] + all_reqs = [long_decode] + shorts + yield from run_until_all_finished(all_reqs, max_steps=DEFAULT_MAX_STEPS * 20) + for r in all_reqs: + assert r.finished + assert long_decode.chunks_done == expected_chunks_done, ( + f"long req must chunk across exactly {expected_chunks_done} chunks; " + f"got chunks_done={long_decode.chunks_done}" + ) + assert len(long_decode.req.output_ids) == long_max_new_tokens, ( + f"ignore_eos long req must decode exactly {long_max_new_tokens} " + f"tokens; got len(output_ids)={len(long_decode.req.output_ids)}" + ) + + def test_engine_stats_monotone_after_each_batch(self): + self.server.execute_script(self._script_engine_stats_monotone_after_each_batch) + + @staticmethod + def _script_engine_stats_monotone_after_each_batch(t: ScriptedContext): + last = None + for _ in range(10): + reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(8)] + yield from run_until_all_finished(reqs) + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + cur = t.engine_stats()["kv_pool_free"] + if last is not None: + assert cur >= last - 1, f"KV pool drifted: {last} -> {cur}" + last = cur + + def test_inflight_middle_chunks_caps_at_one(self): + self.server.execute_script(self._script_inflight_middle_chunks_caps_at_one) + + @staticmethod + def _script_inflight_middle_chunks_caps_at_one(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + running_max = 0 + running_max_post_finish = 0 + post_finish_samples = 0 + for _ in range(DEFAULT_MAX_STEPS): + yield + req = r.req + cur = req.inflight_middle_chunks if req is not None else 0 + running_max = max(running_max, cur) + if r.finished: + running_max_post_finish = max(running_max_post_finish, cur) + post_finish_samples += 1 + if post_finish_samples >= 5: + break + assert r.finished, "req never finished" + assert running_max == 1, ( + f"inflight_middle_chunks must reach exactly 1 across the chunked " + f"lifecycle (the cap from revert e875cd36e4); observed max={running_max}" + ) + assert running_max_post_finish == 0, ( + f"inflight_middle_chunks must be reset to 0 after finish; " + f"observed max post-finish={running_max_post_finish}" + ) + + def test_chunks_done_strictly_increases_no_plateaus(self): + self.server.execute_script( + self._script_chunks_done_strictly_increases_no_plateaus + ) + + @staticmethod + def _script_chunks_done_strictly_increases_no_plateaus(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + prev_chunks_done = r.chunks_done + prev_was_chunking = r.is_chunking + for _ in range(DEFAULT_MAX_STEPS): + yield + if r.finished: + return + cur_chunks_done = r.chunks_done + cur_is_chunking = r.is_chunking + if prev_was_chunking and cur_is_chunking: + assert cur_chunks_done > prev_chunks_done, ( + f"chunks_done plateau between consecutive mid-chunk " + f"yields: {prev_chunks_done} -> {cur_chunks_done}" + ) + prev_chunks_done = cur_chunks_done + prev_was_chunking = cur_is_chunking + raise AssertionError("req never finished") + + def test_output_tokens_len_equals_max_new_tokens_chunked(self): + self.server.execute_script( + self._script_output_tokens_len_equals_max_new_tokens_chunked + ) + + @staticmethod + def _script_output_tokens_len_equals_max_new_tokens_chunked(t: ScriptedContext): + n: int = 8 + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=n, + ignore_eos=True, + ) + yield from run_until_finished(r) + assert r.finished + assert ( + r.chunks_done >= 2 + ), f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}" + assert len(r.req.output_ids) == n, ( + f"ignore_eos=True + max_new_tokens={n} must produce exactly " + f"{n} output tokens; got len(output_tokens)={len(r.req.output_ids)}" + ) + + def test_num_input_tokens_equals_prompt_len_for_chunked(self): + self.server.execute_script( + self._script_num_input_tokens_equals_prompt_len_for_chunked + ) + + @staticmethod + def _script_num_input_tokens_equals_prompt_len_for_chunked(t: ScriptedContext): + prompt_len: int = VERY_LONG_PROMPT_LEN + r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2, ( + f"VERY_LONG_PROMPT_LEN should chunk so this invariant is exercised " + f"on a real multi-chunk prefill; got chunks_done={r.chunks_done}" + ) + assert r.remaining_prompt_tokens == 0, ( + f"the whole prompt must be committed after a chunked finish; " + f"remaining_prompt_tokens={r.remaining_prompt_tokens}" + ) + + def test_chunked_in_flight_count_exactly_zero_after_finish(self): + self.server.execute_script( + self._script_chunked_in_flight_count_exactly_zero_after_finish + ) + + @staticmethod + def _script_chunked_in_flight_count_exactly_zero_after_finish(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + assert (1 if t.scheduler.chunked_req is not None else 0) == 1, ( + f"chunked_in_flight_count should be 1 mid-chunk; got " + f"{(1 if t.scheduler.chunked_req is not None else 0)}" + ) + yield from run_until_finished(r) + for _ in range(3): + yield + assert (1 if t.scheduler.chunked_req is not None else 0) == 0, ( + f"chunked_in_flight_count must be 0 after finish; got " + f"{(1 if t.scheduler.chunked_req is not None else 0)}" + ) + + def test_extend_batch_idx_monotonic_invariant(self): + self.server.execute_script(self._script_extend_batch_idx_monotonic_invariant) + + @staticmethod + def _script_extend_batch_idx_monotonic_invariant(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=64) + observed_regression: bool = False + + yield from run_until( + r, + lambda h: ( + h.req is not None + and not h.req.is_retracted + and h.req.extend_batch_idx > 0 + and h.remaining_prompt_tokens == 0 + and not h.finished + ), + ) + pre_retract_idx: int = r.req.extend_batch_idx + assert not r.req.is_retracted + + t.pause_generation(mode="retract") + retracted = t.find_req_by_rid(r.rid) + assert retracted is not None, "retracted req must stay live in the queue" + if retracted.extend_batch_idx < pre_retract_idx: + observed_regression = True + assert retracted.is_retracted, ( + f"extend_batch_idx regressed without retract flag: " + f"{pre_retract_idx} -> {retracted.extend_batch_idx}" + ) + t.continue_generation() + + prev_extend_batch_idx: int = -1 + regressions: int = 0 + for _ in range(DEFAULT_MAX_STEPS): + req = t.find_req_by_rid(r.rid) + if req is not None: + cur_extend_batch_idx = req.extend_batch_idx + if ( + prev_extend_batch_idx >= 0 + and cur_extend_batch_idx < prev_extend_batch_idx + ): + regressions += 1 + observed_regression = True + assert regressions == 1, ( + f"extend_batch_idx regressed more than once for a single " + f"retract episode: " + f"{prev_extend_batch_idx} -> {cur_extend_batch_idx}" + ) + prev_extend_batch_idx = cur_extend_batch_idx + if r.finished: + break + yield + assert r.finished, "req never finished" + assert observed_regression, ( + "retract must reset extend_batch_idx, producing the regression this " + "test guards" + ) + + def test_inflight_decrement_only_on_final_invariant(self): + self.server.execute_script( + self._script_inflight_decrement_only_on_final_invariant + ) + + @staticmethod + def _script_inflight_decrement_only_on_final_invariant(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + prev_inflight: int = 0 + prev_was_chunked_slot: bool = False + prev_finished: bool = False + observed_decrement: bool = False + for _ in range(DEFAULT_MAX_STEPS): + s = t.scheduler + req = t.find_req_by_rid(r.rid) + cur_inflight = req.inflight_middle_chunks if req is not None else 0 + cur_is_chunked_slot = ( + s.chunked_req is not None and s.chunked_req.rid == r.rid + ) + cur_finished = req.finished() if req is not None else True + if cur_inflight < prev_inflight: + observed_decrement = True + slot_just_released = prev_was_chunked_slot and not cur_is_chunked_slot + finish_just_happened = (not prev_finished) and cur_finished + assert slot_just_released or finish_just_happened, ( + f"inflight_middle_chunks decreased ({prev_inflight} -> " + f"{cur_inflight}) without chunked slot release or req " + f"finish; prev_was_chunked_slot={prev_was_chunked_slot}, " + f"cur_is_chunked_slot={cur_is_chunked_slot}, " + f"prev_finished={prev_finished}, cur_finished={cur_finished}" + ) + prev_inflight = cur_inflight + prev_was_chunked_slot = cur_is_chunked_slot + prev_finished = cur_finished + if r.finished: + break + yield + assert r.finished + assert observed_decrement, ( + "test must observe at least one inflight_middle_chunks decrement " + "across the chunked lifecycle" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_kv_pressure.py b/test/manual/chunked_prefill/test_scripted_kv_pressure.py new file mode 100644 index 000000000..751f1098a --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_kv_pressure.py @@ -0,0 +1,550 @@ +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 ( + BALLAST_MAX_NEW_TOKENS, + DEFAULT_CHUNK_SIZE, + DEFAULT_MAX_STEPS, + SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + SMALL_KV_POOL_BALLAST_PROMPT_LEN, + SMALL_KV_POOL_MAX_TOTAL_TOKENS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +class TestKVPressureBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_lock_refs_tight_concurrent_prefix(self): + self.server.execute_script(self._script_lock_refs_tight_concurrent_prefix) + + @staticmethod + def _script_lock_refs_tight_concurrent_prefix(t: ScriptedContext): + warm_token = 7 + warm_len = DEFAULT_CHUNK_SIZE + r_warm = t.start_req( + prompt_len=warm_len, max_new_tokens=1, prompt_token=warm_token + ) + yield from run_until_finished(r_warm) + r_warm2 = t.start_req(prompt_len=warm_len, max_new_tokens=1, prompt_token=8) + yield from run_until_finished(r_warm2) + assert r_warm.finished + for _ in range(12): + if r_warm.lock_refs == 0: + break + yield + assert r_warm.lock_refs == 0 + + baseline_lock_refs = t.get_all_node_lock_refs() + + t.exhaust_lock_refs(leave_refs=1) + yield + pinned_lock_refs = t.get_all_node_lock_refs() + assert any( + pinned_lock_refs.get(node_id, 0) > baseline_lock_refs.get(node_id, 0) + for node_id in pinned_lock_refs + ), "exhaust_lock_refs(leave_refs=1) must pin at least one warm-prefix node" + + r_long = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + prompt_token=warm_token, + ) + yield from run_until(r_long, lambda h: h.is_chunking) + yield from run_until_finished(r_long, max_steps=2000) + assert r_long.finished + assert r_long.chunks_done >= 2, ( + f"long req must really chunk under pinned cache; got chunks_done=" + f"{r_long.chunks_done}" + ) + assert ( + r_long.lock_refs == 0 + ), f"req {r_long.rid} leaked {r_long.lock_refs} lock_refs after finish" + + t._release_exhausted_pools() + final_lock_refs = t.get_all_node_lock_refs() + for node_id, baseline in baseline_lock_refs.items(): + assert final_lock_refs.get(node_id, 0) == baseline, ( + f"node {node_id} lock_ref leaked: baseline={baseline}, " + f"final={final_lock_refs.get(node_id, 0)}" + ) + + def test_kv_pressure_with_retract_resume(self): + self.server.execute_script(self._script_kv_pressure_with_retract_resume) + + @staticmethod + def _script_kv_pressure_with_retract_resume(t: ScriptedContext): + baseline = t.engine_stats()["kv_pool_free"] + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=60 + ) + yield from run_until(r, lambda h: h.is_chunking) + chunks_before_retract = r.chunks_done + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until( + r, + lambda h: h.chunks_done > chunks_before_retract, + max_steps=2000, + ) + yield from run_until_finished(r, max_steps=2000) + assert r.finished + assert r.kv_pages == 0 + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert final >= baseline, ( + f"KV pool failed to recover after retract+resume: " + f"baseline={baseline}, final={final}" + ) + + def test_chunked_batch_recovers_pools_to_steady_state(self): + self.server.execute_script( + self._script_chunked_batch_recovers_pools_to_steady_state + ) + + @staticmethod + def _script_chunked_batch_recovers_pools_to_steady_state(t: ScriptedContext): + before = t.engine_stats() + reqs = [ + t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE + 1, + max_new_tokens=2, + prompt_token=500 + i, + ) + for i in range(50) + ] + yield from run_until_all_finished(reqs, max_steps=2000) + for r in reqs: + assert r.finished + assert r.kv_pages == 0, f"req {r.rid} kept {r.kv_pages} pages after finish" + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + after = t.engine_stats() + assert after["kv_pool_free"] >= before["kv_pool_free"], ( + f"50 chunked reqs leaked KV: baseline={before['kv_pool_free']}, " + f"final={after['kv_pool_free']}" + ) + assert after["req_pool_free"] >= before["req_pool_free"], ( + f"50 chunked reqs leaked req-pool rows: " + f"baseline={before['req_pool_free']}, final={after['req_pool_free']}" + ) + + def test_chunked_retract_at_chunk_first_mid_last(self): + self.server.execute_script(self._script_chunked_retract_at_chunk_first_mid_last) + + @staticmethod + def _script_chunked_retract_at_chunk_first_mid_last(t: ScriptedContext): + expected_chunks = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + mid_chunk = expected_chunks // 2 + last_minus_one = max(1, expected_chunks - 1) + + r_first = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=31 + ) + yield from run_until(r_first, lambda h: h.is_chunking) + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r_first, max_steps=2000) + assert r_first.finished + assert r_first.kv_pages == 0 + + r_mid = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=32 + ) + yield from run_until(r_mid, lambda h: h.chunks_done >= mid_chunk) + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r_mid, max_steps=2000) + assert r_mid.finished + assert r_mid.kv_pages == 0 + + r_last = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=33 + ) + yield from run_until(r_last, lambda h: h.chunks_done >= last_minus_one) + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r_last, max_steps=2000) + assert r_last.finished + assert r_last.kv_pages == 0 + + def test_flush_cache_during_chunked_in_flight(self): + self.server.execute_script(self._script_flush_cache_during_chunked_in_flight) + + @staticmethod + def _script_flush_cache_during_chunked_in_flight(t: ScriptedContext): + r_warm = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=1, prompt_token=41 + ) + yield from run_until_finished(r_warm) + assert r_warm.finished + + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=42 + ) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + t.flush_cache() + yield from run_until_finished(r, max_steps=2000) + assert r.finished + assert r.kv_pages == 0 + + def test_chunked_oscillation_three_force_retracts(self): + self.server.execute_script( + self._script_chunked_oscillation_three_force_retracts + ) + + @staticmethod + def _script_chunked_oscillation_three_force_retracts(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=70 + ) + yield from run_until(r, lambda h: h.is_chunking) + chunks_at_first = r.chunks_done + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until( + r, + lambda h: h.is_chunking and h.chunks_done >= chunks_at_first, + max_steps=800, + ) + chunks_after_first_resume = r.chunks_done + assert chunks_after_first_resume >= chunks_at_first, ( + f"chunks_done regressed across retract: " + f"before={chunks_at_first}, after={chunks_after_first_resume}" + ) + + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until( + r, + lambda h: h.is_chunking and h.chunks_done >= chunks_after_first_resume, + max_steps=800, + ) + chunks_after_second_resume = r.chunks_done + assert chunks_after_second_resume >= chunks_after_first_resume + + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r, max_steps=2000) + assert r.finished + assert r.chunks_done >= chunks_after_second_resume + + +class TestKVPressureSmallRowPool(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + max_running_requests=8, + ) + + def test_row_pool_tight_admits_after_release(self): + self.server.execute_script(self._script_row_pool_tight_admits_after_release) + + @staticmethod + def _script_row_pool_tight_admits_after_release(t: ScriptedContext): + baseline_rows_used = ( + t.scheduler.req_to_token_pool.size + - t.scheduler.req_to_token_pool.available_size() + ) + row_pool_size = t.scheduler.req_to_token_pool.size + ballast = [ + t.start_req( + prompt_len=1, max_new_tokens=BALLAST_MAX_NEW_TOKENS, ignore_eos=True + ) + for _ in range(row_pool_size) + ] + for _ in range(DEFAULT_MAX_STEPS): + if t.scheduler.req_to_token_pool.available_size() == 0: + break + yield + assert t.scheduler.req_to_token_pool.available_size() == 0, ( + f"ballast must hold every row; " + f"available={t.scheduler.req_to_token_pool.available_size()}" + ) + + reqs = [ + t.start_req(prompt_len=8, max_new_tokens=1, prompt_token=50 + i) + for i in range(5) + ] + for _ in range(6): + yield + for r in reqs: + assert r.status == "waiting", ( + f"fresh req must be unschedulable under a full row pool; " + f"rid={r.rid}, status={r.status}" + ) + + for b in ballast: + t.abort(b) + yield from run_until_all_finished(reqs, max_steps=2000) + for r in reqs: + assert r.finished, f"req {r.rid} did not finish after release" + assert r.kv_pages == 0, ( + f"row-pool pressure must not leave KV held: rid={r.rid}, " + f"kv_pages={r.kv_pages}" + ) + + for _ in range(40): + if t.is_fully_idle: + break + yield + final_rows_used = ( + t.scheduler.req_to_token_pool.size + - t.scheduler.req_to_token_pool.available_size() + ) + assert final_rows_used <= baseline_rows_used, ( + f"row pool leak after admit-after-release: baseline used=" + f"{baseline_rows_used}, final used={final_rows_used}" + ) + + +class TestKVPressureSmallPool(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + max_total_tokens=SMALL_KV_POOL_MAX_TOTAL_TOKENS, + ) + + @staticmethod + def _start_ballast(t: ScriptedContext, *, prompt_token: int): + return t.start_req( + prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN, + max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + prompt_token=prompt_token, + ) + + def test_chunked_completes_when_ballast_retracted(self): + self.server.execute_script( + self._script_chunked_completes_when_ballast_retracted + ) + + @staticmethod + def _script_chunked_completes_when_ballast_retracted(t: ScriptedContext): + baseline = t.engine_stats()["kv_pool_free"] + r_chunk = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=600 + ) + yield from run_until(r_chunk, lambda h: h.is_chunking and h.chunks_done >= 1) + + ballast = TestKVPressureSmallPool._start_ballast(t, prompt_token=601) + + ballast_retracted = False + for _ in range(2000): + if ballast.status == "waiting": + ballast_retracted = True + if r_chunk.finished: + break + yield + assert r_chunk.finished, ( + f"chunked req must complete once the engine retracts the ballast under " + f"real KV pressure; status={r_chunk.status}, kv_pages={r_chunk.kv_pages}" + ) + assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}" + assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}" + + ballast_resolved = ( + ballast_retracted + or ballast.finished + or ballast.status in ("waiting", "finished", "unknown") + ) + assert ballast_resolved, ( + f"ballast must be retracted/aborted under pressure; " + f"status={ballast.status}, retracted={ballast_retracted}" + ) + + t.abort(ballast) + for _ in range(200): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert final >= baseline, ( + f"KV pool not recovered after ballast-retract pressure: " + f"baseline={baseline}, final={final}" + ) + + def test_chunked_completes_under_ballast_then_aborts_chunked(self): + self.server.execute_script( + self._script_chunked_completes_under_ballast_then_aborts_chunked + ) + + @staticmethod + def _script_chunked_completes_under_ballast_then_aborts_chunked( + t: ScriptedContext, + ): + baseline = t.engine_stats()["kv_pool_free"] + r_chunk = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=610 + ) + yield from run_until(r_chunk, lambda h: h.is_chunking and h.chunks_done >= 1) + + ballast = TestKVPressureSmallPool._start_ballast(t, prompt_token=611) + for _ in range(6): + yield + + t.abort(r_chunk) + for _ in range(12): + if ( + r_chunk.kv_pages == 0 + and r_chunk.lock_refs == 0 + and (r_chunk.req is None or r_chunk.req.req_pool_idx is None) + ): + break + yield + assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}" + assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}" + assert r_chunk.req is None or r_chunk.req.req_pool_idx is None + + t.abort(ballast) + for _ in range(200): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert final >= baseline, ( + f"KV pool not recovered after abort-under-pressure: " + f"baseline={baseline}, final={final}" + ) + + def test_kv_recovery_after_full(self): + self.server.execute_script(self._script_kv_recovery_after_full) + + @staticmethod + def _script_kv_recovery_after_full(t: ScriptedContext): + baseline = t.engine_stats()["kv_pool_free"] + b1 = t.start_req( + prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN, + max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + prompt_token=620, + ) + yield from run_until(b1, lambda h: h.status == "running") + + big = t.start_req(prompt_len=2048, max_new_tokens=2, prompt_token=621) + yield from run_until(big, lambda h: h.is_chunking and h.chunks_done >= 2) + + r = t.start_req(prompt_len=16, max_new_tokens=2) + yield + assert r.status == "waiting", ( + f"16-token req must be unschedulable while the ballast and the " + f"mid-chunk req own the pool; status={r.status}" + ) + + t.abort(b1) + yield from run_until(big, lambda h: h.finished, max_steps=3000) + yield from run_until(r, lambda h: h.finished, max_steps=3000) + assert r.finished + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert final >= baseline, ( + f"pool must recover to baseline after release: " + f"baseline={baseline}, final={final}" + ) + + +class TestKVPressurePriority(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + ) + + def test_priority_preempt_multiple_chunked(self): + self.server.execute_script(self._script_priority_preempt_multiple_chunked) + + @staticmethod + def _script_priority_preempt_multiple_chunked(t: ScriptedContext): + baseline = t.engine_stats()["kv_pool_free"] + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=11 + ) + yield from run_until(r1, lambda h: h.is_chunking) + + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + priority=10, + prompt_token=12, + ) + done = {r1.rid: False, r2.rid: False} + for _ in range(DEFAULT_MAX_STEPS * 4): + assert not (r1.is_chunking and r2.is_chunking), ( + f"two reqs cannot share the chunked slot; " + f"r1.is_chunking={r1.is_chunking}, r2.is_chunking={r2.is_chunking}" + ) + done[r1.rid] = done[r1.rid] or r1.finished + done[r2.rid] = done[r2.rid] or r2.finished + if all(done.values()): + break + yield + assert done[r1.rid] and done[r2.rid] + assert r1.kv_pages == 0 + assert r2.kv_pages == 0 + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert final >= baseline, ( + f"KV pool not fully released after preemption: " + f"baseline={baseline}, final={final}" + ) + + +class TestKVPressurePageSize(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + page_size=16, + ) + + def test_strict_mem_check_handles_chunked_tail(self): + self.server.execute_script(self._script_strict_mem_check_handles_chunked_tail) + + @staticmethod + def _script_strict_mem_check_handles_chunked_tail(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 17, max_new_tokens=2) + 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=2000) + assert r.finished + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_lifecycle.py b/test/manual/chunked_prefill/test_scripted_lifecycle.py new file mode 100644 index 000000000..bffc2784c --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_lifecycle.py @@ -0,0 +1,440 @@ +import unittest + +from sglang.srt.managers.schedule_batch import FINISH_ABORT +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, + DEFAULT_MAX_STEPS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_finished, +) + + +def _drain_until_released(t, *handles): + for _ in range(12): + if all( + h.kv_pages == 0 + and h.lock_refs == 0 + and (h.req is None or h.req.req_pool_idx is None) + for h in handles + ): + return + yield + + +class TestLifecycleBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_small_prompt_short_decode(self): + self.server.execute_script(self._script_small_prompt_short_decode) + + @staticmethod + def _script_small_prompt_short_decode(t: ScriptedContext): + r = t.start_req(prompt_len=8, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + assert len(r.req.output_ids) == 2 + + def test_medium_prompt_medium_decode(self): + self.server.execute_script(self._script_medium_prompt_medium_decode) + + @staticmethod + def _script_medium_prompt_medium_decode(t: ScriptedContext): + r = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=16, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + assert len(r.req.output_ids) == 16 + + def test_long_prompt_short_decode(self): + self.server.execute_script(self._script_long_prompt_short_decode) + + @staticmethod + def _script_long_prompt_short_decode(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 8 + assert len(r.req.output_ids) == 2 + + def test_long_prompt_long_decode(self): + self.server.execute_script(self._script_long_prompt_long_decode) + + @staticmethod + def _script_long_prompt_long_decode(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=64, ignore_eos=True + ) + yield from run_until(r, lambda h: h.finished, max_steps=1000) + assert r.finished + assert r.chunks_done == 8 + assert len(r.req.output_ids) == 64 + + def test_tiny_prompt_long_decode(self): + self.server.execute_script(self._script_tiny_prompt_long_decode) + + @staticmethod + def _script_tiny_prompt_long_decode(t: ScriptedContext): + r = t.start_req(prompt_len=1, max_new_tokens=64, ignore_eos=True) + yield from run_until(r, lambda h: h.finished, max_steps=500) + assert r.finished + assert r.chunks_done == 0 + assert len(r.req.output_ids) == 64 + + def test_chunk_size_minus_one_prompt(self): + self.server.execute_script(self._script_chunk_size_minus_one_prompt) + + @staticmethod + def _script_chunk_size_minus_one_prompt(t: ScriptedContext): + r = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE - 1, max_new_tokens=4, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + assert len(r.req.output_ids) == 4 + + def test_chunk_size_plus_two_prompt(self): + self.server.execute_script(self._script_chunk_size_plus_two_prompt) + + @staticmethod + def _script_chunk_size_plus_two_prompt(t: ScriptedContext): + r = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE + 2, max_new_tokens=4, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 2 + assert len(r.req.output_ids) == 4 + + def test_just_over_2x_chunk_size(self): + self.server.execute_script(self._script_just_over_2x_chunk_size) + + @staticmethod + def _script_just_over_2x_chunk_size(t: ScriptedContext): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE + 1, max_new_tokens=4, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 3 + assert len(r.req.output_ids) == 4 + + def test_five_x_chunk_size(self): + self.server.execute_script(self._script_five_x_chunk_size) + + @staticmethod + def _script_five_x_chunk_size(t: ScriptedContext): + r = t.start_req( + prompt_len=5 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 5 + assert len(r.req.output_ids) == 4 + + def test_ten_x_chunk_size(self): + self.server.execute_script(self._script_ten_x_chunk_size) + + @staticmethod + def _script_ten_x_chunk_size(t: ScriptedContext): + r = t.start_req( + prompt_len=10 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, ignore_eos=True + ) + yield from run_until(r, lambda h: h.finished, max_steps=1000) + assert r.finished + assert r.chunks_done == 10 + assert len(r.req.output_ids) == 2 + + def test_status_progression_happy_path(self): + self.server.execute_script(self._script_status_progression_happy_path) + + @staticmethod + def _script_status_progression_happy_path(t: ScriptedContext): + r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + seen = [] + for _ in range(DEFAULT_MAX_STEPS): + seen.append(r.status) + if r.finished: + break + yield + else: + raise AssertionError("req did not finish within DEFAULT_MAX_STEPS") + assert "running" in seen, f"never observed running status; seen={seen}" + assert seen[-1] == "finished", f"final status must be finished; seen={seen}" + finished_idx = seen.index("finished") + assert all( + s in ("finished",) for s in seen[finished_idx:] + ), f"status regressed after finish; seen={seen}" + + def test_long_prompt_only_one_decode(self): + self.server.execute_script(self._script_long_prompt_only_one_decode) + + @staticmethod + def _script_long_prompt_only_one_decode(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=1, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 8 + assert len(r.req.output_ids) == 1 + + def test_kv_pages_consistent_during_run(self): + self.server.execute_script(self._script_kv_pages_consistent_during_run) + + @staticmethod + def _script_kv_pages_consistent_during_run(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, ignore_eos=True + ) + saw_positive = False + for _ in range(DEFAULT_MAX_STEPS): + pages = r.kv_pages + if pages > 0: + saw_positive = True + elif saw_positive and not r.finished: + raise AssertionError( + f"kv_pages collapsed to 0 mid-run before finish; " + f"saw_positive={saw_positive}, status={r.status!r}" + ) + if r.finished: + break + yield + else: + raise AssertionError("req did not finish within DEFAULT_MAX_STEPS") + assert saw_positive + assert r.kv_pages == 0 + assert len(r.req.output_ids) == 4 + + def test_row_idx_recycled_after_finish(self): + self.server.execute_script(self._script_row_idx_recycled_after_finish) + + @staticmethod + def _script_row_idx_recycled_after_finish(t: ScriptedContext): + r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert r.req.req_pool_idx is None + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_two_seq_clean_handoff(self): + self.server.execute_script(self._script_two_seq_clean_handoff) + + @staticmethod + def _script_two_seq_clean_handoff(t: ScriptedContext): + r1 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r1) + yield from _drain_until_released(t, r1) + assert r1.req.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0 + r1_output_len = len(r1.req.output_ids) + + r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r2) + yield from _drain_until_released(t, r2) + assert r1.finished and r2.finished + assert r1_output_len == 2 and len(r2.req.output_ids) == 2 + assert r2.req.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0 + + def test_five_seq_clean(self): + self.server.execute_script(self._script_five_seq_clean) + + @staticmethod + def _script_five_seq_clean(t: ScriptedContext): + reqs = [] + for _ in range(5): + r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + assert r.req.req_pool_idx is None + assert r.kv_pages == 0 + assert r.lock_refs == 0 + reqs.append(r) + for r in reqs: + assert r.finished + + def test_radix_partial_seq(self): + self.server.execute_script(self._script_radix_partial_seq) + + @staticmethod + def _script_radix_partial_seq(t: ScriptedContext): + r1 = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1, ignore_eos=True + ) + yield from run_until_finished(r1) + r2 = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE + 32, max_new_tokens=2, ignore_eos=True + ) + yield from run_until_finished(r2) + assert r1.finished and r2.finished + assert r2.chunks_done == 0 + assert r2.req.cached_tokens > 0, ( + f"r2 must hit r1's radix prefix; got cached_tokens=" + f"{r2.req.cached_tokens}" + ) + assert len(r2.req.output_ids) == 2 + + def test_alternating_short_long_seq(self): + self.server.execute_script(self._script_alternating_short_long_seq) + + @staticmethod + def _script_alternating_short_long_seq(t: ScriptedContext): + for i in range(6): + prompt = 8 if i % 2 == 0 else VERY_LONG_PROMPT_LEN + r = t.start_req( + prompt_len=prompt, + max_new_tokens=2, + ignore_eos=True, + prompt_token=10 + i, + ) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + if prompt == VERY_LONG_PROMPT_LEN: + assert r.chunks_done == 8 + else: + assert r.chunks_done == 0 + + def test_seq_with_growing_prompt(self): + self.server.execute_script(self._script_seq_with_growing_prompt) + + @staticmethod + def _script_seq_with_growing_prompt(t: ScriptedContext): + for idx, L in enumerate([8, 32, 128, 512, 1024]): + r = t.start_req( + prompt_len=L, max_new_tokens=1, ignore_eos=True, prompt_token=10 + idx + ) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 1 + yield from _drain_until_released(t, r) + assert r.req is None or r.req.req_pool_idx is None + assert r.kv_pages == 0 and r.lock_refs == 0 + if L > DEFAULT_CHUNK_SIZE: + assert ( + r.chunks_done == (L + DEFAULT_CHUNK_SIZE - 1) // DEFAULT_CHUNK_SIZE + ) + else: + assert r.chunks_done == 0 + + def test_seq_with_shrinking_prompt(self): + self.server.execute_script(self._script_seq_with_shrinking_prompt) + + @staticmethod + def _script_seq_with_shrinking_prompt(t: ScriptedContext): + for idx, L in enumerate([1024, 512, 128, 32, 8]): + r = t.start_req( + prompt_len=L, max_new_tokens=1, ignore_eos=True, prompt_token=10 + idx + ) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 1 + yield from _drain_until_released(t, r) + assert r.req is None or r.req.req_pool_idx is None + assert r.kv_pages == 0 and r.lock_refs == 0 + if L > DEFAULT_CHUNK_SIZE: + assert ( + r.chunks_done == (L + DEFAULT_CHUNK_SIZE - 1) // DEFAULT_CHUNK_SIZE + ) + else: + assert r.chunks_done == 0 + + def test_seq_with_idle_yields_between(self): + self.server.execute_script(self._script_seq_with_idle_yields_between) + + @staticmethod + def _script_seq_with_idle_yields_between(t: ScriptedContext): + for _ in range(4): + r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + for _ in range(20): + yield + + def test_chunked_then_short_seq(self): + self.server.execute_script(self._script_chunked_then_short_seq) + + @staticmethod + def _script_chunked_then_short_seq(t: ScriptedContext): + seq = [VERY_LONG_PROMPT_LEN, 8, VERY_LONG_PROMPT_LEN, 8] + for idx, L in enumerate(seq): + r = t.start_req( + prompt_len=L, max_new_tokens=2, ignore_eos=True, prompt_token=10 + idx + ) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + if L == VERY_LONG_PROMPT_LEN: + assert r.chunks_done == 8 + else: + assert r.chunks_done == 0 + + def test_seq_engine_stats_stable(self): + self.server.execute_script(self._script_seq_engine_stats_stable) + + @staticmethod + def _script_seq_engine_stats_stable(t: ScriptedContext): + baseline = t.engine_stats()["kv_pool_free"] + for _ in range(5): + r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0 + for _ in range(5): + yield + t.flush_cache() + yield + final = t.engine_stats()["kv_pool_free"] + assert ( + final >= baseline - 1 + ), f"KV pool drift: baseline={baseline}, final={final}" + + def test_abort_all_during_chunked(self): + self.server.execute_script(self._script_abort_all_during_chunked) + + @staticmethod + def _script_abort_all_during_chunked(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.abort_all() + + def _error_message(h): + if h.req is None: + return None + return ( + h.req.finished_reason.message + if isinstance(h.req.finished_reason, FINISH_ABORT) + else None + ) + + for _ in range(DEFAULT_MAX_STEPS): + if r.finished or _error_message(r) is not None: + break + yield + else: + raise AssertionError( + "chunked req did not terminate after abort_all within DEFAULT_MAX_STEPS" + ) + assert r.finished or _error_message(r) is not None + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + assert r.lock_refs == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_lora.py b/test/manual/chunked_prefill/test_scripted_lora.py new file mode 100644 index 000000000..612683937 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_lora.py @@ -0,0 +1,240 @@ +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_all_finished, + run_until_finished, +) + +_LORA_BASE_MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_LORA_ADAPTER = "codelion/Llama-3.2-1B-Instruct-tool-calling-lora" +_LORA_ADAPTER_B = "nicoboss/Llama-3.2-1B-Instruct-Uncensored-Lora" + + +class TestLoRASingleAdapter(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER], + ) + + def test_naive_lora_chunked(self): + self.server.execute_script(self._script_naive_lora_chunked) + + @staticmethod + def _script_naive_lora_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + lora_path=_LORA_ADAPTER, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + assert len(r.req.output_ids) == 4 + + def test_lora_logprob_chunked_pass_idx(self): + self.server.execute_script(self._script_lora_logprob_chunked_pass_idx) + + @staticmethod + def _script_lora_logprob_chunked_pass_idx(t: ScriptedContext): + prompt_len: int = VERY_LONG_PROMPT_LEN + r = t.start_req( + prompt_len=prompt_len, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + return_logprob=True, + logprob_start_len=0, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == prompt_len // DEFAULT_CHUNK_SIZE + assert len(r.req.logprob.input_token_logprobs_val) == prompt_len + + +class TestLoRADrainerBypass(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + max_loras_per_batch=1, + ) + + def test_lora_drainer_does_not_block_chunked_resume(self): + self.server.execute_script( + self._script_lora_drainer_does_not_block_chunked_resume + ) + + @staticmethod + def _script_lora_drainer_does_not_block_chunked_resume(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + yield from run_until(r_a, lambda h: h.is_chunking and h.chunks_done >= 1) + chunks_before = r_a.chunks_done + + r_b = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE // 2, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + + for _ in range(200): + if r_a.chunks_done > chunks_before: + break + yield + else: + raise AssertionError( + f"chunked-resume r_a starved by LoRA drainer; " + f"chunks_done stuck at {chunks_before}" + ) + + yield from run_until_all_finished(handles=[r_a, r_b]) + assert r_a.finished and r_b.finished + + +class TestLoRAAdapterSwitch(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + max_loras_per_batch=2, + ) + + def test_lora_adapter_switch_mid_chunk(self): + self.server.execute_script(self._script_lora_adapter_switch_mid_chunk) + + @staticmethod + def _script_lora_adapter_switch_mid_chunk(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + r_b = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + yield from run_until_all_finished(handles=[r_a, r_b]) + assert r_a.finished and r_b.finished + expected_chunks = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + assert r_a.chunks_done == expected_chunks + assert r_b.chunks_done == expected_chunks + + +class TestLoRAAllDistinctAdapters(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + max_loras_per_batch=2, + max_loaded_loras=4, + ) + + def test_lora_all_distinct_adapters_chunked(self): + self.server.execute_script(self._script_lora_all_distinct_adapters_chunked) + + @staticmethod + def _script_lora_all_distinct_adapters_chunked(t: ScriptedContext): + adapters = [_LORA_ADAPTER, _LORA_ADAPTER_B, _LORA_ADAPTER, _LORA_ADAPTER_B] + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=adapter, + ) + for adapter in adapters + ] + yield from run_until_all_finished(handles=reqs, max_steps=2000) + assert all(r.finished for r in reqs) + for r in reqs: + assert r.kv_pages == 0 + assert r.lock_refs == 0 + expected_first_chunks = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + assert reqs[0].chunks_done == expected_first_chunks + assert reqs[1].chunks_done == expected_first_chunks + assert reqs[2].chunks_done < expected_first_chunks + assert reqs[3].chunks_done < expected_first_chunks + + +class TestLoRAAdapterEviction(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + max_loras_per_batch=1, + max_loaded_loras=2, + ) + + def test_lora_adapter_eviction_between_chunks(self): + self.server.execute_script(self._script_lora_adapter_eviction_between_chunks) + + @staticmethod + def _script_lora_adapter_eviction_between_chunks(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + yield from run_until(r_a, lambda h: h.is_chunking and h.chunks_done >= 1) + + r_b = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE // 2, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + yield from run_until_all_finished(handles=[r_a, r_b], max_steps=800) + assert r_a.finished and r_b.finished + + def test_lora_chunked_abort_during_eviction(self): + self.server.execute_script(self._script_lora_chunked_abort_during_eviction) + + @staticmethod + def _script_lora_chunked_abort_during_eviction(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + yield from run_until(r_a, lambda h: h.is_chunking and h.chunks_done >= 1) + + r_b = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE // 2, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + yield + + t.abort(r_a) + for _ in range(12): + if r_a.kv_pages == 0 and (r_a.req is None or r_a.req.req_pool_idx is None): + break + yield + + assert r_a.status in ("finished", "unknown") + if r_a.req is not None: + assert r_a.kv_pages == 0 + assert r_a.lock_refs == 0 + yield from run_until_finished(r_b) + assert r_b.finished + assert r_b.kv_pages == 0 + assert r_b.lock_refs == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_lora_overlap.py b/test/manual/chunked_prefill/test_scripted_lora_overlap.py new file mode 100644 index 000000000..dc78cc589 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_lora_overlap.py @@ -0,0 +1,191 @@ +import unittest + +from sglang.srt.lora.lora_registry import LoRARef +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_all_finished, + run_until_finished, +) + +_LORA_BASE_MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_LORA_ADAPTER = "codelion/Llama-3.2-1B-Instruct-tool-calling-lora" +_LORA_ADAPTER_B = "nicoboss/Llama-3.2-1B-Instruct-Uncensored-Lora" + + +def _expected_lora_id(adapter_path: str) -> str: + return LoRARef.deterministic_id(adapter_path, adapter_path) + + +class TestLoRAOverlapSingleAdapter(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER], + enable_lora_overlap_loading=True, + max_loras_per_batch=1, + max_loaded_loras=1, + ) + + def test_naive_lora_overlap_chunked(self): + self.server.execute_script(self._script_naive_lora_overlap_chunked) + + @staticmethod + def _script_naive_lora_overlap_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + lora_path=_LORA_ADAPTER, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + + +class TestLoRAOverlapH2dDuringChunk(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + enable_lora_overlap_loading=True, + max_loras_per_batch=2, + max_loaded_loras=2, + ) + + def test_lora_overlap_h2d_during_chunk_admit(self): + self.server.execute_script(self._script_lora_overlap_h2d_during_chunk_admit) + + @staticmethod + def _script_lora_overlap_h2d_during_chunk_admit(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + yield from run_until(r_a, lambda h: h.is_chunking and h.chunks_done >= 1) + + r_b = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + yield from run_until_all_finished(handles=[r_a, r_b], max_steps=1200) + assert r_a.finished and r_b.finished + expected_chunks = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + assert r_a.chunks_done == expected_chunks + assert r_b.chunks_done == expected_chunks + assert r_a.req.lora_id == _expected_lora_id(_LORA_ADAPTER) + assert r_b.req.lora_id == _expected_lora_id(_LORA_ADAPTER_B) + + +class TestLoRAOverlapAbortDuringH2d(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + enable_lora_overlap_loading=True, + max_loras_per_batch=2, + max_loaded_loras=2, + ) + + def test_lora_overlap_abort_during_h2d(self): + self.server.execute_script(self._script_lora_overlap_abort_during_h2d) + + @staticmethod + def _script_lora_overlap_abort_during_h2d(t: ScriptedContext): + r_a = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER, + ) + yield from run_until(r_a, lambda h: h.is_chunking and h.chunks_done >= 1) + + r_b = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=_LORA_ADAPTER_B, + ) + yield + b_loader = t.scheduler.lora_overlap_loader + b_lora_id = _expected_lora_id(_LORA_ADAPTER_B) + assert ( + b_lora_id in b_loader.lora_to_overlap_load_event + or b_lora_id in b_loader.lora_manager.memory_pool.uid_to_buffer_id + ), "adapter B never entered H2D; the abort would not exercise the H2D path" + + t.abort(r_b) + yield + + assert r_b.chunks_done == 0 + assert r_b.req is None + assert r_b.status in ("finished", "unknown") + yield from run_until_finished(r_a, max_steps=800) + assert r_a.finished + assert r_a.req.lora_id == _expected_lora_id(_LORA_ADAPTER) + + +class TestLoRAOverlapAdapterRotation(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + model_path=_LORA_BASE_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_lora=True, + lora_paths=[_LORA_ADAPTER, _LORA_ADAPTER_B], + enable_lora_overlap_loading=True, + max_loras_per_batch=1, + max_loaded_loras=2, + ) + + def test_lora_overlap_back_to_back_adapters_chunked(self): + self.server.execute_script( + self._script_lora_overlap_back_to_back_adapters_chunked + ) + + @staticmethod + def _script_lora_overlap_back_to_back_adapters_chunked(t: ScriptedContext): + adapters = [_LORA_ADAPTER, _LORA_ADAPTER_B, _LORA_ADAPTER, _LORA_ADAPTER_B] + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + lora_path=adapter, + ) + for adapter in adapters + ] + expected_ids = [_expected_lora_id(adapter) for adapter in adapters] + lora_id_by_rid: dict = {} + done = [False] * len(reqs) + for _ in range(2400): + for i, r in enumerate(reqs): + req = r.req + if req is not None: + lora_id_by_rid[r.rid] = req.lora_id + done[i] = done[i] or r.finished + if all(done): + break + yield + assert all(done) + expected_first_chunks = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE + chunk_counts = [r.chunks_done for r in reqs] + a_counts = sorted([reqs[0].chunks_done, reqs[2].chunks_done]) + b_counts = sorted([reqs[1].chunks_done, reqs[3].chunks_done]) + assert a_counts == [0, expected_first_chunks], ( + f"adapter A pair must be one cold full prefill + one full prefix " + f"hit; per-req chunks_done={chunk_counts}" + ) + assert b_counts == [0, expected_first_chunks], ( + f"adapter B pair must be one cold full prefill + one full prefix " + f"hit; per-req chunks_done={chunk_counts}" + ) + for r, expected_id in zip(reqs, expected_ids): + assert lora_id_by_rid.get(r.rid) == expected_id + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_max_new_tokens.py b/test/manual/chunked_prefill/test_scripted_max_new_tokens.py new file mode 100644 index 000000000..f1da3114f --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_max_new_tokens.py @@ -0,0 +1,167 @@ +import unittest +from typing import List + +from sglang.test.scripted_runtime.context import ScriptedContext +from sglang.test.scripted_runtime.scheduler_hook import ScriptedBatchRecord +from sglang.test.scripted_runtime.test_case import ScriptedTestCase +from sglang.test.scripted_runtime_chunked_helpers import ( + DEFAULT_CHUNK_SIZE, + base_engine_kwargs, + run_until, + run_until_finished, +) + + +def _records_for_rid( + batch_log: List[ScriptedBatchRecord], rid: str +) -> List[ScriptedBatchRecord]: + return [rec for rec in batch_log if rid in rec.rids] + + +def _decode_records( + batch_log: List[ScriptedBatchRecord], rid: str +) -> List[ScriptedBatchRecord]: + return [rec for rec in _records_for_rid(batch_log, rid) if rec.mode == "decode"] + + +def _extend_records( + batch_log: List[ScriptedBatchRecord], rid: str +) -> List[ScriptedBatchRecord]: + return [rec for rec in _records_for_rid(batch_log, rid) if rec.mode == "extend"] + + +class TestMaxNewTokensDecodeForwardLaw(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_decode_forward_count_equals_mnt(self): + self.server.execute_script(self._script_decode_forward_count_law) + + @staticmethod + def _script_decode_forward_count_law(t: ScriptedContext): + for max_new_tokens in (1, 2, 3, 4): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, + max_new_tokens=max_new_tokens, + ignore_eos=True, + prompt_token=10 + max_new_tokens, + ) + yield from run_until_finished(r) + assert r.finished + + output_ids = r.req.output_ids + assert len(output_ids) == max_new_tokens, ( + f"max_new_tokens={max_new_tokens} must produce exactly " + f"{max_new_tokens} tokens; got {len(output_ids)}" + ) + + batch_log = t._scheduler_hook._batch_log + decode_records = _decode_records(batch_log, r.rid) + extend_records = _extend_records(batch_log, r.rid) + + assert len(decode_records) == max_new_tokens, ( + f"max_new_tokens={max_new_tokens} expected " + f"{max_new_tokens} decode forward batches, got " + f"{len(decode_records)}" + ) + assert len(extend_records) >= 2, ( + f"max_new_tokens={max_new_tokens} expected >= 2 extend (chunk) " + f"records, got {len(extend_records)}" + ) + + yield + + +class TestMaxNewTokensOneSkipsDecode(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_mnt_one_emits_token_on_prefill_then_one_dead_decode(self): + self.server.execute_script(self._script_mnt_one_skips_decode) + + @staticmethod + def _script_mnt_one_skips_decode(t: ScriptedContext): + r = t.start_req( + prompt_len=3 * DEFAULT_CHUNK_SIZE, + max_new_tokens=1, + ignore_eos=True, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2, ( + f"prompt spanning >=3 chunks should chunk at least twice, got " + f"chunks_done={r.chunks_done}" + ) + assert len(r.req.output_ids) == 1, ( + f"max_new_tokens=1 must produce exactly 1 token, got " + f"{len(r.req.output_ids)}" + ) + + batch_log = t._scheduler_hook._batch_log + decode_records = _decode_records(batch_log, r.rid) + assert len(decode_records) == 1, ( + f"max_new_tokens=1 under overlap launches exactly ONE trailing decode " + f"forward whose token is discarded; got {len(decode_records)}" + ) + + rid_records = _records_for_rid(batch_log, r.rid) + assert rid_records, "expected at least one batch record for the req" + rid_modes = [rec.mode for rec in rid_records] + first_decode_pos = rid_modes.index("decode") + assert first_decode_pos >= 1, ( + f"the lone decode must be preceded by an extend chunk; rid_modes=" + f"{rid_modes}" + ) + assert rid_modes[first_decode_pos - 1] in ("extend", "mixed"), ( + f"the record immediately before the lone decode must be the final " + f"prefill chunk that emits the only token; got " + f"{rid_modes[first_decode_pos - 1]!r}, rid_modes={rid_modes}" + ) + + +class TestMaxNewTokensFirstDecodeAdjacent(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_first_decode_immediately_follows_last_chunk(self): + self.server.execute_script(self._script_first_decode_adjacent) + + @staticmethod + def _script_first_decode_adjacent(t: ScriptedContext): + max_new_tokens = 16 + r = t.start_req( + prompt_len=3 * DEFAULT_CHUNK_SIZE, + max_new_tokens=max_new_tokens, + ignore_eos=True, + ) + yield from run_until(r, lambda h: h.finished, max_steps=400) + assert r.finished + assert len(r.req.output_ids) == max_new_tokens, ( + f"max_new_tokens={max_new_tokens} must produce exactly " + f"{max_new_tokens} tokens; got {len(r.req.output_ids)}" + ) + + batch_log = t._scheduler_hook._batch_log + rid_records = _records_for_rid(batch_log, r.rid) + + decode_records = _decode_records(batch_log, r.rid) + assert len(decode_records) == max_new_tokens, ( + f"expected {max_new_tokens} decode forwards, got " f"{len(decode_records)}" + ) + + rid_modes = [rec.mode for rec in rid_records] + first_decode_pos = rid_modes.index("decode") + assert first_decode_pos >= 1, ( + f"first decode must be preceded by an extend chunk; rid_modes=" + f"{rid_modes}" + ) + assert rid_modes[first_decode_pos - 1] == "extend", ( + f"record immediately before the first decode (in this rid's " + f"subsequence) must be the last extend chunk; got " + f"{rid_modes[first_decode_pos - 1]!r}, rid_modes={rid_modes}" + ) + assert all(mode == "extend" for mode in rid_modes[:first_decode_pos]), ( + f"all records before the first decode must be extend chunks; " + f"rid_modes={rid_modes}" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_multi_req.py b/test/manual/chunked_prefill/test_scripted_multi_req.py new file mode 100644 index 000000000..27db67252 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_multi_req.py @@ -0,0 +1,537 @@ +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 ( + BALLAST_MAX_NEW_TOKENS, + DEFAULT_CHUNK_SIZE, + DEFAULT_MAX_STEPS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +def _drain_flush_then_assert_no_kv_leak(t: ScriptedContext, baseline: dict): + for _ in range(5): + yield + t.flush_cache() + yield + final = t.engine_stats() + assert ( + final["kv_pool_free"] >= baseline["kv_pool_free"] + ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}" + + +class TestMultiReqBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_second_chunked_waits(self): + self.server.execute_script(self._script_second_chunked_waits) + + @staticmethod + def _script_second_chunked_waits(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield + + assert r1.is_chunking, "r1 should still be chunking" + assert ( + not r2.is_chunking + ), "r2 must wait for r1's chunk loop to clear before chunking" + + yield from run_until_all_finished([r1, r2]) + assert r1.finished and r2.finished + + def test_hundred_short_reqs(self): + self.server.execute_script(self._script_hundred_short_reqs) + + @staticmethod + def _script_hundred_short_reqs(t: ScriptedContext): + baseline = t.engine_stats() + reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(100)] + yield from run_until_all_finished(reqs, max_steps=2000) + for r in reqs: + assert r.finished + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_five_hundred_short_reqs(self): + self.server.execute_script(self._script_five_hundred_short_reqs) + + @staticmethod + def _script_five_hundred_short_reqs(t: ScriptedContext): + baseline = t.engine_stats() + reqs = [t.start_req(prompt_len=8, max_new_tokens=1) for _ in range(500)] + yield from run_until_all_finished(reqs, max_steps=20000) + for r in reqs: + assert r.finished + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_mixed_ten_chunked_ten_short(self): + self.server.execute_script(self._script_mixed_ten_chunked_ten_short) + + @staticmethod + def _script_mixed_ten_chunked_ten_short(t: ScriptedContext): + baseline = t.engine_stats() + chunked = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + i + ) + for i in range(10) + ] + shorts = [t.start_req(prompt_len=8, max_new_tokens=2) for _ in range(10)] + all_reqs = chunked + shorts + finished = False + for _ in range(DEFAULT_MAX_STEPS * 20): + assert sum(1 for r in all_reqs if r.is_chunking) <= 1 + if all(r.finished for r in all_reqs): + finished = True + break + yield + if not finished: + raise AssertionError("not all reqs finished") + for r in chunked: + assert r.chunks_done >= 2 + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_submit_during_chunk_mid(self): + self.server.execute_script(self._script_submit_during_chunk_mid) + + @staticmethod + def _script_submit_during_chunk_mid(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield + r3 = t.start_req(prompt_len=16, max_new_tokens=2) + yield + + assert r1.is_chunking, "r1 must still hold the single chunked slot" + assert not r2.is_chunking, "r2 must wait until r1's chunk loop clears" + assert not r3.is_chunking, "r3 must wait until r1's chunk loop clears" + + yield from run_until_all_finished([r1, r2, r3]) + assert r1.finished and r2.finished and r3.finished + + def test_five_identical_prompts(self): + self.server.execute_script(self._script_five_identical_prompts) + + @staticmethod + def _script_five_identical_prompts(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + + others = [ + t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(4) + ] + yield from run_until_all_finished(others) + for r in others: + assert r.finished + assert r.chunks_done == 0 + assert r1.finished + + def test_sibling_shared_prefix(self): + self.server.execute_script(self._script_sibling_shared_prefix) + + @staticmethod + def _script_sibling_shared_prefix(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN + 8, max_new_tokens=2) + yield from run_until_finished(r2) + assert r1.finished and r2.finished + assert ( + r2.chunks_done < r1.chunks_done + ), "r2 reuses r1's cached prefix, so it should chunk fewer times" + + def test_trickle_per_yield_50(self): + self.server.execute_script(self._script_trickle_per_yield_50) + + @staticmethod + def _script_trickle_per_yield_50(t: ScriptedContext): + baseline = t.engine_stats() + reqs = [] + for _ in range(50): + reqs.append(t.start_req(prompt_len=8, max_new_tokens=2)) + yield + yield from run_until_all_finished(reqs, max_steps=2000) + for r in reqs: + assert r.finished + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_submit_then_immediate_abort(self): + self.server.execute_script(self._script_submit_then_immediate_abort) + + @staticmethod + def _script_submit_then_immediate_abort(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + t.abort(r) + for _ in range(5): + yield + assert r.kv_pages == 0 + assert r.req is None or r.req.req_pool_idx is None + + def test_rid_reuse_after_finish(self): + self.server.execute_script(self._script_rid_reuse_after_finish) + + @staticmethod + def _script_rid_reuse_after_finish(t: ScriptedContext): + baseline = t.engine_stats() + r1 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-rid") + yield from run_until_finished(r1) + r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-rid") + yield from run_until_finished(r2) + assert r1.finished and r2.finished + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_submit_pause_n_resubmit_same_rid(self): + self.server.execute_script(self._script_submit_pause_n_resubmit_same_rid) + + @staticmethod + def _script_submit_pause_n_resubmit_same_rid(t: ScriptedContext): + baseline = t.engine_stats() + r1 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-200") + yield from run_until_finished(r1) + for _ in range(200): + assert (1 if t.scheduler.chunked_req is not None else 0) == 0 + yield + r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="reuse-200") + yield from run_until_finished(r2) + assert r2.finished + yield from _drain_flush_then_assert_no_kv_leak(t, baseline) + + def test_submit_during_decode_of_other(self): + self.server.execute_script(self._script_submit_during_decode_of_other) + + @staticmethod + def _script_submit_during_decode_of_other(t: ScriptedContext): + r1 = t.start_req(prompt_len=16, max_new_tokens=16) + yield from run_until(r1, lambda h: h.status == "running") + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r2, lambda h: h.is_chunking) + + comp = t.batch_composition() + chunked_set = set(comp.get("chunked", [])) + prefill_set = set(comp.get("prefill", [])) + decode_set = set(comp.get("decode", [])) + assert r2.rid in chunked_set, f"r2 should be the chunked req; got {comp}" + assert chunked_set.isdisjoint(prefill_set) + assert chunked_set.isdisjoint(decode_set) + + yield from run_until_all_finished([r1, r2], max_steps=DEFAULT_MAX_STEPS * 5) + assert r1.finished and r2.finished + + def test_two_small_parallel(self): + self.server.execute_script(self._script_two_small_parallel) + + @staticmethod + def _script_two_small_parallel(t: ScriptedContext): + r1 = t.start_req(prompt_len=16, max_new_tokens=4) + r2 = t.start_req(prompt_len=16, max_new_tokens=4) + for _ in range(DEFAULT_MAX_STEPS): + assert ( + t.scheduler.chunked_req.rid + if t.scheduler.chunked_req is not None + else None + ) is None + if r1.finished and r2.finished: + break + yield + assert r1.finished and r2.finished + + def test_one_chunked_plus_many_short(self): + self.server.execute_script(self._script_one_chunked_plus_many_short) + + @staticmethod + def _script_one_chunked_plus_many_short(t: ScriptedContext): + chunked = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + shorts = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(5)] + for _ in range(DEFAULT_MAX_STEPS * 5): + assert (1 if t.scheduler.chunked_req is not None else 0) <= 1 + if chunked.finished and all(s.finished for s in shorts): + break + yield + assert chunked.chunks_done >= 2 + assert chunked.finished + for s in shorts: + assert s.finished + + def test_multiple_chunked_staggered(self): + self.server.execute_script(self._script_multiple_chunked_staggered) + + @staticmethod + def _script_multiple_chunked_staggered(t: ScriptedContext): + reqs = [] + for i in range(4): + reqs.append( + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + prompt_token=100 + i, + ) + ) + yield + assert sum(1 for r in reqs if r.is_chunking) <= 1 + yield + for _ in range(DEFAULT_MAX_STEPS * 10): + assert sum(1 for r in reqs if r.is_chunking) <= 1 + if all(r.finished for r in reqs): + break + yield + for r in reqs: + assert r.finished + assert r.chunks_done >= 2 + + def test_eight_concurrent_chunked(self): + self.server.execute_script(self._script_eight_concurrent_chunked) + + @staticmethod + def _script_eight_concurrent_chunked(t: ScriptedContext): + reqs = [ + t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + i + ) + for i in range(8) + ] + finished = False + for _ in range(DEFAULT_MAX_STEPS * 5): + assert sum(1 for r in reqs if r.is_chunking) <= 1 + if all(r.finished for r in reqs): + finished = True + break + yield + if not finished: + raise AssertionError("not all reqs finished") + for r in reqs: + assert r.chunks_done >= 2 + + def test_decode_only_batch(self): + self.server.execute_script(self._script_decode_only_batch) + + @staticmethod + def _script_decode_only_batch(t: ScriptedContext): + reqs = [t.start_req(prompt_len=4, max_new_tokens=8) for _ in range(10)] + for _ in range(DEFAULT_MAX_STEPS * 3): + assert ( + t.scheduler.chunked_req.rid + if t.scheduler.chunked_req is not None + else None + ) is None, "pure decode workload must never populate chunked_req" + assert (1 if t.scheduler.chunked_req is not None else 0) == 0 + if all(r.finished for r in reqs): + return + yield + raise AssertionError("not all reqs finished") + + def test_mixed_prefill_lengths(self): + self.server.execute_script(self._script_mixed_prefill_lengths) + + @staticmethod + def _script_mixed_prefill_lengths(t: ScriptedContext): + lens = [8, 16, 32, 64, 128, 256, 512, 1024] + reqs = [t.start_req(prompt_len=L, max_new_tokens=2) for L in lens] + by_len = dict(zip(lens, reqs)) + for _ in range(DEFAULT_MAX_STEPS * 10): + assert sum(1 for r in reqs if r.is_chunking) <= 1 + if all(r.finished for r in reqs): + break + yield + for r in reqs: + assert r.finished + for L in (8, 16, 32, 64, 128): + assert by_len[L].chunks_done <= 1, f"prompt_len={L} should not chunk" + for L in (512, 1024): + assert by_len[L].chunks_done >= 2, f"prompt_len={L} should chunk >=2" + + def test_chunked_req_exclusive_of_batch_invariant(self): + self.server.execute_script( + self._script_chunked_req_exclusive_of_batch_invariant + ) + + @staticmethod + def _script_chunked_req_exclusive_of_batch_invariant(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(DEFAULT_MAX_STEPS * 5): + s = t.scheduler + chunked = s.chunked_req + running = s.running_batch + if chunked is not None and running is not None: + assert chunked not in running.reqs, ( + f"chunked_req must be exclusive of running_batch.reqs; " + f"chunked.rid={chunked.rid!r} appears in running_batch" + ) + if r1.finished and r2.finished: + return + yield + raise AssertionError("r1 and r2 did not both finish within step budget") + + def test_two_chunked_one_decode(self): + self.server.execute_script(self._script_two_chunked_one_decode) + + @staticmethod + def _script_two_chunked_one_decode(t: ScriptedContext): + chunked1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + chunked2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + short = t.start_req(prompt_len=8, max_new_tokens=4) + all_reqs = [chunked1, chunked2, short] + for _ in range(DEFAULT_MAX_STEPS * 10): + comp = t.batch_composition() + prefill_set = set(comp.get("prefill", [])) + decode_set = set(comp.get("decode", [])) + chunked_set = set(comp.get("chunked", [])) + assert prefill_set.isdisjoint(decode_set) + assert prefill_set.isdisjoint(chunked_set) + assert decode_set.isdisjoint(chunked_set) + if all(r.finished for r in all_reqs): + break + yield + assert chunked1.finished and chunked2.finished and short.finished + + def test_batch_state_query_during_run(self): + self.server.execute_script(self._script_batch_state_query_during_run) + + @staticmethod + def _script_batch_state_query_during_run(t: ScriptedContext): + reqs = [t.start_req(prompt_len=16, max_new_tokens=4) for _ in range(4)] + for _ in range(DEFAULT_MAX_STEPS): + comp = t.batch_composition() + prefill = set(comp.get("prefill", [])) + decode = set(comp.get("decode", [])) + assert prefill & decode == set() + if all(r.finished for r in reqs): + return + yield + raise AssertionError("not all reqs finished") + + +class TestMultiReqPriority(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + ) + + def test_parallel_with_priority(self): + self.server.execute_script(self._script_parallel_with_priority) + + @staticmethod + def _script_parallel_with_priority(t: ScriptedContext): + normal = [ + t.start_req( + prompt_len=16, + max_new_tokens=BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + priority=0, + ) + for _ in range(4) + ] + yield from run_until(normal[-1], lambda h: h.status == "running") + + t.exhaust_kv(leave_pages=0) + high = t.start_req(prompt_len=16, max_new_tokens=2, priority=100) + + preempted = False + for _ in range(DEFAULT_MAX_STEPS * 5): + if any(r.status == "waiting" for r in normal): + preempted = True + if high.finished: + break + yield + assert high.finished, "high-priority req must run to completion" + assert preempted, "a normal req must be preempted back to the waiting queue" + + t.abort_all() + for _ in range(DEFAULT_MAX_STEPS): + if all(r.finished for r in normal): + break + yield + + +class TestMultiReqMixedChunk(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_mixed_chunk=True, + ) + + def test_long_prefill_chunks_more_with_concurrent_decode(self): + self.server.execute_script( + self._script_long_prefill_chunks_more_with_concurrent_decode + ) + + @staticmethod + def _script_long_prefill_chunks_more_with_concurrent_decode(t: ScriptedContext): + decoder = t.start_req(prompt_len=16, max_new_tokens=64) + yield from run_until(decoder, lambda h: h.status == "running") + + long_req = t.start_req(prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until(long_req, lambda h: h.is_chunking) + + saw_mixed_batch: bool = False + for _ in range(DEFAULT_MAX_STEPS * 5): + if t.last_batch_forward_mode == "MIXED": + saw_mixed_batch = True + if long_req.finished and decoder.finished: + break + yield + assert long_req.finished and decoder.finished + + assert saw_mixed_batch, "expected at least one MIXED (prefill+decode) batch" + assert long_req.chunks_done >= 4, ( + f"a 4-chunk prompt must chunk >= 4 times under concurrent decode; " + f"got {long_req.chunks_done}" + ) + + def test_chunked_plus_decode_in_batch(self): + self.server.execute_script(self._script_chunked_plus_decode_in_batch) + + @staticmethod + def _script_chunked_plus_decode_in_batch(t: ScriptedContext): + r2 = t.start_req(prompt_len=8, max_new_tokens=64, ignore_eos=True) + yield from run_until( + r2, lambda h: h.status == "running" and len(h.req.output_ids) >= 1 + ) + + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking) + + r2_out_before = len(r2.req.output_ids) + saw_chunked_r1 = False + saw_r2_decode_during_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + comp = t.batch_composition() + if r1.rid in comp.get("chunked", []): + saw_chunked_r1 = True + if r2.rid in comp.get("decode", []) + comp.get("prefill", []): + saw_r2_decode_during_chunking = True + chunked_set = set(comp.get("chunked", [])) + assert chunked_set.isdisjoint(set(comp.get("prefill", []))) + assert chunked_set.isdisjoint(set(comp.get("decode", []))) + if not r1.is_chunking: + break + yield + assert saw_chunked_r1, "r1 was never observed in the chunked subset" + r2_out_after = len(r2.req.output_ids) if r2.req is not None else 64 + assert r2_out_after > r2_out_before, ( + f"r2's decode must progress while r1 chunks; output_ids stayed at " + f"{r2_out_before}" + ) + assert saw_r2_decode_during_chunking, ( + "r2 was never observed co-batched (decode/mixed) while r1 held " + "the chunked slot" + ) + + t.abort(r2) + yield from run_until_finished(r1) + for _ in range(40): + if t.is_fully_idle: + break + yield + assert r1.finished + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_page_size.py b/test/manual/chunked_prefill/test_scripted_page_size.py new file mode 100644 index 000000000..e023b6f97 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_page_size.py @@ -0,0 +1,106 @@ +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 ( + base_engine_kwargs, + run_until_finished, +) + + +class TestPageSize16ChunkSize16(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=16, page_size=16) + + def test_chunk_size_equals_page_size(self): + self.server.execute_script(self._script_chunk_size_equals_page_size) + + @staticmethod + def _script_chunk_size_equals_page_size(t: ScriptedContext): + r = t.start_req(prompt_len=8 * 16, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 8 + assert r.kv_pages == 0 + + def test_chunk_size_equals_page_size_plus_one(self): + self.server.execute_script(self._script_chunk_size_equals_page_size_plus_one) + + @staticmethod + def _script_chunk_size_equals_page_size_plus_one(t: ScriptedContext): + r = t.start_req(prompt_len=129, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 9 + assert r.kv_pages == 0 + + +class TestPageSize16ChunkSize16NonMultiplePrompt(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=16, page_size=16) + + def test_non_page_multiple_prompt_advances_in_page_steps(self): + self.server.execute_script( + self._script_non_page_multiple_prompt_advances_in_page_steps + ) + + @staticmethod + def _script_non_page_multiple_prompt_advances_in_page_steps(t: ScriptedContext): + r = t.start_req(prompt_len=60, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 4 + assert r.kv_pages == 0 + + +class TestPageSize16ChunkSize64(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=64, page_size=16) + + def test_chunk_spans_multiple_pages_exact(self): + self.server.execute_script(self._script_chunk_spans_multiple_pages_exact) + + @staticmethod + def _script_chunk_spans_multiple_pages_exact(t: ScriptedContext): + r = t.start_req(prompt_len=4 * 64, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 4 + assert r.kv_pages == 0 + + def test_chunk_spans_multiple_pages_plus_one_page(self): + self.server.execute_script( + self._script_chunk_spans_multiple_pages_plus_one_page + ) + + @staticmethod + def _script_chunk_spans_multiple_pages_plus_one_page(t: ScriptedContext): + r = t.start_req(prompt_len=4 * 64 + 16, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 5 + assert r.kv_pages == 0 + + +class TestPageSize16RadixHit(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=16, page_size=16, disable_radix_cache=False + ) + + def test_radix_hit_on_page_boundary(self): + self.server.execute_script(self._script_radix_hit_on_page_boundary) + + @staticmethod + def _script_radix_hit_on_page_boundary(t: ScriptedContext): + r_warm = t.start_req(prompt_len=4 * 16, max_new_tokens=1) + yield from run_until_finished(r_warm, max_steps=400) + assert r_warm.finished + + yield + + r = t.start_req(prompt_len=8 * 16, max_new_tokens=1) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert r.chunks_done == 4 + assert r.kv_pages == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_piecewise_cuda_graph.py b/test/manual/chunked_prefill/test_scripted_piecewise_cuda_graph.py new file mode 100644 index 000000000..a9138d55d --- /dev/null +++ b/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() diff --git a/test/manual/chunked_prefill/test_scripted_pp.py b/test/manual/chunked_prefill/test_scripted_pp.py new file mode 100644 index 000000000..bb5cc9e23 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_pp.py @@ -0,0 +1,246 @@ +import unittest +from typing import Any, Dict + +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_all_finished, + run_until_finished, +) +from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST + + +def _pp_engine_kwargs(*, pp_size: int = 2, **overrides: Any) -> Dict[str, Any]: + return base_engine_kwargs( + model_path=DEFAULT_MODEL_NAME_FOR_TEST, + pp_size=pp_size, + **overrides, + ) + + +def _expected_chunks(prompt_len: int, chunk_size: int) -> int: + if prompt_len <= chunk_size: + return 0 + return (prompt_len + chunk_size - 1) // chunk_size + + +def _drain_until_released(t, *handles): + for _ in range(16): + if all( + h.kv_pages == 0 and (h.req is None or h.req.req_pool_idx is None) + for h in handles + ): + return + yield + + +class TestPPBasic(ScriptedTestCase): + ENGINE_KWARGS = _pp_engine_kwargs() + + def test_pp_abort_during_inflight_chunk(self): + self.server.execute_script(self._script_pp_abort_during_inflight_chunk) + + @staticmethod + def _script_pp_abort_during_inflight_chunk(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.abort(r) + yield from _drain_until_released(t, r) + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_pp_last_chunk_cross_mb_kv_correctness(self): + self.server.execute_script(self._script_pp_last_chunk_cross_mb_kv_correctness) + + @staticmethod + def _script_pp_last_chunk_cross_mb_kv_correctness(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=8, + prompt_token=7, + ignore_eos=True, + ) + yield from run_until_finished(r, max_steps=800) + assert r.finished + expected_chunks = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE) + assert r.chunks_done == expected_chunks, ( + f"cross-mb chunked prefill must complete all {expected_chunks} chunks, " + f"got chunks_done={r.chunks_done}" + ) + + def test_pp_static_chunk_size_predictor_returns_none(self): + self.server.execute_script( + self._script_pp_static_chunk_size_predictor_returns_none + ) + + @staticmethod + def _script_pp_static_chunk_size_predictor_returns_none(t: ScriptedContext): + sched = t.scheduler + assert sched.enable_dynamic_chunking is False + assert sched.predict_next_chunk_size(0) is None + assert sched.predict_next_chunk_size(VERY_LONG_PROMPT_LEN // 2) is None + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r, max_steps=800) + assert r.finished + expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE) + assert r.chunks_done == expected, ( + f"static chunked_prefill_size must produce exactly {expected} chunks, " + f"got {r.chunks_done}" + ) + + def test_pp_multi_microbatch_chunks_done_aggregation(self): + self.server.execute_script( + self._script_pp_multi_microbatch_chunks_done_aggregation + ) + + @staticmethod + def _script_pp_multi_microbatch_chunks_done_aggregation(t: ScriptedContext): + r = t.start_req(prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 4, ( + f"PP=2 multi-chunk req should aggregate >=4 chunks_done across " + f"microbatches, got {r.chunks_done}" + ) + + def test_pp_two_chunked_one_per_mb_simultaneous(self): + self.server.execute_script(self._script_pp_two_chunked_one_per_mb_simultaneous) + + @staticmethod + def _script_pp_two_chunked_one_per_mb_simultaneous(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=1 + ) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=2 + ) + yield from run_until_all_finished(handles=[r1, r2], max_steps=800) + assert r1.finished and r2.finished + assert r1.chunks_done >= 2 and r2.chunks_done >= 2 + assert r1.kv_pages == 0 and r2.kv_pages == 0 + assert r1.lock_refs == 0 and r2.lock_refs == 0 + + def test_pp_retract_chunked_in_middle_mb(self): + self.server.execute_script(self._script_pp_retract_chunked_in_middle_mb) + + @staticmethod + def _script_pp_retract_chunked_in_middle_mb(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + 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 + + def test_pp_chunked_req_to_exclude_pp_context(self): + self.server.execute_script(self._script_pp_chunked_req_to_exclude_pp_context) + + @staticmethod + def _script_pp_chunked_req_to_exclude_pp_context(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + t.abort(r) + r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=2) + yield from run_until_finished(r2, max_steps=400) + assert r2.finished + assert r2.kv_pages == 0 + assert r2.lock_refs == 0 + + +@unittest.skip( + "PD-Multiplexing is mutually exclusive with pipeline parallelism: ServerArgs " + "asserts pp_size == 1 (server_args.py:7284-7287, 'PD-Multiplexing is only " + "supported with pipeline parallelism disabled'). It further requires " + "chunked_prefill_size == -1 and disable_overlap_schedule, both of which the PP " + "chunked scripted-runtime harness mandates the opposite of. enable_pdmux=True " + "with pp_size=2 therefore cannot launch -- the spawned server dies on that " + "assert and setUpClass times out on the handshake. This is a genuine engine " + "config constraint, not a test bug, so the pdmux-under-PP scenario is not a " + "valid combination to exercise here." +) +class TestPPPdmux(ScriptedTestCase): + ENGINE_KWARGS = _pp_engine_kwargs(enable_pdmux=True) + + def test_pp_split_prefill_chunked_no_merge_assert(self): + self.server.execute_script( + self._script_pp_split_prefill_chunked_no_merge_assert + ) + + @staticmethod + def _script_pp_split_prefill_chunked_no_merge_assert(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r, max_steps=800) + assert ( + r.finished + ), "engine died before req finished — merge_batch assert may have tripped" + assert r.chunks_done >= 2, ( + f"pdmux + chunked path must produce >=2 chunks to exercise " + f"split_prefill_batch filter; got chunks_done={r.chunks_done}" + ) + + +class TestPPDynamic(ScriptedTestCase): + ENGINE_KWARGS = _pp_engine_kwargs(enable_dynamic_chunking=True) + + def test_naive_pp_chunked(self): + self.server.execute_script(self._script_naive_pp_chunked) + + @staticmethod + def _script_naive_pp_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + prompt_token=3, + ignore_eos=True, + ) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2, f"expected >=2 chunks, got {r.chunks_done}" + + def test_pp_dynamic_chunk_size_recompute_branch_taken(self): + self.server.execute_script( + self._script_pp_dynamic_chunk_size_recompute_branch_taken + ) + + @staticmethod + def _script_pp_dynamic_chunk_size_recompute_branch_taken(t: ScriptedContext): + sched = t.scheduler + assert sched.enable_dynamic_chunking is True + assert sched.length_predictor is not None + assert sched.length_predictor.is_ready is True + dynamic_size = sched.predict_next_chunk_size(0) + assert dynamic_size is not None + assert isinstance(dynamic_size, int) and dynamic_size > 0 + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=4 + ) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2, f"expected >=2 chunks, got {r.chunks_done}" + + +class TestPPSize4(ScriptedTestCase): + ENGINE_KWARGS = _pp_engine_kwargs(pp_size=4) + + def test_pp_size_4_chunked_completes(self): + self.server.execute_script(self._script_pp_size_4_chunked_completes) + + @staticmethod + def _script_pp_size_4_chunked_completes(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 4 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_priority.py b/test/manual/chunked_prefill/test_scripted_priority.py new file mode 100644 index 000000000..c164342b6 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_priority.py @@ -0,0 +1,309 @@ +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 ( + BALLAST_MAX_NEW_TOKENS, + DEFAULT_CHUNK_SIZE, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +class TestPriorityBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_retract_mid_chunk_releases_kv(self): + self.server.execute_script(self._script_retract_mid_chunk_releases_kv) + + @staticmethod + def _script_retract_mid_chunk_releases_kv(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + pages_before = r.kv_pages + assert pages_before > 0 + + t.pause_generation(mode="retract") + yield + + assert ( + r.status == "waiting" + ), f"force-retracted chunked req must be back in waiting; got {r.status}" + assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}" + + t.continue_generation() + yield from run_until_finished(r) + assert r.finished + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_retract_and_resume(self): + self.server.execute_script(self._script_retract_and_resume) + + @staticmethod + def _script_retract_and_resume(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + t.pause_generation(mode="retract") + yield + assert r.status == "waiting" + assert r.kv_pages == 0 + + t.continue_generation() + yield from run_until_finished(r) + assert r.finished + + def test_force_retract_at_chunk_0(self): + self.server.execute_script(self._script_force_retract_at_chunk_0) + + @staticmethod + def _script_force_retract_at_chunk_0(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done <= 1) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0 + t.continue_generation() + yield from run_until_finished(r, max_steps=800) + assert r.finished + + def test_force_retract_at_chunk_mid(self): + self.server.execute_script(self._script_force_retract_at_chunk_mid) + + @staticmethod + def _script_force_retract_at_chunk_mid(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.chunks_done >= 2 and h.is_chunking) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0 + t.continue_generation() + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.lock_refs == 0 + + def test_force_retract_at_last_chunk(self): + self.server.execute_script(self._script_force_retract_at_last_chunk) + + @staticmethod + def _script_force_retract_at_last_chunk(t: ScriptedContext): + r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4) + yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0 + 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 + + def test_force_retract_then_readmit(self): + self.server.execute_script(self._script_force_retract_then_readmit) + + @staticmethod + def _script_force_retract_then_readmit(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0, "retract must release KV before re-admission" + 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 + + def test_retract_one_admit_one(self): + self.server.execute_script(self._script_retract_one_admit_one) + + @staticmethod + def _script_retract_one_admit_one(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking) + r2 = t.start_req(prompt_len=8, max_new_tokens=2) + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r2) + assert r2.finished + assert r2.kv_pages == 0 + yield from run_until_finished(r1) + assert r1.finished + assert r1.kv_pages == 0 + assert r1.lock_refs == 0 + + def test_retract_during_decode(self): + self.server.execute_script(self._script_retract_during_decode) + + @staticmethod + def _script_retract_during_decode(t: ScriptedContext): + r = t.start_req(prompt_len=8, max_new_tokens=32) + yield from run_until(r, lambda h: h.status == "running") + assert r.kv_pages > 0, "decode-state req must own KV before retract" + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}" + t.continue_generation() + yield from run_until_finished(r) + assert r.finished + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_retract_then_abort_idempotent(self): + self.server.execute_script(self._script_retract_then_abort_idempotent) + + @staticmethod + def _script_retract_then_abort_idempotent(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + t.pause_generation(mode="retract") + t.abort(r) + for _ in range(12): + if ( + r.kv_pages == 0 + and r.lock_refs == 0 + and (r.req is None or r.req.req_pool_idx is None) + ): + break + yield + assert r.kv_pages == 0 + assert r.lock_refs == 0 + assert r.req is None or r.req.req_pool_idx is None + t.continue_generation() + yield + assert r.kv_pages == 0 and r.lock_refs == 0 + + def test_retract_chunked_resume_in_waiting(self): + self.server.execute_script(self._script_retract_chunked_resume_in_waiting) + + @staticmethod + def _script_retract_chunked_resume_in_waiting(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0 + assert r.status == "waiting" + t.continue_generation() + yield from run_until_finished(r, max_steps=800) + assert r.finished + + def test_two_retracts_same_yield(self): + self.server.execute_script(self._script_two_retracts_same_yield) + + @staticmethod + def _script_two_retracts_same_yield(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking) + t.pause_generation(mode="retract") + yield + assert r1.kv_pages == 0 + assert r2.kv_pages == 0 + t.continue_generation() + yield from run_until_all_finished([r1, r2]) + assert r1.finished and r2.finished + assert r1.lock_refs == 0 + assert r2.lock_refs == 0 + + def test_retract_then_re_chunk(self): + self.server.execute_script(self._script_retract_then_re_chunk) + + @staticmethod + def _script_retract_then_re_chunk(t: ScriptedContext): + r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until(r, lambda h: h.chunks_done >= 1) + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0, "retract must release KV" + t.continue_generation() + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.lock_refs == 0 + assert r.kv_pages == 0 + + +class TestPriorityPriority(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + ) + + def test_naive_priority_chunked(self): + self.server.execute_script(self._script_naive_priority_chunked) + + @staticmethod + def _script_naive_priority_chunked(t: ScriptedContext): + low = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, priority=0) + yield from run_until(low, lambda h: h.is_chunking) + + high = t.start_req(prompt_len=8, max_new_tokens=2, priority=10) + + yield from run_until_finished(high) + assert high.finished + assert not low.finished + + yield from run_until_all_finished([low, high]) + assert low.finished and high.finished + + +class TestPriorityPreempt(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + max_running_requests=1, + priority_scheduling_preemption_threshold=0, + ) + + def test_priority_preempt_decode_victim_to_waiting(self): + self.server.execute_script( + self._script_priority_preempt_decode_victim_to_waiting + ) + + @staticmethod + def _script_priority_preempt_decode_victim_to_waiting(t: ScriptedContext): + low = t.start_req( + prompt_len=8, + max_new_tokens=BALLAST_MAX_NEW_TOKENS, + priority=0, + ignore_eos=True, + ) + yield from run_until(low, lambda h: h.status == "running") + assert low.kv_pages > 0 + + high = t.start_req(prompt_len=8, max_new_tokens=2, priority=10) + yield from run_until(low, lambda h: h.status == "waiting") + + assert low.status == "waiting" + assert low.kv_pages == 0 + yield from run_until_finished(high) + assert high.finished + + def test_priority_preempt_release_invariant(self): + self.server.execute_script(self._script_priority_preempt_release_invariant) + + @staticmethod + def _script_priority_preempt_release_invariant(t: ScriptedContext): + r_low = t.start_req( + prompt_len=8, + max_new_tokens=BALLAST_MAX_NEW_TOKENS, + priority=0, + ignore_eos=True, + ) + yield from run_until(r_low, lambda h: h.status == "running") + pages_before = r_low.kv_pages + assert pages_before > 0 + + r_high = t.start_req(prompt_len=8, max_new_tokens=2, priority=10) + yield from run_until(r_low, lambda h: h.status == "waiting") + assert r_low.kv_pages == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_radix.py b/test/manual/chunked_prefill/test_scripted_radix.py new file mode 100644 index 000000000..015a4d090 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_radix.py @@ -0,0 +1,741 @@ +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, + DEFAULT_MAX_STEPS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, +) + + +class TestRadixBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_radix_full_prefix_hit_nine_reqs(self): + self.server.execute_script(self._script_radix_full_prefix_hit_nine_reqs) + + @staticmethod + def _script_radix_full_prefix_hit_nine_reqs(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + + others = [ + t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(9) + ] + yield from run_until_all_finished(others) + for r in others: + assert r.chunks_done == 0 + + def test_radix_hit_full_prefix(self): + self.server.execute_script(self._script_radix_hit_full_prefix) + + @staticmethod + def _script_radix_hit_full_prefix(t: ScriptedContext): + r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r1) + + r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1) + yield from run_until_finished(r2) + assert r2.chunks_done == 0 + + def test_radix_hit_partial_then_chunk_tail(self): + self.server.execute_script(self._script_radix_hit_partial_then_chunk_tail) + + @staticmethod + def _script_radix_hit_partial_then_chunk_tail(t: ScriptedContext): + r1 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r1) + + r2 = t.start_req(prompt_len=3 * DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1) + yield from run_until_finished(r2) + assert r2.chunks_done == 2 + + def test_radix_evict_then_resubmit_rechunks(self): + self.server.execute_script(self._script_radix_evict_then_resubmit_rechunks) + + @staticmethod + def _script_radix_evict_then_resubmit_rechunks(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + assert r1.finished + for _ in range(5): + yield + t.evict_radix(prefix_tokens=None) + yield + + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r2) + assert r2.finished + assert r2.chunks_done >= 2, ( + f"after eviction r2 must re-chunk from scratch; " + f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}" + ) + assert ( + r2.req.cached_tokens == 0 + ), f"eviction must clear r1's prefix; cached_tokens={r2.req.cached_tokens}" + assert r2.kv_pages == 0 + assert r2.lock_refs == 0 + + def test_radix_resume_init_next_round_path(self): + self.server.execute_script(self._script_radix_resume_init_next_round_path) + + @staticmethod + def _script_radix_resume_init_next_round_path(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=1) + yield from run_until_finished(r1) + assert r1.finished + + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN + 2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2 + ) + yield from run_until_finished(r2) + assert r2.finished + assert r2.req.cached_tokens > 0, ( + f"r2 must hit r1's prefix to exercise the partial-hit chunked-" + f"resume branch; got cached_tokens={r2.req.cached_tokens}" + ) + assert r2.chunks_done >= 1, ( + f"residual tail beyond cached prefix should still chunk; got " + f"chunks_done={r2.chunks_done}" + ) + + def test_radix_lock_ref_concurrent_chunked(self): + self.server.execute_script(self._script_radix_lock_ref_concurrent_chunked) + + @staticmethod + def _script_radix_lock_ref_concurrent_chunked(t: ScriptedContext): + r_warm = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1) + yield from run_until_finished(r_warm) + reqs = [ + t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 4 + 8, max_new_tokens=2) + for _ in range(5) + ] + yield from run_until_all_finished(reqs) + for r in reqs: + assert r.finished + assert r.req.cached_tokens > 0 + assert r.lock_refs == 0 + + def test_radix_partial_hit_exact_chunk_boundary(self): + self.server.execute_script(self._script_radix_partial_hit_exact_chunk_boundary) + + @staticmethod + def _script_radix_partial_hit_exact_chunk_boundary(t: ScriptedContext): + r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r1) + r2 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r2) + assert r2.chunks_done == 0, ( + f"residual of exactly chunk_size must not chunk; " + f"chunks_done={r2.chunks_done} cached_tokens={r2.req.cached_tokens}" + ) + + def test_radix_two_distinct_prefixes(self): + self.server.execute_script(self._script_radix_two_distinct_prefixes) + + @staticmethod + def _script_radix_two_distinct_prefixes(t: ScriptedContext): + r_a = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=11 + ) + yield from run_until_finished(r_a) + r_b = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=22 + ) + yield from run_until_finished(r_b) + + r_a2 = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=11 + ) + yield from run_until_finished(r_a2) + assert r_a2.chunks_done == 0 + assert r_a2.req.cached_tokens > 0 + + r_b2 = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=22 + ) + yield from run_until_finished(r_b2) + assert r_b2.chunks_done == 0 + assert r_b2.req.cached_tokens > 0 + + def test_radix_full_prefix_minus_one(self): + self.server.execute_script(self._script_radix_full_prefix_minus_one) + + @staticmethod + def _script_radix_full_prefix_minus_one(t: ScriptedContext): + r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE - 1, max_new_tokens=1) + yield from run_until_finished(r1) + r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=1) + yield from run_until_finished(r2) + assert r2.chunks_done == 0 + + def test_radix_hit_changes_between_chunks(self): + self.server.execute_script(self._script_radix_hit_changes_between_chunks) + + @staticmethod + def _script_radix_hit_changes_between_chunks(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking and h.chunks_done >= 1) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1, max_steps=800) + yield from run_until_finished(r2, max_steps=800) + assert r1.finished and r2.finished + assert r2.chunks_done < r1.chunks_done, ( + f"r2 should hit r1's committed prefix; r2.chunks_done=" + f"{r2.chunks_done} not < r1.chunks_done={r1.chunks_done}" + ) + assert r2.req.cached_tokens > 0 + + def test_radix_evict_during_inflight_chunk(self): + self.server.execute_script(self._script_radix_evict_during_inflight_chunk) + + @staticmethod + def _script_radix_evict_during_inflight_chunk(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + t.evict_radix(prefix_tokens=None) + yield + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_radix_full_hit_no_chunked_path(self): + self.server.execute_script(self._script_radix_full_hit_no_chunked_path) + + @staticmethod + def _script_radix_full_hit_no_chunked_path(t: ScriptedContext): + prompt_len: int = 16 * DEFAULT_CHUNK_SIZE + r_warm = t.start_req(prompt_len=prompt_len, max_new_tokens=1) + yield from run_until_finished(r_warm, max_steps=1200) + assert r_warm.finished + + r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) + yield from run_until_finished(r, max_steps=400) + assert r.finished + assert ( + r.chunks_done == 0 + ), f"full prefix hit must skip chunked path; got chunks_done={r.chunks_done}" + + def test_radix_evict_race_concurrent_chunked_admit(self): + self.server.execute_script( + self._script_radix_evict_race_concurrent_chunked_admit + ) + + @staticmethod + def _script_radix_evict_race_concurrent_chunked_admit(t: ScriptedContext): + warm_len: int = 4 * DEFAULT_CHUNK_SIZE + r_warm = t.start_req(prompt_len=warm_len, max_new_tokens=1) + yield from run_until_finished(r_warm, max_steps=400) + assert r_warm.finished + + for _ in range(5): + yield + t.evict_radix(prefix_tokens=None) + r = t.start_req( + prompt_len=warm_len + DEFAULT_CHUNK_SIZE * 2, + max_new_tokens=2, + ) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.req.cached_tokens == 0, ( + f"eviction must clear the warm prefix; " + f"cached_tokens={r.req.cached_tokens} chunks_done={r.chunks_done}" + ) + assert r.chunks_done >= 2 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_chunked_req_re_chunked_after_resume_same_prefix(self): + self.server.execute_script( + self._script_chunked_req_re_chunked_after_resume_same_prefix + ) + + @staticmethod + def _script_chunked_req_re_chunked_after_resume_same_prefix(t: ScriptedContext): + prompt_len: int = 4 * DEFAULT_CHUNK_SIZE + r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + t.pause_generation(mode="retract") + yield + assert r.kv_pages == 0, f"retract must release KV; got {r.kv_pages}" + + t.continue_generation() + yield from run_until_finished(r, max_steps=800) + assert r.finished + expected_total: int = prompt_len // DEFAULT_CHUNK_SIZE + assert r.chunks_done >= expected_total, ( + f"lifetime chunks_done after retract+resume must cover the " + f"whole prompt; expected >= {expected_total}, got " + f"{r.chunks_done}" + ) + assert r.chunks_done < 2 * expected_total, ( + f"lifetime chunks_done after retract+resume should not double " + f"the prompt's chunk count; expected < {2 * expected_total}, " + f"got {r.chunks_done}" + ) + + +class TestRadixNoTailChunked(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_page_size_one_chunked_has_no_partial_page_tail(self): + self.server.execute_script( + self._script_page_size_one_chunked_has_no_partial_page_tail + ) + + @staticmethod + def _script_page_size_one_chunked_has_no_partial_page_tail(t: ScriptedContext): + s = t.scheduler + prompt_len: int = 4 * DEFAULT_CHUNK_SIZE + r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + observed_mid_chunk: bool = False + for _ in range(800): + req = s.chunked_req + if req is not None and req.rid == r.rid: + observed_mid_chunk = True + prefix_len: int = len(req.prefix_indices) + protected_len: int = req.cache_protected_len + assert prefix_len == protected_len, ( + f"page_size=1 must take the no-tail else branch: " + f"len(prefix_indices)={prefix_len} != " + f"cache_protected_len={protected_len} (a partial-page tail " + f"was appended, which only happens for page_size > 1)" + ) + if r.finished: + break + yield + assert r.finished + assert observed_mid_chunk, ( + "test must observe r as the in-flight chunked_req at least once; the " + "no-tail else branch was never exercised" + ) + assert ( + r.kv_pages == 0 + ), f"finished chunked req must release KV; got {r.kv_pages}" + + +class TestRadixHitCountInvariant(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_chunked_stash_no_hit_count_inflation_invariant(self): + self.server.execute_script( + self._script_chunked_stash_no_hit_count_inflation_invariant + ) + + @staticmethod + def _script_chunked_stash_no_hit_count_inflation_invariant(t: ScriptedContext): + def _snapshot_hit_counts(root) -> dict: + snapshot: dict = {} + stack = [root] + while stack: + node = stack.pop() + snapshot[node.id] = node.hit_count + stack.extend(node.children.values()) + return snapshot + + s = t.scheduler + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + baseline = _snapshot_hit_counts(s.tree_cache.root_node) + + prev_chunks: int = r.chunks_done + observed_chunk_admissions: int = 0 + for _ in range(800): + cur_chunks: int = r.chunks_done + if cur_chunks > prev_chunks: + observed_chunk_admissions += 1 + cur = _snapshot_hit_counts(s.tree_cache.root_node) + for node_id, base_count in baseline.items(): + if node_id in cur: + assert cur[node_id] == base_count, ( + f"_inc_hit_count(chunked=True) inflated existing " + f"node id={node_id} hit_count: baseline={base_count}, " + f"now={cur[node_id]}" + ) + prev_chunks = cur_chunks + if r.finished: + break + yield + assert r.finished + assert observed_chunk_admissions >= 2, ( + f"test must exercise at least 2 chunk admissions after baseline; " + f"observed {observed_chunk_admissions}" + ) + + +class TestRadixHitCountNonChunked(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_non_chunked_prefix_hit_increments_hit_count_by_one(self): + self.server.execute_script( + self._script_non_chunked_prefix_hit_increments_hit_count_by_one + ) + + @staticmethod + def _script_non_chunked_prefix_hit_increments_hit_count_by_one(t: ScriptedContext): + def _snapshot_hit_counts(root) -> dict: + snapshot: dict = {} + stack = [root] + while stack: + node = stack.pop() + snapshot[node.id] = node.hit_count + stack.extend(node.children.values()) + return snapshot + + s = t.scheduler + + r_warm = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=1, prompt_token=11 + ) + yield from run_until_finished(r_warm) + assert r_warm.finished + + baseline = _snapshot_hit_counts(s.tree_cache.root_node) + + r2 = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE + 1, max_new_tokens=1, prompt_token=11 + ) + yield from run_until_finished(r2) + assert r2.finished + assert r2.chunks_done == 0, ( + f"residual past the warm prefix is one token and must not chunk; " + f"got chunks_done={r2.chunks_done}" + ) + assert r2.req.cached_tokens > 0, ( + f"r2 must hit the warm 2-chunk prefix to drive the non-chunked " + f"hit_count increment; got cached_tokens={r2.req.cached_tokens}" + ) + + cur = _snapshot_hit_counts(s.tree_cache.root_node) + incremented_by_one = 0 + for node_id, base_count in baseline.items(): + if node_id not in cur: + continue + delta: int = cur[node_id] - base_count + assert delta in (0, 1), ( + f"non-chunked insert moved existing node id={node_id} hit_count " + f"by {delta} (expected 0 or 1): baseline={base_count}, " + f"now={cur[node_id]}" + ) + if delta == 1: + incremented_by_one += 1 + assert incremented_by_one >= 1, ( + f"_inc_hit_count(chunked=False) at radix_cache.py:672 must bump at " + f"least one matched warm-prefix node by exactly 1; saw " + f"{incremented_by_one} such nodes" + ) + + +class TestRadixPartialPage(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + page_size=4, + ) + + def test_partial_page_tail_no_double_free_invariant(self): + self.server.execute_script( + self._script_partial_page_tail_no_double_free_invariant + ) + + @staticmethod + def _script_partial_page_tail_no_double_free_invariant(t: ScriptedContext): + s = t.scheduler + allocator = s.token_to_kv_pool_allocator + free_before: int = allocator.available_size() + prompt_len: int = 4 * DEFAULT_CHUNK_SIZE + 7 + r = t.start_req(prompt_len=prompt_len, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + for _ in range(800): + req = s.chunked_req + if req is not None and req.rid == r.rid: + prefix_len: int = len(req.prefix_indices) + protected_len: int = req.cache_protected_len + assert prefix_len >= protected_len, ( + f"len(prefix_indices)={prefix_len} dropped below " + f"cache_protected_len={protected_len}: tail was freed " + f"prematurely" + ) + if r.finished: + break + yield + assert r.finished + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + free_after: int = allocator.available_size() + assert free_after == free_before, ( + f"KV pool free count delta on chunked req lifecycle must be 0; " + f"got free_before={free_before}, free_after={free_after} " + f"(double-free or leak of partial-page tail)" + ) + + +class TestRadixFcfs(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + schedule_policy="fcfs", + ) + + def test_naive_radix_chunked(self): + self.server.execute_script(self._script_naive_radix_chunked) + + @staticmethod + def _script_naive_radix_chunked(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + assert r1.finished + + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN + DEFAULT_CHUNK_SIZE * 2, + max_new_tokens=2, + ) + yield from run_until_finished(r2) + assert r2.finished + assert r2.req.cached_tokens > 0 + assert r2.chunks_done >= 1 + + +class TestRadixDisabled(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + disable_radix_cache=True, + kv_canary="none", + kv_canary_real_data="none", + kv_canary_sweep_interval=0, + ) + + def test_radix_disabled_chunks_every_time(self): + self.server.execute_script(self._script_radix_disabled_chunks_every_time) + + @staticmethod + def _script_radix_disabled_chunks_every_time(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r1) + assert r1.chunks_done >= 2 + + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r2) + assert r2.chunks_done >= 2 + + +class TestRadixLpm(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + schedule_policy="lpm", + ) + + def test_radix_lpm_policy_chunked_priority(self): + self.server.execute_script(self._script_radix_lpm_policy_chunked_priority) + + @staticmethod + def _script_radix_lpm_policy_chunked_priority(t: ScriptedContext): + r_warm = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=1 + ) + yield from run_until_finished(r_warm) + assert r_warm.finished + for _ in range(5): + yield + + r_long = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=1 + ) + r_short = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=7 + ) + + first_admitted = None + cached_tokens_by_rid: dict = {} + for _ in range(DEFAULT_MAX_STEPS): + comp = t.batch_composition() + active = ( + comp.get("chunked", []) + + comp.get("prefill", []) + + comp.get("running", []) + ) + if first_admitted is None and active: + first_admitted = active[0] + for r in (r_long, r_short): + if r.req is not None: + cached_tokens_by_rid[r.rid] = r.req.cached_tokens + if r_long.finished and r_short.finished: + break + yield + assert r_long.finished and r_short.finished + assert first_admitted == r_long.rid, ( + f"LPM must admit the longest-prefix-match req first; first admitted " + f"rid was {first_admitted!r}, expected r_long={r_long.rid!r}" + ) + assert cached_tokens_by_rid.get(r_long.rid, 0) > 0 + assert cached_tokens_by_rid.get(r_short.rid, 0) == 0 + for r in (r_long, r_short): + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + +class TestRadixDfsWeight(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + schedule_policy="dfs-weight", + ) + + def test_radix_dfs_weight_policy_chunked(self): + self.server.execute_script(self._script_radix_dfs_weight_policy_chunked) + + @staticmethod + def _script_radix_dfs_weight_policy_chunked(t: ScriptedContext): + warm_a = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=3 + ) + yield from run_until_finished(warm_a) + warm_b = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1, prompt_token=4 + ) + yield from run_until_finished(warm_b) + + heavy = [ + t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=3 + ) + for _ in range(3) + ] + light = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, prompt_token=4 + ) + all_reqs = heavy + [light] + + finish_order: list = [] + cached_tokens_by_rid: dict = {} + for _ in range(DEFAULT_MAX_STEPS * 4): + for r in all_reqs: + if r.req is not None: + cached_tokens_by_rid[r.rid] = r.req.cached_tokens + if r.finished and r.rid not in finish_order: + finish_order.append(r.rid) + if all(r.finished for r in all_reqs): + break + yield + assert all(r.finished for r in all_reqs) + assert finish_order[-1] == light.rid, ( + f"dfs-weight must drain the heavier branch-A subtree before the " + f"lighter branch-B req; finish order was {finish_order!r}, " + f"light={light.rid!r}" + ) + for r in heavy: + assert cached_tokens_by_rid.get(r.rid, 0) > 0 + assert cached_tokens_by_rid.get(light.rid, 0) > 0 + for r in all_reqs: + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + +class TestRadixPriority(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + ) + + def test_radix_prefix_match_with_priority(self): + self.server.execute_script(self._script_radix_prefix_match_with_priority) + + @staticmethod + def _script_radix_prefix_match_with_priority(t: ScriptedContext): + r_warm = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 2, max_new_tokens=1) + yield from run_until_finished(r_warm) + assert r_warm.finished + for _ in range(5): + yield + + r_low = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, priority=0 + ) + r_high = t.start_req( + prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=1, priority=10 + ) + + first_admitted = None + low_done = False + high_done = False + cached_tokens_by_rid: dict = {} + for _ in range(DEFAULT_MAX_STEPS): + comp = t.batch_composition() + active = ( + comp.get("chunked", []) + + comp.get("prefill", []) + + comp.get("running", []) + ) + if first_admitted is None and active: + first_admitted = active[0] + for r in (r_low, r_high): + if r.req is not None: + cached_tokens_by_rid[r.rid] = r.req.cached_tokens + low_done = low_done or r_low.finished + high_done = high_done or r_high.finished + if low_done and high_done: + break + yield + assert low_done and high_done + assert first_admitted == r_high.rid, ( + f"higher-priority req must be admitted first; first admitted rid " + f"was {first_admitted!r}, expected r_high={r_high.rid!r}" + ) + for r in (r_low, r_high): + assert cached_tokens_by_rid.get(r.rid, 0) > 0 + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_radix_calc_priority_skip_chunked_resume(self): + self.server.execute_script(self._script_radix_calc_priority_skip_chunked_resume) + + @staticmethod + def _script_radix_calc_priority_skip_chunked_resume(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, priority=10) + yield from run_until(r1, lambda h: h.is_chunking) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, priority=0) + + prev_chunks_done = r1.chunks_done + r1_fin_step = None + r2_fin_step = None + step = 0 + while not (r1.finished and r2.finished): + assert r1.chunks_done >= prev_chunks_done, ( + f"r1 chunked prefill was preempted by lower-priority r2: " + f"chunks_done regressed {prev_chunks_done} -> {r1.chunks_done}" + ) + prev_chunks_done = r1.chunks_done + if r1.finished and r1_fin_step is None: + r1_fin_step = step + if r2.finished and r2_fin_step is None: + r2_fin_step = step + step += 1 + yield + if r1.finished and r1_fin_step is None: + r1_fin_step = step + if r2.finished and r2_fin_step is None: + r2_fin_step = step + + assert r1.finished and r2.finished + assert r1_fin_step is not None and r2_fin_step is not None + assert r1_fin_step <= r2_fin_step, ( + f"high-priority r1 must finish its chunked prefill no later than " + f"low-priority r2; r1 finished at step {r1_fin_step}, r2 at step " + f"{r2_fin_step} (lower-priority r2 jumped ahead of r1's resume)" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_regression.py b/test/manual/chunked_prefill/test_scripted_regression.py new file mode 100644 index 000000000..cf2ff22a8 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_regression.py @@ -0,0 +1,524 @@ +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, + DEFAULT_MAX_STEPS, + VERY_LONG_PROMPT_LEN, + base_engine_kwargs, + run_until, + run_until_all_finished, + run_until_finished, + warmup_radix, +) + + +def _drain_until_released(t, *handles): + for _ in range(12): + if all( + h.kv_pages == 0 + and h.lock_refs == 0 + and (h.req is None or h.req.req_pool_idx is None) + for h in handles + ): + return + yield + + +class TestRegressionBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_abort_waiting_releases_all(self): + self.server.execute_script(self._script_abort_waiting_releases_all) + + @staticmethod + def _script_abort_waiting_releases_all(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + + t.abort(r) + yield from _drain_until_released(t, r) + + assert r.kv_pages == 0 + assert r.req.req_pool_idx is None + assert r.lock_refs == 0 + assert not r.is_chunking + assert r.req.inflight_middle_chunks == 0 + + def test_pause_covers_waiting_chunked(self): + self.server.execute_script(self._script_pause_covers_waiting_chunked) + + @staticmethod + def _script_pause_covers_waiting_chunked(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + + t.pause_generation(mode="retract") + yield + + assert r.kv_pages == 0 + assert r.req.req_pool_idx is None + assert r.lock_refs == 0 + assert not r.is_chunking + + t.continue_generation() + + def test_inflight_middle_chunks_invariant(self): + self.server.execute_script(self._script_inflight_middle_chunks_invariant) + + @staticmethod + def _script_inflight_middle_chunks_invariant(t: ScriptedContext): + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, ignore_eos=True + ) + + observed_max = 0 + saw_chunking_bump = False + cleared_inflight = False + for _ in range(DEFAULT_MAX_STEPS): + req = r.req + if req is not None: + observed_max = max(observed_max, req.inflight_middle_chunks) + if r.chunks_done >= 1 and r.is_chunking: + saw_chunking_bump = saw_chunking_bump or ( + req.inflight_middle_chunks > 0 + ) + if not r.is_chunking and r.chunks_done >= 2: + cleared_inflight = req.inflight_middle_chunks == 0 + break + if r.finished: + cleared_inflight = True + break + yield + else: + raise AssertionError("chunk loop did not clear within DEFAULT_MAX_STEPS") + + assert saw_chunking_bump, "last-chunk admit must bump inflight_middle_chunks" + assert observed_max == 1, ( + f"e875cd36e4: inflight_middle_chunks must be a 0/1 latch; " + f"observed max={observed_max} (pre-fix bug would bump to 2 " + f"at the last-chunk admit boundary)" + ) + assert ( + cleared_inflight + ), "inflight_middle_chunks should be 0 once the chunk loop clears" + + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + + @unittest.skip( + "dbdcdde245 mamba_pool_idx cleanup-skip is mamba-architecture-specific. " + "The shared test fixture does not configure a mamba model; running this " + "regression against the default transformer model would not exercise the " + "mamba NO_TOKEN cleanup path, so the body would be a pure smoke test " + "with no real protection. Re-enable when a mamba fixture is wired in." + ) + def test_mamba_chunked_resume_no_token(self): + self.server.execute_script(self._script_mamba_chunked_resume_no_token) + + @staticmethod + def _script_mamba_chunked_resume_no_token(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until_finished(r) + assert r.finished + + def test_merge_batch_assert_widened(self): + self.server.execute_script(self._script_merge_batch_assert_widened) + + @staticmethod + def _script_merge_batch_assert_widened(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True) + yield from run_until_all_finished([r1, r2]) + assert r1.finished and r2.finished + assert r1.chunks_done >= 2 and r2.chunks_done == 0 + assert len(r1.req.output_ids) == 2 and len(r2.req.output_ids) == 2 + + def test_chunked_pending_tokens_subtract_prefix(self): + self.server.execute_script(self._script_chunked_pending_tokens_subtract_prefix) + + @staticmethod + def _script_chunked_pending_tokens_subtract_prefix(t: ScriptedContext): + r1 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE * 4, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking and h.chunks_done >= 1) + + r2 = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until(r2, lambda h: h.status == "waiting") + + r1_prefix = len(r1.req.prefix_indices) + assert r1.is_chunking, "r1 must still be the in-flight chunked req" + assert r1_prefix > 0, "r1 must hold a committed prefix as the chunked req" + + observed_pending = t.scheduler.load_inquirer._get_num_pending_tokens() + + r1_seqlen = r1.req.seqlen + r2_seqlen = r2.req.seqlen + expected_post_fix = (r1_seqlen - r1_prefix) + r2_seqlen + pre_fix_bad = r1_seqlen + r2_seqlen + + assert observed_pending == expected_post_fix, ( + f"c79a73bec4: load_inquirer must subtract the chunked req's " + f"prefix_indices_len from its pending-token contribution; " + f"observed={observed_pending}, expected_post_fix=" + f"{expected_post_fix}, pre_fix_bad={pre_fix_bad}" + ) + assert observed_pending < pre_fix_bad, ( + f"c79a73bec4: observed pending tokens did not subtract the chunked " + f"prefix (matches the no-subtraction sum {pre_fix_bad}) — fix is regressed" + ) + + yield from run_until_all_finished([r1, r2]) + assert r1.finished and r2.finished + + def test_chunked_admission_reuse_branch_balanced(self): + self.server.execute_script(self._script_chunked_admission_reuse_branch_balanced) + + @staticmethod + def _script_chunked_admission_reuse_branch_balanced(t: ScriptedContext): + baseline_refs = sum(t.get_all_node_lock_refs().values()) + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + + for target_chunk in (3, 4, 5): + yield from run_until( + r, lambda h: h.chunks_done >= target_chunk and h.is_chunking + ) + assert r.lock_refs == 1, ( + f"reuse branch re-acquired lock_ref on the chunked req's node at " + f"chunk {target_chunk}: r.lock_refs={r.lock_refs} (expected 1)" + ) + per_node = t.get_all_node_lock_refs() + assert max(per_node.values(), default=0) <= 1, ( + f"reuse branch double-locked a single node at chunk " + f"{target_chunk}: per-node lock_refs={per_node}" + ) + + yield from run_until_finished(r) + final_refs = sum(t.get_all_node_lock_refs().values()) + assert final_refs == baseline_refs, ( + f"chunked lifecycle must net to zero lock_ref delta; " + f"baseline={baseline_refs}, final={final_refs}" + ) + assert r.lock_refs == 0 + + def test_multiturn_full_hit_no_reuse_branch(self): + self.server.execute_script(self._script_multiturn_full_hit_no_reuse_branch) + + @staticmethod + def _script_multiturn_full_hit_no_reuse_branch(t: ScriptedContext): + baseline_refs = sum(t.get_all_node_lock_refs().values()) + r1 = t.start_req(prompt_len=64, max_new_tokens=2) + yield from run_until_finished(r1) + + r2 = t.start_req(prompt_len=64, max_new_tokens=2) + yield from run_until_finished(r2) + + assert not r2.is_chunking + assert r2.req.cached_tokens > 0, ( + f"follow-up req must fully hit the warm prefix to exercise the " + f"no-reuse path; got cached_tokens={r2.req.cached_tokens}" + ) + for _ in range(5): + yield + assert sum(t.get_all_node_lock_refs().values()) == baseline_refs, ( + f"full-prefix-hit follow-up must not take the chunked reuse " + f"branch; lock_refs drifted from {baseline_refs} to " + f"{sum(t.get_all_node_lock_refs().values())}" + ) + + def test_abort_chunked_resume_releases_all_resources(self): + self.server.execute_script( + self._script_abort_chunked_resume_releases_all_resources + ) + + @staticmethod + def _script_abort_chunked_resume_releases_all_resources(t: ScriptedContext): + baseline_refs = sum(t.get_all_node_lock_refs().values()) + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + assert r.req.req_pool_idx is not None, "row must be held mid-chunk" + assert r.kv_pages > 0, "committed KV must be held mid-chunk" + assert r.lock_refs >= 1, "radix lock_ref must be held mid-chunk" + + t.abort(r) + yield from _drain_until_released(t, r) + + assert ( + r.req.req_pool_idx is None + ), f"96d4749094: abort must release row; got row_idx={r.req.req_pool_idx!r}" + assert ( + r.kv_pages == 0 + ), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}" + assert ( + r.lock_refs == 0 + ), f"96d4749094: abort must release lock_ref; got lock_refs={r.lock_refs}" + assert not r.is_chunking + assert r.req.inflight_middle_chunks == 0 + assert sum(t.get_all_node_lock_refs().values()) == baseline_refs + + def test_pause_retract_releases_waiting_chunked_resume(self): + self.server.execute_script( + self._script_pause_retract_releases_waiting_chunked_resume + ) + + @staticmethod + def _script_pause_retract_releases_waiting_chunked_resume(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + assert r.req.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1 + + t.pause_generation(mode="retract") + yield + + assert r.req.req_pool_idx is None, ( + f"f38e69f87d: pause(retract) must release waiting " + f"chunked-resume row; got row_idx={r.req.req_pool_idx!r}" + ) + assert r.kv_pages == 0 + assert r.lock_refs == 0 + assert not r.is_chunking + assert r.status == "waiting", ( + f"f38e69f87d: pause(retract) must re-queue the retracted " + f"chunked-resume req; got status={r.status!r}" + ) + assert t.scheduler.chunked_req is None + assert t.scheduler.running_batch.is_empty() + + t.continue_generation() + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2 + + def test_retract_all_clears_batch_with_chunked(self): + self.server.execute_script(self._script_retract_all_clears_batch_with_chunked) + + @staticmethod + def _script_retract_all_clears_batch_with_chunked(t: ScriptedContext): + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + r2 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking and h.chunks_done >= 1) + yield from run_until(r2, lambda h: h.status == "waiting") + + t.pause_generation(mode="retract") + yield + + assert len(t.scheduler.running_batch.reqs) == 0, ( + f"f0388931bf: retract_all must clear batch; got batch_size=" + f"{len(t.scheduler.running_batch.reqs)}" + ) + assert ( + t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None + ) is None + for r in (r1, r2): + assert r.status == "waiting" + assert r.kv_pages == 0 + assert not r.is_chunking + + t.continue_generation() + yield from run_until_all_finished([r1, r2]) + + +class TestRegressionPp(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + tp_size=2, + pp_size=2, + ) + + def test_pp_abort_dedup(self): + self.server.execute_script(self._script_pp_abort_dedup) + + @staticmethod + def _script_pp_abort_dedup(t: ScriptedContext): + r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4) + + yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking) + + t.abort(r) + yield + + rids_after_abort = [req.rid for req in t.scheduler.running_batch.reqs] + occurrences = sum(1 for rid in rids_after_abort if rid == r.rid) + assert occurrences <= 1, ( + f"b823c16e60: batch_rids must dedup across mbs + " + f"waiting_queue; got {occurrences} occurrences of rid=" + f"{r.rid} (pre-fix bug would yield 3)" + ) + yield from _drain_until_released(t, r) + assert r.finished + + +class TestRegressionPriority(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_priority_scheduling=True, + ) + + def test_priority_skips_chunked_in_prefix_match(self): + self.server.execute_script(self._script_priority_skips_chunked_in_prefix_match) + + @staticmethod + def _script_priority_skips_chunked_in_prefix_match(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + priority=0, + ignore_eos=True, + ) + yield from run_until(r1, lambda h: h.is_chunking and h.chunks_done >= 1) + r1_host_hit_before = r1.req.host_hit_length + + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + priority=10, + ignore_eos=True, + ) + prev_chunks = r1.chunks_done + while not (r1.finished and r2.finished): + assert r1.req is None or r1.req.host_hit_length == r1_host_hit_before, ( + f"aaf3752d2b: priority calc re-matched the chunked-resume req; " + f"host_hit_length changed {r1_host_hit_before} -> " + f"{r1.req.host_hit_length}" + ) + assert r1.chunks_done >= prev_chunks, ( + f"r1 chunked prefill was preempted by higher-priority r2: " + f"chunks_done regressed {prev_chunks} -> {r1.chunks_done}" + ) + prev_chunks = r1.chunks_done + yield + + assert r1.finished and r2.finished + assert r1.chunks_done >= 2 + assert len(r1.req.output_ids) == 2 and len(r2.req.output_ids) == 2 + + +class TestRegressionLpm(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + schedule_policy="lpm", + ) + + def test_chunked_resume_priority_in_sort(self): + self.server.execute_script(self._script_chunked_resume_priority_in_sort) + + @staticmethod + def _script_chunked_resume_priority_in_sort(t: ScriptedContext): + r_long = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r_long, lambda h: h.is_chunking) + + shorts = [t.start_req(prompt_len=4, max_new_tokens=1) for _ in range(8)] + + initial_chunks = r_long.chunks_done + for _ in range(200): + if r_long.chunks_done > initial_chunks: + break + yield + else: + raise AssertionError( + f"chunked req starved by short req flood; " + f"chunks_done stuck at {initial_chunks}" + ) + + all_reqs = [r_long, *shorts] + done = {r.rid: False for r in all_reqs} + for _ in range(DEFAULT_MAX_STEPS): + for r in all_reqs: + done[r.rid] = done[r.rid] or r.finished + if all(done.values()): + break + yield + else: + raise AssertionError( + f"reqs did not all finish; done={done}, " + f"long_chunks={r_long.chunks_done}" + ) + assert r_long.finished + assert r_long.chunks_done >= 2 + + def test_lpm_skips_chunked_resume_prefix_match(self): + self.server.execute_script(self._script_lpm_skips_chunked_resume_prefix_match) + + @staticmethod + def _script_lpm_skips_chunked_resume_prefix_match(t: ScriptedContext): + yield from warmup_radix(t, [1] * (2 * DEFAULT_CHUNK_SIZE)) + + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.chunks_done >= 1 and h.is_chunking) + + host_hit_before = r1.req.host_hit_length + + r2 = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + while not (r1.finished and r2.finished): + assert r1.req is None or r1.req.host_hit_length == host_hit_before, ( + f"calc_priority re-matched the chunked-resume req; host_hit_length " + f"changed {host_hit_before} -> {r1.req.host_hit_length}" + ) + yield + + def test_chunked_resume_priority_under_lpm(self): + self.server.execute_script(self._script_chunked_resume_priority_under_lpm) + + @staticmethod + def _script_chunked_resume_priority_under_lpm(t: ScriptedContext): + long_prefix_tokens = [1] * (3 * DEFAULT_CHUNK_SIZE) + yield from warmup_radix(t, long_prefix_tokens) + + r1 = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r1, lambda h: h.is_chunking and h.chunks_done >= 1) + + competitors = [ + t.start_req(prompt_len=3 * DEFAULT_CHUNK_SIZE + 32, max_new_tokens=2) + for _ in range(6) + ] + + baseline_chunks = r1.chunks_done + for _ in range(50): + if r1.chunks_done > baseline_chunks: + break + yield + else: + raise AssertionError( + f"bf5b4e9a10: chunked-resume starved under LPM by " + f"long-prefix competitors; chunks_done stuck at " + f"{baseline_chunks}" + ) + + yield from run_until_all_finished([r1, *competitors]) + + +class TestRegressionGptOss(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + model_path="openai/gpt-oss-20b", + mem_fraction_static=0.70, + disable_piecewise_cuda_graph=True, + ) + + def test_chunked_stash_bounded_by_kv_committed_len(self): + self.server.execute_script( + self._script_chunked_stash_bounded_by_kv_committed_len + ) + + @staticmethod + def _script_chunked_stash_bounded_by_kv_committed_len(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + committed = r.req.kv_committed_len + assert committed > 0 + + assert len(r.req.prefix_indices) <= committed, ( + f"cache_unfinished_req over-read past kv_committed_len: " + f"prefix_indices_len={len(r.req.prefix_indices)}, " + f"kv_committed_len={committed}" + ) + yield from run_until_finished(r) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_sampling.py b/test/manual/chunked_prefill/test_scripted_sampling.py new file mode 100644 index 000000000..11ccee95e --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_sampling.py @@ -0,0 +1,274 @@ +import unittest + +from sglang.srt.managers.schedule_batch import FINISH_LENGTH, FINISH_MATCHED_TOKEN +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 TestSamplingBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_max_new_tokens_zero_prefill_only(self): + self.server.execute_script(self._script_max_new_tokens_zero_prefill_only) + + @staticmethod + def _script_max_new_tokens_zero_prefill_only(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=0) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 1, ( + f"max_new_tokens=0 finishes on the prefill chunk with one sampled " + f"token; got {len(r.req.output_ids)}" + ) + decode_records = [ + rec + for rec in t._scheduler_hook._batch_log + if r.rid in rec.rids and rec.mode == "decode" + ] + assert len(decode_records) == 0, ( + f"max_new_tokens=0 must run zero decode forwards; got " + f"{len(decode_records)}" + ) + + def test_max_new_tokens_one_long_chunked(self): + self.server.execute_script(self._script_max_new_tokens_one_long_chunked) + + @staticmethod + def _script_max_new_tokens_one_long_chunked(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=1) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 1, ( + f"max_new_tokens=1 must produce exactly 1 token, got " + f"{len(r.req.output_ids)}" + ) + + def test_max_new_tokens_1000_long_chunked(self): + self.server.execute_script(self._script_max_new_tokens_1000_long_chunked) + + @staticmethod + def _script_max_new_tokens_1000_long_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=1000, + ignore_eos=True, + ) + yield from run_until(r, lambda h: h.finished, max_steps=2000) + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 1000, ( + f"ignore_eos=True + max_new_tokens=1000 must produce 1000 " + f"output tokens; got {len(r.req.output_ids)}" + ) + + def test_return_logprob_chunked(self): + self.server.execute_script(self._script_return_logprob_chunked) + + @staticmethod + def _script_return_logprob_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + return_logprob=True, + ignore_eos=True, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert r.req.logprob is not None + assert len(r.req.logprob.output_token_logprobs_val) == 4 + + def test_ignore_eos_chunked(self): + self.server.execute_script(self._script_ignore_eos_chunked) + + @staticmethod + def _script_ignore_eos_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=16, ignore_eos=True + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 16 + assert isinstance(r.req.finished_reason, FINISH_LENGTH), ( + f"ignore_eos=True must finish via length cap; got " + f"{r.req.finished_reason!r}" + ) + + def test_return_logprob_top_logprobs_chunked(self): + self.server.execute_script(self._script_return_logprob_top_logprobs_chunked) + + @staticmethod + def _script_return_logprob_top_logprobs_chunked(t: ScriptedContext): + top_k = 5 + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + return_logprob=True, + top_logprobs_num=top_k, + ignore_eos=True, + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + assert r.req.logprob is not None + top = r.req.logprob.output_top_logprobs_val + assert len(top) == 4, ( + f"top logprobs must be reported once per output token; " + f"got {len(top)} entries for 4 tokens" + ) + for step_entries in top: + assert len(step_entries) == top_k, ( + f"each step must carry exactly top_logprobs_num={top_k} " + f"entries; got {len(step_entries)}" + ) + + def test_explicit_rid_chunked(self): + self.server.execute_script(self._script_explicit_rid_chunked) + + @staticmethod + def _script_explicit_rid_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + rid="custom-rid-1", + ignore_eos=True, + ) + yield from run_until_finished(r) + assert r.rid == "custom-rid-1" + assert r.finished + assert r.chunks_done >= 2 + assert len(r.req.output_ids) == 2 + + def test_default_sampling_short(self): + self.server.execute_script(self._script_default_sampling_short) + + @staticmethod + def _script_default_sampling_short(t: ScriptedContext): + r = t.start_req(prompt_len=8, max_new_tokens=2, ignore_eos=True) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done == 0 + assert len(r.req.output_ids) == 2 + + def test_chunked_logprob_input_accumulates_across_chunks(self): + self.server.execute_script( + self._script_chunked_logprob_input_accumulates_across_chunks + ) + + @staticmethod + def _script_chunked_logprob_input_accumulates_across_chunks(t: ScriptedContext): + prompt_len = VERY_LONG_PROMPT_LEN + r = t.start_req( + prompt_len=prompt_len, + max_new_tokens=4, + return_logprob=True, + logprob_start_len=0, + ) + yield from run_until_finished(r) + assert r.finished + assert ( + r.chunks_done >= 2 + ), f"prompt should span multiple chunks, got chunks_done={r.chunks_done}" + assert r.req.logprob is not None + input_lp = r.req.logprob.input_token_logprobs_val + assert len(input_lp) == prompt_len, ( + f"expected {prompt_len} input logprobs (one per prompt token), " + f"got {len(input_lp)}" + ) + + def test_logprob_start_len_inside_chunk_2(self): + self.server.execute_script(self._script_logprob_start_len_inside_chunk_2) + + @staticmethod + def _script_logprob_start_len_inside_chunk_2(t: ScriptedContext): + prompt_len = 4 * DEFAULT_CHUNK_SIZE + start_len = DEFAULT_CHUNK_SIZE + 50 + r = t.start_req( + prompt_len=prompt_len, + max_new_tokens=4, + return_logprob=True, + logprob_start_len=start_len, + ) + yield from run_until_finished(r) + assert r.finished + assert ( + r.chunks_done >= 3 + ), f"prompt should span 3+ chunks, got chunks_done={r.chunks_done}" + assert r.req.logprob is not None + input_lp = r.req.logprob.input_token_logprobs_val + assert len(input_lp) == prompt_len - start_len, ( + f"expected {prompt_len - start_len} input logprobs for tokens " + f">= logprob_start_len={start_len}, got {len(input_lp)}" + ) + + def test_finish_reason_value_eos_vs_length_chunked(self): + self.server.execute_script( + self._script_finish_reason_value_eos_vs_length_chunked + ) + + @staticmethod + def _script_finish_reason_value_eos_vs_length_chunked(t: ScriptedContext): + probe = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=1, + ignore_eos=True, + prompt_token=7, + temperature=0.0, + ) + yield from run_until_finished(probe) + assert probe.finished + first_token = probe.req.output_ids[0] + + for _ in range(40): + if t.is_fully_idle: + break + yield + t.flush_cache() + yield + + r_eos = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=999, + ignore_eos=False, + prompt_token=7, + stop_token_ids=[first_token], + temperature=0.0, + ) + yield from run_until_finished(r_eos, max_steps=2000) + assert r_eos.finished + assert ( + r_eos.chunks_done >= 2 + ), f"scenario 1 should chunk; got chunks_done={r_eos.chunks_done}" + assert isinstance(r_eos.req.finished_reason, FINISH_MATCHED_TOKEN), ( + f"a stop token the model deterministically produces under greedy must " + f"finish via the matched-token path; got {r_eos.req.finished_reason!r}" + ) + + r_length = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=4, + ignore_eos=True, + ) + yield from run_until_finished(r_length) + assert r_length.finished + assert ( + r_length.chunks_done >= 2 + ), f"scenario 2 should chunk; got chunks_done={r_length.chunks_done}" + assert isinstance(r_length.req.finished_reason, FINISH_LENGTH), ( + f"ignore_eos=True + max_new_tokens=4 chunked must finish via " + f"length cap; got {r_length.req.finished_reason!r}" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_spec.py b/test/manual/chunked_prefill/test_scripted_spec.py new file mode 100644 index 000000000..a4b842791 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_spec.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, +) + +_SPEC_MODEL = "Qwen/Qwen3-8B" +_SPEC_DRAFT = "Tengyunw/qwen3_8b_eagle3" + + +def _spec_engine_kwargs(**overrides): + return base_engine_kwargs( + model_path=_SPEC_MODEL, + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + speculative_algorithm="EAGLE3", + speculative_draft_model_path=_SPEC_DRAFT, + speculative_num_steps=6, + speculative_eagle_topk=10, + speculative_num_draft_tokens=32, + kv_canary="none", + kv_canary_real_data="none", + kv_canary_sweep_interval=0, + **overrides, + ) + + +class TestSpecBasic(ScriptedTestCase): + ENGINE_KWARGS = _spec_engine_kwargs() + + def test_spec_chunked_handoff_first_verify(self): + self.server.execute_script(self._script_spec_chunked_handoff_first_verify) + + @staticmethod + def _script_spec_chunked_handoff_first_verify(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=16) + yield from run_until_finished(r, max_steps=800) + assert r.finished + assert r.chunks_done >= 2 + assert r.req.spec_verify_ct >= 1, ( + f"expected >=1 spec verify after chunked handoff, got " + f"{r.req.spec_verify_ct}" + ) + + def test_spec_abort_during_chunked_prepare(self): + self.server.execute_script(self._script_spec_abort_during_chunked_prepare) + + @staticmethod + def _script_spec_abort_during_chunked_prepare(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=16) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + t.abort(r) + for _ in range(40): + if t.is_fully_idle: + break + yield + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/manual/chunked_prefill/test_scripted_special_case.py b/test/manual/chunked_prefill/test_scripted_special_case.py new file mode 100644 index 000000000..d5d4b3c35 --- /dev/null +++ b/test/manual/chunked_prefill/test_scripted_special_case.py @@ -0,0 +1,1253 @@ +import unittest +from typing import Optional + +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, + DEFAULT_MAX_STEPS, + SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + SMALL_KV_POOL_BALLAST_PROMPT_LEN, + SMALL_KV_POOL_MAX_TOTAL_TOKENS, + VERY_LONG_PROMPT_LEN, + advance_to_decode_step, + base_engine_kwargs, + exhaust_row_pool, + run_until, + run_until_finished, +) + + +def _load_inquirer_pending_for_rid(t: ScriptedContext, rid: str) -> int: + s = t.scheduler + chunked = s.chunked_req + if chunked is not None and chunked.rid == rid: + return chunked.seqlen - len(chunked.prefix_indices) + for req in s.waiting_queue: + if req.rid == rid: + return req.seqlen + return 0 + + +class TestSpecialCaseBasic(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_chunked_in_flight_no_idle(self): + self.server.execute_script(self._script_chunked_in_flight_no_idle) + + @staticmethod + def _script_chunked_in_flight_no_idle(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100 + ) + yield from run_until(r, lambda h: h.is_chunking) + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + saw_chunking = True + assert ( + not t.is_idle + ), "scheduler must not idle while chunked_req is in flight" + if r.finished: + break + yield + assert r.finished + assert saw_chunking, "test must observe r mid-chunk at least once" + + def test_add_chunked_req_path(self): + self.server.execute_script(self._script_add_chunked_req_path) + + @staticmethod + def _script_add_chunked_req_path(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=110 + ) + yield from run_until_finished(r) + assert r.finished + assert r.chunks_done >= 2 + + def test_admission_with_chunked_in_flight(self): + self.server.execute_script(self._script_admission_with_chunked_in_flight) + + @staticmethod + def _script_admission_with_chunked_in_flight(t: ScriptedContext): + r_chunk = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=120 + ) + yield from run_until(r_chunk, lambda h: h.is_chunking) + + r_small = t.start_req(prompt_len=4, max_new_tokens=2) + yield + + comp = t.batch_composition() + assert r_chunk.rid in comp.get("chunked", []) + + yield from run_until_finished(r_small) + yield from run_until_finished(r_chunk) + assert r_small.finished and r_chunk.finished + + def test_abort_excludes_chunked_req(self): + self.server.execute_script(self._script_abort_excludes_chunked_req) + + @staticmethod + def _script_abort_excludes_chunked_req(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=130 + ) + yield from run_until(r, lambda h: h.is_chunking) + + t.abort(r) + for _ in range(12): + if t.scheduler.chunked_req is None and r.kv_pages == 0 and r.lock_refs == 0: + break + yield + + assert ( + t.scheduler.chunked_req is None + ), f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}" + assert r.kv_pages == 0 + assert r.lock_refs == 0 + + def test_get_chunked_req_lambda_getter(self): + self.server.execute_script(self._script_get_chunked_req_lambda_getter) + + @staticmethod + def _script_get_chunked_req_lambda_getter(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=140 + ) + yield from run_until(r, lambda h: h.is_chunking) + saw_match = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + cur = ( + t.scheduler.chunked_req.rid + if t.scheduler.chunked_req is not None + else None + ) + assert cur in (None, r.rid), ( + f"getter returned unrelated rid: got {cur!r}, expected " + f"None or {r.rid!r}" + ) + if cur == r.rid: + saw_match = True + if r.finished: + break + yield + assert r.finished + assert saw_match, "getter must return r.rid at least once while r.is_chunking" + assert ( + t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None + ) is None + + @unittest.skip( + "requires real disaggregation prefill/decode split — single-engine " + "ScriptedContext cannot exercise the decode-side waiting_queue " + "KV-hold path. Belongs in test_scripted_disagg.py with D3 counter " + "wiring once disagg topology is available." + ) + def test_disagg_decode_waiting_queue_kv_held(self): + pass + + @unittest.skip( + "requires DLLM model + staging mixin — single-engine cannot drive " + "both DLLM staging AND chunked admission incrementing " + "inflight_middle_chunks from two sources. Belongs in DLLM-specific " + "test file." + ) + def test_dllm_staging_double_inflight_middle_chunks(self): + pass + + @unittest.skip( + "requires real disaggregation topology — single-engine cannot " + "exercise staging_handler chunked path." + ) + def test_staging_handler_chunked(self): + pass + + @unittest.skip( + "requires mooncake KV transport backend — single-engine cannot " + "drive the conn layer chunked path." + ) + def test_mooncake_conn_chunked(self): + pass + + @unittest.skip( + "requires NIXL KV transport backend — single-engine cannot drive " + "the NIXL conn layer chunked path." + ) + def test_nixl_conn_chunked(self): + pass + + def test_filter_batch_exclude_chunked_flag(self): + self.server.execute_script(self._script_filter_batch_exclude_chunked_flag) + + @staticmethod + def _script_filter_batch_exclude_chunked_flag(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=150 + ) + r2 = t.start_req(prompt_len=16, max_new_tokens=2) + saw_r1_chunking = False + for _ in range(DEFAULT_MAX_STEPS * 2): + if r1.is_chunking: + saw_r1_chunking = True + comp = t.batch_composition() + assert r1.rid in comp.get( + "chunked", [] + ), f"mid-chunk r1 must occupy the chunked role; got {comp!r}" + assert r1.rid not in comp.get( + "running", [] + ), f"chunked r1 must be excluded from the running role; got {comp!r}" + if r1.finished and r2.finished: + break + yield + assert r1.finished and r2.finished + assert ( + saw_r1_chunking + ), "r1 must have chunked at some point to exercise the exclude branch" + + @unittest.skip( + "pdmux split_prefill_batch requires the pdmux topology — " + "single-engine ScriptedContext cannot drive the split path. " + "Belongs in a pdmux-specific test once that lane is wired up." + ) + def test_pdmux_split_prefill_batch(self): + pass + + def test_streaming_session_kv_committed_bound(self): + self.server.execute_script(self._script_streaming_session_kv_committed_bound) + + @staticmethod + def _script_streaming_session_kv_committed_bound(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=160 + ) + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + assert len(r.req.prefix_indices) <= r.req.kv_committed_len, ( + f"streaming-session chunked stash must stay bounded by " + f"kv_committed_len; prefix_indices_len={len(r.req.prefix_indices)}, " + f"kv_committed_len={r.req.kv_committed_len}" + ) + if r.finished: + break + yield + assert r.finished + + @unittest.skip( + "mamba_pool_idx cleanup applies only to mamba-class models — " + "single-engine with a non-mamba model cannot drive the NO_TOKEN " + "chunked-resume cleanup-skip branch. Belongs in a mamba-specific " + "test file once mamba scripted coverage is added." + ) + def test_mamba_pool_idx_cleanup_skip_chunked_resume(self): + pass + + def test_pause_retract_clears_chunked_req(self): + self.server.execute_script(self._script_pause_retract_clears_chunked_req) + + @staticmethod + def _script_pause_retract_clears_chunked_req(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=170 + ) + yield from run_until(r, lambda h: h.is_chunking) + + t.pause_generation(mode="retract") + yield + + assert ( + t.scheduler.chunked_req is None + ), f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}" + assert not r.finished, "retract must re-queue r, not finish or abort it" + assert r.status == "waiting", ( + f"retracted chunked req must return to the waiting queue; " + f"got status={r.status!r}" + ) + + t.continue_generation() + yield from run_until_finished(r) + assert ( + r.finished + ), "continue_generation must drive the re-queued req to completion" + + def test_retract_during_gap_inflight_middle_chunks_positive(self): + self.server.execute_script( + self._script_retract_during_gap_inflight_middle_chunks_positive + ) + + @staticmethod + def _script_retract_during_gap_inflight_middle_chunks_positive(t: ScriptedContext): + r = t.start_req( + prompt_len=3 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=180 + ) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + assert r.req.inflight_middle_chunks > 0 + assert r.is_chunking + + t.pause_generation(mode="retract") + yield + + assert r.kv_pages == 0, f"retract must release KV; got kv_pages={r.kv_pages}" + assert not r.finished, "retract must re-queue r, not finish or abort it" + assert not r.is_chunking, "retract must release the chunked slot" + assert r.status == "waiting", ( + f"retracted chunked req must return to the waiting queue; " + f"got status={r.status!r}" + ) + req = t.find_req_by_rid(r.rid) + assert req is not None and req.inflight_middle_chunks == 0, ( + f"retract must reset inflight_middle_chunks; got " + f"{req.inflight_middle_chunks if req is not None else None}" + ) + + t.continue_generation() + yield from run_until_finished(r, max_steps=2000) + assert ( + r.finished + ), "continue_generation must drive the re-queued req to completion" + assert r.kv_pages == 0 + assert len(r.req.output_ids) == 2 + + def test_load_inquirer_pending_tokens_dedup_chunked(self): + self.server.execute_script( + self._script_load_inquirer_pending_tokens_dedup_chunked + ) + + @staticmethod + def _script_load_inquirer_pending_tokens_dedup_chunked(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=190 + ) + yield from run_until(r, lambda h: h.is_chunking) + saw_chunking = False + saw_dedup = False + for _ in range(DEFAULT_MAX_STEPS): + chunked = t.scheduler.chunked_req + if r.is_chunking and chunked is not None and chunked.rid == r.rid: + saw_chunking = True + pending = _load_inquirer_pending_for_rid(t, r.rid) + expected = chunked.seqlen - len(chunked.prefix_indices) + assert pending == expected, ( + f"load_inquirer chunked contribution must equal the prefix-" + f"subtracting formula seqlen - len(prefix_indices) = {expected}; " + f"got {pending}" + ) + assert pending <= chunked.seqlen, ( + f"chunked contribution must never exceed its full seqlen " + f"{chunked.seqlen}; got {pending} — dual-queue dedup violated" + ) + if len(chunked.prefix_indices) > 0: + saw_dedup = pending < chunked.seqlen + if r.finished: + break + yield + assert r.finished + assert ( + saw_chunking + ), "test must observe the dual-queue chunked state at least once" + assert saw_dedup, ( + "test must observe the chunked req with a committed prefix so the " + "dedup subtraction is actually exercised" + ) + + def test_load_inquirer_chunked_contribution_exact_remainder(self): + self.server.execute_script( + self._script_load_inquirer_chunked_contribution_exact_remainder + ) + + @staticmethod + def _script_load_inquirer_chunked_contribution_exact_remainder(t: ScriptedContext): + s = t.scheduler + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=200 + ) + yield from run_until(r, lambda h: h.is_chunking) + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + chunked = s.chunked_req + if r.is_chunking and chunked is not None and chunked.rid == r.rid: + assert len(s.waiting_queue) == 0, ( + "test requires an empty waiting_queue so the chunked req is " + f"the sole pending-token contributor; got {len(s.waiting_queue)}" + ) + expected = chunked.seqlen - len(chunked.prefix_indices) + observed = s.load_inquirer._get_num_pending_tokens() + assert observed == expected, ( + f"chunked contribution must equal remainder " + f"seqlen - len(prefix_indices) = {expected}, got {observed}; " + f"a value of {chunked.seqlen} would mean the committed prefix " + "is being double-counted" + ) + saw_chunking = True + if r.finished: + break + yield + assert r.finished + assert saw_chunking, "test must observe the chunked req mid-flight" + + def test_load_inquirer_chunk_deduct_subtracts_planned_chunk(self): + self.server.execute_script( + self._script_load_inquirer_chunk_deduct_subtracts_planned_chunk + ) + + @staticmethod + def _script_load_inquirer_chunk_deduct_subtracts_planned_chunk(t: ScriptedContext): + s = t.scheduler + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=210 + ) + yield from run_until(r, lambda h: h.is_chunking) + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + chunked = s.chunked_req + if ( + r.is_chunking + and chunked is not None + and chunked.rid == r.rid + and chunked.extend_input_len > 0 + ): + deduct = chunked.extend_input_len + base = s.load_inquirer._get_num_pending_tokens() + deducted = s.load_inquirer._get_num_pending_tokens(chunk_deduct=deduct) + assert deducted == base - deduct, ( + f"chunk_deduct must subtract the planned chunk exactly: " + f"base={base}, deducted={deducted}, expected={base - deduct}" + ) + saw_chunking = True + if r.finished: + break + yield + assert r.finished + assert saw_chunking, "test must observe the chunked req with a planned chunk" + + def test_stage_a_inflight_middle_chunks_sync_invariant(self): + self.server.execute_script( + self._script_stage_a_inflight_middle_chunks_sync_invariant + ) + + @staticmethod + def _script_stage_a_inflight_middle_chunks_sync_invariant(t: ScriptedContext): + def assert_invariant() -> None: + req = t.find_req_by_rid(r.rid) + if req is not None and req.inflight_middle_chunks > 0: + assert r.is_chunking, ( + f"invariant violated: inflight_middle_chunks=" + f"{req.inflight_middle_chunks} but is_chunking={r.is_chunking}" + ) + + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, prompt_token=220 + ) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + assert_invariant() + + t.pause_generation(mode="retract") + yield + assert_invariant() + + t.continue_generation() + for _ in range(DEFAULT_MAX_STEPS): + assert_invariant() + if r.finished: + return + yield + raise AssertionError("req did not finish") + + def test_init_next_round_input_resets_chunk_state(self): + self.server.execute_script( + self._script_init_next_round_input_resets_chunk_state + ) + + @staticmethod + def _script_init_next_round_input_resets_chunk_state(t: ScriptedContext): + r = t.start_req( + prompt_len=3 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=230 + ) + yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking) + saw_mid_chunk = False + for _ in range(DEFAULT_MAX_STEPS): + req = r.req + if ( + r.is_chunking + and r.chunks_done >= 1 + and req is not None + and req.extend_input_len is not None + ): + saw_mid_chunk = True + assert ( + len(req.fill_ids) == len(req.prefix_indices) + req.extend_input_len + ), ( + f"init_next_round_input must rebuild fill_ids to the committed " + f"prefix plus the in-flight chunk; " + f"fill_ids_len={len(req.fill_ids)}, " + f"prefix_indices_len={len(req.prefix_indices)}, " + f"extend_input_len={req.extend_input_len}, " + f"chunks_done={r.chunks_done}" + ) + if r.finished: + break + yield + assert r.finished + assert ( + saw_mid_chunk + ), "test must observe the fill_ids reset boundary at least once" + assert r.finished + + def test_chunked_req_slot_cleared_when_chunk_completes(self): + self.server.execute_script( + self._script_chunked_req_slot_cleared_when_chunk_completes + ) + + @staticmethod + def _script_chunked_req_slot_cleared_when_chunk_completes(t: ScriptedContext): + s = t.scheduler + r = t.start_req( + prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=240 + ) + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + saw_chunking = True + if r.finished: + break + yield + assert r.finished + assert saw_chunking, "req should have occupied the chunked_req slot mid-chunk" + assert ( + s.chunked_req is None + ), f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}" + + def test_second_chunked_admit_blocked_when_chunked_req_set(self): + self.server.execute_script( + self._script_second_chunked_admit_blocked_when_chunked_req_set + ) + + @staticmethod + def _script_second_chunked_admit_blocked_when_chunked_req_set(t: ScriptedContext): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10 + ) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=11 + ) + saw_r1_chunking = False + saw_r2_chunking = False + for _ in range(DEFAULT_MAX_STEPS * 2): + both = r1.is_chunking and r2.is_chunking + assert not both, ( + f"only one chunked-resume slot allowed; r1.is_chunking=" + f"{r1.is_chunking}, r2.is_chunking={r2.is_chunking}" + ) + if r1.is_chunking: + saw_r1_chunking = True + if r2.is_chunking: + saw_r2_chunking = True + if r1.finished and r2.finished: + break + yield + assert r1.finished and r2.finished + assert saw_r1_chunking and saw_r2_chunking, ( + f"both reqs should chunk over their lifetime; saw_r1=" + f"{saw_r1_chunking}, saw_r2={saw_r2_chunking}" + ) + + def test_scheduler_continues_with_only_chunked_req_no_waiting(self): + self.server.execute_script( + self._script_scheduler_continues_with_only_chunked_req_no_waiting + ) + + @staticmethod + def _script_scheduler_continues_with_only_chunked_req_no_waiting( + t: ScriptedContext, + ): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=250 + ) + yield from run_until(r, lambda h: h.is_chunking) + prev_chunks_done = r.chunks_done + progressed = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + assert ( + not t.is_idle + ), "scheduler must not go idle while a chunked req is in flight" + cur_chunks_done = r.chunks_done + if cur_chunks_done > prev_chunks_done: + progressed = True + prev_chunks_done = cur_chunks_done + if r.finished: + break + yield + assert r.finished + assert progressed, ( + "chunks_done must keep advancing without any waiter; pre-fix " + "an empty waiting_queue could cause the loop to skip continuation" + ) + + +class TestSpecialCaseRowPoolExhaustion(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + max_running_requests=8, + ) + + def test_chunked_req_bypasses_req_pool_exhaustion(self): + self.server.execute_script( + self._script_chunked_req_bypasses_req_pool_exhaustion + ) + + @staticmethod + def _script_chunked_req_bypasses_req_pool_exhaustion(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + chunks_before_pressure = r.chunks_done + + yield from exhaust_row_pool(t, leave_rows=0) + + progressed_under_pressure = False + for _ in range(DEFAULT_MAX_STEPS * 2): + if r.chunks_done > chunks_before_pressure: + progressed_under_pressure = True + if r.finished: + break + yield + assert r.finished + assert progressed_under_pressure, ( + "chunked req must advance even when get_num_allocatable_reqs " + "returns 0; pre-fix the bypass would block forever" + ) + assert r.kv_pages == 0 + + +class TestSpecialCaseMixedChunk(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_mixed_chunk=True, + ) + + def test_mix_with_running_chunked_plus_decode(self): + self.server.execute_script(self._script_mix_with_running_chunked_plus_decode) + + @staticmethod + def _script_mix_with_running_chunked_plus_decode(t: ScriptedContext): + decodes = [t.start_req(prompt_len=8, max_new_tokens=16) for _ in range(3)] + for d in decodes: + yield from advance_to_decode_step(d, 1) + + r_chunk = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, prompt_token=300 + ) + yield from run_until(r_chunk, lambda h: h.is_chunking) + + comp = t.batch_composition() + assert r_chunk.rid in comp.get("chunked", []) + + saw_mixed_with_decode = False + all_reqs = [r_chunk, *decodes] + for _ in range(DEFAULT_MAX_STEPS * 2): + comp = t.batch_composition() + batch_rids = ( + set(comp.get("prefill", [])) + | set(comp.get("decode", [])) + | set(comp.get("running", [])) + ) + if ( + t.last_batch_forward_mode == "MIXED" + and r_chunk.rid in comp.get("chunked", []) + and any(d.rid in batch_rids for d in decodes) + ): + saw_mixed_with_decode = True + if all(x.finished for x in all_reqs): + break + yield + assert all(x.finished for x in all_reqs) + assert saw_mixed_with_decode, ( + "enable_mixed_chunk should merge running decode reqs into the chunked " + "prefill iter (MIXED forward_mode) at least once" + ) + + def test_mixed_chunk_with_logprob_falls_back(self): + self.server.execute_script(self._script_mixed_chunk_with_logprob_falls_back) + + @staticmethod + def _script_mixed_chunk_with_logprob_falls_back(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, + max_new_tokens=2, + return_logprob=True, + prompt_token=310, + ) + yield from run_until(r, lambda h: h.is_chunking) + assert ( + t.last_batch_forward_mode != "MIXED" + ), f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}" + yield from run_until_finished(r) + + def test_mixed_chunk_with_running_batch(self): + self.server.execute_script(self._script_mixed_chunk_with_running_batch) + + @staticmethod + def _script_mixed_chunk_with_running_batch(t: ScriptedContext): + r_dec = t.start_req(prompt_len=8, max_new_tokens=32) + yield from run_until(r_dec, lambda h: h.status == "running") + + r_chunk = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=320 + ) + yield + yield from run_until(r_chunk, lambda h: h.is_chunking) + + assert ( + t.last_batch_forward_mode == "MIXED" + ), f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}" + for _ in range(DEFAULT_MAX_STEPS * 2): + if r_chunk.finished and r_dec.finished: + break + yield + assert r_chunk.finished and r_dec.finished + + +class TestSpecialCaseTransformers(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + model_impl="transformers", + ) + + def test_transformers_text_model_still_chunks(self): + self.server.execute_script(self._script_transformers_text_model_still_chunks) + + @staticmethod + def _script_transformers_text_model_still_chunks(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10 + ) + yield from run_until_finished(r) + assert r.finished + expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE) + assert r.chunks_done == expected, ( + f"a text model on the Transformers backend must chunk like the native " + f"backend ({expected} chunks); got chunks_done={r.chunks_done}" + ) + + +class TestSpecialCaseNoChunking(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=-1) + + def test_chunk_size_negative_disables_chunking(self): + self.server.execute_script(self._script_chunk_size_negative_disables_chunking) + + @staticmethod + def _script_chunk_size_negative_disables_chunking(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + for _ in range(DEFAULT_MAX_STEPS): + assert ( + not r.is_chunking + ), "chunked_prefill_size=-1 should disable chunked path" + if r.finished: + return + yield + raise AssertionError("req did not finish under disabled chunking") + + +DETERMINISTIC_ALIGN_SIZE = 4096 + + +class TestSpecialCaseDeterministicFlashInfer(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DETERMINISTIC_ALIGN_SIZE, + page_size=16, + attention_backend="flashinfer", + enable_deterministic_inference=True, + kv_canary="none", + kv_canary_sweep_interval=0, + ) + + def test_chunked_truncation_align_size(self): + self.server.execute_script(self._script_chunked_truncation_align_size) + + @staticmethod + def _script_chunked_truncation_align_size(t: ScriptedContext): + r = t.start_req( + prompt_len=DETERMINISTIC_ALIGN_SIZE + 1024, + max_new_tokens=2, + prompt_token=10, + ) + page_size = 16 + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking and r.req.extend_input_len is not None: + saw_chunking = True + assert r.req.extend_input_len % page_size == 0, ( + f"deterministic chunk boundary must be page-aligned; " + f"got extend_input_len={r.req.extend_input_len}, page_size={page_size}" + ) + if r.finished: + break + yield + assert r.finished, "chunked req did not finish" + assert saw_chunking, "test must observe the req mid-chunk at least once" + + +class TestSpecialCaseHiCache(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_hierarchical_cache=True, + kv_canary="none", + kv_canary_sweep_interval=0, + ) + + def test_hicache_breakdown_only_first_chunk(self): + self.server.execute_script(self._script_hicache_breakdown_only_first_chunk) + + @staticmethod + def _script_hicache_breakdown_only_first_chunk(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=400 + ) + first_chunk_snap = None + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + if r.is_chunking: + saw_chunking = True + if r.is_chunking and r.chunks_done >= 1 and first_chunk_snap is None: + first_chunk_snap = r.req.cached_tokens + if first_chunk_snap is not None and r.is_chunking: + cur = r.req.cached_tokens + assert cur == first_chunk_snap, ( + f"HiCache cached_tokens_* must freeze after first chunk; " + f"first={first_chunk_snap!r}, now={cur!r}" + ) + if r.finished: + break + yield + assert r.finished + assert saw_chunking, "test must observe r mid-chunk at least once" + assert ( + first_chunk_snap is not None + ), "test must snapshot cached_tokens at the first chunk boundary" + + def test_hicache_cached_tokens_set_once_invariant(self): + self.server.execute_script( + self._script_hicache_cached_tokens_set_once_invariant + ) + + @staticmethod + def _script_hicache_cached_tokens_set_once_invariant(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=410 + ) + snap: Optional[tuple] = None + saw_chunking = False + for _ in range(DEFAULT_MAX_STEPS): + req = t.find_req_by_rid(r.rid) + if req is not None and r.is_chunking and r.chunks_done >= 1: + saw_chunking = True + cur = ( + req.cached_tokens_device, + req.cached_tokens_host, + req.cached_tokens_storage, + ) + if snap is None: + snap = cur + else: + assert cur == snap, ( + f"HiCache cached_tokens_* fields must be set exactly " + f"once on first chunk; values changed: snap={snap}, " + f"cur={cur}" + ) + if r.finished: + break + yield + assert r.finished + assert ( + saw_chunking + ), "test must observe the req mid-chunk (chunks_done >= 1) at least once" + assert snap is not None, "test must snapshot the cached_tokens_* breakdown" + + +def _expected_chunks(prompt_len: int, chunk_size: int) -> int: + if prompt_len <= chunk_size: + return 0 + return (prompt_len + chunk_size - 1) // chunk_size + + +class TestSpecialCaseDynamicChunkingPP1(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + enable_dynamic_chunking=True, + ) + + def test_dynamic_chunking_forced_off_on_pp1(self): + self.server.execute_script(self._script_dynamic_chunking_forced_off_on_pp1) + + @staticmethod + def _script_dynamic_chunking_forced_off_on_pp1(t: ScriptedContext): + assert t.scheduler.enable_dynamic_chunking is False, ( + "pp_size==1 must force enable_dynamic_chunking off even when the " + "server arg is True (the 'and ps.pp_size > 1' conjunct)" + ) + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10 + ) + yield from run_until_finished(r) + assert r.finished + expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE) + assert r.chunks_done == expected, ( + f"uniform chunked_prefill_size must yield exactly {expected} chunks " + f"(no dynamic-size prediction on pp1); got {r.chunks_done}" + ) + + +class TestSpecialCaseSmallPool(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + max_total_tokens=SMALL_KV_POOL_MAX_TOTAL_TOKENS, + ) + + @staticmethod + def _start_ballast(t: ScriptedContext, *, prompt_token: int): + return t.start_req( + prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN, + max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS, + ignore_eos=True, + prompt_token=prompt_token, + ) + + @staticmethod + def _run_force_readd_then_complete( + t: ScriptedContext, *, chunk_token: int, ballast_token: int + ): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=chunk_token + ) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + + ballast = TestSpecialCaseSmallPool._start_ballast(t, prompt_token=ballast_token) + + ballast_retracted = False + for _ in range(2000): + if ballast.status == "waiting": + ballast_retracted = True + if r.finished: + break + yield + assert r.finished, ( + "force-re-added chunked resume must complete once the engine retracts " + f"the ballast; status={r.status} kv_pages={r.kv_pages}" + ) + assert r.kv_pages == 0, f"kv_pages={r.kv_pages}" + assert r.lock_refs == 0, f"lock_refs={r.lock_refs}" + assert ( + ballast_retracted + or ballast.finished + or ballast.status in ("waiting", "finished", "unknown") + ), f"ballast must be retracted/aborted under pressure; status={ballast.status}" + + t.abort(ballast) + for _ in range(200): + if t.is_fully_idle: + break + yield + + def test_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens(self): + self.server.execute_script( + self._script_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens + ) + + @staticmethod + def _script_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens( + t: ScriptedContext, + ): + yield from TestSpecialCaseSmallPool._run_force_readd_then_complete( + t, chunk_token=700, ballast_token=701 + ) + + def test_chunked_forced_admission_avoids_leak(self): + self.server.execute_script(self._script_chunked_forced_admission_avoids_leak) + + @staticmethod + def _script_chunked_forced_admission_avoids_leak(t: ScriptedContext): + yield from TestSpecialCaseSmallPool._run_force_readd_then_complete( + t, chunk_token=710, ballast_token=711 + ) + + def test_add_chunked_req_non_swa_forced_admit_on_rem_zero(self): + self.server.execute_script( + self._script_add_chunked_req_non_swa_forced_admit_on_rem_zero + ) + + @staticmethod + def _script_add_chunked_req_non_swa_forced_admit_on_rem_zero(t: ScriptedContext): + yield from TestSpecialCaseSmallPool._run_force_readd_then_complete( + t, chunk_token=720, ballast_token=721 + ) + + +class TestSpecialCaseRetractMerge(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_retract_merges_extend_chunk_batch_before_retract_all(self): + self.server.execute_script( + self._script_retract_merges_extend_chunk_batch_before_retract_all + ) + + @staticmethod + def _script_retract_merges_extend_chunk_batch_before_retract_all( + t: ScriptedContext, + ): + s = t.scheduler + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + + assert s.last_batch is not None, "last_batch must be set while chunking" + assert s.last_batch.forward_mode.is_extend(), ( + f"last_batch must be the extend/chunk batch; got " + f"{s.last_batch.forward_mode!r}" + ) + + t.pause_generation(mode="retract") + yield + + assert ( + s.last_batch is None + ), "retract must clear last_batch after merging the extend chunk batch" + assert len(s.running_batch.reqs) == 0, ( + "the merged extend chunk batch must be retracted out of running_batch, " + f"not stranded; got {len(s.running_batch.reqs)} reqs" + ) + assert ( + r.status == "waiting" + ), f"retracted chunked req must return to the waiting queue; got {r.status!r}" + assert r.kv_pages == 0 + + t.continue_generation() + yield from run_until_finished(r) + assert r.finished + assert len(r.req.output_ids) == 2, ( + f"resumed req must emit exactly max_new_tokens; got " + f"{len(r.req.output_ids)}" + ) + + +class TestSpecialCaseChunkBudgetDefer(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_co_submitted_waiter_deferred_when_chunk_budget_zero(self): + self.server.execute_script( + self._script_co_submitted_waiter_deferred_when_chunk_budget_zero + ) + + @staticmethod + def _script_co_submitted_waiter_deferred_when_chunk_budget_zero( + t: ScriptedContext, + ): + r_chunk = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + r_wait = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE + 8, max_new_tokens=2) + yield from run_until(r_chunk, lambda h: h.is_chunking) + + saw_deferred = False + for _ in range(DEFAULT_MAX_STEPS * 2): + comp = t.batch_composition() + chunk_active = r_chunk.rid in comp.get("chunked", []) + wait_idle = ( + r_wait.status == "waiting" + and r_wait.rid not in comp.get("chunked", []) + and r_wait.rid not in comp.get("running", []) + and r_wait.rid not in comp.get("decode", []) + ) + if chunk_active and wait_idle: + saw_deferred = True + assert r_wait.kv_pages == 0, ( + f"deferred waiter (rem_chunk_tokens<=0 -> OTHER) must hold no KV; " + f"got kv_pages={r_wait.kv_pages}" + ) + if r_chunk.finished and r_wait.finished: + break + yield + assert r_chunk.finished and r_wait.finished + assert saw_deferred, ( + "second waiter must be deferred for >=1 iter while the chunk-token " + "budget was driven to 0 by the in-flight chunked req" + ) + + +class TestSpecialCaseIgnoreEosNoRadix(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs( + chunked_prefill_size=DEFAULT_CHUNK_SIZE, + disable_radix_cache=True, + kv_canary="none", + kv_canary_sweep_interval=0, + ) + + def test_ignore_eos_chunked_truncate_path(self): + self.server.execute_script(self._script_ignore_eos_chunked_truncate_path) + + @staticmethod + def _script_ignore_eos_chunked_truncate_path(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + yield from run_until(r, lambda h: h.is_chunking) + yield from run_until_finished(r) + assert r.finished + expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE) + assert r.chunks_done == expected, ( + f"ignore_eos chunk-truncation path must carry the prompt as new_chunked_req " + f"for {expected} chunks; got {r.chunks_done}" + ) + + def test_ignore_eos_nonchunked_fits_chunk_budget(self): + self.server.execute_script(self._script_ignore_eos_nonchunked_fits_chunk_budget) + + @staticmethod + def _script_ignore_eos_nonchunked_fits_chunk_budget(t: ScriptedContext): + r = t.start_req(prompt_len=8, max_new_tokens=2, ignore_eos=True) + for _ in range(DEFAULT_MAX_STEPS): + assert not r.is_chunking, ( + "a prompt that fits the chunk budget must admit non-chunked via the " + "ignore_eos path, never entering is_chunking" + ) + if r.finished: + break + yield + assert r.finished + assert r.chunks_done == 0, ( + f"prompt fitting the chunk budget completes non-chunked; got " + f"chunks_done={r.chunks_done}" + ) + + def test_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero(self): + self.server.execute_script( + self._script_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero + ) + + @staticmethod + def _script_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero( + t: ScriptedContext, + ): + r1 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + r2 = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + yield from run_until(r1, lambda h: h.is_chunking) + + saw_deferred = False + for _ in range(DEFAULT_MAX_STEPS * 2): + comp = t.batch_composition() + r1_active = r1.rid in comp.get("chunked", []) + r2_idle = ( + r2.status == "waiting" + and r2.rid not in comp.get("chunked", []) + and r2.rid not in comp.get("running", []) + and r2.rid not in comp.get("decode", []) + ) + if r1_active and r2_idle: + saw_deferred = True + assert r2.kv_pages == 0, ( + f"deferred ignore_eos waiter must hold no KV; got " + f"kv_pages={r2.kv_pages}" + ) + if r1.finished and r2.finished: + break + yield + assert r1.finished and r2.finished + assert saw_deferred, ( + "second long ignore_eos waiter must be deferred for >=1 iter while the " + "chunk-token budget was driven to 0 (nested rem_chunk_tokens<=0 OTHER)" + ) + + +class TestSpecialCaseRetractedStain(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_retracted_stain_suppresses_cached_token_recount(self): + self.server.execute_script( + self._script_retracted_stain_suppresses_cached_token_recount + ) + + @staticmethod + def _script_retracted_stain_suppresses_cached_token_recount(t: ScriptedContext): + r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1) + cached_before = r.req.cached_tokens + + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r, max_steps=2000) + assert r.finished + + req = r.req + assert ( + req.retracted_stain is True + ), "retract must set retracted_stain so the cached-token recount is suppressed" + assert req.cached_tokens == cached_before, ( + f"retracted_stain must suppress re-adding pre_len-already_computed on " + f"resume; cached_tokens grew from {cached_before} to {req.cached_tokens}" + ) + + +class TestSpecialCaseResultSkipRetracted(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_result_skip_retracted_emits_no_stray_token(self): + self.server.execute_script( + self._script_result_skip_retracted_emits_no_stray_token + ) + + @staticmethod + def _script_result_skip_retracted_emits_no_stray_token(t: ScriptedContext): + r = t.start_req( + prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True + ) + yield from run_until(r, lambda h: h.is_chunking) + + t.pause_generation(mode="retract") + yield + t.continue_generation() + yield from run_until_finished(r, max_steps=2000) + assert r.finished + assert len(r.req.output_ids) == 2, ( + f"the skipped retracted entry must not append a stray output token; " + f"got output_ids len {len(r.req.output_ids)}" + ) + + +class TestSpecialCaseMiddleChunkNoToken(ScriptedTestCase): + ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE) + + def test_middle_chunk_appends_no_token_no_finish(self): + self.server.execute_script(self._script_middle_chunk_appends_no_token_no_finish) + + @staticmethod + def _script_middle_chunk_appends_no_token_no_finish(t: ScriptedContext): + r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2) + yield from run_until(r, lambda h: h.is_chunking) + + saw_middle_chunk = False + for _ in range(DEFAULT_MAX_STEPS * 2): + if r.is_chunking: + saw_middle_chunk = True + assert len(r.req.output_ids) == 0, ( + f"middle chunk must not append an output token; got " + f"output_ids len {len(r.req.output_ids)}" + ) + assert ( + r.status != "finished" + ), "middle chunk must not finish the req (skip_stream_req)" + if r.finished: + break + yield + assert r.finished + assert saw_middle_chunk, "test must observe r mid-chunk at least once" + assert ( + len(r.req.output_ids) >= 1 + ), "output tokens must appear only after the chunked prefill completes" + + +if __name__ == "__main__": + unittest.main()