diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py new file mode 100644 index 000000000..4da26f497 --- /dev/null +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -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) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 6fae78248..a7fbabdc0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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]