chore: vendor sglang v0.5.10 snapshot

third_party/sglang/test/README.md

@@ -0,0 +1,115 @@
+# 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.
+- `srt/`: Legacy CI setup, to be deprecated.
+
+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 python/sglang/jit_kernel/tests/test_add_constant.py
+
+# Run a suite
+python3 test/run_suite.py --hw cpu --suite stage-a-test-cpu
+python3 test/run_suite.py --hw cuda --suite stage-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 stage-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, suite="stage-b-test-1-gpu-small")
+```
+
+Parameters: `est_time` (seconds), `suite` (target suite), `nightly=True` (nightly-only), `disabled="reason"` (temporarily disable).
+
+Keep `est_time` and `suite` as **literal values** — `run_suite.py` collects them by AST parsing.
+
+JIT kernel files live outside `test/registered/` but still use registration:
+- Correctness tests: `python/sglang/jit_kernel/tests/test_*.py` → `stage-b-kernel-unit-1-gpu-large`
+- Benchmarks: `python/sglang/jit_kernel/benchmark/bench_*.py` → `stage-b-kernel-benchmark-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 | `stage-a-test-cpu` |
+| Small GPU (fits 5090, 32GB) | `stage-b-test-1-gpu-small` (most tests go here) |
+| Large GPU memory or Hopper features | `stage-b-test-1-gpu-large` |
+| JIT kernel correctness | `stage-b-kernel-unit-1-gpu-large` |
+| JIT kernel benchmarks | `stage-b-kernel-benchmark-1-gpu-large` |
+| Multi-GPU (2/4/8) | `stage-b-test-2-gpu-large`, `stage-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/srt/nightly/test_vlms_mmmu_eval.py`.

third_party/sglang/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

third_party/sglang/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

third_party/sglang/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

third_party/sglang/test/manual/ascend/test_ascend_deepseek_mtp.py

@@ -0,0 +1,94 @@
+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_SPEC_V2.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()

third_party/sglang/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()

third_party/sglang/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()

third_party/sglang/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()

third_party/sglang/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()

third_party/sglang/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")

third_party/sglang/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)

third_party/sglang/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)

third_party/sglang/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)

third_party/sglang/test/manual/ep/test_eplb.py

@@ -0,0 +1,154 @@
+import tempfile
+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,
+)
+
+
+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",
+                    "--moe-a2a-backend",
+                    "deepep",
+                    "--deepep-mode",
+                    "normal",
+                    "--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)
+
+    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,
+                moe_a2a_backend="deepep",
+                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,
+            )
+
+            print(f"Action: start engine")
+            envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
+            engine = sgl.Engine(
+                **engine_kwargs,
+                disable_overlap_schedule=True,
+            )
+            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
+
+            print(f"Action: start engine with init_expert_location")
+            engine = sgl.Engine(
+                **engine_kwargs,
+                init_expert_location=str(snapshot_path),
+                port=21000,
+                # TODO auto determine these flags
+                ep_dispatch_algorithm="static",
+            )
+            self._assert_engine_generate_correct(engine)
+            print(f"Action: shutdown engine")
+            engine.shutdown()
+            del engine
+
+    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()

third_party/sglang/test/manual/ep/test_moe_deepep.py

@@ -0,0 +1,116 @@
+import json
+import unittest
+from types import SimpleNamespace
+
+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,
+)
+
+
+class TestPureTP(CustomTestCase):
+    @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",
+                "2",
+                "--moe-a2a-backend",
+                "deepep",
+                "--disable-cuda-graph",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    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 TestDPAttn(unittest.TestCase):
+    @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):
+            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",
+                    "--moe-a2a-backend",
+                    "deepep",
+                    "--deepep-mode",
+                    "normal",
+                    "--disable-cuda-graph",
+                    # Test custom config
+                    "--deepep-config",
+                    json.dumps(
+                        {
+                            "normal_dispatch": {
+                                "num_sms": 20,
+                                "num_max_nvl_chunked_send_tokens": 16,
+                                "num_max_nvl_chunked_recv_tokens": 256,
+                                "num_max_rdma_chunked_send_tokens": 6,
+                                "num_max_rdma_chunked_recv_tokens": 128,
+                            },
+                            "normal_combine": {
+                                "num_sms": 20,
+                                "num_max_nvl_chunked_send_tokens": 6,
+                                "num_max_nvl_chunked_recv_tokens": 256,
+                                "num_max_rdma_chunked_send_tokens": 6,
+                                "num_max_rdma_chunked_recv_tokens": 128,
+                            },
+                        }
+                    ),
+                ],
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    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)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/ep/test_moe_deepep_eval_accuracy_large.py

@@ -0,0 +1,75 @@
+"""
+Usage:
+python -m unittest test_moe_deepep_eval_accuracy_large.TestMoEDeepEPEvalAccuracyLarge.test_mmlu
+"""
+
+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.test_utils import (
+    DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestMoEDeepEPEvalAccuracyLarge(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_DEEPEP_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",
+                "8",
+                "--moe-a2a-backend",
+                "deepep",
+                "--cuda-graph-max-bs",
+                "128",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=64,
+            num_shots=8,
+        )
+        metrics = run_eval(args)
+        print(f"Eval accuracy of GSM8K: {metrics=}")
+
+        self.assertGreater(metrics["score"], 0.93)
+
+    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)
+        print(f"Eval accuracy of MMLU: {metrics=}")
+        self.assertGreater(metrics["score"], 0.87)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/ep/test_mooncake_expert_backup.py

@@ -0,0 +1,143 @@
+import time
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    CustomTestCase,
+    popen_launch_pd_server,
+)
+
+ib_devices = get_rdma_devices_args()
+
+
+class TestBackup(CustomTestCase):
+    extra_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_port = 20000
+        cls.base_url = f"http://127.0.0.1:{cls.base_port}"
+        cls.num_processes = 2
+        # TODO (stage 100): in the future, implement a specified multiprocess launcher
+        cls.processes = [
+            popen_launch_pd_server(
+                cls.model,
+                f"http://127.0.0.1:{cls.base_port + i}",
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--trust-remote-code",
+                    "--tp",
+                    "4",
+                    "--enable-dp-attention",
+                    "--dp",
+                    "4",
+                    "--elastic-ep-backend",
+                    "mooncake",
+                    "--mooncake-ib-device",
+                    ib_devices,
+                    "--moe-a2a-backend",
+                    "mooncake",
+                    "--deepep-mode",
+                    "low_latency",
+                    "--moe-dense-tp-size",
+                    "1",
+                    "--enable-dp-lm-head",
+                    "--enable-two-batch-overlap",
+                    "--disable-custom-all-reduce",
+                    "--enable-elastic-expert-backup",
+                    "--enable-eplb",
+                    "--eplb-rebalance-num-iterations",
+                    "50",
+                    "--chunked-prefill-size",
+                    "512",
+                    "--cuda-graph-max-bs",
+                    "128",
+                    "--max-running-requests",
+                    "512",
+                    "--mem-fraction-static",
+                    "0.5",
+                    "--dist-init-addr",
+                    "127.0.0.1:5000",
+                    "--nnodes",
+                    f"{cls.num_processes}",
+                    "--node-rank",
+                    f"{i}",
+                    "--base-gpu-id",
+                    f"{i * 2}",
+                ],
+            )
+            for i in range(cls.num_processes)
+        ]
+
+        server_ready = [False] * cls.num_processes
+        start_time = time.perf_counter()
+        with requests.Session() as session:
+            while (
+                time.perf_counter() - start_time < DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+                and not all(server_ready)
+            ):
+                for i, process in enumerate(cls.processes):
+                    return_code = process.poll()
+                    if return_code is not None:
+                        # Server failed to start (non-zero exit code) or crashed
+                        raise Exception(
+                            f"Server process exited with code {return_code}. "
+                            "Check server logs for errors."
+                        )
+
+                    try:
+                        headers = {
+                            "Content-Type": "application/json; charset=utf-8",
+                        }
+                        response = session.get(
+                            f"http://127.0.0.1:{cls.base_port + i}/health_generate",
+                            headers=headers,
+                        )
+                        if response.status_code == 200:
+                            server_ready[i] = True
+                    except requests.RequestException:
+                        pass
+
+                    return_code = process.poll()
+                    if return_code is not None:
+                        raise Exception(
+                            f"Server unexpectedly exits ({return_code=}). Usually there will be error logs describing the cause far above this line."
+                        )
+
+                    time.sleep(10)
+        if not all(server_ready):
+            for process in cls.processes:
+                kill_process_tree(process.pid)
+            raise TimeoutError("Server failed to start within the timeout period.")
+
+    @classmethod
+    def tearDownClass(cls):
+        for process in cls.processes:
+            kill_process_tree(process.pid)
+
+    def test_gsm8k(self):
+        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)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/ep/test_nixl_ep.py

@@ -0,0 +1,115 @@
+import os
+import time
+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.server_fixtures.disaggregation_fixture import get_rdma_devices_args
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_MLA,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+TEST_MODEL = os.environ.get("NIXL_EP_TEST_MODEL", DEFAULT_MODEL_NAME_FOR_TEST_MLA)
+os.environ.setdefault("SGLANG_NIXL_EP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", "1024")
+
+ib_devices = get_rdma_devices_args()
+
+NIXL_COMMON = [
+    "--trust-remote-code",
+    "--moe-a2a-backend",
+    "nixl",
+    "--deepep-mode",
+    "low_latency",
+    "--tp",
+    "8",
+    "--mem-fraction-static",
+    "0.78",
+]
+DP_ATTN = ["--dp", "8", "--enable-dp-attention"]
+ELASTIC_NIXL = [
+    "--elastic-ep-backend",
+    "nixl",
+    "--enable-eplb",
+    "--ep-num-redundant-experts",
+    "24",
+]
+ELASTIC_MOONCAKE = [
+    "--elastic-ep-backend",
+    "mooncake",
+    "--mooncake-ib-device",
+    ib_devices,
+    "--enable-eplb",
+    "--ep-num-redundant-experts",
+    "24",
+]
+
+
+class _EPTestBase(CustomTestCase):
+    server_args: list[str] = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = TEST_MODEL
+        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=cls.server_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+        cls.process.wait(timeout=15)
+        time.sleep(2)
+
+    def _run_gsm8k(self):
+        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)
+        return metrics
+
+    def test_gsm8k(self):
+        metrics = self._run_gsm8k()
+        self.assertGreater(metrics["score"], 0.60)
+
+
+class TestNixlEPTP(_EPTestBase):
+    server_args = [*NIXL_COMMON]
+
+
+class TestNixlEPDPAttn(_EPTestBase):
+    server_args = [*NIXL_COMMON, *DP_ATTN]
+
+
+class TestNixlEPElasticEP(_EPTestBase):
+    server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_NIXL]
+
+
+class TestNixlMoeMooncakeElasticEP(_EPTestBase):
+    server_args = [*NIXL_COMMON, *DP_ATTN, *ELASTIC_MOONCAKE]
+
+    pkill_process_1 = "sglang::scheduler_DP1_TP8_EP8"
+
+    def test_gsm8k_fault_1(self):
+        os.system(f"pkill -f {self.pkill_process_1}")
+        metrics = self._run_gsm8k()
+        self.assertGreater(metrics["score"], 0.60)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/hicache/test_disaggregation_hicache.py

@@ -0,0 +1,257 @@
+import os
+import random
+import tempfile
+import unittest
+from typing import Dict
+
+import requests
+
+from sglang.benchmark.utils import get_tokenizer
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+
+class DisaggregationHiCacheBase(PDDisaggregationServerBase):
+    """Base class for disaggregation with HiCache tests"""
+
+    @classmethod
+    def setUpClass(cls):
+        super(DisaggregationHiCacheBase, cls).setUpClass()
+
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        cls.tokenizer = get_tokenizer(cls.model)
+        cls.temp_dir = tempfile.mkdtemp()
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Block until both
+        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 with HiCache enabled
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp-size",
+            "1",
+            "--page-size",
+            "64",
+            "--enable-hierarchical-cache",
+            "--hicache-ratio",
+            "1.2",
+            "--hicache-size",
+            "0",
+            "--hicache-write-policy",
+            "write_through",
+            "--hicache-storage-backend",
+            "file",
+            "--hicache-storage-prefetch-policy",
+            "wait_complete",
+            "--mem-fraction-static",
+            "0.8",
+        ]
+        prefill_args += cls.transfer_backend + cls.rdma_devices
+        env = {
+            **os.environ,
+            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
+        }
+        cls.process_prefill = popen_launch_pd_server(
+            cls.model,
+            cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+            env=env,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        pass
+
+    def gen_prompt(self, token_num: int) -> str:
+        all_available_tokens = list(self.tokenizer.get_vocab().values())
+        selected_tokens = random.choices(all_available_tokens, k=token_num)
+        return self.tokenizer.decode(selected_tokens)
+
+    def send_request(
+        self, prompt: str, max_tokens: int = 100, temperature: float = 0.0
+    ) -> Dict:
+        """Send a generate request and return response"""
+        response = requests.post(
+            f"{self.lb_url}/generate",
+            json={
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": temperature,
+                    "max_new_tokens": max_tokens,
+                    "ignore_eos": True,
+                },
+            },
+            timeout=60,
+        )
+
+        self.assertEqual(
+            response.status_code,
+            200,
+            f"Request failed: {response.status_code} - {response.text}",
+        )
+        return response.json()
+
+    def trigger_offloading_and_flush(self):
+        """Helper method to trigger offloading and flush cache"""
+        # Trigger offloading
+        self.send_request(self.gen_prompt(1), max_tokens=150)
+
+        # Flush device cache to force remote storage access.
+        res = requests.post(
+            f"{self.prefill_url}/flush_cache",
+            params={"timeout": 30},
+            timeout=40,
+        )
+        res.raise_for_status()
+
+
+class TestDisaggregationPrefillWithHiCache(DisaggregationHiCacheBase):
+    """Test disaggregation with HiCache enabled only on Prefill side"""
+
+    @classmethod
+    def start_decode(cls):
+        # Decode without HiCache offload
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp-size",
+            "1",
+            "--page-size",
+            "64",
+            "--mem-fraction-static",
+            "0.8",
+            "--base-gpu-id",
+            "1",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        env = {
+            **os.environ,
+            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
+        }
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+            env=env,
+        )
+
+    def test_prefill_cache_hit(self):
+        """Test that prefill cache works with repeated queries"""
+
+        repeated_prompt = self.gen_prompt(800)
+
+        # First request - should miss cache
+        self.send_request(repeated_prompt, max_tokens=100)
+
+        # Flush cache
+        self.trigger_offloading_and_flush()
+
+        # Second request - should hit cache (faster)
+        response2 = self.send_request(repeated_prompt, max_tokens=100)
+
+        # Assert cached tokens cnt
+        self.assertGreater(response2["meta_info"]["cached_tokens"], 700)
+
+
+class TestDisaggregationDecodeWithHiCache(DisaggregationHiCacheBase):
+    """Test disaggregation with HiCache enabled on both Prefill and Decode sides"""
+
+    @classmethod
+    def start_decode(cls):
+        # Decode with HiCache offload enabled
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp-size",
+            "1",
+            "--page-size",
+            "64",
+            "--mem-fraction-static",
+            "0.8",
+            "--base-gpu-id",
+            "1",
+            "--disaggregation-decode-enable-offload-kvcache",
+            "--hicache-ratio",
+            "1.2",
+            "--hicache-size",
+            "0",
+            "--hicache-storage-backend",
+            "file",
+            "--hicache-storage-prefetch-policy",
+            "wait_complete",
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        env = {
+            **os.environ,
+            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
+        }
+        cls.process_decode = popen_launch_pd_server(
+            cls.model,
+            cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+            env=env,
+        )
+
+    def test_multi_turn_conversation_cache(self):
+        """Test multi-turn conversation scenario with cache hit improvement"""
+
+        print("=== Multi-turn Conversation Cache Test ===")
+
+        # Turn 1
+        initial_prompt = self.gen_prompt(300)
+
+        response1 = self.send_request(initial_prompt, max_tokens=200, temperature=0.1)
+        current_context = initial_prompt + response1["text"]
+
+        # Turns 2-4: Continue generation based on previous context
+        previous_cached_tokens = 0
+
+        for turn in range(2, 5):
+            print(f"\nTurn {turn}: Continuing from previous context")
+
+            response = self.send_request(
+                current_context, max_tokens=200, temperature=0.1
+            )
+            cached_tokens = response["meta_info"]["cached_tokens"]
+
+            print(f"Turn {turn} cached tokens: {cached_tokens}")
+            print(f"Improvement: {cached_tokens - previous_cached_tokens} tokens")
+
+            # Assert cache improvement
+            self.assertGreater(
+                cached_tokens,
+                previous_cached_tokens,
+                f"Turn {turn} should have more cached tokens than turn {turn-1}",
+            )
+
+            # Update context and cached tokens for next iteration
+            current_context += response["text"]
+            previous_cached_tokens = cached_tokens
+
+            # Flush prefill cache
+            self.trigger_offloading_and_flush()
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/hicache/test_pp_with_hicache.py

@@ -0,0 +1,214 @@
+"""
+Usage:
+python3 -m unittest test_pp_with_hicache.TestPPWithHiCache.test_eval_accuracy
+"""
+
+import os
+import subprocess
+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.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    find_available_port,
+    popen_launch_server,
+)
+
+
+class TestPPWithHiCache(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = f"http://127.0.0.1:{find_available_port(23337)}"
+        parsed_url = urlparse(cls.base_url)
+        cls.base_host = parsed_url.hostname
+        cls.base_port = str(parsed_url.port)
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        cls._start_mooncake_services()
+
+        server_args_dict = {
+            "--enable-hierarchical-cache": True,
+            "--mem-fraction-static": 0.6,
+            "--hicache-ratio": 1.2,
+            "--page-size": 64,
+            "--enable-cache-report": True,
+            "--hicache-storage-prefetch-policy": "wait_complete",
+            "--hicache-storage-backend": "mooncake",
+            "--tp-size": 2,
+            "--pp-size": 2,
+            "--chunked-prefill-size": 256,
+            "--hicache-mem-layout": "page_first",
+        }
+
+        final_server_args = []
+        for key, value in server_args_dict.items():
+            final_server_args.append(str(key))
+            if value is not True:
+                final_server_args.append(str(value))
+
+        env_vars = {**os.environ, **cls._mooncake_env()}
+
+        try:
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=final_server_args,
+                env=env_vars,
+            )
+        except Exception:
+            cls._stop_mooncake_services()
+            raise
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process"):
+            kill_process_tree(cls.process.pid)
+        cls._stop_mooncake_services()
+
+    @classmethod
+    def _start_mooncake_services(cls):
+        try:
+            import mooncake.http_metadata_server  # type: ignore  # noqa: F401
+        except Exception as exc:  # pragma: no cover - environment dependent
+            raise unittest.SkipTest(
+                f"Mooncake metadata server module unavailable: {exc}"
+            ) from exc
+
+        cls._mooncake_master_port = find_available_port(50051)
+        cls._mooncake_metadata_port = find_available_port(8080)
+
+        try:
+            cls._mooncake_metadata_process = subprocess.Popen(
+                [
+                    "python3",
+                    "-m",
+                    "mooncake.http_metadata_server",
+                    "--port",
+                    str(cls._mooncake_metadata_port),
+                ],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                preexec_fn=os.setsid,
+            )
+        except (FileNotFoundError, subprocess.SubprocessError) as exc:
+            cls._stop_mooncake_services()
+            raise unittest.SkipTest(
+                f"Could not start Mooncake metadata service: {exc}"
+            ) from exc
+
+        try:
+            cls._mooncake_master_process = subprocess.Popen(
+                ["mooncake_master", "--port", str(cls._mooncake_master_port)],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                preexec_fn=os.setsid,
+            )
+        except (FileNotFoundError, subprocess.SubprocessError) as exc:
+            cls._stop_mooncake_services()
+            raise unittest.SkipTest(f"Could not start mooncake_master: {exc}") from exc
+
+        if not cls._wait_for_mooncake_ready():
+            cls._stop_mooncake_services()
+            raise unittest.SkipTest("Mooncake services did not become ready in time")
+
+    @classmethod
+    def _stop_mooncake_services(cls):
+        for attr in ("_mooncake_metadata_process", "_mooncake_master_process"):
+            proc = getattr(cls, attr, None)
+            if proc:
+                try:
+                    os.killpg(os.getpgid(proc.pid), 9)
+                    proc.wait(timeout=5)
+                except Exception:
+                    pass
+        cls._mooncake_metadata_process = None
+        cls._mooncake_master_process = None
+
+    @classmethod
+    def _mooncake_env(cls):
+        return {
+            "MOONCAKE_MASTER": f"127.0.0.1:{cls._mooncake_master_port}",
+            "MOONCAKE_PROTOCOL": "tcp",
+            "MC_MS_AUTO_DISC": "0",
+            "MOONCAKE_DEVICE": "",
+            "MOONCAKE_TE_META_DATA_SERVER": f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
+            "MOONCAKE_GLOBAL_SEGMENT_SIZE": "4294967296",
+            "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
+        }
+
+    @classmethod
+    def _wait_for_mooncake_ready(cls, timeout: int = 30) -> bool:
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            metadata_ready = False
+            master_ready = False
+
+            if (
+                getattr(cls, "_mooncake_metadata_process", None)
+                and cls._mooncake_metadata_process.poll() is None
+            ):
+                try:
+                    resp = requests.get(
+                        f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
+                        timeout=2,
+                    )
+                    print(resp)
+                    metadata_ready = True
+                except requests.RequestException:
+                    metadata_ready = False
+
+            if (
+                getattr(cls, "_mooncake_master_process", None)
+                and cls._mooncake_master_process.poll() is None
+            ):
+                if time.time() - start_time > 3:
+                    master_ready = True
+
+            if metadata_ready and master_ready:
+                return True
+
+            time.sleep(1.5)
+
+        return False
+
+    def flush_cache(self):
+        res = requests.post(
+            f"{self.base_url}/flush_cache",
+            params={"timeout": 30},
+            timeout=40,
+        )
+        res.raise_for_status()
+
+    def test_eval_accuracy(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=40,
+            num_threads=24,
+        )
+
+        metrics_initial = run_eval(args)
+        self.assertGreater(metrics_initial["score"], 0.6)
+
+        self.flush_cache()
+
+        metrics_cached = run_eval(args)
+        self.assertGreater(metrics_cached["score"], 0.6)
+
+        accuracy_diff = abs(metrics_initial["score"] - metrics_cached["score"])
+        self.assertLess(accuracy_diff, 0.05)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/kv_transfer/test_mooncake_transfer_engine.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+import argparse
+
+try:
+    import mooncake
+
+    BENCH_TOOL_PATH = f"{mooncake.__path__[0]}/transfer_engine_bench"
+    print(f"Mooncake is installed. Bench tool path:\n{BENCH_TOOL_PATH}")
+except ImportError:
+    BENCH_TOOL_PATH = None
+    print("Mooncake is not installed.")
+    exit(0)
+
+
+def run_cmd(args):
+    cmd = [BENCH_TOOL_PATH]
+    if args.initiator:
+        cmd += ["--mode=initiator"]
+    elif args.target:
+        cmd += ["--mode=target"]
+
+    if args.metadata_server:
+        cmd += [f"--metadata_server={args.metadata_server}"]
+    if args.mc_segment_id:
+        cmd += [f"--segment_id={args.mc_segment_id}"]
+    if args.device:
+        cmd += [f"--device_name={args.device}"]
+
+    if args.bench_h2h:
+        cmd += ["--use_vram=false"]
+
+    cmd += ["--auto_discovery"]
+    print(f"Executing command: {' '.join(cmd)}")
+    import subprocess
+
+    try:
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Command failed with error: {e}")
+        exit(1)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("--initiator", action="store_true", help="Run as initiator")
+    group.add_argument("--target", action="store_true", help="Run as target")
+    parser.add_argument("--metadata-server", type=str, default="P2PHANDSHAKE")
+    parser.add_argument("--mc-segment-id", type=str, default=None)
+    parser.add_argument("--bench-h2h", action="store_true")
+    parser.add_argument("--device", type=str, default="mlx5_0")
+    args = parser.parse_args()
+
+    print("Running Mooncake transfer engine benchmark...")
+    if not args.initiator and not args.target:
+        parser.error("Please specify --initiator or --target")
+    if args.initiator and args.mc_segment_id is None:
+        parser.error("Please specify --mc-segment-id for initiator")
+
+    run_cmd(args)
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""
+Test script for validating Mooncake transfer-engine gating and initialization.
+Tests the Mooncake-related branches in the current model-runner flow.
+
+This test verifies:
+1. MooncakeTransferEngine initialization conditions
+2. Different server argument combinations that trigger mooncake TE
+3. Mooncake transfer engine initialization with hostname, gpu_id, and ib_device
+
+Usage:
+    # Run from project root on 2 GPUs
+    CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py
+"""
+
+import argparse
+import multiprocessing
+import os
+import sys
+import time
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Optional
+from unittest.mock import patch
+
+
+@dataclass
+class ServerArgs:
+    """Mock ServerArgs for testing."""
+
+    disaggregation_mode: str = "null"
+    disaggregation_transfer_backend: str = "mooncake"
+    enable_hierarchical_cache: bool = False
+    hicache_storage_backend: str = "mooncake"
+    encoder_only: bool = False
+    language_only: bool = False
+    encoder_transfer_backend: str = "mooncake"
+    enable_elastic_expert_backup: bool = False
+    elastic_ep_backend: Optional[str] = None
+    disaggregation_ib_device: Optional[str] = None
+    mooncake_ib_device: Optional[str] = None
+
+
+def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
+    """
+    Test the condition logic for using MooncakeTransferEngine.
+    """
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+    dummy_runner = SimpleNamespace(server_args=server_args, gpu_id=0)
+    init_called = False
+
+    def _fake_init_mooncake_transfer_engine(*, hostname, gpu_id, ib_device):
+        nonlocal init_called
+        init_called = True
+        return SimpleNamespace(
+            hostname=hostname,
+            gpu_id=gpu_id,
+            ib_device=ib_device,
+        )
+
+    with patch(
+        "sglang.srt.distributed.device_communicators.mooncake_transfer_engine.init_mooncake_transfer_engine",
+        side_effect=_fake_init_mooncake_transfer_engine,
+    ), patch(
+        "sglang.srt.model_executor.model_runner.get_local_ip_auto",
+        return_value="127.0.0.1",
+    ):
+        ModelRunner.init_shared_mooncake_transfer_engine(dummy_runner)
+
+    return init_called
+
+
+def run_mooncake_init(
+    rank: int,
+    world_size: int,
+    master_port: int,
+    args: argparse.Namespace,
+    server_args: ServerArgs,
+):
+    """Worker function for testing mooncake transfer engine initialization."""
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+    os.environ["RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["LOCAL_RANK"] = str(rank)
+
+    # Import before try block to avoid NameError in finally
+    import torch
+    import torch.distributed as dist
+
+    dist_initialized = False
+
+    try:
+        # Initialize distributed environment
+        print(f"[Rank {rank}] Initializing distributed environment...")
+        dist.init_process_group(
+            backend="nccl",
+            world_size=world_size,
+            rank=rank,
+            init_method=f"tcp://127.0.0.1:{master_port}",
+            device_id=rank,
+        )
+        dist_initialized = True
+
+        # Set device
+        torch.cuda.set_device(rank)
+
+        # Sync to ensure all ranks are ready
+        dist.barrier()
+        print(f"[Rank {rank}] Distributed initialization complete.")
+
+        # Test the condition logic
+        use_mooncake_te = test_mooncake_te_condition(server_args)
+        print(f"[Rank {rank}] use_mooncake_te = {use_mooncake_te}")
+
+        if use_mooncake_te:
+            print(f"[Rank {rank}] Attempting to initialize MooncakeTransferEngine...")
+
+            from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+                init_mooncake_transfer_engine,
+            )
+            from sglang.srt.utils import get_local_ip_auto
+
+            ib_device = (
+                server_args.disaggregation_ib_device or server_args.mooncake_ib_device
+            )
+
+            print(f"[Rank {rank}] IB device: {ib_device}")
+
+            # Always actually initialize mooncake
+            engine = init_mooncake_transfer_engine(
+                hostname=get_local_ip_auto(),
+                gpu_id=rank,
+                ib_device=ib_device,
+            )
+            print(f"[Rank {rank}] Session ID: {engine.get_session_id()}")
+            print(f"[Rank {rank}] MooncakeTransferEngine initialized successfully!")
+
+            dist.barrier()
+
+        print(f"[Rank {rank}] Test completed successfully!")
+        sys.exit(0)
+
+    except ImportError as e:
+        print(f"[Rank {rank}] Mooncake not available (ImportError): {e}")
+        sys.exit(1)
+
+    except Exception as e:
+        print(f"[Rank {rank}] Test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+    finally:
+        # Cleanup
+        if dist_initialized and dist.is_initialized():
+            dist.destroy_process_group()
+        print(f"[Rank {rank}] Process group destroyed.")
+
+
+def run_test(args: argparse.Namespace, server_args: ServerArgs) -> bool:
+    """Run the mooncake transfer engine test."""
+    # Set CUDA visible devices
+    cuda_devices = args.cuda_visible_devices.split(",")
+    world_size = len(cuda_devices)
+
+    if world_size < 2:
+        print("ERROR: This test requires at least 2 GPUs.")
+        print(
+            "Usage: CUDA_VISIBLE_DEVICES=0,1 python test/manual/kv_transfer/test_mooncake_transfer_engine_init.py"
+        )
+        sys.exit(1)
+
+    # Check GPU availability
+    import torch
+
+    if not torch.cuda.is_available():
+        print("ERROR: CUDA is not available")
+        sys.exit(1)
+
+    available_gpus = torch.cuda.device_count()
+    if world_size > available_gpus:
+        print(f"ERROR: Requested {world_size} GPUs but only {available_gpus} available")
+        sys.exit(1)
+
+    print(f"Testing with {world_size} GPUs: {cuda_devices}")
+    print()
+
+    # Print server args configuration
+    print("ServerArgs configuration:")
+    for key, value in vars(server_args).items():
+        print(f"  {key}: {value}")
+    print()
+
+    # Check if mooncake should be used
+    use_mooncake_te = test_mooncake_te_condition(server_args)
+    print(f"use_mooncake_te = {use_mooncake_te}")
+    print()
+
+    # Find a free port
+    import socket
+
+    with socket.socket() as s:
+        s.bind(("", 0))
+        master_port = s.getsockname()[1]
+
+    print(f"Using master port: {master_port}")
+
+    # Spawn worker processes
+    ctx = multiprocessing.get_context("spawn")
+    processes = []
+
+    for rank in range(world_size):
+        p = ctx.Process(
+            target=run_mooncake_init,
+            args=(rank, world_size, master_port, args, server_args),
+        )
+        p.start()
+        processes.append(p)
+
+    # Wait for all processes to complete
+    success = True
+    for i, p in enumerate(processes):
+        p.join(timeout=60)
+        if p.exitcode != 0:
+            print(f"Process {i} failed with exit code: {p.exitcode}")
+            success = False
+
+    # Cleanup any remaining processes
+    for p in processes:
+        if p.is_alive():
+            print(f"Process {p.pid} is still alive, terminating...")
+            p.terminate()
+            p.join(timeout=5)
+
+    return success
+
+
+def test_condition_logic():
+    """Test the condition logic for different server argument combinations."""
+    print("=" * 60)
+    print("Testing condition logic for use_mooncake_te")
+    print("=" * 60)
+    print()
+
+    original_hicache_reuse = os.environ.get("SGLANG_HICACHE_MOONCAKE_REUSE_TE")
+    passed = 0
+    failed = 0
+
+    try:
+        test_cases = [
+            # (name, env_value, server_args, expected_result)
+            (
+                "PD disaggregation with mooncake",
+                None,
+                ServerArgs(
+                    disaggregation_mode="prefill",
+                    disaggregation_transfer_backend="mooncake",
+                ),
+                True,
+            ),
+            (
+                "PD disaggregation without mooncake",
+                None,
+                ServerArgs(
+                    disaggregation_mode="prefill",
+                    disaggregation_transfer_backend="other",
+                ),
+                False,
+            ),
+            (
+                "No disaggregation",
+                None,
+                ServerArgs(),
+                False,
+            ),
+            (
+                "HiCache with mooncake (env=False)",
+                "0",
+                ServerArgs(
+                    enable_hierarchical_cache=True,
+                    hicache_storage_backend="mooncake",
+                ),
+                False,
+            ),
+            (
+                "HiCache with mooncake (env=True)",
+                "1",
+                ServerArgs(
+                    enable_hierarchical_cache=True,
+                    hicache_storage_backend="mooncake",
+                ),
+                True,
+            ),
+            (
+                "Encoder only with mooncake",
+                None,
+                ServerArgs(encoder_only=True, encoder_transfer_backend="mooncake"),
+                True,
+            ),
+            (
+                "Language only with mooncake",
+                None,
+                ServerArgs(language_only=True, encoder_transfer_backend="mooncake"),
+                True,
+            ),
+            (
+                "Elastic expert backup with backend",
+                None,
+                ServerArgs(
+                    enable_elastic_expert_backup=True,
+                    elastic_ep_backend="mooncake",
+                ),
+                True,
+            ),
+            (
+                "Elastic expert backup without backend",
+                None,
+                ServerArgs(enable_elastic_expert_backup=True, elastic_ep_backend=None),
+                False,
+            ),
+        ]
+
+        for name, env_value, server_args, expected in test_cases:
+            if env_value is None:
+                os.environ.pop("SGLANG_HICACHE_MOONCAKE_REUSE_TE", None)
+            else:
+                os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = env_value
+
+            result = test_mooncake_te_condition(server_args)
+            status = "PASS" if result == expected else "FAIL"
+
+            if result == expected:
+                passed += 1
+            else:
+                failed += 1
+
+            print(f"{status}: {name}")
+            print(f"       Expected: {expected}, Got: {result}")
+            print()
+    finally:
+        if original_hicache_reuse is None:
+            os.environ.pop("SGLANG_HICACHE_MOONCAKE_REUSE_TE", None)
+        else:
+            os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = original_hicache_reuse
+
+    print(f"Condition logic tests: {passed} passed, {failed} failed")
+    print()
+
+    return failed == 0
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Validate Mooncake transfer-engine gating and initialization"
+    )
+    parser.add_argument(
+        "--cuda-visible-devices",
+        type=str,
+        default="0,1",
+        help="CUDA visible devices (default: 0,1)",
+    )
+    parser.add_argument(
+        "--test-case",
+        type=str,
+        choices=[
+            "pd_disaggregation",
+            "hicache",
+            "encoder_only",
+            "language_only",
+            "elastic_ep",
+        ],
+        default="pd_disaggregation",
+        help="Test case to run",
+    )
+
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("Mooncake Transfer Engine Init Test")
+    print("=" * 60)
+    print()
+
+    # First run condition logic tests
+    condition_passed = test_condition_logic()
+
+    if not condition_passed:
+        print("Condition logic tests failed, skipping distributed test.")
+        sys.exit(1)
+
+    # Configure server args based on test case
+    server_args = ServerArgs()
+
+    if args.test_case == "pd_disaggregation":
+        server_args.disaggregation_mode = "prefill"
+        server_args.disaggregation_transfer_backend = "mooncake"
+    elif args.test_case == "hicache":
+        server_args.enable_hierarchical_cache = True
+        server_args.hicache_storage_backend = "mooncake"
+        os.environ["SGLANG_HICACHE_MOONCAKE_REUSE_TE"] = "1"
+    elif args.test_case == "encoder_only":
+        server_args.encoder_only = True
+        server_args.encoder_transfer_backend = "mooncake"
+    elif args.test_case == "language_only":
+        server_args.language_only = True
+        server_args.encoder_transfer_backend = "mooncake"
+    elif args.test_case == "elastic_ep":
+        server_args.enable_elastic_expert_backup = True
+        server_args.elastic_ep_backend = "mooncake"
+
+    start_time = time.time()
+    success = run_test(args, server_args)
+    elapsed_time = time.time() - start_time
+
+    print()
+    print("=" * 60)
+    if success:
+        print(f"TEST PASSED (elapsed: {elapsed_time:.2f}s)")
+    else:
+        print(f"TEST FAILED (elapsed: {elapsed_time:.2f}s)")
+    print("=" * 60)
+
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/test/manual/lang_frontend/test_bind_cache.py

@@ -0,0 +1,51 @@
+import unittest
+
+import sglang as sgl
+from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
+
+
+class TestBind(CustomTestCase):
+    backend = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.backend = sgl.Runtime(model_path=DEFAULT_MODEL_NAME_FOR_TEST)
+        sgl.set_default_backend(cls.backend)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.backend.shutdown()
+
+    def test_bind(self):
+        @sgl.function
+        def few_shot_qa(s, prompt, question):
+            s += prompt
+            s += "Q: What is the capital of France?\n"
+            s += "A: Paris\n"
+            s += "Q: " + question + "\n"
+            s += "A:" + sgl.gen("answer", stop="\n")
+
+        few_shot_qa_2 = few_shot_qa.bind(
+            prompt="The following are questions with answers.\n\n"
+        )
+
+        tracer = few_shot_qa_2.trace()
+        print(tracer.last_node.print_graph_dfs() + "\n")
+
+    def test_cache(self):
+        @sgl.function
+        def few_shot_qa(s, prompt, question):
+            s += prompt
+            s += "Q: What is the capital of France?\n"
+            s += "A: Paris\n"
+            s += "Q: " + question + "\n"
+            s += "A:" + sgl.gen("answer", stop="\n")
+
+        few_shot_qa_2 = few_shot_qa.bind(
+            prompt="Answer the following questions as if you were a 5-year-old kid.\n\n"
+        )
+        few_shot_qa_2.cache(self.backend)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lang_frontend/test_choices.py

@@ -0,0 +1,91 @@
+import unittest
+
+import numpy as np
+
+from sglang.lang.choices import (
+    greedy_token_selection,
+    token_length_normalized,
+    unconditional_likelihood_normalized,
+)
+from sglang.test.test_utils import CustomTestCase
+
+MOCK_CHOICES_INPUT_DATA = {
+    "choices": [
+        "organ",  # ["organ"]
+        "organism",  # ["organ", "ism"]
+        "antidisestablishmentarianism",  # ["ant", "id", "is", "est", "ablish", "ment", "arian", "ism"]
+    ],
+    "normalized_prompt_logprobs": [-0.1, -0.2, -0.05],
+    "input_token_logprobs": [
+        [[-0.1, 1, None]],
+        [[-0.1, 1, None], [-0.3, 2, None]],
+        [
+            [-0.4, 3, None],
+            [-0.25, 4, None],
+            [-0.1, 5, None],
+            [-0.01, 6, None],
+            [-0.01, 7, None],
+            [-0.01, 8, None],
+            [-0.01, 9, None],
+            [-0.01, 2, None],
+        ],
+    ],
+    "output_token_logprobs": [
+        [[-0.1, 10, None]],
+        [[-0.1, 10, None]],
+        [[-0.1, 10, None]],
+    ],
+    "unconditional_token_logprobs": [
+        [[None, 1, None]],
+        [[None, 1, None], [-1.4, 2, None]],
+        [
+            [None, 3, None],
+            [-0.25, 4, None],
+            [-0.1, 5, None],
+            [-0.01, 6, None],
+            [-0.01, 7, None],
+            [-0.01, 8, None],
+            [-0.01, 9, None],
+            [-0.01, 2, None],
+        ],
+    ],
+}
+
+
+class TestChoices(CustomTestCase):
+
+    def test_token_length_normalized(self):
+        """Confirm 'antidisestablishmentarianism' is selected due to high confidences for
+        its later tokens resulting in highest token length normalized prompt logprob."""
+        decision = token_length_normalized(**MOCK_CHOICES_INPUT_DATA)
+        assert decision.decision == "antidisestablishmentarianism"
+
+    def test_greedy_token_selection(self):
+        """Confirm 'organ' is selected due it having the joint highest initial token
+        logprob, and a higher average logprob than organism's second token."""
+        decision = greedy_token_selection(**MOCK_CHOICES_INPUT_DATA)
+        assert decision.decision == "organ"
+        assert np.allclose(
+            decision.meta_info["greedy_logprob_matrix"],
+            [
+                [-0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1],
+                [-0.1, -0.3, -0.2, -0.2, -0.2, -0.2, -0.2, -0.2],
+                [-0.4, -0.25, -0.1, -0.01, -0.01, -0.01, -0.01, -0.01],
+            ],
+            atol=0.01,
+        )
+
+    def test_unconditional_likelihood_normalized(self):
+        """Confirm 'organism' is selected due to it having the highest average token logprob
+        once normalized by the unconditional token logprobs."""
+        decision = unconditional_likelihood_normalized(**MOCK_CHOICES_INPUT_DATA)
+        assert decision.decision == "organism"
+        assert np.allclose(
+            decision.meta_info["normalized_unconditional_prompt_logprobs"],
+            [-0.1, 0.5, -0.05],
+            atol=0.01,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lang_frontend/test_jump_forward.py

@@ -0,0 +1,138 @@
+import argparse
+from enum import Enum
+
+from pydantic import BaseModel, constr
+
+import sglang as sgl
+from sglang.srt.constrained.outlines_backend import build_regex_from_object
+from sglang.test.test_utils import (
+    add_common_sglang_args_and_parse,
+    select_sglang_backend,
+)
+
+IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
+
+ip_jump_forward = (
+    r"The google's DNS sever address is "
+    + IP_REGEX
+    + r" and "
+    + IP_REGEX
+    + r". "
+    + r"The google's website domain name is "
+    + r"www\.(\w)+\.(\w)+"
+    + r"."
+)
+
+
+# fmt: off
+@sgl.function
+def regex_gen(s):
+    s += "Q: What is the IP address of the Google DNS servers?\n"
+    s += "A: " + sgl.gen(
+        "answer",
+        max_tokens=128,
+        temperature=0,
+        regex=ip_jump_forward,
+    )
+# fmt: on
+
+json_jump_forward = (
+    r"""The information about Hogwarts is in the following JSON format\.\n"""
+    + r"""\n\{\n"""
+    + r"""  "name": "[\w\d\s]*",\n"""
+    + r"""  "country": "[\w\d\s]*",\n"""
+    + r"""  "latitude": [-+]?[0-9]*\.?[0-9]+,\n"""
+    + r"""  "population": [-+]?[0-9]+,\n"""
+    + r"""  "top 3 landmarks": \["[\w\d\s]*", "[\w\d\s]*", "[\w\d\s]*"\],\n"""
+    + r"""\}\n"""
+)
+
+# fmt: off
+@sgl.function
+def json_gen(s):
+    s += sgl.gen(
+        "json",
+        max_tokens=128,
+        temperature=0,
+        regex=json_jump_forward,
+    )
+# fmt: on
+
+
+class Weapon(str, Enum):
+    sword = "sword"
+    axe = "axe"
+    mace = "mace"
+    spear = "spear"
+    bow = "bow"
+    crossbow = "crossbow"
+
+
+class Armor(str, Enum):
+    leather = "leather"
+    chainmail = "chainmail"
+    plate = "plate"
+
+
+class Character(BaseModel):
+    name: constr(max_length=10)
+    age: int
+    armor: Armor
+    weapon: Weapon
+    strength: int
+
+
+@sgl.function
+def character_gen(s):
+    s += "Give me a character description who is a wizard.\n"
+    s += sgl.gen(
+        "character",
+        max_tokens=128,
+        temperature=0,
+        regex=build_regex_from_object(Character),
+    )
+
+
+def main(args):
+    # Select backend
+    backend = select_sglang_backend(args)
+    sgl.set_default_backend(backend)
+
+    state = regex_gen.run(temperature=0)
+
+    print("=" * 20, "IP TEST", "=" * 20)
+    print(state.text())
+
+    state = json_gen.run(temperature=0)
+
+    print("=" * 20, "JSON TEST", "=" * 20)
+    print(state.text())
+
+    state = character_gen.run(temperature=0)
+
+    print("=" * 20, "CHARACTER TEST", "=" * 20)
+    print(state.text())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    args = add_common_sglang_args_and_parse(parser)
+    main(args)
+
+# ==================== IP TEST ====================
+# Q: What is the IP address of the Google DNS servers?
+# A: The google's DNS sever address is 8.8.8.8 and 8.8.4.4. The google's website domain name is www.google.com.
+# ==================== JSON TEST ====================
+# The information about Hogwarts is in the following JSON format.
+
+# {
+#   "name": "Hogwarts School of Witchcraft and Wizardry",
+#   "country": "Scotland",
+#   "latitude": 55.566667,
+#   "population": 1000,
+#   "top 3 landmarks": ["Hogwarts Castle", "The Great Hall", "The Forbidden Forest"],
+# }
+
+# ==================== CHARACTER TEST ====================
+# Give me a character description who is a wizard.
+# { "name" : "Merlin", "age" : 500, "armor" : "chainmail" , "weapon" : "sword" , "strength" : 10 }

third_party/sglang/test/manual/lang_frontend/test_openai_backend.py

@@ -0,0 +1,92 @@
+import unittest
+
+from sglang import OpenAI, set_default_backend
+from sglang.test.test_programs import (
+    test_chat_completion_speculative,
+    test_completion_speculative,
+    test_decode_int,
+    test_decode_json,
+    test_expert_answer,
+    test_few_shot_qa,
+    test_image_qa,
+    test_mt_bench,
+    test_parallel_decoding,
+    test_parallel_encoding,
+    test_react,
+    test_select,
+    test_stream,
+    test_tool_use,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestOpenAIBackend(CustomTestCase):
+    instruct_backend = None
+    chat_backend = None
+    chat_vision_backend = None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.instruct_backend = OpenAI("gpt-3.5-turbo-instruct")
+        cls.chat_backend = OpenAI("gpt-3.5-turbo")
+        cls.chat_vision_backend = OpenAI("gpt-4-turbo")
+
+    def test_few_shot_qa(self):
+        set_default_backend(self.instruct_backend)
+        test_few_shot_qa()
+
+    def test_mt_bench(self):
+        set_default_backend(self.chat_backend)
+        test_mt_bench()
+
+    def test_select(self):
+        set_default_backend(self.instruct_backend)
+        test_select(check_answer=True)
+
+    def test_decode_int(self):
+        set_default_backend(self.instruct_backend)
+        test_decode_int()
+
+    def test_decode_json(self):
+        set_default_backend(self.instruct_backend)
+        test_decode_json()
+
+    def test_expert_answer(self):
+        set_default_backend(self.instruct_backend)
+        test_expert_answer()
+
+    def test_tool_use(self):
+        set_default_backend(self.instruct_backend)
+        test_tool_use()
+
+    def test_react(self):
+        set_default_backend(self.instruct_backend)
+        test_react()
+
+    def test_parallel_decoding(self):
+        set_default_backend(self.instruct_backend)
+        test_parallel_decoding()
+
+    def test_parallel_encoding(self):
+        set_default_backend(self.instruct_backend)
+        test_parallel_encoding()
+
+    def test_image_qa(self):
+        set_default_backend(self.chat_vision_backend)
+        test_image_qa()
+
+    def test_stream(self):
+        set_default_backend(self.instruct_backend)
+        test_stream()
+
+    def test_completion_speculative(self):
+        set_default_backend(self.instruct_backend)
+        test_completion_speculative()
+
+    def test_chat_completion_speculative(self):
+        set_default_backend(self.chat_backend)
+        test_chat_completion_speculative()
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lang_frontend/test_separate_reasoning.py

@@ -0,0 +1,68 @@
+"""
+Tests for the separate_reasoning functionality in sglang.
+
+Usage:
+python3 -m unittest test/lang/test_separate_reasoning.py
+"""
+
+import unittest
+
+from sglang import gen, separate_reasoning
+from sglang.lang.ir import SglExprList, SglSeparateReasoning
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestSeparateReasoning(CustomTestCase):
+    def test_separate_reasoning_creation(self):
+        """Test that SglSeparateReasoning objects are created correctly."""
+        # Test with valid model type and gen expression
+        test_gen = gen("test")
+        expr = separate_reasoning(test_gen, model_type="deepseek-r1")
+        self.assertIsInstance(expr, SglExprList)
+        self.assertEqual(len(expr.expr_list), 2)
+        self.assertEqual(expr.expr_list[0], test_gen)
+        reasoning_expr = expr.expr_list[1]
+        self.assertIsInstance(reasoning_expr, SglSeparateReasoning)
+        self.assertEqual(reasoning_expr.model_type, "deepseek-r1")
+        self.assertEqual(reasoning_expr.name, "test_reasoning_content")
+
+        # Test with another valid model type
+        expr = separate_reasoning(test_gen, model_type="qwen3")
+        self.assertIsInstance(expr, SglExprList)
+        self.assertEqual(expr.expr_list[1].model_type, "qwen3")
+
+    def test_separate_reasoning_name_processing(self):
+        """Test that separate_reasoning correctly processes names."""
+        test_gen = gen("test_var")
+        expr = separate_reasoning(test_gen, model_type="deepseek-r1")
+        reasoning_expr = expr.expr_list[1]
+        self.assertEqual(reasoning_expr.name, "test_var_reasoning_content")
+
+        # Test the process_name_for_reasoning method
+        self.assertEqual(
+            reasoning_expr.process_name_for_reasoning("another_var"),
+            "another_var_reasoning_content",
+        )
+
+    def test_separate_reasoning_repr(self):
+        """Test the string representation of SglSeparateReasoning."""
+        test_gen = gen("test_var")
+        expr = separate_reasoning(test_gen, model_type="deepseek-r1")
+        reasoning_expr = expr.expr_list[1]
+        self.assertEqual(
+            repr(reasoning_expr),
+            "SeparateReasoning(model_type=deepseek-r1, name=test_var_reasoning_content)",
+        )
+
+    def test_separate_reasoning_with_invalid_model_type(self):
+        """Test that separate_reasoning accepts any model type during creation."""
+        # Create with invalid model type
+        test_gen = gen("test")
+        expr = separate_reasoning(test_gen, model_type="invalid-model")
+        self.assertIsInstance(expr, SglExprList)
+        self.assertEqual(expr.expr_list[1].model_type, "invalid-model")
+        # The actual validation happens in the ReasoningParser constructor
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lang_frontend/test_separate_reasoning_execution.py

@@ -0,0 +1,193 @@
+"""
+Tests for the execution of separate_reasoning functionality in sglang.
+
+Usage:
+python3 -m unittest test/lang/test_separate_reasoning_execution.py
+"""
+
+import threading
+import unittest
+from unittest.mock import MagicMock, patch
+
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SglGen, SglSeparateReasoning
+from sglang.test.test_utils import CustomTestCase
+
+
+# Helper function to create events that won't block program exit
+def create_daemon_event():
+    event = threading.Event()
+    return event
+
+
+class MockReasoningParser:
+    def __init__(self, model_type):
+        self.model_type = model_type
+        self.parse_non_stream_called = False
+        self.parse_stream_chunk_called = False
+
+    def parse_non_stream(self, full_text):
+        self.parse_non_stream_called = True
+        # Simulate parsing by adding a prefix to indicate reasoning
+        reasoning = f"[REASONING from {self.model_type}]: {full_text}"
+        normal_text = f"[NORMAL from {self.model_type}]: {full_text}"
+        return reasoning, normal_text
+
+    def parse_stream_chunk(self, chunk_text):
+        self.parse_stream_chunk_called = True
+        # Simulate parsing by adding a prefix to indicate reasoning
+        reasoning = f"[REASONING from {self.model_type}]: {chunk_text}"
+        normal_text = f"[NORMAL from {self.model_type}]: {chunk_text}"
+        return reasoning, normal_text
+
+
+class TestSeparateReasoningExecution(CustomTestCase):
+    def setUp(self):
+        """Set up for the test."""
+        super().setUp()
+        # Store any events created during the test
+        self.events = []
+
+    def tearDown(self):
+        """Clean up any threads that might have been created during the test."""
+        super().tearDown()
+
+        # Set all events to ensure any waiting threads are released
+        for event in self.events:
+            event.set()
+
+    def tearDown(self):
+        super().tearDown()
+        # wake up all threads
+        for ev in self.events:
+            ev.set()
+
+    @patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
+    def test_execute_separate_reasoning(self, mock_parser_class):
+        """Test that _execute_separate_reasoning correctly calls the ReasoningParser."""
+        # Setup mock parser
+        mock_parser = MockReasoningParser("deepseek-r1")
+        mock_parser_class.return_value = mock_parser
+
+        # Create a mock backend to avoid AttributeError in __del__
+        mock_backend = MagicMock()
+
+        # Create a StreamExecutor with necessary setup
+        executor = StreamExecutor(
+            backend=mock_backend,
+            arguments={},
+            default_sampling_para={},
+            chat_template={
+                "role_map": {"user": "user", "assistant": "assistant"}
+            },  # Simple chat template
+            stream=False,
+            use_thread=False,
+        )
+
+        # Set up the executor with a variable and its value
+        var_name = "test_var"
+        reasoning_name = f"{var_name}_reasoning_content"
+        var_value = "Test content"
+        executor.variables = {var_name: var_value}
+
+        # Create events and track them for cleanup
+        var_event = create_daemon_event()
+        reasoning_event = create_daemon_event()
+        self.events.extend([var_event, reasoning_event])
+
+        executor.variable_event = {var_name: var_event, reasoning_name: reasoning_event}
+        executor.variable_event[var_name].set()  # Mark as ready
+
+        # Set up the current role
+        executor.cur_role = "assistant"
+        executor.cur_role_begin_pos = 0
+        executor.text_ = var_value
+
+        # Create a gen expression and a separate_reasoning expression
+        gen_expr = SglGen(var_name)
+        expr = SglSeparateReasoning("deepseek-r1", expr=gen_expr)
+
+        # Execute separate_reasoning
+        executor._execute_separate_reasoning(expr)
+
+        # Verify that the parser was created with the correct model type
+        mock_parser_class.assert_called_once_with("deepseek-r1")
+
+        # Verify that parse_non_stream was called
+        self.assertTrue(mock_parser.parse_non_stream_called)
+
+        # Verify that the variables were updated correctly
+        reasoning_name = f"{var_name}_reasoning_content"
+        self.assertIn(reasoning_name, executor.variables)
+        self.assertEqual(
+            executor.variables[reasoning_name],
+            f"[REASONING from deepseek-r1]: {var_value}",
+        )
+        self.assertEqual(
+            executor.variables[var_name], f"[NORMAL from deepseek-r1]: {var_value}"
+        )
+
+        # Verify that the variable event was set
+        self.assertIn(reasoning_name, executor.variable_event)
+        self.assertTrue(executor.variable_event[reasoning_name].is_set())
+
+        # Verify that the text was updated
+        self.assertEqual(executor.text_, f"[NORMAL from deepseek-r1]: {var_value}")
+
+    @patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
+    def test_reasoning_parser_integration(self, mock_parser_class):
+        """Test the integration between separate_reasoning and ReasoningParser."""
+        # Setup mock parsers for different model types
+        deepseek_parser = MockReasoningParser("deepseek-r1")
+        qwen_parser = MockReasoningParser("qwen3")
+
+        # Configure the mock to return different parsers based on model type
+        def get_parser(model_type):
+            if model_type == "deepseek-r1":
+                return deepseek_parser
+            elif model_type == "qwen3":
+                return qwen_parser
+            else:
+                raise ValueError(f"Unsupported model type: {model_type}")
+
+        mock_parser_class.side_effect = get_parser
+
+        # Test with DeepSeek-R1 model
+        test_text = "This is a test"
+        reasoning, normal_text = deepseek_parser.parse_non_stream(test_text)
+
+        self.assertEqual(reasoning, f"[REASONING from deepseek-r1]: {test_text}")
+        self.assertEqual(normal_text, f"[NORMAL from deepseek-r1]: {test_text}")
+
+        # Test with Qwen3 model
+        reasoning, normal_text = qwen_parser.parse_non_stream(test_text)
+
+        self.assertEqual(reasoning, f"[REASONING from qwen3]: {test_text}")
+        self.assertEqual(normal_text, f"[NORMAL from qwen3]: {test_text}")
+
+    @patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
+    def test_reasoning_parser_invalid_model(self, mock_parser_class):
+        """Test that ReasoningParser raises an error for invalid model types."""
+
+        # Configure the mock to raise an error for invalid model types
+        def get_parser(model_type):
+            if model_type in ["deepseek-r1", "qwen3"]:
+                return MockReasoningParser(model_type)
+            elif model_type is None:
+                raise ValueError("Model type must be specified")
+            else:
+                raise ValueError(f"Unsupported model type: {model_type}")
+
+        mock_parser_class.side_effect = get_parser
+
+        with self.assertRaises(ValueError) as context:
+            mock_parser_class("invalid-model")
+        self.assertIn("Unsupported model type", str(context.exception))
+
+        with self.assertRaises(ValueError) as context:
+            mock_parser_class(None)
+        self.assertIn("Model type must be specified", str(context.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/layers/attention/nsa/test_act_quant_triton.py

@@ -0,0 +1,281 @@
+"""
+Unit tests comparing TileLang and Triton implementations of activation quantization.
+Tests both accuracy and performance.
+"""
+
+import time
+from typing import Tuple
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant
+from sglang.srt.layers.attention.nsa.triton_kernel import act_quant as act_quant_triton
+
+
+def benchmark_kernel(
+    fn,
+    x: torch.Tensor,
+    block_size: int,
+    scale_fmt,
+    warmup: int = 10,
+    repeat: int = 100,
+    use_cuda_graph: bool = True,
+) -> Tuple[float, torch.Tensor, torch.Tensor]:
+    """
+    Benchmark a kernel function.
+
+    Args:
+        fn: Function to benchmark
+        x: Input tensor
+        block_size: Block size for quantization
+        scale_fmt: Scale format
+        warmup: Number of warmup iterations
+        repeat: Number of repeat iterations
+        use_cuda_graph: Whether to use CUDA graphs for more accurate timing
+
+    Returns:
+        Tuple of (avg_time_ms, quantized_output, scales)
+    """
+    # Warmup
+    for _ in range(warmup):
+        y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+    if not x.is_cuda or not use_cuda_graph:
+        # Fallback to regular timing
+        if x.is_cuda:
+            torch.cuda.synchronize()
+
+        start = time.perf_counter()
+        for _ in range(repeat):
+            y, s = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+        if x.is_cuda:
+            torch.cuda.synchronize()
+
+        end = time.perf_counter()
+        avg_time_ms = (end - start) / repeat * 1000
+
+        return avg_time_ms, y, s
+
+    # Use CUDA graph for more accurate timing
+    torch.cuda.synchronize()
+
+    # Allocate output buffers
+    N = x.size(-1)
+    y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
+    s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
+
+    # Capture CUDA graph
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        y_cap, s_cap = fn(x, block_size=block_size, scale_fmt=scale_fmt)
+
+    # Warmup with graph
+    for _ in range(warmup):
+        graph.replay()
+
+    torch.cuda.synchronize()
+
+    # Timing with CUDA graph
+    start_event = torch.cuda.Event(enable_timing=True)
+    end_event = torch.cuda.Event(enable_timing=True)
+
+    start_event.record()
+    for _ in range(repeat):
+        graph.replay()
+    end_event.record()
+
+    torch.cuda.synchronize()
+
+    avg_time_ms = start_event.elapsed_time(end_event) / repeat
+
+    return avg_time_ms, y_cap, s_cap
+
+
+def check_accuracy(
+    y_ref: torch.Tensor,
+    s_ref: torch.Tensor,
+    y_test: torch.Tensor,
+    s_test: torch.Tensor,
+    rtol: float = 1e-2,
+    atol: float = 1e-2,
+) -> Tuple[bool, dict]:
+    """
+    Check accuracy between reference and test outputs.
+
+    Args:
+        y_ref: Reference quantized output
+        s_ref: Reference scales
+        y_test: Test quantized output
+        s_test: Test scales
+        rtol: Relative tolerance
+        atol: Absolute tolerance
+
+    Returns:
+        Tuple of (passed, metrics_dict)
+    """
+    # Convert FP8 to float for comparison
+    y_ref_float = y_ref.float()
+    y_test_float = y_test.float()
+
+    # Compute differences
+    y_diff = torch.abs(y_ref_float - y_test_float)
+    s_diff = torch.abs(s_ref - s_test)
+
+    # Compute metrics
+    y_max_diff = y_diff.max().item()
+    y_mean_diff = y_diff.mean().item()
+    s_max_diff = s_diff.max().item()
+    s_mean_diff = s_diff.mean().item()
+
+    # Check relative and absolute tolerance
+    y_close = torch.allclose(y_ref_float, y_test_float, rtol=rtol, atol=atol)
+    s_close = torch.allclose(s_ref, s_test, rtol=rtol, atol=atol)
+
+    # Compute percentage of matching elements
+    y_match_pct = (y_ref_float == y_test_float).float().mean().item() * 100
+
+    metrics = {
+        "y_max_diff": y_max_diff,
+        "y_mean_diff": y_mean_diff,
+        "y_match_pct": y_match_pct,
+        "s_max_diff": s_max_diff,
+        "s_mean_diff": s_mean_diff,
+        "y_close": y_close,
+        "s_close": s_close,
+    }
+
+    passed = y_close and s_close
+
+    return passed, metrics
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_act_quant_comprehensive_benchmark(scale_fmt=None):
+    """Comprehensive benchmark across multiple sizes with CUDA graphs."""
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    block_size = 128
+
+    shapes = [
+        (128, 512),
+        (256, 1024),
+        (512, 2048),
+        (1024, 4096),
+        (2048, 8192),
+        (4096, 16384),
+    ]
+
+    print("\n" + "=" * 100)
+    print("Comprehensive Performance Benchmark with CUDA Graphs")
+    print("=" * 100)
+    print(
+        f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
+    )
+    print("-" * 100)
+
+    for shape in shapes:
+        torch.manual_seed(42)
+        x = torch.randn(shape, dtype=dtype, device=device)
+
+        try:
+            # Benchmark both with CUDA graphs
+            time_tilelang, y_ref, s_ref = benchmark_kernel(
+                act_quant,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=True,
+            )
+            time_triton, y_triton, s_triton = benchmark_kernel(
+                act_quant_triton,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=True,
+            )
+
+            # Check accuracy
+            passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
+
+            speedup = time_tilelang / time_triton if time_triton > 0 else 0
+            status = "✓ PASS" if passed else "✗ FAIL"
+
+            print(
+                f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
+                f"{speedup:<10.2f} {status}"
+            )
+        except Exception as e:
+            print(f"{str(shape):<20} ERROR: {str(e)}")
+
+    print("=" * 100)
+
+    # Also run without CUDA graphs for comparison
+    print("\n" + "=" * 100)
+    print("Performance Benchmark WITHOUT CUDA Graphs (for comparison)")
+    print("=" * 100)
+    print(
+        f"{'Shape':<20} {'TileLang (ms)':<15} {'Triton (ms)':<15} {'Speedup':<10} {'Status'}"
+    )
+    print("-" * 100)
+
+    for shape in shapes:
+        torch.manual_seed(42)
+        x = torch.randn(shape, dtype=dtype, device=device)
+
+        try:
+            # Benchmark both without CUDA graphs
+            time_tilelang, y_ref, s_ref = benchmark_kernel(
+                act_quant,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=False,
+            )
+            time_triton, y_triton, s_triton = benchmark_kernel(
+                act_quant_triton,
+                x,
+                block_size,
+                scale_fmt,
+                warmup=5,
+                repeat=50,
+                use_cuda_graph=False,
+            )
+
+            # Check accuracy
+            passed, _ = check_accuracy(y_ref, s_ref, y_triton, s_triton)
+
+            speedup = time_tilelang / time_triton if time_triton > 0 else 0
+            status = "✓ PASS" if passed else "✗ FAIL"
+
+            print(
+                f"{str(shape):<20} {time_tilelang:<15.4f} {time_triton:<15.4f} "
+                f"{speedup:<10.2f} {status}"
+            )
+        except Exception as e:
+            print(f"{str(shape):<20} ERROR: {str(e)}")
+
+    print("=" * 100)
+
+
+if __name__ == "__main__":
+    # Run comprehensive benchmark
+    if torch.cuda.is_available():
+        print("\n" + "=" * 80)
+        print("Running Comprehensive Benchmark with scale_fmt=None")
+        print("=" * 80)
+        test_act_quant_comprehensive_benchmark(scale_fmt=None)
+
+        print("\n" + "=" * 80)
+        print("Running Comprehensive Benchmark with scale_fmt!=None")
+        print("=" * 80)
+        test_act_quant_comprehensive_benchmark(scale_fmt="any")
+    else:
+        print("CUDA not available. Skipping tests.")

third_party/sglang/test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py

@@ -0,0 +1,191 @@
+import torch
+
+from sglang.srt.layers.attention.nsa.index_buf_accessor import (
+    _get_k_and_s_triton_kernel,
+)
+
+
+def golden_torch_gen(
+    seq_len_tensor: torch.Tensor,
+    buffer_indexer: torch.Tensor,
+    buffer: torch.Tensor,
+    index_head_dim,
+    page_size,
+):
+    dim_split = page_size * index_head_dim
+    torch_k_out = buffer[:, 0:dim_split]
+    torch_s_out = buffer[:, dim_split:]
+
+    torch_k_out = torch_k_out.reshape(-1, 128)
+    torch_s_out = torch_s_out.reshape(-1, 4)
+
+    batch = seq_len_tensor.shape[0]
+    index_list = []
+    for i in range(batch):
+        seq_len = seq_len_tensor[i].item()
+        buffer_index_ = buffer_indexer[i]
+        align_seq_len = ((seq_len + page_size - 1) / page_size) * page_size
+        needed_block_num = int((seq_len + page_size - 1) / page_size)
+        for j in range(needed_block_num):
+            block_idx = buffer_index_[j].item()
+            start_idx = block_idx * page_size
+            end_idx = 0
+            if j == (needed_block_num - 1):
+                end_idx = block_idx * page_size + (
+                    seq_len - (needed_block_num - 1) * page_size
+                )
+            else:
+                end_idx = (block_idx + 1) * page_size
+
+            index_tensor = (
+                torch.arange(start=start_idx, end=end_idx, step=1)
+                .type(torch.int32)
+                .cuda()
+            )
+            index_list.append(index_tensor)
+
+    index_list_ = torch.cat(index_list, dim=0)
+    torch_k_out = torch.index_select(torch_k_out, dim=0, index=index_list_)
+    torch_s_out = torch.index_select(torch_s_out, dim=0, index=index_list_)
+
+    return torch_k_out, torch_s_out
+
+
+def get_k_and_s_triton():
+    index_head_dim = 128
+    page_size = 64
+    num_page = 128
+    s_offset_in_page = page_size * index_head_dim
+
+    seq_len_tensor = torch.tensor(
+        [256, 267, 215, 32, 129], dtype=torch.int64, device="cuda"
+    )  # 4 + 5 + 3 + 1 + 3 block
+    buffer_indexer = torch.tensor(
+        [
+            [1, 2, 3, 4, 0],
+            [7, 6, 5, 8, 9],
+            [10, 11, 12, 0, 0],
+            [13, 0, 0, 0, 0],
+            [14, 15, 16, 0, 0],
+        ],
+        dtype=torch.int32,
+        device="cuda",
+    )
+    seq_len_sum = seq_len_tensor.sum()
+    batch = seq_len_tensor.shape[0]
+
+    triton_k_out = torch.empty(
+        (seq_len_sum, index_head_dim), dtype=torch.uint8, device="cuda"
+    )
+    triton_s_out = torch.empty((seq_len_sum, 4), dtype=torch.uint8, device="cuda")
+    buffer = torch.randint(
+        0,
+        num_page,
+        (num_page, page_size * index_head_dim + page_size * 4),
+        device="cuda",
+    ).type(torch.uint8)
+
+    _, buf_numel_per_page = buffer.shape
+    _, page_indice_batch_offset = buffer_indexer.shape
+    max_seq_len = seq_len_tensor.max().item()
+
+    BLOCK_SIZE = 256
+    BLOCK_SIZE_K = 128
+
+    num_token_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_k_threads = (index_head_dim + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K
+
+    grid = (batch, num_token_blocks, num_k_threads)
+    seq_num_pow2 = 1
+    while seq_num_pow2 < batch:
+        seq_num_pow2 *= 2
+
+    # acc test =====================
+    _get_k_and_s_triton_kernel[grid](
+        buf_ptr=buffer,
+        page_indices_ptr=buffer_indexer,
+        k_out_ptr=triton_k_out,
+        s_out_ptr=triton_s_out,
+        seq_len_ptr=seq_len_tensor,
+        seq_len_num_pow=seq_num_pow2,
+        page_size=page_size,
+        buf_numel_per_page=buf_numel_per_page,
+        index_head_dim=index_head_dim,
+        s_offset_in_page=s_offset_in_page,
+        page_indice_batch_offset=page_indice_batch_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+    )
+
+    torch_k_out, torch_s_out = golden_torch_gen(
+        seq_len_tensor=seq_len_tensor,
+        buffer_indexer=buffer_indexer,
+        buffer=buffer,
+        index_head_dim=index_head_dim,
+        page_size=page_size,
+    )
+
+    torch.testing.assert_close(
+        triton_k_out, torch_k_out, rtol=0, atol=0, msg="k outputs differ!"
+    )
+    torch.testing.assert_close(
+        triton_s_out, torch_s_out, rtol=0, atol=0, msg="s outputs differ!"
+    )
+    print("_get_k_and_s_triton_kernel test pass")
+
+    # perf test =====================
+    import time
+
+    torch.cuda.synchronize()
+    for _ in range(10):
+        _get_k_and_s_triton_kernel[grid](
+            buf_ptr=buffer,
+            page_indices_ptr=buffer_indexer,
+            k_out_ptr=triton_k_out,
+            s_out_ptr=triton_s_out,
+            seq_len_ptr=seq_len_tensor,
+            seq_len_num_pow=seq_num_pow2,
+            page_size=page_size,
+            buf_numel_per_page=buf_numel_per_page,
+            index_head_dim=index_head_dim,
+            s_offset_in_page=s_offset_in_page,
+            page_indice_batch_offset=page_indice_batch_offset,
+            BLOCK_SIZE=BLOCK_SIZE,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+        )
+
+    torch.cuda.synchronize()
+    start_time = time.perf_counter()
+
+    _get_k_and_s_triton_kernel[grid](
+        buf_ptr=buffer,
+        page_indices_ptr=buffer_indexer,
+        k_out_ptr=triton_k_out,
+        s_out_ptr=triton_s_out,
+        seq_len_ptr=seq_len_tensor,
+        seq_len_num_pow=seq_num_pow2,
+        page_size=page_size,
+        buf_numel_per_page=buf_numel_per_page,
+        index_head_dim=index_head_dim,
+        s_offset_in_page=s_offset_in_page,
+        page_indice_batch_offset=page_indice_batch_offset,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_SIZE_K=BLOCK_SIZE_K,
+    )
+
+    end_time = time.perf_counter()
+    print(
+        f"_get_k_and_s_triton_kernel triton kernel infer time is {((end_time-start_time)*1000):.4f} ms\n"
+    )
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available():
+        print("CUDA not available. Skipping tests.")
+        exit(0)
+
+    print("Start test cases...\n")
+
+    get_k_and_s_triton()
+
+    print("End test cases...\n")

third_party/sglang/test/manual/layers/attention/nsa/test_index_buf_accessor.py

@@ -0,0 +1,591 @@
+"""
+Correctness tests for NSA Indexer K/S Buffer Access with Fused Triton Kernels.
+
+This test verifies that the optimized Triton implementations (GetK, GetS, GetKAndS)
+produce identical results to the torch_fast baseline implementations.
+
+Test coverage:
+- GetK.triton() vs GetK.torch_fast()
+- GetS.triton() vs GetS.torch_fast()
+- GetKAndS.triton() vs separate GetK.torch_fast() + GetS.torch_fast()
+"""
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.nsa.index_buf_accessor import GetK, GetKAndS, GetS
+
+
+class MockNSATokenToKVPool:
+    """Mock pool object that mimics NSATokenToKVPool for testing."""
+
+    def __init__(
+        self,
+        page_size: int = 64,
+        index_head_dim: int = 128,
+        quant_block_size: int = 128,
+        device: str = "cuda",
+    ):
+        self.page_size = page_size
+        self.index_head_dim = index_head_dim
+        self.quant_block_size = quant_block_size
+        self.device = device
+
+
+def create_test_buffer(
+    num_pages: int,
+    page_size: int = 64,
+    index_head_dim: int = 128,
+    device: str = "cuda",
+) -> torch.Tensor:
+    """
+    Create a test buffer mimicking the K/S buffer structure.
+
+    Buffer layout per page:
+    - First page_size * index_head_dim bytes: K data (fp8, stored as uint8)
+    - Next page_size * 4 bytes: S data (fp32 scales, stored as uint8)
+
+    Args:
+        num_pages: Number of pages to allocate
+        page_size: Tokens per page (typically 64)
+        index_head_dim: Dimension of K vectors (typically 128)
+        device: Device to allocate on
+
+    Returns:
+        Buffer of shape (num_pages, page_size * index_head_dim + page_size * 4)
+    """
+    buf_numel_per_page = page_size * index_head_dim + page_size * 4
+    buf = torch.randint(
+        0, 256, (num_pages, buf_numel_per_page), dtype=torch.uint8, device=device
+    )
+    return buf
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGetK:
+    """Test cases for GetK.triton() correctness."""
+
+    @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
+    @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
+    @pytest.mark.parametrize("page_size", [64])
+    @pytest.mark.parametrize("index_head_dim", [128])
+    def test_getk_correctness(self, num_pages, seq_len, page_size, index_head_dim):
+        """Test GetK.triton() produces same output as GetK.torch_fast()."""
+        device = torch.device("cuda")
+
+        # Ensure seq_len doesn't exceed available pages
+        max_seq_len = num_pages * page_size
+        seq_len = min(seq_len, max_seq_len)
+
+        # Create mock pool
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+
+        # Create test buffer
+        buf = create_test_buffer(
+            num_pages=num_pages,
+            page_size=page_size,
+            index_head_dim=index_head_dim,
+            device=device,
+        )
+
+        # Create page indices
+        num_pages_needed = (seq_len + page_size - 1) // page_size
+        page_indices = torch.randint(
+            0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
+        )
+
+        # Run both implementations
+        output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetK.triton(pool, buf, seq_len, page_indices)
+
+        # Verify shapes
+        assert output_torch.shape == (seq_len, index_head_dim)
+        assert output_triton.shape == (seq_len, index_head_dim)
+        assert output_torch.dtype == torch.uint8
+        assert output_triton.dtype == torch.uint8
+
+        # Compare results (should be exact match)
+        torch.testing.assert_close(
+            output_triton, output_torch, rtol=0, atol=0, msg="GetK outputs differ"
+        )
+
+    def test_getk_sequential_pages(self):
+        """Test GetK with sequential page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 10
+        seq_len = 320  # 5 pages
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Sequential page indices [0, 1, 2, 3, 4]
+        page_indices = torch.arange(5, dtype=torch.int32, device=device)
+
+        output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetK.triton(pool, buf, seq_len, page_indices)
+
+        torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
+
+    def test_getk_repeated_pages(self):
+        """Test GetK with repeated page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 5
+        seq_len = 192  # 3 pages
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Repeated page indices [2, 2, 2]
+        page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
+
+        output_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetK.triton(pool, buf, seq_len, page_indices)
+
+        torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGetS:
+    """Test cases for GetS.triton() correctness."""
+
+    @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
+    @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
+    @pytest.mark.parametrize("page_size", [64])
+    @pytest.mark.parametrize("index_head_dim", [128])
+    def test_gets_correctness(self, num_pages, seq_len, page_size, index_head_dim):
+        """Test GetS.triton() produces same output as GetS.torch_fast()."""
+        device = torch.device("cuda")
+
+        # Ensure seq_len doesn't exceed available pages
+        max_seq_len = num_pages * page_size
+        seq_len = min(seq_len, max_seq_len)
+
+        # Create mock pool
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+
+        # Create test buffer
+        buf = create_test_buffer(
+            num_pages=num_pages,
+            page_size=page_size,
+            index_head_dim=index_head_dim,
+            device=device,
+        )
+
+        # Create page indices
+        num_pages_needed = (seq_len + page_size - 1) // page_size
+        page_indices = torch.randint(
+            0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
+        )
+
+        # Run both implementations
+        output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetS.triton(pool, buf, seq_len, page_indices)
+
+        # Verify shapes
+        assert output_torch.shape == (seq_len, 4)
+        assert output_triton.shape == (seq_len, 4)
+        assert output_torch.dtype == torch.uint8
+        assert output_triton.dtype == torch.uint8
+
+        # Compare results (should be exact match)
+        torch.testing.assert_close(
+            output_triton, output_torch, rtol=0, atol=0, msg="GetS outputs differ"
+        )
+
+    def test_gets_sequential_pages(self):
+        """Test GetS with sequential page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 10
+        seq_len = 320  # 5 pages
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Sequential page indices [0, 1, 2, 3, 4]
+        page_indices = torch.arange(5, dtype=torch.int32, device=device)
+
+        output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetS.triton(pool, buf, seq_len, page_indices)
+
+        torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
+
+    def test_gets_repeated_pages(self):
+        """Test GetS with repeated page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 5
+        seq_len = 192  # 3 pages
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Repeated page indices [2, 2, 2]
+        page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
+
+        output_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        output_triton = GetS.triton(pool, buf, seq_len, page_indices)
+
+        torch.testing.assert_close(output_triton, output_torch, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGetKAndS:
+    """Test cases for GetKAndS.triton() correctness."""
+
+    @pytest.mark.parametrize("num_pages", [1, 2, 4, 8, 16])
+    @pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024])
+    @pytest.mark.parametrize("page_size", [64])
+    @pytest.mark.parametrize("index_head_dim", [128])
+    def test_get_k_and_s_correctness(
+        self, num_pages, seq_len, page_size, index_head_dim
+    ):
+        """Test GetKAndS.triton() produces same output as separate torch_fast calls."""
+        device = torch.device("cuda")
+
+        # Ensure seq_len doesn't exceed available pages
+        max_seq_len = num_pages * page_size
+        seq_len = min(seq_len, max_seq_len)
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        # Create mock pool
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+
+        # Create test buffer
+        buf = create_test_buffer(
+            num_pages=num_pages,
+            page_size=page_size,
+            index_head_dim=index_head_dim,
+            device=device,
+        )
+
+        # Create page indices
+        num_pages_needed = (seq_len + page_size - 1) // page_size
+        page_indices = torch.randint(
+            0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
+        )
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Run baseline: separate torch_fast calls
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+
+        # Run fused Triton implementation
+        k_triton, s_triton = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        # Verify shapes
+        assert k_torch.shape == (seq_len, index_head_dim)
+        assert s_torch.shape == (seq_len, 4)
+        assert k_triton.shape == (seq_len, index_head_dim)
+        assert s_triton.shape == (seq_len, 4)
+
+        # Verify dtypes
+        assert k_torch.dtype == torch.uint8
+        assert s_torch.dtype == torch.uint8
+        assert k_triton.dtype == torch.uint8
+        assert s_triton.dtype == torch.uint8
+
+        # Compare K results
+        torch.testing.assert_close(
+            k_triton, k_torch, rtol=0, atol=0, msg="GetKAndS K outputs differ"
+        )
+
+        # Compare S results
+        torch.testing.assert_close(
+            s_triton, s_torch, rtol=0, atol=0, msg="GetKAndS S outputs differ"
+        )
+
+    def test_get_k_and_s_sequential_pages(self):
+        """Test GetKAndS with sequential page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 10
+        seq_len = 320  # 5 pages
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Sequential page indices [0, 1, 2, 3, 4]
+        page_indices = torch.arange(5, dtype=torch.int32, device=device)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Baseline
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+
+        # Fused
+        k_triton, s_triton = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+    def test_get_k_and_s_repeated_pages(self):
+        """Test GetKAndS with repeated page indices."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 5
+        seq_len = 192  # 3 pages
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        # Repeated page indices [2, 2, 2]
+        page_indices = torch.full((3,), 2, dtype=torch.int32, device=device)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Baseline
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+
+        # Fused
+        k_triton, s_triton = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+    def test_get_k_and_s_partial_page(self):
+        """Test GetKAndS when seq_len is not a multiple of page_size."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 5
+        seq_len = 100  # Not a multiple of 64
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        num_pages_needed = (seq_len + page_size - 1) // page_size
+        page_indices = torch.arange(num_pages_needed, dtype=torch.int32, device=device)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Baseline
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+
+        # Fused
+        k_triton, s_triton = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+
+        # Should handle partial pages correctly
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestEdgeCases:
+    """Test edge cases and boundary conditions."""
+
+    def test_single_token(self):
+        """Test with seq_len=1 (single token)."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 2
+        seq_len = 1
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+        page_indices = torch.tensor([0], dtype=torch.int32, device=device)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Test GetK
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        k_triton = GetK.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+
+        # Test GetS
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        s_triton = GetS.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+        # Test GetKAndS
+        k_triton2, s_triton2 = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
+
+    def test_exact_page_boundary(self):
+        """Test when seq_len exactly matches page boundaries."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 5
+        seq_len = 192  # Exactly 3 pages
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+        page_indices = torch.arange(3, dtype=torch.int32, device=device)
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Test GetK
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        k_triton = GetK.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+
+        # Test GetS
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        s_triton = GetS.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+        # Test GetKAndS
+        k_triton2, s_triton2 = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
+
+    def test_large_seq_len(self):
+        """Test with large sequence length."""
+        device = torch.device("cuda")
+        page_size = 64
+        index_head_dim = 128
+        num_pages = 100
+        seq_len = 4096  # 64 pages
+        seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
+
+        pool = MockNSATokenToKVPool(
+            page_size=page_size, index_head_dim=index_head_dim, device=device
+        )
+        buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
+
+        num_pages_needed = (seq_len + page_size - 1) // page_size
+        page_indices = torch.randint(
+            0, num_pages, (num_pages_needed,), dtype=torch.int32, device=device
+        )
+        page_indices_ = page_indices.unsqueeze(0)
+
+        # Test GetK
+        k_torch = GetK.torch_fast(pool, buf, seq_len, page_indices)
+        k_triton = GetK.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(k_triton, k_torch, rtol=0, atol=0)
+
+        # Test GetS
+        s_torch = GetS.torch_fast(pool, buf, seq_len, page_indices)
+        s_triton = GetS.triton(pool, buf, seq_len, page_indices)
+        torch.testing.assert_close(s_triton, s_torch, rtol=0, atol=0)
+
+        # Test GetKAndS
+        k_triton2, s_triton2 = GetKAndS.triton(
+            pool, buf, page_indices_, seq_len_tensor, seq_len, seq_len
+        )
+        torch.testing.assert_close(k_triton2, k_torch, rtol=0, atol=0)
+        torch.testing.assert_close(s_triton2, s_torch, rtol=0, atol=0)
+
+
+def print_test_summary():
+    """Print a summary message about the test suite."""
+    print("\n" + "=" * 80)
+    print("NSA Indexer K/S Buffer Accessor Correctness Tests")
+    print("=" * 80)
+    print("Testing Triton implementations against torch_fast baseline:")
+    print("  - GetK.triton() vs GetK.torch_fast()")
+    print("  - GetS.triton() vs GetS.torch_fast()")
+    print("  - GetKAndS.triton() vs separate GetK/GetS torch_fast() calls")
+    print("=" * 80)
+    print()
+
+
+if __name__ == "__main__":
+    # Run tests manually
+    if not torch.cuda.is_available():
+        print("CUDA not available. Skipping tests.")
+        exit(0)
+
+    print_test_summary()
+
+    # Run a few sample tests
+    print("Running sample correctness tests...\n")
+
+    # Test GetK
+    print("Testing GetK...")
+    test_getk = TestGetK()
+    test_getk.test_getk_correctness(
+        num_pages=4, seq_len=256, page_size=64, index_head_dim=128
+    )
+    test_getk.test_getk_sequential_pages()
+    print("✓ GetK tests passed\n")
+
+    # Test GetS
+    print("Testing GetS...")
+    test_gets = TestGetS()
+    test_gets.test_gets_correctness(
+        num_pages=4, seq_len=256, page_size=64, index_head_dim=128
+    )
+    test_gets.test_gets_sequential_pages()
+    print("✓ GetS tests passed\n")
+
+    # Test GetKAndS
+    print("Testing GetKAndS SeqLen=256...")
+    test_get_k_and_s = TestGetKAndS()
+    test_get_k_and_s.test_get_k_and_s_correctness(
+        num_pages=4, seq_len=256, page_size=64, index_head_dim=128
+    )
+    test_get_k_and_s.test_get_k_and_s_sequential_pages()
+    test_get_k_and_s.test_get_k_and_s_partial_page()
+    print("✓ GetKAndS SeqLen=256 tests passed\n")
+
+    print("Testing GetKAndS SeqLen=128K...")
+    test_get_k_and_s = TestGetKAndS()
+    test_get_k_and_s.test_get_k_and_s_correctness(
+        num_pages=2048, seq_len=131072, page_size=64, index_head_dim=128
+    )
+    test_get_k_and_s.test_get_k_and_s_sequential_pages()
+    test_get_k_and_s.test_get_k_and_s_partial_page()
+    print("✓ GetKAndS SeqLen=128K tests passed\n")
+
+    # Test edge cases
+    print("Testing edge cases...")
+    test_edge = TestEdgeCases()
+    test_edge.test_single_token()
+    test_edge.test_exact_page_boundary()
+    test_edge.test_large_seq_len()
+    print("✓ Edge case tests passed\n")
+
+    print("=" * 80)
+    print("All correctness tests passed successfully!")
+    print("=" * 80)

third_party/sglang/test/manual/layers/moe/test_moe_runners_1gpu.py

@@ -0,0 +1,211 @@
+import os
+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.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
+    DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
+    DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE,
+    DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestMoERunner(CustomTestCase):
+    BASE_URL = DEFAULT_URL_FOR_TEST
+    TIMEOUT = 6000
+    DEFAULT_EVAL_KWARGS = {
+        "eval_name": "mmlu",
+        "num_examples": 5,
+        "num_threads": 1,
+    }
+
+    CONFIGS = {
+        "moe_runner_auto": {
+            "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "triton",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_triton": {
+            "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "triton",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_triton_kernel": {
+            "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "triton_kernel",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_flashinfer_cutlass": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,  # requires model with modelopt_fp4 quantization
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "flashinfer_cutlass",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_deep_gemm": {
+            "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "deep_gemm",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_flashinfer_trtllm": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,  # modelopt_fp4 or fp8 quantization is required for Flashinfer trtllm MOE
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "flashinfer_trtllm",
+            ],
+        },
+        "moe_runner_flashinfer_mxfp4": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "flashinfer_mxfp4",
+                "--quantization",
+                "mxfp4",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_flashinfer_cutedsl": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "flashinfer_cutedsl",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_cutlass": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "cutlass",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+        "moe_runner_cutlass_fp8": {
+            "model": DEFAULT_MODEL_NAME_FOR_TEST_FP8_WITH_MOE,
+            "timeout": 3600,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "cutlass",
+                "--attention-backend",
+                "triton",
+                "--sampling-backend",
+                "pytorch",
+                "--disable-cuda-graph",
+            ],
+        },
+        "moe_runner_speculative": {
+            "model": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "triton",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-draft-model-path",
+                DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+                "--speculative-moe-runner-backend",
+                "triton",
+                "--speculative-num-steps",
+                "2",
+                "--speculative-num-draft-tokens",
+                "4",
+                "--attention-backend",
+                "torch_native",
+                "--sampling-backend",
+                "pytorch",
+            ],
+        },
+    }
+
+    def _run_config(self, config: dict) -> None:
+        model = config["model"]
+        other_args = config.get("other_args", [])
+        eval_kwargs = self.DEFAULT_EVAL_KWARGS
+        env = dict(os.environ)
+        env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1"
+        env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
+        env.update(config.get("env_overrides", {}))
+        timeout = config.get("timeout", self.TIMEOUT)
+
+        process = popen_launch_server(
+            model,
+            self.BASE_URL,
+            timeout=timeout,
+            other_args=other_args,
+            env=env,
+        )
+        try:
+            args = SimpleNamespace(
+                base_url=self.BASE_URL,
+                model=model,
+                **eval_kwargs,
+            )
+            metrics = run_eval(args)
+            print(f"{metrics=}")
+            self.assertGreaterEqual(metrics["score"], 0.48)
+        finally:
+            kill_process_tree(process.pid)
+
+
+for _name, _cfg in TestMoERunner.CONFIGS.items():
+    setattr(
+        TestMoERunner,
+        f"test_{_name}",
+        (lambda self, cfg=_cfg: self._run_config(cfg)),
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/layers/moe/test_moe_runners_4gpu.py

@@ -0,0 +1,116 @@
+import os
+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.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestMoERunner4GPU(CustomTestCase):
+    BASE_URL = DEFAULT_URL_FOR_TEST
+    TIMEOUT = 6000
+    DEFAULT_EVAL_KWARGS = {
+        "eval_name": "mmlu",
+        "num_examples": 5,
+        "num_threads": 1,
+    }
+
+    CONFIGS = {
+        "moe_runner_cutlass_w4a8": {
+            "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8",  # FP8 W8A8 MoE model
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "cutlass",
+                "--attention-backend",
+                "triton",
+                "--sampling-backend",
+                "pytorch",
+                "--tp-size",
+                "4",
+            ],
+        },
+        "moe_runner_cutlass_w4a8_deepep_normal": {
+            "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8",  # FP8 W8A8 MoE model
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "cutlass",
+                "--moe-a2a-backend",
+                "deepep",
+                "--deepep-mode",
+                "normal",
+                "--attention-backend",
+                "triton",
+                "--sampling-backend",
+                "pytorch",
+                "--tp-size",
+                "4",
+            ],
+        },
+        "moe_runner_cutlass_w4a8_deepep_ll": {
+            "model": "tencent/DeepSeek-V3.1-Terminus-W4AFP8",  # FP8 W8A8 MoE model
+            "env_overrides": {"SGLANG_DEEPEP_BF16_DISPATCH": "1"},
+            "other_args": [
+                "--trust-remote-code",
+                "--moe-runner-backend",
+                "cutlass",
+                "--moe-a2a-backend",
+                "deepep",
+                "--deepep-mode",
+                "low_latency",
+                "--attention-backend",
+                "triton",
+                "--sampling-backend",
+                "pytorch",
+                "--tp-size",
+                "4",
+            ],
+        },
+    }
+
+    def _run_config(self, config: dict) -> None:
+        model = config["model"]
+        other_args = config.get("other_args", [])
+        eval_kwargs = self.DEFAULT_EVAL_KWARGS
+        env = dict(os.environ)
+        env["SGLANG_ENABLE_JIT_DEEPGEMM"] = "1"
+        env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0"
+        env.update(config.get("env_overrides", {}))
+        timeout = config.get("timeout", self.TIMEOUT)
+
+        process = popen_launch_server(
+            model,
+            self.BASE_URL,
+            timeout=timeout,
+            other_args=other_args,
+            env=env,
+        )
+        try:
+            args = SimpleNamespace(
+                base_url=self.BASE_URL,
+                model=model,
+                **eval_kwargs,
+            )
+            metrics = run_eval(args)
+            print(f"{metrics=}")
+            self.assertGreaterEqual(metrics["score"], 0.48)
+        finally:
+            kill_process_tree(process.pid)
+
+
+for _name, _cfg in TestMoERunner4GPU.CONFIGS.items():
+    setattr(
+        TestMoERunner4GPU,
+        f"test_{_name}",
+        (lambda self, cfg=_cfg: self._run_config(cfg)),
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lora/test_lora_cuda_graph.py

@@ -0,0 +1,107 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import os
+import unittest
+from typing import List
+
+from sglang.test.lora_utils import (
+    ALL_OTHER_LORA_MODELS,
+    CI_LORA_MODELS,
+    DEFAULT_PROMPTS,
+    TORCH_DTYPES,
+    LoRAModelCase,
+    run_lora_test_by_batch,
+    run_lora_test_one_by_one,
+)
+from sglang.test.test_utils import CustomTestCase, is_in_ci
+
+TEST_CUDA_GRAPH_PADDING_PROMPTS = [
+    "AI is a field of computer science focused on",
+    """
+    ### Instruction:
+    Tell me about llamas and alpacas
+    ### Response:
+    Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids (camels, dromedaries). Llamas live in the Andean mountains of South America where they graze on grasses and shrubs. Alpaca is another name for domesticated llama. The word "alpaca" comes from an Incan language meaning "golden fleece." Alpacas look very similar to llamas but are smaller than their wild relatives. Both species were used by ancient people as pack animals and for meat. Today both llamas and alpacas are raised primarily for their fiber which can be spun into yarn or knitted into clothing.
+    ### Question 2:
+    What do you know about llamas?
+    ### Answer:
+    """,
+    "Computer science is the study of",
+]
+
+
+class TestLoRACudaGraph(CustomTestCase):
+
+    def _run_without_cuda_graph_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        # Since we have already enabled CUDA graph by default in other lora tests,
+        # we only need to run lora tests without CUDA graph here.
+        for model_case in model_cases:
+            # If skip_long_prompt is True, filter out prompts longer than 1000 characters
+            prompts = (
+                DEFAULT_PROMPTS
+                if not model_case.skip_long_prompt
+                else [p for p in DEFAULT_PROMPTS if len(p) < 1000]
+            )
+            for torch_dtype in TORCH_DTYPES:
+                run_lora_test_one_by_one(
+                    prompts,
+                    model_case,
+                    torch_dtype,
+                    max_new_tokens=32,
+                    disable_cuda_graph=True,
+                    test_tag="without_cuda_graph",
+                )
+
+    def _run_cuda_graph_padding_on_model_cases(self, model_cases: List[LoRAModelCase]):
+        for model_case in model_cases:
+            # Run a batch size of 3, which will not be captured by CUDA graph and need padding
+            prompts = TEST_CUDA_GRAPH_PADDING_PROMPTS
+            for torch_dtype in TORCH_DTYPES:
+                run_lora_test_by_batch(
+                    prompts,
+                    model_case,
+                    torch_dtype,
+                    max_new_tokens=32,
+                    disable_cuda_graph=False,
+                    test_tag="cuda_graph_padding",
+                )
+
+    def test_ci_lora_models(self):
+        self._run_without_cuda_graph_on_model_cases(CI_LORA_MODELS)
+        self._run_cuda_graph_padding_on_model_cases(CI_LORA_MODELS)
+
+    def test_all_lora_models(self):
+        if is_in_ci():
+            return
+
+        # Retain ONLY_RUN check here
+        filtered_models = []
+        for model_case in ALL_OTHER_LORA_MODELS:
+            if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base:
+                continue
+            filtered_models.append(model_case)
+
+        self._run_without_cuda_graph_on_model_cases(filtered_models)
+        self._run_cuda_graph_padding_on_model_cases(filtered_models)
+
+
+if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
+    unittest.main(warnings="ignore")

third_party/sglang/test/manual/lora/test_lora_llama4.py

@@ -0,0 +1,63 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+)
+
+MODELS = [
+    SimpleNamespace(
+        model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+        tp_size=8,
+    ),
+]
+
+
+@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
+class TestLlama4LoRA(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+    def test_bringup(self):
+        for model in MODELS:
+            try:
+                process = popen_launch_server(
+                    model.model,
+                    self.base_url,
+                    timeout=3 * DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                    other_args=[
+                        "--enable-lora",
+                        "--max-lora-rank",
+                        "64",
+                        "--lora-target-modules",
+                        "all",
+                        "--tp-size",
+                        str(model.tp_size),
+                        "--context-length",
+                        "262144",
+                        "--attention-backend",
+                        "fa3",
+                    ],
+                )
+            except Exception as e:
+                print(f"Error testing {model.model}: {e}")
+                self.fail(f"Test failed for {model.model}: {e}")
+
+            finally:
+                # Ensure process cleanup happens regardless of success/failure
+                if process is not None and process.poll() is None:
+                    print(f"Cleaning up process {process.pid}")
+                    try:
+                        kill_process_tree(process.pid)
+                    except Exception as e:
+                        print(f"Error killing process: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lora/test_lora_ops.py

@@ -0,0 +1,218 @@
+import random
+import unittest
+
+import torch
+
+from sglang.srt.lora.torch_ops.lora_ops import sgemm_lora_a_fwd, sgemm_lora_b_fwd
+from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestLoraOps(CustomTestCase):
+    def test_sgemm_lora_a_fwd(self):
+        batch_size = 2
+        input_dim = 1024
+        num_loras = 3
+        dtype = torch.float32
+
+        possible_lora_ranks = [8, 16, 32, 64, 128, 256]
+        lora_ranks = random.sample(
+            possible_lora_ranks,
+            counts=[num_loras] * len(possible_lora_ranks),
+            k=num_loras,
+        )
+
+        max_lora_rank = max(lora_ranks)
+
+        possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
+        lora_scaling = random.sample(
+            possible_lora_scaling,
+            counts=[num_loras] * len(possible_lora_scaling),
+            k=num_loras,
+        )
+
+        inputs = torch.randn(batch_size, input_dim, dtype=dtype)
+        lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+        seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
+        lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
+        lora_scaling_tensor = torch.tensor(
+            lora_scaling, dtype=torch.float16, device="cpu"
+        )
+
+        expect_output = reference_sgmv_shrink(
+            inputs,
+            lora_a_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            lora_scaling_tensor,
+        )
+
+        actual_output = sgemm_lora_a_fwd(
+            inputs,
+            lora_a_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            lora_scaling_tensor,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_sgemm_lora_b_fwd(self):
+        batch_size = 2
+        output_dim = 1024
+        num_loras = 3
+        dtype = torch.float32
+
+        possible_lora_ranks = [8, 16, 32, 64, 128, 256]
+        lora_ranks = random.sample(
+            possible_lora_ranks,
+            counts=[num_loras] * len(possible_lora_ranks),
+            k=num_loras,
+        )
+
+        max_lora_rank = max(lora_ranks)
+
+        inputs = torch.randn(batch_size, max_lora_rank, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
+        lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
+        seq_len_tensor = torch.ones(batch_size, dtype=torch.int32, device="cpu")
+        lora_indices_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+        slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
+
+        expect_output = reference_sgmv_expand(
+            inputs,
+            lora_b_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets,
+        )
+
+        actual_output = sgemm_lora_b_fwd(
+            inputs,
+            lora_b_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_sgemm_lora_a_fwd_expand(self):
+        batch_size = 2
+        input_dim = 1024
+        num_loras = 3
+        dtype = torch.float32
+
+        possible_lora_ranks = [8, 16, 32, 64, 128, 256]
+        lora_ranks = random.sample(
+            possible_lora_ranks,
+            counts=[num_loras] * len(possible_lora_ranks),
+            k=num_loras,
+        )
+
+        max_lora_rank = max(lora_ranks)
+
+        possible_lora_scaling = [0.25, 0.5, 1.0, 2.0, 4.0]
+        lora_scaling = random.sample(
+            possible_lora_scaling,
+            counts=[num_loras] * len(possible_lora_scaling),
+            k=num_loras,
+        )
+
+        seq_len_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+
+        seq_len = sum(seq_len_tensor)
+
+        inputs = torch.randn(seq_len, input_dim, dtype=dtype)
+        lora_a_weights = torch.randn(num_loras, max_lora_rank, input_dim, dtype=dtype)
+        lora_indices_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+        lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
+        lora_scaling_tensor = torch.tensor(
+            lora_scaling, dtype=torch.float16, device="cpu"
+        )
+
+        expect_output = reference_sgmv_shrink(
+            inputs,
+            lora_a_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            lora_scaling_tensor,
+        )
+
+        actual_output = sgemm_lora_a_fwd(
+            inputs,
+            lora_a_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            lora_scaling_tensor,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_sgemm_lora_b_fwd_expand(self):
+        batch_size = 2
+        output_dim = 1024
+        num_loras = 3
+        dtype = torch.float32
+
+        possible_lora_ranks = [8, 16, 32, 64, 128, 256]
+        lora_ranks = random.sample(
+            possible_lora_ranks,
+            counts=[num_loras] * len(possible_lora_ranks),
+            k=num_loras,
+        )
+
+        max_lora_rank = max(lora_ranks)
+
+        seq_len_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+
+        seq_len = sum(seq_len_tensor)
+
+        inputs = torch.randn(seq_len, max_lora_rank, dtype=dtype)
+        lora_b_weights = torch.randn(num_loras, output_dim, max_lora_rank, dtype=dtype)
+        lora_ranks_tensor = torch.tensor(lora_ranks, dtype=torch.int32, device="cpu")
+        lora_indices_tensor = torch.randint(
+            num_loras, (batch_size,), dtype=torch.int32, device="cpu"
+        )
+        slice_offsets = torch.tensor([0, output_dim], dtype=torch.int32, device="cpu")
+
+        expect_output = reference_sgmv_expand(
+            inputs,
+            lora_b_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets,
+        )
+
+        actual_output = sgemm_lora_b_fwd(
+            inputs,
+            lora_b_weights,
+            lora_indices_tensor,
+            seq_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets,
+        )
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/lora/test_lora_spec_decoding.py

@@ -0,0 +1,52 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import unittest
+
+from sglang.test.lora_utils import (
+    CI_MULTI_LORA_MODELS,
+    LORA_MODELS_QWEN3,
+    run_lora_multiple_batch_on_model_cases,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestLoRASpecDecoding(CustomTestCase):
+    def test_qwen(self):
+        run_lora_multiple_batch_on_model_cases(
+            LORA_MODELS_QWEN3,
+            attention_backend="triton",
+            use_spec_decoding=True,
+            disable_cuda_graph=True,
+            enable_deterministic_inference=True,
+        )
+
+    def test_llama(self):
+        run_lora_multiple_batch_on_model_cases(
+            CI_MULTI_LORA_MODELS,
+            attention_backend="triton",
+            use_spec_decoding=True,
+            disable_cuda_graph=True,
+            enable_deterministic_inference=True,
+        )
+
+
+if __name__ == "__main__":
+    try:
+        mp.set_start_method("spawn")
+    except RuntimeError:
+        pass
+
+    unittest.main(warnings="ignore")

third_party/sglang/test/manual/lora/test_torch_backend.py

@@ -0,0 +1,246 @@
+import unittest
+
+import torch
+
+from sglang.srt.lora.backend.torch_backend import TorchNativeLoRABackend
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.test.lora_utils import reference_sgmv_expand, reference_sgmv_shrink
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestTorchNativeLoRABackend(CustomTestCase):
+
+    device = "cpu"
+
+    # set duplicate weights to test merging during prepare_lora_batch
+    weight_indices = [0, 0, 1]
+    lora_ranks = [1, 1]
+    scalings = [1.0, 0.5]
+    seq_lens = [1, 1, 1]
+    use_cuda_graph = False
+
+    forward_batch = ForwardBatch(
+        forward_mode=ForwardMode.EXTEND,
+        batch_size=3,
+        input_ids=torch.tensor([[1], [2], [3]], dtype=torch.int32),
+        req_pool_indices=None,
+        seq_lens=None,
+        out_cache_loc=None,
+        seq_lens_sum=3,
+        extend_seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
+        extend_seq_lens_cpu=seq_lens,
+    )
+
+    @classmethod
+    def setUpClass(cls):
+        cls.backend = TorchNativeLoRABackend(max_loras_per_batch=2, device=cls.device)
+        cls.backend.prepare_lora_batch(
+            forward_batch=cls.forward_batch,
+            weight_indices=cls.weight_indices,
+            lora_ranks=cls.lora_ranks,
+            scalings=cls.scalings,
+            use_cuda_graph=cls.use_cuda_graph,
+        )
+
+    def test_run_lora_a_sgemm(self):
+        batch_size = 3
+        input_dim = 4
+        output_dim = 6
+        num_loras = 3
+        dtype = torch.float32
+
+        x = torch.randn(batch_size, input_dim, dtype=dtype)
+        weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+
+        weight_indices_tensor = torch.tensor(
+            self.weight_indices, dtype=torch.int32, device=self.device
+        )
+
+        seg_len_tensor = torch.tensor(
+            self.seq_lens, dtype=torch.int32, device=self.device
+        )
+
+        lora_ranks_tensor = torch.tensor(
+            self.lora_ranks, dtype=torch.int32, device=self.device
+        )
+
+        scalings_tensor = torch.tensor(
+            self.scalings, dtype=torch.float, device=self.device
+        )
+
+        expect_output = reference_sgmv_shrink(
+            x,
+            weights,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            scalings_tensor,
+        )
+
+        actual_output = self.backend.run_lora_a_sgemm(x, weights)
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_run_lora_b_sgemm(self):
+        batch_size = 3
+        input_dim = 6
+        output_dim = 4
+        num_loras = 3
+        dtype = torch.float32
+
+        x = torch.randn(batch_size, input_dim, dtype=dtype)
+        weights = torch.randn(num_loras, output_dim, input_dim, dtype=dtype)
+        _, weight_out_dim, _ = weights.shape
+
+        weight_indices_tensor = torch.tensor(
+            self.weight_indices, dtype=torch.int32, device=self.device
+        )
+
+        seg_len_tensor = torch.tensor(
+            self.seq_lens, dtype=torch.int32, device=self.device
+        )
+
+        lora_ranks_tensor = torch.tensor(
+            self.lora_ranks, dtype=torch.int32, device=self.device
+        )
+
+        expect_output = reference_sgmv_expand(
+            x,
+            weights,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets=torch.tensor(
+                [0, weight_out_dim], dtype=torch.int32, device="cpu"
+            ),
+        )
+
+        actual_output = self.backend.run_lora_b_sgemm(x, weights)
+
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_run_qkv_lora(self):
+        batch_size = 3
+        num_loras = 3
+        input_dim = 6
+        output_offset = [0, 3, 6, 9]
+        output_dim = output_offset[-1]
+        num_slices = len(output_offset) - 1  # 3 slices for Q, K, V
+        max_lora_rank = max(self.lora_ranks)
+        dtype = torch.float32
+
+        x = torch.randn(batch_size, input_dim, dtype=dtype)
+        output_offset_cpu = torch.tensor(output_offset, dtype=torch.int32)
+        qkv_lora_a = torch.randn(
+            num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
+        )
+        qkv_lora_b = torch.randn(
+            num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
+        )
+
+        weight_indices_tensor = torch.tensor(
+            self.weight_indices, dtype=torch.int32, device=self.device
+        )
+
+        seg_len_tensor = torch.tensor(
+            self.seq_lens, dtype=torch.int32, device=self.device
+        )
+
+        lora_ranks_tensor = torch.tensor(
+            self.lora_ranks, dtype=torch.int32, device=self.device
+        )
+
+        scalings_tensor = torch.tensor(
+            self.scalings, dtype=torch.float, device=self.device
+        )
+
+        expect_lora_a_output = reference_sgmv_shrink(
+            x,
+            qkv_lora_a,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            scalings_tensor,
+            num_slices,
+        )
+
+        expect_output = reference_sgmv_expand(
+            expect_lora_a_output,
+            qkv_lora_b,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            output_offset_cpu,
+        )
+
+        actual_output = self.backend.run_qkv_lora(
+            x, qkv_lora_a, qkv_lora_b, None, output_offset_cpu, 0
+        )
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+    def test_run_gate_up_lora(self):
+        batch_size = 3
+        input_dim = 6
+        output_dim = 4
+        num_loras = 3
+        dtype = torch.float32
+
+        max_lora_rank = max(self.lora_ranks)
+
+        num_slices = 2
+
+        x = torch.randn(batch_size, input_dim, dtype=dtype)
+        gate_up_lora_a = torch.randn(
+            num_loras, max_lora_rank * num_slices, input_dim, dtype=dtype
+        )
+        gate_up_lora_b = torch.randn(
+            num_loras, output_dim, max_lora_rank * num_slices, dtype=dtype
+        )
+
+        _, weight_out_dim, _ = gate_up_lora_b.shape
+        slice_size = weight_out_dim // num_slices
+        output_offset = torch.tensor(
+            [0, slice_size, weight_out_dim], dtype=torch.int32, device="cpu"
+        )
+
+        weight_indices_tensor = torch.tensor(
+            self.weight_indices, dtype=torch.int32, device=self.device
+        )
+
+        seg_len_tensor = torch.tensor(
+            self.seq_lens, dtype=torch.int32, device=self.device
+        )
+
+        lora_ranks_tensor = torch.tensor(
+            self.lora_ranks, dtype=torch.int32, device=self.device
+        )
+
+        scalings_tensor = torch.tensor(
+            self.scalings, dtype=torch.float, device=self.device
+        )
+
+        expect_lora_a_output = reference_sgmv_shrink(
+            x,
+            gate_up_lora_a,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            scalings_tensor,
+            num_slices,
+        )
+
+        expect_output = reference_sgmv_expand(
+            expect_lora_a_output,
+            gate_up_lora_b,
+            weight_indices_tensor,
+            seg_len_tensor,
+            lora_ranks_tensor,
+            slice_offsets=output_offset,
+        )
+
+        actual_output = self.backend.run_gate_up_lora(x, gate_up_lora_a, gate_up_lora_b)
+        self.assertTrue(torch.allclose(actual_output, expect_output))
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_clip_models.py

@@ -0,0 +1,78 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import multiprocessing as mp
+import unittest
+
+import torch
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import get_similarities
+
+TEXTS = "two Subway Series sandwiches with meats, cheese, lettuce, tomatoes, and onions on a black background, accompanied by the Subway Series logo, highlighting a new sandwich series."
+IMAGES = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg"
+MODELS = [
+    ("openai/clip-vit-large-patch14-336", 1e-5),
+]
+TORCH_DTYPES = [torch.float16]
+
+
+class TestClipModels(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def assert_close_embeddings(self, model, prefill_tolerance, torch_dtype):
+
+        with HFRunner(
+            model,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+        ) as hf_runner:
+            hf_text_embeds = hf_runner.forward(prompts=TEXTS)
+            hf_image_embeds = hf_runner.forward(image_data=IMAGES)
+
+        with SRTRunner(
+            model,
+            tp_size=1,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+        ) as srt_runner:
+            text_embeds = srt_runner.forward(prompts=TEXTS)
+            image_embeds = srt_runner.forward(prompts="padding", image_data=IMAGES)
+
+        text_similarity = get_similarities(
+            text_embeds.embed_logits[0], hf_text_embeds.embed_logits[0]
+        )
+        image_similarity = get_similarities(
+            image_embeds.embed_logits[0], hf_image_embeds.embed_logits[0]
+        )
+        print("text similarity diff", abs(text_similarity - 1))
+        print("image similarity diff", abs(image_similarity - 1))
+        assert torch.all(
+            abs(text_similarity - 1) < prefill_tolerance
+        ), "embeddings are not all close"
+        assert torch.all(
+            abs(image_similarity - 1) < prefill_tolerance
+        ), "embeddings are not all close"
+
+    def test_accuracy(self):
+        for model, prefill_tolerance in MODELS:
+            for torch_dtype in TORCH_DTYPES:
+                self.assert_close_embeddings(model, prefill_tolerance, torch_dtype)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_falcon_h1_models.py

@@ -0,0 +1,146 @@
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestFalconH1(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "tiiuae/Falcon-H1-0.5B-Instruct"
+        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=[
+                "--tensor-parallel-size",
+                "1",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.74)
+
+
+class TestFalconH1TP4(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "tiiuae/Falcon-H1-0.5B-Instruct"
+        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=[
+                "--tensor-parallel-size",
+                "4",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.74)
+
+
+class TestFalconH1NoGatedRMS(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "tiiuae/Falcon-H1-1.5B-Instruct"
+        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=[
+                "--tensor-parallel-size",
+                "1",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.74)
+
+
+class TestFalconH1NoGatedTP4(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "tiiuae/Falcon-H1-1.5B-Instruct"
+        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=[
+                "--tensor-parallel-size",
+                "4",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.74)

third_party/sglang/test/manual/models/test_gme_qwen_models.py

@@ -0,0 +1,85 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+import multiprocessing as mp
+import unittest
+
+import torch
+
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase, get_similarities
+
+TEXTS = "two Subway Series sandwiches with meats, cheese, lettuce, tomatoes, and onions on a black background, accompanied by the Subway Series logo, highlighting a new sandwich series."
+IMAGES = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg"
+
+
+MODELS = [
+    ("Alibaba-NLP/gme-Qwen2-VL-2B-Instruct", 1e-3),
+]
+TORCH_DTYPES = [torch.float16]
+
+
+class TestQmeQwenModels(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def assert_close_embeddings(self, model, prefill_tolerance, torch_dtype):
+
+        prompts_no_image = f"<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n{TEXTS}<|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
+        prompts_with_image = f"<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>assistant\n<|endoftext|>"
+        with HFRunner(
+            model,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+        ) as hf_runner:
+            hf_text_embeddings = hf_runner.forward(prompts=[prompts_no_image])
+            hf_image_embeddings = hf_runner.forward(
+                prompts=[prompts_with_image], image_data=[IMAGES]
+            )
+        with SRTRunner(
+            model,
+            tp_size=1,
+            torch_dtype=torch_dtype,
+            model_type="embedding",
+        ) as srt_runner:
+            srt_text_embeddings = srt_runner.forward(prompts=prompts_no_image)
+            srt_image_embeddings = srt_runner.forward(
+                prompts=prompts_with_image, image_data=IMAGES
+            )
+
+        similarity = get_similarities(
+            hf_text_embeddings.embed_logits[0], srt_text_embeddings.embed_logits[0]
+        )
+        print("texts similarity diff", abs(similarity - 1))
+        assert torch.all(
+            abs(similarity - 1) < prefill_tolerance
+        ), "embeddings are not all close"
+        similarity = get_similarities(
+            hf_image_embeddings.embed_logits[0], srt_image_embeddings.embed_logits[0]
+        )
+        print("images similarity diff", abs(similarity - 1))
+        assert torch.all(
+            abs(similarity - 1) < prefill_tolerance
+        ), "embeddings are not all close"
+
+    def test_accuracy(self):
+        for model, prefill_tolerance in MODELS:
+            for torch_dtype in TORCH_DTYPES:
+                self.assert_close_embeddings(model, prefill_tolerance, torch_dtype)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_grok_models.py

@@ -0,0 +1,53 @@
+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.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestGrok(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "lmzheng/grok-1"
+        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=[
+                "--load-format",
+                "dummy",
+                "--json-model-override-args",
+                '{"num_hidden_layers": 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",
+            api="completion",
+            max_tokens=512,
+            num_examples=64,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        # It is dummy weights so we only assert the output throughput instead of accuracy.
+        self.assertGreater(metrics["output_throughput"], 1000)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_kimi_k2_models.py

@@ -0,0 +1,70 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval 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,
+    write_github_step_summary,
+)
+
+
+class TestKimiK2Thinking(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "moonshotai/Kimi-K2-Thinking"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--tp",
+            "8",
+            "--trust-remote-code",
+            "--tool-call-parser",
+            "kimi_k2",
+            "--reasoning-parser",
+            "kimi_k2",
+            "--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,
+    ):
+        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(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (Kimi-K2-Thinking)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.95)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_llama4_models.py

@@ -0,0 +1,72 @@
+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.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+MODELS = [
+    SimpleNamespace(
+        model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
+        accuracy=0.9,
+        tp_size=4,
+    ),
+]
+
+
+class TestLlama4(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+    def test_gsm8k(self):
+
+        for model in MODELS:
+            try:
+                process = popen_launch_server(
+                    model.model,
+                    self.base_url,
+                    timeout=3 * DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                    other_args=[
+                        "--chat-template",
+                        "llama-4",
+                        "--tp-size",
+                        str(model.tp_size),
+                        "--mem-fraction-static",
+                        "0.8",
+                        "--context-length",
+                        "8192",
+                    ],
+                )
+                args = SimpleNamespace(
+                    base_url=self.base_url,
+                    eval_name="gsm8k",
+                    api="completion",
+                    max_tokens=512,
+                    num_examples=200,
+                    num_threads=128,
+                )
+                metrics = run_eval(args)
+                print(f"{metrics=}")
+                self.assertGreaterEqual(metrics["score"], model.accuracy)
+            except Exception as e:
+                print(f"Error testing {model.model}: {e}")
+                self.fail(f"Test failed for {model.model}: {e}")
+
+            finally:
+                # Ensure process cleanup happens regardless of success/failure
+                if process is not None and process.poll() is None:
+                    print(f"Cleaning up process {process.pid}")
+                    try:
+                        kill_process_tree(process.pid)
+                    except Exception as e:
+                        print(f"Error killing process: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_mistral_large3_basic.py

@@ -0,0 +1,88 @@
+import os
+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,
+)
+
+MISTRAL_LARGE3_MODEL_PATH = "mistralai/Mistral-Large-3-675B-Instruct-2512"
+
+
+class TestMistralLarge3Basic(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        # Set environment variable to disable JIT DeepGemm
+        os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"] = "0"
+
+        cls.model = MISTRAL_LARGE3_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--tp",
+            "8",
+            "--attention-backend",
+            "trtllm_mla",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true}',
+            "--chat-template",
+            "mistral",
+        ]
+        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)
+        # Clean up environment variable
+        if "SGLANG_ENABLE_JIT_DEEPGEMM" in os.environ:
+            del os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"]
+
+    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 (mistral-large-3)\n" f'{metrics["score"]=:.3f}\n'
+            )
+            self.assertGreater(metrics["score"], 0.90)
+
+    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 (mistral-large-3)\n" f"{speed=:.2f} token/s\n"
+            )
+            self.assertGreater(speed, 50)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_mtp_models.py

@@ -0,0 +1,58 @@
+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.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestMiMoMTP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "XiaomiMiMo/MiMo-7B-RL"
+        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",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-num-steps",
+                "1",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "2",
+                "--mem-fraction-static",
+                "0.5",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.7)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/models/test_unsloth_models.py

@@ -0,0 +1,213 @@
+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.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestUnslothPhi4(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/phi-4"
+        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=[],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.78)
+
+
+class TestUnslothPhi4Bnb4bit(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/phi-4-bnb-4bit"
+        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=[
+                "--load-format",
+                "bitsandbytes",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.75)
+
+
+class TestUnslothPhi4UnslothBnb4bit(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/phi-4-unsloth-bnb-4bit"
+        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=[
+                "--load-format",
+                "bitsandbytes",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.75)
+
+
+class TestUnslothPhi4MiniInstruct(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/Phi-4-mini-instruct"
+        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=[],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.65)
+
+
+class TestUnslothPhi4MiniBnb4bit(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/Phi-4-mini-instruct-bnb-4bit"
+        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=[
+                "--load-format",
+                "bitsandbytes",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.6)
+
+
+class TestUnslothPhi4MiniUnslothBnb4bit(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "unsloth/Phi-4-mini-instruct-unsloth-bnb-4bit"
+        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=[
+                "--load-format",
+                "bitsandbytes",
+            ],
+        )
+
+    @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",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreater(metrics["score"], 0.6)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_deepseek_v31_perf.py

@@ -0,0 +1,86 @@
+import unittest
+
+from sglang.test.nightly_utils import NightlyBenchmarkRunner
+from sglang.test.test_utils import DEFAULT_URL_FOR_TEST, _parse_int_list_env
+
+DEEPSEEK_V31_MODEL_PATH = "deepseek-ai/DeepSeek-V3.1"
+PROFILE_DIR = "performance_profiles_deepseek_v31"
+
+
+class TestNightlyDeepseekV31Performance(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEEPSEEK_V31_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.batch_sizes = [1, 1, 8, 16, 64]
+        cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
+        cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
+
+        # Define variant configurations
+        cls.variants = [
+            {
+                "name": "basic",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+            {
+                "name": "mtp",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--speculative-algorithm",
+                    "EAGLE",
+                    "--speculative-num-steps",
+                    "3",
+                    "--speculative-eagle-topk",
+                    "1",
+                    "--speculative-num-draft-tokens",
+                    "4",
+                    "--mem-frac",
+                    "0.7",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+        ]
+
+        cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
+        cls.runner.setup_profile_directory()
+
+    def test_bench_one_batch(self):
+        failed_variants = []
+
+        try:
+            for variant_config in self.variants:
+                with self.subTest(variant=variant_config["name"]):
+                    results, success = self.runner.run_benchmark_for_model(
+                        model_path=self.model,
+                        batch_sizes=self.batch_sizes,
+                        input_lens=self.input_lens,
+                        output_lens=self.output_lens,
+                        other_args=variant_config["other_args"],
+                        variant=variant_config["name"],
+                    )
+
+                    if not success:
+                        failed_variants.append(variant_config["name"])
+
+                    self.runner.add_report(results, variant=variant_config["name"])
+        finally:
+            self.runner.write_final_report()
+
+        if failed_variants:
+            raise AssertionError(
+                f"Benchmark failed for {self.model} with the following variants: "
+                f"{', '.join(failed_variants)}"
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_deepseek_v32_perf.py

@@ -0,0 +1,127 @@
+import unittest
+
+from sglang.test.nightly_utils import NightlyBenchmarkRunner
+from sglang.test.test_utils import DEFAULT_URL_FOR_TEST, _parse_int_list_env
+
+DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2"
+PROFILE_DIR = "performance_profiles_deepseek_v32"
+
+
+class TestNightlyDeepseekV32Performance(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEEPSEEK_V32_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.batch_sizes = [1, 1, 8, 16, 64]
+        cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
+        cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
+
+        # Define variant configurations
+        cls.variants = [
+            {
+                "name": "basic",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--dp",
+                    "8",
+                    "--enable-dp-attention",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+            {
+                "name": "mtp",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--dp",
+                    "8",
+                    "--enable-dp-attention",
+                    "--speculative-algorithm",
+                    "EAGLE",
+                    "--speculative-num-steps",
+                    "3",
+                    "--speculative-eagle-topk",
+                    "1",
+                    "--speculative-num-draft-tokens",
+                    "4",
+                    "--mem-frac",
+                    "0.7",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+            {
+                "name": "nsa",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--dp",
+                    "8",
+                    "--enable-dp-attention",
+                    "--attention-backend",
+                    "nsa",
+                    "--nsa-prefill-backend",
+                    "flashmla_sparse",
+                    "--nsa-decode-backend",
+                    "flashmla_kv",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+            {
+                "name": "pure_tp",
+                "other_args": [
+                    "--trust-remote-code",
+                    "--tp",
+                    "8",
+                    "--attention-backend",
+                    "nsa",
+                    "--nsa-prefill-backend",
+                    "flashmla_sparse",
+                    "--nsa-decode-backend",
+                    "flashmla_kv",
+                    "--model-loader-extra-config",
+                    '{"enable_multithread_load": true}',
+                ],
+            },
+        ]
+
+        cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
+        cls.runner.setup_profile_directory()
+
+    def test_bench_one_batch(self):
+        failed_variants = []
+
+        try:
+            for variant_config in self.variants:
+                with self.subTest(variant=variant_config["name"]):
+                    results, success = self.runner.run_benchmark_for_model(
+                        model_path=self.model,
+                        batch_sizes=self.batch_sizes,
+                        input_lens=self.input_lens,
+                        output_lens=self.output_lens,
+                        other_args=variant_config["other_args"],
+                        variant=variant_config["name"],
+                    )
+
+                    if not success:
+                        failed_variants.append(variant_config["name"])
+
+                    self.runner.add_report(results, variant=variant_config["name"])
+        finally:
+            self.runner.write_final_report()
+
+        if failed_variants:
+            raise AssertionError(
+                f"Benchmark failed for {self.model} with the following variants: "
+                f"{', '.join(failed_variants)}"
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_text_models_gsm8k_eval.py

@@ -0,0 +1,124 @@
+import json
+import unittest
+import warnings
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1,
+    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2,
+    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1,
+    DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    ModelLaunchSettings,
+    check_evaluation_test_results,
+    parse_models,
+    popen_launch_server,
+    write_results_to_json,
+)
+
+MODEL_SCORE_THRESHOLDS = {
+    "meta-llama/Llama-3.1-8B-Instruct": 0.82,
+    "mistralai/Mistral-7B-Instruct-v0.3": 0.58,
+    "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": 0.85,
+    "google/gemma-2-27b-it": 0.91,
+    "meta-llama/Llama-3.1-70B-Instruct": 0.95,
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": 0.616,
+    "Qwen/Qwen2-57B-A14B-Instruct": 0.86,
+    "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8": 0.83,
+    "neuralmagic/Mistral-7B-Instruct-v0.3-FP8": 0.54,
+    "neuralmagic/DeepSeek-Coder-V2-Lite-Instruct-FP8": 0.835,
+    "zai-org/GLM-4.5-Air-FP8": 0.75,
+    # The threshold of neuralmagic/gemma-2-2b-it-FP8 should be 0.6, but this model has some accuracy regression.
+    # The fix is tracked at https://github.com/sgl-project/sglang/issues/4324, we set it to 0.50, for now, to make CI green.
+    "neuralmagic/gemma-2-2b-it-FP8": 0.50,
+    "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8": 0.94,
+    "neuralmagic/Mixtral-8x7B-Instruct-v0.1-FP8": 0.65,
+    "neuralmagic/Qwen2-72B-Instruct-FP8": 0.94,
+    "neuralmagic/Qwen2-57B-A14B-Instruct-FP8": 0.82,
+}
+
+
+# Do not use `CustomTestCase` since `test_mgsm_en_all_models` does not want retry
+class TestNightlyGsm8KEval(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.models = []
+        models_tp1 = parse_models(
+            DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1
+        ) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1)
+        for model_path in models_tp1:
+            cls.models.append(ModelLaunchSettings(model_path, tp_size=1))
+
+        models_tp2 = parse_models(
+            DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2
+        ) + parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2)
+        for model_path in models_tp2:
+            cls.models.append(ModelLaunchSettings(model_path, tp_size=2))
+
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+    def test_mgsm_en_all_models(self):
+        warnings.filterwarnings(
+            "ignore", category=ResourceWarning, message="unclosed.*socket"
+        )
+        is_first = True
+        all_results = []
+        for model_setup in self.models:
+            with self.subTest(model=model_setup.model_path):
+                other_args = list(model_setup.extra_args)
+
+                if model_setup.model_path == "meta-llama/Llama-3.1-70B-Instruct":
+                    other_args.extend(["--mem-fraction-static", "0.9"])
+
+                process = popen_launch_server(
+                    model=model_setup.model_path,
+                    other_args=other_args,
+                    base_url=self.base_url,
+                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                )
+
+                try:
+                    args = SimpleNamespace(
+                        base_url=self.base_url,
+                        model=model_setup.model_path,
+                        eval_name="mgsm_en",
+                        num_examples=None,
+                        num_threads=1024,
+                    )
+
+                    metrics = run_eval(args)
+                    print(
+                        f"{'=' * 42}\n{model_setup.model_path} - metrics={metrics} score={metrics['score']}\n{'=' * 42}\n"
+                    )
+
+                    write_results_to_json(
+                        model_setup.model_path, metrics, "w" if is_first else "a"
+                    )
+                    is_first = False
+
+                    # 0.0 for empty latency
+                    all_results.append((model_setup.model_path, metrics["score"], 0.0))
+                finally:
+                    kill_process_tree(process.pid)
+
+        try:
+            with open("results.json", "r") as f:
+                print("\nFinal Results from results.json:")
+                print(json.dumps(json.load(f), indent=2))
+        except Exception as e:
+            print(f"Error reading results.json: {e}")
+
+        # Check all scores after collecting all results
+        check_evaluation_test_results(
+            all_results,
+            self.__class__.__name__,
+            model_accuracy_thresholds=MODEL_SCORE_THRESHOLDS,
+            model_count=len(self.models),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_text_models_perf.py

@@ -0,0 +1,59 @@
+import unittest
+
+from sglang.test.nightly_utils import NightlyBenchmarkRunner
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    ModelLaunchSettings,
+    _parse_int_list_env,
+    parse_models,
+)
+
+PROFILE_DIR = "performance_profiles_text_models"
+
+
+class TestNightlyTextModelsPerformance(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.models = []
+        # TODO: replace with DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1 or other model lists
+        for model_path in parse_models("meta-llama/Llama-3.1-8B-Instruct"):
+            cls.models.append(ModelLaunchSettings(model_path, tp_size=1))
+        for model_path in parse_models("Qwen/Qwen2-57B-A14B-Instruct"):
+            cls.models.append(ModelLaunchSettings(model_path, tp_size=2))
+        # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP1), False, False),
+        # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_TP2), False, True),
+        # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP1), True, False),
+        # (parse_models(DEFAULT_MODEL_NAME_FOR_NIGHTLY_EVAL_FP8_TP2), True, True),
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.batch_sizes = [1, 1, 8, 16, 64]
+        cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_INPUT_LENS", "4096"))
+        cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_OUTPUT_LENS", "512"))
+        cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
+        cls.runner.setup_profile_directory()
+
+    def test_bench_one_batch(self):
+        all_model_succeed = True
+
+        for model_setup in self.models:
+            with self.subTest(model=model_setup.model_path):
+                results, success = self.runner.run_benchmark_for_model(
+                    model_path=model_setup.model_path,
+                    batch_sizes=self.batch_sizes,
+                    input_lens=self.input_lens,
+                    output_lens=self.output_lens,
+                    other_args=model_setup.extra_args,
+                )
+
+                if not success:
+                    all_model_succeed = False
+
+                self.runner.add_report(results)
+
+        self.runner.write_final_report()
+
+        if not all_model_succeed:
+            raise AssertionError("Some models failed the perf tests.")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_vlms_mmmu_eval.py

@@ -0,0 +1,127 @@
+import json
+import unittest
+import warnings
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    ModelEvalMetrics,
+    ModelLaunchSettings,
+    check_evaluation_test_results,
+    popen_launch_server,
+    write_results_to_json,
+)
+
+MODEL_THRESHOLDS = {
+    # Conservative thresholds on 100 MMMU samples, especially for latency thresholds
+    ModelLaunchSettings("deepseek-ai/deepseek-vl2-small"): ModelEvalMetrics(
+        0.330, 56.1
+    ),
+    ModelLaunchSettings("deepseek-ai/Janus-Pro-7B"): ModelEvalMetrics(0.285, 40.3),
+    ModelLaunchSettings("Efficient-Large-Model/NVILA-8B-hf"): ModelEvalMetrics(
+        0.270, 56.7
+    ),
+    ModelLaunchSettings("Efficient-Large-Model/NVILA-Lite-2B-hf"): ModelEvalMetrics(
+        0.270, 23.8
+    ),
+    ModelLaunchSettings("google/gemma-3-4b-it"): ModelEvalMetrics(0.360, 10.9),
+    ModelLaunchSettings("google/gemma-3n-E4B-it"): ModelEvalMetrics(0.360, 17.7),
+    ModelLaunchSettings("mistral-community/pixtral-12b"): ModelEvalMetrics(0.360, 16.6),
+    ModelLaunchSettings("moonshotai/Kimi-VL-A3B-Instruct"): ModelEvalMetrics(
+        0.330, 22.3
+    ),
+    ModelLaunchSettings("openbmb/MiniCPM-o-2_6"): ModelEvalMetrics(0.330, 29.3),
+    ModelLaunchSettings("openbmb/MiniCPM-v-2_6"): ModelEvalMetrics(0.259, 36.3),
+    ModelLaunchSettings("OpenGVLab/InternVL2_5-2B"): ModelEvalMetrics(0.300, 17.0),
+    ModelLaunchSettings("Qwen/Qwen2-VL-7B-Instruct"): ModelEvalMetrics(0.310, 83.3),
+    ModelLaunchSettings("Qwen/Qwen2.5-VL-7B-Instruct"): ModelEvalMetrics(0.340, 31.9),
+    ModelLaunchSettings(
+        "Qwen/Qwen3-VL-30B-A3B-Instruct", extra_args=["--tp=2"]
+    ): ModelEvalMetrics(0.29, 37.0),
+    ModelLaunchSettings(
+        "unsloth/Mistral-Small-3.1-24B-Instruct-2503"
+    ): ModelEvalMetrics(0.310, 16.7),
+    ModelLaunchSettings("XiaomiMiMo/MiMo-VL-7B-RL"): ModelEvalMetrics(0.28, 32.0),
+    ModelLaunchSettings("zai-org/GLM-4.1V-9B-Thinking"): ModelEvalMetrics(0.280, 30.4),
+}
+
+
+class TestNightlyVLMMmmuEval(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.models = list(MODEL_THRESHOLDS.keys())
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+    def test_mmmu_vlm_models(self):
+        warnings.filterwarnings(
+            "ignore", category=ResourceWarning, message="unclosed.*socket"
+        )
+        is_first = True
+        all_results = []
+
+        for model in self.models:
+            model_path = model.model_path
+            with self.subTest(model=model_path):
+                process = popen_launch_server(
+                    model=model_path,
+                    base_url=self.base_url,
+                    other_args=model.extra_args,
+                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                )
+                try:
+                    args = SimpleNamespace(
+                        base_url=self.base_url,
+                        model=model_path,
+                        eval_name="mmmu",
+                        num_examples=100,
+                        num_threads=64,
+                        max_tokens=30,
+                    )
+
+                    args.return_latency = True
+
+                    metrics, latency = run_eval(args)
+
+                    metrics["score"] = round(metrics["score"], 4)
+                    metrics["latency"] = round(latency, 4)
+                    print(
+                        f"{'=' * 42}\n{model_path} - metrics={metrics} score={metrics['score']}\n{'=' * 42}\n"
+                    )
+
+                    write_results_to_json(model_path, metrics, "w" if is_first else "a")
+                    is_first = False
+
+                    all_results.append(
+                        (model_path, metrics["score"], metrics["latency"])
+                    )
+                finally:
+                    kill_process_tree(process.pid)
+
+        try:
+            with open("results.json", "r") as f:
+                print("\nFinal Results from results.json:")
+                print(json.dumps(json.load(f), indent=2))
+        except Exception as e:
+            print(f"Error reading results: {e}")
+
+        model_accuracy_thresholds = {
+            model.model_path: threshold.accuracy
+            for model, threshold in MODEL_THRESHOLDS.items()
+        }
+        model_latency_thresholds = {
+            model.model_path: threshold.eval_time
+            for model, threshold in MODEL_THRESHOLDS.items()
+        }
+        check_evaluation_test_results(
+            all_results,
+            self.__class__.__name__,
+            model_accuracy_thresholds=model_accuracy_thresholds,
+            model_latency_thresholds=model_latency_thresholds,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_vlms_perf.py

@@ -0,0 +1,87 @@
+import os
+import unittest
+import warnings
+
+from sglang.test.nightly_utils import NightlyBenchmarkRunner
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    ModelLaunchSettings,
+    _parse_int_list_env,
+    parse_models,
+)
+
+PROFILE_DIR = "performance_profiles_vlms"
+
+MODEL_DEFAULTS = [
+    # Keep conservative defaults. Can be overridden by env NIGHTLY_VLM_MODELS
+    ModelLaunchSettings(
+        "Qwen/Qwen2.5-VL-7B-Instruct",
+        extra_args=["--mem-fraction-static=0.7"],
+    ),
+    ModelLaunchSettings(
+        "google/gemma-3-27b-it",
+    ),
+    ModelLaunchSettings("Qwen/Qwen3-VL-30B-A3B-Instruct", extra_args=["--tp=2"]),
+    # "OpenGVLab/InternVL2_5-2B",
+    # buggy in official transformers impl
+    # "openbmb/MiniCPM-V-2_6",
+]
+
+
+class TestNightlyVLMModelsPerformance(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        warnings.filterwarnings(
+            "ignore", category=ResourceWarning, message="unclosed.*socket"
+        )
+
+        nightly_vlm_models_str = os.environ.get("NIGHTLY_VLM_MODELS")
+        if nightly_vlm_models_str:
+            cls.models = []
+            model_paths = parse_models(nightly_vlm_models_str)
+            for model_path in model_paths:
+                cls.models.append(ModelLaunchSettings(model_path))
+        else:
+            cls.models = MODEL_DEFAULTS
+
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+        cls.batch_sizes = _parse_int_list_env("NIGHTLY_VLM_BATCH_SIZES", "1,1,2,8,16")
+        cls.input_lens = tuple(_parse_int_list_env("NIGHTLY_VLM_INPUT_LENS", "4096"))
+        cls.output_lens = tuple(_parse_int_list_env("NIGHTLY_VLM_OUTPUT_LENS", "512"))
+        cls.runner = NightlyBenchmarkRunner(PROFILE_DIR, cls.__name__, cls.base_url)
+        cls.runner.setup_profile_directory()
+
+    def test_bench_one_batch(self):
+        all_model_succeed = True
+
+        for model_setup in self.models:
+            with self.subTest(model=model_setup.model_path):
+                # VLMs need additional benchmark args for dataset and trust-remote-code
+                extra_bench_args = [
+                    "--trust-remote-code",
+                    "--dataset-name=mmmu",
+                ]
+
+                results, success = self.runner.run_benchmark_for_model(
+                    model_path=model_setup.model_path,
+                    batch_sizes=self.batch_sizes,
+                    input_lens=self.input_lens,
+                    output_lens=self.output_lens,
+                    other_args=model_setup.extra_args,
+                    extra_bench_args=extra_bench_args,
+                )
+
+                if not success:
+                    all_model_succeed = False
+
+                self.runner.add_report(results)
+
+        self.runner.write_final_report()
+
+        if not all_model_succeed:
+            raise AssertionError("Some models failed the perf tests.")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/nightly/test_vlms_piecewise_cuda_graph.py

@@ -0,0 +1,262 @@
+import argparse
+import glob
+import json
+import os
+import random
+import sys
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+)
+
+MODELS = [
+    SimpleNamespace(model="Qwen/Qwen2.5-VL-7B-Instruct", mmmu_accuracy=0.60),
+]
+
+
+# Set default mem_fraction_static to 0.8
+DEFAULT_MEM_FRACTION_STATIC = 0.8
+
+
+class TestVLMPiecewiseCudaGraph(CustomTestCase):
+    parsed_args = None  # Class variable to store args
+
+    @classmethod
+    def setUpClass(cls):
+        # Removed argument parsing from here
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+
+        if cls.parsed_args is None:
+            cls.parsed_args = SimpleNamespace(
+                mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
+            )
+
+        # Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
+        os.environ["OPENAI_API_KEY"] = cls.api_key
+        os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
+
+    def run_mmmu_eval(
+        self,
+        model_version: str,
+        output_path: str,
+        *,
+        env: dict | None = None,
+    ):
+        """
+        Evaluate a VLM on the MMMU validation set with lmms‑eval.
+        Only `model_version` (checkpoint) and `chat_template` vary;
+        We are focusing only on the validation set due to resource constraints.
+        """
+        # -------- fixed settings --------
+        model = "openai_compatible"
+        tp = 1
+        tasks = "mmmu_val"
+        batch_size = 32
+        log_suffix = "openai_compatible"
+        os.makedirs(output_path, exist_ok=True)
+
+        # -------- compose --model_args --------
+        model_args = f'model_version="{model_version}",' f"tp={tp}"
+
+        # -------- build command list --------
+        cmd = [
+            "python3",
+            "-m",
+            "lmms_eval",
+            "--model",
+            model,
+            "--model_args",
+            model_args,
+            "--tasks",
+            tasks,
+            "--batch_size",
+            str(batch_size),
+            "--output_path",
+            str(output_path),
+        ]
+
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
+
+    def _run_vlm_mmmu_test(
+        self,
+        model,
+        output_path,
+        test_name="",
+        custom_env=None,
+        log_level="info",
+        capture_output=False,
+    ):
+        """
+        Common method to run VLM MMMU benchmark test.
+        Args:
+            model: Model to test
+            output_path: Path for output logs
+            test_name: Optional test name for logging
+            custom_env: Optional custom environment variables
+            log_level: Log level for server (default: "info")
+            capture_output: Whether to capture server stdout/stderr
+        """
+        print(f"\nTesting model: {model.model}{test_name}")
+
+        process = None
+        mmmu_accuracy = 0  # Initialize to handle potential exceptions
+        server_output = ""
+
+        try:
+            # Prepare environment variables
+            process_env = os.environ.copy()
+            if custom_env:
+                process_env.update(custom_env)
+            # if test vlm with cuda_ipc feature, open this env_var
+            process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
+
+            # Prepare stdout/stderr redirection if needed
+            stdout_file = None
+            stderr_file = None
+            if capture_output:
+                stdout_file = open("/tmp/server_stdout.log", "w")
+                stderr_file = open("/tmp/server_stderr.log", "w")
+
+            # Launch server for testing
+            process = popen_launch_server(
+                model.model,
+                base_url=self.base_url,
+                timeout=self.time_out,
+                api_key=self.api_key,
+                other_args=[
+                    "--trust-remote-code",
+                    "--piecewise-cuda-graph-max-tokens",
+                    "8192",
+                    "--enforce-piecewise-cuda-graph",
+                    "--tp=8",
+                    "--piecewise-cuda-graph-compiler=eager",
+                    "--disable-radix-cache",
+                    "--log-level",
+                    log_level,
+                ],
+                env=process_env,
+                return_stdout_stderr=(
+                    (stdout_file, stderr_file) if capture_output else None
+                ),
+            )
+
+            # Run evaluation
+            self.run_mmmu_eval(model.model, output_path)
+
+            # Get the result file
+            # Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
+            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
+            if not result_files:
+                result_files = glob.glob(f"{output_path}/*.json")
+
+            if not result_files:
+                raise FileNotFoundError(f"No JSON result files found in {output_path}")
+
+            result_file_path = result_files[0]
+
+            with open(result_file_path, "r") as f:
+                result = json.load(f)
+                print(f"Result{test_name}\n: {result}")
+
+            # Process the result
+            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
+            print(
+                f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
+            )
+
+            # Capture server output if requested
+            if capture_output and process:
+                server_output = self._read_output_from_files()
+
+            # Assert performance meets expected threshold
+            self.assertGreaterEqual(
+                mmmu_accuracy,
+                model.mmmu_accuracy,
+                f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
+            )
+
+            return server_output
+
+        except Exception as e:
+            print(f"Error testing {model.model}{test_name}: {e}")
+            self.fail(f"Test failed for {model.model}{test_name}: {e}")
+
+        finally:
+            # Ensure process cleanup happens regardless of success/failure
+            if process is not None and process.poll() is None:
+                print(f"Cleaning up process {process.pid}")
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process: {e}")
+
+            # clean up temporary files
+            if capture_output:
+                if stdout_file:
+                    stdout_file.close()
+                if stderr_file:
+                    stderr_file.close()
+                for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
+                    try:
+                        if os.path.exists(filename):
+                            os.remove(filename)
+                    except Exception as e:
+                        print(f"Error removing {filename}: {e}")
+
+    def _read_output_from_files(self):
+        output_lines = []
+
+        log_files = [
+            ("/tmp/server_stdout.log", "[STDOUT]"),
+            ("/tmp/server_stderr.log", "[STDERR]"),
+        ]
+        for filename, tag in log_files:
+            try:
+                if os.path.exists(filename):
+                    with open(filename, "r") as f:
+                        for line in f:
+                            output_lines.append(f"{tag} {line.rstrip()}")
+            except Exception as e:
+                print(f"Error reading {tag.lower()} file: {e}")
+
+        return "\n".join(output_lines)
+
+    def test_vlm_mmmu_benchmark(self):
+        """Test VLM models against MMMU benchmark."""
+        models_to_test = MODELS
+
+        if is_in_ci():
+            models_to_test = [random.choice(MODELS)]
+
+        for model in models_to_test:
+            self._run_vlm_mmmu_test(model, "./logs")
+
+
+if __name__ == "__main__":
+    # Define and parse arguments here, before unittest.main
+    parser = argparse.ArgumentParser(description="Test VLM models")
+    parser.add_argument(
+        "--mem-fraction-static",
+        type=float,
+        help="Static memory fraction for the model",
+        default=DEFAULT_MEM_FRACTION_STATIC,
+    )
+
+    # Parse args intended for unittest
+    args = parser.parse_args()
+
+    # Store the parsed args object on the class
+    TestVLMPiecewiseCudaGraph.parsed_args = args
+
+    # Pass args to unittest
+    unittest.main(argv=[sys.argv[0]])

third_party/sglang/test/manual/nightly/test_vlms_vit_cuda_graph.py

@@ -0,0 +1,268 @@
+import argparse
+import glob
+import json
+import os
+import random
+import sys
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+)
+
+MODELS = [
+    SimpleNamespace(model="Qwen/Qwen2.5-VL-7B-Instruct", mmmu_accuracy=0.60),
+    SimpleNamespace(model="Qwen/Qwen3-VL-8B-Instruct", mmmu_accuracy=0.60),
+]
+
+
+# Set default mem_fraction_static to 0.8
+DEFAULT_MEM_FRACTION_STATIC = 0.8
+
+
+class TestVLMViTCudaGraph(CustomTestCase):
+    parsed_args = None  # Class variable to store args
+
+    @classmethod
+    def setUpClass(cls):
+        # Removed argument parsing from here
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+        cls.enable_vit_cuda_graph = "1"
+
+        if cls.parsed_args is None:
+            cls.parsed_args = SimpleNamespace(
+                mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
+            )
+
+        # Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
+        os.environ["OPENAI_API_KEY"] = cls.api_key
+        os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
+        os.environ["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = cls.enable_vit_cuda_graph
+
+    def run_mmmu_eval(
+        self,
+        model_version: str,
+        output_path: str,
+        *,
+        env: dict | None = None,
+    ):
+        """
+        Evaluate a VLM on the MMMU validation set with lmms‑eval.
+        Only `model_version` (checkpoint) and `chat_template` vary;
+        We are focusing only on the validation set due to resource constraints.
+        """
+        # -------- fixed settings --------
+        model = "openai_compatible"
+        tp = 1
+        tasks = "mmmu_val"
+        batch_size = 32
+        log_suffix = "openai_compatible"
+        os.makedirs(output_path, exist_ok=True)
+
+        # -------- compose --model_args --------
+        model_args = f'model_version="{model_version}",' f"tp={tp}"
+
+        # -------- build command list --------
+        cmd = [
+            "python3",
+            "-m",
+            "lmms_eval",
+            "--model",
+            model,
+            "--model_args",
+            model_args,
+            "--tasks",
+            tasks,
+            "--batch_size",
+            str(batch_size),
+            "--output_path",
+            str(output_path),
+        ]
+
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
+
+    def _run_vlm_mmmu_test(
+        self,
+        model,
+        output_path,
+        test_name="",
+        custom_env=None,
+        log_level="info",
+        capture_output=False,
+    ):
+        """
+        Common method to run VLM MMMU benchmark test.
+        Args:
+            model: Model to test
+            output_path: Path for output logs
+            test_name: Optional test name for logging
+            custom_env: Optional custom environment variables
+            log_level: Log level for server (default: "info")
+            capture_output: Whether to capture server stdout/stderr
+        """
+        print(f"\nTesting model: {model.model}{test_name}")
+
+        process = None
+        mmmu_accuracy = 0  # Initialize to handle potential exceptions
+        server_output = ""
+
+        try:
+            # Prepare environment variables
+            process_env = os.environ.copy()
+            if custom_env:
+                process_env.update(custom_env)
+            # if test vlm with cuda_ipc feature, open this env_var
+            process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
+            process_env["SGLANG_VIT_ENABLE_CUDA_GRAPH"] = "1"
+
+            # Prepare stdout/stderr redirection if needed
+            stdout_file = None
+            stderr_file = None
+            if capture_output:
+                stdout_file = open("/tmp/server_stdout.log", "w")
+                stderr_file = open("/tmp/server_stderr.log", "w")
+
+            # Launch server for testing
+            process = popen_launch_server(
+                model.model,
+                base_url=self.base_url,
+                timeout=self.time_out,
+                api_key=self.api_key,
+                other_args=[
+                    "--mm-attention-backend",
+                    "fa3",
+                    "--enforce-piecewise-cuda-graph",
+                    "--piecewise-cuda-graph-max-tokens",
+                    "8192",
+                    "--chunked-prefill-size",
+                    "8192",
+                    "--disable-radix-cache",
+                    "--disable-overlap-schedule",
+                    "--piecewise-cuda-graph-compiler",
+                    "eager",
+                ],
+                env=process_env,
+                return_stdout_stderr=(
+                    (stdout_file, stderr_file) if capture_output else None
+                ),
+            )
+
+            # Run evaluation
+            self.run_mmmu_eval(model.model, output_path)
+
+            # Get the result file
+            # Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
+            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
+            if not result_files:
+                result_files = glob.glob(f"{output_path}/*.json")
+
+            if not result_files:
+                raise FileNotFoundError(f"No JSON result files found in {output_path}")
+
+            result_file_path = result_files[0]
+
+            with open(result_file_path, "r") as f:
+                result = json.load(f)
+                print(f"Result{test_name}\n: {result}")
+
+            # Process the result
+            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
+            print(
+                f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
+            )
+
+            # Capture server output if requested
+            if capture_output and process:
+                server_output = self._read_output_from_files()
+
+            # Assert performance meets expected threshold
+            self.assertGreaterEqual(
+                mmmu_accuracy,
+                model.mmmu_accuracy,
+                f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
+            )
+
+            return server_output
+
+        except Exception as e:
+            print(f"Error testing {model.model}{test_name}: {e}")
+            self.fail(f"Test failed for {model.model}{test_name}: {e}")
+
+        finally:
+            # Ensure process cleanup happens regardless of success/failure
+            if process is not None and process.poll() is None:
+                print(f"Cleaning up process {process.pid}")
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process: {e}")
+
+            # clean up temporary files
+            if capture_output:
+                if stdout_file:
+                    stdout_file.close()
+                if stderr_file:
+                    stderr_file.close()
+                for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
+                    try:
+                        if os.path.exists(filename):
+                            os.remove(filename)
+                    except Exception as e:
+                        print(f"Error removing {filename}: {e}")
+
+    def _read_output_from_files(self):
+        output_lines = []
+
+        log_files = [
+            ("/tmp/server_stdout.log", "[STDOUT]"),
+            ("/tmp/server_stderr.log", "[STDERR]"),
+        ]
+        for filename, tag in log_files:
+            try:
+                if os.path.exists(filename):
+                    with open(filename, "r") as f:
+                        for line in f:
+                            output_lines.append(f"{tag} {line.rstrip()}")
+            except Exception as e:
+                print(f"Error reading {tag.lower()} file: {e}")
+
+        return "\n".join(output_lines)
+
+    def test_vlm_mmmu_benchmark(self):
+        """Test VLM models against MMMU benchmark."""
+        models_to_test = MODELS
+
+        if is_in_ci():
+            models_to_test = [random.choice(MODELS)]
+
+        for model in models_to_test:
+            self._run_vlm_mmmu_test(model, "./logs")
+
+
+if __name__ == "__main__":
+    # Define and parse arguments here, before unittest.main
+    parser = argparse.ArgumentParser(description="Test VLM models")
+    parser.add_argument(
+        "--mem-fraction-static",
+        type=float,
+        help="Static memory fraction for the model",
+        default=DEFAULT_MEM_FRACTION_STATIC,
+    )
+
+    # Parse args intended for unittest
+    args = parser.parse_args()
+
+    # Store the parsed args object on the class
+    TestVLMViTCudaGraph.parsed_args = args
+
+    # Pass args to unittest
+    unittest.main(argv=[sys.argv[0]])

third_party/sglang/test/manual/nightly/test_vlms_vit_flashinfer_cudnn.py

@@ -0,0 +1,258 @@
+import argparse
+import glob
+import json
+import os
+import random
+import sys
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.mmmu_vlm_kit import _run_lmms_eval_with_retry
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+)
+
+MODELS = [
+    SimpleNamespace(model="Qwen/Qwen3-VL-30B-A3B-Instruct", mmmu_accuracy=0.51),
+]
+
+
+# Set default mem_fraction_static to 0.8
+DEFAULT_MEM_FRACTION_STATIC = 0.8
+
+
+class TestVLMViTFlashinferCudnn(CustomTestCase):
+    parsed_args = None  # Class variable to store args
+
+    @classmethod
+    def setUpClass(cls):
+        # Removed argument parsing from here
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.time_out = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+
+        if cls.parsed_args is None:
+            cls.parsed_args = SimpleNamespace(
+                mem_fraction_static=DEFAULT_MEM_FRACTION_STATIC
+            )
+
+        # Set OpenAI API key and base URL environment variables. Needed for lmm-evals to work.
+        os.environ["OPENAI_API_KEY"] = cls.api_key
+        os.environ["OPENAI_API_BASE"] = f"{cls.base_url}/v1"
+
+    def run_mmmu_eval(
+        self,
+        model_version: str,
+        output_path: str,
+        *,
+        env: dict | None = None,
+    ):
+        """
+        Evaluate a VLM on the MMMU validation set with lmms‑eval.
+        Only `model_version` (checkpoint) and `chat_template` vary;
+        We are focusing only on the validation set due to resource constraints.
+        """
+        # -------- fixed settings --------
+        model = "openai_compatible"
+        tp = 1
+        tasks = "mmmu_val"
+        batch_size = 32
+        log_suffix = "openai_compatible"
+        os.makedirs(output_path, exist_ok=True)
+
+        # -------- compose --model_args --------
+        model_args = f'model_version="{model_version}",' f"tp={tp}"
+
+        # -------- build command list --------
+        cmd = [
+            "python3",
+            "-m",
+            "lmms_eval",
+            "--model",
+            model,
+            "--model_args",
+            model_args,
+            "--tasks",
+            tasks,
+            "--batch_size",
+            str(batch_size),
+            "--output_path",
+            str(output_path),
+        ]
+
+        _run_lmms_eval_with_retry(cmd, timeout=3600)
+
+    def _run_vlm_mmmu_test(
+        self,
+        model,
+        output_path,
+        test_name="",
+        custom_env=None,
+        log_level="info",
+        capture_output=False,
+    ):
+        """
+        Common method to run VLM MMMU benchmark test.
+        Args:
+            model: Model to test
+            output_path: Path for output logs
+            test_name: Optional test name for logging
+            custom_env: Optional custom environment variables
+            log_level: Log level for server (default: "info")
+            capture_output: Whether to capture server stdout/stderr
+        """
+        print(f"\nTesting model: {model.model}{test_name}")
+
+        process = None
+        mmmu_accuracy = 0  # Initialize to handle potential exceptions
+        server_output = ""
+
+        try:
+            # Prepare environment variables
+            process_env = os.environ.copy()
+            if custom_env:
+                process_env.update(custom_env)
+            # if test vlm with cuda_ipc feature, open this env_var
+            process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1"
+
+            # Prepare stdout/stderr redirection if needed
+            stdout_file = None
+            stderr_file = None
+            if capture_output:
+                stdout_file = open("/tmp/server_stdout.log", "w")
+                stderr_file = open("/tmp/server_stderr.log", "w")
+
+            # Launch server for testing
+            process = popen_launch_server(
+                model.model,
+                base_url=self.base_url,
+                timeout=self.time_out,
+                api_key=self.api_key,
+                other_args=[
+                    "--mm-attention-backend",
+                    "flashinfer_cudnn",
+                    "--chunked-prefill-size",
+                    "8192",
+                    "--disable-radix-cache",
+                ],
+                env=process_env,
+                return_stdout_stderr=(
+                    (stdout_file, stderr_file) if capture_output else None
+                ),
+            )
+
+            # Run evaluation
+            self.run_mmmu_eval(model.model, output_path)
+
+            # Get the result file
+            # Search recursively for JSON result files (lmms-eval v0.4.1+ creates subdirectories)
+            result_files = glob.glob(f"{output_path}/**/*.json", recursive=True)
+            if not result_files:
+                result_files = glob.glob(f"{output_path}/*.json")
+
+            if not result_files:
+                raise FileNotFoundError(f"No JSON result files found in {output_path}")
+
+            result_file_path = result_files[0]
+
+            with open(result_file_path, "r") as f:
+                result = json.load(f)
+                print(f"Result{test_name}\n: {result}")
+
+            # Process the result
+            mmmu_accuracy = result["results"]["mmmu_val"]["mmmu_acc,none"]
+            print(
+                f"Model {model.model} achieved accuracy{test_name}: {mmmu_accuracy:.4f}"
+            )
+
+            # Capture server output if requested
+            if capture_output and process:
+                server_output = self._read_output_from_files()
+
+            # Assert performance meets expected threshold
+            self.assertGreaterEqual(
+                mmmu_accuracy,
+                model.mmmu_accuracy,
+                f"Model {model.model} accuracy ({mmmu_accuracy:.4f}) below expected threshold ({model.mmmu_accuracy:.4f}){test_name}",
+            )
+
+            return server_output
+
+        except Exception as e:
+            print(f"Error testing {model.model}{test_name}: {e}")
+            self.fail(f"Test failed for {model.model}{test_name}: {e}")
+
+        finally:
+            # Ensure process cleanup happens regardless of success/failure
+            if process is not None and process.poll() is None:
+                print(f"Cleaning up process {process.pid}")
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process: {e}")
+
+            # clean up temporary files
+            if capture_output:
+                if stdout_file:
+                    stdout_file.close()
+                if stderr_file:
+                    stderr_file.close()
+                for filename in ["/tmp/server_stdout.log", "/tmp/server_stderr.log"]:
+                    try:
+                        if os.path.exists(filename):
+                            os.remove(filename)
+                    except Exception as e:
+                        print(f"Error removing {filename}: {e}")
+
+    def _read_output_from_files(self):
+        output_lines = []
+
+        log_files = [
+            ("/tmp/server_stdout.log", "[STDOUT]"),
+            ("/tmp/server_stderr.log", "[STDERR]"),
+        ]
+        for filename, tag in log_files:
+            try:
+                if os.path.exists(filename):
+                    with open(filename, "r") as f:
+                        for line in f:
+                            output_lines.append(f"{tag} {line.rstrip()}")
+            except Exception as e:
+                print(f"Error reading {tag.lower()} file: {e}")
+
+        return "\n".join(output_lines)
+
+    def test_vlm_mmmu_benchmark(self):
+        """Test VLM models against MMMU benchmark."""
+        models_to_test = MODELS
+
+        if is_in_ci():
+            models_to_test = [random.choice(MODELS)]
+
+        for model in models_to_test:
+            self._run_vlm_mmmu_test(model, "./logs")
+
+
+if __name__ == "__main__":
+    # Define and parse arguments here, before unittest.main
+    parser = argparse.ArgumentParser(description="Test VLM models")
+    parser.add_argument(
+        "--mem-fraction-static",
+        type=float,
+        help="Static memory fraction for the model",
+        default=DEFAULT_MEM_FRACTION_STATIC,
+    )
+
+    # Parse args intended for unittest
+    args = parser.parse_args()
+
+    # Store the parsed args object on the class
+    TestVLMViTFlashinferCudnn.parsed_args = args
+
+    # Pass args to unittest
+    unittest.main(argv=[sys.argv[0]])

third_party/sglang/test/manual/openai_server/features/test_cache_report.py

@@ -0,0 +1,289 @@
+import unittest
+
+import openai
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestCacheReport(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.min_cached = 5
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=300,
+            other_args=[
+                "--chunked-prefill-size=40",
+                "--enable-cache-report",
+            ],
+        )
+        cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+        cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+
+        usage = cls.run_openai(cls, "1").usage
+        # we can assume that our request is of size 1, plus the total template size
+        # ideally we would like to know the begin size / end size of the template to be more precise
+        total_template_size = usage.prompt_tokens - 1
+        print(f"template size: {total_template_size}")
+        usage2 = cls.run_openai(cls, "2").usage
+        assert usage2.prompt_tokens_details.cached_tokens <= total_template_size
+        cls.min_cached = max(
+            usage2.prompt_tokens_details.cached_tokens,
+            total_template_size - usage2.prompt_tokens_details.cached_tokens,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_decode(self, return_logprob=False, top_logprobs_num=0, n=1):
+        response = requests.post(
+            self.base_url + "/generate",
+            # we use an uncommon start to minimise the chance that the cache is hit by chance
+            json={
+                "text": "_ The capital of France is",
+                "sampling_params": {
+                    "temperature": 0 if n == 1 else 0.5,
+                    "max_new_tokens": 128,
+                    "n": n,
+                    "stop_token_ids": [119690],
+                },
+                "stream": False,
+                "return_logprob": return_logprob,
+                "top_logprobs_num": top_logprobs_num,
+                "logprob_start_len": 0,
+            },
+        )
+        return response
+
+    def run_openai(self, message):
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=[
+                # {"role": "system", "content": "You are a helpful AI assistant"},
+                {"role": "user", "content": message},
+            ],
+            temperature=0,
+            max_tokens=100,
+        )
+        return response
+
+    async def run_openai_async(self, message):
+        response = await self.aclient.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "user", "content": message},
+            ],
+            temperature=0,
+            max_tokens=100,
+        )
+        return response
+
+    def cache_report_openai(self, message):
+        response = self.run_openai(message)
+        print(
+            f"openai first request cached_tokens: {int(response.usage.prompt_tokens_details.cached_tokens)}"
+        )
+        first_cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
+        # assert int(response.usage.cached_tokens) == 0
+        assert first_cached_tokens <= self.min_cached
+        response = self.run_openai(message)
+        cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
+        print(f"openai second request cached_tokens: {cached_tokens}")
+        assert cached_tokens > 0
+        assert cached_tokens == int(response.usage.prompt_tokens) - 1
+        return first_cached_tokens
+
+    async def cache_report_openai_async(self, message):
+        response = await self.run_openai_async(message)
+        cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
+        prompt_tokens = int(response.usage.prompt_tokens)
+        return cached_tokens, prompt_tokens
+
+    def test_generate(self):
+        print("=" * 100)
+        response = self.run_decode()
+        # print(response.json())
+        cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
+        print(f"sglang first request cached_tokens: {cached_tokens}")
+        print(
+            f"sglang first request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
+        )
+        # can't assure to be 0: depends on the initialisation request / if a template is used with the model
+        assert cached_tokens < self.min_cached
+        response = self.run_decode()
+        cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
+        print(f"sglang second request cached_tokens: {cached_tokens}")
+        print(
+            f"sglang second request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
+        )
+        assert cached_tokens == int(response.json()["meta_info"]["prompt_tokens"]) - 1
+
+    def test_cache_split_prefill_openai(self):
+        print("=" * 100)
+        self.cache_report_openai(
+            "€ This is a very long and unique text that should not be already cached, the twist is"
+            " that it should be longer than the chunked-prefill-size, so it should be split among"
+            " several prefill requests. Still, it shouldn't be cached"
+        )
+
+    def test_cache_report_openai(self):
+        print("=" * 100)
+        # warm up the cache, for the template
+        self.run_openai("Introduce the capital of France.")
+
+        first_cached_tokens_1 = self.run_openai(
+            "How many sparrow do you need to lift a coconut?"
+        ).usage.prompt_tokens_details.cached_tokens
+
+        usage_2 = self.run_openai("* sing something about cats").usage
+        first_cached_tokens_2 = usage_2.prompt_tokens_details.cached_tokens
+        # first request may not have 0 cached tokens, but if they only have the template in common they
+        # should be the same once the cache is warmed up
+        assert first_cached_tokens_1 == first_cached_tokens_2
+
+        resp = self.run_openai("* sing something about cats and dogs")
+        print(resp.usage)
+
+        resp = self.run_openai("* sing something about cats, please")
+        print(resp.usage)
+        assert (
+            resp.usage.prompt_tokens_details.cached_tokens
+            >= usage_2.prompt_tokens - self.min_cached
+        )
+
+    # TODO: flaky test
+    # def test_cache_report_openai_async(self):
+    #     print("=" * 100)
+
+    #     async def run_test():
+    #         task0 = asyncio.create_task(
+    #             self.cache_report_openai_async(
+    #                 "first request, to start the inference and let the next two request be started in the same batch"
+    #             )
+    #         )
+    #         await asyncio.sleep(1)  # to force the first request to be started first
+    #         task1 = asyncio.create_task(
+    #             self.cache_report_openai_async(
+    #                 "> can the same batch parallel request use the cache?"
+    #             )
+    #         )
+    #         task2 = asyncio.create_task(
+    #             self.cache_report_openai_async(
+    #                 "> can the same batch parallel request use the cache?"
+    #             )
+    #         )
+    #         result0, result1, result2 = await asyncio.gather(task0, task1, task2)
+
+    #         cached_tokens0, prompt_tokens0 = result0
+    #         cached_tokens1, prompt_tokens1 = result1
+    #         cached_tokens2, prompt_tokens2 = result2
+
+    #         print(
+    #             f"Async request 0 - Cached tokens: {cached_tokens0}, Prompt tokens: {prompt_tokens0}"
+    #         )
+    #         print(
+    #             f"Async request 1 - Cached tokens: {cached_tokens1}, Prompt tokens: {prompt_tokens1}"
+    #         )
+    #         print(
+    #             f"Async request 2 - Cached tokens: {cached_tokens2}, Prompt tokens: {prompt_tokens2}"
+    #         )
+
+    #         # Assert that no requests used the cache (because first is alone, and the next two are in the same batch)
+    #         # If a new optimisation limiting starting request with same prefix at the same time was added
+    #         # to maximise the cache hit, this would not be true
+    #         assert cached_tokens1 == cached_tokens2 == cached_tokens0
+
+    #     asyncio.run(run_test())
+
+    def test_cache_salt_effectiveness(self):
+        print("=" * 100)
+        print("Testing cache_salt effectiveness")
+
+        # Use a unique message to avoid interference with other tests
+        test_message = "What is the capital of Japan?"
+
+        # First request with cache_salt "salt1"
+        response1 = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": test_message}],
+            temperature=0,
+            max_tokens=10,
+            extra_body={"cache_salt": "salt1"},
+        )
+        cached_tokens_1_first = int(response1.usage.prompt_tokens_details.cached_tokens)
+        prompt_tokens_1 = int(response1.usage.prompt_tokens)
+        print(
+            f"First request with salt1 - cached_tokens: {cached_tokens_1_first}, prompt_tokens: {prompt_tokens_1}"
+        )
+
+        # Second request with same cache_salt "salt1" - should get cache hit
+        response2 = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": test_message}],
+            temperature=0,
+            max_tokens=10,
+            extra_body={"cache_salt": "salt1"},
+        )
+        cached_tokens_1_second = int(
+            response2.usage.prompt_tokens_details.cached_tokens
+        )
+        print(
+            f"Second request with salt1 - cached_tokens: {cached_tokens_1_second}, prompt_tokens: {prompt_tokens_1}"
+        )
+
+        # Verify cache hit for same salt
+        assert (
+            cached_tokens_1_second > cached_tokens_1_first
+        ), "Should have cache hit with same cache_salt"
+        assert (
+            cached_tokens_1_second == prompt_tokens_1 - 1
+        ), "Should cache all prompt tokens except the last one"
+
+        # Third request with different cache_salt "salt2" - should not get cache hit
+        response3 = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": test_message}],
+            temperature=0,
+            max_tokens=10,
+            extra_body={"cache_salt": "salt2"},
+        )
+        cached_tokens_2_first = int(response3.usage.prompt_tokens_details.cached_tokens)
+        print(f"First request with salt2 - cached_tokens: {cached_tokens_2_first}")
+
+        # Verify no cache hit for different salt (should be similar to first request with salt1)
+        assert (
+            cached_tokens_2_first <= cached_tokens_1_first + self.min_cached
+        ), "Different cache_salt should not share cache"
+
+        # Fourth request with same cache_salt "salt2" - should now get cache hit
+        response4 = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": test_message}],
+            temperature=0,
+            max_tokens=10,
+            extra_body={"cache_salt": "salt2"},
+        )
+        cached_tokens_2_second = int(
+            response4.usage.prompt_tokens_details.cached_tokens
+        )
+        print(f"Second request with salt2 - cached_tokens: {cached_tokens_2_second}")
+
+        # Verify cache hit for salt2
+        assert (
+            cached_tokens_2_second == cached_tokens_2_first
+        ), "Should have cache hit with same cache_salt for salt2"
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/openai_server/features/test_continuous_usage_stats.py

@@ -0,0 +1,105 @@
+import asyncio
+import unittest
+
+import openai
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestContinuousUsageStats(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(cls.model, cls.base_url, timeout=300)
+        cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+        cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_continuous_usage_stats_enabled(self):
+        stream = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is machine learning?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": True},
+        )
+
+        chunks_with_usage = 0
+        chunks_with_content = 0
+        last_usage = None
+
+        for chunk in stream:
+            has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
+            if chunk.usage:
+                chunks_with_usage += 1
+                last_usage = chunk.usage
+            if has_content:
+                chunks_with_content += 1
+
+        assert chunks_with_content > 0
+        assert chunks_with_usage >= chunks_with_content
+        assert last_usage.prompt_tokens > 0
+        assert last_usage.completion_tokens > 0
+        assert (
+            last_usage.total_tokens
+            == last_usage.prompt_tokens + last_usage.completion_tokens
+        )
+
+    async def test_continuous_usage_stats_async(self):
+        stream = await self.aclient.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is deep learning?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": True},
+        )
+
+        chunks_with_usage = 0
+        chunks_with_content = 0
+
+        async for chunk in stream:
+            has_content = len(chunk.choices) > 0 and chunk.choices[0].delta.content
+            if chunk.usage:
+                chunks_with_usage += 1
+            if has_content:
+                chunks_with_content += 1
+
+        assert chunks_with_content > 0
+        assert chunks_with_usage >= chunks_with_content
+
+    def test_continuous_usage_stats_disabled(self):
+        stream = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": "What is AI?"}],
+            stream=True,
+            max_tokens=30,
+            temperature=0,
+            stream_options={"include_usage": True, "continuous_usage_stats": False},
+        )
+
+        usage_chunks = []
+        for chunk in stream:
+            if chunk.usage:
+                usage_chunks.append(chunk)
+
+        assert len(usage_chunks) == 1
+        assert len(usage_chunks[0].choices) == 0
+
+    def test_async_runner(self):
+        asyncio.run(self.test_continuous_usage_stats_async())
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/openai_server/features/test_structural_tag.py

@@ -0,0 +1,126 @@
+"""
+python3 -m unittest test.srt.openai_server.features.test_structural_tag
+"""
+
+import json
+import unittest
+from typing import Any
+
+import openai
+
+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,
+)
+
+
+def setup_class(cls, backend: str):
+    cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+    cls.base_url = DEFAULT_URL_FOR_TEST
+
+    other_args = [
+        "--max-running-requests",
+        "10",
+        "--grammar-backend",
+        backend,
+    ]
+
+    cls.process = popen_launch_server(
+        cls.model,
+        cls.base_url,
+        timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+        other_args=other_args,
+    )
+
+
+class TestStructuralTagXGrammarBackend(CustomTestCase):
+    model: str
+    base_url: str
+    process: Any
+
+    @classmethod
+    def setUpClass(cls):
+        setup_class(cls, backend="xgrammar")
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_stag_constant_str_openai(self):
+        client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
+
+        # even when the answer is ridiculous, the model should follow the instruction
+        answer = "The capital of France is Berlin."
+
+        response = client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {
+                    "role": "user",
+                    "content": "Introduce the capital of France. Return in a JSON format.",
+                },
+            ],
+            temperature=0,
+            max_tokens=128,
+            response_format={
+                "type": "structural_tag",
+                "format": {
+                    "type": "const_string",
+                    "value": answer,
+                },
+            },
+        )
+
+        text = response.choices[0].message.content
+        self.assertEqual(text, answer)
+
+    def test_stag_json_schema_openai(self):
+        client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1")
+        json_schema = {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "pattern": "^[\\w]+$"},
+                "population": {"type": "integer"},
+            },
+            "required": ["name", "population"],
+            "additionalProperties": False,
+        }
+
+        response = client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {
+                    "role": "user",
+                    "content": "Introduce the capital of France. Return in a JSON format.",
+                },
+            ],
+            temperature=0,
+            max_tokens=128,
+            response_format={
+                "type": "structural_tag",
+                "format": {
+                    "type": "json_schema",
+                    "json_schema": json_schema,
+                },
+            },
+        )
+
+        text = response.choices[0].message.content
+        try:
+            js_obj = json.loads(text)
+        except (TypeError, json.decoder.JSONDecodeError):
+            print("JSONDecodeError", text)
+            raise
+
+        self.assertIsInstance(js_obj["name"], str)
+        self.assertIsInstance(js_obj["population"], int)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/piecewise_cudagraph/test_disaggregation_piecewise_cuda_graph.py

@@ -0,0 +1,88 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.test.run_eval import run_eval
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+
+class TestDisaggregationPiecewiseCudaGraph(PDDisaggregationServerBase):
+    """Test piecewise CUDA graph support in disaggregation prefill server"""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+
+        # Start servers
+        cls.start_prefill()
+        cls.start_decode()
+
+        # Wait for both to be ready
+        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",
+            "--tp",
+            "1",
+            "--enforce-piecewise-cuda-graph",
+        ]
+        prefill_args += 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,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+        ]
+        decode_args += 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,
+        )
+
+    def test_gsm8k_accuracy(self):
+        """Verify that piecewise cuda graph works correctly in prefill server"""
+        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(f"GSM8K accuracy with piecewise cuda graph: {metrics['score']:.3f}")
+
+        self.assertGreater(metrics["score"], 0.62)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/quant/kv_cache_scales_llama3_1_8b.json

@@ -0,0 +1,42 @@
+{
+    "model_type": "llama",
+    "kv_cache": {
+        "dtype": "float8_e4m3fn",
+        "scaling_factor": {
+            "0": {
+                "0": 1,
+                "1": 1,
+                "2": 1,
+                "3": 1,
+                "4": 1,
+                "5": 1,
+                "6": 1,
+                "7": 1,
+                "8": 1,
+                "9": 1,
+                "10": 1,
+                "11": 1,
+                "12": 1,
+                "13": 1,
+                "14": 1,
+                "15": 1,
+                "16": 1,
+                "17": 1,
+                "18": 1,
+                "19": 1,
+                "20": 1,
+                "21": 1,
+                "22": 1,
+                "23": 1,
+                "24": 1,
+                "25": 1,
+                "26": 1,
+                "27": 1,
+                "28": 1,
+                "29": 1,
+                "30": 1,
+                "31": 1
+            }
+        }
+    }
+}

third_party/sglang/test/manual/quant/kv_cache_scales_llama3_8b.json

@@ -0,0 +1,42 @@
+{
+    "model_type": "llama",
+    "kv_cache": {
+        "dtype": "float8_e4m3fn",
+        "scaling_factor": {
+            "0": {
+                "0": 0.0408,
+                "1": 0.0503,
+                "2": 0.0667,
+                "3": 0.0909,
+                "4": 0.1135,
+                "5": 0.127,
+                "6": 0.1768,
+                "7": 0.1488,
+                "8": 0.1135,
+                "9": 0.1203,
+                "10": 0.1013,
+                "11": 0.0842,
+                "12": 0.1231,
+                "13": 0.1096,
+                "14": 0.1221,
+                "15": 0.1013,
+                "16": 0.1067,
+                "17": 0.0952,
+                "18": 0.0899,
+                "19": 0.097,
+                "20": 0.087,
+                "21": 0.0994,
+                "22": 0.0904,
+                "23": 0.1013,
+                "24": 0.1019,
+                "25": 0.1053,
+                "26": 0.1,
+                "27": 0.0894,
+                "28": 0.1013,
+                "29": 0.1488,
+                "30": 0.0766,
+                "31": 0.0821
+            }
+        }
+    }
+}

third_party/sglang/test/manual/quant/kv_cache_scales_qwen2_1_5b.json

@@ -0,0 +1,38 @@
+{
+    "model_type": "qwen",
+    "kv_cache": {
+        "dtype": "float8_e4m3fn",
+        "scaling_factor": {
+            "0": {
+                "0": 0.9846,
+                 "1": 0.0645,
+                 "2": 0.0731,
+                 "3": 0.0800,
+                 "4": 0.0748,
+                 "5": 0.0780,
+                 "6": 0.0702,
+                 "7": 0.0894,
+                 "8": 0.0410,
+                 "9": 0.0758,
+                 "10": 0.0556,
+                 "11": 0.0731,
+                 "12": 0.0899,
+                 "13": 0.0780,
+                 "14": 0.1441,
+                 "15": 0.0914,
+                 "16": 0.5614,
+                 "17": 0.1067,
+                 "18": 0.0537,
+                 "19": 0.0658,
+                 "20": 0.0523,
+                 "21": 0.0533,
+                 "22": 0.0699,
+                 "23": 0.0635,
+                 "24": 0.0588,
+                 "25": 0.0884,
+                 "26": 0.0947,
+                 "27": 0.1032
+            }
+        }
+    }
+}

third_party/sglang/test/manual/quant/test_fp8_kvcache.py

@@ -0,0 +1,114 @@
+import os
+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.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestFp8KvcacheBase(CustomTestCase):
+    model_config = None
+
+    @classmethod
+    def setUpClass(cls):
+        if cls.model_config is None:
+            raise NotImplementedError("model_config must be specified in subclass")
+
+        cls.model = cls.model_config["model_name"]
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        dirpath = os.path.dirname(__file__)
+        config_file = os.path.join(dirpath, cls.model_config["config_filename"])
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--kv-cache-dtype",
+                "fp8_e4m3",
+                "--quantization-param-path",
+                config_file,
+            ],
+        )
+
+
+class TestFp8KvcacheLlama(TestFp8KvcacheBase):
+    model_config = {
+        "model_name": DEFAULT_MODEL_NAME_FOR_TEST,
+        "config_filename": "kv_cache_scales_llama3_8b.json",
+    }
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_mgsm_en(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mgsm_en",
+            num_examples=None,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        self.assertGreater(metrics["score"], 0.80)
+
+    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.assertGreaterEqual(metrics["score"], 0.65)
+
+
+class TestFp8KvcacheQwen(TestFp8KvcacheBase):
+    model_config = {
+        "model_name": DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+        "config_filename": "kv_cache_scales_qwen2_1_5b.json",
+    }
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_mgsm_en(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mgsm_en",
+            num_examples=None,
+            num_threads=1024,
+        )
+
+        metrics = run_eval(args)
+        self.assertGreater(metrics["score"], 0.01)
+
+    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.assertGreaterEqual(metrics["score"], 0.3)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_async_dynamic_batch_tokenizer.py

@@ -0,0 +1,296 @@
+"""
+Unit tests for AsyncDynamicbatchTokenizer.
+
+Tests the async dynamic batching functionality for tokenization,
+including batch efficiency, timeout handling, and error cases.
+"""
+
+import asyncio
+import logging
+import sys
+import time
+from unittest.mock import Mock
+
+import pytest
+from transformers import AutoTokenizer
+
+from sglang.srt.managers.async_dynamic_batch_tokenizer import AsyncDynamicbatchTokenizer
+
+
+class TestAsyncDynamicbatchTokenizer:
+    """Test suite for AsyncDynamicbatchTokenizer."""
+
+    @pytest.fixture
+    def mock_tokenizer(self):
+        """Create a mock tokenizer that behaves like HuggingFace tokenizer."""
+
+        def mock_encode(texts, **kwargs):
+            is_single = isinstance(texts, str)
+            if is_single:
+                texts = [texts]
+
+            # Simulate tokenization - convert text to mock token ids
+            input_ids = []
+            token_type_ids = []
+
+            for text in texts:
+                # Simple mock: text length determines number of tokens
+                tokens = [i for i in range(len(text.split()))]
+                input_ids.append(tokens)
+
+                if kwargs.get("return_token_type_ids", False):
+                    token_type_ids.append([0] * len(tokens))
+
+            result = {"input_ids": input_ids}
+            if kwargs.get("return_token_type_ids", False):
+                result["token_type_ids"] = token_type_ids
+
+            # For single inputs, return individual result (not wrapped in a list)
+            if is_single:
+                result = {"input_ids": input_ids[0]}
+                if kwargs.get("return_token_type_ids", False):
+                    result["token_type_ids"] = token_type_ids[0]
+
+            # Create a proper BatchEncoding-like object that supports dict operations
+            class MockBatchEncoding(dict):
+                def __init__(self, data):
+                    super().__init__(data)
+                    for key, value in data.items():
+                        setattr(self, key, value)
+
+            return MockBatchEncoding(result)
+
+        # Return the function directly - the AsyncDynamicbatchTokenizer will call it
+        return mock_encode
+
+    @pytest.fixture
+    def async_tokenizer(self, mock_tokenizer):
+        """Create AsyncDynamicbatchTokenizer instance."""
+        return AsyncDynamicbatchTokenizer(
+            tokenizer=mock_tokenizer, max_batch_size=4, batch_wait_timeout_s=0.01
+        )
+
+    @pytest.mark.asyncio
+    async def test_single_request(self, async_tokenizer):
+        """Test tokenizing a single request."""
+        text = "hello world"
+        result = await async_tokenizer.encode(text)
+
+        assert "input_ids" in result
+        assert result["input_ids"] == [0, 1]  # 2 words -> 2 tokens
+
+    @pytest.mark.asyncio
+    async def test_single_request_with_token_type_ids(self, async_tokenizer):
+        """Test tokenizing with token type IDs."""
+        text = "hello world"
+        result = await async_tokenizer.encode(text, return_token_type_ids=True)
+
+        assert "input_ids" in result
+        assert "token_type_ids" in result
+        assert result["input_ids"] == [0, 1]
+        assert result["token_type_ids"] == [0, 0]
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests_same_kwargs(self, async_tokenizer):
+        """Test that concurrent requests with same kwargs get batched."""
+        texts = ["hello world", "how are you", "fine thanks", "good morning"]
+
+        # Start all requests concurrently
+        tasks = [async_tokenizer.encode(text) for text in texts]
+        results = await asyncio.gather(*tasks)
+
+        # Verify all results
+        assert len(results) == 4
+        for i, result in enumerate(results):
+            assert "input_ids" in result
+            expected_tokens = list(range(len(texts[i].split())))
+            assert result["input_ids"] == expected_tokens
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests_different_kwargs(self, async_tokenizer):
+        """Test that requests with different kwargs are processed individually."""
+        text1 = "hello world"
+        text2 = "how are you"
+
+        # One with token_type_ids, one without
+        task1 = async_tokenizer.encode(text1, return_token_type_ids=True)
+        task2 = async_tokenizer.encode(text2)
+
+        result1, result2 = await asyncio.gather(task1, task2)
+
+        # First result should have token_type_ids
+        assert "input_ids" in result1
+        assert "token_type_ids" in result1
+        assert result1["input_ids"] == [0, 1]
+        assert result1["token_type_ids"] == [0, 0]
+
+        # Second result should not have token_type_ids
+        assert "input_ids" in result2
+        assert "token_type_ids" not in result2
+        assert result2["input_ids"] == [0, 1, 2]
+
+    @pytest.mark.asyncio
+    async def test_batch_timeout(self, async_tokenizer):
+        """Test that batching respects timeout."""
+        # Send first request
+        task1 = asyncio.create_task(async_tokenizer.encode("hello world"))
+
+        # Wait longer than batch timeout
+        await asyncio.sleep(0.02)  # Longer than 0.01s timeout
+
+        # Send second request
+        task2 = asyncio.create_task(async_tokenizer.encode("how are you"))
+
+        results = await asyncio.gather(task1, task2)
+
+        # Both should complete successfully
+        assert len(results) == 2
+        assert results[0]["input_ids"] == [0, 1]
+        assert results[1]["input_ids"] == [0, 1, 2]
+
+    @pytest.mark.asyncio
+    async def test_max_batch_size_limit(self, async_tokenizer):
+        """Test that batching respects max_batch_size."""
+        # Send more requests than max_batch_size (4)
+        texts = [f"text {i}" for i in range(6)]
+        tasks = [async_tokenizer.encode(text) for text in texts]
+
+        results = await asyncio.gather(*tasks)
+
+        # All should complete successfully
+        assert len(results) == 6
+        for i, result in enumerate(results):
+            assert "input_ids" in result
+            assert result["input_ids"] == [0, 1]  # "text i" -> 2 tokens
+
+    @pytest.mark.asyncio
+    async def test_callable_interface(self, async_tokenizer):
+        """Test that the tokenizer is callable."""
+        text = "hello world"
+        result = await async_tokenizer(text)
+
+        assert "input_ids" in result
+        assert result["input_ids"] == [0, 1]
+
+    @pytest.mark.asyncio
+    async def test_lazy_initialization(self, mock_tokenizer):
+        """Test that initialization happens lazily."""
+        tokenizer = AsyncDynamicbatchTokenizer(mock_tokenizer)
+
+        # Should not be initialized yet
+        assert not tokenizer._initialized
+
+        # First encode should initialize
+        await tokenizer.encode("hello")
+
+        # Should now be initialized
+        assert tokenizer._initialized
+
+    @pytest.mark.asyncio
+    async def test_error_handling_in_tokenizer(self, mock_tokenizer):
+        """Test error handling when tokenizer fails."""
+
+        # Create a new async tokenizer with a failing tokenizer
+        def failing_tokenizer(*args, **kwargs):
+            raise ValueError("Tokenizer error")
+
+        async_tokenizer = AsyncDynamicbatchTokenizer(
+            tokenizer=failing_tokenizer, max_batch_size=4, batch_wait_timeout_s=0.01
+        )
+
+        with pytest.raises(ValueError, match="Tokenizer error"):
+            await async_tokenizer.encode("hello world")
+
+    @pytest.mark.asyncio
+    async def test_batch_processing_logs(self, async_tokenizer, caplog):
+        """Test that batch processing logs are generated."""
+        caplog.set_level(logging.DEBUG)
+
+        # Send multiple requests to trigger batching
+        tasks = [
+            async_tokenizer.encode("hello world"),
+            async_tokenizer.encode("how are you"),
+        ]
+
+        await asyncio.gather(*tasks)
+
+        # Should have batch processing log
+        assert any(
+            "Processing dynamic batch of size" in record.message
+            for record in caplog.records
+        )
+
+    @pytest.mark.asyncio
+    async def test_empty_queue_immediate_processing(self, async_tokenizer):
+        """Test that single requests are processed immediately when queue is empty."""
+        start_time = time.time()
+        result = await async_tokenizer.encode("hello world")
+        end_time = time.time()
+
+        # Should complete quickly (much less than batch timeout)
+        assert end_time - start_time < 0.005  # 5ms should be plenty
+        assert result["input_ids"] == [0, 1]
+
+    @pytest.mark.asyncio
+    async def test_real_tokenizer_integration(self):
+        """Test with a real HuggingFace tokenizer."""
+        try:
+            # Use a small, fast tokenizer for testing
+            real_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+            async_tokenizer = AsyncDynamicbatchTokenizer(
+                tokenizer=real_tokenizer, max_batch_size=2, batch_wait_timeout_s=0.01
+            )
+
+            text = "Hello, world!"
+            result = await async_tokenizer.encode(text)
+
+            # Should get actual token IDs
+            assert "input_ids" in result
+            assert isinstance(result["input_ids"], list)
+            assert len(result["input_ids"]) > 0
+            assert all(isinstance(token_id, int) for token_id in result["input_ids"])
+
+        except Exception as e:
+            pytest.skip(f"Real tokenizer test skipped: {e}")
+
+    @pytest.mark.asyncio
+    async def test_concurrent_mixed_requests(self, async_tokenizer):
+        """Test mixing single and batched requests."""
+        # Start some requests
+        task1 = asyncio.create_task(async_tokenizer.encode("hello"))
+        task2 = asyncio.create_task(async_tokenizer.encode("world"))
+
+        # Wait a bit
+        await asyncio.sleep(0.005)
+
+        # Start more requests
+        task3 = asyncio.create_task(async_tokenizer.encode("how are"))
+        task4 = asyncio.create_task(async_tokenizer.encode("you doing"))
+
+        results = await asyncio.gather(task1, task2, task3, task4)
+
+        # All should complete successfully
+        assert len(results) == 4
+        for result in results:
+            assert "input_ids" in result
+            assert isinstance(result["input_ids"], list)
+
+    def test_cleanup_on_destruction(self, mock_tokenizer):
+        """Test that resources are cleaned up properly."""
+        tokenizer = AsyncDynamicbatchTokenizer(mock_tokenizer)
+
+        # Mock the executor and task
+        tokenizer._executor = Mock()
+        tokenizer._batcher_task = Mock()
+        tokenizer._batcher_task.done.return_value = False
+
+        # Call destructor
+        tokenizer.__del__()
+
+        # Should cancel task and shutdown executor
+        tokenizer._batcher_task.cancel.assert_called_once()
+        tokenizer._executor.shutdown.assert_called_once_with(wait=False)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/test/manual/test_config_integration.py

@@ -0,0 +1,166 @@
+"""
+Test script to verify SGLang config file integration.
+"""
+
+import argparse
+import os
+import sys
+import tempfile
+
+import pytest
+import yaml
+
+from sglang.srt.server_args import ServerArgs, prepare_server_args
+from sglang.srt.server_args_config_parser import ConfigArgumentMerger
+
+
+@pytest.fixture
+def merger():
+    """Fixture providing a ConfigArgumentMerger instance."""
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    return ConfigArgumentMerger(parser)
+
+
+def test_server_args_config_parser(merger):
+    """Test the config parser functionality."""
+    # Create a temporary config file
+    config_data = {
+        "model-path": "microsoft/DialoGPT-medium",
+        "host": "0.0.0.0",
+        "port": 30000,
+        "tensor-parallel-size": 2,
+        "trust-remote-code": False,
+        "enable-metrics": True,
+        "incremental-streaming-output": True,
+        "skip-server-warmup": False,
+        "log-requests": True,
+        "show-time-cost": True,
+        "is-embedding": False,
+    }
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(config_data, f)
+        config_file = f.name
+
+    try:
+        # Test config parser directly
+        config_args = merger._parse_yaml_config(config_file)
+
+        # Test merging with CLI args
+        cli_args = ["--config", config_file, "--max-running-requests", "128"]
+        merged_args = merger.merge_config_with_args(cli_args)
+
+        # Verify the merged args contain both config and CLI values
+        assert "--model-path" in merged_args
+        assert "microsoft/DialoGPT-medium" in merged_args
+        assert "--host" in merged_args
+        assert "0.0.0.0" in merged_args
+        assert "--port" in merged_args
+        assert "30000" in merged_args
+        assert "--tensor-parallel-size" in merged_args
+        assert "2" in merged_args
+        assert "--max-running-requests" in merged_args
+        assert "128" in merged_args
+
+        # Test boolean arguments
+        assert "--enable-metrics" in merged_args  # True boolean
+        assert "--incremental-streaming-output" in merged_args  # True boolean
+        assert "--log-requests" in merged_args  # True boolean
+        assert "--show-time-cost" in merged_args  # True boolean
+        # False booleans should not be present (only add flag if True)
+        assert "--trust-remote-code" not in merged_args  # False boolean
+        assert "--skip-server-warmup" not in merged_args  # False boolean
+        assert "--is-embedding" not in merged_args  # False boolean
+
+    finally:
+        os.unlink(config_file)
+
+
+def test_server_args_integration():
+    """Test the integration with server args."""
+    # Create a temporary config file
+    config_data = {
+        "model-path": "microsoft/DialoGPT-medium",
+        "host": "0.0.0.0",
+        "port": 30000,
+        "tensor-parallel-size": 1,
+        "max-running-requests": 256,
+    }
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(config_data, f)
+        config_file = f.name
+
+    try:
+        # Test with config file
+        argv = ["--config", config_file]
+        server_args = prepare_server_args(argv)
+
+        # Verify that config values were loaded
+        assert server_args.model_path == "microsoft/DialoGPT-medium"
+        assert server_args.host == "0.0.0.0"
+        assert server_args.port == 30000
+        assert server_args.tp_size == 1
+        assert server_args.max_running_requests == 256
+
+    finally:
+        os.unlink(config_file)
+
+
+def test_cli_override():
+    """Test that CLI arguments override config file values."""
+    # Create a temporary config file
+    config_data = {
+        "model-path": "microsoft/DialoGPT-medium",
+        "port": 30000,
+        "tensor-parallel-size": 1,
+    }
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        yaml.dump(config_data, f)
+        config_file = f.name
+
+    try:
+        # Test CLI override (CLI should take precedence)
+        argv = [
+            "--config",
+            config_file,
+            "--port",
+            "40000",
+            "--tensor-parallel-size",
+            "2",
+        ]
+        server_args = prepare_server_args(argv)
+
+        # Verify that CLI values override config values
+        assert server_args.model_path == "microsoft/DialoGPT-medium"  # From config
+        assert server_args.port == 40000  # From CLI (overrides config)
+        assert server_args.tp_size == 2  # From CLI (overrides config)
+
+    finally:
+        os.unlink(config_file)
+
+
+def test_error_handling():
+    """Test error handling for invalid config files."""
+    # Test non-existent config file
+    with pytest.raises(ValueError, match="Config file not found"):
+        argv = ["--config", "non-existent.yaml"]
+        prepare_server_args(argv)
+
+    # Test invalid YAML file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+        f.write("invalid: yaml: content: [")
+        invalid_yaml_file = f.name
+
+    try:
+        with pytest.raises(Exception):
+            argv = ["--config", invalid_yaml_file]
+            prepare_server_args(argv)
+    finally:
+        os.unlink(invalid_yaml_file)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

third_party/sglang/test/manual/test_cross_node_scheduler_info_sync.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+"""
+Test cross-node scheduler_infos synchronization for remote weight loading.
+
+Simulates multi-node setups on a single machine using different GPU subsets.
+Validates that scheduler_infos are correctly synced across nodes via Gloo.
+
+IMPORTANT: For multi-node tests, start both nodes within a few seconds of each
+other to avoid port binding conflicts (they share the same network namespace).
+
+Test cases:
+  - tp4_nodes2: TP=4 across 2 nodes, validates basic cross-node sync
+  - dp2_single_node: DP=2 with dp_attention on single node
+  - dp2_tp2_nodes2: DP=2, TP=4 across 2 nodes with dp_attention
+
+Usage (multi-node):
+    Terminal 1: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --node-rank 0
+    Terminal 2: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --node-rank 1
+    Terminal 3: python test_cross_node_scheduler_info_sync.py --test-case tp4_nodes2 --test-only
+
+Usage (single-node):
+    Terminal 1: python test_cross_node_scheduler_info_sync.py --test-case dp2_single_node --node-rank 0
+    Terminal 2: python test_cross_node_scheduler_info_sync.py --test-case dp2_single_node --test-only
+
+Requirements: 4 GPUs on single machine
+"""
+
+import argparse
+import socket
+import subprocess
+import sys
+import time
+from dataclasses import dataclass
+from typing import List
+
+import requests
+
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+)
+
+
+@dataclass
+class TestCase:
+    name: str
+    tp_size: int
+    dp_size: int
+    nnodes: int
+    gpus_per_node: int
+    expected_ranks: int
+    extra_args: List[str]
+
+
+TEST_CASES = {
+    "tp4_nodes2": TestCase(
+        name="tp4_nodes2",
+        tp_size=4,
+        dp_size=1,
+        nnodes=2,
+        gpus_per_node=2,
+        expected_ranks=4,
+        extra_args=[],
+    ),
+    "dp2_single_node": TestCase(
+        name="dp2_single_node",
+        tp_size=2,
+        dp_size=2,
+        nnodes=1,
+        gpus_per_node=2,
+        expected_ranks=2,
+        extra_args=["--enable-dp-attention", "--dp", "2", "--attention-backend", "fa3"],
+    ),
+    "dp2_tp2_nodes2": TestCase(
+        name="dp2_tp2_nodes2",
+        tp_size=4,
+        dp_size=2,
+        nnodes=2,
+        gpus_per_node=2,
+        expected_ranks=4,
+        extra_args=["--enable-dp-attention", "--dp", "2", "--attention-backend", "fa3"],
+    ),
+}
+
+TEST_CASE_MODELS = {
+    "tp4_nodes2": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+    "dp2_single_node": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+    "dp2_tp2_nodes2": DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
+}
+
+
+def get_local_ip() -> str:
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except Exception:
+        return "127.0.0.1"
+    finally:
+        s.close()
+
+
+def launch_node(
+    test_case: TestCase, node_rank: int, model_path: str, dist_init_addr: str
+):
+    cmd = [
+        sys.executable,
+        "-m",
+        "sglang.launch_server",
+        "--model-path",
+        model_path,
+        "--tp",
+        str(test_case.tp_size),
+        "--port",
+        str(30000 + node_rank * 100),
+        "--host",
+        "0.0.0.0",
+        "--remote-instance-weight-loader-start-seed-via-transfer-engine",
+    ]
+    if test_case.nnodes > 1:
+        cmd.extend(
+            [
+                "--nnodes",
+                str(test_case.nnodes),
+                "--node-rank",
+                str(node_rank),
+                "--dist-init-addr",
+                dist_init_addr,
+                "--base-gpu-id",
+                str(node_rank * test_case.gpus_per_node),
+            ]
+        )
+    cmd.extend(test_case.extra_args)
+    print(f"[Node {node_rank}] {' '.join(cmd)}")
+    subprocess.run(cmd)
+
+
+def test_api(test_case: TestCase) -> bool:
+    base_url = "http://127.0.0.1:30000"
+    print(f"Testing {test_case.name}: expecting {test_case.expected_ranks} ranks")
+
+    for _ in range(60):
+        try:
+            if requests.get(f"{base_url}/health", timeout=2).status_code == 200:
+                break
+        except Exception:
+            pass
+        time.sleep(2)
+    else:
+        print("ERROR: Server not ready")
+        return False
+
+    all_passed = True
+    for rank in range(test_case.expected_ranks):
+        try:
+            resp = requests.get(
+                f"{base_url}/get_remote_instance_transfer_engine_info",
+                params={"rank": rank},
+                timeout=5,
+            )
+            status = "✓" if resp.status_code == 200 else "✗"
+            print(f"{status} Rank {rank}: {resp.status_code}")
+            if resp.status_code != 200:
+                all_passed = False
+        except Exception as e:
+            print(f"✗ Rank {rank}: {e}")
+            all_passed = False
+
+    print("PASSED" if all_passed else "FAILED")
+    return all_passed
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--test-case", type=str, choices=list(TEST_CASES.keys()), required=True
+    )
+    parser.add_argument("--node-rank", type=int, choices=[0, 1])
+    parser.add_argument("--model-path", type=str, default=None)
+    parser.add_argument("--dist-init-addr", type=str, default=None)
+    parser.add_argument("--test-only", action="store_true")
+    args = parser.parse_args()
+
+    test_case = TEST_CASES[args.test_case]
+    model_path = args.model_path or TEST_CASE_MODELS.get(
+        args.test_case, DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
+    )
+
+    if args.test_only:
+        sys.exit(0 if test_api(test_case) else 1)
+
+    if test_case.nnodes == 1:
+        launch_node(test_case, 0, model_path, "")
+        return
+
+    if args.node_rank is None:
+        print(f"Usage: --node-rank 0 or 1, then --test-only in another terminal")
+        sys.exit(0)
+
+    dist_init_addr = args.dist_init_addr or f"{get_local_ip()}:20000"
+    launch_node(test_case, args.node_rank, model_path, dist_init_addr)
+
+
+if __name__ == "__main__":
+    main()

third_party/sglang/test/manual/test_custom_allreduce.py

@@ -0,0 +1,182 @@
+import os
+import random
+import socket
+import unittest
+from typing import Any
+
+import ray
+import torch
+import torch.distributed as dist
+
+from sglang.srt.distributed import init_distributed_environment
+from sglang.srt.distributed.communication_op import (  # noqa
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.distributed.parallel_state import (
+    get_tensor_model_parallel_group,
+    graph_capture,
+    initialize_model_parallel,
+)
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.test_utils import CustomTestCase
+
+
+def get_open_port() -> int:
+    # try ipv4
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+    except OSError:
+        # try ipv6
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+
+def multi_process_parallel(
+    world_size: int,
+    cls: Any,
+    test_target: Any,
+) -> None:
+
+    # Using ray helps debugging the error when it failed
+    # as compared to multiprocessing.
+    # NOTE: We need to set working_dir for distributed tests,
+    # otherwise we may get import errors on ray workers
+
+    ray.init(log_to_driver=True)
+
+    distributed_init_port = get_open_port()
+    refs = []
+    for rank in range(world_size):
+        refs.append(test_target.remote(cls, world_size, rank, distributed_init_port))
+    ray.get(refs)
+
+    ray.shutdown()
+
+
+class TestCustomAllReduce(CustomTestCase):
+    TEST_SIZES = [
+        512,
+        4096,
+        32768,
+        262144,
+        2097152,
+        16777216,
+        33554432,
+        67108864,
+    ]  # 512B...32MB
+    WORLD_SIZES = [2, 4, 6, 8]
+    TEST_LOOP = 10
+
+    @classmethod
+    def setUpClass(cls):
+        random.seed(42)  # keep the deterministic seed
+
+    def test_graph_allreduce(self):
+        for world_size in self.WORLD_SIZES:
+            if world_size > torch.cuda.device_count():
+                continue
+            multi_process_parallel(world_size, self, self.graph_allreduce)
+
+    def test_eager_allreduce(self):
+        for world_size in self.WORLD_SIZES:
+            if world_size > torch.cuda.device_count():
+                continue
+            multi_process_parallel(world_size, self, self.eager_allreduce)
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def graph_allreduce(self, world_size, rank, distributed_init_port):
+        del os.environ["CUDA_VISIBLE_DEVICES"]
+        device = torch.device(f"cuda:{rank}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        # Set global server args to avoid "Global server args is not set yet!" error
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+        # A small all_reduce for warmup.
+        # this is needed because device communicators might be created lazily
+        # (e.g. NCCL). This will ensure that the communicator is initialized
+        # before any communication happens, so that this group can be used for
+        # graph capture immediately.
+        data = torch.zeros(1)
+        data = data.to(device=device)
+        torch.distributed.all_reduce(data, group=group)
+        torch.cuda.synchronize()
+        del data
+
+        for sz in self.TEST_SIZES:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                for _ in range(self.TEST_LOOP):
+                    with graph_capture() as graph_capture_context:
+                        # use integers so result matches NCCL exactly
+                        inp1 = torch.randint(
+                            1,
+                            16,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        inp2 = torch.randint(
+                            1,
+                            16,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        torch.cuda.synchronize()
+                        graph = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(
+                            graph, stream=graph_capture_context.stream
+                        ):
+                            out1 = tensor_model_parallel_all_reduce(inp1)
+                            # the input buffer is immediately modified to test
+                            # synchronization
+                            dist.all_reduce(inp1, group=group)
+                            out2 = tensor_model_parallel_all_reduce(inp2)
+                            dist.all_reduce(inp2, group=group)
+                    graph.replay()
+                    torch.testing.assert_close(out1, inp1)
+                    torch.testing.assert_close(out2, inp2)
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def eager_allreduce(self, world_size, rank, distributed_init_port):
+        del os.environ["CUDA_VISIBLE_DEVICES"]
+        device = torch.device(f"cuda:{rank}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        # Set global server args to avoid "Global server args is not set yet!" error
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+        for sz in self.TEST_SIZES:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                for _ in range(self.TEST_LOOP):
+                    inp1 = torch.randint(
+                        1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
+                    )
+                    out1 = tensor_model_parallel_all_reduce(inp1)
+                    dist.all_reduce(inp1, group=group)
+                    torch.testing.assert_close(out1, inp1)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_deepseek_chat_templates.py

@@ -0,0 +1,318 @@
+"""
+Unit tests for DeepSeek chat template tool call handling.
+
+Tests verify that the DeepSeek chat templates (v3, v3.1, v3.2) correctly handle
+both dict and string types for tool['function']['arguments'] without double-escaping,
+addressing issue #11700.
+"""
+
+import os
+import unittest
+
+from jinja2 import Template
+
+
+class TestDeepSeekChatTemplateToolCalls(unittest.TestCase):
+    """Test DeepSeek chat templates handle tool calls correctly."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Load all DeepSeek chat templates."""
+        base_path = os.path.join(
+            os.path.dirname(__file__), "..", "..", "examples", "chat_template"
+        )
+
+        cls.templates = {}
+        template_files = {
+            "v3": "tool_chat_template_deepseekv3.jinja",
+            "v3.1": "tool_chat_template_deepseekv31.jinja",
+            "v3.2": "tool_chat_template_deepseekv32.jinja",
+        }
+
+        for version, filename in template_files.items():
+            template_path = os.path.join(base_path, filename)
+            with open(template_path, "r") as f:
+                template_content = f.read()
+            cls.templates[version] = Template(template_content)
+
+    def _render_template(
+        self, version, messages, tools=None, add_generation_prompt=True
+    ):
+        """Helper method to render a template with given messages and tools."""
+        template = self.templates[version]
+
+        # Common template variables
+        context = {
+            "messages": messages,
+            "add_generation_prompt": add_generation_prompt,
+            "bos_token": "<｜begin▁of▁sentence｜>",
+        }
+
+        if tools is not None:
+            context["tools"] = tools
+
+        return template.render(**context)
+
+    def test_tool_arguments_as_dict(self):
+        """Test that tool arguments as dict are properly JSON-encoded (normal case)."""
+        # This tests the normal case where arguments come from OpenAI API as dict
+
+        for version in ["v3", "v3.1", "v3.2"]:
+            with self.subTest(version=version):
+                messages = [
+                    {"role": "user", "content": "What's the weather in NYC?"},
+                    {
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": {
+                                        "city": "New York",
+                                        "unit": "celsius",
+                                    },  # Dict
+                                },
+                            }
+                        ],
+                    },
+                ]
+
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "description": "Get weather information",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "city": {"type": "string"},
+                                    "unit": {"type": "string"},
+                                },
+                            },
+                        },
+                    }
+                ]
+
+                output = self._render_template(version, messages, tools)
+
+                # Should contain properly formatted JSON (not double-escaped)
+                self.assertIn('"city"', output, f"{version}: Should contain city key")
+                self.assertIn(
+                    '"New York"', output, f"{version}: Should contain city value"
+                )
+
+                # Should NOT contain double-escaped quotes
+                self.assertNotIn(
+                    '\\"city\\"', output, f"{version}: Should not double-escape"
+                )
+                self.assertNotIn(
+                    '\\\\"', output, f"{version}: Should not have escaped backslashes"
+                )
+
+    def test_tool_arguments_as_string(self):
+        """Test that tool arguments as string are used as-is (multi-round case)."""
+        # This tests the multi-round function calling case from issue #11700
+        # where arguments might already be JSON strings from previous model output
+
+        for version in ["v3", "v3.1", "v3.2"]:
+            with self.subTest(version=version):
+                messages = [
+                    {"role": "user", "content": "What's the stock price of NVDA?"},
+                    {
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_stock_info",
+                                    "arguments": '{"symbol": "NVDA"}',  # Already a JSON string
+                                },
+                            }
+                        ],
+                    },
+                ]
+
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_stock_info",
+                            "description": "Get stock information",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {
+                                    "symbol": {"type": "string"},
+                                },
+                            },
+                        },
+                    }
+                ]
+
+                output = self._render_template(version, messages, tools)
+
+                # Should contain the JSON string as-is
+                self.assertIn(
+                    '{"symbol": "NVDA"}',
+                    output,
+                    f"{version}: Should contain JSON as-is",
+                )
+
+                # Should NOT double-escape (the bug from issue #11700)
+                # Bad output would look like: "{\"symbol\": \"NVDA\"}" or "{\\"symbol\\": \\"NVDA\\"}"
+                self.assertNotIn(
+                    '{\\"symbol\\"', output, f"{version}: Should not double-escape"
+                )
+                self.assertNotIn(
+                    '"{\\"symbol', output, f"{version}: Should not wrap and escape"
+                )
+
+                # Verify it's not triple-quoted or escaped
+                self.assertNotIn(
+                    '""{"', output, f"{version}: Should not have extra quotes"
+                )
+
+    def test_multiple_tool_calls_mixed_types(self):
+        """Test multiple tool calls with mixed dict and string argument types."""
+        # This tests a complex scenario with multiple tools, some with dict args, some with string
+
+        for version in ["v3", "v3.1", "v3.2"]:
+            with self.subTest(version=version):
+                messages = [
+                    {"role": "user", "content": "Get weather and stock info"},
+                    {
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": {"city": "Boston"},  # Dict
+                                },
+                            },
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_stock_info",
+                                    "arguments": '{"symbol": "TSLA"}',  # String
+                                },
+                            },
+                        ],
+                    },
+                ]
+
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "description": "Get weather",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {"city": {"type": "string"}},
+                            },
+                        },
+                    },
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_stock_info",
+                            "description": "Get stock info",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {"symbol": {"type": "string"}},
+                            },
+                        },
+                    },
+                ]
+
+                output = self._render_template(version, messages, tools)
+
+                # First tool (dict) should be properly JSON-encoded
+                self.assertIn(
+                    '"city"', output, f"{version}: First tool should have city key"
+                )
+                self.assertIn(
+                    '"Boston"',
+                    output,
+                    f"{version}: First tool should have Boston value",
+                )
+
+                # Second tool (string) should be used as-is
+                self.assertIn(
+                    '{"symbol": "TSLA"}',
+                    output,
+                    f"{version}: Second tool should use string as-is",
+                )
+
+                # Neither should be double-escaped
+                self.assertNotIn(
+                    '\\"city\\"',
+                    output,
+                    f"{version}: First tool should not double-escape",
+                )
+                self.assertNotIn(
+                    '\\"symbol\\"',
+                    output,
+                    f"{version}: Second tool should not double-escape",
+                )
+
+    def test_tool_call_with_content(self):
+        """Test tool calls that also include content text."""
+        # Some models include explanatory text along with tool calls
+
+        for version in ["v3", "v3.1", "v3.2"]:
+            with self.subTest(version=version):
+                messages = [
+                    {"role": "user", "content": "What's the weather?"},
+                    {
+                        "role": "assistant",
+                        "content": "Let me check the weather for you.",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "function": {
+                                    "name": "get_weather",
+                                    "arguments": {"city": "Seattle"},
+                                },
+                            }
+                        ],
+                    },
+                ]
+
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "description": "Get weather",
+                            "parameters": {
+                                "type": "object",
+                                "properties": {"city": {"type": "string"}},
+                            },
+                        },
+                    }
+                ]
+
+                output = self._render_template(version, messages, tools)
+
+                # Should contain both the content and the tool call
+                self.assertIn(
+                    "Let me check the weather",
+                    output,
+                    f"{version}: Should include content",
+                )
+                self.assertIn(
+                    '"city"', output, f"{version}: Should include tool arguments"
+                )
+                self.assertNotIn(
+                    '\\"city\\"', output, f"{version}: Should not double-escape"
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_double_sparsity.py

@@ -0,0 +1,65 @@
+import os
+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.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestDoubleSparsity(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        dirpath = os.path.dirname(__file__)
+        config_file = os.path.join(
+            dirpath, "double-sparsity-config-Llama-3.1-8B-Instruct.json"
+        )
+        # NOTE: Generate the config file by running https://github.com/andy-yang-1/DoubleSparse/blob/main/evaluation/group_channel_config.py
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--enable-double-sparsity",
+                "--ds-channel-config-path",
+                config_file,
+                "--ds-heavy-channel-num",
+                "32",
+                "--ds-heavy-channel-type",
+                "k",
+                "--ds-heavy-token-num",
+                "512",
+                "--ds-sparse-decode-threshold",
+                "0",
+                "--max-total-tokens",
+                "200000",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    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.assertGreaterEqual(metrics["score"], 0.65)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_expert_distribution.py

@@ -0,0 +1,101 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import requests
+import torch
+
+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_launch_server,
+)
+
+
+class TestExpertDistribution(CustomTestCase):
+    def test_expert_distribution_record(self):
+        # TODO: Add tests for DeepEP gatherer (currently our CI cannot run that)
+        for info in [
+            dict(model_path="deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct"),
+            dict(model_path="Qwen/Qwen1.5-MoE-A2.7B"),
+            dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", tp_size=2),
+            dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", mode="per_pass"),
+            dict(model_path="Qwen/Qwen1.5-MoE-A2.7B", mode="per_token"),
+        ]:
+            with self.subTest(info=info):
+                self._execute_core(**info)
+
+    def _execute_core(self, model_path: str, mode: str = "stat", tp_size: int = 1):
+        """Test expert distribution record endpoints"""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
+
+            process = popen_launch_server(
+                model_path,
+                DEFAULT_URL_FOR_TEST,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=[
+                    "--trust-remote-code",
+                    "--tp-size",
+                    str(tp_size),
+                    "--expert-distribution-recorder-mode",
+                    mode,
+                    "--disable-cuda-graph",
+                    "--disable-overlap-schedule",
+                ],
+            )
+
+            try:
+                # Start recording
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/start_expert_distribution_record"
+                )
+                self.assertEqual(response.status_code, 200)
+
+                # Make some requests to generate expert distribution data
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/generate",
+                    json={
+                        "text": "The capital of France is",
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 32,
+                        },
+                    },
+                )
+                self.assertEqual(response.status_code, 200)
+
+                # Stop recording
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/stop_expert_distribution_record"
+                )
+                self.assertEqual(response.status_code, 200)
+
+                # Dump the recorded data
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/dump_expert_distribution_record"
+                )
+                self.assertEqual(response.status_code, 200)
+
+                # Check data rows
+                data = torch.load(
+                    list(Path(tmp_dir).glob("*.pt"))[0], weights_only=True
+                )
+                print(f"{data=}")
+
+                if mode in ["per_pass", "per_token"]:
+                    self.assertGreater(len(data), 0, "Should contain data rows")
+                else:
+                    logical_count = data["logical_count"]
+                    print(f"{logical_count.sum()=} {logical_count=}")
+                    self.assertTrue(logical_count.sum() > 0)
+
+            finally:
+                kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_expert_location_updater.py

@@ -0,0 +1,258 @@
+import os
+import traceback
+import unittest
+from dataclasses import dataclass
+from typing import List
+
+import torch
+import torch.distributed
+import torch.multiprocessing as mp
+from torch.multiprocessing import Process
+
+from sglang.srt.eplb import expert_location_updater
+from sglang.srt.utils import get_device
+from sglang.test.test_utils import CustomTestCase, find_available_port
+from sglang.utils import is_in_ci
+
+
+@dataclass
+class _TestInfo:
+    nnodes: int
+    num_logical_experts: int
+    num_physical_experts: int
+    num_repeat: int = 5000
+
+
+class TestExpertLocationUpdater(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        mp.set_start_method("spawn", force=True)
+
+    def test_cpu(self):
+        self._test_common(device="cpu")
+        self._test_core(
+            num_gpus=32,
+            device="cpu",
+            infos=[
+                _TestInfo(
+                    nnodes=4,
+                    num_logical_experts=256,
+                    num_physical_experts=288,
+                    num_repeat=10000,
+                )
+            ],
+        )
+
+    def test_cpu_slow(self):
+        if is_in_ci():
+            return
+        self._test_core(
+            num_gpus=144,
+            device="cpu",
+            infos=[
+                _TestInfo(
+                    nnodes=18,
+                    num_logical_experts=256,
+                    num_physical_experts=288,
+                    num_repeat=10000,
+                )
+            ],
+        )
+
+    def test_gpu(self):
+        if is_in_ci():
+            return
+        self._test_common(device=get_device())
+
+    def _test_common(self, device):
+        infos = []
+
+        for nnodes in [1, 2, 4]:
+            for num_logical_experts in [2, 5, 20, 256]:
+                for num_physical_experts in [8, 16, 256, 288]:
+                    if num_logical_experts > num_physical_experts:
+                        continue
+                    infos.append(
+                        _TestInfo(
+                            nnodes=nnodes,
+                            num_logical_experts=num_logical_experts,
+                            num_physical_experts=num_physical_experts,
+                        )
+                    )
+
+        self._test_core(num_gpus=8, device=device, infos=infos)
+
+    def _test_core(
+        self,
+        num_gpus: int,
+        device: str,
+        infos: List[_TestInfo],
+    ):
+        master_port = find_available_port(23456)
+
+        processes = []
+        output_reader, output_writer = mp.Pipe(duplex=False)
+        for rank in range(num_gpus):
+            p = Process(
+                target=_run_subprocess,
+                kwargs=dict(
+                    rank=rank,
+                    num_gpus=num_gpus,
+                    output_writer=output_writer,
+                    master_port=master_port,
+                    device=device,
+                    infos=infos,
+                ),
+            )
+            p.start()
+            processes.append(p)
+
+        for _ in range(num_gpus):
+            self.assertTrue(
+                output_reader.recv(), f"Subprocess has error, please see logs above."
+            )
+
+        for p in processes:
+            p.join()
+
+
+def _run_subprocess(
+    rank: int,
+    num_gpus: int,
+    master_port: int,
+    device: str,
+    infos: List[_TestInfo],
+    output_writer,
+):
+    try:
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(master_port)
+
+        torch.random.manual_seed(42)
+        torch.distributed.init_process_group(
+            rank=rank,
+            world_size=num_gpus,
+            backend={"cpu": "gloo", "cuda": None}[device],
+        )
+        if device == "cuda":
+            torch.cuda.set_device(f"cuda:{rank}")
+        if device == "xpu":
+            torch.xpu.set_device(f"xpu:{rank}")
+
+        for info in infos:
+            _execute_test(info, rank=rank, num_gpus=num_gpus, device=device)
+
+        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()
+
+
+def _execute_test(info: _TestInfo, rank: int, num_gpus: int, device: str):
+    if rank == 0:
+        print(f"Test: {num_gpus=} {info=}", flush=True)
+
+    assert info.num_physical_experts % num_gpus == 0
+    num_local_physical_experts = info.num_physical_experts // num_gpus
+    assert num_gpus % info.nnodes == 0
+    num_gpu_per_node = num_gpus // info.nnodes
+
+    def _create_routed_experts_weights(physical_to_logical_map):
+        local_logical_expert_ids = physical_to_logical_map[
+            rank * num_local_physical_experts : (rank + 1) * num_local_physical_experts
+        ].cpu()
+        return [
+            local_logical_expert_ids.to(device).clone(),
+            torch.tensor(
+                [
+                    [local_logical_expert_id * 10, local_logical_expert_id * 100]
+                    for local_logical_expert_id in local_logical_expert_ids.tolist()
+                ],
+                device=device,
+            ),
+        ]
+
+    def _create_physical_to_logical_map():
+        if rank == 0:
+            ans = torch.concat(
+                [
+                    torch.arange(0, info.num_logical_experts),
+                    torch.randint(
+                        0,
+                        info.num_logical_experts,
+                        (info.num_physical_experts - info.num_logical_experts,),
+                    ),
+                ]
+            )
+            ans = ans[torch.randperm(ans.shape[0])]
+        else:
+            ans = torch.empty((info.num_physical_experts,), dtype=torch.int64)
+
+        assert ans.dtype == torch.int64 and ans.shape == (info.num_physical_experts,)
+        ans = ans.to(device)
+        torch.distributed.broadcast(ans, src=0)
+
+        return ans.cpu()
+
+    physical_to_logical_map = _create_physical_to_logical_map()
+    routed_experts_weights = _create_routed_experts_weights(physical_to_logical_map)
+
+    for i in range(info.num_repeat):
+        if rank == 0 and ((i % 500 == 0) or (i == info.num_repeat - 1)):
+            print(f"Step {i}/{info.num_repeat}", flush=True)
+
+        new_physical_to_logical_map = _create_physical_to_logical_map()
+        expect_new_weights = _create_routed_experts_weights(new_physical_to_logical_map)
+
+        output_logs = expert_location_updater.update_expert_weights_single_layer(
+            routed_experts_weights=routed_experts_weights,
+            temp_buffers=expert_location_updater.create_temp_buffers(
+                routed_experts_weights
+            ),
+            old_physical_to_logical_map=physical_to_logical_map.tolist(),
+            new_physical_to_logical_map=new_physical_to_logical_map.tolist(),
+            num_local_physical_experts=num_local_physical_experts,
+            num_gpu_per_node=num_gpu_per_node,
+            rank=rank,
+            debug=True,
+        )
+
+        local_has_error = not all(
+            torch.all(x == y)
+            for x, y in zip(routed_experts_weights, expect_new_weights, strict=True)
+        )
+        global_has_error = torch.tensor(local_has_error, device=device)
+        torch.distributed.all_reduce(
+            global_has_error, op=torch.distributed.ReduceOp.MAX
+        )
+
+        if global_has_error.cpu().item():
+            output_logs_str = "\n".join(output_logs)
+            local_message = (
+                f"===================== rank {rank} ============================\n"
+                f"{num_gpus=} {info=}\n"
+                f"{routed_experts_weights[0].tolist()=}\n"
+                f"{expect_new_weights[0].tolist()=}\n"
+                f"{physical_to_logical_map.tolist()=}\n"
+                f"{new_physical_to_logical_map.tolist()=}\n"
+                f"===logs===\n"
+                f"{output_logs_str}\n"
+                f"==============================================================\n"
+            )
+
+            global_messages = ([None] * num_gpus) if rank == 0 else None
+            torch.distributed.gather_object(local_message, global_messages, dst=0)
+
+            if rank == 0:
+                print("\n\n".join(global_messages), flush=True)
+            raise AssertionError(f"Error happens, see logs above")
+
+        physical_to_logical_map = new_physical_to_logical_map
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_fim_completion.py

@@ -0,0 +1,72 @@
+import unittest
+
+import openai
+
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestFimCompletion(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "deepseek-ai/deepseek-coder-1.3b-base"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        other_args = ["--completion-template", "deepseek_coder"]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=cls.api_key,
+            other_args=other_args,
+        )
+        cls.base_url += "/v1"
+        cls.tokenizer = get_tokenizer(cls.model)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_fim_completion(self, number_of_completion):
+        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
+        prompt = "function sum(a: number, b: number): number{\n"
+        suffix = "}"
+
+        prompt_input = self.tokenizer.encode(prompt) + self.tokenizer.encode(suffix)
+        num_prompt_tokens = len(prompt_input) + 2
+
+        response = client.completions.create(
+            model=self.model,
+            prompt=prompt,
+            suffix=suffix,
+            temperature=0.3,
+            max_tokens=32,
+            stream=False,
+            n=number_of_completion,
+        )
+
+        print(response)
+        print(len(response.choices))
+        assert len(response.choices) == number_of_completion
+        assert response.id
+        assert response.created
+        assert response.object == "text_completion"
+        assert (
+            response.usage.prompt_tokens == num_prompt_tokens
+        ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}"
+        assert response.usage.completion_tokens > 0
+        assert response.usage.total_tokens > 0
+
+    def test_fim_completion(self):
+        for number_of_completion in [1, 3]:
+            self.run_fim_completion(number_of_completion)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_forward_split_prefill.py

@@ -0,0 +1,307 @@
+"""
+Test forward_split_prefill functionality.
+
+Usage:
+python3 -m unittest test_forward_split_prefill.TestForwardSplitPrefill
+or
+python3 test_forward_split_prefill.py
+"""
+
+import unittest
+
+import numpy as np
+import torch
+
+from sglang.bench_one_batch import TreeCacheNamespace
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.utils import get_device
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
+
+
+class TestForwardSplitPrefill(CustomTestCase):
+    """Test cases for forward_split_prefill functionality."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Set up the test environment once for all tests."""
+        cls.model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.tp_size = 1
+        cls.device = get_device()
+
+        # Initialize server args
+        cls.server_args = ServerArgs(
+            model_path=cls.model_path,
+            tokenizer_path=cls.model_path,
+            host="127.0.0.1",
+            disable_cuda_graph=True,  # Disable CUDA graph for testing split prefill
+            disable_hybrid_swa_memory=True,
+            port=30000,
+            tp_size=cls.tp_size,
+            mem_fraction_static=0.8,
+            trust_remote_code=True,
+        )
+
+        cls.port_args = PortArgs.init_new(cls.server_args)
+
+        # Load model and tokenizer
+        cls.model_config = ModelConfig.from_server_args(cls.server_args)
+        cls.model_runner = ModelRunner(
+            model_config=cls.model_config,
+            mem_fraction_static=cls.server_args.mem_fraction_static,
+            gpu_id=0,
+            tp_rank=0,
+            tp_size=cls.tp_size,
+            pp_rank=0,
+            pp_size=1,
+            nccl_port=cls.port_args.nccl_port,
+            server_args=cls.server_args,
+            moe_ep_rank=0,
+            moe_ep_size=1,
+        )
+
+        cls.tokenizer = get_tokenizer(
+            cls.server_args.tokenizer_path,
+            tokenizer_mode=cls.server_args.tokenizer_mode,
+            trust_remote_code=cls.server_args.trust_remote_code,
+        )
+
+        print(
+            f"Test with model: {cls.model_path}, num_hidden_layers: {cls.model_config.num_hidden_layers}"
+        )
+
+    def prepare_test_batch(self, batch_size=2, input_len=128, is_split_prefill=True):
+        """Prepare a test batch for split prefill testing."""
+        # Create synthetic input
+        input_ids = np.random.randint(10, 1000, (batch_size, input_len), dtype=np.int32)
+
+        sampling_params = SamplingParams(
+            temperature=0.0,
+            max_new_tokens=8,
+        )
+
+        reqs = []
+        for i in range(batch_size):
+            req = Req(
+                rid=i,
+                origin_input_text="",
+                origin_input_ids=list(input_ids[i]),
+                sampling_params=sampling_params,
+            )
+            req.fill_ids = req.origin_input_ids
+            req.logprob_start_len = -1
+            req.set_extend_input_len(len(req.fill_ids) - len(req.prefix_indices))
+            reqs.append(req)
+
+        # Create dummy tree_cache for tests (no prefix caching, just allocation)
+        dummy_tree_cache = TreeCacheNamespace(
+            page_size=1,
+            device=self.model_runner.device,
+            token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
+        )
+
+        batch = ScheduleBatch.init_new(
+            reqs=reqs,
+            req_to_token_pool=self.model_runner.req_to_token_pool,
+            token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
+            tree_cache=dummy_tree_cache,
+            model_config=self.model_config,
+            enable_overlap=False,
+            spec_algorithm=SpeculativeAlgorithm.NONE,
+        )
+        if is_split_prefill:
+            batch.prepare_for_split_prefill()
+        else:
+            batch.prepare_for_extend()
+
+        # Create forward batch
+        model_worker_batch = batch.get_model_worker_batch()
+        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
+
+        return forward_batch
+
+    def test_split_prefill_functionality(self):
+        """Test that split prefill can complete successfully."""
+        print("\n=== Testing split prefill functionality ===")
+
+        forward_batch = self.prepare_test_batch(batch_size=2, input_len=64)
+
+        # Reset split index
+        forward_batch.split_index = 0
+
+        # Test split prefill in chunks
+        num_layers = self.model_config.num_hidden_layers
+        chunk_size = max(1, num_layers // 4)  # Split into 4 chunks
+
+        results = []
+        split_count = 0
+
+        while forward_batch.split_index < num_layers:
+            print(
+                f"Processing split {split_count}, split_index: {forward_batch.split_index}"
+            )
+
+            result = self.model_runner.forward_split_prefill(
+                forward_batch=forward_batch,
+                reinit_attn_backend=(split_count == 0),
+                forward_count=chunk_size,
+            )
+
+            results.append(result)
+            split_count += 1
+
+            # Verify split_index is updated correctly
+            expected_next_index = min(split_count * chunk_size, num_layers)
+            self.assertEqual(forward_batch.split_index, expected_next_index)
+
+        # The last result should contain logits
+        self.assertIsNotNone(results[-1], "Final split should return logits")
+        print(f"Split prefill completed in {split_count} splits")
+
+    def test_split_prefill_vs_normal_prefill(self):
+        """Test that split prefill produces the same results as normal prefill."""
+        print("\n=== Testing split prefill vs normal prefill consistency ===")
+
+        forward_batch_normal = self.prepare_test_batch(
+            batch_size=2, input_len=128, is_split_prefill=False
+        )
+        forward_batch_split = self.prepare_test_batch(
+            batch_size=2, input_len=128, is_split_prefill=True
+        )
+
+        # Ensure same input
+        forward_batch_split.input_ids = forward_batch_normal.input_ids.clone()
+        forward_batch_split.positions = forward_batch_normal.positions.clone()
+
+        # Method 1: Normal extend (prefill)
+        print("Running normal extend (prefill)...")
+        normal_result = self.model_runner.forward_extend(forward_batch_normal)
+
+        # Method 2: Split prefill
+        print("Running split prefill...")
+        num_layers = self.model_config.num_hidden_layers
+        chunk_size = max(1, num_layers // 3)  # Split into 3 chunks
+
+        split_result = None
+
+        while forward_batch_split.split_index < num_layers:
+            result = self.model_runner.forward_split_prefill(
+                forward_batch=forward_batch_split,
+                forward_count=chunk_size,
+            )
+            if result is not None:
+                split_result = result
+
+        # Compare results
+        self.assertIsNotNone(normal_result, "Normal prefill should return result")
+        self.assertIsNotNone(split_result, "Split prefill should return result")
+
+        # Compare logits shapes
+        self.assertEqual(
+            normal_result.next_token_logits.shape,
+            split_result.next_token_logits.shape,
+            "Logits shapes should match",
+        )
+
+        # Compare logits values (should be very close due to same computation)
+        # Use a larger tolerance for numerical differences in split computation
+        torch.testing.assert_close(
+            normal_result.next_token_logits,
+            split_result.next_token_logits,
+            rtol=1e-3,
+            atol=1e-3,
+            msg="Split prefill and normal prefill should produce similar logits",
+        )
+
+        print("✓ Split prefill and normal prefill produce consistent results")
+
+    def test_split_prefill_different_chunk_sizes(self):
+        """Test split prefill with different chunk sizes."""
+        print("\n=== Testing split prefill with different chunk sizes ===")
+
+        num_layers = self.model_config.num_hidden_layers
+        chunk_sizes = [1, 2, max(1, num_layers // 2), num_layers]
+
+        # Prepare identical batches for each test
+        base_batch = self.prepare_test_batch(batch_size=1, input_len=16)
+        base_input_ids = base_batch.input_ids.clone()
+        base_positions = base_batch.positions.clone()
+
+        results = []
+
+        for chunk_size in chunk_sizes:
+            if chunk_size > num_layers:
+                continue
+
+            print(f"Testing chunk size: {chunk_size}")
+
+            # Prepare fresh batch
+            forward_batch = self.prepare_test_batch(batch_size=1, input_len=16)
+            forward_batch.input_ids = base_input_ids.clone()
+            forward_batch.positions = base_positions.clone()
+            forward_batch.split_index = 0
+
+            # Run split prefill
+            split_result = None
+
+            while forward_batch.split_index < num_layers:
+                result = self.model_runner.forward_split_prefill(
+                    forward_batch=forward_batch,
+                    forward_count=chunk_size,
+                )
+                if result is not None:
+                    split_result = result
+
+            self.assertIsNotNone(
+                split_result,
+                f"Split prefill should succeed with chunk_size={chunk_size}",
+            )
+            results.append(split_result)
+
+        # Compare all results should be identical (same input, same computation)
+        if len(results) > 1:
+            for i, result in enumerate(results[1:], 1):
+                torch.testing.assert_close(
+                    results[0].next_token_logits,
+                    result.next_token_logits,
+                    rtol=1e-3,
+                    atol=1e-3,
+                    msg=f"Results with different chunk sizes should be identical (chunk_size {chunk_sizes[i]})",
+                )
+
+        print("✓ All chunk sizes produce consistent results")
+
+    def test_split_prefill_edge_cases(self):
+        """Test edge cases for split prefill."""
+        print("\n=== Testing split prefill edge cases ===")
+
+        # Test with single layer chunks
+        forward_batch = self.prepare_test_batch(batch_size=1, input_len=8)
+
+        # Process one layer at a time
+        num_layers = self.model_config.num_hidden_layers
+        for layer_idx in range(num_layers):
+            result = self.model_runner.forward_split_prefill(
+                forward_batch=forward_batch,
+                reinit_attn_backend=(layer_idx == 0),
+                forward_count=1,  # One layer at a time
+            )
+
+            if layer_idx == num_layers - 1:
+                # Last layer should return result
+                self.assertIsNotNone(result, "Last layer should return logits")
+            else:
+                # Intermediate layers should return None
+                self.assertIsNone(result, f"Layer {layer_idx} should return None")
+
+        print("✓ Single layer processing works correctly")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_get_weights_by_name.py

@@ -0,0 +1,185 @@
+import gc
+import unittest
+
+import numpy as np
+import requests
+from transformers import AutoModelForCausalLM
+
+import sglang as sgl
+from sglang.srt.utils import get_device
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    empty_gpu_cache,
+    get_gpu_count,
+    is_in_ci,
+    popen_launch_server,
+)
+from sglang.utils import terminate_process
+
+
+def _process_return(ret):
+    if isinstance(ret, list) and len(ret) == 2:
+        print(f"running assert_allclose on data parallel")
+        np.testing.assert_allclose(ret[0], ret[1])
+        return np.array(ret[0])
+    return np.array(ret)
+
+
+class TestGetWeightsByName(CustomTestCase):
+
+    def init_hf_model(self, model_name, tie_word_embeddings):
+        self.hf_model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype="bfloat16", tie_word_embeddings=tie_word_embeddings
+        ).to(get_device())
+
+    def init_backend(self, backend, dp, tp, model_name):
+        self.backend = backend
+        self.dp = dp
+        self.tp = tp
+        if backend == "Engine":
+            self.engine = sgl.Engine(
+                model_path=model_name,
+                random_seed=42,
+                tp_size=tp,
+                dp_size=dp,
+            )
+        else:
+            self.process = popen_launch_server(
+                model_name,
+                DEFAULT_URL_FOR_TEST,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=(
+                    "--tp-size",
+                    str(tp),
+                    "--dp-size",
+                    str(dp),
+                ),
+            )
+
+    def clean_up(self):
+        del self.hf_model
+        gc.collect()
+        empty_gpu_cache()
+        if self.backend == "Engine":
+            self.engine.shutdown()
+        else:
+            terminate_process(self.process)
+
+    def assert_tie_word_embeddings(self, truncate_size):
+        print("assert_tie_word_embeddings")
+        if self.backend == "Engine":
+            backend_ret = _process_return(
+                self.engine.get_weights_by_name("lm_head.weight", truncate_size)
+            )
+        else:
+            backend_ret = _process_return(
+                requests.get(
+                    f"{DEFAULT_URL_FOR_TEST}/get_weights_by_name",
+                    json={"name": "lm_head.weight", "truncate_size": truncate_size},
+                ).json()
+            )
+        print("assert_tie_word_embeddings of hf and backend")
+        assert np.allclose(
+            self.hf_model.get_parameter("model.embed_tokens.weight")
+            .cpu()
+            .detach()
+            .float()
+            .numpy()[:truncate_size],
+            backend_ret,
+        )
+        assert np.allclose(
+            self.hf_model.get_parameter("lm_head.weight")
+            .cpu()
+            .detach()
+            .float()
+            .numpy()[:truncate_size],
+            self.hf_model.get_parameter("model.embed_tokens.weight")
+            .cpu()
+            .detach()
+            .float()
+            .numpy()[:truncate_size],
+        )
+
+    def assert_weights_all_close(self, param_name, truncate_size):
+        print(
+            f"param_name: {param_name}, backend: {self.backend}, dp: {self.dp}, tp: {self.tp}"
+        )
+        param = self.hf_model.get_parameter(param_name)[:truncate_size]
+        param_np = param.cpu().detach().float().numpy()
+
+        if self.backend == "Engine":
+            engine_ret = self.engine.get_weights_by_name(param_name, truncate_size)
+            engine_ret = _process_return(engine_ret)
+            np.testing.assert_allclose(engine_ret, param_np, rtol=1e-5, atol=1e-5)
+
+        if self.backend == "Runtime":
+            runtime_ret = requests.get(
+                f"{DEFAULT_URL_FOR_TEST}/get_weights_by_name",
+                json={"name": param_name, "truncate_size": truncate_size},
+            ).json()
+            runtime_ret = _process_return(runtime_ret)
+            np.testing.assert_allclose(runtime_ret, param_np, rtol=1e-5, atol=1e-5)
+
+    def test_get_weights_by_name(self):
+        if is_in_ci():
+            test_suits = [
+                ("Engine", 1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
+            ]
+        else:
+            test_suits = [
+                ("Runtime", 1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
+                ("Engine", 1, 1, DEFAULT_MODEL_NAME_FOR_TEST),
+            ]
+            if get_gpu_count() >= 2:
+                test_suits.append(("Engine", 1, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST))
+                test_suits.append(("Runtime", 2, 1, DEFAULT_MODEL_NAME_FOR_TEST))
+
+            if get_gpu_count() >= 4:
+                test_suits.extend(
+                    [
+                        ("Engine", 2, 2, DEFAULT_SMALL_MODEL_NAME_FOR_TEST),
+                        ("Runtime", 2, 2, DEFAULT_MODEL_NAME_FOR_TEST),
+                    ]
+                )
+
+        parameters = [
+            "model.embed_tokens.weight",
+            "model.layers.0.input_layernorm.weight",
+            "model.layers.1.self_attn.q_proj.weight",
+            "model.layers.2.self_attn.k_proj.weight",
+            "model.layers.3.self_attn.v_proj.weight",
+            "model.layers.4.self_attn.o_proj.weight",
+            "model.layers.5.mlp.gate_proj.weight",
+            "model.layers.6.mlp.up_proj.weight",
+            "model.layers.7.mlp.down_proj.weight",
+            "model.layers.8.post_attention_layernorm.weight",
+            "model.norm.weight",
+            "lm_head.weight",
+        ]
+
+        truncate_size = 100
+
+        for test_suit in test_suits:
+            if test_suit[-1] == DEFAULT_MODEL_NAME_FOR_TEST:
+                tie_word_embeddings = False
+            else:
+                tie_word_embeddings = True
+
+            self.init_hf_model(test_suit[-1], tie_word_embeddings)
+            self.init_backend(*test_suit)
+
+            for param_name in parameters:
+                self.assert_weights_all_close(param_name, truncate_size)
+
+            if tie_word_embeddings:
+                self.assert_tie_word_embeddings(truncate_size)
+
+            self.clean_up()
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_health_check.py

@@ -0,0 +1,28 @@
+import unittest
+
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestHealthCheck(CustomTestCase):
+    def test_health_check(self):
+        """Test that metrics endpoint returns data when enabled"""
+        with self.assertRaises(TimeoutError):
+            popen_launch_server(
+                DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+                DEFAULT_URL_FOR_TEST,
+                timeout=60,
+                other_args=[
+                    "--disable-cuda-graph",
+                    "--json-model-override-args",
+                    '{"architectures": ["LlamaForCausalLMForHealthTest"]}',
+                ],
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_kv_events.py

@@ -0,0 +1,292 @@
+import time
+import unittest
+
+import requests
+import zmq
+from msgspec.msgpack import Decoder
+
+from sglang.srt.disaggregation.kv_events import (
+    AllBlocksCleared,
+    BlockRemoved,
+    BlockStored,
+    KVEventBatch,
+)
+from sglang.srt.utils import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestKvEvents(CustomTestCase):
+    def test_kv_events_enabled(self):
+        """Test that kv events are sent and received by subscriber data when enabled"""
+
+        # Launch kv events subscriber
+        decoder = Decoder(type=KVEventBatch)
+        context = zmq.Context()
+        sub = context.socket(zmq.SUB)
+        sub.connect("tcp://localhost:5557")
+        topic = "kv-events"
+        sub.setsockopt_string(zmq.SUBSCRIBE, topic)
+
+        # Launch sglang server
+        process = popen_launch_server(
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--kv-events-config",
+                '{"publisher": "zmq", "topic": "kv-events"}',
+                "--max-total-tokens",
+                32,
+                "--cuda-graph-max-bs",
+                2,
+                "--enable-dp-attention",
+                "--dp-size",
+                1,
+            ],
+        )
+
+        try:
+            # Make some requests to generate some metrics
+            response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
+            self.assertEqual(response.status_code, 200)
+
+            response = requests.post(
+                f"{DEFAULT_URL_FOR_TEST}/generate",
+                json={
+                    "text": "The capital of France is",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": 32,
+                    },
+                },
+            )
+            response = requests.post(
+                f"{DEFAULT_URL_FOR_TEST}/generate",
+                json={
+                    "text": "The capital of Spain is",
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": 32,
+                    },
+                },
+            )
+
+            # Get events
+            events = []
+            start = time.time()
+            max_wait_s = 5
+            min_events_expected = 5  # Expect at least some events
+
+            while (
+                len(events) < min_events_expected and (time.time() - start) < max_wait_s
+            ):
+                if sub.poll(timeout=100):  # 100ms timeout
+                    _, seq_bytes, payload = sub.recv_multipart()
+                    event_batch = decoder.decode(payload)
+                    for event in event_batch.events:
+                        events.append(event)
+
+            # Verify we received events
+            self.assertGreater(
+                len(events), 0, "Should have received at least one KV cache event"
+            )
+
+            # Track which blocks were stored and removed
+            stored_blocks = {}  # hash -> BlockStored event
+            removed_hashes = set()
+
+            # Validate event structure and relationships
+            for event in events:
+                self.assertIsInstance(
+                    event,
+                    (BlockStored, BlockRemoved, AllBlocksCleared),
+                    f"Event should be a KV cache event, got {type(event)}",
+                )
+
+                if isinstance(event, BlockStored):
+                    # Validate BlockStored structure
+                    self.assertIsInstance(event.block_hashes, list)
+                    self.assertEqual(
+                        len(event.block_hashes), 1, "Should have one hash per block"
+                    )
+                    self.assertIsInstance(event.token_ids, list)
+                    self.assertEqual(
+                        event.block_size,
+                        len(event.token_ids),
+                        "block_size should match token_ids length",
+                    )
+                    self.assertIsNone(
+                        event.lora_id, "lora_id should be None for basic test"
+                    )
+
+                    # Store this block for later validation
+                    block_hash = event.block_hashes[0]
+                    stored_blocks[block_hash] = event
+
+                    # If parent_block_hash is set, verify it was stored earlier
+                    if event.parent_block_hash is not None:
+                        # Parent should either be in stored_blocks or could be from a previous request
+                        pass  # Don't strictly enforce this as root blocks may have synthetic parents
+
+                elif isinstance(event, BlockRemoved):
+                    # Validate BlockRemoved structure
+                    self.assertIsInstance(event.block_hashes, list)
+                    self.assertEqual(
+                        len(event.block_hashes), 1, "Should have one hash per block"
+                    )
+                    removed_hashes.add(event.block_hashes[0])
+
+            # Verify we got both BlockStored and BlockRemoved events
+            self.assertGreater(
+                len(stored_blocks), 0, "Should have at least one BlockStored event"
+            )
+            # BlockRemoved events may not always occur in this short test, so just check if they do occur
+            # that they reference previously stored blocks
+            for removed_hash in removed_hashes:
+                # It's OK if the removed block wasn't in our stored_blocks
+                # (it could have been stored before we started listening)
+                pass
+
+        finally:
+            sub.close()
+            context.term()
+            kill_process_tree(process.pid)
+
+    def test_kv_events_attn_dp(self):
+        """Test that kv events are properly tagged with DP rank in attention DP mode"""
+
+        # Launch multiple subscribers for different DP ranks
+        decoder = Decoder(type=KVEventBatch)
+        context = zmq.Context()
+
+        # Subscribe to both DP rank endpoints
+        sub_dp0 = context.socket(zmq.SUB)
+        sub_dp0.connect("tcp://localhost:5557")  # DP rank 0
+        topic = "kv-events"
+        sub_dp0.setsockopt_string(zmq.SUBSCRIBE, topic)
+
+        sub_dp1 = context.socket(zmq.SUB)
+        sub_dp1.connect("tcp://localhost:5558")  # DP rank 1 (offset by rank)
+        sub_dp1.setsockopt_string(zmq.SUBSCRIBE, topic)
+
+        # Launch sglang server with DP attention enabled
+        process = popen_launch_server(
+            DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--kv-events-config",
+                '{"publisher": "zmq", "topic": "kv-events"}',
+                "--max-total-tokens",
+                64,
+                "--cuda-graph-max-bs",
+                4,
+                "--enable-dp-attention",
+                "--dp-size",
+                2,
+                "--tp-size",
+                2,
+            ],
+        )
+
+        try:
+            # Make requests to generate events
+            response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
+            self.assertEqual(response.status_code, 200)
+
+            # Send multiple requests to trigger events from both DP ranks
+            for i in range(4):
+                response = requests.post(
+                    f"{DEFAULT_URL_FOR_TEST}/generate",
+                    json={
+                        "text": f"Request {i}: The capital of country {i} is",
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 16,
+                        },
+                    },
+                )
+
+            # Collect events from both DP ranks
+            events_dp0 = []
+            events_dp1 = []
+            start = time.time()
+            max_wait_s = 10
+            min_events_per_rank = 3  # Expect at least a few events from each rank
+
+            while (time.time() - start) < max_wait_s and (
+                len(events_dp0) < min_events_per_rank
+                or len(events_dp1) < min_events_per_rank
+            ):
+                # Check DP rank 0
+                if sub_dp0.poll(timeout=100):  # 100ms timeout
+                    _, seq_bytes, payload = sub_dp0.recv_multipart()
+                    event_batch = decoder.decode(payload)
+                    print(
+                        f"DP Rank 0 - EventBatch: ts={event_batch.ts}, attn_dp_rank={event_batch.attn_dp_rank}"
+                    )
+                    self.assertEqual(
+                        event_batch.attn_dp_rank,
+                        0,
+                        "DP rank 0 events should have attn_dp_rank=0",
+                    )
+                    for event in event_batch.events:
+                        print(f"  DP0 - {event}")
+                        events_dp0.append(event)
+
+                # Check DP rank 1
+                if sub_dp1.poll(timeout=100):  # 100ms timeout
+                    _, seq_bytes, payload = sub_dp1.recv_multipart()
+                    event_batch = decoder.decode(payload)
+                    print(
+                        f"DP Rank 1 - EventBatch: ts={event_batch.ts}, attn_dp_rank={event_batch.attn_dp_rank}"
+                    )
+                    self.assertEqual(
+                        event_batch.attn_dp_rank,
+                        1,
+                        "DP rank 1 events should have attn_dp_rank=1",
+                    )
+                    for event in event_batch.events:
+                        print(f"  DP1 - {event}")
+                        events_dp1.append(event)
+
+            # Verify we got events from both DP ranks
+            print(f"Collected {len(events_dp0)} events from DP rank 0")
+            print(f"Collected {len(events_dp1)} events from DP rank 1")
+
+            self.assertGreaterEqual(
+                len(events_dp0),
+                min_events_per_rank,
+                f"Expected at least {min_events_per_rank} events from DP rank 0",
+            )
+            self.assertGreaterEqual(
+                len(events_dp1),
+                min_events_per_rank,
+                f"Expected at least {min_events_per_rank} events from DP rank 1",
+            )
+
+            # Verify event types are as expected
+            for events in [events_dp0, events_dp1]:
+                for event in events:
+                    self.assertIsInstance(
+                        event,
+                        (BlockStored, BlockRemoved, AllBlocksCleared),
+                        f"Event should be a KV cache event, got {type(event)}",
+                    )
+
+        finally:
+            sub_dp0.close()
+            sub_dp1.close()
+            context.term()
+            kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_logprobs.py

@@ -0,0 +1,527 @@
+"""
+Logprobs Accuracy Test for SGLang
+
+======================
+With deterministic/batch invariant kernels, we can ensure that SGLang produces exactly the same
+logprobs results for identical inputs. However, logprobs are highly sensitive to GPU hardware,
+kernels, torch versions, and other factors, so we cannot maintain a unified logprobs baseline
+across different machines.
+
+This test is designed to be run locally by contributors to verify logprobs accuracy
+before making changes to related code.
+When submitting changes that affect logprobs computation, please:
+1. Generate baseline
+2. Run test
+3. Submit results
+
+We really appreciate your effort and contribution to SGLang!
+
+======================
+What does this test do?
+This test fetches 1000 samples from the ShareGPT dataset, generates logprobs for each sample,
+and saves them as a baseline. Then, by running the test mode, it validates the accuracy of
+logprobs by comparing them against the baseline.
+
+This test ensures that:
+- the boundary of log probs requests are correct, eg, the index for tokens that required log probs are strictly followed
+- logprobs remain invariant between test runs, and also before and after your code changes;
+
+======================
+Usage
+
+Step 1: Generate Baseline (Before Code Changes)
+```bash
+python test/manual/test_logprobs.py gen
+```
+
+Step 2: Test Against Baseline (After Code Changes)
+```bash
+python test/manual/test_logprobs.py test
+```
+This tests your changes against the locally generated baseline from Step 1.
+The test passes if the maximum and mean differences are within the tolerance thresholds.
+======================
+"""
+
+import argparse
+import json
+import os
+import pickle
+import random
+import unittest
+
+import numpy as np
+import requests
+import torch
+from transformers import AutoTokenizer
+
+import sglang as sgl
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+# Configuration
+DENSE_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+SHAREGPT_URL = (
+    "https://huggingface.co/datasets/anon8231489123/"
+    "ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
+)
+
+# Hardware-specific configuration
+if torch.version.cuda is not None:
+    print("Running on NVIDIA CUDA GPU")
+    DENSE_TOLERANCE_MAX_DIFF = 1e-5
+    DENSE_TOLERANCE_MEAN_DIFF = 1e-5
+else:
+    print("No GPU backend (CPU only)")
+    raise ValueError("No GPU backend (CPU only)")
+
+# Common configuration
+TOP_K = 20
+NUM_SAMPLES = 1000
+LOGPROB_SAMPLE_RATIO = 0.5
+TEMPERATURE = 1.0
+MAX_LEN = 20000
+
+# Default output files
+DEFAULT_BASELINE_PKL = "sglang_baseline_local.pkl"
+DEFAULT_META_JSON = "baseline_meta_preview.json"
+
+# Default engine configuration
+DEFAULT_ENGINE_CONFIG = {
+    "model_path": DENSE_MODEL_NAME,
+    "random_seed": 42,
+    "skip_tokenizer_init": True,
+    "mem_fraction_static": 0.8,
+    "enable_deterministic_inference": True,
+    "attention_backend": "flashinfer",
+}
+
+
+def generate_baseline(
+    baseline_file=DEFAULT_BASELINE_PKL,
+    meta_file=DEFAULT_META_JSON,
+    num_samples=NUM_SAMPLES,
+):
+    """Generate a local baseline for logprobs testing.
+
+    Args:
+        baseline_file: Path to save the baseline pickle file
+        meta_file: Path to save the metadata preview JSON file
+        num_samples: Number of samples to generate
+    """
+    print(f"SGLang version: {sgl.__version__}")
+    print("Downloading ShareGPT dataset...")
+
+    # Download ShareGPT dataset
+    try:
+        response = requests.get(SHAREGPT_URL, timeout=30)
+        response.raise_for_status()
+        data = response.json()
+        print(f"Dataset size: {len(data)}")
+    except requests.exceptions.RequestException as e:
+        raise Exception(f"Failed to download ShareGPT dataset: {e}") from e
+
+    # Filter and prepare texts
+    texts = []
+    for s in data:
+        if "conversations" in s and len(s["conversations"]) > 0:
+            try:
+                text = s["conversations"][0]["value"]
+                if isinstance(text, str) and len(text) <= MAX_LEN and len(text) >= 5500:
+                    texts.append(text)
+                    if len(texts) >= num_samples * 40:  # Get more samples for filtering
+                        break
+            except (KeyError, IndexError, TypeError) as e:
+                print(f"Warning: Skipping invalid conversation data: {e}")
+                continue
+
+    if not texts:
+        raise ValueError("No valid texts found in the dataset")
+
+    print(f"Loading tokenizer for {DENSE_MODEL_NAME}...")
+    tokenizer = AutoTokenizer.from_pretrained(DENSE_MODEL_NAME, use_fast=True)
+
+    rng = np.random.default_rng(42)
+
+    print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
+    engine = sgl.Engine(
+        model_path=DENSE_MODEL_NAME,
+        attention_backend="flashinfer",
+        enable_deterministic_inference=True,
+        random_seed=42,
+        skip_tokenizer_init=True,
+        mem_fraction_static=0.8,
+        max_running_requests=1,
+    )
+
+    records = []
+    prompt_lengths = []
+
+    try:
+        for i, text in enumerate(texts):
+            if len(records) >= num_samples:
+                break
+
+            try:
+                ids = tokenizer.encode(text, add_special_tokens=False)
+                if len(ids) < 5:
+                    continue
+
+                start_pos = int(rng.integers(0, max(1, len(ids) - 3)))
+
+                outputs = engine.generate(
+                    input_ids=[ids],
+                    sampling_params={
+                        "temperature": 1.0,
+                        "top_p": 1.0,
+                        "top_k": TOP_K,
+                        "max_new_tokens": 1,
+                    },
+                    return_logprob=True,
+                    logprob_start_len=start_pos,
+                    top_logprobs_num=TOP_K,
+                )
+                meta = outputs[0]["meta_info"]
+
+                records.append(
+                    dict(id=i, text=text, ids=ids, start_pos=start_pos, meta=meta)
+                )
+                prompt_lengths.append(len(ids))
+
+                if (i + 1) % 50 == 0:
+                    print(f"Processed {len(records)}/{num_samples} samples")
+
+            except Exception as e:
+                print(f"Warning: Failed to process sample {i}: {e}")
+                continue
+
+        if not records:
+            raise RuntimeError(
+                "Failed to generate any baseline records. Please check the warnings above for errors."
+            )
+
+        # Save baseline files
+        with open(baseline_file, "wb") as f:
+            pickle.dump(records, f)
+        with open(meta_file, "w", encoding="utf-8") as f:
+            json.dump(records[:2], f, ensure_ascii=False, indent=2)
+
+        print(f"✅ Saved {len(records)} samples to {baseline_file}")
+        print(f"✅ Meta preview saved to {meta_file}")
+
+        if prompt_lengths:
+            avg_prompt_length = sum(prompt_lengths) / len(prompt_lengths)
+            print(f"📊 Average prompt length: {avg_prompt_length:.2f} tokens")
+
+    finally:
+        engine.shutdown()
+        torch.cuda.empty_cache()
+
+
+class TestLogprobsDense(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        """Set up the test class - initialize the engine once for all tests."""
+        print(f"Launching SGLang Engine with {DENSE_MODEL_NAME}...")
+        cls.engine = sgl.Engine(**DEFAULT_ENGINE_CONFIG)
+
+    @classmethod
+    def tearDownClass(cls):
+        """Clean up after all tests - shutdown the engine."""
+        cls.engine.shutdown()
+        torch.cuda.empty_cache()
+
+    @classmethod
+    def restart_engine_with_config(cls, **kwargs):
+        """Create engine with custom configuration"""
+        # Safely shutdown existing engine
+        cls.engine.shutdown()
+        torch.cuda.empty_cache()
+
+        # Set chunk size
+        chunk_size = kwargs.pop("chunk_size", None)
+        if chunk_size is not None:
+            print(f"Setting chunk size to {chunk_size}")
+            os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True"
+            os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size)
+        else:
+            os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False"
+
+        # Create engine with merged configuration
+        engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
+        cls.engine = sgl.Engine(**engine_config)
+
+    def load_test_data(self, baseline_file=None):
+        """Load test data from local baseline file. In test mode, only local baseline is supported."""
+        if not baseline_file:
+            raise ValueError("baseline_file is required in test mode")
+
+        if not os.path.exists(baseline_file):
+            raise FileNotFoundError(
+                f"Baseline file not found: {baseline_file}. Please run 'gen' mode first to generate the baseline."
+            )
+
+        print(f"Loading local baseline from {baseline_file}...")
+        try:
+            with open(baseline_file, "rb") as f:
+                records = pickle.load(f)
+            print(f"Successfully loaded {len(records)} records from local baseline")
+            return records
+        except (IOError, pickle.PickleError) as e:
+            raise Exception(f"Failed to load local baseline: {e}") from e
+
+    def compare_meta(self, baseline_meta, sglang_meta):
+        """Compare metadata between two outputs and return max and mean differences."""
+        diffs = []
+        for key in ["input_top_logprobs", "output_top_logprobs"]:
+            baseline_logprobs, sglang_logprobs = baseline_meta[key], sglang_meta[key]
+            self.assertEqual(
+                len(baseline_logprobs),
+                len(sglang_logprobs),
+                f"Length of {key} is not equal, sglang did not return the correct number of log probs(should be top 20)",
+            )
+            for baseline_entry, sglang_entry in zip(baseline_logprobs, sglang_logprobs):
+                if not baseline_entry or not sglang_entry:
+                    continue
+                baseline_token_map = {tid: lp for lp, tid, _ in baseline_entry}
+                sglang_token_map = {tid: lp for lp, tid, _ in sglang_entry}
+                common_tokens = baseline_token_map.keys() & sglang_token_map.keys()
+                self.assertGreaterEqual(
+                    len(common_tokens),
+                    TOP_K,
+                    f"there are only {len(common_tokens)} common topk tokens that matches",
+                )
+                for token_id in common_tokens:
+                    diffs.append(
+                        abs(baseline_token_map[token_id] - sglang_token_map[token_id])
+                    )
+        if not diffs:
+            return 0.0, 0.0
+        return max(diffs), float(np.mean(diffs))
+
+    def test_logprobs_comparison(self, baseline_file=None):
+        """Test the logprobs comparison functionality with different parameter combinations."""
+        # Load test data with retry mechanism
+        records = self.load_test_data(baseline_file)
+
+        # Fast configs for CI
+        test_configs = [
+            {"num_samples": NUM_SAMPLES},
+            {"num_samples": 42, "chunk_size": 1, "max_running_requests": 16},
+            {"num_samples": 42, "chunk_size": 2, "max_running_requests": 16},
+            {"num_samples": 42, "chunk_size": 3, "max_running_requests": 16},
+            {"num_samples": NUM_SAMPLES, "chunk_size": 16, "max_running_requests": 128},
+            {"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 16},
+            {"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 8},
+            {"num_samples": NUM_SAMPLES, "chunk_size": 128, "max_running_requests": 32},
+            {
+                "num_samples": NUM_SAMPLES,
+                "chunk_size": 128,
+                "max_running_requests": 128,
+            },
+            {"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 8},
+            {"num_samples": NUM_SAMPLES, "chunk_size": 256, "max_running_requests": 32},
+            {
+                "num_samples": NUM_SAMPLES,
+                "chunk_size": 256,
+                "max_running_requests": 128,
+            },
+        ]
+
+        # Run tests
+        for config in test_configs:
+            with self.subTest(config=config):
+                print(f"Testing with config: {config}")
+
+                # Sample records for this config
+                num_samples = config.get("num_samples", NUM_SAMPLES)
+                test_records = random.sample(records, k=min(num_samples, len(records)))
+                random.shuffle(test_records)
+
+                # Calculate how many samples should return logprobs
+                logprob_count = int(len(test_records) * LOGPROB_SAMPLE_RATIO)
+                print(
+                    f"Testing with {len(test_records)} samples, temperature={TEMPERATURE}"
+                )
+                print(
+                    f"Will return logprobs for {logprob_count} samples (ratio: {LOGPROB_SAMPLE_RATIO})"
+                )
+
+                all_max, all_mean = [], []
+                logprob_returned_count = 0
+
+                # Process all records at once
+                input_ids = [rec["ids"] for rec in test_records]
+                logprob_start_lens = [rec["start_pos"] for rec in test_records]
+
+                # Determine which samples should return logprobs (randomly selected)
+                logprob_indices = set(
+                    random.sample(range(len(test_records)), logprob_count)
+                )
+                return_logprob_array = [
+                    sample_idx in logprob_indices
+                    for sample_idx in range(len(test_records))
+                ]
+
+                # Sampling param per request
+                sampling_params = [
+                    {
+                        "temperature": TEMPERATURE,
+                        "top_p": 1.0,
+                        "top_k": TOP_K,
+                        "max_new_tokens": 1,
+                    }
+                    for _ in test_records
+                ]
+
+                # Some configs must restart the engine to take effect
+                chunk_size = config.get("chunk_size", None)
+                max_running_requests = config.get("max_running_requests", None)
+                if chunk_size is not None or max_running_requests is not None:
+                    self.restart_engine_with_config(
+                        chunk_size=chunk_size,
+                        max_running_requests=max_running_requests,
+                    )
+
+                outputs = self.engine.generate(
+                    input_ids=input_ids,
+                    sampling_params=sampling_params,
+                    return_logprob=return_logprob_array,
+                    logprob_start_len=logprob_start_lens,
+                    top_logprobs_num=TOP_K,
+                )
+
+                for sample_idx, (rec, output) in enumerate(zip(test_records, outputs)):
+                    # Only compare logprobs for samples that should have them
+                    if sample_idx in logprob_indices:
+                        # Safe access to meta_info and input_top_logprobs
+                        meta_info = output.get("meta_info")
+                        input_top_logprobs = (
+                            meta_info.get("input_top_logprobs") if meta_info else None
+                        )
+
+                        self.assertIsNotNone(
+                            input_top_logprobs,
+                            f"return_logprob enabled on this sample, but input_top_logprobs is None (length: {len(input_top_logprobs) if input_top_logprobs is not None else 'N/A'})",
+                        )
+                        baseline_meta = rec["meta"]
+                        sglang_meta = meta_info
+
+                        max_diff, mean_diff = self.compare_meta(
+                            baseline_meta, sglang_meta
+                        )
+                        all_max.append(max_diff)
+                        all_mean.append(mean_diff)
+                        logprob_returned_count += 1
+                    else:
+                        # Verify that logprobs were not returned for this sample
+                        meta_info = output.get("meta_info")
+                        input_top_logprobs = (
+                            meta_info.get("input_top_logprobs") if meta_info else None
+                        )
+                        output_token_ids_logprobs = (
+                            meta_info.get("output_token_ids_logprobs")
+                            if meta_info
+                            else None
+                        )
+
+                        self.assertFalse(
+                            input_top_logprobs,
+                            f"return_logprob is disabled on this sample, Sample {sample_idx} should not have logprobs, content: {output_token_ids_logprobs}",
+                        )
+
+                max_of_max = max(all_max) if all_max else 0.0
+                mean_of_mean = np.mean(all_mean) if all_mean else 0.0
+
+                print(f"max Δ={max_of_max:.6g}")
+                print(f"mean Δ={mean_of_mean:.6g}")
+                print(
+                    f"logprobs returned for {logprob_returned_count} samples (expected: {logprob_count})"
+                )
+
+                # Verify correct number of logprobs returned
+                self.assertEqual(
+                    logprob_returned_count,
+                    logprob_count,
+                    f"Expected {logprob_count} samples with logprobs, got {logprob_returned_count}",
+                )
+
+                # Basic validation
+                self.assertIsInstance(all_max, list)
+                self.assertIsInstance(all_mean, list)
+                self.assertGreater(
+                    len(all_max),
+                    0,
+                    f"No test samples processed for config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}}",
+                )
+
+                # Tolerance checks with clear error messages
+                failed_samples = []
+                for sample_idx, (max_diff, mean_diff) in enumerate(
+                    zip(all_max, all_mean)
+                ):
+                    if max_diff > DENSE_TOLERANCE_MAX_DIFF:
+                        failed_samples.append(
+                            f"Sample {sample_idx}: max_diff={max_diff:.6g} > {DENSE_TOLERANCE_MAX_DIFF}"
+                        )
+                    if mean_diff > DENSE_TOLERANCE_MEAN_DIFF:
+                        failed_samples.append(
+                            f"Sample {sample_idx}: mean_diff={mean_diff:.6g} > {DENSE_TOLERANCE_MEAN_DIFF}"
+                        )
+
+                if failed_samples:
+                    self.fail(
+                        f"Config {{'num_samples': {NUM_SAMPLES}, 'logprob_sample_ratio': {LOGPROB_SAMPLE_RATIO}, 'temperature': {TEMPERATURE}}} - Tolerance exceeded in {len(failed_samples)} samples:\n"
+                        + "\n".join(failed_samples[:5])
+                    )
+
+
+def main():
+    """Main function to handle command line arguments and run either generation or testing."""
+    parser = argparse.ArgumentParser(
+        description="SGLang Logprobs Test and Baseline Generation"
+    )
+    parser.add_argument(
+        "mode",
+        choices=["gen", "test"],
+        help="Mode to run: 'gen' to generate baseline, 'test' to run tests",
+    )
+
+    args = parser.parse_args()
+
+    if args.mode == "gen":
+        print("🚀 Generating baseline...")
+        generate_baseline()
+        print(f"\n✅ Baseline generation complete!")
+        print(f"📁 Baseline saved to: {DEFAULT_BASELINE_PKL}")
+        print(f"📁 Metadata preview saved to: {DEFAULT_META_JSON}")
+        print(f"\n💡 Next steps:")
+        print(f"   1. Make your code changes")
+        print(f"   2. Run: python {__file__} test")
+
+    elif args.mode == "test":
+        print("🧪 Running logprobs test...")
+        if not os.path.exists(DEFAULT_BASELINE_PKL):
+            print(f"❌ Baseline file not found: {DEFAULT_BASELINE_PKL}")
+            print(f"💡 Generate baseline first by running:")
+            print(f"   python {__file__} gen")
+            print(f"   This will download ShareGPT data and generate a local baseline.")
+            return 1
+
+        # Set environment variable for testing
+        os.environ["RETURN_ORIGINAL_LOGPROB"] = "True"
+
+        # Create test instance and run
+        test_instance = TestLogprobsDense()
+        test_instance.setUpClass()
+        try:
+            test_instance.test_logprobs_comparison(baseline_file=DEFAULT_BASELINE_PKL)
+            print("\n✅ Test completed successfully!")
+        finally:
+            test_instance.tearDownClass()
+
+    return 0
+
+
+if __name__ == "__main__":
+    exit(main())

third_party/sglang/test/manual/test_mla_tp.py

@@ -0,0 +1,66 @@
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestDeepseekTP2(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "lmsys/sglang-ci-dsv3-test"
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = ["--trust-remote-code"]
+        if torch.cuda.is_available() and torch.version.cuda:
+            other_args.extend(
+                ["--tp", "2", "--enable-torch-compile", "--cuda-graph-max-bs", "2"]
+            )
+        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):
+        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)
+        self.assertGreater(metrics["score"], 0.62)
+
+    def test_gsm8k_bs1(self):
+        # test torch compile accuracy for bs=1
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=10,
+            num_threads=1,
+        )
+        metrics = run_eval(args)
+        self.assertGreater(metrics["score"], 0.62)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_modelopt.py

@@ -0,0 +1,58 @@
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8,
+    DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8_REVISION,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestEvalFP8ModelOptQuantAccuracy(CustomTestCase):
+
+    def _run_test(self, model, other_args, expected_score):
+        base_url = DEFAULT_URL_FOR_TEST
+        other_args = other_args or []
+
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                model=model,
+                eval_name="mmlu",
+                num_examples=64,
+                num_threads=32,
+                temperature=0.1,
+            )
+
+            metrics = run_eval(args)
+            self.assertGreaterEqual(metrics["score"], expected_score)
+        finally:
+            kill_process_tree(process.pid)
+
+    @unittest.skipIf(
+        torch.version.hip is not None, "modelopt quantization unsupported on ROCm"
+    )
+    def test_mmlu_offline_only(self):
+        """Test with offline quantization only."""
+        self._run_test(
+            model=DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8,
+            other_args=[
+                "--revision",
+                DEFAULT_MODEL_NAME_FOR_MODELOPT_QUANT_ACCURACY_TEST_FP8_REVISION,
+            ],
+            expected_score=0.64,
+        )

third_party/sglang/test/manual/test_modelopt_fp8kvcache.py

@@ -0,0 +1,29 @@
+import unittest
+
+from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp8Config,
+    ModelOptFp8KVCacheMethod,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestModelOptFp8KVCacheMethod(CustomTestCase):
+    def test_kv_cache_method_initialization(self):
+        """Test that ModelOptFp8KVCacheMethod can be instantiated and
+        inherits from BaseKVCacheMethod."""
+        # Create a ModelOptFp8Config object
+        quant_config = ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
+
+        # Instantiate the KV cache method
+        kv_cache_method = ModelOptFp8KVCacheMethod(quant_config)
+
+        # Check inheritance
+        self.assertIsInstance(kv_cache_method, BaseKVCacheMethod)
+
+        # Check that the quant_config is stored
+        self.assertEqual(kv_cache_method.quant_config, quant_config)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_models_from_modelscope.py

@@ -0,0 +1,40 @@
+import os
+import shutil
+import subprocess
+import unittest
+from unittest import mock
+
+from sglang.srt.utils import prepare_model_and_tokenizer
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestDownloadFromModelScope(CustomTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "iic/nlp_lstmcrf_word-segmentation_chinese-news"
+        stat, output = subprocess.getstatusoutput("pip install modelscope")
+
+        cls.with_modelscope_environ = {k: v for k, v in os.environ.items()}
+        cls.with_modelscope_environ["SGLANG_USE_MODELSCOPE"] = "True"
+
+    @classmethod
+    def tearDownClass(cls):
+        pass
+
+    def test_prepare_model_and_tokenizer(self):
+        from modelscope.utils.file_utils import get_model_cache_root
+
+        model_cache_root = get_model_cache_root()
+        if os.path.exists(model_cache_root):
+            shutil.rmtree(model_cache_root)
+        with mock.patch.dict(os.environ, self.with_modelscope_environ, clear=True):
+            model_path, tokenizer_path = prepare_model_and_tokenizer(
+                self.model, self.model
+            )
+            assert os.path.exists(os.path.join(model_path, "pytorch_model.bin"))
+            assert os.path.exists(os.path.join(tokenizer_path, "config.json"))
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_mori_transfer_engine_e2e.py

@@ -0,0 +1,293 @@
+import os
+import subprocess
+import unittest
+
+import requests
+
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_pd_server,
+)
+
+
+class TestMoriTransferEngineE2E(PDDisaggregationServerBase):
+    """
+    Run:
+        SGLANG_MORI_MANUAL_E2E=1 python3 test/manual/test_mori_transfer_engine_e2e.py
+
+    Optional:
+    - SGLANG_MORI_E2E_TEST_MODEL: override model (defaults to a small test model)
+    - SGLANG_TEST_PD_DISAGG_DEVICES: RDMA devices string, e.g. "mlx5_roce0,mlx5_roce4"
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"):
+            raise unittest.SkipTest(
+                "Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test."
+            )
+
+        try:
+            import torch
+
+            if not torch.cuda.is_available():
+                raise unittest.SkipTest("torch.cuda is not available.")
+        except Exception as e:
+            raise unittest.SkipTest(f"torch is not available/usable: {e}")
+
+        # Force the disaggregation fixture to use MORI backend in local/manual runs.
+        os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori"
+
+        super().setUpClass()
+
+        cls.model = os.environ.get(
+            "SGLANG_MORI_E2E_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        )
+
+        cls.start_prefill()
+        cls.start_decode()
+
+        cls.wait_server_ready(
+            cls.prefill_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_prefill,
+        )
+        cls.wait_server_ready(
+            cls.decode_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_decode,
+        )
+
+        cls.launch_lb()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None)
+        super().tearDownClass()
+
+    @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 = subprocess.Popen(lb_command, stdout=None, stderr=None)
+        cls.wait_server_ready(
+            cls.lb_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_lb,
+        )
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "1",
+        ]
+        prefill_args += 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,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+        ]
+        decode_args += 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,
+        )
+
+    def test_generate_smoke(self):
+        resp = requests.post(
+            self.lb_url + "/generate",
+            json={
+                "text": "Hello",
+                "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+            },
+            timeout=120,
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        out = resp.json()
+        self.assertIn("text", out)
+        self.assertIsInstance(out["text"], str)
+        self.assertGreater(len(out["text"]), 0)
+
+
+class TestMoriTransferEngineTPMismatchE2E(PDDisaggregationServerBase):
+    """Manual MORI PD-disaggregation E2E with TP mismatch.
+
+    Scenario:
+    - prefill: tp=2 (GPU 0-1)
+    - decode:  tp=4 (GPU 2-5)
+
+    Manual-only and requires >= 6 visible GPUs.
+    """
+
+    _PORT_DELTA = 10
+
+    @classmethod
+    def setUpClass(cls):
+        if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"):
+            raise unittest.SkipTest(
+                "Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test."
+            )
+
+        try:
+            import torch
+
+            if not torch.cuda.is_available():
+                raise unittest.SkipTest("torch.cuda is not available.")
+            if torch.cuda.device_count() < 6:
+                raise unittest.SkipTest(
+                    "TP-mismatch test requires >= 6 visible GPUs (prefill tp=2 + decode tp=4)."
+                )
+        except Exception as e:
+            raise unittest.SkipTest(f"torch is not available/usable: {e}")
+
+        os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori"
+        super().setUpClass()
+
+        # Shift ports to avoid clashing with TestMoriTransferEngineE2E.
+        cls.lb_port = str(int(cls.lb_port) + cls._PORT_DELTA)
+        cls.prefill_port = str(int(cls.prefill_port) + cls._PORT_DELTA)
+        cls.decode_port = str(int(cls.decode_port) + cls._PORT_DELTA)
+        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}"
+
+        cls.model = os.environ.get(
+            "SGLANG_MORI_E2E_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        )
+
+        cls.start_prefill()
+        cls.start_decode()
+
+        cls.wait_server_ready(
+            cls.prefill_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_prefill,
+        )
+        cls.wait_server_ready(
+            cls.decode_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_decode,
+        )
+        cls.launch_lb()
+
+    @classmethod
+    def tearDownClass(cls):
+        os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None)
+        super().tearDownClass()
+
+    @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 = subprocess.Popen(lb_command, stdout=None, stderr=None)
+        cls.wait_server_ready(
+            cls.lb_url + "/health",
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            process=cls.process_lb,
+        )
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "2",
+        ]
+        prefill_args += 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,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp",
+            "4",
+            "--base-gpu-id",
+            "2",
+        ]
+        decode_args += 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,
+        )
+
+    def test_generate_smoke_tp_mismatch(self):
+        resp = requests.post(
+            self.lb_url + "/generate",
+            json={
+                "text": "Hello",
+                "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+            },
+            timeout=120,
+        )
+        self.assertEqual(resp.status_code, 200, resp.text)
+        out = resp.json()
+        self.assertIn("text", out)
+        self.assertIsInstance(out["text"], str)
+        self.assertGreater(len(out["text"]), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_mscclpp.py

@@ -0,0 +1,196 @@
+"""For Now, MSCCL is only supported on TP16 and TP8 case
+
+if [[ $RANK -eq 0  ]]; then
+    ray start --block --head --port=6379 &
+    python3 test_mscclpp.py;
+else
+    ray start --block --address=${MASTER_ADDR}:6379;
+fi
+"""
+
+import os
+import random
+import socket
+import unittest
+from typing import Any
+
+import ray
+import torch
+import torch.distributed as dist
+
+from sglang.srt.distributed import init_distributed_environment
+from sglang.srt.distributed.communication_op import (  # noqa
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.distributed.parallel_state import (
+    get_tensor_model_parallel_group,
+    graph_capture,
+    initialize_model_parallel,
+    set_custom_all_reduce,
+    set_mscclpp_all_reduce,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+def get_open_port() -> int:
+    # try ipv4
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+    except OSError:
+        # try ipv6
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+
+def multi_process_parallel(
+    world_size: int,
+    master_addr: str,
+    cls: Any,
+    test_target: Any,
+) -> None:
+
+    # Using ray helps debugging the error when it failed
+    # as compared to multiprocessing.
+    # NOTE: We need to set working_dir for distributed tests,
+    # otherwise we may get import errors on ray workers
+
+    ray.init(log_to_driver=True)
+
+    distributed_init_port = get_open_port()
+    refs = []
+    for rank in range(world_size):
+        refs.append(
+            test_target.remote(
+                cls, world_size, master_addr, rank, distributed_init_port
+            )
+        )
+    ray.get(refs)
+
+    ray.shutdown()
+
+
+class TestMSCCLAllReduce(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        random.seed(42)
+        # 1KB to 1MB
+        cls.test_sizes = [512, 4096, 32768, 262144, 524288]
+        cls.world_sizes = [8]
+        TEST_TP16 = int(os.getenv("SGL_MSCCLPP_TEST_TP16", "0"))
+        if TEST_TP16:
+            cls.world_sizes = [16]
+        cls.test_loop = 10
+
+    def test_graph_allreduce(self):
+        TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
+        for world_size in self.world_sizes:
+            if world_size not in [8, 16]:
+                continue
+            multi_process_parallel(
+                world_size, TEST_MASTER_ADDR, self, self.graph_allreduce
+            )
+
+    def test_eager_allreduce(self):
+        TEST_MASTER_ADDR = os.getenv("SGL_MSCCLPP_TEST_MASTER_ADDR", "localhost")
+        for world_size in self.world_sizes:
+            if world_size not in [8, 16]:
+                continue
+            multi_process_parallel(
+                world_size, TEST_MASTER_ADDR, self, self.eager_allreduce
+            )
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def graph_allreduce(self, world_size, master_addr, rank, distributed_init_port):
+        del os.environ["CUDA_VISIBLE_DEVICES"]
+        device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
+        set_mscclpp_all_reduce(True)
+        set_custom_all_reduce(False)
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank % torch.cuda.device_count(),
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        # A small all_reduce for warmup.
+        # this is needed because device communicators might be created lazily
+        # (e.g. NCCL). This will ensure that the communicator is initialized
+        # before any communication happens, so that this group can be used for
+        # graph capture immediately.
+        data = torch.zeros(1)
+        data = data.to(device=device)
+        torch.distributed.all_reduce(data, group=group)
+        torch.cuda.synchronize()
+        del data
+
+        for sz in self.test_sizes:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                for _ in range(self.test_loop):
+                    with graph_capture() as graph_capture_context:
+                        # use integers so result matches NCCL exactly
+                        inp1 = torch.randint(
+                            1,
+                            16,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        inp2 = torch.randint(
+                            1,
+                            16,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        torch.cuda.synchronize()
+                        graph = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(
+                            graph, stream=graph_capture_context.stream
+                        ):
+                            out1 = tensor_model_parallel_all_reduce(inp1)
+                            # the input buffer is immediately modified to test
+                            # synchronization
+                            dist.all_reduce(inp1, group=group)
+                            out2 = tensor_model_parallel_all_reduce(inp2)
+                            dist.all_reduce(inp2, group=group)
+                    graph.replay()
+                    torch.testing.assert_close(out1, inp1)
+                    torch.testing.assert_close(out2, inp2)
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def eager_allreduce(self, world_size, master_addr, rank, distributed_init_port):
+        del os.environ["CUDA_VISIBLE_DEVICES"]
+        device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://{master_addr}:{distributed_init_port}"
+        set_mscclpp_all_reduce(True)
+        set_custom_all_reduce(False)
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        for sz in self.test_sizes:
+            for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+                for _ in range(self.test_loop):
+                    inp1 = torch.randint(
+                        1, 16, (sz,), dtype=dtype, device=torch.cuda.current_device()
+                    )
+                    out1 = tensor_model_parallel_all_reduce(inp1)
+                    dist.all_reduce(inp1, group=group)
+                    torch.testing.assert_close(out1, inp1)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_quick_allreduce.py

@@ -0,0 +1,303 @@
+import multiprocessing
+import os
+import random
+import socket
+import unittest
+from typing import Any
+
+import ray
+import torch
+import torch.distributed as dist
+
+import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
+from sglang.srt.distributed import init_distributed_environment
+from sglang.srt.distributed.communication_op import (  # noqa
+    tensor_model_parallel_all_reduce,
+)
+from sglang.srt.distributed.device_communicators.quick_all_reduce import (
+    qr_rocm_arch_available,
+)
+from sglang.srt.distributed.parallel_state import (
+    get_tensor_model_parallel_group,
+    graph_capture,
+    initialize_model_parallel,
+)
+from sglang.test.test_utils import CustomTestCase
+
+torch.manual_seed(42)
+random.seed(44)  # keep the deterministic seed
+
+
+def get_open_port() -> int:
+    # try ipv4
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+    except OSError:
+        # try ipv6
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+
+def multi_process_parallel(
+    world_size: int, cls: Any, test_target: Any, quant_mode: str
+) -> None:
+
+    # Using ray helps debugging the error when it failed
+    # as compared to multiprocessing.
+    # NOTE: We need to set working_dir for distributed tests,
+    # otherwise we may get import errors on ray workers
+
+    ray.init(log_to_driver=True)
+
+    distributed_init_port = get_open_port()
+    refs = []
+    for rank in range(world_size):
+        refs.append(
+            test_target.remote(cls, world_size, rank, distributed_init_port, quant_mode)
+        )
+    ray.get(refs)
+
+    ray.shutdown()
+
+
+class TestQuickAllReduce(CustomTestCase):
+    TEST_SIZES = [
+        2 * 1024 * 1024,
+        4 * 1024 * 1024,
+        8 * 1024 * 1024,
+        16 * 1024 * 1024,
+        32 * 1024 * 1024,
+    ]
+    TEST_LOOP = 5
+    # Too many configurations can lead to a test grid that is too large
+    # The tp takes too long to boot,let's just choose 4 out of 12 configurations
+    # WORLD_SIZES = [2, 4, 8]
+    # QUANT_MODE = ["FP", "INT8", "INT6", "INT4"]
+    QUANT_MODE_WORLD_SIZE_PART = [["FP", 8], ["INT4", 4], ["INT8", 2], ["INT6", 2]]
+
+    @unittest.skipIf(
+        not qr_rocm_arch_available(),
+        "Only test Quick AllReduce on ROCm architectures >= gfx94*",
+    )
+    def test_graph_allreduce(self):
+        for quant_mode_world_size_part in self.QUANT_MODE_WORLD_SIZE_PART:
+            quant_mode = quant_mode_world_size_part[0]
+            world_size = quant_mode_world_size_part[1]
+            if world_size > torch.cuda.device_count():
+                continue
+            multi_process_parallel(world_size, self, self.graph_allreduce, quant_mode)
+
+    @unittest.skipIf(
+        not qr_rocm_arch_available(),
+        "Only test Quick AllReduce on ROCm architectures >= gfx94*",
+    )
+    def test_eager_allreduce(self):
+        for quant_mode_world_size_part in self.QUANT_MODE_WORLD_SIZE_PART:
+            quant_mode = quant_mode_world_size_part[0]
+            world_size = quant_mode_world_size_part[1]
+            if world_size > torch.cuda.device_count():
+                continue
+            multi_process_parallel(world_size, self, self.eager_allreduce, quant_mode)
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def graph_allreduce(self, world_size, rank, distributed_init_port, quant_mode):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
+        os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
+        device = torch.device(f"cuda:{rank}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        # A small all_reduce for warmup.
+        # this is needed because device communicators might be created lazily
+        # (e.g. NCCL). This will ensure that the communicator is initialized
+        # before any communication happens, so that this group can be used for
+        # graph capture immediately.
+        data = torch.zeros(1)
+        data = data.to(device=device)
+        torch.distributed.all_reduce(data, group=group)
+        torch.cuda.synchronize()
+        del data
+
+        for sz in self.TEST_SIZES:
+            for dtype in [torch.float16, torch.bfloat16]:
+                for _ in range(self.TEST_LOOP):
+                    with graph_capture() as graph_capture_context:
+                        # use integers so result matches NCCL exactly
+                        inp1 = torch.randint(
+                            1,
+                            23,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        inp2 = torch.randint(
+                            -23,
+                            1,
+                            (sz,),
+                            dtype=dtype,
+                            device=torch.cuda.current_device(),
+                        )
+                        torch.cuda.synchronize()
+                        graph = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(
+                            graph, stream=graph_capture_context.stream
+                        ):
+                            out1 = tensor_model_parallel_all_reduce(inp1)
+                            # the input buffer is immediately modified to test
+                            # synchronization
+                            dist.all_reduce(inp1, group=group)
+                            out2 = tensor_model_parallel_all_reduce(inp2)
+                            dist.all_reduce(inp2, group=group)
+                    graph.replay()
+                    atol = 1.25 * world_size
+                    rtol = 0.5 * world_size
+                    for inp, out in [[inp1, out1], [inp2, out2]]:
+                        torch.testing.assert_close(out, inp, atol=atol, rtol=rtol)
+                        # try:
+                        #     torch.testing.assert_close(out, inp, atol=atol, rtol=rtol)
+                        # except AssertionError as e:
+                        #     print("Max abs diff:", (out - inp).abs().max())
+                        #     print("Max rel diff:", ((out - inp).abs() / inp.abs().clamp(min=1e-5)).max())
+
+    @ray.remote(num_gpus=1, max_calls=1)
+    def eager_allreduce(self, world_size, rank, distributed_init_port, quant_mode):
+        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+        os.environ["ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_mode
+        os.environ["ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "0"
+        device = torch.device(f"cuda:{rank}")
+        torch.cuda.set_device(device)
+        distributed_init_method = f"tcp://localhost:{distributed_init_port}"
+        init_distributed_environment(
+            world_size=world_size,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            local_rank=rank,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=world_size)
+        group = get_tensor_model_parallel_group().device_group
+
+        for sz in self.TEST_SIZES:
+            for dtype in [torch.float16, torch.bfloat16]:
+                for _ in range(self.TEST_LOOP):
+                    inp1 = torch.randint(
+                        1,
+                        23,
+                        (sz,),
+                        dtype=dtype,
+                        device=torch.cuda.current_device(),
+                    )
+                    out1 = tensor_model_parallel_all_reduce(inp1)
+                    dist.all_reduce(inp1, group=group)
+                    atol = 1.25 * world_size
+                    rtol = 0.5 * world_size
+                    torch.testing.assert_close(out1, inp1, atol=atol, rtol=rtol)
+                    # try:
+                    #     torch.testing.assert_close(out1, inp1, atol=atol, rtol=rtol)
+                    # except AssertionError as e:
+                    #     print("Max abs diff:", (out1 - inp1).abs().max())
+                    #     print("Max rel diff:", ((out1 - inp1).abs() / inp1.abs().clamp(min=1e-5)).max())
+
+
+def qr_variable_input(rank, world_size):
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+    qr_max_size = None  # MB
+    _ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
+    ranks = []
+    for i in range(world_size):
+        ranks.append(i)
+    dist.init_process_group(
+        backend="nccl",
+        init_method="tcp://127.0.0.1:29500",
+        rank=rank,
+        world_size=world_size,
+    )
+    cpu_group = torch.distributed.new_group(ranks, backend="nccl")
+
+    handle = ops.qr_get_handle(_ptr)
+    world_size = dist.get_world_size(group=cpu_group)
+    handles = [None] * world_size
+    dist.all_gather_object(handles, handle, group=cpu_group)
+    ops.qr_open_handles(_ptr, handles)
+
+    num = 1
+    s1 = 1024
+    while num < 50000:  # 50000 is sufficient to identify issues.
+        dtype = torch.float16
+        if num % 2 == 0:
+            s2 = 1024
+            inp1 = torch.zeros(
+                (s1, s2), dtype=dtype, device=torch.cuda.current_device()
+            )
+        else:
+            s2 = 2048
+            inp1 = torch.ones((s1, s2), dtype=dtype, device=torch.cuda.current_device())
+        result = torch.empty_like(inp1)
+        # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4
+        ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
+        try:
+            if inp1[0, 0] == 0:
+                assert torch.all(result == 0)
+            else:
+                assert torch.all(result == world_size)
+        except AssertionError:
+            print("Assertion failed! Allreduce results are incorrect.")
+            raise
+        num += 1
+
+
+class TestQuickreduceVariableInput(CustomTestCase):
+    """
+    When the tensor parallelism is set to 4 or 8, frequent changes
+    in the input shape can cause QuickReduce to hang (this issue
+    has been observed with the gpt_oss model).
+    """
+
+    TP_SIZES = [4, 8]
+
+    @unittest.skipIf(
+        not qr_rocm_arch_available(),
+        "Only test Quick AllReduce on ROCm architectures >= gfx94*",
+    )
+    def test_custom_quick_allreduce_variable_input(self):
+        for tp_size in self.TP_SIZES:
+            world_size = tp_size
+            if world_size > torch.cuda.device_count():
+                return
+
+            multiprocessing.set_start_method("spawn", force=True)
+            # 90s is enough
+            timeout = 90
+            processes = []
+            for rank in range(tp_size):
+                p = multiprocessing.Process(
+                    target=qr_variable_input, args=(rank, tp_size)
+                )
+                p.start()
+                processes.append((rank, p))
+            for rank, p in processes:
+                p.join(timeout=timeout)
+                if p.is_alive():
+                    for r, proc in processes:
+                        if proc.is_alive():
+                            proc.terminate()
+                            proc.join()
+                    raise RuntimeError(
+                        f"QuickReduce hang detected after {timeout} seconds!"
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_ray_engine.py

@@ -0,0 +1,461 @@
+"""Integration tests for RayEngine and Ray HTTP server (requires GPU + Ray).
+
+Tests the Ray actor scheduler backend:
+  - Offline inference via Engine(use_ray=True) inside a Ray actor on a placement group
+  - Error paths in RayEngine._launch_scheduler_processes()
+  - HTTP server launched via --use-ray flag
+
+Usage:
+    # 1-GPU tests
+    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP1 -v -s
+    python -m pytest test/manual/test_ray_engine.py::TestRayEngineErrors -v -s
+    python -m pytest test/manual/test_ray_engine.py::TestRayHTTPServerTP1 -v -s
+
+    # 2-GPU tests
+    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP2 -v -s
+    python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflinePP2 -v -s
+"""
+
+from __future__ import annotations
+
+import os
+import time
+import unittest
+
+import torch
+
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+# Allow overriding the model via env var for environments without gated access
+_MODEL = os.environ.get("SGLANG_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+
+try:
+    import ray
+    from ray.runtime_env import RuntimeEnv
+    from ray.util.placement_group import placement_group
+    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+    # Prevent Ray from overriding CUDA_VISIBLE_DEVICES so that all GPUs
+    # remain visible inside actors regardless of num_gpus allocation.
+    _env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
+    if os.environ.get("HF_TOKEN"):
+        _env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
+    _RAY_RUNTIME_ENV = RuntimeEnv(env_vars=_env_vars)
+    _has_ray = True
+except ImportError:
+    _has_ray = False
+    _RAY_RUNTIME_ENV = None
+
+
+_NUM_GPUS = torch.cuda.device_count()
+
+_SAMPLING_PARAMS = {"max_new_tokens": 32, "temperature": 0.0}
+
+_PROMPTS = [
+    "The capital of France is",
+    "Explain quantum computing in simple terms:",
+    "Write a haiku about programming:",
+    "What is 2 + 2?",
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _create_engine_on_pg(tp_size, pp_size=1, model=_MODEL, extra_kwargs=None):
+    """Create an EngineActor on a placement group and wait for it to be ready.
+
+    Returns (engine_actor, placement_group).
+    """
+
+    @ray.remote
+    class EngineActor:
+        def __init__(self, **kwargs):
+            from sglang.srt.ray.engine import RayEngine
+
+            self.engine = RayEngine(**kwargs)
+
+        def is_ready(self):
+            return True
+
+        def generate(self, prompt, sampling_params):
+            return self.engine.generate(prompt=prompt, sampling_params=sampling_params)
+
+        def shutdown(self):
+            if self.engine:
+                self.engine.shutdown()
+                self.engine = None
+
+    total_gpus = tp_size * pp_size
+    pg = placement_group(
+        [{"CPU": 1, "GPU": total_gpus}],
+        strategy="STRICT_PACK",
+    )
+    ray.get(pg.ready())
+
+    kwargs = dict(
+        model_path=model,
+        tp_size=tp_size,
+        pp_size=pp_size,
+    )
+    if extra_kwargs:
+        kwargs.update(extra_kwargs)
+
+    actor = EngineActor.options(
+        num_cpus=1,
+        num_gpus=0,
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_bundle_index=0,
+        ),
+    ).remote(**kwargs)
+
+    ray.get(actor.is_ready.remote(), timeout=600)
+    return actor, pg
+
+
+def _cleanup(actor, pg):
+    """Shutdown engine actor and remove placement group."""
+    try:
+        ray.get(actor.shutdown.remote(), timeout=60)
+    except Exception:
+        pass
+    try:
+        ray.util.remove_placement_group(pg)
+    except Exception:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Tests: Offline TP=1
+# ---------------------------------------------------------------------------
+
+
+@unittest.skipUnless(_has_ray, "ray is not installed")
+@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
+class TestRayEngineOfflineTP1(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        if not ray.is_initialized():
+            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
+        cls.actor, cls.pg = _create_engine_on_pg(tp_size=1)
+
+    @classmethod
+    def tearDownClass(cls):
+        _cleanup(cls.actor, cls.pg)
+        ray.shutdown()
+
+    def test_offline_generate(self):
+        result = ray.get(
+            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
+        )
+        self.assertIn("text", result)
+        self.assertGreater(len(result["text"]), 0)
+        print(f"Generated: {result['text'][:200]}")
+
+    def test_batch_generate(self):
+        for prompt in _PROMPTS:
+            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
+            self.assertIn("text", result)
+            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
+
+    def test_deterministic(self):
+        prompt = "The meaning of life is"
+        r1 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
+        r2 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
+        self.assertEqual(r1["text"], r2["text"])
+
+
+# ---------------------------------------------------------------------------
+# Tests: Offline TP=2
+# ---------------------------------------------------------------------------
+
+
+@unittest.skipUnless(_has_ray, "ray is not installed")
+@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
+class TestRayEngineOfflineTP2(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        if not ray.is_initialized():
+            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
+        cls.actor, cls.pg = _create_engine_on_pg(tp_size=2)
+
+    @classmethod
+    def tearDownClass(cls):
+        _cleanup(cls.actor, cls.pg)
+        ray.shutdown()
+
+    def test_offline_generate_tp2(self):
+        result = ray.get(
+            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
+        )
+        self.assertIn("text", result)
+        self.assertGreater(len(result["text"]), 0)
+        print(f"Generated (TP=2): {result['text'][:200]}")
+
+    def test_batch_generate_tp2(self):
+        for prompt in _PROMPTS:
+            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
+            self.assertIn("text", result)
+            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
+
+
+# ---------------------------------------------------------------------------
+# Tests: Offline PP=2
+# ---------------------------------------------------------------------------
+
+
+@unittest.skipUnless(_has_ray, "ray is not installed")
+@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
+class TestRayEngineOfflinePP2(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        if not ray.is_initialized():
+            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
+        cls.actor, cls.pg = _create_engine_on_pg(tp_size=1, pp_size=2)
+
+    @classmethod
+    def tearDownClass(cls):
+        _cleanup(cls.actor, cls.pg)
+        ray.shutdown()
+
+    def test_offline_generate_pp2(self):
+        result = ray.get(
+            self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS)
+        )
+        self.assertIn("text", result)
+        self.assertGreater(len(result["text"]), 0)
+        print(f"Generated (PP=2): {result['text'][:200]}")
+
+    def test_batch_generate_pp2(self):
+        for prompt in _PROMPTS:
+            result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS))
+            self.assertIn("text", result)
+            self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}")
+
+
+# ---------------------------------------------------------------------------
+# Tests: Error paths
+# ---------------------------------------------------------------------------
+
+
+@unittest.skipUnless(_has_ray, "ray is not installed")
+@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
+class TestRayEngineErrors(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        if not ray.is_initialized():
+            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown()
+
+    def test_dp_greater_than_1_raises(self):
+        """RayEngine with dp_size > 1 should raise NotImplementedError."""
+
+        @ray.remote
+        class _BadActor:
+            def try_create(self):
+                from sglang.srt.ray.engine import RayEngine
+
+                try:
+                    RayEngine(
+                        model_path=_MODEL,
+                        tp_size=1,
+                        dp_size=2,
+                        use_ray=True,
+                    )
+                    return None
+                except (NotImplementedError, RuntimeError) as e:
+                    return str(e)
+
+        pg = placement_group([{"CPU": 1, "GPU": 1}], strategy="STRICT_PACK")
+        ray.get(pg.ready())
+
+        actor = _BadActor.options(
+            num_cpus=1,
+            num_gpus=0,
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg,
+                placement_group_bundle_index=0,
+            ),
+        ).remote()
+
+        try:
+            error_msg = ray.get(actor.try_create.remote(), timeout=120)
+            self.assertIsNotNone(error_msg, "Expected error but RayEngine created OK")
+            self.assertIn("dp_size", error_msg.lower())
+        finally:
+            ray.util.remove_placement_group(pg)
+
+    def test_missing_placement_group_raises(self):
+        """RayEngine without a placement group should raise RuntimeError."""
+
+        @ray.remote(num_gpus=1)
+        def _try_create_without_pg():
+            from sglang.srt.ray.engine import RayEngine
+
+            try:
+                RayEngine(
+                    model_path=_MODEL,
+                    tp_size=1,
+                    use_ray=True,
+                )
+                return None
+            except RuntimeError as e:
+                return str(e)
+
+        error_msg = ray.get(_try_create_without_pg.remote(), timeout=120)
+        self.assertIsNotNone(
+            error_msg, "Expected RuntimeError but RayEngine created OK"
+        )
+        self.assertIn("placement group", error_msg.lower())
+
+
+# ---------------------------------------------------------------------------
+# Tests: HTTP server
+# ---------------------------------------------------------------------------
+
+
+@unittest.skipUnless(_has_ray, "ray is not installed")
+@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
+class TestRayHTTPServerTP1(unittest.TestCase):
+    """Test the Ray HTTP server path (launch_server.py --use-ray).
+
+    Launches the server inside a Ray task on a placement group (mirrors
+    examples/anyscale/driver_online.py) and sends HTTP requests to it.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        import requests as req_lib
+
+        if not ray.is_initialized():
+            ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
+
+        cls.port = 30100
+        cls.pg = placement_group(
+            [{"CPU": 1, "GPU": 1}],
+            strategy="STRICT_PACK",
+        )
+        ray.get(cls.pg.ready())
+
+        pg_strategy = PlacementGroupSchedulingStrategy(
+            placement_group=cls.pg,
+            placement_group_bundle_index=0,
+        )
+
+        # Resolve the node IP where the server will run
+        @ray.remote(num_cpus=0, num_gpus=0)
+        def _get_ip():
+            return ray.util.get_node_ip_address()
+
+        cls.node_ip = ray.get(_get_ip.options(scheduling_strategy=pg_strategy).remote())
+        cls.base_url = f"http://{cls.node_ip}:{cls.port}"
+
+        # Launch server as a Ray task (blocks until server exits)
+        @ray.remote
+        def _launch(**kwargs):
+            from sglang.srt.ray.http_server import launch_server
+            from sglang.srt.server_args import ServerArgs
+
+            launch_server(ServerArgs(**kwargs))
+
+        cls.server_ref = _launch.options(
+            num_cpus=1,
+            num_gpus=0,
+            scheduling_strategy=pg_strategy,
+        ).remote(
+            model_path=_MODEL,
+            tp_size=1,
+            port=cls.port,
+            host="0.0.0.0",
+            use_ray=True,
+        )
+
+        # Wait for health check
+        t0 = time.time()
+        timeout = 600
+        healthy = False
+        while time.time() - t0 < timeout:
+            ready, _ = ray.wait([cls.server_ref], timeout=0)
+            if ready:
+                try:
+                    ray.get(cls.server_ref)
+                except Exception as e:
+                    raise RuntimeError(f"Server task crashed: {e}") from e
+                raise RuntimeError("Server task exited before becoming healthy")
+            try:
+                if req_lib.get(f"{cls.base_url}/health", timeout=5).status_code == 200:
+                    healthy = True
+                    break
+            except req_lib.exceptions.RequestException:
+                pass
+            time.sleep(3)
+
+        if not healthy:
+            ray.cancel(cls.server_ref, force=True)
+            raise RuntimeError(f"Server did not become healthy within {timeout}s")
+
+    @classmethod
+    def tearDownClass(cls):
+        try:
+            ray.cancel(cls.server_ref, force=True)
+        except Exception:
+            pass
+        try:
+            ray.util.remove_placement_group(cls.pg)
+        except Exception:
+            pass
+        ray.shutdown()
+
+    def test_health_endpoint(self):
+        import requests
+
+        resp = requests.get(f"{self.base_url}/health", timeout=10)
+        self.assertEqual(resp.status_code, 200)
+
+    def test_generate_endpoint(self):
+        import requests
+
+        resp = requests.post(
+            f"{self.base_url}/generate",
+            json={
+                "text": "The capital of France is",
+                "sampling_params": _SAMPLING_PARAMS,
+            },
+            timeout=60,
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        self.assertIn("text", data)
+        self.assertGreater(len(data["text"]), 0)
+        print(f"HTTP response: {data['text'][:200]}")
+
+    def test_generate_multiple(self):
+        import requests
+
+        for prompt in _PROMPTS:
+            resp = requests.post(
+                f"{self.base_url}/generate",
+                json={
+                    "text": prompt,
+                    "sampling_params": _SAMPLING_PARAMS,
+                },
+                timeout=60,
+            )
+            resp.raise_for_status()
+            data = resp.json()
+            self.assertIn("text", data)
+            self.assertGreater(len(data["text"]), 0, f"Empty output for: {prompt}")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_sagemaker_server.py

@@ -0,0 +1,183 @@
+"""
+python3 -m unittest test_sagemaker_server.TestSageMakerServer.test_chat_completion
+"""
+
+import json
+import unittest
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+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 TestSageMakerServer(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-123456"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            api_key=cls.api_key,
+        )
+        cls.tokenizer = get_tokenizer(DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_chat_completion(self, logprobs, parallel_sample_num):
+        data = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {
+                    "role": "user",
+                    "content": "What is the capital of France? Answer in a few words.",
+                },
+            ],
+            "temperature": 0,
+            "logprobs": logprobs is not None and logprobs > 0,
+            "top_logprobs": logprobs,
+            "n": parallel_sample_num,
+        }
+
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        response = requests.post(
+            f"{self.base_url}/invocations", json=data, headers=headers
+        ).json()
+
+        if logprobs:
+            assert isinstance(
+                response["choices"][0]["logprobs"]["content"][0]["top_logprobs"][0][
+                    "token"
+                ],
+                str,
+            )
+
+            ret_num_top_logprobs = len(
+                response["choices"][0]["logprobs"]["content"][0]["top_logprobs"]
+            )
+            assert (
+                ret_num_top_logprobs == logprobs
+            ), f"{ret_num_top_logprobs} vs {logprobs}"
+
+        assert len(response["choices"]) == parallel_sample_num
+        assert response["choices"][0]["message"]["role"] == "assistant"
+        assert isinstance(response["choices"][0]["message"]["content"], str)
+        assert response["id"]
+        assert response["created"]
+        assert response["usage"]["prompt_tokens"] > 0
+        assert response["usage"]["completion_tokens"] > 0
+        assert response["usage"]["total_tokens"] > 0
+
+    def run_chat_completion_stream(self, logprobs, parallel_sample_num=1):
+        data = {
+            "model": self.model,
+            "messages": [
+                {"role": "system", "content": "You are a helpful AI assistant"},
+                {
+                    "role": "user",
+                    "content": "What is the capital of France? Answer in a few words.",
+                },
+            ],
+            "temperature": 0,
+            "logprobs": logprobs is not None and logprobs > 0,
+            "top_logprobs": logprobs,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+            "n": parallel_sample_num,
+        }
+
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        response = requests.post(
+            f"{self.base_url}/invocations", json=data, stream=True, headers=headers
+        )
+
+        is_firsts = {}
+        for line in response.iter_lines():
+            line = line.decode("utf-8").replace("data: ", "")
+            if len(line) < 1 or line == "[DONE]":
+                continue
+            print(f"value: {line}")
+            line = json.loads(line)
+            usage = line.get("usage")
+            if usage is not None:
+                assert usage["prompt_tokens"] > 0
+                assert usage["completion_tokens"] > 0
+                assert usage["total_tokens"] > 0
+                continue
+
+            index = line.get("choices")[0].get("index")
+            data = line.get("choices")[0].get("delta")
+
+            if is_firsts.get(index, True):
+                assert data["role"] == "assistant"
+                is_firsts[index] = False
+                continue
+
+            # Skip chunks that are just empty placeholders, usually at stream end/stop
+            if data.get("content") is None:
+                continue
+
+            if logprobs:
+                assert line.get("choices")[0].get("logprobs")
+                assert isinstance(
+                    line.get("choices")[0]
+                    .get("logprobs")
+                    .get("content")[0]
+                    .get("top_logprobs")[0]
+                    .get("token"),
+                    str,
+                )
+                assert isinstance(
+                    line.get("choices")[0]
+                    .get("logprobs")
+                    .get("content")[0]
+                    .get("top_logprobs"),
+                    list,
+                )
+                ret_num_top_logprobs = len(
+                    line.get("choices")[0]
+                    .get("logprobs")
+                    .get("content")[0]
+                    .get("top_logprobs")
+                )
+                assert (
+                    ret_num_top_logprobs == logprobs
+                ), f"{ret_num_top_logprobs} vs {logprobs}"
+
+            assert isinstance(data["content"], str)
+            assert line["id"]
+            assert line["created"]
+
+        for index in [i for i in range(parallel_sample_num)]:
+            assert not is_firsts.get(
+                index, True
+            ), f"index {index} is not found in the response"
+
+    def test_chat_completion(self):
+        for logprobs in [None, 5]:
+            for parallel_sample_num in [1, 2]:
+                self.run_chat_completion(logprobs, parallel_sample_num)
+
+    def test_chat_completion_stream(self):
+        for logprobs in [None, 5]:
+            for parallel_sample_num in [1, 2]:
+                self.run_chat_completion_stream(logprobs, parallel_sample_num)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_schedule_policy.py

@@ -0,0 +1,314 @@
+import unittest
+
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.managers.schedule_policy import (
+    CacheAgnosticPolicy,
+    CacheAwarePolicy,
+    SchedulePolicy,
+)
+from sglang.srt.mem_cache.radix_cache import RadixCache
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestSchedulePolicy(CustomTestCase):
+
+    def setUp(self):
+        self.tree_cache = RadixCache.create_simulated()
+
+    def test_init_with_cache_aware_policy(self):
+        policy = SchedulePolicy(
+            policy="lpm",
+            tree_cache=self.tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        self.assertEqual(policy.policy, CacheAwarePolicy.LPM)
+
+    def test_init_with_cache_agnostic_policy(self):
+        policy = SchedulePolicy(
+            policy="fcfs",
+            tree_cache=self.tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        self.assertEqual(policy.policy, CacheAgnosticPolicy.FCFS)
+
+    def test_init_with_unknown_policy(self):
+        with self.assertRaises(ValueError):
+            SchedulePolicy(
+                policy="invalid",
+                tree_cache=self.tree_cache,
+                enable_hierarchical_cache=True,
+                enable_priority_scheduling=False,
+                schedule_low_priority_values_first=False,
+            )
+
+    def test_init_with_disabled_cache(self):
+        tree_cache = RadixCache.create_simulated(disable=True)
+        policy = SchedulePolicy(
+            policy="lpm",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        self.assertEqual(policy.policy, CacheAgnosticPolicy.FCFS)
+
+    def test_calc_priority_fcfs(self):
+        tree_cache = RadixCache.create_simulated()
+        waiting_queue = [
+            Req(1, "a b", [1, 2], SamplingParams()),
+            Req(3, "a b c", [1, 2, 3], SamplingParams()),
+            Req(2, "a", [1], SamplingParams()),
+        ]
+
+        policy = SchedulePolicy(
+            policy="fcfs",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue)
+        # Check if FCFS keeps the original order
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 3)
+        self.assertEqual(waiting_queue[2].rid, 2)
+
+    def test_calc_priority_priority_enabled_fcfs_scheduling(self):
+        tree_cache = RadixCache.create_simulated()
+        r1 = Req(1, "a b", [1, 2], SamplingParams())
+        r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
+        r3 = Req(2, "a", [1], SamplingParams())
+        r1.priority, r1.time_stats.wait_queue_entry_time = 1, 1
+        r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
+        r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
+
+        waiting_queue = [r1, r2, r3]
+
+        policy = SchedulePolicy(
+            policy="fcfs",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=True,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue)
+
+        # Check if priority enabled fcfs ordering is applied.
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 2)
+        self.assertEqual(waiting_queue[2].rid, 3)
+
+    def test_calc_priority_priority_enabled_fcfs_scheduling_with_low_priority_values_first(
+        self,
+    ):
+        tree_cache = RadixCache.create_simulated()
+        r1 = Req(1, "a b", [1, 2], SamplingParams())
+        r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
+        r3 = Req(2, "a", [1], SamplingParams())
+        r1.priority, r1.time_stats.wait_queue_entry_time = -1, 1
+        r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
+        r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
+
+        waiting_queue = [r1, r2, r3]
+
+        policy = SchedulePolicy(
+            policy="fcfs",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=True,
+            schedule_low_priority_values_first=True,
+        )
+        policy.calc_priority(waiting_queue)
+        # Check if priority enabled fcfs ordering is applied.
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 2)
+        self.assertEqual(waiting_queue[2].rid, 3)
+
+    def test_calc_priority_longest_output_first_scheduling(self):
+        tree_cache = RadixCache.create_simulated()
+
+        waiting_queue = [
+            Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)),
+            Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)),
+            Req(2, "a", [1], SamplingParams(max_new_tokens=100)),
+        ]
+
+        policy = SchedulePolicy(
+            policy="lof",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue)
+        # Check if priority enabled fcfs ordering is applied.
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 2)
+        self.assertEqual(waiting_queue[2].rid, 3)
+
+    def test_calc_priority_priority_enabled_longest_output_first_scheduling(self):
+        tree_cache = RadixCache.create_simulated()
+
+        waiting_queue = [
+            Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=1),
+            Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=0),
+            Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=0),
+        ]
+
+        policy = SchedulePolicy(
+            policy="lof",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=True,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue)
+        # Check if priority enabled fcfs ordering is applied.
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 2)
+        self.assertEqual(waiting_queue[2].rid, 3)
+
+    def test_calc_priority_priority_enabled_longest_output_first_scheduling_with_low_priority_values_first(
+        self,
+    ):
+        tree_cache = RadixCache.create_simulated()
+
+        waiting_queue = [
+            Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=0),
+            Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=1),
+            Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=1),
+        ]
+
+        policy = SchedulePolicy(
+            policy="lof",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=True,
+            enable_priority_scheduling=True,
+            schedule_low_priority_values_first=True,
+        )
+        policy.calc_priority(waiting_queue)
+        # Check if priority enabled fcfs ordering is applied.
+        self.assertEqual(waiting_queue[0].rid, 1)
+        self.assertEqual(waiting_queue[1].rid, 2)
+        self.assertEqual(waiting_queue[2].rid, 3)
+
+    def test_calc_priority_routing_key_scheduling(self):
+        """Test routing-key policy: prioritize by routing key frequency in running batch."""
+        tree_cache = RadixCache.create_simulated()
+
+        running_reqs = [
+            Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
+            Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
+            Req("r3", "c", [3], SamplingParams(), routing_key="key_b"),
+        ]
+        running_batch = ScheduleBatch(reqs=running_reqs)
+
+        waiting_queue = [
+            Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
+            Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
+            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+        ]
+
+        policy = SchedulePolicy(
+            policy="routing-key",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=False,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue, running_batch)
+
+        self.assertEqual(waiting_queue[0].rid, "w2")
+        self.assertEqual(waiting_queue[1].rid, "w1")
+        self.assertEqual(waiting_queue[2].rid, "w3")
+
+    def test_calc_priority_routing_key_tie_break_by_lexicographic_order(self):
+        """Test routing-key policy: tie-break by lexicographic order."""
+        tree_cache = RadixCache.create_simulated()
+
+        running_reqs = [
+            Req("r1", "a", [1], SamplingParams(), routing_key="key_b"),
+            Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
+        ]
+        running_batch = ScheduleBatch(reqs=running_reqs)
+
+        waiting_queue = [
+            Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
+            Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
+        ]
+
+        policy = SchedulePolicy(
+            policy="routing-key",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=False,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue, running_batch)
+
+        self.assertEqual(waiting_queue[0].rid, "w2")
+        self.assertEqual(waiting_queue[1].rid, "w1")
+
+    def test_calc_priority_routing_key_no_match_deprioritized(self):
+        """Test routing-key policy: requests without matching routing keys are deprioritized."""
+        tree_cache = RadixCache.create_simulated()
+
+        running_reqs = [
+            Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
+            Req("r2", "b", [2], SamplingParams(), routing_key="key_b"),
+            Req("r3", "c", [3], SamplingParams(), routing_key="key_c"),
+        ]
+        running_batch = ScheduleBatch(reqs=running_reqs)
+
+        waiting_queue = [
+            Req("w1", "d", [4], SamplingParams(), routing_key="key_d"),
+            Req("w2", "e", [5], SamplingParams(), routing_key="key_e"),
+            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+        ]
+
+        policy = SchedulePolicy(
+            policy="routing-key",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=False,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue, running_batch)
+
+        self.assertEqual(waiting_queue[0].rid, "w3")
+        self.assertEqual(waiting_queue[1].rid, "w1")
+        self.assertEqual(waiting_queue[2].rid, "w2")
+
+    def test_calc_priority_routing_key_empty_running_batch(self):
+        """Test routing-key policy: empty running batch keeps original order."""
+        tree_cache = RadixCache.create_simulated()
+
+        running_batch = ScheduleBatch(reqs=[])
+
+        waiting_queue = [
+            Req("w1", "d", [4], SamplingParams(), routing_key="key_a"),
+            Req("w2", "e", [5], SamplingParams(), routing_key="key_b"),
+            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+        ]
+
+        policy = SchedulePolicy(
+            policy="routing-key",
+            tree_cache=tree_cache,
+            enable_hierarchical_cache=False,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+        policy.calc_priority(waiting_queue, running_batch)
+
+        self.assertEqual(waiting_queue[0].rid, "w1")
+        self.assertEqual(waiting_queue[1].rid, "w2")
+        self.assertEqual(waiting_queue[2].rid, "w3")
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_srt_engine_with_quant_args.py

@@ -0,0 +1,60 @@
+import unittest
+
+import sglang as sgl
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
+
+
+class TestSRTEngineWithQuantArgs(CustomTestCase):
+
+    def test_1_quantization_args(self):
+
+        # we only test fp8 because other methods are currently dependent on vllm. We can add other methods back to test after vllm dependency is resolved.
+        quantization_args_list = [
+            # "awq",
+            "fp8",
+            # "gptq",
+            # "marlin",
+            # "gptq_marlin",
+            # "awq_marlin",
+            # "bitsandbytes",
+            # "gguf",
+        ]
+
+        prompt = "Today is a sunny day and I like"
+        model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+        sampling_params = {"temperature": 0, "max_new_tokens": 8}
+
+        for quantization_args in quantization_args_list:
+            engine = sgl.Engine(
+                model_path=model_path, random_seed=42, quantization=quantization_args
+            )
+            engine.generate(prompt, sampling_params)
+            engine.shutdown()
+
+    def test_2_torchao_args(self):
+
+        # we don't test int8dq because currently there is conflict between int8dq and capture cuda graph
+        torchao_args_list = [
+            # "int8dq",
+            "int8wo",
+            "fp8wo",
+            "fp8dq-per_tensor",
+            "fp8dq-per_row",
+        ] + [f"int4wo-{group_size}" for group_size in [32, 64, 128, 256]]
+
+        prompt = "Today is a sunny day and I like"
+        model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+        sampling_params = {"temperature": 0, "max_new_tokens": 8}
+
+        for torchao_config in torchao_args_list:
+            engine = sgl.Engine(
+                model_path=model_path, random_seed=42, torchao_config=torchao_config
+            )
+            engine.generate(prompt, sampling_params)
+            engine.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_tokenizer_batch_encode.py

@@ -0,0 +1,121 @@
+"""
+Unit tests for enable_tokenizer_batch_encode feature.
+
+This tests the batch tokenization functionality which allows processing
+multiple text inputs in a single batch for improved performance.
+
+Usage:
+python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncode.test_batch_validation_constraints
+python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncodeUnit.test_batch_tokenize_and_process_logic
+python3 -m unittest test_tokenizer_batch_encode.TestTokenizerBatchEncodeLogic.test_batch_processing_path
+"""
+
+import unittest
+from unittest.mock import Mock, patch
+
+from sglang.srt.managers.io_struct import GenerateReqInput
+from sglang.srt.managers.tokenizer_manager import TokenizerManager
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+
+class TestTokenizerBatchEncode(unittest.TestCase):
+    """Test cases for tokenizer batch encoding validation and setup."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.server_args = ServerArgs(
+            model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            enable_tokenizer_batch_encode=True,
+        )
+        self.port_args = PortArgs.init_new(self.server_args)
+
+        with patch("zmq.asyncio.Context"), patch(
+            "sglang.srt.utils.get_zmq_socket"
+        ), patch(
+            "sglang.srt.utils.hf_transformers_utils.get_tokenizer"
+        ) as mock_tokenizer:
+
+            mock_tokenizer.return_value = Mock(vocab_size=32000)
+            self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
+
+    def test_batch_encode_enabled(self):
+        """Test that batch encoding is enabled when configured."""
+        self.assertTrue(self.server_args.enable_tokenizer_batch_encode)
+
+    def test_batch_encode_disabled(self):
+        """Test that batch encoding can be disabled."""
+        server_args_disabled = ServerArgs(
+            model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            enable_tokenizer_batch_encode=False,
+        )
+        self.assertFalse(server_args_disabled.enable_tokenizer_batch_encode)
+
+    def test_multimodal_input_validation(self):
+        """Test that multimodal inputs are rejected in batch mode."""
+        req = GenerateReqInput(text="test", image_data=["dummy"])
+        req.contains_mm_input = Mock(return_value=True)
+
+        batch_obj = Mock()
+        batch_obj.__getitem__ = lambda self, i: req
+
+        self.tokenizer_manager.is_generation = True
+
+        with self.assertRaises(ValueError) as cm:
+            self.tokenizer_manager._validate_batch_tokenization_constraints(
+                1, batch_obj
+            )
+
+        self.assertIn("multimodal", str(cm.exception))
+
+    def test_pretokenized_input_validation(self):
+        """Test that pre-tokenized inputs are rejected in batch mode."""
+        req = GenerateReqInput(input_ids=[1, 2, 3])
+
+        batch_obj = Mock()
+        batch_obj.__getitem__ = lambda self, i: req
+
+        with self.assertRaises(ValueError) as cm:
+            self.tokenizer_manager._validate_batch_tokenization_constraints(
+                1, batch_obj
+            )
+
+        self.assertIn("pre-tokenized", str(cm.exception))
+
+    def test_input_embeds_validation(self):
+        """Test that input embeds are rejected in batch mode."""
+        req = GenerateReqInput(input_embeds=[0.1, 0.2])
+
+        batch_obj = Mock()
+        batch_obj.__getitem__ = lambda self, i: req
+
+        with self.assertRaises(ValueError) as cm:
+            self.tokenizer_manager._validate_batch_tokenization_constraints(
+                1, batch_obj
+            )
+
+        self.assertIn("input_embeds", str(cm.exception))
+
+    def test_valid_text_only_requests_pass_validation(self):
+        """Test that valid text-only requests pass validation."""
+        # Create valid requests (text-only)
+        requests = []
+        for i in range(3):
+            req = GenerateReqInput(text=f"test text {i}")
+            req.contains_mm_input = Mock(return_value=False)
+            requests.append(req)
+
+        batch_obj = Mock()
+        batch_obj.__getitem__ = Mock(side_effect=lambda i: requests[i])
+
+        # Should not raise any exception
+        try:
+            self.tokenizer_manager._validate_batch_tokenization_constraints(
+                3, batch_obj
+            )
+        except Exception as e:
+            self.fail(f"Validation failed for valid text-only requests: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

third_party/sglang/test/manual/test_tokenizer_manager.py

@@ -0,0 +1,408 @@
+"""
+Unit tests for TokenizerManager helper methods.
+
+This tests the refactored tokenization functionality including input format detection,
+tokenizer input preparation, and result extraction logic.
+
+Usage:
+python3 -m unittest test_tokenizer_manager.TestInputFormatDetection
+python3 -m unittest test_tokenizer_manager.TestTokenizerInputPreparation
+python3 -m unittest test_tokenizer_manager.TestTokenizerResultExtraction
+python3 -m unittest test_tokenizer_manager.TestTokenizerManagerIntegration
+"""
+
+import unittest
+from unittest.mock import Mock, patch
+
+from sglang.srt.managers.tokenizer_manager import InputFormat, TokenizerManager
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+
+
+class TestInputFormatDetection(unittest.TestCase):
+    """Test cases for _detect_input_format method."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        with patch("sglang.srt.utils.get_device", return_value="cpu"):
+            self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+            self.port_args = PortArgs.init_new(self.server_args)
+
+        with patch("zmq.asyncio.Context"), patch(
+            "sglang.srt.utils.get_zmq_socket"
+        ), patch(
+            "sglang.srt.utils.hf_transformers_utils.get_tokenizer"
+        ) as mock_tokenizer:
+            mock_tokenizer.return_value = Mock(vocab_size=32000)
+            self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
+
+    def test_detect_single_string(self):
+        """Test detection of single string input."""
+        text = "Hello world"
+        result = self.tokenizer_manager._detect_input_format(
+            text, is_cross_encoder=False
+        )
+        self.assertEqual(result, InputFormat.SINGLE_STRING)
+
+    def test_detect_single_string_cross_encoder_disabled(self):
+        """Test single string with cross_encoder disabled still returns single_string."""
+        text = "Hello world"
+        result = self.tokenizer_manager._detect_input_format(
+            text, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.SINGLE_STRING)
+
+    def test_detect_batch_strings(self):
+        """Test detection of batch string inputs."""
+        texts = ["Hello", "World", "How are you?"]
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=False
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+    def test_detect_batch_strings_cross_encoder_disabled(self):
+        """Test batch strings with cross_encoder disabled."""
+        texts = ["Hello", "World"]
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+    def test_detect_cross_encoder_single_pair(self):
+        """Test detection of cross-encoder single pair."""
+        texts = [["query text", "document text"]]
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.CROSS_ENCODER_PAIRS)
+
+    def test_detect_cross_encoder_multiple_pairs(self):
+        """Test detection of cross-encoder multiple pairs."""
+        texts = [["q1", "d1"], ["q2", "d2"], ["q3", "d3"]]
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.CROSS_ENCODER_PAIRS)
+
+    def test_detect_cross_encoder_disabled_with_pairs(self):
+        """Test pairs with cross_encoder disabled should return batch_strings."""
+        texts = [["query", "document"]]
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=False
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+    def test_detect_empty_list(self):
+        """Test detection with empty list."""
+        texts = []
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+    def test_detect_malformed_cross_encoder_pairs(self):
+        """Test malformed cross-encoder pairs (not length 2)."""
+        texts = [["query only"]]  # Single element, not a pair
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+        texts = [["query", "doc", "extra"]]  # Three elements, not a pair
+        result = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(result, InputFormat.BATCH_STRINGS)
+
+
+class TestTokenizerInputPreparation(unittest.TestCase):
+    """Test cases for _prepare_tokenizer_input method."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        with patch("sglang.srt.utils.get_device", return_value="cpu"):
+            self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+            self.port_args = PortArgs.init_new(self.server_args)
+
+        with patch("zmq.asyncio.Context"), patch(
+            "sglang.srt.utils.get_zmq_socket"
+        ), patch(
+            "sglang.srt.utils.hf_transformers_utils.get_tokenizer"
+        ) as mock_tokenizer:
+            mock_tokenizer.return_value = Mock(vocab_size=32000)
+            self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
+
+    def test_prepare_single_string_input(self):
+        """Test preparation of single string input."""
+        text = "Hello world"
+        result = self.tokenizer_manager._prepare_tokenizer_input(
+            text, InputFormat.SINGLE_STRING
+        )
+        self.assertEqual(result, ["Hello world"])
+
+    def test_prepare_batch_strings_input(self):
+        """Test preparation of batch strings input."""
+        texts = ["Hello", "World", "Test"]
+        result = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, InputFormat.BATCH_STRINGS
+        )
+        self.assertEqual(result, ["Hello", "World", "Test"])
+
+    def test_prepare_cross_encoder_pairs_input(self):
+        """Test preparation of cross-encoder pairs input."""
+        texts = [["query1", "doc1"], ["query2", "doc2"]]
+        result = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, InputFormat.CROSS_ENCODER_PAIRS
+        )
+        self.assertEqual(result, [["query1", "doc1"], ["query2", "doc2"]])
+
+    def test_prepare_cross_encoder_single_pair_input(self):
+        """Test preparation of single cross-encoder pair."""
+        texts = [["query text", "document text"]]
+        result = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, InputFormat.CROSS_ENCODER_PAIRS
+        )
+        self.assertEqual(result, [["query text", "document text"]])
+
+    def test_prepare_batch_strings_input_format_passthrough(self):
+        """Batch strings should pass through unchanged."""
+        texts = ["test"]
+        result = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, InputFormat.BATCH_STRINGS
+        )
+        self.assertEqual(result, ["test"])
+
+
+class TestTokenizerResultExtraction(unittest.TestCase):
+    """Test cases for _extract_tokenizer_results method."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        with patch("sglang.srt.utils.get_device", return_value="cpu"):
+            self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+            self.port_args = PortArgs.init_new(self.server_args)
+
+        with patch("zmq.asyncio.Context"), patch(
+            "sglang.srt.utils.get_zmq_socket"
+        ), patch(
+            "sglang.srt.utils.hf_transformers_utils.get_tokenizer"
+        ) as mock_tokenizer:
+            mock_tokenizer.return_value = Mock(vocab_size=32000)
+            self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
+
+    def test_extract_single_string_results(self):
+        """Test extraction for single string input."""
+        input_ids = [[101, 2129, 102]]
+        token_type_ids = [[0, 0, 0]]
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.SINGLE_STRING,
+                original_batch_size=1,
+            )
+        )
+
+        self.assertEqual(result_input_ids, [101, 2129, 102])
+        self.assertEqual(result_token_type_ids, [0, 0, 0])
+
+    def test_extract_single_cross_encoder_results(self):
+        """Test extraction for single cross-encoder pair."""
+        input_ids = [[101, 2129, 102, 4068, 102]]
+        token_type_ids = [[0, 0, 0, 1, 1]]
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.CROSS_ENCODER_PAIRS,
+                original_batch_size=1,
+            )
+        )
+
+        self.assertEqual(result_input_ids, [101, 2129, 102, 4068, 102])
+        self.assertEqual(result_token_type_ids, [0, 0, 0, 1, 1])
+
+    def test_extract_batch_results(self):
+        """Test extraction for batch inputs."""
+        input_ids = [[101, 2129, 102], [101, 4068, 102]]
+        token_type_ids = [[0, 0, 0], [0, 0, 0]]
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.BATCH_STRINGS,
+                original_batch_size=2,
+            )
+        )
+
+        self.assertEqual(result_input_ids, [[101, 2129, 102], [101, 4068, 102]])
+        self.assertEqual(result_token_type_ids, [[0, 0, 0], [0, 0, 0]])
+
+    def test_extract_multiple_cross_encoder_results(self):
+        """Test extraction for multiple cross-encoder pairs."""
+        input_ids = [[101, 2129, 102, 4068, 102], [101, 7592, 102, 2088, 102]]
+        token_type_ids = [[0, 0, 0, 1, 1], [0, 0, 0, 1, 1]]
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.CROSS_ENCODER_PAIRS,
+                original_batch_size=2,
+            )
+        )
+
+        self.assertEqual(
+            result_input_ids, [[101, 2129, 102, 4068, 102], [101, 7592, 102, 2088, 102]]
+        )
+        self.assertEqual(result_token_type_ids, [[0, 0, 0, 1, 1], [0, 0, 0, 1, 1]])
+
+    def test_extract_empty_results(self):
+        """Test extraction with empty results."""
+        input_ids = []
+        token_type_ids = None
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.SINGLE_STRING,
+                original_batch_size=1,
+            )
+        )
+
+        self.assertEqual(result_input_ids, [])
+        self.assertIsNone(result_token_type_ids)
+
+    def test_extract_with_none_token_type_ids(self):
+        """Test extraction when token_type_ids is None."""
+        input_ids = [[101, 2129, 102]]
+        token_type_ids = None
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                input_ids,
+                token_type_ids,
+                InputFormat.SINGLE_STRING,
+                original_batch_size=1,
+            )
+        )
+
+        self.assertEqual(result_input_ids, [101, 2129, 102])
+        self.assertIsNone(result_token_type_ids)
+
+
+class TestTokenizerManagerIntegration(unittest.TestCase):
+    """Integration tests combining multiple helper methods."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        with patch("sglang.srt.utils.get_device", return_value="cpu"):
+            self.server_args = ServerArgs(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+            self.port_args = PortArgs.init_new(self.server_args)
+
+        with patch("zmq.asyncio.Context"), patch(
+            "sglang.srt.utils.get_zmq_socket"
+        ), patch(
+            "sglang.srt.utils.hf_transformers_utils.get_tokenizer"
+        ) as mock_tokenizer:
+            mock_tokenizer.return_value = Mock(vocab_size=32000)
+            self.tokenizer_manager = TokenizerManager(self.server_args, self.port_args)
+
+    def test_full_workflow_single_string(self):
+        """Test complete workflow for single string input."""
+        text = "Hello world"
+
+        # Step 1: Detect format
+        input_format = self.tokenizer_manager._detect_input_format(
+            text, is_cross_encoder=False
+        )
+        self.assertEqual(input_format, InputFormat.SINGLE_STRING)
+
+        # Step 2: Prepare input
+        tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
+            text, input_format
+        )
+        self.assertEqual(tokenizer_input, ["Hello world"])
+
+        # Step 3: Extract results (simulated tokenizer output)
+        mock_input_ids = [[101, 2129, 4248, 102]]
+        mock_token_type_ids = None
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                mock_input_ids, mock_token_type_ids, input_format, original_batch_size=1
+            )
+        )
+
+        self.assertEqual(result_input_ids, [101, 2129, 4248, 102])
+        self.assertIsNone(result_token_type_ids)
+
+    def test_full_workflow_cross_encoder_pairs(self):
+        """Test complete workflow for cross-encoder pairs."""
+        texts = [
+            ["How many people live in Berlin?", "Berlin is well known for its museums."]
+        ]
+
+        # Step 1: Detect format
+        input_format = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=True
+        )
+        self.assertEqual(input_format, InputFormat.CROSS_ENCODER_PAIRS)
+
+        # Step 2: Prepare input
+        tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, input_format
+        )
+        self.assertEqual(tokenizer_input, texts)
+
+        # Step 3: Extract results (simulated tokenizer output for cross-encoder)
+        mock_input_ids = [[101, 2129, 2116, 102, 4068, 2003, 102]]
+        mock_token_type_ids = [[0, 0, 0, 0, 1, 1, 1]]
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                mock_input_ids, mock_token_type_ids, input_format, original_batch_size=1
+            )
+        )
+
+        self.assertEqual(result_input_ids, [101, 2129, 2116, 102, 4068, 2003, 102])
+        self.assertEqual(result_token_type_ids, [0, 0, 0, 0, 1, 1, 1])
+
+    def test_full_workflow_batch_strings(self):
+        """Test complete workflow for batch strings."""
+        texts = ["Hello", "World", "Test"]
+
+        # Step 1: Detect format
+        input_format = self.tokenizer_manager._detect_input_format(
+            texts, is_cross_encoder=False
+        )
+        self.assertEqual(input_format, InputFormat.BATCH_STRINGS)
+
+        # Step 2: Prepare input
+        tokenizer_input = self.tokenizer_manager._prepare_tokenizer_input(
+            texts, input_format
+        )
+        self.assertEqual(tokenizer_input, ["Hello", "World", "Test"])
+
+        # Step 3: Extract results (simulated tokenizer output)
+        mock_input_ids = [[101, 7592, 102], [101, 2088, 102], [101, 2774, 102]]
+        mock_token_type_ids = None
+
+        result_input_ids, result_token_type_ids = (
+            self.tokenizer_manager._extract_tokenizer_results(
+                mock_input_ids, mock_token_type_ids, input_format, original_batch_size=3
+            )
+        )
+
+        self.assertEqual(
+            result_input_ids, [[101, 7592, 102], [101, 2088, 102], [101, 2774, 102]]
+        )
+        self.assertIsNone(result_token_type_ids)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

third_party/sglang/test/manual/test_torch_flex_attention_backend.py

@@ -0,0 +1,49 @@
+"""
+Usage:
+python3 -m unittest test_torch_flex_attention_backend.TestTorchFlexAttnBackend.test_gsm8k
+"""
+
+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.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestTorchFlexAttnBackend(CustomTestCase):
+    def test_gsm8k(self):
+        model = DEFAULT_MODEL_NAME_FOR_TEST
+        base_url = DEFAULT_URL_FOR_TEST
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--attention-backend", "flex_attention"],
+        )
+
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                eval_name="gsm8k",
+                api="completion",
+                max_tokens=512,
+                num_examples=100,
+                num_threads=10,
+                num_shots=8,
+            )
+            metrics = run_eval(args)
+            print(f"{metrics=}")
+            self.assertGreater(metrics["score"], 0.62)
+        finally:
+            kill_process_tree(process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_torch_tp.py

@@ -0,0 +1,30 @@
+import unittest
+
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    run_bench_offline_throughput,
+)
+
+
+class TestTorchTP(CustomTestCase):
+    def test_torch_native_llama(self):
+        output_throughput = run_bench_offline_throughput(
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            [
+                "--tp",
+                "2",
+                # This cannot run anymore with the new torch version.
+                # "--json-model-override-args",
+                # '{"architectures": ["TorchNativeLlamaForCausalLM"]}',
+                "--disable-cuda-graph",
+            ],
+        )
+
+        if is_in_ci():
+            self.assertGreater(output_throughput, 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_triton_attention_rocm_mla.py

@@ -0,0 +1,259 @@
+import random
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.triton_ops.decode_attention import (
+    decode_attention_fwd_grouped,
+)
+from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import (
+    decode_attention_fwd_grouped_rope,
+)
+from sglang.srt.layers.rotary_embedding import DeepseekScalingRotaryEmbedding
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestTritonAttentionMLA(CustomTestCase):
+
+    def _set_all_seeds(self, seed):
+        """Set all random seeds for reproducibility."""
+        random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    def setUp(self):
+        # Set seeds before each test method
+        self._set_all_seeds(42)
+
+    def preprocess_kv_cache(self, kv_cache, kv_lora_rank):
+        latent_cache = kv_cache
+        v_input = latent_cache[..., :kv_lora_rank]
+        v_input = v_input.contiguous().unsqueeze(1)
+        k_input = latent_cache.unsqueeze(1)
+        k_input[..., :kv_lora_rank] = v_input
+
+        return k_input, v_input
+
+    def input_helper(
+        self,
+        B,
+        H,
+        S,
+        kv_lora_rank,
+        rotary_dim,
+        qk_rope_head_dim,
+        num_kv_splits,
+        dtype,
+        device,
+        rope_base=10,
+        rope_max_seq_len=16384,
+        rope_scaling=1.0,
+        is_neox_style=False,
+    ):
+        q = torch.randn(
+            B, H, kv_lora_rank + qk_rope_head_dim, device=device, dtype=dtype
+        )
+        kv_cache = torch.randn(
+            B * S, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device
+        )
+        kv_indptr = torch.arange(B + 1, device=device) * S
+        kv_indices = torch.arange(B * S, device=device)
+        attn_logits = torch.empty(
+            B, H, num_kv_splits, kv_lora_rank + 1, dtype=dtype, device=device
+        )
+        rotary_emb = DeepseekScalingRotaryEmbedding(
+            qk_rope_head_dim,
+            rotary_dim,
+            rope_max_seq_len,
+            rope_base,
+            is_neox_style,
+            rope_scaling,
+            q.dtype,
+            device="cpu",
+        ).cuda()
+        positions = torch.tensor([S], device=device).unsqueeze(0).repeat(B, 1)
+
+        return kv_indptr, kv_indices, q, kv_cache, attn_logits, rotary_emb, positions
+
+    def ref_compute_full_fwd(
+        self,
+        q,
+        k_input,
+        v_input,
+        kv_lora_rank,
+        kv_indptr,
+        kv_indices,
+        num_kv_splits,
+        sm_scale,
+        logit_cap,
+        rotary_emb,
+        positions,
+        use_rope,
+        device="cuda",
+    ):
+
+        B, H = q.shape[0], q.shape[1]
+        S = kv_indptr[1].item()
+        qk_rope_head_dim = k_input.shape[-1] - kv_lora_rank
+
+        q_input = torch.empty(B, H, kv_lora_rank + qk_rope_head_dim, dtype=q.dtype).to(
+            device
+        )
+        q_nope_out, q_pe = q.split([kv_lora_rank, qk_rope_head_dim], dim=-1)
+        k_pe_t = k_input.view(B, 1, S, -1)[:, :, -1:, kv_lora_rank:]
+
+        if use_rope:
+            q_pe, k_pe_t = rotary_emb(positions, q_pe.unsqueeze(2), k_pe_t)
+            q_pe = q_pe.squeeze()
+
+        k_input.view(B, 1, S, -1)[:, :, -1:, kv_lora_rank:] = k_pe_t
+
+        q_input[..., :kv_lora_rank] = q_nope_out
+        q_input[..., kv_lora_rank:] = q_pe
+
+        B, H = q_input.shape[0], q_input.shape[1]
+        kv_lora_rank = v_input.shape[-1]
+        device = q_input.device
+
+        attn_logits = torch.empty(
+            B, H, num_kv_splits, kv_lora_rank + 1, dtype=q_input.dtype, device=device
+        )
+        o = torch.empty(B, H, kv_lora_rank, dtype=q_input.dtype, device=device)
+
+        decode_attention_fwd_grouped(
+            q_input,
+            k_input,
+            v_input,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            num_kv_splits,
+            sm_scale,
+            logit_cap,
+        )
+
+        return attn_logits, o, k_pe_t.squeeze()
+
+    def _test_rocm_fused_mla_kernel(
+        self,
+        B,
+        H,
+        S,
+        kv_lora_rank,
+        qk_rope_head_dim,
+        rotary_dim,
+        dtype,
+        use_rope,
+        is_neox_style,
+        num_kv_splits=2,
+        sm_scale=1.0,
+        logit_cap=0.0,
+        device="cuda",
+    ):
+        kv_indptr, kv_indices, q, kv_cache, attn_logits, rotary_emb, positions = (
+            self.input_helper(
+                B,
+                H,
+                S,
+                kv_lora_rank,
+                rotary_dim,
+                qk_rope_head_dim,
+                num_kv_splits,
+                dtype,
+                device=device,
+                is_neox_style=is_neox_style,
+            )
+        )
+
+        k_input, v_input = self.preprocess_kv_cache(kv_cache, kv_lora_rank)
+        k_pe_tokens = torch.empty(
+            B, qk_rope_head_dim, dtype=kv_cache.dtype, device=device
+        )
+        tri_o = torch.empty(B, H, kv_lora_rank, dtype=kv_cache.dtype, device=device)
+
+        decode_attention_fwd_grouped_rope(
+            q,
+            k_input,
+            v_input,
+            tri_o,
+            kv_indptr,
+            kv_indices,
+            k_pe_tokens if use_rope else None,
+            kv_lora_rank,
+            rotary_dim if use_rope else None,
+            rotary_emb.cos_sin_cache if use_rope else None,
+            positions if use_rope else None,
+            attn_logits,
+            num_kv_splits,
+            sm_scale,
+            logit_cap,
+            use_rope,
+            is_neox_style,
+        )
+
+        tri_logits = attn_logits
+
+        # reference
+        ref_logits, ref_o, ref_k_pe_tokens = self.ref_compute_full_fwd(
+            q,
+            k_input,
+            v_input,
+            kv_lora_rank,
+            kv_indptr,
+            kv_indices,
+            num_kv_splits,
+            sm_scale,
+            logit_cap,
+            rotary_emb,
+            positions,
+            use_rope,
+            device="cuda",
+        )
+
+        if use_rope:
+            torch.testing.assert_close(
+                ref_k_pe_tokens, k_pe_tokens.squeeze(), atol=1e-2, rtol=1e-2
+            )
+        torch.testing.assert_close(ref_logits, tri_logits, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(ref_o, tri_o, atol=1e-2, rtol=1e-2)
+
+    def test_grouped_rocm_fused_mla(self):
+        configs = [
+            (1, 128, 2048, 512, 64, 64),
+            (1, 128, 2048, 512, 128, 64),
+            (1, 128, 2048, 512, 127, 64),
+            (1, 128, 2050, 512, 127, 64),
+            (1, 128, 2050, 512, 128, 64),
+            (8, 128, 2048, 512, 64, 64),
+            (8, 128, 2048, 512, 128, 64),
+            (8, 128, 2048, 512, 127, 64),
+            (8, 128, 2050, 512, 127, 64),
+            (8, 128, 2050, 512, 128, 64),
+        ]
+        dtypes = [torch.bfloat16, torch.float32]
+        use_rope_list = [True, False]
+        is_neox_style_list = [True, False]
+
+        for B, H, S, kv_lora_rank, qk_rope_head_dim, rotary_dim in configs:
+            for dtype in dtypes:
+                for use_rope in use_rope_list:
+                    for is_neox_style in is_neox_style_list:
+                        self._test_rocm_fused_mla_kernel(
+                            B,
+                            H,
+                            S,
+                            kv_lora_rank,
+                            qk_rope_head_dim,
+                            rotary_dim,
+                            dtype,
+                            use_rope,
+                            is_neox_style,
+                        )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_triton_moe_wna16.py

@@ -0,0 +1,252 @@
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_moe
+from sglang.srt.layers.moe.topk import TopKConfig, select_experts
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.srt.utils import get_device
+
+NUM_EXPERTS = [8, 64]
+TOP_KS = [2, 6]
+
+
+def quantize_weights(
+    w: torch.Tensor,
+    quant_type: str,
+    group_size: Optional[int],
+    zero_points: bool = False,
+    ref_zero_points_after_scales: bool = False,
+):
+    assert quant_type in ["w4a16", "w4a16b8", "w8a16", "w8a16b128"]
+    assert not zero_points or group_size is not None, (
+        "to have group zero points, group_size must be provided "
+        "(-1 group_size is channelwise)"
+    )
+
+    orig_device = w.device
+    orig_type = w.dtype
+    size_k, size_n = w.shape
+
+    assert w.is_floating_point(), "w must be float"
+
+    if group_size == -1:
+        group_size = size_k
+
+    # Reshape to [groupsize, -1]
+    if group_size is not None and group_size < size_k:
+        w = w.reshape((-1, group_size, size_n))
+        w = w.permute(1, 0, 2)
+        w = w.reshape((group_size, -1))
+
+    # Compute scale for each group
+    max_val = torch.max(w, 0, keepdim=True).values
+    min_val = torch.min(w, 0, keepdim=True).values
+
+    if quant_type == "w4a16":
+        max_q_val = 15
+        min_q_val = 0
+    elif quant_type == "w4a16b8":
+        max_q_val = 7
+        min_q_val = -1
+    elif quant_type == "w8a16":
+        max_q_val = 255
+        min_q_val = 0
+    elif quant_type == "w8a16b128":
+        max_q_val = 127
+        min_q_val = -128
+
+    w_s = torch.Tensor([1.0]).to(w.device)  # unscaled case
+    maybe_w_zp = None
+    if group_size is not None:
+        if zero_points:
+            w_s = (max_val - min_val).clamp(min=1e-5) / max_q_val
+            maybe_w_zp = (
+                torch.round(torch.abs(min_val / w_s)).clamp(min_q_val, max_q_val).int()
+            )
+        else:
+            # If the bias is such that there are no possible negative/positive
+            #  values, set the max value to inf to avoid divide by 0
+            w_s = torch.max(
+                abs(max_val / (max_q_val if max_q_val != 0 else torch.inf)),
+                abs(min_val / (min_q_val if min_q_val != 0 else torch.inf)),
+            )
+
+    # Quantize
+    w_q = torch.round(w / w_s).int() + (maybe_w_zp if zero_points else 0)
+    w_q = torch.clamp(w_q, min_q_val, max_q_val)
+
+    # Compute ref (dequantized)
+    # For some kernels (namely Machete) the zero-points are applied after the
+    # scales are applied, for this case computing the reference in similar way
+    # allows us to use tighter error tolerances in our unit tests.
+    if ref_zero_points_after_scales and maybe_w_zp is not None:
+        w_ref = w_q.to(orig_type) * w_s - maybe_w_zp.to(orig_type) * w_s
+    else:
+        w_ref = (w_q - (maybe_w_zp if zero_points else 0)).to(orig_type) * w_s
+
+    if quant_type == "w4a16b8":
+        w_q += 8
+    elif quant_type == "w8a16b128":
+        w_q += 128
+
+    # Restore original shapes
+    if group_size is not None and group_size < size_k:
+
+        def reshape_w(w):
+            w = w.reshape((group_size, -1, size_n))
+            w = w.permute(1, 0, 2)
+            w = w.reshape((size_k, size_n)).contiguous()
+            return w
+
+        w_q = reshape_w(w_q)
+        w_ref = reshape_w(w_ref)
+        w_s = w_s.reshape((-1, size_n)).contiguous()
+
+    if maybe_w_zp is not None:
+        maybe_w_zp = maybe_w_zp.reshape((-1, size_n)).contiguous()
+        maybe_w_zp = maybe_w_zp.to(device=orig_device)
+
+    return (
+        w_ref.to(device=orig_device),
+        w_q.to(device=orig_device),
+        w_s if group_size is not None else None,
+        maybe_w_zp,
+    )
+
+
+def torch_moe(a, w1, w2, score, topk):
+    B, D = a.shape
+    a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
+    out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_weight = topk_weight.view(-1)
+    topk_ids = topk_ids.view(-1)
+    for i in range(w1.shape[0]):
+        mask = topk_ids == i
+        if mask.sum():
+            out[mask] = SiluAndMul()(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(
+                0, 1
+            )
+    return (
+        out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
+    ).sum(dim=1)
+
+
+# fork from https://github.com/vllm-project/vllm/blob/main/tests/kernels/test_moe.py
+@pytest.mark.parametrize("m", [1, 32, 222])
+@pytest.mark.parametrize("n", [128, 1024, 2048])
+@pytest.mark.parametrize("k", [128, 1024])
+@pytest.mark.parametrize("e", NUM_EXPERTS)
+@pytest.mark.parametrize("topk", TOP_KS)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("group_size", [64, 128])
+@pytest.mark.parametrize("has_zp", [True, False])
+@pytest.mark.parametrize("weight_bits", [8])  # [4, 8])
+def test_fused_moe_wn16(
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    topk: int,
+    dtype: torch.dtype,
+    group_size: int,
+    has_zp: bool,
+    weight_bits: int,
+):
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+    print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
+    a = torch.randn((m, k), device=get_device(), dtype=dtype) / 10
+    w1 = torch.randn((e, 2 * n, k), device=get_device(), dtype=dtype) / 10
+    w2 = torch.randn((e, k, n), device=get_device(), dtype=dtype) / 10
+    score = torch.randn((m, e), device=get_device(), dtype=dtype)
+
+    if weight_bits == 4:
+        pack_factor = 2
+        quant_type = "w4a16" if has_zp else "w4a16b8"
+    elif weight_bits == 8:
+        pack_factor = 1
+        quant_type = "w8a16" if has_zp else "w8a16b128"
+
+    w1_ref = w1.clone()
+    w2_ref = w2.clone()
+    w1_qweight = torch.empty(
+        (e, 2 * n, k // pack_factor), device=get_device(), dtype=torch.uint8
+    )
+    w2_qweight = torch.empty(
+        (e, k, n // pack_factor), device=get_device(), dtype=torch.uint8
+    )
+    w1_scales = torch.empty(
+        (e, 2 * n, k // group_size), device=get_device(), dtype=dtype
+    )
+    w2_scales = torch.empty((e, k, n // group_size), device=get_device(), dtype=dtype)
+    w1_qzeros = torch.empty(
+        (e, 2 * n // pack_factor, k // group_size),
+        device=get_device(),
+        dtype=torch.uint8,
+    )
+    w2_qzeros = torch.empty(
+        (e, k // pack_factor, n // group_size), device=get_device(), dtype=torch.uint8
+    )
+
+    for i in range(e * 2):
+        expert_id = i % e
+        if i // e == 0:
+            w, w_ref, w_qweight, w_scales, w_qzeros = (
+                w1,
+                w1_ref,
+                w1_qweight,
+                w1_scales,
+                w1_qzeros,
+            )
+        else:
+            w, w_ref, w_qweight, w_scales, w_qzeros = (
+                w2,
+                w2_ref,
+                w2_qweight,
+                w2_scales,
+                w2_qzeros,
+            )
+        weight, qweight, scales, qzeros = quantize_weights(
+            w[expert_id].T, quant_type, group_size, has_zp, False
+        )
+        weight = weight.T
+        qweight = qweight.T.contiguous().to(torch.uint8)
+        scales = scales.T
+        if has_zp:
+            qzeros = qzeros.T.contiguous().to(torch.uint8)
+        if weight_bits == 4:
+            qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]
+            if has_zp:
+                qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]
+
+        w_ref[expert_id] = weight
+        w_qweight[expert_id] = qweight
+        w_scales[expert_id] = scales
+        if has_zp:
+            w_qzeros[expert_id] = qzeros
+
+    topk_output = select_experts(
+        hidden_states=a,
+        router_logits=score,
+        topk_config=TopKConfig(top_k=topk),
+    )
+
+    triton_output = fused_moe(
+        a,
+        w1_qweight,
+        w2_qweight,
+        topk_output,
+        use_int4_w4a16=weight_bits == 4,
+        use_int8_w8a16=weight_bits == 8,
+        w1_scale=w1_scales,
+        w2_scale=w2_scales,
+        w1_zp=w1_qzeros if has_zp else None,
+        w2_zp=w2_qzeros if has_zp else None,
+        block_shape=[0, group_size],
+    )
+    torch_output = torch_moe(a, w1_ref, w2_ref, score, topk)
+    torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)

third_party/sglang/test/manual/test_trtllm_fp8_kv_kernel.py

@@ -0,0 +1,481 @@
+"""
+Unit tests for TRTLLM FP8 KV cache fusion kernel.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
+    fused_fp8_set_kv_buffer,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestTRTLLMFP8KVKernel(CustomTestCase):
+    """Test fused FP8 KV cache write kernel correctness."""
+
+    @classmethod
+    def setUpClass(cls):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA not available")
+
+        if torch.cuda.get_device_capability()[0] < 9:
+            raise unittest.SkipTest("FP8 requires compute capability >= 9.0")
+
+    def _test_kernel_correctness(
+        self,
+        num_tokens,
+        num_kv_heads,
+        head_dim,
+        page_size,
+        use_scale,
+        input_ndim,
+        cache_ndim,
+    ):
+        """Compare Triton kernel output against naive implementation."""
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+
+        # Create input tensors
+        if input_ndim == 3:
+            k = torch.randn(
+                num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype
+            )
+            v = torch.randn(
+                num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype
+            )
+        else:
+            k = torch.randn(
+                num_tokens, num_kv_heads * head_dim, device=device, dtype=dtype
+            )
+            v = torch.randn(
+                num_tokens, num_kv_heads * head_dim, device=device, dtype=dtype
+            )
+
+        # Create cache tensors (use FP8 to match real runtime behavior)
+        num_pages = 128
+        total_slots = num_pages * page_size
+        cache_dtype = torch.float8_e4m3fn
+        if cache_ndim == 3:
+            k_cache_triton = torch.zeros(
+                total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
+            )
+            v_cache_triton = torch.zeros(
+                total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
+            )
+            k_cache_naive = torch.zeros(
+                total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
+            )
+            v_cache_naive = torch.zeros(
+                total_slots, num_kv_heads, head_dim, device=device, dtype=cache_dtype
+            )
+        else:
+            k_cache_triton = torch.zeros(
+                num_pages,
+                page_size,
+                num_kv_heads,
+                head_dim,
+                device=device,
+                dtype=cache_dtype,
+            )
+            v_cache_triton = torch.zeros(
+                num_pages,
+                page_size,
+                num_kv_heads,
+                head_dim,
+                device=device,
+                dtype=cache_dtype,
+            )
+            k_cache_naive = torch.zeros(
+                num_pages,
+                page_size,
+                num_kv_heads,
+                head_dim,
+                device=device,
+                dtype=cache_dtype,
+            )
+            v_cache_naive = torch.zeros(
+                num_pages,
+                page_size,
+                num_kv_heads,
+                head_dim,
+                device=device,
+                dtype=cache_dtype,
+            )
+
+        # Create cache locations (ensure unique indices to avoid race conditions)
+        cache_loc = torch.randperm(total_slots, device=device, dtype=torch.int32)[
+            :num_tokens
+        ]
+
+        # Optional scales
+        k_scale = 0.5 if use_scale else None
+        v_scale = 0.75 if use_scale else None
+
+        # Run Triton kernel
+        fused_fp8_set_kv_buffer(
+            k.clone(),
+            v.clone(),
+            k_cache_triton,
+            v_cache_triton,
+            cache_loc,
+            k_scale,
+            v_scale,
+            page_size,
+            use_triton=True,
+        )
+
+        # Run naive fallback
+        fused_fp8_set_kv_buffer(
+            k.clone(),
+            v.clone(),
+            k_cache_naive,
+            v_cache_naive,
+            cache_loc,
+            k_scale,
+            v_scale,
+            page_size,
+            use_triton=False,
+        )
+
+        # Compare results (bit-exact match expected)
+        self.assertTrue(
+            torch.equal(k_cache_triton, k_cache_naive),
+            "K cache mismatch between Triton and naive",
+        )
+        self.assertTrue(
+            torch.equal(v_cache_triton, v_cache_naive),
+            "V cache mismatch between Triton and naive",
+        )
+
+    def test_basic_3d_input_3d_cache(self):
+        """Test basic case: 3D input, 3D cache, no scale."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=False,
+            input_ndim=3,
+            cache_ndim=3,
+        )
+
+    def test_basic_3d_input_4d_cache(self):
+        """Test basic case: 3D input, 4D cache, no scale."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=False,
+            input_ndim=3,
+            cache_ndim=4,
+        )
+
+    def test_with_scale_3d_cache(self):
+        """Test with scale: 3D input, 3D cache."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=True,
+            input_ndim=3,
+            cache_ndim=3,
+        )
+
+    def test_with_scale_4d_cache(self):
+        """Test with scale: 3D input, 4D cache."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=True,
+            input_ndim=3,
+            cache_ndim=4,
+        )
+
+    def test_2d_input_3d_cache(self):
+        """Test 2D input (flattened): 2D input, 3D cache."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=False,
+            input_ndim=2,
+            cache_ndim=3,
+        )
+
+    def test_2d_input_4d_cache(self):
+        """Test 2D input (flattened): 2D input, 4D cache."""
+        self._test_kernel_correctness(
+            num_tokens=16,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=False,
+            input_ndim=2,
+            cache_ndim=4,
+        )
+
+    def test_single_token(self):
+        """Test edge case: single token."""
+        self._test_kernel_correctness(
+            num_tokens=1,
+            num_kv_heads=8,
+            head_dim=128,
+            page_size=16,
+            use_scale=True,
+            input_ndim=3,
+            cache_ndim=3,
+        )
+
+    def test_large_batch(self):
+        """Test larger batch size."""
+        self._test_kernel_correctness(
+            num_tokens=128,
+            num_kv_heads=16,
+            head_dim=64,
+            page_size=16,
+            use_scale=True,
+            input_ndim=3,
+            cache_ndim=4,
+        )
+
+    def test_different_head_dims(self):
+        """Test different head dimensions."""
+        for head_dim in [64, 128]:
+            self._test_kernel_correctness(
+                num_tokens=16,
+                num_kv_heads=8,
+                head_dim=head_dim,
+                page_size=16,
+                use_scale=False,
+                input_ndim=3,
+                cache_ndim=3,
+            )
+
+    def test_empty_input(self):
+        """Test edge case: empty input (0 tokens)."""
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+        num_kv_heads = 8
+        head_dim = 128
+        page_size = 16
+        num_tokens = 0
+
+        # Empty inputs
+        k = torch.randn(num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
+        v = torch.randn(num_tokens, num_kv_heads, head_dim, device=device, dtype=dtype)
+
+        # Cache (use FP8 to match real runtime behavior)
+        total_slots = 128
+        k_cache = torch.zeros(
+            total_slots,
+            num_kv_heads,
+            head_dim,
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        v_cache = torch.zeros(
+            total_slots,
+            num_kv_heads,
+            head_dim,
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+
+        # Empty cache locations
+        cache_loc = torch.empty(num_tokens, device=device, dtype=torch.int32)
+
+        # Should not crash
+        fused_fp8_set_kv_buffer(
+            k,
+            v,
+            k_cache,
+            v_cache,
+            cache_loc,
+            k_scale=None,
+            v_scale=None,
+            page_size=page_size,
+        )
+
+    def test_fp8_kv_kernel_accepts_tensor_scales(self):
+        """
+        Regression test for B200 Triton compilation issue.
+
+        This test ensures that fused_fp8_set_kv_buffer correctly handles
+        k_scale/v_scale when they are 0-dimensional tensors (torch.nn.Parameter).
+
+        Previously, Triton would treat 0-D tensor arguments as pointers,
+        causing a type error when performing "1.0 / k_scale" inside the kernel.
+        The fix converts tensor scales to Python floats in the wrapper.
+        """
+        device = torch.device("cuda")
+
+        num_tokens = 4
+        num_kv_heads = 2
+        head_dim = 64
+        page_size = 16
+        total_slots = page_size
+
+        k = torch.randn(
+            num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
+        )
+        v = torch.randn_like(k)
+
+        k_cache = torch.empty(
+            total_slots,
+            num_kv_heads,
+            head_dim,
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        v_cache = torch.empty_like(k_cache)
+
+        cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
+
+        # Use 0D tensor form of scale to reproduce the original bug scenario
+        k_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
+        v_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
+
+        # Old code would trigger Triton's IncompatibleTypeError here
+        # New code should handle this gracefully by converting to float
+        fused_fp8_set_kv_buffer(
+            k,
+            v,
+            k_cache,
+            v_cache,
+            cache_loc,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            page_size=page_size,
+            use_triton=True,
+        )
+
+        # If we get here without exception, the regression is fixed
+
+    def test_fp8_kv_kernel_cuda_graph_compatible(self):
+        """
+        Regression test for CUDA graph capture compatibility.
+
+        This test ensures that fused_fp8_set_kv_buffer works correctly within
+        CUDA graph capture, which is used in production for performance.
+
+        Previously, float(k_scale) caused GPU→CPU synchronization, triggering
+        cudaErrorStreamCaptureUnsupported during graph capture. The fix computes
+        inverse scales purely on GPU using tensor operations.
+        """
+        device = torch.device("cuda")
+
+        num_tokens = 4
+        num_kv_heads = 2
+        head_dim = 64
+        page_size = 16
+        total_slots = page_size
+
+        k = torch.randn(
+            num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
+        )
+        v = torch.randn_like(k)
+
+        k_cache = torch.empty(
+            total_slots,
+            num_kv_heads,
+            head_dim,
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        v_cache = torch.empty_like(k_cache)
+
+        cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
+
+        # Use 0D tensor scales (like nn.Parameter) to reproduce production scenario
+        k_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
+        v_scale = torch.tensor(1.0, device=device, dtype=torch.float32)
+
+        # Test that kernel works under CUDA graph capture
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            # Old code would fail here with cudaErrorStreamCaptureUnsupported
+            # New code should succeed because all operations stay on GPU
+            fused_fp8_set_kv_buffer(
+                k,
+                v,
+                k_cache,
+                v_cache,
+                cache_loc,
+                k_scale=k_scale,
+                v_scale=v_scale,
+                page_size=page_size,
+                use_triton=True,
+            )
+
+        # Replay the graph to verify it works
+        graph.replay()
+
+        # If we get here without exception, CUDA graph compatibility is confirmed
+
+    def test_fp8_kv_kernel_cuda_graph_compatible_no_scale(self):
+        """
+        Regression test for CUDA graph capture compatibility without scales.
+
+        This test ensures that fused_fp8_set_kv_buffer works correctly within
+        CUDA graph capture when k_scale/v_scale are None (use_provided_scale=False).
+
+        Previously, the code created new GPU tensors (torch.tensor(1.0, device=...))
+        during graph capture, triggering cudaErrorStreamCaptureUnsupported.
+        The fix passes dummy pointers when use_provided_scale=False, as the kernel
+        uses constant 1.0 and Triton optimizes away the pointer loads.
+        """
+        device = torch.device("cuda")
+
+        num_tokens = 4
+        num_kv_heads = 2
+        head_dim = 64
+        page_size = 16
+        total_slots = page_size
+
+        k = torch.randn(
+            num_tokens, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16
+        )
+        v = torch.randn_like(k)
+
+        k_cache = torch.empty(
+            total_slots,
+            num_kv_heads,
+            head_dim,
+            device=device,
+            dtype=torch.float8_e4m3fn,
+        )
+        v_cache = torch.empty_like(k_cache)
+
+        cache_loc = torch.arange(num_tokens, device=device, dtype=torch.int32)
+
+        # Test that kernel works under CUDA graph capture WITHOUT scales
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            # No k_scale/v_scale provided - use_provided_scale=False branch
+            # Old code would fail here with cudaErrorStreamCaptureUnsupported
+            # New code should succeed by using dummy pointers
+            fused_fp8_set_kv_buffer(
+                k,
+                v,
+                k_cache,
+                v_cache,
+                cache_loc,
+                page_size=page_size,
+                use_triton=True,
+            )
+
+        # Replay the graph to verify it works
+        graph.replay()
+
+        # If we get here without exception, no-scale CUDA graph compatibility is confirmed
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_two_batch_overlap.py

@@ -0,0 +1,152 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.batch_overlap.two_batch_overlap import (
+    compute_split_seq_index,
+    compute_split_token_index,
+)
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_ENABLE_THINKING_MODEL_NAME_FOR_TEST,
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    popen_launch_server,
+)
+
+
+class TestTwoBatchOverlap(unittest.TestCase):
+    @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):
+            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",
+                    "--moe-a2a-backend",
+                    "deepep",
+                    "--deepep-mode",
+                    "normal",
+                    "--disable-cuda-graph",  # DeepEP normal does not support CUDA Graph
+                    "--enable-two-batch-overlap",
+                ],
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_generate_single_prompt(self):
+        response = requests.post(
+            self.base_url + "/generate",
+            # we use an uncommon start to minimise the chance that the cache is hit by chance
+            json={
+                "text": "_ 1+1=2, 1+2=3, 1+3=4, 1+4=",
+                "sampling_params": {"temperature": 0, "max_new_tokens": 8},
+            },
+        )
+        print(f"{response.json()=}")
+        self.assertEqual(response.json()["text"], "5, 1+5=6")
+
+    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 TestTwoBatchOverlapUnitTest(unittest.TestCase):
+    def test_compute_split_seq_and_token_index(self):
+        for num_tokens, expect in [
+            (0, 0),
+            (100, 50),
+            (99, 49),
+        ]:
+            actual = compute_split_seq_index(
+                forward_mode=ForwardMode.DECODE,
+                num_tokens=num_tokens,
+                extend_lens=None,
+                token_num_per_seq=1,
+            )
+            self.assertEqual(actual, expect)
+
+        for extend_lens, expect in [
+            ([], (0, 0)),
+            ([42], (0, 21)),
+            ([42, 999], (1, 520)),
+            ([999, 42], (0, 520)),
+            ([498, 502], (1, 498)),
+            ([4096, 4096, 4096, 4096], (2, 8192)),
+            ([4095, 4096, 4096, 4096, 1], (2, 8191)),
+            ([1, 4095, 4096, 4096, 4096], (3, 8192)),
+            ([4097, 4096, 4096, 4095, 1], (2, 8193)),
+            ([1, 1, 1, 1, 99999], (4, 50001)),
+            ([99999, 1, 1, 1, 1], (0, 50001)),
+        ]:
+            actual_seq_idx = compute_split_seq_index(
+                forward_mode=ForwardMode.EXTEND,
+                num_tokens=None,
+                extend_lens=extend_lens,
+                token_num_per_seq=None,
+            )
+            actual_token_idx = compute_split_token_index(
+                split_seq_index=actual_seq_idx,
+                forward_mode=ForwardMode.EXTEND,
+                extend_seq_lens=extend_lens,
+                token_num_per_seq=None,
+            )
+            actual = (actual_seq_idx, actual_token_idx)
+            print(f"{extend_lens=} {expect=} {actual=}")
+            self.assertEqual(actual, expect)
+
+
+class TestQwen3TwoBatchOverlap(TestTwoBatchOverlap):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_ENABLE_THINKING_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.api_key = "sk-1234"
+        with 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=[
+                    "--trust-remote-code",
+                    "--tp",
+                    "2",
+                    "--dp",
+                    "2",
+                    "--enable-dp-attention",
+                    "--moe-a2a-backend",
+                    "deepep",
+                    "--deepep-mode",
+                    "normal",
+                    "--disable-cuda-graph",  # DeepEP normal does not support CUDA Graph
+                    "--enable-two-batch-overlap",
+                ],
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_vertex_endpoint.py

@@ -0,0 +1,64 @@
+"""
+python3 -m unittest test_vertex_endpoint.TestVertexEndpoint.test_vertex_generate
+"""
+
+import unittest
+from http import HTTPStatus
+
+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 TestVertexEndpoint(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_SMALL_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=["--cuda-graph-max-bs", 2],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def run_generate(self, parameters):
+        data = {
+            "instances": [
+                {"text": "The capital of France is"},
+                {"text": "The capital of China is"},
+            ],
+            "parameters": parameters,
+        }
+        response = requests.post(self.base_url + "/vertex_generate", json=data)
+        response_json = response.json()
+        assert len(response_json["predictions"]) == len(data["instances"])
+        return response_json
+
+    def test_vertex_generate(self):
+        for parameters in [None, {"sampling_params": {"max_new_tokens": 4}}]:
+            self.run_generate(parameters)
+
+    def test_vertex_generate_fail(self):
+        data = {
+            "instances": [
+                {"prompt": "The capital of France is"},
+            ],
+        }
+        response = requests.post(self.base_url + "/vertex_generate", json=data)
+        assert response.status_code == HTTPStatus.BAD_REQUEST
+
+
+if __name__ == "__main__":
+    unittest.main()

third_party/sglang/test/manual/test_vlm_accuracy.py

@@ -0,0 +1,320 @@
+""" """
+
+import unittest
+from typing import List, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers import AutoModel, AutoProcessor, AutoTokenizer
+
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
+from sglang.srt.managers.mm_utils import embed_mm_inputs, init_mm_embedding_cache
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputs,
+)
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
+from sglang.srt.parser.conversation import generate_chat_conv
+from sglang.srt.server_args import ServerArgs
+from sglang.test.test_utils import download_image_with_retry
+
+
+# Test the logits output between HF and SGLang
+class VisionLLMLogitsBase(unittest.IsolatedAsyncioTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.image_url = "https://github.com/sgl-project/sglang/blob/main/examples/assets/example_image.png?raw=true"
+        cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        cls.model_path = ""
+        cls.chat_template = ""
+        cls.processor = ""
+        cls.main_image = download_image_with_retry(cls.image_url)
+
+    def compare_outputs(self, sglang_output: torch.Tensor, hf_output: torch.Tensor):
+        # Convert to float32 for numerical stability if needed
+        hf = hf_output.float()
+        sg = sglang_output.float()
+
+        # Basic shape and dtype comparison
+        print("\n=== Basic Properties ===")
+        print(f"Shapes match: {hf.shape == sg.shape}")
+        print(f"HF shape: {hf.shape}, SGLang shape: {sg.shape}")
+        print(f"HF dtype: {hf.dtype}, SGLang dtype: {sg.dtype}")
+
+        # Move tensors to CPU for numpy operations
+        hf_np = hf.cpu().numpy()
+        sg_np = sg.cpu().numpy()
+
+        # Statistical metrics
+        print("\n=== Statistical Metrics ===")
+        print(f"Mean absolute difference: {torch.mean(torch.abs(hf - sg)).item():.6f}")
+        print(f"Max absolute difference: {torch.max(torch.abs(hf - sg)).item():.6f}")
+        print(f"Mean squared error: {torch.mean((hf - sg) ** 2).item():.6f}")
+        print(
+            f"Root mean squared error: {torch.sqrt(torch.mean((hf - sg) ** 2)).item():.6f}"
+        )
+
+        # Cosine similarity (across feature dimension)
+        cos_sim = F.cosine_similarity(hf, sg)
+        print(f"Mean cosine similarity: {torch.mean(cos_sim).item():.6f}")
+        print(f"Min cosine similarity: {torch.min(cos_sim).item():.6f}")
+
+        # Find largest absolute differences
+        print("\n=== Largest Absolute Differences ===")
+        diffs = torch.abs(hf - sg)
+        flat_diffs = diffs.flatten()
+
+        # Get indices of top 10 differences
+        top_k = 10
+        top_values, top_flat_indices = torch.topk(flat_diffs, top_k)
+
+        # Convert flat indices to multidimensional indices
+        top_indices = np.unravel_index(top_flat_indices.cpu().numpy(), diffs.shape)
+
+        print(f"\nTop {top_k} largest absolute differences:")
+        print(
+            "Index".ljust(30)
+            + "Difference".ljust(15)
+            + "HF Value".ljust(15)
+            + "SGLang Value"
+        )
+        print("-" * 75)
+
+        for i in range(top_k):
+            # Get the index tuple for this difference
+            idx = tuple(dim[i] for dim in top_indices)
+        diff_val = top_values[i].item()
+        hf_val = hf[idx].item()
+        sg_val = sg[idx].item()
+
+        # Format the index tuple and values
+        idx_str = str(idx)
+        print(f"{idx_str:<30}{diff_val:<15.6f}{hf_val:<15.6f}{sg_val:.6f}")
+
+        np.testing.assert_allclose(hf_np, sg_np)
+
+    def get_completion_request(self) -> ChatCompletionRequest:
+        json_str = f"""
+        {{
+  "model": "{self.model_path}",
+  "messages": [
+    {{
+      "role": "user",
+      "content": [
+        {{
+          "type": "image_url",
+          "image_url": {{
+            "url": "{self.image_url}"
+          }}
+        }},
+        {{
+          "type": "text",
+          "text": "What's in this picture?"
+        }}
+      ]
+    }}
+  ]
+}}
+        """
+
+        return ChatCompletionRequest.model_validate_json(json_str)
+
+    def get_processor_output(self, req: Optional[ChatCompletionRequest] = None):
+        if req is None:
+            req = self.get_completion_request()
+        conv = generate_chat_conv(req, template_name=self.chat_template)
+        text = conv.get_prompt()
+
+        # Process inputs using processor
+        # FIXME: the formal arguments may differ
+        inputs = self.processor(
+            text=[text],
+            images=[self.main_image],
+            return_tensors="pt",
+        ).to(self.device)
+
+        return inputs
+
+    def get_sglang_model(self):
+        self.model_runner = ModelRunner(
+            model_config=ModelConfig(self.model_path, model_override_args="{}"),
+            mem_fraction_static=0.8,
+            gpu_id=0,
+            tp_rank=0,
+            tp_size=1,
+            pp_rank=0,
+            pp_size=1,
+            nccl_port=12435,
+            server_args=ServerArgs(
+                model_path=self.model_path,
+                disable_cuda_graph=True,
+            ),
+        )
+        return self.model_runner.model
+
+
+class TestMiniCPMV2_6Logits(VisionLLMLogitsBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model_path = "openbmb/MiniCPM-V-2_6"
+        cls.tokenizer = AutoTokenizer.from_pretrained(
+            cls.model_path, trust_remote_code=True
+        )
+        cls.processor = AutoProcessor.from_pretrained(
+            cls.model_path, trust_remote_code=True
+        )
+        cls.chat_template = "minicpmv"
+
+        cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        cls.hf_model = (
+            AutoModel.from_pretrained(
+                cls.model_path, torch_dtype=torch.bfloat16, trust_remote_code=True
+            )
+            .eval()
+            .to(cls.device)
+        )
+        init_mm_embedding_cache()
+
+    async def test_vlm_embedding_output(self):
+        """
+        Compares the embedding output of vlm
+        """
+        inputs = self.get_processor_output()
+
+        with torch.no_grad():
+            # hf
+            model_inputs = {
+                "input_ids": inputs.input_ids,
+                "image_bound": inputs.image_bound,
+                "pixel_values": inputs.pixel_values,
+                "tgt_sizes": inputs.tgt_sizes,
+            }
+            hf_output, _ = self.hf_model.get_vllm_embedding(
+                model_inputs,
+            )
+            hf_output = hf_output.squeeze(0)
+
+            # sglang
+            model = self.get_sglang_model()
+            input_ids = inputs["input_ids"].to(self.device).flatten()
+
+            pixel_values = inputs["pixel_values"]
+            tgt_sizes = inputs["tgt_sizes"]
+            pixel_values_flat: List[torch.Tensor] = []
+            tgt_sizes_flat: List[torch.Tensor] = []
+            for pixel_b, tgt_b in zip(pixel_values, tgt_sizes):
+                # per image
+                if len(pixel_b) != len(tgt_b):
+                    raise ValueError(
+                        "Inconsistent N lengths, found: "
+                        f"{len(pixel_b)} vs {len(tgt_b)}"
+                    )
+                for pixel_n, tgt_n in zip(pixel_b, tgt_b):
+                    pixel_values_flat += [pixel_n]
+                    tgt_sizes_flat += [tgt_n]
+
+            im_start_id, im_end_id = (
+                self.tokenizer.im_start_id,
+                self.tokenizer.im_end_id,
+            )
+            slice_start_id, slice_end_id = (
+                self.tokenizer.slice_start_id,
+                self.tokenizer.slice_end_id,
+            )
+
+            image_offsets = BaseMultimodalProcessor.get_mm_items_offset_by_pair(
+                input_ids=input_ids, mm_start_id=im_start_id, mm_end_id=im_end_id
+            )
+            slice_offsets = BaseMultimodalProcessor.get_mm_items_offset_by_pair(
+                input_ids=input_ids, mm_start_id=slice_start_id, mm_end_id=slice_end_id
+            )
+            image_offsets.extend(slice_offsets)
+            image_offsets = sorted(image_offsets)
+
+            sglang_output = embed_mm_inputs(
+                mm_inputs_list=[
+                    MultimodalInputs(
+                        mm_items=[
+                            MultimodalDataItem(
+                                feature=pixel_values_flat,
+                                offsets=image_offsets,
+                                tgt_size=tgt_sizes_flat,
+                                modality=Modality.IMAGE,
+                                pad_value=self.processor.tokenizer.unk_token_id,
+                            )
+                        ]
+                    ),
+                ],
+                extend_prefix_lens=[0],
+                extend_seq_lens=[input_ids.shape[0]],
+                input_ids=input_ids,
+                input_embedding=model.get_input_embeddings(),
+                multimodal_model=model,
+                placeholder_tokens={
+                    Modality.IMAGE: self.processor.tokenizer.unk_token_id,
+                },
+            )
+
+        self.compare_outputs(sglang_output, hf_output)
+
+
+class TestMiniCPMV4Logits(VisionLLMLogitsBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model_path = "openbmb/MiniCPM-V-4"
+        cls.tokenizer = AutoTokenizer.from_pretrained(
+            cls.model_path, trust_remote_code=True
+        )
+        cls.processor = AutoProcessor.from_pretrained(
+            cls.model_path, trust_remote_code=True
+        )
+        cls.chat_template = "minicpmv"
+
+        cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        cls.hf_model = (
+            AutoModel.from_pretrained(
+                cls.model_path, torch_dtype=torch.bfloat16, trust_remote_code=True
+            )
+            .eval()
+            .to(cls.device)
+        )
+        init_mm_embedding_cache()
+
+    async def test_vlm_embedding_output(self):
+        """
+        Compares the embedding output of vlm
+        """
+        inputs = self.get_processor_output()
+
+        with torch.no_grad():
+            # hf
+            model_inputs = {
+                "input_ids": inputs.input_ids,
+                "image_bound": inputs.image_bound,
+                "pixel_values": inputs.pixel_values,
+                "tgt_sizes": inputs.tgt_sizes,
+            }
+            hf_output = self.hf_model.get_input_embeddings()(inputs.input_ids)
+
+            # sglang
+            model = self.get_model()
+            sglang_output = self.vlm_func(
+                model,
+                input_ids=inputs.input_ids.to(self.device),
+                pixel_values=inputs.pixel_values,
+                image_bound=inputs.image_bound.to(self.device),
+                tgt_sizes=inputs.tgt_sizes.to(self.device),
+                input_embedding=model.get_input_embeddings(),
+                multimodal_model=model,
+                placeholder_tokens={
+                    Modality.IMAGE: self.processor.tokenizer.unk_token_id,
+                },
+            )
+
+        self.compare_outputs(sglang_output, hf_output)