[Feature] Optimize DeepSeek's DeepEP on Ascend NPU (#8355)
Co-authored-by: ronnie_zheng <zl19940307@163.com> Co-authored-by: Hexq0210 <hexq0809521@gmail.com>
This commit is contained in:
co-authored by
ronnie_zheng
Hexq0210
parent
52e1f52f32
commit
137e75daa1
@@ -250,10 +250,11 @@ class TopK(CustomOp):
|
||||
|
||||
# NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern
|
||||
if global_num_experts == 256:
|
||||
router_logits = router_logits.to(torch.float32)
|
||||
return torch_npu.npu_moe_gating_top_k(
|
||||
router_logits,
|
||||
k=self.top_k,
|
||||
bias=self.correction_bias,
|
||||
bias=self.correction_bias.to(torch.float32),
|
||||
k_group=self.topk_group,
|
||||
group_count=self.num_expert_group,
|
||||
group_select_mode=1,
|
||||
|
||||
Reference in New Issue
Block a user