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

This commit is contained in:
Even Zhou
2026-02-25 20:49:59 +08:00
committed by GitHub
parent 9840cd3f68
commit cdc411160b

View File

@@ -546,7 +546,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
# x.shape = [B*S, H]
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
original_dtype = x.dtype
@@ -554,10 +556,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
topk_weights = topk_weights.to(x.dtype)
topk_ids = topk_ids.to(torch.int32)
num_experts = layer.num_experts
top_k = layer.top_k
w13 = layer.w13_weight
w2 = layer.w2_weight
top_k = layer.top_k or topk_ids.shape[1] # in case layer.top_k is not set
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(