[3/n jit_kernel restruct] Clean up benchmark naming and benchmarking helpers (#20250)

python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py

@@ -6,6 +6,7 @@ import triton
 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
 
 try:
     from sgl_kernel import awq_dequantize as aot_awq_dequantize
@@ -19,7 +20,6 @@ IS_CI = (
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-# CI environment uses simplified parameters
 if IS_CI:
     qweight_row_range = [128]
     qweight_cols_range = [16]
@@ -107,8 +107,6 @@ def benchmark(qweight_row, qweight_col, provider):
         device=device,
     )
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
     elif provider == "aot":
@@ -116,8 +114,7 @@ def benchmark(qweight_row, qweight_col, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py

@@ -9,6 +9,7 @@ from sgl_kernel.scalar_type import scalar_types
 from sglang.jit_kernel.awq_marlin_repack import (
     awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
 )
+from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
 
 try:
@@ -23,7 +24,6 @@ IS_CI = (
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-# Fixed parameters
 NUM_BITS = 4
 GROUP_SIZE = 128
 SIZE_N = 4096
@@ -111,8 +111,6 @@ def benchmark(num_experts, size_k, size_n, num_bits, provider):
         num_experts, size_k, size_n, num_bits, group_size
     )
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: jit_awq_marlin_moe_repack(
             b_q_weight, perm, size_k, size_n, num_bits
@@ -124,8 +122,7 @@ def benchmark(num_experts, size_k, size_n, num_bits, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py

@@ -9,6 +9,7 @@ from sgl_kernel.scalar_type import scalar_types
 from sglang.jit_kernel.awq_marlin_repack import (
     awq_marlin_repack as jit_awq_marlin_repack,
 )
+from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
 
 try:
@@ -23,7 +24,6 @@ IS_CI = (
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-# Fixed problem dimensions
 SIZE_K = 4096
 SIZE_N = 4096
 NUM_BITS = 4
@@ -43,7 +43,6 @@ def awq_pack(q_w, num_bits, size_k, size_n):
     return pack_cols(q_w, num_bits, size_k, size_n)
 
 
-# Quantize weights once
 _b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
 _w_ref, _q_w, _s, _zp = quantize_weights(
     _b_weight, scalar_types.uint4, GROUP_SIZE, zero_points=True
@@ -99,8 +98,6 @@ def benchmark(size_k, size_n, num_bits, provider):
     )
     q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
     elif provider == "aot":
@@ -108,8 +105,7 @@ def benchmark(size_k, size_n, num_bits, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_concat_mla.py

@@ -6,13 +6,12 @@ import triton.testing
 from sgl_kernel import concat_mla_absorb_q as aot_absorb_q
 from sgl_kernel import concat_mla_k as aot_k
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, 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
 
 IS_CI = is_in_ci()
 
-# Constants
 NUM_LOCAL_HEADS = 128
 QK_NOPE_HEAD_DIM = 128
 QK_ROPE_HEAD_DIM = 64
@@ -109,9 +108,7 @@ def bench_concat_mla_k(num_tokens: int, provider: str):
         "torch": torch_concat_mla_k,
     }
     fn = lambda: FN_MAP[provider](k, k_nope, k_rope)
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if IS_CI:
@@ -153,9 +150,7 @@ def bench_concat_mla_absorb_q(dim_0: int, dim_1: int, provider: str):
         }
         fn = lambda: FN_MAP[provider](a, b)
 
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py

@@ -5,7 +5,7 @@ import triton
 import triton.testing
 from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
 from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
 
 IS_CI = is_in_ci()
@@ -33,7 +33,7 @@ else:
     BS_LIST = [2**n for n in range(0, 14)]
     HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
 
-LINE_VALS = ["jit", "fi"]
+LINE_VALS = ["jit", "flashinfer"]
 LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
 STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
 
@@ -59,14 +59,12 @@ def benchmark(hidden_size: int, batch_size: int, provider: str):
     weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
     FN_MAP = {
         "jit": sglang_jit_fused_add_rmsnorm,
-        "fi": flashinfer_fused_add_rmsnorm,
+        "flashinfer": flashinfer_fused_add_rmsnorm,
     }
     fn = lambda: FN_MAP[provider](
         input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps
     )
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py

@@ -8,7 +8,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNormScaleShift,
     RMSNormScaleShift,
@@ -78,9 +78,7 @@ def bench_fused_norm_scale_shift(
     else:
         fn = lambda: layer.forward_cuda(x, shift, scale)
 
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark_no_cudagraph(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms  # convert to us
 
 
 # ============================================================================
@@ -117,9 +115,7 @@ def bench_fused_scale_residual_norm_scale_shift(
     else:
         fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale)
 
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark_no_cudagraph(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms  # convert to us
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py

@@ -5,6 +5,7 @@ import triton
 import triton.testing
 from sgl_kernel.scalar_type import scalar_types
 
+from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm as jit_gptq_marlin_gemm
 from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
 from sglang.test.test_marlin_utils import marlin_quantize
@@ -21,13 +22,11 @@ IS_CI = (
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-# Fixed problem dimensions
 SIZE_K = 4096
 SIZE_N = 4096
 GROUP_SIZE = 128
 QUANT_TYPE = scalar_types.uint4b8
 
-# Quantize weights once
 _b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
 _w_ref, _marlin_q_w, _marlin_s, _g_idx, _sort_indices, _ = marlin_quantize(
     _b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False
@@ -100,8 +99,6 @@ def benchmark(size_m, provider):
     device = torch.device("cuda")
     a = torch.randn((size_m, SIZE_K), dtype=torch.float16, device=device)
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: _run_gemm(jit_gptq_marlin_gemm, a)
     elif provider == "aot":
@@ -109,8 +106,7 @@ def benchmark(size_m, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py

@@ -5,6 +5,7 @@ import triton
 import triton.testing
 from sgl_kernel.scalar_type import scalar_types
 
+from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn
 from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows
 
@@ -20,13 +21,11 @@ IS_CI = (
     or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
 )
 
-# Fixed problem dimensions
 SIZE_N = 4096
 NUM_BITS = 4
 QUANT_TYPE = scalar_types.uint4b8
 GROUP_SIZE = 128
 
-# Pre-compute quantized weight for each size_k in the sweep
 _cache = {}
 
 
@@ -86,8 +85,6 @@ else:
 def benchmark(size_k, provider):
     q_w_gptq, sort_indices = _get_inputs(size_k)
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
     elif provider == "aot":
@@ -95,8 +92,7 @@ def benchmark(size_k, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_hicache.py

@@ -32,9 +32,6 @@ from sglang.jit_kernel.hicache import (
     transfer_hicache_one_layer,
 )
 
-# NOTE: Adjustable hyperparameters for better benchmark stability
-
-# NOTE: torch impl is too slow in benchmark
 DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
 PAGE_SIZE = 1
 ENABLE_SORT = True
@@ -193,7 +190,7 @@ ELEMENT_SIZE_RANGE = get_benchmark_range(
     ci_range=[1024],
 )
 
-LINE_VALS = ["aot", "jit", "pytorch"]
+LINE_VALS = ["aot", "jit", "torch"]
 LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
 STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
 
@@ -228,12 +225,10 @@ def benchmark_one_layer_h2d(
     v_cache_src = cache_local.v_cache_host
     k_cache_dst = cache_local.k_cache_cuda
     v_cache_dst = cache_local.v_cache_cuda
-    # to avoid fluctutation, we set the seed as const
     torch.manual_seed(batch_size * 65536 + element_size)
     indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
     indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
 
-    # sort by host indices to improve host access performance
     if ENABLE_SORT:
         indices_src_gpu, mapping = indices_src_gpu.sort()
         indices_dst_gpu = indices_dst_gpu[mapping]
@@ -267,7 +262,7 @@ def benchmark_one_layer_h2d(
             )
             for i in range(NUM_LAYERS)
         ],
-        "pytorch": lambda: [
+        "torch": lambda: [
             pytorch_transfer(
                 k_cache_dst[i],
                 v_cache_dst[i],
@@ -283,7 +278,7 @@ def benchmark_one_layer_h2d(
     if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
         return (float("nan"), float("nan"), float("nan"))
 
-    if DISABLE_TORCH and provider in ["pytorch"]:
+    if DISABLE_TORCH and provider in ["torch"]:
         return (float("nan"), float("nan"), float("nan"))
 
     ms, min_ms, max_ms = triton.testing.do_bench(  # type: ignore
@@ -333,12 +328,10 @@ def benchmark_all_layer_d2h(
     v_caches_src = cache_local.v_cache_cuda
     k_caches_dst = cache_local.k_cache_host
     v_caches_dst = cache_local.v_cache_host
-    # to avoid fluctutation, we set the seed as const
     torch.manual_seed(batch_size * 65536 + element_size)
 
     indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
     indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
-    # sort by host indices to improve host access performance
     if ENABLE_SORT:
         indices_dst_gpu, mapping = indices_dst_gpu.sort()
         indices_src_gpu = indices_src_gpu[mapping]
@@ -373,7 +366,7 @@ def benchmark_all_layer_d2h(
             element_bytes,
             element_bytes,
         ),
-        "pytorch": lambda: [
+        "torch": lambda: [
             pytorch_transfer(
                 k_caches_dst[i],
                 v_caches_dst[i],
@@ -389,7 +382,7 @@ def benchmark_all_layer_d2h(
     if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
         return (float("nan"), float("nan"), float("nan"))
 
-    if DISABLE_TORCH and provider in ["pytorch"]:
+    if DISABLE_TORCH and provider in ["torch"]:
         return (float("nan"), float("nan"), float("nan"))
 
     ms, min_ms, max_ms = triton.testing.do_bench(  # type: ignore

python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py

@@ -5,6 +5,7 @@ import triton
 import triton.testing
 from sgl_kernel.scalar_type import scalar_types
 
+from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm as jit_fn
 from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
 from sglang.test.test_marlin_utils import marlin_quantize
@@ -27,7 +28,6 @@ def stack_and_dev(tensors):
     return torch.stack(tensors, dim=0).to(dev)
 
 
-# Fixed problem dimensions
 E = 8
 SIZE_K = 4096
 SIZE_N = 4096
@@ -37,7 +37,6 @@ QUANT_TYPE = scalar_types.uint4b8
 DTYPE = torch.float16
 BLOCK_SIZE_M = 64
 
-# Quantize weights once (per-expert)
 torch.manual_seed(0)
 _qweight_l, _scales_l, _w_ref_l = [], [], []
 for i in range(E):
@@ -215,8 +214,6 @@ def benchmark(size_m, provider):
         _make_inputs(size_m)
     )
 
-    quantiles = [0.5, 0.2, 0.8]
-
     if provider == "jit":
         fn = lambda: _run_jit(
             a,
@@ -242,8 +239,7 @@ def benchmark(size_m, provider):
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return run_benchmark(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_norm.py

@@ -6,7 +6,7 @@ import triton.testing
 from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
 from flashinfer.norm import rmsnorm as fi_rmsnorm
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, 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
 
@@ -25,7 +25,7 @@ else:
     BS_LIST = [2**n for n in range(0, 14)]
     HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
 
-LINE_VALS = ["jit", "fi"]
+LINE_VALS = ["jit", "flashinfer"]
 LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
 STYLES = [("blue", "--"), ("green", "-.")]
 
@@ -50,12 +50,10 @@ def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
     weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
     FN_MAP = {
         "jit": lambda: jit_rmsnorm(input.clone(), weight),
-        "fi": lambda: fi_rmsnorm(input.clone(), weight, out=input.clone()),
+        "flashinfer": lambda: fi_rmsnorm(input.clone(), weight, out=input.clone()),
     }
     fn = FN_MAP[provider]
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 @triton.testing.perf_report(
@@ -79,14 +77,12 @@ def benchmark_fused_add_rmsnorm(hidden_size: int, batch_size: int, provider: str
         "jit": lambda: jit_fused_add_rmsnorm(
             input.clone(), residual.clone(), weight, torch.finfo(DTYPE).eps
         ),
-        "fi": lambda: fi_fused_add_rmsnorm(
+        "flashinfer": lambda: fi_fused_add_rmsnorm(
             input.clone(), residual.clone(), weight, eps=torch.finfo(DTYPE).eps
         ),
     }
     fn = FN_MAP[provider]
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_qknorm.py

@@ -91,7 +91,7 @@ HEAD_DIM_RANGE = get_benchmark_range(
     ci_range=[128],
 )
 
-LINE_VALS = ["aot", "jit", "fi", "torch"]
+LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
 LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
 STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
 
@@ -126,7 +126,7 @@ def benchmark(
     FN_MAP = {
         "aot": sglang_aot_qknorm,
         "jit": sglang_jit_qknorm,
-        "fi": flashinfer_qknorm,
+        "flashinfer": flashinfer_qknorm,
         "torch": torch_impl_qknorm,
     }
     fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)

python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py

@@ -6,7 +6,7 @@ import triton
 import triton.testing
 from sgl_kernel import rmsnorm
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
 from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
 from sglang.srt.utils import get_current_device_stream_fast
 
@@ -78,7 +78,7 @@ else:
     BS_RANGE = [2**n for n in range(0, 14)]
     HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192]
 
-LINE_VALS = ["jit", "aot", "fi", "torch"]
+LINE_VALS = ["jit", "aot", "flashinfer", "torch"]
 LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"]
 STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")]
 
@@ -108,13 +108,11 @@ def benchmark(
     FN_MAP = {
         "jit": sglang_jit_qknorm_across_heads,
         "aot": sglang_aot_qknorm_across_heads,
-        "fi": flashinfer_qknorm_across_heads,
+        "flashinfer": flashinfer_qknorm_across_heads,
         "torch": torch_impl_qknorm_across_heads,
     }
     fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)
-    quantiles = [0.5, 0.2, 0.8]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_renorm.py

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.testing
 
-from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph
 
 
 def torch_top_k_renorm_probs(probs, top_k):
@@ -171,7 +171,7 @@ configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range))
         x_vals=configs_k,
         line_arg="provider",
         line_vals=["torch", "sglang"],
-        line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
+        line_names=["Torch Reference", "SGL Kernel"],
         styles=[("red", "-"), ("green", "-")],
         ylabel="us",
         plot_name="top-k-renorm-probs-performance",
@@ -195,8 +195,7 @@ def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
     elif provider == "sglang":
         fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor)
 
-    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return run_benchmark_no_cudagraph(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 @triton.testing.perf_report(
@@ -205,7 +204,7 @@ def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
         x_vals=configs_p,
         line_arg="provider",
         line_vals=["torch", "sglang"],
-        line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
+        line_names=["Torch Reference", "SGL Kernel"],
         styles=[("red", "-"), ("blue", "-")],
         ylabel="us",
         plot_name="top-p-renorm-probs-performance",
@@ -225,8 +224,7 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
     elif provider == "sglang":
         fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor)
 
-    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return run_benchmark_no_cudagraph(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 @triton.testing.perf_report(
@@ -235,7 +233,7 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
         x_vals=configs_k,
         line_arg="provider",
         line_vals=["torch", "sglang"],
-        line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
+        line_names=["Torch Reference", "SGL Kernel"],
         styles=[("red", "-"), ("orange", "-")],
         ylabel="us",
         plot_name="top-k-mask-logits-performance",
@@ -258,8 +256,7 @@ def benchmark_top_k_mask(batch_size, vocab_size, k, provider):
     elif provider == "sglang":
         fn = lambda: sgl_kernel.top_k_mask_logits(logits.clone(), top_k_tensor)
 
-    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
+    return run_benchmark_no_cudagraph(fn)
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

python/sglang/jit_kernel/benchmark/bench_rmsnorm.py

@@ -56,7 +56,7 @@ HIDDEN_SIZE_LIST = get_benchmark_range(
     ci_range=[512, 2048],
 )
 
-LINE_VALS = ["aot", "jit", "fi", "torch"]
+LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
 LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
 STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
 
@@ -84,7 +84,7 @@ def benchmark(hidden_size: int, batch_size: int, provider: str):
     FN_MAP = {
         "aot": sglang_aot_rmsnorm,
         "jit": sglang_jit_rmsnorm,
-        "fi": flashinfer_rmsnorm,
+        "flashinfer": flashinfer_rmsnorm,
         "torch": torch_impl_rmsnorm,
     }
     fn = lambda: FN_MAP[provider](input.clone(), weight)

python/sglang/jit_kernel/benchmark/bench_rope.py

@@ -64,7 +64,7 @@ def flashinfer_rope(
     )
 
 
-def sglang_rope_v0(
+def sglang_pos_enc_rope(
     q: torch.Tensor,
     k: torch.Tensor,
     positions: torch.Tensor,
@@ -83,7 +83,7 @@ def sglang_rope_v0(
     )
 
 
-def sglang_rope_v1(
+def sgl_kernel_rope(
     q: torch.Tensor,
     k: torch.Tensor,
     positions: torch.Tensor,
@@ -102,7 +102,7 @@ def sglang_rope_v1(
     )
 
 
-def sglang_rope_v2(
+def sglang_fused_rope(
     q: torch.Tensor,
     k: torch.Tensor,
     positions: torch.Tensor,
@@ -118,7 +118,7 @@ def sglang_rope_v2(
 # ---------------------------------------------------------------------------
 
 
-def rope_v0_store(
+def jit_rope_then_store(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
@@ -150,7 +150,7 @@ def rope_v0_store(
     )
 
 
-def rope_v1_store(
+def sgl_kernel_fused_rope_store(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
@@ -181,7 +181,7 @@ def rope_v1_store(
     )
 
 
-def rope_v2_store(
+def jit_fused_rope_store(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
@@ -221,8 +221,13 @@ IS_NEOX_RANGE = get_benchmark_range(
 # Benchmark 1: RoPE only
 # ---------------------------------------------------------------------------
 
-ROPE_LINE_VALS = ["fi", "rope_v0", "rope_v1", "rope_v2"]
+ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "sgl_kernel", "jit_fused_rope"]
-ROPE_LINE_NAMES = ["FlashInfer", "SGL RoPE v0", "SGL RoPE v1", "SGL RoPE v2"]
+ROPE_LINE_NAMES = [
+    "FlashInfer",
+    "SGL JIT PosEnc",
+    "sgl-kernel",
+    "SGL JIT Fused RoPE",
+]
 ROPE_STYLES = [("green", "-."), ("red", "-"), ("orange", "-"), ("blue", "--")]
 
 rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@@ -263,10 +268,10 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str)
     torch.cuda.synchronize()
 
     FN_MAP = {
-        "fi": flashinfer_rope,
+        "flashinfer": flashinfer_rope,
-        "rope_v0": sglang_rope_v0,
+        "jit_pos_enc": sglang_pos_enc_rope,
-        "rope_v1": sglang_rope_v1,
+        "sgl_kernel": sgl_kernel_rope,
-        "rope_v2": sglang_rope_v2,
+        "jit_fused_rope": sglang_fused_rope,
     }
     fn = lambda: FN_MAP[provider](q, k, positions, is_neox)
     return run_benchmark(fn)
@@ -276,8 +281,12 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str)
 # Benchmark 2: RoPE + KV cache store
 # ---------------------------------------------------------------------------
 
-STORE_LINE_VALS = ["rope_v0_store", "rope_v1_store", "rope_v2_store"]
+STORE_LINE_VALS = ["jit_rope_then_store", "sgl_kernel_fused_store", "jit_fused_store"]
-STORE_LINE_NAMES = ["SGL RoPE v0 + Store", "SGL RoPE v1 + Store", "SGL RoPE v2 + Store"]
+STORE_LINE_NAMES = [
+    "SGL JIT RoPE + Store",
+    "sgl-kernel Fused RoPE + Store",
+    "SGL JIT Fused RoPE + Store",
+]
 STORE_STYLES = [("red", "-"), ("orange", "-"), ("blue", "--")]
 
 store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@@ -333,9 +342,9 @@ def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider
     torch.cuda.synchronize()
 
     FN_MAP = {
-        "rope_v0_store": rope_v0_store,
+        "jit_rope_then_store": jit_rope_then_store,
-        "rope_v1_store": rope_v1_store,
+        "sgl_kernel_fused_store": sgl_kernel_fused_rope_store,
-        "rope_v2_store": rope_v2_store,
+        "jit_fused_store": jit_fused_rope_store,
     }
     fn = lambda: FN_MAP[provider](
         q, k, v, k_cache, v_cache, positions, out_loc, is_neox

python/sglang/jit_kernel/benchmark/utils.py

@@ -40,3 +40,11 @@ def run_benchmark(
     quantiles = quantiles or DEFAULT_QUANTILES
     ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
     return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+def run_benchmark_no_cudagraph(
+    fn: Callable, quantiles: List[float] = None
+) -> Tuple[float, float, float]:
+    quantiles = quantiles or DEFAULT_QUANTILES
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms