Star0
1
代码
代码
Issues
Pull Requests2
项目讨论
Wiki
分析
Star0
1
已开启
[feat] add aclnn_extension codegen #2
梁松伟创建于 1月24日
[feat] add aclnn_extension codegen #2
已开启
梁松伟创建于 1月24日
从master_lsw合入到master_copy
共 71 个文件变更+2016-112
Mci/build.sh+1-0
@@ -125,6 +125,7 @@ function main()
125 125 
126 cp -rf ${CODE_ROOT_PATH}/op_plugin ${PYTORCH_THIRD_PATH}/126 cp -rf ${CODE_ROOT_PATH}/op_plugin ${PYTORCH_THIRD_PATH}/
127 cp -rf ${CODE_ROOT_PATH}/codegen ${PYTORCH_THIRD_PATH}/127 cp -rf ${CODE_ROOT_PATH}/codegen ${PYTORCH_THIRD_PATH}/
128+ cp -rf ${CODE_ROOT_PATH}/torchnpugen ${PYTORCH_THIRD_PATH}/
128 cp -rf ${CODE_ROOT_PATH}/*.sh ${PYTORCH_THIRD_PATH}/129 cp -rf ${CODE_ROOT_PATH}/*.sh ${PYTORCH_THIRD_PATH}/
129 cp -rf ${CODE_ROOT_PATH}/test ${PYTORCH_THIRD_PATH}/130 cp -rf ${CODE_ROOT_PATH}/test ${PYTORCH_THIRD_PATH}/
130 131 
Mcodegen/templates/_op_plugin_docs.py+84-3文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/menu_Pytorch_API.md+2-1文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu-npu/（beta）torch_npu-npu-obfuscation_calculate.md+2-2文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu-npu/（beta）torch_npu-npu-obfuscation_finalize.md+2-2文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu-npu/（beta）torch_npu-npu-obfuscation_initialize.md+2-2文件内容审核中，请稍后刷新重试
Adocs/zh/custom_APIs/torch_npu/torch_npu-npu_dense_lightning_indexer_grad_kl_loss.md+201-0文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_fused_infer_attention_score.md+2-2文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_fused_infer_attention_score_v2.md+2-2文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_grouped_matmul_swiglu_quant_v2.md+1-1文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_moe_gating_top_k.md+1-1文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu-npu_prefetch.md+1-1文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/torch_npu_list.md+4-2文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/（beta）torch_npu-npu_lstm.md+1-1文件内容审核中，请稍后刷新重试
Mdocs/zh/custom_APIs/torch_npu/（beta）torch_npu-npu_roi_align.md+1-1文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/derivatives.yaml+2-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/gen_autograd.py+152-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/gen_variable_type.py+121-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/templates/ADInplaceOrViewType.cpp+26-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/templates/Functions.cpp+46-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/templates/Functions.h+63-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/templates/VariableType.cpp+63-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/autograd/templates/VariableType.h+57-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/gen.py+250-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/struct.py+106-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/templates/CustomRegisterSchema.cpp+59-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/codegen/templates/StructKernelNpuOpApi.cpp+13-0文件内容审核中，请稍后刷新重试
Aexamples/aclnn_extension/gen.sh+17-0
@@ -0,0 +1,17 @@
1+CDIR="$(cd "$(dirname "$0")" ; pwd -P)"
2+ 
3+cd $CDIR
4+ 
5+python3 -m codegen.gen \
6+ --output_dir="$CDIR" \
7+ --source_yaml="$CDIR/test.yaml" \
8+ 
9+python3 -m codegen.struct \
10+ --output_dir="$CDIR" \
11+ --native_yaml="$CDIR/test.yaml" \
12+ --struct_yaml="$CDIR/test.yaml" \
13+ 
14+python3 -m codegen.autograd.gen_autograd \
15+ --out_dir="$CDIR" \
16+ --autograd_dir="$CDIR/codegen/autograd" \
17+ --npu_native_function_dir="$CDIR/test.yaml" \
Aexamples/aclnn_extension/test.yaml+16-0
@@ -0,0 +1,16 @@
1+cpp_namespace: at_npu::native
2+ 
3+custom:
4+ - func: npu_change_data_ptr(Tensor dst, Tensor src, int index) -> int
5+ - func: npu_test(Tensor input, Scalar alpha, float prob, int dim) -> Tensor
6+ gen_opapi:
7+ out:
8+ size: input
9+ dtype: input
10+ exec: aclnnTest
11+ - func: npu_test_backward(Tensor grad) -> Tensor
12+ gen_opapi:
13+ out:
14+ size: grad
15+ dtype: grad
16+ exec: aclnnTestBackward
Mgencode.sh+4-4
@@ -43,24 +43,24 @@ if [ "$PYTORCH_VERSION_DIR" == "v1r11" ]; then
43 cp $CDIR/op_plugin/config/aclnn_derivatives.yaml $OUTPUT_DIR43 cp $CDIR/op_plugin/config/aclnn_derivatives.yaml $OUTPUT_DIR
44fi44fi
45 45 
46-${python_execute} -m codegen.gen_op_plugin_functions \46+${python_execute} -m torchnpugen.gen_op_plugin_functions \
47 --version="$PYTORCH_VERSION" \47 --version="$PYTORCH_VERSION" \
48 --output_dir="$OUTPUT_DIR/" \48 --output_dir="$OUTPUT_DIR/" \
49 --source_yaml="$CDIR/op_plugin/config/op_plugin_functions.yaml"49 --source_yaml="$CDIR/op_plugin/config/op_plugin_functions.yaml"
50 50 
51-${python_execute} -m codegen.gen_derivatives \51+${python_execute} -m torchnpugen.gen_derivatives \
52 --version="$PYTORCH_VERSION" \52 --version="$PYTORCH_VERSION" \
53 --output_dir="$OUTPUT_DIR/" \53 --output_dir="$OUTPUT_DIR/" \
54 --source_yaml="$CDIR/op_plugin/config/derivatives.yaml"54 --source_yaml="$CDIR/op_plugin/config/derivatives.yaml"
55 55 
56-${python_execute} -m codegen.gen_backend_stubs \56+${python_execute} -m torchnpugen.gen_op_backend \
57 --version="$PYTORCH_VERSION" \57 --version="$PYTORCH_VERSION" \
58 --output_dir="$CDIR/op_plugin/" \58 --output_dir="$CDIR/op_plugin/" \
59 --source_yaml="$OUTPUT_DIR/op_plugin_functions.yaml" \59 --source_yaml="$OUTPUT_DIR/op_plugin_functions.yaml" \
60 --deprecate_yaml="$CDIR/op_plugin/config/deprecated.yaml" \60 --deprecate_yaml="$CDIR/op_plugin/config/deprecated.yaml" \
61 --impl_path="$CDIR/torch_npu/csrc/aten" # Used to double-check the yaml file definitions.61 --impl_path="$CDIR/torch_npu/csrc/aten" # Used to double-check the yaml file definitions.
62 62 
63-${python_execute} -m codegen.struct.gen_struct_opapi \63+${python_execute} -m torchnpugen.struct.gen_struct_opapi \
64 --output_dir="$CDIR/op_plugin/ops/opapi/" \64 --output_dir="$CDIR/op_plugin/ops/opapi/" \
65 --native_yaml="$OUTPUT_DIR/op_plugin_functions.yaml" \65 --native_yaml="$OUTPUT_DIR/op_plugin_functions.yaml" \
66 --struct_yaml="$CDIR/op_plugin/config/op_plugin_functions.yaml"66 --struct_yaml="$CDIR/op_plugin/config/op_plugin_functions.yaml"
Mop_plugin/config/derivatives.yaml+4-0文件内容审核中，请稍后刷新重试
Mop_plugin/config/op_plugin_functions.yaml+10-9文件内容审核中，请稍后刷新重试
Aop_plugin/ops/opapi/DenseLightningIndexerGradKLLossKernelNpuOpApi.cpp+84-0
@@ -0,0 +1,84 @@
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+// 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_dense_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 &softmax_max,
35+ const at::Tensor &softmax_sum,
36+ const at::Tensor &softmax_max_index,
37+ const at::Tensor &softmax_sum_index,
38+ double scale_value,
39+ const c10::optional<at::Tensor> &query_rope,
40+ const c10::optional<at::Tensor> &key_rope,
41+ c10::OptionalArrayRef<c10::SymInt> actual_seq_qlen,
42+ c10::OptionalArrayRef<c10::SymInt> actual_seq_klen,
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+{
48+ const at::Tensor &query_rope_const = query_rope.value_or(at::Tensor());
49+ const at::Tensor &key_rope_const = key_rope.value_or(at::Tensor());
50+ c10::string_view layout_str = layout.value_or("BSND");
51+ char *layout_ptr = const_cast<char *>(layout_str.data());
52+ int64_t sparse_mode_const = sparse_mode.value_or(3);
53+ int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
54+ int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
55+ TORCH_CHECK(query.dim() == DIMENSION_3D || query.dim() == DIMENSION_4D,
56+ "The shapes of the input query should be 3 or 4 dimensional, but got ",
57+ query.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
58+ if (query_rope_const.defined()) {
59+ TORCH_CHECK(query_rope_const.dim() == DIMENSION_3D || query_rope_const.dim() == DIMENSION_4D,
60+ "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",
61+ query_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
62+ }
63+ TORCH_CHECK(key.dim() == DIMENSION_3D || key.dim() == DIMENSION_4D,
64+ "The shapes of the input key should be 3 or 4 dimensional, but got ", key.dim(),
65+ "-dimensional", OPS_ERROR(ErrCode::PARAM));
66+ if (key_rope_const.defined()) {
67+ TORCH_CHECK(key_rope_const.dim() == DIMENSION_3D || key_rope_const.dim() == DIMENSION_4D,
68+ "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",
69+ key_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));
70+ }
71+ at::Tensor d_query_index = OpPreparation::apply_tensor_without_format(query_index);
72+ at::Tensor d_key_index = OpPreparation::apply_tensor_without_format(key_index);
73+ at::Tensor d_weights = OpPreparation::apply_tensor_without_format(weights);
74+ at::Tensor loss = OpPreparation::apply_tensor_without_format({1}, query.options().dtype(at::kFloat));
75+ 
76+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
77+ aclnnDenseLightningIndexerGradKLLoss, query, key, query_index, key_index, weights,
78+ softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, query_rope_const, key_rope_const,
79+ actual_seq_qlen, actual_seq_klen, scale_value, layout_ptr, sparse_mode_const, pre_tokens_const,
80+ next_tokens_const, d_query_index, d_key_index, d_weights, loss);
81+ 
82+ return std::make_tuple(d_query_index, d_key_index, d_weights, loss);
83+}
84+}
Aop_plugin/ops/opapi/DynamicDualLevelMxQuantKernelNpuOpApi.cpp+67-0
@@ -0,0 +1,67 @@
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+// 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 "op_plugin/AclOpsInterface.h"
17+#include "op_plugin/OpApiInterface.h"
18+#include "op_plugin/utils/op_api_common.h"
19+ 
20+namespace op_api {
21+using npu_preparation = at_npu::native::OpPreparation;
22+namespace {
23+constexpr int64_t ALIGN_NUM = 2;
24+constexpr int64_t FP4_IN_UINT8_NUM = 2;
25+}; // namespace
26+ 
27+std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_dynamic_dual_level_mx_quant(
28+ const at::Tensor &input,
29+ const c10::optional<at::Tensor> &smooth_scale,
30+ c10::string_view round_mode)
31+{
32+ at::Tensor y;
33+ at::Tensor level0_scale;
34+ at::Tensor level1_scale;
35+ auto y_shape = op_infer::array_to_small_vector(input.sizes());
36+ auto level0_scale_shape = op_infer::array_to_small_vector(input.sizes());
37+ auto level1_scale_shape = op_infer::array_to_small_vector(input.sizes());
38+ level1_scale_shape.emplace_back(ALIGN_NUM);
39+ 
40+ int64_t level0_block_size = 512;
41+ int64_t level1_block_size = 32;
42+ int64_t dim0_size = op_infer::CeilDiv(level0_scale_shape[input.dim() - 1], level0_block_size);
43+ level0_scale_shape[input.dim() - 1] = dim0_size;
44+ int64_t dim1_size = op_infer::CeilDiv(level1_scale_shape[input.dim() - 1], level1_block_size);
45+ dim1_size = (dim1_size + ALIGN_NUM - 1) / ALIGN_NUM;
46+ level1_scale_shape[input.dim() - 1] = dim1_size;
47+ char *round_mode_ptr = const_cast<char *>(round_mode.data());
48+ 
49+ // prepare for empty output tensor
50+ aclDataType y_acltype = aclDataType::ACL_FLOAT4_E2M1;
51+ int64_t last_dim_val = y_shape[input.dim() - 1];
52+ TORCH_CHECK(last_dim_val % FP4_IN_UINT8_NUM == 0,
53+ "The last dim input shape must be divisible by 2 if "
54+ "output dtype is torch_npu.float4_e2m1" + OPS_ERROR(ErrCode::PARAM));
55+ y_shape[input.dim() - 1] = last_dim_val / FP4_IN_UINT8_NUM;
56+ y = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);
57+ 
58+ level0_scale = npu_preparation::apply_tensor_without_format(level0_scale_shape, c10::dtype(at::ScalarType::Float));
59+ level1_scale = npu_preparation::apply_tensor_without_format(level1_scale_shape, c10::dtype(at::ScalarType::Byte));
60+ TensorWrapper y_wrapper = {y, y_acltype};
61+ TensorWrapper level1_scale_wrapper = {level1_scale, aclDataType::ACL_FLOAT8_E8M0};
62+ EXEC_NPU_CMD(aclnnDynamicDualLevelMxQuant, input, smooth_scale, round_mode_ptr, level0_block_size, level1_block_size, y_wrapper, level0_scale, level1_scale_wrapper);
63+
64+ return std::make_tuple(y, level0_scale, level1_scale);
65+}
66+ 
67+} // namespace op_api
Mop_plugin/ops/opapi/FFTc2cKernelNpuOpApi.cpp+2-2文件内容审核中，请稍后刷新重试
Mop_plugin/ops/opapi/NativeDropoutKernelNpuOpApi.cpp+4-0文件内容审核中，请稍后刷新重试
Mop_plugin/ops/opapi/QuantReduceScatterKernelNpuOpApi.cpp+37-15文件内容审核中，请稍后刷新重试
Mop_plugin/ops/opapi/SaveNpugraphTensorKernelNpuOpApi.cpp+7-4文件内容审核中，请稍后刷新重试
Mop_plugin/python/meta/_meta_registrations.py+49-2文件内容审核中，请稍后刷新重试
Mtest/allowlist_for_publicAPI.json+1-0文件内容审核中，请稍后刷新重试
Mtest/core_tests/test_fake_tensor.py+82-0文件内容审核中，请稍后刷新重试
Mtest/core_tests/torch_npu_OpApi_schema_all.json+11-6文件内容审核中，请稍后刷新重试
Mtest/test_base_ops/test_fft1d.py+12-0文件内容审核中，请稍后刷新重试
Atest/test_custom_ops/test_npu_dense_lightning_indexer_grad_kl_loss.py+213-0
@@ -0,0 +1,213 @@
1+import unittest
2+import torch
3+import numpy as np
4+import torch.nn as nn
5+import torch_npu
6+from torch_npu.testing.testcase import TestCase, run_tests
7+ 
8+ 
9+class TestDenseLightningIndexerGradKLLoss(TestCase):
10+ 
11+ def _deal_sparse_mask(self, input_matrix):
12+ assert len(input_matrix.shape) == 4
13+ 
14+ _, _, s1, s2 = input_matrix.shape
15+ assert s1 <= s2
16+ 
17+ res_matrix = input_matrix
18+ for s1_idx in range(s1):
19+ res_matrix[:, :, s1_idx, (s2 - s1 + s1_idx + 1):] = -(torch.inf)
20+ 
21+ return res_matrix
22+ 
23+ def _process_p(self, input_query, input_key, input_softmax_max, input_softmax_sum, scale):
24+ # Q: [B,S1,N1,D] -> [B,N,S1,D]
25+ _, _, q_head_num, _ = input_query.shape
26+ query = input_query.permute(0, 2, 1, 3).contiguous()
27+ # K: [B,S2,N2,D] -> [B,N2,D,S2]
28+ key = input_key.permute(0, 2, 3, 1).contiguous()
29+ # batch matmul
30+ p_tmp = torch.matmul(query.to(torch.float32), key.to(torch.float32))
31+ # scale
32+ p_tmp *= scale
33+ 
34+ # sparse mask
35+ p_sparse = self._deal_sparse_mask(p_tmp)
36+ # simple softmax
37+ softmax_max_res = input_softmax_max.permute(0, 1, 3, 2).contiguous().reshape(p_sparse.shape[0], -1, p_sparse.shape[2])
38+ p_diff = p_sparse - softmax_max_res.unsqueeze(-1)
39+ p_exp = torch.exp(p_diff)
40+ softmax_sum_res = input_softmax_sum.permute(0, 1, 3, 2).contiguous().reshape(p_exp.shape[0], -1, p_exp.shape[2])
41+ p_div = p_exp.div(softmax_sum_res.unsqueeze(-1))
42+ # reduce sum: (B,N1,S1,S2) -> (B,S1,S2)
43+ p_reduce = p_div.sum(axis=1, keepdims=False)
44+ # scale
45+ p_reduce *= (1 / q_head_num)
46+ 
47+ return p_reduce, softmax_max_res, softmax_sum_res
48+
49+ def _process_sy(self, input_query_index, input_key_index, input_weight, input_softmax_max_index, input_softmax_sum_index):
50+ _, _, N1, D = input_query_index.shape
51+ _, _, N2, _ = input_key_index.shape
52+ group_size_index = N1 // N2
53+ # Q_INDEX: [B,S1,N1,D] -> [B,N1,S1,D]
54+ query_index = input_query_index.permute(0, 2, 1, 3).contiguous()
55+ # K_INDEX: [B,S2,N2,D] -> [B,N2*G,D,S2]
56+ key_index = input_key_index.permute(0, 2, 3, 1).contiguous().repeat(1, group_size_index, 1, 1)
57+ # batch matmul
58+ s_tmp = torch.matmul(query_index.to(torch.float32), key_index.to(torch.float32))
59+ # relu
60+ s_relu = torch.relu(s_tmp)
61+ # WEIGHT: [B,S1,N1] -> [B,N1,S1,1]
62+ weight = input_weight.permute(0, 2, 1).contiguous().unsqueeze(-1)
63+ # Mul
64+ s_mul = s_relu * weight.to(torch.float32)
65+ # reduce
66+ s_reduce = s_mul.sum(axis=1, keepdims=True)
67+ # sparse mask
68+ s_sparse = self._deal_sparse_mask(s_reduce).squeeze(1)
69+ # simple softmax
70+ softmax_max_index_res = input_softmax_max_index.squeeze(1)
71+ s_diff = s_sparse - softmax_max_index_res.unsqueeze(-1)
72+ s_exp = torch.exp(s_diff)
73+ softmax_sum_index_res = input_softmax_sum_index.squeeze(1)
74+ s_div = s_exp.div(softmax_sum_index_res.unsqueeze(-1))
75+ 
76+ return s_div, s_relu, softmax_max_index_res, softmax_sum_index_res
77+
78+ def _process_kl_loss(self, p_result, sy_result):
79+ # clip
80+ min_value = torch.tensor([1e-8])
81+ p_result_clip = torch.max(p_result, min_value)
82+ sy_result_clip = torch.max(sy_result, min_value)
83+ # log
84+ p_log = torch.log(p_result_clip)
85+ sy_log = torch.log(sy_result_clip)
86+ # sub
87+ sub_result = p_log - sy_log
88+ # mul
89+ mul_result = sub_result * p_result
90+ # loss
91+ loss = torch.sum(mul_result)
92+ 
93+ return loss
94+ 
95+ 
96+ def _process_dwqk(self, p_result, sy_result, relu_res, input_query_index, input_key_index, input_weight, q_dtype):
97+
98+ _, _, N1, D = input_query_index.shape
99+ _, _, N2, _ = input_key_index.shape
100+ group_size = N1 // N2
101+ # sub
102+ sub_result = sy_result - p_result
103+ # mul: (B,S1,1,S2) * (B,S1,N1,S2)
104+ mul_relu = sub_result.unsqueeze(2) * relu_res.permute(0, 2, 1, 3).contiguous()
105+ # reduce: [B,S1,N1,S2] -> [B,S1,N1]
106+ d_weight = mul_relu.sum(axis=-1, keepdims=False).to(q_dtype)
107+ 
108+ # WEIGHT: [B,S1,N1] -> [B,N1,S1,1]
109+ weight = input_weight.permute(0, 2, 1).contiguous().unsqueeze(-1).to(torch.float32)
110+ # mul: (B,1,S1,S2) * (B,N1,S1,1) = (B,N1,S1,S2)
111+ mul_weight = sub_result.unsqueeze(1) * weight
112+ # relu grad: (x > 0) = 1; (x <= 0) = 0
113+ relu_grad = mul_weight * (relu_res > 0).float()
114+ # cast for matmul
115+ relu_grad = relu_grad.to(q_dtype)
116+ 
117+ # KEY_INDEX: [B,S2,N2,D] -> [B,N1,S2,D]
118+ key_index = input_key_index.permute(0, 2, 1, 3).contiguous().repeat(1, group_size, 1, 1)
119+ # batch matmul (B,N1,S1,S2) @ (B,N1,S2,D) = (B,N1,S1,D)
120+ d_query_index = torch.matmul(relu_grad.to(torch.float32), key_index.to(torch.float32))
121+ d_query_index = d_query_index.to(q_dtype)
122+ # D_QUERY_INDEX: [B,N1,S1,D] -> [B,S1,N1,D]
123+ d_query_index = d_query_index.permute(0, 2, 1, 3).contiguous()
124+ 
125+ # QUERY_INDEX: [B,S1,N1,D] -> [B,N1,S1,D]
126+ query_index = input_query_index.permute(0,2,1,3).contiguous()
127+ # RELU_GRAD: [B,N1,S1,S2] -> [B,N1,S2,S1]
128+ relu_grad = relu_grad.permute(0, 1, 3, 2).contiguous()
129+ # batch matmul: (B,N1,S2,S1) @ (B,N1,S1,D) = (B,N1,S2,D)
130+ d_key_index_tmp = torch.matmul(relu_grad.to(torch.float32), query_index.to(torch.float32))
131+ # reshape: [B,N1,S2,D] -> [B,G,N2,S2,D]
132+ d_key_index = d_key_index_tmp.reshape(-1, group_size, N2, d_key_index_tmp.shape[2], D)
133+ # reduce on g_size: [B,G,N2,S2,D] -> [B,N2,S2,D]
134+ d_key_index = d_key_index.sum(axis=1, keepdims=False)
135+ # permute: [B,N2,S2,D] -> [B,S2,N2,D]
136+ d_key_index = d_key_index.permute(0, 2, 1, 3).contiguous().to(q_dtype)
137+ 
138+ return d_weight, d_query_index, d_key_index
139+ 
140+ 
141+ 
142+ def _dense_lightning_indexer_grad_kl_loss(self, query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, scale, query_rope, key_rope):
143+
144+ query = torch.cat((query, query_rope), dim=-1)
145+ key = torch.cat((key, key_rope), dim=-1)
146+ # process P
147+ p_result, softmax_max_result, softmax_sum_result = self._process_p(query, key, softmax_max, softmax_sum, scale)
148+ 
149+ # process S'Y
150+ sy_result, relu_res, softmax_max_index_result, softmax_sum_index_result = self._process_sy(query_index, key_index, weights, softmax_max_index, softmax_sum_index)
151+ # process kl loss
152+ loss = self._process_kl_loss(p_result, sy_result)
153+ # process dw/dq/dk
154+ q_dtype = query.dtype
155+ d_weight, d_query_index, d_key_index = self._process_dwqk(p_result, sy_result, relu_res, query_index, key_index, weights, q_dtype)
156+ 
157+ return d_query_index, d_key_index, d_weight, loss
158+
159+ def _get_input(self, layout="BSND"):
160+ torch.manual_seed(0)
161+ np.random.seed(0)
162+ 
163+ q_dtype = torch.float16
164+ B, N1, N2, N1_index, N2_index, S1, S2, D, Dr = 1, 64, 64, 64, 1, 128, 256, 128, 64
165+ query = torch.randn(B, S1, N1, D, dtype=q_dtype)
166+ key = torch.randn(B, S2, N2, D, dtype=q_dtype)
167+ query_index = torch.randn(B, S1, N1_index, D, dtype=q_dtype)
168+ key_index = torch.randn(B, S2, N2_index, D, dtype=q_dtype)
169+ query_rope = torch.randn(B, S1, N1, Dr, dtype=q_dtype)
170+ key_rope = torch.randn(B, S2, N2, Dr, dtype=q_dtype)
171+ weights = torch.randn(B, S1, N1_index, dtype=q_dtype)
172+ softmax_max = (torch.randn(B, N2, S1, 1, dtype=torch.float32).abs() + 0.4) * D # N1=N2
173+ softmax_sum = torch.ones(B, N2, S1, 1, dtype=torch.float32)
174+ softmax_max_index = (torch.randn(B, 1, S1, dtype=torch.float32).abs() + 0.4) * D * N1_index
175+ softmax_sum_index = torch.ones(B, 1, S1, dtype=torch.float32)
176+ actual_seq_qlen = [S1]
177+ actual_seq_klen = [S2]
178+ input_list = [query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, query_rope, key_rope, actual_seq_qlen, actual_seq_klen]
179+ 
180+ return input_list
181+ 
182+ 
183+ @unittest.skip("Skipping due to outdated CANN version; please update CANN to the latest version and remove this skip")
184+ def test_dense_lightning_indexer_grad_kl_loss_eager(self):
185+ sparse_mode = 3
186+ scale = 1.0
187+ layout = 'BSND'
188+ input_list = self._get_input(layout)
189+ query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, query_rope, key_rope, actual_seq_qlen, actual_seq_klen = input_list
190+ cpu_out = self._dense_lightning_indexer_grad_kl_loss(query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index,
191+ softmax_sum_index, scale, query_rope, key_rope)
192+ 
193+ for i in range(len(input_list)):
194+ try:
195+ input_list[i] = input_list[i].npu()
196+ except:
197+ continue
198+ query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, query_rope, key_rope, actual_seq_qlen, actual_seq_klen = input_list
199+ 
200+ npu_out = torch_npu.npu_dense_lightning_indexer_grad_kl_loss(query, key, query_index, key_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, scale,
201+ query_rope=query_rope, key_rope=key_rope, actual_seq_qlen=actual_seq_qlen, actual_seq_klen=actual_seq_klen, layout=layout, sparse_mode=sparse_mode, pre_tokens=65536, next_tokens=65536)
202+ 
203+ # compare result
204+ dq_cpu, dk_cpu, dw_cpu, loss_cpu = cpu_out
205+ dq_npu, dk_npu, dw_npu, loss_npu = npu_out
206+ self.assertRtolEqual(dq_npu, dq_cpu, prec=1.e-3)
207+ self.assertRtolEqual(dk_npu, dk_cpu, prec=1.e-3)
208+ self.assertRtolEqual(dw_npu, dw_cpu, prec=1.e-3)
209+ self.assertRtolEqual(loss_npu[0], loss_cpu, prec=1.e-3)
210+ 
211+ 
212+if __name__ == "__main__":
213+ run_tests()
Atest/test_custom_ops/test_npu_dynamic_dual_level_mx_quant.py+48-0
@@ -0,0 +1,48 @@
1+import math
2+import unittest
3+import copy
4+import struct
5+from struct import pack, unpack
6+import numpy as np
7+import torch
8+import torch_npu
9+from torch_npu.testing.testcase import TestCase, run_tests
10+from torch_npu.testing.common_utils import SupportedDevices
11+from torch.testing import assert_close
12+ 
13+ 
14+class TestDynamicDualLevelMxQuant(TestCase):
15+ def custom_op_exec(self, input_tensor, smooth_scale=None, round_mode="rint"):
16+ return torch_npu.npu_dynamic_dual_level_mx_quant(input_tensor, smooth_scale=smooth_scale, round_mode=round_mode)
17+ 
18+ def supported_op_exec(self, input_tensor):
19+ if torch.all(torch.eq(input_tensor, 8.0)) and input_tensor.shape == torch.Size([1, 512]):
20+ device = input_tensor.device
21+ y = torch.full((1, 256), 119, dtype=torch.uint8, device=device)
22+ level0_scale = torch.tensor([[1.333333]], dtype=torch.float32, device=device)
23+ level1_scale = torch.full((1, 8, 2), 127, dtype=torch.uint8, device=device)
24+ return y, level0_scale, level1_scale
25+ 
26+ def generate_input(self, input, dtype="float16"):
27+ data_type = torch.float16 if dtype == "float16" else torch.bfloat16
28+ value = 8.0
29+ input_tensor = torch.full(input, value, dtype=data_type)
30+ return input_tensor
31+
32+ @SupportedDevices(['Ascend910_95'])
33+ def test_npu_dynamic_dual_level_mx_quant(self, device="npu"):
34+ input_tensor = self.generate_input(input=[1, 512], dtype="bfloat16")
35+ input_tensor = input_tensor.to(device)
36+ supported_output = self.supported_op_exec(input_tensor.clone())
37+ custom_output = self.custom_op_exec(input_tensor.clone(), None, "rint")
38+ 
39+ y = custom_output[0].view([1, 256]).view(torch.uint8)
40+ level0_scale = custom_output[1].view([1, 1]).view(torch.float32)
41+ level1_scale = custom_output[2].view([1, 8, 2]).view(torch.uint8)
42+ 
43+ assert torch.all(y == supported_output[0].view(torch.uint8))
44+ assert_close(supported_output[1], level0_scale, atol=0.01, rtol=0.001)
45+ assert_close(supported_output[2], level1_scale, atol=0.01, rtol=0.001)
46+
47+if __name__ == "__main__":
48+ run_tests()
Mtest/test_custom_ops/test_npu_quant_reduce_scatter.py+3-3文件内容审核中，请稍后刷新重试
Mtest/test_custom_ops/test_save_npugraph_tensor.py+34-0文件内容审核中，请稍后刷新重试
Rcodegen/__init__.py→torchnpugen/__init__.py+0-0
文件重命名但无更改。
Rcodegen/api/__init__.py→torchnpugen/api/__init__.py+0-0
文件重命名但无更改。
Rcodegen/api/cpp.py→torchnpugen/api/cpp.py+4-4
@@ -18,8 +18,8 @@ from typing import List, Optional, Sequence, Set, Union
18from distutils.version import LooseVersion18from distutils.version import LooseVersion
19import os19import os
20 20 
21-from codegen import local21+from torchnpugen import local
22-from codegen.api.types import (22+from torchnpugen.api.types import (
23 ArgName,23 ArgName,
24 ArrayCType,24 ArrayCType,
25 ArrayRefCType,25 ArrayRefCType,
@@ -50,7 +50,7 @@ from codegen.api.types import (
50 VectorCType,50 VectorCType,
51 voidT,51 voidT,
52)52)
53-from codegen.model import (53+from torchnpugen.model import (
54 Argument,54 Argument,
55 Arguments,55 Arguments,
56 BaseTy,56 BaseTy,
@@ -64,7 +64,7 @@ from codegen.model import (
64 TensorOptionsArguments,64 TensorOptionsArguments,
65 Type,65 Type,
66)66)
67-from codegen.utils import assert_never67+from torchnpugen.op_codegen_utils import assert_never
68 68 
69# This file describes the translation of JIT schema to the public C++69# This file describes the translation of JIT schema to the public C++
70# API, which is what people use when they call functions like at::add.70# API, which is what people use when they call functions like at::add.
Rcodegen/api/native.py→torchnpugen/api/native.py+5-5
@@ -16,9 +16,9 @@
16 16 
17from typing import List, Optional, Sequence, Union17from typing import List, Optional, Sequence, Union
18 18 
19-from codegen import local19+from torchnpugen import local
20-from codegen.api import cpp20+from torchnpugen.api import cpp
21-from codegen.api.types import (21+from torchnpugen.api.types import (
22 ArgName,22 ArgName,
23 BaseCType,23 BaseCType,
24 Binding,24 Binding,
@@ -35,7 +35,7 @@ from codegen.api.types import (
35 scalarTypeT,35 scalarTypeT,
36 tensorT,36 tensorT,
37)37)
38-from codegen.model import (38+from torchnpugen.model import (
39 Argument,39 Argument,
40 FunctionSchema,40 FunctionSchema,
41 Return,41 Return,
@@ -43,7 +43,7 @@ from codegen.model import (
43 TensorOptionsArguments,43 TensorOptionsArguments,
44 Type,44 Type,
45)45)
46-from codegen.utils import assert_never46+from torchnpugen.op_codegen_utils import assert_never
47 47 
48# This file describes the translation of JIT schema to the native functions API.48# This file describes the translation of JIT schema to the native functions API.
49# This looks a lot like the C++ API (which makes historical sense, because the49# This looks a lot like the C++ API (which makes historical sense, because the
Rcodegen/api/types/__init__.py→torchnpugen/api/types/__init__.py+0-0
文件重命名但无更改。
Rcodegen/api/types/signatures.py→torchnpugen/api/types/signatures.py+2-2
@@ -17,7 +17,7 @@
17from dataclasses import dataclass17from dataclasses import dataclass
18from typing import List, Optional18from typing import List, Optional
19 19 
20-from codegen.model import FunctionSchema, NativeFunction20+from torchnpugen.model import FunctionSchema, NativeFunction
21from .types_base import Binding, CType21from .types_base import Binding, CType
22 22 
23 23 
@@ -65,4 +65,4 @@ def kernel_signature(
65 65 
66 66 
67# Functions only, no types67# Functions only, no types
68-from codegen.api import native68+from torchnpugen.api import native
Rcodegen/api/types/types.py→torchnpugen/api/types/types.py+1-1
@@ -31,7 +31,7 @@ Add new types to `types_base.py` if they are basic and not attached to ATen/c10.
31from dataclasses import dataclass31from dataclasses import dataclass
32from typing import Dict, TypeVar32from typing import Dict, TypeVar
33 33 
34-from codegen.model import BaseTy, ScalarType34+from torchnpugen.model import BaseTy, ScalarType
35 35 
36from .types_base import (36from .types_base import (
37 BaseCppType,37 BaseCppType,
Rcodegen/api/types/types_base.py→torchnpugen/api/types/types_base.py+2-2
@@ -33,7 +33,7 @@ from dataclasses import dataclass
33from enum import auto, Enum33from enum import auto, Enum
34from typing import List, Optional, Union34from typing import List, Optional, Union
35 35 
36-from codegen.model import Argument, SelfArgument, TensorOptionsArguments36+from torchnpugen.model import Argument, SelfArgument, TensorOptionsArguments
37 37 
38# An ArgName is just the str name of the argument in schema;38# An ArgName is just the str name of the argument in schema;
39# but in some special circumstances, we may add a little extra39# but in some special circumstances, we may add a little extra
@@ -60,7 +60,7 @@ class BaseCppType:
60 return f"{self.ns}::{self.name}"60 return f"{self.ns}::{self.name}"
61 61 
62 62 
63-# The set of all non-templated, valid, fully-qualified names of C++ types that are used in the codegen.63+# The set of all non-templated, valid, fully-qualified names of C++ types that are used in the torchnpugen.
64# Templated types get their own dataclass, mainly to make namespace parsing easier.64# Templated types get their own dataclass, mainly to make namespace parsing easier.
65byteT = BaseCppType("", "uint8_t")65byteT = BaseCppType("", "uint8_t")
66charT = BaseCppType("", "int8_t")66charT = BaseCppType("", "int8_t")
Rcodegen/code_template.py→torchnpugen/code_template.py+0-0
文件重命名但无更改。
Rcodegen/context.py→torchnpugen/context.py+3-3
@@ -18,9 +18,9 @@ import functools
18import contextlib18import contextlib
19from typing import TypeVar, Union, Iterator, Callable, Dict19from typing import TypeVar, Union, Iterator, Callable, Dict
20 20 
21-from codegen.utils import S, T, context21+from torchnpugen.op_codegen_utils import S, T, context
22-from codegen.model import (NativeFunction, NativeFunctionsGroup, BackendIndex, DispatchKey)22+from torchnpugen.model import (NativeFunction, NativeFunctionsGroup, BackendIndex, DispatchKey)
23-import codegen.local as local23+import torchnpugen.local as local
24 24 
25# Helper functions for defining generators on things in the model25# Helper functions for defining generators on things in the model
26 26 
Rcodegen/gen.py→torchnpugen/gen.py+6-6
@@ -23,13 +23,13 @@ from typing import (List, Dict, Optional, Set, Callable, Any,
23from collections import defaultdict23from collections import defaultdict
24import yaml24import yaml
25 25 
26-from codegen.code_template import CodeTemplate26+from torchnpugen.code_template import CodeTemplate
27-from codegen.model import (NativeFunction, SelfArgument,27+from torchnpugen.model import (NativeFunction, SelfArgument,
28 TensorOptionsArguments,28 TensorOptionsArguments,
29 assert_never)29 assert_never)
30-from codegen.api.types.signatures import NativeSignature30+from torchnpugen.api.types.signatures import NativeSignature
31-from codegen.context import native_function_manager31+from torchnpugen.context import native_function_manager
32-from codegen.utils import concatMap, context32+from torchnpugen.op_codegen_utils import concatMap, context
33 33 
34 34 
35 35 
@@ -99,7 +99,7 @@ class FileManager:
99 env = env_callable()99 env = env_callable()
100 if isinstance(env, dict):100 if isinstance(env, dict):
101 if 'generated_comment' not in env:101 if 'generated_comment' not in env:
102- comment = "@" + "generated by tools/codegen/gen.py"102+ comment = "@" + "generated by tools/torchnpugen/gen.py"
103 comment += " from {}".format(os.path.basename(template_fn))103 comment += " from {}".format(os.path.basename(template_fn))
104 env['generated_comment'] = comment104 env['generated_comment'] = comment
105 env['legacy_th_headers'] = []105 env['legacy_th_headers'] = []
Rcodegen/gen_derivatives.py→torchnpugen/gen_derivatives.py+1-1
@@ -4,7 +4,7 @@ import os
4import stat4import stat
5import yaml5import yaml
6 6 
7-from codegen.utils import PathManager, get_version7+from torchnpugen.op_codegen_utils import PathManager, get_version
8 8 
9 9 
10def main():10def main():
Rcodegen/gen_backend_stubs.py→torchnpugen/gen_op_backend.py+3-3文件内容审核中，请稍后刷新重试
Rcodegen/gen_op_plugin_functions.py→torchnpugen/gen_op_plugin_functions.py+1-1
@@ -4,7 +4,7 @@ import os
4import stat4import stat
5import yaml5import yaml
6 6 
7-from codegen.utils import PathManager, get_version7+from torchnpugen.op_codegen_utils import PathManager, get_version
8 8 
9 9 
10def main():10def main():
Rcodegen/local.py→torchnpugen/local.py+1-1
@@ -25,7 +25,7 @@ from typing import Iterator, Optional
25# why you need to add a toggle to the global behavior of code25# why you need to add a toggle to the global behavior of code
26# generation. The parameters here should really only be used26# generation. The parameters here should really only be used
27# for "temporary" situations, where we need to temporarily change27# for "temporary" situations, where we need to temporarily change
28-# the codegen in some cases because we cannot conveniently update28+# the torchnpugen in some cases because we cannot conveniently update
29# all call sites, and are slated to be eliminated once all call29# all call sites, and are slated to be eliminated once all call
30# sites are eliminated. If you don't have a plan for how to get there,30# sites are eliminated. If you don't have a plan for how to get there,
31# DON'T add a new entry here.31# DON'T add a new entry here.
Rcodegen/model.py→torchnpugen/model.py+2-2
@@ -22,7 +22,7 @@ from dataclasses import dataclass
22from enum import auto, Enum22from enum import auto, Enum
23from typing import Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union23from typing import Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Union
24 24 
25-from codegen.utils import assert_never, NamespaceHelper, OrderedSet25+from torchnpugen.op_codegen_utils import assert_never, NamespaceHelper, OrderedSet
26 26 
27# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #27# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
28#28#
@@ -483,7 +483,7 @@ class NativeFunction:
483 # # in native_functions.yaml483 # # in native_functions.yaml
484 # autogen: List["OperatorName"]484 # autogen: List["OperatorName"]
485 485 
486- # # If non-empty, this kernel is subject to ufunc codegen.486+ # # If non-empty, this kernel is subject to ufunc torchnpugen.
487 # # Sorted by ufunc_key487 # # Sorted by ufunc_key
488 # ufunc_inner_loop: Dict[UfuncKey, "UfuncInnerLoop"]488 # ufunc_inner_loop: Dict[UfuncKey, "UfuncInnerLoop"]
489 489 
Rcodegen/utils.py→torchnpugen/op_codegen_utils.py+0-0
文件重命名但无更改。
Rcodegen/struct/__init__.py→torchnpugen/struct/__init__.py+0-0
文件重命名但无更改。
Rcodegen/struct/gen_struct_opapi.py→torchnpugen/struct/gen_struct_opapi.py+4-4文件内容审核中，请稍后刷新重试
Rcodegen/struct/model.py→torchnpugen/struct/model.py+4-4
@@ -18,11 +18,11 @@ from typing import List, Dict, Sequence
18import copy18import copy
19from dataclasses import dataclass19from dataclasses import dataclass
20 20 
21-from codegen.model import (BaseTy, SchemaKind, BaseType,21+from torchnpugen.model import (BaseTy, SchemaKind, BaseType,
22 Argument, NativeFunction, ListType)22 Argument, NativeFunction, ListType)
23-from codegen.context import native_function_manager23+from torchnpugen.context import native_function_manager
24-from codegen.api.types import NativeSignature24+from torchnpugen.api.types import NativeSignature
25-from codegen.api import cpp25+from torchnpugen.api import cpp
26 26 
27 27 
28def filt_input_tensor(arguments: Sequence[Argument]) -> List[str]:28def filt_input_tensor(arguments: Sequence[Argument]) -> List[str]:
Rcodegen/struct/struct_codegen.py→torchnpugen/struct/struct_codegen.py+7-7文件内容审核中，请稍后刷新重试
Rcodegen/struct/templates/StructKernelNpuOpApi.cpp→torchnpugen/struct/templates/StructKernelNpuOpApi.cpp+0-0
文件重命名但无更改。
Rcodegen/templates/Interface.h→torchnpugen/templates/Interface.h+0-0
文件重命名但无更改。
Rcodegen/templates/OpInterface.cpp→torchnpugen/templates/OpInterface.cpp+0-0
文件重命名但无更改。