描述

将 FillTriple / ProcessToken / TokenToExpert 中循环内的 (k < axisK_) && hasExpertScalesFlag_ 运行时判断外提：SendToMoeExpert 拆分为入口 + 循环体，按 hasExpertScalesFlag_ 以字面量 true / false 分流传参，共享专家路径固定传 false，消除未启用 expertScales 场景逐 token 热循环中的分支开销。同步修复 arch22 / arch35 FullMesh 变体中的同款问题，行为与基线严格等价。

关联的Issue

关联Issue #5029

测试

本地验证，二级冒烟

文档更新

不涉及

类型标签

• 🐛 Bug修复
• ✨ 新特性
• ⚡ 性能优化
• ♻️ 重构
• 🧪 测试
• 📦 构建/CI
• 🔧 配置变更
• 📝 文档更新
• ⬆️ 依赖升级
• 🔒 安全修复
• 🧹 代码清理
• ❓ 其他，请描述：

PR #11394 代码检视报告

项目	内容
PR 标题	将专家权重判断提至moe发送循环外，消除可选topk weight引入的性能劣化
PR 地址	https://gitcode.com/cann/ops-transformer/pull/11394
作者 / 来源分支	zhong-zixin / fix_topkw_loop → master
检视基线	merge-base 33d8246e0f，PR head 1e37cfe17（单提交）
变更范围	3 个文件（+93 / -46）：moe_distribute_dispatch_v2.h、arch22 / arch35 两个 FullMesh 变体
模块	mc2
检视结论	可通过（无阻塞问题；两项非阻塞观察项见第 5 节）

---

1. 整体概述

1.1 背景与动机

提交 2a0d633137（PR !7350，"Support optional topk weights for MTE low latency"）为 moe_distribute_dispatch_v2 引入了可选的 topK 权重输入 expertScales。该特性在逐 token 发送热循环内以运行时条件 (k < axisK_) && hasExpertScalesFlag_ 控制权重写入，对未启用 expertScales 的存量场景（绝大多数调用方）造成性能劣化。

1.2 修复方式

将"是否写专家权重"的判断提升到循环外，通过函数参数显式传递：

1. FillTriple、ProcessToken 新增 bool writeExpertScale 参数；
2. SendToMoeExpert 拆分为入口 + SendToMoeExpertLoop 循环体，入口按 hasExpertScalesFlag_ 以字面量 true/false 分流；
3. SendToSharedExpert（共享专家路径）固定传 false。

该修复同时覆盖标准 kernel 与两个 FullMesh 变体，回归在三种部署形态下全部消除。

---

2. 主头文件变更解析

• FillTriple / ProcessToken 新增 writeExpertScale 参数，内部判断 if (writeExpertScale)；
• SendToMoeExpert 拆分为入口 + SendToMoeExpertLoop，入口按 hasExpertScalesFlag_ 以字面量分流；
• SendToSharedExpert 固定传 false；
• 两类调用路径（MoE：k = expertIdx % axisK_ 恒小于 axisK_；共享专家：k = axisK_ + idx 恒不小于 axisK_）行为与基线严格等价；
• expertScaleAlign_ 仅在 hasExpertScalesFlag_ == true 时初始化，expertScalesTensor_ 仅在同条件下加载且先于发送，无读未初始化风险；
• 调用点全部同步，v3/v4 经头文件复用自动覆盖。

---

3. FullMesh 两文件变更解析

3.1 改动模式

两个 FullMesh 文件复刻主头文件的修复思路，但因函数族不同（TokenToExpert / TokenToExpertInQuant 直接持有 FillTriple 调用）适配方式略有差异：

1. FillTriple 新增 bool writeExpertScale，内部判断由 (k < axisK_) && hasExpertScalesFlag_ 改为 (k < axisK_) && writeExpertScale（注意：与主头文件不同，k < axisK_ 被保留）；
2. TokenToExpert / TokenToExpertInQuant 新增 bool writeExpertScale 形参并透传给 FillTriple；
3. SendToMoeExpert 拆分为入口 + SendToMoeExpertLoop，入口按 hasExpertScalesFlag_ 以字面量 true/false 分流；
4. SendToSharedExpert 的两处调用固定传 false。

3.2 arch22（moe_distribute_dispatch_v2_full_mesh.h）等价性论证

调用路径逐一验证：

路径	k 取值来源	旧行为	新行为	结论
共享专家（SendToSharedExpert:607/609）	axisK_ + toSharedExpertIndex，恒 >= axisK_	(k < axisK_) 恒假 → 不写	字面量 false → 不写	等价
MoE 专家（SendToMoeExpertLoop:707/710）	topKIndex = calExpertIdsIdx % axisK_（:700），恒 < axisK_	条件退化为 hasExpertScalesFlag_	writeExpertScale = hasExpertScalesFlag_（入口 :666-670 分流）	等价

拆分完整性：入口保留 SplitExpertNumToCore() + CalExpertSendNum() + 分流；calExpertIdsIdx / maskN64Num / dstWinGMTensor / expertMaskTensorU64 四个声明（均无副作用）与双重循环体逐语句移入 SendToMoeExpertLoop，循环体与旧版逐行一致（仅两处调用增加实参）。

BS 模式旁证：SendToMoeExpertByBS（:836）不调用 TokenToExpert / FillTriple，不受签名变更影响，无需改动。

3.3 arch35（moe_distribute_dispatch_v2_a5_full_mesh.h）等价性论证

路径	k 取值来源	旧行为	新行为	结论
共享专家（SendToSharedExpert:626/628）	axisK_ + toSharedExpertIndex，恒 >= axisK_	恒不写	字面量 false → 不写	等价
MoE 专家（SendToMoeExpertLoop:695/697）	topKId = index % axisK_（:671），恒 < axisK_	条件退化为 hasExpertScalesFlag_	writeExpertScale = hasExpertScalesFlag_（入口 :651-655 分流）	等价

拆分完整性：与旧版 SendToMoeExpert 逐语句对比确认——前奏（Duplicate、validTokenNum 计算、syncFlagId_ = 0、SetFlag 循环）保留在入口；dstWinGMTensor / dstTokenIdx 声明（无副作用）与 token 循环移入 SendToMoeExpertLoop；尾部 WaitFlag 循环保留在入口 if/else 之后。执行顺序与旧单函数完全一致（SetFlag 全部置位 → 逐 token Wait/Set → 收尾 Wait）。

3.4 arch22 与 arch35 实现一致性

实现模式一致（FillTriple 判断式、字面量分流、共享专家传 false 均相同）。存在两处结构差异，均为两文件既有循环组织差异所致，非本次引入：

• SendToMoeExpertLoop 签名不同：arch35 额外传 validTokenNum（其循环按 token 索引跨核步进，validTokenNum 在前奏中计算）；arch22 按专家分核，循环上界用成员 sendNum_，无需传参。
• 拆分边界不同：arch22 入口只保留分核与计数两步；arch35 入口还保留 Duplicate/SetFlag 前奏与 WaitFlag 尾声。各自与旧版单函数语义一一对应。

3.5 状态一致性

• expertScaleAlign_ 仅在 Init 中 hasExpertScalesFlag_ == true 时初始化（两文件同款守卫，arch22:384 / arch35:384）；新代码仅在 writeExpertScale == true（即 hasExpertScalesFlag_ == true）时读取，无读未初始化成员风险。
• expertScalesTensor_ 在 ExpIdsCopyAndMaskCal 尾部按 hasExpertScalesFlag_ 守卫加载（arch22:1966-1970 / arch35:1816-1820），先于 AllToAllDispatchA3/A5 分流执行；共享专家核虽也执行该函数，但路径固定传 false 不读 scales，与旧代码（靠 k < axisK_ 恒假跳过）行为一致。

3.6 调用点完整性

全库检索确认：

• 两文件内 FillTriple 全部调用（定义体内各 2 处）、TokenToExpert / TokenToExpertInQuant 全部调用（共享专家各 2 处 + MoE 循环各 2 处）实参数目与新签名一一匹配；
• 四者为类私有成员，无外部调用者；
• 实例化点 moe_distribute_dispatch_v2 / moe_distribute_dispatch_v3 的 arch22 *_a3.cpp 与 arch35 *_apt.cpp 共 4 个 cpp 仅调用公共入口（Init / Process），私有方法签名变更随模板惰性实例化自动传播，无需改动；
• examples/fast_kernel_launch_example/ 下的 full_mesh 副本是独立的旧版拷贝（独立命名空间、未被本 PR 修改），不受影响。

3.7 格式检查

以仓库 .clang-format（clang-format 22.1.8）对 3 个文件 PR 版本全文执行：

• 两个 FullMesh 文件：无任何偏差；
• 主头文件：仅第 458、1065 行两处偏差（本 PR 未触碰的历史行，注释对齐，疑为 clang-format 版本差异），无新增偏差。

---

4. 性能修复有效性

• 主头文件：无权重场景热循环中权重写入、成员加载、比较与分支全部消失；有权重场景仅保留无条件写入。字面量实参 + 内联 + 常量传播可生成两份无内部分支的特化循环；即使编译器未完成折叠，循环内分支也退化为寄存器参数判断，不会劣于旧代码。
• FullMesh：无权重场景（回归场景）(k < axisK_) && false 经内联 + 常量折叠完全消除，回归修复覆盖 FullMesh 部署形态；有权重场景保留一次 k < axisK_ 运行期比较（见观察项 5.2），仍优于旧代码（成员加载 + 比较 + 分支）。
• 相比把 hasExpertScales 提升为模板参数的方案（需翻倍实例化、增加 tiling key 复杂度），当前运行时分流 + 编译期折叠的折中是合理的。

---

5. 检视发现

阻塞问题

无。

观察项（非阻塞）

[观察-1] FullMesh 修复未覆盖 BS 模式路径（arch22）
arch22 的 SendToMoeExpertByBS（:836，canUseBSMode 分支）不经过 TokenToExpert / FillTriple，本就不含 (k < axisK_) && hasExpertScalesFlag_ 判断，无回归残留。仅作完整性说明，无需改动。

[观察-2] FullMesh FillTriple 保留了冗余的 k < axisK_ 判断（可选清理）
两文件 FillTriple 的判断为 (k < axisK_) && writeExpertScale，而两条调用路径上该条件已可静态判定：共享专家路径 writeExpertScale == false 使 k < axisK_ 成为死条件；MoE 路径 k = xxx % axisK_ 恒真。主头文件的对应修改已将判断简化为 if (writeExpertScale)。FullMesh 保留 k < axisK_ 的后果：无权重特化路径（回归场景）经常量折叠完全消除，不受影响；有权重特化路径每次 FillTriple 多一次运行期比较，属微小冗余。保留亦可视为防御性写法，不要求修改；如后续清理，建议与主头文件风格统一。

---

6. 综合判定

维度	判定
正确性	通过。三类 kernel 变体（主、arch22 FullMesh、arch35 FullMesh）两类调用路径与基线严格等价，状态初始化/加载时序一致
一致性	通过。arch22 / arch35 实现模式一致，结构差异源于既有循环组织差异
调用点	通过。全库无私有成员外部引用遗漏，4 个实例化 cpp 无需改动
性能	通过。回归场景（无 expertScales）热循环开销在三种 kernel 变体中全部消除
格式/规范	通过。两 FullMesh 文件 clang-format 干净；主头文件仅存两处历史行偏差（本 PR 未触碰）
综合判定	可通过（观察-2 可选清理）