minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[NPU] Fix a corner case where FusedMoE.top_k is not explicitly declared (#19287)

Browse Source
This commit is contained in:
Even Zhou
2026-02-25 20:49:59 +08:00
committed by GitHub
parent 9840cd3f68
commit cdc411160b
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/unquant.py
+3 -4
View File
@@ -546,7 +546,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
# x.shape = [B*S, H]
x = dispatch_output.hidden_states x = dispatch_output.hidden_states
# topk_weights.shape = [B*S, K]; topk_ids.shape = [B*S, K]
topk_weights, topk_ids, _ = dispatch_output.topk_output topk_weights, topk_ids, _ = dispatch_output.topk_output
original_dtype = x.dtype original_dtype = x.dtype
@@ -554,10 +556,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
topk_weights = topk_weights.to(x.dtype) topk_weights = topk_weights.to(x.dtype)
topk_ids = topk_ids.to(torch.int32) topk_ids = topk_ids.to(torch.int32)
num_experts = layer.num_experts num_experts = layer.num_experts
top_k = layer.top_k top_k = layer.top_k or topk_ids.shape[1] # in case layer.top_k is not set
w13 = layer.w13_weight
w2 = layer.w2_weight
hidden_states, expanded_row_idx, expert_tokens, _ = ( hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2( torch.ops.npu.npu_moe_init_routing_v2(