chore: import upstream snapshot with attribution

test/README.md

@@ -0,0 +1,114 @@
+# Test and Continuous Integration (CI) System in SGLang
+
+This page covers principles and essentials: folder layout, how to run tests, registration, and suite selection. For complete references, see the skill guides:
+
+- **Writing tests** — templates, fixtures, model selection, complete suite tables, checklist: [`.claude/skills/write-sglang-test/SKILL.md`](../.claude/skills/write-sglang-test/SKILL.md)
+- **CI pipeline internals** — stage flow diagrams, fast-fail layers, gating, partitioning, execution modes, debugging failures: [`.claude/skills/ci-workflow-guide/SKILL.md`](../.claude/skills/ci-workflow-guide/SKILL.md)
+
+## CI Pipeline Overview
+
+The CI pipeline runs in three sequential stages: **A** (pre-flight, ~3 min) → **B** (basic, ~30 min) → **C** (advanced, ~30 min). Kernel and multimodal-gen tests run in parallel with stage B. For details on stage gating, fast-fail mechanisms, execution modes (PR vs scheduled vs `/rerun-stage`), and debugging CI failures, see the [CI workflow guide](../.claude/skills/ci-workflow-guide/SKILL.md).
+
+## Folder Organization
+
+- `registered/`: CI test files, auto-discovered by `run_suite.py`. Most tests live here. JIT kernel tests are an exception (see below).
+- `manual/`: Non-CI tests for local debugging or special setups.
+- `run_suite.py`: CI runner — scans `registered/` and JIT kernel directories.
+
+The system supports both [unittest](https://docs.python.org/3/library/unittest.html) and [pytest](https://docs.pytest.org/en/stable/). The launcher runs `python filename.py -f` with **failfast enabled by default**.
+
+Make sure your file ends with **exactly** one of:
+
+```python
+# for unittest
+if __name__ == "__main__":
+    unittest.main()
+```
+
+```python
+# for pytest
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main([__file__]))
+```
+
+Do not add custom `argparse` or modify `sys.argv` before these calls — the CI runner appends `-f` for failfast.
+
+## Run Tests Locally
+
+```bash
+# Single file
+python3 test/registered/core/test_srt_endpoint.py
+
+# Single test method
+python3 test/registered/core/test_srt_endpoint.py TestSRTEndpoint.test_simple_decode
+
+# Single JIT kernel test
+python3 test/registered/jit/test_add_constant.py
+
+# Run a suite
+python3 test/run_suite.py --hw cpu --suite base-a-test-cpu
+python3 test/run_suite.py --hw cuda --suite base-a-test-1-gpu-small
+
+# Nightly tests
+python3 test/run_suite.py --hw cuda --suite nightly-1-gpu --nightly
+
+# With auto-partitioning (for parallel CI jobs)
+python3 test/run_suite.py --hw cuda --suite base-b-test-1-gpu-small \
+    --auto-partition-id 0 --auto-partition-size 4
+```
+
+## CI Registration
+
+Every CI-discovered test file must call a registration function at module level:
+
+```python
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=80, stage="base-b", runner_config="1-gpu-small")
+```
+
+Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and runner pool from `scripts/ci/runner_configs.yml`), `nightly=True` (nightly-only), `disabled="reason"` (temporarily disable).
+
+Keep `est_time`, `stage`, `runner_config` as **literal values** — `run_suite.py` collects them by AST parsing.
+
+JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path):
+- Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
+- Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
+
+## Choosing a Suite
+
+Use the lightest suite that meets your test's needs. Full suite tables are in the [write-sglang-test skill](../.claude/skills/write-sglang-test/SKILL.md#all-ci-suites).
+
+| Need | Suite |
+|------|-------|
+| No GPU required | `base-a-test-cpu` |
+| Small GPU (fits 5090, 32GB) | `base-b-test-1-gpu-small` (most tests go here) |
+| Large GPU memory or Hopper features | `base-b-test-1-gpu-large` |
+| JIT kernel correctness | `base-b-kernel-unit-test-1-gpu-large` |
+| JIT kernel benchmarks | `base-b-kernel-benchmark-test-1-gpu-large` |
+| Multi-GPU (2/4/8) | `base-b-test-2-gpu-large`, `base-c-test-*` |
+| Long-running or experimental | `nightly-*` suites |
+
+## Steps for Adding a Test
+
+See the [write-sglang-test skill](../.claude/skills/write-sglang-test/SKILL.md) for templates, fixtures, model selection, and a complete checklist.
+
+## Multi-Hardware Backends
+
+This README mostly describes the NVIDIA GPU CI pipeline. Other hardware backends (AMD, NPU) follow the same practices and use the multi-backend registry system. A scheduled job summarizes test coverage across all backends; [here is an example run](https://github.com/sgl-project/sglang/actions/runs/23424304300).
+
+## Tips
+
+- Learn from existing examples in [test/registered](https://github.com/sgl-project/sglang/tree/main/test/registered).
+- Reuse servers — launching is expensive. Share one server across many test methods via `setUpClass`.
+- Use as few GPUs as possible. Prefer 1-GPU runners.
+- Each test file should take < 500 seconds; split if longer.
+- Each GitHub Actions job should take < 30 minutes; split if longer.
+- If tests are too slow for per-commit, consider nightly suites.
+
+## Other Notes
+
+### Adding New Models to Nightly CI
+- **Text models**: Extend the [global model list variables](https://github.com/sgl-project/sglang/blob/85c1f7937781199203b38bb46325a2840f353a04/python/sglang/test/test_utils.py#L104) in `test_utils.py`.
+- **VLMs**: Extend the `MODEL_THRESHOLDS` dictionary in `test/registered/eval/test_vlms_mmmu_eval.py`.

test/lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml

@@ -0,0 +1,13 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.847
+  - name: "exact_match,flexible-extract"
+    value: 0.556
+limit: 1319
+num_concurrent: 128
+num_fewshot: 5
+apply_chat_template: false
+fewshot_as_multiturn: true

test/lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

@@ -0,0 +1,13 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.847
+  - name: "exact_match,flexible-extract"
+    value: 0.556
+limit: 1319
+num_concurrent: 128
+num_fewshot: 5
+apply_chat_template: false
+fewshot_as_multiturn: true

test/lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4.yaml

@@ -0,0 +1,13 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.847
+  - name: "exact_match,flexible-extract"
+    value: 0.444
+limit: 1319
+num_concurrent: 128
+num_fewshot: 5
+apply_chat_template: false
+fewshot_as_multiturn: true

test/lm_eval_configs/Qwen3.5-397B-A17B.yaml

@@ -0,0 +1,13 @@
+model_name: "Qwen/Qwen3.5-397B-A17B"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.9704
+  - name: "exact_match,flexible-extract"
+    value: 0.9697
+limit: 1319
+num_concurrent: 256
+num_fewshot: 5
+gen_kwargs: "max_gen_toks=2048"
+rtol: 0.05

test/manual/4-gpu-models/test_qwen35_fp4_triton.py

@@ -0,0 +1,68 @@
+import unittest
+
+from sglang.test.accuracy_test_runner import AccuracyTestParams
+
+# This eval harness applies the chat_template, which is critical for qwen3.5
+# to get good accuracy on gsm8k
+from sglang.test.run_combined_tests import run_combined_tests
+from sglang.test.test_utils import (
+    CustomTestCase,
+    ModelLaunchSettings,
+)
+
+QWEN35_FP4_MODEL = "nvidia/Qwen3.5-397B-A17B-NVFP4"
+ACC_THRESHOLDS = {QWEN35_FP4_MODEL: {"gsm8k": 0.95}}
+
+
+class TestQwen35FP4(CustomTestCase):
+    def test_gsm8k(self):
+        base_args = [
+            "--tp-size",
+            "4",
+            "--chunked-prefill-size",
+            "2048",
+            "--mamba-scheduler-strategy",
+            "extra_buffer",
+            "--mamba-track-interval",
+            "128",
+            "--mamba-ssm-dtype",
+            "bfloat16",
+            "--max-running-requests",
+            "128",
+            "--reasoning-parser",
+            "qwen3",
+            "--attention-backend",
+            "trtllm_mha",
+            "--quantization",
+            "modelopt_fp4",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true,"num_threads": 64}',
+        ]
+
+        variants = [
+            ModelLaunchSettings(
+                QWEN35_FP4_MODEL,
+                extra_args=base_args,
+                variant="Triton",
+            ),
+        ]
+
+        run_combined_tests(
+            models=variants,
+            test_name="Qwen3.5-397B-A17B-NVFP4",
+            accuracy_params=AccuracyTestParams(
+                dataset="gsm8k",
+                baseline_accuracy=ACC_THRESHOLDS[QWEN35_FP4_MODEL]["gsm8k"],
+                num_examples=200,
+                num_threads=128,
+                max_tokens=16000,
+                thinking_mode="qwen3",
+                temperature=0.6,
+                top_p=0.95,
+                top_k=20,
+            ),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/4-gpu-models/test_qwen35_models_archived.py

@@ -0,0 +1,168 @@
+"""Archived test classes split out of test/registered/4-gpu-models/test_qwen35_models.py.
+
+Originally registered with `register_cuda_ci(...)`. Moved here as part of
+the per-commit pruning effort to keep the code reachable manually.
+Run with `python3 test/manual/4-gpu-models/test_qwen35_models_archived.py`.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.accuracy_test_runner import AccuracyTestParams
+from sglang.test.kits.reasoning_kit import ReasoningTokenUsageMixin
+
+# This eval harness applies the chat_template, which is critical for qwen3.5
+# to get good accuracy on gsm8k
+from sglang.test.run_combined_tests import run_combined_tests
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    ModelLaunchSettings,
+    popen_launch_server,
+)
+
+QWEN35_FP4_MODEL = "nvidia/Qwen3.5-397B-A17B-NVFP4"
+ACC_THRESHOLDS = {QWEN35_FP4_MODEL: {"gsm8k": 0.95}}
+
+
+class TestQwen35FP4(CustomTestCase):
+    def test_gsm8k(self):
+        base_args = [
+            "--tp-size",
+            "4",
+            "--chunked-prefill-size",
+            "2048",
+            "--mamba-scheduler-strategy",
+            "extra_buffer",
+            "--mamba-track-interval",
+            "128",
+            "--mamba-ssm-dtype",
+            "bfloat16",
+            "--max-running-requests",
+            "128",
+            "--reasoning-parser",
+            "qwen3",
+            "--attention-backend",
+            "trtllm_mha",
+            "--quantization",
+            "modelopt_fp4",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true,"num_threads": 64}',
+        ]
+
+        variants = [
+            ModelLaunchSettings(
+                QWEN35_FP4_MODEL,
+                extra_args=base_args,
+                variant="Triton",
+            ),
+            # TODO: Fix this and re-enable it
+            # ModelLaunchSettings(
+            #     QWEN35_FP4_MODEL,
+            #     extra_args=base_args + ["--linear-attn-decode-backend", "flashinfer"],
+            #     variant="FlashInfer",
+            # ),
+        ]
+
+        run_combined_tests(
+            models=variants,
+            test_name="Qwen3.5-397B-A17B-NVFP4",
+            accuracy_params=AccuracyTestParams(
+                dataset="gsm8k",
+                baseline_accuracy=ACC_THRESHOLDS[QWEN35_FP4_MODEL]["gsm8k"],
+                num_examples=200,
+                num_threads=128,
+                max_tokens=16000,
+                thinking_mode="qwen3",
+                temperature=0.6,
+                top_p=0.95,
+                top_k=20,
+            ),
+        )
+
+
+class TestQwen35FP4MTP(ReasoningTokenUsageMixin, CustomTestCase):
+    reasoning_parser_name = "qwen3"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = QWEN35_FP4_MODEL
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.init_reasoning_token_verifier()
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                "4",
+                "--chunked-prefill-size",
+                "2048",
+                "--mamba-scheduler-strategy",
+                "extra_buffer",
+                "--mamba-track-interval",
+                "128",
+                "--mamba-ssm-dtype",
+                "bfloat16",
+                "--max-running-requests",
+                "128",
+                "--reasoning-parser",
+                "qwen3",
+                "--attention-backend",
+                "trtllm_mha",
+                "--quantization",
+                "modelopt_fp4",
+                "--speculative-algorithm",
+                "NEXTN",
+                "--speculative-num-steps",
+                "3",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "4",
+                "--mem-fraction-static",
+                "0.8",
+                "--model-loader-extra-config",
+                '{"enable_multithread_load": true,"num_threads": 64}',
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            model=self.model,
+            eval_name="gsm8k",
+            num_shots=5,
+            num_examples=200,
+            max_tokens=16000,
+            num_threads=128,
+            repeat=1,
+            temperature=0.6,
+            top_p=0.95,
+            top_k=20,
+            base_url=self.base_url,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"])
+
+        server_info = requests.get(self.base_url + "/server_info")
+        avg_spec_accept_length = server_info.json()["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        self.assertGreater(avg_spec_accept_length, 3.3)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/4-gpu-models/test_qwen3_next_models.py

@@ -0,0 +1,62 @@
+import unittest
+
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
+from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
+
+
+class TestQwen3Next(
+    GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    model = QWEN3_NEXT_MODEL
+    cache_chunk_size = 64
+    gsm8k_accuracy_thres = 0.93
+    kl_div_thres = 0.0025
+    other_args = [
+        "--tp-size",
+        "4",
+        "--chunked-prefill-size",
+        "1024",
+        "--mamba-scheduler-strategy",
+        "extra_buffer",
+        "--mamba-track-interval",
+        "2",
+        "--page-size",
+        "1",
+        "--attention-backend",
+        "triton",
+        "--moe-runner-backend",
+        "triton",
+    ]
+
+
+class TestQwen3NextLazyExtraBuffer(
+    GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    model = QWEN3_NEXT_MODEL
+    cache_chunk_size = 64
+    gsm8k_accuracy_thres = 0.93
+    kl_div_thres = 0.0025
+    other_args = [
+        "--tp-size",
+        "4",
+        "--chunked-prefill-size",
+        "1024",
+        "--mamba-scheduler-strategy",
+        "extra_buffer_lazy",
+        "--mamba-track-interval",
+        "2",
+        "--page-size",
+        "1",
+        "--attention-backend",
+        "triton",
+        "--moe-runner-backend",
+        "triton",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/4-gpu-models/test_qwen3_next_models_mtp_archived.py

@@ -0,0 +1,44 @@
+"""Archived test classes split out of test/registered/4-gpu-models/test_qwen3_next_models_mtp.py.
+
+Originally registered with `register_cuda_ci(...)`. Moved here as part of
+the per-commit pruning effort to keep the code reachable manually.
+Run with `python3 test/manual/4-gpu-models/test_qwen3_next_models_mtp_archived.py`.
+"""
+
+import unittest
+
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
+
+
+class TestQwen3NextMTP(GSM8KMixin, KLDivergenceMixin, DefaultServerBase):
+    model = QWEN3_NEXT_MODEL
+    gsm8k_accuracy_thres = 0.93
+    kl_div_thres = 0.0025
+    other_args = [
+        "--trust-remote-code",
+        "--speculative-algorithm",
+        "NEXTN",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--mem-fraction-static",
+        "0.8",
+        "--tp",
+        "4",
+        "--chunked-prefill-size",
+        "2048",
+        "--mamba-scheduler-strategy",
+        "no_buffer",
+        "--disable-radix-cache",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/8-gpu-models/test_deepseek_v3_basic.py

@@ -0,0 +1,82 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.send_one import BenchArgs, send_one_prompt
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_amd_ci,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
+
+
+class TestDeepseekV3Basic(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = FULL_DEEPSEEK_V3_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1400,
+            num_threads=1400,
+            num_shots=8,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (deepseek-v3)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.935)
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = send_one_prompt(args)
+
+        print(f"{speed=:.2f}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed (deepseek-v3)\n" f"{speed=:.2f} token/s\n"
+            )
+            if is_in_amd_ci():
+                self.assertGreater(speed, 12)
+            else:
+                self.assertGreater(speed, 75)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/8-gpu-models/test_dsa_models_basic.py

@@ -0,0 +1,259 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.send_one import BenchArgs, send_one_prompt
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2"
+GLM5_MODEL_PATH = "zai-org/GLM-5-FP8"
+
+
+class TestDeepseekV32DP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEEPSEEK_V32_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--dp",
+            "8",
+            "--enable-dp-attention",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1400,
+            num_threads=1400,
+            num_shots=20,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (deepseek-v32)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.935)
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = send_one_prompt(args)
+
+        print(f"{speed=:.2f}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
+            )
+            self.assertGreater(speed, 50)
+
+
+class TestDeepseekV32TP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEEPSEEK_V32_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1400,
+            num_threads=1400,
+            num_shots=20,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (deepseek-v32)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.935)
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = send_one_prompt(args)
+
+        print(f"{speed=:.2f}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed (deepseek-v32)\n" f"{speed=:.2f} token/s\n"
+            )
+            self.assertGreater(speed, 80)
+
+
+class TestGLM5DP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = GLM5_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--dp",
+            "8",
+            "--enable-dp-attention",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1400,
+            num_threads=1400,
+            num_shots=20,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (glm-5)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.935)
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = send_one_prompt(args)
+
+        print(f"{speed=:.2f}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed (glm-5)\n" f"{speed=:.2f} token/s\n"
+            )
+            self.assertGreater(speed, 40)
+
+
+class TestGLM5TP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = GLM5_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "8",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(
+        self,
+    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1400,
+            num_threads=1400,
+            num_shots=20,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (glm-5)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.935)
+
+    def test_bs_1_speed(self):
+        args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
+        acc_length, speed = send_one_prompt(args)
+
+        print(f"{speed=:.2f}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_bs_1_speed (glm-5)\n" f"{speed=:.2f} token/s\n"
+            )
+            self.assertGreater(speed, 60)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/disaggregation_utils.py

@@ -0,0 +1,153 @@
+import logging
+import os
+import time
+import warnings
+from urllib.parse import urlparse
+
+import requests
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_with_error_check,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TestDisaggregationBase(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        parsed_url = urlparse(DEFAULT_URL_FOR_TEST)
+        cls.base_host = parsed_url.hostname
+        base_port = str(parsed_url.port)
+        cls.lb_port = base_port
+        cls.prefill_port = f"{int(base_port) + 100}"
+        cls.decode_port = f"{int(base_port) + 200}"
+        cls.prefill_url = f"http://{cls.base_host}:{cls.prefill_port}"
+        cls.decode_url = f"http://{cls.base_host}:{cls.decode_port}"
+        cls.lb_url = f"http://{cls.base_host}:{cls.lb_port}"
+        print(f"{cls.base_host=} {cls.lb_port=} {cls.prefill_port=} {cls.decode_port=}")
+        cls.process_lb, cls.process_decode, cls.process_prefill = None, None, None
+
+        # config transfer backend and rdma devices
+        cls.transfer_backend = [
+            "--disaggregation-transfer-backend",
+            envs.SGLANG_TEST_PD_DISAGG_BACKEND.get(),
+        ]
+        cls.rdma_devices = [
+            "--disaggregation-ib-device",
+            envs.SGLANG_TEST_PD_DISAGG_DEVICES.get(),
+        ]
+        if cls.rdma_devices[1] is None:
+            cls.rdma_devices = []
+            msg = "No RDMA devices specified for disaggregation test, using default settings."
+            warnings.warn(msg)
+
+    @classmethod
+    def launch_lb(cls):
+        lb_command = [
+            "python3",
+            "-m",
+            "sglang_router.launch_router",
+            "--pd-disaggregation",
+            "--mini-lb",
+            "--prefill",
+            cls.prefill_url,
+            "--decode",
+            cls.decode_url,
+            "--host",
+            cls.base_host,
+            "--port",
+            cls.lb_port,
+        ]
+        print("Starting load balancer:", " ".join(lb_command))
+        cls.process_lb = popen_with_error_check(lb_command)
+        cls.wait_server_ready(cls.lb_url + "/health")
+
+    @classmethod
+    def wait_server_ready(cls, url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH):
+        start_time = time.perf_counter()
+        while True:
+            try:
+                response = requests.get(url)
+                if response.status_code == 200:
+                    print(f"Server {url} is ready")
+                    return
+            except Exception:
+                pass
+
+            if time.perf_counter() - start_time > timeout:
+                raise RuntimeError(f"Server {url} failed to start in {timeout}s")
+            time.sleep(1)
+
+    @classmethod
+    def tearDownClass(cls):
+        for process in [cls.process_lb, cls.process_decode, cls.process_prefill]:
+            if process:
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process {process.pid}: {e}")
+
+        # wait for 5 seconds
+        time.sleep(5)
+
+
+def get_rdma_devices_args():
+    def _parse_list_env(var_name: str):
+        val = os.getenv(var_name)
+        if not val:
+            return None
+        items = [x.strip() for x in val.split(",") if x.strip()]
+        return items or None
+
+    def _pick_default_pair(rdma_all_devices):
+        return [rdma_all_devices[0], rdma_all_devices[len(rdma_all_devices) // 2]]
+
+    rdma_all_devices = _parse_list_env("SGLANG_CI_RDMA_ALL_DEVICES") or [
+        f"mlx5_roce{i}" for i in range(8)
+    ]
+    logger.info("Resolved rdma_all_devices=%s", rdma_all_devices)
+
+    n_rdma = len(rdma_all_devices)
+
+    # 1. Get visible GPU indices
+    cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if not cuda_visible_devices:
+        warnings.warn("CUDA_VISIBLE_DEVICES is not set. Using default RDMA devices.")
+        return ",".join(_pick_default_pair(rdma_all_devices))
+
+    try:
+        # Convert to list of integers (handling possible spaces and empty strings)
+        gpu_indices = [
+            int(idx.strip()) for idx in cuda_visible_devices.split(",") if idx.strip()
+        ]
+        if not gpu_indices or len(gpu_indices) > 4:
+            return ",".join(_pick_default_pair(rdma_all_devices))
+    except ValueError:
+        warnings.warn(f"Invalid CUDA_VISIBLE_DEVICES format: {cuda_visible_devices}")
+        return ",".join(_pick_default_pair(rdma_all_devices))
+
+    # 2. Calculate base RDMA index group (each group of 4 GPUs uses consecutive devices)
+    base_rdma_group = (min(gpu_indices) // 4) * 4
+    for gpu_idx in gpu_indices:
+        if not (base_rdma_group <= gpu_idx < base_rdma_group + 4):
+            warnings.warn(
+                f"GPU index {gpu_idx} is outside expected group "
+                f"{base_rdma_group}-{base_rdma_group+3}"
+            )
+
+    # 3. Generate RDMA device names
+    rdma_devices = []
+    for gpu_idx in gpu_indices:
+        nic_index = gpu_idx // (8 // n_rdma)
+        rdma_devices.append(rdma_all_devices[nic_index])
+
+    if not rdma_devices:
+        return ",".join(_pick_default_pair(rdma_all_devices))
+
+    return ",".join(rdma_devices)

test/manual/ascend/llm_models/test_npu_gemma_4_26b_a4b_it_llm.py

@@ -0,0 +1,30 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_26B_A4B_IT_WEIGHTS_PATH
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestGemma426BA4BIt(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-26B-A4B-it model on the GSM8K dataset is no less than 0.35.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-26B-A4B-it
+    """
+
+    model = GEMMA_4_26B_A4B_IT_WEIGHTS_PATH
+    accuracy = 0.35
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.7",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "2",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/llm_models/test_npu_gemma_4_31b_llm.py

@@ -0,0 +1,30 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_31B_WEIGHTS_PATH
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestGemma431B(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-31B-it model on the GSM8K dataset is no less than 0.70.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-31B-it
+    """
+
+    model = GEMMA_4_31B_WEIGHTS_PATH
+    accuracy = 0.70
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.7",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "2",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/llm_models/test_npu_gemma_4_e2b_llm.py

@@ -0,0 +1,30 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_E2B_WEIGHTS_PATH
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestGemma4E2B(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-E2B-it model on the GSM8K dataset is no less than 0.05.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-E2B-it
+    """
+
+    model = GEMMA_4_E2B_WEIGHTS_PATH
+    accuracy = 0.05
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.7",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "1",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/llm_models/test_npu_gemma_4_e4b_llm.py

@@ -0,0 +1,30 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_E4B_WEIGHTS_PATH
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestGemma4E4B(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-E4B-it model on the GSM8K dataset is no less than 0.60.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-E4B-it
+    """
+
+    model = GEMMA_4_E4B_WEIGHTS_PATH
+    accuracy = 0.60
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.7",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "1",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/llm_models/test_npu_mimo_v2_flash.py

@@ -0,0 +1,39 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ascend.test_ascend_utils import MIMO_V2_FLASH_WEIGHTS_PATH
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestMiMoV2FlashGraphWithMTP(GSM8KAscendMixin, CustomTestCase):
+    """Testcase: Verify the inference accuracy of MiMo-V2-Flash on GSM8K with cuda graph and MTP (speculative decoding).
+
+    [Test Category] Model
+    [Test Target] XiaomiMiMo/MiMo-V2-Flash
+    [Test Config] Prefill+Decode, cuda graph enabled, EAGLE speculative decoding
+    """
+
+    model = MIMO_V2_FLASH_WEIGHTS_PATH
+    accuracy = 0.9
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.8",
+        "--attention-backend",
+        "ascend",
+        "--tp-size",
+        "16",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--enable-multi-layer-eagle",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/test_ascend_deepseek_mtp.py

@@ -0,0 +1,93 @@
+import unittest
+from types import SimpleNamespace
+from urllib.parse import urlparse
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+TEST_MODEL_MATRIX = {
+    "/root/.cache/modelscope/hub/models/vllm-ascend/DeepSeek-R1-0528-W8A8": {
+        "accuracy": 0.95,
+        "latency": 1000,
+        "output_throughput": 6,
+    },
+}
+
+
+class TestAscendDeepSeekMTP(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.models = TEST_MODEL_MATRIX.keys()
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.url = urlparse(DEFAULT_URL_FOR_TEST)
+
+        cls.common_args = [
+            "--trust-remote-code",
+            "--attention-backend",
+            "ascend",
+            "--mem-fraction-static",
+            0.8,
+            "--disable-radix-cache",
+            "--chunked-prefill-size",
+            32768,
+            "--tp-size",
+            16,
+            "--dp-size",
+            2,
+            "--enable-dp-attention",
+            "--speculative-algorithm",
+            "NEXTN",
+            "--speculative-num-steps",
+            1,
+            "--speculative-eagle-topk",
+            1,
+            "--speculative-num-draft-tokens",
+            2,
+        ]
+
+        envs.SGLANG_NPU_USE_MLAPO.set(True)
+        envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.set(True)
+
+    def test_a_gsm8k(self):
+        for model in self.models:
+            with self.subTest(model=model):
+                print(f"##=== Testing accuracy: {model} ===##")
+
+                process = popen_launch_server(
+                    model,
+                    self.base_url,
+                    timeout=2400,
+                    other_args=[
+                        *self.common_args,
+                    ],
+                )
+
+                try:
+                    args = SimpleNamespace(
+                        num_shots=5,
+                        data_path=None,
+                        num_questions=1319,
+                        max_new_tokens=512,
+                        parallel=128,
+                        host=f"http://{self.url.hostname}",
+                        port=int(self.url.port),
+                    )
+
+                    metrics = run_eval_few_shot_gsm8k(args)
+                    self.assertGreaterEqual(
+                        metrics["accuracy"],
+                        TEST_MODEL_MATRIX[model]["accuracy"],
+                    )
+                finally:
+                    kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/test_ascend_vocab_mask.py

@@ -0,0 +1,81 @@
+import math
+
+import pytest
+import torch
+
+from sglang.srt.constrained import xgrammar_backend as xb
+
+
+def _pack_mask(allowed_ids, vocab_size, batch_size=1):
+    nwords = math.ceil(vocab_size / 32)
+    m = torch.zeros((batch_size, nwords), dtype=torch.int32)
+    for b in range(batch_size):
+        for tid in allowed_ids[b]:
+            m[b, tid // 32] |= 1 << (tid % 32)
+    return m
+
+
+def _apply_ref_cpu(logits, vocab_mask):
+    vocab_size = logits.shape[-1]
+    token_ids = torch.arange(vocab_size, device="cpu", dtype=torch.int64)
+    word_idx = token_ids // 32
+    bit_idx = (token_ids % 32).to(torch.int32)
+    words = vocab_mask.cpu()[:, word_idx].to(torch.int32)
+    allowed = ((words >> bit_idx) & 1).bool().to(logits.device)
+    out = logits.clone()
+    out.masked_fill_(~allowed, float("-inf"))
+    return out
+
+
+@pytest.mark.skipif(
+    not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
+)
+def test_mask_blocks_disallowed_token_on_npu():
+    device = "npu:0"
+    vocab_size = 64
+
+    logits = torch.zeros((1, vocab_size), device=device, dtype=torch.float32)
+    logits[0, 16] = 22.125
+    logits[0, 5] = 10.0
+
+    allowed = [[5, 6, 7, 8]]
+    vocab_mask = _pack_mask(allowed, vocab_size).to(device=device, dtype=torch.int32)
+
+    g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
+    out = logits.clone()
+    g.apply_vocab_mask(out, vocab_mask)
+
+    assert not torch.isfinite(out[0, 16])
+    assert int(torch.argmax(out[0]).item()) != 16
+
+
+@pytest.mark.skipif(
+    not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
+)
+def test_npu_path_matches_reference_random():
+    device = "npu:0"
+    B, V = 4, 257
+    torch.manual_seed(0)
+
+    logits = torch.randn(B, V, device=device, dtype=torch.float32)
+
+    allowed = []
+    for _ in range(B):
+        ids = torch.randperm(V)[: V // 4].tolist()
+        allowed.append(ids)
+    vocab_mask = _pack_mask(allowed, V, B).to(device=device, dtype=torch.int32)
+
+    g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
+    out_npu = logits.clone()
+    g.apply_vocab_mask(out_npu, vocab_mask)
+
+    out_ref = _apply_ref_cpu(logits, vocab_mask)
+
+    assert torch.equal(torch.isfinite(out_npu), torch.isfinite(out_ref))
+    diff = (
+        torch.nan_to_num(out_npu - out_ref, nan=0.0, posinf=0.0, neginf=0.0)
+        .abs()
+        .max()
+        .item()
+    )
+    assert diff < 1e-5

test/manual/ascend/test_ascend_w8a8_quantization.py

@@ -0,0 +1,102 @@
+"""
+Usage:
+python3 -m unittest test_ascend_w8a8_quantization.TestAscendW8A8.test_gsm8k
+"""
+
+import os
+import time
+import unittest
+from types import SimpleNamespace
+from urllib.parse import urlparse
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+)
+
+if "ASCEND_RT_VISIBLE_DEVICES" not in os.environ:
+    os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0,1"
+DEFAULT_PORT_FOR_SRT_TEST_RUNNER = (
+    7000 + int(os.environ.get("ASCEND_RT_VISIBLE_DEVICES", "0")[0]) * 100
+)
+DEFAULT_URL_FOR_TEST = f"http://127.0.0.1:{DEFAULT_PORT_FOR_SRT_TEST_RUNNER + 1000}"
+
+
+class TestAscendW8A8(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "vllm-ascend/Qwen2.5-0.5B-Instruct-w8a8"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--disable-cuda-graph",
+                "--device",
+                "npu",
+                "--attention-backend",
+                "ascend",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        base_url = DEFAULT_URL_FOR_TEST
+        url = urlparse(base_url)
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host=f"http://{url.hostname}",
+            port=int(url.port),
+        )
+        metrics = run_eval(args)
+        print(metrics)
+
+        self.assertGreaterEqual(metrics["accuracy"], 0.25)
+        self.assertGreaterEqual(metrics["output_throughput"], 1000)
+
+    def run_decode(self, max_new_tokens):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "The capital of France is",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": max_new_tokens,
+                },
+                "ignore_eos": True,
+            },
+        )
+        return response.json()
+
+    def test_throughput(self):
+        max_tokens = 256
+
+        tic = time.perf_counter()
+        res = self.run_decode(max_tokens)
+        tok = time.perf_counter()
+        print(res["text"])
+        throughput = max_tokens / (tok - tic)
+        print(f"Throughput: {throughput} tokens/s")
+
+        if is_in_ci():
+            self.assertGreaterEqual(throughput, 25)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/test_mindspore_models.py

@@ -0,0 +1,64 @@
+"""
+Usage:
+python3 -m unittest test_mindspore_models.TestMindSporeQwen3.test_gsm8k
+"""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestMindSporeQwen3(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-8B"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--device",
+                "npu",
+                "--model-impl",
+                "mindspore",
+                "--attention-backend",
+                "ascend",
+                "--tp-size",
+                "1",
+                "--dp-size",
+                "1",
+                "--mem-fraction-static",
+                0.8,
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["accuracy"], 0.78)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/vlm_models/test_npu_gemma_4_26b_a4b_it.py

@@ -0,0 +1,22 @@
+import unittest
+
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_26B_A4B_IT_WEIGHTS_PATH
+from sglang.test.ascend.vlm_utils import TestVLMModels
+
+
+class TestGemma426BA4BIt(TestVLMModels):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-26B-A4B-it model on the MMMU dataset is no less than 0.40.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-26B-A4B-it
+    """
+
+    model = GEMMA_4_26B_A4B_IT_WEIGHTS_PATH
+    mmmu_accuracy = 0.40
+
+    def test_vlm_mmmu_benchmark(self):
+        self._run_vlm_mmmu_test()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/vlm_models/test_npu_gemma_4_31b.py

@@ -0,0 +1,22 @@
+import unittest
+
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_31B_WEIGHTS_PATH
+from sglang.test.ascend.vlm_utils import TestVLMModels
+
+
+class TestGemma431B(TestVLMModels):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-31B-it model on the MMMU dataset is no less than 0.50.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-31B-it
+    """
+
+    model = GEMMA_4_31B_WEIGHTS_PATH
+    mmmu_accuracy = 0.50
+
+    def test_vlm_mmmu_benchmark(self):
+        self._run_vlm_mmmu_test()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/vlm_models/test_npu_gemma_4_e2b.py

@@ -0,0 +1,22 @@
+import unittest
+
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_E2B_WEIGHTS_PATH
+from sglang.test.ascend.vlm_utils import TestVLMModels
+
+
+class TestGemma4E2B(TestVLMModels):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-E2B-it model on the MMMU dataset is no less than 0.15.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-E2B-it
+    """
+
+    model = GEMMA_4_E2B_WEIGHTS_PATH
+    mmmu_accuracy = 0.15
+
+    def test_vlm_mmmu_benchmark(self):
+        self._run_vlm_mmmu_test()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ascend/vlm_models/test_npu_gemma_4_e4b.py

@@ -0,0 +1,22 @@
+import unittest
+
+from sglang.test.ascend.test_ascend_utils import GEMMA_4_E4B_WEIGHTS_PATH
+from sglang.test.ascend.vlm_utils import TestVLMModels
+
+
+class TestGemma4E4B(TestVLMModels):
+    """Testcase: Verify that the inference accuracy of the google/gemma-4-E4B-it model on the MMMU dataset is no less than 0.30.
+
+    [Test Category] Model
+    [Test Target] google/gemma-4-E4B-it
+    """
+
+    model = GEMMA_4_E4B_WEIGHTS_PATH
+    mmmu_accuracy = 0.30
+
+    def test_vlm_mmmu_benchmark(self):
+        self._run_vlm_mmmu_test()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/attention/test_fa3.py

@@ -0,0 +1,263 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_device_sm, kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_DRAFT_MODEL_EAGLE3,
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+# FlashAttention3 integration tests (requires SM 90+ / H100)
+# Multiple test classes: FA3, FA3+MLA, FA3+SpecDecode variants
+GSM_DATASET_PATH = None
+
+# In case of some machine lack internet connection, we can set OFFLINE_MODE to True.
+OFFLINE_MODE = False
+
+# Change the path below when OFFLINE_MODE is True.
+OFFLINE_PATH_DICT = {
+    DEFAULT_MODEL_NAME_FOR_TEST: "/shared/public/elr-models/meta-llama/Meta-Llama-3.1-8B-Instruct",
+    DEFAULT_DRAFT_MODEL_EAGLE3: "/shared/public/elr-models/jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B",
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA: "/shared/public/sharing/deepseek/dsv3-test/snapshots/",
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN: "/shared/public/sharing/deepseek/dsv3-test-NextN/snapshots/",
+    GSM_DATASET_PATH: "/shared/public/data/gsm8k/test.jsonl",
+}
+
+if OFFLINE_MODE:
+    DEFAULT_MODEL_NAME_FOR_TEST = OFFLINE_PATH_DICT[DEFAULT_MODEL_NAME_FOR_TEST]
+    DEFAULT_DRAFT_MODEL_EAGLE3 = OFFLINE_PATH_DICT[DEFAULT_DRAFT_MODEL_EAGLE3]
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA = OFFLINE_PATH_DICT[DEFAULT_MODEL_NAME_FOR_TEST_MLA]
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN = OFFLINE_PATH_DICT[
+        DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN
+    ]
+    GSM_DATASET_PATH = OFFLINE_PATH_DICT[GSM_DATASET_PATH]
+
+# Default server arguments shared across all tests
+DEFAULT_SERVER_ARGS = [
+    "--trust-remote-code",
+    "--cuda-graph-max-bs-decode",
+    "8",
+    "--attention-backend",
+    "fa3",
+]
+
+"""
+Integration test for python/sglang/srt/layers/attention/flashattention_backend.py
+"""
+
+
+@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
+class BaseFlashAttentionTest(CustomTestCase):
+    """Base class for testing FlashAttention3."""
+
+    model = DEFAULT_MODEL_NAME_FOR_TEST
+    base_url = DEFAULT_URL_FOR_TEST
+    accuracy_threshold = 0.65  # derived tests need to override this
+    speculative_decode = False
+    spec_decode_threshold = 1.0  # derived spec decoding tests need to override this
+
+    @classmethod
+    def get_server_args(cls):
+        """Return the arguments for the server launch. Override in subclasses."""
+        return DEFAULT_SERVER_ARGS
+
+    @classmethod
+    def setUpClass(cls):
+        # disable deep gemm precompile to make launch server faster
+        # please don't do this if you want to make your inference workload faster
+        with (
+            envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.override(False),
+            envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False),
+        ):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=cls.get_server_args(),
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=100,
+            num_threads=128,
+            num_shots=4,
+            gsm8k_data_path=GSM_DATASET_PATH,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        # Use the appropriate metric key based on the test class
+        metric_key = "score"
+        self.assertGreater(metrics[metric_key], self.accuracy_threshold)
+
+        if self.speculative_decode:
+            server_info = requests.get(self.base_url + "/server_info").json()
+            avg_spec_accept_length = server_info["internal_states"][0][
+                "avg_spec_accept_length"
+            ]
+            print(f"{avg_spec_accept_length=}")
+            self.assertGreater(avg_spec_accept_length, self.spec_decode_threshold)
+
+
+class TestFlashAttention3MLA(BaseFlashAttentionTest):
+    """Test FlashAttention3 with MLA, e.g. deepseek v3 test model"""
+
+    accuracy_threshold = 0.60
+    model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+
+    @classmethod
+    def get_server_args(cls):
+        return DEFAULT_SERVER_ARGS
+
+
+class TestFlashAttention3SpeculativeDecode(BaseFlashAttentionTest):
+    """Test FlashAttention3 with speculative decode enabled with Llama 3.1 8B and its eagle3 model"""
+
+    model = DEFAULT_MODEL_NAME_FOR_TEST
+    accuracy_threshold = 0.65
+    speculative_decode = True
+    spec_decode_threshold = 1.5
+
+    @classmethod
+    def get_server_args(cls):
+        args = DEFAULT_SERVER_ARGS
+        args.extend(
+            [
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--speculative-algorithm",
+                "EAGLE3",
+                "--speculative-draft-model-path",
+                DEFAULT_DRAFT_MODEL_EAGLE3,
+                "--speculative-num-steps",
+                "3",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "4",
+                "--dtype",
+                "float16",
+            ]
+        )
+        return args
+
+
+class TestFlashAttention3SpeculativeDecodeTopk(BaseFlashAttentionTest):
+    """Tests FlashAttention3 with enhanced speculative decoding using Llama 3.1 8B and EAGLE3.
+    This test will be using top-k value > 1 which would verify the other branches of the FA3 code
+    """
+
+    model = DEFAULT_MODEL_NAME_FOR_TEST
+    accuracy_threshold = 0.65
+    speculative_decode = True
+    spec_decode_threshold = 1.6
+
+    @classmethod
+    def get_server_args(cls):
+        args = DEFAULT_SERVER_ARGS
+        args.extend(
+            [
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--speculative-algorithm",
+                "EAGLE3",
+                "--speculative-draft-model-path",
+                DEFAULT_DRAFT_MODEL_EAGLE3,
+                "--speculative-num-steps",
+                "5",
+                "--speculative-eagle-topk",
+                "4",
+                "--speculative-num-draft-tokens",
+                "8",
+                "--dtype",
+                "float16",
+            ]
+        )
+        return args
+
+
+class TestFlashAttention3MLASpeculativeDecode(BaseFlashAttentionTest):
+    """Test FlashAttention3 with speculative decode enabled with deepseek v3 test model and its nextN model"""
+
+    model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+    accuracy_threshold = 0.60
+    speculative_decode = True
+    spec_decode_threshold = 2.5
+
+    @classmethod
+    def get_server_args(cls):
+        args = DEFAULT_SERVER_ARGS
+        args.extend(
+            [
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-draft-model-path",
+                DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+                "--speculative-num-steps",
+                "3",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "4",
+            ]
+        )
+        return args
+
+
+class TestFlashAttention3MLASpeculativeDecodeTopk(BaseFlashAttentionTest):
+    """Test FlashAttention3 with speculative decode enabled with deepseek v3 test model and its nextN model
+    This test will be using top-k value > 1 which would verify the other branches of the FA3 code
+    """
+
+    model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+    accuracy_threshold = 0.60
+    speculative_decode = True
+    spec_decode_threshold = 2.95
+
+    @classmethod
+    def get_server_args(cls):
+        args = DEFAULT_SERVER_ARGS
+        args.extend(
+            [
+                "--cuda-graph-max-bs-decode",
+                "4",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-draft-model-path",
+                DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+                "--speculative-num-steps",
+                "5",
+                "--speculative-eagle-topk",
+                "4",
+                "--speculative-num-draft-tokens",
+                "8",
+            ]
+        )
+        return args
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/attention/test_flashattn_backend.py

@@ -0,0 +1,588 @@
+import unittest
+
+import torch
+
+from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.layers.attention.flashattention_backend import (
+    FlashAttentionBackend,
+    draft_decode_set_expand_metadata,
+)
+from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    set_forward_context,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class MockModelRunner:
+    def __init__(
+        self,
+        page_size=1,
+        num_heads=2,
+        head_dim=8,
+    ):
+        self.device = "cuda"
+        self.dtype = torch.float16
+        self.kv_cache_dtype = torch.float16
+        self.is_hybrid_swa = False
+        self.attention_chunk_size = None
+        attention_arch = AttentionArch.MHA
+        # Max batch size for the test.
+        max_batch_size = 160
+        # Total tokens(prefix + extend + decode) in the test should not exceed this length.
+        max_context_len = 2048
+        self.model_config = type(
+            "ModelConfig",
+            (),
+            {
+                "context_len": max_context_len,
+                "is_multimodal": False,
+                "attention_arch": attention_arch,
+                "is_encoder_decoder": False,
+                "is_local_attention_model": False,
+            },
+        )()
+        self.sliding_window_size = None
+        self.kv_cache_dtype = (
+            self.dtype
+        )  # torch dtype, required by FlashAttentionBackend
+
+        # server_args is still needed for string-based config (kv_cache_dtype_str)
+        self.server_args = type(
+            "ServerArgs",
+            (),
+            {
+                "kv_cache_dtype": "auto",  # string version for kv_cache_dtype_str
+                "speculative_eagle_topk": None,
+                "speculative_num_draft_tokens": 0,
+                "enable_deterministic_inference": False,
+            },
+        )
+        self.attn_cp_size = 1
+        # Create a large enough req_to_token_pool to fit the test usage.
+        self.req_to_token_pool = type(
+            "TokenPool",
+            (),
+            {
+                # A typical max_bs * max_context_len for cuda graph decode
+                "size": max_batch_size,
+                # Add req_to_token attribute
+                "req_to_token": torch.zeros(
+                    max_batch_size,
+                    max_context_len,
+                    dtype=torch.int32,
+                    device=self.device,
+                ),
+            },
+        )()
+        self.page_size = page_size
+        max_total_num_tokens = max_batch_size * max_context_len
+        self.token_to_kv_pool = MHATokenToKVPool(
+            size=max_total_num_tokens,
+            page_size=page_size,
+            dtype=self.dtype,
+            head_num=num_heads,
+            head_dim=head_dim,
+            layer_num=1,  # only consider layer=1 for unit test
+            device=self.device,
+            enable_memory_saver=False,
+        )
+        self.hisparse_coordinator = None
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFlashAttentionBackend(CustomTestCase):
+    def setUp(self):
+        # Test parameters
+        self.batch_size = 2
+        self.seq_len = 256
+        self.num_heads = 2
+        self.head_dim = 8
+        self.device = "cuda"
+        self.dtype = torch.float16
+
+    def _init_model_runner(self, page_size=1):
+        self.model_runner = MockModelRunner(
+            page_size=page_size,
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+        )
+        self.backend = FlashAttentionBackend(self.model_runner)
+        self.ref_backend = TorchNativeAttnBackend(self.model_runner)
+        self.model_runner.model_config.num_attention_heads = self.num_heads
+        # Publish the backend for any RadixAttention.forward path the tests
+        # exercise; tearDown is unnecessary here since each test re-inits.
+        set_forward_context(ForwardContext(attn_backend=self.backend))
+
+    def _mock_write_to_req_to_token_pool(self, batch_size, seq_len, page_size):
+        # if page_size > 1, the token pool stores the index to the page.
+        # so we need to multiply the index by page_size.
+        self.req_to_token = (
+            torch.arange(0, batch_size, dtype=torch.int32, device=self.device)[:, None]
+            * seq_len
+            + torch.arange(0, seq_len, dtype=torch.int32, device=self.device)[None, :]
+            + page_size
+        )
+        self.model_runner.req_to_token_pool.req_to_token[:batch_size, :seq_len] = (
+            self.req_to_token
+        )
+
+    def _create_attention_layer(self):
+        """Create attention layer for testing."""
+        return RadixAttention(
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            scaling=1.0,
+            num_kv_heads=self.num_heads,
+            layer_id=0,
+        )
+
+    def _create_qkv_tensors(self, tokens_len):
+        """Create q, k, v tensors for testing."""
+        shape = (tokens_len, self.num_heads, self.head_dim)
+        return (
+            torch.randn(shape, dtype=self.dtype, device=self.device),
+            torch.randn(shape, dtype=self.dtype, device=self.device),
+            torch.randn(shape, dtype=self.dtype, device=self.device),
+        )
+
+    def _run_reference_forward(
+        self, mode, q, k, v, layer, forward_batch, expected_shape
+    ):
+        """Run reference forward pass using native backend."""
+        if mode == ForwardMode.EXTEND:
+            output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
+        else:  # ForwardMode.DECODE
+            output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
+        return output.view(expected_shape)
+
+    def _verify_output(self, output, expected_shape, output_ref=None):
+        """Verify output tensor shape, dtype, and values."""
+        self.assertEqual(
+            output.shape,
+            expected_shape,
+            f"Expected shape {expected_shape}, got {output.shape}",
+        )
+        self.assertEqual(output.dtype, self.dtype)
+        self.assertEqual(output.device.type, "cuda")
+        self.assertEqual(
+            torch.isnan(output).sum().item(), 0, "Output contains NaN values"
+        )
+
+        if output_ref is not None:
+            if not torch.allclose(output, output_ref, atol=1e-1, rtol=0.0):
+                # Check where the values differ beyond the given tolerances
+                diff_mask = ~torch.isclose(output, output_ref, atol=1e-1, rtol=0.0)
+
+                # Find the first index where the difference occurs
+                if diff_mask.any():
+                    first_mismatch_idx = diff_mask.nonzero()[0]
+                    print(
+                        "First mismatch at index:", tuple(first_mismatch_idx.tolist())
+                    )
+                    print("output:", output[tuple(first_mismatch_idx.tolist())])
+                    print("output_ref:", output_ref[tuple(first_mismatch_idx.tolist())])
+                raise AssertionError(
+                    "Attention output is not close to the torch native backend output"
+                )
+
+    def _create_forward_batch(
+        self, mode, q_len=None, prefix_len=0, page_size=1, attn_cp_size=1
+    ):
+        """Create a forward batch for testing based on mode and lengths."""
+        self._init_model_runner(page_size=page_size)
+
+        # Default to self.seq_len if not specified
+        q_len = q_len or self.seq_len
+
+        if mode == ForwardMode.EXTEND:
+            total_len = prefix_len + q_len
+            out_cache_start = prefix_len * self.batch_size
+            out_cache_end = total_len * self.batch_size
+
+            forward_batch = ForwardBatch(
+                batch_size=self.batch_size,
+                input_ids=torch.randint(
+                    0, 100, (self.batch_size, q_len), device=self.device
+                ),
+                out_cache_loc=torch.arange(
+                    out_cache_start, out_cache_end, device=self.device
+                ),
+                seq_lens_sum=self.batch_size * total_len,
+                forward_mode=mode,
+                req_pool_indices=torch.arange(self.batch_size, device=self.device),
+                seq_lens=torch.tensor(
+                    [total_len] * self.batch_size, device=self.device
+                ),
+                seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
+                extend_prefix_lens=torch.tensor(
+                    [prefix_len] * self.batch_size, device=self.device
+                ),
+                extend_prefix_lens_cpu=[prefix_len] * self.batch_size,
+                extend_seq_lens=torch.tensor(
+                    [q_len] * self.batch_size, device=self.device
+                ),
+                extend_seq_lens_cpu=[q_len] * self.batch_size,
+            )
+            if attn_cp_size > 1:
+                forward_batch.attn_cp_metadata = type(
+                    "AttnCPMetadata",
+                    (),
+                    {
+                        "kv_len_prev_tensor": torch.tensor(
+                            [q_len // 2] * self.batch_size,
+                            dtype=torch.int32,
+                            device=self.device,
+                        ),
+                        "kv_len_next_tensor": torch.tensor(
+                            [q_len] * self.batch_size,
+                            dtype=torch.int32,
+                            device=self.device,
+                        ),
+                        "actual_seq_q_prev": q_len // 2,
+                        "actual_seq_q_next": q_len // 2,
+                    },
+                )
+        else:  # ForwardMode.DECODE
+            decode_len = q_len  # Assuming 1 for decode testing
+            total_len = self.seq_len + decode_len
+            if mode == ForwardMode.DECODE and page_size > 1:
+                # Get next page_size multiple of self.seq_len
+                out_cache_start = (
+                    self.batch_size * self.seq_len // page_size + 1
+                ) * page_size
+                # out_cache_end is the start of the next block
+                out_cache_end = out_cache_start + decode_len * page_size
+            else:
+                out_cache_start = self.batch_size * self.seq_len
+                out_cache_end = self.batch_size * total_len
+
+            forward_batch = ForwardBatch(
+                batch_size=self.batch_size,
+                input_ids=torch.randint(
+                    0, 100, (self.batch_size, decode_len), device=self.device
+                ),
+                out_cache_loc=torch.tensor(
+                    [out_cache_start, out_cache_end], device=self.device
+                ),
+                seq_lens_sum=self.batch_size * total_len,
+                forward_mode=mode,
+                req_pool_indices=torch.arange(self.batch_size, device=self.device),
+                seq_lens=torch.tensor(
+                    [total_len] * self.batch_size, device=self.device
+                ),
+                seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
+            )
+
+        # Pool refs are resolved via the active ForwardContext (published in
+        # setUp). Write the test fixture's req_to_token mapping.
+        self._mock_write_to_req_to_token_pool(self.batch_size, total_len, page_size)
+
+        return forward_batch
+
+    def _setup_kv_cache(self, forward_batch, layer, cache_len):
+        # Create constant values for the prefix cache for easy debugging
+        cache_k = torch.ones(
+            self.batch_size * cache_len,
+            self.num_heads,
+            self.head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        cache_v = (
+            torch.ones(
+                self.batch_size * cache_len,
+                self.num_heads,
+                self.head_dim,
+                dtype=self.dtype,
+                device=self.device,
+            )
+            * 2
+        )
+
+        # Set the prefix KV cache
+        self.model_runner.token_to_kv_pool.set_kv_buffer(
+            layer,
+            torch.arange(self.batch_size * cache_len, device=self.device),
+            cache_k,
+            cache_v,
+            layer.k_scale,
+            layer.v_scale,
+        )
+
+    def _run_attention_test(self, mode, q_len, prefix_len=0, page_size=1):
+        """
+            Run an attention test with the specified parameters.
+        Args:
+            mode: ForwardMode.EXTEND or ForwardMode.DECODE
+            q_len: Length of the query sequence. For decode mode, q_len is 1.
+            prefix_len: Length of the prefix sequence for extend mode
+            page_size: Page size for the KV cache
+        """
+        layer = self._create_attention_layer()
+
+        # Create forward batch and set up
+        forward_batch = self._create_forward_batch(mode, q_len, prefix_len, page_size)
+
+        # Create QKV tensors for the input
+        q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
+
+        # KV cache for prefixed extend is prefix_len
+        # KV cache for decode is same as seq_len
+        # No KV cache for extend without prefix
+        if mode == ForwardMode.EXTEND:
+            if prefix_len > 0:
+                self._setup_kv_cache(forward_batch, layer, prefix_len)
+        else:
+            self._setup_kv_cache(forward_batch, layer, self.seq_len)
+
+        self.backend.init_forward_metadata(forward_batch)
+
+        if mode == ForwardMode.EXTEND:
+            expected_shape = (
+                self.batch_size * q_len,
+                self.num_heads * self.head_dim,
+            )
+            output = self.backend.forward_extend(q, k, v, layer, forward_batch)
+        else:
+            expected_shape = (self.batch_size, self.num_heads * self.head_dim)
+            output = self.backend.forward_decode(q, k, v, layer, forward_batch)
+
+        output_ref = self._run_reference_forward(
+            mode, q, k, v, layer, forward_batch, expected_shape
+        )
+
+        self._verify_output(output, expected_shape, output_ref)
+
+        return output
+
+    def _run_attention_cp_test(self, mode, q_len, prefix_len=0, page_size=1):
+        layer = self._create_attention_layer()
+
+        # Create forward batch and set up
+        forward_batch = self._create_forward_batch(
+            mode, q_len, prefix_len, page_size, attn_cp_size=2
+        )
+        self.backend.attn_cp_size = 2
+
+        # Create QKV tensors for the input
+        q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
+
+        # KV cache for prefixed extend is prefix_len
+        # KV cache for decode is same as seq_len
+        # No KV cache for extend without prefix
+        # Setup KV cache for CP testing - need KV cache to have actual values
+        # For extend with CP, we need KV cache populated so attention has something to attend to
+        self._setup_kv_cache(forward_batch, layer, q_len)
+
+        self.backend.init_forward_metadata(forward_batch)
+
+        # if mode == ForwardMode.EXTEND:
+        expected_shape = (
+            self.batch_size * q_len,
+            self.num_heads * self.head_dim,
+        )
+
+        output = self.backend.forward_extend(q, k, v, layer, forward_batch)
+        # else:
+        #     expected_shape = (self.batch_size, self.num_heads * self.head_dim)
+        #     output = self.backend.forward_decode(q, k, v, layer, forward_batch)
+
+        output_ref = self._run_reference_forward(
+            mode, q, k, v, layer, forward_batch, expected_shape
+        )
+
+        self._verify_output(output, expected_shape, output_ref)
+
+        return output
+
+    def test_forward_extend_cp(self):
+        """Test the standard extend operation with context parallel."""
+        self._run_attention_cp_test(ForwardMode.EXTEND, q_len=self.seq_len)
+
+    # def test_forward_extend_cp_with_prefix(self):
+    #     """Test the standard extend operation with context parallel and prefix."""
+    #     prefix_len = self.seq_len // 2
+    #     extend_len = self.seq_len - prefix_len
+    #     self._run_attention_cp_test(ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len)
+
+    # def test_forward_extend(self):
+    #     """Test the standard extend operation."""
+    #     self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
+
+    # def test_forward_decode(self):
+    #     """Test the decode operation with cached tokens."""
+    #     self._run_attention_test(ForwardMode.DECODE, q_len=1)
+
+    # def test_forward_extend_with_prefix(self):
+    #     """Test extending from cached prefix tokens."""
+    #     prefix_len = self.seq_len // 2
+    #     extend_len = self.seq_len - prefix_len
+    #     self._run_attention_test(
+    #         ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
+    #     )
+
+    # def test_forward_extend_with_page_size_greater_than_1(self):
+    #     """Test extending from cached prefix tokens with page size greater than 1."""
+    #     self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len, page_size=64)
+
+    # def test_forward_decode_with_page_size_greater_than_1(self):
+    #     """Test decode operation with page size greater than 1."""
+    #     self._run_attention_test(ForwardMode.DECODE, q_len=1, page_size=64)
+
+
+class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
+    """
+    All the test cases examples have 1 additional cache location than the decode length.
+    This is to align with the current allocation logic. It does not affect the correctness.
+    """
+
+    def test_draft_decode_set_expand_metadata(self):
+        bs, topk, page_size = 1, 2, 4
+
+        cases = [
+            (
+                torch.tensor(
+                    [
+                        [23, 24],
+                        [31, 32],
+                    ],
+                    dtype=torch.int32,
+                ),
+                torch.tensor(
+                    [
+                        [5, 6],
+                        [7, 8],
+                    ],
+                    dtype=torch.int32,
+                ),
+                1,
+            ),
+            # Decode span multiple pages:
+            # duplicated kv cache: 24, 25, 26
+            # decode locations: 27, 28, 29, 30, 31, 32
+            # We need 3 pages in total.
+            (
+                torch.tensor(
+                    [
+                        [27, 28, 29, 30, 31, 32],
+                        [35, 36, 37, 38, 39, 40],
+                    ],
+                    dtype=torch.int32,
+                ),
+                torch.tensor(
+                    [
+                        [6, 7, 8, 0, 0, 0],
+                        [8, 9, 10, 0, 0, 0],
+                    ],
+                    dtype=torch.int32,
+                ),
+                5,
+            ),
+        ]
+
+        last_page_lens = torch.tensor([3], dtype=torch.int32)
+        for cache_loc, expected_page_table, decode_length in cases:
+            cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
+            page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
+            draft_decode_set_expand_metadata(
+                cache_seqlens_int32=cache_seqlens_int32,
+                page_table=page_table,
+                last_page_lens=last_page_lens,
+                decode_length=decode_length,
+                cache_loc=cache_loc,
+                topk=topk,
+                page_size=page_size,
+            )
+
+            expected_cache_seqlens = torch.tensor(
+                [decode_length + 3, decode_length + 3], dtype=torch.int32
+            )
+            self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
+            self.assertTrue(torch.equal(page_table, expected_page_table))
+
+    def test_draft_decode_set_expand_metadata_page_crossing(self):
+        """
+        Regression for fa3 EAGLE draft decode with topk > 1 and page_size > 1.
+        cache_loc arrives num_steps-wide; callers pre-slice it to `decode_length`
+        (the live draft tokens) before this helper runs, so the dedup'd scatter
+        never writes past the (decode_length + 1)-wide expand page_table row even
+        when consecutive draft tokens land on distinct pages.
+        """
+        bs, topk, page_size = 1, 2, 4
+        decode_length = 2
+        last_page_lens = torch.tensor([3], dtype=torch.int32)
+        # 2 live draft tokens per (batch, topk) crossing into distinct pages.
+        cache_loc = torch.tensor([[23, 28], [31, 36]], dtype=torch.int32)
+        cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
+        # page_table is (decode_length + 1) wide (extra slot for the last partial
+        # page); the trailing column must stay zero.
+        page_table = torch.zeros(bs * topk, decode_length + 1, dtype=torch.int32)
+        draft_decode_set_expand_metadata(
+            cache_seqlens_int32=cache_seqlens_int32,
+            page_table=page_table,
+            last_page_lens=last_page_lens,
+            decode_length=decode_length,
+            cache_loc=cache_loc,
+            topk=topk,
+            page_size=page_size,
+        )
+        expected_page_table = torch.tensor([[5, 7, 0], [7, 9, 0]], dtype=torch.int32)
+        expected_cache_seqlens = torch.tensor([5, 5], dtype=torch.int32)
+        self.assertTrue(torch.equal(page_table, expected_page_table))
+        self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
+
+    def test_update_draft_decode_set_expand_metadata_multi_batch(self):
+        """
+        Ensure expand metadata works when batch size > 1 and last pages differ.
+        """
+        bs, topk, decode_length, page_size = 3, 2, 3, 4
+        cache_loc = torch.tensor(
+            [
+                # First batch: last page duplicate is 1, consecutive pages
+                [1, 2, 3, 4],
+                [6, 7, 8, 9],
+                # Second batch: last page duplicate is 3, non-consecutive pages
+                [3, 8, 9, 10],
+                [14, 15, 16, 17],
+                # Third batch: last page duplicate is 0, consecutive pages
+                [0, 1, 2, 3],
+                [4, 5, 6, 7],
+            ],
+            dtype=torch.int32,
+        )
+        cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
+        last_page_lens = torch.tensor([1, 3, 0], dtype=torch.int32)
+        page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
+        draft_decode_set_expand_metadata(
+            cache_seqlens_int32=cache_seqlens_int32,
+            page_table=page_table,
+            last_page_lens=last_page_lens,
+            decode_length=decode_length,
+            cache_loc=cache_loc,
+            topk=topk,
+            page_size=page_size,
+        )
+
+        expected_cache_seqlens = torch.tensor([4, 4, 6, 6, 3, 3], dtype=torch.int32)
+        expected_page_table = torch.tensor(
+            [
+                [0, 1, 0, 0],
+                [1, 2, 0, 0],
+                [0, 2, 0, 0],
+                [3, 4, 0, 0],
+                [0, 0, 0, 0],
+                [1, 0, 0, 0],
+            ],
+            dtype=torch.int32,
+        )
+        self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
+        self.assertTrue(torch.equal(page_table, expected_page_table))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/attention/test_flashattn_mla_backend.py

@@ -0,0 +1,328 @@
+import unittest
+
+import torch
+
+from sglang.srt.configs.model_config import AttentionArch
+from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_context import (
+    ForwardContext,
+    set_forward_context,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class MockModelRunner:
+    def __init__(
+        self,
+        kv_lora_rank,
+        qk_rope_head_dim,
+    ):
+        attention_arch = AttentionArch.MLA
+        self.device = "cuda"
+        self.dtype = torch.float16
+        self.is_hybrid_swa = False
+        context_len = 2048
+        self.model_config = type(
+            "ModelConfig",
+            (),
+            {
+                "context_len": context_len,
+                "attention_arch": attention_arch,
+                "is_encoder_decoder": False,
+                "is_local_attention_model": False,
+            },
+        )
+        self.sliding_window_size = None
+        # Add server_args attribute
+        self.server_args = type(
+            "ServerArgs",
+            (),
+            {
+                "kv_cache_dtype": torch.float16,
+                "speculative_eagle_topk": None,
+                "speculative_num_draft_tokens": 0,
+                "enable_deterministic_inference": False,
+            },
+        )
+        self.kv_cache_dtype = self.server_args.kv_cache_dtype
+
+        batch_size = 160
+        # Create a proper req_to_token_pool with the req_to_token attribute
+        self.req_to_token_pool = type(
+            "TokenPool",
+            (),
+            {
+                # A typical max_bs * max_context_len for cuda graph decode
+                "size": batch_size,
+                # Add req_to_token attribute
+                "req_to_token": torch.zeros(
+                    batch_size, context_len, dtype=torch.int32, device=self.device
+                ),
+            },
+        )
+        self.page_size = 1
+        max_total_num_tokens = batch_size * context_len
+        self.token_to_kv_pool = MLATokenToKVPool(
+            size=max_total_num_tokens,
+            page_size=self.page_size,
+            dtype=self.kv_cache_dtype,
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            layer_num=1,  # only consider layer=1 for unit test
+            device=self.device,
+            enable_memory_saver=False,
+        )
+        self.hisparse_coordinator = None
+
+
+class MockReqToTokenPool:
+    def __init__(self, batch_size, seq_len, device):
+        self.req_to_token = (
+            torch.arange(batch_size * seq_len, device=device)
+            .reshape(batch_size, seq_len)
+            .to(torch.int32)
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFlashAttentionMLABackend(CustomTestCase):
+    def setUp(self):
+        # MLA with different V headdim requires Hopper architecture (compute capability >= 9.0)
+        if torch.cuda.is_available():
+            compute_capability = torch.cuda.get_device_capability()
+            if compute_capability[0] < 9:
+                self.skipTest(
+                    f"MLA requires Hopper GPU (compute capability >= 9.0), "
+                    f"but found compute capability {compute_capability[0]}.{compute_capability[1]}"
+                )
+
+        # Test parameters
+        self.batch_size = 2
+        self.seq_len = 360
+        self.num_heads = 2
+        self.device = "cuda"
+        self.dtype = torch.float16
+        self.kv_lora_rank = 512
+        self.q_lora_rank = 128
+        self.qk_rope_head_dim = 64
+        self.qk_head_dim = self.qk_rope_head_dim + self.kv_lora_rank
+        # Assume no rope scaling
+        self.scaling = self.qk_head_dim**-0.5
+        # Initialize model runner and backend
+        self._init_model_runner()
+        self.backend = FlashAttentionBackend(self.model_runner)
+        self.ref_backend = TorchNativeAttnBackend(self.model_runner)
+        self.num_local_heads = 2
+        # Publish the backend so RadixAttention.forward resolves correctly.
+        set_forward_context(ForwardContext(attn_backend=self.backend))
+
+    def _init_model_runner(self):
+        self.model_runner = MockModelRunner(
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+        )
+
+    def _create_attention_layer(self):
+        """Create attention layer for testing."""
+        self.attn_mqa = RadixAttention(
+            num_heads=self.num_local_heads,
+            head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
+            scaling=self.scaling,
+            num_kv_heads=1,
+            layer_id=0,
+            v_head_dim=self.kv_lora_rank,
+            prefix="attn_mqa",
+        )
+        return self.attn_mqa
+
+    def _run_reference_forward(
+        self, mode, q, k, v, layer, forward_batch, expected_shape
+    ):
+        """Run reference forward pass using native backend."""
+        if mode == ForwardMode.EXTEND:
+            output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
+        else:  # ForwardMode.DECODE
+            output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
+        return output.view(expected_shape)
+
+    def _verify_output(self, output, expected_shape):
+        """Verify output tensor shape, dtype, and values."""
+        self.assertEqual(
+            output.shape,
+            expected_shape,
+            f"Expected shape {expected_shape}, got {output.shape}",
+        )
+        self.assertEqual(output.dtype, self.dtype)
+        self.assertEqual(output.device.type, "cuda")
+        self.assertEqual(
+            torch.isnan(output).sum().item(), 0, "Output contains NaN values"
+        )
+
+    def _create_forward_batch(self, mode, q_len=None, prefix_len=0):
+        """Create a forward batch for testing based on mode and lengths."""
+        # Default to self.seq_len if not specified
+        q_len = q_len or self.seq_len
+
+        if mode == ForwardMode.EXTEND:
+            total_len = prefix_len + q_len
+            out_cache_start = prefix_len * self.batch_size
+            out_cache_end = total_len * self.batch_size
+
+            forward_batch = ForwardBatch(
+                batch_size=self.batch_size,
+                input_ids=torch.randint(
+                    0, 100, (self.batch_size, q_len), device=self.device
+                ),
+                out_cache_loc=torch.arange(
+                    out_cache_start, out_cache_end, device=self.device
+                ),
+                seq_lens_sum=self.batch_size * total_len,
+                forward_mode=mode,
+                req_pool_indices=torch.arange(self.batch_size, device=self.device),
+                seq_lens=torch.tensor(
+                    [total_len] * self.batch_size, device=self.device
+                ),
+                seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
+                extend_prefix_lens=torch.tensor(
+                    [prefix_len] * self.batch_size, device=self.device
+                ),
+                extend_prefix_lens_cpu=[prefix_len] * self.batch_size,
+                extend_seq_lens=torch.tensor(
+                    [q_len] * self.batch_size, device=self.device
+                ),
+                extend_seq_lens_cpu=[q_len] * self.batch_size,
+            )
+
+        else:  # ForwardMode.DECODE
+            decode_len = q_len  # typically 1 for decode mode
+            total_len = self.seq_len + decode_len
+            out_cache_start = self.batch_size * self.seq_len
+            out_cache_end = self.batch_size * total_len
+
+            forward_batch = ForwardBatch(
+                batch_size=self.batch_size,
+                input_ids=torch.randint(
+                    0, 100, (self.batch_size, decode_len), device=self.device
+                ),
+                out_cache_loc=torch.arange(
+                    out_cache_start, out_cache_end, device=self.device
+                ),
+                seq_lens_sum=self.batch_size * total_len,
+                forward_mode=mode,
+                req_pool_indices=torch.arange(self.batch_size, device=self.device),
+                seq_lens=torch.tensor(
+                    [total_len] * self.batch_size, device=self.device
+                ),
+                seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
+            )
+
+        # Pool refs are resolved via the active ForwardContext (published in
+        # setUp); the fixture no longer needs to attach them to forward_batch.
+        return forward_batch
+
+    def _setup_kv_cache(self, forward_batch, layer, cache_len):
+        """Set up KV cache with prefix tokens."""
+        if cache_len <= 0:
+            return
+
+        # For MLA, create separate nope and rope caches
+        cache_k_nope = torch.ones(
+            self.batch_size * cache_len,
+            1,  # latent cache has only one head in MQA
+            self.kv_lora_rank,
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+        cache_k_rope = torch.ones(
+            self.batch_size * cache_len,
+            1,  # latent cache has only one head in MQA
+            self.qk_rope_head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+        # Set the prefix KV cache using MLA-specific method
+        self.model_runner.token_to_kv_pool.set_mla_kv_buffer(
+            layer,
+            torch.arange(self.batch_size * cache_len, device=self.device),
+            cache_k_nope,
+            cache_k_rope,
+        )
+
+    def _run_attention_test(self, mode, q_len, prefix_len=0):
+        """
+            Run an attention test with the specified parameters.
+        Args:
+            mode: ForwardMode.EXTEND or ForwardMode.DECODE
+            q_len: Length of the query sequence. For decode mode, q_len is 1.
+            prefix_len: Length of the prefix sequence for extend mode
+        """
+        layer = self._create_attention_layer()
+
+        # Create forward batch and set up
+        forward_batch = self._create_forward_batch(mode, q_len, prefix_len)
+
+        # Create q, kv_compressed for testing
+        q_shape = (self.batch_size * q_len, self.num_heads, self.qk_head_dim)
+        kv_shape = (self.batch_size * q_len, self.qk_head_dim)
+        q = torch.randn(q_shape, dtype=self.dtype, device=self.device)
+        kv_compressed = torch.randn(kv_shape, dtype=self.dtype, device=self.device)
+
+        # For MLA, split kv_compressed into k_nope and k_rope
+        # k_nope has dimension kv_lora_rank, k_rope has dimension qk_rope_head_dim
+        k_nope = kv_compressed[:, : self.kv_lora_rank]
+        k_rope = kv_compressed[:, self.kv_lora_rank :]
+
+        # k_nope needs to be unsqueezed for the num_heads dimension
+        k = k_nope.unsqueeze(1)
+        # k_rope also needs to be unsqueezed
+        k_rope = k_rope.unsqueeze(1)
+
+        # v is not used for mqa
+        v = torch.randn((1), dtype=self.dtype, device=self.device)
+
+        self._setup_kv_cache(forward_batch, layer, prefix_len)
+
+        self.backend.init_forward_metadata(forward_batch)
+
+        expected_shape = (
+            self.batch_size * q_len,
+            self.num_heads * self.kv_lora_rank,
+        )
+
+        if mode == ForwardMode.EXTEND:
+            output = self.backend.forward_extend(
+                q, k, v, layer, forward_batch, k_rope=k_rope
+            )
+        else:
+            output = self.backend.forward_decode(
+                q, k, v, layer, forward_batch, k_rope=k_rope
+            )
+
+        self._verify_output(output, expected_shape)
+        return output
+
+    def test_forward_extend(self):
+        """Test the standard extend operation."""
+        self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
+
+    def test_forward_decode(self):
+        """Test the decode operation with cached tokens."""
+        self._run_attention_test(ForwardMode.DECODE, q_len=1)
+
+    def test_forward_extend_with_prefix(self):
+        """Test extending from cached prefix tokens."""
+        prefix_len = self.seq_len // 2
+        extend_len = self.seq_len - prefix_len
+        self._run_attention_test(
+            ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/attention/test_local_attn.py

@@ -0,0 +1,74 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import get_device_sm, kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_LOCAL_ATTENTION,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+# Local attention with FA3 (requires SM 90+ / H100, tp=4)
+@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
+class TestFlashAttention3LocalAttn(CustomTestCase):
+    model = DEFAULT_MODEL_NAME_FOR_TEST_LOCAL_ATTENTION
+    base_url = DEFAULT_URL_FOR_TEST
+    accuracy_threshold = 0.90
+
+    @classmethod
+    def get_server_args(cls):
+        return [
+            "--cuda-graph-max-bs-decode",
+            "2",
+            "--attention-backend",
+            "fa3",
+            "--tp",
+            "4",
+            "--context-length",
+            "1000000",
+        ]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=cls.get_server_args(),
+            env=os.environ,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=100,
+            num_threads=128,
+            num_shots=4,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        # Use the appropriate metric key based on the test class
+        metric_key = "score"
+        self.assertGreater(metrics[metric_key], self.accuracy_threshold)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/attention/test_prefix_chunk_info.py

@@ -0,0 +1,238 @@
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.utils.common import get_device
+from sglang.test.test_utils import CustomTestCase
+
+TEST_CASES = [
+    # Sequence with same prefix lens
+    {
+        "batch_size": 3,
+        "prefix_lens": [64, 64, 64],
+        "max_chunk_capacity": 48,
+        "prefix_chunk_len": 16,
+        "num_prefix_chunks": 4,
+        "prefix_chunk_starts": torch.tensor(
+            [
+                [0, 0, 0],
+                [16, 16, 16],
+                [32, 32, 32],
+                [48, 48, 48],
+            ],
+            dtype=torch.int32,
+        ),
+        "prefix_chunk_seq_lens": torch.tensor(
+            [
+                [16, 16, 16],
+                [16, 16, 16],
+                [16, 16, 16],
+                [16, 16, 16],
+            ],
+            dtype=torch.int32,
+        ),
+    },
+    # Sequence with different prefix lens
+    {
+        "batch_size": 4,
+        "prefix_lens": [16, 32, 48, 64],
+        "max_chunk_capacity": 64,
+        "prefix_chunk_len": 16,
+        "num_prefix_chunks": 4,
+        "prefix_chunk_starts": torch.tensor(
+            [
+                [0, 0, 0, 0],
+                [16, 16, 16, 16],
+                [32, 32, 32, 32],
+                [48, 48, 48, 48],
+            ],
+            dtype=torch.int32,
+        ),
+        "prefix_chunk_seq_lens": torch.tensor(
+            [
+                [16, 16, 16, 16],
+                [0, 16, 16, 16],
+                [0, 0, 16, 16],
+                [0, 0, 0, 16],
+            ],
+            dtype=torch.int32,
+        ),
+    },
+    # Sequence with irregular shapes
+    {
+        "batch_size": 2,
+        "prefix_lens": [1, 64],
+        "max_chunk_capacity": 31,
+        "prefix_chunk_len": 15,
+        "num_prefix_chunks": 5,
+        "prefix_chunk_starts": torch.tensor(
+            [
+                [0, 0],
+                [15, 15],
+                [30, 30],
+                [45, 45],
+                [60, 60],
+            ],
+            dtype=torch.int32,
+        ),
+        "prefix_chunk_seq_lens": torch.tensor(
+            [
+                [1, 15],
+                [0, 15],
+                [0, 15],
+                [0, 15],
+                [0, 4],
+            ],
+            dtype=torch.int32,
+        ),
+    },
+]
+
+
+class MockForwardBatch(ForwardBatch):
+    def __init__(self, max_chunk_capacity: int, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_chunk_capacity = max_chunk_capacity
+
+    def get_max_chunk_capacity(self):
+        return self.max_chunk_capacity
+
+
+class MockReqToTokenPool:
+    def __init__(self, batch_size, seq_len, device):
+        self.req_to_token = (
+            torch.arange(batch_size * seq_len, device=device)
+            .reshape(batch_size, seq_len)
+            .to(torch.int32)
+        )
+
+
+# Test correctness of triton kernel for computing kv indices
+def check_kv_indices(forward_batch, req_to_token_pool):
+    for i in range(forward_batch.num_prefix_chunks):
+        computed_kv_indices = forward_batch.prefix_chunk_kv_indices[i]
+        req_to_token = req_to_token_pool.req_to_token[: forward_batch.batch_size, :]
+        ref_kv_indices = torch.empty(
+            forward_batch.prefix_chunk_num_tokens[i],
+            dtype=torch.int32,
+            device=computed_kv_indices.device,
+        )
+        running_ptr = 0
+        for j in range(forward_batch.batch_size):
+            seq_start = forward_batch.prefix_chunk_starts[i, j].item()
+            seq_len = forward_batch.prefix_chunk_seq_lens[i, j].item()
+            ref_kv_indices[running_ptr : running_ptr + seq_len].copy_(
+                req_to_token[j, seq_start : seq_start + seq_len]
+            )
+            running_ptr += seq_len
+        assert torch.allclose(computed_kv_indices, ref_kv_indices)
+
+
+@unittest.skipIf(
+    not (torch.cuda.is_available() or torch.xpu.is_available()),
+    "Test requires CUDA or XPU",
+)
+class TestPrefixChunkInfo(CustomTestCase):
+    def setUp(self):
+        # Common test parameters
+        self.num_local_heads = 128
+        self.kv_lora_rank = 512
+        self.qk_rope_head_dim = 64
+        self.device = get_device()
+        self.dtype = torch.bfloat16
+        self.extend_len = 64
+        self.max_bs = 4
+        self.max_seq_len = 128
+
+        # req_to_token_pool
+        self.req_to_token_pool = MockReqToTokenPool(
+            self.max_bs,
+            self.max_seq_len,
+            self.device,
+        )
+
+        # token_to_kv_pool
+        self.token_to_kv_pool = MLATokenToKVPool(
+            size=self.max_bs * self.max_seq_len,
+            page_size=1,  # only consider page=1 for unit test
+            dtype=self.dtype,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            layer_num=1,  # only consider layer=1 for unit test
+            device=self.device,
+            enable_memory_saver=False,
+        )
+
+    def test_prefix_chunk_info(self):
+        """Test the standard extend operation."""
+
+        for test_case in TEST_CASES:
+            print(
+                f"Test case with batch_size={test_case['batch_size']}, prefix_lens={test_case['prefix_lens']}, max_chunk_capacity={test_case['max_chunk_capacity']}"
+            )
+            batch_size = test_case["batch_size"]
+            prefix_lens_cpu = test_case["prefix_lens"]
+            assert len(prefix_lens_cpu) == batch_size
+            prefix_lens = torch.tensor(prefix_lens_cpu, device=self.device)
+            max_chunk_capacity = test_case["max_chunk_capacity"]
+            seq_lens_cpu = [
+                self.extend_len + prefix_lens_cpu[i] for i in range(batch_size)
+            ]
+            seq_lens = torch.tensor(seq_lens_cpu, device=self.device)
+
+            # Create forward batch
+            # input_ids and out_cache_loc are dummy tensors in this test
+            forward_batch = MockForwardBatch(
+                max_chunk_capacity=max_chunk_capacity,
+                batch_size=batch_size,
+                input_ids=torch.randint(
+                    0, 100, (batch_size, self.extend_len), device=self.device
+                ),
+                out_cache_loc=torch.arange(
+                    self.max_bs * self.max_seq_len - batch_size * self.extend_len,
+                    self.max_bs * self.max_seq_len,
+                    device=self.device,
+                ),
+                seq_lens_sum=sum(seq_lens_cpu),
+                forward_mode=ForwardMode.EXTEND,
+                req_pool_indices=torch.arange(batch_size, device=self.device),
+                seq_lens=seq_lens,
+                seq_lens_cpu=seq_lens_cpu,
+                extend_prefix_lens=prefix_lens,
+                extend_prefix_lens_cpu=prefix_lens_cpu,
+            )
+            # Pool refs are resolved via the active ForwardContext; mock an
+            # attn_backend that carries the pools (Pattern A invariant).
+            from types import SimpleNamespace
+
+            from sglang.srt.model_executor.forward_context import (
+                ForwardContext,
+                set_forward_context,
+            )
+
+            mock_backend = SimpleNamespace(
+                req_to_token_pool=self.req_to_token_pool,
+                token_to_kv_pool=self.token_to_kv_pool,
+            )
+            set_forward_context(ForwardContext(attn_backend=mock_backend))
+
+            forward_batch.prepare_chunked_prefix_cache_info(self.device)
+            assert forward_batch.get_max_chunk_capacity() == max_chunk_capacity
+            assert forward_batch.prefix_chunk_len == test_case["prefix_chunk_len"]
+            assert forward_batch.num_prefix_chunks == test_case["num_prefix_chunks"]
+            assert torch.allclose(
+                forward_batch.prefix_chunk_starts,
+                test_case["prefix_chunk_starts"].to(self.device),
+            )
+            assert torch.allclose(
+                forward_batch.prefix_chunk_seq_lens,
+                test_case["prefix_chunk_seq_lens"].to(self.device),
+            )
+
+            check_kv_indices(forward_batch, self.req_to_token_pool)
+
+
+if __name__ == "__main__":
+    unittest.main()

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 check "
+        "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()

test/manual/core/test_dynamic_grad_mode.py

@@ -0,0 +1,58 @@
+import unittest
+
+import torch
+
+from sglang.srt.utils import DynamicGradMode
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestDynamicGradMode(CustomTestCase):
+    def test_inference(self):
+        # Test inference_mode
+        DynamicGradMode.set_inference_mode(True)
+
+        @DynamicGradMode()
+        def create_tensor_x():
+            return torch.empty(0)
+
+        X = create_tensor_x()
+        self.assertTrue(not X.requires_grad and X.is_inference())
+
+    def test_no_grad(self):
+        # Test no_grad
+        DynamicGradMode.set_inference_mode(False)
+
+        @DynamicGradMode()
+        def create_tensor_y():
+            return torch.empty(0)
+
+        Y = create_tensor_y()
+        self.assertTrue(not Y.requires_grad and not Y.is_inference())
+
+    def test_nested_inference(self):
+        # Test no_grad nested inference_mode, inference_mode should has higher priority
+        DynamicGradMode.set_inference_mode(False)
+
+        @DynamicGradMode()
+        def create_tensor_z():
+            with torch.inference_mode():
+                return torch.empty(0)
+
+        Z = create_tensor_z()
+        self.assertTrue(not Z.requires_grad and Z.is_inference())
+
+    def test_nested_no_grad(self):
+        # Test inference_mode nested no_grad, inference_mode should has higher priority
+        DynamicGradMode.set_inference_mode(True)
+
+        @DynamicGradMode()
+        def create_tensor_w():
+            with torch.no_grad():
+                return torch.empty(0)
+
+        W = create_tensor_w()
+        self.assertTrue(not W.requires_grad and W.is_inference())
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

test/manual/core/test_gpt_oss_1gpu.py

@@ -0,0 +1,31 @@
+import unittest
+
+from sglang.test.gpt_oss_common import BaseTestGptOss
+
+
+class TestGptOss1Gpu(BaseTestGptOss):
+    def test_mxfp4_20b(self):
+        self.run_test(
+            model_variant="20b",
+            quantization="mxfp4",
+            expected_score_of_reasoning_effort={
+                "low": 0.34,
+                "medium": 0.34,
+                "high": 0.27,  # TODO investigate
+            },
+        )
+
+    def test_bf16_20b(self):
+        self.run_test(
+            model_variant="20b",
+            quantization="bf16",
+            expected_score_of_reasoning_effort={
+                "low": 0.34,
+                "medium": 0.34,
+                "high": 0.27,  # TODO investigate
+            },
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/core/test_session_radix_cache.py

@@ -0,0 +1,174 @@
+"""Manual test for the session radix cache (--enable-session-radix-cache).
+Run directly: python test/manual/core/test_session_radix_cache.py
+"""
+
+import unittest
+from array import array
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import (
+    EvictParams,
+    InsertParams,
+    MatchPrefixParams,
+)
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
+
+
+class TestSessionRadixCache(unittest.TestCase):
+    def setUp(self):
+        dtype = torch.float16
+        kv = MHATokenToKVPool(
+            size=64,
+            page_size=1,
+            dtype=dtype,
+            head_num=2,
+            head_dim=8,
+            layer_num=1,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+        allocator = TokenToKVPoolAllocator(
+            size=64, dtype=dtype, device="cpu", kvcache=kv, need_sort=False
+        )
+        req_to_token_pool = ReqToTokenPool(
+            size=8, max_context_len=1024, device="cpu", enable_memory_saver=False
+        )
+        self.cache = RadixCache(
+            CacheInitParams(
+                disable=False,
+                req_to_token_pool=req_to_token_pool,
+                token_to_kv_pool_allocator=allocator,
+                page_size=1,
+                eviction_policy="lru",
+                enable_kv_cache_events=False,
+                enable_session_radix_cache=True,
+            )
+        )
+
+    def _insert(self, toks):
+        idx = self.cache.token_to_kv_pool_allocator.alloc(len(toks))
+        self.cache.insert(
+            InsertParams(key=RadixKey(array("q", toks)), value=idx.to(torch.int64))
+        )
+
+    def _tag(self, toks, sid):
+        self.cache._tag_session_leaf(
+            SimpleNamespace(session_id=sid),
+            RadixKey(array("q", toks)),
+            node=self._leaf(toks),
+        )
+
+    def _cached(self, toks):
+        return int(
+            self.cache.match_prefix(
+                MatchPrefixParams(key=RadixKey(array("q", toks)))
+            ).device_indices.numel()
+        )
+
+    def _leaf(self, toks):
+        return self.cache.match_prefix(
+            MatchPrefixParams(key=RadixKey(array("q", toks)))
+        ).last_device_node
+
+    def test_tag_with_known_node_skips_match_prefix(self):
+        self._insert([1, 2, 3, 4])
+        leaf = self._leaf([1, 2, 3, 4])
+        orig_match_prefix = self.cache.match_prefix
+
+        def fail_match_prefix(_params):
+            raise AssertionError("match_prefix should not run when node is supplied")
+
+        self.cache.match_prefix = fail_match_prefix
+        try:
+            self.cache._tag_session_leaf(
+                SimpleNamespace(session_id="S"),
+                RadixKey(array("q", [1, 2, 3, 4])),
+                node=leaf,
+            )
+        finally:
+            self.cache.match_prefix = orig_match_prefix
+        self.assertEqual(getattr(leaf, "session_ids", None), {"S"})
+        self.assertIn(leaf, self.cache._session_leaves["S"])
+
+    def test_disabled_cache_does_not_tag_session_kv(self):
+        self.cache.enable_session_radix_cache = False
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "S")
+        self.assertIsNone(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None))
+
+    def test_shared_prefix_frees_only_unique_tail(self):
+        # A/B share prefix [1,2]; close(A) frees only A's tail, B + shared stay.
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "A")
+        self._insert([1, 2, 5, 6])
+        self._tag([1, 2, 5, 6], "B")
+        self.assertGreater(self.cache.release_radix_session("A"), 0)
+        self.assertEqual(self._cached([1, 2, 3, 4]), 2)  # only shared [1,2] left
+        self.assertEqual(self._cached([1, 2, 5, 6]), 4)  # B intact
+
+    def test_same_leaf_freed_only_on_last_holder(self):
+        # Identical content -> one leaf held by {A,B}; freed only on last close.
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "A")
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "B")
+        self.assertEqual(
+            getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A", "B"}
+        )
+        self.assertEqual(self.cache.release_radix_session("A"), 0)  # B still holds
+        self.assertEqual(self._cached([1, 2, 3, 4]), 4)
+        self.assertEqual(self.cache.release_radix_session("B"), 1)  # last holder frees
+        self.assertEqual(self._cached([1, 2, 3, 4]), 0)
+
+    def test_legacy_release_does_not_release_radix_session(self):
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "S")
+        self.cache.release_session("S")
+        self.assertEqual(self._cached([1, 2, 3, 4]), 4)
+        self.assertEqual(self.cache.release_radix_session("S"), 1)
+
+    def test_tag_is_lru_neutral_not_pinned(self):
+        # The tag must add no lock/pin: a tagged, never-closed node is evictable.
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "S")
+        leaf = self._leaf([1, 2, 3, 4])
+        self.assertEqual(leaf.lock_ref, 0)
+        self.assertEqual(self.cache.protected_size(), 0)
+        self.assertIn(leaf, self.cache.evictable_leaves)
+        self.cache.evict(EvictParams(num_tokens=4))  # LRU reclaims it while open
+        self.assertEqual(self._cached([1, 2, 3, 4]), 0)
+        self.assertNotIn("S", self.cache._session_leaves)
+        self.assertEqual(
+            self.cache.release_radix_session("S"), 0
+        )  # late close is a no-op
+
+    def test_close_tombstone_blocks_late_finish(self):
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "S")
+        self.assertEqual(self.cache.release_radix_session("S"), 1)
+
+        self._insert([5, 6, 7, 8])
+        self._tag([5, 6, 7, 8], "S")  # simulates a finish racing after close
+        self.assertIsNone(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None))
+
+    def test_tombstoned_shared_holder_cannot_retag_after_last_holder_close(self):
+        self._insert([1, 2, 3, 4])
+        self._tag([1, 2, 3, 4], "A")
+        self._tag([1, 2, 3, 4], "B")
+
+        self.assertEqual(self.cache.release_radix_session("B"), 0)
+        self.assertEqual(getattr(self._leaf([1, 2, 3, 4]), "session_ids", None), {"A"})
+        self.assertEqual(self.cache.release_radix_session("A"), 1)
+
+        self._insert([5, 6, 7, 8])
+        self._tag([5, 6, 7, 8], "B")
+        self.assertIsNone(getattr(self._leaf([5, 6, 7, 8]), "session_ids", None))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/cpu/test_comm.py

@@ -0,0 +1,116 @@
+import copy
+import multiprocessing
+import os
+import traceback
+import unittest
+from multiprocessing import Process
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from sglang.test.test_utils import CustomTestCase, find_available_port
+
+
+def run_distributed_test(rank, world_size, master_port, output_writer, fn):
+    try:
+        os.environ["RANK"] = str(rank)
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+        os.environ["LOCAL_SIZE"] = str(world_size)
+
+        dist.init_process_group("gloo", rank=rank, world_size=world_size)
+        torch.ops.sgl_kernel.initialize(world_size, rank)
+
+        fn(rank, world_size)
+
+        execution_ok = True
+    except Exception as e:
+        print(f"subprocess[{rank=}] has error: {e}", flush=True)
+        traceback.print_exc()
+        execution_ok = False
+
+    output_writer.send(execution_ok)
+    output_writer.close()
+
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def all_reduce_fn(rank, world_size):
+    op = dist.ReduceOp.SUM
+    for dtype in [torch.float32, torch.bfloat16, torch.float16]:
+        tensor = torch.randn(2, 10, dtype=dtype)
+        tensor_shm = copy.deepcopy(tensor)
+
+        dist.all_reduce(tensor, op=op)
+        torch.ops.sgl_kernel.shm_allreduce(tensor_shm, op)
+
+        torch.testing.assert_close(tensor, tensor_shm)
+
+
+def all_gather_fn(rank, world_size):
+    dim = -1
+
+    for dtype in [torch.float32, torch.bfloat16, torch.float16]:
+        tensor = torch.randn(2, 10, dtype=dtype)
+
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += tensor.dim()
+
+        input_size = tensor.size()
+        output_size = (input_size[0] * world_size,) + input_size[1:]
+        output_tensor = torch.empty(
+            output_size, dtype=tensor.dtype, device=tensor.device
+        )
+        dist.all_gather_into_tensor(output_tensor, tensor)
+        output_tensor = output_tensor.reshape((world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim] + (world_size * input_size[dim],) + input_size[dim + 1 :]
+        )
+
+        output_shm = torch.ops.sgl_kernel.shm_allgather(tensor, dim)
+
+        torch.testing.assert_close(output_tensor, output_shm)
+
+
+class TestComm(CustomTestCase):
+    def _spawn_and_check(self, fn, world_size=2):
+        mp.set_start_method("spawn", force=True)
+        master_port = find_available_port(23456)
+
+        processes = []
+        output_reader, output_writer = multiprocessing.Pipe(duplex=False)
+
+        for rank in range(world_size):
+            p = Process(
+                target=run_distributed_test,
+                kwargs=dict(
+                    rank=rank,
+                    world_size=world_size,
+                    master_port=master_port,
+                    output_writer=output_writer,
+                    fn=fn,
+                ),
+            )
+            p.start()
+            processes.append(p)
+
+        for _ in range(world_size):
+            self.assertTrue(output_reader.recv(), "Subprocess fail. Check logs above.")
+
+        for p in processes:
+            p.join()
+
+    def test_all_reduce(self):
+        self._spawn_and_check(all_reduce_fn)
+
+    def test_all_gather(self):
+        self._spawn_and_check(all_gather_fn)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/debug_utils/get_logits_ut.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+
+
+class DummyModel(nn.Module):
+    def __init__(self, d_in=2048, n_heads=128, softmax_scale=0.5):
+        super().__init__()
+        self.weights_proj = nn.Linear(d_in, 1024)
+        self.n_heads = n_heads
+        self.softmax_scale = softmax_scale
+
+    def _get_logits_head_gate_orig(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        weights = weights * self.n_heads**-0.5
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
+        return weights
+
+    def _get_logits_head_gate_opt(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        scale_const = self.n_heads**-0.5 * q_scale * self.softmax_scale  # (B,1,1)
+        weights = weights.unsqueeze(-1) * scale_const  # (B,1024,1)
+        return weights
+
+
+def main():
+    torch.manual_seed(0)
+    model = DummyModel(d_in=2048, n_heads=128, softmax_scale=0.5)
+    x = torch.randn(128, 2048)  # batch=128, d_in=2048
+    q_scale = torch.randn(128, 1)
+
+    import time
+
+    start = time.time()
+    for _ in range(1000):
+        out_orig = model._get_logits_head_gate_orig(x, q_scale)
+    print("Original version time:", time.time() - start)
+
+    start = time.time()
+    for _ in range(1000):
+        out_opt = model._get_logits_head_gate_opt(x, q_scale)
+    print("Optimized version time:", time.time() - start)
+
+    print("Difference:", (out_orig - out_opt).abs().max().item())
+    assert torch.allclose(out_orig, out_opt), "Mismatch between original and optimized"
+
+
+if __name__ == "__main__":
+    main()
+
+
+"""
+Original version time: 0.49235057830810547
+Optimized version time: 0.4087331295013428
+Difference: 1.4901161193847656e-08
+"""

test/manual/debug_utils/run_with_retry.py

@@ -0,0 +1,53 @@
+"""
+Simple wrapper to run a test file with retry logic.
+
+Usage:
+    python3 -m sglang.test.ci.run_with_retry test_file.py [--max-attempts 2] [--retry-wait 60]
+"""
+
+import argparse
+import sys
+
+from sglang.test.ci.ci_utils import TestFile, run_unittest_files
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run a test file with retry logic")
+    parser.add_argument("test_file", help="The test file to run")
+    parser.add_argument(
+        "--max-attempts",
+        type=int,
+        default=2,
+        help="Maximum number of attempts (default: 2)",
+    )
+    parser.add_argument(
+        "--retry-wait",
+        type=int,
+        default=60,
+        help="Seconds to wait between retries (default: 60)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=1200,
+        help="Timeout per attempt in seconds (default: 1200)",
+    )
+    args = parser.parse_args()
+
+    # Create a TestFile with a reasonable estimated time
+    test_file = TestFile(name=args.test_file, estimated_time=args.timeout)
+
+    exit_code = run_unittest_files(
+        files=[test_file],
+        timeout_per_file=args.timeout,
+        continue_on_error=False,
+        enable_retry=True,
+        max_attempts=args.max_attempts,
+        retry_wait_seconds=args.retry_wait,
+    )
+
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

test/manual/debug_utils/test_dump_metric.py

@@ -0,0 +1,132 @@
+"""Unit tests for dump_metric() function."""
+
+import json
+import os
+import tempfile
+import unittest
+from pathlib import Path
+
+from sglang.test.test_utils import dump_metric
+
+
+class TestDumpMetric(unittest.TestCase):
+    """Test suite for dump_metric() function."""
+
+    _ENV_KEYS_TO_CLEAN = ["SGLANG_TEST_METRICS_OUTPUT", "PYTEST_CURRENT_TEST"]
+
+    def setUp(self):
+        """Clean up env vars before each test."""
+        for key in self._ENV_KEYS_TO_CLEAN:
+            os.environ.pop(key, None)
+
+    def tearDown(self):
+        """Clean up env vars after each test."""
+        for key in self._ENV_KEYS_TO_CLEAN:
+            os.environ.pop(key, None)
+
+    def test_writes_valid_jsonl(self):
+        """Test that dump_metric writes one valid JSON line when env is set."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_path = os.path.join(tmpdir, "metrics")
+            os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
+
+            dump_metric("test_accuracy", 0.95, labels={"model": "llama"})
+
+            # Check file exists with PID suffix
+            pid = os.getpid()
+            jsonl_path = f"{base_path}.{pid}.jsonl"
+            self.assertTrue(os.path.exists(jsonl_path))
+
+            # Read and validate
+            with open(jsonl_path, encoding="utf-8") as f:
+                lines = f.readlines()
+
+            self.assertEqual(len(lines), 1)
+            record = json.loads(lines[0])
+
+            # Validate required fields
+            self.assertIn("filename", record)
+            self.assertIn("test_case", record)
+            self.assertEqual(record["metric_name"], "test_accuracy")
+            self.assertEqual(record["value"], 0.95)
+
+            # Validate optional fields
+            self.assertIn("ts", record)
+            self.assertIsInstance(record["ts"], (int, float))
+            self.assertEqual(record["labels"], {"model": "llama"})
+
+    def test_no_env_no_file(self):
+        """Test that dump_metric doesn't create file when env var not set."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Don't set env var
+            dump_metric("test_metric", 42)
+
+            # Verify no files created
+            files = list(Path(tmpdir).glob("*.jsonl"))
+            self.assertEqual(len(files), 0)
+
+    def test_labels_not_serializable_stringified(self):
+        """Test that non-serializable labels are stringified."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_path = os.path.join(tmpdir, "metrics")
+            os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
+
+            # Non-serializable label
+            class NonSerializable:
+                pass
+
+            dump_metric("test_metric", 100, labels={"obj": NonSerializable()})
+
+            pid = os.getpid()
+            jsonl_path = f"{base_path}.{pid}.jsonl"
+            with open(jsonl_path, encoding="utf-8") as f:
+                lines = f.readlines()
+
+            record = json.loads(lines[0])
+            self.assertIn("labels", record)
+            self.assertIsInstance(record["labels"], str)
+
+    def test_bool_to_int(self):
+        """Test that bool values are converted to int."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_path = os.path.join(tmpdir, "metrics")
+            os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
+
+            dump_metric("bool_true", True)
+            dump_metric("bool_false", False)
+
+            pid = os.getpid()
+            jsonl_path = f"{base_path}.{pid}.jsonl"
+            with open(jsonl_path, encoding="utf-8") as f:
+                lines = f.readlines()
+
+            self.assertEqual(len(lines), 2)
+            record1 = json.loads(lines[0])
+            record2 = json.loads(lines[1])
+
+            self.assertEqual(record1["value"], 1)  # True -> 1
+            self.assertEqual(record2["value"], 0)  # False -> 0
+
+    def test_pytest_current_test_parsing(self):
+        """Test PYTEST_CURRENT_TEST parsing for test_case."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            base_path = os.path.join(tmpdir, "metrics")
+            os.environ["SGLANG_TEST_METRICS_OUTPUT"] = base_path
+            os.environ["PYTEST_CURRENT_TEST"] = (
+                "test/srt/test_example.py::TestClass::test_method (call)"
+            )
+
+            dump_metric("pytest_metric", 123)
+
+            pid = os.getpid()
+            jsonl_path = f"{base_path}.{pid}.jsonl"
+            with open(jsonl_path, encoding="utf-8") as f:
+                lines = f.readlines()
+
+            record = json.loads(lines[0])
+            # Only assert test_case parsing, not filename
+            self.assertEqual(record["test_case"], "TestClass.test_method")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/debug_utils/test_log_parser.py

@@ -0,0 +1,28 @@
+import json
+import unittest
+
+from sglang.srt.debug_utils import log_parser
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestLogParser(CustomTestCase):
+    def test_log_parser(self):
+        lines = """
+(SGLangEngine pid=35555) [2025-10-31 03:45:20 TP0] Decode batch [51341], #running-req: 317, #token: 1094261, token usage: 0.67, cuda graph: True, gen throughput (token/s): 14806.57, #queue-req: 0,
+(SGLangEngine pid=111711, ip=10.15.36.1) [2025-10-31 03:45:20 TP0] Decode batch [39913], #running-req: 78, #token: 432100, token usage: 0.27, cuda graph: True, gen throughput (token/s): 7269.16, #queue-req: 0,
+[2025-11-03 14:31:10 DP6 TP6 EP6] Decode batch, #running-req: 251, #token: 2811200, token usage: 1.00, cuda graph: True, gen throughput (token/s): 2055.94, #queue-req: 655,
+"""
+        expect_rows = json.loads(
+            """[{"line":"(SGLangEngine pid=35555) [2025-10-31 03:45:20 TP0] Decode batch [51341], #running-req: 317, #token: 1094261, token usage: 0.67, cuda graph: True, gen throughput (token/s): 14806.57, #queue-req: 0,","1":"(SGLangEngine pid=35555)","pid":35555,"ip":null,"time":"2025-10-31 03:45:20","dp_rank":null,"tp_rank":0,"ep_rank":null,"pp_rank":null,"9":" [51341]","num_running_req":317,"num_token":1094261,"token_usage":0.67,"gen_throughput":14806.57,"queue_req":0},{"line":"(SGLangEngine pid=111711, ip=10.15.36.1) [2025-10-31 03:45:20 TP0] Decode batch [39913], #running-req: 78, #token: 432100, token usage: 0.27, cuda graph: True, gen throughput (token/s): 7269.16, #queue-req: 0,","1":"(SGLangEngine pid=111711, ip=10.15.36.1)","pid":111711,"ip":"10.15.36.1","time":"2025-10-31 03:45:20","dp_rank":null,"tp_rank":0,"ep_rank":null,"pp_rank":null,"9":" [39913]","num_running_req":78,"num_token":432100,"token_usage":0.27,"gen_throughput":7269.16,"queue_req":0},{"line":"[2025-11-03 14:31:10 DP6 TP6 EP6] Decode batch, #running-req: 251, #token: 2811200, token usage: 1.00, cuda graph: True, gen throughput (token/s): 2055.94, #queue-req: 655,","1":null,"pid":null,"ip":null,"time":"2025-11-03 14:31:10","dp_rank":6,"tp_rank":6,"ep_rank":6,"pp_rank":null,"9":null,"num_running_req":251,"num_token":2811200,"token_usage":1.0,"gen_throughput":2055.94,"queue_req":655}]""",
+        )
+
+        df = log_parser.parse(lines)
+        print(df)
+        print(df.write_json())
+
+        assert len(df) == len(lines.strip().splitlines()), f"{len(df)=}"
+        self.assertEqual(json.loads(df.write_json()), expect_rows)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/distributed/test_dp_attention_archived.py

@@ -0,0 +1,101 @@
+"""Archived test classes split out of test/registered/distributed/test_dp_attention.py.
+
+Originally registered with `register_cuda_ci(...)`. Moved here as part of
+the per-commit pruning effort to keep the code reachable manually.
+Run with `python3 test/manual/distributed/test_dp_attention_archived.py`.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.ebnf_constrained_kit import EBNFConstrainedMixin
+from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
+from sglang.test.kits.regex_constrained_kit import RegexConstrainedMixin
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_amd_ci,
+    popen_launch_server,
+)
+
+
+class TestDPAttentionDP2TP2DeepseekV3MTP(
+    CustomTestCase,
+    JSONConstrainedMixin,
+    EBNFConstrainedMixin,
+    RegexConstrainedMixin,
+):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--disable-radix",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "2",
+            "--speculative-eagle-topk",
+            "4",
+            "--speculative-num-draft-tokens",
+            "4",
+            "--speculative-draft-model-path",
+            DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+            "--tp-size",
+            "2",
+            "--enable-dp-attention",
+            "--dp-size",
+            "2",
+        ]
+        if not is_in_amd_ci():
+            other_args += ["--mem-frac", "0.7"]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+
+        self.assertGreater(metrics["score"], 0.60)
+
+        server_info = requests.get(self.base_url + "/server_info")
+        avg_spec_accept_length = server_info.json()["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(
+            f"###test_gsm8k (deepseek-v3 mtp + dp):\n"
+            f"accuracy={metrics['score']=:.3f}\n"
+            f"{avg_spec_accept_length=:.3f}\n"
+        )
+        self.assertGreater(avg_spec_accept_length, 2.5)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/distributed/test_dp_attention_large.py

@@ -0,0 +1,188 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.ebnf_constrained_kit import EBNFConstrainedMixin
+from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
+from sglang.test.kits.regex_constrained_kit import RegexConstrainedMixin
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_IMAGE_URL,
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_amd_ci,
+    popen_launch_server,
+)
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "DeepSeek MLA forward_mla NameError on AMD (batched_gemm not defined)",
+)
+class TestDPAttentionDP2TP4(
+    CustomTestCase,
+    JSONConstrainedMixin,
+    EBNFConstrainedMixin,
+    RegexConstrainedMixin,
+):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp=4",
+                "--enable-dp-attention",
+                "--dp=2",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=None,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.8)
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "DeepSeek MTP forward_mla NameError on AMD + needs 8 GPUs",
+)
+class TestDPAttentionDP2TP2DeepseekV3MTP(
+    CustomTestCase,
+    JSONConstrainedMixin,
+    EBNFConstrainedMixin,
+    RegexConstrainedMixin,
+):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--disable-radix",
+            "--speculative-algorithm=EAGLE",
+            "--speculative-num-steps=2",
+            "--speculative-eagle-topk=4",
+            "--speculative-num-draft-tokens=4",
+            "--speculative-draft-model-path",
+            DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
+            "--tp-size=4",
+            "--enable-dp-attention",
+            "--dp-size=2",
+        ]
+        if not is_in_amd_ci():
+            other_args += ["--mem-frac", "0.7"]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(metrics)
+
+        self.assertGreater(metrics["score"], 0.60)
+
+        server_info = requests.get(self.base_url + "/server_info")
+        avg_spec_accept_length = server_info.json()["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(
+            f"###test_gsm8k (deepseek-v3 mtp + dp):\n"
+            f"accuracy={metrics['score']=:.3f}\n"
+            f"{avg_spec_accept_length=:.3f}\n"
+        )
+        self.assertGreater(avg_spec_accept_length, 2.5)
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "Qwen3-VL-30B-A3B-Instruct OOMs at TP=4 DP=2 on MI325 4-GPU runners",
+)
+class TestDPAttentionDP2TP4VLM(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "Qwen/Qwen3-VL-30B-A3B-Instruct"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.image_url = DEFAULT_IMAGE_URL
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--enable-dp-attention",
+                "--dp",
+                "2",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_vlm_generate(self):
+        chat_template = get_chat_template_by_model_path(self.model)
+        prompt = f"{chat_template.image_token}What is in this image?"
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": prompt,
+                "image_data": [self.image_url],
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 16,
+                },
+            },
+        )
+        response.raise_for_status()
+        response_json = response.json()
+        print(response_json)
+        self.assertIn("output_ids", response_json)
+        self.assertGreater(len(response_json["output_ids"]), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/_common.py

@@ -0,0 +1,317 @@
+"""Shared fixture for DeepSeek-V4 cookbook launch-command tests.
+
+Each sibling ``test_<hardware>_<model_size>.py`` declares ONE
+``hardware x model_size`` cell from the cookbook (e.g. B200 x Flash)
+and contains one ``CustomTestCase`` subclass per recipe
+(Low-Latency / Balanced / Max-Throughput / CP, where supported).
+
+Each subclass launches the server with the cookbook's exact flags and
+runs two sgl-eval evaluations (https://github.com/sgl-project/sgl-eval):
+- ``test_smoke_gsm8k`` — short, cheap GSM8K pass to verify the server
+  can produce coherent math answers at all (sanity gate).
+- ``test_aime25`` — full AIME25 accuracy run (heavy; 16 repeats default).
+
+Cookbook reference:
+    https://docs.sglang.io/cookbook/autoregressive/DeepSeek/DeepSeek-V4
+
+These are MANUAL tests (not CI). ``sgl-eval`` must be on PATH.
+
+Per-variant defaults (set on the Flash/Pro intermediate base classes):
+    Flash recipes -> AIME25 score threshold 0.93
+    Pro   recipes -> AIME25 score threshold 0.95
+GSM8K sanity threshold (0.93) is shared across Flash and Pro.
+
+AIME25 knobs (env vars):
+    DSV4_AIME25_NUM_REPEATS       (default 16    -> --n-repeats)
+    DSV4_AIME25_TEMPERATURE       (default 1.0   -> --temperature)
+    DSV4_AIME25_TOP_P             (default 1.0   -> --top-p)
+    DSV4_AIME25_MAX_TOKENS        (default 65536 -> --max-tokens)
+    DSV4_AIME25_NUM_THREADS       (default 512   -> --num-threads)
+    DSV4_AIME25_SCORE_METRIC      (default "score"; sgl-eval JSON key under "aggregate")
+    DSV4_AIME25_SCORE_THRESHOLD   (default 0; >0 overrides per-variant default)
+
+GSM8K sanity knobs (env vars):
+    DSV4_GSM8K_NUM_EXAMPLES       (default 50    -> --num-examples)
+    DSV4_GSM8K_N_REPEATS          (default 1     -> --n-repeats)
+    DSV4_GSM8K_TEMPERATURE        (default 0.6   -> --temperature)
+    DSV4_GSM8K_TOP_P              (default 0.95  -> --top-p)
+    DSV4_GSM8K_MAX_TOKENS         (default 8192  -> --max-tokens)
+    DSV4_GSM8K_NUM_THREADS        (default 64    -> --num-threads)
+    DSV4_GSM8K_SCORE_METRIC       (default "score"; sgl-eval JSON key under "aggregate")
+    DSV4_GSM8K_SCORE_THRESHOLD    (default 0.93; set to 0 to skip the assertion)
+
+Shared knobs:
+    DSV4_SGL_EVAL_OUT_DIR         (default /tmp/sgl-eval-out -> --out-dir)
+    DSV4_SGL_EVAL_BIN             (default "sgl-eval"; override path to the CLI)
+    DSV4_SERVER_LAUNCH_TIMEOUT    (default 3600s; the sglang 600s default is
+                                   too short for DSV4 model load + DeepGEMM
+                                   warmup. 1800s is also tight for the heavier
+                                   recipes (DP-attn + DeepEP); 3600s is the
+                                   safe default. Bump again for first-run
+                                   model downloads if needed.)
+
+Multi-node knobs (only consumed by multi-node test classes; if either
+is unset, those classes ``SkipTest``):
+    DSV4_NODE_RANK                (per-node rank for --node-rank)
+    DSV4_DIST_INIT_ADDR           (e.g. 10.0.0.1:20000 for --dist-init-addr)
+
+Always-on env (set by the base class for every recipe; per-recipe EXTRA_ENV
+wins on key conflict):
+    SGLANG_JIT_DEEPGEMM_FAST_WARMUP=1   skip the slow DeepGEMM warmup grid
+"""
+
+import json
+import os
+import shutil
+import subprocess
+import unittest
+from pathlib import Path
+from typing import ClassVar, Dict, List, Optional
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+SGL_EVAL_BIN = os.environ.get("DSV4_SGL_EVAL_BIN", "sgl-eval")
+SGL_EVAL_OUT_DIR = os.environ.get("DSV4_SGL_EVAL_OUT_DIR", "/tmp/sgl-eval-out")
+
+# DSV4 server launch needs more than the 600s sglang default: model load alone
+# can take 5+ min and DeepGEMM warmup another ~5 min. First-run model download
+# adds ~10-30 min on top. 1800s covers steady-state; bump via env for downloads.
+SERVER_LAUNCH_TIMEOUT = int(os.environ.get("DSV4_SERVER_LAUNCH_TIMEOUT", "3600"))
+
+# Defaults applied to every recipe's EXTRA_ENV. Per-recipe EXTRA_ENV wins on key
+# conflict.
+BASE_ENV: Dict[str, str] = {
+    # Skip the slow exhaustive DeepGEMM warmup grid; covers the shapes DSV4
+    # actually hits and shaves several minutes off server startup.
+    "SGLANG_JIT_DEEPGEMM_FAST_WARMUP": "1",
+}
+
+AIME25_NUM_REPEATS = int(os.environ.get("DSV4_AIME25_NUM_REPEATS", "16"))
+AIME25_TEMPERATURE = float(os.environ.get("DSV4_AIME25_TEMPERATURE", "1.0"))
+AIME25_TOP_P = float(os.environ.get("DSV4_AIME25_TOP_P", "1.0"))
+AIME25_MAX_TOKENS = int(os.environ.get("DSV4_AIME25_MAX_TOKENS", "65536"))
+AIME25_NUM_THREADS = int(os.environ.get("DSV4_AIME25_NUM_THREADS", "512"))
+AIME25_SCORE_METRIC = os.environ.get("DSV4_AIME25_SCORE_METRIC", "score")
+AIME25_SCORE_THRESHOLD = float(os.environ.get("DSV4_AIME25_SCORE_THRESHOLD", "0.0"))
+
+GSM8K_NUM_EXAMPLES = int(os.environ.get("DSV4_GSM8K_NUM_EXAMPLES", "50"))
+GSM8K_N_REPEATS = int(os.environ.get("DSV4_GSM8K_N_REPEATS", "1"))
+GSM8K_TEMPERATURE = float(os.environ.get("DSV4_GSM8K_TEMPERATURE", "0.6"))
+GSM8K_TOP_P = float(os.environ.get("DSV4_GSM8K_TOP_P", "0.95"))
+GSM8K_MAX_TOKENS = int(os.environ.get("DSV4_GSM8K_MAX_TOKENS", "8192"))
+GSM8K_NUM_THREADS = int(os.environ.get("DSV4_GSM8K_NUM_THREADS", "64"))
+GSM8K_SCORE_METRIC = os.environ.get("DSV4_GSM8K_SCORE_METRIC", "score")
+GSM8K_SCORE_THRESHOLD = float(os.environ.get("DSV4_GSM8K_SCORE_THRESHOLD", "0.93"))
+
+# DeepEP "large SMS" config — appears as `--deepep-config '{...}'` in every
+# DeepEP recipe except multi-node ones (where it is gated off in the JSX).
+DEEPEP_LARGE_SMS_CONFIG = (
+    '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+)
+
+
+def multinode_args(nnodes: int) -> List[str]:
+    """Return CLI args for a multi-node launch, or skip the test.
+
+    Reads DSV4_NODE_RANK and DSV4_DIST_INIT_ADDR from the env. Raises
+    ``unittest.SkipTest`` when either is missing — call from inside
+    ``setUpClass`` so the whole class skips cleanly.
+    """
+    rank = os.environ.get("DSV4_NODE_RANK")
+    addr = os.environ.get("DSV4_DIST_INIT_ADDR")
+    if rank is None or addr is None:
+        raise unittest.SkipTest(
+            "multi-node test requires DSV4_NODE_RANK and DSV4_DIST_INIT_ADDR"
+        )
+    return [
+        "--nnodes",
+        str(nnodes),
+        "--node-rank",
+        rank,
+        "--dist-init-addr",
+        addr,
+    ]
+
+
+class DSV4Aime25TestBase(CustomTestCase):
+    """Subclass via ``DSV4FlashAime25TestBase`` or ``DSV4ProAime25TestBase``,
+    not directly. Per-recipe subclasses set MODEL / OTHER_ARGS / EXTRA_ENV.
+
+    SCORE_THRESHOLD is set by the Flash/Pro intermediate base classes:
+    Flash 0.93, Pro 0.95.
+    """
+
+    MODEL: ClassVar[str] = ""
+    OTHER_ARGS: ClassVar[List[str]] = []
+    EXTRA_ENV: ClassVar[Dict[str, str]] = {}
+
+    SCORE_THRESHOLD: ClassVar[float] = 0.0
+
+    _BASE_CLASSES: ClassVar[set] = set()
+
+    @classmethod
+    def setUpClass(cls):
+        if cls in cls._BASE_CLASSES:
+            raise unittest.SkipTest("base class; subclass to run")
+        if not cls.MODEL or not cls.OTHER_ARGS:
+            raise unittest.SkipTest(f"{cls.__name__}: MODEL and OTHER_ARGS must be set")
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        env: Optional[Dict[str, str]] = {**BASE_ENV, **(cls.EXTRA_ENV or {})}
+        cls.process = popen_launch_server(
+            cls.MODEL,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=list(cls.OTHER_ARGS),
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def test_smoke_gsm8k(self):
+        """Quick GSM8K pass to verify the server is producing math answers."""
+        self._run_sgl_eval(
+            eval_name="gsm8k",
+            n_repeats=GSM8K_N_REPEATS,
+            temperature=GSM8K_TEMPERATURE,
+            top_p=GSM8K_TOP_P,
+            max_tokens=GSM8K_MAX_TOKENS,
+            num_threads=GSM8K_NUM_THREADS,
+            num_examples=GSM8K_NUM_EXAMPLES,
+            metric=GSM8K_SCORE_METRIC,
+            threshold=GSM8K_SCORE_THRESHOLD,
+        )
+
+    def test_aime25(self):
+        """Full AIME25 accuracy run; threshold gated by Flash vs Pro base."""
+        threshold = (
+            AIME25_SCORE_THRESHOLD
+            if AIME25_SCORE_THRESHOLD > 0
+            else self.SCORE_THRESHOLD
+        )
+        self._run_sgl_eval(
+            eval_name="aime25",
+            n_repeats=AIME25_NUM_REPEATS,
+            temperature=AIME25_TEMPERATURE,
+            top_p=AIME25_TOP_P,
+            max_tokens=AIME25_MAX_TOKENS,
+            num_threads=AIME25_NUM_THREADS,
+            num_examples=None,
+            metric=AIME25_SCORE_METRIC,
+            threshold=threshold,
+        )
+
+    def _run_sgl_eval(
+        self,
+        eval_name,
+        n_repeats,
+        temperature,
+        top_p,
+        max_tokens,
+        num_threads,
+        num_examples,
+        metric,
+        threshold,
+    ):
+        if shutil.which(SGL_EVAL_BIN) is None:
+            self.skipTest(f"{SGL_EVAL_BIN!r} not found on PATH")
+
+        out_dir = Path(SGL_EVAL_OUT_DIR)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        glob_pattern = f"sgl_eval_{eval_name}_*.json"
+        before = set(out_dir.glob(glob_pattern))
+
+        cmd = [
+            SGL_EVAL_BIN,
+            "run",
+            eval_name,
+            "--base-url",
+            f"{self.base_url}/v1",
+            "--n-repeats",
+            str(n_repeats),
+            "--temperature",
+            str(temperature),
+            "--top-p",
+            str(top_p),
+            "--max-tokens",
+            str(max_tokens),
+            "--num-threads",
+            str(num_threads),
+            "--out-dir",
+            str(out_dir),
+        ]
+        if num_examples is not None:
+            cmd += ["--num-examples", str(num_examples)]
+
+        print(f"[{type(self).__name__}] + {' '.join(cmd)}", flush=True)
+        subprocess.run(cmd, check=True)
+
+        new = sorted(set(out_dir.glob(glob_pattern)) - before)
+        if not new:
+            self.fail(f"sgl-eval produced no new {eval_name} JSON in {out_dir}")
+        result_path = new[-1]
+        with open(result_path) as f:
+            result = json.load(f)
+        print(
+            f"[{type(self).__name__}] sgl-eval {eval_name} result "
+            f"({result_path.name}): {json.dumps(result, indent=2)}",
+            flush=True,
+        )
+
+        score = self._extract_score(result, metric)
+        if threshold > 0:
+            self.assertGreaterEqual(
+                score,
+                threshold,
+                f"{eval_name} {metric}={score} below threshold {threshold}",
+            )
+
+    @staticmethod
+    def _extract_score(result, metric):
+        """Find ``metric`` (e.g. "pass@1") anywhere in the sgl-eval JSON tree."""
+
+        def walk(o):
+            if isinstance(o, dict):
+                if metric in o and isinstance(o[metric], (int, float)):
+                    return float(o[metric])
+                for v in o.values():
+                    s = walk(v)
+                    if s is not None:
+                        return s
+            elif isinstance(o, list):
+                for v in o:
+                    s = walk(v)
+                    if s is not None:
+                        return s
+            return None
+
+        score = walk(result)
+        if score is None:
+            raise AssertionError(f"metric {metric!r} not found in sgl-eval result JSON")
+        return score
+
+
+class DSV4FlashAime25TestBase(DSV4Aime25TestBase):
+    """Base for DeepSeek-V4-Flash recipes: AIME25 threshold 0.93."""
+
+    SCORE_THRESHOLD = 0.93
+
+
+class DSV4ProAime25TestBase(DSV4Aime25TestBase):
+    """Base for DeepSeek-V4-Pro recipes: AIME25 threshold 0.95."""
+
+    SCORE_THRESHOLD = 0.95
+
+
+DSV4Aime25TestBase._BASE_CLASSES = {
+    DSV4Aime25TestBase,
+    DSV4FlashAime25TestBase,
+    DSV4ProAime25TestBase,
+}

test/manual/dsv4/test_b200_flash.py

@@ -0,0 +1,108 @@
+"""B200 (FP4) x DeepSeek-V4-Flash.
+
+Covers the four cookbook recipes for this hardware x model_size cell:
+Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+
+
+class TestB200FlashLowLatency(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestB200FlashBalanced(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestB200FlashMaxThroughput(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestB200FlashCP(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--max-running-requests",
+        "1024",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_b200_pro.py

@@ -0,0 +1,124 @@
+"""B200 (FP4) x DeepSeek-V4-Pro.
+
+Covers the four cookbook recipes for this hardware x model_size cell:
+Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Pro"
+
+
+class TestB200ProLowLatency(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestB200ProBalanced(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--dp",
+        "8",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--mem-fraction-static",
+        "0.82",
+        "--cuda-graph-max-bs-decode",
+        "64",
+        "--max-running-requests",
+        "128",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestB200ProMaxThroughput(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--dp",
+        "8",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--mem-fraction-static",
+        "0.82",
+        "--cuda-graph-max-bs-decode",
+        "64",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestB200ProCP(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--cuda-graph-max-bs-decode",
+        "256",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_b300_flash.py

@@ -0,0 +1,110 @@
+"""B300 x DeepSeek-V4-Flash.
+
+The cookbook generator aliases B300 to B200, so the launch flags
+are identical to the B200(FP4) Flash cell. Kept as a separate file
+because the hardware target (and therefore the runtime environment)
+is different. Covers Low-Latency, Balanced, Max-Throughput, CP.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+
+
+class TestB300FlashLowLatency(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestB300FlashBalanced(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestB300FlashMaxThroughput(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestB300FlashCP(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--max-running-requests",
+        "1024",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_b300_pro.py

@@ -0,0 +1,126 @@
+"""B300 x DeepSeek-V4-Pro.
+
+The cookbook generator aliases B300 to B200, so the launch flags
+are identical to the B200(FP4) Pro cell. Kept as a separate file
+because the hardware target (and therefore the runtime environment)
+is different. Covers Low-Latency, Balanced, Max-Throughput, CP.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Pro"
+
+
+class TestB300ProLowLatency(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestB300ProBalanced(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--dp",
+        "8",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--mem-fraction-static",
+        "0.82",
+        "--cuda-graph-max-bs-decode",
+        "64",
+        "--max-running-requests",
+        "128",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestB300ProMaxThroughput(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--dp",
+        "8",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--mem-fraction-static",
+        "0.82",
+        "--cuda-graph-max-bs-decode",
+        "64",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestB300ProCP(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--cuda-graph-max-bs-decode",
+        "256",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_flash_mtp_dp4.py

@@ -0,0 +1,175 @@
+"""DSV4 Flash MTP test using EAGLE speculative algorithm.
+
+DSV4 Flash MTP shares the EAGLE wire path: EAGLE algo + NextN head built
+into the target model weights. No separate draft model is needed (sglang
+auto-falls back `--speculative-draft-model-path` to the target model).
+
+Test matrix mirrors test_eagle_infer_b.TestEAGLEServerBasic to maximize
+cuda-graph + buffer-pool coverage on the DSV4 path:
+  - test_gsm8k         (accuracy + spec path full forward)
+  - test_max_token_one (degenerate spec step, still cuda-graph captured)
+  - test_request_abort (cuda-graph buffer pool survives abort+restart)
+
+Server launch matches `run_flash_dp4.sh`: tp=4, dp=4, deepep MoE backend,
+DSV4 FP8 (FP4 experts disabled).
+"""
+
+import random
+import threading
+import time
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.few_shot_gsm8k import run_eval as run_gsm8k_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+DSV4_FLASH_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+    # MTP runs ~num_draft_tokens forward passes per step, so the deepep
+    # dispatch input size scales by that factor. Default 256 (used by the
+    # plain server) overflows once cuda-graph-max-bs-decode * num_draft_tokens
+    # > 256. 1024 covers bs=128 * 4 draft tokens with headroom.
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+}
+
+DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+
+PROMPTS = [
+    "[INST] You are a helpful assistant.\\nWhere are you from? [/INST]",
+    "[INST] You are a helpful assistant.\\nSummarize gradient descent in 2 sentences. [/INST]",
+    "[INST] You are a helpful assistant.\\nWhat is 17*23? [/INST]",
+    "[INST] You are a helpful assistant.\\nList three primary colors. [/INST]",
+]
+
+
+class DSV4FlashMTPServerBase(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DSV4_FLASH_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "4",
+            "--dp",
+            "4",
+            "--enable-dp-attention",
+            "--moe-a2a-backend",
+            "deepep",
+            "--cuda-graph-max-bs-decode",
+            "128",
+            "--max-running-requests",
+            "256",
+            "--deepep-config",
+            DEEPEP_CONFIG,
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            "--mem-fraction-static",
+            "0.7",
+        ]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+            env=DSV4_FLASH_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def send_request(self):
+        time.sleep(random.uniform(0, 2))
+        for prompt in PROMPTS:
+            resp = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": prompt,
+                    "sampling_params": {"temperature": 0, "max_new_tokens": 256},
+                },
+            )
+            assert resp.status_code == 200
+
+    def send_requests_abort(self):
+        for prompt in PROMPTS:
+            try:
+                time.sleep(random.uniform(0, 2))
+                requests.post(
+                    self.base_url + "/generate",
+                    json={
+                        "text": prompt,
+                        "sampling_params": {"temperature": 0, "max_new_tokens": 256},
+                    },
+                    timeout=0.5,
+                )
+            except requests.exceptions.Timeout:
+                pass
+
+
+class TestDSV4FlashMTPBasic(DSV4FlashMTPServerBase):
+    def test_gsm8k(self):
+        """Accuracy + spec path full forward."""
+        requests.get(self.base_url + "/flush_cache")
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_gsm8k_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["accuracy"], 0.95)
+
+    def test_max_token_one(self):
+        """Degenerate spec step (still cuda-graph captured)."""
+        requests.get(self.base_url + "/flush_cache")
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=100,
+            max_new_tokens=1,
+            parallel=128,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_gsm8k_eval(args)
+        self.assertGreater(metrics["output_throughput"], 50)
+
+    def test_request_abort(self):
+        """Cuda-graph buffer pool must survive abort+restart cycles."""
+        concurrency = 4
+        threads = [
+            threading.Thread(target=self.send_request) for _ in range(concurrency)
+        ] + [
+            threading.Thread(target=self.send_requests_abort)
+            for _ in range(concurrency)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+        self.assertIsNone(self.process.poll())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_flash_mtp_tp8.py

@@ -0,0 +1,122 @@
+"""DSV4-Flash 285B MTP performance tests on H200 TP=8.
+
+Manual test (8× H200, 285B FP8 weights). Not registered in CI.
+"""
+
+import os
+import tempfile
+import unittest
+
+import requests
+
+from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
+from sglang.bench_one_batch_server import run_benchmark as run_one_batch_benchmark
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+DSV4_FLASH_BASE_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+}
+
+DSV4_FLASH_SERVER_ARGS = [
+    "--trust-remote-code",
+    "--tp",
+    "8",
+    "--speculative-algorithm",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+    "--max-running-requests",
+    "8",
+]
+
+
+def _launch_dsv4_flash_server(extra_env=None):
+    env = dict(DSV4_FLASH_BASE_ENV)
+    if extra_env:
+        env.update(extra_env)
+    return popen_launch_server(
+        DSV4_FLASH_MODEL_PATH,
+        DEFAULT_URL_FOR_TEST,
+        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 4,
+        other_args=DSV4_FLASH_SERVER_ARGS,
+        env=env,
+    )
+
+
+class TestDSV4FlashMTPSimulatedAcc(CustomTestCase):
+    """bs=1 latency at isl=4096 / 900000 with `SGLANG_SIMULATE_ACC_LEN=3`.
+
+    Reference (H200 Flash TP8):
+      - isl=4096   → output 258.1 tok/s, accept 2.94
+      - isl=900000 → output 222.9 tok/s, accept 2.90
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch_dsv4_flash_server(
+            extra_env={"SGLANG_SIMULATE_ACC_LEN": "3"}
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def _run_one_batch(self, input_len):
+        requests.get(self.base_url + "/flush_cache")
+        server_args = ServerArgs(model_path=DSV4_FLASH_MODEL_PATH)
+        bench_args = OneBatchBenchArgs(
+            run_name=f"dsv4_flash_simacc_isl{input_len}",
+            batch_size=(1,),
+            input_len=(input_len,),
+            output_len=(1024,),
+            base_url=self.base_url,
+            skip_warmup=True,
+            result_filename=os.path.join(
+                tempfile.gettempdir(), f"dsv4_flash_simacc_isl{input_len}.jsonl"
+            ),
+            append_to_github_summary=False,
+        )
+        results, _ = run_one_batch_benchmark(server_args, bench_args)
+        self.assertTrue(results, "bench_one_batch_server returned no results")
+        return results[0]
+
+    def test_isl_4096(self):
+        r = self._run_one_batch(4096)
+        print(
+            f"[flash simacc isl=4096] output_throughput={r.output_throughput:.2f} tok/s "
+            f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
+            f"acc_length={r.acc_length:.2f}"
+        )
+        # Reference 258.1 tok/s / acc=2.94.
+        self.assertGreater(r.output_throughput, 232.0)
+        self.assertGreater(r.acc_length, 2.85)
+
+    def test_isl_900k(self):
+        r = self._run_one_batch(900_000)
+        print(
+            f"[flash simacc isl=900k] output_throughput={r.output_throughput:.2f} tok/s "
+            f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
+            f"acc_length={r.acc_length:.2f}"
+        )
+        # Reference 222.9 tok/s / acc=2.90.
+        self.assertGreater(r.output_throughput, 200.0)
+        self.assertGreater(r.acc_length, 2.85)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_flash_sanity_dp4.py

@@ -0,0 +1,160 @@
+"""DSV4-Flash 4-GPU server sanity matrix (TP4 variants)."""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+DSV4_FLASH_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+}
+
+DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+
+
+def _launch(other_args, env_extra=None, timeout_mult=1):
+    env = dict(DSV4_FLASH_ENV)
+    if env_extra:
+        env.update(env_extra)
+    return popen_launch_server(
+        DSV4_FLASH_MODEL_PATH,
+        DEFAULT_URL_FOR_TEST,
+        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * timeout_mult,
+        other_args=other_args,
+        env=env,
+    )
+
+
+_EAGLE_SPEC_ARGS = [
+    "--speculative-algorithm",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+]
+
+
+class TestDSV4FlashTP4DP4(
+    BasicDecodeCorrectnessMixin,
+    CustomTestCase,
+):
+    """TP4 + DP4 + deepep + EAGLE MTP."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch(
+            [
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--enable-dp-attention",
+                "--moe-a2a-backend",
+                "deepep",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "256",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+                "--mem-fraction-static",
+                "0.7",
+                *_EAGLE_SPEC_ARGS,
+            ]
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+class TestDSV4FlashTP4EP(
+    BasicDecodeCorrectnessMixin,
+    CustomTestCase,
+):
+    """TP attn + EP MoE (no DP attn) — exercises the DeepEP + TP-attn path."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch(
+            [
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--ep",
+                "4",
+                # No --enable-dp-attention by design: covers TP-attn path.
+                "--moe-a2a-backend",
+                "deepep",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "64",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+                "--mem-fraction-static",
+                "0.7",
+                *_EAGLE_SPEC_ARGS,
+            ]
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+class TestDSV4FlashTP4DP4ChunkedPrefillLarge(
+    BasicDecodeCorrectnessMixin,
+    CustomTestCase,
+):
+    """TP4 + DP4 with --chunked-prefill-size 16384 — large chunked prefill."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch(
+            [
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--enable-dp-attention",
+                "--moe-a2a-backend",
+                "deepep",
+                "--chunked-prefill-size",
+                "16384",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "256",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+                "--mem-fraction-static",
+                "0.7",
+                *_EAGLE_SPEC_ARGS,
+            ]
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_flash_sanity_tp8.py

@@ -0,0 +1,53 @@
+"""DSV4-Flash 8-GPU server sanity (TP8, no spec decoding)."""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+DSV4_FLASH_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+}
+
+
+class TestDSV4FlashTP8NoSpec(
+    BasicDecodeCorrectnessMixin,
+    CustomTestCase,
+):
+    """TP8, no spec decoding."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            DSV4_FLASH_MODEL_PATH,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "8",
+                "--max-running-requests",
+                "8",
+                "--mem-fraction-static",
+                "0.85",
+            ],
+            env=DSV4_FLASH_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_pd_disagg_nixl.py

@@ -0,0 +1,148 @@
+"""DSV4 Flash PD-disagg with NIXL backend. Both sides run dp-attention
++ deepep + EAGLE MTP so attn_tp_size and the V4 state pool layout are
+fully symmetric: same SWA item_len under matching attn_tp, and same
+DSA c4/c128 indexer ring buffer size under matching spec status. nixl
+`send_state` is page-by-index and has no V4 TP-slice / spec-asymmetric
+path, so any layout mismatch would trip the item_len assert in
+`nixl/conn.py`."""
+
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.few_shot_gsm8k import run_eval as run_gsm8k_eval
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+DSV4_FLASH_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+    # MTP num_draft_tokens=4 scales dispatch by ~4x; 256 overflows at bs=128.
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+}
+
+DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+
+
+class TestDSV4FlashPDDisaggNIXL(PDDisaggregationServerBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.transfer_backend = ["--disaggregation-transfer-backend", "nixl"]
+        cls.rdma_devices = []
+        cls.model = DSV4_FLASH_MODEL_PATH
+
+        cls.start_prefill()
+        cls.start_decode()
+
+        cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
+        cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
+        cls.launch_lb()
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--base-gpu-id",
+            "0",
+            "--tp",
+            "4",
+            "--dp",
+            "4",
+            "--enable-dp-attention",
+            "--moe-a2a-backend",
+            "deepep",
+            "--deepep-config",
+            DEEPEP_CONFIG,
+            "--cuda-graph-max-bs-decode",
+            "128",
+            "--max-running-requests",
+            "256",
+            "--mem-fraction-static",
+            "0.7",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            *cls.transfer_backend,
+            *cls.rdma_devices,
+        ]
+        cls.process_prefill = popen_launch_pd_server(
+            cls.model,
+            cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+            env=DSV4_FLASH_ENV,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--base-gpu-id",
+            "4",
+            "--tp",
+            "4",
+            "--dp",
+            "4",
+            "--enable-dp-attention",
+            "--moe-a2a-backend",
+            "deepep",
+            "--deepep-config",
+            DEEPEP_CONFIG,
+            "--cuda-graph-max-bs-decode",
+            "128",
+            "--max-running-requests",
+            "256",
+            "--mem-fraction-static",
+            "0.7",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            *cls.transfer_backend,
+            *cls.rdma_devices,
+        ]
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+            env=DSV4_FLASH_ENV,
+        )
+
+    def test_gsm8k(self):
+        """End-to-end PD-disagg accuracy through the LB."""
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=200,
+            max_new_tokens=512,
+            parallel=64,
+            host=f"http://{self.base_host}",
+            port=int(self.lb_port),
+        )
+        metrics = run_gsm8k_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["accuracy"], 0.95)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_pro_mtp.py

@@ -0,0 +1,278 @@
+"""DSV4-Pro 1.6T MTP performance tests on B200 TP=8.
+
+1. TestDSV4ProMTPSimulatedAcc — `SGLANG_SIMULATE_ACC_LEN=3` pins EAGLE accept
+   length so latency comparisons are apples-to-apples. Runs `bench_one_batch_server`
+   at bs=1 for isl=4096 and isl=900000 (osl=1024).
+
+2. TestDSV4ProMTPHongloumeng — real EAGLE accept (no SIMULATE) on Chinese
+   long-context input (`hongloumeng.txt`, ~627k DSV4 tokens). Builds a one-line
+   custom JSONL dataset on the fly and drives `bench_serving --dataset-name custom`
+   with one short slice (30k tokens) and the full long prompt.
+
+Manual test (8× B200, 1.6T weights). Not registered in CI.
+"""
+
+import json
+import os
+import tempfile
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.bench_one_batch_server import BenchArgs as OneBatchBenchArgs
+from sglang.bench_one_batch_server import run_benchmark as run_one_batch_benchmark
+from sglang.benchmark.serving import run_benchmark as run_serving_benchmark
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_PRO_MODEL_PATH = "deepseek-ai/DeepSeek-V4-Pro"
+
+HONGLOUMENG_PATH = os.environ.get(
+    "SGLANG_HONGLOUMENG_PATH",
+    os.path.join(os.path.dirname(__file__), "hongloumeng.txt"),
+)
+
+DSV4_PRO_BASE_ENV = {
+    "SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2": "1",
+}
+
+DSV4_PRO_SERVER_ARGS = [
+    "--trust-remote-code",
+    "--tp",
+    "8",
+    "--moe-runner-backend",
+    "flashinfer_mxfp4",
+    "--speculative-algorithm",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+    "--chunked-prefill-size",
+    "4096",
+    "--disable-flashinfer-autotune",
+    "--mem-fraction-static",
+    "0.82",
+    "--max-running-requests",
+    "8",
+]
+
+
+def _launch_dsv4_pro_server(extra_env=None):
+    env = dict(DSV4_PRO_BASE_ENV)
+    if extra_env:
+        env.update(extra_env)
+    return popen_launch_server(
+        DSV4_PRO_MODEL_PATH,
+        DEFAULT_URL_FOR_TEST,
+        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 4,
+        other_args=DSV4_PRO_SERVER_ARGS,
+        env=env,
+    )
+
+
+class TestDSV4ProMTPSimulatedAcc(CustomTestCase):
+    """bs=1 latency at isl=4096 / 900000 with `SGLANG_SIMULATE_ACC_LEN=3`.
+
+    Reference (B200 Pro TP8):
+      - isl=4096   → output 194.6 tok/s, accept 2.96
+      - isl=900000 → output 174.6 tok/s, accept 2.93
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch_dsv4_pro_server(
+            extra_env={"SGLANG_SIMULATE_ACC_LEN": "3"}
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def _run_one_batch(self, input_len):
+        requests.get(self.base_url + "/flush_cache")
+        server_args = ServerArgs(model_path=DSV4_PRO_MODEL_PATH)
+        bench_args = OneBatchBenchArgs(
+            run_name=f"dsv4_pro_simacc_isl{input_len}",
+            batch_size=(1,),
+            input_len=(input_len,),
+            output_len=(1024,),
+            base_url=self.base_url,
+            skip_warmup=True,
+            result_filename=os.path.join(
+                tempfile.gettempdir(), f"dsv4_pro_simacc_isl{input_len}.jsonl"
+            ),
+            append_to_github_summary=False,
+        )
+        results, _ = run_one_batch_benchmark(server_args, bench_args)
+        self.assertTrue(results, "bench_one_batch_server returned no results")
+        return results[0]
+
+    def test_isl_4096(self):
+        r = self._run_one_batch(4096)
+        print(
+            f"[pro simacc isl=4096] output_throughput={r.output_throughput:.2f} tok/s "
+            f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
+            f"acc_length={r.acc_length:.2f}"
+        )
+        # Reference 194.6 tok/s / acc=2.96 — give 10% throughput margin and a
+        # generous accept-length floor to absorb run-to-run jitter.
+        self.assertGreater(r.output_throughput, 175.0)
+        self.assertGreater(r.acc_length, 2.85)
+
+    def test_isl_900k(self):
+        r = self._run_one_batch(900_000)
+        print(
+            f"[pro simacc isl=900k] output_throughput={r.output_throughput:.2f} tok/s "
+            f"latency={r.latency:.2f}s last_ttft={r.last_ttft:.2f}s "
+            f"acc_length={r.acc_length:.2f}"
+        )
+        # Reference 174.6 tok/s / acc=2.93.
+        self.assertGreater(r.output_throughput, 155.0)
+        self.assertGreater(r.acc_length, 2.85)
+
+
+def _build_hongloumeng_jsonl(num_tokens, tokenizer, out_path):
+    """Slice the first `num_tokens` DSV4 tokens of hongloumeng.txt into a
+    one-line CustomDataset JSONL. Pass num_tokens=None to keep the full text.
+    """
+    with open(HONGLOUMENG_PATH, "r", encoding="utf-8") as f:
+        text = f.read()
+    if num_tokens is not None:
+        ids = tokenizer.encode(text)
+        text = tokenizer.decode(ids[:num_tokens])
+    with open(out_path, "w", encoding="utf-8") as f:
+        f.write(
+            json.dumps(
+                {"conversations": [{"value": text}, {"value": "x"}]},
+                ensure_ascii=False,
+            )
+            + "\n"
+        )
+    return out_path
+
+
+class TestDSV4ProMTPHongloumeng(CustomTestCase):
+    """Real EAGLE accept on Chinese long-context (hongloumeng.txt).
+
+    Reference (B200 Pro TP8, no SIMULATE):
+      - isl=30000  → output 124.4 tok/s, decode peak 184 tok/s, accept 2.47
+      - isl=627059 → output 125.7 tok/s, decode peak 179 tok/s, accept 2.52
+    """
+
+    SHORT_TOKENS = 30_000
+    LONG_TOKENS = None  # full file (~627k DSV4 tokens)
+    OUTPUT_TOKENS = 4096
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = _launch_dsv4_pro_server()
+
+        # Resolve tokenizer once; the server reports its own tokenizer path so
+        # on-the-fly token-level slicing matches what the server will see.
+        info = requests.get(cls.base_url + "/server_info", timeout=60).json()
+        tokenizer_path = info.get("tokenizer_path") or DSV4_PRO_MODEL_PATH
+        from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+
+        cls.tokenizer = get_tokenizer(tokenizer_path)
+
+        cls.tmpdir = tempfile.mkdtemp(prefix="dsv4_hongloumeng_")
+        cls.short_jsonl = _build_hongloumeng_jsonl(
+            cls.SHORT_TOKENS,
+            cls.tokenizer,
+            os.path.join(cls.tmpdir, "hongloumeng_30k.jsonl"),
+        )
+        cls.long_jsonl = _build_hongloumeng_jsonl(
+            cls.LONG_TOKENS,
+            cls.tokenizer,
+            os.path.join(cls.tmpdir, "hongloumeng_full.jsonl"),
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def _run_custom_bench(self, dataset_path):
+        requests.get(self.base_url + "/flush_cache")
+        args = SimpleNamespace(
+            backend="sglang",
+            base_url=self.base_url,
+            host=None,
+            port=None,
+            dataset_name="custom",
+            dataset_path=dataset_path,
+            model=None,
+            tokenizer=None,
+            num_prompts=1,
+            sharegpt_output_len=self.OUTPUT_TOKENS,
+            sharegpt_context_len=None,
+            random_input_len=4096,
+            random_output_len=2048,
+            random_range_ratio=0.0,
+            request_rate=float("inf"),
+            max_concurrency=1,
+            warmup_requests=0,
+            flush_cache=True,
+            multi=None,
+            output_file=None,
+            disable_tqdm=False,
+            disable_stream=False,
+            return_logprob=False,
+            return_routed_experts=False,
+            seed=0,
+            disable_ignore_eos=False,
+            extra_request_body=None,
+            apply_chat_template=False,
+            profile=None,
+            lora_name=None,
+            lora_request_distribution="uniform",
+            lora_zipf_alpha=1.5,
+            prompt_suffix="",
+            device="cuda",
+            pd_separated=False,
+            ready_check_timeout_sec=0,
+        )
+        return run_serving_benchmark(args)
+
+    def test_short_30k(self):
+        res = self._run_custom_bench(self.short_jsonl)
+        print(
+            f"[hongloumeng 30k] output_throughput={res['output_throughput']:.2f} tok/s "
+            f"accept_length={res['accept_length']:.2f} "
+            f"mean_ttft_ms={res['mean_ttft_ms']:.0f} "
+            f"mean_tpot_ms={res['mean_tpot_ms']:.2f}"
+        )
+        # Reference 124 tok/s / accept 2.47.
+        self.assertGreater(res["output_throughput"], 105.0)
+        self.assertGreater(res["accept_length"], 2.30)
+
+    def test_long_full(self):
+        res = self._run_custom_bench(self.long_jsonl)
+        print(
+            f"[hongloumeng full] output_throughput={res['output_throughput']:.2f} tok/s "
+            f"accept_length={res['accept_length']:.2f} "
+            f"mean_ttft_ms={res['mean_ttft_ms']:.0f} "
+            f"mean_tpot_ms={res['mean_tpot_ms']:.2f}"
+        )
+        # Reference 125 tok/s / accept 2.52. Cold prefill takes ~85s on 627k
+        # tokens so the run is dominated by prefill, but decode steady-state
+        # accept_length is the metric we care about.
+        self.assertGreater(res["output_throughput"], 105.0)
+        self.assertGreater(res["accept_length"], 2.30)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_dsv4_swa_radix_retract.py

@@ -0,0 +1,165 @@
+"""DSV4 stress test for SWA radix cache + tombstone + retract interaction.
+
+Reproduces the assert in `swa_radix_cache.cache_unfinished_req`:
+    assert old_prefix_len <= len(new_indices)
+
+Trip conditions (all required):
+  1. Fork-only SWA leaf early-release on (`SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW=1`)
+  2. Multiple requests share a long prefix (so one req's tombstoned leaf
+     poisons match_prefix for others walking the same radix path).
+  3. Memory pressure forces retract while at least one req has tombstoned
+     its leaf (decode_batch_idx >= sliding_window_size at retract time).
+
+After main #19427 changed `old_prefix_len = req.cache_protected_len`
+(stable), tombstone-induced shrinks in match's `best_value_len` across
+chunked-prefill rounds can make stale `cache_protected_len` exceed
+current matchable length -> assert trips.
+
+Test passes iff the scheduler does not crash under this stress workload.
+"""
+
+import random
+import threading
+import time
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+DSV4_FLASH_MODEL_PATH = "sgl-project/DeepSeek-V4-Flash-FP8"
+
+# Long shared prefix forces multi-chunk prefill and ensures cross-request
+# prefix-cache hits so one req's tombstone affects later reqs.
+SHARED_PREFIX_BLOCK = (
+    "You are a careful, expert assistant. Answer concisely.\n"
+    "Context: " + ("the quick brown fox jumps over the lazy dog. " * 600)
+)
+
+QUESTION_TAILS = [
+    " Q: What is 17*23?\n",
+    " Q: List three primary colors.\n",
+    " Q: Where is Mount Everest?\n",
+    " Q: Summarize gradient descent in two sentences.\n",
+    " Q: Name two bodies of water in Africa.\n",
+    " Q: What language is spoken in Brazil?\n",
+]
+
+
+class TestDSV4FlashSWARadixRetract(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DSV4_FLASH_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            "4",
+            "--dp",
+            "4",
+            "--enable-dp-attention",
+            "--moe-a2a-backend",
+            "deepep",
+            "--cuda-graph-max-bs-decode",
+            "128",
+            "--max-running-requests",
+            "256",
+            "--deepep-config",
+            '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}',
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            # Tight static memory so SWA pool fills up under load and
+            # retract is forced.
+            "--mem-fraction-static",
+            "0.7",
+        ]
+        env = {
+            "SGLANG_DSV4_FP4_EXPERTS": "0",
+            "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+            "SGLANG_OPT_SWA_RADIX_CACHE_COMPACT": "0",
+            "SGLANG_TEST_RETRACT": "1",
+            "SGLANG_TEST_RETRACT_INTERVAL": "3",
+        }
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def _send_req(self, prompt: str, max_new_tokens: int):
+        try:
+            resp = requests.post(
+                self.base_url + "/generate",
+                json={
+                    "text": prompt,
+                    "sampling_params": {
+                        # Vary outputs slightly so reqs don't share decode
+                        # paths perfectly; we want some to finish, some to
+                        # be retracted under pressure.
+                        "temperature": 0.7,
+                        "max_new_tokens": max_new_tokens,
+                    },
+                },
+                timeout=600,
+            )
+            # Per-request success is not the gate; some requests are
+            # expected to be retracted/aborted under heavy pressure.
+            return resp.status_code == 200
+        except Exception:
+            return False
+
+    def test_swa_tombstone_retract_does_not_crash(self):
+        """Stress: 64 concurrent long-prompt reqs with long generation force
+        retract under SWA pool pressure. Reqs share a 30k+ token prefix so
+        tombstoned leaves from retracted reqs are on the radix path of new
+        reqs. Scheduler must not crash on the swa_radix_cache assert."""
+
+        random.seed(0)
+        concurrency = 64
+        # Long enough generation to push past sliding_window_size -> fires
+        # `dec_swa_lock_only` -> tombstones leaves. Combined with SWA pool
+        # pressure this guarantees retract while tombstones are live.
+        max_new_tokens = 1024
+
+        threads = []
+        for i in range(concurrency):
+            tail = QUESTION_TAILS[i % len(QUESTION_TAILS)]
+            # Add a small per-req suffix so reqs don't dedup at radix root
+            # but still share the bulk of the prefix.
+            prompt = SHARED_PREFIX_BLOCK + tail + f"(seed={i})"
+            t = threading.Thread(target=self._send_req, args=(prompt, max_new_tokens))
+            threads.append(t)
+            t.start()
+            # Stagger so requests enter prefill in waves; some are still in
+            # decode (and have tombstoned leaves) when later waves of
+            # chunked-prefill reqs walk the same radix path.
+            time.sleep(0.05)
+
+        for t in threads:
+            t.join(timeout=600)
+
+        # The only invariant: scheduler survived. Per-request completion is
+        # best-effort under retract pressure.
+        self.assertIsNone(self.process.poll())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_fused_compress_attn_hip.py

@@ -0,0 +1,465 @@
+"""Unit tests for the fused compressor attention Triton kernel on HIP.
+
+Validates numerical parity between the fused single-kernel path (plan-driven
+Triton) and the reference per-seq Python implementation.
+
+Usage:
+    python -m pytest test/manual/dsv4/test_fused_compress_attn_hip.py -v
+    # or directly:
+    python test/manual/dsv4/test_fused_compress_attn_hip.py
+"""
+
+import unittest
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass
+class FusedCompressPlan:
+    compress_plan_gpu: torch.Tensor
+    write_plan_gpu: torch.Tensor
+    num_compress: int
+    num_write: int
+
+
+def write_current_token_to_state(
+    kv_score_input: torch.Tensor,
+    write_plan: torch.Tensor,
+    state_pool_buffer: torch.Tensor,
+    head_dim: int,
+    overlap: bool,
+    ratio: int,
+) -> None:
+    """Reference write path used by this manual test.
+
+    Plan row layout: [ragged_id, batch_id, position, window_len, state_base].
+    """
+    del head_dim  # layout is already encoded in kv_score_input/state_pool_buffer shape.
+    state_size = (2 if overlap else 1) * ratio
+    plan_cpu = write_plan.cpu()
+    for row in plan_cpu:
+        ragged_id = int(row[0].item())
+        position = int(row[2].item())
+        state_base = int(row[4].item())
+        if ragged_id < 0 or position < 0:
+            continue
+        dst = state_base + (position % state_size)
+        if (
+            0 <= dst < state_pool_buffer.shape[0]
+            and 0 <= ragged_id < kv_score_input.shape[0]
+        ):
+            state_pool_buffer[dst] = kv_score_input[ragged_id]
+
+
+def fused_compress_attn(
+    state_pool_buffer: torch.Tensor,
+    plan: torch.Tensor,
+    ape: torch.Tensor,
+    rms_weight: torch.Tensor,
+    rms_eps: float,
+    freqs_cis_real: torch.Tensor,
+    head_dim: int,
+    rope_head_dim: int,
+    overlap: bool,
+    ratio: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    """Reference compress path for manual parity tests.
+
+    This keeps the test runnable after removing `fused_compress_kernel.py`.
+    """
+    freqs_cis = torch.view_as_complex(
+        freqs_cis_real.view(freqs_cis_real.shape[0], -1, 2).contiguous()
+    )
+    result = _ref_compress(
+        kv_score_input=torch.empty(
+            0, device=state_pool_buffer.device, dtype=torch.float32
+        ),
+        state_pool=state_pool_buffer,
+        plan=plan,
+        ape=ape,
+        rms_weight=rms_weight,
+        rms_eps=rms_eps,
+        freqs_cis=freqs_cis,
+        head_dim=head_dim,
+        rope_head_dim=rope_head_dim,
+        overlap=overlap,
+        ratio=ratio,
+        num_compress=plan.shape[0],
+    )
+    out.copy_(result)
+    return out
+
+
+def _make_plan_from_params(
+    extend_lens: list[int],
+    seq_lens: list[int],
+    ratio: int,
+    overlap: bool,
+    state_bases: list[int],
+    device: torch.device,
+) -> FusedCompressPlan:
+    """Build a test plan without requiring real SWA / req_to_token tables."""
+    bs = len(extend_lens)
+    ext = np.array(extend_lens, dtype=np.int32)
+    seq = np.array(seq_lens, dtype=np.int32)
+    total = int(ext.sum())
+
+    state_size = (2 if overlap else 1) * ratio
+    K = state_size
+
+    batch_ids = np.repeat(np.arange(bs, dtype=np.int32), ext)
+    ragged_ids = np.arange(total, dtype=np.int32)
+    cu_extend = np.empty(bs + 1, dtype=np.int32)
+    cu_extend[0] = 0
+    np.cumsum(ext, out=cu_extend[1:])
+    j_in_seq = ragged_ids - cu_extend[batch_ids]
+    prefix_lens = seq - ext
+    positions = prefix_lens[batch_ids] + j_in_seq
+
+    window_lens = np.maximum(0, K - np.minimum(j_in_seq + 1, K)).astype(np.int32)
+    state_base_arr = np.array(state_bases, dtype=np.int32)
+    state_base_per_token = state_base_arr[batch_ids]
+
+    plan_rows = np.stack(
+        [ragged_ids, batch_ids, positions, window_lens, state_base_per_token],
+        axis=1,
+    ).astype(np.int32)
+
+    compress_mask = (positions + 1) % ratio == 0
+    compress_plan = plan_rows[compress_mask]
+
+    write_starts = np.maximum(0, seq - K).astype(np.int32)
+    write_mask = positions >= write_starts[batch_ids]
+    write_plan = plan_rows[write_mask]
+
+    n_compress = int(compress_plan.shape[0]) if compress_plan.size > 0 else 0
+    n_write = int(write_plan.shape[0]) if write_plan.size > 0 else 0
+
+    compress_gpu = (
+        torch.from_numpy(np.ascontiguousarray(compress_plan)).to(device)
+        if n_compress > 0
+        else torch.empty((0, 5), dtype=torch.int32, device=device)
+    )
+    write_gpu = (
+        torch.from_numpy(np.ascontiguousarray(write_plan)).to(device)
+        if n_write > 0
+        else torch.empty((0, 5), dtype=torch.int32, device=device)
+    )
+
+    return FusedCompressPlan(
+        compress_plan_gpu=compress_gpu,
+        write_plan_gpu=write_gpu,
+        num_compress=n_compress,
+        num_write=n_write,
+    )
+
+
+def _make_freqs_cis(max_seq: int, rope_dim: int, device: torch.device) -> torch.Tensor:
+    """Create test freqs_cis as complex64 [max_seq, rope_dim/2], matching production."""
+    half = rope_dim // 2
+    angles = torch.randn(max_seq, half, device=device, dtype=torch.float32) * 0.1
+    return torch.polar(torch.ones_like(angles), angles)
+
+
+def _freqs_to_real(freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Convert complex64 freqs to float32 [max_seq, rope_dim] interleaved."""
+    return torch.view_as_real(freqs_cis).flatten(-2).contiguous()
+
+
+def _ref_compress(
+    kv_score_input: torch.Tensor,
+    state_pool: torch.Tensor,
+    plan: torch.Tensor,
+    ape: torch.Tensor,
+    rms_weight: torch.Tensor,
+    rms_eps: float,
+    freqs_cis: torch.Tensor,
+    head_dim: int,
+    rope_head_dim: int,
+    overlap: bool,
+    ratio: int,
+    num_compress: int,
+) -> torch.Tensor:
+    """Pure-PyTorch reference matching SGLang compress_decode_paged semantics.
+
+    State already has current tokens written (no APE).
+    APE is added to ALL K scores at compress time.
+    """
+    if num_compress == 0:
+        return torch.empty(0, head_dim, dtype=torch.float32, device=state_pool.device)
+
+    coff = 2 if overlap else 1
+    half_dim = coff * head_dim
+    state_size = coff * ratio
+    K = state_size
+
+    plan_cpu = plan[:num_compress].cpu()
+    out = torch.empty(
+        num_compress, head_dim, dtype=torch.float32, device=state_pool.device
+    )
+
+    for pid in range(num_compress):
+        position = int(plan_cpu[pid, 2].item())
+        state_base = int(plan_cpu[pid, 4].item())
+
+        if position < 0:
+            continue
+
+        kv_rows = []
+        score_rows = []
+        for k in range(K):
+            s = position - K + 1 + k
+            col_off = (head_dim if k >= ratio else 0) if overlap else 0
+            ape_row = k % ratio
+            d_slice = slice(col_off, col_off + head_dim)
+
+            if s < 0:
+                kv_rows.append(
+                    torch.zeros(head_dim, dtype=torch.float32, device=state_pool.device)
+                )
+                score_rows.append(
+                    torch.full(
+                        (head_dim,),
+                        float("-inf"),
+                        dtype=torch.float32,
+                        device=state_pool.device,
+                    )
+                )
+            else:
+                ring = s % state_size
+                row = state_pool[state_base + ring]
+                kv_rows.append(row[d_slice].float())
+                # APE added to ALL scores
+                score_rows.append(
+                    row[half_dim + col_off : half_dim + col_off + head_dim].float()
+                    + ape[ape_row, d_slice].float()
+                )
+
+        kv_stack = torch.stack(kv_rows, dim=0)
+        sc_stack = torch.stack(score_rows, dim=0)
+        weights = torch.softmax(sc_stack, dim=0)
+        compressed = (weights * kv_stack).sum(dim=0)
+
+        var = (compressed * compressed).mean()
+        normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
+
+        comp_pos = (position // ratio) * ratio
+        rope_seg = normed[-rope_head_dim:].clone()
+        freqs_row = torch.view_as_real(freqs_cis[comp_pos]).flatten()
+        cos_v = freqs_row[0::2].float()
+        sin_v = freqs_row[1::2].float()
+
+        even = rope_seg[0::2]
+        odd = rope_seg[1::2]
+        normed[-rope_head_dim:] = torch.stack(
+            [even * cos_v - odd * sin_v, odd * cos_v + even * sin_v], dim=-1
+        ).flatten()
+
+        out[pid] = normed
+
+    return out
+
+
+class TestFusedCompressAttn(unittest.TestCase):
+
+    def _run_test(
+        self,
+        ratio: int,
+        overlap: bool,
+        bs: int,
+        extend_lens: list[int],
+        prefix_lens: list[int],
+        head_dim: int = 512,
+        rope_head_dim: int = 64,
+    ):
+        device = torch.device("cuda")
+        torch.manual_seed(42)
+        coff = 2 if overlap else 1
+        half_dim = coff * head_dim
+        last_dim = 2 * half_dim
+        state_size = coff * ratio
+
+        seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
+        total_tokens = sum(extend_lens)
+        max_seq = max(seq_lens) + 128
+
+        kv_score_input = torch.randn(
+            total_tokens, last_dim, device=device, dtype=torch.float32
+        )
+
+        pool_size = bs * state_size + 2
+        state_pool = torch.randn(
+            pool_size, last_dim, device=device, dtype=torch.float32
+        )
+        state_pool[:, half_dim:] *= 0.5  # reasonable score magnitudes
+
+        state_bases = [i * state_size for i in range(bs)]
+        ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1
+        rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32)
+        rms_eps = 1e-6
+        freqs_cis = _make_freqs_cis(max_seq, rope_head_dim, device)
+        freqs_real = _freqs_to_real(freqs_cis)
+
+        plan = _make_plan_from_params(
+            extend_lens, seq_lens, ratio, overlap, state_bases, device
+        )
+        if plan.num_compress == 0:
+            return
+
+        # Step 1: write current tokens to state (same for both paths)
+        state_triton = state_pool.clone()
+        state_ref = state_pool.clone()
+
+        write_current_token_to_state(
+            kv_score_input=kv_score_input,
+            write_plan=plan.write_plan_gpu,
+            state_pool_buffer=state_triton,
+            head_dim=head_dim,
+            overlap=overlap,
+            ratio=ratio,
+        )
+        # Reference: same write
+        write_current_token_to_state(
+            kv_score_input=kv_score_input,
+            write_plan=plan.write_plan_gpu,
+            state_pool_buffer=state_ref,
+            head_dim=head_dim,
+            overlap=overlap,
+            ratio=ratio,
+        )
+
+        # Step 2a: Triton fused compress
+        out_triton = torch.empty(
+            plan.num_compress, head_dim, device=device, dtype=torch.float32
+        )
+        fused_compress_attn(
+            state_pool_buffer=state_triton,
+            plan=plan.compress_plan_gpu,
+            ape=ape,
+            rms_weight=rms_weight,
+            rms_eps=rms_eps,
+            freqs_cis_real=freqs_real,
+            head_dim=head_dim,
+            rope_head_dim=rope_head_dim,
+            overlap=overlap,
+            ratio=ratio,
+            out=out_triton,
+        )
+
+        # Step 2b: reference compress
+        out_ref = _ref_compress(
+            kv_score_input=kv_score_input,
+            state_pool=state_ref,
+            plan=plan.compress_plan_gpu,
+            ape=ape,
+            rms_weight=rms_weight,
+            rms_eps=rms_eps,
+            freqs_cis=freqs_cis,
+            head_dim=head_dim,
+            rope_head_dim=rope_head_dim,
+            overlap=overlap,
+            ratio=ratio,
+            num_compress=plan.num_compress,
+        )
+
+        torch.testing.assert_close(out_triton, out_ref, atol=1e-3, rtol=1e-3)
+
+    def test_hca_single(self):
+        self._run_test(
+            ratio=128, overlap=False, bs=1, extend_lens=[128], prefix_lens=[0]
+        )
+
+    def test_hca_multi(self):
+        self._run_test(
+            ratio=128, overlap=False, bs=2, extend_lens=[128, 256], prefix_lens=[0, 128]
+        )
+
+    def test_csa_single(self):
+        self._run_test(ratio=4, overlap=True, bs=1, extend_lens=[16], prefix_lens=[8])
+
+    def test_csa_multi(self):
+        self._run_test(
+            ratio=4, overlap=True, bs=3, extend_lens=[8, 12, 16], prefix_lens=[4, 8, 0]
+        )
+
+    def test_csa_small_dim(self):
+        self._run_test(
+            ratio=4,
+            overlap=True,
+            bs=2,
+            extend_lens=[8, 8],
+            prefix_lens=[4, 0],
+            head_dim=256,
+        )
+
+
+class TestStateOrdering(unittest.TestCase):
+
+    def test_write_then_compress(self):
+        """Verify write-first, compress-second matches reference."""
+        device = torch.device("cuda")
+        torch.manual_seed(123)
+        ratio, overlap = 4, True
+        coff = 2
+        head_dim, rope_head_dim = 128, 64
+        half_dim = coff * head_dim
+        last_dim = 2 * half_dim
+        state_size = coff * ratio
+
+        pool_size = state_size + 2
+        state_pool = torch.randn(
+            pool_size, last_dim, device=device, dtype=torch.float32
+        )
+        state_pool[:, half_dim:] *= 0.5
+
+        kv_score_input = torch.randn(8, last_dim, device=device, dtype=torch.float32)
+        ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1
+        rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32)
+        freqs_cis = _make_freqs_cis(64, rope_head_dim, device)
+
+        plan = _make_plan_from_params([8], [8], ratio, overlap, [0], device)
+        if plan.num_compress == 0:
+            return
+
+        state_before = state_pool.clone()
+
+        # Write first
+        write_current_token_to_state(
+            kv_score_input=kv_score_input,
+            write_plan=plan.write_plan_gpu,
+            state_pool_buffer=state_pool,
+            head_dim=head_dim,
+            overlap=overlap,
+            ratio=ratio,
+        )
+
+        # State should now be different (tokens written)
+        self.assertFalse(torch.allclose(state_pool, state_before))
+
+        # Compress
+        out = torch.empty(
+            plan.num_compress, head_dim, device=device, dtype=torch.float32
+        )
+        fused_compress_attn(
+            state_pool_buffer=state_pool,
+            plan=plan.compress_plan_gpu,
+            ape=ape,
+            rms_weight=rms_weight,
+            rms_eps=1e-6,
+            freqs_cis_real=_freqs_to_real(freqs_cis),
+            head_dim=head_dim,
+            rope_head_dim=rope_head_dim,
+            overlap=overlap,
+            ratio=ratio,
+            out=out,
+        )
+
+        self.assertFalse(torch.any(torch.isnan(out)).item())
+        self.assertFalse(torch.any(torch.isinf(out)).item())
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_gb300_flash.py

@@ -0,0 +1,108 @@
+"""GB300 x DeepSeek-V4-Flash.
+
+Single-node TP=4 path on the deepseek-ai MXFP4 repo. Covers
+Low-Latency, Balanced, Max-Throughput, Context-Parallel (CP).
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+
+
+class TestGB300FlashLowLatency(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestGB300FlashBalanced(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestGB300FlashMaxThroughput(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024"}
+
+
+class TestGB300FlashCP(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--max-running-requests",
+        "1024",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_gb300_pro.py

@@ -0,0 +1,126 @@
+"""GB300 x DeepSeek-V4-Pro.
+
+Single-node TP=4 path. Note that GB300 Pro CP bumps
+mem-fraction-static to 0.88 (1.6T weights at TP=4 on 273 GB don't
+fit at the default 0.78). Covers Low-Latency, Balanced,
+Max-Throughput, Context-Parallel (CP).
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4ProAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Pro"
+
+
+class TestGB300ProLowLatency(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "flashinfer_mxfp4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--chunked-prefill-size",
+        "4096",
+        "--disable-flashinfer-autotune",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestGB300ProBalanced(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--mem-fraction-static",
+        "0.9",
+        "--cuda-graph-max-bs-decode",
+        "128",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestGB300ProMaxThroughput(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--mem-fraction-static",
+        "0.9",
+        "--cuda-graph-max-bs-decode",
+        "128",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256"}
+
+
+class TestGB300ProCP(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.88",
+        "--cuda-graph-max-bs-decode",
+        "256",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_h200_fp4_flash.py

@@ -0,0 +1,71 @@
+"""H200 (FP4 / Marlin) x DeepSeek-V4-Flash.
+
+The cookbook disables Context-Parallel for the H200 FP4 (Marlin)
+hardware, so this file only covers Low-Latency, Balanced, and
+Max-Throughput.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DSV4FlashAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+
+
+class TestH200Fp4FlashLowLatency(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "marlin",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestH200Fp4FlashBalanced(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "marlin",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestH200Fp4FlashMaxThroughput(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-runner-backend",
+        "marlin",
+    ]
+    EXTRA_ENV = {}
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_h200_fp4_pro.py

@@ -0,0 +1,77 @@
+"""H200 (FP4 / Marlin) x DeepSeek-V4-Pro.
+
+The cookbook disables Context-Parallel for the H200 FP4 (Marlin)
+hardware, so this file only covers Low-Latency, Balanced, and
+Max-Throughput.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DSV4ProAime25TestBase
+
+MODEL = "deepseek-ai/DeepSeek-V4-Pro"
+
+
+class TestH200Fp4ProLowLatency(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-runner-backend",
+        "marlin",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestH200Fp4ProBalanced(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-runner-backend",
+        "marlin",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+class TestH200Fp4ProMaxThroughput(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "8",
+        "--moe-runner-backend",
+        "marlin",
+        "--mem-fraction-static",
+        "0.88",
+    ]
+    EXTRA_ENV = {}
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_h200_fp8_flash.py

@@ -0,0 +1,120 @@
+"""H200 (FP8) x DeepSeek-V4-Flash.
+
+Uses the FP8-repackaged repo (sgl-project/DeepSeek-V4-Flash-FP8) and
+the SGLANG_DSV4_FP4_EXPERTS=0 env that the cookbook generator emits
+for H200 FP8 cells. Covers Low-Latency, Balanced, Max-Throughput, CP.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DEEPEP_LARGE_SMS_CONFIG, DSV4FlashAime25TestBase
+
+MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
+H200_FP8_ENV = {"SGLANG_DSV4_FP4_EXPERTS": "0"}
+
+
+class TestH200Fp8FlashLowLatency(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "3",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "4",
+    ]
+    EXTRA_ENV = dict(H200_FP8_ENV)
+
+
+class TestH200Fp8FlashBalanced(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--speculative-algorithm",
+        "EAGLE",
+        "--speculative-num-steps",
+        "1",
+        "--speculative-eagle-topk",
+        "1",
+        "--speculative-num-draft-tokens",
+        "2",
+        "--cuda-graph-max-bs-decode",
+        "128",
+        "--max-running-requests",
+        "128",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        **H200_FP8_ENV,
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
+    }
+
+
+class TestH200Fp8FlashMaxThroughput(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--dp",
+        "4",
+        "--enable-dp-attention",
+        "--moe-a2a-backend",
+        "deepep",
+        "--cuda-graph-max-bs-decode",
+        "128",
+        "--max-running-requests",
+        "256",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        **H200_FP8_ENV,
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
+    }
+
+
+class TestH200Fp8FlashCP(DSV4FlashAime25TestBase):
+    MODEL = MODEL
+    OTHER_ARGS = [
+        "--trust-remote-code",
+        "--tp",
+        "4",
+        "--moe-a2a-backend",
+        "deepep",
+        "--enable-dsa-prefill-context-parallel",
+        "--dsa-prefill-cp-mode",
+        "round-robin-split",
+        "--chunked-prefill-size",
+        "16384",
+        "--mem-fraction-static",
+        "0.78",
+        "--max-running-requests",
+        "1024",
+        "--deepep-config",
+        DEEPEP_LARGE_SMS_CONFIG,
+    ]
+    EXTRA_ENV = {
+        **H200_FP8_ENV,
+        "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+    }
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/dsv4/test_h200_fp8_pro.py

@@ -0,0 +1,126 @@
+"""H200 (FP8) x DeepSeek-V4-Pro.
+
+The cookbook ships this cell as a multi-node (2 nodes, TP=16) launch
+using the FP8-repackaged repo (sgl-project/DeepSeek-V4-Pro-FP8).
+Each test class skips itself unless DSV4_NODE_RANK and
+DSV4_DIST_INIT_ADDR are exported. Runtime expectation:
+
+    On every node:
+        DSV4_NODE_RANK=<0 or 1> \\
+        DSV4_DIST_INIT_ADDR=<head-node-ip>:20000 \\
+        python test/manual/models/dsv4/test_h200_fp8_pro.py
+
+Context-Parallel is marked TBD in the cookbook for this cell, so it
+is intentionally omitted.
+"""
+
+import os
+import sys
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _common import DSV4ProAime25TestBase, multinode_args
+
+MODEL = "sgl-project/DeepSeek-V4-Pro-FP8"
+H200_FP8_PRO_ENV = {
+    "SGLANG_DSV4_FP4_EXPERTS": "0",
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "128",
+}
+
+
+class TestH200Fp8ProLowLatency(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    EXTRA_ENV = dict(H200_FP8_PRO_ENV)
+
+    @classmethod
+    def setUpClass(cls):
+        cls.OTHER_ARGS = [
+            "--trust-remote-code",
+            "--tp",
+            "16",
+            "--dp",
+            "16",
+            "--enable-dp-attention",
+            *multinode_args(2),
+            "--moe-a2a-backend",
+            "deepep",
+            "--cuda-graph-max-bs-decode",
+            "8",
+            "--max-running-requests",
+            "32",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            "--mem-fraction-static",
+            "0.88",
+        ]
+        super().setUpClass()
+
+
+class TestH200Fp8ProBalanced(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    EXTRA_ENV = dict(H200_FP8_PRO_ENV)
+
+    @classmethod
+    def setUpClass(cls):
+        cls.OTHER_ARGS = [
+            "--trust-remote-code",
+            "--tp",
+            "16",
+            "--dp",
+            "16",
+            "--enable-dp-attention",
+            *multinode_args(2),
+            "--moe-a2a-backend",
+            "deepep",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "1",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "2",
+            "--mem-fraction-static",
+            "0.88",
+            "--cuda-graph-max-bs-decode",
+            "8",
+            "--max-running-requests",
+            "32",
+        ]
+        super().setUpClass()
+
+
+class TestH200Fp8ProMaxThroughput(DSV4ProAime25TestBase):
+    MODEL = MODEL
+    EXTRA_ENV = dict(H200_FP8_PRO_ENV)
+
+    @classmethod
+    def setUpClass(cls):
+        cls.OTHER_ARGS = [
+            "--trust-remote-code",
+            "--tp",
+            "16",
+            "--dp",
+            "16",
+            "--enable-dp-attention",
+            *multinode_args(2),
+            "--moe-a2a-backend",
+            "deepep",
+            "--mem-fraction-static",
+            "0.88",
+            "--cuda-graph-max-bs-decode",
+            "128",
+            "--max-running-requests",
+            "256",
+        ]
+        super().setUpClass()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/entrypoints/http_server/test_abort_request.py

@@ -0,0 +1,205 @@
+"""
+Integration test for abort_request functionality with a SGLang server.
+
+Run with:
+    python -m unittest sglang.test.srt.entrypoints.http_server.test_abort_request -v
+"""
+
+import threading
+import time
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestAbortRequest(CustomTestCase):
+    """Integration test class for abort request functionality."""
+
+    model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+    base_url = DEFAULT_URL_FOR_TEST
+
+    @classmethod
+    def setUpClass(cls):
+        """Launch the server."""
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--disable-cuda-graph"],
+        )
+
+        cls.completion_url = f"{cls.base_url}/generate"
+        cls.abort_url = f"{cls.base_url}/abort_request"
+        cls.health_url = f"{cls.base_url}/health"
+
+        print(f"Server started at {cls.base_url}")
+
+    @classmethod
+    def tearDownClass(cls):
+        """Clean up the server."""
+        kill_process_tree(cls.process.pid)
+
+    def _send_completion_request(
+        self,
+        text: str,
+        request_id: str,
+        max_tokens: int = 50,
+        temperature: float = 0.8,
+        stream: bool = True,
+    ) -> requests.Response:
+        """Send a completion request to the server."""
+        payload = {
+            "text": text,
+            "sampling_params": {
+                "max_new_tokens": max_tokens,
+                "temperature": temperature,
+            },
+            "stream": stream,
+            "rid": request_id,
+        }
+
+        response = requests.post(
+            self.completion_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=30,
+            stream=stream,
+        )
+
+        return response
+
+    def _send_abort_request(self, request_id: str) -> requests.Response:
+        """Send an abort request."""
+        payload = {"rid": request_id}
+        return requests.post(self.abort_url, json=payload, timeout=10)
+
+    def _check_server_health(self) -> bool:
+        """Check if server is healthy."""
+        try:
+            response = requests.get(self.health_url, timeout=5)
+            return response.status_code == 200
+        except:
+            return False
+
+    def test_abort_during_non_streaming_generation(self):
+        """Test aborting a non-streaming request during generation."""
+        self.assertTrue(self._check_server_health(), "Server should be healthy")
+
+        request_id = "test_abort_non_streaming"
+        completion_result = {}
+
+        def run_completion():
+            response = self._send_completion_request(
+                "Write a detailed essay about artificial intelligence",
+                max_tokens=500,
+                temperature=1,
+                request_id=request_id,
+                stream=False,
+            )
+
+            if response.status_code == 200:
+                result = response.json()
+                completion_result["text"] = result.get("text", "")
+                completion_result["finish_reason"] = result.get("meta_info", {}).get(
+                    "finish_reason"
+                )
+
+        completion_thread = threading.Thread(target=run_completion)
+        completion_thread.start()
+        time.sleep(0.1)
+
+        abort_response = self._send_abort_request(request_id)
+        completion_thread.join()
+
+        self.assertEqual(abort_response.status_code, 200)
+        self.assertIsNotNone(completion_result, "Should have completion result")
+        if completion_result:
+            finish_reason_obj = completion_result.get("finish_reason")
+            self.assertIsNotNone(finish_reason_obj, "Should have finish_reason")
+            if finish_reason_obj:
+                self.assertEqual(
+                    finish_reason_obj.get("type"), "abort", "Should be aborted"
+                )
+
+    def test_batch_requests_with_selective_abort(self):
+        """Test multiple concurrent requests with selective abort of one request."""
+        self.assertTrue(self._check_server_health(), "Server should be healthy")
+
+        request_ids = ["batch_test_0", "batch_test_1", "batch_test_2"]
+        abort_target_id = "batch_test_1"
+        completion_results = {}
+        threads = []
+
+        def run_completion(req_id, prompt):
+            response = self._send_completion_request(
+                f"Write a story about {prompt}",
+                max_tokens=100,
+                temperature=0.8,
+                request_id=req_id,
+                stream=False,
+            )
+
+            if response.status_code == 200:
+                result = response.json()
+                completion_results[req_id] = {
+                    "text": result.get("text", ""),
+                    "finish_reason": result.get("meta_info", {}).get("finish_reason"),
+                }
+
+        # Start all requests
+        prompts = ["a knight's adventure", "a space discovery", "a chef's restaurant"]
+        for i, req_id in enumerate(request_ids):
+            thread = threading.Thread(target=run_completion, args=(req_id, prompts[i]))
+            threads.append(thread)
+            thread.start()
+
+        # Abort one request
+        time.sleep(0.1)
+        abort_response = self._send_abort_request(abort_target_id)
+
+        # Wait for completion
+        for thread in threads:
+            thread.join(timeout=30)
+
+        # Verify results
+        self.assertEqual(abort_response.status_code, 200)
+
+        # Check aborted request
+        aborted_result = completion_results.get(abort_target_id)
+        self.assertIsNotNone(
+            aborted_result, f"Aborted request {abort_target_id} should have result"
+        )
+        if aborted_result:
+            aborted_finish_reason = aborted_result.get("finish_reason")
+            self.assertIsNotNone(
+                aborted_finish_reason, "Aborted request should have finish_reason"
+            )
+            if aborted_finish_reason:
+                self.assertEqual(aborted_finish_reason.get("type"), "abort")
+
+        # Check other requests completed normally
+        normal_completions = 0
+        for req_id in request_ids:
+            if req_id != abort_target_id and req_id in completion_results:
+                result = completion_results[req_id]
+                if result:
+                    finish_reason = result.get("finish_reason")
+                    if finish_reason and finish_reason.get("type") == "length":
+                        normal_completions += 1
+
+        self.assertEqual(
+            normal_completions, 2, "Other 2 requests should complete normally"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2, warnings="ignore")

test/manual/ep/test_deepep_internode.py

@@ -0,0 +1,445 @@
+# Copy from deepseek-ai/DeepEP/tests/test_internode.py
+
+import os
+import time
+
+# noinspection PyUnresolvedReferences
+import deep_ep
+
+# Test compatibility with low latency functions
+import test_deepep_low_latency
+import torch
+import torch.distributed as dist
+
+from sglang.test.test_deepep_utils import (
+    bench,
+    calc_diff,
+    create_grouped_scores,
+    init_dist,
+    inplace_unique,
+    per_token_cast_back,
+    per_token_cast_to_fp8,
+)
+
+
+def test_main(
+    num_sms: int,
+    local_rank: int,
+    num_local_ranks: int,
+    num_ranks: int,
+    num_nodes: int,
+    rank: int,
+    buffer: deep_ep.Buffer,
+    group: dist.ProcessGroup,
+):
+    # Settings
+    num_tokens, hidden, num_topk_groups, num_topk, num_experts = (
+        4096,
+        7168,
+        min(num_nodes, 4),
+        8,
+        (256 // num_ranks) * num_ranks,
+    )
+    assert num_experts % num_ranks == 0 and num_local_ranks == 8
+    if local_rank == 0:
+        print(
+            f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}",
+            flush=True,
+        )
+
+    # Random data
+    x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
+    x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+    x_e4m3 = per_token_cast_to_fp8(x)
+    scores = (
+        torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+        + 1
+    )
+    group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1)
+    group_idx = torch.topk(
+        group_scores, k=num_topk_groups, dim=-1, sorted=False
+    ).indices
+    masked_scores = create_grouped_scores(scores, group_idx, num_nodes)
+    topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[
+        1
+    ]
+    topk_weights = (
+        torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
+    )
+    topk_weights_pure_rand = torch.randn(
+        (num_tokens, num_topk), dtype=torch.float32, device="cuda"
+    )
+    rank_idx = topk_idx // (num_experts // num_ranks)
+    rank_idx.masked_fill_(topk_idx == -1, -1)
+    inplace_unique(rank_idx, num_ranks)
+    rdma_rank_idx = rank_idx // num_local_ranks
+    rdma_rank_idx.masked_fill_(rank_idx == -1, -1)
+    inplace_unique(rdma_rank_idx, num_nodes)
+
+    # RDMA dispatch counts
+    rdma_idx = topk_idx // (num_experts // num_nodes)
+    rdma_idx.masked_fill_(topk_idx == -1, -1)
+    inplace_unique(rdma_idx, num_nodes)
+    num_rdma_token_sent = rdma_idx.ne(-1).sum().item()
+
+    # Expert meta
+    num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
+    for i in range(num_experts):
+        num_tokens_per_expert[i] = (topk_idx == i).sum()
+    gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
+    dist.all_reduce(gbl_num_tokens_per_expert, group=group)
+
+    # Rank layout meta
+    num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
+    num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda")
+    token_idx_in_rank = torch.full(
+        (num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
+    )
+    for i in range(num_ranks):
+        num_tokens_per_rank[i] = (rank_idx == i).sum()
+        token_sel = (rank_idx == i).max(dim=-1)[0]
+        count = token_sel.sum().item()
+        tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
+        tokens[:count] = torch.sort(tokens[:count])[0]
+        token_idx_in_rank[i][tokens[:count]] = torch.arange(
+            count, dtype=torch.long, device="cuda"
+        )
+    for i in range(num_nodes):
+        num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum()
+    token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
+    is_token_in_rank = token_idx_in_rank >= 0
+    gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
+    dist.all_reduce(gbl_num_tokens_per_rank, group=group)
+
+    (
+        ref_num_tokens_per_rank,
+        ref_num_tokens_per_rdma_rank,
+        ref_num_tokens_per_expert,
+        ref_is_token_in_rank,
+        _,
+    ) = buffer.get_dispatch_layout(topk_idx, num_experts)
+    assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
+    assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank)
+    assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
+    assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
+    t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
+    if local_rank == 0:
+        print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
+        print("", flush=True)
+    group.barrier()
+    time.sleep(1)
+
+    # Config
+    rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512)
+    config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size)
+
+    # Test dispatch
+    # noinspection PyShadowingNames
+    def check_data(check_x, recv_gbl_rank_prefix_sum):
+        assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
+        check_start = 0
+        for i in range(num_ranks):
+            check_end = recv_gbl_rank_prefix_sum[i].item()
+            assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
+            check_start = check_end
+
+    for previous_mode in (False, True):
+        for async_mode in (False, True):
+            for current_x in (x_pure_rand, x, x_e4m3):
+                for with_topk in (False, True):
+                    if local_rank == 0:
+                        print(
+                            f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
+                            flush=True,
+                            end="",
+                        )
+                    dispatch_args = {
+                        "x": current_x,
+                        "num_tokens_per_rank": num_tokens_per_rank,
+                        "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
+                        "is_token_in_rank": is_token_in_rank,
+                        "num_tokens_per_expert": num_tokens_per_expert,
+                        "config": config,
+                        "async_finish": async_mode,
+                    }
+                    if with_topk:
+                        dispatch_args.update(
+                            {
+                                "topk_idx": topk_idx,
+                                "topk_weights": (
+                                    topk_weights_pure_rand
+                                    if current_x is x_pure_rand
+                                    else topk_weights
+                                ),
+                            }
+                        )
+                    if previous_mode:
+                        dispatch_args.update({"previous_event": buffer.capture()})
+                    (
+                        recv_x,
+                        recv_topk_idx,
+                        recv_topk_weights,
+                        recv_num_tokens_per_expert_list,
+                        handle,
+                        event,
+                    ) = buffer.dispatch(**dispatch_args)
+                    event.current_stream_wait() if async_mode else ()
+                    recv_x = (
+                        per_token_cast_back(*recv_x)
+                        if isinstance(recv_x, tuple)
+                        else recv_x
+                    )
+
+                    # Checks
+                    recv_gbl_rank_prefix_sum = handle[-4]
+                    assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
+                        0
+                    ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+                    assert (
+                        gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
+                        == recv_num_tokens_per_expert_list
+                    )
+                    if current_x is not x_pure_rand:
+                        check_data(recv_x, recv_gbl_rank_prefix_sum)
+                    if with_topk:
+                        # Check `topk_idx`
+                        assert (
+                            recv_topk_idx.eq(-1)
+                            | (
+                                (recv_topk_idx >= 0)
+                                & (recv_topk_idx < (num_experts // num_ranks))
+                            )
+                        ).sum().item() == recv_topk_idx.numel()
+                        for i, count in enumerate(recv_num_tokens_per_expert_list):
+                            assert recv_topk_idx.eq(i).sum().item() == count
+
+                        # Check `topk_weights`
+                        if current_x is not x_pure_rand:
+                            recv_topk_weights[recv_topk_idx.eq(-1)] = (
+                                recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
+                                    recv_topk_weights
+                                )[recv_topk_idx.eq(-1)]
+                            )
+                            check_data(recv_topk_weights, recv_gbl_rank_prefix_sum)
+
+                    # Test cached dispatch (must without top-k staffs)
+                    if not with_topk:
+                        dispatch_args = {
+                            "x": current_x,
+                            "handle": handle,
+                            "config": config,
+                            "async_finish": async_mode,
+                        }
+                        if previous_mode:
+                            dispatch_args.update({"previous_event": buffer.capture()})
+                        recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
+                        event.current_stream_wait() if async_mode else ()
+                        recv_x = (
+                            per_token_cast_back(*recv_x)
+                            if isinstance(recv_x, tuple)
+                            else recv_x
+                        )
+                        if current_x is not x_pure_rand:
+                            check_data(recv_x, recv_gbl_rank_prefix_sum)
+
+                    # Test combine
+                    combine_args = {
+                        "x": recv_x,
+                        "handle": handle,
+                        "config": config,
+                        "async_finish": async_mode,
+                    }
+                    if with_topk:
+                        combine_args.update({"topk_weights": recv_topk_weights})
+                    if previous_mode:
+                        dispatch_args.update({"previous_event": buffer.capture()})
+                    combined_x, combined_topk_weights, event = buffer.combine(
+                        **combine_args
+                    )
+                    event.current_stream_wait() if async_mode else ()
+                    check_x = combined_x.float() / is_token_in_rank.sum(
+                        dim=1
+                    ).unsqueeze(1)
+                    ref_x = x_pure_rand if current_x is x_pure_rand else x
+                    assert calc_diff(check_x, ref_x) < 5e-6
+                    if with_topk:
+                        check_topk_weights = (
+                            combined_topk_weights
+                            if (current_x is x_pure_rand)
+                            else (
+                                combined_topk_weights
+                                / is_token_in_rank.sum(dim=1).unsqueeze(1)
+                            )
+                        )
+                        ref_topk_weights = (
+                            topk_weights_pure_rand
+                            if current_x is x_pure_rand
+                            else topk_weights
+                        )
+                        assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
+
+                    # For later tuning
+                    dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2
+                    dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
+                    combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
+                    combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes
+
+                    if local_rank == 0:
+                        print(" passed", flush=True)
+    if local_rank == 0:
+        print("", flush=True)
+
+    # Tune dispatch performance
+    best_dispatch_results = None
+    fp8_factor = (1 + 4 / 128) / 2
+    for current_x in (x_e4m3, x):
+        best_time, best_results = 1e10, None
+        rdma_send_bytes = (
+            (dispatch_bf16_rdma_send_bytes * fp8_factor)
+            if isinstance(current_x, tuple)
+            else dispatch_bf16_rdma_send_bytes
+        )
+        nvl_recv_bytes = (
+            (dispatch_bf16_nvl_recv_bytes * fp8_factor)
+            if isinstance(current_x, tuple)
+            else dispatch_bf16_nvl_recv_bytes
+        )
+        for nvl_chunk_size in range(4, 33, 4):
+            for rdma_chunk_size in range(4, 33, 4):
+                config = deep_ep.Config(
+                    num_sms,
+                    nvl_chunk_size,
+                    nvl_buffer_size,
+                    rdma_chunk_size,
+                    rdma_buffer_size,
+                )
+                tune_args = {"x": current_x, "handle": handle, "config": config}
+                t = bench(lambda: buffer.dispatch(**tune_args))[0]
+                if t < best_time:
+                    best_time, best_results = t, (
+                        num_sms,
+                        nvl_chunk_size,
+                        rdma_chunk_size,
+                    )
+                if local_rank == 0:
+                    print(
+                        f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
+                        flush=True,
+                    )
+        if local_rank == 0:
+            print(
+                f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
+                flush=True,
+            )
+            print("", flush=True)
+
+        if isinstance(current_x, tuple):
+            # Gather FP8 the best config from rank 0
+            best_dispatch_results = torch.tensor(
+                [best_results[0], best_results[1], best_results[2]],
+                dtype=torch.int32,
+                device="cuda",
+            )
+            all_best_fp8_results_list = [
+                torch.zeros_like(best_dispatch_results)
+                for _ in range(torch.distributed.get_world_size())
+            ]
+            dist.all_gather(
+                all_best_fp8_results_list, best_dispatch_results, group=group
+            )
+            best_dispatch_results = all_best_fp8_results_list[0].tolist()
+    dispatch_config = deep_ep.Config(
+        best_dispatch_results[0],
+        best_dispatch_results[1],
+        nvl_buffer_size,
+        best_dispatch_results[2],
+        rdma_buffer_size,
+    )
+
+    dispatch_args = {
+        "x": x,
+        "num_tokens_per_rank": num_tokens_per_rank,
+        "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
+        "is_token_in_rank": is_token_in_rank,
+        "num_tokens_per_expert": num_tokens_per_expert,
+        "config": dispatch_config if dispatch_config is not None else config,
+    }
+    recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
+
+    # Tune combine performance
+    best_time, best_results = 1e10, None
+    for nvl_chunk_size in range(1, 5, 1):
+        for rdma_chunk_size in range(8, 33, 4):
+            config = deep_ep.Config(
+                num_sms,
+                nvl_chunk_size,
+                nvl_buffer_size,
+                rdma_chunk_size,
+                rdma_buffer_size,
+            )
+            tune_args = {"x": recv_x, "handle": handle, "config": config}
+            t = bench(lambda: buffer.combine(**tune_args))[0]
+            if local_rank == 0:
+                print(
+                    f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
+                    flush=True,
+                )
+                if t < best_time:
+                    best_time, best_results = t, (
+                        num_sms,
+                        nvl_chunk_size,
+                        rdma_chunk_size,
+                    )
+
+    if local_rank == 0:
+        print(
+            f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
+            flush=True,
+        )
+        print("", flush=True)
+
+
+# noinspection PyUnboundLocalVariable
+def test_loop(local_rank: int, num_local_ranks: int):
+    num_nodes = int(os.getenv("WORLD_SIZE", 1))
+    rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
+    test_ll_compatibility = False
+    if test_ll_compatibility:
+        ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
+
+    buffer = deep_ep.Buffer(
+        group,
+        int(1e9),
+        int(1e9),
+        low_latency_mode=test_ll_compatibility,
+        num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
+    )
+    assert num_local_ranks == 8 and num_ranks > 8
+    torch.manual_seed(rank)
+
+    for i in (24,):
+        test_main(
+            i, local_rank, num_local_ranks, num_ranks, num_nodes, rank, buffer, group
+        )
+        if local_rank == 0:
+            print("", flush=True)
+
+    # Test compatibility with low latency functions
+    if test_ll_compatibility:
+        buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
+        test_deepep_low_latency.test_main(
+            ll_num_tokens,
+            ll_hidden,
+            ll_num_experts,
+            ll_num_topk,
+            rank,
+            num_ranks,
+            group,
+            buffer,
+            seed=1,
+        )
+
+
+if __name__ == "__main__":
+    num_processes = 8
+    torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

test/manual/ep/test_deepep_intranode.py

@@ -0,0 +1,378 @@
+# Copy from deepseek-ai/DeepEP/tests/test_intranode.py
+
+import time
+
+# noinspection PyUnresolvedReferences
+import deep_ep
+
+# Test compatibility with low latency functions
+import test_deepep_low_latency
+import torch
+import torch.distributed as dist
+
+from sglang.test.test_deepep_utils import (
+    bench,
+    calc_diff,
+    init_dist,
+    inplace_unique,
+    per_token_cast_back,
+    per_token_cast_to_fp8,
+)
+
+
+def test_main(
+    num_sms: int,
+    local_rank: int,
+    num_ranks: int,
+    rank: int,
+    buffer: deep_ep.Buffer,
+    group: dist.ProcessGroup,
+):
+    # Settings
+    num_tokens, hidden, num_topk, num_experts = (
+        4096,
+        7168,
+        8,
+        (256 // num_ranks) * num_ranks,
+    )
+    assert num_experts % num_ranks == 0
+    if local_rank == 0:
+        print(
+            f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk={num_topk}",
+            flush=True,
+        )
+
+    # Random data
+    x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
+    x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+    x_e4m3 = per_token_cast_to_fp8(x)
+    scores = (
+        torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+        + 1
+    )
+    topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1]
+    topk_weights = (
+        torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
+    )
+    topk_weights_pure_rand = torch.randn(
+        (num_tokens, num_topk), dtype=torch.float32, device="cuda"
+    )
+    rank_idx = topk_idx // (num_experts // num_ranks)
+    rank_idx.masked_fill_(topk_idx == -1, -1)
+    inplace_unique(rank_idx, num_ranks)
+
+    # Expert meta
+    num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
+    for i in range(num_experts):
+        num_tokens_per_expert[i] = (topk_idx == i).sum()
+    gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
+    dist.all_reduce(gbl_num_tokens_per_expert, group=group)
+
+    # Rank layout meta
+    num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
+    token_idx_in_rank = torch.full(
+        (num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
+    )
+    for i in range(num_ranks):
+        num_tokens_per_rank[i] = (rank_idx == i).sum()
+        token_sel = (rank_idx == i).max(dim=-1)[0]
+        count = token_sel.sum().item()
+        tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
+        tokens[:count] = torch.sort(tokens[:count])[0]
+        token_idx_in_rank[i][tokens[:count]] = torch.arange(
+            count, dtype=torch.long, device="cuda"
+        )
+    token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
+    is_token_in_rank = token_idx_in_rank >= 0
+    gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
+    dist.all_reduce(gbl_num_tokens_per_rank, group=group)
+
+    ref_num_tokens_per_rank, _, ref_num_tokens_per_expert, ref_is_token_in_rank, _ = (
+        buffer.get_dispatch_layout(topk_idx, num_experts)
+    )
+    assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
+    assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
+    assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
+    t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
+    if local_rank == 0:
+        print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
+        print("", flush=True)
+    group.barrier()
+    time.sleep(1)
+
+    # Config
+    nvl_buffer_size = 256
+    config = deep_ep.Config(num_sms, 8, nvl_buffer_size)
+
+    # Test dispatch
+    # noinspection PyShadowingNames
+    def check_data(check_x, rank_prefix_matrix):
+        assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
+        check_start = 0
+        for i in range(num_ranks):
+            check_end = rank_prefix_matrix[i][rank].item()
+            assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
+            check_start = check_end
+
+    for previous_mode in (False, True):
+        for async_mode in (False, True):
+            for current_x in (x_pure_rand, x, x_e4m3):
+                for with_topk in (False, True):
+                    if local_rank == 0:
+                        print(
+                            f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
+                            flush=True,
+                            end="",
+                        )
+                    dispatch_args = {
+                        "x": current_x,
+                        "num_tokens_per_rank": num_tokens_per_rank,
+                        "is_token_in_rank": is_token_in_rank,
+                        "num_tokens_per_expert": num_tokens_per_expert,
+                        "config": config,
+                        "async_finish": async_mode,
+                    }
+                    if with_topk:
+                        dispatch_args.update(
+                            {
+                                "topk_idx": topk_idx,
+                                "topk_weights": (
+                                    topk_weights_pure_rand
+                                    if current_x is x_pure_rand
+                                    else topk_weights
+                                ),
+                            }
+                        )
+                    if previous_mode:
+                        dispatch_args.update({"previous_event": buffer.capture()})
+                    (
+                        recv_x,
+                        recv_topk_idx,
+                        recv_topk_weights,
+                        recv_num_tokens_per_expert_list,
+                        handle,
+                        event,
+                    ) = buffer.dispatch(**dispatch_args)
+                    event.current_stream_wait() if async_mode else ()
+                    recv_x = (
+                        per_token_cast_back(*recv_x)
+                        if isinstance(recv_x, tuple)
+                        else recv_x
+                    )
+
+                    # Checks
+                    rank_prefix_matrix = handle[0]
+                    assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
+                        0
+                    ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+                    assert (
+                        gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
+                        == recv_num_tokens_per_expert_list
+                    )
+                    if current_x is not x_pure_rand:
+                        check_data(recv_x, rank_prefix_matrix)
+                    if with_topk:
+                        # Check `topk_idx`
+                        assert (
+                            recv_topk_idx.eq(-1)
+                            | (
+                                (recv_topk_idx >= 0)
+                                & (recv_topk_idx < (num_experts // num_ranks))
+                            )
+                        ).sum().item() == recv_topk_idx.numel()
+                        for i, count in enumerate(recv_num_tokens_per_expert_list):
+                            assert recv_topk_idx.eq(i).sum().item() == count
+
+                        # Check `topk_weights`
+                        if current_x is not x_pure_rand:
+                            recv_topk_weights[recv_topk_idx.eq(-1)] = (
+                                recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
+                                    recv_topk_weights
+                                )[recv_topk_idx.eq(-1)]
+                            )
+                            check_data(recv_topk_weights, rank_prefix_matrix)
+
+                    # Test cached dispatch (must without top-k staffs)
+                    if not with_topk:
+                        dispatch_args = {
+                            "x": current_x,
+                            "handle": handle,
+                            "config": config,
+                            "async_finish": async_mode,
+                        }
+                        if previous_mode:
+                            dispatch_args.update({"previous_event": buffer.capture()})
+                        recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
+                        event.current_stream_wait() if async_mode else ()
+                        recv_x = (
+                            per_token_cast_back(*recv_x)
+                            if isinstance(recv_x, tuple)
+                            else recv_x
+                        )
+                        if current_x is not x_pure_rand:
+                            check_data(recv_x, rank_prefix_matrix)
+
+                    # Test combine
+                    combine_args = {
+                        "x": recv_x,
+                        "handle": handle,
+                        "config": config,
+                        "async_finish": async_mode,
+                    }
+                    if with_topk:
+                        combine_args.update({"topk_weights": recv_topk_weights})
+                    if previous_mode:
+                        dispatch_args.update({"previous_event": buffer.capture()})
+                    combined_x, combined_topk_weights, event = buffer.combine(
+                        **combine_args
+                    )
+                    event.current_stream_wait() if async_mode else ()
+                    check_x = combined_x.float() / is_token_in_rank.sum(
+                        dim=1
+                    ).unsqueeze(1)
+                    ref_x = x_pure_rand if current_x is x_pure_rand else x
+                    assert calc_diff(check_x, ref_x) < 5e-6
+                    if with_topk:
+                        check_topk_weights = (
+                            combined_topk_weights
+                            if (current_x is x_pure_rand)
+                            else (
+                                combined_topk_weights
+                                / is_token_in_rank.sum(dim=1).unsqueeze(1)
+                            )
+                        )
+                        ref_topk_weights = (
+                            topk_weights_pure_rand
+                            if current_x is x_pure_rand
+                            else topk_weights
+                        )
+                        assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
+
+                    # For later tuning
+                    dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
+                    combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
+
+                    if local_rank == 0:
+                        print(" passed", flush=True)
+    if local_rank == 0:
+        print("", flush=True)
+
+    # Tune dispatch performance
+    best_dispatch_results = None
+    fp8_factor = (1 + 4 / 128) / 2
+    for current_x in (x_e4m3, x):
+        best_time, best_results = 1e10, None
+        nvl_recv_bytes = (
+            (dispatch_bf16_nvl_recv_bytes * fp8_factor)
+            if isinstance(current_x, tuple)
+            else dispatch_bf16_nvl_recv_bytes
+        )
+        for nvl_chunk_size in range(4, 33, 4):
+            config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
+            tune_args = {"x": current_x, "handle": handle, "config": config}
+            t = bench(lambda: buffer.dispatch(**tune_args))[0]
+            if t < best_time:
+                best_time, best_results = t, (num_sms, nvl_chunk_size)
+            if local_rank == 0:
+                print(
+                    f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
+                    flush=True,
+                )
+        if local_rank == 0:
+            print(
+                f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
+                flush=True,
+            )
+            print("", flush=True)
+
+        if isinstance(current_x, tuple):
+            # Gather FP8 the best config from rank 0
+            best_dispatch_results = torch.tensor(
+                [best_results[0], best_results[1]], dtype=torch.int32, device="cuda"
+            )
+            all_best_fp8_results_list = [
+                torch.zeros_like(best_dispatch_results)
+                for _ in range(torch.distributed.get_world_size())
+            ]
+            dist.all_gather(
+                all_best_fp8_results_list, best_dispatch_results, group=group
+            )
+            best_dispatch_results = all_best_fp8_results_list[0].tolist()
+    dispatch_config = deep_ep.Config(
+        best_dispatch_results[0], best_dispatch_results[1], nvl_buffer_size
+    )
+
+    dispatch_args = {
+        "x": x,
+        "num_tokens_per_rank": num_tokens_per_rank,
+        "is_token_in_rank": is_token_in_rank,
+        "num_tokens_per_expert": num_tokens_per_expert,
+        "config": dispatch_config if dispatch_config is not None else config,
+    }
+    recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
+
+    # Tune combine performance
+    best_time, best_results = 1e10, None
+    for nvl_chunk_size in range(1, 7, 1):
+        config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
+        tune_args = {"x": recv_x, "handle": handle, "config": config}
+        t = bench(lambda: buffer.combine(**tune_args))[0]
+        if local_rank == 0:
+            print(
+                f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
+                flush=True,
+            )
+            if t < best_time:
+                best_time, best_results = t, (num_sms, nvl_chunk_size)
+
+    if local_rank == 0:
+        print(
+            f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}: {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
+            flush=True,
+        )
+        print("", flush=True)
+
+
+# noinspection PyUnboundLocalVariable
+def test_loop(local_rank: int, num_local_ranks: int):
+    rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
+    test_ll_compatibility, num_rdma_bytes = False, 0
+    if test_ll_compatibility:
+        ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
+        num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
+            ll_num_tokens, ll_hidden, num_ranks, ll_num_experts
+        )
+
+    buffer = deep_ep.Buffer(
+        group,
+        int(1e9),
+        num_rdma_bytes,
+        low_latency_mode=test_ll_compatibility,
+        num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
+    )
+    torch.manual_seed(rank)
+
+    for i in (24,):
+        test_main(i, local_rank, num_ranks, rank, buffer, group)
+        if local_rank == 0:
+            print("", flush=True)
+
+    # Test compatibility with low latency functions
+    if test_ll_compatibility:
+        buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
+        test_deepep_low_latency.test_main(
+            ll_num_tokens,
+            ll_hidden,
+            ll_num_experts,
+            ll_num_topk,
+            rank,
+            num_ranks,
+            group,
+            buffer,
+            seed=1,
+        )
+
+
+if __name__ == "__main__":
+    num_processes = 8
+    torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

test/manual/ep/test_deepep_low_latency.py

@@ -0,0 +1,325 @@
+# Copy from deepseek-ai/DeepEP/tests/test_low_latency.py
+
+import random
+from functools import partial
+
+import deep_ep
+import torch
+import torch.distributed as dist
+
+from sglang.test.test_deepep_utils import (
+    bench,
+    bench_kineto,
+    calc_diff,
+    hash_tensor,
+    init_dist,
+    per_token_cast_back,
+)
+
+
+def test_main(
+    num_tokens: int,
+    hidden: int,
+    num_experts: int,
+    num_topk: int,
+    rank: int,
+    num_ranks: int,
+    group: dist.ProcessGroup,
+    buffer: deep_ep.Buffer,
+    seed: int = 0,
+):
+    torch.manual_seed(seed + rank)
+    random.seed(seed + rank)
+
+    assert num_experts % num_ranks == 0
+    num_local_experts = num_experts // num_ranks
+
+    # NOTES: the integers greater than 256 exceeds the BF16 precision limit
+    rank_offset = 128
+    assert (
+        num_ranks - rank_offset < 257
+    ), "Too many ranks (exceeding test precision limit)"
+
+    x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (
+        rank - rank_offset
+    )
+    x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1)
+    scores = (
+        torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
+        + 1
+    )
+    topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1]
+    topk_weights = torch.randn(
+        (num_tokens, num_topk), dtype=torch.float32, device="cuda"
+    ).abs()
+
+    # Randomly mask some positions
+    for i in range(10):
+        topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = (
+            -1
+        )
+
+    # Check dispatch correctness
+    do_check = True
+    hash_value, num_times = 0, 0
+    for return_recv_hook in (False, True):
+        for dispatch_use_fp8 in (False, True):
+            num_times += 1
+            for i in range((num_times % 2) + 1):
+                packed_recv_x, packed_recv_count, handle, event, hook = (
+                    buffer.low_latency_dispatch(
+                        x,
+                        topk_idx,
+                        num_tokens,
+                        num_experts,
+                        use_fp8=dispatch_use_fp8,
+                        async_finish=not return_recv_hook,
+                        return_recv_hook=return_recv_hook,
+                    )
+                )
+                hook() if return_recv_hook else event.current_stream_wait()
+            packed_recv_x = (
+                (packed_recv_x[0], packed_recv_x[1].contiguous())
+                if dispatch_use_fp8
+                else packed_recv_x
+            )
+            simulated_gemm_x = (
+                per_token_cast_back(
+                    packed_recv_x[0].view(-1, hidden),
+                    packed_recv_x[1].view(-1, hidden // 128),
+                ).view(packed_recv_x[0].shape)
+                if dispatch_use_fp8
+                else packed_recv_x.clone()
+            )
+            all_topk_idx = torch.empty(
+                (num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda"
+            )
+            dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group)
+            for i in range(num_local_experts if do_check else 0):
+                expert_id = rank * num_local_experts + i
+                recv_x = (
+                    per_token_cast_back(packed_recv_x[0][i], packed_recv_x[1][i])
+                    if dispatch_use_fp8
+                    else packed_recv_x[i]
+                )
+                recv_count, recv_src_info, recv_layout_range = (
+                    packed_recv_count[i],
+                    handle[0][i],
+                    handle[1][i],
+                )
+
+                # Check expert indices
+                int_mask = (2**32) - 1
+                num_valid_tokens = recv_count.item()
+                assert (
+                    num_valid_tokens == (recv_layout_range & int_mask).sum().item()
+                ), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()"
+                assert (
+                    num_valid_tokens == (all_topk_idx == expert_id).sum().item()
+                ), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
+
+                # Check received data
+                recv_x = recv_x[:num_valid_tokens]
+                recv_x_amin = recv_x[:, :-128].amin(dim=-1)
+                recv_src_info = recv_src_info[:num_valid_tokens]
+                assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1))
+                assert (
+                    recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens
+                ).sum().item() == 0
+                for j in range(num_ranks):
+                    begin_idx, count = (recv_layout_range[j] >> 32).item(), (
+                        recv_layout_range[j] & int_mask
+                    ).item()
+                    assert (recv_x_amin == j - rank_offset).sum().item() == (
+                        all_topk_idx[j] == expert_id
+                    ).sum().item()
+                    assert (
+                        recv_x[begin_idx : begin_idx + count][:-128] - j
+                    ).sum().item() == 0
+                if dispatch_use_fp8:
+                    hash_value ^= hash_tensor(packed_recv_x[0][i, :num_valid_tokens])
+                    hash_value ^= hash_tensor(packed_recv_x[1][i, :num_valid_tokens])
+                else:
+                    hash_value ^= hash_tensor(packed_recv_x[i, :num_valid_tokens])
+
+            # Check combine correctness
+            for zero_copy in (False, True):
+                if zero_copy:
+                    buffer.get_next_low_latency_combine_buffer(handle)[
+                        :, :, :
+                    ] = simulated_gemm_x
+                out = torch.empty(
+                    (num_tokens, hidden), dtype=torch.bfloat16, device="cuda"
+                )
+                combined_x, event, hook = buffer.low_latency_combine(
+                    simulated_gemm_x,
+                    topk_idx,
+                    topk_weights,
+                    handle,
+                    async_finish=not return_recv_hook,
+                    zero_copy=zero_copy,
+                    return_recv_hook=return_recv_hook,
+                    out=out,
+                )
+                hook() if return_recv_hook else event.current_stream_wait()
+                if do_check:
+                    diff = calc_diff(
+                        x
+                        * topk_weights.masked_fill(topk_idx == -1, 0)
+                        .sum(dim=1)
+                        .view(-1, 1),
+                        combined_x,
+                    )
+                    assert torch.isnan(combined_x).sum().item() == 0
+                    assert diff < 1e-5, f"Error: {diff=}, {zero_copy=}"
+                    hash_value ^= hash_tensor(combined_x)
+
+    def create_test_cast_with_outliers(num_outliers):
+        tmp = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
+        tmp /= tmp.abs().amax(dim=1).view(-1, 1)
+        assert tmp.abs().amax().item() <= 1
+
+        # Create some amax outliers
+        for i in range(num_outliers):
+            tmp[random.randint(0, num_tokens - 1)] *= 1e3
+        return tmp
+
+    # noinspection PyShadowingNames
+    def large_gemm_with_hook(hook):
+        mat_0 = torch.randn((8192, 8192), dtype=torch.float)
+        mat_1 = torch.randn((8192, 8192), dtype=torch.float)
+        mat_0 @ mat_1
+        hook()
+
+    # noinspection PyShadowingNames
+    def test_func(zero_copy: bool, return_recv_hook: bool):
+        recv_x, recv_count, handle, event, hook = buffer.low_latency_dispatch(
+            x,
+            topk_idx,
+            num_tokens,
+            num_experts,
+            async_finish=False,
+            return_recv_hook=return_recv_hook,
+        )
+        large_gemm_with_hook(hook) if return_recv_hook else None
+        if zero_copy:
+            buffer.get_next_low_latency_combine_buffer(handle)[
+                :, :, :
+            ] = simulated_gemm_x
+        combined_x, event, hook = buffer.low_latency_combine(
+            simulated_gemm_x,
+            topk_idx,
+            topk_weights,
+            handle,
+            zero_copy=zero_copy,
+            return_recv_hook=return_recv_hook,
+        )
+        large_gemm_with_hook(hook) if return_recv_hook else None
+
+    # Calculate bandwidth
+    num_fp8_bytes, num_bf16_bytes = (hidden + hidden / 128 * 4 + 16), hidden * 2
+    num_dispatch_comm_bytes, num_combine_comm_bytes = 0, 0
+    for i in range(num_tokens):
+        num_selections = (topk_idx[i] != -1).sum().item()
+        num_dispatch_comm_bytes += num_fp8_bytes * num_selections
+        num_combine_comm_bytes += num_bf16_bytes * num_selections
+
+    # Dispatch + combine testing
+    avg_t, min_t, max_t = bench(
+        partial(test_func, zero_copy=False, return_recv_hook=False)
+    )
+    print(
+        f"[rank {rank}] Dispatch + combine bandwidth: {(num_dispatch_comm_bytes + num_combine_comm_bytes) / 1e9 / avg_t:.2f} GB/s, "
+        f"avg_t={avg_t * 1e6:.2f} us, min_t={min_t * 1e6:.2f} us, max_t={max_t * 1e6:.2f} us",
+        flush=True,
+    )
+
+    # Separate profiling
+    for return_recv_hook in (False, True):
+        group.barrier()
+        dispatch_t, combine_t = bench_kineto(
+            partial(test_func, zero_copy=True, return_recv_hook=return_recv_hook),
+            kernel_names=("dispatch", "combine"),
+            barrier_comm_profiling=True,
+            suppress_kineto_output=True,
+        )
+        if not return_recv_hook:
+            print(
+                f"[rank {rank}] Dispatch bandwidth: {num_dispatch_comm_bytes / 1e9 / dispatch_t:.2f} GB/s, avg_t={dispatch_t * 1e6:.2f} us | "
+                f"Combine bandwidth: {num_combine_comm_bytes / 1e9 / combine_t:.2f} GB/s, avg_t={combine_t * 1e6:.2f} us",
+                flush=True,
+            )
+        else:
+            print(
+                f"[rank {rank}] Dispatch send/recv time: {dispatch_t * 2 * 1e6:.2f} us | "
+                f"Combine send/recv time: {combine_t * 2 * 1e6:.2f} us",
+                flush=True,
+            )
+
+    return hash_value
+
+
+# noinspection PyUnboundLocalVariable
+def test_loop(local_rank: int, num_local_ranks: int):
+    rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
+    num_tokens, hidden, num_topk, num_experts = 128, 7168, 8, 288
+
+    num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
+        num_tokens, hidden, num_ranks, num_experts
+    )
+    if local_rank == 0:
+        print(f"Allocating buffer size: {num_rdma_bytes / 1e6} MB ...", flush=True)
+    buffer = deep_ep.Buffer(
+        group,
+        num_rdma_bytes=num_rdma_bytes,
+        low_latency_mode=True,
+        num_qps_per_rank=num_experts // num_ranks,
+    )
+    test_main(
+        num_tokens,
+        hidden,
+        num_experts,
+        num_topk,
+        rank,
+        num_ranks,
+        group,
+        buffer,
+        seed=1,
+    )
+
+    do_pressure_test = False
+    for seed in range(int(1e9) if do_pressure_test else 0):
+        if local_rank == 0:
+            print(f"Testing with seed {seed} ...", flush=True)
+        ref_hash = test_main(
+            num_tokens,
+            hidden,
+            num_experts,
+            num_topk,
+            rank,
+            num_ranks,
+            group,
+            buffer,
+            seed=seed,
+        )
+        for i in range(20):
+            assert (
+                test_main(
+                    num_tokens,
+                    hidden,
+                    num_experts,
+                    num_topk,
+                    rank,
+                    num_ranks,
+                    group,
+                    buffer,
+                    seed=seed,
+                )
+                == ref_hash
+            ), f"Error: seed={seed}"
+
+
+if __name__ == "__main__":
+    # TODO: you may modify NUMA binding for less CPU overhead
+    num_processes = 8
+    torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)

test/manual/ep/test_eplb.py

@@ -0,0 +1,182 @@
+import os
+import tempfile
+import time
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+import sglang as sgl
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+def get_a2a_backend_config():
+    # On Blackwell or machines where DeepEP is hard to compile, set
+    # SGLANG_EPLB_TEST_MOE_A2A_BACKEND=flashinfer to use FlashInfer A2A.
+    moe_a2a_backend = os.environ.get("SGLANG_EPLB_TEST_MOE_A2A_BACKEND", "deepep")
+    args = ["--moe-a2a-backend", moe_a2a_backend]
+    kwargs = {"moe_a2a_backend": moe_a2a_backend}
+    if moe_a2a_backend == "deepep":
+        args.extend(["--deepep-mode", "normal"])
+        kwargs["deepep_mode"] = "normal"
+    elif moe_a2a_backend == "flashinfer":
+        args.extend(["--moe-runner-backend", "flashinfer_cutlass"])
+        kwargs["moe_runner_backend"] = "flashinfer_cutlass"
+    return args, kwargs
+
+
+def get_a2a_backend_args():
+    return get_a2a_backend_config()[0]
+
+
+def get_a2a_backend_kwargs():
+    return get_a2a_backend_config()[1]
+
+
+class _BaseTestDynamicEPLB(CustomTestCase):
+    extra_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        with (
+            envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False),
+            envs.SGLANG_EXPERT_LOCATION_UPDATER_CANARY.override(True),
+        ):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--trust-remote-code",
+                    "--tp",
+                    "2",
+                    "--dp",
+                    "2",
+                    "--enable-dp-attention",
+                    *get_a2a_backend_args(),
+                    "--disable-cuda-graph",
+                    "--enable-eplb",
+                    "--ep-num-redundant-experts",
+                    "4",
+                    "--eplb-rebalance-num-iterations",
+                    "50",
+                    "--expert-distribution-recorder-buffer-size",
+                    "50",
+                    # TODO pr-chain: enable later
+                    # "--enable-expert-distribution-metrics",
+                    # TODO auto determine these flags
+                    "--expert-distribution-recorder-mode",
+                    "stat",
+                    "--ep-dispatch-algorithm",
+                    "static",
+                    *cls.extra_args,
+                ],
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        time.sleep(5)
+
+    def test_mmlu(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=64,
+            num_threads=32,
+        )
+
+        metrics = run_eval(args)
+        self.assertGreater(metrics["score"], 0.5)
+
+
+class TestDynamicEPLBSimple(_BaseTestDynamicEPLB):
+    pass
+
+
+class TestDynamicEPLBMultiChunk(_BaseTestDynamicEPLB):
+    extra_args = ["--eplb-rebalance-layers-per-chunk", "1"]
+
+
+class TestStaticEPLB(CustomTestCase):
+    def test_save_expert_distribution_and_init_expert_location(self):
+        envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            engine_kwargs = dict(
+                model_path=DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+                trust_remote_code=True,
+                ep_num_redundant_experts=4,
+                enable_dp_attention=True,
+                disable_cuda_graph=True,
+                expert_distribution_recorder_mode="stat",
+                tp_size=2,
+                dp_size=2,
+                log_level="info",
+                # TODO pr-chain: enable later
+                # enable_expert_distribution_metrics=True,
+            )
+            engine_kwargs.update(get_a2a_backend_kwargs())
+
+            print(f"Action: start engine")
+            envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
+            engine = sgl.Engine(
+                **engine_kwargs,
+                disable_overlap_schedule=True,
+                port=11000,
+            )
+            engine.start_expert_distribution_record()
+            self._assert_engine_generate_correct(engine)
+
+            print(f"Action: dump_expert_distribution_record")
+            engine.dump_expert_distribution_record()
+            snapshot_path = list(Path(tmp_dir).glob("*.pt"))[0]
+            assert snapshot_path is not None
+            print(f"{snapshot_path=}")
+
+            print(f"Action: shutdown engine")
+            engine.shutdown()
+            del engine
+            engine = None
+            time.sleep(5)
+
+            print(f"Action: start engine with init_expert_location")
+            engine = sgl.Engine(
+                **engine_kwargs,
+                init_expert_location=str(snapshot_path),
+                # TODO auto determine these flags
+                ep_dispatch_algorithm="static",
+                port=12000,
+            )
+            self._assert_engine_generate_correct(engine)
+            print(f"Action: shutdown engine")
+            engine.shutdown()
+            del engine
+            engine = None
+            time.sleep(5)
+
+    def _assert_engine_generate_correct(self, engine: sgl.Engine):
+        output = engine.generate(
+            prompt=["1+1=2, 2+2=4", "One plus one is two, two plus two is four"],
+            sampling_params=dict(max_new_tokens=8, temperature=0.0),
+        )
+        print(f"engine.generate {output=}")
+        self.assertEqual(
+            [x["text"] for x in output],
+            [", 4+4=8,", ", four plus four is eight, eight"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/manual/ep/test_eplb_mori.py

@@ -0,0 +1,242 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.server_args import ZMQ_TCP_PORT_DELTA
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.network import is_port_available
+from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
+from sglang.test.test_utils import (
+    DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+def wait_all_ports_release(base_url, timeout_s=60):
+    import time
+
+    port = int(base_url.split(":")[-1])
+    offsets = [
+        0,
+        ZMQ_TCP_PORT_DELTA,
+        ZMQ_TCP_PORT_DELTA + 1,
+        ZMQ_TCP_PORT_DELTA + 2,
+        ZMQ_TCP_PORT_DELTA + 3,
+        ZMQ_TCP_PORT_DELTA + 4,
+    ]
+    for _ in range(timeout_s):
+        if all(is_port_available(port + off) for off in offsets):
+            return
+        time.sleep(1)
+    print(f"Warning: some ports still occupied after {timeout_s}s")
+
+
+mori_env = {
+    **os.environ,
+    "SGLANG_USE_AITER": "1",
+    "SGLANG_MORI_DISPATCH_DTYPE": "bf16",
+    "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "4096",
+    "SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE": "32",
+    "MORI_SHMEM_MODE": "ISOLATION",
+}
+
+common_args = [
+    "--tp-size",
+    "8",
+    "--ep-size",
+    "8",
+    "--dp-size",
+    "8",
+    "--enable-dp-attention",
+    "--moe-a2a-backend",
+    "mori",
+    "--trust-remote-code",
+    "--load-balance-method",
+    "round_robin",
+    "--moe-dense-tp-size",
+    "1",
+    "--enable-dp-lm-head",
+    "--mem-fraction-static",
+    "0.6",
+    "--chunked-prefill-size",
+    "32768",
+    "--max-running-requests",
+    "128",
+    "--context-length",
+    "12288",
+    "--attention-backend",
+    "aiter",
+    "--cuda-graph-max-bs-decode",
+    "32",
+]
+
+eplb_args = [
+    "--enable-eplb",
+    "--ep-num-redundant-experts",
+    "32",
+    "--eplb-rebalance-num-iterations",
+    "50",
+    "--expert-distribution-recorder-buffer-size",
+    "50",
+    "--ep-dispatch-algorithm",
+    "static",
+]
+
+
+mtp_args = [
+    "--speculative-algo",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+]
+
+
+class TestEPLBMoriStat(CustomTestCase):
+    """EPLB with mori backend, stat mode (on_select_experts path)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+class TestEPLBMoriStatApprox(CustomTestCase):
+    """EPLB with mori backend, stat_approx mode (local_expert_count kernel)."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat_approx",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+class TestEPLBMoriMultiChunk(CustomTestCase):
+    """EPLB with mori backend, chunked layer updates."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = (
+            common_args
+            + eplb_args
+            + [
+                "--deepep-mode",
+                "normal",
+                "--expert-distribution-recorder-mode",
+                "stat",
+                "--eplb-rebalance-layers-per-chunk",
+                "1",
+            ]
+        )
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5,
+            other_args=other_args,
+            env=mori_env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        wait_all_ports_release(cls.base_url)
+
+    def test_gsm8k(self):
+        args = SimpleNamespace(
+            num_shots=5,
+            data_path=None,
+            num_questions=1209,
+            max_new_tokens=512,
+            parallel=1209,
+            host="http://127.0.0.1",
+            port=int(self.base_url.split(":")[-1]),
+        )
+        metrics = run_eval_few_shot_gsm8k(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["accuracy"], 0.9)
+
+
+if __name__ == "__main__":
+    unittest.main()