Merged

feat: add torch.compile support(take engram_fused_weight as example) #47

ChaoyangJiAdd on 20 days agoMerged fromChaoyangJi/ops-tilelang： master into cann/ops-tilelang： master

feat: add torch.compile support(take engram_fused_weight as example) #47

Merged

ChaoyangJiAdd on 20 days ago

Merged fromChaoyangJi/ops-tilelang masterintocann/ops-tilelang master

ChaoyangJiCollaborators

20 days ago

当前PR是否有AI参与:

[x] 否

[ ] 是
__1. AI Agent 平台:
__2. AI 模型:
__3. Prompt上下文 :

PR功能描述 / 为什么需要这个合入**:

   以engram_fused_weight为例，增加torch.compile的支持，可以支持aclgraph的图接入；

   另外，修改：
   1.  engram_fused_weight通过返回值返回tensor，不通过原来的参数out返回；
   2.  增加对pytorch版本的要求
   3. 同时支持eager和compile两种方式的调用和测试

该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

Pull Request has been successfully merged by CANN-robot.

（Thank you for ChaoyangJi's contribution.）

ChaoyangJiCollaborators

20 days ago Created pull request, commit 16cfec00

CANN-robotCollaborators

20 days ago Added label: cann-cla/yes

CANN-robotCollaborators

20 days ago 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, jiawei_zhuang (2/2)	✅ kevin365 (1/1)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

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

CANN-robotCollaborators

20 days ago Set kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver

CANN-robotCollaborators

20 days ago Set kevin365,monologue815,jiang-lirui,Allan_Yu,jiawei_zhuang,liudan12,mabing1118 as approver

ChaoyangJiCollaborators

20 days ago Comments：

compile

ChaoyangJiCollaborators

20 days ago Pre merge succeeded (commit_id: 63c2596b3788a9333b5ecfcc41911658a1a135f1)

CANN-robotCollaborators

20 days ago Comments：

🚀 CI 流水线已启动
📋 执行详情: 点击查看流水线

CANN-robotCollaborators

20 days ago Added label: ci-pipeline-running

CANN-robotCollaborators

20 days ago Removed label: ci-pipeline-running

CANN-robotCollaborators

20 days ago Added label: ci-pipeline-failed

CANN-robotCollaborators

20 days ago Comments：

流水线分析报告 — ops-tilelang run#109

流水线: https://gitcode.com/CANN/ops-tilelang/actions/runs/0d173d3d924b4301b15b60b44e9b5887 | PR#47 | 2026-10-23 20:07:38 ~ 2026-10-23 20:12:25 | 失败 Job 1 个

定界结论: 环境问题 / 设备NPU异常 — Runner 上 NPU 设备 3 启动失败(err 507033)

失败类型: 环境问题(设备/NPU异常, 可重试)
根因: ST_test 以 6 个 xdist worker 运行于 NPU 设备 2,3, 分配到设备 3 的 3 个 worker(gw1/gw3/gw5)在首个 NPU 调用处瞬时失败: TsdOpen/rtSetDevice 报 device retain error(错误码 507033, Failed to start the device); 跑在健康设备上的用例全部通过
处置建议: ① 直接重新触发流水线(设备瞬时占用/异常, 大概率恢复)
② 重试仍失败 → 联系 CI/环境维护者检查该 Runner 的 NPU 设备 3(npu-smi 确认设备健康状态与占用, 排除同机并发任务占用)
排除项: 排除代码问题 — 两轮 pytest session 失败用例集合不同(hidden=3072 第一轮通过/第二轮失败, hidden=2560 反之), 失败仅与 worker/设备绑定相关, 与用例内容无关; 失败均为设备初始化错误(非断言/编译/逻辑错误), PR 变更用例在健康 worker 上全部执行通过

失败根因分析

根因 #1: 环境问题 / 设备NPU异常 — NPU devId=3 启动失败(TsdOpen/rtSetDevice, err 507033)

判据: 设备启动失败: rtSetDevice execution failed(device retain error) + TsdOpen failed devId=3, 错误码 507033; 失败按 worker 而非按用例聚簇
定位步骤: 3_run_st_test.log
影响 Job: ST_test
关键上下文:
- Running regular pytest cases on NPU 2,3 with 6 workers(6 个 worker 共用设备 2,3)
- 第一轮(TILELANG_DEFAULT_TARGET=auto): 3 failed, 4 passed; 失败 worker 为 gw1/gw3/gw5(用例 6144/2560/7168, 均在 session 启动后 1 秒内瞬时失败)
- 第二轮(TILELANG_DEFAULT_TARGET=pto): 3 failed, 4 passed; 失败 worker 仍为 gw1/gw3/gw5(用例变为 3072/6144/7168) — 用例分配随调度变化, 失败 worker 恒定
- 跑在 gw0/gw2/gw4(健康设备)上的用例(两轮共 4 个不同参数用例及 compile 用例)全部通过
根因证据(日志原文):

[2026-10-23 20:10:27.817+08:00] Running regular pytest cases on NPU 2,3 with 6 workers
[2026-10-23 20:10:46.552+08:00] [gw5] [ 14%] FAILED tests/engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=6144,]   ← session 启动后 <1s 瞬时失败
E       RuntimeError: LazySetDevice:../torch_npu/csrc/core/npu/NPUFunctions.cpp:342 NPU function error: err, error code is 507033
E       [ERROR] 2026-10-23-20-10-46 (PID:1607796, Device:0, RankID:-1) ERR00100 PTA call acl api failed
E       [Error]: Failed to start the device.
E               TsdOpen failed. devId=3, tdt error=1.[FUNC:PrintfTsdError][FILE:runtime.cc][LINE:2194]
E               Check param failed, dev can not be NULL![FUNC:DeviceRetain][FILE:runtime.cc][LINE:3238]
E               rtSetDevice execution failed, reason=device retain error[FUNC:FuncErrorReason][FILE:error_message_manage.cc][LINE:76]

结构化定界摘要

[{"repo": "ops-tilelang", "run_number": 109, "Failed_Type": "环境问题",
  "Failed_Sub_Type": "设备/NPU异常",
  "Failed_Reason": "ST_test 的 6 个 xdist worker 运行于 NPU 设备 2,3, 分配到设备 3 的 gw1/gw3/gw5 在首个 NPU 调用处瞬时失败: TsdOpen/rtSetDevice device retain error(错误码 507033, Failed to start the device); 两轮 session 失败用例集合随调度变化而失败 worker 恒定, 健康设备上的用例全部通过",
  "evidence": "TsdOpen failed. devId=3, tdt error=1.[FUNC:PrintfTsdError][FILE:runtime.cc][LINE:2194] Check param failed, dev can not be NULL! rtSetDevice execution failed, reason=device retain error",
  "suggestion": "重新触发流水线; 若复现联系 CI/环境维护者检查该 Runner 的 NPU 设备 3 健康状态与占用",
  "jobs": "ST_test"}]

由门禁失败定界分析自动生成 | repo: ops-tilelang | run: 0d173d3d924b4301b15b60b44e9b5887 | PR#47

wangliweiCollaborators19 days agoStart code review2

tests/engram/test_engram_fused_weight.py

3636

3737

38+@torch.inference_mode()

39+def test_engram_fused_weight_compile():

wangliwei19 days agoComments：

test_engram_fused_weight_compile 当前仅调用一次 compiled 函数。建议复用同一个 compiled 函数，覆盖原地修改输入、传入同形状新 Tensor 后的多次调用，并逐次与 FP32 参考结果比较，确认图重放能够读取最新输入，避免旧结果导致误判。

ChaoyangJiCollaborators

17 days ago Comments：

已修改，对同一个compiled kernel，分别进行初始输入，修改输入，同形状的新tensor的测试验证。

wangliweiCollaborators19 days agoStart code review2

src/cann_ops_tilelang/engram/engram_fused_weight.py

65+def engram_fused_weight(

66+ weight_hidden: torch.Tensor,

67+ weight_embed: torch.Tensor,

68+) -> torch.Tensor:

wangliwei19 days agoComments：

This code comment section is from+65To+68

删除 out= 后，已有 engram_fused_weight(a, b, out=buffer) 调用会直接报 TypeError。建议保留 eager 路径的可选 out= 支持；如果确定统一为返回值接口，请在文档中明确兼容性变化及迁移方式。

ChaoyangJiCollaborators

17 days ago Comments：

去掉out参数，和原来tile_kernel开发的时候对齐，这个out参数是后面ops-tilelang迁移的时候引入的，这里去掉，不引入兼容性问题。

ChaoyangJiCollaborators

17 days ago Pre merge succeeded (commit_id: 63a362b72e288ae3e02401b9c37900109bc524c3)

ChaoyangJiCollaborators

17 days ago Pushed 1 commits:844483a7-fix: enhance test case to cover different input of the compiled kernel

ChaoyangJiCollaborators

17 days ago Pre merge succeeded (commit_id: e1158f1668d4b9b95fb17d518336f0ff35709fd5)

CANN-robotCollaborators

17 days ago Removed label: cann-cla/yes

CANN-robotCollaborators

17 days ago Added label: stat/needs-squash

CANN-robotCollaborators

17 days ago Added label: cann-cla/yes

ChaoyangJiCollaborators

17 days ago Comments：

compile

ChaoyangJiCollaborators

17 days ago Pre merge succeeded (commit_id: 23c22a052c02ae255641cf88215040287d670864)

CANN-robotCollaborators

17 days ago Comments：

🚀 CI 流水线已启动
📋 执行详情: 点击查看流水线

CANN-robotCollaborators

17 days ago Removed label: ci-pipeline-failed

CANN-robotCollaborators

17 days ago Added label: ci-pipeline-running

CANN-robotCollaborators

17 days ago Removed label: ci-pipeline-running

CANN-robotCollaborators

17 days ago Added label: ci-pipeline-failed

CANN-robotCollaborators

17 days ago Comments：

流水线分析报告 — ops-tilelang run#114

流水线: https://gitcode.com/CANN/ops-tilelang/actions/runs/917120e5ff3440cc825dfdf9afe5b85d | PR#47 | 2026-10-26 11:39:05 ~ 2026-10-26 11:44:37 | 失败 Job 1 个

定界结论: 环境问题 / 设备/NPU异常 — runner 侧 NPU 设备3启动失败(TsdOpen/rtSetDevice, 507033)

失败类型: 环境问题(设备/NPU异常)
根因: ST_test 以 6 个 pytest-xdist worker 共用 NPU 设备 2、3 执行 tests/engram/test_engram_fused_weight.py, 落到 devId=3 的 worker 在首个 NPU 算子(torch.randn(..., device=get_device()), tests/engram/test_engram_fused_weight.py:19)处设备启动失败——TsdOpen failed. devId=3, tdt error=31 → Check param failed, dev can not be NULL![DeviceRetain] → rtSetDevice execution failed, reason=device retain error(torch_npu 错误码 507033, "Failed to start the device")。失败集合跟随 worker→设备分配而非用例身份(同一用例在 auto/pto 两轮目标模式下成败互换), 是设备侧故障不是代码缺陷。
处置建议: ① 直接重新触发流水线(设备瞬时状态异常/占用, 大概率恢复); ② 重试仍失败 → 转 CI/环境维护者检查该 runner 的 NPU 设备 3(npu-smi info 健康状态/占用进程/驱动日志), 修复后重跑。
排除项(为何不是代码问题):
- PR#47 新增的 test_engram_fused_weight_compile(新 torch.library.custom_op + npugraph_ex 编译路径)两轮均 PASSED——PR 新代码在健康设备上工作正常;
- 失败的参数化用例是 PR 未修改的既有 eager 用例, 失败点在设备初始化(测试逻辑执行之前), 非断言/逻辑错误;
- 同一参数化用例结果随模式翻转(hidden=4096: auto 轮 PASSED / pto 轮 FAILED; hidden=7168: auto 轮 FAILED / pto 轮 PASSED), 失败由设备分配决定, 与用例代码无关;
- 本 run 其余 11 个 Job(含 CodeCheck/precommit/SCA/Antipoison)全部 COMPLETED, 唯 ST_test 失败。

失败根因分析

根因 #1: 环境问题 / 设备/NPU异常 — NPU 设备3启动失败: TsdOpen devId=3 tdt error=31, rtSetDevice device retain error(507033)

判据: rtSetDevice execution failed, reason=device retain error / TsdOpen failed. devId=3 / [Error]: Failed to start the device. → 环境问题/设备/NPU异常(device); 表层 pytest FAILED 条目由设备初始化失败传导, 非用例断言失败
定位步骤: 3_run_st_test.log(Job ST_test, 唯一失败 Job)
影响 Job: ST_test
失败用例(两轮共 6 个 FAILED 条目, 去重 4 个参数化用例):
- tests/engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=6144,](auto/pto 两轮均败)
- tests/engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=2560,](auto/pto 两轮均败)
- tests/engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=7168,](仅 auto 轮败)
- tests/engram/test_engram_fused_weight.py::test_engram_fused_weight[hc=4,hidden=4096,](仅 pto 轮败)

两轮对照(同一份代码, 结果随设备分配翻转):

参数化用例	auto 轮(TILELANG_DEFAULT_TARGET=auto, 11:43)	pto 轮(TILELANG_DEFAULT_TARGET=pto, 11:44)
hidden=6144	FAILED(gw5)	FAILED(gw3)
hidden=2560	FAILED(gw1)	FAILED(gw1)
hidden=7168	FAILED(gw3)	PASSED(gw0)
hidden=4096	PASSED(gw4)	FAILED(gw5)
hidden=2048 / 3072	PASSED	PASSED
test_engram_fused_weight_compile(PR 新增)	PASSED	PASSED

两轮失败测试全部落在 gw1/gw3/gw5 三个 worker, 且所有错误块均指向 devId=3——失败因子是该设备, 而非用例代码。

环境上下文: Running regular pytest cases on NPU 2,3 with 6 workers; 两轮汇总均为 3 failed, 4 passed(auto 轮 53.53s, pto 轮 33.27s)。
根因证据(日志原文):

[2026-10-26 11:43:47.412+08:00] E       RuntimeError: LazySetDevice:../torch_npu/csrc/core/npu/NPUFunctions.cpp:342 NPU function error: err, error code is 507033
[2026-10-26 11:43:47.412+08:00] E       [ERROR] 2026-10-26-11:43:07 (PID:771415, Device:0, RankID:-1) ERR00100 PTA call acl api failed
[2026-10-26 11:43:47.412+08:00] E       [Error]: Failed to start the device. 
[2026-10-26 11:43:47.412+08:00] E               Rectify the fault based on the error information in the ascend log.
[2026-10-26 11:43:47.412+08:00] E               TsdOpen failed. devId=3, tdt error=31.[FUNC:PrintfTsdError][FILE:runtime.cc][LINE:2194]
[2026-10-26 11:43:47.412+08:00] E               Check param failed, dev can not be NULL![FUNC:DeviceRetain][FILE:runtime.cc][LINE:3238]
[2026-10-26 11:43:47.412+08:00] E               rtSetDevice execution failed, reason=device retain error[FUNC:FuncErrorReason][FILE:error_message_manage.cc][LINE:76]

触发点(首个 NPU 算子, 测试逻辑执行之前):

    def generate_test_data(params):
        hc_mult = params['hc']
        hidden_size = params['hidden']
>       wh_data = torch.randn(hc_mult, hidden_size, dtype=torch.bfloat16, device=get_device())
E       RuntimeError: LazySetDevice:../torch_npu/csrc/core/npu/NPUFunctions.cpp:342 NPU function error: err, error code is 507033
tests/engram/test_engram_fused_weight.py:19: RuntimeError

结构化定界摘要

[{"repo": "ops-tilelang", "run_number": 114, "Failed_Type": "环境问题",
  "Failed_Sub_Type": "设备/NPU异常",
  "Failed_Reason": "ST_test 6个pytest worker共用NPU 2,3, 落到devId=3的用例在首个NPU算子处设备启动失败: TsdOpen failed devId=3 tdt error=31, rtSetDevice device retain error(507033, Failed to start the device); 同一用例在auto/pto两轮成败互换, 与代码无关",
  "evidence": "RuntimeError: LazySetDevice:../torch_npu/csrc/core/npu/NPUFunctions.cpp:342 NPU function error: err, error code is 507033 | TsdOpen failed. devId=3, tdt error=31.[FUNC:PrintfTsdError][FILE:runtime.cc][LINE:2194] | rtSetDevice execution failed, reason=device retain error[FUNC:FuncErrorReason][FILE:error_message_manage.cc][LINE:76]",
  "suggestion": "先重新触发流水线(设备瞬时异常大概率恢复); 重试仍失败则转CI/环境维护者检查该runner的NPU设备3(npu-smi健康状态/占用进程/驱动日志)",
  "jobs": "ST_test"}]

由门禁失败定界分析自动生成 | repo: ops-tilelang | run: 917120e5ff3440cc825dfdf9afe5b85d | PR#47

ChaoyangJiCollaborators

17 days ago Comments：

/compile

ChaoyangJiCollaborators

17 days ago Pre merge succeeded (commit_id: 208e2258c7df2b58b466d4b8b46a2df7fb6bb9be)

CANN-robotCollaborators

17 days ago Comments：

🚀 CI 流水线已启动
📋 执行详情: 点击查看流水线

CANN-robotCollaborators

17 days ago Removed label: ci-pipeline-failed

CANN-robotCollaborators

17 days ago Added label: ci-pipeline-running

CANN-robotCollaborators

17 days ago Removed label: ci-pipeline-running

CANN-robotCollaborators

17 days ago Added label: ci-pipeline-passed

wangkaiCollaborators

16 days ago Comments：

/approve
/lgtm

CANN-robotCollaborators

16 days ago Added label: approved

wangliweiCollaborators16 days agoStart code review2

docs/architecture.md

122122

123+## Engram fused weight 的 torch.compile 接入

124+

125+`engram_fused_weight` 支持固定形状推理的 `npugraph_ex` 编译路径：

wangliwei16 days agoComments：

非固定形状能支持吗？

ChaoyangJiCollaborators

16 days ago Comments：

支持。图这层是支持的接收不同的shape，同一个compiled的kernel是支持不同的shape输入的。已验证。
和算子层的动态shape的支持是解耦的关系（和算子实现有关）。

W wangliweiCollaborators

16 days ago Resolved final issue

wangliweiCollaborators

15 days ago Comments：

/lgtm

jiawei_zhuangCollaborators

15 days ago Comments：

/lgtm

CANN-robotCollaborators

15 days ago Comments：

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
/lgtm: wangliwei12

CANN-robotCollaborators

15 days ago Added label: lgtm

CANN-robotCollaborators

15 days ago Merged pull request