Refactor JIT kernel CI to use run_suite.py registration system (#21239)

python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py

@@ -6,8 +6,11 @@ import triton.testing
 
 from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
 from sglang.jit_kernel.benchmark.utils import run_benchmark
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 try:
     from sgl_kernel import awq_dequantize as aot_awq_dequantize
 

python/sglang/jit_kernel/benchmark/bench_clamp_position.py

@@ -11,6 +11,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.jit_kernel.clamp_position import clamp_position_cuda
 from sglang.srt.utils import get_compiler_backend
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],

python/sglang/jit_kernel/benchmark/bench_concat_mla.py

@@ -9,8 +9,11 @@ from sgl_kernel import concat_mla_k as aot_k
 from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q
 from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 IS_CI = is_in_ci()
 
 NUM_LOCAL_HEADS = 128

python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py

@@ -22,6 +22,13 @@ import torch
 import torch.distributed as dist
 
 from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=120,
+    suite="stage-b-kernel-benchmark-1-gpu-large",
+    disabled="requires multi-GPU, self-skips in CI",
+)
 
 DTYPE_MAP = {
     "float16": torch.float16,

python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py

@@ -7,8 +7,11 @@ from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm
 
 from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 IS_CI = is_in_ci()
 
 

python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py

@@ -15,8 +15,11 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     ScaleResidualLayerNormScaleShift,
     ScaleResidualRMSNormScaleShift,
 )
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=17, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 if is_in_ci():
     B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024]
 else:

python/sglang/jit_kernel/benchmark/bench_hadamard.py

@@ -14,6 +14,9 @@ from sglang.jit_kernel.benchmark.utils import (
     run_benchmark,
 )
 from sglang.jit_kernel.hadamard import hadamard_transform
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 # AOT kernel: might not be available in all environments.
 # This is used for performance baseline comparison.

python/sglang/jit_kernel/benchmark/bench_hicache.py

@@ -31,6 +31,9 @@ from sglang.jit_kernel.hicache import (
     transfer_hicache_all_layer,
     transfer_hicache_one_layer,
 )
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=29, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
 PAGE_SIZE = 1

python/sglang/jit_kernel/benchmark/bench_norm.py

@@ -9,8 +9,11 @@ from flashinfer.norm import rmsnorm as fi_rmsnorm
 from sglang.jit_kernel.benchmark.utils import run_benchmark
 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.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 IS_CI = is_in_ci()
 
 DTYPE = torch.bfloat16

python/sglang/jit_kernel/benchmark/bench_norm_impls.py

@@ -21,8 +21,15 @@ 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.utils import is_in_ci
 
+register_cuda_ci(
+    est_time=120,
+    suite="stage-b-kernel-benchmark-1-gpu-large",
+    disabled="self-skips in CI, standalone tool",
+)
+
 os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
 
 REPO_ROOT = KERNEL_PATH.parents[2]

python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py

@@ -13,6 +13,9 @@ from sglang.jit_kernel.nvfp4 import (
     scaled_fp4_quant,
 )
 from sglang.srt.utils import is_sm100_supported
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py

@@ -8,6 +8,9 @@ import triton
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
 from sglang.srt.utils import is_sm100_supported
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py

@@ -8,6 +8,9 @@ import triton
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from sglang.srt.utils import is_sm100_supported
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py

@@ -6,6 +6,9 @@ import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 try:
     from vllm import _custom_ops as ops

python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py

@@ -17,8 +17,11 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 )
 from sglang.srt.utils import is_hip
 from sglang.srt.utils.bench_utils import bench_kineto
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 IS_CI = is_in_ci()
 
 _is_hip = is_hip()

python/sglang/jit_kernel/benchmark/bench_qknorm.py

@@ -13,6 +13,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.jit_kernel.norm import fused_inplace_qknorm
 from sglang.srt.utils import get_current_device_stream_fast
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 alt_stream = torch.cuda.Stream()
 

python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py

@@ -9,8 +9,11 @@ from sgl_kernel import rmsnorm
 from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
 from sglang.srt.utils import get_current_device_stream_fast
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=12, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 IS_CI = is_in_ci()
 
 alt_stream = torch.cuda.Stream()

python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py

@@ -9,8 +9,11 @@ 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.utils import is_in_ci
 
+register_cuda_ci(est_time=13, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 if is_in_ci():
     B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
 else:

python/sglang/jit_kernel/benchmark/bench_renorm.py

@@ -6,8 +6,11 @@ import triton
 import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
+from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 
+register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
+
 
 def torch_top_k_renorm_probs(probs, top_k):
     """Vectorized PyTorch implementation of top-k renormalization."""

python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py

@@ -11,6 +11,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda
 from sglang.srt.utils import get_compiler_backend
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],  # 16 … 32K elements

python/sglang/jit_kernel/benchmark/bench_rmsnorm.py

@@ -13,6 +13,9 @@ from sglang.jit_kernel.benchmark.utils import (
     run_benchmark,
 )
 from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=21, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 
 def sglang_aot_rmsnorm(

python/sglang/jit_kernel/benchmark/bench_rope.py

@@ -10,6 +10,9 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 MAX_SEQ_LEN = 131072
 ROPE_BASE = 10000.0

python/sglang/jit_kernel/benchmark/bench_store_cache.py

@@ -12,6 +12,9 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
 )
 from sglang.jit_kernel.kvcache import store_cache
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=9, suite="stage-b-kernel-benchmark-1-gpu-large")
 
 
 def sglang_jit_store_cache(