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

test/registered/rl/test_fp32_lm_head.py

@@ -1,9 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=15, suite="stage-b-test-small-1-gpu-amd")
-
 import unittest
 from types import SimpleNamespace
 from unittest.mock import patch
@@ -18,6 +12,10 @@ from sglang.srt.server_args import (
     get_global_server_args,
     set_global_server_args_for_scheduler,
 )
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=15, suite="stage-b-test-small-1-gpu-amd")
 
 
 class LMHeadStub(nn.Module):

test/registered/rl/test_lora_load_from_tensor.py

@@ -1,9 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=90, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=54, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=90, suite="stage-b-test-small-1-gpu")
-
 import json
 import os
 import unittest
@@ -13,8 +7,12 @@ from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
 
 import sglang as sgl
+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=90, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=90, suite="stage-b-test-small-1-gpu")
+
 MODEL_PATH = "Qwen/Qwen3-0.6B"
 LORA_REPO = "charent/self_cognition_Alice"
 TEST_PROMPT = "Hello, my name is"

test/registered/rl/test_patch_torch.py

@@ -1,10 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=19, suite="stage-b-test-large-2-gpu")
-register_amd_ci(
-    est_time=19, suite="stage-b-test-large-2-gpu-amd", disabled="see #11127"
-)
-
 import os
 import traceback
 import unittest
@@ -14,6 +7,12 @@ import torch
 import torch.multiprocessing as mp
 
 from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_amd_ci(
+    est_time=19, suite="stage-b-test-large-2-gpu-amd", disabled="see #11127"
+)
+register_cuda_ci(est_time=19, suite="stage-b-test-large-2-gpu")
 
 
 class TestReleaseMemoryOccupation(unittest.TestCase):

test/registered/rl/test_return_routed_experts.py

@@ -1,7 +1,3 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=180, suite="stage-c-test-large-4-gpu")
-
 import asyncio
 import logging
 import unittest
@@ -16,6 +12,7 @@ from sglang.srt.layers.moe.routed_experts_capturer import (
     extract_routed_experts_from_meta_info,
 )
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -24,6 +21,8 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_cuda_ci(est_time=180, suite="stage-c-test-large-4-gpu")
+
 SHAREGPT_URL = (
     "https://huggingface.co/datasets/anon8231489123/"
     "ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"

test/registered/rl/test_update_weights_from_disk.py

@@ -1,10 +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", disabled="see #14021")
-register_amd_ci(
-    est_time=210, suite="stage-b-test-small-1-gpu-amd", disabled="see #14021"
-)
-
 import json
 import random
 import time
@@ -15,6 +8,7 @@ import requests
 
 import sglang as sgl
 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import (
     DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -24,6 +18,11 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
+register_amd_ci(
+    est_time=210, suite="stage-b-test-small-1-gpu-amd", disabled="see #14021"
+)
+register_cuda_ci(est_time=210, suite="stage-b-test-large-1-gpu", disabled="see #14021")
+
 
 ###############################################################################
 # Engine Mode Tests (Single-configuration)

test/registered/rl/test_update_weights_from_distributed.py

@@ -1,8 +1,5 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
-register_cuda_ci(est_time=103, suite="stage-b-test-large-2-gpu")
-register_amd_ci(est_time=103, suite="stage-b-test-large-2-gpu-amd")
-
 """Test distributed weight updates.
 
 This test suite simulates a distributed training environment to ensure
@@ -45,6 +42,9 @@ from sglang.test.test_utils import (
 )
 from sglang.utils import terminate_process
 
+register_cuda_ci(est_time=103, suite="stage-b-test-large-2-gpu")
+register_amd_ci(est_time=103, suite="stage-b-test-large-2-gpu-amd")
+
 mp.set_start_method("spawn", force=True)
 
 

test/registered/rl/test_update_weights_from_tensor.py

@@ -2,7 +2,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=195, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=195, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=164, suite="stage-b-test-small-1-gpu-5090")
 
 import gc
 import json