Add more testing for chunked prefill (#27506)

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()

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)

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())

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,
     )
 
 

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:

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,

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)

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):

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

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()

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()