chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Benchmark] use flashinfer bench_gpu_time instead of triton do_bench (#20305)

Browse Source
This commit is contained in:
Mook
2026-03-11 21:04:30 -07:00
committed by GitHub
parent ae7c2397b9
commit abc672e717
13 changed files with 85 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
View File

@@ -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(),

7
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
View File

@@ -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,

7
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
View File

@@ -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),

8
benchmark/kernels/elementwise/benchmark_concat_mla.py
View File

@@ -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

5
benchmark/kernels/fused_moe_triton/benchmark_sglang_fused_moe_triton.py
View File

@@ -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

5
benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py
View File

@@ -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,

5
benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py
View File

@@ -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,

9
benchmark/kernels/quantization/bench_fp4_quant.py
View File

@@ -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,

9
benchmark/kernels/quantization/bench_int8_quant.py
View File

@@ -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,
)

10
benchmark/kernels/scheduler_batch/benchmark_get_last_loc_triton.py
View File

@@ -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

14
benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py
View File

@@ -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

15
benchmark/kernels/sliding_window_attention_triton/bench_triton_swa_kernel.py
View File

@@ -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)

23
python/sglang/benchmark/bench_utils.py Normal file
View File

@@ -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)