已合并
fix(moe): 补齐 MoeFinalizeRouting / MoeReRouting 的结构化输入契约 #251
fix(moe): 补齐 MoeFinalizeRouting / MoeReRouting 的结构化输入契约 #251
已合并
Deng Pan创建于 14 天前
Deng Pan
Deng Pan成员
14 天前

两个 MoE 算子的输入都有单区间 value_range 表达不了的结构约束,通用生成器与现有 get_input 都没兜住。


1. MoeFinalizeRouting:expanded_src_to_dst_row 应为单射映射

该输入是 MoeInitRouting 产出的"源行 → 展开缓冲区行"映射,每个源行 (token, k) 占据展开缓冲区中互不相同的一行:

  • drop_less(mode 0/2):[0, NK) 的完整置换
  • drop_pad(mode 1/3):到 [0, E*C) 的单射,容量溢出的源行取 -1 表示丢弃

而通用生成器按 randint 独立有放回采样,重复量很大(约 37% 的条目落在已被占用的目的行上):

case 修复前重复 case 修复前重复
c1 3014 / 8192 c9 24158 / 65536
c2 12083 / 32768 c15 96284 / 262144
c10 (drop_pad) 14 / 99 c11 (drop_pad) 93 / 399

这不是"当前判错",但有两个实质问题

golden 与真实 MoeFinalizeRoutingV2 都做 gather,重复索引下数值仍然可比(现网用例即如此通过)。然而:

  1. 访存模式失真 —— 真实场景每行恰好读一次,重复采样把它变成随机重复读,L2 命中率虚高,perf 数据不具代表性。
  2. 契约站在候选一边 —— 任何采用 scatter 方向实现的候选(遍历展开行写回目的行,与 gather 等价当且仅当映射单射)都会与 golden 分叉。

修法

新增 get_input:按 case 实际形状推导目的空间 num_dst(drop_less 下 == NK,drop_pad 下 == E*C),为非 -1 位置分配 [0, num_dst) 的互异值;-1 的位置原样保留,维持 drop_pad 用例既有的丢弃覆盖与随种子可复现的行为。其余 6 个输入原样透传。

未做的部分已在注释中写明:drop_pad 的完整契约还要求目的行落在该源行所属专家的容量块 [e*C, (e+1)*C) 内。golden 不校验这一点,且按专家分配会让丢弃率降到近乎 0(当前用例 E*C 远大于 NK),反而削弱 -1 路径覆盖,故只做单射重建;专家对齐应由用例设计连同容量一并调整。


2. MoeReRouting:get_input 自身会静默违反契约

proto.yaml:39 声明 expert_token_num_per_rank 的元素必须大于 0,而 get_inputbase_value = A // (N*E) 平铺 —— A < N*Ebase_value 为 0,除最后一格外全是 0,静默产出违反契约的输入。候选 kernel 在"某张卡的某个专家分到 0 个 token"上的行为未定义,失败会表现为莫名的精度不符而非配置错误。

当前用例集最小 base_value = 4(c20: A=64, NE=16),不会触发;这是为后续新增小 A 用例设的护栏:base_value < 1 时直接抛 ValueError 并指明 A 与 NE 的具体数值。不改既有的余数分配方式,20 个用例行为完全不变。


验证

  • MoeFinalizeRouting 20 个用例:重复数全部归零、无越界、-1 位置与数量保持不变、golden 输出 shape 不变。
  • MoeReRouting 20 个用例:sum == Amin > 0,与修复前逐位一致。
  • 新增 tests/ut/test_moe_get_input.py(15 项):drop_less 满置换、drop_pad 非 -1 项互异、-1 位置保留、目的空间取值域、shape/dtype 契约、跨调用可复现、其余输入透传、抽屉原理兜底、golden 可执行;以及 ReRouting 的可行/边界/不可行三态与透传。
  • tests/ut 全量 1069 passed / 6 skipped

顺带记录一个用例设计层面的覆盖缺口(本次不改)

MoeReRouting 20 个用例里 18 个expert_token_num_per_rank 是完全均匀分布,专家负载不均的场景没有覆盖。这属于用例设计范畴,建议由算子 owner 评估是否补充。

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 Deng Pan 的贡献)
Deng PanDeng Pan成员
14 天前 创建了 pull request,commit 776e5b60
atomgit-bot
atomgit-bot
14 天前 评论:

变更摘要

此 PR 为 MoeFinalizeRoutingMoeReRouting 两个 MoE 算子补齐结构化输入契约:MoeFinalizeRouting 新增 get_input 函数,将 expanded_src_to_dst_row 从通用生成器产生的有放回随机采样重建为符合单射约束的映射;MoeReRoutingget_input 中新增 A < N*E 时的显式报错护栏,防止静默产出违反 proto.yaml 契约(元素必须大于 0)的输入。同时新增 15 项单元测试覆盖两类算子的契约校验与边界行为。

主要改动

  • tasks/level3/moe_finalize_routing/golden.py 新增 get_input 函数:根据 expanded_permuted_rows 形状推导目的空间 num_dst,利用 argsort(rand) 为非 -1 位置分配互异的目的行,确保 expanded_src_to_dst_row 满足 drop_less 下的完整置换和 drop_pad 下的单射约束,同时原样保留 -1 位置以维持既有丢弃覆盖行为。
  • tasks/level3/moe_re_routing/golden.py 新增 base_value < 1 防护:在原有 base_value = A // (N*E) 平铺逻辑之后、生成新 expert_token_num_per_rank 之前,检查当 A < N*E 时直接抛出包含 A 与 N*E 具体数值的 ValueError,避免后续用例因 token 不足而静默产出全零输入。
  • 新增 tests/ut/test_moe_get_input.py 测试文件:包含 TestFinalizeRoutingInjectivity(9 项:drop_less 满置换、drop_pad 非 -1 项互异、-1 位置保留、目的空间取值域、shape/dtype 契约、跨调用可复现、其余输入透传、抽屉原理兜底、golden 可执行)和 TestReRoutingGuard(5 项:可行/边界/不可行三态及总和校验、余数吸收、per_token_scales 透传),共 14 项测试。
likedislike
atomgit-bot
atomgit-bot
14 天前 评论:

代码审查

✅ 未发现问题

likedislike
CANN-robotCANN-robot成员
14 天前 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
14 天前 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
repo-cann/cann-bench gxj1123, fanyuwei_math (2/2) gxj1123 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

Deng_Pan, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
fanyuwei成员
11 天前 评论:

compile

likedislike
Deng PanDeng Pan成员
11 天前 virtual merging failed, update merge request[project_id: 9587699, iid: 251, target_commit_sha: e46a42e1bc5edfb5cab69d41d84ddb478515fa91], message: You are not allowed to merge to this ref
CANN-robotCANN-robot成员
11 天前 添加了label:ci-pipeline-running
CANN-robot
CANN-robot成员
11 天前 评论:

流水线任务触发成功
任务链接 [6a1ea73b25ab4d0f9bcfebf19424eedb][流水线指导]

任务名称状态日志下载链接
SCA ✅ SUCCESS >>>>>
antipoison ✅ SUCCESS >>>>>
Check_Pr ✅ SUCCESS >>>>>
codecheck_style ✅ SUCCESS >>>>>
pre-commit ✅ SUCCESS
UT_Test ✅ SUCCESS >>>>>
pre_comment ✅ SUCCESS >>>>>
PreSmoke_A900 ✅ SUCCESS >>>>>

[2026-08-10 14:36:18]    CI执行结束

likedislike
CANN-robotCANN-robot成员
11 天前 删除了label:ci-pipeline-running
CANN-robotCANN-robot成员
11 天前 添加了label:ci-pipeline-passed
gxj1123成员
11 天前 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
11 天前 添加了label:approved
fanyuwei成员
10 天前 评论:

/lgtm

likedislike
CANN-robotCANN-robot成员
10 天前 添加了label:lgtm
CANN-robotCANN-robot成员
10 天前 合入了pull request