251 lines
7.6 KiB
Python
251 lines
7.6 KiB
Python
import argparse
|
|
import os
|
|
from typing import List
|
|
|
|
import torch
|
|
import triton
|
|
from flashinfer.gemm import mm_M1_16_K7168_N256
|
|
from sgl_kernel import dsv3_router_gemm
|
|
|
|
N = 256
|
|
K = 7168
|
|
|
|
|
|
def create_benchmark_configs(tp_sizes: List[int]):
|
|
configs = []
|
|
for tp_size in tp_sizes:
|
|
for m in range(1, 17):
|
|
configs.append((m, N, K, tp_size))
|
|
return configs
|
|
|
|
|
|
def dsv3_router_gemm_flashinfer(
|
|
hidden_states: torch.Tensor,
|
|
router_weights: torch.Tensor,
|
|
):
|
|
"""Flashinfer implementation of dsv3 router gemm"""
|
|
output = torch.empty(
|
|
hidden_states.shape[0],
|
|
router_weights.shape[0],
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
)
|
|
mm_M1_16_K7168_N256(
|
|
hidden_states, router_weights.t(), output, launch_with_pdl=args.use_pdl
|
|
)
|
|
return output
|
|
|
|
|
|
def dsv3_router_gemm_sgl(
|
|
hidden_states: torch.Tensor,
|
|
router_weights: torch.Tensor,
|
|
):
|
|
"""SGLang implementation of dsv3 router gemm"""
|
|
output = dsv3_router_gemm(
|
|
hidden_states,
|
|
router_weights,
|
|
out_dtype=torch.float32,
|
|
)
|
|
return output
|
|
|
|
|
|
def check_accuracy(a, b, atol, rtol, percent):
|
|
"""Unified accuracy checking function with detailed error reporting."""
|
|
if not torch.isfinite(a).all():
|
|
print("Non-finite values in reference output")
|
|
return False
|
|
if not torch.isfinite(b).all():
|
|
print("Non-finite values in actual output")
|
|
return False
|
|
assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}"
|
|
|
|
close = torch.isclose(a, b, atol=atol, rtol=rtol)
|
|
match_ratio = close.float().mean()
|
|
if match_ratio >= percent:
|
|
return True
|
|
|
|
mismatch_percent = 1.0 - match_ratio.item()
|
|
if mismatch_percent > 1 - percent:
|
|
print(
|
|
f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
|
|
f"(threshold: {1 - percent:.4f})"
|
|
)
|
|
return False
|
|
|
|
|
|
def calculate_diff(m: int, n: int, k: int):
|
|
hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
|
|
router_weights = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
|
|
|
|
out_flashinfer = dsv3_router_gemm_flashinfer(
|
|
hidden_states.clone(memory_format=torch.contiguous_format),
|
|
router_weights.clone(memory_format=torch.contiguous_format),
|
|
)
|
|
|
|
out_sgl = dsv3_router_gemm_sgl(
|
|
hidden_states.clone(memory_format=torch.contiguous_format),
|
|
router_weights.clone(memory_format=torch.contiguous_format),
|
|
)
|
|
|
|
print(f"Shape m={m}, n={n}, k={k}:")
|
|
print(f"Using PDL={args.use_pdl}")
|
|
print(f"Flashinfer output: {out_flashinfer[0, 0:5]}")
|
|
print(f"SGLang output: {out_sgl[0, 0:5]}")
|
|
|
|
flashinfer_sgl_match = check_accuracy(out_flashinfer, out_sgl, 0.1, 0.6, 0.95)
|
|
print("Correctness check:")
|
|
print(f" - Flashinfer vs SGLang: {'✅' if flashinfer_sgl_match else '❌'}")
|
|
|
|
|
|
def _benchmark(m, n, k, tp_size, provider):
|
|
print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}")
|
|
hidden_states = torch.randn(
|
|
(m, k), device="cuda", dtype=torch.bfloat16
|
|
).contiguous()
|
|
router_weights = torch.randn(
|
|
(n, k), device="cuda", dtype=torch.bfloat16
|
|
).contiguous()
|
|
|
|
quantiles = [0.5, 0.2, 0.8]
|
|
|
|
if provider == "sglang":
|
|
ms, min_ms, max_ms = triton.testing.do_bench(
|
|
lambda: dsv3_router_gemm_sgl(
|
|
hidden_states.clone(memory_format=torch.contiguous_format),
|
|
router_weights.clone(memory_format=torch.contiguous_format),
|
|
),
|
|
quantiles=quantiles,
|
|
)
|
|
elif provider == "flashinfer":
|
|
ms, min_ms, max_ms = triton.testing.do_bench(
|
|
lambda: dsv3_router_gemm_flashinfer(
|
|
hidden_states.clone(memory_format=torch.contiguous_format),
|
|
router_weights.clone(memory_format=torch.contiguous_format),
|
|
),
|
|
quantiles=quantiles,
|
|
)
|
|
|
|
# Calculate TFLOPS
|
|
flops = 2 * m * n * k # multiply-adds
|
|
tflops = flops / (ms * 1e-3) / 1e12
|
|
|
|
# Print shape-specific results with TFLOPS
|
|
print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}")
|
|
return ms, max_ms, min_ms
|
|
|
|
|
|
def get_benchmark_plot_friendly(tp_sizes):
|
|
all_configs = create_benchmark_configs(tp_sizes)
|
|
x_vals = list(range(len(all_configs)))
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=["cfg_id"],
|
|
x_vals=x_vals,
|
|
line_arg="provider",
|
|
line_vals=["sglang", "flashinfer"],
|
|
line_names=["SGLang", "Flashinfer"],
|
|
styles=[("blue", "-"), ("red", "-")],
|
|
ylabel="us",
|
|
plot_name=f"fp8-gemm-performance-comparison-tp-{'-'.join(str(tp) for tp in tp_sizes)}",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark(cfg_id, provider):
|
|
m, n, k, tp_size = all_configs[cfg_id]
|
|
ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider)
|
|
return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
|
|
|
|
return benchmark
|
|
|
|
|
|
def get_benchmark(tp_sizes):
|
|
all_configs = create_benchmark_configs(tp_sizes)
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=[
|
|
"m",
|
|
"n",
|
|
"k",
|
|
"tp_size",
|
|
],
|
|
x_vals=[list(config) for config in all_configs],
|
|
line_arg="provider",
|
|
line_vals=["sglang", "flashinfer"],
|
|
line_names=["SGLang", "Flashinfer"],
|
|
styles=[("blue", "-"), ("red", "-")],
|
|
ylabel="us",
|
|
plot_name=f"fp8-gemm-performance-comparison-tp-{'-'.join(str(tp) for tp in tp_sizes)}",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark(m, n, k, tp_size, provider):
|
|
ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider)
|
|
return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms
|
|
|
|
return benchmark
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
|
|
print("Skipping benchmark because the device is not supported")
|
|
exit(0)
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--save-path",
|
|
type=str,
|
|
default="./configs/benchmark_ops/dsv3_router_gemm/",
|
|
help="Path to save dsv3 router gemm benchmark results",
|
|
)
|
|
parser.add_argument(
|
|
"--run-correctness",
|
|
action="store_true",
|
|
default=True,
|
|
help="Whether to run correctness test",
|
|
)
|
|
parser.add_argument(
|
|
"--tp-sizes",
|
|
type=int,
|
|
nargs="+",
|
|
default=[1],
|
|
help="List of tensor parallelism sizes to benchmark",
|
|
)
|
|
parser.add_argument(
|
|
"--plot-friendly",
|
|
action="store_true",
|
|
default=False,
|
|
help="Plot x axis as the config index instead of the m",
|
|
)
|
|
parser.add_argument(
|
|
"--use-pdl",
|
|
action="store_true",
|
|
default=False,
|
|
help="Use PDL if true.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Set random seed for reproducibility
|
|
torch.manual_seed(0)
|
|
torch.cuda.manual_seed(0)
|
|
|
|
if args.use_pdl:
|
|
os.environ["TRTLLM_ENABLE_PDL"] = "1"
|
|
|
|
# Run correctness tests on a few examples
|
|
if args.run_correctness:
|
|
print("Running correctness tests...")
|
|
for m, n, k, _ in create_benchmark_configs(args.tp_sizes):
|
|
calculate_diff(m, n, k)
|
|
|
|
# Get the benchmark function with the specified tp_size
|
|
benchmark = (
|
|
get_benchmark_plot_friendly(args.tp_sizes)
|
|
if args.plot_friendly
|
|
else get_benchmark(args.tp_sizes)
|
|
)
|
|
|
|
print(f"Running performance benchmark for TP sizes = {args.tp_sizes}...")
|
|
benchmark.run(print_data=True, save_path=args.save_path)
|