【950】op-plugin回合主仓：moe_re_routing/moe_init_routing_v2/scaled_masked_softmax/group_norm_silu/moe_gating_top_k_softmax

Mcodegen/templates/_op_plugin_docs.py+88-23

	@@ -3981,7 +3981,7 @@ result: Tensor类型, 代表量化matmul的计算结果.
3981	该接口支持图模式.	3981	该接口支持图模式.
3982	传入的x1、x2、scale不能是空.	3982	传入的x1、x2、scale不能是空.
3983	x1、x2、bias、scale、offset、pertoken_scale、output_dtype的数据类型和数据格式需要在支持的范围之内.	3983	x1、x2、bias、scale、offset、pertoken_scale、output_dtype的数据类型和数据格式需要在支持的范围之内.
3984	-x1与x2最后一维的shape大小不能超过65535.	3984	+当x1的数据类型为float8_e4m3fn, x2_dtype为torch_npu.float4_e2m1或torch_npu.float4_e1m2的情况下, x1、x2的k值必须是64的倍数并且大小不能超过65535, x2的n值大小不能超过65535. 其他情况, x1与x2最后一维的shape大小不能超过65535.
3985	目前输出int8或float16且无pertoken_scale情况下, 图模式不支持scale直接传入float32数据类型.	3985	目前输出int8或float16且无pertoken_scale情况下, 图模式不支持scale直接传入float32数据类型.
3986	如果在PyTorch图模式中使用本接口, 且环境变量ENABLE_ACLNN=false, 则在调用接口前需要对shape为(n, k//8)的x2数据进行转置, 转置过程应写在图中.	3986	如果在PyTorch图模式中使用本接口, 且环境变量ENABLE_ACLNN=false, 则在调用接口前需要对shape为(n, k//8)的x2数据进行转置, 转置过程应写在图中.
3987	支持将x2转为昇腾亲和的数据排布以提高搬运效率. 需要调用torch_npu.npu_format_cast完成输入x2(weight)为昇腾亲和的数据排布功能.	3987	支持将x2转为昇腾亲和的数据排布以提高搬运效率. 需要调用torch_npu.npu_format_cast完成输入x2(weight)为昇腾亲和的数据排布功能.
	@@ -9277,6 +9277,65 @@ print(res)
9277	"""	9277	"""
9278	)	9278	)
9279		9279
		9280	+_add_torch_npu_docstr(
		9281	+ "npu_moe_gating_top_k_softmax_v2",
		梁松伟2025年12月1日新增接口需要在test/allowlist_for_publicAPI.json补充下，会导致ut不通过
		9282	+ """
		9283	+接口原型：
		9284	+npu_moe_gating_top_k_softmax_v2(Tensor x, *, int k=1, Tensor? finished=None, int? renorm=0, bool? output_softmax=False) -> (Tensor, Tensor, Tensor)
		9285	+
		9286	+功能描述
		9287	+MoE计算中，当renorm参数设置为0时，对输入x做softmax操作，再做topk操作；当renorm参数设置为1时，对x做topk操作，后做softmax操作。
		9288	+
		9289	+参数说明
		9290	+x: Tensor类型, 必选输入, 表示待计算的输入要求是一个2D/3D的Tensor, 数据类型支持float16、bfloat16、float32, 数据格式要求为ND.
		9291	+finished: Tensor类型, 可选输入, 表示输入中需要参与计算的行, 要求是一个1D/2D的Tensor, 数据类型支持bool, shape为x[:-1], 数据格式要求为ND.
		9292	+k: Host侧的int类型, 表示topk的k值, 大小为0<k<=x的-1轴大小, k<=1024.
		9293	+renorm: int类型，可选输入，表示先计算softmax还是先计算topk。
		9294	+output_softmax：bool类型，可选输入，表示是否输出softmax的结果，取值true和false。true表示输出softmax的结果，false表示不输出。
		9295	+
		9296	+输出说明
		9297	+y: Tensor类型, 对x做softmax后取的topk值, 要求是一个2D/3D的Tensor, 数据类型与x需要保持一致, 其非-1轴要求与x的对应轴大小一致, 其-1轴要求其大小同k值. 数据格式要求为ND.
		9298	+expert_idx: Tensor类型, 对x做softmax后取topk值的索引, 即专家的序号. shape要求与y一致, 数据类型支持int32, 数据格式要求为ND.
		9299	+row_idx: Tensor类型, 指示每个位置对应的原始行位置, shape要求与y一致, 数据类型支持int32, 数据格式要求为ND.
		9300	+
		9301	+约束说明
		9302	+该接口支持推理场景下使用.
		9303	+该接口支持图模式(PyTorch 2.1版本).
		9304	+
		9305	+支持的型号
		9306	+Atlas A2 训练系列产品/Atlas 800I A2 推理产品
		9307	+Atlas A3 训练系列产品/Atlas A3 推理系列产品
		9308	+调用示例
		9309	+单算子模式调用
		9310	+import torch
		9311	+import torch_npu
		9312	+x = torch.rand((3, 3), dtype=torch.float32).to("npu")
		9313	+finished = torch.randint(2, size=(3,), dtype=torch.bool).to("npu")
		9314	+y, expert_idx, row_idx = torch_npu.npu_moe_gating_top_k_softmax(
		9315	+ x, finished, k=2)
		9316	+图模式调用
		9317	+import torch
		9318	+import torch_npu
		9319	+import torchair as tng
		9320	+from torchair.configs.compiler_config import CompilerConfig
		9321	+torch_npu.npu.set_compile_mode(jit_compile=True)
		9322	+config = CompilerConfig()
		9323	+npu_backend = tng.get_npu_backend(compiler_config=config)
		9324	+device=torch.device(f'npu:0')
		9325	+torch_npu.npu.set_device(device)
		9326	+class MoeGatingTopkSoftmaxModelV2(torch.nn.Module):
		9327	+ def __init__(self):
		9328	+ super().__init__()
		9329	+ def forward(self, x, finished, k, renorm, output_softmax):
		9330	+ res = torch_npu.npu_moe_gating_top_k_softmax_v2(x = x, finished = finished, k = k, renorm = renorm, output_softmax = output_softmax)
		9331	+ return res
		9332	+x = torch.randn((2, 4, 6),device='npu',dtype=torch.float16).npu()
		9333	+moe_gating_topk_softmax_model_v2 = MoeGatingTopkSoftmaxV2Model().npu()
		9334	+moe_gating_topk_softmax_model_v2 = torch.compile(moe_gating_topk_softmax_model_v2, backend=npu_backend, dynamic=True)
		9335	+res = moe_gating_topk_softmax_model_v2(x = x, finished = None, k = 2, renorm = 0, output_softmax = output_softmax = True)
		9336	+"""
		9337	+)
		9338	+
9280	_add_torch_npu_docstr(	9339	_add_torch_npu_docstr(
9281	"npu_moe_gating_top_k",	9340	"npu_moe_gating_top_k",
9282	"""	9341	"""
	@@ -9478,37 +9537,43 @@ torch_npu.npu_moe_init_routing_v2(Tensor x, Tensor expert_idx, *, Tensor? scale=
9478	动态quant场景下不输入。	9537	动态quant场景下不输入。
9479	active_num：int类型，表示总的最大处理row数，输出expanded_x只有这么多行是有效的，当前入参校验需大于等于0。当前未使用，校验需等于NUM_ROWS*K。	9538	active_num：int类型，表示总的最大处理row数，输出expanded_x只有这么多行是有效的，当前入参校验需大于等于0。当前未使用，校验需等于NUM_ROWS*K。
9480	expert_capacity：int类型，表示每个专家能够处理的tokens数，取值范围大于等于0。当前未使用，仅校验非空。	9539	expert_capacity：int类型，表示每个专家能够处理的tokens数，取值范围大于等于0。当前未使用，仅校验非空。
9481	- expert_num：int类型，表示专家数。expert_tokens_num_type为key_value模式时，取值范围为[0, 5120]；其他模式取值范围为[0, 10240]。	9540	+ expert_num：int类型，表示专家数，要求大于0。expert_tokens_num_type为key_value模式时，取值范围为(0, 5120]；其他模式取值范围为(0, 10240]。
9482	drop_pad_mode：int类型，表示是否为drop_pad场景，取值为0和1。0表示dropless场景，该场景下不校验expert_capacity。1表示drop_pad场景。当前仅支持0。	9541	drop_pad_mode：int类型，表示是否为drop_pad场景，取值为0和1。0表示dropless场景，该场景下不校验expert_capacity。1表示drop_pad场景。当前仅支持0。
9483	expert_tokens_num_type：int类型，取值为0、1和2。0表示cumsum模式；1表示count模式，即输出的值为各个专家处理的token数量的累计值；2表示key_value模式，即输出的值为专家和对应专家处理token数量的累计值。当前仅支持1和2。	9542	expert_tokens_num_type：int类型，取值为0、1和2。0表示cumsum模式；1表示count模式，即输出的值为各个专家处理的token数量的累计值；2表示key_value模式，即输出的值为专家和对应专家处理token数量的累计值。当前仅支持1和2。
9484	expert_tokens_num_flag：bool类型，表示是否输出expert_token_cumsum_or_count，默认False表示不输出。当前仅支持True。	9543	expert_tokens_num_flag：bool类型，表示是否输出expert_token_cumsum_or_count，默认False表示不输出。当前仅支持True。
9485	quant_mode：int类型，表示量化模式，支持取值为0、1、-1。0表示静态量化，-1表示不量化场景；1表示动态quant场景。当前仅支持-1和1。x数据类型为int8时仅支持-1，不可再量化。	9544	quant_mode：int类型，表示量化模式，支持取值为0、1、-1。0表示静态量化，-1表示不量化场景；1表示动态quant场景。当前仅支持-1和1。x数据类型为int8时仅支持-1，不可再量化。
9486	- active_expert_range：int类型长度为2的数组，表示活跃expert的范围。数组内值为[expert_start, expert_end]，表示活跃的expert范围在expert_start到expert_end之间，左闭右开。要求值大于等于0，并且expert_end不大于expert_num。	9545	+ active_expert_range：int类型长度为2的数组，表示活跃expert的范围。数组内值为[expert_start, expert_end]，表示活跃的expert范围在[expert_start, expert_end)区间内，左闭右开。要求数组内的值大于等于0，并且expert_end不大于expert_num。
9487	- row_idx_type：int类型，表示输出expanded_row_idx使用的索引类型，支持取值0和1，默认值0。0表示gather类型的索引；1表示scatter类型的索引。性能模板下仅支持1。	9546	+ row_idx_type：int类型，表示输出expanded_row_idx使用的索引类型，支持取值0和1，默认值0。0表示gather类型的索引；1表示scatter类型的索引。
9488	输出说明	9547	输出说明
9489	- expanded_x：Tensor类型，根据expert_idx进行扩展过的特征，要求是2D的Tensor，shape为(NUM_ROWSK, H)。非量化场景下数据类型同x；量化场景下数据类型支持int8。数据格式要求为ND。前available_idx_numH个元素为有效数据，其余由row_idx_type决定。其中available_idx_num为expert_idx中active_expert_range范围的元素的个数。量化场景下，当x的数据类型为int8时，输出值未定义。	9548	+ expanded_x：Tensor类型，根据expert_idx进行扩展过的特征，要求是2D的Tensor，shape为(NUM_ROWSK, H)。非量化场景下数据类型同x；量化场景下数据类型支持int8。数据格式要求为ND。前available_idx_numH个元素为有效数据，其余由row_idx_type决定。其中available_idx_num为expert_idx中在active_expert_range范围的元素的个数。量化场景下，当x的数据类型为int8时，输出值未定义。
9490	- expanded_row_idx：Tensor类型，expanded_x和x的映射关系，要求是1D的Tensor，shape为(NUM_ROWS*K, )，数据类型支持int32，数据格式要求为ND。前available_idx_num个元素为有效数据，其余无效数据由row_idx_type决定，其中available_idx_num为expert_idx中active_expert_range范围的元素的个数。row_idx_type为0时，无效数据由-1填充；row_idx_type为1时，无效数据未初始化。	9549	+ expanded_row_idx：Tensor类型，expanded_x和x的映射关系，要求是1D的Tensor，shape为(NUM_ROWS*K,)，数据类型支持int32，数据格式要求为ND。当row_idx_type为0时，有效元素与无效元素共存，其中无效元素由-1填充；当row_idx_type为1时，前available_idx_num个元素有效，其余无效元素值未定义。其中available_idx_num为expert_idx中在active_expert_range范围的元素的个数。量化场景下，当x的数据类型为int8时，输出值未定义。
9491	- expert_token_cumsum_or_count：Tensor类型。在expert_tokens_num_type为1的场景下，要求是1D的Tensor，表示active_expert_range范围内expert对应的处理token的总数。shape为(expert_end-expert_start, )；在expert_tokens_num_type为2的场景下，要求是2D的Tensor，shape为(expert_num, 2)，表示active_expert_range范围内token总数为非0的expert，以及对应expert处理token的总数；expert id在active_expert_range范围且剔除对应expert处理token为0的元素对为有效元素对，存放于Tensor头部并保持原序。数据类型支持int64，数据格式要求为ND。	9550	+ expert_token_cumsum_or_count：Tensor类型，数据类型支持int64，数据格式要求为ND。
9492	- expanded_scale：Tensor类型，数据类型支持float32，数据格式要求为ND。令available_idx_num为active_expert_range范围的元素的个数。	9551	+ 在expert_tokens_num_type为1的场景下，要求是1D的Tensor，表示active_expert_range范围内每个expert对应的处理token的总数，shape为(expert_end-expert_start,)。
9493	- 非量化场景下，即quant_mode为-1，shape为(NUM_ROWSK, )。当scale未输入时，输出值未定义；当scale输入时，输出表示一个1D的Tensor，前available_idx_numH个元素为有效数据，其余为无效数据。	9552	+ 在expert_tokens_num_type为2的场景下，要求是2D的Tensor，shape为(expert_num, 2)，表示active_expert_range范围内的每个expert的expert_idx及其对应处理的token总数。有效元素对是指expert_idx在active_expert_range范围内，且处理的token数不为0的元素对，这些有效元素对按原顺序存放在Tensor头部。如果有效元素对的数量少于expert_num，其后会跟一对元素对(0,0)以表示有效元素对的结束。
9494	- 动态quant场景下，即quant_mode为1，输出量化计算过程中scale的中间值，shape为(NUM_ROWS*K)。当scale未输入时，输出值未定义；当scale输入时，输出表示一个1D的Tensor，前available_idx_num个元素为有效数据，其余为无效数据，若x的输入类型为int8，输出值未定义。	9553	+ expanded_scale：Tensor类型，数据类型支持float32，数据格式要求为ND。令available_idx_num为expert_idx中在active_expert_range范围的元素的个数。
		9554	+ 非量化场景下，即quant_mode为-1，shape为(NUM_ROWS*K,)。当scale未输入时，输出值未定义；当scale输入时，输出表示一个1D的Tensor，前available_idx_num个元素为有效数据，其余为无效数据。
		9555	+ 动态量化场景下，即quant_mode为1，输出量化计算过程中scale的中间值，shape为(NUM_ROWS*K,)。输出表示一个1D的Tensor，前available_idx_num个元素为有效数据，其余为无效数据，若x的输入类型为int8，输出值未定义。
9495	约束说明	9556	约束说明
9496	该接口支持推理场景下使用。	9557	该接口支持推理场景下使用。
9497	该接口支持图模式。	9558	该接口支持图模式。
9498	不支持静态量化模式。	9559	不支持静态量化模式。
9499	- 该算子支持两种性能模板，进入两种性能模板需要分别额外满足以下条件，不满足条件则进入通用模板：	9560	+ 该接口在部分产品型号下，支持两种性能模板。进入两种性能模板需要分别额外满足以下条件，不满足条件则进入通用模板：
9500	- 进入低时延性能模板需要同时满足以下条件：	9561	+ 支持性能模板的产品型号：
9501	- x、expert_idx、scale输入Shape要求分别为：(1, 7168)、(1, 8)、(256, 7168)	9562	+ Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件
9502	- x数据类型要求：bfloat16	9563	+ Atlas A3 训练系列产品/Atlas A3 推理系列产品
9503	- 属性要求：active_expert_range=[0,256]、 quant_mode=1、expert_tokens_num_type=2、expert_num=256	9564	+ 性能模板的约束条件：
9504	- 进入大batch性能模板需要同时满足以下条件：	9565	+ 进入低时延性能模板需要同时满足以下条件：
9505	- NUM_ROWS范围为[384, 8192]	9566	+ x、expert_idx、scale输入Shape要求分别为：(1, 7168)、(1, 8)、(256, 7168)
9506	- K=8	9567	+ x数据类型要求：bfloat16
9507	- expert_num=256	9568	+ 属性要求：active_expert_range=[0, 256]、 quant_mode=1、expert_tokens_num_type=2、expert_num=256
9508	- expert_end-expert_start<=32	9569	+ 进入大batch性能模板需要同时满足以下条件：
9509	- quant_mode=-1	9570	+ NUM_ROWS范围为[384, 8192]
9510	- row_idx_type=1	9571	+ K=8
9511	- expert_tokens_num_type=1	9572	+ expert_num=256
		9573	+ expert_end-expert_start<=32
		9574	+ quant_mode=-1
		9575	+ row_idx_type=1
		9576	+ expert_tokens_num_type=1
9512		9577
9513	支持的PyTorch版本	9578	支持的PyTorch版本
9514	PyTorch 2.6	9579	PyTorch 2.6

Mop_plugin/config/op_plugin_functions.yaml+8-2

	@@ -6006,10 +6006,10 @@ custom:
6006	dtype: input	6006	dtype: input
6007	out1:	6007	out1:
6008	size: '{input.size(0), group}'	6008	size: '{input.size(0), group}'
6009	- dtype: input	6009	+ dtype: npu_group_norm_silu_dst_type(input, weight, bias)
6010	out2:	6010	out2:
6011	size: '{input.size(0), group}'	6011	size: '{input.size(0), group}'
6012	- dtype: input	6012	+ dtype: npu_group_norm_silu_dst_type(input, weight, bias)
6013	exec: aclnnGroupNormSilu	6013	exec: aclnnGroupNormSilu
6014		6014
6015	- func: npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, Tensor[]? offset=None, Tensor[]? antiquant_scale=None, Tensor[]? antiquant_offset=None, Tensor[]? per_token_scale=None, Tensor? group_list=None, Tensor[]? activation_input=None, Tensor[]? activation_quant_scale=None, Tensor[]? activation_quant_offset=None, int? split_item=0, int? group_type=None, int? group_list_type=0, int? act_type=0, int[]? tuning_config=None, int? output_dtype=None, int? x_dtype=None, int? weight_dtype=None, int? scale_dtype=None, int? per_token_scale_dtype=None) -> Tensor[]	6015	- func: npu_grouped_matmul(Tensor[] x, Tensor[] weight, *, Tensor[]? bias=None, Tensor[]? scale=None, Tensor[]? offset=None, Tensor[]? antiquant_scale=None, Tensor[]? antiquant_offset=None, Tensor[]? per_token_scale=None, Tensor? group_list=None, Tensor[]? activation_input=None, Tensor[]? activation_quant_scale=None, Tensor[]? activation_quant_offset=None, int? split_item=0, int? group_type=None, int? group_list_type=0, int? act_type=0, int[]? tuning_config=None, int? output_dtype=None, int? x_dtype=None, int? weight_dtype=None, int? scale_dtype=None, int? per_token_scale_dtype=None) -> Tensor[]
	@@ -6197,6 +6197,10 @@ custom:
6197	op_api: all_version	6197	op_api: all_version
6198	exposed: all_version	6198	exposed: all_version
6199		6199
		6200	+ - func: npu_moe_gating_top_k_softmax_v2(Tensor x, *, int k=1, Tensor? finished=None, int? renorm=0, bool? output_softmax=False) -> (Tensor, Tensor, Tensor)
		6201	+ op_api: all_version
		6202	+ exposed: all_version
		6203	+
6200	- func: npu_moe_gating_top_k(Tensor x, int k, *, Tensor? bias=None, int k_group=1, int group_count=1, int group_select_mode=0, int renorm=0, int norm_type=0, bool out_flag=False, float routed_scaling_factor=1.0, float eps=1e-20) -> (Tensor, Tensor, Tensor)	6204	- func: npu_moe_gating_top_k(Tensor x, int k, *, Tensor? bias=None, int k_group=1, int group_count=1, int group_select_mode=0, int renorm=0, int norm_type=0, bool out_flag=False, float routed_scaling_factor=1.0, float eps=1e-20) -> (Tensor, Tensor, Tensor)
6201	op_api: all_version	6205	op_api: all_version
6202		6206
	@@ -6478,10 +6482,12 @@ custom:
6478	exposed: all_version	6482	exposed: all_version
6479		6483
6480	- func: npu_scaled_masked_softmax(Tensor x, Tensor mask, Scalar scale=1, bool fixed_triu_mask=False) -> Tensor	6484	- func: npu_scaled_masked_softmax(Tensor x, Tensor mask, Scalar scale=1, bool fixed_triu_mask=False) -> Tensor
		6485	+ op_api: all_version
6481	acl_op: all_version	6486	acl_op: all_version
6482	exposed: all_version	6487	exposed: all_version
6483		6488
6484	- func: npu_scaled_masked_softmax_backward(Tensor y_grad, Tensor y, Tensor mask, Scalar scale, bool fixed_triu_mask) -> Tensor	6489	- func: npu_scaled_masked_softmax_backward(Tensor y_grad, Tensor y, Tensor mask, Scalar scale, bool fixed_triu_mask) -> Tensor
		6490	+ op_api: all_version
6485	acl_op: all_version	6491	acl_op: all_version
6486		6492
6487	- func: npu_scatter(Tensor self, Tensor indices, Tensor updates, int dim) -> Tensor	6493	- func: npu_scatter(Tensor self, Tensor indices, Tensor updates, int dim) -> Tensor

Aop_plugin/ops/opapi/MoeGatingTopKSoftmaxV2KernelNpuOpApi.cpp+98-0

@@ -0,0 +1,98 @@
	1	+// Copyright (c) 2025 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 related link.
	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	+const int DIM_TWO = 2;
	19	+const int DIM_THREE = 3;
	20	+
	21	+namespace op_api {
	22	+ using npu_preparation = at_npu::native::OpPreparation;
	23	+ using tensor_list = std::tuple<at::Tensor, at::Tensor, at::Tensor>;
	24	+
	25	+ tensor_list npu_moe_gating_top_k_softmax_v2 (const at::Tensor &x, int64_t k, const c10::optional<at::Tensor> &finished_opt,
	26	+ const c10::optional<int64_t> renorm_opt, const c10::optional<bool> softmax_flag_opt)
	27	+ {
	28	+ // check x's shape
	29	+ TORCH_CHECK(x.dim() == DIM_TWO or x.dim() == DIM_THREE, "The x's shape should be 2D or 3D", OPS_ERROR(ErrCode::PARAM));
	30	+ // check x's datatype
	31	+ TORCH_CHECK(x.scalar_type() == at::kHalf \|\| x.scalar_type() == at::kFloat \|\| x.scalar_type() == at::kBFloat16,
	32	+ "float16, float32 or bfloat16 tensor expected but got a tensor with dtype: ",
	33	+ x.scalar_type(), OPS_ERROR(ErrCode::PARAM));
	34	+
	35	+ // check k's datatype
	36	+ auto x_size = x.sizes();
	37	+ TORCH_CHECK(k >= 0 and k <= x_size[x.dim() - 1], "The k's shape should be in [0, ", x_size[x.dim() - 1], "]", OPS_ERROR(ErrCode::PARAM));
	38	+
	39	+ // renorm optional
	40	+ int64_t renorm = c10::value_or_else(renorm_opt, [] {return 0;});
	41	+ TORCH_CHECK(renorm == 0 \|\| renorm == 1, "renorm must be 0 or 1, but got: ", renorm, OPS_ERROR(ErrCode::PARAM));
	42	+
	43	+ bool softmax_result_flag = c10::value_or_else(softmax_flag_opt, [] {return false; });
	44	+
	45	+ // finished Tensor
	46	+ const at::Tensor &finished = c10::value_or_else(finished_opt, [] {return at::Tensor(); });
	47	+ if (finished.defined()) {
	48	+ TORCH_CHECK(finished.scalar_type() == at::kBool, "bool tensor expected but got a tensor with dtype: ", finished.scalar_type(), OPS_ERROR(ErrCode::PARAM));
	49	+ auto finished_size = finished.sizes();
	50	+ TORCH_CHECK((x.dim() - 1) == finished.dim(), "x.dim() should be 1 more than finished.dim().", OPS_ERROR(ErrCode::PARAM));
	51	+ TORCH_CHECK(x_size[0] == finished_size[0], "Input rows should be same as finished rows.", OPS_ERROR(ErrCode::PARAM));
	52	+ if (x.dim() == DIM_THREE) {
	53	+ TORCH_CHECK(x_size[1] == finished_size[1], "Input rows should be same as finished rows.", OPS_ERROR(ErrCode::PARAM));
	54	+ }
	55	+ }
	56	+
	57	+ at::Tensor y;
	58	+ at::Tensor expert_idx;
	59	+ at::Tensor softmax_result; // Optional output
	60	+
	61	+ if (x.dim() == DIM_THREE) {
	62	+ y = npu_preparation::apply_tensor_without_format({x_size[0], x_size[1], k}, x.options());
	63	+ expert_idx = npu_preparation::apply_tensor_without_format(
	64	+ {x_size[0], x_size[1], k}, x.options().dtype(at::kInt));
	65	+ } else {
	66	+ y = npu_preparation::apply_tensor_without_format({x_size[0], k}, x.options());
	67	+ expert_idx = npu_preparation::apply_tensor_without_format(
	68	+ {x_size[0], k}, x.options().dtype(at::kInt));
	69	+ }
	70	+
	71	+ bool softmaxFlag = (renorm == 0) && softmax_result_flag;
	72	+ if (softmaxFlag) {
	73	+ if (x.dim() == DIM_THREE) {
	74	+ softmax_result = npu_preparation::apply_tensor_without_format(
	75	+ {x_size[0], x_size[1], x_size[2]}, x.options().dtype(at::kFloat));
	76	+ } else {
	77	+ softmax_result = npu_preparation::apply_tensor_without_format(
	78	+ {x_size[0], x_size[1]}, x.options().dtype(at::kFloat));
	79	+ }
	80	+ } else {
	81	+ softmax_result = npu_preparation::apply_tensor_without_format({0}, x.options().dtype(at::kFloat));
	82	+ }
	83	+
	84	+ if (k == 0) {
	85	+ return std::tie(y, expert_idx, softmax_result);
	86	+ }
	87	+
	88	+ for (int32_t i = 0; i < x.dim(); i++) {
	89	+ if (x_size[i] == 0) {
	90	+ return std::tie(y, expert_idx, softmax_result);
	91	+ }
	92	+ }
	93	+
	94	+ EXEC_NPU_CMD(aclnnMoeGatingTopKSoftmaxV2, x, finished, k, renorm, softmaxFlag, y, expert_idx, softmax_result);
	95	+
	96	+ return std::tie(y, expert_idx, softmax_result);
	97	+ }
	98	+}

Mop_plugin/ops/opapi/MoeInitRoutingV2KernelNpuOpApi.cpp+62-26

	@@ -17,25 +17,32 @@
17	#include "op_plugin/OpApiInterface.h"	17	#include "op_plugin/OpApiInterface.h"
18	#include "op_plugin/utils/op_api_common.h"	18	#include "op_plugin/utils/op_api_common.h"
19		19
20	-// expert_tokens_num_types	20	+namespace {
21	-// 'expert_tokens_count_or_cumsum' is suggested to be reverted into 'expert_tokens_count'	21	+constexpr int64_t DIM_X = 2;
22	-static const int64_t CUMSUM = 0;	22	+constexpr int64_t DIM_EXPERT_IDX = 2;
23	-static const int64_t COUNT = 1;	23	+constexpr int64_t LENGTH_ACTIVE_EXPERT_RANGE = 2;
24	-static const int64_t KEY_VALUE = 2;	24	+constexpr int64_t EXPERT_TOKENS_COUNT = 1;
		25	+constexpr int64_t EXPERT_TOKENS_KEY_VALUE = 2;
		26	+constexpr int64_t QUANT_MODE_UNQUANT = -1;
		27	+constexpr int64_t QUANT_MODE_DYNAMIC_QUANT = 1;
25		28
26	-static const int64_t EXPERT_NUM_V2 = 128;	29	+constexpr int64_t EXPERT_NUM_V2 = 128;
27	-static const int64_t EXPERT_NUM_MIN_V2 = 0;	30	+constexpr int64_t EXPERT_NUM_MIN_V2 = 0;
28	-static const int64_t EXPERT_NUM_MAX_V2 = 128;	31	+constexpr int64_t EXPERT_NUM_MAX_V2 = 128;
29	-static const int64_t HIDDEN_DIM_VAL_V2 = 2048;	32	+constexpr int64_t HIDDEN_DIM_VAL_V2 = 2048;
		33	+
		34	+}; // namespace
30		35
31	namespace op_api {	36	namespace op_api {
32	using npu_preparation = at_npu::native::OpPreparation;	37	using npu_preparation = at_npu::native::OpPreparation;
33	using npu_utils = at_npu::native::NpuUtils;	38	using npu_utils = at_npu::native::NpuUtils;
34	using tensor_list = std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>;	39	using tensor_list = std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>;
35		40
36	-static bool CheckV2Case(int hidden_dim, int64_t expert_num, at::IntArrayRef active_expert_range, int64_t expert_tokens_num_type, int64_t quant_mode)	41	+static bool CheckV2Case(int hidden_dim, int64_t expert_num, at::IntArrayRef active_expert_range,
		42	+ int64_t expert_tokens_num_type, int64_t quant_mode)
37	{	43	{
38	- if (expert_num == EXPERT_NUM_V2 && active_expert_range[0] == EXPERT_NUM_MIN_V2 && active_expert_range[1] == EXPERT_NUM_MAX_V2 && hidden_dim == HIDDEN_DIM_VAL_V2) {	44	+ if (expert_num == EXPERT_NUM_V2 && active_expert_range[0] == EXPERT_NUM_MIN_V2 &&
		45	+ active_expert_range[1] == EXPERT_NUM_MAX_V2 && hidden_dim == HIDDEN_DIM_VAL_V2) {
39	if (quant_mode == -1 && expert_tokens_num_type == 1) {	46	if (quant_mode == -1 && expert_tokens_num_type == 1) {
40	return true;	47	return true;
41	}	48	}
	@@ -48,6 +55,33 @@ tensor_list npu_moe_init_routing_v2(const at::Tensor &x, const at::Tensor &exper
48	int64_t expert_capacity, int64_t expert_num, int64_t drop_pad_mode, int64_t expert_tokens_num_type,	55	int64_t expert_capacity, int64_t expert_num, int64_t drop_pad_mode, int64_t expert_tokens_num_type,
49	bool expert_tokens_num_flag, int64_t quant_mode, at::IntArrayRef active_expert_range, int64_t row_idx_type)	56	bool expert_tokens_num_flag, int64_t quant_mode, at::IntArrayRef active_expert_range, int64_t row_idx_type)
50	{	57	{
		58	+ int64_t x_dim = x.dim();
		59	+ TORCH_CHECK(x_dim == DIM_X,
		60	+ "The x should be ",
		61	+ DIM_X,
		62	+ "-Dimension, current is ",
		63	+ x_dim,
		64	+ "-Dimension.",
		65	+ OPS_ERROR(ErrCode::PARAM));
		66	+
		67	+ int64_t expert_idx_dim = expert_idx.dim();
		68	+ TORCH_CHECK(expert_idx_dim == DIM_EXPERT_IDX,
		69	+ "The expert_idx should be ",
		70	+ DIM_EXPERT_IDX,
		71	+ "-Dimension, current is ",
		72	+ expert_idx,
		73	+ "-Dimension.",
		74	+ OPS_ERROR(ErrCode::PARAM));
		75	+
		76	+ int64_t active_expert_range_length = active_expert_range.size();
		77	+ TORCH_CHECK(active_expert_range_length == LENGTH_ACTIVE_EXPERT_RANGE,
		78	+ "The length of list active_expert_range should be ",
		79	+ LENGTH_ACTIVE_EXPERT_RANGE,
		80	+ ", current is ",
		81	+ active_expert_range_length,
		82	+ ".",
		83	+ OPS_ERROR(ErrCode::PARAM));
		84	+
51	int expert_length = active_expert_range[1] - active_expert_range[0];	85	int expert_length = active_expert_range[1] - active_expert_range[0];
52	auto x_size = x.sizes();	86	auto x_size = x.sizes();
53	auto expert_idx_size = expert_idx.sizes();	87	auto expert_idx_size = expert_idx.sizes();
	@@ -62,30 +96,31 @@ tensor_list npu_moe_init_routing_v2(const at::Tensor &x, const at::Tensor &exper
62	bool using_v2 = CheckV2Case(h, expert_num, active_expert_range, expert_tokens_num_type, quant_mode);	96	bool using_v2 = CheckV2Case(h, expert_num, active_expert_range, expert_tokens_num_type, quant_mode);
63		97
64	at::Tensor expanded_x;	98	at::Tensor expanded_x;
65	- if (quant_mode == -1) {	99	+ if (quant_mode == QUANT_MODE_DYNAMIC_QUANT) {
66	- expanded_x = npu_preparation::apply_tensor_without_format(x, {bs * k, h});
67	- } else {
68	expanded_x = npu_preparation::apply_tensor_without_format({bs * k, h}, x.options().dtype(at::kChar));	100	expanded_x = npu_preparation::apply_tensor_without_format({bs * k, h}, x.options().dtype(at::kChar));
		101	+ } else { // quant_mode == QUANT_MODE_UNQUANT
		102	+ expanded_x = npu_preparation::apply_tensor_without_format(x, {bs * k, h});
69	}	103	}
70	at::Tensor expanded_row_idx = npu_preparation::apply_tensor_without_format(expert_idx, {bs * k});	104	at::Tensor expanded_row_idx = npu_preparation::apply_tensor_without_format(expert_idx, {bs * k});
71	at::Tensor expert_tokens_count_or_cumsum;	105	at::Tensor expert_tokens_count_or_cumsum;
72	- if (expert_tokens_num_type >= CUMSUM && expert_tokens_num_type <= COUNT) {	106	+ if (expert_tokens_num_type == EXPERT_TOKENS_KEY_VALUE) {
73	- // expert_tokens_count_or_cumsum in [end-start, ]
74	- if (using_v2) {
75	- expert_tokens_count_or_cumsum =
76	- npu_preparation::apply_tensor_without_format({expert_num}, x.options().dtype(at::kInt));
77	- } else {
78	- expert_tokens_count_or_cumsum =
79	- npu_preparation::apply_tensor_without_format({expert_length}, x.options().dtype(at::kLong));
80	- }
81	- } else if (expert_tokens_num_type == KEY_VALUE) {
82	// key_value in [2, end-start]	107	// key_value in [2, end-start]
83	expert_tokens_count_or_cumsum =	108	expert_tokens_count_or_cumsum =
84	npu_preparation::apply_tensor_without_format({expert_num, 2}, x.options().dtype(at::kLong));	109	npu_preparation::apply_tensor_without_format({expert_num, 2}, x.options().dtype(at::kLong));
		110	+ } else { // expert_tokens_num_type == EXPERT_TOKENS_COUNT
		111	+ // expert_tokens_count_or_cumsum in [end-start, ]
		112	+ if (using_v2) {
		113	+ expert_tokens_count_or_cumsum =
		114	+ npu_preparation::apply_tensor_without_format({expert_num}, x.options().dtype(at::kInt));
		115	+ } else {
		116	+ expert_tokens_count_or_cumsum =
		117	+ npu_preparation::apply_tensor_without_format({expert_length}, x.options().dtype(at::kLong));
		118	+ }
85	}	119	}
86		120
87	if (using_v2) {	121	if (using_v2) {
88	- at::Tensor expert_tokens_before_capacity = npu_preparation::apply_tensor_without_format({expert_num}, x.options().dtype(at::kInt));	122	+ at::Tensor expert_tokens_before_capacity =
		123	+ npu_preparation::apply_tensor_without_format({expert_num}, x.options().dtype(at::kInt));
89	expert_capacity = 0;	124	expert_capacity = 0;
90	drop_pad_mode = 0;	125	drop_pad_mode = 0;
91	int64_t expert_tokens_count_or_cumsum_flag = 2;	126	int64_t expert_tokens_count_or_cumsum_flag = 2;
	@@ -105,7 +140,8 @@ tensor_list npu_moe_init_routing_v2(const at::Tensor &x, const at::Tensor &exper
105	expert_tokens_before_capacity);	140	expert_tokens_before_capacity);
106	return std::tie(expanded_x, expanded_row_idx, expert_tokens_count_or_cumsum, expert_tokens_before_capacity);	141	return std::tie(expanded_x, expanded_row_idx, expert_tokens_count_or_cumsum, expert_tokens_before_capacity);
107	} else {	142	} else {
108	- at::Tensor expanded_scale = npu_preparation::apply_tensor_without_format({bs * k}, x.options().dtype(at::kFloat));	143	+ at::Tensor expanded_scale =
		144	+ npu_preparation::apply_tensor_without_format({bs * k}, x.options().dtype(at::kFloat));
109	EXEC_NPU_CMD(aclnnMoeInitRoutingV3,	145	EXEC_NPU_CMD(aclnnMoeInitRoutingV3,
110	x,	146	x,
111	expert_idx,	147	expert_idx,

Aop_plugin/ops/opapi/ScaledMaskedSoftmaxKernelNpuOpApi.cpp+67-0

@@ -0,0 +1,67 @@
	1	+// Copyright (c) 2025 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 related link.
	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 <ATen/AccumulateType.h>
	15	+#include "op_plugin/AclOpsInterface.h"
	16	+#include "op_plugin/OpApiInterface.h"
	17	+#include "op_plugin/utils/op_api_common.h"
	18	+
	19	+namespace op_api {
	20	+using npu_preparation = at_npu::native::OpPreparation;
	21	+
	22	+at::Tensor npu_scaled_masked_softmax(const at::Tensor& x, const at::Tensor& mask, const at::Scalar& scale,
	23	+ bool fixed_triu_mask)
	24	+{
	25	+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend910_95) {
	26	+ return acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask);
	27	+ }
	28	+ DO_COMPATIBILITY(aclnnScaledMaskedSoftmax, acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask));
	29	+ double scale_value = 1.0;
	30	+ if (scale.isFloatingPoint()) {
	31	+ scale_value = scale.to<double>();
	32	+ } else if (scale.isIntegral(true)) {
	33	+ scale_value = static_cast<double>(scale.to<int64_t>());
	34	+ } else {
	35	+ TORCH_CHECK(false, "scaled_masked_softmax expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));
	36	+ }
	37	+
	38	+ at::Tensor result = npu_preparation::apply_tensor_without_format(x.sizes(), x.options());
	39	+ EXEC_NPU_CMD(aclnnScaledMaskedSoftmax, x, mask, scale_value, fixed_triu_mask, result);
	40	+ return result;
	41	+}
	42	+
	43	+at::Tensor npu_scaled_masked_softmax_backward(
	44	+ const at::Tensor& y_grad,
	45	+ const at::Tensor& y,
	46	+ const at::Tensor& mask,
	47	+ const at::Scalar& scale,
	48	+ bool fixed_triu_mask)
	49	+{
	50	+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend910_95) {
	51	+ return acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask);
	52	+ }
	53	+ DO_COMPATIBILITY(aclnnScaledMaskedSoftmaxBackward, acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask));
	54	+ double scale_value = 1.0;
	55	+ if (scale.isFloatingPoint()) {
	56	+ scale_value = scale.to<double>();
	57	+ } else if (scale.isIntegral(true)) {
	58	+ scale_value = static_cast<double>(scale.to<int64_t>());
	59	+ } else {
	60	+ TORCH_CHECK(false, "scaled_masked_softmax_backward expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));
	61	+ }
	62	+
	63	+ at::Tensor result = npu_preparation::apply_tensor_without_format(y_grad.sizes(), y_grad.options());
	64	+ EXEC_NPU_CMD(aclnnScaledMaskedSoftmaxBackward, y_grad, y, mask, scale_value, fixed_triu_mask, result);
	65	+ return result;
	66	+}
	67	+}

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

	@@ -907,6 +907,33 @@ def npu_moe_gating_top_k_softmax_meta(x, finished=None, k=1):
907	x.new_empty(tuple(row_idx_dim_list), dtype=torch.int32))	907	x.new_empty(tuple(row_idx_dim_list), dtype=torch.int32))
908		908
909		909
		910	+@impl(m, "npu_moe_gating_top_k_softmax_v2")
		911	+def npu_moe_gating_top_k_softmax_v2_meta(x, *, k=1, finished=None, renorm=0, output_softmax=False):
		912	+ x_dim = x.dim()
		913	+ torch._check(
		914	+ x_dim == 2 or x_dim == 3,
		915	+ lambda: "the x shape support only 2d and 3d)" + ops_error(ErrCode.VALUE),
		916	+ )
		917	+ if x_dim == 3:
		918	+ y_dim_list = [x.size(0), x.size(1), k]
		919	+ expert_idx_dim_list = [x.size(0), x.size(1), k]
		920	+ else:
		921	+ y_dim_list = [x.size(0), k]
		922	+ expert_idx_dim_list = [x.size(0), k]
		923	+
		924	+ if renorm == 0 and output_softmax:
		925	+ if x.dim == 3:
		926	+ softmax_result_dim_list = [x.size(0), x.size(1), x.size(2)]
		927	+ else:
		928	+ softmax_result_dim_list = [x.size(0), x.size(1)]
		929	+ else:
		930	+ softmax_result_dim_list = [0, ]
		931	+
		932	+ return (x.new_empty(tuple(y_dim_list), dtype=x.dtype),
		933	+ x.new_empty(tuple(expert_idx_dim_list), dtype=torch.int32),
		934	+ x.new_empty(tuple(softmax_result_dim_list), dtype=torch.float32))
		935	+
		936	+
910	@impl(m, "npu_moe_gating_top_k")	937	@impl(m, "npu_moe_gating_top_k")
911	def npu_moe_gating_top_k_meta(x, k=1, bias=None, k_group=1, group_count=1, group_select_mode=0, renorm=0, norm_type=0, out_flag=False, routed_scaling_factor=1.0, eps=1e-20):	938	def npu_moe_gating_top_k_meta(x, k=1, bias=None, k_group=1, group_count=1, group_select_mode=0, renorm=0, norm_type=0, out_flag=False, routed_scaling_factor=1.0, eps=1e-20):
912	x_dim = x.dim()	939	x_dim = x.dim()

Mop_plugin/utils/KernelNpuOutputDtype.cpp+17-0

	@@ -44,6 +44,23 @@ at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle)
44	return high_type;	44	return high_type;
45	}	45	}
46		46
		47	+at::ScalarType npu_group_norm_silu_dst_type(const at::Tensor& input, const c10::optional<at::Tensor>& weight,
		48	+ const c10::optional<at::Tensor>& bias)
		49	+{
		50	+ at::native::ResultTypeState state = {};
		51	+ state = at::native::update_result_type_state(input, state);
		52	+
		53	+ if (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910_95) {
		54	+ if (weight.has_value()) {
		55	+ state = at::native::update_result_type_state(weight.value(), state);
		56	+ } else if (bias.has_value()) {
		57	+ state = at::native::update_result_type_state(bias.value(), state);
		58	+ }
		59	+ }
		60	+
		61	+ return at::native::result_type(state);
		62	+}
		63	+
47	at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype)	64	at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype)
48	{	65	{
49	at::ScalarType dst_type = c10::value_or_else(dst_dtype, [] {return at::ScalarType::Char;});	66	at::ScalarType dst_type = c10::value_or_else(dst_dtype, [] {return at::ScalarType::Char;});

Mop_plugin/utils/KernelNpuOutputDtype.h+1-0

	@@ -24,6 +24,7 @@ namespace op_infer {
24		24
25	OP_PLUGIN_HIDDEN at::ScalarType angle_out_dtype(const at::Tensor& self);	25	OP_PLUGIN_HIDDEN at::ScalarType angle_out_dtype(const at::Tensor& self);
26	OP_PLUGIN_HIDDEN at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle);	26	OP_PLUGIN_HIDDEN at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle);
		27	+OP_PLUGIN_HIDDEN at::ScalarType npu_group_norm_silu_dst_type(const at::Tensor& input, const c10::optional<at::Tensor>& weight, const c10::optional<at::Tensor>& bias);
27	OP_PLUGIN_HIDDEN at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype);	28	OP_PLUGIN_HIDDEN at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype);
28	OP_PLUGIN_HIDDEN at::ScalarType clamp_out_dtype(const at::Tensor& self, const c10::optional<at::Tensor>& min, const c10::optional<at::Tensor>& max);	29	OP_PLUGIN_HIDDEN at::ScalarType clamp_out_dtype(const at::Tensor& self, const c10::optional<at::Tensor>& min, const c10::optional<at::Tensor>& max);
29	OP_PLUGIN_HIDDEN at::ScalarType clamp_scalar_out_dtype(const at::Tensor& self, const c10::optional<at::Scalar>& min, const c10::optional<at::Scalar>& max);	30	OP_PLUGIN_HIDDEN at::ScalarType clamp_scalar_out_dtype(const at::Tensor& self, const c10::optional<at::Scalar>& min, const c10::optional<at::Scalar>& max);

Mtest/core_tests/torch_npu_OpApi_schema_all.json+9-0

	@@ -17,6 +17,12 @@
17	"op_api: torch_npu.npu_moe_gating_top_k_softmax(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {	17	"op_api: torch_npu.npu_moe_gating_top_k_softmax(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
18	"version": ["v2.8", "newest"]	18	"version": ["v2.8", "newest"]
19	},	19	},
		20	+ "op_api: torch_npu.npu_moe_gating_top_k_softmax_v2(args, *kwargs)": {
		21	+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
		22	+ },
		23	+ "op_api: torch_npu.npu_moe_gating_top_k_softmax_v2(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
		24	+ "version": ["v2.8", "newest"]
		25	+ },
20	"op_api: torch_npu.npu_geglu(args, *kwargs)": {	26	"op_api: torch_npu.npu_geglu(args, *kwargs)": {
21	"version": ["v2.1"]	27	"version": ["v2.1"]
22	},	28	},
	@@ -1277,6 +1283,9 @@
1277	"func: npu_moe_gating_top_k_softmax(Tensor x, Tensor? finished=None, int k=1) -> (Tensor, Tensor, Tensor)": {	1283	"func: npu_moe_gating_top_k_softmax(Tensor x, Tensor? finished=None, int k=1) -> (Tensor, Tensor, Tensor)": {
1278	"version": ["all_version"]	1284	"version": ["all_version"]
1279	},	1285	},
		1286	+ "func: npu_moe_gating_top_k_softmax_v2(Tensor x, *, int k=1, Tensor? finished=None, int? renorm=0, bool? output_softmax=False) -> (Tensor, Tensor, Tensor)": {
		1287	+ "version": ["all_version"]
		1288	+ },
1280	"func: npu_ffn(Tensor x, Tensor weight1, Tensor weight2, str activation, *, int[]? expert_tokens=None, int[]? expert_tokens_index=None, Tensor? bias1=None, Tensor? bias2=None, Tensor? scale=None, Tensor? offset=None, Tensor? deq_scale1=None, Tensor? deq_scale2=None, Tensor? antiquant_scale1=None, Tensor? antiquant_scale2=None, Tensor? antiquant_offset1=None, Tensor? antiquant_offset2=None, int? inner_precise=None, ScalarType? output_dtype=None) -> Tensor": {	1289	"func: npu_ffn(Tensor x, Tensor weight1, Tensor weight2, str activation, *, int[]? expert_tokens=None, int[]? expert_tokens_index=None, Tensor? bias1=None, Tensor? bias2=None, Tensor? scale=None, Tensor? offset=None, Tensor? deq_scale1=None, Tensor? deq_scale2=None, Tensor? antiquant_scale1=None, Tensor? antiquant_scale2=None, Tensor? antiquant_offset1=None, Tensor? antiquant_offset2=None, int? inner_precise=None, ScalarType? output_dtype=None) -> Tensor": {
1281	"version": ["all_version"]	1290	"version": ["all_version"]
1282	},	1291	},

Mtest/test_custom_ops/test_group_norm_silu.py+13-0

	@@ -32,6 +32,19 @@ class TestGroupNormSilu(TestCase):
32	custom_output = self.custom_op_exec(x, gama, beta, group, eps)	32	custom_output = self.custom_op_exec(x, gama, beta, group, eps)
33	self.assertRtolEqual(supported_output, custom_output)	33	self.assertRtolEqual(supported_output, custom_output)
34		34
		35	+ @SupportedDevices(['Ascend910_95'])
		36	+ def test_npu_910_95_(self, device="npu"):
		37	+ x = torch.randn(24, 320, 48, 48, dtype=torch.float16).npu()
		38	+ gama = torch.randn(320, dtype=torch.float32).npu()
		39	+ beta = torch.randn(320, dtype=torch.float32).npu()
		40	+
		41	+ group = 32
		42	+ eps = 0.0001
		43	+
		44	+ supported_output = self.supported_op_exec(x, gama, beta, group, eps)
		45	+ custom_output = self.custom_op_exec(x, gama, beta, group, eps)
		46	+ self.assertRtolEqual(supported_output, custom_output)
		47	+
35		48
36	if __name__ == "__main__":	49	if __name__ == "__main__":
37	run_tests()	50	run_tests()

Atest/test_custom_ops/test_npu_moe_gating_top_k_softmax_v2.py+99-0

@@ -0,0 +1,99 @@
	1	+# Copyright (c) 2020, Huawei Technologies.All rights reserved.
	2	+#
	3	+# Licensed under the BSD 3-Clause License (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at relate links.
	6	+#
	7	+# Unless required by applicable law or agreed to in writing, software
	8	+# distributed under the License is distributed on an "AS IS" BASIS,
	9	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	10	+# See the License for the specific language governing permissions and
	11	+# limitations under the License.
	12	+
	13	+import itertools
	14	+import numpy as np
	15	+import torch
	16	+import torch_npu
	17	+from torch_npu.testing.testcase import TestCase, run_tests
	18	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
	19	+
	20	+
	21	+class TestNpuMoeGatingTopKSoftmaxV2(TestCase):
	22	+
	23	+ def softmax_func(self, x, axis=None):
	24	+ m = np.max(x, axis=axis, keepdims=True)
	25	+ e = np.exp(x - m)
	26	+ return e / np.sum(e, axis=axis, keepdims=True)
	27	+
	28	+ def cpu_op_exec(self, x, k, finished, renorm, output_softmax):
	29	+ num_expert = x.shape[-1]
	30	+ leading_shape = x.shape[:-1]
	31	+ x_2d = x
	32	+
	33	+ if x_2d.ndim == 3:
	34	+ x_2d = x_2d.reshape(-1, num_expert)
	35	+ if finished is not None:
	36	+ finished = finished.flatten()
	37	+ else:
	38	+ if finished is not None and finished.ndim > 1:
	39	+ finished = finished.reshape(-1)
	40	+
	41	+ softmax_full_f32 = self.softmax_func(x_2d.astype(np.float32), axis=-1)
	42	+
	43	+ if renorm == 1:
	44	+ indices = np.argsort(-x_2d, axis=-1, kind='stable')[:, :k]
	45	+ values = np.take_along_axis(x_2d, indices, axis=-1)
	46	+ out = self.softmax_func(values.astype(np.float32), axis=-1).astype(x.dtype)
	47	+ else:
	48	+ indices = np.argsort(-softmax_full_f32, axis=-1, kind='stable')[:, :k]
	49	+ out = np.take_along_axis(softmax_full_f32, indices, axis=-1).astype(x.dtype)
	50	+
	51	+ indices = indices.astype(np.int32)
	52	+ if finished is not None:
	53	+ finished_expanded = np.tile(finished.reshape(-1, 1), (1, k))
	54	+ indices = np.where(finished_expanded, num_expert, indices)
	55	+
	56	+ out = out.reshape(*leading_shape, k)
	57	+ indices = indices.reshape(*leading_shape, k)
	58	+
	59	+ if renorm == 0 and output_softmax:
	60	+ softmax_full = softmax_full_f32.reshape(*leading_shape, num_expert)
	61	+ return out, indices, softmax_full
	62	+ else:
	63	+ empty_softmax = np.array([], dtype=softmax_full_f32.dtype)
	64	+ return out, indices, empty_softmax
	65	+
	66	+ def npu_op_exec(self, x, k, finished, renorm, output_softmax):
	67	+ y, indices, softmax_result = torch_npu.npu_moe_gating_top_k_softmax_v2(
	68	+ x=x, k=k, finished=finished, renorm=renorm, output_softmax=output_softmax)
	69	+ return y, indices, softmax_result
	70	+
	71	+ @SupportedDevices(['Ascend910B'])
	72	+ def test_npu_noe_gating_top_k_softmax_v2(self, device="npu"):
	73	+ n_list = [10, 430, 520]
	74	+ k_list = [2, 4, 5, 9]
	75	+ col_list = [200, 1256, 5120]
	76	+ flag_list = [True, False]
	77	+ renorm_list = [0, 1]
	78	+ dtype_list = [np.float16, np.float32]
	79	+ for n, k, col, flag, renorm, dtype in itertools.product(n_list, k_list, col_list, flag_list, renorm_list, dtype_list):
	80	+ x = np.random.uniform(-1, 1, size=(n, col)).astype(dtype)
	81	+ finished = np.random.uniform(-1, 1, size=(n,)).astype(bool)
	82	+ x_npu = torch.from_numpy(x).npu()
	83	+ finished_npu = torch.from_numpy(finished).npu()
	84	+
	85	+ try:
	86	+ y, expert_idx, output_softmax = self.cpu_op_exec(x, k, finished, renorm, flag)
	87	+ y_npu, expert_idx_npu, output_softmax_npu = self.npu_op_exec(x_npu, k, finished_npu, renorm, flag)
	88	+ self.assertRtolEqual(y, y_npu.cpu().numpy())
	89	+ self.assertRtolEqual(expert_idx, expert_idx_npu.cpu().numpy())
	90	+ self.assertRtolEqual(output_softmax, output_softmax_npu.cpu().numpy())
	91	+
	92	+ except Exception as e:
	93	+ raise AssertionError(f"Task failed unecpectedly: {e}") from e
	94	+
	95	+
	96	+
	97	+
	98	+if __name__ == "__main__":
	99	+ run_tests()

Mtest/test_custom_ops/test_npu_moe_init_routing_v2.py+2-0

	@@ -1,3 +1,4 @@
		1	+import unittest
1	import itertools	2	import itertools
2	from os import scandir	3	from os import scandir
3	import numpy as np	4	import numpy as np
	@@ -231,6 +232,7 @@ class TestNpuMoeInitRoutingV2(TestCase):
231	return	232	return
232	self.assertRtolEqual(expanded_scale, local_expanded_scale_npu.numpy())	233	self.assertRtolEqual(expanded_scale, local_expanded_scale_npu.numpy())
233		234
		235	+ @unittest.skip("Skipping test_npu_moe_init_routing_static_quant due to unsupported quant_mode=0 in current CANN version")
234	@SupportedDevices(['Ascend910B'])	236	@SupportedDevices(['Ascend910B'])
235	def test_npu_moe_init_routing_static_quant(self):	237	def test_npu_moe_init_routing_static_quant(self):
236	bs_list = [4]	238	bs_list = [4]