minke.yu/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-12 04:04:30 +00: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
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py
+5 -4
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(),
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py
+4 -3
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,
benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py
+4 -3
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),