[Benchmark] use flashinfer bench_gpu_time instead of triton do_bench (#20305)
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user