[CI] Reorganize stage-b 1-GPU tests for 5090 compatibility (#16826)

test/registered/moe/test_cutedsl_moe.py

@@ -1,8 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=300, suite="stage-c-test-large-4-gpu-b200")
-
 import unittest
 from typing import Callable
 
@@ -14,6 +10,9 @@ from torch.nn import functional as F
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.flashinfer_cutedsl_moe import flashinfer_cutedsl_moe_masked
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=300, suite="stage-c-test-large-4-gpu-b200")
 
 SKIP_TEST = torch.cuda.get_device_capability() < (10, 0)
 SKIP_REASON = "Nvfp4 Requires compute capability of 10 or above."

test/registered/moe/test_fused_moe.py

@@ -1,8 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=80, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=30, suite="stage-b-test-small-1-gpu-amd")
-
 import unittest
 
 import torch
@@ -15,8 +10,12 @@ from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import is_hip
+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=80, suite="stage-b-test-large-1-gpu")
+register_amd_ci(est_time=30, suite="stage-b-test-small-1-gpu-amd")
+
 _is_hip = is_hip()
 _is_fp8_fnuz = is_fp8_fnuz()
 

test/registered/moe/test_glm4_moe_models.py

@@ -1,11 +1,8 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=100, suite="stage-b-test-large-2-gpu")
-
 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.few_shot_gsm8k import run_eval
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -14,6 +11,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=100, suite="stage-b-test-large-2-gpu")
+
 
 class TestGLM4MoE(CustomTestCase):
     @classmethod

test/registered/moe/test_moe_ep.py

@@ -1,11 +1,8 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=140, suite="stage-b-test-large-2-gpu")
-
 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_MLA_MODEL_NAME_FOR_TEST,
@@ -15,6 +12,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=140, suite="stage-b-test-large-2-gpu")
+
 
 class TestEp(CustomTestCase):
     @classmethod

test/registered/moe/test_torch_compile_moe.py

@@ -1,8 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=210, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=1400, suite="stage-b-test-small-1-gpu-amd")
-
 import time
 import unittest
 from types import SimpleNamespace
@@ -10,6 +5,7 @@ from types import SimpleNamespace
 import requests
 
 from sglang.srt.utils import is_cuda, kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
@@ -19,6 +15,9 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=210, suite="stage-b-test-large-1-gpu")
+register_amd_ci(est_time=1400, suite="stage-b-test-small-1-gpu-amd")
+
 
 class TestTorchCompileMoe(CustomTestCase):
     @classmethod

test/registered/moe/test_triton_fused_moe.py

@@ -1,7 +1,3 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=89, suite="stage-b-test-small-1-gpu")
-
 import unittest
 
 import torch
@@ -13,8 +9,11 @@ from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsQuantIn
 from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
 from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
+register_cuda_ci(est_time=89, suite="stage-b-test-large-1-gpu")
+
 
 class TestFusedMOE(CustomTestCase):
     NUM_EXPERTS = [8, 64]

test/registered/moe/test_triton_moe_channel_fp8_kernel.py

@@ -1,7 +1,3 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=16, suite="stage-b-test-small-1-gpu")
-
 import itertools
 import unittest
 
@@ -12,8 +8,11 @@ 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.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
+register_cuda_ci(est_time=16, suite="stage-b-test-large-1-gpu")
+
 
 def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
     """Matrix multiplication function that supports per-token input quantization and per-column weight quantization"""