feat: add new api (add)rms_norm_dynamic_mx

Mop_plugin/config/op_plugin_functions.yaml+6-0

	@@ -5624,6 +5624,9 @@ custom:
5624	- func: npu_rms_norm_quant(Tensor x, Tensor gamma, Tensor beta, Tensor scale, Tensor offset, float epsilon=1e-06, *, int? dst_dtype=None) -> Tensor	5624	- func: npu_rms_norm_quant(Tensor x, Tensor gamma, Tensor beta, Tensor scale, Tensor offset, float epsilon=1e-06, *, int? dst_dtype=None) -> Tensor
5625	op_api: all_version	5625	op_api: all_version
5626		5626
		5627	+ - func: npu_rms_norm_dynamic_mx_quant(Tensor x, Tensor gamma, *, Tensor? beta=None, float epsilon=1e-06, int scale_alg=0, str round_mode="rint", int dst_type=296) -> (Tensor, Tensor, Tensor)
		5628	+ op_api: all_version
		5629	+
5627	- func: npu_add_rms_norm_cast(Tensor x1, Tensor x2, Tensor gamma, float epsilon=1e-06) -> (Tensor, Tensor, Tensor, Tensor)	5630	- func: npu_add_rms_norm_cast(Tensor x1, Tensor x2, Tensor gamma, float epsilon=1e-06) -> (Tensor, Tensor, Tensor, Tensor)
5628	op_api: all_version	5631	op_api: all_version
5629	gen_opapi:	5632	gen_opapi:
	@@ -6390,6 +6393,9 @@ custom:
6390	dtype: at::kFloat	6393	dtype: at::kFloat
6391	exec: aclnnAddRmsNormDynamicQuantV2, x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon, output_mask, y1, y2, x_out, scale1, scale2	6394	exec: aclnnAddRmsNormDynamicQuantV2, x1, x2, gamma, smooth_scale1, smooth_scale2, beta, epsilon, output_mask, y1, y2, x_out, scale1, scale2
6392		6395
		6396	+ - func: npu_add_rms_norm_dynamic_mx_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? beta=None, float epsilon=1e-06, int scale_alg=0, str round_mode="rint", int dst_type=296) -> (Tensor, Tensor, Tensor, Tensor)
		6397	+ op_api: all_version
		6398	+
6393	- func: npu_attention_worker_combine(Tensor schedule_context, Tensor expert_scales, Tensor layer_id, int hidden_size, *, int token_dtype=0, int need_schedule=0) -> (Tensor, Tensor)	6399	- func: npu_attention_worker_combine(Tensor schedule_context, Tensor expert_scales, Tensor layer_id, int hidden_size, *, int token_dtype=0, int need_schedule=0) -> (Tensor, Tensor)
6394	op_api: all_version	6400	op_api: all_version
6395		6401

Aop_plugin/ops/opapi/AddRmsNormDynamicMxQuantKernelNpuOpApi.cpp+107-0

@@ -0,0 +1,107 @@
	1	+// Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+// All rights reserved.
	3	+//
	4	+// Licensed under the BSD 3-Clause License (the "License");
	5	+// you may not use this file except in compliance with the License.
	6	+// You may obtain a copy of the License at
	7	+//
	8	+// Unless required by applicable law or agreed to in writing, software
	9	+// distributed under the License is distributed on an "AS IS" BASIS,
	10	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+// See the License for the specific language governing permissions and
	12	+// limitations under the License.
	13	+
	14	+#include "op_plugin/OpApiInterface.h"
	15	+#include "op_plugin/utils/op_api_common.h"
	16	+
	17	+namespace op_api {
	18	+ using npu_preparation = at_npu::native::OpPreparation;
	19	+ using tensor_list = std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>;
	20	+ using namespace op_infer;
	21	+ namespace {
	22	+ constexpr int64_t BLOCK_SIZE_BASE_NUM = 32;
	23	+ constexpr int64_t ALIGN_NUM = 2;
	24	+ constexpr int64_t FP4_IN_UINT8_NUM = 2;
	25	+ constexpr int64_t MIN_INPUT_DIM = 1;
	26	+ constexpr int64_t MAX_INPUT_DIM = 7;
	27	+ } // namespace
	28	+
	29	+ tensor_list npu_add_rms_norm_dynamic_mx_quant(const at::Tensor &x1, const at::Tensor &x2, const at::Tensor &gamma,
	30	+ const c10::optional<at::Tensor> &beta, double epsilon, int64_t scale_alg,
	31	+ c10::string_view round_mode, int64_t dst_type)
	32	+ {
	33	+ // 输出Tensor准备
	34	+ at::Tensor y;
	35	+ at::Tensor x_out;
	36	+ at::Tensor mxscale;
	37	+ at::Tensor rstd;
	38	+
	39	+ // 参数检查
	40	+ TORCH_CHECK(x1.dim() >= MIN_INPUT_DIM && x1.dim() <= MAX_INPUT_DIM, "The x1 should be in 1~7D" + OPS_ERROR(ErrCode::PARAM));
	41	+ TORCH_CHECK(x2.dim() >= MIN_INPUT_DIM && x2.dim() <= MAX_INPUT_DIM, "The x2 should be in 1~7D" + OPS_ERROR(ErrCode::PARAM));
	42	+ TORCH_CHECK(x1.sizes() == x2.sizes(), "The shape of x1 and x2 must be the same" + OPS_ERROR(ErrCode::PARAM));
	43	+ TORCH_CHECK(x1.requires_grad() == x2.requires_grad(),
	44	+ "The requires_grad of x1 and x2 must be consistent" + OPS_ERROR(ErrCode::PARAM));
	45	+
	46	+ static const bool is_available = check_aclnn_kernel_available("aclnnAddRmsNormDynamicMxQuant");
	47	+ TORCH_CHECK(is_available,
	48	+ "Current CANN version do not support this api. Please try to update the version of CANN."
	49	+ + OPS_ERROR(ErrCode::PARAM));
	50	+
	51	+ // 类型推断
	52	+ auto y_shape = array_to_small_vector(x1.sizes());
	53	+ auto mxscale_shape = array_to_small_vector(x1.sizes());
	54	+ mxscale_shape.emplace_back(ALIGN_NUM);
	55	+
	56	+ // y shape&dtype 推导
	57	+ aclDataType y_acltype;
	58	+ bool special_output_type = (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) \|\|
	59	+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
	60	+ ASCEND_LOGI("[npu_add_rms_norm_dynamic_mx_quant]: Getting aclTensor y dtype by Parameter(dst_type): %ld", dst_type);
	61	+ if (special_output_type) {
	62	+ int64_t y_last_dim_val = y_shape[x1.dim() - 1];
	63	+ TORCH_CHECK(y_last_dim_val % FP4_IN_UINT8_NUM == 0,
	64	+ "The last dim input shape must be divisible by 2 if "
	65	+ "y dtype is torch_npu.float4_e2m1fn_x2 or torch_npu.float4_e1m2" + OPS_ERROR(ErrCode::PARAM));
	66	+ // Pytorch2.8之前最小单位是8位，不支持FP4类型，就将两个FP4合成一个FP8
	67	+ y_shape[x1.dim() - 1] = y_last_dim_val / FP4_IN_UINT8_NUM;
	68	+ y = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);
	69	+ y_acltype = c10_npu::GetAclDataType(dst_type);
	70	+ } else {
	71	+ y_acltype = c10_npu::GetAclDataType(dst_type);
	72	+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
	73	+ y = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));
	74	+ }
	75	+
	76	+ // x_out shape 推导
	77	+ auto x_out_shape = x1.sizes();
	78	+ auto x_out_dtype = x1.scalar_type();
	79	+ x_out = npu_preparation::apply_tensor_without_format(x_out_shape, x1.options().dtype(x_out_dtype));
	80	+
	81	+ // mxscale shape 推导
	82	+ int64_t last_axis_change = x1.dim() - 1;
	83	+ int64_t last_dim_size = CeilDiv(mxscale_shape[last_axis_change], BLOCK_SIZE_BASE_NUM);
	84	+ last_dim_size = (last_dim_size + ALIGN_NUM - 1) / ALIGN_NUM;
	85	+ mxscale_shape[last_axis_change] = last_dim_size;
	86	+ mxscale = npu_preparation::apply_tensor_without_format(mxscale_shape, c10::dtype(at::ScalarType::Byte));
	87	+
	88	+ // rstd shape 推导
	89	+ bool output_rstd = x1.requires_grad() && x2.requires_grad();
	90	+ if (output_rstd) {
	91	+ auto output_size = rms_norm_npu_output_size(x1, gamma);
	92	+ rstd = npu_preparation::apply_tensor_without_format(output_size[1], x1.options().dtype(at::kFloat));
	93	+ } else {
	94	+ rstd = at::empty({0}, x1.options().dtype(at::kFloat));
	95	+ }
	96	+
	97	+ // 调用NPU原生算子执行
	98	+ char round_mode_ptr = const_cast<char >(round_mode.data());
	99	+ TensorWrapper y_wrapper = {y, y_acltype};
	100	+ TensorWrapper mxscale_wrapper = {mxscale, aclDataType::ACL_FLOAT8_E8M0};
	101	+
	102	+ EXEC_NPU_CMD(aclnnAddRmsNormDynamicMxQuant, x1, x2, gamma, beta, epsilon, scale_alg,
	W W wuyi_513月23日接口兼容性判断要有，下面那个算子也是刘琦 3月24日评论： 3月24日评论：
	103	+ round_mode_ptr, y_acltype, output_rstd, y_wrapper, x_out, mxscale_wrapper, rstd);
	104	+
	105	+ return std::make_tuple(y, x_out, mxscale, rstd);
	106	+ }
	107	+} // namespace op_api

Aop_plugin/ops/opapi/RmsNormDynamicMxQuantKernelNpuOpApi.cpp+97-0

@@ -0,0 +1,97 @@
	1	+// Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+// All rights reserved.
	3	+//
	4	+// Licensed under the BSD 3-Clause License (the "License");
	5	+// you may not use this file except in compliance with the License.
	6	+// You may obtain a copy of the License at
	7	+//
	8	+// Unless required by applicable law or agreed to in writing, software
	9	+// distributed under the License is distributed on an "AS IS" BASIS,
	10	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+// See the License for the specific language governing permissions and
	12	+// limitations under the License.
	13	+
	14	+#include "op_plugin/AclOpsInterface.h"
	15	+#include "op_plugin/OpApiInterface.h"
	16	+#include "op_plugin/utils/op_api_common.h"
	17	+
	18	+namespace op_api {
	19	+using npu_preparation = at_npu::native::OpPreparation;
	20	+namespace {
	21	+constexpr int64_t BLOCK_SIZE = 32;
	22	+constexpr int64_t ALIGN_NUM = 2;
	23	+constexpr int64_t FP4_IN_UINT8_NUM = 2;
	24	+constexpr int64_t MIN_INPUT_DIM = 1;
	25	+constexpr int64_t MAX_INPUT_DIM = 7;
	26	+}; // namespace
	27	+
	28	+std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_rms_norm_dynamic_mx_quant(
	29	+ const at::Tensor &x,
	30	+ const at::Tensor &gamma,
	31	+ const c10::optional<at::Tensor> &beta,
	32	+ double epsilon,
	33	+ const int64_t scale_alg,
	34	+ c10::string_view round_mode,
	35	+ int64_t dst_type)
	36	+{
	37	+ // output
	38	+ at::Tensor y;
	39	+ at::Tensor mxscale;
	40	+ at::Tensor rstd;
	41	+
	42	+ // check params
	43	+ TORCH_CHECK(x.dim() >= MIN_INPUT_DIM && x.dim() <= MAX_INPUT_DIM, "The x should be in 1~7D" + OPS_ERROR(ErrCode::PARAM));
	44	+
	45	+ static const bool is_available = check_aclnn_kernel_available("aclnnRmsNormDynamicMxQuant");
	46	+ TORCH_CHECK(is_available,
	47	+ "Current CANN version do not support this api. Please try to update the version of CANN."
	48	+ + OPS_ERROR(ErrCode::PARAM));
	49	+
	50	+ // y
	51	+ auto y_shape = op_infer::array_to_small_vector(x.sizes());
	52	+ aclDataType y_acltype = c10_npu::GetAclDataType(dst_type);
	53	+ bool special_output_type = (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) \|\|
	54	+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
	55	+ ASCEND_LOGI("[npu_rms_norm_dynamic_mx_quant]: Getting aclTensor y dtype by Parameter(dst_type): %ld", dst_type);
	56	+ if (special_output_type) {
	57	+ int64_t last_dim_val = y_shape[x.dim() - 1];
	58	+ TORCH_CHECK(last_dim_val % FP4_IN_UINT8_NUM == 0,
	59	+ "The last dim x shape must be divisible by 2 if "
	60	+ "output dtype is torch_npu.float4_e2m1 or torch_npu.float4_e1m2" + OPS_ERROR(ErrCode::PARAM));
	61	+ y_shape[x.dim() - 1] = last_dim_val / FP4_IN_UINT8_NUM;
	62	+ y = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);
	63	+ } else {
	64	+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
	65	+ y = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));
	66	+ }
	67	+
	68	+ // mxscale
	69	+ auto mxscale_shape = op_infer::array_to_small_vector(x.sizes());
	70	+ mxscale_shape.emplace_back(ALIGN_NUM);
	71	+ int64_t last_axis_change = x.dim() - 1;
	72	+ int64_t dim_size = op_infer::CeilDiv(mxscale_shape[last_axis_change], BLOCK_SIZE);
	73	+ dim_size = (dim_size + ALIGN_NUM - 1) / ALIGN_NUM;
	74	+ mxscale_shape[last_axis_change] = dim_size;
	75	+ mxscale = npu_preparation::apply_tensor_without_format(mxscale_shape, c10::dtype(at::ScalarType::Byte));
	76	+
	77	+ // rstd
	78	+ bool output_rstd = x.requires_grad();
	79	+ if (output_rstd) {
	80	+ auto output_size = op_infer::rms_norm_npu_output_size(x, gamma);
	81	+ rstd = npu_preparation::apply_tensor_with_format(output_size[1], x.options().dtype(at::kFloat), ACL_FORMAT_ND);
	82	+ } else {
	83	+ rstd = at::empty({0}, x.options().dtype(at::kFloat));
	84	+ }
	85	+
	86	+ // call aclnn
	87	+ char round_mode_ptr = const_cast<char >(round_mode.data());
	88	+ TensorWrapper y_wrapper = {y, y_acltype};
	89	+ TensorWrapper mxscale_wrapper = {mxscale, aclDataType::ACL_FLOAT8_E8M0};
	90	+
	91	+ EXEC_NPU_CMD(aclnnRmsNormDynamicMxQuant, x, gamma, beta, epsilon, scale_alg,
	92	+ round_mode_ptr, y_acltype, output_rstd, y_wrapper, mxscale_wrapper, rstd);
	93	+
	94	+ return std::tuple<at::Tensor, at::Tensor, at::Tensor>(y, mxscale, rstd);
	95	+}
	96	+
	97	+} // namespace op_api

Mop_plugin/python/meta/_meta_registrations.py+102-0

	@@ -2230,6 +2230,54 @@ def npu_rms_norm_quant_meta(x, gamma, beta, scale, offset, epsilon=1e-06, dst_dt
2230	return torch.empty(x.size(), dtype=dst_torch_dtype, device=x.device)	2230	return torch.empty(x.size(), dtype=dst_torch_dtype, device=x.device)
2231		2231
2232		2232
		2233	+@impl(m, "npu_rms_norm_dynamic_mx_quant")
		2234	+def npu_rms_norm_dynamic_mx_quant_meta(x, gamma, *, beta=None, epsilon=1e-06, scale_alg=0, round_mode='rint', dst_type=296):
		2235	+ if scale_alg not in [0, 1]:
		2236	+ raise RuntimeError(f"Invalid scale_alg value: {scale_alg}. Expected 0 or 1." +
		2237	+ ops_error(ErrCode.PARAM))
		2238	+
		2239	+ # 以下变量为本函数局部使用
		2240	+ align_num = 2
		2241	+ mxscale_block_size = 32
		2242	+
		2243	+ dst_torch_dtype = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP.get(dst_type, torch.int8)
		2244	+ if dst_torch_dtype == torch.float8_e5m2 or dst_type == 291:
		2245	+ y = torch.empty_like(x, dtype=torch.float8_e5m2)
		2246	+ elif dst_torch_dtype == torch.float8_e4m3fn or dst_type == 292:
		2247	+ y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
		2248	+ else: # float4_e2m1, float4_e1m2
		2249	+ if x.size(x.dim() - 1) % 2:
		2250	+ raise RuntimeError("If output dtype is float4_e2m1 or float4_e1m2, " \
		2251	+ "the last dim of input must be divisible by 2, " +
		2252	+ ops_error(ErrCode.PARAM))
		2253	+ y_shape = []
		2254	+ for dim in range(x.dim() - 1):
		2255	+ y_shape.append(x.size(dim))
		2256	+ y_shape.append(x.size(x.dim() - 1) // align_num)
		2257	+ y = x.new_empty(y_shape, dtype=torch.uint8)
		2258	+
		2259	+ # mxscale
		2260	+ mxscale_shape = []
		2261	+ for dim in range(x.dim()):
		2262	+ mxscale_shape.append(x.size(dim))
		2263	+ mxscale_shape.append(2)
		2264	+ last_axis_change = x.dim() - 1
		2265	+ last_dim_size = int(math.ceil(mxscale_shape[last_axis_change] / (mxscale_block_size * align_num)))
		2266	+ mxscale_shape[last_axis_change] = last_dim_size
		2267	+ mxscale = x.new_empty(mxscale_shape, dtype=torch.uint8)
		2268	+
		2269	+ # rstd
		2270	+ rstd_dim = x.dim() - gamma.dim()
		2271	+ ret = []
		2272	+ for dim in range(x.dim()):
		2273	+ if dim < rstd_dim:
		2274	+ ret.append(x.size(dim))
		2275	+ else:
		2276	+ ret.append(1)
		2277	+ rstd = torch.empty(ret, dtype=torch.float32, device='meta')
		2278	+ return (y, mxscale, rstd)
		2279	+
		2280	+
2233	@impl(m, "npu_add_rms_norm_cast")	2281	@impl(m, "npu_add_rms_norm_cast")
2234	def npu_add_rms_norm_cast_meta(x1, x2, gamma, epsilon=1e-6):	2282	def npu_add_rms_norm_cast_meta(x1, x2, gamma, epsilon=1e-6):
2235	rstd_dim = x1.dim() - gamma.dim()	2283	rstd_dim = x1.dim() - gamma.dim()
	@@ -2263,6 +2311,60 @@ def npu_add_rms_norm_dynamic_quant_meta(x1, x2, gamma, *, smooth_scale1=None, sm
2263	torch.empty(x1.size()[:-1], dtype=torch.float32, device=x1.device))	2311	torch.empty(x1.size()[:-1], dtype=torch.float32, device=x1.device))
2264		2312
2265		2313
		2314	+@impl(m, "npu_add_rms_norm_dynamic_mx_quant")
		2315	+def npu_add_rms_norm_dynamic_mx_quant_meta(x1, x2, gamma, *, beta=None, epsilon=1e-6, scale_alg=0, round_mode='rint', dst_type=296):
		2316	+ if scale_alg not in [0, 1]:
		2317	+ raise RuntimeError(f"Invalid scale_alg value: {scale_alg}. Expected 0 or 1." +
		2318	+ ops_error(ErrCode.PARAM))
		2319	+
		2320	+ dim_num = x1.dim()
		2321	+
		2322	+ # 以下变量为本函数局部使用
		2323	+ align_num = 2
		2324	+ mxscale_block_size = 32
		2325	+
		2326	+ # y
		2327	+ torch_dtype = TORCH_DTYPE_ENUM_VALUE_TO_SCALAR_TYPE_MAP.get(dst_type, torch.int8)
		2328	+ if torch_dtype == torch.float8_e5m2 or dst_type == torch_npu.float8_e5m2:
		2329	+ y = torch.empty_like(x1, dtype=torch.float8_e5m2)
		2330	+ elif torch_dtype == torch.float8_e4m3fn or dst_type == torch_npu.float8_e4m3fn:
		2331	+ y = torch.empty_like(x1, dtype=torch.float8_e4m3fn)
		2332	+ else: # float4_e2m1, float4_e1m2
		2333	+ if x1.size(dim_num - 1) % 2:
		2334	+ raise RuntimeError("If output dtype is float4_e2m1 or float4_e1m2, " \
		2335	+ "the last dim of input must be divisible by 2, " +
		2336	+ ops_error(ErrCode.PARAM))
		2337	+ y_shape = []
		2338	+ for dim in range(dim_num - 1):
		2339	+ y_shape.append(x1.size(dim))
		2340	+ y_shape.append(x1.size(dim_num - 1) // align_num)
		2341	+ y = x1.new_empty(y_shape, dtype=torch.uint8)
		2342	+
		2343	+ # x_out
		2344	+ x_out = torch.empty_like(x1, dtype=x1.dtype)
		2345	+
		2346	+ # mxscale
		2347	+ mxscale_shape = []
		2348	+ for dim in range(dim_num):
		2349	+ mxscale_shape.append(x1.size(dim))
		2350	+ mxscale_shape.append(2)
		2351	+ last_axis_change = dim_num - 1
		2352	+ last_dim_size = int(math.ceil(mxscale_shape[last_axis_change] / (mxscale_block_size * align_num)))
		2353	+ mxscale_shape[last_axis_change] = last_dim_size
		2354	+ mxscale = x1.new_empty(mxscale_shape, dtype=torch.uint8)
		2355	+
		2356	+ # rstd
		2357	+ rstd_dim = dim_num - gamma.dim()
		2358	+ ret = []
		2359	+ for i in range(dim_num):
		2360	+ if i < rstd_dim:
		2361	+ ret.append(x1.size(i))
		2362	+ else:
		2363	+ ret.append(1)
		2364	+ rstd = torch.empty(ret, dtype=torch.float32, device='meta')
		2365	+ return (y, x_out, mxscale, rstd)
		2366	+
		2367	+
2266	@impl(m, "npu_rms_norm_backward")	2368	@impl(m, "npu_rms_norm_backward")
2267	def npu_rms_norm_backward_meta(dy, self, gamma, rstd):	2369	def npu_rms_norm_backward_meta(dy, self, gamma, rstd):
2268	return (torch.empty_like(self, dtype=self.dtype), torch.empty_like(gamma, dtype=torch.float32))	2370	return (torch.empty_like(self, dtype=self.dtype), torch.empty_like(gamma, dtype=torch.float32))

Mtest/core_tests/test_fake_tensor.py+37-0

             self.assertTrue(y.dtype == torch.float8_e5m2)
++class TestRmsNormDynamicMxQuant(TestCase):
++    def test_npu_rms_norm_dynamic_mx_quant_meta(self):
++        with FakeTensorMode():
++            x = torch.randn([8, 64], dtype=torch.float16, device='npu')
++            gamma = torch.ones([64, ], dtype=torch.float16, device='npu')
++            beta = torch.zeros([64, ], dtype=torch.float16, device='npu')
++            y_npu, mxscale_npu, rstd_npu = torch_npu.npu_rms_norm_dynamic_mx_quant(
++                x, gamma, beta=beta, epsilon=1e-6, scale_alg=0, round_mode="rint", dst_type=torch_npu.float8_e5m2
++            )
++            self.assertEqual(y_npu.shape, x.shape)
++            self.assertEqual(y_npu.dtype, torch.float8_e5m2)
++            self.assertEqual(mxscale_npu.shape, torch.Size([8, 1, 2]))
++            self.assertEqual(mxscale_npu.dtype, torch.uint8)
++            self.assertEqual(rstd_npu.shape, torch.Size([8, 1]))
++            self.assertEqual(rstd_npu.dtype, torch.float32)
++
++
 class TestNpuRmsNorm(TestCase):
     def test_npu_rms_norm(self):
         with FakeTensorMode():
             self.assertEqual(y2_npu.dtype, torch.int32)
++class TestAddRmsNormDynamicMxQuant(TestCase):
++    def test_npu_add_rms_norm_dynamic_mx_quant_meta(self):
++        with FakeTensorMode():
++            x1 = torch.randn([8, 64], dtype=torch.float16, device='npu')
++            x2 = torch.randn([8, 64], dtype=torch.float16, device='npu')
++            gamma = torch.ones([64, ], dtype=torch.float16, device='npu')
++            beta = torch.zeros([64, ], dtype=torch.float16, device='npu')
++            y_npu, x_out_npu, mxscale_npu, rstd_npu = torch_npu.npu_add_rms_norm_dynamic_mx_quant(
++                x1, x2, gamma, beta=beta, epsilon=1e-6, scale_alg=0, round_mode="rint", dst_type=torch_npu.float8_e5m2
++            )
++            self.assertEqual(y_npu.shape, x1.shape)
++            self.assertEqual(y_npu.dtype, torch.float8_e5m2)
++            self.assertEqual(x_out_npu.shape, x1.shape)
++            self.assertEqual(x_out_npu.dtype, x1.dtype)
++            self.assertEqual(mxscale_npu.shape, torch.Size([8, 1, 2]))
++            self.assertEqual(mxscale_npu.dtype, torch.uint8)
++            self.assertEqual(rstd_npu.shape, torch.Size([8, 1]))
++            self.assertEqual(rstd_npu.dtype, torch.float32)
++
++
 class TestMoeUpdateExpert(TestCase):
     def test_moe_update_expert(self):
         with FakeTensorMode():

Mtest/core_tests/torch_npu_OpApi_schema_all.json+6-0

	@@ -206,6 +206,9 @@
206	"func: npu_add_rms_norm_dynamic_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? smooth_scale1=None, Tensor? smooth_scale2=None, Tensor? beta=None, float epsilon=1e-6, bool[2] output_mask=[], ScalarType? y_dtype=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {	206	"func: npu_add_rms_norm_dynamic_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? smooth_scale1=None, Tensor? smooth_scale2=None, Tensor? beta=None, float epsilon=1e-6, bool[2] output_mask=[], ScalarType? y_dtype=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {
207	"version": ["all_version"]	207	"version": ["all_version"]
208	},	208	},
		209	+ "func: npu_add_rms_norm_dynamic_mx_quant(Tensor x1, Tensor x2, Tensor gamma, *, Tensor? beta=None, float epsilon=1e-06, int scale_alg=0, str round_mode=\"rint\", int dst_type=296) -> (Tensor, Tensor, Tensor, Tensor)": {
		210	+ "version": ["all_version"]
		211	+ },
209	"op_api: torch_npu.attention_worker_scheduler_(args, *kwargs)": {	212	"op_api: torch_npu.attention_worker_scheduler_(args, *kwargs)": {
210	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]	213	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]
211	},	214	},
	@@ -1081,6 +1084,9 @@
1081	"func: npu_rms_norm_backward(Tensor dy, Tensor self, Tensor gamma, Tensor rstd) -> (Tensor, Tensor)": {	1084	"func: npu_rms_norm_backward(Tensor dy, Tensor self, Tensor gamma, Tensor rstd) -> (Tensor, Tensor)": {
1082	"version": ["all_version"]	1085	"version": ["all_version"]
1083	},	1086	},
		1087	+ "func: npu_rms_norm_dynamic_mx_quant(Tensor x, Tensor gamma, *, Tensor? beta=None, float epsilon=1e-06, int scale_alg=0, str round_mode=\"rint\", int dst_type=296) -> (Tensor, Tensor, Tensor)": {
		1088	+ "version": ["all_version"]
		1089	+ },
1084	"func: npu_swiglu(Tensor self, int dim=-1) -> Tensor": {	1090	"func: npu_swiglu(Tensor self, int dim=-1) -> Tensor": {
1085	"version": ["all_version"]	1091	"version": ["all_version"]
1086	},	1092	},

Atest/test_custom_ops/test_npu_add_rms_norm_dynamic_mx_quant.py+80-0

@@ -0,0 +1,80 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+# All rights reserved.
	3	+
	4	+# Licensed under the BSD 3-Clause License (the "License");
	5	+# you may not use this file except in compliance with the License.
	6	+# You may obtain a copy of the License at
	7	+
	8	+# Unless required by applicable law or agreed to in writing, software
	9	+# distributed under the License is distributed on an "AS IS" BASIS,
	10	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+# See the License for the specific language governing permissions and
	12	+# limitations under the License.
	13	+
	14	+import unittest
	15	+import math
	Y Y yang-binrong3月23日缺少文件头
	16	+import numpy as np
	17	+import torch
	18	+import torch_npu
	19	+
	20	+from torch_npu.testing.testcase import TestCase, run_tests
	21	+from torch_npu.testing.common_utils import SupportedDevices
	22	+from torch.testing import assert_close
	23	+
	24	+class TestAddRmsNormDynamicMxQuant(TestCase):
	25	+
	26	+ # 此处传递的是已经.npu()后的
	27	+ def npu_op_exec(self, x1, x2, gamma, beta=None, epsilon=1e-6, scale_alg=0, round_mode="rint", dst_type=torch_npu.float8_e5m2):
	28	+ return torch_npu.npu_add_rms_norm_dynamic_mx_quant(x1, x2, gamma, beta=beta, epsilon=epsilon,
	29	+ scale_alg=scale_alg, round_mode=round_mode, dst_type=dst_type)
	30	+
	31	+ def golden_op_exec(self, input_tensor):
	32	+ if torch.all(torch.eq(input_tensor, 0.0)) and input_tensor.shape == torch.Size([1, 2, 2]):
	33	+ device = input_tensor.device
	34	+ y = torch.tensor([[[0, 0], [0, 0]]], dtype=torch.float8_e5m2, device=device)
	35	+ x = torch.tensor([[[0, 0], [0, 0]]], dtype=torch.float8_e5m2, device=device)
	36	+ mxscale = torch.tensor([[[[0, 0]], [[0, 0]]]], dtype=torch.uint8, device=device)
	37	+ rstd = torch.tensor([[[1000], [1000]]], dtype=torch.float32, device=device)
	38	+ return y, x, mxscale, rstd
	39	+
	40	+ def generate_input(self, input, value, dtype="float16"):
	41	+ if dtype == "float32":
	42	+ data_type = torch.float32
	43	+ elif dtype == "float16":
	44	+ data_type = torch.float16
	45	+ elif dtype == "bfloat16":
	46	+ data_type = torch.bfloat16
	47	+ input_tensor = torch.full(input, value, dtype=data_type)
	48	+ return input_tensor
	49	+
	50	+ @SupportedDevices(['Ascend950'])
	51	+ def test_npu_add_rms_norm_quant_float8_e5m2_with_rstd(self, device="npu"):
	52	+ x1 = self.generate_input(input=[1, 2, 2], value=0.0, dtype="bfloat16")
	53	+ x2 = self.generate_input(input=[1, 2, 2], value=0.0, dtype="bfloat16")
	54	+ C = x1.shape[-1]
	55	+ gamma = self.generate_input(input=[C], value=0.0, dtype="float32")
	56	+ beta = self.generate_input(input=[C], value=0.0, dtype="float32")
	57	+ x1 = x1.to(device).requires_grad_(True)
	58	+ x2 = x2.to(device).requires_grad_(True)
	59	+ gamma = gamma.to(device)
	60	+ beta = beta.to(device)
	61	+ eps = 1e-6
	62	+ scale_alg = 0
	63	+ round_mode = "rint"
	64	+ out_dtype = 23
	65	+
	66	+ golden_output = self.golden_op_exec(x1.clone().detach())
	67	+ npu_output = self.npu_op_exec(x1, x2, gamma, beta=beta, epsilon=eps, scale_alg=scale_alg, round_mode=round_mode, dst_type=out_dtype)
	68	+ y = npu_output[0].view([1, 2, 2]).view(torch.uint8)
	69	+ x = npu_output[1].view([1, 2, 2]).view(torch.bfloat16)
	70	+ mxscale = npu_output[2].view([1, 2, 1, 2]).to(torch.uint8)
	71	+ rstd = npu_output[3].view([1, 2, 1]).to(torch.float32)
	72	+
	73	+ assert torch.all(y == golden_output[0].view(torch.uint8))
	74	+ assert torch.all(x == golden_output[1].view(torch.bfloat16))
	75	+ assert_close(golden_output[2], mxscale, atol=0.01, rtol=0.001)
	76	+ assert torch.all(rstd == golden_output[3].view(torch.float32))
	77	+
	78	+
	79	+if __name__ == "__main__":
	80	+ run_tests()

Atest/test_custom_ops/test_npu_rms_norm_dynamic_mx_quant.py+76-0

@@ -0,0 +1,76 @@
	1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd
	2	+# All rights reserved.
	3	+
	4	+# Licensed under the BSD 3-Clause License (the "License");
	5	+# you may not use this file except in compliance with the License.
	6	+# You may obtain a copy of the License at
	7	+
	8	+# Unless required by applicable law or agreed to in writing, software
	9	+# distributed under the License is distributed on an "AS IS" BASIS,
	10	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	11	+# See the License for the specific language governing permissions and
	12	+# limitations under the License.
	13	+
	14	+import unittest
	15	+import math
	Y Y yang-binrong3月23日缺少文件头
	16	+import numpy as np
	17	+import torch
	18	+import torch_npu
	19	+
	20	+from torch_npu.testing.testcase import TestCase, run_tests
	21	+from torch_npu.testing.common_utils import SupportedDevices
	22	+from torch.testing import assert_close
	23	+
	24	+class TestRmsNormDynamicMxQuant(TestCase):
	25	+
	26	+ # 此处传递的是已经.npu()后的
	27	+ def npu_op_exec(self, x, gamma, beta=None, epsilon=1e-6, scale_alg=0, round_mode="rint", dst_type=torch_npu.float8_e5m2):
	28	+ return torch_npu.npu_rms_norm_dynamic_mx_quant(x, gamma, beta=beta, epsilon=epsilon, scale_alg=scale_alg,
	29	+ round_mode=round_mode, dst_type=dst_type)
	30	+
	31	+ def golden_op_exec(self, input_tensor):
	32	+ if torch.all(torch.eq(input_tensor, 0.0)) and input_tensor.shape == torch.Size([1, 2, 2]):
	33	+ device = input_tensor.device
	34	+ y = torch.tensor([[[0, 0], [0, 0]]], dtype=torch.float8_e5m2, device=device)
	35	+ mxscale = torch.tensor([[[[0, 0]], [[0, 0]]]], dtype=torch.uint8, device=device)
	36	+ rstd = torch.tensor([[[1000], [1000]]], dtype=torch.float32, device=device)
	37	+ return y, mxscale, rstd
	38	+
	39	+ def generate_input(self, input, value, dtype="float16"):
	40	+ if dtype == "float32":
	41	+ data_type = torch.float32
	42	+ elif dtype == "float16":
	43	+ data_type = torch.float16
	44	+ elif dtype == "bfloat16":
	45	+ data_type = torch.bfloat16
	46	+ input_tensor = torch.full(input, value, dtype=data_type)
	47	+ return input_tensor
	48	+
	49	+ @SupportedDevices(['Ascend950'])
	50	+ def test_npu_rms_norm_quant_float8_e5m2_with_rstd(self, device="npu"):
	51	+ x = self.generate_input(input=[1, 2, 2], value=0.0, dtype="bfloat16")
	52	+ C = x.shape[-1]
	53	+ gamma = self.generate_input(input=[C], value=0.0, dtype="float32")
	54	+ beta = self.generate_input(input=[C], value=0.0, dtype="float32")
	55	+ x = x.to(device).requires_grad_(True)
	56	+
	57	+ gamma = gamma.to(device)
	58	+ beta = beta.to(device)
	59	+ eps = 1e-6
	60	+ scale_alg=0
	61	+ round_mode="rint"
	62	+ out_dtype = 23
	63	+
	64	+ golden_output = self.golden_op_exec(x)
	65	+ npu_output = self.npu_op_exec(x, gamma, beta=beta, epsilon=eps, scale_alg=scale_alg, round_mode=round_mode, dst_type=out_dtype)
	66	+ y = npu_output[0].view([1, 2, 2]).view(torch.uint8)
	67	+ mxscale = npu_output[1].view([1, 2, 1, 2]).to(torch.uint8)
	68	+ rstd = npu_output[2].view([1, 2, 1]).to(torch.float32)
	69	+
	70	+ assert torch.all(y == golden_output[0].view(torch.uint8))
	71	+ assert_close(golden_output[1], mxscale, atol=0.01, rtol=0.001)
	72	+ assert torch.all(rstd == golden_output[2].view(torch.float32))
	73	+
	74	+
	75	+if __name__ == "__main__":
	76	+ run_tests()