[AMD] Optimize MiniMax-M2.5 - use aiter biased_grouped_topk for sigmoid scoring in MoE routing (#23611)
This commit is contained in:
@@ -538,13 +538,24 @@ def fused_topk(
|
|||||||
renormalize,
|
renormalize,
|
||||||
)
|
)
|
||||||
elif scoring_func == "sigmoid":
|
elif scoring_func == "sigmoid":
|
||||||
topk_sigmoid(
|
if _use_aiter and correction_bias is not None:
|
||||||
topk_weights,
|
aiter_biased_grouped_topk(
|
||||||
topk_ids,
|
gating_output,
|
||||||
gating_output,
|
correction_bias.to(dtype=gating_output.dtype),
|
||||||
renormalize,
|
topk_weights,
|
||||||
correction_bias,
|
topk_ids,
|
||||||
)
|
num_expert_group=1,
|
||||||
|
topk_group=1,
|
||||||
|
need_renorm=renormalize,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
topk_sigmoid(
|
||||||
|
topk_weights,
|
||||||
|
topk_ids,
|
||||||
|
gating_output,
|
||||||
|
renormalize,
|
||||||
|
correction_bias,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Invalid scoring function: {scoring_func}")
|
raise ValueError(f"Invalid scoring function: {scoring_func}")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user