Add deepseek style fused moe group gate selection kernel (#4530)

2025-03-29 11:51:45 -07:00
parent ddf8981d91
commit 45dcfc2e76
9 changed files with 616 additions and 1 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -36,7 +36,7 @@ from sgl_kernel.gemm import (
    sgl_per_token_group_quant_int8,
    sgl_per_token_quant_fp8,
 )
-from sgl_kernel.moe import moe_align_block_size, topk_softmax
+from sgl_kernel.moe import moe_align_block_size, moe_fused_gate, topk_softmax
 from sgl_kernel.sampling import (
    min_p_sampling_from_probs,
    top_k_renorm_prob,