[fix] Use bundled AscendC backend for torch.compile-torchtitan-npu-AtomGit

M.agents/skills/torchtitan-npu-env-setup/SKILL.md+18-6

	@@ -153,15 +153,27 @@ pip install -e .
153	> pip uninstall torchtitan_npu	153	> pip uninstall torchtitan_npu
154	> ```	154	> ```
155		155
156	-### Step 4：可选特性 — 算子自动融合支持	156	+### Step 4：可选特性 — torch.compile 支持
157		157
158	-如需启用 `torch.compile` 编译链路下的 NPU Codegen 后端，按文档执行：	158	+兼容版本的 `torch_npu` 已内置 AscendC Codegen 后端，无需单独安装 `torchair` 或
		159	+`inductor-npu-ext`。通过 `torchtitan_npu.entry` 启动训练时，使用以下命令启用
		160	+`torch.compile`，本仓会自动选择 AscendC 后端：
159		161
160	```bash	162	```bash
161	-git clone https://gitcode.com/Ascend/torchair.git	163	+export TORCHINDUCTOR_SIZE_ASSERTS=0
162	-cd torchair/experimental/_inductor_npu_ext/	164	+bash scripts/run_train.sh --compile.enable
163	-pip3 install -e ./python/	165	+```
164	-cd -	166	+
		167	+在独立代码中直接调用 `torch.compile` 时，通过 `options` 显式选择 AscendC 后端：
		168	+
		169	+```python
		170	+import torch
		171	+import torch_npu # noqa: F401
		172	+
		173	+compiled_fn = torch.compile(
		174	+ fn,
		175	+ options={"npu_backend": "ascendc"},
		176	+)
165	```	177	```
166		178
167	功能说明参考：	179	功能说明参考：

M.ci/smoke_test.sh+4-13

	@@ -22,15 +22,16 @@ INTEGRATION_REPORT_DIR="${PROJECT_ROOT}/test_reports/integration_tests"
22	TORCHTITAN_BRANCH="main"	22	TORCHTITAN_BRANCH="main"
23	TORCHTITAN_COMMIT="ac13e536c84e7f6647b14fa9375c3c8a8a2b8578"	23	TORCHTITAN_COMMIT="ac13e536c84e7f6647b14fa9375c3c8a8a2b8578"
24	TORCHTITAN_DIR="${PROJECT_ROOT}/third_party/torchtitan"	24	TORCHTITAN_DIR="${PROJECT_ROOT}/third_party/torchtitan"
25	-TORCHAIR_COMMIT="3c9418c2"
26	-TORCHAIR_DIR="${PROJECT_ROOT}/third_party/torchair"
27	DEEPSEEK_TOKENIZER_REPO="${DEEPSEEK_TOKENIZER_REPO:-https://gitcode.com/hitwdy/deepseekv4.git}"	25	DEEPSEEK_TOKENIZER_REPO="${DEEPSEEK_TOKENIZER_REPO:-https://gitcode.com/hitwdy/deepseekv4.git}"
28	DEEPSEEK_V4_TOKENIZER_DIR="${PROJECT_ROOT}/tests/assets/tokenizer/deepseekv4_tokenizer"	26	DEEPSEEK_V4_TOKENIZER_DIR="${PROJECT_ROOT}/tests/assets/tokenizer/deepseekv4_tokenizer"
29	DEEPSEEK_V32_TOKENIZER_DIR="${PROJECT_ROOT}/tests/assets/tokenizer/deepseekv32_tokenizer"	27	DEEPSEEK_V32_TOKENIZER_DIR="${PROJECT_ROOT}/tests/assets/tokenizer/deepseekv32_tokenizer"
30	TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-300}	28	TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-300}
31	SMOKE_STEPS=${SMOKE_STEPS:-1}	29	SMOKE_STEPS=${SMOKE_STEPS:-1}
32	# Known false-positive patterns to exclude from error detection	30	# Known false-positive patterns to exclude from error detection
33	-ERROR_EXCLUDE_PATTERNS=("TORCH_NCCL_ASYNC_ERROR_HANDLING")	31	+ERROR_EXCLUDE_PATTERNS=(
		32	+ "TORCH_NCCL_ASYNC_ERROR_HANDLING"
		33	+ "HCCL_ASYNC_ERROR_HANDLING is deprecated"
		34	+)
34		35
35	export PYTHONPATH="${PROJECT_ROOT}/tests/smoke_tests/npu_bypass_triton_codegen:${PYTHONPATH:-}"	36	export PYTHONPATH="${PROJECT_ROOT}/tests/smoke_tests/npu_bypass_triton_codegen:${PYTHONPATH:-}"
36		37
	@@ -43,16 +44,6 @@ _setup_env() {
43	python3 -m pip install -e .	44	python3 -m pip install -e .
44	fi	45	fi
45		46
46	- # Ensure the torch.compile NPU codegen extension is installed.
47	- echo "Installing inductor_npu_ext from torchair ${TORCHAIR_COMMIT}..."
48	- mkdir -p third_party
49	- if [[ ! -d "$TORCHAIR_DIR/.git" ]]; then
50	- git clone https://gitcode.com/Ascend/torchair.git "$TORCHAIR_DIR"
51	- fi
52	- git -C "$TORCHAIR_DIR" fetch origin
53	- git -C "$TORCHAIR_DIR" checkout "$TORCHAIR_COMMIT"
54	- pip3 install -e "$TORCHAIR_DIR/experimental/_inductor_npu_ext/python/"
55	-
56	# Clone torchtitan source if not exists	47	# Clone torchtitan source if not exists
57	if [[ ! -d "$TORCHTITAN_DIR/.git" ]]; then	48	if [[ ! -d "$TORCHTITAN_DIR/.git" ]]; then
58	echo "Cloning torchtitan source..."	49	echo "Cloning torchtitan source..."

Mdocs/community/versioning_policy.md+1-1

	@@ -9,7 +9,7 @@
9		9
10	\| torchtitan-npu \| torchtitan \| Python \| Stable CANN \| PyTorch/torch_npu \| triton-ascend \|	10	\| torchtitan-npu \| torchtitan \| Python \| Stable CANN \| PyTorch/torch_npu \| triton-ascend \|
11	\| --- \| --- \| --- \| --- \| --- \| --- \|	11	\| --- \| --- \| --- \| --- \| --- \| --- \|
12	-\| master \| main (ac13e536c84e7f6647b14fa9375c3c8a8a2b8578) \| 3.11.x \| 9.0.0 \| 2.12.0 / 2.12.0rc1 \| 3.2.1 \|	12	+\| master \| main (ac13e536c84e7f6647b14fa9375c3c8a8a2b8578) \| 3.11.x \| 9.0.0 \| 2.12.0 / 2.12.0.dev20260808 \| 3.2.1 \|
13	\| v0.2.2-dev \| v0.2.2 (73a0e6979dd10b6b1904098eb3c8f62c18ab87ce) \| 3.11.x \| 9.0.0 \| 2.10.0 / 2.10.0 \|3.2.1 \|	13	\| v0.2.2-dev \| v0.2.2 (73a0e6979dd10b6b1904098eb3c8f62c18ab87ce) \| 3.11.x \| 9.0.0 \| 2.10.0 / 2.10.0 \|3.2.1 \|
14		14
15	对于活跃开发分支，请始终以 `分支同步表` 为准。	15	对于活跃开发分支，请始终以 `分支同步表` 为准。

Mdocs/feature_guides/muon_optimizer.md+0-2

	@@ -38,8 +38,6 @@ $$X_{k+1} = a \cdot X_k + b \cdot X_k X_k^T X_k + c \cdot (X_k X_k^T)^2 X_k$$
38		38
39	Muon 支持只对 Newton-Schulz 张量函数启用 `torch.compile`，编译边界只覆盖 `zeropower_via_newtonschulz5()`，其余逻辑仍保持 eager 执行。	39	Muon 支持只对 Newton-Schulz 张量函数启用 `torch.compile`，编译边界只覆盖 `zeropower_via_newtonschulz5()`，其余逻辑仍保持 eager 执行。
40		40
41	-当前验证环境版本：`torch_npu 2.12.0.20260805`、`CANN 9.2.0 (20260730)`。当前版本的代码暂不支持同时开启 DSV4 的 `npu_smla` converter。
42	-
43	启用方式是打开全局 compile，并在 `components` 中包含 `"muon"`：	41	启用方式是打开全局 compile，并在 `components` 中包含 `"muon"`：
44		42
45	```python	43	```python

Mdocs/feature_guides/torch_compile.md+15-9

	@@ -9,23 +9,27 @@ torch.compile 是 PyTorch 2.0 的核心特性。通过 JIT （即时编译），
9	<img src="../assets/include_npu_ext.png" style="width:80%; max-width: 1200px" >	9	<img src="../assets/include_npu_ext.png" style="width:80%; max-width: 1200px" >
10	</p>	10	</p>
11		11
12	-为了在 NPU 平台上充分利用 `torch.compile` 原生的编译能力，`torchtitan_npu` 在保留 Dynamo 与 Inductor 既有编译流程的基础上，接入了 Codegen 后端 [`inductor-npu-ext`](https://gitcode.com/Ascend/torchair/blob/master/experimental/_inductor_npu_ext/README.md)。该后端借助 [AutoFuse](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/900beta1/graph/graphguide/autofuse_1_0001.html) 的自动融合能力，从 Inductor IR 生成 AscendC 融合 Kernel。	12	+为了在 NPU 平台上充分利用 `torch.compile` 原生的编译能力，`torchtitan_npu` 在保留 Dynamo 与 Inductor 既有编译流程的基础上，使用 [`torch_npu` 内置的 AscendC Codegen 后端](https://gitcode.com/Ascend/torchair/blob/master/experimental/_inductor_npu_ext/README.md)。该后端借助 [AutoFuse](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/900beta1/graph/graphguide/autofuse_1_0001.html) 的自动融合能力，从 Inductor IR 生成 AscendC 融合 Kernel。
13		13
14	## 支持范围	14	## 支持范围
15	torchtitan-npu 当前支持 `DeepSeek-V3、DeepSeek-V3.2、DeepSeek-V4` 模型的全流程编译。	15	torchtitan-npu 当前支持 `DeepSeek-V3、DeepSeek-V3.2、DeepSeek-V4` 模型的全流程编译。
16		16
17	## torch.compile 示例	17	## torch.compile 示例
18		18
19	-### 1. 安装 inductor_npu_ext	19	+### 1. 使用 AscendC Codegen 后端
		X X xuyujun10 天前这个用法也提交一个PR到override分支吧，作为默认用法
20		20
21	-inductor_npu_ext 需要从源码安装。在运行环境内执行以下命令：	21	+AscendC Codegen 后端已随兼容版本的 `torch_npu` 打包在 `torch_npu/_inductor/ascendc` 中，无需单独安装。通过 `torchtitan_npu.entry` 启动训练并开启 `torch.compile` 时，本仓会自动选择该后端。
22		22
23	-```bash	23	+在独立代码中直接调用 `torch.compile` 时，可通过 `options` 显式选择：
24	-git clone https://gitcode.com/Ascend/torchair.git	24	+
25	-cd torchair	25	+```python
26	-git checkout 3c9418c2	26	+import torch
27	-pip3 install -e experimental/_inductor_npu_ext/python/	27	+import torch_npu # noqa: F401
28	-cd -	28	+
		29	+compiled_fn = torch.compile(
		30	+ fn,
		31	+ options={"npu_backend": "ascendc"},
		X X xuyujun10 天前 .agents/skills/torchtitan-npu-env-setup/SKILL.md 未更新，需要更新下里面的torch.compile用法张伟 9 天前评论： 9 天前评论：
		32	+)
29	```	33	```
30		34
31	### 2. 配置 compile	35	### 2. 配置 compile
	@@ -39,6 +43,8 @@ compile = CompileConfig(
39	enable=True,	43	enable=True,
40	# 编译完整模型	44	# 编译完整模型
41	components=["model", "loss", "muon"],	45	components=["model", "loss", "muon"],
		46	+ # Dynamo 使用 Inductor；AscendC 是 Inductor 内部的 NPU Codegen 后端
		47	+ backend="inductor",
42	)	48	)
43	```	49	```
44		50

Mdocs/lint_guide.md+1-1

	@@ -9,7 +9,7 @@ pre-commit run --all-files
9	为了正确运行 pre-commit 检查，需要安装开发依赖：	9	为了正确运行 pre-commit 检查，需要安装开发依赖：
10		10
11	```bash	11	```bash
12	-pip install -r requirements.txt -r requirements_dev.txt	12	+python3 -m pip install -r requirements.txt -r requirements_dev.txt
13	# 下载依赖并运行检查	13	# 下载依赖并运行检查
14	pre-commit run --all-files	14	pre-commit run --all-files
15	```	15	```

Mdocs/test_guides/test_guide.md+1-1

	@@ -63,7 +63,7 @@ OverrideDefinitions(
63	#### 新增测试用例步骤	63	#### 新增测试用例步骤
64		64
65	1. 打开 `tests/smoke_tests/integration_test.py`	65	1. 打开 `tests/smoke_tests/integration_test.py`
66	-2. 在 `generate_smoke_tests()` 中，按用例类型将新配置添加到对应的分组函数（`_fake_backend_tests()` / `_base_tests()` / `_tp_tests()` / `_ep_tests()` / `_checkpoint_tests()`）的返回列表中（例如下面的 TP 用例加到 `_tp_tests()`）：	66	+2. 在 `generate_smoke_tests()` 中，按用例类型将新配置添加到对应的分组函数（`_fake_backend_tests()` / `_base_tests()` / `_tp_tests()` / `_ep_tests()`）的返回列表中（例如下面的 TP 用例加到 `_tp_tests()`）：
67	```python	67	```python
68	OverrideDefinitions(	68	OverrideDefinitions(
69	[	69	[

Mdocs/user-guides/installation.md+6-17

	@@ -6,7 +6,7 @@ torchtitan-npu支持Atlas 800T A3等昇腾训练硬件形态。软件版本配
6		6
7	\| torchtitan-npu版本 \| torchtitan版本 \| PyTorch版本 \| torch_npu版本 \| CANN版本 \| Python版本 \| Triton Ascend \|	7	\| torchtitan-npu版本 \| torchtitan版本 \| PyTorch版本 \| torch_npu版本 \| CANN版本 \| Python版本 \| Triton Ascend \|
8	\|------------------------\|-------------\|--------------\|-------------\|---------\|----------------------------------------\|--------------\|	8	\|------------------------\|-------------\|--------------\|-------------\|---------\|----------------------------------------\|--------------\|
9	-\| master \| main `ac13e536c84e7f6647b14fa9375c3c8a8a2b8578` \| 2.12.0 \| 2.12.0rc1 \| [9.2.0_daily0730](https://ascend.devcloud.huaweicloud.com/artifactory/cann-run-mirror/software/master/20260730143326087/) \| Python3.11.x \| 3.2.1 \|	9	+\| master \| main `ac13e536c84e7f6647b14fa9375c3c8a8a2b8578` \| 2.12.0 \| 2.12 daily版本（详见requirements.txt文件） \| [9.2.0_daily0805](https://ascend.devcloud.huaweicloud.com/artifactory/cann-run-mirror/software/legacy/20260805101249091/) \| Python3.11.x \| 3.2.1 \|
10	\| v0.2.2-dev \| v0.2.2 `73a0e6979dd10b6b1904098eb3c8f62c18ab87ce` \| 2.10.0 \| 2.10.0 \| 9.0.0 \| Python3.11.x \| 3.2.1	10	\| v0.2.2-dev \| v0.2.2 `73a0e6979dd10b6b1904098eb3c8f62c18ab87ce` \| 2.10.0 \| 2.10.0 \| 9.0.0 \| Python3.11.x \| 3.2.1
11		11
12	## 源码安装	12	## 源码安装
	@@ -52,8 +52,6 @@ torchtitan-npu支持Atlas 800T A3等昇腾训练硬件形态。软件版本配
52		52
53	> 注：安装 NNAL（Ascend Transformer Boost加速库）前，请先执行 `source /usr/local/Ascend/cann/set_env.sh` 配置 CANN 环境变量。	53	> 注：安装 NNAL（Ascend Transformer Boost加速库）前，请先执行 `source /usr/local/Ascend/cann/set_env.sh` 配置 CANN 环境变量。
54		54
55	-> 注：单独安装 torch_npu 2.12.0rc1 时，请同时安装 numpy 和 PyYAML，避免导入 torch_npu 时缺少依赖。
56	-
57	### 2. 下载torchtitan-npu源码master分支（请注意下列命令的大小写）	55	### 2. 下载torchtitan-npu源码master分支（请注意下列命令的大小写）
58		56
59		57
	@@ -65,8 +63,8 @@ git clone https://gitcode.com/cann/torchtitan-npu.git
65		63
66	```shell	64	```shell
67	cd torchtitan-npu	65	cd torchtitan-npu
68	-pip install -r requirements.txt	66	+python3 -m pip install -r requirements.txt
69	-pip install -e .	67	+python3 -m pip install -e .
70	```	68	```
71		69
72	> 注：如有旧版本torchtitan-npu，请先[卸载](#卸载)，再进行安装	70	> 注：如有旧版本torchtitan-npu，请先[卸载](#卸载)，再进行安装
	@@ -74,20 +72,11 @@ pip install -e .
74	> 注：`requirements.txt` 已为 `triton-ascend==3.2.1` 配置 `--extra-index-url https://triton-ascend.osinfra.cn/pypi/simple`；默认 PyPI 源当前不提供该版本，单独安装时请使用 `pip install --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple triton-ascend==3.2.1`。`triton-ascend` 是安装包名，实际 Python 模块名为 `triton`，代码中请使用 `import triton`，不要使用 `import triton_ascend`。	72	> 注：`requirements.txt` 已为 `triton-ascend==3.2.1` 配置 `--extra-index-url https://triton-ascend.osinfra.cn/pypi/simple`；默认 PyPI 源当前不提供该版本，单独安装时请使用 `pip install --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple triton-ascend==3.2.1`。`triton-ascend` 是安装包名，实际 Python 模块名为 `triton`，代码中请使用 `import triton`，不要使用 `import triton_ascend`。
75		73
76		74
77	-### 4. 安装inductor-npu-ext（可选，使能torch.compile时需要）	75	+### 4. 配置torch.compile（可选）
78		76
79	-为了在 NPU 平台上充分利用 `torch.compile` 原生的编译能力，`torchtitan_npu` 在保留 Dynamo 与 Inductor 既有编译流程的基础上，接入了 Codegen 后端 [`inductor-npu-ext`](https://gitcode.com/Ascend/torchair/blob/master/experimental/_inductor_npu_ext/README.md)。该后端借助 [AutoFuse](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/900beta1/graph/graphguide/autofuse_1_0001.html) 的自动融合能力，从 Inductor IR 生成 AscendC 融合 Kernel。	77	+AscendC Codegen 后端已随兼容版本的 `torch_npu` 打包在 `torch_npu/_inductor/ascendc` 中，无需单独安装。通过 `torchtitan_npu.entry` 启动训练并开启 `torch.compile` 时，本仓会自动选择该后端。
80		78
81	-inductor_npu_ext 需要从源码安装。在运行环境内执行以下命令：	79	+具体配置和启动方式请参考 [torch.compile 说明文档](../feature_guides/torch_compile.md)。
82	-```bash
83	-git clone https://gitcode.com/Ascend/torchair.git
84	-cd torchair/experimental/_inductor_npu_ext/
85	-pip3 install -e ./python/
86	-cd -
87	-```
88	-
89	-
90	-> 注：具体torch.compile特性文档请参考 [说明文档](https://gitcode.com/cann/torchtitan-npu/blob/master/docs/feature_guides/torch_compile.md)
91		80
92		81
93	## PyPI安装	82	## PyPI安装

Mrequirements.txt+2-1

	@@ -1,6 +1,7 @@
1	--extra-index-url https://download.pytorch.org/whl/cpu	1	--extra-index-url https://download.pytorch.org/whl/cpu
		2	+--find-links https://cann-ai.obs.cn-north-4.myhuaweicloud.com/torchtitan-npu/simple/torch-npu/index.html
2	torch==2.12.0+cpu	3	torch==2.12.0+cpu
3	-torch_npu==2.12.0rc1	4	+torch_npu==2.12.0.dev20260808
		Z Z zzyyjj01212 天前已过期 pip install -r requirements.txt 没办法安装这个包张伟 11 天前评论： 11 天前评论： titan-npu-reviewer11 天前已过期一般(S2): 这里直接 pin 到 `torch_npu==2.12.0.dev20260808`，但安装文档已说明该 daily wheel 默认 pip 源无法直接解析；clean 环境按 AGENTS 里的 `pip install -r requirements.txt -r requirements_dev.txt` 或用户直接安装 requirements 时仍会失败。建议统一安装契约：改为可解析的直接 wheel/约束入口，或从 requirements 中移除该 pin 并把脚本前置要求同步到 AGENTS/所有安装路径。
4	numpy	5	numpy
5	PyYAML	6	PyYAML
6	torchtitan @ git+https://gitcode.com/GitHub_Trending/to/torchtitan.git@ac13e536c84e7f6647b14fa9375c3c8a8a2b8578	7	torchtitan @ git+https://gitcode.com/GitHub_Trending/to/torchtitan.git@ac13e536c84e7f6647b14fa9375c3c8a8a2b8578

Mtests/smoke_tests/features/test_quant.py+4-1

	@@ -24,7 +24,10 @@ def test_quant_linear_mxfp8(npu_device):
24	except RuntimeError as error:	24	except RuntimeError as error:
25	skip_on_runtime_unsupported(	25	skip_on_runtime_unsupported(
26	error,	26	error,
27	- ("does not support opType [DynamicMxQuant]",),	27	+ (
		28	+ "does not support opType [DynamicMxQuant]",
		29	+ "SoC version ascend910b verification failed",
		30	+ ),
28	"DynamicMxQuant is not supported on the current Ascend SOC",	31	"DynamicMxQuant is not supported on the current Ascend SOC",
29	)	32	)
30		33

Mtests/smoke_tests/integration_test.py+1-27

	@@ -25,7 +25,6 @@ Usage:
25	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v3_fake_backend_ep	25	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v3_fake_backend_ep
26	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v4_fake_backend_ep	26	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v4_fake_backend_ep
27	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v32_fake_backend_ep	27	python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v32_fake_backend_ep
28	- python tests/smoke_tests/integration_test.py ./outputs --test_name deepseek_v32_checkpoint_async_with_pinned_mem
29	python tests/smoke_tests/integration_test.py ./outputs --ngpu 4	28	python tests/smoke_tests/integration_test.py ./outputs --ngpu 4
30	"""	29	"""
31		30
	@@ -377,33 +376,8 @@ def _fake_backend_tests() -> list[OverrideDefinitions]:
377	]	376	]
378		377
379		378
380	-def _checkpoint_tests() -> list[OverrideDefinitions]:
381	- """Checkpoint integration tests."""
382	- return [
383	- OverrideDefinitions(
384	- [
385	- [
386	- f"--module {_DEEPSEEK_V32_MODULE}",
387	- f"--config {_DEEPSEEK_V32_CONFIG}",
388	- "--optimizer.swap-optimizer",
389	- "--training.steps 2",
390	- "--checkpoint.enable",
391	- "--checkpoint.no-load-only",
392	- "--checkpoint.interval 1",
393	- "--checkpoint.folder checkpoint_async_with_pinned_mem",
394	- "--checkpoint.async-mode async_with_pinned_mem",
395	- "--checkpoint.no-sync-files",
396	- ]
397	- ],
398	- "DeepSeek V3.2 checkpoint async_with_pinned_mem",
399	- "deepseek_v32_checkpoint_async_with_pinned_mem",
400	- ngpu=2,
401	- ),
402	- ]
403	-
404	-
405	def generate_smoke_tests() -> list[OverrideDefinitions]:	379	def generate_smoke_tests() -> list[OverrideDefinitions]:
406	- return _tp_compile_tests() + _fake_backend_tests() + _base_tests() + _tp_tests() + _ep_tests() + _checkpoint_tests()	380	+ return _tp_compile_tests() + _fake_backend_tests() + _base_tests() + _tp_tests() + _ep_tests()
407		381
408		382
409	# ============================================================================	383	# ============================================================================

Mtests/smoke_tests/npu_bypass_triton_codegen/npu_bypass_triton_codegen.py+58-18

	@@ -4,12 +4,23 @@
4	# LICENSE file in the root directory of this source tree.	4	# LICENSE file in the root directory of this source tree.
5		5
6		6
7	-from functools import wraps	7	+import importlib
8	import logging	8	import logging
		9	+from functools import wraps
9	from typing import Any	10	from typing import Any
10		11
11	logger = logging.getLogger(__name__)	12	logger = logging.getLogger(__name__)
12	_PATCHED = False	13	_PATCHED = False
		14	+_ASCENDC_LOWERING_GUARD = "_LoweringGuard"
		15	+_ASCENDC_LOWERING_REGISTRY = "_data"
		16	+_INDUCTOR_LOWERING_MODULE = "torch._inductor.lowering"
		17	+_TORCH_NPU_ASCENDC_LOWERING_MODULE = "torch_npu._inductor.ascendc.lowering.common"
		18	+_TORCH_NPU_DYNAMO_MODULE = "torch_npu.utils._dynamo"
		19	+_TORCH_NPU_LOWERING_OP_LIST_MODULE = "torch_npu._inductor.lowering_op_list"
		20	+_NPU_BACKEND_SCOPE = "_NpuBackendScope"
		21	+_GENERATE_LIST = "GENERATE_LIST"
		22	+_LOWERING_OVERLOAD_OP = "LOWERING_OVERLOAD_OP"
		23	+_FALLBACK_LIST = "FALLBACK_LIST"
13		24
14	_LOWERINGS_TO_KEEP = (	25	_LOWERINGS_TO_KEEP = (
15	("copy_", "default"),	26	("copy_", "default"),
	@@ -39,22 +50,32 @@ def _enable_implicit_fallbacks() -> None:
39		50
40		51
41	def _prepare_inductor_bypass() -> None:	52	def _prepare_inductor_bypass() -> None:
		53	+ # Importing torch_npu._inductor can lazily register AscendC lowerings.
		54	+ _fix_torch_npu_inductor_lowering()
42	preserved = _collect_lowerings_to_keep()	55	preserved = _collect_lowerings_to_keep()
43	_clear_inductor_tables(preserved)	56	_clear_inductor_tables(preserved)
44	- _fix_torch_npu_inductor_lowering()
45	_restore_lowerings(preserved)	57	_restore_lowerings(preserved)
46		58
47		59
48	-def _collect_lowerings_to_keep() -> dict[Any, Any]:	60	+def _collect_lowering_targets_to_keep() -> tuple[Any, ...]:
49	import torch	61	import torch
50	- from torch._inductor.lowering import lowerings
51		62
52	- preserved = {}	63	+ targets = []
53	for packet_name, overload_name in _LOWERINGS_TO_KEEP:	64	for packet_name, overload_name in _LOWERINGS_TO_KEEP:
54	packet = getattr(torch.ops.aten, packet_name, None)	65	packet = getattr(torch.ops.aten, packet_name, None)
55	if packet is None:	66	if packet is None:
56	continue	67	continue
57	target = getattr(packet, overload_name, None)	68	target = getattr(packet, overload_name, None)
		69	+ if target is not None:
		70	+ targets.append(target)
		71	+ return tuple(targets)
		72	+
		73	+
		74	+def _collect_lowerings_to_keep() -> dict[Any, Any]:
		75	+ lowerings = importlib.import_module(_INDUCTOR_LOWERING_MODULE).lowerings
		76	+
		77	+ preserved = {}
		78	+ for target in _collect_lowering_targets_to_keep():
58	if target in lowerings:	79	if target in lowerings:
59	preserved[target] = lowerings[target]	80	preserved[target] = lowerings[target]
60	return preserved	81	return preserved
	@@ -77,27 +98,46 @@ def _restore_lowerings(preserved: dict[Any, Any]) -> None:
77		98
78	def _fix_torch_npu_inductor_lowering() -> None:	99	def _fix_torch_npu_inductor_lowering() -> None:
79	try:	100	try:
80	- from torch_npu._inductor.lowering import _init_set	101	+ lowering_op_list = importlib.import_module(_TORCH_NPU_LOWERING_OP_LIST_MODULE)
81	- from torch_npu._inductor.lowering_op_list import (
82	- FALLBACK_LIST,
83	- GENERATE_LIST,
84	- LOWERING_OVERLOAD_OP,
85	- )
86	except Exception as exc:	102	except Exception as exc:
87	- logger.debug("Skip torch_npu lowering cleanup for bypass smoke test: %r", exc)	103	+ logger.debug("Skip legacy torch_npu lowering cleanup for bypass smoke test: %r", exc)
88	- return	104	+ else:
		105	+ getattr(lowering_op_list, _GENERATE_LIST).clear()
		106	+ getattr(lowering_op_list, _LOWERING_OVERLOAD_OP).clear()
		107	+ getattr(lowering_op_list, _FALLBACK_LIST).clear()
89		108
90	- _init_set(GENERATE_LIST, set())	109	+ try:
91	- _init_set(LOWERING_OVERLOAD_OP, set())	110	+ ascendc_lowering = importlib.import_module(_TORCH_NPU_ASCENDC_LOWERING_MODULE)
92	- FALLBACK_LIST.clear()	111	+ except Exception as exc:
		112	+ logger.debug("Skip AscendC lowering cleanup for bypass smoke test: %r", exc)
		113	+ else:
		114	+ lowering_guard = getattr(
		115	+ ascendc_lowering,
		116	+ _ASCENDC_LOWERING_GUARD,
		117	+ )
		118	+ lowering_registry = getattr(
		119	+ lowering_guard,
		120	+ _ASCENDC_LOWERING_REGISTRY,
		121	+ )
		122	+ preserved_registry = {
		123	+ target: lowering_registry[target]
		124	+ for target in _collect_lowering_targets_to_keep()
		125	+ if target in lowering_registry
		126	+ }
		127	+ lowering_registry.clear()
		128	+ lowering_registry.update(preserved_registry)
93		129
94		130
95	def _npu_bypass_backend(gm: Any, example_inputs: Any):	131	def _npu_bypass_backend(gm: Any, example_inputs: Any):
96	"""Test-only compile backend: clear NPU codegen lowerings, then call Inductor."""	132	"""Test-only compile backend: clear NPU codegen lowerings, then call Inductor."""
97	- _prepare_inductor_bypass()
98	from torch._dynamo.backends.registry import lookup_backend	133	from torch._dynamo.backends.registry import lookup_backend
99		134
100	- return lookup_backend("inductor")(gm, example_inputs)	135	+ inductor_backend = lookup_backend("inductor")
		136	+ torch_npu_dynamo = importlib.import_module(_TORCH_NPU_DYNAMO_MODULE)
		137	+ npu_backend_scope = getattr(torch_npu_dynamo, _NPU_BACKEND_SCOPE)
		138	+ with npu_backend_scope("ascendc"):
		139	+ _prepare_inductor_bypass()
		140	+ return inductor_backend(gm, example_inputs)
101		141
102		142
103	def _wrap_torch_compile(torch_module) -> None:	143	def _wrap_torch_compile(torch_module) -> None:

Mtorchtitan_npu/entry.py+8-11

	@@ -24,6 +24,12 @@ if TYPE_CHECKING:
24	_SKIP_FLEX_TO_SDPA_REWRITE_MODELS = {"vlm"}	24	_SKIP_FLEX_TO_SDPA_REWRITE_MODELS = {"vlm"}
25		25
26		26
		27	+def _configure_compile_backend() -> None:
		28	+ from torch._inductor import config as inductor_config
		29	+
		30	+ inductor_config.npu_backend = "ascendc" # pyrefly: ignore [missing-attribute]
		31	+
		32	+
27	def main() -> None:	33	def main() -> None:
28	"""Main entry point for NPU training with new config system."""	34	"""Main entry point for NPU training with new config system."""
29	init_logger()	35	init_logger()
	@@ -70,6 +76,8 @@ def main() -> None:
70	logger.warning("There might be performance issues with activation checkpointing and torch.compile enabled!")	76	logger.warning("There might be performance issues with activation checkpointing and torch.compile enabled!")
71		77
72	if config.compile.enable:	78	if config.compile.enable:
		79	+ _configure_compile_backend()
		80	+
73	if model_name == "deepseek_v3":	81	if model_name == "deepseek_v3":
74	# MLA performs shape inference according to the value tensor;	82	# MLA performs shape inference according to the value tensor;
75	# patch the meta registration so dynamo traces the right shapes.	83	# patch the meta registration so dynamo traces the right shapes.
	@@ -86,17 +94,6 @@ def main() -> None:
86	)	94	)
87		95
88	original_meta_func.__code__ = npu_fusion_attention_forward.__code__	96	original_meta_func.__code__ = npu_fusion_attention_forward.__code__
89	-
90	- try:
91	- # pyrefly: ignore [missing-import]
92	- import inductor_npu_ext # noqa: F401
93	- except Exception as e:
94	- raise RuntimeError(
95	- f"compile.enable is True for {model_name} model but inductor_npu_ext is not available. "
96	- "Please install inductor_npu_ext before enabling compile. "
97	- "See docs/torch_compile.md for installation instructions."
98	- ) from e
99	-
100	if model_name in ("deepseek_v32", "deepseek_v4"):	97	if model_name in ("deepseek_v32", "deepseek_v4"):
101	from torchtitan_npu.train import (	98	from torchtitan_npu.train import (
102	_patch_init_for_dsa_set_loss_scale,	99	_patch_init_for_dsa_set_loss_scale,

Mtorchtitan_npu/experiments/ao_npu/tests/ops/test_compile_ops.py+19-24

	@@ -11,17 +11,16 @@ quantization paths (MX FP8, MX FP4, Block FP8). The key concern is that
11	omitted for standard FP8 types, where the tensor's native dtype is	11	omitted for standard FP8 types, where the tensor's native dtype is
12	sufficient, but passed for FP4 where the tensor is stored as ``uint8``.	12	sufficient, but passed for FP4 where the tensor is stored as ``uint8``.
13		13
14	-The compilation uses ``backend="inductor"`` with ``inductor_npu_ext``,	14	+The compilation uses ``backend="inductor"`` and selects the bundled AscendC
15	-which is the standard NPU compile path used by torchtitan-npu.	15	+codegen with the ``npu_backend`` compile option.
16		16
17	First invocation can be slow due to TBE kernel compilation (~minutes);	17	First invocation can be slow due to TBE kernel compilation (~minutes);
18	subsequent runs with the same shape use the cached graph.	18	subsequent runs with the same shape use the cached graph.
19	"""	19	"""
20		20
21	-# Extends the inductor backend for NPU; import order matters (before torch.compile).
22	-import inductor_npu_ext # noqa: F401
23	import pytest	21	import pytest
24	import torch	22	import torch
		23	+import torch_npu # noqa: F401
25		24
26	from torchtitan_npu.experiments.ao_npu.torchao_npu.ops.block_ops import (	25	from torchtitan_npu.experiments.ao_npu.torchao_npu.ops.block_ops import (
27	to_block_fp8_then_mm,	26	to_block_fp8_then_mm,
	@@ -37,6 +36,18 @@ def _npu_available():
37	return hasattr(torch, "npu") and torch.npu.is_available()	36	return hasattr(torch, "npu") and torch.npu.is_available()
38		37
39		38
		39	+def _compile_and_assert_output(model, lhs, rhs, *, expected_shape):
		40	+ compiled = torch.compile(
		41	+ model,
		42	+ backend="inductor",
		43	+ dynamic=False,
		44	+ options={"npu_backend": "ascendc"},
		45	+ )
		46	+ output = compiled(lhs, rhs)
		47	+ assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}"
		48	+ assert output.dtype == torch.bfloat16
		49	+
		50	+
40	# ============================================================================	51	# ============================================================================
41	# MX quantized matmul (to_mx_then_mm)	52	# MX quantized matmul (to_mx_then_mm)
42	# ============================================================================	53	# ============================================================================
	@@ -68,11 +79,7 @@ def test_mx_fp8_matmul_compile(M, K, N):
68	config = MXQuantizeConfig() # default: float8_e4m3fn	79	config = MXQuantizeConfig() # default: float8_e4m3fn
69		80
70	model = MXMMModel(config, config).npu()	81	model = MXMMModel(config, config).npu()
71	- compiled = torch.compile(model, backend="inductor", dynamic=False)	82	+ _compile_and_assert_output(model, A, B, expected_shape=(M, N))
		Z Z zzyyjj0129 天前测一下低精度的场景？张伟 9 天前评论： 9 天前评论：
72	-
73	- Y = compiled(A, B)
74	- assert Y.shape == (M, N), f"Expected ({M}, {N}), got {Y.shape}"
75	- assert Y.dtype == torch.bfloat16
76		83
77		84
78	@pytest.mark.skipif(not _npu_available(), reason="NPU not available")	85	@pytest.mark.skipif(not _npu_available(), reason="NPU not available")
	@@ -102,11 +109,7 @@ def test_mx_fp4_matmul_compile(M, K, N):
102	config = MXQuantizeConfig(elem_dtype=torch.float4_e2m1fn_x2)	109	config = MXQuantizeConfig(elem_dtype=torch.float4_e2m1fn_x2)
103		110
104	model = MXMMModel(config, config).npu()	111	model = MXMMModel(config, config).npu()
105	- compiled = torch.compile(model, backend="inductor", dynamic=False)	112	+ _compile_and_assert_output(model, A, B, expected_shape=(M, N))
106	-
107	- Y = compiled(A, B)
108	- assert Y.shape == (M, N), f"Expected ({M}, {N}), got {Y.shape}"
109	- assert Y.dtype == torch.bfloat16
110		113
111		114
112	# ============================================================================	115	# ============================================================================
	@@ -141,11 +144,7 @@ def test_block_fp8_without_mxfp4_compile(M, K, N):
141	config_B = BlockQuantizeConfig() # no mxfp4_fake_quantize_config	144	config_B = BlockQuantizeConfig() # no mxfp4_fake_quantize_config
142		145
143	model = BlockMMModel(config_A, config_B).npu()	146	model = BlockMMModel(config_A, config_B).npu()
144	- compiled = torch.compile(model, backend="inductor", dynamic=False)	147	+ _compile_and_assert_output(model, A, B, expected_shape=(M, N))
145	-
146	- Y = compiled(A, B)
147	- assert Y.shape == (M, N), f"Expected ({M}, {N}), got {Y.shape}"
148	- assert Y.dtype == torch.bfloat16
149		148
150		149
151	@pytest.mark.skipif(not _npu_available(), reason="NPU not available")	150	@pytest.mark.skipif(not _npu_available(), reason="NPU not available")
	@@ -167,8 +166,4 @@ def test_block_fp8_with_mxfp4_compile(M, K, N):
167	)	166	)
168		167
169	model = BlockMMModel(config_A, config_B).npu()	168	model = BlockMMModel(config_A, config_B).npu()
170	- compiled = torch.compile(model, backend="inductor", dynamic=False)	169	+ _compile_and_assert_output(model, A, B, expected_shape=(M, N))
171	-
172	- Y = compiled(A, B)
173	- assert Y.shape == (M, N), f"Expected ({M}, {N}), got {Y.shape}"
174	- assert Y.dtype == torch.bfloat16