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

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,

sgl-kernel/python/sgl_kernel/moe.py

@@ -32,3 +32,15 @@ def topk_softmax(
     torch.ops.sgl_kernel.topk_softmax.default(
         topk_weights, topk_ids, token_expert_indices, gating_output
     )
+
+
+def moe_fused_gate(input_tensor, bias, num_expert_group, topk_group, topk):
+    # This fused kernel function is used to select topk expert in a hierarchical 2-layer fashion
+    # it split group of expert into num_expert_group, and use top2 expert weight sum in each group
+    # as the group weight to select exerpt groups and then select topk experts within the selected groups
+    # the #experts is decided by the input tensor shape and we currently only support power of 2 #experts
+    # and #experts should be divisible by num_expert_group. #expert/num_expert_group <= 32 is limitted for now.
+    # for non-supported case, we suggestion to use the biased_grouped_topk func in sglang.srt.layers.moe.topk
+    return torch.ops.sgl_kernel.moe_fused_gate(
+        input_tensor, bias, num_expert_group, topk_group, topk
+    )