Merged
feat(moe): migrate gate operators #9
wangliweiAdd on Sep 17
feat(moe): migrate gate operators #9
Merged
Pull Request has been successfully merged by CANN-robot.
（Thank you for wangliwei's contribution.）Sep 17 Added label: cann-cla/yes
CANN-robot
Sep 17 Comments：
Sep 17 Comments：
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.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-cann/ops-tilelang | ✅ kevin365, ChaoyangJi (2/2) | ✅ kevin365 (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
wangliwei12, thanks for your pull request. All authors of the commits have signed the CLA. 👍
CANN-robotSep 17 Set ChaoyangJi,kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver
Sep 17 Set ChaoyangJi,kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver
CANN-robotSep 17 Set ChaoyangJi,kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver
Sep 17 Set ChaoyangJi,kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver
CANN-robot
Sep 17 Comments：
Sep 17 Comments：
| 🚀 CI 流水线已启动 |
|---|
| 📋 执行详情: 点击查看流水线 |
Sep 17 Added label: ci-pipeline-running
Sep 17 Removed label: ci-pipeline-running
Sep 17 Added label: ci-pipeline-failed
CANN-robot
Sep 17 Comments：
Sep 17 Comments：
| 🚀 CI 流水线已启动 |
|---|
| 📋 执行详情: 点击查看流水线 |
Sep 17 Removed label: ci-pipeline-failed
Sep 17 Added label: ci-pipeline-running
Sep 18 Removed label: ci-pipeline-running
Sep 18 Added label: ci-pipeline-failed
zhangwei503
Sep 18 Comments：
Sep 18 Comments：
compile
Sep 18 virtual merging failed, update merge request[project_id: 10728021, iid: 9, target_commit_sha: cea6d8ce1a46d9ffba2f4015c684319dd48f71f2], message: You are not allowed to merge to this ref
Sep 18 virtual merging failed, update merge request[project_id: 10728021, iid: 9, target_commit_sha: cea6d8ce1a46d9ffba2f4015c684319dd48f71f2], message: You are not allowed to merge to this ref
CANN-robot
Sep 18 Comments：
Sep 18 Comments：
| 🚀 CI 流水线已启动 |
|---|
| 📋 执行详情: 点击查看流水线 |
Sep 18 Removed label: ci-pipeline-failed
7 event messages are collapsed here。 See more
Sep 20 Added label: cann-cla/yes
CANN-robot
Sep 20 Comments：
Sep 20 Comments：
| 🚀 CI 流水线已启动 |
|---|
| 📋 执行详情: 点击查看流水线 |
Sep 20 Removed label: ci-pipeline-failed
Sep 20 Added label: ci-pipeline-running
Sep 20 Removed label: ci-pipeline-running
Sep 20 Added label: ci-pipeline-passed
Sep 20 Added label: approved
README.md
@@ -31,29 +33,30 @@
3133└── pyproject.toml # 构建、依赖和工具配置
@@ -115,6 +117,24 @@ row_norm, col_sumsq = engram_sinkhorn_step(matrix, row_scale, col_scale, eps=1e-
115117
116118三个输入必须位于同一 NPU、连续且为 FP32，并且不能共享存储。`matrix` 的形状为 `(num_rows, head_dim)`，`head_dim` 必须是正的 256 倍数；调用方需保证 `row_scale` 包含非负的 Sinkhorn 缩放系数。已验证的典型维度为 256 和 5120。由于 `SimtVF` 需预留线程上下文，wrapper 会拒绝共享 UB 估算超过 216 KiB 的 shape。`num_rows=0` 时返回空的 `row_norm` 和全零的 `col_sumsq`。
117119
120+MoE 路由算子接收连续的 NPU FP32 分数。`topk_gate` 稳定选择每个 token 的专家索引，分数相同时选择较小索引；`moe_topk_gate` 融合打分、偏置、掩码、固定路由、随机路由和逻辑到物理专家映射；`normalize_weight` 返回每行分母及归一化权重：
121+
122+```python
123+from cann_ops_tilelang import moe_topk_gate, normalize_weight, topk_gate
124+
125+topk_idx = topk_gate(scores, num_topk=6)
126+denominator, normalized_weights = normalize_weight(topk_weights)
127+topk_idx, topk_weights = moe_topk_gate(
128+ logits,
129+ num_topk=6,
130+ use_shared_as_routed=False,
131+ num_shared_experts=1,
132+ routed_scaling_factor=1.0,
133+ ep_rank=0,
134+ scoring_func='sigmoid',
135+)
136+```
137+
118138## 测试
119139
120140```bash
@@ -135,6 +155,11 @@ pytest tests/engram/test_engram_fused_weight.py -m 'not benchmark' --refresh-mem
135155pytest tests/engram -m benchmark --run-benchmark --refresh-baseline
136156pytest tests/engram -m 'not benchmark' --refresh-mem
137157
158+# 运行或刷新 MoE 算子族
这里只是举例，不是所有的算子都在这里列一下，前面已经有Engram举例了，所以这里的moe就不用加了。
README.md
@@ -4,7 +4,7 @@
44
55> **硬件支持范围**：本代码仓仅支持 Ascend A5 及以上架构，不保证在更早的 Ascend 架构上可编译或正确运行。
66
7-当前仓库提供基础工程框架：Engram 算子族包含 `engram_fused_weight`、`engram_gate_fwd`、`engram_gate_bwd`、`engram_grad_w_reduce` 和 `engram_sinkhorn_step`，覆盖前向、反向、分片梯度归并及 Sinkhorn 行归一化。所有算子均接入分级参数生成、确定性随机种子、NPU 显存画像、A5 性能回归框架、正则筛选、xdist 多卡绑定以及分阶段测试调度。
7+当前仓库提供基础工程框架：Engram 算子族覆盖前向、反向、分片梯度归并及 Sinkhorn 行归一化；MoE 算子族提供 `topk_gate`、`moe_topk_gate` 和 `normalize_weight`。算子测试接入分级参数生成、确定性随机种子、NPU 显存画像、A5 性能回归框架、正则筛选、xdist 多卡绑定以及分阶段测试调度。
88
99## 设计原则
1010
@@ -21,11 +21,13 @@ ops-tilelang/
2121├── src/cann_ops_tilelang/
2222│ ├── config.py # Ascend 环境与运行配置
2323│ ├── utils.py # 通用纯 Python 工具
24-│ └── engram/ # Engram 前向、反向和梯度归并算子
24+│ ├── engram/ # Engram 前向、反向和梯度归并算子
25+│ └── moe/ # MoE 路由和权重归一化算子
2526├── tests/
2627│ ├── testing/ # 参数生成、数值断言和 benchmark 工具
2728│ ├── pytest_*_plugin.py # benchmark/显存/seed/xdist/warning 插件
28-│ └── engram/ # 算子测试、PyTorch reference、benchmark 与显存画像
29+│ ├── engram/ # Engram 测试与 PyTorch reference
30+│ └── moe/ # MoE 测试与 PyTorch reference
2931├── docs/ # 架构和新增算子指南
3032├── scripts/ # 本地/CI 执行入口
3133└── pyproject.toml # 构建、依赖和工具配置
@@ -115,6 +117,24 @@ row_norm, col_sumsq = engram_sinkhorn_step(matrix, row_scale, col_scale, eps=1e-
115117
116118三个输入必须位于同一 NPU、连续且为 FP32，并且不能共享存储。`matrix` 的形状为 `(num_rows, head_dim)`，`head_dim` 必须是正的 256 倍数；调用方需保证 `row_scale` 包含非负的 Sinkhorn 缩放系数。已验证的典型维度为 256 和 5120。由于 `SimtVF` 需预留线程上下文，wrapper 会拒绝共享 UB 估算超过 216 KiB 的 shape。`num_rows=0` 时返回空的 `row_norm` 和全零的 `col_sumsq`。
117119
120+MoE 路由算子接收连续的 NPU FP32 分数。`topk_gate` 稳定选择每个 token 的专家索引，分数相同时选择较小索引；`moe_topk_gate` 融合打分、偏置、掩码、固定路由、随机路由和逻辑到物理专家映射；`normalize_weight` 返回每行分母及归一化权重：
这里只是举例，不是所有的算子都在这里列一下，前面已经有Engram举例了，所以这里的moe就不用加了。
ChaoyangJi
Sep 20 Comments：
Sep 20 Comments：
/lgtm
Sep 20 Added label: lgtm
CANN-robot
Sep 20 Comments：
Sep 20 Comments：
The MR can not be merged, because of CodeReview discussion not resolved
If you want to solve this problem, you can click here to do it in the FAQs.
Sep 20 Resolved final issue
Sep 20 Merged pull request
当前PR是否有AI参与:
[x] 否
[ ] 是
__1. AI Agent 平台:
__2. AI 模型:
__3. Prompt上下文 :
PR功能描述 / 为什么需要这个合入**:
该PR关联的issue
(格式为fixes #<issue号>, 或者resolves #<issue号>): fixes #
希望检视人员了解:
改动类型 / Change Type
Bug 修复 / Bug Fix
[ ✅] 新功能 / New Feature
性能优化 / Performance
代码重构 / Refactoring
文档更新 / Documentation
测试相关 / Test
其它 / Other
测试信息 / Testing
[ ✅] 单元测试通过 / UT passed
[ ✅] 集成测试通过 / ST passed
[ ]✅ 人工验证通过 / Manual verified
检查清单 / Checklist
[ ✅] 代码符合规范 / Code follows style guide
[ ✅] 测试添加并通过 / Tests added and passed
[ ✅] 文档已更新 / Docs updated if needed
[ ✅] 无硬编码敏感信息 / No secrets hardcoded
[ ✅] 提交信息符合规范 / Commit message follows convention