[feat] npu_mhc_pre新增inner_precise支持使能HF32模式-op-plugin-AtomGit

Mcodegen/templates/_op_plugin_docs.py+2-1

	@@ -4003,7 +4003,7 @@ _add_torch_npu_docstr(
4003	"npu_mhc_pre",	4003	"npu_mhc_pre",
4004	"""	4004	"""
4005	接口原型：	4005	接口原型：
4006	torch_npu.npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)	4006	torch_npu.npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
4007		4007
4008	功能描述：	4008	功能描述：
4009	通过一系列计算，可得到 MHC (流形约束超连接)架构中 hidden 层对应的投影矩阵 Hres 和 Hpost，以及作为 Atten 或 MLP 层输入的矩阵 Hin。	4009	通过一系列计算，可得到 MHC (流形约束超连接)架构中 hidden 层对应的投影矩阵 Hres 和 Hpost，以及作为 Atten 或 MLP 层输入的矩阵 Hin。
	@@ -4017,6 +4017,7 @@ gamma: Tensor类型，可选输入，表示进行 RMSNorm 计算时的缩放因
4017	norm_eps: Float类型，可选输入，RMSNorm 的防除零参数。	4017	norm_eps: Float类型，可选输入，RMSNorm 的防除零参数。
4018	hc_eps: Float类型，可选输入，H_pre 经过 sigmoid 运算后的 eps 参数。	4018	hc_eps: Float类型，可选输入，H_pre 经过 sigmoid 运算后的 eps 参数。
4019	out_flag: Int类型，可选输入，表示是否输出中间结果标识，默认值为0(仅输出最终变换结果)。	4019	out_flag: Int类型，可选输入，表示是否输出中间结果标识，默认值为0(仅输出最终变换结果)。
		4020	inner_precise: Int类型，可选输入，指定MhcPre算子的计算模式，默认值为0。0表示在Cube中使用FP32模式计算；1表示在Cube中使用HF32模式计算。
4020	n：shape 中的 n 常取 4、6、8。	4021	n：shape 中的 n 常取 4、6、8。
4021		4022
4022	输出说明:	4023	输出说明:

Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_mhc_pre.md+3-1

	@@ -30,7 +30,7 @@ $$
30	## 函数原型	30	## 函数原型
31		31
32	```python	32	```python
33	torch_npu.npu_mhc_pre(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=1e-6, out_flag=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)	33	torch_npu.npu_mhc_pre(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=1e-6, out_flag=0, inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
34	```	34	```
35		35
36	## 参数说明	36	## 参数说明
	@@ -44,6 +44,7 @@ torch_npu.npu_mhc_pre(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=
44	- norm\_eps（`float`）：可选参数，RmsNorm的防除零参数，默认值是1e-6，数据类型为`float32`。	44	- norm\_eps（`float`）：可选参数，RmsNorm的防除零参数，默认值是1e-6，数据类型为`float32`。
45	- hc\_eps（`float`）：可选参数，H<sub>pre</sub>的sigmoid后的eps参数，默认值是1e-6，数据类型支持`float32`。	45	- hc\_eps（`float`）：可选参数，H<sub>pre</sub>的sigmoid后的eps参数，默认值是1e-6，数据类型支持`float32`。
46	- out\_flag（`int`）：可选参数，表示是否输出h\_mix/inv\_rms/h\_pre，默认为0表示不输出，为1表示全输出。	46	- out\_flag（`int`）：可选参数，表示是否输出h\_mix/inv\_rms/h\_pre，默认为0表示不输出，为1表示全输出。
		47	- inner\_precise（`int`）：可选参数，指定MhcPre算子的计算模式，默认值为0。0表示在Cube中使用FP32模式计算；1表示在Cube中使用HF32模式计算。
47		48
48	## 返回值说明	49	## 返回值说明
49		50
	@@ -63,6 +64,7 @@ torch_npu.npu_mhc_pre(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=
63	- D：支持1\~16384，需满足D为16对齐。	64	- D：支持1\~16384，需满足D为16对齐。
64	- 当alpha=\[3\]时，支持 h\_res 输出，必须满足以下条件：输入 phi=$n<sup>2</sup>+2n, nD$，bias=$n<sup>2</sup>+2n$，输出 h\_mix=$B, S ,n<sup>2</sup>+2n$ 或 $T, n<sup>2</sup>+2n$；当alpha=\[2\]时，h\_res 输出为0，必须满足以下条件：输入 phi=$2n, nD$，bias=$2n$，输出 h\_mix=$B, S ,2n$ 或 $T, 2n$。	65	- 当alpha=\[3\]时，支持 h\_res 输出，必须满足以下条件：输入 phi=$n<sup>2</sup>+2n, nD$，bias=$n<sup>2</sup>+2n$，输出 h\_mix=$B, S ,n<sup>2</sup>+2n$ 或 $T, n<sup>2</sup>+2n$；当alpha=\[2\]时，h\_res 输出为0，必须满足以下条件：输入 phi=$2n, nD$，bias=$2n$，输出 h\_mix=$B, S ,2n$ 或 $T, 2n$。
65	- 可选输出 inv\_rms、h\_mix、h\_pre 为互存关系，需同时输出或全部不输出，不支持仅返回其中部分。	66	- 可选输出 inv\_rms、h\_mix、h\_pre 为互存关系，需同时输出或全部不输出，不支持仅返回其中部分。
		67	- inner_precise：目前支持0、1。
66		68
67	## 调用示例	69	## 调用示例
68		70

Mop_plugin/config/derivatives.yaml+1-1

	@@ -356,7 +356,7 @@ backward:
356	x2: _cdist_backward(grad.mT().contiguous(), x2, x1, p, result.mT().contiguous())	356	x2: _cdist_backward(grad.mT().contiguous(), x2, x1, p, result.mT().contiguous())
357	version: all_version	357	version: all_version
358		358
359	- name: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)	359	- name: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
360	output_differentiability: [true, true, true, false, false, false]	360	output_differentiability: [true, true, true, false, false, false]
361	x, phi, alpha, bias, gamma: 'npu_mhc_pre_backward(x, phi, alpha, grads[0], grads[1], grads[2], result3, result4, result5, result1, gamma, hc_eps, c10::nullopt)'	361	x, phi, alpha, bias, gamma: 'npu_mhc_pre_backward(x, phi, alpha, grads[0], grads[1], grads[2], result3, result4, result5, result1, gamma, hc_eps, c10::nullopt)'
		ffmh11 天前当inner_precise传入1（HF32）时，反向是否要做特殊处理？ L liweijian16 11 天前评论： 11 天前评论：
362	version: all_version	362	version: all_version

Mop_plugin/config/op_plugin_functions.yaml+1-1

	@@ -5616,7 +5616,7 @@ custom:
5616	- func: attention_worker_scheduler(Tensor self) -> Tensor	5616	- func: attention_worker_scheduler(Tensor self) -> Tensor
5617	op_api: [v2.1, newest]	5617	op_api: [v2.1, newest]
5618		5618
5619	- func: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)	5619	- func: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
5620	op_api: all_version	5620	op_api: all_version
5621	internal_format_opapi: all_version	5621	internal_format_opapi: all_version
5622		5622

Mop_plugin/ops/opapi/MhcPreNpuOpapi.cpp+38-17

	@@ -27,21 +27,28 @@ constexpr int64_t REMOVE_ONE_DIM = 1;
27	constexpr int64_t REMOVE_TWO_DIMS = 2;	27	constexpr int64_t REMOVE_TWO_DIMS = 2;
28	constexpr int64_t ALPHA_NUMEL = 3;	28	constexpr int64_t ALPHA_NUMEL = 3;
29	constexpr int64_t ALPHA_NUMEL_HY = 2;	29	constexpr int64_t ALPHA_NUMEL_HY = 2;
		30	constexpr int64_t CUBE_USE_FP32 = 0;
		31	constexpr int64_t CUBE_USE_HF32 = 1;
30		32
31	// aclnnMhcPre 依赖 CANN 9.0.0 及以上版本，此次校验避免旧版本环境出现找不到算子的兼容性问题。	33	// aclnnMhcPre 依赖 CANN 9.0.0 及以上版本，此次校验避免旧版本环境出现找不到算子的兼容性问题。
32	inline void check_mhc_pre_supported()	34	// Prefer V2 so inner_precise reaches the kernel. FP32 may fall back to V1 for compatibility.
		35	// Returns true when the caller should dispatch aclnnMhcPreV2.
		36	inline bool check_mhc_pre_supported(int64_t inner_precise)
33	{	37	{
34	static const bool is_cann_ready = op_plugin::utils::is_gte_cann_version_900();	38	static const bool is_cann_ready = op_plugin::utils::is_gte_cann_version_900();
35	static const bool is_aclnn_kernel_available = check_aclnn_kernel_available("aclnnMhcPre");	39	static const bool is_mhc_pre_available = check_aclnn_kernel_available("aclnnMhcPre");
		40	static const bool is_mhc_pre_v2_available = check_aclnn_kernel_available("aclnnMhcPreV2");
36	TORCH_CHECK(	41	TORCH_CHECK(
37	is_cann_ready && is_aclnn_kernel_available,	42	is_cann_ready && (is_mhc_pre_v2_available \|\|
38	"torch_npu.npu_mhc_pre requires CANN >= 9.0.0 and aclnnMhcPre support. "	43	(inner_precise == CUBE_USE_FP32 && is_mhc_pre_available)),
		44	"torch_npu.npu_mhc_pre requires CANN >= 9.0.0. inner_precise=1 additionally requires aclnnMhcPreV2. "
39	"Please upgrade CANN.",	45	"Please upgrade CANN.",
40	OPS_ERROR(ErrCode::NOT_SUPPORT));	46	OPS_ERROR(ErrCode::NOT_SUPPORT));
		47	return is_mhc_pre_v2_available;
41	}	48	}
42		49
43	/**	50	/**
44	* @brief 构造 aclnnMhcPre 所需的输出张量。	51	* @brief 构造 aclnnMhcPreV2 所需的输出张量。
45	*	52	*
46	* 该函数根据输入张量 x 的维度布局，并结合 phi 的第 0 维大小，	53	* 该函数根据输入张量 x 的维度布局，并结合 phi 的第 0 维大小，
47	* 预先创建算子执行所需的各个输出张量。	54	* 预先创建算子执行所需的各个输出张量。
	@@ -185,24 +192,33 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
185	namespace op_api {	192	namespace op_api {
186	std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mhc_pre(	193	std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mhc_pre(
187	const at::Tensor &x, const at::Tensor &phi, const at::Tensor &alpha, const at::Tensor &bias,	194	const at::Tensor &x, const at::Tensor &phi, const at::Tensor &alpha, const at::Tensor &bias,
188	const c10::optional<at::Tensor> &gamma, double norm_eps, double hc_eps, int64_t out_flag)	195	const c10::optional<at::Tensor> &gamma, double norm_eps, double hc_eps, int64_t out_flag, int64_t inner_precise)
189	{	196	{
190	TORCH_CHECK(x.numel() > 0, "Input x should not be empty.");	197	TORCH_CHECK(x.numel() > 0, "Input x should not be empty.", OPS_ERROR(ErrCode::VALUE));
191	TORCH_CHECK(phi.numel() > 0, "Input phi should not be empty.");	198	TORCH_CHECK(phi.numel() > 0, "Input phi should not be empty.", OPS_ERROR(ErrCode::VALUE));
192	TORCH_CHECK(alpha.numel() == ALPHA_NUMEL \|\| alpha.numel() == ALPHA_NUMEL_HY,	199	TORCH_CHECK(alpha.numel() == ALPHA_NUMEL \|\| alpha.numel() == ALPHA_NUMEL_HY,
193	"Input alpha must have 3 or 2 elements, but got ", alpha.numel(), ".");	200	"Input alpha must have 3 or 2 elements, but got ", alpha.numel(), ".", OPS_ERROR(ErrCode::VALUE));
194	TORCH_CHECK(bias.numel() > 0, "Input bias should not be empty.");	201	TORCH_CHECK(bias.numel() > 0, "Input bias should not be empty.", OPS_ERROR(ErrCode::VALUE));
195		202
196	TORCH_CHECK(x.dim() == TND_DIMS \|\| x.dim() == BSND_DIMS, "Input x must be 3D or 4D, but got ", x.dim(), "D.");	203	TORCH_CHECK(x.dim() == TND_DIMS \|\| x.dim() == BSND_DIMS, "Input x must be 3D or 4D, but got ", x.dim(), "D.",
		204	OPS_ERROR(ErrCode::VALUE));
197		205
198	// out_flag 用于控制 aclnnMhcPre 是否实际写出全部输出：	206	// out_flag 用于控制 aclnnMhcPreV2 是否实际写出全部输出：
199	// - out_flag == 1：写出全部 6 个结果	207	// - out_flag == 1：写出全部 6 个结果
200	// (outHin, outHpost, outHres, outInvRms, outHmix, outHpre)	208	// (outHin, outHpost, outHres, outInvRms, outHmix, outHpre)
201	// - out_flag == 0：仅需要 outHin / outHpost / outHres，	209	// - out_flag == 0：仅需要 outHin / outHpost / outHres，
202	// 后 3 个输出在调用 aclnnMhcPre 时使用 nullTensor 占位，不参与实际写出	210	// 后 3 个输出在调用 aclnnMhcPreV2 时使用 nullTensor 占位，不参与实际写出
203	TORCH_CHECK(out_flag == 0 \|\| out_flag == 1, "out_flag must be 0 or 1, but got ", out_flag, ".");	211	TORCH_CHECK(
		212	inner_precise == CUBE_USE_FP32 \|\| inner_precise == CUBE_USE_HF32,
		213	"Input inner_precise must be 0 (Cube uses FP32) or 1 (Cube uses HF32), but got ",
		214	inner_precise,
		215	".",
		216	OPS_ERROR(ErrCode::VALUE));
204		217
205	check_mhc_pre_supported();	218	TORCH_CHECK(out_flag == 0 \|\| out_flag == 1, "Input out_flag must be 0 or 1, but got ", out_flag, ".",
		219	OPS_ERROR(ErrCode::VALUE));
		220
		221	bool useMhcPreV2 = check_mhc_pre_supported(inner_precise);
206		222
207	bool hasResi = (alpha.numel() == ALPHA_NUMEL);	223	bool hasResi = (alpha.numel() == ALPHA_NUMEL);
208	auto mhcPreOutput = construct_mhc_pre_outputs(x, phi, out_flag, hasResi);	224	auto mhcPreOutput = construct_mhc_pre_outputs(x, phi, out_flag, hasResi);
	@@ -214,8 +230,13 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tenso
214	at::Tensor outHmix = std::get<4>(mhcPreOutput);	230	at::Tensor outHmix = std::get<4>(mhcPreOutput);
215	at::Tensor outHpre = std::get<5>(mhcPreOutput);	231	at::Tensor outHpre = std::get<5>(mhcPreOutput);
216		232
217	EXEC_NPU_CMD(aclnnMhcPre, x, phi, alpha, bias, gamma, norm_eps, hc_eps, outHin, outHpost,	233	if (useMhcPreV2) {
218	outHres, outInvRms, outHmix, outHpre);	234	EXEC_NPU_CMD(aclnnMhcPreV2, x, phi, alpha, bias, gamma, norm_eps, hc_eps, inner_precise, outHin, outHpost,
		235	outHres, outInvRms, outHmix, outHpre);
		236	} else {
		237	EXEC_NPU_CMD(aclnnMhcPre, x, phi, alpha, bias, gamma, norm_eps, hc_eps, outHin, outHpost,
		238	outHres, outInvRms, outHmix, outHpre);
		239	}
219	return std::make_tuple(outHin, outHpost, outHres, outInvRms, outHmix, outHpre);	240	return std::make_tuple(outHin, outHpost, outHres, outInvRms, outHmix, outHpre);
220	}	241	}
221	} // namespace op_api	242	} // namespace op_api

Mop_plugin/python/meta/_meta_registrations.py+8-1

	@@ -183,7 +183,7 @@ if os.getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") != "1":
183		183
184		184
185	@impl(m, "npu_mhc_pre")	185	@impl(m, "npu_mhc_pre")
186	def npu_mhc_pre_meta(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=1e-6, out_flag=0):	186	def npu_mhc_pre_meta(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=1e-6, out_flag=0, inner_precise=0):
187	torch._check(	187	torch._check(
188	x.numel() > 0,	188	x.numel() > 0,
189	lambda: "Input x should not be empty." + ops_error(ErrCode.VALUE),	189	lambda: "Input x should not be empty." + ops_error(ErrCode.VALUE),
	@@ -209,6 +209,13 @@ def npu_mhc_pre_meta(x, phi, alpha, bias, *, gamma=None, norm_eps=1e-6, hc_eps=1
209	out_flag == 0 or out_flag == 1,	209	out_flag == 0 or out_flag == 1,
210	lambda: f"Input out_flag must be 0 or 1, but got {out_flag}." + ops_error(ErrCode.VALUE),	210	lambda: f"Input out_flag must be 0 or 1, but got {out_flag}." + ops_error(ErrCode.VALUE),
211	)	211	)
		212	torch._check(
		213	inner_precise == 0 or inner_precise == 1,
		214	lambda: (
		215	f"Input inner_precise must be 0 (Cube uses FP32) or 1 (Cube uses HF32), but got {inner_precise}."
		216	+ ops_error(ErrCode.VALUE)
		217	),
		218	)
212		219
213	has_resi = alpha.numel() == 3	220	has_resi = alpha.numel() == 3
214		221

Mtest/core_tests/torch_npu_OpApi_schema_all.json+1-1

	@@ -922,7 +922,7 @@
922	"func: npu_mhc_post_backward(Tensor grad_y, Tensor x, Tensor? h_res, Tensor h_out, Tensor h_post) -> (Tensor, Tensor, Tensor, Tensor)": {	922	"func: npu_mhc_post_backward(Tensor grad_y, Tensor x, Tensor? h_res, Tensor h_out, Tensor h_post) -> (Tensor, Tensor, Tensor, Tensor)": {
923	"version": ["all_version"]	923	"version": ["all_version"]
924	},	924	},
925	"func: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)": {	925	"func: npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)": {
926	"version": ["all_version"]	926	"version": ["all_version"]
927	},	927	},
928	"func: npu_mhc_pre_backward(Tensor x, Tensor phi, Tensor alpha, Tensor grad_h_in, Tensor grad_h_post, Tensor grad_h_res, Tensor inv_rms, Tensor h_mix, Tensor h_pre, Tensor h_post, Tensor? gamma=None, float hc_eps=1e-6, Tensor? grad_x_post=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {	928	"func: npu_mhc_pre_backward(Tensor x, Tensor phi, Tensor alpha, Tensor grad_h_in, Tensor grad_h_post, Tensor grad_h_res, Tensor inv_rms, Tensor h_mix, Tensor h_pre, Tensor h_post, Tensor? gamma=None, float hc_eps=1e-6, Tensor? grad_x_post=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {

Mtest/test_custom_ops/test_npu_mhc_pre.py+28-12

	@@ -6,6 +6,13 @@ from torch_npu.testing.common_utils import SupportedDevices
6		6
7		7
8	class TestNpuMhcPre(TestCase):	8	class TestNpuMhcPre(TestCase):
		9	@staticmethod
		10	def fp32_to_hf32(value):
		11	value_bits = value.contiguous().view(torch.int32)
		12	hf32_bits = torch.bitwise_right_shift(torch.bitwise_right_shift(value_bits, 12) + 1, 1)
		13	hf32_bits = torch.bitwise_left_shift(hf32_bits, 13)
		14	return hf32_bits.view(torch.float32)
		15
9	def cpu_op_exec(	16	def cpu_op_exec(
10	self,	17	self,
11	x: torch.Tensor,	18	x: torch.Tensor,
	@@ -15,6 +22,7 @@ class TestNpuMhcPre(TestCase):
15	gamma: torch.Tensor = None,	22	gamma: torch.Tensor = None,
16	norm_eps: float = 1e-6,	23	norm_eps: float = 1e-6,
17	hc_eps: float = 1e-6,	24	hc_eps: float = 1e-6,
		25	inner_precise: int = 0,
18	):	26	):
19	T, N, D = x.shape	27	T, N, D = x.shape
20	ND = N * D	28	ND = N * D
	@@ -23,13 +31,16 @@ class TestNpuMhcPre(TestCase):
23	inv_rms = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps)	31	inv_rms = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps)
24		32
25	if gamma is not None:	33	if gamma is not None:
26	gamma = gamma.reshape(ND).float()	34	matmul_x = x * gamma.reshape(ND).float()
27	h_mix = F.linear(x * gamma, phi.float())
28	weight = h_mix * inv_rms
29	else:	35	else:
30	h_mix = F.linear(x, phi.float())	36	matmul_x = x
31	weight = h_mix * inv_rms	37	matmul_phi = phi.float()
		38	if inner_precise == 1:
		39	matmul_x = self.fp32_to_hf32(matmul_x)
		40	matmul_phi = self.fp32_to_hf32(matmul_phi)
32		41
		42	h_mix = F.linear(matmul_x, matmul_phi)
		43	weight = h_mix * inv_rms
33	h_pre, h_post, h_res = weight.split([N, N, N * N], dim=-1)	44	h_pre, h_post, h_res = weight.split([N, N, N * N], dim=-1)
34	h_res = h_res.unflatten(-1, (N, N))	45	h_res = h_res.unflatten(-1, (N, N))
35	h_pre = torch.sigmoid(h_pre * alpha[0] + bias[:N].unsqueeze(0)) + hc_eps	46	h_pre = torch.sigmoid(h_pre * alpha[0] + bias[:N].unsqueeze(0)) + hc_eps
	@@ -76,14 +87,15 @@ class TestNpuMhcPre(TestCase):
76		87
77	return (h_in, h_post, inv_rms[:, 0], h_mix, h_pre)	88	return (h_in, h_post, inv_rms[:, 0], h_mix, h_pre)
78		89
79	def custom_op_exec(self, x, phi, alpha, bias, gamma, out_flag):	90	def custom_op_exec(self, x, phi, alpha, bias, gamma, out_flag, inner_precise=0):
80	return torch_npu.npu_mhc_pre(	91	return torch_npu.npu_mhc_pre(
81	x,	92	x,
82	phi,	93	phi,
83	alpha,	94	alpha,
84	bias,	95	bias,
85	gamma=gamma,	96	gamma=gamma,
86	out_flag=out_flag	97	out_flag=out_flag,
		98	inner_precise=inner_precise
87	)	99	)
88		100
89	def build_input_tensors(self, T, n, D):	101	def build_input_tensors(self, T, n, D):
	@@ -111,13 +123,16 @@ class TestNpuMhcPre(TestCase):
111		123
112	return x, phi, alpha, bias, gamma	124	return x, phi, alpha, bias, gamma
113		125
114	def run_and_check(self, T, n, D, out_flag, output_names, tol_map):	126	def run_and_check(self, T, n, D, out_flag, output_names, tol_map, inner_precise=0):
115	with torch.no_grad():	127	with torch.no_grad():
116	x, phi, alpha, bias, gamma = self.build_input_tensors(T, n, D)	128	x, phi, alpha, bias, gamma = self.build_input_tensors(T, n, D)
117		129
118	expected_output = self.cpu_op_exec(x, phi, alpha, bias, gamma)[:len(output_names)]	130	expected_output = self.cpu_op_exec(
		131	x, phi, alpha, bias, gamma, inner_precise=inner_precise
		132	)[:len(output_names)]
119	actual_output = self.custom_op_exec(	133	actual_output = self.custom_op_exec(
120	x.npu(), phi.npu(), alpha.npu(), bias.npu(), gamma.npu(), out_flag=out_flag	134	x.npu(), phi.npu(), alpha.npu(), bias.npu(), gamma.npu(), out_flag=out_flag,
		135	inner_precise=inner_precise
121	)[:len(output_names)]	136	)[:len(output_names)]
122		137
123	for name, exp, act in zip(output_names, expected_output, actual_output):	138	for name, exp, act in zip(output_names, expected_output, actual_output):
	@@ -130,7 +145,7 @@ class TestNpuMhcPre(TestCase):
130	except AssertionError as e:	145	except AssertionError as e:
131	raise AssertionError(	146	raise AssertionError(
132	f"Output {name} compare failed for shape (T={T}, n={n}, D={D}), "	147	f"Output {name} compare failed for shape (T={T}, n={n}, D={D}), "
133	f"out_flag={out_flag}: {e}"	148	f"out_flag={out_flag}, inner_precise={inner_precise}: {e}"
134	)	149	)
135		150
136	def run_hy_and_check(self, T, n, D, out_flag, output_names, tol_map):	151	def run_hy_and_check(self, T, n, D, out_flag, output_names, tol_map):
	@@ -197,7 +212,8 @@ class TestNpuMhcPre(TestCase):
197	"h_post": 1e-3,	212	"h_post": 1e-3,
198	"h_res": 1e-3,	213	"h_res": 1e-3,
199	}	214	}
200	self.run_and_check(T, n, D, out_flag, output_names, tol_map)	215	for inner_precise in (0, 1):
		216	self.run_and_check(T, n, D, out_flag, output_names, tol_map, inner_precise=inner_precise)
201		217
202	@SupportedDevices(['Ascend950'])	218	@SupportedDevices(['Ascend950'])
203	def test_npu_mhc_pre_prefill_hy(self, device="npu"):	219	def test_npu_mhc_pre_prefill_hy(self, device="npu"):