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

2026-01-15 15:23:35 -08:00
parent 8b22deef5b
commit 146b5fcc84
136 changed files with 277 additions and 404 deletions
@@ -595,110 +595,13 @@ jobs:
          ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
        )
      )
-    runs-on: 1-gpu-runner
-    env:
-      RUNNER_LABELS: 1-gpu-runner
-    strategy:
-      fail-fast: false
-      max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel) }}
-      matrix:
-        partition: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
-
-      - name: Download artifacts
-        if: needs.check-changes.outputs.sgl_kernel == 'true'
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 10
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
-
-      - name: Run test
-        timeout-minutes: 30
-        run: |
-          cd test/
-          CONTINUE_ON_ERROR_FLAG=""
-          if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
-            CONTINUE_ON_ERROR_FLAG="--continue-on-error"
-          fi
-          python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 14 $CONTINUE_ON_ERROR_FLAG
-
-  stage-b-test-large-1-gpu:
-    needs: [check-changes, call-gate, stage-a-test-1, sgl-kernel-build-wheels]
-    if: |
-      always() &&
-      (
-        (inputs.target_stage == 'stage-b-test-large-1-gpu') ||
-        (
-          !inputs.target_stage &&
-          (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
-          ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
-        )
-      )
-    runs-on: 1-gpu-runner
-    env:
-      RUNNER_LABELS: 1-gpu-runner
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
-
-      - name: Download artifacts
-        if: needs.check-changes.outputs.sgl_kernel == 'true'
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 10
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
-
-      - name: Run test
-        timeout-minutes: 30
-        run: |
-          cd test/
-          CONTINUE_ON_ERROR_FLAG=""
-          if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
-            CONTINUE_ON_ERROR_FLAG="--continue-on-error"
-          fi
-          python3 run_suite.py --hw cuda --suite stage-b-test-large-1-gpu $CONTINUE_ON_ERROR_FLAG
-
-  # 5090 dry run stage - only runs on scheduled CI or when explicitly targeted via /rerun-stage
-  # Used to validate 5090 compatibility before full integration
-  stage-b-test-small-1-gpu-5090:
-    needs: [check-changes, call-gate, stage-a-test-1, sgl-kernel-build-wheels]
-    if: |
-      always() &&
-      (
-        (inputs.target_stage == 'stage-b-test-small-1-gpu-5090') ||
-        (
-          !inputs.target_stage &&
-          github.event_name == 'schedule' &&
-          (!failure() && !cancelled()) &&
-          ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
-        )
-      )
    runs-on: 1-gpu-5090
-    continue-on-error: true
    env:
      RUNNER_LABELS: 1-gpu-5090
      IS_BLACKWELL: "1"
    strategy:
      fail-fast: false
-      max-parallel: 8
+      max-parallel: 4
      matrix:
        partition: [0, 1, 2, 3, 4, 5, 6, 7]
    steps:
@@ -726,7 +629,61 @@ jobs:
        run: |
          source /etc/profile.d/sglang-ci.sh
          cd test/
-          python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu-5090 --auto-partition-id ${{ matrix.partition }} --auto-partition-size 8 --continue-on-error
+          CONTINUE_ON_ERROR_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
+            CONTINUE_ON_ERROR_FLAG="--continue-on-error"
+          fi
+          python3 run_suite.py --hw cuda --suite stage-b-test-small-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 8 $CONTINUE_ON_ERROR_FLAG
+
+  # Runs on H200 (80GB, SM90) - tests that don't pass on 5090 (FA3, FP8, high VRAM, etc.)
+  stage-b-test-large-1-gpu:
+    needs: [check-changes, call-gate, stage-a-test-1, sgl-kernel-build-wheels]
+    if: |
+      always() &&
+      (
+        (inputs.target_stage == 'stage-b-test-large-1-gpu') ||
+        (
+          !inputs.target_stage &&
+          (github.event_name == 'schedule' || (!failure() && !cancelled())) &&
+          ((needs.check-changes.outputs.main_package == 'true') || (needs.check-changes.outputs.sgl_kernel == 'true'))
+        )
+      )
+    runs-on: 1-gpu-runner
+    env:
+      RUNNER_LABELS: 1-gpu-runner
+    strategy:
+      fail-fast: false
+      max-parallel: ${{ fromJson(needs.check-changes.outputs.max_parallel) }}
+      matrix:
+        partition: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
+
+      - name: Download artifacts
+        if: needs.check-changes.outputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 10
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
+
+      - name: Run test
+        timeout-minutes: 30
+        run: |
+          cd test/
+          CONTINUE_ON_ERROR_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.continue_on_error }}" == "true" ]]; then
+            CONTINUE_ON_ERROR_FLAG="--continue-on-error"
+          fi
+          python3 run_suite.py --hw cuda --suite stage-b-test-large-1-gpu --auto-partition-id ${{ matrix.partition }} --auto-partition-size 10 $CONTINUE_ON_ERROR_FLAG

  stage-b-test-large-2-gpu:
    needs: [check-changes, call-gate, stage-a-test-1, sgl-kernel-build-wheels]
@@ -227,7 +227,6 @@ def handle_rerun_stage(
        "stage-a-cpu-only",
        "stage-b-test-small-1-gpu",
        "stage-b-test-large-1-gpu",
-        "stage-b-test-small-1-gpu-5090",
        "stage-b-test-large-2-gpu",
        "stage-c-test-large-4-gpu",
        "stage-c-test-large-4-gpu-b200",
@@ -9,7 +9,6 @@ from sglang.test.test_utils import CustomTestCase

 # Triton kernel unit test for KV indices creation
 register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")


@@ -16,7 +16,6 @@ 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=9, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")


 class TestMamba(unittest.TestCase):
@@ -15,7 +15,6 @@ from sglang.test.test_utils import (

 # RadixAttention server integration tests
 register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=79, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")


@@ -22,7 +22,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 # CPU-based unit test, runs quickly on any GPU runner
 register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=5, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")

 import time
 import unittest
@@ -9,7 +9,7 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllo
 from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=8, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")


@@ -19,7 +19,6 @@ from sglang.test.test_utils import (

 # Torch native attention backend integration test with MMLU eval
 register_cuda_ci(est_time=169, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=120, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=150, suite="stage-b-test-small-1-gpu-amd")


@@ -20,7 +20,7 @@ from sglang.test.test_utils import (
 )

 # Triton attention backend integration test with latency benchmark and MMLU eval
-register_cuda_ci(est_time=200, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=200, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=1110, suite="stage-b-test-small-1-gpu-amd")


@@ -23,7 +23,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase, is_in_amd_ci

 # Triton attention kernel unit tests (decode, extend, prefill)
-register_cuda_ci(est_time=30, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=30, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=30, suite="stage-b-test-small-1-gpu-amd")


@@ -15,7 +15,7 @@ from sglang.test.test_utils import (
 )

 # Sliding window attention with Triton backend (Gemma-3 model)
-register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")


@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=144, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=144, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=1100, suite="stage-b-test-small-1-gpu-amd")


@@ -1,12 +1,7 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=111, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=76, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=179, suite="stage-b-test-small-1-gpu-amd")
-
 import unittest

 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kits.ebnf_constrained_kit import TestEBNFConstrainedMixin
 from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin
 from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin
@@ -18,6 +13,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=111, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=179, suite="stage-b-test-small-1-gpu-amd")
+

 class ServerWithGrammar(CustomTestCase):
    backend = "xgrammar"
@@ -15,7 +15,7 @@ from sglang.test.test_deterministic_utils import (
    TestDeterministicBase,
 )

-register_cuda_ci(est_time=278, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=278, suite="stage-b-test-large-1-gpu")


 class TestFlashinferDeterministic(TestDeterministicBase):
@@ -3,7 +3,7 @@ import unittest
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.gpt_oss_common import BaseTestGptOss

-register_cuda_ci(est_time=519, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=519, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=750, suite="stage-b-test-small-1-gpu-amd-mi35x")


@@ -10,7 +10,6 @@ from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTest

 register_cuda_ci(est_time=55, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=55, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=31, suite="stage-b-test-small-1-gpu-5090")

 _is_hip = is_hip()
 if _is_hip:
@@ -17,7 +17,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=38, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=40, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=38, suite="stage-b-test-small-1-gpu-amd")


@@ -9,7 +9,7 @@ from sglang.test.test_utils import (
    CustomTestCase,
 )

-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=8, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=8, suite="stage-b-test-small-1-gpu-amd")


@@ -10,7 +10,6 @@ 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=6, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")

 HOOK_CALLS = []
@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=62, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd")


@@ -18,7 +18,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=47, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=49, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=70, suite="stage-b-test-small-1-gpu-amd")


@@ -7,7 +7,7 @@ from sglang.srt.entrypoints.engine import Engine
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase

-register_cuda_ci(est_time=260, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=260, suite="stage-b-test-large-1-gpu")

 TEST_MODEL_NAME = DEFAULT_SMALL_MODEL_NAME_FOR_TEST

@@ -7,7 +7,6 @@ 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=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=1, suite="stage-b-test-small-1-gpu-amd")


@@ -28,7 +28,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=127, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=122, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=130, suite="stage-b-test-small-1-gpu-amd")


@@ -22,7 +22,7 @@ from sglang.test.test_utils import (
    CustomTestCase,
 )

-register_cuda_ci(est_time=252, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=252, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=261, suite="stage-b-test-small-1-gpu-amd")


@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
 )

 # CI Registration - Small 1-GPU tests (24GB GPU sufficient)
-register_cuda_ci(est_time=539, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=539, suite="stage-b-test-large-1-gpu")


 class TestPiecewiseCudaGraphCorrectness(CustomTestCase):
@@ -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=7, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=15, suite="stage-b-test-small-1-gpu-amd")
-
 import unittest

 import torch
@@ -21,6 +15,18 @@ from sglang.srt.layers.linear import LinearBase
 from sglang.srt.models.qwen2 import Qwen2MLP
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import add_prefix
+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",
+    disabled="Test uses pytest-style function without TestCase class - see #17145",
+)
+register_amd_ci(
+    est_time=15,
+    suite="stage-b-test-small-1-gpu-amd",
+    disabled="Test uses pytest-style function without TestCase class - see #17145",
+)

 TEST_HIDDEN_SIZE = 32

@@ -1,6 +1,6 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=181, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=181, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=330, suite="stage-b-test-small-1-gpu-amd")

 import unittest
@@ -1,7 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci
-
-register_amd_ci(est_time=520, suite="stage-b-test-small-1-gpu-amd")
-
 """
 Test LLaDA2 (Diffusion Language Model) on AMD GPUs.

@@ -12,6 +8,7 @@ import unittest
 from types import SimpleNamespace

 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.send_one import BenchArgs, send_one_prompt
 from sglang.test.test_utils import (
@@ -23,6 +20,8 @@ from sglang.test.test_utils import (
    write_github_step_summary,
 )

+register_amd_ci(est_time=520, suite="stage-b-test-small-1-gpu-amd")
+

 class TestLLaDA2MiniAMD(CustomTestCase):
    @classmethod
@@ -2,7 +2,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=96, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=300, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=96, suite="stage-b-test-small-1-gpu-5090")

 import time
 import unittest
@@ -1,6 +1,6 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=524, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=524, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=524, suite="stage-b-test-small-1-gpu-amd")
 """
 Consolidated HiCache variant tests.
@@ -4,10 +4,6 @@ from typing import Callable
 import pytest
 import torch
 from flashinfer import fp4_quantize, scaled_fp4_grouped_quantize
-
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=300, suite="nightly-4-gpu-b200", nightly=True)
 from flashinfer.fused_moe import cutlass_fused_moe as flashinfer_cutlass_fused_moe
 from sgl_kernel import scaled_fp4_quant, silu_and_mul
 from torch.nn import functional as F
@@ -15,6 +11,9 @@ from torch.nn import functional as F
 from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
 from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType
 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="nightly-4-gpu-b200", nightly=True)

 if torch.cuda.get_device_capability() < (10, 0):
    pytest.skip(
@@ -4,11 +4,8 @@ from unittest.mock import MagicMock, patch

 import torch

-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=2, suite="nightly-1-gpu", nightly=True)
-
 from sglang.srt.layers import dp_attention as _dp_attn
+from sglang.test.ci.ci_register import register_cuda_ci

 # Patch DP-attention globals before importing backends
 _dp_attn.get_attention_tp_size = lambda: 1  # TP size = 1 for unit test
@@ -27,6 +24,8 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.test_utils import CustomTestCase

+register_cuda_ci(est_time=2, suite="nightly-1-gpu", nightly=True)
+
 # Global configuration for all indexer tests
 DEFAULT_CONFIG = {
    "device": "cuda",
@@ -2,7 +2,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=25, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=25, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=3, suite="stage-b-test-small-1-gpu-5090")

 # Adapted from https://github.com/vllm-project/vllm/blob/main/tests/kernels/mamba/test_causal_conv1d.py

@@ -1,7 +1,3 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=50, suite="stage-b-test-large-2-gpu")
-
 # Adapted from https://github.com/vllm-project/vllm/blob/2c58742dff8613a3bd7496f2008ce927e18d38d1/tests/kernels/mamba/test_mamba_mixer2.py


@@ -17,6 +13,9 @@ from sglang.srt.distributed.parallel_state import (
    init_distributed_environment,
    initialize_model_parallel,
 )
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=50, suite="stage-b-test-large-2-gpu")

 NUM_GPUS = 2

@@ -2,7 +2,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=20, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu-5090")

 # Adapted from https://github.com/vllm-project/vllm/blob/633f943e30a4444d890d26b81850f7217736f840/tests/kernels/mamba/test_mamba_ssm_ssd.py

@@ -2,7 +2,6 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=13, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=30, suite="stage-b-test-small-1-gpu-amd")
-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu-5090")

 # Adapted from https://github.com/vllm-project/vllm/blob/633f943e30a4444d890d26b81850f7217736f840/tests/kernels/mamba/test_mamba_ssm_ssd.py

@@ -30,7 +30,6 @@ from sglang.test.lora_utils import (
 from sglang.test.test_utils import CustomTestCase, is_in_ci

 register_cuda_ci(est_time=200, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=137, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(
    est_time=200,
    suite="stage-b-test-small-1-gpu-amd",
@@ -24,7 +24,6 @@ from sglang.test.runners import SRTRunner
 from sglang.test.test_utils import CustomTestCase

 register_cuda_ci(est_time=224, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=215, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=224, suite="stage-b-test-small-1-gpu-amd")

 PROMPTS = [
@@ -36,6 +36,7 @@ import torch

 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase

 register_cuda_ci(
    est_time=150,
@@ -45,13 +46,6 @@ register_amd_ci(
    est_time=250,
    suite="stage-b-test-small-1-gpu-amd",
 )
-register_cuda_ci(
-    est_time=76,
-    suite="stage-b-test-small-1-gpu-5090",
-)
-
-from sglang.test.test_utils import DEFAULT_PORT_FOR_SRT_TEST_RUNNER, CustomTestCase
-
 # Test configuration constants
 LORA_BACKEND = "triton"
 DISABLE_CUDA_GRAPH = False
@@ -17,12 +17,8 @@ import unittest

 import openai

-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=150, suite="nightly-1-gpu", nightly=True)
-register_amd_ci(est_time=150, suite="nightly-amd-1-gpu", nightly=True)
-
 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,
@@ -31,6 +27,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=150, suite="nightly-1-gpu", nightly=True)
+register_amd_ci(est_time=150, suite="nightly-amd-1-gpu", nightly=True)
+

 def get_real_lora_adapter() -> str:
    """Use a real LoRA adapter from Hugging Face."""
@@ -20,15 +20,14 @@ from sglang.test.lora_utils import (
    LORA_MODELS_QWEN3,
    run_lora_multiple_batch_on_model_cases,
 )
+from sglang.test.test_utils import CustomTestCase

-register_cuda_ci(est_time=97, suite="nightly-1-gpu", nightly=True)
 register_amd_ci(
    est_time=30,
    suite="stage-b-test-small-1-gpu-amd",
    disabled="see https://github.com/sgl-project/sglang/issues/13107",
 )
-
-from sglang.test.test_utils import CustomTestCase
+register_cuda_ci(est_time=97, suite="nightly-1-gpu", nightly=True)


 class TestLoRAQwen3(CustomTestCase):
@@ -34,7 +34,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=487, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=487, suite="stage-b-test-large-1-gpu")

 PROMPTS = [
    "SGL is a",
@@ -25,7 +25,7 @@ from sglang.test.lora_utils import (
 )
 from sglang.test.test_utils import CustomTestCase, is_in_ci

-register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")


@@ -2,12 +2,6 @@ import unittest
 from typing import Dict, List

 import requests
-
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-register_cuda_ci(est_time=32, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=29, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=32, suite="stage-b-test-small-1-gpu-amd")
 from prometheus_client.parser import text_string_to_metric_families
 from prometheus_client.samples import Sample

@@ -17,6 +11,7 @@ from sglang.srt.metrics.collector import (
    compute_routing_key_stats,
 )
 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_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
@@ -25,6 +20,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=32, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=32, suite="stage-b-test-small-1-gpu-amd")
+
 _MODEL_NAME = "Qwen/Qwen3-0.6B"


@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
 )

 # FlashMLA attention backend tests with MTP speculative decoding
-register_cuda_ci(est_time=284, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=284, suite="stage-b-test-large-1-gpu")


 class TestFlashMLAAttnBackend(unittest.TestCase):
@@ -13,7 +13,7 @@ from sglang.test.test_utils import (
 )

 # MLA attention test with MGSM evaluation
-register_cuda_ci(est_time=194, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=194, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=1100, suite="stage-b-test-small-1-gpu-amd")


@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
 )

 # DeepSeek-V3 MLA tests with torch compile, FA3, and MTP speculative decoding
-register_cuda_ci(est_time=442, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=442, suite="stage-b-test-large-1-gpu")
 register_amd_ci(
    est_time=221,
    suite="stage-b-test-small-1-gpu-amd",
@@ -15,7 +15,7 @@ from sglang.test.test_utils import (
 )

 # FlashInfer MLA backend tests with MTP speculative decoding
-register_cuda_ci(est_time=302, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=302, suite="stage-b-test-large-1-gpu")


 class TestFlashinferMLA(CustomTestCase):
@@ -13,7 +13,7 @@ from sglang.test.test_utils import (
 )

 # MLA FP8 KV cache test with MGSM evaluation
-register_cuda_ci(est_time=77, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=77, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=360, suite="stage-b-test-small-1-gpu-amd")


@@ -16,7 +16,7 @@ from sglang.test.test_utils import (
 )

 # DeepSeek-V3 INT8 quantization tests (channel and block INT8)
-register_cuda_ci(est_time=341, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=341, suite="stage-b-test-large-1-gpu")


 class TestMLADeepseekV3ChannelInt8(CustomTestCase):
@@ -6,7 +6,6 @@ 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=30, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=26, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=45, suite="stage-b-test-small-1-gpu-amd")


@@ -19,7 +19,7 @@ from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.model_loader.loader import ModelOptModelLoader
 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=9, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=9, suite="stage-b-test-small-1-gpu-amd")

 # Note: PYTHONPATH=python should be set when running tests
@@ -26,7 +26,6 @@ CALIBRATION_NUM_SAMPLES = 512
 DEFAULT_DEVICE = "cuda:0"

 register_cuda_ci(est_time=11, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu-5090")


 class TestModelOptModelLoader(CustomTestCase):
@@ -12,7 +12,7 @@ from sglang.srt.weight_sync.utils import update_weights
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST

-register_cuda_ci(est_time=29, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=29, suite="stage-b-test-large-1-gpu")


 class AsyncEngine(Engine):
@@ -1,13 +1,10 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
 # Model tests for compressed tensors (FP8)
-register_cuda_ci(est_time=42, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=42, suite="stage-b-test-small-1-gpu-amd")

 import unittest
 from types import SimpleNamespace

 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval
 from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -16,6 +13,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=42, suite="stage-b-test-large-1-gpu")
+register_amd_ci(est_time=42, suite="stage-b-test-small-1-gpu-amd")
+

 class TestCompressedTensorsLlama3FP8(CustomTestCase):
    @classmethod
@@ -1,19 +1,19 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-# Cross encoder model tests
-register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=105, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=150, suite="stage-b-test-small-1-gpu-amd")
-
 import multiprocessing as mp
 import random
 import unittest

 import torch

+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.runners import TEST_RERANK_QUERY_DOCS, HFRunner, SRTRunner
 from sglang.test.test_utils import CustomTestCase, is_in_ci

+# Cross encoder model tests
+
+
+register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=150, suite="stage-b-test-small-1-gpu-amd")
+
 MODELS = [
    ("cross-encoder/ms-marco-MiniLM-L6-v2", 1, 1e-2),
    ("BAAI/bge-reranker-v2-m3", 1, 1e-2),
@@ -1,14 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-# Embedding model tests
-register_cuda_ci(est_time=73, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=58, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(
-    est_time=73,
-    suite="stage-b-test-small-1-gpu-amd",
-    disabled="see https://github.com/sgl-project/sglang/issues/11127",
-)
-
 # 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.
@@ -31,6 +20,7 @@ from typing import Optional
 import torch
 from transformers import AutoConfig, AutoTokenizer

+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
 from sglang.test.test_utils import (
    CustomTestCase,
@@ -39,6 +29,14 @@ from sglang.test.test_utils import (
    is_in_ci,
 )

+# Embedding model tests
+register_amd_ci(
+    est_time=73,
+    suite="stage-b-test-small-1-gpu-amd",
+    disabled="see https://github.com/sgl-project/sglang/issues/11127",
+)
+register_cuda_ci(est_time=73, suite="stage-b-test-small-1-gpu")
+
 MODEL_TO_CONFIG = {
    "Alibaba-NLP/gte-Qwen2-1.5B-instruct": (1, 1e-5),
    "intfloat/e5-mistral-7b-instruct": (1, 1e-5),
@@ -1,8 +1,16 @@
+import multiprocessing as mp
+import random
+import time
+import unittest
+
+import torch
+from transformers import AutoConfig, AutoTokenizer
+
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase, get_similarities, is_in_ci

 # Encoder embedding model tests (CUDA only)
-register_cuda_ci(est_time=270, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=144, suite="stage-b-test-small-1-gpu-5090")

 # Copyright 2023-2024 SGLang Team
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,16 +28,8 @@ register_cuda_ci(est_time=144, suite="stage-b-test-small-1-gpu-5090")

 # python -m unittest test_encoder_embedding_models.TestEncoderEmbeddingModels.test_prefill_logits

-import multiprocessing as mp
-import random
-import time
-import unittest

-import torch
-from transformers import AutoConfig, AutoTokenizer
-
-from sglang.test.runners import DEFAULT_PROMPTS, HFRunner, SRTRunner
-from sglang.test.test_utils import CustomTestCase, get_similarities, is_in_ci
+register_cuda_ci(est_time=270, suite="stage-b-test-small-1-gpu")

 MODELS = [("BAAI/bge-small-en", 1, 1e-5), ("BAAI/bge-m3", 1, 1e-5)]

@@ -1,7 +1,7 @@
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 # Generation model tests (CUDA only)
-register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=103, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=106, suite="stage-b-test-small-1-gpu-amd")

 # Copyright 2023-2024 SGLang Team
@@ -1,16 +1,17 @@
-from sglang.test.ci.ci_register import register_cuda_ci
-
-# NVIDIA Nemotron Nano V2 VL model tests (CUDA only)
-# GSM8k + MMMU evaluation
-register_cuda_ci(est_time=214, suite="stage-b-test-small-1-gpu")
-
 import unittest

+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.gsm8k_accuracy_kit import GSM8KMixin
 from sglang.test.kits.mmmu_vlm_kit import MMMUMixin
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase

+# NVIDIA Nemotron Nano V2 VL model tests (CUDA only)
+# GSM8k + MMMU evaluation
+
+
+register_cuda_ci(est_time=214, suite="stage-b-test-large-1-gpu")
+
 MODEL = "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16"


@@ -1,14 +1,10 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
 # Qwen model tests
-register_cuda_ci(est_time=90, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=84, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=130, suite="stage-b-test-small-1-gpu-amd")

 import unittest
 from types import SimpleNamespace

 from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.few_shot_gsm8k import run_eval
 from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -17,6 +13,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=90, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=130, suite="stage-b-test-small-1-gpu-amd")
+

 class TestQwen2(CustomTestCase):
    @classmethod
@@ -1,9 +1,13 @@
+import multiprocessing as mp
+import unittest
+
+import torch
+
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.runners import HFRunner, SRTRunner
+from sglang.test.test_utils import CustomTestCase

 # Reward model tests
-register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=69, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=132, suite="stage-b-test-small-1-gpu-amd")

 # Copyright 2023-2024 SGLang Team
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,13 +23,9 @@ register_amd_ci(est_time=132, suite="stage-b-test-small-1-gpu-amd")
 # 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
+register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=132, suite="stage-b-test-small-1-gpu-amd")

 MODELS = [
    ("LxzGordon/URM-LLaMa-3.1-8B", 1, 4e-2),
@@ -1,9 +1,4 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
 # Transformers fallback model tests
-register_cuda_ci(est_time=245, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=210, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=320, suite="stage-b-test-small-1-gpu-amd")

 import dataclasses
 import multiprocessing as mp
@@ -14,6 +9,7 @@ from typing import List
 import torch

 from sglang.srt.utils import is_hip, kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.runners import DEFAULT_PROMPTS, SRTRunner, check_close_model_outputs
 from sglang.test.test_utils import (
    DEFAULT_MODEL_NAME_FOR_TEST,
@@ -24,6 +20,9 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

+register_cuda_ci(est_time=245, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=320, suite="stage-b-test-small-1-gpu-amd")
+

 class TestTransformersFallbackEndpoint(CustomTestCase):
    @classmethod
@@ -1,9 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-# VLM (Vision Language Model) tests
-register_cuda_ci(est_time=228, suite="stage-b-test-small-1-gpu")
-register_amd_ci(est_time=420, suite="stage-b-test-small-1-gpu-amd")
-
 import argparse
 import random
 import sys
@@ -11,12 +5,19 @@ import unittest
 from types import SimpleNamespace

 from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kits.mmmu_vlm_kit import (
    DEFAULT_MEM_FRACTION_STATIC,
    MMMUMultiModelTestBase,
 )
 from sglang.test.test_utils import is_in_ci

+# VLM (Vision Language Model) tests
+
+
+register_cuda_ci(est_time=228, suite="stage-b-test-large-1-gpu")
+register_amd_ci(est_time=420, suite="stage-b-test-small-1-gpu-amd")
+
 _is_hip = is_hip()
 # VLM models for testing
 if _is_hip:
@@ -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."
@@ -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()

@@ -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
@@ -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
@@ -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
@@ -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]
@@ -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"""
@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=70, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=48, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=141, suite="stage-b-test-small-1-gpu-amd")


@@ -29,7 +29,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=184, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=130, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=149, suite="stage-b-test-small-1-gpu-amd")


@@ -31,7 +31,6 @@ from sglang.srt.entrypoints.openai.protocol import (
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=3, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=2, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")


@@ -24,7 +24,6 @@ from sglang.srt.utils import get_or_create_event_loop
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=10, 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=10, suite="stage-b-test-small-1-gpu-amd")


@@ -14,7 +14,6 @@ from sglang.srt.managers.tokenizer_manager import TokenizerManager
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")


@@ -16,7 +16,7 @@ from sglang.srt.entrypoints.openai.serving_embedding import OpenAIServingEmbeddi
 from sglang.srt.managers.io_struct import EmbeddingReqInput
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=10, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")


@@ -6,7 +6,7 @@ from sglang.srt.entrypoints.openai.protocol import V1RerankReqInput
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 # Keep consistent with other openai_server/basic unit tests.
-register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=10, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu")

 try:
@@ -21,7 +21,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=103, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=200, suite="stage-b-test-small-1-gpu-amd")


@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=109, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=84, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=180, suite="stage-b-test-small-1-gpu-amd")


@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=27, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=20, suite="stage-b-test-small-1-gpu-amd")


@@ -17,7 +17,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=186, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=134, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(
    est_time=186,
    suite="stage-b-test-small-1-gpu-amd",
@@ -24,7 +24,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=89, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=77, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=89, suite="stage-b-test-small-1-gpu-amd")


@@ -14,7 +14,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=73, suite="stage-b-test-small-1-gpu-amd")


@@ -23,7 +23,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=120, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=109, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=258, suite="stage-b-test-small-1-gpu-amd")


@@ -22,7 +22,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=41, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=41, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=41, suite="stage-b-test-small-1-gpu-amd")


@@ -12,7 +12,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=40, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=39, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd")


@@ -12,7 +12,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=6, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=37, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=47, suite="stage-b-test-small-1-gpu-amd")


@@ -12,7 +12,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=38, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=38, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=31, suite="stage-b-test-small-1-gpu-amd")


@@ -1,10 +1,3 @@
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-
-# Ops - Repeat Interleave tests (1-GPU)
-register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")
-register_amd_ci(est_time=75, suite="stage-b-test-small-1-gpu-amd")
-
 import time
 from typing import Tuple

@@ -13,6 +6,13 @@ import pytest
 import torch

 from sglang.srt.models.utils import compute_cu_seqlens_from_grid_numpy as cpu_numpy_impl
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+# Ops - Repeat Interleave tests (1-GPU)
+
+
+register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=75, suite="stage-b-test-small-1-gpu-amd")


 def torch_ref_impl(grid_thw: torch.Tensor) -> torch.Tensor:
@@ -18,7 +18,6 @@ from sglang.srt.utils.profile_merger import ProfileMerger
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

 register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=7, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=8, suite="stage-b-test-small-1-gpu-amd")


@@ -5,7 +5,6 @@ from sglang.srt.managers.io_struct import ProfileReqInput
 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=6, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=9, suite="stage-b-test-small-1-gpu-amd")


@@ -30,7 +30,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=41, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=39, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd")

 OUTPUT_DIR = "./profiler_dir"
@@ -18,7 +18,6 @@ from sglang.test.test_utils import (
 )

 register_cuda_ci(est_time=77, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=76, suite="stage-b-test-small-1-gpu-5090")


 class TestAutoRound(CustomTestCase):
@@ -11,7 +11,6 @@ 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=44, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=25, suite="stage-b-test-small-1-gpu-5090")
 register_amd_ci(est_time=22, suite="stage-b-test-small-1-gpu-amd")


@@ -14,7 +14,7 @@ from sglang.test.test_utils import (
    popen_launch_server,
 )

-register_cuda_ci(est_time=250, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=250, suite="stage-b-test-large-1-gpu")
 register_amd_ci(est_time=303, suite="stage-b-test-small-1-gpu-amd")


@@ -9,7 +9,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

-register_cuda_ci(est_time=132, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=132, suite="stage-b-test-large-1-gpu")


 class TestFP8Base(CustomTestCase):
@@ -10,7 +10,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

-register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
+register_cuda_ci(est_time=9, suite="stage-b-test-large-1-gpu")


 class TestInverseTransformScaleUe8m0(CustomTestCase):
@@ -12,7 +12,6 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

 register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
-register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu-5090")


 def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
--- a/Show More
+++ b/Show More