已合并
[feat][compilation]按torch版本兼容rms_norm/rope pattern的dtype cast分解 #519
zqxu创建于 23 天前
[feat][compilation]按torch版本兼容rms_norm/rope pattern的dtype cast分解 #519
已合并
zqxu创建于 23 天前
zqxu
23 天前

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),按旧分解结构编写的 RMSNormPatternRopePattern 在 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 的修改:

  1. 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.defaultdtype 为关键字参数)
  2. 修复 rope_pattern.pyx_out.type_as(x) 未赋值的 no-op bug,并补充真实图中存在的末尾 .to(x.dtype) cast 节点(diffusers apply_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 对照为 0

Test Report

  • 环境:torch 2.9 / torch_npu 2.9 / 910B
  • RMSNorm:PatternMatchPass replace 190 → 342(+152),profiling 中 npu_rms_norm 调用 152×steps 次(4 步采样 608 次),融合真实生效
  • RoPE:profiling 中出现 RotaryPositionEmbedding 融合算子(456 次),eager 对照为 0 次
  • torch 2.10 源码级验证:rms_norm_symint / rms_norm_composite / _to_copy decomposition 与 2.9 完全一致(layer_norm.cpp diff 为空),本修复在 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%

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 zqxu 的贡献)
Zzqxu
23 天前 创建了 pull request,commit 23010e1b
atomgit-bot
atomgit-bot
23 天前 评论:

变更摘要

本 PR 修复 torch_npu 2.9+ 环境下 RMSNormPatternRopePattern 编译 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.pyrope_pattern.py 中新增 IS_TORCH_GE_29 版本判断,按 torch.__version__ 分支选择 dtype cast 算子:旧版本保持 _npu_dtype_cast 不变,2.9+ 版本改用 torch.ops.aten._to_copy.defaultdtype 为关键字参数);同时修复 rope_pattern.pyx_out.type_as(x) 未赋值的 no-op bug,并补上真实图中存在的末尾 .to(x.dtype) cast 节点,使 RoPE 融合在 torch 2.9+ 上重新匹配生效。

主要改动

  • 版本分支选择 dtype cast 算子:在 rms_norm_pattern.pyrope_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 与 diffusers apply_rotary_emb 真实图结构一致,保证末尾 .to(x.dtype) 节点能够被 pattern 匹配。
  • 恢复 RMSNorm/RoPE 融合匹配:修复后 PatternMatchPass 的替换数由 190 提升至 342(RMSNorm +152)、494 总量(含 RoPE +152),profiling 验证 npu_rms_normRotaryPositionEmbedding 融合算子真实执行,并在 torch 2.10 上同样有效。
likedislike
不准确?
Zzqxu
23 天前 关联了看板:MindIE-SD
atomgit-bot
atomgit-bot
23 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
Xxiangjie10成员
23 天前 添加了label:pr-audit-failed
此处折叠了91条消息 查看更多
ascend-robotascend-robot成员
21 天前 添加了label:ci-pipeline-passed
ascend-robot
ascend-robot成员
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 >>>
此流水线已支持下列评论快捷指令,仅PR创建者和白名单成员评论有效
  • compile : 运行流水线
  • retry : 重试流水线所有失败子任务
  • retry <任务名> : 仅重试指定失败子任务
  • stop : 停止流水线
likedislike
ascend-robotascend-robot成员
21 天前 关闭了关联的issue
ascend-robotascend-robot成员
21 天前 合入了pull request
ascend-robot
ascend-robot成员
21 天前 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike