From 7fbe91c6ea3c0ea0061cf313f682a66a54ca0f1a Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:00:00 -0700 Subject: [PATCH] [AMD] register 8 JIT kernel benchmarks to jit-kernel-benchmark-test-amd (#31492) --- test/registered/jit/benchmark/bench_custom_all_reduce.py | 3 ++- test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py | 3 ++- test/registered/jit/benchmark/bench_post_reorder_deepgemm.py | 3 ++- test/registered/jit/benchmark/bench_symm_mem_all_gather.py | 3 ++- test/registered/jit/benchmark/bench_tp_qknorm.py | 3 ++- .../jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py | 3 ++- .../jit/benchmark/diffusion/bench_group_norm_silu.py | 3 ++- test/registered/jit/benchmark/diffusion/bench_norm_impls.py | 3 ++- 8 files changed, 16 insertions(+), 8 deletions(-) diff --git a/test/registered/jit/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py index e0ea022d3..4bc187693 100644 --- a/test/registered/jit/benchmark/bench_custom_all_reduce.py +++ b/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") # --------------------------------------------------------------------------- diff --git a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py b/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py index dcab974f1..fb603bfa0 100644 --- a/test/registered/jit/benchmark/bench_fp8_blockwise_gemm.py +++ b/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"): diff --git a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py b/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py index 185bf5c19..3964f7e34 100644 --- a/test/registered/jit/benchmark/bench_post_reorder_deepgemm.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py b/test/registered/jit/benchmark/bench_symm_mem_all_gather.py index 9abb934ff..86c31ff84 100644 --- a/test/registered/jit/benchmark/bench_symm_mem_all_gather.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py index a63fdb6ba..290f5b71c 100644 --- a/test/registered/jit/benchmark/bench_tp_qknorm.py +++ b/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") # --------------------------------------------------------------------------- diff --git a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py b/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py index f4e6ae892..318e98f57 100644 --- a/test/registered/jit/benchmark/diffusion/bench_causal_conv3d_cat_pad.py +++ b/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 diff --git a/test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py b/test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py index 6875a111c..d116cd02c 100644 --- a/test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py +++ b/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 diff --git a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py index 59db976af..489f10693 100644 --- a/test/registered/jit/benchmark/diffusion/bench_norm_impls.py +++ b/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")