EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Qingquan Song
2025-03-29 11:51:45 -07:00
committed by GitHub
parent ddf8981d91
commit 45dcfc2e76
9 changed files with 616 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
sgl-kernel/python/sgl_kernel/moe.py
View File

@@ -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
)