[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:
Even Zhou
2025-08-09 01:35:00 -07:00
committed by GitHub
co-authored by ronnie_zheng Hexq0210
parent 52e1f52f32
commit 137e75daa1
7 changed files with 210 additions and 61 deletions
+2 -1
View File
@@ -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,