文件最后提交记录最后更新时间
7 个月前
7 个月前
7 个月前
2 个月前
7 个月前
README

ScoreNormalize

产品支持情况

产品 是否支持
Atlas A2 训练系列产品

功能说明

  • 算子功能:MoE路由专家得分进行归一化。
  • 计算公式:

    x=x∑i=1kxi×2.5x = \frac{x}{\sum_{i=1}^{k} x_i} \times 2.5

  • 这里xx是一个二维张量,表示路由专家的分数,shape为[rowsrows, kk],每行表示一个token选择的topk个专家的分数。计算公式对每个token选择的所有专家分数进行归一化,除以所有分数的总和再乘2.5。

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
x 输入/输出 公式中的输入/输出张量x,shape为(rows, k) BFLOAT16 ND
rows 输入 输入/输出x的第一维 INT -
k 输入 输入/输出x的第二维 INT -

约束说明

  • x仅支持BFLOAT16类型。
  • k <= 16。

调用说明

torch.ops.npu_ops_transformer_ext.score_normalize(x, rows, k)