已合并
[feat][compilation]按torch版本兼容rms_norm/rope pattern的dtype cast分解 #519
zqxu创建于 23 天前
[feat][compilation]按torch版本兼容rms_norm/rope pattern的dtype cast分解 #519
已合并
Pull Request已成功合入, 合并人@ascend-robot
(感谢 zqxu 的贡献)atomgit-bot
23 天前 评论:
23 天前 评论:
变更摘要
本 PR 修复 torch_npu 2.9+ 环境下 RMSNormPattern 与 RopePattern 编译 pattern 全部失配的问题。根因是 torch_npu 2.9 起(MR 30358)移除了 Tensor.to 的 NPU 过适配,dtype cast 在 AOT 分解中由 torch.ops.npu._npu_dtype_cast 变为 torch.ops.aten._to_copy,而 PatternMatchPass 基于文本级算子序列匹配,首个节点不匹配会导致整条 pattern 链放弃。本 PR 在 rms_norm_pattern.py 和 rope_pattern.py 中新增 IS_TORCH_GE_29 版本判断,按 torch.__version__ 分支选择 dtype cast 算子:旧版本保持 _npu_dtype_cast 不变,2.9+ 版本改用 torch.ops.aten._to_copy.default(dtype 为关键字参数);同时修复 rope_pattern.py 中 x_out.type_as(x) 未赋值的 no-op bug,并补上真实图中存在的末尾 .to(x.dtype) cast 节点,使 RoPE 融合在 torch 2.9+ 上重新匹配生效。
主要改动
- 版本分支选择 dtype cast 算子:在
rms_norm_pattern.py与rope_pattern.py中新增基于torch.__version__解析的IS_TORCH_GE_29判断,torch >= 2.9时_dtype_cast_func改为lambda x, dtype: torch.ops.aten._to_copy.default(x, dtype=dtype),torch < 2.9保持原有torch.ops.npu._npu_dtype_cast.default,确保旧版本行为完全不变。 - 修复 RoPE pattern 的 no-op 与缺失 cast 节点:将
rope_pattern.py中未赋值的x_out.type_as(x)替换为return torch.ops.aten._to_copy.default(x_out, dtype=x.dtype),使输出 cast 与 diffusersapply_rotary_emb真实图结构一致,保证末尾.to(x.dtype)节点能够被 pattern 匹配。 - 恢复 RMSNorm/RoPE 融合匹配:修复后
PatternMatchPass的替换数由 190 提升至 342(RMSNorm +152)、494 总量(含 RoPE +152),profiling 验证npu_rms_norm与RotaryPositionEmbedding融合算子真实执行,并在 torch 2.10 上同样有效。


不准确?
atomgit-bot
23 天前 评论:
23 天前 评论:
23 天前 添加了label:pr-audit-failed
此处折叠了91条消息 查看更多
21 天前 添加了label:ci-pipeline-passed
ascend-robot
21 天前 评论:
21 天前 评论:
流水线 PR-pipeline_MindIE-SD_gitcode#1758 [ commitID:69117ba1 ] 已完成
>>>代码风格自动修复执行失败,具体请查看日志,不影响流水线执行及PR合入
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_linux_x86_abi1 | ✅ | >>> |
| Build_linux_arm_abi1 | ✅ | >>> | |
| Build_linux_arm_abi0 | ✅ | >>> | |
| 恶意代码检查 | Antipoison | ✅ | >>> |
| 开源片段检查 | SCA | ✅ | >>> |
| 开发者测试 | UT_linux_arm | ✅ | >>> |
| 流水线 | PR-pipeline_MindIE-SD_gitcode | ✅ | >>> |
- compile : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线


21 天前 关闭了关联的issue
21 天前 合入了pull request
ascend-robot
21 天前 评论:
21 天前 评论:
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


Which issue(s) this PR fixes or accomplishes
Fixes #298
Purpose
torch_npu 2.9 起(MR 30358,commit
5c2817cd85,2026-02)移除了Tensor.to的 NPU 过适配:dtype cast 在编译图(AOT 分解)中由torch.ops.npu._npu_dtype_cast变为torch.ops.aten._to_copy。由于 compilation pattern 的匹配是文本级算子序列匹配(
PatternMatchPass基于torch._inductor.pattern_matcher),按旧分解结构编写的RMSNormPattern与RopePattern在 torch 2.9+ 环境全部失配:PatternMatchPass replace 190 patterns(仅 AdaLayerNorm 114 + GELU 76,RMSNorm/RoPE 均 0 匹配)rms_norm经 core decomposition 表展开为_to_copy → pow → mean → add → rsqrt → mul → mul → _to_copy,第一个 dtype cast 节点从_npu_dtype_cast变为_to_copy,第一个节点不匹配则整链放弃本 PR 的修改:
rms_norm_pattern.py/rope_pattern.py中按torch.__version__版本分支选择 dtype cast 算子:torch < 2.9:保持原有_npu_dtype_cast匹配逻辑,行为完全不变torch >= 2.9:pattern 适配为torch.ops.aten._to_copy.default(dtype为关键字参数)rope_pattern.py中x_out.type_as(x)未赋值的 no-op bug,并补充真实图中存在的末尾.to(x.dtype)cast 节点(diffusersapply_rotary_emb源码固有结构)Test Plan
环境要求:torch 2.9 / torch_npu 2.9(或更新)、cache-dit 源码安装、FLUX.1-dev 权重。
# 1. 安装本 PR 分支 python3 -m pip install -e . # 2. 编译匹配数验证(cache-dit 标准 CLI,预期 494) MINDIE_LOG_LEVEL=debug python3 -m cache_dit.generate flux \ --model-path /data/weights/FLUX.1-dev \ --compile --warmup 2 --repeat 2 2>&1 | grep "PatternMatchPass replace" # 预期: 494 = 190(AdaLayerNorm+GELU) + 152(RMSNorm) + 152(RoPE) # 3. profiling 验证融合算子真实执行(torch_npu profiler 采集) # 在推理脚本中用 torch_npu.profiler.profile 包裹 generate 调用,导出后: grep -ci "rmsnorm\|rotary" ./profiling/*/kernel_details.csv # 预期: 出现 aclnnRmsNorm / RotaryPositionEmbedding,eager 对照为 0Test Report
PatternMatchPass replace190 → 342(+152),profiling 中npu_rms_norm调用 152×steps 次(4 步采样 608 次),融合真实生效RotaryPositionEmbedding融合算子(456 次),eager 对照为 0 次rms_norm_symint/rms_norm_composite/_to_copydecomposition 与 2.9 完全一致(layer_norm.cppdiff 为空),本修复在 2.10 上继续有效Summary
Name Stmts Miss Branch BrPart Cover Missing
mindiesd/env.py 10 0 0 0 100%
mindiesd/eplb/init.py 4 0 0 0 100%
mindiesd/eplb/collector.py 28 19 2 0 30% 21-29, 32-36, 39-44, 47
mindiesd/eplb/dispatcher.py 71 53 0 0 25% 22-34, 37-43, 46-47, 50, 53, 56, 59, 64-79, 82, 85-88, 91-109, 112-127
mindiesd/eplb/eplb_scheduler.py 128 45 14 2 61% 53-64, 127-136, 176-184, 194-195, 200-208, 222-243, 247-248
mindiesd/eplb/greedy_algorithm.py 321 197 108 8 35% 97, 113-123, 131, 138, 180-182, 223-225, 256-339, 371, 374-469, 505-512, 519-628, 632-653, 699, 721-724
...... .... ... .... .. ..
mindiesd/quantization/utils.py 175 26 66 13 80% 36-48, 56, 72, 78, 123, 146, 156, 162, 183, 255, 259, 274, 285, 301
mindiesd/share_memory.py 106 106 26 0 0% 13-202
mindiesd/utils/init.py 3 0 0 0 100%
mindiesd/utils/env.py 114 17 52 19 78% 54, 132, 135, 138, 141, 144, 147, 150, 153, 156, 159, 162, 166->172, 169, 174-175, 178->184, 181, 185->exit, 188
..........................................................................................
TOTAL 7152 2476 2326 369 63%