From abc672e7177b3ec366c2791845447e3f121fdfc9 Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Wed, 11 Mar 2026 21:04:30 -0700 Subject: [PATCH] [Benchmark] use flashinfer bench_gpu_time instead of triton do_bench (#20305) --- .../deepseek/benchmark_deepgemm_fp8_gemm.py | 9 ++++---- .../benchmark_deepgemm_fp8_gemm_blackwell.py | 7 +++--- .../benchmark_deepgemm_fp8_group_gemm.py | 7 +++--- .../elementwise/benchmark_concat_mla.py | 8 +++---- .../benchmark_sglang_fused_moe_triton.py | 5 ++-- .../benchmark_torch_compile_fused_moe.py | 5 ++-- ...nchmark_vllm_vs_sglang_fused_moe_triton.py | 5 ++-- .../kernels/quantization/bench_fp4_quant.py | 9 ++++---- .../kernels/quantization/bench_int8_quant.py | 9 ++++---- .../benchmark_get_last_loc_triton.py | 10 ++++---- ...enchmark_write_req_to_token_pool_triton.py | 14 +++++------ .../bench_triton_swa_kernel.py | 15 ++++++++++-- python/sglang/benchmark/bench_utils.py | 23 +++++++++++++++++++ 13 files changed, 85 insertions(+), 41 deletions(-) create mode 100644 python/sglang/benchmark/bench_utils.py diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py index bd02e2aee..0b18e3bad 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import ( w8a8_block_fp8_matmul as vllm_w8a8_block_fp8_matmul, ) +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.fp8_kernel import ( w8a8_block_fp8_matmul_deepgemm as w8a8_block_fp8_matmul, ) @@ -303,10 +304,10 @@ def get_benchmark(tp_size): y_fp8, y_scale = per_block_cast_to_fp8(y) x_scale_col_major = get_mn_major_tma_aligned_tensor(x_scale.clone()) - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_deepgemm( x_fp8.clone(), x_scale_col_major.clone(), @@ -319,7 +320,7 @@ def get_benchmark(tp_size): quantiles=quantiles, ) elif provider == "sglang": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_sglang( x_fp8.clone(), x_scale.clone(), @@ -334,7 +335,7 @@ def get_benchmark(tp_size): else: # tilelang tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: tilelang_kernel( x_fp8.clone(), x_scale.clone(), diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py index de14bd90e..3257da7b3 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py @@ -6,6 +6,7 @@ import triton from deep_gemm import ceil_div from flashinfer.gemm import gemm_fp8_nt_groupwise +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, w8a8_block_fp8_matmul_deepgemm, @@ -195,10 +196,10 @@ def _benchmark(m, n, k, tp_size, provider): y_fp8, y_scale, [BLOCK_SIZE, BLOCK_SIZE] ) - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_deepgemm_blackwell( dg_x_fp8, dg_x_scale, @@ -208,7 +209,7 @@ def _benchmark(m, n, k, tp_size, provider): quantiles=quantiles, ) elif provider == "flashinfer": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_flashinfer( x_fp8, x_scale, diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py index b2cea0705..8b1be7b88 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -8,6 +8,7 @@ from deep_gemm import calc_diff from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # Import shared functionality from the regular GEMM benchmark +from sglang.benchmark.bench_utils import run_bench from sglang.benchmark.kernels.deepseek.benchmark_deepgemm_fp8_gemm import ( per_block_cast_to_fp8, per_token_cast_to_fp8, @@ -397,10 +398,10 @@ def get_benchmark(tp_size): .view(-1) ) - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) if provider == "deepgemm": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_group_deepgemm( x_fp8_grouped, y_fp8_grouped, @@ -420,7 +421,7 @@ def get_benchmark(tp_size): M, _ = a.shape _, N = b.shape c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: fp8_gemm_group_triton( (a, a_scale), (b, b_scale), diff --git a/benchmark/kernels/elementwise/benchmark_concat_mla.py b/benchmark/kernels/elementwise/benchmark_concat_mla.py index c4d7bb1c8..7bc51d3da 100644 --- a/benchmark/kernels/elementwise/benchmark_concat_mla.py +++ b/benchmark/kernels/elementwise/benchmark_concat_mla.py @@ -3,6 +3,8 @@ import triton import triton.language as tl from sgl_kernel import concat_mla_k as concat_mla_k_cuda +from sglang.benchmark.bench_utils import run_bench + DEVICE = triton.runtime.driver.active.get_active_torch_device() num_local_heads = 128 @@ -179,7 +181,7 @@ if not torch.all(output_ref == output_exp): ) def benchmark(num_tokens, provider): data = create_data(num_tokens=num_tokens) - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) fn = { "torch": fn_torch, "torch_compiled": fn_torch_compiled, @@ -187,9 +189,7 @@ def benchmark(num_tokens, provider): "hack_non_strided": fn_hack_non_strided, "cuda": fn_cuda, }[provider] - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: fn(**data), quantiles=quantiles - ) + ms, min_ms, max_ms = run_bench(lambda: fn(**data), quantiles=quantiles) return ms, min_ms, max_ms diff --git a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py index df2952b29..eaec7712a 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py +++ b/benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py @@ -5,6 +5,7 @@ import torch import triton from common_utils import get_model_config +from sglang.benchmark.bench_utils import run_bench from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, @@ -181,8 +182,8 @@ def benchmark( else: bench_lambda = lambda: api_func(**api_kwargs) - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench(bench_lambda, quantiles=quantiles) + quantiles = (0.5, 0.2, 0.8) + ms, min_ms, max_ms = run_bench(bench_lambda, quantiles=quantiles) return ms, min_ms, max_ms diff --git a/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py index 2b4faa24b..e94c66f4d 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py +++ b/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py @@ -6,6 +6,7 @@ import triton from torch.nn import functional as F from transformers import AutoConfig +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( fused_moe as fused_moe_triton, ) @@ -258,8 +259,8 @@ def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): ) torch.cuda.synchronize() - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench( + quantiles = (0.5, 0.2, 0.8) + ms, min_ms, max_ms = run_bench( lambda: api_func( x, w1, diff --git a/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py index 206ee2a86..24fb580dc 100644 --- a/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py +++ b/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py @@ -5,6 +5,7 @@ import torch import triton from vllm.model_executor.layers.fused_moe.fused_moe import fused_moe as fused_moe_vllm +from sglang.benchmark.bench_utils import run_bench from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, @@ -190,8 +191,8 @@ def benchmark(batch_size, provider, model_config, use_fp8_w8a8=False): ) torch.cuda.synchronize() - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench( + quantiles = (0.5, 0.2, 0.8) + ms, min_ms, max_ms = run_bench( lambda: api_func( x, w1, diff --git a/benchmark/kernels/quantization/bench_fp4_quant.py b/benchmark/kernels/quantization/bench_fp4_quant.py index afc12dd8d..9baedf407 100644 --- a/benchmark/kernels/quantization/bench_fp4_quant.py +++ b/benchmark/kernels/quantization/bench_fp4_quant.py @@ -9,6 +9,7 @@ from flashinfer import ( ) from sgl_kernel.elementwise import silu_and_mul +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_post_quant_fwd @@ -75,9 +76,9 @@ def benchmark(M, K, provider): dtype=torch.float32, ) - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) if provider == "triton_fp8": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = run_bench( lambda: silu_and_mul_masked_post_quant_fwd( x, fp8_out, @@ -89,7 +90,7 @@ def benchmark(M, K, provider): quantiles=quantiles, ) if provider == "cuda_unfused_fp4": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = run_bench( lambda: scaled_fp4_grouped_quantize( silu_and_mul(x), masks, @@ -98,7 +99,7 @@ def benchmark(M, K, provider): quantiles=quantiles, ) if provider == "cuda_fused_fp4": - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + ms, min_ms, max_ms = run_bench( lambda: silu_and_mul_scaled_nvfp4_experts_quantize( x, masks, diff --git a/benchmark/kernels/quantization/bench_int8_quant.py b/benchmark/kernels/quantization/bench_int8_quant.py index 94b795690..d40458ed9 100644 --- a/benchmark/kernels/quantization/bench_int8_quant.py +++ b/benchmark/kernels/quantization/bench_int8_quant.py @@ -4,6 +4,7 @@ import torch import triton from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 @@ -59,19 +60,19 @@ def benchmark(batch_size, provider): M, K = batch_size, 16384 x = torch.randn(M, K, dtype=torch.float16, device="cuda") * 1000 - quantiles = [0.5, 0.2, 0.8] + quantiles = (0.5, 0.2, 0.8) if provider == "vllm op": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: vllm_scaled_int8_quant(x, symmetric=True), quantiles=quantiles, ) if provider == "triton": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: per_token_quant_int8(x), quantiles=quantiles, ) if provider == "torch.compile": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: torch_int8_quant(x), quantiles=quantiles, ) diff --git a/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py b/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py index 3e17205e7..911cdb827 100644 --- a/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py +++ b/benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py @@ -4,6 +4,8 @@ import torch import triton import triton.language as tl +from sglang.benchmark.bench_utils import run_bench + @torch.compile(dynamic=True) def get_last_loc_torch( @@ -124,14 +126,14 @@ def get_benchmark(): quantiles = [0.5, 0.2, 0.8] if provider == "reference": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: get_last_loc_torch(req_to_token, req_pool_indices, pre_lens), - quantiles=quantiles, + quantiles=tuple(quantiles), ) elif provider == "triton": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: get_last_loc_triton(req_to_token, req_pool_indices, pre_lens), - quantiles=quantiles, + quantiles=tuple(quantiles), ) return 1000 * ms, 1000 * max_ms, 1000 * min_ms diff --git a/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py b/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py index 1ce43c8ba..561ff88ee 100644 --- a/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py +++ b/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py @@ -5,6 +5,8 @@ import torch import triton import triton.language as tl +from sglang.benchmark.bench_utils import run_bench + @triton.jit def write_req_to_token_pool_triton( @@ -263,7 +265,7 @@ def get_benchmark(): quantiles = [0.5, 0.2, 0.8] if provider == "reference": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: write_req_to_token_pool_reference( req_to_token.clone(), req_pool_indices, @@ -272,10 +274,10 @@ def get_benchmark(): extend_lens, out_cache_loc, ), - quantiles=quantiles, + quantiles=tuple(quantiles), ) elif provider == "triton": - ms, min_ms, max_ms = triton.testing.do_bench( + ms, min_ms, max_ms = run_bench( lambda: write_req_to_token_pool_triton[(batch_size,)]( req_to_token.clone(), req_pool_indices, @@ -285,7 +287,7 @@ def get_benchmark(): out_cache_loc, max_context_len, ), - quantiles=quantiles, + quantiles=tuple(quantiles), ) else: @@ -303,9 +305,7 @@ def get_benchmark(): BLOCK_SIZE=block_size, ) - ms, min_ms, max_ms = triton.testing.do_bench( - run_optimized, quantiles=quantiles - ) + ms, min_ms, max_ms = run_bench(run_optimized, quantiles=tuple(quantiles)) return 1000 * ms, 1000 * max_ms, 1000 * min_ms diff --git a/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py index 98144d470..9fd42fb12 100644 --- a/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py +++ b/benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py @@ -4,6 +4,7 @@ import torch import torch.nn.functional as F import triton.testing as tt +from sglang.benchmark.bench_utils import run_bench from sglang.srt.layers.attention.triton_ops.extend_attention import extend_attention_fwd @@ -270,9 +271,19 @@ def bench( raise AssertionError("Mismatch between triton and torch reference.") if provider == "triton": - ms = tt.do_bench(lambda: _run_triton(inputs), warmup=warmup, rep=rep) + ms = run_bench( + lambda: _run_triton(inputs), + quantiles=None, + warmup_ms=warmup, + rep_ms=rep, + )[0] elif provider == "torch": - ms = tt.do_bench(lambda: _run_torch_ref(inputs), warmup=warmup, rep=rep) + ms = run_bench( + lambda: _run_torch_ref(inputs), + quantiles=None, + warmup_ms=warmup, + rep_ms=rep, + )[0] else: raise ValueError(provider) diff --git a/python/sglang/benchmark/bench_utils.py b/python/sglang/benchmark/bench_utils.py new file mode 100644 index 000000000..958b93709 --- /dev/null +++ b/python/sglang/benchmark/bench_utils.py @@ -0,0 +1,23 @@ +"""Triton do_bench/do_bench_cudagraph compatible wrapper using flashinfer.testing.bench_gpu_time.""" + +import numpy as np +from flashinfer.testing import bench_gpu_time + + +def run_bench( + fn, + use_cuda_graph: bool = True, + quantiles=(0.5, 0.2, 0.8), + warmup_ms: int = 25, + rep_ms: int = 100, +): + """Returns (ms, min_ms, max_ms) or (median,) when quantiles=None.""" + times = bench_gpu_time( + fn=fn, + use_cuda_graph=use_cuda_graph, + dry_run_time_ms=warmup_ms, + repeat_time_ms=rep_ms, + ) + if quantiles is None: + return (float(np.median(times)),) + return tuple(float(np.percentile(times, q * 100)) for q in quantiles)