[AMD] register 8 JIT kernel benchmarks to jit-kernel-benchmark-test-amd (#31492)

test/registered/jit/benchmark/bench_custom_all_reduce.py

@@ -14,7 +14,7 @@ from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main
 from sglang.jit_kernel.mp import register_comm_cleanup
 from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl
-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=120,
@@ -22,6 +22,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="requires multi-GPU, self-skips in CI",
 )
+register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 # ---------------------------------------------------------------------------

test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py

@@ -8,13 +8,14 @@ import triton
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
 from sglang.srt.utils import is_sm120_supported
-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=5,
     stage="base-b-kernel-benchmark",
     runner_config="1-gpu-large",
 )
+register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 def _make_inputs(m: int, n: int, k: int, device: str = "cuda"):

test/registered/jit/benchmark/bench_post_reorder_deepgemm.py

@@ -5,11 +5,12 @@ from sglang.kernels.ops.moe.ep_moe_kernels import (
     post_reorder_deepgemm,
     post_reorder_triton_kernel,
 )
-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=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=8, stage="jit-kernel-benchmark", runner_config="amd")
 
 HIDDEN = 6144
 NUM_EXPERTS = 129

test/registered/jit/benchmark/bench_symm_mem_all_gather.py

@@ -33,7 +33,7 @@ from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
     all_gather_inner,
     create_state,
 )
-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=120,
@@ -41,6 +41,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="requires multi-GPU, self-skips in CI",
 )
+register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd")
 
 # ---------------------------------------------------------------------------
 # Sweep parameters

test/registered/jit/benchmark/bench_tp_qknorm.py

@@ -36,7 +36,7 @@ from sglang.jit_kernel.utils import cache_once, get_ci_test_range
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
-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=120,
@@ -44,6 +44,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="requires multi-GPU, self-skips in CI",
 )
+register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 # ---------------------------------------------------------------------------

test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py

@@ -9,7 +9,7 @@ from sglang.jit_kernel.diffusion.causal_conv3d_cat_pad import (
 from sglang.jit_kernel.diffusion.triton.causal_conv3d_pad import (
     fused_causal_conv3d_cat_pad as fused_causal_conv3d_cat_pad_triton,
 )
-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=20,
@@ -17,6 +17,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="standalone benchmark",
 )
+register_amd_ci(est_time=20, stage="jit-kernel-benchmark", runner_config="amd")
 
 DEVICE = "cuda"
 DTYPE = torch.bfloat16

test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py

@@ -11,7 +11,7 @@ import torch.nn.functional as F
 import triton.testing
 
 from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
-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.utils import is_in_ci
 
 register_cuda_ci(
@@ -20,6 +20,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="standalone benchmark",
 )
+register_amd_ci(est_time=45, stage="jit-kernel-benchmark", runner_config="amd")
 
 DEVICE = "cuda"
 EPS = 1e-5

test/registered/jit/benchmark/diffusion/bench_norm_impls.py

@@ -19,7 +19,7 @@ from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_r
 from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
 from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
 from sglang.jit_kernel.utils import KERNEL_PATH
-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.utils import is_in_ci
 
 register_cuda_ci(
@@ -28,6 +28,7 @@ register_cuda_ci(
     runner_config="1-gpu-large",
     disabled="self-skips in CI, standalone tool",
 )
+register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd")
 
 os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")