feat: add torch.compile support(take engram_fused_weight as example)-ops-tilelang-AtomGit

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

Total: 5 files changed+108-32

Mdocs/architecture.md+43-1

	@@ -77,7 +77,7 @@ cann_ops_tilelang.engram_fused_weight
77	-> validate BF16/NPU/2-D/contiguous inputs	77	-> validate BF16/NPU/2-D/contiguous inputs
78	-> choose block size and persistent scheduling from total elements	78	-> choose block size and persistent scheduling from total elements
79	-> get_engram_fused_weight_kernel(hidden_size, hc_mult)	79	-> get_engram_fused_weight_kernel(hidden_size, hc_mult)
80	- -> allocate/reuse an FP32 output	80	+ -> allocate an FP32 output
81	-> launch flattened tensors	81	-> launch flattened tensors
82	-> compare with the FP32 PyTorch reference	82	-> compare with the FP32 PyTorch reference
83	```	83	```
	@@ -120,6 +120,48 @@ mhc_post_backward(...)
120		120
121	`mhc_head_compute_mix_forward` 和 `mhc_head_compute_mix_backward` 处理 head mixing 系数，`mhc_multilayer_recompute` 融合多层 pre/post mix 重计算。`mhc_pre_big_fuse` 融合 reduce/RMSNorm、split mixes、Sinkhorn 和 pre apply mix 的前向路径。	121	`mhc_head_compute_mix_forward` 和 `mhc_head_compute_mix_backward` 处理 head mixing 系数，`mhc_multilayer_recompute` 融合多层 pre/post mix 重计算。`mhc_pre_big_fuse` 融合 reduce/RMSNorm、split mixes、Sinkhorn 和 pre apply mix 的前向路径。
122		122
		123	+## Engram fused weight 的 torch.compile 接入
		124	+
		125	+`engram_fused_weight` 支持固定形状推理的 `npugraph_ex` 编译路径：
		W W wangliwei16 days ago 非固定形状能支持吗？ ChaoyangJi 16 days ago Comments： 16 days ago Comments：
		126	+
		127	+```python
		128	+import torch
		129	+import torch_npu
		130	+
		131	+from cann_ops_tilelang import engram_fused_weight
		132	+
		133	+
		134	+def model(weight_hidden, weight_embed):
		135	+ return engram_fused_weight(-weight_hidden, weight_embed) + 1
		136	+
		137	+
		138	+compiled_model = torch.compile(
		139	+ model, backend='npugraph_ex', fullgraph=True, dynamic=False,
		140	+)
		141	+with torch.inference_mode():
		142	+ a = torch.randn(4, 4096, device='npu', dtype=torch.bfloat16)
		143	+ b = torch.randn_like(a)
		144	+ result = compiled_model(a, b)
		145	+```
		146	+
		147	+编译时，公共接口调用 `cann_ops_tilelang::engram_fused_weight` 自定义算子；其 fake 实现只校验
		148	+shape、dtype、device 和布局，并推导同形状的 FP32 输出。实际执行仍使用原有 TileLang 内核。
		149	+公共接口只接收两个输入张量，直接返回 FP32 结果，不修改输入。普通 eager 调用保持原有路径。
		150	+
		151	+本仓库要求 PyTorch >= 2.10；此接入使用 `torch.library.custom_op`，并需要提供
		152	+`npugraph_ex` 的配套 `torch_npu`。实机验证环境为 PyTorch / torch_npu 2.10。
		153	+输入仍须为连续的二维 BF16 NPU 张量，非空元素数须能被 256 整除；空输入也受支持。
		154	+这里按形状编译，形状变化可能触发重新编译，未注册 autograd 公式。
		155	+首次执行包含 TileLang JIT 和图准备开销。
		156	+
		157	+以下测试包含原有 eager 用例和一个固定形状的 compile 基本功能用例，验证
		158	+`npugraph_ex` 编译执行的结果与 FP32 参考结果一致：
		159	+
		160	+```bash
		161	+OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OPS_TILELANG_TEST_LEVEL=0 \
		162	+ python -m pytest tests/engram/test_engram_fused_weight.py -m 'not benchmark' -v
		163	+```
		164	+
123	## 扩展方向	165	## 扩展方向
124		166
125	- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；	167	- `src/cann_ops_tilelang/quant/`：公共 scaling-factor layout 和 FP8/FP4 类型系统；

Mpyproject.toml+1-1

	@@ -11,7 +11,7 @@ requires-python = ">=3.10"
11	license = "MIT"	11	license = "MIT"
12	authors = [{name = "ops-tilelang contributors"}]	12	authors = [{name = "ops-tilelang contributors"}]
13	dependencies = [	13	dependencies = [
14	- "torch>=2.1",	14	+ "torch>=2.10",
15	"tilelang>=0.1.9",	15	"tilelang>=0.1.9",
16	]	16	]
17	classifiers = [	17	classifiers = [

Mrequirements.txt+1-1

	@@ -9,7 +9,7 @@
9	# The tilelang version constraint alone does not guarantee Ascend support.	9	# The tilelang version constraint alone does not guarantee Ascend support.
10		10
11	# Runtime	11	# Runtime
12	-torch>=2.1	12	+torch>=2.10
13	tilelang>=0.1.9	13	tilelang>=0.1.9
14		14
15	# Tests (NumPy is also required by Tensor.numpy() in test references)	15	# Tests (NumPy is also required by Tensor.numpy() in test references)

Msrc/cann_ops_tilelang/engram/engram_fused_weight.py+38-29

	@@ -9,14 +9,9 @@ from cann_ops_tilelang.config import require_ascend, should_print_kernel_source
9	from .engram_fused_weight_kernel import get_engram_fused_weight_kernel	9	from .engram_fused_weight_kernel import get_engram_fused_weight_kernel
10		10
11		11
12	-def _validate_inputs(	12	+def _validate_inputs(weight_hidden: torch.Tensor, weight_embed: torch.Tensor) -> None:
13	- weight_hidden: torch.Tensor,
14	- weight_embed: torch.Tensor,
15	- out: torch.Tensor \| None,
16	-) -> None:
17	if not isinstance(weight_hidden, torch.Tensor) or not isinstance(weight_embed, torch.Tensor):	13	if not isinstance(weight_hidden, torch.Tensor) or not isinstance(weight_embed, torch.Tensor):
18	raise TypeError('weight_hidden and weight_embed must be torch.Tensor instances.')	14	raise TypeError('weight_hidden and weight_embed must be torch.Tensor instances.')
19	- require_ascend()
20	if weight_hidden.device.type != 'npu' or weight_embed.device.type != 'npu':	15	if weight_hidden.device.type != 'npu' or weight_embed.device.type != 'npu':
21	raise ValueError(	16	raise ValueError(
22	f'weight_hidden and weight_embed must be NPU tensors, got {weight_hidden.device} and {weight_embed.device}.'	17	f'weight_hidden and weight_embed must be NPU tensors, got {weight_hidden.device} and {weight_embed.device}.'
	@@ -37,33 +32,14 @@ def _validate_inputs(
37	if weight_hidden.numel() and weight_hidden.numel() % 256:	32	if weight_hidden.numel() and weight_hidden.numel() % 256:
38	raise ValueError('the total number of input elements must be divisible by 256.')	33	raise ValueError('the total number of input elements must be divisible by 256.')
39		34
40	- if out is None:
41	- return
42	- if not isinstance(out, torch.Tensor):
43	- raise TypeError('out must be a torch.Tensor when provided.')
44	- if out.shape != weight_hidden.shape:
45	- raise ValueError('out must have the same shape as weight_hidden.')
46	- if out.dtype != torch.float32 or out.device != weight_hidden.device:
47	- raise ValueError('out must be float32 and on the same NPU device as weight_hidden.')
48	- if not out.is_contiguous():
49	- raise ValueError('out must be contiguous.')
50		35
51	-	36	+def _engram_fused_weight_impl(
52	-def engram_fused_weight(
53	weight_hidden: torch.Tensor,	37	weight_hidden: torch.Tensor,
54	weight_embed: torch.Tensor,	38	weight_embed: torch.Tensor,
55	- *,
56	- out: torch.Tensor \| None = None,
57	) -> torch.Tensor:	39	) -> torch.Tensor:
58	- """Multiply two BF16 Engram weights and return an FP32 tensor.	40	+ _validate_inputs(weight_hidden, weight_embed)
59	-	41	+ require_ascend()
60	- The inputs must be contiguous NPU tensors with shape	42	+ result = torch.empty_like(weight_hidden, dtype=torch.float32)
61	- ``(hc_mult, hidden_size)``. The Ascend kernel converts both operands to
62	- FP32 before multiplication.
63	- """
64	-
65	- _validate_inputs(weight_hidden, weight_embed, out)
66	- result = torch.empty_like(weight_hidden, dtype=torch.float32) if out is None else out
67	if weight_hidden.numel() == 0:	43	if weight_hidden.numel() == 0:
68	return result	44	return result
69		45
	@@ -73,3 +49,36 @@ def engram_fused_weight(
73	print(kernel.get_kernel_source())	49	print(kernel.get_kernel_source())
74	kernel(weight_hidden.view(-1), weight_embed.view(-1), result.view(-1))	50	kernel(weight_hidden.view(-1), weight_embed.view(-1), result.view(-1))
75	return result	51	return result
		52	+
		53	+
		54	+@torch.library.custom_op('cann_ops_tilelang::engram_fused_weight', mutates_args=())
		55	+def _compiled_fused_weight(weight_hidden: torch.Tensor, weight_embed: torch.Tensor) -> torch.Tensor:
		56	+ return _engram_fused_weight_impl(weight_hidden, weight_embed)
		57	+
		58	+
		59	+@_compiled_fused_weight.register_fake
		60	+def _fused_weight_fake(weight_hidden, weight_embed):
		61	+ _validate_inputs(weight_hidden, weight_embed)
		62	+ return torch.empty_like(weight_hidden, dtype=torch.float32)
		63	+
		64	+
		65	+def engram_fused_weight(
		66	+ weight_hidden: torch.Tensor,
		67	+ weight_embed: torch.Tensor,
		68	+) -> torch.Tensor:
		W W wangliwei19 days ago This code comment section is from+65To+68 删除 out= 后，已有 engram_fused_weight(a, b, out=buffer) 调用会直接报 TypeError。建议保留 eager 路径的可选 out= 支持；如果确定统一为返回值接口，请在文档中明确兼容性变化及迁移方式。 ChaoyangJi 17 days ago Comments： 17 days ago Comments：
		69	+ """Multiply two BF16 Engram weights and return an FP32 tensor.
		70	+
		71	+ The inputs must be contiguous NPU tensors with shape
		72	+ ``(hc_mult, hidden_size)``. The Ascend kernel converts both operands to
		73	+ FP32 before multiplication and returns the result without modifying
		74	+ either input.
		75	+
		76	+ With a compatible torch_npu, supports inference with
		77	+ ``torch.compile(..., backend='npugraph_ex', fullgraph=True, dynamic=False)``.
		78	+ The compiled path uses an opaque custom operator with a fake implementation;
		79	+ tracing does not compile or launch the TileLang kernel. Shapes specialize
		80	+ at compilation. No autograd formula is registered for this operator.
		81	+ """
		82	+ if torch.compiler.is_compiling():
		83	+ return _compiled_fused_weight(weight_hidden, weight_embed)
		84	+ return _engram_fused_weight_impl(weight_hidden, weight_embed)

Mtests/engram/test_engram_fused_weight.py+25-0

	@@ -35,6 +35,31 @@ def test_engram_fused_weight(params):
35	assert_equal(out, ref)	35	assert_equal(out, ref)
36		36
37		37
		38	+@torch.inference_mode()
		39	+def test_engram_fused_weight_compile():
		W W wangliwei19 days ago test_engram_fused_weight_compile 当前仅调用一次 compiled 函数。建议复用同一个 compiled 函数，覆盖原地修改输入、传入同形状新 Tensor 后的多次调用，并逐次与 FP32 参考结果比较，确认图重放能够读取最新输入，避免旧结果导致误判。 ChaoyangJi 17 days ago Comments： 17 days ago Comments：
		40	+ wh_data, we_data = generate_test_data({'hc': 4, 'hidden': 4096})
		41	+ compiled = torch.compile(
		42	+ cann_ops_tilelang.engram_fused_weight,
		43	+ backend='npugraph_ex',
		44	+ fullgraph=True,
		45	+ dynamic=False,
		46	+ )
		47	+ for step in range(3):
		48	+ if step == 1:
		49	+ # Update values while keeping the input storage unchanged.
		50	+ wh_data.add_(1)
		51	+ we_data.neg_()
		52	+ elif step == 2:
		53	+ # Replace both inputs with new tensors of the same shape.
		54	+ wh_data, we_data = generate_test_data({'hc': 4, 'hidden': 4096})
		55	+
		56	+ ref = wh_data.float() * we_data.float()
		57	+ out = compiled(wh_data, we_data)
		58	+ assert_equal(out, ref)
		59	+ # Graph execution may reuse output storage on the next call.
		60	+ del out
		61	+
		62	+
38	@pytest.mark.benchmark	63	@pytest.mark.benchmark
39	@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)	64	@pytest.mark.parametrize('params', generate_test_params(0), ids=make_param_id)
40	def test_engram_fused_weight_benchmark(benchmark_timer, benchmark_record, params):	65	def test_engram_fused_weight_benchmark(benchmark_timer, benchmark_record, params):