Add dsv3 router gemm kernel (#7627)

2025-06-29 23:31:55 -07:00
parent 22352d47a9
commit 7248272ccc
8 changed files with 398 additions and 0 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -200,6 +200,7 @@ void bmm_fp8(
    at::Tensor workspace_buffer,
    int64_t cublas_handle,
    int64_t cuda_stream);
+void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, const torch::Tensor& mat_b);

 void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b);