From 3484058f97ab664798436950d242922e9716215e Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sun, 25 Jan 2026 10:10:37 +0000 Subject: [PATCH 01/23] Add accuracy test for dsv3_router_gemm kernel --- ...ark_deepgemm_dsv3_router_gemm_blackwell.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py 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 000000000000..d00a2bd8fbc5 --- /dev/null +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -0,0 +1,136 @@ +import argparse + +import torch +from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 + +from sgl_kernel import dsv3_router_gemm + +def dsv3_router_gemm_flashinfer( + hidden_states: torch.Tensor, + router_weights: torch.Tensor, + launch_with_pdl=False +): + """Flashinfer implementation of dsv3 router gemm""" + num_tokens, num_experts = hidden_states.shape[0], router_weights.shape[0] + # output = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.bfloat16) + output = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32).contiguous() + + print(f"hidden_states.shape: {hidden_states.shape}") + print(f"num_tokens: {num_tokens}, num_experts: {num_experts}") + print(f"router_weights.shape: {router_weights.shape}") + print(f"output.shape: {output.shape}") + + mm_M1_16_K7168_N256( + hidden_states, + router_weights.t(), + output, + launch_with_pdl=launch_with_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(num_tokens: int, num_experts: int, hidden_dim: int): + hidden_states = torch.randn((num_tokens, hidden_dim), device="cuda", dtype=torch.bfloat16).contiguous() + router_weights = torch.randn((num_experts, hidden_dim), device="cuda", dtype=torch.bfloat16).contiguous() + + out_flashinfer = dsv3_router_gemm_flashinfer( + hidden_states, + router_weights, + False, + ) + + out_sgl = dsv3_router_gemm_sgl( + hidden_states, + router_weights, + ) + + print(f"Shape m={num_tokens}, n={num_experts}, k={hidden_dim}:") + print(f"Flashinfer output: {out_flashinfer[0, 0:5]}") + print(f"DeepGEMM output: {out_sgl[0, 0:5]}") + + flashinfer_deepgemm_match = check_accuracy( + out_flashinfer, out_sgl, 0.1, 0.6, 0.95 + ) + print("Correctness check:") + print(f" - Flashinfer vs DeepGEMM: {'✅' if flashinfer_deepgemm_match else '❌'}") + + +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-size", + type=int, + default=1, + help="Tensor parallelism size to benchmark (default: 1)", + ) + parser.add_argument( + "--plot-friendly", + action="store_true", + default=False, + help="Plot x axis as the config index instead of the m", + ) + args = parser.parse_args() + + # Set random seed for reproducibility + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # Run correctness tests on a few examples + if args.run_correctness: + print("Running correctness tests...") + calculate_diff(1, 256, 7168) # Small test + calculate_diff(8, 256, 7168) # Medium test + calculate_diff(16, 256, 7168) # Large test + From d1b7e7215e3cfb195f28b9996547f5df0c632356 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sun, 25 Jan 2026 12:26:48 +0000 Subject: [PATCH 02/23] Add performance benchmark for dsv3_rotuer_gemm --- ...ark_deepgemm_dsv3_router_gemm_blackwell.py | 151 ++++++++++++++---- 1 file changed, 124 insertions(+), 27 deletions(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index d00a2bd8fbc5..a3f3e43ea2c9 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -1,24 +1,29 @@ import argparse -import torch +import torch +import triton from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 -from sgl_kernel import dsv3_router_gemm +from sgl_kernel import dsv3_router_gemm as dsv3_router_gemm + +N = 256 +K = 7168 + +def create_benchmark_configs(tp_size): + configs = [] + for launch_with_pdl in [False, True]: + for m in range(1, 17): + configs.append((m, N, K, tp_size, launch_with_pdl)) + return configs + def dsv3_router_gemm_flashinfer( hidden_states: torch.Tensor, router_weights: torch.Tensor, - launch_with_pdl=False + launch_with_pdl: bool, ): """Flashinfer implementation of dsv3 router gemm""" - num_tokens, num_experts = hidden_states.shape[0], router_weights.shape[0] - # output = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.bfloat16) - output = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32).contiguous() - - print(f"hidden_states.shape: {hidden_states.shape}") - print(f"num_tokens: {num_tokens}, num_experts: {num_experts}") - print(f"router_weights.shape: {router_weights.shape}") - print(f"output.shape: {output.shape}") + output = torch.randn(hidden_states.shape[0], router_weights.shape[0], device="cuda", dtype=torch.float32).contiguous() mm_M1_16_K7168_N256( hidden_states, @@ -65,30 +70,114 @@ def check_accuracy(a, b, atol, rtol, percent): ) return False -def calculate_diff(num_tokens: int, num_experts: int, hidden_dim: int): - hidden_states = torch.randn((num_tokens, hidden_dim), device="cuda", dtype=torch.bfloat16).contiguous() - router_weights = torch.randn((num_experts, hidden_dim), device="cuda", dtype=torch.bfloat16).contiguous() + +def calculate_diff(m: int, n: int, k: int, launch_with_pdl: bool): + 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, - router_weights, - False, + hidden_states.clone(memory_format=torch.contiguous_format), + router_weights.clone(memory_format=torch.contiguous_format), + launch_with_pdl, ) out_sgl = dsv3_router_gemm_sgl( - hidden_states, - router_weights, + hidden_states.clone(memory_format=torch.contiguous_format), + router_weights.clone(memory_format=torch.contiguous_format), ) - print(f"Shape m={num_tokens}, n={num_experts}, k={hidden_dim}:") + print(f"Shape m={m}, n={n}, k={k}:") + print(f"Using launch_with_pdl={launch_with_pdl} for flashinfer") print(f"Flashinfer output: {out_flashinfer[0, 0:5]}") - print(f"DeepGEMM output: {out_sgl[0, 0:5]}") + print(f"SGLang output: {out_sgl[0, 0:5]}") - flashinfer_deepgemm_match = check_accuracy( + flashinfer_sgl_match = check_accuracy( out_flashinfer, out_sgl, 0.1, 0.6, 0.95 ) print("Correctness check:") - print(f" - Flashinfer vs DeepGEMM: {'✅' if flashinfer_deepgemm_match else '❌'}") + print(f" - Flashinfer vs SGLang: {'✅' if flashinfer_sgl_match else '❌'}") + + +def _benchmark(m, n, k, tp_size, launch_with_pdl, provider): + print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), launch_with_pdl={launch_with_pdl}, 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), + launch_with_pdl, + ), + 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_size): + all_configs = create_benchmark_configs(tp_size) + 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{tp_size}", + args={}, + ) + ) + def benchmark(cfg_id, provider): + m, n, k, tp_size, launch_with_pdl= all_configs[cfg_id] + ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, provider) + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark + + +def get_benchmark(tp_size): + all_configs = create_benchmark_configs(tp_size) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k", "tp_size", "launch_with_pdl", ], + 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{tp_size}", + args={}, + ) + ) + def benchmark(m, n, k, tp_size, launch_with_pdl, provider): + ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, provider) + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark if __name__ == "__main__": @@ -129,8 +218,16 @@ def calculate_diff(num_tokens: int, num_experts: int, hidden_dim: int): # Run correctness tests on a few examples if args.run_correctness: - print("Running correctness tests...") - calculate_diff(1, 256, 7168) # Small test - calculate_diff(8, 256, 7168) # Medium test - calculate_diff(16, 256, 7168) # Large test + print("Running correctness tests...") + for m, n, k, _, launch_with_pdl in create_benchmark_configs(args.tp_size): + calculate_diff(m, n, k, launch_with_pdl) + + # Get the benchmark function with the specified tp_size + benchmark = ( + get_benchmark_plot_friendly(args.tp_size) + if args.plot_friendly + else get_benchmark(args.tp_size) + ) + print(f"Running performance benchmark for TP size = {args.tp_size}...") + benchmark.run(print_data=True, save_path=args.save_path) From dd79d06714d4a02373dfc8e0ec26a639b530ce38 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sun, 25 Jan 2026 12:58:08 +0000 Subject: [PATCH 03/23] Update the script to allow iterating tp size list --- ...ark_deepgemm_dsv3_router_gemm_blackwell.py | 35 ++++++++++--------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index a3f3e43ea2c9..2544b45db6e4 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -1,4 +1,5 @@ import argparse +from typing import List import torch import triton @@ -9,11 +10,12 @@ N = 256 K = 7168 -def create_benchmark_configs(tp_size): +def create_benchmark_configs(tp_sizes: List[int]): configs = [] for launch_with_pdl in [False, True]: - for m in range(1, 17): - configs.append((m, N, K, tp_size, launch_with_pdl)) + for tp_size in tp_sizes: + for m in range(1, 17): + configs.append((m, N, K, tp_size, launch_with_pdl)) return configs @@ -132,8 +134,8 @@ def _benchmark(m, n, k, tp_size, launch_with_pdl, provider): return ms, max_ms, min_ms -def get_benchmark_plot_friendly(tp_size): - all_configs = create_benchmark_configs(tp_size) +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( @@ -145,7 +147,7 @@ def get_benchmark_plot_friendly(tp_size): line_names=["SGLang", "Flashinfer"], styles=[("blue", "-"), ("red", "-")], ylabel="us", - plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", + plot_name=f"fp8-gemm-performance-comparison-tp-{"-".join(str(tp) for tp in tp_sizes)}", args={}, ) ) @@ -157,8 +159,8 @@ def benchmark(cfg_id, provider): return benchmark -def get_benchmark(tp_size): - all_configs = create_benchmark_configs(tp_size) +def get_benchmark(tp_sizes): + all_configs = create_benchmark_configs(tp_sizes) @triton.testing.perf_report( triton.testing.Benchmark( @@ -169,7 +171,7 @@ def get_benchmark(tp_size): line_names=["SGLang", "Flashinfer"], styles=[("blue", "-"), ("red", "-")], ylabel="us", - plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", + plot_name=f"fp8-gemm-performance-comparison-tp-{"-".join(str(tp) for tp in tp_sizes)}", args={}, ) ) @@ -199,10 +201,11 @@ def benchmark(m, n, k, tp_size, launch_with_pdl, provider): help="Whether to run correctness test", ) parser.add_argument( - "--tp-size", + "--tp-sizes", type=int, - default=1, - help="Tensor parallelism size to benchmark (default: 1)", + nargs='+', + default=[1], + help="List of tensor parallelism sizes to benchmark", ) parser.add_argument( "--plot-friendly", @@ -219,15 +222,15 @@ def benchmark(m, n, k, tp_size, launch_with_pdl, provider): # Run correctness tests on a few examples if args.run_correctness: print("Running correctness tests...") - for m, n, k, _, launch_with_pdl in create_benchmark_configs(args.tp_size): + for m, n, k, _, launch_with_pdl in create_benchmark_configs(args.tp_sizes): calculate_diff(m, n, k, launch_with_pdl) # Get the benchmark function with the specified tp_size benchmark = ( - get_benchmark_plot_friendly(args.tp_size) + get_benchmark_plot_friendly(args.tp_sizes) if args.plot_friendly - else get_benchmark(args.tp_size) + else get_benchmark(args.tp_sizes) ) - print(f"Running performance benchmark for TP size = {args.tp_size}...") + print(f"Running performance benchmark for TP sizes = {args.tp_sizes}...") benchmark.run(print_data=True, save_path=args.save_path) From 8691db2fe6cc8fd95a4e4e31c3f44f4ca4483c01 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sun, 25 Jan 2026 13:14:51 +0000 Subject: [PATCH 04/23] Updat dsv3_router_gemm benchmark script --- ...ark_deepgemm_dsv3_router_gemm_blackwell.py | 48 ++++++++++++------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index 2544b45db6e4..d2756760554f 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -3,13 +3,13 @@ import torch import triton -from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 - +from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 from sgl_kernel import dsv3_router_gemm as dsv3_router_gemm N = 256 K = 7168 + def create_benchmark_configs(tp_sizes: List[int]): configs = [] for launch_with_pdl in [False, True]: @@ -25,13 +25,15 @@ def dsv3_router_gemm_flashinfer( launch_with_pdl: bool, ): """Flashinfer implementation of dsv3 router gemm""" - output = torch.randn(hidden_states.shape[0], router_weights.shape[0], device="cuda", dtype=torch.float32).contiguous() + output = torch.randn( + hidden_states.shape[0], + router_weights.shape[0], + device="cuda", + dtype=torch.float32, + ).contiguous() mm_M1_16_K7168_N256( - hidden_states, - router_weights.t(), - output, - launch_with_pdl=launch_with_pdl + hidden_states, router_weights.t(), output, launch_with_pdl=launch_with_pdl ) return output @@ -93,17 +95,21 @@ def calculate_diff(m: int, n: int, k: int, launch_with_pdl: bool): 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 - ) + 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, launch_with_pdl, provider): - print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), launch_with_pdl={launch_with_pdl}, 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() + print( + f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), launch_with_pdl={launch_with_pdl}, 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] @@ -152,7 +158,7 @@ def get_benchmark_plot_friendly(tp_sizes): ) ) def benchmark(cfg_id, provider): - m, n, k, tp_size, launch_with_pdl= all_configs[cfg_id] + m, n, k, tp_size, launch_with_pdl = all_configs[cfg_id] ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, provider) return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms @@ -164,7 +170,13 @@ def get_benchmark(tp_sizes): @triton.testing.perf_report( triton.testing.Benchmark( - x_names=["m", "n", "k", "tp_size", "launch_with_pdl", ], + x_names=[ + "m", + "n", + "k", + "tp_size", + "launch_with_pdl", + ], x_vals=[list(config) for config in all_configs], line_arg="provider", line_vals=["sglang", "flashinfer"], @@ -203,7 +215,7 @@ def benchmark(m, n, k, tp_size, launch_with_pdl, provider): parser.add_argument( "--tp-sizes", type=int, - nargs='+', + nargs="+", default=[1], help="List of tensor parallelism sizes to benchmark", ) @@ -221,9 +233,9 @@ def benchmark(m, n, k, tp_size, launch_with_pdl, provider): # Run correctness tests on a few examples if args.run_correctness: - print("Running correctness tests...") + print("Running correctness tests...") for m, n, k, _, launch_with_pdl in create_benchmark_configs(args.tp_sizes): - calculate_diff(m, n, k, launch_with_pdl) + calculate_diff(m, n, k, launch_with_pdl) # Get the benchmark function with the specified tp_size benchmark = ( From fe0dc0acb39cb1ea17ca95da83cc06f1fa1e7fb9 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Wed, 4 Feb 2026 10:54:20 +0000 Subject: [PATCH 05/23] Update to use PDL for sglang-kernel using env var --- ...ark_deepgemm_dsv3_router_gemm_blackwell.py | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index d2756760554f..4423a2346ad5 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -1,4 +1,5 @@ import argparse +import os from typing import List import torch @@ -12,28 +13,25 @@ def create_benchmark_configs(tp_sizes: List[int]): configs = [] - for launch_with_pdl in [False, True]: - for tp_size in tp_sizes: - for m in range(1, 17): - configs.append((m, N, K, tp_size, launch_with_pdl)) + 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, - launch_with_pdl: bool, ): """Flashinfer implementation of dsv3 router gemm""" - output = torch.randn( + output = torch.empty( hidden_states.shape[0], router_weights.shape[0], device="cuda", dtype=torch.float32, - ).contiguous() - + ) mm_M1_16_K7168_N256( - hidden_states, router_weights.t(), output, launch_with_pdl=launch_with_pdl + hidden_states, router_weights.t(), output, launch_with_pdl=args.use_pdl ) return output @@ -75,14 +73,13 @@ def check_accuracy(a, b, atol, rtol, percent): return False -def calculate_diff(m: int, n: int, k: int, launch_with_pdl: bool): +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), - launch_with_pdl, ) out_sgl = dsv3_router_gemm_sgl( @@ -91,7 +88,7 @@ def calculate_diff(m: int, n: int, k: int, launch_with_pdl: bool): ) print(f"Shape m={m}, n={n}, k={k}:") - print(f"Using launch_with_pdl={launch_with_pdl} for flashinfer") + 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]}") @@ -100,10 +97,8 @@ def calculate_diff(m: int, n: int, k: int, launch_with_pdl: bool): print(f" - Flashinfer vs SGLang: {'✅' if flashinfer_sgl_match else '❌'}") -def _benchmark(m, n, k, tp_size, launch_with_pdl, provider): - print( - f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), launch_with_pdl={launch_with_pdl}, Provider: {provider}" - ) +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() @@ -126,7 +121,6 @@ def _benchmark(m, n, k, tp_size, launch_with_pdl, provider): lambda: dsv3_router_gemm_flashinfer( hidden_states.clone(memory_format=torch.contiguous_format), router_weights.clone(memory_format=torch.contiguous_format), - launch_with_pdl, ), quantiles=quantiles, ) @@ -158,8 +152,8 @@ def get_benchmark_plot_friendly(tp_sizes): ) ) def benchmark(cfg_id, provider): - m, n, k, tp_size, launch_with_pdl = all_configs[cfg_id] - ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, 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 @@ -175,7 +169,6 @@ def get_benchmark(tp_sizes): "n", "k", "tp_size", - "launch_with_pdl", ], x_vals=[list(config) for config in all_configs], line_arg="provider", @@ -187,8 +180,8 @@ def get_benchmark(tp_sizes): args={}, ) ) - def benchmark(m, n, k, tp_size, launch_with_pdl, provider): - ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, provider) + 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 @@ -225,17 +218,26 @@ def benchmark(m, n, k, tp_size, launch_with_pdl, provider): 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, _, launch_with_pdl in create_benchmark_configs(args.tp_sizes): - calculate_diff(m, n, k, launch_with_pdl) + 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 = ( From 42bc2e99ee1c9bef73582d793e644c8ad1f16097 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sat, 7 Feb 2026 08:38:15 +0000 Subject: [PATCH 06/23] Update to use flashinfer mm_M1_16_K7168_N256 for deepseekv2 --- python/sglang/srt/models/deepseek_v2.py | 25 +++++++++++++++++-------- sglang | 1 + 2 files changed, 18 insertions(+), 8 deletions(-) create mode 160000 sglang diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 988ed91e831d..b3a9f0da10f0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -177,7 +177,8 @@ ) if _is_cuda: - from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm + from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm + from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: @@ -344,9 +345,19 @@ def forward( ): # router gemm output float32 - logits = dsv3_router_gemm( - hidden_states, self.weight, out_dtype=torch.float32 + logits = torch.empty( + hidden_states.shape[0], + self.weight.shape[0], + device="cuda", + dtype=torch.float32, ) + mm_M1_16_K7168_N256( + hidden_states, + self.weight.t(), + logits, + launch_with_pdl=os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0" + ) + elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator @@ -2844,7 +2855,7 @@ def determine_num_fused_shared_experts( or self.config.n_routed_experts != 256 or self.config.n_shared_experts != 1 ): - disable_reason = "Config does not support fused shared expert(s)." + disable_reason = "Config not support fused shared expert(s)." elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): @@ -2856,10 +2867,8 @@ def determine_num_fused_shared_experts( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism." - elif disable_reason is None and ( - get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori() - ): - disable_reason = "Deepseek V3/R1 cannot use shared experts fusion optimization under deepep expert parallelism." + elif disable_reason is None and get_moe_a2a_backend().is_deepep(): + disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under deepep expert parallelism." elif self.quant_config and self.quant_config.get_name() == "w4afp8": disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts." diff --git a/sglang b/sglang new file mode 160000 index 000000000000..1552aab741bf --- /dev/null +++ b/sglang @@ -0,0 +1 @@ +Subproject commit 1552aab741bf89c241b10e06297ee657597bd758 From 1c3d9877aa776b3fd3d0123e90e6a840ae60ebf5 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sat, 7 Feb 2026 08:39:30 +0000 Subject: [PATCH 07/23] Revert unrelated changes --- python/sglang/srt/models/deepseek_v2.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b3a9f0da10f0..8613acde367b 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2855,7 +2855,7 @@ def determine_num_fused_shared_experts( or self.config.n_routed_experts != 256 or self.config.n_shared_experts != 1 ): - disable_reason = "Config not support fused shared expert(s)." + disable_reason = "Config does not support fused shared expert(s)." elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): @@ -2867,8 +2867,10 @@ def determine_num_fused_shared_experts( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism." - elif disable_reason is None and get_moe_a2a_backend().is_deepep(): - disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under deepep expert parallelism." + elif disable_reason is None and ( + get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori() + ): + disable_reason = "Deepseek V3/R1 cannot use shared experts fusion optimization under deepep expert parallelism." elif self.quant_config and self.quant_config.get_name() == "w4afp8": disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts." From bfcbaeb40b1e7ef8eaa3f788a60d5c26bb42d454 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Sat, 7 Feb 2026 08:47:49 +0000 Subject: [PATCH 08/23] formatting --- python/sglang/srt/models/deepseek_v2.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8613acde367b..f8e36715594a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -356,8 +356,7 @@ def forward( self.weight.t(), logits, launch_with_pdl=os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0" - ) - + ) elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator From 9d2797a23c9605e427cb1a661fe1f5780040b4f2 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Wed, 11 Feb 2026 16:39:28 +0000 Subject: [PATCH 09/23] Update the kernel usage based on weight tensor shape --- python/sglang/srt/models/deepseek_v2.py | 37 ++++++++++++++----------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 4696c937202a..adecf1a823ac 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -177,7 +177,7 @@ ) if _is_cuda: - from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm + from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 elif _is_cpu and _is_cpu_amx_available: pass @@ -340,23 +340,28 @@ def forward( _is_cuda and hidden_states.shape[0] <= 16 and hidden_states.shape[1] == 7168 - and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384) and _device_sm >= 90 ): - - # router gemm output float32 - logits = torch.empty( - hidden_states.shape[0], - self.weight.shape[0], - device="cuda", - dtype=torch.float32, - ) - mm_M1_16_K7168_N256( - hidden_states, - self.weight.t(), - logits, - launch_with_pdl=os.environ.get("TRTLLM_ENABLE_PDL", "1") != "0" - ) + if self.weight.shape[0] == 384: + logits = dsv3_router_gemm( + hidden_states, + self.weight, + out_dtype=torch.float32 + ) + elif 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, + ) + # Note: PDL is disabled by default. + mm_M1_16_K7168_N256( + hidden_states, + self.weight.t(), + logits, + ) elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator From 41b9cbde702cbceccbbeba468573d39da0d47e01 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Wed, 11 Feb 2026 23:17:32 +0000 Subject: [PATCH 10/23] Update --- python/sglang/srt/models/deepseek_v2.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index adecf1a823ac..b1251b65a96b 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -342,13 +342,7 @@ def forward( and hidden_states.shape[1] == 7168 and _device_sm >= 90 ): - if self.weight.shape[0] == 384: - logits = dsv3_router_gemm( - hidden_states, - self.weight, - out_dtype=torch.float32 - ) - elif self.weight.shape[0] == 256: + if self.weight.shape[0] == 256: # router gemm output float32 logits = torch.empty( hidden_states.shape[0], @@ -362,6 +356,12 @@ def forward( self.weight.t(), logits, ) + elif self.weight.shape[0] == 384: + logits = dsv3_router_gemm( + hidden_states, + self.weight, + out_dtype=torch.float32 + ) elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator From f668bc5f66dcca2853073bec0be03c2ff1dfdf28 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Wed, 11 Feb 2026 23:29:04 +0000 Subject: [PATCH 11/23] remove sglang subfolder --- sglang | 1 - 1 file changed, 1 deletion(-) delete mode 160000 sglang diff --git a/sglang b/sglang deleted file mode 160000 index 1552aab741bf..000000000000 --- a/sglang +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 1552aab741bf89c241b10e06297ee657597bd758 From 77ff69f76434819ab279bc33f353f81e4d83bbe6 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Mon, 23 Feb 2026 03:45:01 +0000 Subject: [PATCH 12/23] Turn on PDL by default with sglang enviorn --- python/sglang/srt/environ.py | 1 + python/sglang/srt/models/deepseek_v2.py | 1 + 2 files changed, 2 insertions(+) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2308b212ccf7..13c7513c3d1c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -328,6 +328,7 @@ class Envs: # Default to the pick from flashinfer SGLANG_FLASHINFER_FP4_GEMM_BACKEND = EnvStr("") SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) + SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL = EnvBool(True) # Triton SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index a94fdc7868e9..9612fa9fd611 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -362,6 +362,7 @@ def forward( hidden_states, self.weight.t(), logits, + launch_with_pdl=envs.SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL.get() ) elif self.weight.shape[0] == 384: logits = dsv3_router_gemm( From 8356ed3747c2ea75010bc1336fea0a404ff6f34d Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Mon, 23 Feb 2026 03:45:41 +0000 Subject: [PATCH 13/23] Remove oudated comment --- python/sglang/srt/models/deepseek_v2.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 9612fa9fd611..b6fc4bab0b21 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -357,7 +357,6 @@ def forward( device=hidden_states.device, dtype=torch.float32, ) - # Note: PDL is disabled by default. mm_M1_16_K7168_N256( hidden_states, self.weight.t(), From 6f299252a227d42c21c3a7512a525693a48eb413 Mon Sep 17 00:00:00 2001 From: harrisonlimh Date: Mon, 23 Feb 2026 03:50:49 +0000 Subject: [PATCH 14/23] lint --- python/sglang/srt/models/deepseek_v2.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b6fc4bab0b21..1280f41f64fd 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -180,8 +180,8 @@ ) if _is_cuda: - from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 + from sgl_kernel import bmm_fp8, dsv3_fused_a_gemm, dsv3_router_gemm elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: @@ -361,14 +361,12 @@ def forward( hidden_states, self.weight.t(), logits, - launch_with_pdl=envs.SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL.get() - ) + launch_with_pdl=envs.SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL.get(), + ) elif self.weight.shape[0] == 384: logits = dsv3_router_gemm( - hidden_states, - self.weight, - out_dtype=torch.float32 - ) + hidden_states, self.weight, out_dtype=torch.float32 + ) elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator From 7d45a2ce510d5384a1bbf1e3a8543d2b6013e6d0 Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:19:06 -0800 Subject: [PATCH 15/23] Update environ.py --- python/sglang/srt/environ.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 13c7513c3d1c..2308b212ccf7 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -328,7 +328,6 @@ class Envs: # Default to the pick from flashinfer SGLANG_FLASHINFER_FP4_GEMM_BACKEND = EnvStr("") SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) - SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL = EnvBool(True) # Triton SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False) From 52e64d8166e1c1ec6e437521e78d939f8f5b2a9d Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:19:34 -0800 Subject: [PATCH 16/23] Update deepseek_v2.py --- python/sglang/srt/models/deepseek_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 1280f41f64fd..786d18c03773 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -361,7 +361,7 @@ def forward( hidden_states, self.weight.t(), logits, - launch_with_pdl=envs.SGLANG_FLASHINFER_DSV3_ROTUER_GEMM_LAUNCH_WITH_PDL.get(), + launch_with_pdl=True, ) elif self.weight.shape[0] == 384: logits = dsv3_router_gemm( From 78da480b9bc08b4c8f4f4b03652f53328e840d04 Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:23:01 -0800 Subject: [PATCH 17/23] Update benchmark_deepgemm_dsv3_router_gemm_blackwell.py --- .../deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index 4423a2346ad5..7bce73ca2c3b 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -5,7 +5,7 @@ import torch import triton from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 -from sgl_kernel import dsv3_router_gemm as dsv3_router_gemm +from sgl_kernel import dsv3_router_gemm N = 256 K = 7168 From 5062d3cb39146359e294b6ee99bee2e7ed25e534 Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Tue, 3 Mar 2026 17:24:48 -0800 Subject: [PATCH 18/23] Update deepseek_v2.py --- python/sglang/srt/models/deepseek_v2.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index f69b98691461..0e48341e0006 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -176,7 +176,6 @@ else: pass - logger = logging.getLogger(__name__) From 977f80604c5152be2cbcfd3764d40b3c20942e75 Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Mon, 9 Mar 2026 16:57:01 -0700 Subject: [PATCH 19/23] Update deepseek_v2.py --- python/sglang/srt/models/deepseek_v2.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 1616f4512f4d..b13d181569ee 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -318,9 +318,10 @@ def forward( _is_cuda and hidden_states.shape[0] <= 16 and hidden_states.shape[1] == 7168 + and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384) and _device_sm >= 90 ): - if self.weight.shape[0] == 256: + if _device_sm >= 100 and self.weight.shape[0] == 256: # router gemm output float32 logits = torch.empty( hidden_states.shape[0], @@ -334,7 +335,7 @@ def forward( logits, launch_with_pdl=True, ) - elif self.weight.shape[0] == 384: + else: logits = dsv3_router_gemm( hidden_states, self.weight, out_dtype=torch.float32 ) From 59033349b8008b3fbe1f1c15f731ee36c94fe78f Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Thu, 2 Apr 2026 17:42:00 -0700 Subject: [PATCH 20/23] Update benchmark_deepgemm_dsv3_router_gemm_blackwell.py --- .../deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py index 7bce73ca2c3b..4da26f497846 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py @@ -4,7 +4,7 @@ import torch import triton -from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 +from flashinfer.gemm import mm_M1_16_K7168_N256 from sgl_kernel import dsv3_router_gemm N = 256 From 4a353303535884f9d54b74c9f1de0b91e8c00abb Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Thu, 2 Apr 2026 17:42:23 -0700 Subject: [PATCH 21/23] Update deepseek_v2.py --- python/sglang/srt/models/deepseek_v2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 97e9a51bab77..c67e004a7d7a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -165,7 +165,7 @@ pass if _is_cuda: - from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 + from flashinfer.gemm import mm_M1_16_K7168_N256 from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm elif _is_npu: from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import ( From 72b2f4c40172f9dbf8fa980957ccc22253088abc Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Thu, 2 Apr 2026 18:24:24 -0700 Subject: [PATCH 22/23] bypass torch compiler with decorator --- python/sglang/srt/models/deepseek_v2.py | 27 ++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c67e004a7d7a..0cb7ff0c3348 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -165,7 +165,8 @@ pass if _is_cuda: - from flashinfer.gemm import mm_M1_16_K7168_N256 + from flashinfer.gemm import mm_M1_16_K7168_N256 as _raw_dsv3_router_gemm + from sglang.srt.utils.custom_op import register_custom_op from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm elif _is_npu: from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import ( @@ -333,12 +334,7 @@ def forward( device=hidden_states.device, dtype=torch.float32, ) - mm_M1_16_K7168_N256( - hidden_states, - self.weight.t(), - logits, - launch_with_pdl=True, - ) + flashinfer_dsv3_router_gemm(logits, hidden_states, self.weight) else: logits = dsv3_router_gemm( hidden_states, self.weight, out_dtype=torch.float32 @@ -2274,4 +2270,21 @@ 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] From f40fde4c24100446505eaf630f54e50b1e824624 Mon Sep 17 00:00:00 2001 From: harrisonlimh <97203667+harrisonlimh@users.noreply.github.com> Date: Thu, 2 Apr 2026 18:25:23 -0700 Subject: [PATCH 23/23] bypass torch compiler with decorator --- python/sglang/srt/models/deepseek_v2.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 0cb7ff0c3348..a7fbabdc05ff 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -166,8 +166,9 @@ if _is_cuda: from flashinfer.gemm import mm_M1_16_K7168_N256 as _raw_dsv3_router_gemm - from sglang.srt.utils.custom_op import register_custom_op 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, @@ -2273,7 +2274,7 @@ class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): @register_custom_op( op_name="flashinfer_dsv3_router_gemm", mutates_args=[], - fake_impl=lambda logits, hidden_states, weight: None + fake_impl=lambda logits, hidden_states, weight: None, ) def flashinfer_dsv3_router_gemm( logits: torch.Tensor, @@ -2287,4 +2288,5 @@ def flashinfer_dsv3_router_gemm( launch_with_pdl=True, ) + EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]