Closed
[质量加固]: moe_gating_top_k arch35 hash模式与group-normal路径耦合在同一kernel模板,拆分为独立的MoeGatingTopKHashRegbase模板 #7201
CPZhouAdd on  3 days agoClosed at  2 days ago
CPZhou
CPZhouCollaborators
3 days ago Create

问题描述

moe_gating_top_k 的 arch35 register-based kernel 模板 MoeGatingTopKRegbase 同时承担两条互斥的执行路径:

  1. group-normal TopK 路径(TILING_KEY_REGBASE):按 bias 加权后排序选 TopK;
  2. hash 路径(TILING_KEY_HASH_{INT32,INT64}_{INT32,INT64} 共 4 种 tiling key):由 inputIds + tid2eid 哈希表直接映射选专家。

两者共用同一模板导致:

  • 模板需要同时携带 IDXT/EIDXT 两个仅 hash 路径使用的模板参数,group-normal 路径被迫以默认参数实例化,分支与模板参数耦合,可读性差;
  • hash 专属逻辑(inputIds/tid2eid 拷入、哈希查表、按 tid 行 gather)与 group-normal 的排序/归一化逻辑交织在同一实现内,单文件臃肿(重构前约 508 行),维护与后续调优互相牵制;
  • 两条路径无法独立演进(如 hash 路径后续要做数据类型扩展或性能优化会波及 group-normal 路径)。

期望改进

将 hash 模式拆分为独立的 kernel 模板 MoeGatingTopKHashRegbase<T, IDXT, EIDXT>(继承精简后的 MoeGatingTopKRegbase<T>,复用输入拷贝/归一化/CopyOut 公共通路),moe_gating_top_k_apt.cpp 按 tiling key 分别实例化,实现两条路径解耦、各自独立维护。

likedislike
周云飞
周云飞Collaborators
3 days ago Comments:

/assign @ZhouChuping

likedislike
CANN-robotCANN-robotCollaborators
3 days ago Assign ZhouChuping as responsible person
周云飞
周云飞Collaborators
3 days ago Comments:

问题说明

这是模板职责划分的重构建议,目前没有证据表明存在计算错误。MoeGatingTopKRegbase 确实带有两个默认整型模板参数,初始化接口也接收 inputIds、tid2eid。按报告描述,拆出 hash 路径有助于隔离类型扩展,但继承方案仍需确认公共逻辑与状态的边界。

处理建议

可以考虑新增 MoeGatingTopKHashRegbase<T, IDXT, EIDXT>,将查表与 gather 逻辑移入其中,在 moe_gating_top_k_apt.cpp 按 tiling key 分别实例化。输入拷贝、归一化和输出仅复用真正共通的部分,避免基类继续依赖 hash 参数;也可评估抽取公共辅助函数,减少继承耦合。重构应保持原有接口与 tiling key 行为,重点验证四种 hash 类型组合及 group-normal 路径的输出一致性。

likedislike
CANN-robotCANN-robotCollaborators
2 days ago Close issue
CANN-robotCANN-robotCollaborators
2 days ago changed custom state from 进行中 to 已完成
CANN-robotCANN-robotCollaborators
2 days ago Added label: resolved