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


3 days ago Assign ZhouChuping as responsible person
周云飞
3 days ago Comments:
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 路径的输出一致性。


2 days ago Close issue
2 days ago changed custom state from 进行中 to 已完成
2 days ago Added label: resolved
问题描述
moe_gating_top_k的 arch35 register-based kernel 模板MoeGatingTopKRegbase同时承担两条互斥的执行路径:TILING_KEY_REGBASE):按 bias 加权后排序选 TopK;TILING_KEY_HASH_{INT32,INT64}_{INT32,INT64}共 4 种 tiling key):由inputIds+tid2eid哈希表直接映射选专家。两者共用同一模板导致:
IDXT/EIDXT两个仅 hash 路径使用的模板参数,group-normal 路径被迫以默认参数实例化,分支与模板参数耦合,可读性差;inputIds/tid2eid拷入、哈希查表、按 tid 行 gather)与 group-normal 的排序/归一化逻辑交织在同一实现内,单文件臃肿(重构前约 508 行),维护与后续调优互相牵制;期望改进
将 hash 模式拆分为独立的 kernel 模板
MoeGatingTopKHashRegbase<T, IDXT, EIDXT>(继承精简后的MoeGatingTopKRegbase<T>,复用输入拷贝/归一化/CopyOut 公共通路),moe_gating_top_k_apt.cpp按 tiling key 分别实例化,实现两条路径解耦、各自独立维护。