[AMD] add dedicated jit-kernel-benchmark-test-amd stage + register portable JIT benches (#30307)

test/registered/jit/benchmark/bench_activation.py

@@ -10,11 +10,12 @@ from sglang.jit_kernel.activation import relu2 as relu2_jit
 from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
-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=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=30, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @torch.compile

test/registered/jit/benchmark/bench_add_constant.py

@@ -8,11 +8,12 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-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-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 CONSTANT = 7
 SIZE_LIST = get_benchmark_range(

test/registered/jit/benchmark/bench_clamp_position.py

@@ -16,7 +16,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 register_cuda_ci(
     est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
-register_amd_ci(est_time=16, stage="jit-kernel-unit", runner_config="amd")
+register_amd_ci(est_time=16, stage="jit-kernel-benchmark", runner_config="amd")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],

test/registered/jit/benchmark/bench_dsv3_router_gemm.py

@@ -16,11 +16,12 @@ from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
 from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm
 from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
-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")
 
 # sgl_kernel AOT kernel is specialized for hidden_dim=7168 only.
 SGL_KERNEL_HIDDEN_DIM = 7168

test/registered/jit/benchmark/bench_fused_eh_norm.py

@@ -4,11 +4,12 @@ import torch
 
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.fused_eh_norm import fused_eh_norm
-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=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=6, stage="jit-kernel-benchmark", runner_config="amd")
 
 EPS = 1e-6
 

test/registered/jit/benchmark/bench_hicache.py

@@ -31,11 +31,12 @@ from sglang.jit_kernel.hicache import (
     transfer_hicache_all_layer,
     transfer_hicache_one_layer,
 )
-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=29, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=29, stage="jit-kernel-benchmark", runner_config="amd")
 
 DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
 PAGE_SIZE = 1

test/registered/jit/benchmark/bench_hisparse.py

@@ -7,11 +7,12 @@ import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
 from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
-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=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=12, stage="jit-kernel-benchmark", runner_config="amd")
 
 DEVICE = DEFAULT_DEVICE
 DTYPE = DEFAULT_DTYPE

test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py

@@ -18,11 +18,12 @@ from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import (
     mla_kv_pack_quantize_fp8 as hybrid_pack,
 )
 from sglang.jit_kernel.utils import 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=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @triton.jit

test/registered/jit/benchmark/bench_ngram_compute_decode.py

@@ -11,11 +11,12 @@ from sglang.jit_kernel.ngram_embedding import (
     compute_n_gram_ids,
     compute_n_gram_ids_decode,
 )
-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-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 NE_N = 8
 NE_K = 2

test/registered/jit/benchmark/bench_ngram_update_token_table.py

@@ -11,11 +11,12 @@ from sglang.jit_kernel.ngram_embedding import (
     update_token_table,
     update_token_table_decode,
 )
-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-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 MAX_CONTEXT_LEN = 4096
 BATCH_SIZE_LIST = get_benchmark_range(

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -16,11 +16,12 @@ from sglang.jit_kernel.benchmark.utils import (
     run_benchmark_no_cudagraph,
 )
 from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module
-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=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=10, stage="jit-kernel-benchmark", runner_config="amd")
 
 HEAD_DIM = 512
 STATE_DIM = HEAD_DIM * 3

test/registered/jit/benchmark/bench_resolve_future_token_ids.py

@@ -16,7 +16,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 register_cuda_ci(
     est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
-register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
+register_amd_ci(est_time=10, stage="jit-kernel-benchmark", runner_config="amd")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],  # 16 … 32K elements

test/registered/jit/benchmark/bench_store_cache.py

@@ -7,11 +7,12 @@ from sglang.jit_kernel.benchmark.utils import (
     create_random,
 )
 from sglang.jit_kernel.kvcache import store_cache
-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=9, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=9, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @torch.compile()

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

@@ -9,12 +9,13 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
 )
-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(
     est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=13, stage="jit-kernel-benchmark", runner_config="amd")
 
 if is_in_ci():
     B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]

test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py

@@ -15,11 +15,12 @@ from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx i
     _topk_index_merge_kernel,
     _topk_index_partial_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")
 
 BLOCK_SIZE = 128
 TOPK = 16

test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py

@@ -5,11 +5,12 @@ import torch
 
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
-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=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=6, stage="jit-kernel-benchmark", runner_config="amd")
 
 HEAD_DIM = 128
 NUM_KV_HEADS = 1

test/run_suite.py

@@ -44,6 +44,7 @@ PER_COMMIT_SUITES = {
         "stage-b-test-1-gpu-large-amd",
         "stage-b-test-2-gpu-large-amd",
         "jit-kernel-unit-test-amd",
+        "jit-kernel-benchmark-test-amd",
         "sgl-kernel-unit-test-2-gpu-amd",
         "stage-c-test-4-gpu-amd",
         "stage-c-test-large-8-gpu-amd",