Add dsv3 router gemm benchmark on blackwell (#17707)

2026-04-04 01:18:01 -07:00
parent 82ea4906cf
commit 9fa12d605a
2 changed files with 284 additions and 4 deletions
@@ -0,0 +1,250 @@
+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)
@@ -165,7 +165,10 @@ if _use_aiter:
    pass

 if _is_cuda:
+    from flashinfer.gemm import mm_M1_16_K7168_N256 as _raw_dsv3_router_gemm
    from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm
+
+    from sglang.srt.utils.custom_op import register_custom_op
 elif _is_npu:
    from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
        forward_dsa_core_npu,
@@ -324,11 +327,20 @@ class MoEGate(nn.Module):
                and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
                and _device_sm >= 90
            ):
+                if _device_sm >= 100 and self.weight.shape[0] == 256:
+                    # router gemm output float32
+                    logits = torch.empty(
+                        hidden_states.shape[0],
+                        self.weight.shape[0],
+                        device=hidden_states.device,
+                        dtype=torch.float32,
+                    )
+                    flashinfer_dsv3_router_gemm(logits, hidden_states, self.weight)
+                else:
+                    logits = dsv3_router_gemm(
+                        hidden_states, self.weight, out_dtype=torch.float32
+                    )

-                # router gemm output float32
-                logits = dsv3_router_gemm(
-                    hidden_states, self.weight, out_dtype=torch.float32
-                )
            elif _use_aiter:
                logits = aiter_dsv3_router_gemm(hidden_states, self.weight)
            else:
@@ -2259,4 +2271,22 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
    pass


+@register_custom_op(
+    op_name="flashinfer_dsv3_router_gemm",
+    mutates_args=[],
+    fake_impl=lambda logits, hidden_states, weight: None,
+)
+def flashinfer_dsv3_router_gemm(
+    logits: torch.Tensor,
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+) -> None:
+    _raw_dsv3_router_gemm(
+        hidden_states,
+        weight.t(),
+        logits,
+        launch_with_pdl=True,
+    )
+
+
 EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]