[CI] Ascend NPU CI enhancement (#8294)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Even Zhou
2025-08-04 13:16:38 +08:00
committed by GitHub
parent f57d2dc162
commit fee0ab0fba
9 changed files with 415 additions and 189 deletions

View File

@@ -398,8 +398,12 @@ def grouped_topk_gpu(
.reshape(num_token, -1)
) # [n, e]
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
# TODO: NPU can't support directly evaluating a comparison for now
topk_weights, topk_ids = torch.topk(
tmp_scores, k=topk, dim=-1, sorted=num_fused_shared_experts > 0
tmp_scores,
k=topk,
dim=-1,
sorted=(True if num_fused_shared_experts > 0 else False),
)
if num_fused_shared_experts:
topk_ids[:, -1] = torch.randint(
@@ -489,8 +493,12 @@ def biased_grouped_topk_impl(
tmp_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")
) # [n, e]
# TODO: NPU can't support directly evaluating a comparison for now
_, topk_ids = torch.topk(
tmp_scores, k=topk, dim=-1, sorted=num_fused_shared_experts > 0
tmp_scores,
k=topk,
dim=-1,
sorted=(True if num_fused_shared_experts > 0 else False),
)
topk_weights = scores.gather(1, topk_ids)