add dsa-op-plugin-AtomGit

已合并

add dsa #3529

daifu1234567创建于 2025年11月19日从daifu1234567/op-plugin： 2511dsa 合入到 Ascend/op-plugin： master

add dsa #3529

已合并

daifu1234567创建于 2025年11月19日

从daifu1234567/op-plugin 2511dsa合入到Ascend/op-plugin master

共 10 个文件变更+460-22

Mcodegen/templates/_op_plugin_docs.py+144-0

	@@ -12529,5 +12529,149 @@ mode = "none"
12529	input = torch.rand(shape, dtype=torch.float16).npu()	12529	input = torch.rand(shape, dtype=torch.float16).npu()
12530	output = torch_npu.npu_gelu_mul(input, approximate=mode)	12530	output = torch_npu.npu_gelu_mul(input, approximate=mode)
12531		12531
		12532	+"""
		12533	+)
		12534	+
		12535	+
		12536	+_add_torch_npu_docstr(
		12537	+ "npu_sparse_lightning_indexer_grad_kl_loss",
		12538	+ """
		12539	+接口原型:
		12540	+npu_sparse_lightning_indexer_grad_kl_loss(Tensor query, Tensor key, Tensor query_index, Tensor key_index, Tensor weights, Tensor sparse_indices, Tensor softmax_max, Tensor softmax_sum, float scale_value=1, *, Tensor? query_rope=None, Tensor? key_rope=None, SymInt[]? actual_seq_qlen=None, SymInt[]? actual_seq_klen=None, str? layout='BSND', int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor, Tensor)
		12541	+
		12542	+功能描述:
		12543	+该接口实现了npu_lightning_indexer的反向功能，再额外融合了Loss的计算功能。npu_lightning_indexer将Attention的query和Attention的key之间的最高内在联系的top_k个筛选出来，存放在sparse_indices中，从而减少长序列场景下Attention的计算量，加速长序列的网络的推理和训练的性能。
		12544	+
		12545	+参数说明:
		12546	+query（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1, D)、(T1, N1, D)。
		12547	+key（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S2, N2, D)、(T2, N2, D)。
		12548	+query_index（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1index, D)、(T1, N1index, D)。
		12549	+key_index（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1index, D)、(T1, N1index, D)。
		12550	+weights（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1)、(T1, N1)。
		12551	+sparse_indices（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, topK)、(T1, topK)。
		12552	+softmax_max（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1, G)、(T1, N1, G)。
		12553	+softmax_sum（Tensor）：必选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1, G)、(T1, N1, G)。
		12554	+scale_value（float）：必选参数，表示缩放系数，数据类型支持FLOAT。
		12555	+query_rope（Tensor）：可选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S1, N1, Dr)、(T1, N1, Dr)。
		12556	+key_rope（Tensor）：可选参数，数据格式支持ND，数据类型支持BFLOAT16、FLOAT16。支持输入shape(B, S2, N2, Dr)、(T2, N2, Dr)。
		12557	+actual_seq_qlen（int[]）：可选参数，int类型数组，TND场景时需传入此参数。表示query每个S的累加和长度，数据类型支持INT64，数据格式支持ND，默认值为None。
		12558	+actual_seq_klen（int[]）：可选参数，int类型数组，TND场景时需传入此参数。表示key每个S的累加和长度，数据类型支持INT64，数据格式支持ND，默认值为None。
		12559	+layout（str）：可选参数，用于标识输入query的数据排布格式，数据类型支持str。当前支持BSND、TND，默认值为"BSND"。
		12560	+sparse_mode（int）：可选参数，表示sparse的模式，数据类型支持INT32.默认值为3。
		12561	+pre_tokens（int）：必选参数，数据类型支持INT64，默认值9223372036854775807。
		12562	+next_tokens（int）：必选参数，数据类型支持INT64，默认值9223372036854775807。
		12563	+
		12564	+输出说明:
		12565	+d_query_index（Tensor）：表示query_index的梯度，数据类型支持BFLOAT16、FLOAT16。
		12566	+d_key_index（Tensor）：表示key_index的梯度，数据类型支持BFLOAT16、FLOAT16。
		12567	+d_weights（Tensor）：表示weights的梯度，数据类型支持BFLOAT16、FLOAT16。
		12568	+loss（Tensor）：表示网络正向输出和golden值的差异，数据类型支持FLOAT。
		12569	+
		12570	+支持版本:
		12571	+PyTorch 2.1
		12572	+PyTorch 2.5及更高版本
		12573	+
		12574	+支持的型号:
		12575	+Atlas A2训练系列产品
		12576	+Atlas A3训练系列产品
		12577	+
		12578	+调用示例:
		12579	+import torch
		12580	+import torch_npu
		12581	+
		12582	+def gen_inputs(seqlens_list_array, seqlens_list_kv_array, isTnd):
		12583	+ B = 1
		12584	+ NQuery = 64
		12585	+ NQueryIndex = 64
		12586	+ N2 = 1
		12587	+ S1 = 128
		12588	+ S2 = 128
		12589	+ topK = 2048
		12590	+ D = 512
		12591	+ DIndex = 128
		12592	+ DR = 64
		12593	+ output_dtype = torch.float16
		12594	+ q = torch.randn(B, S1, NQuery, D, dtype=output_dtype, device=torch.device('npu'))
		12595	+ k = torch.randn(B, S2, N2, D, dtype=output_dtype, device=torch.device('npu'))
		12596	+
		12597	+ q_index = torch.randn(B, S1, NQueryIndex, DIndex, dtype=output_dtype, device=torch.device('npu'))
		12598	+ k_index = torch.randn(B, S2, N2, DIndex, dtype=output_dtype, device=torch.device('npu'))
		12599	+ if DR != 0:
		12600	+ q_rope = torch.randn(B, S1, NQuery, DR, dtype=output_dtype, device=torch.device('npu'))
		12601	+ k_rope = torch.randn(B, S2, N2, DR, dtype=output_dtype, device=torch.device('npu'))
		12602	+ else:
		12603	+ q_rope = None
		12604	+ k_rope = None
		12605	+ weights = torch.randn(B, S1, NQueryIndex, dtype=output_dtype, device=torch.device('npu'))
		12606	+ a = -0.05 # 最小值
		12607	+ b = 0.05 # 最大值
		12608	+ kk = 3.0 # 控制分布范围（3σ 覆盖绝大多数值）
		12609	+ scale = (b - a) / (2 * kk)
		12610	+ shift = (a + b) / 2
		12611	+ weights = weights * scale + shift
		12612	+ if isTnd:
		12613	+ sparse_indices = torch.zeros(S1, N2, topK).to(torch.int32).npu()
		12614	+ tIdx = 0
		12615	+ for bIdx in range(B):
		12616	+ for s1Idx in range(seqlens_list_array[bIdx]):
		12617	+ s2RealSize = (int)((seqlens_list_kv_array[bIdx] - seqlens_list_array[bIdx]) + s1Idx + 1)
		12618	+ if s2RealSize <= 0:
		12619	+ s2RealSize = seqlens_list_kv_array[bIdx]
		12620	+
		12621	+ if s2RealSize > topK:
		12622	+ s2RealLen = topK
		12623	+ else:
		12624	+ s2RealLen = s2RealSize
		12625	+ #处理S2无效行场景，把对应的sparse indices置为-1
		12626	+ sparse_indices[tIdx, :, 0 : s2RealLen] = (torch.randint(0, s2RealSize, (s2RealLen,)).to(torch.int32)).npu()
		12627	+ sparse_indices[tIdx, :, s2RealLen : topK] = -1
		12628	+ tIdx = tIdx + 1
		12629	+ q_tnd = q.squeeze(dim=0)
		12630	+ k_tnd = k.squeeze(dim=0)
		12631	+ q_index_tnd = q_index.squeeze(dim=0)
		12632	+ k_index_tnd = k_index.squeeze(dim=0)
		12633	+ if q_rope is not None:
		12634	+ q_rope_tnd = q_rope.squeeze(dim=0)
		12635	+ k_rope_tnd = k_rope.squeeze(dim=0)
		12636	+ else :
		12637	+ q_rope_tnd = None
		12638	+ k_rope_tnd = None
		12639	+ weights_tnd = weights.squeeze(dim=0)
		12640	+
		12641	+ softmax_max = torch.randn(N2, S1, NQueryIndex, dtype=torch.float, device=torch.device('npu'))
		12642	+ softmax_sum = torch.randn(N2, S1, NQueryIndex, dtype=torch.float, device=torch.device('npu'))
		12643	+ return q_tnd, k_tnd, q_index_tnd, k_index_tnd, q_rope_tnd, k_rope_tnd, weights_tnd, sparse_indices, softmax_max, softmax_sum
		12644	+ else :
		12645	+ sparse_indices = torch.zeros(B, S1, N2, topK).to(torch.int32).npu()
		12646	+ for s1Idx in range(S1):
		12647	+ s2RealSize = (int)(S2 - S1 + s1Idx + 1)
		12648	+ if s2RealSize <= 0:
		12649	+ s2RealSize = S2
		12650	+
		12651	+ if s2RealSize > topK:
		12652	+ s2RealLen = topK
		12653	+ else:
		12654	+ s2RealLen = s2RealSize
		12655	+ sparse_indices[:, s1Idx, 0, 0 : s2RealLen] = (torch.randint(0, s2RealSize, (s2RealLen,)).to(torch.int32)).npu()
		12656	+ sparse_indices[:, s1Idx, 0, s2RealLen : topK] = -1
		12657	+
		12658	+ softmax_max = torch.randn(B, N2, S1, NQueryIndex, dtype=torch.float, device=torch.device('npu'))
		12659	+ softmax_sum = torch.randn(B, N2, S1, NQueryIndex, dtype=torch.float, device=torch.device('npu'))
		12660	+ return q, k, q_index, k_index, q_rope, k_rope, weights, sparse_indices, softmax_max, softmax_sum
		12661	+
		12662	+
		12663	+actual_seq_qlen = [128]
		12664	+actual_seq_kvlen = [128]
		12665	+input_layout = 'TND'
		12666	+isTnd = True
		12667	+sparse_mode = 3
		12668	+scale = 1.0
		12669	+q, k, q_index, k_index, q_rope, k_rope, weights, sparse_indices, softmax_max, softmax_sum = gen_inputs(actual_seq_qlen, actual_seq_kvlen, isTnd)
		12670	+
		12671	+torch_npu.npu_sparse_lightning_indexer_grad_kl_loss(
		12672	+ q, k, q_index, k_index, weights, sparse_indices, softmax_max, softmax_sum, scale,
		12673	+ query_rope=q_rope, key_rope=k_rope, actual_seq_qlen=actual_seq_qlen, actual_seq_klen=actual_seq_kvlen, layout=input_layout, sparse_mode=sparse_mode, pre_tokens=65536, next_tokens=65536
		12674	+ )
		12675	+
12532	"""	12676	"""
12533	)	12677	)

Mop_plugin/config/derivatives.yaml+9-0

	@@ -285,3 +285,12 @@ backward:
285	- name: npu_moe_token_unpermute(Tensor permuted_tokens, Tensor sorted_indices, Tensor? probs=None, bool padded_mode=False, int[]? restore_shape=None) -> Tensor	285	- name: npu_moe_token_unpermute(Tensor permuted_tokens, Tensor sorted_indices, Tensor? probs=None, bool padded_mode=False, int[]? restore_shape=None) -> Tensor
286	permuted_tokens, probs: npu_moe_token_unpermute_grad(permuted_tokens, grad, sorted_indices, probs, padded_mode, restore_shape)	286	permuted_tokens, probs: npu_moe_token_unpermute_grad(permuted_tokens, grad, sorted_indices, probs, padded_mode, restore_shape)
287	version: [v2.1, newest]	287	version: [v2.1, newest]
		288	+
		289	+- name: npu_sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, float scale_value, *, Tensor? block_table=None, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_kv=None, Tensor? query_rope=None, Tensor? key_rope=None, int sparse_block_size=1, str layout_query="BSND", str layout_kv="BSND", int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, int attention_mode=0, bool return_softmax_lse=False) -> (Tensor, Tensor, Tensor)
		290	+ query, key, value, query_rope, key_rope: npu_sparse_flash_attention_grad(query, key, value, sparse_indices, grad, result0, result1, result2, scale_value, sparse_block_size, query_rope, key_rope, actual_seq_lengths_query, actual_seq_lengths_kv, layout_query, sparse_mode, pre_tokens, next_tokens, attention_mode)
		291	+ version: all_version
		292	+
		293	+- name: npu_lightning_indexer(Tensor query, Tensor key, Tensor weights, *, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, str layout_query="BSND", str layout_key="BSND", int sparse_count=2048, int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, bool return_value=False) -> (Tensor, Tensor)
		294	+ output_differentiability: [false, true]
		295	+ query, key, weights: npu_lightning_indexer_grad(query, key, grad, result0, weights, actual_seq_lengths_query, actual_seq_lengths_key, layout_query, sparse_mode, pre_tokens, next_tokens)
		296	+ version: all_version

Mop_plugin/config/op_plugin_functions.yaml+13-0

	@@ -6083,6 +6083,9 @@ custom:
6083	op_api: all_version	6083	op_api: all_version
6084	exposed: all_version	6084	exposed: all_version
6085		6085
		6086	+ - func: npu_lightning_indexer_grad(Tensor query, Tensor key, Tensor dy, Tensor sparse_indices, Tensor weights, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, str? layout="BSND", int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor)
		6087	+ op_api: all_version
		6088	+
6086	- func: npu_linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor	6089	- func: npu_linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor
6087	acl_op: all_version	6090	acl_op: all_version
6088	op_api: all_version	6091	op_api: all_version
	@@ -7005,6 +7008,13 @@ custom:
7005	op_api: all_version	7008	op_api: all_version
7006	exposed: all_version	7009	exposed: all_version
7007		7010
		7011	+ - func: npu_sparse_flash_attention_grad(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, Tensor d_out, Tensor out, Tensor softmax_max, Tensor softmax_sum, float scale_value, int sparse_block_size, Tensor? query_rope=None, Tensor? key_rope=None, Tensor? actual_seq_qlen=None, Tensor? actual_seq_kvlen=None, str? layout="BSND", int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807, int? attention_mode=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor)
		7012	+ op_api: all_version
		7013	+
		7014	+ - func: npu_sparse_lightning_indexer_grad_kl_loss(Tensor query, Tensor key, Tensor query_index, Tensor key_index, Tensor weights, Tensor sparse_indices, Tensor softmax_max, Tensor softmax_sum, float scale_value=1, *, Tensor? query_rope=None, Tensor? key_rope=None, SymInt[]? actual_seq_qlen=None, SymInt[]? actual_seq_klen=None, str? layout='BSND', int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor, Tensor)
		7015	+ op_api: all_version
		7016	+ exposed: all_version
		7017	+
7008	- func: npu_moe_token_unpermute_grad(Tensor permuted_tokens, Tensor grad_unpermuted_tokens, Tensor sorted_indices, Tensor? probs=None, bool padded_mode=False, int[]? restore_shape=None) -> (Tensor, Tensor)	7018	- func: npu_moe_token_unpermute_grad(Tensor permuted_tokens, Tensor grad_unpermuted_tokens, Tensor sorted_indices, Tensor? probs=None, bool padded_mode=False, int[]? restore_shape=None) -> (Tensor, Tensor)
7009	op_api: [v2.1, newest]	7019	op_api: [v2.1, newest]
7010	gen_opapi:	7020	gen_opapi:
	@@ -7153,6 +7163,9 @@ symint:
7153	- func: npu_add_quant_gmm(Tensor self, Tensor x1, Tensor x2, Tensor x2_scale, Tensor group_list, *, Tensor? x1_scale=None, int? group_list_type=0, SymInt[]? group_sizes=None, int? x1_dtype=None, int? x2_dtype=None, int? x1_scale_dtype=None, int? x2_scale_dtype=None) -> Tensor	7163	- func: npu_add_quant_gmm(Tensor self, Tensor x1, Tensor x2, Tensor x2_scale, Tensor group_list, *, Tensor? x1_scale=None, int? group_list_type=0, SymInt[]? group_sizes=None, int? x1_dtype=None, int? x2_dtype=None, int? x1_scale_dtype=None, int? x2_scale_dtype=None) -> Tensor
7154	op_api: all_version	7164	op_api: all_version
7155		7165
		7166	+ - func: npu_sparse_lightning_indexer_grad_kl_loss(Tensor query, Tensor key, Tensor query_index, Tensor key_index, Tensor weights, Tensor sparse_indices, Tensor softmax_max, Tensor softmax_sum, float scale_value=1, *, Tensor? query_rope=None, Tensor? key_rope=None, SymInt[]? actual_seq_qlen=None, SymInt[]? actual_seq_klen=None, str? layout='BSND', int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor, Tensor)
		7167	+ op_api: all_version
		7168	+
7156	quant:	7169	quant:
7157	- func: q_scale(Tensor self) -> float	7170	- func: q_scale(Tensor self) -> float
7158	acl_op: v2.2	7171	acl_op: v2.2

Aop_plugin/ops/opapi/LightningIndexerGradKernelNpuOpApi.cpp+62-0

@@ -0,0 +1,62 @@
	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
	7	+//
	8	+// https://opensource.org/licenses/BSD-3-Clause
	9	+//
	10	+// Unless required by applicable law or agreed to in writing, software
	11	+// distributed under the License is distributed on an "AS IS" BASIS,
	12	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	13	+// See the License for the specific language governing permissions and
	14	+// limitations under the License.
	15	+
	16	+#include <cstring>
	17	+
	18	+#include "torch_npu/csrc/framework/utils/RandomOpAdapter.h"
	19	+#include "torch_npu/csrc/aten/CustomFunctions.h"
	20	+#include "op_plugin/OpApiInterface.h"
	21	+#include "op_plugin/utils/op_api_common.h"
	22	+
	23	+namespace op_api {
	24	+using namespace at_npu::native;
	25	+
	26	+std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_lightning_indexer_grad(
	27	+ const at::Tensor &query,
	28	+ const at::Tensor &key,
	29	+ const at::Tensor &dy,
	30	+ const at::Tensor &sparse_indices,
	31	+ const at::Tensor &weights,
	32	+ const c10::optional<at::Tensor> &actual_seq_lengths_query,
	33	+ const c10::optional<at::Tensor> &actual_seq_lengths_key,
	34	+ c10::optional<c10::string_view> layout,
	35	+ c10::optional<int64_t> sparse_mode,
	36	+ c10::optional<int64_t> pre_tokens,
	37	+ c10::optional<int64_t> next_tokens)
	38	+{
	39	+ const at::Tensor &actual_seq_lengths_query_const = actual_seq_lengths_query.value_or(at::Tensor());
	40	+ const at::Tensor &actual_seq_lengths_key_const = actual_seq_lengths_key.value_or(at::Tensor());
	41	+
	42	+ at::Tensor d_query = OpPreparation::apply_tensor_without_format(query);
	43	+ at::Tensor d_key = OpPreparation::apply_tensor_without_format(key);
	44	+ at::Tensor d_weights = OpPreparation::apply_tensor_without_format(weights);
	45	+
	46	+ c10::string_view layout_str_view = layout.value_or("BSND");
	47	+ char layout_ptr = const_cast<char >(layout_str_view.data());
	48	+ const int64_t sparse_mode_const = sparse_mode.value_or(0);
	49	+ const int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
	50	+ const int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
	51	+ const int64_t head_num = 0;
	52	+ const bool deterministic = true;
	53	+
	54	+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
	55	+ aclnnLightningIndexerGrad, query, key, dy, sparse_indices, weights,
	56	+ actual_seq_lengths_query_const, actual_seq_lengths_key_const,
	57	+ head_num, layout_ptr, sparse_mode_const, pre_tokens_const, next_tokens_const, deterministic,
	58	+ d_query, d_key, d_weights);
	59	+
	60	+ return std::make_tuple(d_query, d_key, d_weights);
	61	+}
	62	+}

Mop_plugin/ops/opapi/LightningIndexerKernelNpuOpApi.cpp+3-9

Aop_plugin/ops/opapi/SparseFlashAttentionGradKernelNpuOpApi.cpp+104-0

@@ -0,0 +1,104 @@
	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
	7	+//
	8	+// https://opensource.org/licenses/BSD-3-Clause
	9	+//
	10	+// Unless required by applicable law or agreed to in writing, software
	11	+// distributed under the License is distributed on an "AS IS" BASIS,
	12	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	13	+// See the License for the specific language governing permissions and
	14	+// limitations under the License.
	15	+
	16	+#include <cstring>
	17	+
	18	+#include "torch_npu/csrc/framework/utils/RandomOpAdapter.h"
	19	+#include "torch_npu/csrc/aten/CustomFunctions.h"
	20	+#include "op_plugin/OpApiInterface.h"
	21	+#include "op_plugin/utils/op_api_common.h"
	22	+
	23	+namespace op_api {
	24	+using namespace at_npu::native;
	25	+const int DIMENSION_3D = 3;
	26	+const int DIMENSION_4D = 4;
	27	+
	28	+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention_grad(
	29	+ const at::Tensor &query,
	30	+ const at::Tensor &key,
	31	+ const at::Tensor &value,
	32	+ const at::Tensor &sparse_indices,
	33	+ const at::Tensor &d_out,
	34	+ const at::Tensor &out,
	35	+ const at::Tensor &softmax_max,
	36	+ const at::Tensor &softmax_sum,
	37	+ double scale_value,
	38	+ int64_t sparse_block_size,
	39	+ const c10::optional<at::Tensor> &query_rope,
	40	+ const c10::optional<at::Tensor> &key_rope,
	41	+ const c10::optional<at::Tensor> &actual_seq_qlen,
	42	+ const c10::optional<at::Tensor> &actual_seq_kvlen,
	43	+ c10::optional<c10::string_view> layout,
	44	+ c10::optional<int64_t> sparse_mode,
	45	+ c10::optional<int64_t> pre_tokens,
	46	+ c10::optional<int64_t> next_tokens,
	47	+ c10::optional<int64_t> attention_mode)
	48	+{
	49	+ const at::Tensor &query_rope_const = query_rope.value_or(at::Tensor());
	50	+ const at::Tensor &key_rope_const = key_rope.value_or(at::Tensor());
	51	+ const at::Tensor &ac_seq_qlen = actual_seq_qlen.value_or(at::Tensor());
	52	+ const at::Tensor &ac_seq_kvlen = actual_seq_kvlen.value_or(at::Tensor());
	53	+ TORCH_CHECK(query.dim() == DIMENSION_3D \|\| query.dim() == DIMENSION_4D,
	54	+ "The shapes of the input query should be 3 or 4 dimensional, but got ",
	55	+ query.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	56	+ if (query_rope_const.defined()) {
	57	+ TORCH_CHECK(query_rope_const.dim() == DIMENSION_3D \|\| query_rope_const.dim() == DIMENSION_4D,
	58	+ "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",
	59	+ query_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	60	+ }
	61	+ TORCH_CHECK(key.dim() == DIMENSION_3D \|\| key.dim() == DIMENSION_4D,
	62	+ "The shapes of the input key should be 3 or 4 dimensional, but got ", key.dim(),
	63	+ "-dimensional", OPS_ERROR(ErrCode::PARAM));
	64	+ if (key_rope_const.defined()) {
	65	+ TORCH_CHECK(key_rope_const.dim() == DIMENSION_3D \|\| key_rope_const.dim() == DIMENSION_4D,
	66	+ "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",
	67	+ key_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	68	+ }
	69	+ TORCH_CHECK(value.dim() == DIMENSION_3D \|\| value.dim() == DIMENSION_4D,
	70	+ "The shapes of the input value should be 3 or 4 dimensional, but got ",
	71	+ value.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	72	+ at::Tensor d_query = OpPreparation::apply_tensor_without_format(query);
	73	+ at::Tensor d_key = OpPreparation::apply_tensor_without_format(key);
	74	+ at::Tensor d_value = OpPreparation::apply_tensor_without_format(value);
	75	+ at::Tensor d_query_rope;
	76	+ at::Tensor d_key_rope;
	77	+ if (query_rope_const.defined()) {
	78	+ d_query_rope = OpPreparation::apply_tensor_without_format(query_rope_const);
	79	+ } else {
	80	+ d_query_rope = at::empty({0}, query.options());
	81	+ }
	82	+ if (key_rope_const.defined()) {
	83	+ d_key_rope = OpPreparation::apply_tensor_without_format(key_rope_const);
	84	+ } else {
	85	+ d_key_rope = at::empty({0}, key.options());
	86	+ }
	87	+
	88	+ c10::string_view layout_str_view = layout.value_or("BSND");
	89	+ char layout_ptr = const_cast<char >(layout_str_view.data());
	O O openLiBingCI2025年11月19日指针与引用安全: 代码中使用const_cast<char>移除了layout_str_view.data()的const限定符，然后将其赋值给layout_ptr。虽然layout_str_view.data()返回的是const char，但这里通过const_cast强制转换后，如果后续代码通过layout_ptr修改了字符串内容，会导致未定义行为。layout_str_view是一个c10::string_view，应该保持其不可变性。问题类型: 指针与引用安全文件路径: `op_plugin/ops/opapi/SparseFlashAttentionGradKernelNpuOpApi.cpp` 行号: 114 问题代码: `char layout_ptr = const_cast<char >(layout_str_view.data());` 修改建议: `如果底层API确实需要非const的char参数，应该考虑复制字符串数据到新分配的缓冲区。或者重新评估底层API的设计，看是否可以使用const char。建议改为：std::string layout_str = std::string(layout_str_view); char layout_ptr = layout_str.data();` 此评论由代码审查工具自动生成*
	90	+
	91	+ const int64_t sparse_mode_const = sparse_mode.value_or(3);
	92	+ const int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
	93	+ const int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
	94	+ const bool deterministic_const = true;
	95	+
	96	+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
	97	+ aclnnSparseFlashAttentionGrad, query, key, value, sparse_indices, d_out,
	98	+ out, softmax_max, softmax_sum, ac_seq_qlen, ac_seq_kvlen, query_rope_const, key_rope_const,
	99	+ scale_value, sparse_block_size, layout_ptr, sparse_mode_const, pre_tokens_const, next_tokens_const,
	100	+ deterministic_const, d_query, d_key, d_value, d_query_rope, d_key_rope);
	101	+
	102	+ return std::make_tuple(d_query, d_key, d_value, d_query_rope, d_key_rope);
	103	+}
	104	+}

Mop_plugin/ops/opapi/SparseFlashAttentionKernelNpuOpApi.cpp+9-13

	@@ -24,11 +24,10 @@ const static int64_t DIM_0 = 0;
24	const static int64_t DIM_1 = 1;	24	const static int64_t DIM_1 = 1;
25	const static int64_t DIM_2 = 2;	25	const static int64_t DIM_2 = 2;
26	const static int64_t DIM_3 = 3;	26	const static int64_t DIM_3 = 3;
27	-const static int64_t DIM_NUM_3 = 3;
28		27
29	-using namespace at_npu::native;
30	using npu_preparation = at_npu::native::OpPreparation;	28	using npu_preparation = at_npu::native::OpPreparation;
31		29
		30	+namespace {
32	at::Tensor construct_sparse_flash_attention_output_tensor(	31	at::Tensor construct_sparse_flash_attention_output_tensor(
33	const at::Tensor& query, std::string layout)	32	const at::Tensor& query, std::string layout)
34	{	33	{
	@@ -46,6 +45,8 @@ at::Tensor construct_sparse_flash_attention_output_tensor(
46		45
47	return output;	46	return output;
48	}	47	}
		48	+}
		49	+
49		50
50	std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention(	51	std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention(
51	const at::Tensor &query, const at::Tensor &key, const at::Tensor &value,	52	const at::Tensor &query, const at::Tensor &key, const at::Tensor &value,
	@@ -66,23 +67,18 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention(
66	std::string layout_kv_str = std::string(layout_kv);	67	std::string layout_kv_str = std::string(layout_kv);
67		68
68	// construct the output tensor	69	// construct the output tensor
69	- at::Tensor sparse_flash_attention_output = op_api::construct_sparse_flash_attention_output_tensor(	70	+ at::Tensor sparse_flash_attention_output = construct_sparse_flash_attention_output_tensor(
70	query, layout_query_str);	71	query, layout_query_str);
71	at::Tensor softmax_max;	72	at::Tensor softmax_max;
72	at::Tensor softmax_sum;	73	at::Tensor softmax_sum;
73	at::SmallVector<int64_t, SIZE> softmax_max_size;	74	at::SmallVector<int64_t, SIZE> softmax_max_size;
74	at::SmallVector<int64_t, SIZE> softmax_sum_size;	75	at::SmallVector<int64_t, SIZE> softmax_sum_size;
75	- if (return_softmax_lse) {	76	+ if (query.dim() == DIM_3) {
76	- if (query.dim() == DIM_NUM_3) {	77	+ softmax_max_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
77	- softmax_max_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};	78	+ softmax_sum_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
78	- softmax_sum_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
79	- } else {
80	- softmax_max_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
81	- softmax_sum_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
82	- }
83	} else {	79	} else {
84	- softmax_max_size = {0};	80	+ softmax_max_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
85	- softmax_sum_size = {0};	81	+ softmax_sum_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
86	}	82	}
87	softmax_max = at::empty(softmax_max_size, query.options().dtype(at::kFloat));	83	softmax_max = at::empty(softmax_max_size, query.options().dtype(at::kFloat));
88	softmax_sum = at::empty(softmax_sum_size, query.options().dtype(at::kFloat));	84	softmax_sum = at::empty(softmax_sum_size, query.options().dtype(at::kFloat));

Aop_plugin/ops/opapi/SparseLightningIndexerGradKLLossKernelNpuOpApi.cpp+86-0

@@ -0,0 +1,86 @@
	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
	7	+//
	8	+// https://opensource.org/licenses/BSD-3-Clause
	9	+//
	10	+// Unless required by applicable law or agreed to in writing, software
	11	+// distributed under the License is distributed on an "AS IS" BASIS,
	12	+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	13	+// See the License for the specific language governing permissions and
	14	+// limitations under the License.
	15	+
	16	+#include <cstring>
	17	+
	18	+#include "torch_npu/csrc/framework/utils/RandomOpAdapter.h"
	19	+#include "torch_npu/csrc/aten/CustomFunctions.h"
	20	+#include "op_plugin/OpApiInterface.h"
	21	+#include "op_plugin/utils/op_api_common.h"
	22	+
	23	+namespace op_api {
	24	+using namespace at_npu::native;
	25	+const int DIMENSION_3D = 3;
	26	+const int DIMENSION_4D = 4;
	27	+
	28	+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_lightning_indexer_grad_kl_loss_symint(
	29	+ const at::Tensor &query,
	30	+ const at::Tensor &key,
	31	+ const at::Tensor &query_index,
	32	+ const at::Tensor &key_index,
	33	+ const at::Tensor &weights,
	34	+ const at::Tensor &sparse_indices,
	35	+ const at::Tensor &softmax_max,
	36	+ const at::Tensor &softmax_sum,
	37	+ double scale_value,
	38	+ const c10::optional<at::Tensor> &query_rope,
	39	+ const c10::optional<at::Tensor> &key_rope,
	40	+ c10::OptionalArrayRef<c10::SymInt> actual_seq_qlen,
	41	+ c10::OptionalArrayRef<c10::SymInt> actual_seq_klen,
	42	+ c10::optional<c10::string_view> layout,
	43	+ c10::optional<int64_t> sparse_mode,
	44	+ c10::optional<int64_t> pre_tokens,
	45	+ c10::optional<int64_t> next_tokens)
	46	+{
	47	+ const at::Tensor &query_rope_const = query_rope.value_or(at::Tensor());
	48	+ const at::Tensor &key_rope_const = key_rope.value_or(at::Tensor());
	49	+ auto actual_seq_qlen_const = actual_seq_qlen.has_value() ? c10::asIntArrayRefUnchecked(actual_seq_qlen.value()) : at::IntArrayRef{};
	50	+ auto actual_seq_klen_const = actual_seq_klen.has_value() ? c10::asIntArrayRefUnchecked(actual_seq_klen.value()) : at::IntArrayRef{};
	51	+ c10::string_view layout_str = layout.value_or("BSND");
	52	+ char layout_ptr = const_cast<char >(layout_str.data());
	53	+ int64_t sparse_mode_const = sparse_mode.value_or(3);
	54	+ int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
	55	+ int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
	56	+ bool deterministic_const = true;
	57	+ TORCH_CHECK(query.dim() == DIMENSION_3D \|\| query.dim() == DIMENSION_4D,
	58	+ "The shapes of the input query should be 3 or 4 dimensional, but got ",
	59	+ query.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	60	+ if (query_rope_const.defined()) {
	61	+ TORCH_CHECK(query_rope_const.dim() == DIMENSION_3D \|\| query_rope_const.dim() == DIMENSION_4D,
	62	+ "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",
	63	+ query_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	64	+ }
	65	+ TORCH_CHECK(key.dim() == DIMENSION_3D \|\| key.dim() == DIMENSION_4D,
	66	+ "The shapes of the input key should be 3 or 4 dimensional, but got ", key.dim(),
	67	+ "-dimensional", OPS_ERROR(ErrCode::PARAM));
	68	+ if (key_rope_const.defined()) {
	69	+ TORCH_CHECK(key_rope_const.dim() == DIMENSION_3D \|\| key_rope_const.dim() == DIMENSION_4D,
	70	+ "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",
	71	+ key_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
	72	+ }
	73	+ at::Tensor d_query_index = OpPreparation::apply_tensor_without_format(query_index);
	74	+ at::Tensor d_key_index = OpPreparation::apply_tensor_without_format(key_index);
	75	+ at::Tensor d_weights = OpPreparation::apply_tensor_without_format(weights);
	76	+ at::Tensor loss = OpPreparation::apply_tensor_without_format({1}, query.options().dtype(at::kFloat));
	77	+
	78	+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
	79	+ aclnnSparseLightningIndexerGradKLLoss, query, key, query_index, key_index, weights,
	80	+ sparse_indices, softmax_max, softmax_sum, query_rope_const, key_rope_const, actual_seq_qlen_const,
	81	+ actual_seq_klen_const, scale_value, layout_ptr, sparse_mode_const, pre_tokens_const, next_tokens_const, deterministic_const,
	82	+ d_query_index, d_key_index, d_weights, loss);
	83	+
	84	+ return std::make_tuple(d_query_index, d_key_index, d_weights, loss);
	85	+}
	86	+}

Mtest/allowlist_for_publicAPI.json+3-0

	@@ -37,6 +37,9 @@
37	"npu_quant_lightning_indexer",	37	"npu_quant_lightning_indexer",
38	"npu_lightning_indexer",	38	"npu_lightning_indexer",
39	"npu_sparse_flash_attention",	39	"npu_sparse_flash_attention",
		40	+ "npu_lightning_indexer_grad",
		41	+ "npu_sparse_flash_attention_grad",
		42	+ "npu_sparse_lightning_indexer_grad_kl_loss",
40	"npu_kv_quant_sparse_flash_attention",	43	"npu_kv_quant_sparse_flash_attention",
41	"npu_convert_weight_to_int4pack",	44	"npu_convert_weight_to_int4pack",
42	"npu_ffn",	45	"npu_ffn",

Mtest/core_tests/torch_npu_OpApi_schema_all.json+27-0

	@@ -41,9 +41,18 @@
41	"op_api: torch_npu.npu_lightning_indexer(args, *kwargs)": {	41	"op_api: torch_npu.npu_lightning_indexer(args, *kwargs)": {
42	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]	42	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]
43	},	43	},
		44	+ "op_api: torch_npu.npu_lightning_indexer_grad(args, *kwargs)": {
		45	+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
		46	+ },
44	"op_api: torch_npu.npu_sparse_flash_attention(args, *kwargs)": {	47	"op_api: torch_npu.npu_sparse_flash_attention(args, *kwargs)": {
45	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]	48	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]
46	},	49	},
		50	+ "op_api: torch_npu.npu_sparse_flash_attention_grad(args, *kwargs)": {
		51	+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
		52	+ },
		53	+ "op_api: torch_npu.npu_sparse_lightning_indexer_grad_kl_loss(args, *kwargs)": {
		54	+ "version": ["v2.1", "v2.5", "v2.6", "v2.7"]
		55	+ },
47	"op_api: torch_npu.npu_kv_quant_sparse_flash_attention(args, *kwargs)": {	56	"op_api: torch_npu.npu_kv_quant_sparse_flash_attention(args, *kwargs)": {
48	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]	57	"version": ["v2.1", "v2.5", "v2.6", "v2.7"]
49	},	58	},
	@@ -380,9 +389,18 @@
380	"op_api: torch_npu.npu_lightning_indexer(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {	389	"op_api: torch_npu.npu_lightning_indexer(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
381	"version": ["v2.8", "newest"]	390	"version": ["v2.8", "newest"]
382	},	391	},
		392	+ "op_api: torch_npu.npu_lightning_indexer_grad(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
		393	+ "version": ["v2.8", "newest"]
		394	+ },
383	"op_api: torch_npu.npu_sparse_flash_attention(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {	395	"op_api: torch_npu.npu_sparse_flash_attention(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
384	"version": ["v2.8", "newest"]	396	"version": ["v2.8", "newest"]
385	},	397	},
		398	+ "op_api: torch_npu.npu_sparse_flash_attention_grad(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
		399	+ "version": ["v2.8", "newest"]
		400	+ },
		401	+ "op_api: torch_npu.npu_sparse_lightning_indexer_grad_kl_loss(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
		402	+ "version": ["v2.8", "newest"]
		403	+ },
386	"op_api: torch_npu.npu_kv_quant_sparse_flash_attention(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {	404	"op_api: torch_npu.npu_kv_quant_sparse_flash_attention(args: _P.args, *kwargs: _P.kwargs) -> ~_T": {
387	"version": ["v2.8", "newest"]	405	"version": ["v2.8", "newest"]
388	},	406	},
	@@ -1199,9 +1217,18 @@
1199	"func: npu_lightning_indexer(Tensor query, Tensor key, Tensor weights, *, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, str layout_query=\"BSND\", str layout_key=\"BSND\", int sparse_count=2048, int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, bool return_value=False) -> (Tensor, Tensor)": {	1217	"func: npu_lightning_indexer(Tensor query, Tensor key, Tensor weights, *, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, str layout_query=\"BSND\", str layout_key=\"BSND\", int sparse_count=2048, int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, bool return_value=False) -> (Tensor, Tensor)": {
1200	"version": ["all_version"]	1218	"version": ["all_version"]
1201	},	1219	},
		1220	+ "func: npu_lightning_indexer_grad(Tensor query, Tensor key, Tensor dy, Tensor sparse_indices, Tensor weights, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, str? layout=\"BSND\", int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor)": {
		1221	+ "version": ["all_version"]
		1222	+ },
1202	"func: npu_sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, float scale_value, *, Tensor? block_table=None, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_kv=None, Tensor? query_rope=None, Tensor? key_rope=None, int sparse_block_size=1, str layout_query=\"BSND\", str layout_kv=\"BSND\", int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, int attention_mode=0, bool return_softmax_lse=False) -> (Tensor, Tensor, Tensor)": {	1223	"func: npu_sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, float scale_value, *, Tensor? block_table=None, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_kv=None, Tensor? query_rope=None, Tensor? key_rope=None, int sparse_block_size=1, str layout_query=\"BSND\", str layout_kv=\"BSND\", int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, int attention_mode=0, bool return_softmax_lse=False) -> (Tensor, Tensor, Tensor)": {
1203	"version": ["all_version"]	1224	"version": ["all_version"]
1204	},	1225	},
		1226	+ "func: npu_sparse_flash_attention_grad(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, Tensor d_out, Tensor out, Tensor softmax_max, Tensor softmax_sum, float scale_value, int sparse_block_size, Tensor? query_rope=None, Tensor? key_rope=None, Tensor? actual_seq_qlen=None, Tensor? actual_seq_kvlen=None, str? layout=\"BSND\", int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807, int? attention_mode=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor)": {
		1227	+ "version": ["all_version"]
		1228	+ },
		1229	+ "func: npu_sparse_lightning_indexer_grad_kl_loss(Tensor query, Tensor key, Tensor query_index, Tensor key_index, Tensor weights, Tensor sparse_indices, Tensor softmax_max, Tensor softmax_sum, float scale_value=1, *, Tensor? query_rope=None, Tensor? key_rope=None, SymInt[]? actual_seq_qlen=None, SymInt[]? actual_seq_klen=None, str? layout='BSND', int? sparse_mode=3, int? pre_tokens=9223372036854775807, int? next_tokens=9223372036854775807) -> (Tensor, Tensor, Tensor, Tensor)": {
		1230	+ "version": ["all_version"]
		1231	+ },
1205	"func: npu_kv_quant_sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, float scale_value, int key_quant_mode, int value_quant_mode, *, Tensor? key_dequant_scale=None, Tensor? value_dequant_scale=None, Tensor? block_table=None, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_kv=None, int sparse_block_size=1, str layout_query=\"BSND\", str layout_kv=\"BSND\", int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, int attention_mode=0, int quant_scale_repo_mode=1, int tile_size=128, int rope_head_dim=64) -> Tensor": {	1232	"func: npu_kv_quant_sparse_flash_attention(Tensor query, Tensor key, Tensor value, Tensor sparse_indices, float scale_value, int key_quant_mode, int value_quant_mode, *, Tensor? key_dequant_scale=None, Tensor? value_dequant_scale=None, Tensor? block_table=None, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_kv=None, int sparse_block_size=1, str layout_query=\"BSND\", str layout_kv=\"BSND\", int sparse_mode=3, int pre_tokens=9223372036854775807, int next_tokens=9223372036854775807, int attention_mode=0, int quant_scale_repo_mode=1, int tile_size=128, int rope_head_dim=64) -> Tensor": {
1206	"version": ["all_version"]	1233	"version": ["all_version"]
1207	},	1234	},

	@@ -23,13 +23,12 @@ namespace op_api {
23	const static int64_t DIM_0 = 0;	23	const static int64_t DIM_0 = 0;
24	const static int64_t DIM_1 = 1;	24	const static int64_t DIM_1 = 1;
25	const static int64_t DIM_2 = 2;	25	const static int64_t DIM_2 = 2;
26	-const static int64_t DIM_3 = 3;
27	using namespace at_npu::native;	26	using namespace at_npu::native;
28	using npu_preparation = at_npu::native::OpPreparation;	27	using npu_preparation = at_npu::native::OpPreparation;
29		28
30	std::tuple<at::Tensor, at::Tensor> construct_lightning_indexer_output_tensor(const at::Tensor& query,	29	std::tuple<at::Tensor, at::Tensor> construct_lightning_indexer_output_tensor(const at::Tensor& query,
31	const at::Tensor& key, const c10::optional<at::Tensor> &actual_seq_lengths_query, int64_t sparse_count,	30	const at::Tensor& key, const c10::optional<at::Tensor> &actual_seq_lengths_query, int64_t sparse_count,
32	- std::string query_layout_str, std::string key_layout_str, bool return_value)	31	+ std::string query_layout_str, std::string key_layout_str)
33	{	32	{
34	at::SmallVector<int64_t, SIZE> output_size;	33	at::SmallVector<int64_t, SIZE> output_size;
35	for (size_t i = 0; i < query.sizes().size(); i++) {	34	for (size_t i = 0; i < query.sizes().size(); i++) {
	@@ -49,12 +48,7 @@ std::tuple<at::Tensor, at::Tensor> construct_lightning_indexer_output_tensor(con
49	output_size = {query.size(DIM_0), key.size(n_dim_index), sparse_count};	48	output_size = {query.size(DIM_0), key.size(n_dim_index), sparse_count};
50	}	49	}
51	at::Tensor sparse_indices_out = npu_preparation::apply_tensor_without_format(output_size, at::kInt);	50	at::Tensor sparse_indices_out = npu_preparation::apply_tensor_without_format(output_size, at::kInt);
52	- at::Tensor sparse_values_out;	51	+ at::Tensor sparse_values_out = npu_preparation::apply_tensor_without_format(output_size, query.dtype());
53	- if (return_value) {
54	- sparse_values_out = npu_preparation::apply_tensor_without_format(output_size, query.dtype());
55	- } else {
56	- sparse_values_out = npu_preparation::apply_tensor_without_format({0}, query.dtype());
57	- }
58		52
59	return std::tuple<at::Tensor, at::Tensor>(sparse_indices_out, sparse_values_out);	53	return std::tuple<at::Tensor, at::Tensor>(sparse_indices_out, sparse_values_out);
60	}	54	}
	@@ -75,7 +69,7 @@ std::tuple<at::Tensor, at::Tensor> npu_lightning_indexer(
75		69
76	// construct the output tensor	70	// construct the output tensor
77	std::tuple<at::Tensor, at::Tensor> lightning_indexer_output = op_api::construct_lightning_indexer_output_tensor(	71	std::tuple<at::Tensor, at::Tensor> lightning_indexer_output = op_api::construct_lightning_indexer_output_tensor(
78	- query, key, actual_seq_lengths_query, sparse_count, query_layout_str, key_layout_str, return_value);	72	+ query, key, actual_seq_lengths_query, sparse_count, query_layout_str, key_layout_str);
79	at::Tensor sparse_indices_out = std::get<0>(lightning_indexer_output);	73	at::Tensor sparse_indices_out = std::get<0>(lightning_indexer_output);
80	at::Tensor sparse_values_out = std::get<1>(lightning_indexer_output);	74	at::Tensor sparse_values_out = std::get<1>(lightning_indexer_output);
81	// convert str	75	// convert str