From a2aa51c818dbf5ec32d113a93bc3da057bdb4d6d Mon Sep 17 00:00:00 2001
From: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Date: Tue, 16 Jun 2026 21:58:34 -0700
Subject: [PATCH] [AMD] register 4 2-gpu tests to stage-b-test-2-gpu-large-amd
 (#28344)

---
 .../attention/test_gemma4_swa_triton_oob_regression.py         | 3 ++-
 .../unified_radix_tree/test_unified_radix_cache_kl_full.py     | 3 ++-
 test/registered/rl/test_patch_torch.py                         | 3 ++-
 test/registered/scheduler/test_load_snapshot_server.py         | 3 ++-
 4 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/test/registered/attention/test_gemma4_swa_triton_oob_regression.py b/test/registered/attention/test_gemma4_swa_triton_oob_regression.py
index 6245df7e5..675c7d6f9 100644
--- a/test/registered/attention/test_gemma4_swa_triton_oob_regression.py
+++ b/test/registered/attention/test_gemma4_swa_triton_oob_regression.py
@@ -15,7 +15,7 @@ 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.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,
@@ -24,6 +24,7 @@ from sglang.test.test_utils import (
 )
 
 register_cuda_ci(est_time=107, stage="base-b", runner_config="2-gpu-large")
+register_amd_ci(est_time=160, suite="stage-b-test-2-gpu-large-amd")
 
 
 PROMPT = (
diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
index 1f66f6d9c..03f988448 100644
--- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
+++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_full.py
@@ -1,7 +1,7 @@
 import unittest
 
 from sglang.srt.utils import kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
 from sglang.test.kl_multiturn_utils import get_input_ids
 from sglang.test.test_utils import (
@@ -12,6 +12,7 @@ from sglang.test.test_utils import (
 )
 
 register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
+register_amd_ci(est_time=400, suite="stage-b-test-2-gpu-large-amd")
 
 FULL_MODEL = "Qwen/Qwen3-32B"
 
diff --git a/test/registered/rl/test_patch_torch.py b/test/registered/rl/test_patch_torch.py
index fd5dc9412..b28869e8a 100644
--- a/test/registered/rl/test_patch_torch.py
+++ b/test/registered/rl/test_patch_torch.py
@@ -7,9 +7,10 @@ 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_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=15, stage="base-b", runner_config="2-gpu-large")
+register_amd_ci(est_time=30, suite="stage-b-test-2-gpu-large-amd")
 
 
 class TestReleaseMemoryOccupation(unittest.TestCase):
diff --git a/test/registered/scheduler/test_load_snapshot_server.py b/test/registered/scheduler/test_load_snapshot_server.py
index b3ec25e1c..9dd701801 100644
--- a/test/registered/scheduler/test_load_snapshot_server.py
+++ b/test/registered/scheduler/test_load_snapshot_server.py
@@ -10,7 +10,7 @@ import unittest
 import urllib.request
 
 from sglang.srt.utils import kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
+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,
@@ -20,6 +20,7 @@ from sglang.test.test_utils import (
 )
 
 register_cuda_ci(est_time=300, stage="base-b", runner_config="2-gpu-large")
+register_amd_ci(est_time=450, suite="stage-b-test-2-gpu-large-amd")
 
 
 def _query_loads(base_url, retries=5, interval=2.0):