Backgroud（背景信息）

使用AscendC 实现对MoE专家得分进行归一化的算子ScoreNormalize。ScoreNormalize算子将MoE模型中的路由专家分数进行归一化，输入x（BFloat16）是shape为[rows，k]的张量，即路由专家的分数，其中rows表示序列长度，k表示topk。

ScoreNormalize将分数以最后一维执行归一化，计算公式为：

$$x = \frac{x}{\sum_{i=1}^{k}x_i} \times 2.5$$

Benefit / Necessity （价值/作用）

将通过多个aclnn适配的归一化操作融合为一个算子实现，优化效率。

Design（设计方案）

Tilling 策略

• 分核策略：
  • 将输入[num_rows, h] 按 num_rows分核， 结合UB大小计算得到每个核能存放的最大的rows_per_core，再根据数据总量计算每个核需要被调用的次数loop_count

Kernel 侧设计

进行Init和Process两个阶段，其中Process包括数据搬入（CopyIn）、计算（Compute）、数据搬出（CopyOut）三个阶段，实现专家得分归一化计算。

• 初始化（Init）
  • 计算rows_per_core以及loop_count
  • 建立GM Tensor映射，初始化队列缓冲区，初始化临时缓冲区
• 计算流程（Compute）
  • 调用WholeReduceSum对得分求和
  • 调用Broadcast将得分广播
  • 调用Div将广播后的得分和求和的得分做除法
  • 调用Muls将做除法后的得分乘2.5
• 数据搬入（CopyIn）搬出（CopyOut）
  • CopyIn: 将score数据从inGM中搬入到inQueue
  • CopyOut: 将归一化后的score数据从outQueue搬出到outGM
• 流程调度（Process）
  • 遍历loop_count，按CopyIn->Compute->CopyOut的顺序调度