[CI] Trim the base-c 4-gpu-h100 stage from 5 shards to 4 (#35407)

python/sglang/test/kits/pd_parity_kit.py

@@ -0,0 +1,70 @@
+import time
+
+import requests
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.server_fixtures.disaggregation_fixture import assert_process_healthy
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    popen_launch_server,
+)
+
+
+class PDLogprobParityMixin:
+    # Mix in before the PD server fixture, which owns the P/D launches.
+    reference_parallel_args = []
+    baseline_args = []
+
+    @staticmethod
+    def generate(base_url):
+        response = requests.post(
+            base_url + "/generate",
+            json={
+                "input_ids": [1] + [100 + i % 1000 for i in range(256)],
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 4,
+                    "ignore_eos": True,
+                },
+                "return_logprob": True,
+                "top_logprobs_num": 5,
+            },
+            timeout=120,
+        )
+        response.raise_for_status()
+        return response.json()["meta_info"]
+
+    def test_logprob_parity(self):
+        baseline = popen_launch_server(
+            self.model,
+            self.lb_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=self.reference_parallel_args
+            + ["--trust-remote-code"]
+            + self.baseline_args,
+            env=self.extra_prefill_env,
+        )
+        try:
+            reference = self.generate(self.lb_url)
+        finally:
+            kill_process_tree(baseline.pid, wait_timeout=60)
+        time.sleep(5)
+
+        self.launch_all()
+        disaggregated = self.generate(self.lb_url)
+
+        reference_logprobs = reference["output_token_logprobs"]
+        disaggregated_logprobs = disaggregated["output_token_logprobs"]
+        self.assertEqual(
+            [item[1] for item in reference_logprobs],
+            [item[1] for item in disaggregated_logprobs],
+        )
+        self.assertEqual(len(reference_logprobs), 4)
+        for reference_item, disaggregated_item in zip(
+            reference_logprobs, disaggregated_logprobs
+        ):
+            self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
+
+        assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
+        assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
+        assert_process_healthy(self, "decode", self.process_decode, self.decode_url)

test/registered/disaggregation/test_disaggregation_kimi_linear.py

@@ -1,24 +1,17 @@
-import time
 import unittest
 
-import requests
-
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
-    assert_process_healthy,
-)
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    popen_launch_server,
 )
 
 register_cuda_ci(est_time=480, stage="base-c", runner_config="4-gpu-h100")
 
 KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
 SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
-SERVER_ARGS = [
+
+DETERMINISTIC_ARGS = [
     "--skip-tokenizer-init",
     "--random-seed",
     "1",
@@ -34,74 +27,19 @@ SERVER_ARGS = [
 ]
 
 
-class TestKimiLinearHeterogeneousTPDisaggregation(PDDisaggregationServerBase):
+class TestKimiLinearHeterogeneousTPDisaggregation(
+    PDLogprobParityMixin, PDDisaggregationServerBase
+):
+    model = KIMI_LINEAR_MODEL
+    extra_prefill_env = SERVER_ENV
+    extra_decode_env = SERVER_ENV
     prefill_tp_size = 2
     decode_tp_size = 1
     decode_base_gpu_id = 2
     reference_parallel_args = ["--tp-size", "2"]
-    extra_prefill_args = SERVER_ARGS
+    baseline_args = DETERMINISTIC_ARGS
-    extra_decode_args = SERVER_ARGS
+    extra_prefill_args = DETERMINISTIC_ARGS
-    extra_prefill_env = SERVER_ENV
+    extra_decode_args = DETERMINISTIC_ARGS
-    extra_decode_env = SERVER_ENV
-
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        cls.model = KIMI_LINEAR_MODEL
-
-    @staticmethod
-    def generate(base_url):
-        response = requests.post(
-            base_url + "/generate",
-            json={
-                "input_ids": [1] + [100 + i % 1000 for i in range(256)],
-                "sampling_params": {
-                    "temperature": 0,
-                    "max_new_tokens": 4,
-                    "ignore_eos": True,
-                },
-                "return_logprob": True,
-                "top_logprobs_num": 5,
-            },
-            timeout=120,
-        )
-        response.raise_for_status()
-        return response.json()["meta_info"]
-
-    def test_logprob_parity(self):
-        baseline = popen_launch_server(
-            self.model,
-            self.lb_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=self.reference_parallel_args
-            + ["--trust-remote-code"]
-            + SERVER_ARGS,
-            env=SERVER_ENV,
-        )
-        try:
-            reference = self.generate(self.lb_url)
-        finally:
-            kill_process_tree(baseline.pid, wait_timeout=60)
-        time.sleep(5)
-
-        self.launch_all()
-        disaggregated = self.generate(self.lb_url)
-
-        reference_logprobs = reference["output_token_logprobs"]
-        disaggregated_logprobs = disaggregated["output_token_logprobs"]
-        self.assertEqual(
-            [item[1] for item in reference_logprobs],
-            [item[1] for item in disaggregated_logprobs],
-        )
-        self.assertEqual(len(reference_logprobs), 4)
-        for reference_item, disaggregated_item in zip(
-            reference_logprobs, disaggregated_logprobs
-        ):
-            self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
-
-        assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
-        assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
-        assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
 
 
 class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggregation):
@@ -109,7 +47,7 @@ class TestKimiLinearPipelineDisaggregation(TestKimiLinearHeterogeneousTPDisaggre
     decode_tp_size = 1
     decode_base_gpu_id = 2
     reference_parallel_args = ["--tp-size", "1", "--pp-size", "2"]
-    extra_prefill_args = SERVER_ARGS + ["--pp-size", "2"]
+    extra_prefill_args = DETERMINISTIC_ARGS + ["--pp-size", "2"]
 
 
 if __name__ == "__main__":

test/registered/disaggregation/test_disaggregation_unified_memory.py

@@ -1,53 +1,24 @@
-"""PD disaggregation with --enable-unified-memory (MLA hybrid-Mamba).
-
-Guards the unified-memory PD transfer scheme end to end: whole page-envelope
-KV registration (`UnifiedMLATokenToKVPool.get_contiguous_buf_infos`), whole
-slot-envelope KDA/mamba state transfer, virtual->physical index translation at
-the prefill send / decode prealloc sites, and the compaction move gate. A
-regression in any of them shifts the decode-side KV/state bytes and breaks
-logprob parity with the non-PD unified-memory reference.
-
-`--attention-backend` is deliberately NOT pinned, matching
-`models_e2e/test_kimi_linear_unified_memory.py`, which documents that pinning
-hides defects reachable only under the resolved default. The transferred bytes
-are backend-independent, so the default (fa3 on this suite's H100 runner) covers
-this file's subject either way. The linear-attn/Mamba backends stay pinned to
-triton -- the page-major layout requires them.
-
-`--enable-deterministic-inference` is deliberately NOT set. It would only guard
-against batch-shape-dependent kernel variation, and the reference and P+D paths
-run the same shapes: measured, two fresh servers on separate GPUs produce
-bit-identical logits without it. Setting it would narrow the test to the
-batch-invariant op set and a non-default sampling backend -- a less
-representative config -- and couple a PD-transfer test to the deterministic code
-path, so a defect there would fail this file for an unrelated reason.
-"""
-
-import time
 import unittest
 
-import requests
-
-from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.pd_parity_kit import PDLogprobParityMixin
 from sglang.test.server_fixtures.disaggregation_fixture import (
     PDDisaggregationServerBase,
-    assert_process_healthy,
-)
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    popen_launch_server,
 )
 
-register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=900, stage="base-b", runner_config="2-gpu-large")
 
 KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random"
 SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"}
-SERVER_ARGS = [
+
+# --attention-backend and --enable-deterministic-inference are deliberately
+# absent: bytes are backend-independent and both paths run identical shapes.
+UNIFIED_MEMORY_ARGS = [
     "--skip-tokenizer-init",
     "--random-seed",
     "1",
     "--enable-unified-memory",
+    # Page-major layout requires triton for both.
     "--linear-attn-backend",
     "triton",
     "--mamba-backend",
@@ -63,75 +34,19 @@ SERVER_ARGS = [
 ]
 
 
-class TestUnifiedMemoryDisaggregation(PDDisaggregationServerBase):
+class TestUnifiedMemoryDisaggregation(PDLogprobParityMixin, PDDisaggregationServerBase):
     """1 prefill + 1 decode, both with --enable-unified-memory, vs a non-PD
     unified-memory reference server."""
 
+    model = KIMI_LINEAR_MODEL
+    extra_prefill_env = SERVER_ENV
+    extra_decode_env = SERVER_ENV
     prefill_tp_size = 1
     decode_tp_size = 1
     decode_base_gpu_id = 1
-    extra_prefill_args = SERVER_ARGS
+    baseline_args = UNIFIED_MEMORY_ARGS
-    extra_decode_args = SERVER_ARGS
+    extra_prefill_args = UNIFIED_MEMORY_ARGS
-    extra_prefill_env = SERVER_ENV
+    extra_decode_args = UNIFIED_MEMORY_ARGS
-    extra_decode_env = SERVER_ENV
-    baseline_args = SERVER_ARGS
-
-    @classmethod
-    def setUpClass(cls):
-        super().setUpClass()
-        cls.model = KIMI_LINEAR_MODEL
-
-    @staticmethod
-    def generate(base_url):
-        response = requests.post(
-            base_url + "/generate",
-            json={
-                "input_ids": [1] + [100 + i % 1000 for i in range(256)],
-                "sampling_params": {
-                    "temperature": 0,
-                    "max_new_tokens": 4,
-                    "ignore_eos": True,
-                },
-                "return_logprob": True,
-                "top_logprobs_num": 5,
-            },
-            timeout=120,
-        )
-        response.raise_for_status()
-        return response.json()["meta_info"]
-
-    def test_logprob_parity(self):
-        baseline = popen_launch_server(
-            self.model,
-            self.lb_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--trust-remote-code"] + self.baseline_args,
-            env=SERVER_ENV,
-        )
-        try:
-            reference = self.generate(self.lb_url)
-        finally:
-            kill_process_tree(baseline.pid, wait_timeout=60)
-        time.sleep(5)
-
-        self.launch_all()
-        disaggregated = self.generate(self.lb_url)
-
-        reference_logprobs = reference["output_token_logprobs"]
-        disaggregated_logprobs = disaggregated["output_token_logprobs"]
-        self.assertEqual(
-            [item[1] for item in reference_logprobs],
-            [item[1] for item in disaggregated_logprobs],
-        )
-        self.assertEqual(len(reference_logprobs), 4)
-        for reference_item, disaggregated_item in zip(
-            reference_logprobs, disaggregated_logprobs
-        ):
-            self.assertAlmostEqual(reference_item[0], disaggregated_item[0], delta=0.05)
-
-        assert_process_healthy(self, "load balancer", self.process_lb, self.lb_url)
-        assert_process_healthy(self, "prefill", self.process_prefill, self.prefill_url)
-        assert_process_healthy(self, "decode", self.process_decode, self.decode_url)
 
 
 class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregation):
@@ -142,10 +57,10 @@ class TestUnifiedMemoryDisaggregationChunkedPrefill(TestUnifiedMemoryDisaggregat
     chunk size so any parity break isolates to the PD transfer.
     """
 
-    _chunked_args = SERVER_ARGS + ["--chunked-prefill-size", "64"]
+    _chunked_args = UNIFIED_MEMORY_ARGS + ["--chunked-prefill-size", "64"]
+    baseline_args = _chunked_args
     extra_prefill_args = _chunked_args
     extra_decode_args = _chunked_args
-    baseline_args = _chunked_args
 
 
 if __name__ == "__main__":

test/registered/ep/test_deepep_small.py

@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=478, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=270, stage="base-c", runner_config="4-gpu-h100")
 
 
 class TestPureDP(CustomTestCase):
@@ -66,51 +66,6 @@ class TestPureDP(CustomTestCase):
         self.assertGreater(metrics["score"], 0.60)
 
 
-class TestHybridDPTP(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--trust-remote-code",
-                "--tp",
-                "4",
-                "--enable-dp-attention",
-                "--dp",
-                "2",
-                "--moe-a2a-backend",
-                "deepep",
-                "--cuda-graph-max-bs-decode",
-                "128",
-                "--max-running-requests",
-                "256",
-            ],
-        )
-
-    @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(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-
 class TestTP(CustomTestCase):
     @classmethod
     def setUpClass(cls):
@@ -153,55 +108,6 @@ class TestTP(CustomTestCase):
         self.assertGreater(metrics["score"], 0.60)
 
 
-@unittest.skip("covered in test_deepep_large.py")
-class TestNoGatherdBuffer(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--trust-remote-code",
-                "--tp",
-                "4",
-                "--enable-dp-attention",
-                "--dp",
-                "4",
-                "--moe-dense-tp-size",
-                "1",
-                "--enable-dp-lm-head",
-                "--moe-a2a-backend",
-                "deepep",
-                "--cuda-graph-max-bs-decode",
-                "32",
-                "--max-running-requests",
-                "512",
-            ],
-        )
-
-    @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(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-
 class TestTBO(CustomTestCase):
     @classmethod
     def setUpClass(cls):
@@ -254,176 +160,6 @@ class TestTBO(CustomTestCase):
         self.assertGreater(metrics["score"], 0.60)
 
 
-class TestTBOWithTPAttn(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--trust-remote-code",
-                "--tp",
-                "4",
-                "--moe-a2a-backend",
-                "deepep",
-                "--enable-two-batch-overlap",
-                "--cuda-graph-max-bs-decode",
-                "128",
-                "--max-running-requests",
-                "512",
-                "--mem-fraction-static",  # temp fix as DeepEP buffer is too large.
-                "0.7",
-            ],
-            env={
-                **os.environ,
-                "SGLANG_TBO_DEBUG": "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(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-
-# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case.
-# @unittest.skip("covered in TestMTPWithTPAttnAndTBO")
-class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--trust-remote-code",
-                "--tp",
-                "4",
-                "--moe-dense-tp-size",
-                "1",
-                "--moe-a2a-backend",
-                "deepep",
-                "--enable-two-batch-overlap",
-                "--cuda-graph-max-bs-decode",
-                "128",
-                "--max-running-requests",
-                "512",
-                "--mem-fraction-static",  # temp fix as DeepEP buffer is too large.
-                "0.7",
-            ],
-            env={
-                **os.environ,
-                "SGLANG_TBO_DEBUG": "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(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-
-@unittest.skip("covered in TestMTPWithTBO")
-class TestMTP(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--trust-remote-code",
-                "--tp",
-                "4",
-                "--enable-dp-attention",
-                "--dp",
-                "2",
-                "--enable-dp-lm-head",
-                "--moe-a2a-backend",
-                "deepep",
-                "--speculative-algo",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
-                "--speculative-num-steps",
-                "2",
-                "--speculative-eagle-topk",
-                "3",
-                "--speculative-num-draft-tokens",
-                "3",
-                "--cuda-graph-max-bs-decode",
-                "32",
-                "--max-running-requests",
-                "64",
-            ],
-        )
-
-    @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(metrics)
-
-        self.assertGreater(metrics["score"], 0.60)
-
-        server_info = requests.get(self.base_url + "/server_info")
-        avg_spec_accept_length = server_info.json()["internal_states"][0][
-            "avg_spec_accept_length"
-        ]
-        print(
-            f"###test_gsm8k (deepseek-v3 mtp + dp + tbo):\n"
-            f"accuracy={metrics['score']=:.3f}\n"
-            f"{avg_spec_accept_length=:.3f}\n"
-        )
-        self.assertGreater(avg_spec_accept_length, 2.1)
-
-
 class TestMTPWithTBO(CustomTestCase):
     @classmethod
     def setUpClass(cls):

test/registered/ep/test_deepep_small_extra.py

@@ -0,0 +1,167 @@
+import os
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+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,
+)
+
+register_cuda_ci(est_time=210, stage="extra-b", runner_config="4-gpu-h100")
+
+
+class TestHybridDPTP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--enable-dp-attention",
+                "--dp",
+                "2",
+                "--moe-a2a-backend",
+                "deepep",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "256",
+            ],
+        )
+
+    @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(metrics)
+
+        self.assertGreater(metrics["score"], 0.60)
+
+
+class TestTBOWithTPAttn(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--moe-a2a-backend",
+                "deepep",
+                "--enable-two-batch-overlap",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "512",
+                "--mem-fraction-static",  # temp fix as DeepEP buffer is too large.
+                "0.7",
+            ],
+            env={
+                **os.environ,
+                "SGLANG_TBO_DEBUG": "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(metrics)
+
+        self.assertGreater(metrics["score"], 0.60)
+
+
+# There exists bug when using MTP + TBO + attn_tp_size > 1, currently skip that case.
+# @unittest.skip("covered in TestMTPWithTPAttnAndTBO")
+class TestTBOWithTPAttnAndDenseDP(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "4",
+                "--moe-dense-tp-size",
+                "1",
+                "--moe-a2a-backend",
+                "deepep",
+                "--enable-two-batch-overlap",
+                "--cuda-graph-max-bs-decode",
+                "128",
+                "--max-running-requests",
+                "512",
+                "--mem-fraction-static",  # temp fix as DeepEP buffer is too large.
+                "0.7",
+            ],
+            env={
+                **os.environ,
+                "SGLANG_TBO_DEBUG": "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(metrics)
+
+        self.assertGreater(metrics["score"], 0.60)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/hicache/test_hicache_storage_3fs_backend.py

@@ -13,7 +13,7 @@ from test_hicache_storage_file_backend import HiCacheStorageBaseMixin
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=300, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=300, stage="base-b", runner_config="2-gpu-large")
 register_amd_ci(est_time=300, suite="base-b-test-2-gpu-large")
 
 

test/registered/models_e2e/test_qwen3_next_models.py

@@ -1,4 +1,3 @@
-import os
 import unittest
 
 from sglang.test.ci.ci_register import register_cuda_ci
@@ -7,7 +6,7 @@ from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
 from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 
-register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=260, stage="base-c", runner_config="4-gpu-h100")
 
 QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 
@@ -54,45 +53,5 @@ class TestQwen3NextLazyExtraBufferLargePage(
     other_args = _make_args(page_size=2, track_interval=2)
 
 
-class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase):
-    model = QWEN3_NEXT_MODEL
-    cache_chunk_size = 64
-    kl_div_thres = 0.002
-    other_args = _make_args(page_size=1, track_interval=2)
-
-    @classmethod
-    def setUpClass(cls):
-        os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
-        os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
-        super().setUpClass()
-
-    @classmethod
-    def tearDownClass(cls):
-        super().tearDownClass()
-        os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
-        os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
-
-
-class TestQwen3NextLazyExtraBufferLargePageAllocFail(
-    KLDivergenceMixin, DefaultServerBase
-):
-    model = QWEN3_NEXT_MODEL
-    cache_chunk_size = 64
-    kl_div_thres = 0.002
-    other_args = _make_args(page_size=2, track_interval=2)
-
-    @classmethod
-    def setUpClass(cls):
-        os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
-        os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
-        super().setUpClass()
-
-    @classmethod
-    def tearDownClass(cls):
-        super().tearDownClass()
-        os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
-        os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
-
-
 if __name__ == "__main__":
     unittest.main()

test/registered/models_e2e/test_qwen3_next_models_extra.py

@@ -0,0 +1,76 @@
+import os
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-h100")
+
+QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
+
+_COMMON_ARGS = [
+    "--trust-remote-code",
+    "--tp-size",
+    "4",
+    "--chunked-prefill-size",
+    "2048",
+    "--mamba-scheduler-strategy",
+    "extra_buffer_lazy",
+    "--attention-backend",
+    "triton",
+]
+
+
+def _make_args(*, page_size=1, track_interval=2):
+    return [
+        *_COMMON_ARGS,
+        "--mamba-track-interval",
+        str(track_interval),
+        "--page-size",
+        str(page_size),
+    ]
+
+
+class TestQwen3NextLazyExtraBufferAllocFail(KLDivergenceMixin, DefaultServerBase):
+    model = QWEN3_NEXT_MODEL
+    cache_chunk_size = 64
+    kl_div_thres = 0.002
+    other_args = _make_args(page_size=1, track_interval=2)
+
+    @classmethod
+    def setUpClass(cls):
+        os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
+        os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
+        super().setUpClass()
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
+        os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
+
+
+class TestQwen3NextLazyExtraBufferLargePageAllocFail(
+    KLDivergenceMixin, DefaultServerBase
+):
+    model = QWEN3_NEXT_MODEL
+    cache_chunk_size = 64
+    kl_div_thres = 0.002
+    other_args = _make_args(page_size=2, track_interval=2)
+
+    @classmethod
+    def setUpClass(cls):
+        os.environ["SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL"] = "1"
+        os.environ["SGLANG_TEST_SKIP_CACHE_HIT_ASSERT"] = "1"
+        super().setUpClass()
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        os.environ.pop("SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL", None)
+        os.environ.pop("SGLANG_TEST_SKIP_CACHE_HIT_ASSERT", None)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/pp/test_pp_gemma4.py

@@ -0,0 +1,153 @@
+import time
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP,
+    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    is_in_amd_ci,
+    is_in_ci,
+    popen_launch_server,
+)
+
+# tp=1 pp=2 -- two GPUs.
+register_cuda_ci(est_time=220, stage="base-b", runner_config="2-gpu-large")
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "Gemma4 PP not yet validated on AMD",
+)
+class TestGemma4PPAccuracy(unittest.TestCase):
+    """End-to-end PP=2 accuracy gate for Gemma4 multimodal.
+
+    Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so
+    sglang auto-selects the triton attention backend; no manual flag needed.
+    The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an
+    H100's 80 GB.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP
+        cls.base_url = "http://127.0.0.1:23333"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                1,
+                "--pp-size",
+                2,
+                "--trust-remote-code",
+                "--enable-multimodal",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        # Gemma4 is instruction-tuned and doesn't follow few-shot completion
+        # prompts well — use the chat API (default in run_eval), which scores
+        # ~0.98 on this model vs ~0.44 with api="completion".
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=200,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        # Chat-API baseline ~0.98; gate well below to absorb sample-noise
+        # without missing a real PP-routing regression (pre-PP-fix the model
+        # produced garbage outputs scoring ≈ 0).
+        self.assertGreaterEqual(metrics["score"], 0.90)
+        # Wait a little bit so that the memory check happens.
+        time.sleep(4)
+
+    @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
+    def test_mmmu(self):
+        # Multimodal accuracy gate covering the vision_tower → embed_vision
+        # (first rank) → PP-proxy handoff → LM tail (last rank) chain.
+        # Measured 0.71 on 200 examples; full eval (~900 questions) takes
+        # ~5-7 min on H100 so this is manual-only.
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmmu",
+            num_examples=None,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in
+        # 0.69-0.73.  Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary
+        # samples ≈ 0.015) while still catching mid-grade vision/PP
+        # regressions, not just complete breakage.
+        self.assertGreater(metrics["score"], 0.65)
+
+
+@unittest.skipIf(
+    is_in_amd_ci(),
+    "Gemma4 PP not yet validated on AMD",
+)
+class TestGemma4PLEPPAccuracy(unittest.TestCase):
+    """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path).
+
+    26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP
+    test never crosses the PLE branch.  Cuda graph + PLE corrupts outputs
+    (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema
+    drops ``per_layer_inputs``), so this test pins the eager configuration.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP
+        cls.base_url = "http://127.0.0.1:23339"
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--tp-size",
+                1,
+                "--pp-size",
+                2,
+                "--trust-remote-code",
+                "--enable-multimodal",
+                # Required for PLE under PP — see Gemma4TextModel guard.
+                "--disable-cuda-graph",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage
+        # (corruption collapses score to ~0).
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            num_examples=100,
+            num_threads=32,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+        self.assertGreaterEqual(metrics["score"], 0.80)
+        time.sleep(4)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/pp/test_pp_single_node.py

@@ -23,18 +23,15 @@ from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_MLA_MODEL_NAME_FOR_TEST,
     DEFAULT_MODEL_NAME_FOR_TEST,
-    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP,
-    DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
     is_in_amd_ci,
-    is_in_ci,
     popen_launch_server,
     run_bench_one_batch_server,
 )
 
-register_cuda_ci(est_time=500, stage="base-c", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=280, stage="base-c", runner_config="4-gpu-h100")
 register_amd_ci(est_time=500, suite="stage-c-test-4-gpu-amd")
 
 
@@ -145,137 +142,6 @@ class TestDPAttentionDP2PP2(CustomTestCase):
         self.assertGreater(metrics["score"], 0.8)
 
 
-@unittest.skipIf(
-    is_in_amd_ci(),
-    "Gemma4 PP not yet validated on AMD",
-)
-class TestGemma4PPAccuracy(unittest.TestCase):
-    """End-to-end PP=2 accuracy gate for Gemma4 multimodal.
-
-    Gemma4 has full-attention layers with head_dim=512 (FA's max is 256), so
-    sglang auto-selects the triton attention backend; no manual flag needed.
-    The 26B BF16 model splits to ~26 GB per stage under PP=2, well within an
-    H100's 80 GB.
-    """
-
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PP
-        cls.base_url = "http://127.0.0.1:23333"
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--tp-size",
-                1,
-                "--pp-size",
-                2,
-                "--trust-remote-code",
-                "--enable-multimodal",
-            ],
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        # Gemma4 is instruction-tuned and doesn't follow few-shot completion
-        # prompts well — use the chat API (default in run_eval), which scores
-        # ~0.98 on this model vs ~0.44 with api="completion".
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            num_examples=200,
-            num_threads=32,
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-
-        # Chat-API baseline ~0.98; gate well below to absorb sample-noise
-        # without missing a real PP-routing regression (pre-PP-fix the model
-        # produced garbage outputs scoring ≈ 0).
-        self.assertGreaterEqual(metrics["score"], 0.90)
-        # Wait a little bit so that the memory check happens.
-        time.sleep(4)
-
-    @unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
-    def test_mmmu(self):
-        # Multimodal accuracy gate covering the vision_tower → embed_vision
-        # (first rank) → PP-proxy handoff → LM tail (last rank) chain.
-        # Measured 0.71 on 200 examples; full eval (~900 questions) takes
-        # ~5-7 min on H100 so this is manual-only.
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="mmmu",
-            num_examples=None,
-            num_threads=32,
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-        # Measured 0.72 on this setup; published Gemma-4-26B MMMU lies in
-        # 0.69-0.73.  Gate 0.65 leaves ~5 SE of headroom (SE on 900 binary
-        # samples ≈ 0.015) while still catching mid-grade vision/PP
-        # regressions, not just complete breakage.
-        self.assertGreater(metrics["score"], 0.65)
-
-
-@unittest.skipIf(
-    is_in_amd_ci(),
-    "Gemma4 PP not yet validated on AMD",
-)
-class TestGemma4PLEPPAccuracy(unittest.TestCase):
-    """PP=2 coverage for Gemma4 PLE variants (per_layer_inputs proxy path).
-
-    26B-A4B has ``hidden_size_per_layer_input=0`` so the default Gemma4 PP
-    test never crosses the PLE branch.  Cuda graph + PLE corrupts outputs
-    (the runner's hardcoded ``{hidden_states, residual}`` PP-proxy schema
-    drops ``per_layer_inputs``), so this test pins the eager configuration.
-    """
-
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST_GEMMA4_PLE_PP
-        cls.base_url = "http://127.0.0.1:23339"
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--tp-size",
-                1,
-                "--pp-size",
-                2,
-                "--trust-remote-code",
-                "--enable-multimodal",
-                # Required for PLE under PP — see Gemma4TextModel guard.
-                "--disable-cuda-graph",
-            ],
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        # Eager-path baseline ~0.92; gate 0.80 catches PLE breakage
-        # (corruption collapses score to ~0).
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            num_examples=100,
-            num_threads=32,
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-        self.assertGreaterEqual(metrics["score"], 0.80)
-        time.sleep(4)
-
-
 class TestPPMixedChunk(CustomTestCase):
     @classmethod
     def setUpClass(cls):

test/registered/rl/test_release_memory_occupation.py

@@ -1,477 +0,0 @@
-"""Test memory release and resume operations for SGLang engine in hybrid RL training.
-
-This test suite evaluates the SGLang engine's memory management capabilities, focusing
-on releasing and resuming memory occupation for KV cache and model weights. It simulates
-an RL workflow where the SGLang engine acts as a rollout engine for experience collection.
-The process involves initializing the engine, sending a small number of requests to simulate
-rollout, releasing memory to mimic offloading during RL training, resuming memory occupation,
-updating weights with a trained HuggingFace model, and verifying the updated weights.
-
-Detailed in our proposal (https://github.com/sgl-project/sglang/pull/7099), two test cases
-are included:
-
-1. Basic Release and Resume: Uses a lower mem_fraction_static (0.6) to control memory allocation
-and avoid OOM errors carefully. This test simulates a scenario without multi-stage memory management,
-ensuring the engine can release and resume memory occupation while maintaining functionality after
-weight updates.
-
-2. Multi-Stage Release and Resume: Employs a higher mem_fraction_static (0.85) to simulate higher
-memory pressure, leveraging multi-stage memory management. It sequentially releases and resumes
-KV cache and model weights, verifying memory deallocation and reallocation at each stage, and
-ensuring correct weight updates and text generation.
-
-3. Tensor Parallel Tests: Tests memory release and resume operations with different tensor parallel
-configurations (tp=1, tp=2) to ensure proper memory management in distributed settings. For different
-data parallel size, we test it in verl.
-
-NOTE: This test is temporarily disabled.
-"""
-
-import os
-import time
-import unittest
-
-from transformers import AutoModelForCausalLM
-
-import sglang as sgl
-from sglang.srt.constants import (
-    GPU_MEMORY_TYPE_CUDA_GRAPH,
-    GPU_MEMORY_TYPE_KV_CACHE,
-    GPU_MEMORY_TYPE_WEIGHTS,
-)
-from sglang.srt.utils import get_device
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import (
-    DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
-    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
-    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
-    DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
-    DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
-    CustomTestCase,
-    empty_gpu_cache,
-    get_gpu_count,
-    get_gpu_memory_gb,
-)
-
-register_cuda_ci(
-    est_time=200,
-    stage="base-c",
-    runner_config="4-gpu-h100",
-    disabled="Temporarily disabled - needs investigation",
-)
-
-# (temporarily) set to true to observe memory usage in nvidia-smi more clearly
-_DEBUG_EXTRA = False
-
-
-class TestReleaseMemoryOccupation(CustomTestCase):
-    def _setup_engine(
-        self,
-        model_name,
-        mem_fraction_static=0.8,
-        tp_size=1,
-        ep_size=1,
-        enable_weights_cpu_backup=False,
-    ):
-        """Common setup for engine and HF model."""
-
-        os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
-        engine = sgl.Engine(
-            model_path=model_name,
-            random_seed=42,
-            enable_memory_saver=True,
-            mem_fraction_static=mem_fraction_static,
-            tp_size=tp_size,
-            ep_size=ep_size,
-            enable_weights_cpu_backup=enable_weights_cpu_backup,
-            # disable_cuda_graph=True,  # for debugging only
-        )
-
-        return engine
-
-    def _common_test_params(self):
-        """Common test parameters."""
-        return {
-            "prompt": "Today is a sunny day and I like",
-            "sampling_params": {"temperature": 0, "max_new_tokens": 8},
-            "expect_output_before_update_weights": " to spend it outdoors. I decided to",
-            "expect_output_after_update_weights": " to go for a walk. I like",
-            "prompt_moe": "The weather is nice today, and I want to",
-            "sampling_params_moe": {"temperature": 0, "max_new_tokens": 16},
-            "expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a",
-            "expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I",
-            "prompt_hybrid_mamba": "The weather is nice today, and I want to",
-            "sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16},
-            "expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
-            "expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
-        }
-
-    def _test_initial_generation(
-        self, engine, prompt, sampling_params, expect_output_before_update_weights
-    ):
-        """Test initial generation and memory allocation."""
-        print("generate (#1)")
-        outputs = engine.generate(prompt, sampling_params)["text"]
-        self.assertEqual(outputs, expect_output_before_update_weights)
-
-        if _DEBUG_EXTRA:
-            time.sleep(3)
-
-    def test_release_and_resume_occupation(self):
-        # Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM
-        model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
-        assert get_gpu_count() >= 2, "Need at least 2 GPUs for tensor parallel tests"
-
-        for tp_size in [1, 2]:
-
-            print(f"Testing tp_size={tp_size} for test_release_and_resume_occupation")
-            engine = self._setup_engine(
-                model_name=model_name, mem_fraction_static=0.6, tp_size=tp_size
-            )
-            params = self._common_test_params()
-
-            self._test_initial_generation(
-                engine,
-                params["prompt"],
-                params["sampling_params"],
-                params["expect_output_before_update_weights"],
-            )
-
-            t = time.perf_counter()
-            gpu_memory_usage_before_release = get_gpu_memory_gb()
-            engine.release_memory_occupation()
-            gpu_memory_usage_after_release = get_gpu_memory_gb()
-
-            self.assertLess(
-                gpu_memory_usage_after_release,
-                gpu_memory_usage_before_release,
-            )
-
-            print(
-                f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
-            )
-
-            if _DEBUG_EXTRA:
-                time.sleep(3)
-
-            t = time.perf_counter()
-            engine.resume_memory_occupation()
-            print(
-                f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
-            )
-
-            hf_model_new = AutoModelForCausalLM.from_pretrained(
-                DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
-                torch_dtype="bfloat16",
-                device_map=get_device(),
-            )
-            engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
-
-            # destroy the hf model
-            del hf_model_new
-            empty_gpu_cache()
-
-            print("generate (#2)")
-            outputs = engine.generate(params["prompt"], params["sampling_params"])[
-                "text"
-            ]
-            self.assertEqual(outputs, params["expect_output_after_update_weights"])
-            engine.shutdown()
-
-    def test_release_and_resume_occupation_with_weights_cpu_backup(self):
-        # Test release and resume occupation with weights CPU backup
-        model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
-
-        print("Testing test_release_and_resume_occupation_with_weights_cpu_backup")
-        engine = self._setup_engine(
-            model_name=model_name,
-            mem_fraction_static=0.6,
-            enable_weights_cpu_backup=True,
-        )
-        params = self._common_test_params()
-
-        self._test_initial_generation(
-            engine,
-            params["prompt"],
-            params["sampling_params"],
-            params["expect_output_before_update_weights"],
-        )
-
-        t = time.perf_counter()
-        gpu_memory_usage_before_release = get_gpu_memory_gb()
-        engine.release_memory_occupation()
-        gpu_memory_usage_after_release = get_gpu_memory_gb()
-
-        self.assertLess(
-            gpu_memory_usage_after_release,
-            gpu_memory_usage_before_release,
-        )
-
-        print(
-            f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
-        )
-
-        if _DEBUG_EXTRA:
-            time.sleep(3)
-
-        t = time.perf_counter()
-        engine.resume_memory_occupation()
-        print(
-            f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
-        )
-
-        print("generate post resume")
-        outputs = engine.generate(params["prompt"], params["sampling_params"])["text"]
-        self.assertEqual(outputs, params["expect_output_before_update_weights"])
-        engine.shutdown()
-
-    def test_multi_stage_release_and_resume(self):
-        # With multi-stage release and resume, we can set the memory fraction to 0.85 without concern of OOM
-        model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
-
-        for tp_size in [1, 2]:
-            if tp_size == 2 and get_gpu_count() < 2:
-                continue
-
-            print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume")
-            os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
-            engine = sgl.Engine(
-                model_path=model_name,
-                random_seed=42,
-                enable_memory_saver=True,
-                mem_fraction_static=0.85,  # Higher memory pressure
-                tp_size=tp_size,
-            )
-            params = self._common_test_params()
-
-            self._test_initial_generation(
-                engine,
-                params["prompt"],
-                params["sampling_params"],
-                params["expect_output_before_update_weights"],
-            )
-
-            t = time.perf_counter()
-            gpu_memory_usage_before_release = get_gpu_memory_gb()
-            engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
-
-            gpu_memory_usage_after_release_kv_cache = get_gpu_memory_gb()
-
-            self.assertLess(
-                gpu_memory_usage_after_release_kv_cache,
-                gpu_memory_usage_before_release,
-            )
-
-            engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
-            gpu_memory_usage_after_release_weights = get_gpu_memory_gb()
-
-            self.assertLess(
-                gpu_memory_usage_after_release_weights,
-                gpu_memory_usage_after_release_kv_cache,
-            )
-
-            engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
-            gpu_memory_usage_after_release_cuda_graph = get_gpu_memory_gb()
-
-            self.assertLess(
-                gpu_memory_usage_after_release_cuda_graph,
-                gpu_memory_usage_after_release_weights,
-            )
-
-            print(f"Release took {time.perf_counter() - t:.2f}s")
-            print(
-                f"Memory: {gpu_memory_usage_before_release:.1f} → {gpu_memory_usage_after_release_kv_cache:.1f} → {gpu_memory_usage_after_release_weights:.1f} → {gpu_memory_usage_after_release_cuda_graph:.1f} GB"
-            )
-
-            if _DEBUG_EXTRA:
-                time.sleep(3)
-
-            t = time.perf_counter()
-            gpu_memory_usage_before_resume = get_gpu_memory_gb()
-
-            # gpu_memory_usage_after_release_weights and gpu_memory_usage_before_resume should be close
-
-            self.assertAlmostEqual(
-                gpu_memory_usage_after_release_weights,
-                gpu_memory_usage_before_resume,
-                delta=3.0,
-            )
-            print(f"Resume weights took {time.perf_counter() - t:.2f}s")
-
-            engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
-            gpu_memory_usage_after_resume_cuda_graph = get_gpu_memory_gb()
-
-            self.assertGreater(
-                gpu_memory_usage_after_resume_cuda_graph,
-                gpu_memory_usage_before_resume,
-            )
-
-            engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
-            gpu_memory_usage_after_resume_weights = get_gpu_memory_gb()
-
-            self.assertGreater(
-                gpu_memory_usage_after_resume_weights,
-                gpu_memory_usage_after_resume_cuda_graph,
-            )
-
-            # Update weights from a trained model to serving engine, and then destroy the trained model
-            hf_model_new = AutoModelForCausalLM.from_pretrained(
-                DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
-                torch_dtype="bfloat16",
-                device_map=get_device(),
-            )
-            gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb()
-            engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
-
-            # destroy the hf model
-            del hf_model_new
-            empty_gpu_cache()
-            engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
-
-            gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb()
-            self.assertGreater(
-                gpu_memory_usage_after_resume_kv_cache,
-                gpu_memory_usage_after_resume_weights,
-            )
-
-            print(f"Resume + update took {time.perf_counter() - t:.2f}s")
-            print(
-                f"Memory: {gpu_memory_usage_before_resume:.1f} → {gpu_memory_usage_after_resume_cuda_graph:.1f} → {gpu_memory_usage_after_resume_weights:.1f} → {gpu_memory_usage_after_loaded_hf_model:.1f} → {gpu_memory_usage_after_resume_kv_cache:.1f} GB"
-            )
-
-            print("generate (#2)")
-            outputs = engine.generate(params["prompt"], params["sampling_params"])[
-                "text"
-            ]
-            self.assertEqual(outputs, params["expect_output_after_update_weights"])
-            engine.shutdown()
-
-    def test_moe_model_release_and_resume(self):
-        # Test with MoE model
-        model_name = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
-
-        tp_size = ep_size = 2
-
-        print(
-            f"Testing tp_size={tp_size} and ep_size={ep_size} for test_moe_model_release_and_resume"
-        )
-        engine = sgl.Engine(
-            model_path=model_name,
-            random_seed=42,
-            enable_memory_saver=True,
-            mem_fraction_static=0.5,
-            tp_size=tp_size,
-            ep_size=ep_size,
-        )
-        params = self._common_test_params()
-
-        self._test_initial_generation(
-            engine,
-            params["prompt_moe"],
-            params["sampling_params_moe"],
-            params["expect_output_before_update_weights_moe"],
-        )
-
-        t = time.perf_counter()
-        gpu_memory_usage_before_release = get_gpu_memory_gb()
-        engine.release_memory_occupation()
-        gpu_memory_usage_after_release = get_gpu_memory_gb()
-        self.assertLess(
-            gpu_memory_usage_after_release,
-            gpu_memory_usage_before_release,
-        )
-
-        print(
-            f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
-        )
-
-        if _DEBUG_EXTRA:
-            time.sleep(3)
-
-        t = time.perf_counter()
-        engine.resume_memory_occupation()
-        print(
-            f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
-        )
-
-        hf_model_new = AutoModelForCausalLM.from_pretrained(
-            DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
-            torch_dtype="bfloat16",
-            device_map=get_device(),
-        )
-        engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
-
-        # destroy the hf model
-        del hf_model_new
-        empty_gpu_cache()
-
-        print("generate (#2)")
-        outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[
-            "text"
-        ]
-        self.assertEqual(outputs, params["expect_output_after_update_weights_moe"])
-        engine.shutdown()
-
-    def test_hybrid_mamba_model_release_and_resume(self):
-        # Test with Hybrid Mamba model
-        model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
-
-        tp_size = 4
-
-        print(
-            f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume"
-        )
-        engine = sgl.Engine(
-            model_path=model_name,
-            random_seed=42,
-            enable_memory_saver=True,
-            tp_size=tp_size,
-        )
-        params = self._common_test_params()
-
-        self._test_initial_generation(
-            engine,
-            params["prompt_hybrid_mamba"],
-            params["sampling_params_hybrid_mamba"],
-            params["expect_output_before_update_weights_hybrid_mamba"],
-        )
-
-        t = time.perf_counter()
-        gpu_memory_usage_before_release = get_gpu_memory_gb()
-        engine.release_memory_occupation()
-        gpu_memory_usage_after_release = get_gpu_memory_gb()
-        self.assertLess(
-            gpu_memory_usage_after_release,
-            gpu_memory_usage_before_release,
-        )
-
-        print(
-            f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
-        )
-
-        if _DEBUG_EXTRA:
-            time.sleep(3)
-
-        t = time.perf_counter()
-        engine.resume_memory_occupation()
-        print(
-            f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
-        )
-
-        engine.update_weights_from_disk(model_name)
-
-        # destroy the hf model
-        empty_gpu_cache()
-
-        print("generate (#2)")
-        outputs = engine.generate(
-            params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"]
-        )["text"]
-        self.assertEqual(
-            outputs, params["expect_output_after_update_weights_hybrid_mamba"]
-        )
-        engine.shutdown()
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/unit/layers/test_flashinfer_comm_fusion.py

@@ -10,9 +10,9 @@ from sglang.srt.runtime_context import get_parallel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100")
+# Collectives are mocked and world_size is a plain int, so the world_size=4
-register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-b200")
+# cases need one real CUDA device.
-register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-gb300")
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
 
 
 class _FakeWorkspace: