[CI] Ascend NPU CI enhancement (#8294)
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user