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

This commit is contained in:
Mook
2026-03-12 04:04:30 +00:00
committed by GitHub
parent ae7c2397b9
commit abc672e717
13 changed files with 85 additions and 41 deletions
@@ -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,
@@ -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,
)