Star186
854
代码
代码
Issues178
Pull Requests333
项目讨论
Wiki
分析
Star186
854
已合并
[master][Fix] Fix static check errors detected by clang-format #5599
thickhair创建于 7月29日
[master][Fix] Fix static check errors detected by clang-format #5599
已合并
thickhair创建于 7月29日
从thickhair/op-pluginclang-format-12合入到Ascend/op-pluginmaster
共 52 个文件变更+3446-2998
Mop_plugin/ops/opapi/ScaledMaskedSoftmaxKernelNpuOpApi.cpp+39-36
@@ -19,25 +19,27 @@
19namespace op_api {19namespace op_api {
20using npu_preparation = at_npu::native::OpPreparation;20using npu_preparation = at_npu::native::OpPreparation;
21 21 
22-at::Tensor npu_scaled_masked_softmax(const at::Tensor& x, const at::Tensor& mask, const at::Scalar& scale,22+at::Tensor npu_scaled_masked_softmax(
23- bool fixed_triu_mask)23+ const at::Tensor& x,
24-{24+ const at::Tensor& mask,
25- if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {25+ const at::Scalar& scale,
26- return acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask);26+ bool fixed_triu_mask) {
27- }27+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {
28- DO_COMPATIBILITY(aclnnScaledMaskedSoftmax, acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask));28+ return acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask);
29- double scale_value = 1.0;29+ }
30- if (scale.isFloatingPoint()) {30+ DO_COMPATIBILITY(aclnnScaledMaskedSoftmax, acl_op::npu_scaled_masked_softmax(x, mask, scale, fixed_triu_mask));
31- scale_value = scale.to<double>();31+ double scale_value = 1.0;
32- } else if (scale.isIntegral(true)) {32+ if (scale.isFloatingPoint()) {
33- scale_value = static_cast<double>(scale.to<int64_t>());33+ scale_value = scale.to<double>();
34- } else {34+ } else if (scale.isIntegral(true)) {
35- TORCH_CHECK(false, "scaled_masked_softmax expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));35+ scale_value = static_cast<double>(scale.to<int64_t>());
36- }36+ } else {
37+ TORCH_CHECK(false, "scaled_masked_softmax expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));
38+ }
37 39 
38- at::Tensor result = npu_preparation::apply_tensor_without_format(x.sizes(), x.options());40+ 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);41+ EXEC_NPU_CMD(aclnnScaledMaskedSoftmax, x, mask, scale_value, fixed_triu_mask, result);
40- return result;42+ return result;
41}43}
42 44 
43at::Tensor npu_scaled_masked_softmax_backward(45at::Tensor npu_scaled_masked_softmax_backward(
@@ -45,23 +47,24 @@ at::Tensor npu_scaled_masked_softmax_backward(
45 const at::Tensor& y,47 const at::Tensor& y,
46 const at::Tensor& mask,48 const at::Tensor& mask,
47 const at::Scalar& scale,49 const at::Scalar& scale,
48- bool fixed_triu_mask)50+ bool fixed_triu_mask) {
49-{51+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {
50- if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {52+ return acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask);
51- return acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask);53+ }
52- }54+ DO_COMPATIBILITY(
53- DO_COMPATIBILITY(aclnnScaledMaskedSoftmaxBackward, acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask));55+ aclnnScaledMaskedSoftmaxBackward,
54- double scale_value = 1.0;56+ acl_op::npu_scaled_masked_softmax_backward(y_grad, y, mask, scale, fixed_triu_mask));
55- if (scale.isFloatingPoint()) {57+ double scale_value = 1.0;
56- scale_value = scale.to<double>();58+ if (scale.isFloatingPoint()) {
57- } else if (scale.isIntegral(true)) {59+ scale_value = scale.to<double>();
58- scale_value = static_cast<double>(scale.to<int64_t>());60+ } else if (scale.isIntegral(true)) {
59- } else {61+ scale_value = static_cast<double>(scale.to<int64_t>());
60- TORCH_CHECK(false, "scaled_masked_softmax_backward expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));62+ } else {
61- }63+ TORCH_CHECK(false, "scaled_masked_softmax_backward expects scale to be float or int", OPS_ERROR(ErrCode::TYPE));
64+ }
62 65 
63- at::Tensor result = npu_preparation::apply_tensor_without_format(y_grad.sizes(), y_grad.options());66+ 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);67+ EXEC_NPU_CMD(aclnnScaledMaskedSoftmaxBackward, y_grad, y, mask, scale_value, fixed_triu_mask, result);
65- return result;68+ return result;
66}69}
67-}70+} // namespace op_api
Mop_plugin/ops/opapi/ScaledMmKernelNpuOpApi.cpp+441-355
@@ -18,7 +18,6 @@
18#include "op_plugin/utils/op_api_common.h"18#include "op_plugin/utils/op_api_common.h"
19#include "torch_npu/csrc/framework/utils/InternalFormatOpAdapter.h"19#include "torch_npu/csrc/framework/utils/InternalFormatOpAdapter.h"
20 20 
21- 
22namespace op_api {21namespace op_api {
23constexpr size_t LAST_SECOND_DIM_INDEX = 2;22constexpr size_t LAST_SECOND_DIM_INDEX = 2;
24constexpr int64_t PERGROUP_DIM_NUM = 2;23constexpr int64_t PERGROUP_DIM_NUM = 2;
@@ -28,414 +27,501 @@ static const size_t A8W4_GROUP_DIM = 3;
28static const size_t A8W4_INPUT_DIM = 2;27static const size_t A8W4_INPUT_DIM = 2;
29using npu_preparation = at_npu::native::OpPreparation;28using npu_preparation = at_npu::native::OpPreparation;
30 29 
31-bool static is_transpose_last_two_dims(const at::Tensor &tensor)30+bool static is_transpose_last_two_dims(const at::Tensor& tensor) {
32-{31+ if (tensor.dim() < 2 || tensor.dim() > 6) {
33- if (tensor.dim() < 2 || tensor.dim() > 6) {32+ return false;
33+ }
34+ int64_t dim1 = tensor.dim() - 1;
35+ int64_t dim2 = tensor.dim() - 2;
36+ if (tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2)) {
37+ int64_t tmpNxD = tensor.size(dim1) * tensor.size(dim2);
38+ for (int64_t batchDim = tensor.dim() - 3; batchDim >= 0; batchDim--) {
39+ if (tensor.stride(batchDim) != tmpNxD) {
34 return false;40 return false;
41+ }
42+ tmpNxD *= tensor.size(batchDim);
35 }43 }
36- int64_t dim1 = tensor.dim() - 1;44+ if (tensor.size(dim1) == 1 && tensor.size(dim2) == 1) {
37- int64_t dim2 = tensor.dim() - 2;45+ return false;
38- if (tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2)) {
39- int64_t tmpNxD = tensor.size(dim1) * tensor.size(dim2);
40- for (int64_t batchDim = tensor.dim() - 3; batchDim >= 0; batchDim--) {
41- if (tensor.stride(batchDim) != tmpNxD) {
42- return false;
43- }
44- tmpNxD *= tensor.size(batchDim);
45- }
46- if (tensor.size(dim1) == 1 && tensor.size(dim2) == 1) {
47- return false;
48- }
49- return true;
50 }46 }
51- return false;47+ return true;
48+ }
49+ return false;
52}50}
53 51 
54-static bool is_transpose_certain_two_dims(const at::Tensor &tensor, int64_t dim)52+static bool is_transpose_certain_two_dims(const at::Tensor& tensor, int64_t dim) {
55-{53+ return tensor.stride(dim + 1) == tensor.stride(dim) * tensor.size(dim);
56- return tensor.stride(dim + 1) == tensor.stride(dim) * tensor.size(dim);
57}54}
58 55 
59-static bool is_x_scale_same_transpose(const at::Tensor &x, const at::Tensor &scale, int64_t dim_x, int64_t dim_scale)56+static bool is_x_scale_same_transpose(const at::Tensor& x, const at::Tensor& scale, int64_t dim_x, int64_t dim_scale) {
60-{57+ if (x.dim() < dim_x + 2 || scale.dim() < dim_scale + 2) { // make sure dims after the start dim are no less than 2
61- if (x.dim() < dim_x + 2 || scale.dim() < dim_scale + 2) { // make sure dims after the start dim are no less than 258+ return true;
62- return true;59+ }
63- }60+ if (x.size(dim_x) == 1 && x.size(dim_x + 1) == 1) {
64- if (x.size(dim_x) == 1 && x.size(dim_x + 1)== 1) {61+ return true;
65- return true;62+ }
66- }63+ if (scale.size(dim_scale) == 1 && scale.size(dim_scale + 1) == 1) {
67- if (scale.size(dim_scale) == 1 && scale.size(dim_scale + 1)== 1) {64+ return true;
68- return true;65+ }
69- }66+ bool x_trans = is_transpose_certain_two_dims(x, dim_x);
70- bool x_trans = is_transpose_certain_two_dims(x, dim_x);67+ bool scale_trans = is_transpose_certain_two_dims(scale, dim_scale);
71- bool scale_trans = is_transpose_certain_two_dims(scale, dim_scale);68+ if (x_trans == scale_trans) {
72- if (x_trans == scale_trans) {69+ return true;
73- return true;70+ }
74- }71+ return false;
75- return false;
76}72}
77-static bool is_nz_format(const at::Tensor& x2)73+static bool is_nz_format(const at::Tensor& x2) {
78-{74+ const torch_npu::NPUStorageDesc& tensor_desc = torch_npu::NPUBridge::GetNpuStorageImpl(x2)->npu_desc_;
79- const torch_npu::NPUStorageDesc &tensor_desc =75+ return tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ || tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_4 ||
80- torch_npu::NPUBridge::GetNpuStorageImpl(x2)->npu_desc_;76+ tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_16;
81- return tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ ||
82- tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_4 ||
83- tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_16;
84}77}
85 78 
86-static uint64_t infer_out_batch_shape(const at::Tensor &x1, const at::Tensor &x2, std::vector<uint64_t> &batch_record)79+static uint64_t infer_out_batch_shape(const at::Tensor& x1, const at::Tensor& x2, std::vector<uint64_t>& batch_record) {
87-{80+ TORCH_CHECK(
88- TORCH_CHECK(at_npu::native::FormatHelper::IsBaseFormatType(x2) || is_nz_format(x2),81+ at_npu::native::FormatHelper::IsBaseFormatType(x2) || is_nz_format(x2),
89- "x2 should be in the original image format or nz format, but it is ",82+ "x2 should be in the original image format or nz format, but it is ",
90- npu_preparation::get_tensor_npu_format(x2), OPS_ERROR(ErrCode::PARAM));83+ npu_preparation::get_tensor_npu_format(x2),
91- uint64_t batch_val = 1;84+ OPS_ERROR(ErrCode::PARAM));
92- auto x1_dim_num = x1.dim();85+ uint64_t batch_val = 1;
93- auto x2_dim_num = x2.dim();86+ auto x1_dim_num = x1.dim();
94- auto out_dim_num = std::max(x1_dim_num, x2_dim_num);87+ auto x2_dim_num = x2.dim();
95- auto &shape_long = x1_dim_num > x2_dim_num ? x1 : x2;88+ auto out_dim_num = std::max(x1_dim_num, x2_dim_num);
96- auto &shape_short = x1_dim_num > x2_dim_num ? x2 : x1;89+ auto& shape_long = x1_dim_num > x2_dim_num ? x1 : x2;
97- int64_t vaild_offset = out_dim_num - std::min(x1_dim_num, x2_dim_num);90+ auto& shape_short = x1_dim_num > x2_dim_num ? x2 : x1;
98- for (int64_t i = 0; i < out_dim_num - LAST_SECOND_DIM_INDEX; i++) {91+ int64_t vaild_offset = out_dim_num - std::min(x1_dim_num, x2_dim_num);
99- auto short_dim = i < vaild_offset ? 1 : shape_short.size(i - vaild_offset);92+ for (int64_t i = 0; i < out_dim_num - LAST_SECOND_DIM_INDEX; i++) {
100- auto long_dim = shape_long.size(i);93+ auto short_dim = i < vaild_offset ? 1 : shape_short.size(i - vaild_offset);
101- TORCH_CHECK(!(short_dim > 1 && long_dim > 1 && short_dim != long_dim),94+ auto long_dim = shape_long.size(i);
102- "the x1 shape and x2 shape not supported for broadcast, the short_dim is ",95+ TORCH_CHECK(
103- short_dim, " and the long_dim is ", long_dim, OPS_ERROR(ErrCode::PARAM));96+ !(short_dim > 1 && long_dim > 1 && short_dim != long_dim),
104- uint64_t cur_batch_value = static_cast<uint64_t>(std::max(short_dim, long_dim));97+ "the x1 shape and x2 shape not supported for broadcast, the short_dim is ",
105- batch_val = batch_val * cur_batch_value;98+ short_dim,
106- batch_record.push_back(cur_batch_value);99+ " and the long_dim is ",
107- }100+ long_dim,
108- return batch_val;101+ OPS_ERROR(ErrCode::PARAM));
102+ uint64_t cur_batch_value = static_cast<uint64_t>(std::max(short_dim, long_dim));
103+ batch_val = batch_val * cur_batch_value;
104+ batch_record.push_back(cur_batch_value);
105+ }
106+ return batch_val;
109}107}
110 108 
111-static int64_t check_and_get_groups(at::IntArrayRef group_size_list)109+static int64_t check_and_get_groups(at::IntArrayRef group_size_list) {
112-{110+ int64_t groups = 0;
113- int64_t groups = 0;111+ if (group_size_list.empty()) {
114- if (group_size_list.empty()) {
115- return groups;
116- }
117- size_t group_dim = group_size_list.size();
118- TORCH_CHECK(group_dim == A8W4_GROUP_DIM, "group_sizes only support input with three elements, but got ",
119- group_dim, OPS_ERROR(ErrCode::PARAM));
120- int64_t group_m = static_cast<int64_t>(group_size_list[0]);
121- int64_t group_n = static_cast<int64_t>(group_size_list[1]);
122- int64_t group_k = static_cast<int64_t>(group_size_list[2]);
123- bool invalid_group_param = ((group_m <= GROUP_MAX && group_m >= 0)
124- && (group_n <= GROUP_MAX && group_n >= 0)
125- && (group_k <= GROUP_MAX && group_k >= 0));
126- TORCH_CHECK(invalid_group_param, "group param value must conform to range [0, 65535]", OPS_ERROR(ErrCode::VALUE));
127- groups = static_cast<int64_t>((static_cast<uint64_t>(group_m) << 32) + (static_cast<uint64_t>(group_n) << 16) +
128- (static_cast<uint64_t>(group_k)));
129 return groups;112 return groups;
113+ }
114+ size_t group_dim = group_size_list.size();
115+ TORCH_CHECK(
116+ group_dim == A8W4_GROUP_DIM,
117+ "group_sizes only support input with three elements, but got ",
118+ group_dim,
119+ OPS_ERROR(ErrCode::PARAM));
120+ int64_t group_m = static_cast<int64_t>(group_size_list[0]);
121+ int64_t group_n = static_cast<int64_t>(group_size_list[1]);
122+ int64_t group_k = static_cast<int64_t>(group_size_list[2]);
123+ bool invalid_group_param =
124+ ((group_m <= GROUP_MAX && group_m >= 0) && (group_n <= GROUP_MAX && group_n >= 0) &&
125+ (group_k <= GROUP_MAX && group_k >= 0));
126+ TORCH_CHECK(invalid_group_param, "group param value must conform to range [0, 65535]", OPS_ERROR(ErrCode::VALUE));
127+ groups = static_cast<int64_t>(
128+ (static_cast<uint64_t>(group_m) << 32) + (static_cast<uint64_t>(group_n) << 16) +
129+ (static_cast<uint64_t>(group_k)));
130+ return groups;
130}131}
131 132 
132static const std::map<c10::ScalarType, int64_t> SCALAR_TO_INT_TYPE_MAP = {133static const std::map<c10::ScalarType, int64_t> SCALAR_TO_INT_TYPE_MAP = {
133- {c10::ScalarType::Char, static_cast<int64_t>(at::kChar)}, /**< int8 */134+ {c10::ScalarType::Char, static_cast<int64_t>(at::kChar)}, /**< int8 */
134- {c10::ScalarType::Int, static_cast<int64_t>(at::kInt)}, /**< int32 */135+ {c10::ScalarType::Int, static_cast<int64_t>(at::kInt)}, /**< int32 */
135- {c10::ScalarType::BFloat16, static_cast<int64_t>(at::kBFloat16)}, /**< bfp16 */136+ {c10::ScalarType::BFloat16, static_cast<int64_t>(at::kBFloat16)}, /**< bfp16 */
136- {c10::ScalarType::Half, static_cast<int64_t>(at::kHalf)}, /**< fp16 */137+ {c10::ScalarType::Half, static_cast<int64_t>(at::kHalf)}, /**< fp16 */
137- {c10::ScalarType::Float, static_cast<int64_t>(at::kFloat)}, /**< fp32 */138+ {c10::ScalarType::Float, static_cast<int64_t>(at::kFloat)}, /**< fp32 */
138 139 
139- {c10::ScalarType::Float8_e4m3fn, static_cast<int64_t>(at::kFloat8_e4m3fn)}, /**< fp8e4m3 */140+ {c10::ScalarType::Float8_e4m3fn, static_cast<int64_t>(at::kFloat8_e4m3fn)}, /**< fp8e4m3 */
140- {c10::ScalarType::Float8_e5m2, static_cast<int64_t>(at::kFloat8_e5m2)}, /**< fp8e5m2 */141+ {c10::ScalarType::Float8_e5m2, static_cast<int64_t>(at::kFloat8_e5m2)}, /**< fp8e5m2 */
141- {c10::ScalarType::Byte, static_cast<int64_t>(c10_npu::DType::HIFLOAT8)} /**< hif8 */142+ {c10::ScalarType::Byte, static_cast<int64_t>(c10_npu::DType::HIFLOAT8)} /**< hif8 */
142};143};
143 144 
144-static c10::optional<int64_t> ToIntType(const std::optional<c10::ScalarType> &torchType) {145+static c10::optional<int64_t> ToIntType(const std::optional<c10::ScalarType>& torchType) {
145- c10::optional<int64_t> int_type = c10::nullopt;146+ c10::optional<int64_t> int_type = c10::nullopt;
146- if (torchType.has_value()) {147+ if (torchType.has_value()) {
147- const auto &it = SCALAR_TO_INT_TYPE_MAP.find(torchType.value());148+ const auto& it = SCALAR_TO_INT_TYPE_MAP.find(torchType.value());
148- if (it != SCALAR_TO_INT_TYPE_MAP.cend()) {149+ if (it != SCALAR_TO_INT_TYPE_MAP.cend()) {
149- int_type = c10::make_optional(it->second);150+ int_type = c10::make_optional(it->second);
150- }
151 }151 }
152- return int_type;152+ }
153+ return int_type;
153}154}
154 155 
155-at::Tensor _scaled_mm(const at::Tensor &mat_a,156+at::Tensor _scaled_mm(
156- const at::Tensor &mat_b,157+ const at::Tensor& mat_a,
157- const at::Tensor &scale_a,158+ const at::Tensor& mat_b,
158- const at::Tensor &scale_b,159+ const at::Tensor& scale_a,
160+ const at::Tensor& scale_b,
159 const std::optional<at::Tensor>& bias,161 const std::optional<at::Tensor>& bias,
160 const std::optional<at::Tensor>& scale_result, // 对result做scale，仅当output为float8时才有用，162 const std::optional<at::Tensor>& scale_result, // 对result做scale，仅当output为float8时才有用，
161- std::optional<c10::ScalarType> out_dtype, // 当前不支持float8163+ std::optional<c10::ScalarType> out_dtype, // 当前不支持float8
162 bool use_fast_accum)164 bool use_fast_accum)
163 165 
164{166{
167+ // check A5
168+ TORCH_CHECK(
169+ c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,
170+ "_scaled_mm is supported only on the Ascend950 platform and after.",
171+ OPS_ERROR(ErrCode::PARAM));
165 172 
166- // check A5173+ // Check data types: mat_a and mat_b must be float8 type
167- TORCH_CHECK(c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,174+ TORCH_CHECK(
168- "_scaled_mm is supported only on the Ascend950 platform and after.", OPS_ERROR(ErrCode::PARAM));175+ mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn || mat_a.scalar_type() == c10::ScalarType::Float8_e5m2,
176+ "mat_a must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ",
177+ mat_a.scalar_type(),
178+ OPS_ERROR(ErrCode::TYPE));
179+ TORCH_CHECK(
180+ mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn || mat_b.scalar_type() == c10::ScalarType::Float8_e5m2,
181+ "mat_b must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ",
182+ mat_b.scalar_type(),
183+ OPS_ERROR(ErrCode::TYPE));
169 184 
170- // Check data types: mat_a and mat_b must be float8 type185+ // Check multiplication of two Float8_e5m2 is not supported (reference: _scaled_mm_out_cuda)
171- TORCH_CHECK(mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn ||186+ TORCH_CHECK(
172- mat_a.scalar_type() == c10::ScalarType::Float8_e5m2,187+ mat_a.scalar_type() != c10::ScalarType::Float8_e5m2 || mat_b.scalar_type() != c10::ScalarType::Float8_e5m2,
173- "mat_a must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ", mat_a.scalar_type(),188+ "Multiplication of two Float8_e5m2 matrices is not supported",
174- OPS_ERROR(ErrCode::TYPE));189+ OPS_ERROR(ErrCode::TYPE));
175- TORCH_CHECK(mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn ||
176- mat_b.scalar_type() == c10::ScalarType::Float8_e5m2,
177- "mat_b must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ", mat_b.scalar_type(),
178- OPS_ERROR(ErrCode::TYPE));
179 190 
180- // Check multiplication of two Float8_e5m2 is not supported (reference: _scaled_mm_out_cuda)191+ // Check data types: scale_a and scale_b must be float32 or float8_e8m0 type
181- TORCH_CHECK(mat_a.scalar_type() != c10::ScalarType::Float8_e5m2 ||192+ TORCH_CHECK(
182- mat_b.scalar_type() != c10::ScalarType::Float8_e5m2,193+ scale_a.scalar_type() == c10::ScalarType::Float ||
183- "Multiplication of two Float8_e5m2 matrices is not supported",194+ scale_a.scalar_type() ==
184- OPS_ERROR(ErrCode::TYPE));195+ npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),
196+ "scale_a must be float32 or float8_e8m0 type, but got ",
197+ scale_a.scalar_type(),
198+ OPS_ERROR(ErrCode::TYPE));
199+ TORCH_CHECK(
200+ scale_b.scalar_type() == c10::ScalarType::Float ||
201+ scale_b.scalar_type() ==
202+ npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),
203+ "scale_b must be float32 or float8_e8m0 type, but got ",
204+ scale_b.scalar_type(),
205+ OPS_ERROR(ErrCode::TYPE));
185 206 
186- // Check data types: scale_a and scale_b must be float32 or float8_e8m0 type207+ // Check bias (reference: _scaled_mm_out_cuda)
187- TORCH_CHECK(scale_a.scalar_type() == c10::ScalarType::Float ||208+ if (bias.has_value()) {
188- scale_a.scalar_type() == npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),209+ TORCH_CHECK(
189- "scale_a must be float32 or float8_e8m0 type, but got ", scale_a.scalar_type(),210+ bias->numel() == mat_b.sizes()[1],
211+ "Bias must be size ",
212+ mat_b.sizes()[1],
213+ " but got ",
214+ bias->numel(),
215+ OPS_ERROR(ErrCode::PARAM));
216+ // Check out_dtype vs bias compatibility
217+ auto out_dtype_value = out_dtype.value_or(c10::ScalarType::BFloat16);
218+ TORCH_CHECK(
219+ out_dtype_value != c10::ScalarType::Float,
220+ "Bias is not supported when out_dtype is set to Float32",
190 OPS_ERROR(ErrCode::TYPE));221 OPS_ERROR(ErrCode::TYPE));
191- TORCH_CHECK(scale_b.scalar_type() == c10::ScalarType::Float ||222+ TORCH_CHECK(
192- scale_b.scalar_type() == npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),223+ bias->scalar_type() == c10::ScalarType::BFloat16 || bias->scalar_type() == c10::ScalarType::Half,
193- "scale_b must be float32 or float8_e8m0 type, but got ", scale_b.scalar_type(),224+ "Bias must be BFloat16 or Half, but got ",
225+ bias->scalar_type(),
194 OPS_ERROR(ErrCode::TYPE));226 OPS_ERROR(ErrCode::TYPE));
227+ TORCH_CHECK(
228+ (out_dtype_value != c10::ScalarType::Float && out_dtype_value != c10::ScalarType::BFloat16) ||
229+ bias->scalar_type() == c10::ScalarType::BFloat16,
230+ "Bias must be BFloat16 to compute ",
231+ out_dtype_value,
232+ " output, but got ",
233+ bias->scalar_type(),
234+ OPS_ERROR(ErrCode::TYPE));
235+ TORCH_CHECK(
236+ out_dtype_value != c10::ScalarType::Half || bias->scalar_type() == c10::ScalarType::Half,
237+ "Bias must be Float16 to compute ",
238+ out_dtype_value,
239+ " output, but got ",
240+ bias->scalar_type(),
241+ OPS_ERROR(ErrCode::TYPE));
242+ }
195 243 
196- // Check bias (reference: _scaled_mm_out_cuda)244+ // Check scale_result: currently only supports null/empty because output float8 is not supported
197- if (bias.has_value()) {245+ TORCH_CHECK(
198- TORCH_CHECK(bias->numel() == mat_b.sizes()[1],246+ !scale_result.has_value() || scale_result == c10::nullopt,
199- "Bias must be size ", mat_b.sizes()[1], " but got ", bias->numel(),247+ "scale_result is not supported currently, as output float8 type is not enabled. "
200- OPS_ERROR(ErrCode::PARAM));248+ "scale_result is only utilized when output is float8 type.",
201- // Check out_dtype vs bias compatibility249+ OPS_ERROR(ErrCode::NOT_SUPPORT));
202- auto out_dtype_value = out_dtype.value_or(c10::ScalarType::BFloat16);250+ 
203- TORCH_CHECK(out_dtype_value != c10::ScalarType::Float,251+ // Check out_dtype: currently only supports Float32, BFloat16, and Float16
204- "Bias is not supported when out_dtype is set to Float32",252+ if (out_dtype.has_value()) {
205- OPS_ERROR(ErrCode::TYPE));253+ TORCH_CHECK(
206- TORCH_CHECK(bias->scalar_type() == c10::ScalarType::BFloat16 ||254+ out_dtype.value() == c10::ScalarType::Float || out_dtype.value() == c10::ScalarType::BFloat16 ||
207- bias->scalar_type() == c10::ScalarType::Half,255+ out_dtype.value() == c10::ScalarType::Half,
208- "Bias must be BFloat16 or Half, but got ", bias->scalar_type(),256+ "out_dtype must be Float32, BFloat16, or Float16, but got ",
209- OPS_ERROR(ErrCode::TYPE));257+ out_dtype.value(),
210- TORCH_CHECK((out_dtype_value != c10::ScalarType::Float &&258+ OPS_ERROR(ErrCode::TYPE));
211- out_dtype_value != c10::ScalarType::BFloat16) ||259+ }
212- bias->scalar_type() == c10::ScalarType::BFloat16,260+ 
213- "Bias must be BFloat16 to compute ", out_dtype_value,261+ // Check sizes
214- " output, but got ", bias->scalar_type(),262+ TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a matrix, please check mat_a dim num.", OPS_ERROR(ErrCode::PARAM));
215- OPS_ERROR(ErrCode::TYPE));263+ TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a matrix, please check mat_b dim num.", OPS_ERROR(ErrCode::PARAM));
216- TORCH_CHECK(out_dtype_value != c10::ScalarType::Half ||264+ TORCH_CHECK(
217- bias->scalar_type() == c10::ScalarType::Half,265+ mat_a.sizes()[1] == mat_b.sizes()[0],
218- "Bias must be Float16 to compute ", out_dtype_value,266+ "mat_a and mat_b shapes cannot be multiplied (",
219- " output, but got ", bias->scalar_type(),267+ mat_a.sizes()[0],
220- OPS_ERROR(ErrCode::TYPE));268+ "x",
269+ mat_b.sizes()[1],
270+ " and ",
271+ mat_b.sizes()[0],
272+ "x",
273+ mat_b.sizes()[1],
274+ ")",
275+ OPS_ERROR(ErrCode::PARAM));
276+ 
277+ //////////////rowwise transfer//////////////
278+ at::Tensor processed_scale_a = scale_a;
279+ at::Tensor processed_scale_b = scale_b;
280+ if (mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn && mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn &&
281+ scale_a.scalar_type() == c10::ScalarType::Float && scale_b.scalar_type() == c10::ScalarType::Float) {
282+ if (scale_a.dim() == 2 && scale_a.sizes()[1] == 1 && scale_b.dim() == 2 && scale_b.sizes()[0] == 1) {
283+ TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be contiguous in last dim");
284+ TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be contiguous in first dim");
285+ processed_scale_a = scale_a.squeeze(-1).contiguous();
286+ processed_scale_b = scale_b.squeeze(0);
221 }287 }
288+ }
222 289 
290+ //////////////parameters transfer//////////////
291+ at::Tensor x1 = mat_a;
292+ at::Tensor x2 = mat_b;
293+ c10::optional<at::Tensor> pertoken_scale = processed_scale_a;
294+ at::Tensor scale = processed_scale_b;
295+ c10::optional<int64_t> scale_dtype = c10::nullopt; // scale_b.scalar_type()
223 296 
224- // Check scale_result: currently only supports null/empty because output float8 is not supported297+ c10::optional<int64_t> output_dtype = ToIntType(out_dtype);
225- TORCH_CHECK(!scale_result.has_value() || scale_result == c10::nullopt,298+ c10::optional<at::Tensor> offset = c10::nullopt;
226- "scale_result is not supported currently, as output float8 type is not enabled. "299+ c10::optional<int64_t> x1_dtype = c10::nullopt;
227- "scale_result is only utilized when output is float8 type.",300+ c10::optional<int64_t> x2_dtype = c10::nullopt;
228- OPS_ERROR(ErrCode::NOT_SUPPORT));301+ c10::optional<int64_t> pertoken_scale_dtype = c10::nullopt;
302+ c10::OptionalIntArrayRef group_sizes = c10::nullopt;
303+ c10::optional<at::Tensor> y_scale = c10::nullopt;
229 304 
230- // Check out_dtype: currently only supports Float32, BFloat16, and Float16305+ ///////////////////////////////////////////npu_quant_matmul//////////////////////////////////////
231- if (out_dtype.has_value()) {306+ if (is_nz_format(x2)) {
232- TORCH_CHECK(out_dtype.value() == c10::ScalarType::Float ||307+ static const bool is_quant_matmul_weight_nz_available = check_aclnn_kernel_available("aclnnQuantMatmulWeightNz");
233- out_dtype.value() == c10::ScalarType::BFloat16 ||308+ TORCH_CHECK(
234- out_dtype.value() == c10::ScalarType::Half,309+ is_quant_matmul_weight_nz_available,
235- "out_dtype must be Float32, BFloat16, or Float16, but got ", out_dtype.value(),310+ "Get aclnnQuantMatmulWeightNz or aclnnQuantMatmulWeightNzGetWorkspaceSize failed, only "
236- OPS_ERROR(ErrCode::TYPE));311+ "aclnnQuantMatmulWeightNz support X2's format is nz, please upgrade CANN.",
237- }312+ OPS_ERROR(ErrCode::PARAM));
238- 313+ } else {
239- 314+ static const bool is_quant_matmul_v5_available = check_aclnn_kernel_available("aclnnQuantMatmulV5");
240- // Check sizes315+ TORCH_CHECK(
241- TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a matrix, please check mat_a dim num." ,OPS_ERROR(ErrCode::PARAM));316+ is_quant_matmul_v5_available,
242- TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a matrix, please check mat_b dim num." ,OPS_ERROR(ErrCode::PARAM));317+ "Get aclnnQuantMatmulV5 or aclnnQuantMatmulV5 failed, only "
243- TORCH_CHECK(mat_a.sizes()[1] == mat_b.sizes()[0], "mat_a and mat_b shapes cannot be multiplied (",mat_a.sizes()[0],318+ "aclnnQuantMatmulV5 support A8W4, please upgrade CANN.",
244- "x", mat_b.sizes()[1], " and ", mat_b.sizes()[0], "x", mat_b.sizes()[1], ")", OPS_ERROR(ErrCode::PARAM));319+ OPS_ERROR(ErrCode::TYPE));
245- 320+ }
246- 321+ bool is_a8W4_int = x1.dtype() == at::kChar && x2.dtype() == at::kInt;
247- //////////////rowwise transfer//////////////322+ bool is_a8W4_float = x1.dtype() == at::kFloat8_e4m3fn && x2.dtype() == at::kFloat;
248- at::Tensor processed_scale_a = scale_a;323+ at::IntArrayRef group_size_list = group_sizes.value_or(at::IntArrayRef{});
249- at::Tensor processed_scale_b = scale_b;324+ int64_t group_size = check_and_get_groups(group_size_list);
250- if (mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn325+ bool is_a4w4 = x1.dtype() == at::kInt && x2.dtype() == at::kInt;
251- && mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn326+ bool trans_x1 = is_transpose_last_two_dims(x1);
252- && scale_a.scalar_type() == c10::ScalarType::Float327+ bool trans_x2 = is_transpose_last_two_dims(x2);
253- && scale_b.scalar_type() == c10::ScalarType::Float) {328+ auto x1_dim_num = x1.dim();
254- if (scale_a.dim() == 2 && scale_a.sizes()[1] == 1 && scale_b.dim() == 2 && scale_b.sizes()[0] == 1) {329+ auto x2_dim_num = x2.dim();
255- TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be contiguous in last dim");330+ auto x2_n_dim = (is_a4w4 && !trans_x2) ? x2.size(x2_dim_num - 1) * INT4_NUMS_IN_INT32 : x2.size(x2_dim_num - 1);
256- TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be contiguous in first dim");
257- processed_scale_a = scale_a.squeeze(-1).contiguous();
258- processed_scale_b = scale_b.squeeze(0);
259- }
260- }
261- 
262- //////////////parameters transfer//////////////
263- at::Tensor x1 = mat_a;
264- at::Tensor x2 = mat_b;
265- c10::optional<at::Tensor> pertoken_scale = processed_scale_a;
266- at::Tensor scale = processed_scale_b;
267- c10::optional<int64_t> scale_dtype = c10::nullopt; // scale_b.scalar_type()
268- 
269- c10::optional<int64_t> output_dtype = ToIntType(out_dtype);
270- c10::optional<at::Tensor> offset = c10::nullopt;
271- c10::optional<int64_t> x1_dtype = c10::nullopt;
272- c10::optional<int64_t> x2_dtype = c10::nullopt;
273- c10::optional<int64_t> pertoken_scale_dtype = c10::nullopt;
274- c10::OptionalIntArrayRef group_sizes = c10::nullopt;
275- c10::optional<at::Tensor> y_scale = c10::nullopt;
276- 
277- 
278-///////////////////////////////////////////npu_quant_matmul//////////////////////////////////////
279- if (is_nz_format(x2)) {
280- static const bool is_quant_matmul_weight_nz_available = check_aclnn_kernel_available("aclnnQuantMatmulWeightNz");
281- TORCH_CHECK(is_quant_matmul_weight_nz_available,
282- "Get aclnnQuantMatmulWeightNz or aclnnQuantMatmulWeightNzGetWorkspaceSize failed, only "
283- "aclnnQuantMatmulWeightNz support X2's format is nz, please upgrade CANN.",
284- OPS_ERROR(ErrCode::PARAM));
285- } else {
286- static const bool is_quant_matmul_v5_available = check_aclnn_kernel_available("aclnnQuantMatmulV5");
287- TORCH_CHECK(is_quant_matmul_v5_available,
288- "Get aclnnQuantMatmulV5 or aclnnQuantMatmulV5 failed, only "
289- "aclnnQuantMatmulV5 support A8W4, please upgrade CANN.",
290- OPS_ERROR(ErrCode::TYPE));
291- }
292- bool is_a8W4_int = x1.dtype() == at::kChar && x2.dtype() == at::kInt;
293- bool is_a8W4_float = x1.dtype() == at::kFloat8_e4m3fn && x2.dtype() == at::kFloat;
294- at::IntArrayRef group_size_list = group_sizes.value_or(at::IntArrayRef{});
295- int64_t group_size = check_and_get_groups(group_size_list);
296- bool is_a4w4 = x1.dtype() == at::kInt && x2.dtype() == at::kInt;
297- bool trans_x1 = is_transpose_last_two_dims(x1);
298- bool trans_x2 = is_transpose_last_two_dims(x2);
299- auto x1_dim_num = x1.dim();
300- auto x2_dim_num = x2.dim();
301- auto x2_n_dim = (is_a4w4 && !trans_x2) ? x2.size(x2_dim_num - 1) * INT4_NUMS_IN_INT32 : x2.size(x2_dim_num - 1);
302 331 
303#if VERSION_BETWEEN(V2R1, V2R7)332#if VERSION_BETWEEN(V2R1, V2R7)
304- bool mxfp4_valid = x1_dtype.has_value() && x2_dtype.has_value() &&333+ bool mxfp4_valid = x1_dtype.has_value() && x2_dtype.has_value() &&
305- x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) &&334+ x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) &&
306- x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);335+ x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
307#endif336#endif
308#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)337#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)
309- bool mxfp4_valid = false;338+ bool mxfp4_valid = false;
310- if (x1_dtype.has_value()) {339+ if (x1_dtype.has_value()) {
311- mxfp4_valid = x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);340+ mxfp4_valid = x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
312- } else {341+ } else {
313- mxfp4_valid = x1.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;342+ mxfp4_valid = x1.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
314- }343+ }
315- if (x2_dtype.has_value()) {344+ if (x2_dtype.has_value()) {
316- mxfp4_valid = mxfp4_valid && x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);345+ mxfp4_valid = mxfp4_valid && x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
317- } else {346+ } else {
318- mxfp4_valid = mxfp4_valid && x2.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;347+ mxfp4_valid = mxfp4_valid && x2.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
319- }348+ }
320#endif349#endif
321 350 
322- c10::SmallVector<int64_t, SIZE> output_size;351+ c10::SmallVector<int64_t, SIZE> output_size;
323- if (is_a8W4_int) {352+ if (is_a8W4_int) {
324- output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};353+ output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
325- } else if (is_a8W4_float) {354+ } else if (is_a8W4_float) {
326- if (trans_x2) {355+ if (trans_x2) {
327- output_size = {x1.sizes()[0], x2.sizes()[1]};356+ output_size = {x1.sizes()[0], x2.sizes()[1]};
328- } else {
329- output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
330- }
331 } else {357 } else {
332- std::vector<uint64_t> batch_record;358+ output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
333- uint64_t batch_val = infer_out_batch_shape(x1, x2, batch_record);
334- const at::Tensor long_tensor = x1_dim_num > x2_dim_num ? x1 : x2;
335- output_size = op_infer::array_to_small_vector(long_tensor.sizes());
336- if (mxfp4_valid) {
337- TORCH_CHECK(x1.dim() >= 2 && x1.dim() <= 6,
338- "x1 dim num should be 2 ~ 6, please check x1 dim num. Actual x1 dim = ", x1.dim(),
339- OPS_ERROR(ErrCode::PARAM));
340- TORCH_CHECK(x2.dim() >= 2 && x2.dim() <= 6,
341- "x2 dim num should be 2 ~ 6, please check x2 dim num. Actual x2 dim = ", x2.dim(),
342- OPS_ERROR(ErrCode::PARAM));
343- int64_t x1_size_last_second = x1.sizes()[x1_dim_num - LAST_SECOND_DIM_INDEX];
344- int64_t x2_size_last = x2.sizes()[x2_dim_num - 1];
345- int64_t real_m = !trans_x1 ? x1_size_last_second : x1_size_last_second * FP4_IN_INT8;
346- int64_t real_n = trans_x2 ? x2_size_last : x2_size_last * FP4_IN_INT8;
347- output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = real_m;
348- output_size[long_tensor.dim() - 1] = real_n;
349- } else {
350- output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = x1.size(x1_dim_num - LAST_SECOND_DIM_INDEX);
351- output_size[long_tensor.dim() - 1] = x2_n_dim;
352- }
353- for (int64_t i = 0; i < long_tensor.dim() - LAST_SECOND_DIM_INDEX; i++) {
354- output_size[i] = static_cast<int64_t>(batch_record[i]);
355- }
356 }359 }
357- c10::TensorOptions options;360+ } else {
358- aclDataType output_acltype = ACL_INT8;361+ std::vector<uint64_t> batch_record;
359- if (!output_dtype.has_value()) {362+ uint64_t batch_val = infer_out_batch_shape(x1, x2, batch_record);
360- options = x1.options().dtype(at::kChar);363+ const at::Tensor long_tensor = x1_dim_num > x2_dim_num ? x1 : x2;
364+ output_size = op_infer::array_to_small_vector(long_tensor.sizes());
365+ if (mxfp4_valid) {
366+ TORCH_CHECK(
367+ x1.dim() >= 2 && x1.dim() <= 6,
368+ "x1 dim num should be 2 ~ 6, please check x1 dim num. Actual x1 dim = ",
369+ x1.dim(),
370+ OPS_ERROR(ErrCode::PARAM));
371+ TORCH_CHECK(
372+ x2.dim() >= 2 && x2.dim() <= 6,
373+ "x2 dim num should be 2 ~ 6, please check x2 dim num. Actual x2 dim = ",
374+ x2.dim(),
375+ OPS_ERROR(ErrCode::PARAM));
376+ int64_t x1_size_last_second = x1.sizes()[x1_dim_num - LAST_SECOND_DIM_INDEX];
377+ int64_t x2_size_last = x2.sizes()[x2_dim_num - 1];
378+ int64_t real_m = !trans_x1 ? x1_size_last_second : x1_size_last_second * FP4_IN_INT8;
379+ int64_t real_n = trans_x2 ? x2_size_last : x2_size_last * FP4_IN_INT8;
380+ output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = real_m;
381+ output_size[long_tensor.dim() - 1] = real_n;
361 } else {382 } else {
362- output_acltype = c10_npu::GetAclDataType(output_dtype.value());383+ output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = x1.size(x1_dim_num - LAST_SECOND_DIM_INDEX);
363- options = x1.options().dtype(npu_preparation::convert_to_scalar_type(output_acltype));384+ output_size[long_tensor.dim() - 1] = x2_n_dim;
364 }385 }
365- at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);386+ for (int64_t i = 0; i < long_tensor.dim() - LAST_SECOND_DIM_INDEX; i++) {
387+ output_size[i] = static_cast<int64_t>(batch_record[i]);
388+ }
389+ }
390+ c10::TensorOptions options;
391+ aclDataType output_acltype = ACL_INT8;
392+ if (!output_dtype.has_value()) {
393+ options = x1.options().dtype(at::kChar);
394+ } else {
395+ output_acltype = c10_npu::GetAclDataType(output_dtype.value());
396+ options = x1.options().dtype(npu_preparation::convert_to_scalar_type(output_acltype));
397+ }
398+ at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);
366 399 
367- const at::Tensor &offset_real = offset.value_or(at::Tensor());400+ const at::Tensor& offset_real = offset.value_or(at::Tensor());
368- const at::Tensor &pertoken_scale_real = pertoken_scale.value_or(at::Tensor());401+ const at::Tensor& pertoken_scale_real = pertoken_scale.value_or(at::Tensor());
369- const at::Tensor &bias_real = bias.value_or(at::Tensor());402+ const at::Tensor& bias_real = bias.value_or(at::Tensor());
370- bool transpose1 = false;403+ bool transpose1 = false;
371- bool transpose2 = false;404+ bool transpose2 = false;
372 405 
373- TensorWrapper x1_wrapper = make_wrapper(x1, x1_dtype);406+ TensorWrapper x1_wrapper = make_wrapper(x1, x1_dtype);
374- TensorWrapper x2_wrapper = make_wrapper(x2, x2_dtype);407+ TensorWrapper x2_wrapper = make_wrapper(x2, x2_dtype);
375- TensorWrapper x1_scale_wrapper = make_wrapper(pertoken_scale_real, pertoken_scale_dtype);408+ TensorWrapper x1_scale_wrapper = make_wrapper(pertoken_scale_real, pertoken_scale_dtype);
376- TensorWrapper x2_scale_wrapper = make_wrapper(scale, scale_dtype);409+ TensorWrapper x2_scale_wrapper = make_wrapper(scale, scale_dtype);
377- TensorWrapper result_wrapper = make_wrapper(result, output_dtype);410+ TensorWrapper result_wrapper = make_wrapper(result, output_dtype);
378- at::Tensor x1_offset = at::empty({0}, options);411+ at::Tensor x1_offset = at::empty({0}, options);
379- at::Tensor x2_offset = at::Tensor();412+ at::Tensor x2_offset = at::Tensor();
380- at::Tensor y_offset = at::empty({0}, options);413+ at::Tensor y_offset = at::empty({0}, options);
381- if (is_a8W4_int) { // Only A8W4 int needs y_offset414+ if (is_a8W4_int) { // Only A8W4 int needs y_offset
382- y_offset = offset_real;415+ y_offset = offset_real;
416+ } else {
417+ x2_offset = offset_real;
418+ }
419+ 
420+ bool use_aclnn_v5 =
421+ x1_dtype.has_value() || (x1.dtype() != at::kInt && x1.dtype() != at::kChar) || is_a8W4_float || is_a8W4_int;
422+ 
423+ aclDataType pertoken_scale_dtype_real = pertoken_scale_dtype.has_value()
424+ ? c10_npu::GetAclDataType(pertoken_scale_dtype.value())
425+ : (pertoken_scale.has_value() ? c10_npu::GetAclDataType(static_cast<int64_t>(pertoken_scale_real.scalar_type()))
426+ : aclDataType::ACL_INT8);
427+ bool need_check_trans = pertoken_scale.has_value() &&
428+ (((pertoken_scale_real.dim() == x1.dim() && scale.dim() == x2.dim()) ||
429+ pertoken_scale_dtype_real == aclDataType::ACL_FLOAT8_E8M0) &&
430+ (pertoken_scale_real.dim() >= 2 && scale.dim() >= 2)) &&
431+ !(is_a8W4_float || is_a8W4_int);
432+ if (need_check_trans) {
433+ int64_t dim_x1 = x1.dim() - 2; // check the last 2 dim
434+ int64_t dim_x2 = x2.dim() - 2; // check the last 2 dim
435+ int64_t dim_x1_scale = 0;
436+ int64_t dim_x2_scale = 0;
437+ if (pertoken_scale_dtype_real != aclDataType::ACL_FLOAT8_E8M0) {
438+ dim_x1_scale = pertoken_scale_real.dim() - 2; // check the last 2 dim in GB/BB
439+ dim_x2_scale = scale.dim() - 2; // check the last 2 dim in GB/BB
440+ }
441+ TORCH_CHECK(
442+ is_x_scale_same_transpose(x1, pertoken_scale_real, dim_x1, dim_x1_scale),
443+ "Input x1 tensor and pertoken_scale tensor's transpose are not same, please check input.",
444+ OPS_ERROR(ErrCode::PARAM));
445+ TORCH_CHECK(
446+ is_x_scale_same_transpose(x2, scale, dim_x2, dim_x2_scale),
447+ "Input x2 tensor and scale tensor's transpose are not same, please check input.",
448+ OPS_ERROR(ErrCode::PARAM));
449+ }
450+ 
451+ bool use_trans_quant_param = scale.dtype() == at::kFloat && !pertoken_scale.has_value() &&
452+ (output_acltype != ACL_BF16 || use_aclnn_v5) && output_acltype != ACL_INT32;
453+ if (use_trans_quant_param) {
454+ const at::Tensor quant_param = op_api::npu_trans_quant_param(scale, offset);
455+ if (is_nz_format(x2)) {
456+ EXEC_NPU_CMD(
457+ aclnnQuantMatmulWeightNz,
458+ x1_wrapper,
459+ x2_wrapper,
460+ pertoken_scale_real,
461+ quant_param,
462+ y_scale,
463+ x1_offset,
464+ x2_offset,
465+ y_offset,
466+ bias_real,
467+ transpose1,
468+ transpose2,
469+ group_size,
470+ result_wrapper);
383 } else {471 } else {
384- x2_offset = offset_real;472+ EXEC_NPU_CMD(
473+ aclnnQuantMatmulV5,
474+ x1_wrapper,
475+ x2_wrapper,
476+ pertoken_scale_real,
477+ quant_param,
478+ y_scale,
479+ x1_offset,
480+ x2_offset,
481+ y_offset,
482+ bias_real,
483+ transpose1,
484+ transpose2,
485+ group_size,
486+ result_wrapper);
385 }487 }
386- 488+ } else {
387- bool use_aclnn_v5 = x1_dtype.has_value() || (x1.dtype() != at::kInt && x1.dtype() != at::kChar) ||489+ if (!is_a4w4 && is_nz_format(x2)) {
388- is_a8W4_float || is_a8W4_int;490+ EXEC_NPU_CMD(
389- 491+ aclnnQuantMatmulWeightNz,
390- aclDataType pertoken_scale_dtype_real = pertoken_scale_dtype.has_value()492+ x1_wrapper,
391- ? c10_npu::GetAclDataType(pertoken_scale_dtype.value())493+ x2_wrapper,
392- : (pertoken_scale.has_value()494+ x1_scale_wrapper,
393- ? c10_npu::GetAclDataType(static_cast<int64_t>(pertoken_scale_real.scalar_type()))495+ x2_scale_wrapper,
394- : aclDataType::ACL_INT8);496+ y_scale,
395- bool need_check_trans = pertoken_scale.has_value()497+ x1_offset,
396- && (((pertoken_scale_real.dim() == x1.dim() && scale.dim() == x2.dim())498+ x2_offset,
397- || pertoken_scale_dtype_real == aclDataType::ACL_FLOAT8_E8M0)499+ y_offset,
398- && (pertoken_scale_real.dim() >= 2 && scale.dim() >= 2))500+ bias_real,
399- && !(is_a8W4_float || is_a8W4_int);501+ transpose1,
400- if (need_check_trans) {502+ transpose2,
401- int64_t dim_x1 = x1.dim() - 2; // check the last 2 dim503+ group_size,
402- int64_t dim_x2 = x2.dim() - 2; // check the last 2 dim504+ result_wrapper);
403- int64_t dim_x1_scale = 0;
404- int64_t dim_x2_scale = 0;
405- if (pertoken_scale_dtype_real != aclDataType::ACL_FLOAT8_E8M0) {
406- dim_x1_scale = pertoken_scale_real.dim() - 2; // check the last 2 dim in GB/BB
407- dim_x2_scale = scale.dim() - 2; // check the last 2 dim in GB/BB
408- }
409- TORCH_CHECK(is_x_scale_same_transpose(x1, pertoken_scale_real, dim_x1, dim_x1_scale),
410- "Input x1 tensor and pertoken_scale tensor's transpose are not same, please check input.",
411- OPS_ERROR(ErrCode::PARAM));
412- TORCH_CHECK(is_x_scale_same_transpose(x2, scale, dim_x2, dim_x2_scale),
413- "Input x2 tensor and scale tensor's transpose are not same, please check input.",
414- OPS_ERROR(ErrCode::PARAM));
415- }
416- 
417- bool use_trans_quant_param = scale.dtype() == at::kFloat && !pertoken_scale.has_value() &&
418- (output_acltype != ACL_BF16 || use_aclnn_v5) && output_acltype != ACL_INT32;
419- if (use_trans_quant_param) {
420- const at::Tensor quant_param = op_api::npu_trans_quant_param(scale, offset);
421- if (is_nz_format(x2)) {
422- EXEC_NPU_CMD(aclnnQuantMatmulWeightNz, x1_wrapper, x2_wrapper, pertoken_scale_real, quant_param, y_scale,
423- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);
424- } else {
425- EXEC_NPU_CMD(aclnnQuantMatmulV5, x1_wrapper, x2_wrapper, pertoken_scale_real, quant_param, y_scale,
426- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);
427- }
428 } else {505 } else {
429- if (!is_a4w4 && is_nz_format(x2)) {506+ EXEC_NPU_CMD(
430- EXEC_NPU_CMD(aclnnQuantMatmulWeightNz, x1_wrapper, x2_wrapper, x1_scale_wrapper, x2_scale_wrapper, y_scale,507+ aclnnQuantMatmulV5,
431- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);508+ x1_wrapper,
432- } else {509+ x2_wrapper,
433- EXEC_NPU_CMD(aclnnQuantMatmulV5, x1_wrapper, x2_wrapper, x1_scale_wrapper, x2_scale_wrapper, y_scale,510+ x1_scale_wrapper,
434- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);511+ x2_scale_wrapper,
435- }512+ y_scale,
513+ x1_offset,
514+ x2_offset,
515+ y_offset,
516+ bias_real,
517+ transpose1,
518+ transpose2,
519+ group_size,
520+ result_wrapper);
436 }521 }
522+ }
437 523 
438- return result;524+ return result;
439}525}
440 526 
441-}527+} // namespace op_api
Mop_plugin/ops/opapi/ScaledMmV2KernelNpuOpApi.cpp+640-497
@@ -24,19 +24,19 @@
24// Helper functions for ceil_div and round_up24// Helper functions for ceil_div and round_up
25template <typename T>25template <typename T>
26inline T ceil_div(T a, T b) {26inline T ceil_div(T a, T b) {
27- if (b == 0) return T{};27+ if (b == 0)
28- return (a + b - 1) / b;28+ return T{};
29+ return (a + b - 1) / b;
29}30}
30 31 
31template <typename T>32template <typename T>
32inline T round_up(T value, T alignment) {33inline T round_up(T value, T alignment) {
33- if (alignment == 0) {34+ if (alignment == 0) {
34- return value;35+ return value;
35- }36+ }
36- return (value + alignment - 1) / alignment * alignment;37+ return (value + alignment - 1) / alignment * alignment;
37}38}
38 39 
39- 
40namespace op_api {40namespace op_api {
41constexpr size_t LAST_SECOND_DIM_INDEX = 2;41constexpr size_t LAST_SECOND_DIM_INDEX = 2;
42constexpr int64_t PERGROUP_DIM_NUM = 2;42constexpr int64_t PERGROUP_DIM_NUM = 2;
@@ -46,170 +46,189 @@ static const size_t A8W4_GROUP_DIM = 3;
46static const size_t A8W4_INPUT_DIM = 2;46static const size_t A8W4_INPUT_DIM = 2;
47using npu_preparation = at_npu::native::OpPreparation;47using npu_preparation = at_npu::native::OpPreparation;
48 48 
49-bool static is_transpose_last_two_dims(const at::Tensor &tensor)49+bool static is_transpose_last_two_dims(const at::Tensor& tensor) {
50-{50+ if (tensor.dim() < 2 || tensor.dim() > 6) {
51- if (tensor.dim() < 2 || tensor.dim() > 6) {51+ return false;
52+ }
53+ int64_t dim1 = tensor.dim() - 1;
54+ int64_t dim2 = tensor.dim() - 2;
55+ if (tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2)) {
56+ int64_t tmpNxD = tensor.size(dim1) * tensor.size(dim2);
57+ for (int64_t batchDim = tensor.dim() - 3; batchDim >= 0; batchDim--) {
58+ if (tensor.stride(batchDim) != tmpNxD) {
52 return false;59 return false;
60+ }
61+ tmpNxD *= tensor.size(batchDim);
53 }62 }
54- int64_t dim1 = tensor.dim() - 1;63+ if (tensor.size(dim1) == 1 && tensor.size(dim2) == 1) {
55- int64_t dim2 = tensor.dim() - 2;64+ return false;
56- if (tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2)) {
57- int64_t tmpNxD = tensor.size(dim1) * tensor.size(dim2);
58- for (int64_t batchDim = tensor.dim() - 3; batchDim >= 0; batchDim--) {
59- if (tensor.stride(batchDim) != tmpNxD) {
60- return false;
61- }
62- tmpNxD *= tensor.size(batchDim);
63- }
64- if (tensor.size(dim1) == 1 && tensor.size(dim2) == 1) {
65- return false;
66- }
67- return true;
68 }65 }
69- return false;66+ return true;
67+ }
68+ return false;
70}69}
71 70 
72-static bool is_transpose_certain_two_dims(const at::Tensor &tensor, int64_t dim)71+static bool is_transpose_certain_two_dims(const at::Tensor& tensor, int64_t dim) {
73-{72+ return tensor.stride(dim + 1) == tensor.stride(dim) * tensor.size(dim);
74- return tensor.stride(dim + 1) == tensor.stride(dim) * tensor.size(dim);
75}73}
76 74 
77-static bool is_x_scale_same_transpose(const at::Tensor &x, const at::Tensor &scale, int64_t dim_x, int64_t dim_scale)75+static bool is_x_scale_same_transpose(const at::Tensor& x, const at::Tensor& scale, int64_t dim_x, int64_t dim_scale) {
78-{76+ if (x.dim() < dim_x + 2 || scale.dim() < dim_scale + 2) { // make sure dims after the start dim are no less than 2
79- if (x.dim() < dim_x + 2 || scale.dim() < dim_scale + 2) { // make sure dims after the start dim are no less than 277+ return true;
80- return true;78+ }
81- }79+ if (x.size(dim_x) == 1 && x.size(dim_x + 1) == 1) {
82- if (x.size(dim_x) == 1 && x.size(dim_x + 1)== 1) {80+ return true;
83- return true;81+ }
84- }82+ if (scale.size(dim_scale) == 1 && scale.size(dim_scale + 1) == 1) {
85- if (scale.size(dim_scale) == 1 && scale.size(dim_scale + 1)== 1) {83+ return true;
86- return true;84+ }
87- }85+ bool x_trans = is_transpose_certain_two_dims(x, dim_x);
88- bool x_trans = is_transpose_certain_two_dims(x, dim_x);86+ bool scale_trans = is_transpose_certain_two_dims(scale, dim_scale);
89- bool scale_trans = is_transpose_certain_two_dims(scale, dim_scale);87+ if (x_trans == scale_trans) {
90- if (x_trans == scale_trans) {88+ return true;
91- return true;89+ }
92- }90+ return false;
93- return false;
94}91}
95-static bool is_nz_format(const at::Tensor& x2)92+static bool is_nz_format(const at::Tensor& x2) {
96-{93+ const torch_npu::NPUStorageDesc& tensor_desc = torch_npu::NPUBridge::GetNpuStorageImpl(x2)->npu_desc_;
97- const torch_npu::NPUStorageDesc &tensor_desc =94+ return tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ || tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_4 ||
98- torch_npu::NPUBridge::GetNpuStorageImpl(x2)->npu_desc_;95+ tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_16;
99- return tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ ||
100- tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_4 ||
101- tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_16;
102}96}
103 97 
104-static uint64_t infer_out_batch_shape(const at::Tensor &x1, const at::Tensor &x2, std::vector<uint64_t> &batch_record)98+static uint64_t infer_out_batch_shape(const at::Tensor& x1, const at::Tensor& x2, std::vector<uint64_t>& batch_record) {
105-{99+ TORCH_CHECK(
106- TORCH_CHECK(at_npu::native::FormatHelper::IsBaseFormatType(x2) || is_nz_format(x2),100+ at_npu::native::FormatHelper::IsBaseFormatType(x2) || is_nz_format(x2),
107- "x2 should be in the original image format or nz format, but it is ",101+ "x2 should be in the original image format or nz format, but it is ",
108- npu_preparation::get_tensor_npu_format(x2), OPS_ERROR(ErrCode::PARAM));102+ npu_preparation::get_tensor_npu_format(x2),
109- uint64_t batch_val = 1;103+ OPS_ERROR(ErrCode::PARAM));
110- auto x1_dim_num = x1.dim();104+ uint64_t batch_val = 1;
111- auto x2_dim_num = x2.dim();105+ auto x1_dim_num = x1.dim();
112- auto out_dim_num = std::max(x1_dim_num, x2_dim_num);106+ auto x2_dim_num = x2.dim();
113- auto &shape_long = x1_dim_num > x2_dim_num ? x1 : x2;107+ auto out_dim_num = std::max(x1_dim_num, x2_dim_num);
114- auto &shape_short = x1_dim_num > x2_dim_num ? x2 : x1;108+ auto& shape_long = x1_dim_num > x2_dim_num ? x1 : x2;
115- int64_t vaild_offset = out_dim_num - std::min(x1_dim_num, x2_dim_num);109+ auto& shape_short = x1_dim_num > x2_dim_num ? x2 : x1;
116- for (int64_t i = 0; i < out_dim_num - LAST_SECOND_DIM_INDEX; i++) {110+ int64_t vaild_offset = out_dim_num - std::min(x1_dim_num, x2_dim_num);
117- auto short_dim = i < vaild_offset ? 1 : shape_short.size(i - vaild_offset);111+ for (int64_t i = 0; i < out_dim_num - LAST_SECOND_DIM_INDEX; i++) {
118- auto long_dim = shape_long.size(i);112+ auto short_dim = i < vaild_offset ? 1 : shape_short.size(i - vaild_offset);
119- TORCH_CHECK(!(short_dim > 1 && long_dim > 1 && short_dim != long_dim),113+ auto long_dim = shape_long.size(i);
120- "the x1 shape and x2 shape not supported for broadcast, the short_dim is ",114+ TORCH_CHECK(
121- short_dim, " and the long_dim is ", long_dim, OPS_ERROR(ErrCode::PARAM));115+ !(short_dim > 1 && long_dim > 1 && short_dim != long_dim),
122- uint64_t cur_batch_value = static_cast<uint64_t>(std::max(short_dim, long_dim));116+ "the x1 shape and x2 shape not supported for broadcast, the short_dim is ",
123- batch_val = batch_val * cur_batch_value;117+ short_dim,
124- batch_record.push_back(cur_batch_value);118+ " and the long_dim is ",
125- }119+ long_dim,
126- return batch_val;120+ OPS_ERROR(ErrCode::PARAM));
121+ uint64_t cur_batch_value = static_cast<uint64_t>(std::max(short_dim, long_dim));
122+ batch_val = batch_val * cur_batch_value;
123+ batch_record.push_back(cur_batch_value);
124+ }
125+ return batch_val;
127}126}
128 127 
129-static int64_t check_and_get_groups(at::IntArrayRef group_size_list)128+static int64_t check_and_get_groups(at::IntArrayRef group_size_list) {
130-{129+ int64_t groups = 0;
131- int64_t groups = 0;130+ if (group_size_list.empty()) {
132- if (group_size_list.empty()) {
133- return groups;
134- }
135- size_t group_dim = group_size_list.size();
136- TORCH_CHECK(group_dim == A8W4_GROUP_DIM, "group_sizes only support input with three elements, but got ",
137- group_dim, OPS_ERROR(ErrCode::PARAM));
138- int64_t group_m = static_cast<int64_t>(group_size_list[0]);
139- int64_t group_n = static_cast<int64_t>(group_size_list[1]);
140- int64_t group_k = static_cast<int64_t>(group_size_list[2]);
141- bool invalid_group_param = ((group_m <= GROUP_MAX && group_m >= 0)
142- && (group_n <= GROUP_MAX && group_n >= 0)
143- && (group_k <= GROUP_MAX && group_k >= 0));
144- TORCH_CHECK(invalid_group_param, "group param value must conform to range [0, 65535]", OPS_ERROR(ErrCode::VALUE));
145- groups = static_cast<int64_t>((static_cast<uint64_t>(group_m) << 32) + (static_cast<uint64_t>(group_n) << 16) +
146- (static_cast<uint64_t>(group_k)));
147 return groups;131 return groups;
132+ }
133+ size_t group_dim = group_size_list.size();
134+ TORCH_CHECK(
135+ group_dim == A8W4_GROUP_DIM,
136+ "group_sizes only support input with three elements, but got ",
137+ group_dim,
138+ OPS_ERROR(ErrCode::PARAM));
139+ int64_t group_m = static_cast<int64_t>(group_size_list[0]);
140+ int64_t group_n = static_cast<int64_t>(group_size_list[1]);
141+ int64_t group_k = static_cast<int64_t>(group_size_list[2]);
142+ bool invalid_group_param =
143+ ((group_m <= GROUP_MAX && group_m >= 0) && (group_n <= GROUP_MAX && group_n >= 0) &&
144+ (group_k <= GROUP_MAX && group_k >= 0));
145+ TORCH_CHECK(invalid_group_param, "group param value must conform to range [0, 65535]", OPS_ERROR(ErrCode::VALUE));
146+ groups = static_cast<int64_t>(
147+ (static_cast<uint64_t>(group_m) << 32) + (static_cast<uint64_t>(group_n) << 16) +
148+ (static_cast<uint64_t>(group_k)));
149+ return groups;
148}150}
149 151 
150static const std::map<c10::ScalarType, int64_t> SCALAR_TO_INT_TYPE_MAP = {152static const std::map<c10::ScalarType, int64_t> SCALAR_TO_INT_TYPE_MAP = {
151- {c10::ScalarType::Char, static_cast<int64_t>(at::kChar)}, /**< int8 */153+ {c10::ScalarType::Char, static_cast<int64_t>(at::kChar)}, /**< int8 */
152- {c10::ScalarType::Int, static_cast<int64_t>(at::kInt)}, /**< int32 */154+ {c10::ScalarType::Int, static_cast<int64_t>(at::kInt)}, /**< int32 */
153- {c10::ScalarType::BFloat16, static_cast<int64_t>(at::kBFloat16)}, /**< bfp16 */155+ {c10::ScalarType::BFloat16, static_cast<int64_t>(at::kBFloat16)}, /**< bfp16 */
154- {c10::ScalarType::Half, static_cast<int64_t>(at::kHalf)}, /**< fp16 */156+ {c10::ScalarType::Half, static_cast<int64_t>(at::kHalf)}, /**< fp16 */
155- {c10::ScalarType::Float, static_cast<int64_t>(at::kFloat)}, /**< fp32 */157+ {c10::ScalarType::Float, static_cast<int64_t>(at::kFloat)}, /**< fp32 */
156 158 
157- {c10::ScalarType::Float8_e4m3fn, static_cast<int64_t>(at::kFloat8_e4m3fn)}, /**< fp8e4m3 */159+ {c10::ScalarType::Float8_e4m3fn, static_cast<int64_t>(at::kFloat8_e4m3fn)}, /**< fp8e4m3 */
158- {c10::ScalarType::Float8_e5m2, static_cast<int64_t>(at::kFloat8_e5m2)}, /**< fp8e5m2 */160+ {c10::ScalarType::Float8_e5m2, static_cast<int64_t>(at::kFloat8_e5m2)}, /**< fp8e5m2 */
159- {c10::ScalarType::Byte, static_cast<int64_t>(c10_npu::DType::HIFLOAT8)} /**< hif8 */161+ {c10::ScalarType::Byte, static_cast<int64_t>(c10_npu::DType::HIFLOAT8)} /**< hif8 */
160};162};
161 163 
162-static c10::optional<int64_t> ToIntType(const std::optional<c10::ScalarType> &torchType) {164+static c10::optional<int64_t> ToIntType(const std::optional<c10::ScalarType>& torchType) {
163- c10::optional<int64_t> int_type = c10::nullopt;165+ c10::optional<int64_t> int_type = c10::nullopt;
164- if (torchType.has_value()) {166+ if (torchType.has_value()) {
165- const auto &it = SCALAR_TO_INT_TYPE_MAP.find(torchType.value());167+ const auto& it = SCALAR_TO_INT_TYPE_MAP.find(torchType.value());
166- if (it != SCALAR_TO_INT_TYPE_MAP.cend()) {168+ if (it != SCALAR_TO_INT_TYPE_MAP.cend()) {
167- int_type = c10::make_optional(it->second);169+ int_type = c10::make_optional(it->second);
168- }
169 }170 }
170- return int_type;171+ }
172+ return int_type;
171}173}
172 174 
173-using acceptance_fn = std::function<bool(c10::ScalarType, std::vector<ScalingType>&, c10::ArrayRef<at::Tensor>&, c10::ScalarType, std::vector<ScalingType>&, c10::ArrayRef<at::Tensor>&)>;175+using acceptance_fn = std::function<bool(
176+ c10::ScalarType,
177+ std::vector<ScalingType>&,
178+ c10::ArrayRef<at::Tensor>&,
179+ c10::ScalarType,
180+ std::vector<ScalingType>&,
181+ c10::ArrayRef<at::Tensor>&)>;
174 182 
175// Namespace for scaled_blas check functions183// Namespace for scaled_blas check functions
176namespace scaled_blas {184namespace scaled_blas {
177 185 
178// Check deepseek recipe (BlockWise 1x128 and 128x128 scaling)186// Check deepseek recipe (BlockWise 1x128 and 128x128 scaling)
179-bool check_deepseek_recipe(c10::ScalarType a_type, std::vector<ScalingType>& a_recipe, c10::ArrayRef<at::Tensor>& a_scale,187+bool check_deepseek_recipe(
180- c10::ScalarType b_type, std::vector<ScalingType>& b_recipe, c10::ArrayRef<at::Tensor>& b_scale,188+ c10::ScalarType a_type,
181- ScalingType expected_a_type, ScalingType expected_b_type) {189+ std::vector<ScalingType>& a_recipe,
182- if (a_type != c10::ScalarType::Float8_e4m3fn && a_type != c10::ScalarType::Float8_e5m2) {190+ c10::ArrayRef<at::Tensor>& a_scale,
183- return false;191+ c10::ScalarType b_type,
184- }192+ std::vector<ScalingType>& b_recipe,
185- if (b_type != c10::ScalarType::Float8_e4m3fn && b_type != c10::ScalarType::Float8_e5m2) {193+ c10::ArrayRef<at::Tensor>& b_scale,
186- return false;194+ ScalingType expected_a_type,
187- }195+ ScalingType expected_b_type) {
188- if (a_recipe.empty() || b_recipe.empty()) {196+ if (a_type != c10::ScalarType::Float8_e4m3fn && a_type != c10::ScalarType::Float8_e5m2) {
189- return false;197+ return false;
190- }198+ }
191- if (a_scale.empty() || b_scale.empty()) {199+ if (b_type != c10::ScalarType::Float8_e4m3fn && b_type != c10::ScalarType::Float8_e5m2) {
192- return false;200+ return false;
193- }201+ }
194- return (a_recipe[0] == expected_a_type && b_recipe[0] == expected_b_type);202+ if (a_recipe.empty() || b_recipe.empty()) {
203+ return false;
204+ }
205+ if (a_scale.empty() || b_scale.empty()) {
206+ return false;
207+ }
208+ return (a_recipe[0] == expected_a_type && b_recipe[0] == expected_b_type);
195}209}
196 210 
197// Check mxfp8 recipe211// Check mxfp8 recipe
198-static bool check_mxfp8_recipe(c10::ScalarType a_type, std::vector<ScalingType>& a_recipe, c10::ArrayRef<at::Tensor>& a_scale,212+static bool check_mxfp8_recipe(
199- c10::ScalarType b_type, std::vector<ScalingType>& b_recipe, c10::ArrayRef<at::Tensor>& b_scale) {213+ c10::ScalarType a_type,
200- if (a_type != c10::ScalarType::Float8_e4m3fn && a_type != c10::ScalarType::Float8_e5m2) {214+ std::vector<ScalingType>& a_recipe,
201- return false;215+ c10::ArrayRef<at::Tensor>& a_scale,
202- }216+ c10::ScalarType b_type,
203- if (b_type != c10::ScalarType::Float8_e4m3fn && b_type != c10::ScalarType::Float8_e5m2) {217+ std::vector<ScalingType>& b_recipe,
204- return false;218+ c10::ArrayRef<at::Tensor>& b_scale) {
205- }219+ if (a_type != c10::ScalarType::Float8_e4m3fn && a_type != c10::ScalarType::Float8_e5m2) {
206- if (a_recipe.empty() || b_recipe.empty()) {220+ return false;
207- return false;221+ }
208- }222+ if (b_type != c10::ScalarType::Float8_e4m3fn && b_type != c10::ScalarType::Float8_e5m2) {
209- if (a_scale.empty() || b_scale.empty()) {223+ return false;
210- return false;224+ }
211- }225+ if (a_recipe.empty() || b_recipe.empty()) {
212- return true;226+ return false;
227+ }
228+ if (a_scale.empty() || b_scale.empty()) {
229+ return false;
230+ }
231+ return true;
213}232}
214 233 
215// Check tensorwise recipe234// Check tensorwise recipe
@@ -233,10 +252,14 @@ bool check_tensorwise_recipe(
233 return false;252 return false;
234 }253 }
235 // Need {Tensorwise, e8m0} for A & B254 // Need {Tensorwise, e8m0} for A & B
236- if (recipe_a[0] != ScalingType::TensorWise) return false;255+ if (recipe_a[0] != ScalingType::TensorWise)
237- if (scales_a[0].scalar_type() != c10::ScalarType::Float) return false;256+ return false;
238- if (recipe_b[0] != ScalingType::TensorWise) return false;257+ if (scales_a[0].scalar_type() != c10::ScalarType::Float)
239- if (scales_b[0].scalar_type() != c10::ScalarType::Float) return false;258+ return false;
259+ if (recipe_b[0] != ScalingType::TensorWise)
260+ return false;
261+ if (scales_b[0].scalar_type() != c10::ScalarType::Float)
262+ return false;
240 263 
241 return true;264 return true;
242}265}
@@ -263,10 +286,14 @@ static bool check_rowwise_recipe(
263 }286 }
264 287 
265 // Need {RowWise, fp32} for A & B288 // Need {RowWise, fp32} for A & B
266- if (recipe_a[0] != ScalingType::RowWise) return false;289+ if (recipe_a[0] != ScalingType::RowWise)
267- if (scales_a[0].scalar_type() != c10::ScalarType::Float) return false;290+ return false;
268- if (recipe_b[0] != ScalingType::RowWise) return false;291+ if (scales_a[0].scalar_type() != c10::ScalarType::Float)
269- if (scales_b[0].scalar_type() != c10::ScalarType::Float) return false;292+ return false;
293+ if (recipe_b[0] != ScalingType::RowWise)
294+ return false;
295+ if (scales_b[0].scalar_type() != c10::ScalarType::Float)
296+ return false;
270 297 
271 return true;298 return true;
272}299}
@@ -274,37 +301,34 @@ static bool check_rowwise_recipe(
274} // namespace scaled_blas301} // namespace scaled_blas
275 302 
276// Wrapper for check_deepseek_recipe that captures expected ScalingTypes303// Wrapper for check_deepseek_recipe that captures expected ScalingTypes
277-std::function<bool(c10::ScalarType, std::vector<ScalingType>&, c10::ArrayRef<at::Tensor>&, c10::ScalarType, std::vector<ScalingType>&, c10::ArrayRef<at::Tensor>&)>304+std::function<bool(
305+ c10::ScalarType,
306+ std::vector<ScalingType>&,
307+ c10::ArrayRef<at::Tensor>&,
308+ c10::ScalarType,
309+ std::vector<ScalingType>&,
310+ c10::ArrayRef<at::Tensor>&)>
278make_deepseek_checker(ScalingType expected_a, ScalingType expected_b) {311make_deepseek_checker(ScalingType expected_a, ScalingType expected_b) {
279- return [expected_a, expected_b](c10::ScalarType a_type, std::vector<ScalingType>& a_recipe, c10::ArrayRef<at::Tensor>& a_scale,312+ return [expected_a, expected_b](
280- c10::ScalarType b_type, std::vector<ScalingType>& b_recipe, c10::ArrayRef<at::Tensor>& b_scale) {313+ c10::ScalarType a_type,
281- return scaled_blas::check_deepseek_recipe(a_type, a_recipe, a_scale, b_type, b_recipe, b_scale, expected_a, expected_b);314+ std::vector<ScalingType>& a_recipe,
282- };315+ c10::ArrayRef<at::Tensor>& a_scale,
316+ c10::ScalarType b_type,
317+ std::vector<ScalingType>& b_recipe,
318+ c10::ArrayRef<at::Tensor>& b_scale) {
319+ return scaled_blas::check_deepseek_recipe(
320+ a_type, a_recipe, a_scale, b_type, b_recipe, b_scale, expected_a, expected_b);
321+ };
283}322}
284 323 
285using KernelDispatchItem = std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>;324using KernelDispatchItem = std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>;
286-const std::array<KernelDispatchItem, 4> scale_kernel_dispatch_table = {{325+const std::array<KernelDispatchItem, 4> scale_kernel_dispatch_table = {
287- {326+ {{"tensorwise_tensorwise", scaled_blas::check_tensorwise_recipe, ScaledGemmImplementation::TENSORWISE_TENSORWISE},
288- "tensorwise_tensorwise",327+ {"rowwise_rowwise", scaled_blas::check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
289- scaled_blas::check_tensorwise_recipe,328+ {"deepseek_block_1x128_1x128",
290- ScaledGemmImplementation::TENSORWISE_TENSORWISE329+ make_deepseek_checker(ScalingType::BlockWise1x128, ScalingType::BlockWise1x128),
291- },330+ ScaledGemmImplementation::BLOCK_1x128_1x128},
292- {331+ {"mxfp8_double_channel", scaled_blas::check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8}}};
293- "rowwise_rowwise",
294- scaled_blas::check_rowwise_recipe,
295- ScaledGemmImplementation::ROWWISE_ROWWISE
296- },
297- {
298- "deepseek_block_1x128_1x128",
299- make_deepseek_checker(ScalingType::BlockWise1x128, ScalingType::BlockWise1x128),
300- ScaledGemmImplementation::BLOCK_1x128_1x128
301- },
302- {
303- "mxfp8_double_channel",
304- scaled_blas::check_mxfp8_recipe,
305- ScaledGemmImplementation::MXFP8_MXFP8
306- }
307-}};
308 332 
309// Signature note: scale_a / scale_b use at::ITensorListRef so this impl matches333// Signature note: scale_a / scale_b use at::ITensorListRef so this impl matches
310// the wrapper generated by torch_npu's codegen on PyTorch >= 2.14, where334// the wrapper generated by torch_npu's codegen on PyTorch >= 2.14, where
@@ -313,7 +337,9 @@ const std::array<KernelDispatchItem, 4> scale_kernel_dispatch_table = {{
313// const at::ITensorListRef&. IListRef is implicitly constructible from337// const at::ITensorListRef&. IListRef is implicitly constructible from
314// ArrayRef, so on 2.13 the wrapper still passes at::TensorList and it converts338// ArrayRef, so on 2.13 the wrapper still passes at::TensorList and it converts
315// transparently -- no version gating needed. See docs/issue_scaled_mm_v2_structured.md.339// transparently -- no version gating needed. See docs/issue_scaled_mm_v2_structured.md.
316-at::Tensor _scaled_mm_v2(const at::Tensor& mat_a, const at::Tensor& mat_b,340+at::Tensor _scaled_mm_v2(
341+ const at::Tensor& mat_a,
342+ const at::Tensor& mat_b,
317 const at::ITensorListRef& scale_a_ilist,343 const at::ITensorListRef& scale_a_ilist,
318 at::IntArrayRef scale_recipe_a,344 at::IntArrayRef scale_recipe_a,
319 at::IntArrayRef swizzle_a,345 at::IntArrayRef swizzle_a,
@@ -323,361 +349,478 @@ at::Tensor _scaled_mm_v2(const at::Tensor& mat_a, const at::Tensor& mat_b,
323 const std::optional<at::Tensor>& bias,349 const std::optional<at::Tensor>& bias,
324 const std::optional<c10::ScalarType> out_dtype, // 当前cpp中不支持float8350 const std::optional<c10::ScalarType> out_dtype, // 当前cpp中不支持float8
325 at::IntArrayRef contraction_dim,351 at::IntArrayRef contraction_dim,
326- bool use_fast_accum)352+ bool use_fast_accum) {
327-{353+ // Materialize the ilist inputs into local TensorList (ArrayRef<Tensor>)
328- // Materialize the ilist inputs into local TensorList (ArrayRef<Tensor>)354+ // views so the rest of the function body -- which does .empty() / .size()
329- // views so the rest of the function body -- which does .empty() / .size()355+ // and scale_a[i] indexing -- needs no further change. Cost: shallow
330- // and scale_a[i] indexing -- needs no further change. Cost: shallow356+ // per-Tensor refcount bump into a local vector; vector stays alive for
331- // per-Tensor refcount bump into a local vector; vector stays alive for357+ // the whole call, so the ArrayRef view never dangles.
332- // the whole call, so the ArrayRef view never dangles.358+ auto scale_a_materialized = scale_a_ilist.materialize();
333- auto scale_a_materialized = scale_a_ilist.materialize();359+ auto scale_b_materialized = scale_b_ilist.materialize();
334- auto scale_b_materialized = scale_b_ilist.materialize();360+ std::vector<at::Tensor> scale_a_owned(scale_a_materialized.begin(), scale_a_materialized.end());
335- std::vector<at::Tensor> scale_a_owned(scale_a_materialized.begin(), scale_a_materialized.end());361+ std::vector<at::Tensor> scale_b_owned(scale_b_materialized.begin(), scale_b_materialized.end());
336- std::vector<at::Tensor> scale_b_owned(scale_b_materialized.begin(), scale_b_materialized.end());362+ at::TensorList scale_a(scale_a_owned);
337- at::TensorList scale_a(scale_a_owned);363+ at::TensorList scale_b(scale_b_owned);
338- at::TensorList scale_b(scale_b_owned);
339 364 
340- // check A5365+ // check A5
341- TORCH_CHECK(c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,366+ TORCH_CHECK(
342- "_scaled_mm is supported only on the Ascend950 platform and after.", OPS_ERROR(ErrCode::PARAM));367+ c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,
368+ "_scaled_mm is supported only on the Ascend950 platform and after.",
369+ OPS_ERROR(ErrCode::PARAM));
343 370 
344- // Check data types: mat_a and mat_b must be float8 type371+ // Check data types: mat_a and mat_b must be float8 type
345- TORCH_CHECK(mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn ||372+ TORCH_CHECK(
346- mat_a.scalar_type() == c10::ScalarType::Float8_e5m2,373+ mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn || mat_a.scalar_type() == c10::ScalarType::Float8_e5m2,
347- "mat_a must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ", mat_a.scalar_type(),374+ "mat_a must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ",
348- OPS_ERROR(ErrCode::TYPE));375+ mat_a.scalar_type(),
349- TORCH_CHECK(mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn ||376+ OPS_ERROR(ErrCode::TYPE));
350- mat_b.scalar_type() == c10::ScalarType::Float8_e5m2,377+ TORCH_CHECK(
351- "mat_b must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ", mat_b.scalar_type(),378+ mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn || mat_b.scalar_type() == c10::ScalarType::Float8_e5m2,
352- OPS_ERROR(ErrCode::TYPE));379+ "mat_b must be float8 type (Float8_e4m3fn or Float8_e5m2), but got ",
380+ mat_b.scalar_type(),
381+ OPS_ERROR(ErrCode::TYPE));
353 382 
354- // Check multiplication of two Float8_e5m2 is not supported (reference: _scaled_mm_out_cuda)383+ // Check multiplication of two Float8_e5m2 is not supported (reference: _scaled_mm_out_cuda)
355- TORCH_CHECK(mat_a.scalar_type() != c10::ScalarType::Float8_e5m2 ||384+ TORCH_CHECK(
356- mat_b.scalar_type() != c10::ScalarType::Float8_e5m2,385+ mat_a.scalar_type() != c10::ScalarType::Float8_e5m2 || mat_b.scalar_type() != c10::ScalarType::Float8_e5m2,
357- "Multiplication of two Float8_e5m2 matrices is not supported",386+ "Multiplication of two Float8_e5m2 matrices is not supported",
358- OPS_ERROR(ErrCode::TYPE));387+ OPS_ERROR(ErrCode::TYPE));
359 388 
360- // Check sizes389+ // Check sizes
361- TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a matrix, please check mat_a dim num." ,OPS_ERROR(ErrCode::PARAM));390+ TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a matrix, please check mat_a dim num.", OPS_ERROR(ErrCode::PARAM));
362- TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a matrix, please check mat_b dim num." ,OPS_ERROR(ErrCode::PARAM));391+ TORCH_CHECK(mat_b.dim() == 2, "mat_b must be a matrix, please check mat_b dim num.", OPS_ERROR(ErrCode::PARAM));
363- TORCH_CHECK(mat_a.sizes()[1] == mat_b.sizes()[0], "mat_a and mat_b shapes cannot be multiplied (",mat_a.sizes()[0],392+ TORCH_CHECK(
364- "x", mat_b.sizes()[1], " and ", mat_b.sizes()[0], "x", mat_b.sizes()[1], ")", OPS_ERROR(ErrCode::PARAM));393+ mat_a.sizes()[1] == mat_b.sizes()[0],
394+ "mat_a and mat_b shapes cannot be multiplied (",
395+ mat_a.sizes()[0],
396+ "x",
397+ mat_b.sizes()[1],
398+ " and ",
399+ mat_b.sizes()[0],
400+ "x",
401+ mat_b.sizes()[1],
402+ ")",
403+ OPS_ERROR(ErrCode::PARAM));
365 404 
366- // Check data types: scale_a and scale_b must be float32 or float8_e8m0 type405+ // Check data types: scale_a and scale_b must be float32 or float8_e8m0 type
367- // Check scale_a and scale_b are not empty406+ // Check scale_a and scale_b are not empty
368- TORCH_CHECK(!scale_a.empty(), "scale_a must not be empty", OPS_ERROR(ErrCode::PARAM));407+ TORCH_CHECK(!scale_a.empty(), "scale_a must not be empty", OPS_ERROR(ErrCode::PARAM));
369- TORCH_CHECK(!scale_b.empty(), "scale_b must not be empty", OPS_ERROR(ErrCode::PARAM));408+ TORCH_CHECK(!scale_b.empty(), "scale_b must not be empty", OPS_ERROR(ErrCode::PARAM));
370 409 
371- TORCH_CHECK(scale_a[0].scalar_type() == c10::ScalarType::Float ||410+ TORCH_CHECK(
372- scale_a[0].scalar_type() == npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),411+ scale_a[0].scalar_type() == c10::ScalarType::Float ||
373- "scale_a must be float32 or float8_e8m0 type, but got ", scale_a[0].scalar_type(),412+ scale_a[0].scalar_type() ==
374- OPS_ERROR(ErrCode::TYPE));413+ npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),
375- TORCH_CHECK(scale_b[0].scalar_type() == c10::ScalarType::Float ||414+ "scale_a must be float32 or float8_e8m0 type, but got ",
376- scale_b[0].scalar_type() == npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),415+ scale_a[0].scalar_type(),
377- "scale_b must be float32 or float8_e8m0 type, but got ", scale_b[0].scalar_type(),416+ OPS_ERROR(ErrCode::TYPE));
378- OPS_ERROR(ErrCode::TYPE));417+ TORCH_CHECK(
418+ scale_b[0].scalar_type() == c10::ScalarType::Float ||
419+ scale_b[0].scalar_type() ==
420+ npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(c10_npu::DType::FLOAT8_E8M0)),
421+ "scale_b must be float32 or float8_e8m0 type, but got ",
422+ scale_b[0].scalar_type(),
423+ OPS_ERROR(ErrCode::TYPE));
379 424 
380- // Cast integer type IDs to corresponding enumeration values425+ // Cast integer type IDs to corresponding enumeration values
381- auto scaling_type_a = convert_int_to_enum<ScalingType>(scale_recipe_a);426+ auto scaling_type_a = convert_int_to_enum<ScalingType>(scale_recipe_a);
382- auto scaling_type_b = convert_int_to_enum<ScalingType>(scale_recipe_b);427+ auto scaling_type_b = convert_int_to_enum<ScalingType>(scale_recipe_b);
383 428 
384- bool matched_valid_impl = false;429+ bool matched_valid_impl = false;
385- ScaledGemmImplementation selected_kernel = ScaledGemmImplementation::NONE;430+ ScaledGemmImplementation selected_kernel = ScaledGemmImplementation::NONE;
386 431 
387- // Iterate registered kernel implementations to find matching one432+ // Iterate registered kernel implementations to find matching one
388- for (const auto& entry : scale_kernel_dispatch_table) {433+ for (const auto& entry : scale_kernel_dispatch_table) {
389- auto const& kernel_id = std::get<0>(entry);434+ auto const& kernel_id = std::get<0>(entry);
390- auto const& matcher_fn = std::get<1>(entry);435+ auto const& matcher_fn = std::get<1>(entry);
391- auto const& kernel_impl = std::get<2>(entry);436+ auto const& kernel_impl = std::get<2>(entry);
392 437 
393- bool config_ok = matcher_fn(438+ bool config_ok =
394- mat_a.scalar_type(),439+ matcher_fn(mat_a.scalar_type(), scaling_type_a, scale_a, mat_b.scalar_type(), scaling_type_b, scale_b);
395- scaling_type_a,
396- scale_a,
397- mat_b.scalar_type(),
398- scaling_type_b,
399- scale_b
400- );
401 440 
402- if (config_ok) {441+ if (config_ok) {
403- selected_kernel = kernel_impl;442+ selected_kernel = kernel_impl;
404- matched_valid_impl = true;443+ matched_valid_impl = true;
405- break;444+ break;
406- }
407 }445 }
446+ }
447+ 
448+ TORCH_CHECK(
449+ matched_valid_impl,
450+ "Unsupported scaling configuration.\n"
451+ "- TensorWise: mat_a/mat_b must be float8; scales float32; scale_a and scale_b are singletons (scalar).\n"
452+ "- RowWise: mat_a/mat_b float8; scales float32; scale_a shape (",
453+ mat_a.size(0),
454+ ", 1); scale_b shape (1, ",
455+ mat_b.size(1),
456+ "); both contiguous.\n"
457+ "- BlockWise 1x128: mat_a/mat_b must be float8; scales float32; scale_a shape (",
458+ mat_a.size(0),
459+ ", ",
460+ ceil_div<int64_t>(mat_a.size(1), 128),
461+ "); scale_b shape (",
462+ ceil_div<int64_t>(mat_b.size(0), 128),
463+ ", ",
464+ mat_b.size(1),
465+ "); outer-dim-major.\n"
466+ "- Blockwise 1x32: mat_a/mat_b float8; scales float8_e8m0fnu; scale_a elements=",
467+ round_up<int64_t>(mat_a.size(0), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_a.size(1), 32), 4),
468+ "; scale_b elements=",
469+ round_up<int64_t>(mat_b.size(1), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_b.size(0), 32), 4),
470+ "; contiguous only.\n"
471+ "Current inputs: mat_a.dtype=",
472+ mat_a.scalar_type(),
473+ ", scale_a.dtype=",
474+ scale_a[0].scalar_type(),
475+ ", scale_a.shape=",
476+ scale_a[0].sizes(),
477+ ", mat_b.dtype=",
478+ mat_b.scalar_type(),
479+ ", scale_b.dtype=",
480+ scale_b[0].scalar_type(),
481+ ", scale_b.shape=",
482+ scale_b[0].sizes());
483+ 
484+ // Check bias (reference: _scaled_mm_out_cuda)
485+ if (bias.has_value()) {
486+ TORCH_CHECK(
487+ bias->numel() == mat_b.sizes()[1],
488+ "Bias must be size ",
489+ mat_b.sizes()[1],
490+ " but got ",
491+ bias->numel(),
492+ OPS_ERROR(ErrCode::PARAM));
493+ // Check out_dtype vs bias compatibility
494+ auto out_dtype_value = out_dtype.value_or(c10::ScalarType::BFloat16);
495+ TORCH_CHECK(
496+ out_dtype_value != c10::ScalarType::Float,
497+ "Bias is not supported when out_dtype is set to Float32",
498+ OPS_ERROR(ErrCode::TYPE));
499+ TORCH_CHECK(
500+ bias->scalar_type() == c10::ScalarType::BFloat16 || bias->scalar_type() == c10::ScalarType::Half,
501+ "Bias must be BFloat16 or Half, but got ",
502+ bias->scalar_type(),
503+ OPS_ERROR(ErrCode::TYPE));
504+ TORCH_CHECK(
505+ (out_dtype_value != c10::ScalarType::Float && out_dtype_value != c10::ScalarType::BFloat16) ||
506+ bias->scalar_type() == c10::ScalarType::BFloat16,
507+ "Bias must be BFloat16 to compute ",
508+ out_dtype_value,
509+ " output, but got ",
510+ bias->scalar_type(),
511+ OPS_ERROR(ErrCode::TYPE));
512+ TORCH_CHECK(
513+ out_dtype_value != c10::ScalarType::Half || bias->scalar_type() == c10::ScalarType::Half,
514+ "Bias must be Float16 to compute ",
515+ out_dtype_value,
516+ " output, but got ",
517+ bias->scalar_type(),
518+ OPS_ERROR(ErrCode::TYPE));
519+ }
520+ 
521+ // Check out_dtype: currently only supports Float32, BFloat16, and Float16
522+ if (out_dtype.has_value()) {
523+ TORCH_CHECK(
524+ out_dtype.value() == c10::ScalarType::Float || out_dtype.value() == c10::ScalarType::BFloat16 ||
525+ out_dtype.value() == c10::ScalarType::Half,
526+ "out_dtype must be Float32, BFloat16, or Float16, but got ",
527+ out_dtype.value(),
528+ OPS_ERROR(ErrCode::TYPE));
529+ }
530+ 
531+ bool has_contraction = !contraction_dim.empty();
532+ if (has_contraction) {
533+ int64_t c_dim_size = static_cast<int64_t>(contraction_dim.size());
534+ TORCH_CHECK(c_dim_size == 2, "contraction_dim must have exactly 2 elements", OPS_ERROR(ErrCode::PARAM));
535+ 
536+ int64_t a_contract_idx = contraction_dim[0];
537+ int64_t b_contract_idx = contraction_dim[1];
538+ int64_t a_dim_size = mat_a.size(a_contract_idx);
539+ int64_t b_dim_size = mat_b.size(b_contract_idx);
408 540 
409 TORCH_CHECK(541 TORCH_CHECK(
410- matched_valid_impl,542+ a_dim_size == b_dim_size,
411- "Unsupported scaling configuration.\n"543+ "dimension mismatch for matrix multiplication (",
412- "- TensorWise: mat_a/mat_b must be float8; scales float32; scale_a and scale_b are singletons (scalar).\n"544+ "A shape=",
413- "- RowWise: mat_a/mat_b float8; scales float32; scale_a shape (", mat_a.size(0), ", 1); scale_b shape (1, ", mat_b.size(1), "); both contiguous.\n"545+ mat_a.size(0),
414- "- BlockWise 1x128: mat_a/mat_b must be float8; scales float32; scale_a shape (", mat_a.size(0), ", ", ceil_div<int64_t>(mat_a.size(1), 128), "); scale_b shape (", ceil_div<int64_t>(mat_b.size(0), 128), ", ", mat_b.size(1), "); outer-dim-major.\n"546+ "x",
415- "- Blockwise 1x32: mat_a/mat_b float8; scales float8_e8m0fnu; scale_a elements=", round_up<int64_t>(mat_a.size(0), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_a.size(1), 32), 4), "; scale_b elements=", round_up<int64_t>(mat_b.size(1), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_b.size(0), 32), 4), "; contiguous only.\n"547+ mat_a.size(1),
416- "Current inputs: mat_a.dtype=", mat_a.scalar_type(), ", scale_a.dtype=", scale_a[0].scalar_type(), ", scale_a.shape=", scale_a[0].sizes(),548+ ", B shape=",
417- ", mat_b.dtype=", mat_b.scalar_type(), ", scale_b.dtype=", scale_b[0].scalar_type(), ", scale_b.shape=", scale_b[0].sizes()549+ mat_b.size(0),
418- );550+ "x",
551+ mat_b.size(1),
552+ ")",
553+ " contraction axis A: ",
554+ a_contract_idx,
555+ ", contraction axis B: ",
556+ b_contract_idx,
557+ OPS_ERROR(ErrCode::PARAM));
558+ } else {
559+ int64_t a_col = mat_a.size(1);
560+ int64_t b_row = mat_b.size(0);
561+ TORCH_CHECK(
562+ a_col == b_row,
563+ "matrix multiplication dimension mismatch (",
564+ "A: ",
565+ mat_a.size(0),
566+ "x",
567+ mat_a.size(1),
568+ ", B: ",
569+ mat_b.size(0),
570+ "x",
571+ mat_b.size(1),
572+ ")",
573+ OPS_ERROR(ErrCode::PARAM));
574+ }
419 575 
420- // Check bias (reference: _scaled_mm_out_cuda)576+ //////////////rowwise transfer//////////////
421- if (bias.has_value()) {577+ at::Tensor processed_scale_a = scale_a[0];
422- TORCH_CHECK(bias->numel() == mat_b.sizes()[1],578+ at::Tensor processed_scale_b = scale_b[0];
423- "Bias must be size ", mat_b.sizes()[1], " but got ", bias->numel(),579+ if (mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn && mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn &&
424- OPS_ERROR(ErrCode::PARAM));580+ scale_a[0].scalar_type() == c10::ScalarType::Float && scale_b[0].scalar_type() == c10::ScalarType::Float) {
425- // Check out_dtype vs bias compatibility581+ if (scale_a[0].dim() == 2 && scale_a[0].sizes()[1] == 1 && scale_b[0].dim() == 2 && scale_b[0].sizes()[0] == 1) {
426- auto out_dtype_value = out_dtype.value_or(c10::ScalarType::BFloat16);582+ TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be contiguous in last dim");
427- TORCH_CHECK(out_dtype_value != c10::ScalarType::Float,583+ TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be contiguous in first dim");
428- "Bias is not supported when out_dtype is set to Float32",584+ processed_scale_a = scale_a[0].squeeze(-1).contiguous();
429- OPS_ERROR(ErrCode::TYPE));585+ processed_scale_b = scale_b[0].squeeze(0);
430- TORCH_CHECK(bias->scalar_type() == c10::ScalarType::BFloat16 ||
431- bias->scalar_type() == c10::ScalarType::Half,
432- "Bias must be BFloat16 or Half, but got ", bias->scalar_type(),
433- OPS_ERROR(ErrCode::TYPE));
434- TORCH_CHECK((out_dtype_value != c10::ScalarType::Float &&
435- out_dtype_value != c10::ScalarType::BFloat16) ||
436- bias->scalar_type() == c10::ScalarType::BFloat16,
437- "Bias must be BFloat16 to compute ", out_dtype_value,
438- " output, but got ", bias->scalar_type(),
439- OPS_ERROR(ErrCode::TYPE));
440- TORCH_CHECK(out_dtype_value != c10::ScalarType::Half ||
441- bias->scalar_type() == c10::ScalarType::Half,
442- "Bias must be Float16 to compute ", out_dtype_value,
443- " output, but got ", bias->scalar_type(),
444- OPS_ERROR(ErrCode::TYPE));
445 }586 }
587+ }
446 588 
447- // Check out_dtype: currently only supports Float32, BFloat16, and Float16589+ //////////////parameters transfer//////////////
448- if (out_dtype.has_value()) {590+ at::Tensor x1 = mat_a;
449- TORCH_CHECK(out_dtype.value() == c10::ScalarType::Float ||591+ at::Tensor x2 = mat_b;
450- out_dtype.value() == c10::ScalarType::BFloat16 ||592+ c10::optional<at::Tensor> pertoken_scale = scale_a.size() > 0 ? c10::make_optional(processed_scale_a) : c10::nullopt;
451- out_dtype.value() == c10::ScalarType::Half,593+ at::Tensor scale = scale_b.size() > 0 ? processed_scale_b : at::Tensor();
452- "out_dtype must be Float32, BFloat16, or Float16, but got ", out_dtype.value(),594+ c10::optional<int64_t> scale_dtype = c10::nullopt; // scale_b.scalar_type()
453- OPS_ERROR(ErrCode::TYPE));
454- }
455 595 
596+ c10::optional<int64_t> output_dtype = ToIntType(out_dtype);
597+ c10::optional<at::Tensor> offset = c10::nullopt;
598+ c10::optional<int64_t> x1_dtype = c10::nullopt;
599+ c10::optional<int64_t> x2_dtype = c10::nullopt;
600+ c10::optional<int64_t> pertoken_scale_dtype = c10::nullopt;
601+ c10::OptionalIntArrayRef group_sizes = c10::nullopt;
602+ c10::optional<at::Tensor> y_scale = c10::nullopt;
456 603 
457- bool has_contraction = !contraction_dim.empty();604+ ///////////////////////////////////////////npu_quant_matmul//////////////////////////////////////
458- if (has_contraction) {605+ if (is_nz_format(x2)) {
459- int64_t c_dim_size = static_cast<int64_t>(contraction_dim.size());606+ static const bool is_quant_matmul_weight_nz_available = check_aclnn_kernel_available("aclnnQuantMatmulWeightNz");
460- TORCH_CHECK(c_dim_size == 2, "contraction_dim must have exactly 2 elements", OPS_ERROR(ErrCode::PARAM));607+ TORCH_CHECK(
461- 608+ is_quant_matmul_weight_nz_available,
462- int64_t a_contract_idx = contraction_dim[0];609+ "Get aclnnQuantMatmulWeightNz or aclnnQuantMatmulWeightNzGetWorkspaceSize failed, only "
463- int64_t b_contract_idx = contraction_dim[1];610+ "aclnnQuantMatmulWeightNz support X2's format is nz, please upgrade CANN.",
464- int64_t a_dim_size = mat_a.size(a_contract_idx);611+ OPS_ERROR(ErrCode::PARAM));
465- int64_t b_dim_size = mat_b.size(b_contract_idx);612+ } else {
466- 613+ static const bool is_quant_matmul_v5_available = check_aclnn_kernel_available("aclnnQuantMatmulV5");
467- TORCH_CHECK(614+ TORCH_CHECK(
468- a_dim_size == b_dim_size,615+ is_quant_matmul_v5_available,
469- "dimension mismatch for matrix multiplication (",616+ "Get aclnnQuantMatmulV5 or aclnnQuantMatmulV5 failed, only "
470- "A shape=", mat_a.size(0), "x", mat_a.size(1),617+ "aclnnQuantMatmulV5 support A8W4, please upgrade CANN.",
471- ", B shape=", mat_b.size(0), "x", mat_b.size(1), ")",618+ OPS_ERROR(ErrCode::TYPE));
472- " contraction axis A: ", a_contract_idx,619+ }
473- ", contraction axis B: ", b_contract_idx,620+ bool is_a8W4_int = x1.dtype() == at::kChar && x2.dtype() == at::kInt;
474- OPS_ERROR(ErrCode::PARAM));621+ bool is_a8W4_float = x1.dtype() == at::kFloat8_e4m3fn && x2.dtype() == at::kFloat;
475- } else {622+ at::IntArrayRef group_size_list = group_sizes.value_or(at::IntArrayRef{});
476- int64_t a_col = mat_a.size(1);623+ int64_t group_size = check_and_get_groups(group_size_list);
477- int64_t b_row = mat_b.size(0);624+ bool is_a4w4 = x1.dtype() == at::kInt && x2.dtype() == at::kInt;
478- TORCH_CHECK(625+ bool trans_x1 = is_transpose_last_two_dims(x1);
479- a_col == b_row,626+ bool trans_x2 = is_transpose_last_two_dims(x2);
480- "matrix multiplication dimension mismatch (",627+ auto x1_dim_num = x1.dim();
481- "A: ", mat_a.size(0), "x", mat_a.size(1),628+ auto x2_dim_num = x2.dim();
482- ", B: ", mat_b.size(0), "x", mat_b.size(1), ")",629+ auto x2_n_dim = (is_a4w4 && !trans_x2) ? x2.size(x2_dim_num - 1) * INT4_NUMS_IN_INT32 : x2.size(x2_dim_num - 1);
483- OPS_ERROR(ErrCode::PARAM)
484- );
485- }
486- 
487- 
488- 
489- //////////////rowwise transfer//////////////
490- at::Tensor processed_scale_a = scale_a[0];
491- at::Tensor processed_scale_b = scale_b[0];
492- if (mat_a.scalar_type() == c10::ScalarType::Float8_e4m3fn
493- && mat_b.scalar_type() == c10::ScalarType::Float8_e4m3fn
494- && scale_a[0].scalar_type() == c10::ScalarType::Float
495- && scale_b[0].scalar_type() == c10::ScalarType::Float) {
496- if (scale_a[0].dim() == 2 && scale_a[0].sizes()[1] == 1 && scale_b[0].dim() == 2 && scale_b[0].sizes()[0] == 1) {
497- TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be contiguous in last dim");
498- TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be contiguous in first dim");
499- processed_scale_a = scale_a[0].squeeze(-1).contiguous();
500- processed_scale_b = scale_b[0].squeeze(0);
501- }
502- }
503- 
504- //////////////parameters transfer//////////////
505- at::Tensor x1 = mat_a;
506- at::Tensor x2 = mat_b;
507- c10::optional<at::Tensor> pertoken_scale = scale_a.size() > 0 ? c10::make_optional(processed_scale_a) : c10::nullopt;
508- at::Tensor scale = scale_b.size() > 0 ? processed_scale_b : at::Tensor();
509- c10::optional<int64_t> scale_dtype = c10::nullopt; // scale_b.scalar_type()
510- 
511- c10::optional<int64_t> output_dtype = ToIntType(out_dtype);
512- c10::optional<at::Tensor> offset = c10::nullopt;
513- c10::optional<int64_t> x1_dtype = c10::nullopt;
514- c10::optional<int64_t> x2_dtype = c10::nullopt;
515- c10::optional<int64_t> pertoken_scale_dtype = c10::nullopt;
516- c10::OptionalIntArrayRef group_sizes = c10::nullopt;
517- c10::optional<at::Tensor> y_scale = c10::nullopt;
518- 
519- 
520- ///////////////////////////////////////////npu_quant_matmul//////////////////////////////////////
521- if (is_nz_format(x2)) {
522- static const bool is_quant_matmul_weight_nz_available = check_aclnn_kernel_available("aclnnQuantMatmulWeightNz");
523- TORCH_CHECK(is_quant_matmul_weight_nz_available,
524- "Get aclnnQuantMatmulWeightNz or aclnnQuantMatmulWeightNzGetWorkspaceSize failed, only "
525- "aclnnQuantMatmulWeightNz support X2's format is nz, please upgrade CANN.",
526- OPS_ERROR(ErrCode::PARAM));
527- } else {
528- static const bool is_quant_matmul_v5_available = check_aclnn_kernel_available("aclnnQuantMatmulV5");
529- TORCH_CHECK(is_quant_matmul_v5_available,
530- "Get aclnnQuantMatmulV5 or aclnnQuantMatmulV5 failed, only "
531- "aclnnQuantMatmulV5 support A8W4, please upgrade CANN.",
532- OPS_ERROR(ErrCode::TYPE));
533- }
534- bool is_a8W4_int = x1.dtype() == at::kChar && x2.dtype() == at::kInt;
535- bool is_a8W4_float = x1.dtype() == at::kFloat8_e4m3fn && x2.dtype() == at::kFloat;
536- at::IntArrayRef group_size_list = group_sizes.value_or(at::IntArrayRef{});
537- int64_t group_size = check_and_get_groups(group_size_list);
538- bool is_a4w4 = x1.dtype() == at::kInt && x2.dtype() == at::kInt;
539- bool trans_x1 = is_transpose_last_two_dims(x1);
540- bool trans_x2 = is_transpose_last_two_dims(x2);
541- auto x1_dim_num = x1.dim();
542- auto x2_dim_num = x2.dim();
543- auto x2_n_dim = (is_a4w4 && !trans_x2) ? x2.size(x2_dim_num - 1) * INT4_NUMS_IN_INT32 : x2.size(x2_dim_num - 1);
544 630 
545#if VERSION_BETWEEN(V2R1, V2R7)631#if VERSION_BETWEEN(V2R1, V2R7)
546- bool mxfp4_valid = x1_dtype.has_value() && x2_dtype.has_value() &&632+ bool mxfp4_valid = x1_dtype.has_value() && x2_dtype.has_value() &&
547- x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) &&633+ x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) &&
548- x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);634+ x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
549#endif635#endif
550#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)636#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)
551- bool mxfp4_valid = false;637+ bool mxfp4_valid = false;
552- if (x1_dtype.has_value()) {638+ if (x1_dtype.has_value()) {
553- mxfp4_valid = x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);639+ mxfp4_valid = x1_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
554- } else {640+ } else {
555- mxfp4_valid = x1.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;641+ mxfp4_valid = x1.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
556- }642+ }
557- if (x2_dtype.has_value()) {643+ if (x2_dtype.has_value()) {
558- mxfp4_valid = mxfp4_valid && x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);644+ mxfp4_valid = mxfp4_valid && x2_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1);
559- } else {645+ } else {
560- mxfp4_valid = mxfp4_valid && x2.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;646+ mxfp4_valid = mxfp4_valid && x2.scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
561- }647+ }
562#endif648#endif
563 649 
564- c10::SmallVector<int64_t, SIZE> output_size;650+ c10::SmallVector<int64_t, SIZE> output_size;
565- if (is_a8W4_int) {651+ if (is_a8W4_int) {
566- output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};652+ output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
567- } else if (is_a8W4_float) {653+ } else if (is_a8W4_float) {
568- if (trans_x2) {654+ if (trans_x2) {
569- output_size = {x1.sizes()[0], x2.sizes()[1]};655+ output_size = {x1.sizes()[0], x2.sizes()[1]};
570- } else {
571- output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
572- }
573 } else {656 } else {
574- std::vector<uint64_t> batch_record;657+ output_size = {x1.sizes()[0], x2.sizes()[1] * INT4_NUMS_IN_INT32};
575- uint64_t batch_val = infer_out_batch_shape(x1, x2, batch_record);
576- const at::Tensor long_tensor = x1_dim_num > x2_dim_num ? x1 : x2;
577- output_size = op_infer::array_to_small_vector(long_tensor.sizes());
578- if (mxfp4_valid) {
579- TORCH_CHECK(x1.dim() >= 2 && x1.dim() <= 6,
580- "x1 dim num should be 2 ~ 6, please check x1 dim num. Actual x1 dim = ", x1.dim(),
581- OPS_ERROR(ErrCode::PARAM));
582- TORCH_CHECK(x2.dim() >= 2 && x2.dim() <= 6,
583- "x2 dim num should be 2 ~ 6, please check x2 dim num. Actual x2 dim = ", x2.dim(),
584- OPS_ERROR(ErrCode::PARAM));
585- int64_t x1_size_last_second = x1.sizes()[x1_dim_num - LAST_SECOND_DIM_INDEX];
586- int64_t x2_size_last = x2.sizes()[x2_dim_num - 1];
587- int64_t real_m = !trans_x1 ? x1_size_last_second : x1_size_last_second * FP4_IN_INT8;
588- int64_t real_n = trans_x2 ? x2_size_last : x2_size_last * FP4_IN_INT8;
589- output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = real_m;
590- output_size[long_tensor.dim() - 1] = real_n;
591- } else {
592- output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = x1.size(x1_dim_num - LAST_SECOND_DIM_INDEX);
593- output_size[long_tensor.dim() - 1] = x2_n_dim;
594- }
595- for (int64_t i = 0; i < long_tensor.dim() - LAST_SECOND_DIM_INDEX; i++) {
596- output_size[i] = static_cast<int64_t>(batch_record[i]);
597- }
598 }658 }
599- c10::TensorOptions options;659+ } else {
600- aclDataType output_acltype = ACL_INT8;660+ std::vector<uint64_t> batch_record;
601- if (!output_dtype.has_value()) {661+ uint64_t batch_val = infer_out_batch_shape(x1, x2, batch_record);
602- options = x1.options().dtype(at::kChar);662+ const at::Tensor long_tensor = x1_dim_num > x2_dim_num ? x1 : x2;
663+ output_size = op_infer::array_to_small_vector(long_tensor.sizes());
664+ if (mxfp4_valid) {
665+ TORCH_CHECK(
666+ x1.dim() >= 2 && x1.dim() <= 6,
667+ "x1 dim num should be 2 ~ 6, please check x1 dim num. Actual x1 dim = ",
668+ x1.dim(),
669+ OPS_ERROR(ErrCode::PARAM));
670+ TORCH_CHECK(
671+ x2.dim() >= 2 && x2.dim() <= 6,
672+ "x2 dim num should be 2 ~ 6, please check x2 dim num. Actual x2 dim = ",
673+ x2.dim(),
674+ OPS_ERROR(ErrCode::PARAM));
675+ int64_t x1_size_last_second = x1.sizes()[x1_dim_num - LAST_SECOND_DIM_INDEX];
676+ int64_t x2_size_last = x2.sizes()[x2_dim_num - 1];
677+ int64_t real_m = !trans_x1 ? x1_size_last_second : x1_size_last_second * FP4_IN_INT8;
678+ int64_t real_n = trans_x2 ? x2_size_last : x2_size_last * FP4_IN_INT8;
679+ output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = real_m;
680+ output_size[long_tensor.dim() - 1] = real_n;
603 } else {681 } else {
604- output_acltype = c10_npu::GetAclDataType(output_dtype.value());682+ output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = x1.size(x1_dim_num - LAST_SECOND_DIM_INDEX);
605- options = x1.options().dtype(npu_preparation::convert_to_scalar_type(output_acltype));683+ output_size[long_tensor.dim() - 1] = x2_n_dim;
606 }684 }
607- at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);685+ for (int64_t i = 0; i < long_tensor.dim() - LAST_SECOND_DIM_INDEX; i++) {
686+ output_size[i] = static_cast<int64_t>(batch_record[i]);
687+ }
688+ }
689+ c10::TensorOptions options;
690+ aclDataType output_acltype = ACL_INT8;
691+ if (!output_dtype.has_value()) {
692+ options = x1.options().dtype(at::kChar);
693+ } else {
694+ output_acltype = c10_npu::GetAclDataType(output_dtype.value());
695+ options = x1.options().dtype(npu_preparation::convert_to_scalar_type(output_acltype));
696+ }
697+ at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options);
608 698 
609- const at::Tensor &offset_real = offset.value_or(at::Tensor());699+ const at::Tensor& offset_real = offset.value_or(at::Tensor());
610- const at::Tensor &pertoken_scale_real = pertoken_scale.value_or(at::Tensor());700+ const at::Tensor& pertoken_scale_real = pertoken_scale.value_or(at::Tensor());
611- const at::Tensor &bias_real = bias.value_or(at::Tensor());701+ const at::Tensor& bias_real = bias.value_or(at::Tensor());
612- bool transpose1 = false;702+ bool transpose1 = false;
613- bool transpose2 = false;703+ bool transpose2 = false;
614 704 
615- TensorWrapper x1_wrapper = make_wrapper(x1, x1_dtype);705+ TensorWrapper x1_wrapper = make_wrapper(x1, x1_dtype);
616- TensorWrapper x2_wrapper = make_wrapper(x2, x2_dtype);706+ TensorWrapper x2_wrapper = make_wrapper(x2, x2_dtype);
617- TensorWrapper x1_scale_wrapper = make_wrapper(pertoken_scale_real, pertoken_scale_dtype);707+ TensorWrapper x1_scale_wrapper = make_wrapper(pertoken_scale_real, pertoken_scale_dtype);
618- TensorWrapper x2_scale_wrapper = make_wrapper(scale, scale_dtype);708+ TensorWrapper x2_scale_wrapper = make_wrapper(scale, scale_dtype);
619- TensorWrapper result_wrapper = make_wrapper(result, output_dtype);709+ TensorWrapper result_wrapper = make_wrapper(result, output_dtype);
620- at::Tensor x1_offset = at::empty({0}, options);710+ at::Tensor x1_offset = at::empty({0}, options);
621- at::Tensor x2_offset = at::Tensor();711+ at::Tensor x2_offset = at::Tensor();
622- at::Tensor y_offset = at::empty({0}, options);712+ at::Tensor y_offset = at::empty({0}, options);
623- if (is_a8W4_int) { // Only A8W4 int needs y_offset713+ if (is_a8W4_int) { // Only A8W4 int needs y_offset
624- y_offset = offset_real;714+ y_offset = offset_real;
715+ } else {
716+ x2_offset = offset_real;
717+ }
718+ 
719+ bool use_aclnn_v5 =
720+ x1_dtype.has_value() || (x1.dtype() != at::kInt && x1.dtype() != at::kChar) || is_a8W4_float || is_a8W4_int;
721+ 
722+ aclDataType pertoken_scale_dtype_real = pertoken_scale_dtype.has_value()
723+ ? c10_npu::GetAclDataType(pertoken_scale_dtype.value())
724+ : (pertoken_scale.has_value() ? c10_npu::GetAclDataType(static_cast<int64_t>(pertoken_scale_real.scalar_type()))
725+ : aclDataType::ACL_INT8);
726+ bool need_check_trans = pertoken_scale.has_value() &&
727+ (((pertoken_scale_real.dim() == x1.dim() && scale.dim() == x2.dim()) ||
728+ pertoken_scale_dtype_real == aclDataType::ACL_FLOAT8_E8M0) &&
729+ (pertoken_scale_real.dim() >= 2 && scale.dim() >= 2)) &&
730+ !(is_a8W4_float || is_a8W4_int);
731+ if (need_check_trans) {
732+ int64_t dim_x1 = x1.dim() - 2; // check the last 2 dim
733+ int64_t dim_x2 = x2.dim() - 2; // check the last 2 dim
734+ int64_t dim_x1_scale = 0;
735+ int64_t dim_x2_scale = 0;
736+ if (pertoken_scale_dtype_real != aclDataType::ACL_FLOAT8_E8M0) {
737+ dim_x1_scale = pertoken_scale_real.dim() - 2; // check the last 2 dim in GB/BB
738+ dim_x2_scale = scale.dim() - 2; // check the last 2 dim in GB/BB
739+ }
740+ TORCH_CHECK(
741+ is_x_scale_same_transpose(x1, pertoken_scale_real, dim_x1, dim_x1_scale),
742+ "Input x1 tensor and pertoken_scale tensor's transpose are not same, please check input.",
743+ OPS_ERROR(ErrCode::PARAM));
744+ TORCH_CHECK(
745+ is_x_scale_same_transpose(x2, scale, dim_x2, dim_x2_scale),
746+ "Input x2 tensor and scale tensor's transpose are not same, please check input.",
747+ OPS_ERROR(ErrCode::PARAM));
748+ }
749+ 
750+ bool use_trans_quant_param = scale.dtype() == at::kFloat && !pertoken_scale.has_value() &&
751+ (output_acltype != ACL_BF16 || use_aclnn_v5) && output_acltype != ACL_INT32;
752+ if (use_trans_quant_param) {
753+ const at::Tensor quant_param = op_api::npu_trans_quant_param(scale, offset);
754+ if (is_nz_format(x2)) {
755+ EXEC_NPU_CMD(
756+ aclnnQuantMatmulWeightNz,
757+ x1_wrapper,
758+ x2_wrapper,
759+ pertoken_scale_real,
760+ quant_param,
761+ y_scale,
762+ x1_offset,
763+ x2_offset,
764+ y_offset,
765+ bias_real,
766+ transpose1,
767+ transpose2,
768+ group_size,
769+ result_wrapper);
625 } else {770 } else {
626- x2_offset = offset_real;771+ EXEC_NPU_CMD(
772+ aclnnQuantMatmulV5,
773+ x1_wrapper,
774+ x2_wrapper,
775+ pertoken_scale_real,
776+ quant_param,
777+ y_scale,
778+ x1_offset,
779+ x2_offset,
780+ y_offset,
781+ bias_real,
782+ transpose1,
783+ transpose2,
784+ group_size,
785+ result_wrapper);
627 }786 }
628- 787+ } else {
629- bool use_aclnn_v5 = x1_dtype.has_value() || (x1.dtype() != at::kInt && x1.dtype() != at::kChar) ||788+ if (!is_a4w4 && is_nz_format(x2)) {
630- is_a8W4_float || is_a8W4_int;789+ EXEC_NPU_CMD(
631- 790+ aclnnQuantMatmulWeightNz,
632- aclDataType pertoken_scale_dtype_real = pertoken_scale_dtype.has_value()791+ x1_wrapper,
633- ? c10_npu::GetAclDataType(pertoken_scale_dtype.value())792+ x2_wrapper,
634- : (pertoken_scale.has_value()793+ x1_scale_wrapper,
635- ? c10_npu::GetAclDataType(static_cast<int64_t>(pertoken_scale_real.scalar_type()))794+ x2_scale_wrapper,
636- : aclDataType::ACL_INT8);795+ y_scale,
637- bool need_check_trans = pertoken_scale.has_value()796+ x1_offset,
638- && (((pertoken_scale_real.dim() == x1.dim() && scale.dim() == x2.dim())797+ x2_offset,
639- || pertoken_scale_dtype_real == aclDataType::ACL_FLOAT8_E8M0)798+ y_offset,
640- && (pertoken_scale_real.dim() >= 2 && scale.dim() >= 2))799+ bias_real,
641- && !(is_a8W4_float || is_a8W4_int);800+ transpose1,
642- if (need_check_trans) {801+ transpose2,
643- int64_t dim_x1 = x1.dim() - 2; // check the last 2 dim802+ group_size,
644- int64_t dim_x2 = x2.dim() - 2; // check the last 2 dim803+ result_wrapper);
645- int64_t dim_x1_scale = 0;
646- int64_t dim_x2_scale = 0;
647- if (pertoken_scale_dtype_real != aclDataType::ACL_FLOAT8_E8M0) {
648- dim_x1_scale = pertoken_scale_real.dim() - 2; // check the last 2 dim in GB/BB
649- dim_x2_scale = scale.dim() - 2; // check the last 2 dim in GB/BB
650- }
651- TORCH_CHECK(is_x_scale_same_transpose(x1, pertoken_scale_real, dim_x1, dim_x1_scale),
652- "Input x1 tensor and pertoken_scale tensor's transpose are not same, please check input.",
653- OPS_ERROR(ErrCode::PARAM));
654- TORCH_CHECK(is_x_scale_same_transpose(x2, scale, dim_x2, dim_x2_scale),
655- "Input x2 tensor and scale tensor's transpose are not same, please check input.",
656- OPS_ERROR(ErrCode::PARAM));
657- }
658- 
659- bool use_trans_quant_param = scale.dtype() == at::kFloat && !pertoken_scale.has_value() &&
660- (output_acltype != ACL_BF16 || use_aclnn_v5) && output_acltype != ACL_INT32;
661- if (use_trans_quant_param) {
662- const at::Tensor quant_param = op_api::npu_trans_quant_param(scale, offset);
663- if (is_nz_format(x2)) {
664- EXEC_NPU_CMD(aclnnQuantMatmulWeightNz, x1_wrapper, x2_wrapper, pertoken_scale_real, quant_param, y_scale,
665- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);
666- } else {
667- EXEC_NPU_CMD(aclnnQuantMatmulV5, x1_wrapper, x2_wrapper, pertoken_scale_real, quant_param, y_scale,
668- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);
669- }
670 } else {804 } else {
671- if (!is_a4w4 && is_nz_format(x2)) {805+ EXEC_NPU_CMD(
672- EXEC_NPU_CMD(aclnnQuantMatmulWeightNz, x1_wrapper, x2_wrapper, x1_scale_wrapper, x2_scale_wrapper, y_scale,806+ aclnnQuantMatmulV5,
673- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);807+ x1_wrapper,
674- } else {808+ x2_wrapper,
675- EXEC_NPU_CMD(aclnnQuantMatmulV5, x1_wrapper, x2_wrapper, x1_scale_wrapper, x2_scale_wrapper, y_scale,809+ x1_scale_wrapper,
676- x1_offset, x2_offset, y_offset, bias_real, transpose1, transpose2, group_size, result_wrapper);810+ x2_scale_wrapper,
677- }811+ y_scale,
812+ x1_offset,
813+ x2_offset,
814+ y_offset,
815+ bias_real,
816+ transpose1,
817+ transpose2,
818+ group_size,
819+ result_wrapper);
678 }820 }
821+ }
679 822 
680- return result;823+ return result;
681}824}
682 825 
683-}826+} // namespace op_api
Mop_plugin/ops/opapi/ScatterAddKernelNpuOpApi.cpp+15-31
@@ -17,45 +17,29 @@
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- 
21namespace op_api {20namespace op_api {
22using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
23 22 
24-at::Tensor scatter_add(23+at::Tensor scatter_add(const at::Tensor& self, int64_t dim, const at::Tensor& index, const at::Tensor& src) {
25- const at::Tensor& self,24+ DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add(self, dim, index, src));
26- int64_t dim,25+ auto selfClone = self.clone(at::MemoryFormat::Contiguous);
27- const at::Tensor& index,26+ npu_preparation::CheckMemory({selfClone, index, src}, {selfClone});
28- const at::Tensor& src)27+ EXEC_NPU_CMD(aclnnScatterAdd, selfClone, dim, index, src, selfClone);
29-{28+ return selfClone;
30- DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add(self, dim, index, src));
31- auto selfClone = self.clone(at::MemoryFormat::Contiguous);
32- npu_preparation::CheckMemory({selfClone, index, src}, {selfClone});
33- EXEC_NPU_CMD(aclnnScatterAdd, selfClone, dim, index, src, selfClone);
34- return selfClone;
35}29}
36 30 
37-at::Tensor& scatter_add_(31+at::Tensor& scatter_add_(at::Tensor& self, int64_t dim, const at::Tensor& index, const at::Tensor& src) {
38- at::Tensor& self,32+ DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add_(self, dim, index, src));
39- int64_t dim,33+ npu_preparation::CheckMemory({self, index, src}, {self});
40- const at::Tensor& index,34+ EXEC_NPU_CMD(aclnnScatterAdd, self, dim, index, src, self);
41- const at::Tensor& src)35+ return self;
42-{
43- DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add_(self, dim, index, src));
44- npu_preparation::CheckMemory({self, index, src}, {self});
45- EXEC_NPU_CMD(aclnnScatterAdd, self, dim, index, src, self);
46- return self;
47}36}
48 37 
49#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)38#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)
50-at::Tensor scatter_add(39+at::Tensor scatter_add(const at::Tensor& self, at::Dimname dim, const at::Tensor& index, const at::Tensor& src) {
51- const at::Tensor& self,40+ DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add(self, dim, index, src));
52- at::Dimname dim,41+ return op_api::scatter_add(self, dimname_to_position(self, dim), index, src);
53- const at::Tensor& index,
54- const at::Tensor& src)
55-{
56- DO_COMPATIBILITY(aclnnScatterAdd, acl_op::scatter_add(self, dim, index, src));
57- return op_api::scatter_add(self, dimname_to_position(self, dim), index, src);
58}42}
59#endif43#endif
60 44 
61-}45+} // namespace op_api
Mop_plugin/ops/opapi/ScatterKernelNpuOpApi.cpp+68-58
@@ -21,76 +21,86 @@ namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23// reduce value must be "add" or "multiply"23// reduce value must be "add" or "multiply"
24-static inline bool reduce_valid(c10::string_view reduce)24+static inline bool reduce_valid(c10::string_view reduce) {
25-{25+ return (reduce == "add" || reduce == "multiply");
26- return (reduce == "add" || reduce == "multiply");
27}26}
28 27 
29-static int64_t get_reduce(c10::string_view reduce)28+static int64_t get_reduce(c10::string_view reduce) {
30-{29+ if (reduce == "add") {
31- if (reduce == "add") {30+ return 1;
32- return 1;31+ } else if (reduce == "multiply") {
33- } else if (reduce == "multiply") {32+ return 2;
34- return 2;33+ }
35- }34+ return 0;
36- return 0;
37}35}
38 36 
39-at::Tensor& scatter_out(const at::Tensor& self, int64_t dim, const at::Tensor& index,37+at::Tensor& scatter_out(
40- const at::Tensor& src, at::Tensor& out)38+ const at::Tensor& self,
41-{39+ int64_t dim,
42- DO_COMPATIBILITY(aclnnScatter, acl_op::scatter_out(self, dim, index, src, out));40+ const at::Tensor& index,
43- npu_preparation::check_tensor({self, src, index}, out, self);41+ const at::Tensor& src,
44- int64_t reduction = 0;42+ at::Tensor& out) {
45- EXEC_NPU_CMD(aclnnScatter, self, dim, index, src, reduction, out);43+ DO_COMPATIBILITY(aclnnScatter, acl_op::scatter_out(self, dim, index, src, out));
46- return out;44+ npu_preparation::check_tensor({self, src, index}, out, self);
45+ int64_t reduction = 0;
46+ EXEC_NPU_CMD(aclnnScatter, self, dim, index, src, reduction, out);
47+ return out;
47}48}
48 49 
49-at::Tensor& scatter_out(const at::Tensor& self, int64_t dim, const at::Tensor& index,50+at::Tensor& scatter_out(
50- const at::Tensor& src, c10::string_view reduce, at::Tensor& out)51+ const at::Tensor& self,
51-{52+ int64_t dim,
52- npu_preparation::check_tensor({self, src, index}, out, self);53+ const at::Tensor& index,
53- TORCH_CHECK(reduce_valid(reduce), "Reduce should be either add or multiply", OPS_ERROR(ErrCode::PARAM));54+ const at::Tensor& src,
54- int64_t reduction = get_reduce(reduce);55+ c10::string_view reduce,
55- EXEC_NPU_CMD(aclnnScatter, self, dim, index, src, reduction, out);56+ at::Tensor& out) {
56- return out;57+ npu_preparation::check_tensor({self, src, index}, out, self);
58+ TORCH_CHECK(reduce_valid(reduce), "Reduce should be either add or multiply", OPS_ERROR(ErrCode::PARAM));
59+ int64_t reduction = get_reduce(reduce);
60+ EXEC_NPU_CMD(aclnnScatter, self, dim, index, src, reduction, out);
61+ return out;
57}62}
58 63 
59-at::Tensor& scatter_out(const at::Tensor& self, int64_t dim, const at::Tensor& index,64+at::Tensor& scatter_out(
60- const at::Scalar& value, at::Tensor& out)65+ const at::Tensor& self,
61-{66+ int64_t dim,
62- DO_COMPATIBILITY(aclnnScatterValue, acl_op::scatter_out(self, dim, index, value, out));67+ const at::Tensor& index,
63- npu_preparation::check_tensor({self, index}, out, self);68+ const at::Scalar& value,
64- int64_t reduction = 0;69+ at::Tensor& out) {
65- EXEC_NPU_CMD(aclnnScatterValue, self, dim, index, value, reduction, out);70+ DO_COMPATIBILITY(aclnnScatterValue, acl_op::scatter_out(self, dim, index, value, out));
66- return out;71+ npu_preparation::check_tensor({self, index}, out, self);
72+ int64_t reduction = 0;
73+ EXEC_NPU_CMD(aclnnScatterValue, self, dim, index, value, reduction, out);
74+ return out;
67}75}
68 76 
69-at::Tensor& scatter_out(const at::Tensor& self, int64_t dim, const at::Tensor& index,77+at::Tensor& scatter_out(
70- const at::Scalar& value, c10::string_view reduce, at::Tensor& out)78+ const at::Tensor& self,
71-{79+ int64_t dim,
72- npu_preparation::check_tensor({self, index}, out, self);80+ const at::Tensor& index,
73- TORCH_CHECK(reduce_valid(reduce), "Reduce should be either add or multiply", OPS_ERROR(ErrCode::PARAM));81+ const at::Scalar& value,
74- int64_t reduction = get_reduce(reduce);82+ c10::string_view reduce,
75- EXEC_NPU_CMD(aclnnScatterValue, self, dim, index, value, reduction, out);83+ at::Tensor& out) {
76- return out;84+ npu_preparation::check_tensor({self, index}, out, self);
85+ TORCH_CHECK(reduce_valid(reduce), "Reduce should be either add or multiply", OPS_ERROR(ErrCode::PARAM));
86+ int64_t reduction = get_reduce(reduce);
87+ EXEC_NPU_CMD(aclnnScatterValue, self, dim, index, value, reduction, out);
88+ return out;
77}89}
78 90 
79-at::Tensor &scatter_(at::Tensor &self, int64_t dim, const at::Tensor &index, const at::Tensor &src)91+at::Tensor& scatter_(at::Tensor& self, int64_t dim, const at::Tensor& index, const at::Tensor& src) {
80-{92+ DO_COMPATIBILITY(aclnnInplaceScatter, acl_op::scatter_(self, dim, index, src));
81- DO_COMPATIBILITY(aclnnInplaceScatter, acl_op::scatter_(self, dim, index, src));93+ npu_preparation::check_tensor({self, src, index}, self, self);
82- npu_preparation::check_tensor({self, src, index}, self, self);94+ int64_t reduction = 0;
83- int64_t reduction = 0;95+ EXEC_NPU_CMD(aclnnInplaceScatter, self, dim, index, src, reduction);
84- EXEC_NPU_CMD(aclnnInplaceScatter, self, dim, index, src, reduction);96+ return self;
85- return self;
86}97}
87 98 
88-at::Tensor &scatter_(at::Tensor &self, int64_t dim, const at::Tensor &index, const at::Scalar& value)99+at::Tensor& scatter_(at::Tensor& self, int64_t dim, const at::Tensor& index, const at::Scalar& value) {
89-{100+ DO_COMPATIBILITY(aclnnInplaceScatterValue, acl_op::scatter_(self, dim, index, value));
90- DO_COMPATIBILITY(aclnnInplaceScatterValue, acl_op::scatter_(self, dim, index, value));101+ npu_preparation::check_tensor({self, index}, self, self);
91- npu_preparation::check_tensor({self, index}, self, self);102+ int64_t reduction = 0;
92- int64_t reduction = 0;103+ EXEC_NPU_CMD(aclnnInplaceScatterValue, self, dim, index, value, reduction);
93- EXEC_NPU_CMD(aclnnInplaceScatterValue, self, dim, index, value, reduction);104+ return self;
94- return self;
95-}
96}105}
106+} // namespace op_api
Mop_plugin/ops/opapi/ScatterListKernelNpuOpApi.cpp+33-35
@@ -18,42 +18,40 @@
18#include "op_plugin/utils/op_api_common.h"18#include "op_plugin/utils/op_api_common.h"
19 19 
20namespace op_api {20namespace op_api {
21- using npu_preparation = at_npu::native::OpPreparation;21+using npu_preparation = at_npu::native::OpPreparation;
22 22 
23- std::vector<at::Tensor> npu_scatter_list(23+std::vector<at::Tensor> npu_scatter_list(
24- at::TensorList self,24+ at::TensorList self,
25- const at::Tensor &indice,25+ const at::Tensor& indice,
26- const at::Tensor &updates,26+ const at::Tensor& updates,
27- const c10::optional<at::Tensor> &mask,27+ const c10::optional<at::Tensor>& mask,
28- c10::string_view reduce,28+ c10::string_view reduce,
29- int64_t axis)29+ int64_t axis) {
30- {30+ std::string reduce_str = std::string(reduce);
31- std::string reduce_str = std::string(reduce);31+ char* reduce_ptr = const_cast<char*>(reduce_str.c_str());
32- char *reduce_ptr = const_cast<char *>(reduce_str.c_str());32+ // The attribute 'reduce' of ScatterList only supports setting it to 'update'.
33- // The attribute 'reduce' of ScatterList only supports setting it to 'update'.33+ std::vector<at::Tensor> result;
34- std::vector<at::Tensor> result;34+ for (const at::Tensor& tensor : self) {
35- for (const at::Tensor &tensor : self) {35+ result.push_back(tensor.clone());
36- result.push_back(tensor.clone());36+ }
37- }37+ at::TensorList result_ = at::TensorList(result);
38- at::TensorList result_ = at::TensorList(result);
39 38 
40- EXEC_NPU_CMD(aclnnScatterList, result_, indice, updates, mask, reduce_ptr, axis);39+ EXEC_NPU_CMD(aclnnScatterList, result_, indice, updates, mask, reduce_ptr, axis);
41- 
42- return result;
43- }
44- 
45- void npu_scatter_list_(
46- at::TensorList self,
47- const at::Tensor &indice,
48- const at::Tensor &updates,
49- const c10::optional<at::Tensor> &mask,
50- c10::string_view reduce,
51- int64_t axis)
52- {
53- std::string reduce_str = std::string(reduce);
54- char *reduce_ptr = const_cast<char *>(reduce_str.c_str());
55- EXEC_NPU_CMD(aclnnScatterList, self, indice, updates, mask, reduce_ptr, axis);
56- return;
57- }
58 40 
41+ return result;
59}42}
43+ 
44+void npu_scatter_list_(
45+ at::TensorList self,
46+ const at::Tensor& indice,
47+ const at::Tensor& updates,
48+ const c10::optional<at::Tensor>& mask,
49+ c10::string_view reduce,
50+ int64_t axis) {
51+ std::string reduce_str = std::string(reduce);
52+ char* reduce_ptr = const_cast<char*>(reduce_str.c_str());
53+ EXEC_NPU_CMD(aclnnScatterList, self, indice, updates, mask, reduce_ptr, axis);
54+ return;
55+}
56+ 
57+} // namespace op_api
Mop_plugin/ops/opapi/ScatterNdUpdateKernelNpuOpApi.cpp+9-17
@@ -18,25 +18,17 @@
18#include "op_plugin/utils/op_api_common.h"18#include "op_plugin/utils/op_api_common.h"
19 19 
20namespace op_api {20namespace op_api {
21- using npu_preparation = at_npu::native::OpPreparation;21+using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor npu_scatter_nd_update(23+at::Tensor npu_scatter_nd_update(const at::Tensor& self, const at::Tensor& indices, const at::Tensor& updates) {
24- const at::Tensor &self,24+ at::Tensor result = self.clone();
25- const at::Tensor &indices,25+ EXEC_NPU_CMD(aclnnScatterNdUpdate, result, indices, updates);
26- const at::Tensor &updates)26+ return result;
27-{
28- at::Tensor result = self.clone();
29- EXEC_NPU_CMD(aclnnScatterNdUpdate, result, indices, updates);
30- return result;
31}27}
32 28 
33-at::Tensor &npu_scatter_nd_update_(29+at::Tensor& npu_scatter_nd_update_(at::Tensor& self, const at::Tensor& indices, const at::Tensor& updates) {
34- at::Tensor &self,30+ EXEC_NPU_CMD(aclnnScatterNdUpdate, self, indices, updates);
35- const at::Tensor &indices,31+ return self;
36- const at::Tensor &updates)
37-{
38- EXEC_NPU_CMD(aclnnScatterNdUpdate, self, indices, updates);
39- return self;
40}32}
41 33 
42-}34+} // namespace op_api
Mop_plugin/ops/opapi/ScatterPaKvCacheNpuOpApi.cpp+46-24
@@ -29,19 +29,30 @@ std::tuple<at::Tensor, at::Tensor> npu_scatter_pa_kv_cache_functional(
29 const c10::optional<at::Tensor>& compress_lens,29 const c10::optional<at::Tensor>& compress_lens,
30 const c10::optional<at::Tensor>& compress_seq_offsets,30 const c10::optional<at::Tensor>& compress_seq_offsets,
31 const c10::optional<at::Tensor>& seq_lens,31 const c10::optional<at::Tensor>& seq_lens,
32- c10::optional<c10::string_view> cache_mode)32+ c10::optional<c10::string_view> cache_mode) {
33-{33+ char* cache_mode_ptr = cache_mode.has_value() ? const_cast<char*>(cache_mode.value().data()) : nullptr;
34- char* cache_mode_ptr = cache_mode.has_value() ? const_cast<char *>(cache_mode.value().data()) : nullptr;34+ char* scatter_mode = "None";
35- char* scatter_mode = "None";35+ c10::SmallVector<int64_t, op_infer::SIZE> strides_size = {1, 1};
36- c10::SmallVector<int64_t, op_infer::SIZE> strides_size = {1, 1};36+ at::IntArrayRef strides = at::IntArrayRef(strides_size);
37- at::IntArrayRef strides = at::IntArrayRef(strides_size);37+ c10::SmallVector<int64_t, op_infer::SIZE> offsets_size = {0, 0};
38- c10::SmallVector<int64_t, op_infer::SIZE> offsets_size = {0, 0};38+ at::IntArrayRef offsets = at::IntArrayRef(offsets_size);
39- at::IntArrayRef offsets = at::IntArrayRef(offsets_size);39+ auto keyCacheClone = key_cache.clone(at::MemoryFormat::Contiguous);
40- auto keyCacheClone = key_cache.clone(at::MemoryFormat::Contiguous);40+ auto valueCacheClone = value_cache.clone(at::MemoryFormat::Contiguous);
41- auto valueCacheClone = value_cache.clone(at::MemoryFormat::Contiguous);41+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
42- EXEC_NPU_NO_FORMAT_CHECK_CMD(aclnnScatterPaKvCache, key, keyCacheClone, slot_mapping, value,42+ aclnnScatterPaKvCache,
43- valueCacheClone, compress_lens, compress_seq_offsets, seq_lens, cache_mode_ptr, scatter_mode, strides, offsets);43+ key,
44- return std::make_tuple(keyCacheClone, valueCacheClone);44+ keyCacheClone,
45+ slot_mapping,
46+ value,
47+ valueCacheClone,
48+ compress_lens,
49+ compress_seq_offsets,
50+ seq_lens,
51+ cache_mode_ptr,
52+ scatter_mode,
53+ strides,
54+ offsets);
55+ return std::make_tuple(keyCacheClone, valueCacheClone);
45}56}
46 57 
47void npu_scatter_pa_kv_cache(58void npu_scatter_pa_kv_cache(
@@ -53,17 +64,28 @@ void npu_scatter_pa_kv_cache(
53 const c10::optional<at::Tensor>& compress_lens,64 const c10::optional<at::Tensor>& compress_lens,
54 const c10::optional<at::Tensor>& compress_seq_offsets,65 const c10::optional<at::Tensor>& compress_seq_offsets,
55 const c10::optional<at::Tensor>& seq_lens,66 const c10::optional<at::Tensor>& seq_lens,
56- c10::optional<c10::string_view> cache_mode)67+ c10::optional<c10::string_view> cache_mode) {
57-{68+ char* cache_mode_ptr = cache_mode.has_value() ? const_cast<char*>(cache_mode.value().data()) : nullptr;
58- char* cache_mode_ptr = cache_mode.has_value() ? const_cast<char *>(cache_mode.value().data()) : nullptr;69+ char* scatter_mode = "None";
59- char* scatter_mode = "None";70+ c10::SmallVector<int64_t, op_infer::SIZE> strides_size = {1, 1};
60- c10::SmallVector<int64_t, op_infer::SIZE> strides_size = {1, 1};71+ at::IntArrayRef strides = at::IntArrayRef(strides_size);
61- at::IntArrayRef strides = at::IntArrayRef(strides_size);72+ c10::SmallVector<int64_t, op_infer::SIZE> offsets_size = {0, 0};
62- c10::SmallVector<int64_t, op_infer::SIZE> offsets_size = {0, 0};73+ at::IntArrayRef offsets = at::IntArrayRef(offsets_size);
63- at::IntArrayRef offsets = at::IntArrayRef(offsets_size);
64 74 
65- EXEC_NPU_NO_FORMAT_CHECK_CMD(aclnnScatterPaKvCache, key, key_cache, slot_mapping, value, value_cache,75+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
66- compress_lens, compress_seq_offsets, seq_lens, cache_mode_ptr, scatter_mode, strides, offsets);76+ aclnnScatterPaKvCache,
77+ key,
78+ key_cache,
79+ slot_mapping,
80+ value,
81+ value_cache,
82+ compress_lens,
83+ compress_seq_offsets,
84+ seq_lens,
85+ cache_mode_ptr,
86+ scatter_mode,
87+ strides,
88+ offsets);
67}89}
68 90 
69-}91+} // namespace op_api
Mop_plugin/ops/opapi/ScatterReduceKernelNpuOpApi.cpp+68-79
@@ -23,12 +23,12 @@ namespace op_api {
23using npu_preparation = at_npu::native::OpPreparation;23using npu_preparation = at_npu::native::OpPreparation;
24namespace {24namespace {
25enum class ScatterReduceType : int64_t {25enum class ScatterReduceType : int64_t {
26- REDUCE_NONE = 0,26+ REDUCE_NONE = 0,
27- REDUCE_ADD = 1,27+ REDUCE_ADD = 1,
28- REDUCE_MUL = 2,28+ REDUCE_MUL = 2,
29- REDUCE_MAX = 3,29+ REDUCE_MAX = 3,
30- REDUCE_MIN = 4,30+ REDUCE_MIN = 4,
31- REDUCE_MEAN = 5,31+ REDUCE_MEAN = 5,
32};32};
33 33 
34const std::pair<c10::string_view, ScatterReduceType> REDUCE_TYPE_MAP[] = {34const std::pair<c10::string_view, ScatterReduceType> REDUCE_TYPE_MAP[] = {
@@ -44,17 +44,18 @@ const std::pair<c10::string_view, ScatterReduceType> REDUCE_TYPE_MAP[] = {
44 {"mean", ScatterReduceType::REDUCE_MEAN},44 {"mean", ScatterReduceType::REDUCE_MEAN},
45};45};
46 46 
47-int64_t get_reduce(c10::string_view reduce, const char* op_name)47+int64_t get_reduce(c10::string_view reduce, const char* op_name) {
48-{48+ for (const auto& reduce_pair : REDUCE_TYPE_MAP) {
49- for (const auto& reduce_pair : REDUCE_TYPE_MAP) {49+ if (reduce == reduce_pair.first) {
50- if (reduce == reduce_pair.first) {50+ return static_cast<int64_t>(reduce_pair.second);
51- return static_cast<int64_t>(reduce_pair.second);
52- }
53 }51 }
54- TORCH_CHECK(52+ }
55- false, op_name,53+ TORCH_CHECK(
56- ": expected reduce to be one of none, sum, add, amin, min, amax, max, mul, prod or mean, but got ", reduce,54+ false,
57- OPS_ERROR(ErrCode::PARAM));55+ op_name,
56+ ": expected reduce to be one of none, sum, add, amin, min, amax, max, mul, prod or mean, but got ",
57+ reduce,
58+ OPS_ERROR(ErrCode::PARAM));
58}59}
59 60 
60at::Tensor scatter_reduce_cpu_fallback(61at::Tensor scatter_reduce_cpu_fallback(
@@ -63,12 +64,11 @@ at::Tensor scatter_reduce_cpu_fallback(
63 const at::Tensor& index,64 const at::Tensor& index,
64 const at::Tensor& src,65 const at::Tensor& src,
65 c10::string_view reduce,66 c10::string_view reduce,
66- bool include_self)67+ bool include_self) {
67-{68+ at::Tensor self_cpu = self.cpu();
68- at::Tensor self_cpu = self.cpu();69+ at::Tensor index_cpu = index.cpu();
69- at::Tensor index_cpu = index.cpu();70+ at::Tensor src_cpu = src.cpu();
70- at::Tensor src_cpu = src.cpu();71+ return at::scatter_reduce(self_cpu, dim, index_cpu, src_cpu, reduce, include_self).to(self.options());
71- return at::scatter_reduce(self_cpu, dim, index_cpu, src_cpu, reduce, include_self).to(self.options());
72}72}
73 73 
74at::Tensor& scatter_reduce_out_cpu_fallback(74at::Tensor& scatter_reduce_out_cpu_fallback(
@@ -78,15 +78,14 @@ at::Tensor& scatter_reduce_out_cpu_fallback(
78 const at::Tensor& src,78 const at::Tensor& src,
79 c10::string_view reduce,79 c10::string_view reduce,
80 bool include_self,80 bool include_self,
81- at::Tensor& out)81+ at::Tensor& out) {
82-{82+ at::Tensor self_cpu = self.cpu();
83- at::Tensor self_cpu = self.cpu();83+ at::Tensor index_cpu = index.cpu();
84- at::Tensor index_cpu = index.cpu();84+ at::Tensor src_cpu = src.cpu();
85- at::Tensor src_cpu = src.cpu();85+ at::Tensor out_cpu = out.cpu();
86- at::Tensor out_cpu = out.cpu();86+ at::scatter_reduce_out(out_cpu, self_cpu, dim, index_cpu, src_cpu, reduce, include_self);
87- at::scatter_reduce_out(out_cpu, self_cpu, dim, index_cpu, src_cpu, reduce, include_self);87+ out.copy_(out_cpu);
88- out.copy_(out_cpu);88+ return out;
89- return out;
90}89}
91 90 
92at::Tensor& scatter_reduce_inplace_cpu_fallback(91at::Tensor& scatter_reduce_inplace_cpu_fallback(
@@ -95,14 +94,13 @@ at::Tensor& scatter_reduce_inplace_cpu_fallback(
95 const at::Tensor& index,94 const at::Tensor& index,
96 const at::Tensor& src,95 const at::Tensor& src,
97 c10::string_view reduce,96 c10::string_view reduce,
98- bool include_self)97+ bool include_self) {
99-{98+ at::Tensor self_cpu = self.cpu();
100- at::Tensor self_cpu = self.cpu();99+ at::Tensor index_cpu = index.cpu();
101- at::Tensor index_cpu = index.cpu();100+ at::Tensor src_cpu = src.cpu();
102- at::Tensor src_cpu = src.cpu();101+ self_cpu.scatter_reduce_(dim, index_cpu, src_cpu, reduce, include_self);
103- self_cpu.scatter_reduce_(dim, index_cpu, src_cpu, reduce, include_self);102+ self.copy_(self_cpu);
104- self.copy_(self_cpu);103+ return self;
105- return self;
106}104}
107} // namespace105} // namespace
108 106 
@@ -112,23 +110,19 @@ at::Tensor scatter_reduce(
112 const at::Tensor& index,110 const at::Tensor& index,
113 const at::Tensor& src,111 const at::Tensor& src,
114 c10::string_view reduce,112 c10::string_view reduce,
115- bool include_self)113+ bool include_self) {
116-{114+ if (include_self && (reduce == "sum" || reduce == "add") &&
117- if (include_self && (reduce == "sum" || reduce == "add") &&115+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950 && self.scalar_type() == at::kFloat) {
118- c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950 &&116+ DO_COMPATIBILITY_COMMON(
119- self.scalar_type() == at::kFloat) {117+ aclnnScatterReduce, scatter_reduce_cpu_fallback(self, dim, index, src, reduce, include_self));
120- DO_COMPATIBILITY_COMMON(118+ auto result = self.clone(at::MemoryFormat::Contiguous);
121- aclnnScatterReduce,119+ npu_preparation::CheckMemory({result, index, src}, {result});
122- scatter_reduce_cpu_fallback(self, dim, index, src, reduce, include_self)120+ int64_t reduction = get_reduce(reduce, "scatter_reduce()");
123- );121+ EXEC_NPU_CMD(aclnnScatterReduce, result, dim, index, src, reduction, include_self, result);
124- auto result = self.clone(at::MemoryFormat::Contiguous);122+ return result;
125- npu_preparation::CheckMemory({result, index, src}, {result});123+ }
126- int64_t reduction = get_reduce(reduce, "scatter_reduce()");
127- EXEC_NPU_CMD(aclnnScatterReduce, result, dim, index, src, reduction, include_self, result);
128- return result;
129- }
130 124 
131- return scatter_reduce_cpu_fallback(self, dim, index, src, reduce, include_self);125+ return scatter_reduce_cpu_fallback(self, dim, index, src, reduce, include_self);
132}126}
133 127 
134at::Tensor& scatter_reduce_out(128at::Tensor& scatter_reduce_out(
@@ -138,14 +132,13 @@ at::Tensor& scatter_reduce_out(
138 const at::Tensor& src,132 const at::Tensor& src,
139 c10::string_view reduce,133 c10::string_view reduce,
140 bool include_self,134 bool include_self,
141- at::Tensor& out)135+ at::Tensor& out) {
142-{136+ DO_COMPATIBILITY_COMMON(
143- DO_COMPATIBILITY_COMMON(137+ aclnnScatterReduce, scatter_reduce_out_cpu_fallback(self, dim, index, src, reduce, include_self, out));
144- aclnnScatterReduce, scatter_reduce_out_cpu_fallback(self, dim, index, src, reduce, include_self, out));138+ npu_preparation::CheckMemory({self, index, src}, {out});
145- npu_preparation::CheckMemory({self, index, src}, {out});139+ int64_t reduction = get_reduce(reduce, "scatter_reduce_out()");
146- int64_t reduction = get_reduce(reduce, "scatter_reduce_out()");140+ EXEC_NPU_CMD(aclnnScatterReduce, self, dim, index, src, reduction, include_self, out);
147- EXEC_NPU_CMD(aclnnScatterReduce, self, dim, index, src, reduction, include_self, out);141+ return out;
148- return out;
149}142}
150 143 
151at::Tensor& scatter_reduce_(144at::Tensor& scatter_reduce_(
@@ -154,20 +147,16 @@ at::Tensor& scatter_reduce_(
154 const at::Tensor& index,147 const at::Tensor& index,
155 const at::Tensor& src,148 const at::Tensor& src,
156 c10::string_view reduce,149 c10::string_view reduce,
157- bool include_self)150+ bool include_self) {
158-{151+ if (include_self && (reduce == "sum" || reduce == "add") &&
159- if (include_self && (reduce == "sum" || reduce == "add") &&152+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950 && self.scalar_type() == at::kFloat) {
160- c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950 &&153+ DO_COMPATIBILITY_COMMON(
161- self.scalar_type() == at::kFloat) {154+ aclnnInplaceScatterReduce, scatter_reduce_inplace_cpu_fallback(self, dim, index, src, reduce, include_self));
162- DO_COMPATIBILITY_COMMON(155+ npu_preparation::CheckMemory({self, index, src}, {self});
163- aclnnInplaceScatterReduce,156+ int64_t reduction = get_reduce(reduce, "scatter_reduce_()");
164- scatter_reduce_inplace_cpu_fallback(self, dim, index, src, reduce, include_self)157+ EXEC_NPU_CMD(aclnnInplaceScatterReduce, self, dim, index, src, reduction, include_self);
165- );158+ return self;
166- npu_preparation::CheckMemory({self, index, src}, {self});159+ }
167- int64_t reduction = get_reduce(reduce, "scatter_reduce_()");160+ return scatter_reduce_inplace_cpu_fallback(self, dim, index, src, reduce, include_self);
168- EXEC_NPU_CMD(aclnnInplaceScatterReduce, self, dim, index, src, reduction, include_self);
169- return self;
170- }
171- return scatter_reduce_inplace_cpu_fallback(self, dim, index, src, reduce, include_self);
172-}
173}161}
162+} // namespace op_api
Mop_plugin/ops/opapi/ScatterUpdateKernelNpuOpApi.cpp+10-20
@@ -17,26 +17,16 @@
17#include "op_plugin/utils/op_api_common.h"17#include "op_plugin/utils/op_api_common.h"
18 18 
19namespace op_api {19namespace op_api {
20-at::Tensor scatter_update(20+at::Tensor scatter_update(const at::Tensor& self, const at::Tensor& indices, const at::Tensor& updates, int64_t axis) {
21- const at::Tensor &self,21+ // The attribute 'reduce' of Scatter only supports setting it to 'update'.
22- const at::Tensor &indices,22+ at::Tensor result = self.clone();
23- const at::Tensor &updates,23+ EXEC_NPU_CMD(aclnnInplaceScatterUpdate, result, indices, updates, axis);
24- int64_t axis)24+ return result;
25-{
26- // The attribute 'reduce' of Scatter only supports setting it to 'update'.
27- at::Tensor result = self.clone();
28- EXEC_NPU_CMD(aclnnInplaceScatterUpdate, result, indices, updates, axis);
29- return result;
30}25}
31 26 
32-at::Tensor &scatter_update_(27+at::Tensor& scatter_update_(at::Tensor& self, const at::Tensor& indices, const at::Tensor& updates, int64_t axis) {
33- at::Tensor &self,28+ // The attribute 'reduce' of Scatter only supports setting it to 'update'.
34- const at::Tensor &indices,29+ EXEC_NPU_CMD(aclnnInplaceScatterUpdate, self, indices, updates, axis);
35- const at::Tensor &updates,30+ return self;
36- int64_t axis)
37-{
38- // The attribute 'reduce' of Scatter only supports setting it to 'update'.
39- EXEC_NPU_CMD(aclnnInplaceScatterUpdate, self, indices, updates, axis);
40- return self;
41}31}
42-} // namespace op_api32+} // namespace op_api
Mop_plugin/ops/opapi/SignKernelNpuOpApi.cpp+9-11
@@ -20,22 +20,20 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor &sgn_out(const at::Tensor &self, at::Tensor &result)23+at::Tensor& sgn_out(const at::Tensor& self, at::Tensor& result) {
24-{
25 DO_COMPATIBILITY(aclnnSign, acl_op::sgn_out(self, result));24 DO_COMPATIBILITY(aclnnSign, acl_op::sgn_out(self, result));
26 npu_preparation::check_tensor({self}, result, self);25 npu_preparation::check_tensor({self}, result, self);
27 EXEC_NPU_CMD(aclnnSign, self, result);26 EXEC_NPU_CMD(aclnnSign, self, result);
28 return result;27 return result;
29}28}
30 29 
31-at::Tensor sgn(const at::Tensor &self)30+at::Tensor sgn(const at::Tensor& self) {
32-{31+ DO_COMPATIBILITY(aclnnSign, acl_op::sgn(self));
33- DO_COMPATIBILITY(aclnnSign, acl_op::sgn(self));32+ auto outputSize = op_infer::input_same_output_size(self);
34- auto outputSize = op_infer::input_same_output_size(self);33+ at::Tensor result = npu_preparation::apply_tensor_without_format(outputSize, self.options());
35- at::Tensor result = npu_preparation::apply_tensor_without_format(outputSize, self.options());34+ EXEC_NPU_CMD(aclnnSign, self, result);
36- EXEC_NPU_CMD(aclnnSign, self, result);35+ at::namedinference::propagate_names(result, self);
37- at::namedinference::propagate_names(result, self);36+ return result;
38- return result;
39}37}
40 38 
41-}39+} // namespace op_api
Mop_plugin/ops/opapi/SilentCheckV2KernelNpuOpApi.cpp+29-32
@@ -20,45 +20,42 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor aclnnNotFound(const at::Tensor &val)23+at::Tensor aclnnNotFound(const at::Tensor& val) {
24-{24+ TORCH_CHECK(false, "aclnnSilentCheck not found. " + OPS_ERROR(ErrCode::INTERNAL));
25- TORCH_CHECK(false, "aclnnSilentCheck not found. " + OPS_ERROR(ErrCode::INTERNAL));25+ return val;
26- return val;
27}26}
28 27 
29at::Tensor _npu_silent_check_v2(28at::Tensor _npu_silent_check_v2(
30- const at::Tensor &val,29+ const at::Tensor& val,
31- at::Tensor &input_grad,30+ at::Tensor& input_grad,
32- at::Tensor &sfda,31+ at::Tensor& sfda,
33- at::Tensor &step,32+ at::Tensor& step,
34 int64_t c_min_steps,33 int64_t c_min_steps,
35 double c_thresh_l1,34 double c_thresh_l1,
36 double c_coeff_l1,35 double c_coeff_l1,
37 double c_thresh_l2,36 double c_thresh_l2,
38 double c_coeff_l2,37 double c_coeff_l2,
39- int64_t npu_asd_detect)38+ int64_t npu_asd_detect) {
40-{39+ DO_COMPATIBILITY(aclnnSilentCheck, aclnnNotFound(val));
41- DO_COMPATIBILITY(aclnnSilentCheck, aclnnNotFound(val));40+ at::Tensor result = npu_preparation::apply_tensor_without_format(step.sizes(), step.options().dtype(at::kInt));
42- at::Tensor result = npu_preparation::apply_tensor_without_format(step.sizes(), step.options().dtype(at::kInt));41+ int32_t c_min_steps_cast = static_cast<int32_t>(c_min_steps);
43- int32_t c_min_steps_cast = static_cast<int32_t>(c_min_steps);42+ float c_thresh_l1_cast = static_cast<float>(c_thresh_l1);
44- float c_thresh_l1_cast = static_cast<float>(c_thresh_l1);43+ float c_coeff_l1_cast = static_cast<float>(c_coeff_l1);
45- float c_coeff_l1_cast = static_cast<float>(c_coeff_l1);44+ float c_thresh_l2_cast = static_cast<float>(c_thresh_l2);
46- float c_thresh_l2_cast = static_cast<float>(c_thresh_l2);45+ float c_coeff_l2_cast = static_cast<float>(c_coeff_l2);
47- float c_coeff_l2_cast = static_cast<float>(c_coeff_l2);46+ EXEC_NPU_CMD(
48- EXEC_NPU_CMD(47+ aclnnSilentCheck,
49- aclnnSilentCheck,48+ val,
50- val,49+ input_grad,
51- input_grad,50+ sfda,
52- sfda,51+ step,
53- step,52+ c_min_steps_cast,
54- c_min_steps_cast,53+ c_thresh_l1_cast,
55- c_thresh_l1_cast,54+ c_coeff_l1_cast,
56- c_coeff_l1_cast,55+ c_thresh_l2_cast,
57- c_thresh_l2_cast,56+ c_coeff_l2_cast,
58- c_coeff_l2_cast,57+ npu_asd_detect,
59- npu_asd_detect,58+ result);
60- result);59+ return result;
61- return result;
62}60}
63} // namespace op_api61} // namespace op_api
64- 
Mop_plugin/ops/opapi/SilentCheckV3KernelNpuOpApi.cpp+34-24
@@ -20,35 +20,45 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor aclnnNotFoundV2(const at::Tensor &val)23+at::Tensor aclnnNotFoundV2(const at::Tensor& val) {
24-{24+ TORCH_CHECK(false, "aclnnSilentCheckV2 not found. " + OPS_ERROR(ErrCode::INTERNAL));
25- TORCH_CHECK(false, "aclnnSilentCheckV2 not found. " + OPS_ERROR(ErrCode::INTERNAL));25+ return val;
26- return val;
27}26}
28 27 
29at::Tensor _npu_silent_check_v3(28at::Tensor _npu_silent_check_v3(
30- const at::Tensor &val,29+ const at::Tensor& val,
31- at::Tensor &input_grad,30+ at::Tensor& input_grad,
32- at::Tensor &step,31+ at::Tensor& step,
33- at::Tensor &max,32+ at::Tensor& max,
34- at::Tensor &avg,33+ at::Tensor& avg,
35 double c_thresh_l1,34 double c_thresh_l1,
36 double c_thresh_l2,35 double c_thresh_l2,
37 double beta1,36 double beta1,
38- int64_t npu_asd_detect)37+ int64_t npu_asd_detect) {
39-{38+ DO_COMPATIBILITY(aclnnSilentCheckV2, aclnnNotFoundV2(val));
40- DO_COMPATIBILITY(aclnnSilentCheckV2, aclnnNotFoundV2(val));39+ at::Tensor result = npu_preparation::apply_tensor_without_format(step.sizes(), step.options().dtype(at::kInt));
41- at::Tensor result = npu_preparation::apply_tensor_without_format(step.sizes(), step.options().dtype(at::kInt));40+ float c_thresh_l1_cast = static_cast<float>(c_thresh_l1);
42- float c_thresh_l1_cast = static_cast<float>(c_thresh_l1);41+ float c_thresh_l2_cast = static_cast<float>(c_thresh_l2);
43- float c_thresh_l2_cast = static_cast<float>(c_thresh_l2);42+ float beta1_cast = static_cast<float>(beta1);
44- float beta1_cast = static_cast<float>(beta1);43+ auto dst_size = input_grad.sizes();
45- auto dst_size = input_grad.sizes();44+ auto dst_stride = input_grad.strides();
46- auto dst_stride = input_grad.strides();45+ c10::SmallVector<int64_t, N> offset = {input_grad.storage_offset()};
47- c10::SmallVector<int64_t, N> offset = {input_grad.storage_offset()};46+ auto dst_offset = at::IntArrayRef(offset);
48- auto dst_offset = at::IntArrayRef(offset);47+ EXEC_NPU_CMD(
49- EXEC_NPU_CMD(aclnnSilentCheckV2, val, max, avg, input_grad, step, dst_size, dst_stride, dst_offset,48+ aclnnSilentCheckV2,
50- c_thresh_l1_cast, c_thresh_l2_cast, beta1_cast, npu_asd_detect, result);49+ val,
51- return result;50+ max,
51+ avg,
52+ input_grad,
53+ step,
54+ dst_size,
55+ dst_stride,
56+ dst_offset,
57+ c_thresh_l1_cast,
58+ c_thresh_l2_cast,
59+ beta1_cast,
60+ npu_asd_detect,
61+ result);
62+ return result;
52}63}
53} // namespace op_api64} // namespace op_api
54- 
Mop_plugin/ops/opapi/SiluBackwardKernelNpuOpApi.cpp+23-24
@@ -20,32 +20,31 @@
20#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"20#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"
21 21 
22namespace op_api {22namespace op_api {
23-at::Tensor& silu_backward_out(const at::Tensor& grad_output, const at::Tensor& self, at::Tensor& result)23+at::Tensor& silu_backward_out(const at::Tensor& grad_output, const at::Tensor& self, at::Tensor& result) {
24-{24+ DO_COMPATIBILITY(aclnnSiluBackward, acl_op::silu_backward_out(grad_output, self, result));
25- DO_COMPATIBILITY(aclnnSiluBackward, acl_op::silu_backward_out(grad_output, self, result));25+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {
26- if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {26+ at_npu::native::OpPreparation::check_tensor({grad_output, self}, result, grad_output);
27- at_npu::native::OpPreparation::check_tensor({grad_output, self}, result, grad_output);27+ }
28- }28+ EXEC_NPU_CMD(aclnnSiluBackward, grad_output, self, result);
29- EXEC_NPU_CMD(aclnnSiluBackward, grad_output, self, result);29+ return result;
30- return result;
31}30}
32 31 
33-at::Tensor silu_backward(const at::Tensor& grad_output, const at::Tensor& self)32+at::Tensor silu_backward(const at::Tensor& grad_output, const at::Tensor& self) {
34-{33+ DO_COMPATIBILITY(aclnnSiluBackward, acl_op::silu_backward(grad_output, self));
35- DO_COMPATIBILITY(aclnnSiluBackward, acl_op::silu_backward(grad_output, self));34+ at::Tensor grad_input;
36- at::Tensor grad_input;35+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {
37- if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950) {36+ grad_input = at_npu::native::OpPreparation::apply_tensor_without_format(grad_output);
38- grad_input = at_npu::native::OpPreparation::apply_tensor_without_format(grad_output);37+ } else {
39- } else {38+ at::ScalarType output_dtype = grad_output.scalar_type();
40- at::ScalarType output_dtype = grad_output.scalar_type();39+ if (grad_output.scalar_type() != self.scalar_type()) {
41- if (grad_output.scalar_type() != self.scalar_type()) {40+ output_dtype = at::kFloat;
42- output_dtype = at::kFloat;
43- }
44- auto output_size = op_infer::broadcast_ops_npu_output_size(grad_output, self);
45- grad_input = at_npu::native::OpPreparation::apply_tensor_without_format(output_size, self.options().dtype(output_dtype));
46 }41 }
47- EXEC_NPU_CMD(aclnnSiluBackward, grad_output, self, grad_input);42+ auto output_size = op_infer::broadcast_ops_npu_output_size(grad_output, self);
48- return grad_input;43+ grad_input =
44+ at_npu::native::OpPreparation::apply_tensor_without_format(output_size, self.options().dtype(output_dtype));
45+ }
46+ EXEC_NPU_CMD(aclnnSiluBackward, grad_output, self, grad_input);
47+ return grad_input;
49}48}
50 49 
51-}50+} // namespace op_api
Mop_plugin/ops/opapi/SimExponentialNpuOpApi.cpp+20-20
@@ -20,26 +20,26 @@
20 20 
21namespace op_api {21namespace op_api {
22 22 
23-at::Tensor& npu_sim_exponential_(at::Tensor& self, double lambd, c10::optional<at::Generator> generator)23+at::Tensor& npu_sim_exponential_(at::Tensor& self, double lambd, c10::optional<at::Generator> generator) {
24-{24+ TORCH_CHECK(
25- TORCH_CHECK(lambd > 0.0, "npu_sim_exponential_ expects lambd > 0.0, but found lambd=",25+ lambd > 0.0, "npu_sim_exponential_ expects lambd > 0.0, but found lambd=", lambd, OPS_ERROR(ErrCode::PARAM));
26- lambd, OPS_ERROR(ErrCode::PARAM));26+ if (std::isinf(lambd)) {
27- if (std::isinf(lambd)) {27+ self.zero_();
28- self.zero_();
29- return self;
30- }
31- 
32- auto gen = at::get_generator_or_default<at_npu::NPUGeneratorImpl>(generator, at_npu::detail::getDefaultNPUGenerator());
33- // Remove false after aclnnSimThreadExponential supports aclnnSetPytorchRandom.
34- auto counter_offset = op_plugin::utils::calc_final_counter_offset(self, false);
35- auto pair = gen->philox_engine_inputs(counter_offset);
36- int64_t seed = static_cast<int64_t>(pair.first);
37- int64_t offset = static_cast<int64_t>(pair.second);
38- int64_t count = self.numel();
39- ASCEND_LOGI("count:%lld, lambd:%lf, seed:%lld, offset:%lld", count, lambd, seed, offset);
40- 
41- EXEC_NPU_CMD(aclnnSimThreadExponential, self, count, lambd, seed, offset);
42 return self;28 return self;
29+ }
30+ 
31+ auto gen =
32+ at::get_generator_or_default<at_npu::NPUGeneratorImpl>(generator, at_npu::detail::getDefaultNPUGenerator());
33+ // Remove false after aclnnSimThreadExponential supports aclnnSetPytorchRandom.
34+ auto counter_offset = op_plugin::utils::calc_final_counter_offset(self, false);
35+ auto pair = gen->philox_engine_inputs(counter_offset);
36+ int64_t seed = static_cast<int64_t>(pair.first);
37+ int64_t offset = static_cast<int64_t>(pair.second);
38+ int64_t count = self.numel();
39+ ASCEND_LOGI("count:%lld, lambd:%lf, seed:%lld, offset:%lld", count, lambd, seed, offset);
40+ 
41+ EXEC_NPU_CMD(aclnnSimThreadExponential, self, count, lambd, seed, offset);
42+ return self;
43}43}
44 44 
45-} // namespace op_api45+} // namespace op_api
Mop_plugin/ops/opapi/SlogdetKernelNpuOpApi.cpp+18-16
@@ -20,21 +20,23 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-std::tuple<at::Tensor, at::Tensor> slogdet(const at::Tensor& self)23+std::tuple<at::Tensor, at::Tensor> slogdet(const at::Tensor& self) {
24-{24+ DO_COMPATIBILITY(aclnnSlogdet, acl_op::slogdet(self));
25- DO_COMPATIBILITY(aclnnSlogdet, acl_op::slogdet(self));25+ // input dimension at least 2
26- // input dimension at least 226+ TORCH_CHECK(
27- TORCH_CHECK(self.ndimension() >= 2, "Expected nonempty least 2D tensor, but got a tensor with sizes ", self.dim(),27+ self.ndimension() >= 2,
28- OPS_ERROR(ErrCode::PARAM));28+ "Expected nonempty least 2D tensor, but got a tensor with sizes ",
29- // calculate the output size29+ self.dim(),
30- auto outputSize = op_infer::array_to_small_vector(self.sizes());30+ OPS_ERROR(ErrCode::PARAM));
31- outputSize.erase(outputSize.end() - 2, outputSize.end());31+ // calculate the output size
32- // construct the output tensor of the NPU32+ auto outputSize = op_infer::array_to_small_vector(self.sizes());
33- at::Tensor sign = npu_preparation::apply_tensor(self, outputSize);33+ outputSize.erase(outputSize.end() - 2, outputSize.end());
34- at::Tensor log = npu_preparation::apply_tensor(self, outputSize);34+ // construct the output tensor of the NPU
35- // calculate the output result of the NPU35+ at::Tensor sign = npu_preparation::apply_tensor(self, outputSize);
36- EXEC_NPU_CMD(aclnnSlogdet, self, sign, log);36+ at::Tensor log = npu_preparation::apply_tensor(self, outputSize);
37+ // calculate the output result of the NPU
38+ EXEC_NPU_CMD(aclnnSlogdet, self, sign, log);
37 39 
38- return std::tie(sign, log);40+ return std::tie(sign, log);
39-}
40}41}
42+} // namespace op_api
Mop_plugin/ops/opapi/SmoothL1LossBackwrdKernelNpuOpApi.cpp+7-9
@@ -26,10 +26,10 @@ at::Tensor& smooth_l1_loss_backward_out(
26 const at::Tensor& target,26 const at::Tensor& target,
27 int64_t reduction,27 int64_t reduction,
28 double beta,28 double beta,
29- at::Tensor& grad_input)29+ at::Tensor& grad_input) {
30-{30+ DO_COMPATIBILITY(
31- DO_COMPATIBILITY(aclnnSmoothL1LossBackward,31+ aclnnSmoothL1LossBackward,
32- acl_op::smooth_l1_loss_backward_out(grad_out, self, target, reduction, beta, grad_input));32+ acl_op::smooth_l1_loss_backward_out(grad_out, self, target, reduction, beta, grad_input));
33 auto mid_shape = op_infer::broadcast_ops_npu_output_size(self.sizes(), target.sizes());33 auto mid_shape = op_infer::broadcast_ops_npu_output_size(self.sizes(), target.sizes());
34 auto output_size = op_infer::broadcast_ops_npu_output_size(mid_shape, grad_out.sizes());34 auto output_size = op_infer::broadcast_ops_npu_output_size(mid_shape, grad_out.sizes());
35 npu_preparation::check_tensor({grad_out, self, target}, grad_input, grad_input.scalar_type(), output_size);35 npu_preparation::check_tensor({grad_out, self, target}, grad_input, grad_input.scalar_type(), output_size);
@@ -43,10 +43,8 @@ at::Tensor smooth_l1_loss_backward(
43 const at::Tensor& self,43 const at::Tensor& self,
44 const at::Tensor& target,44 const at::Tensor& target,
45 int64_t reduction,45 int64_t reduction,
46- double beta)46+ double beta) {
47-{47+ DO_COMPATIBILITY(aclnnSmoothL1LossBackward, acl_op::smooth_l1_loss_backward(grad_out, self, target, reduction, beta));
48- DO_COMPATIBILITY(aclnnSmoothL1LossBackward,
49- acl_op::smooth_l1_loss_backward(grad_out, self, target, reduction, beta));
50 auto mid_shape = op_infer::broadcast_ops_npu_output_size(self.sizes(), target.sizes());48 auto mid_shape = op_infer::broadcast_ops_npu_output_size(self.sizes(), target.sizes());
51 auto output_size = op_infer::broadcast_ops_npu_output_size(mid_shape, grad_out.sizes());49 auto output_size = op_infer::broadcast_ops_npu_output_size(mid_shape, grad_out.sizes());
52 at::Tensor grad_input = npu_preparation::apply_tensor_without_format(self, output_size);50 at::Tensor grad_input = npu_preparation::apply_tensor_without_format(self, output_size);
@@ -55,4 +53,4 @@ at::Tensor smooth_l1_loss_backward(
55 return grad_input;53 return grad_input;
56}54}
57 55 
58-}56+} // namespace op_api
Mop_plugin/ops/opapi/SmoothL1LossKernelNpuOpApi.cpp+22-20
@@ -20,26 +20,28 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor &smooth_l1_loss_out(const at::Tensor &self, const at::Tensor &target, int64_t reduction, double beta,23+at::Tensor& smooth_l1_loss_out(
24- at::Tensor &out)24+ const at::Tensor& self,
25-{25+ const at::Tensor& target,
26- DO_COMPATIBILITY(aclnnSmoothL1Loss, acl_op::smooth_l1_loss_out(self, target, reduction, beta, out));26+ int64_t reduction,
27- auto outputSize = op_infer::smooth_l1_loss_npu_output_size(self, reduction);27+ double beta,
28- npu_preparation::check_tensor({self, target}, out, out.scalar_type(), outputSize);28+ at::Tensor& out) {
29- npu_preparation::check_memory({self, target}, {out});29+ DO_COMPATIBILITY(aclnnSmoothL1Loss, acl_op::smooth_l1_loss_out(self, target, reduction, beta, out));
30- float sigma = static_cast<float>(beta);30+ auto outputSize = op_infer::smooth_l1_loss_npu_output_size(self, reduction);
31- EXEC_NPU_CMD(aclnnSmoothL1Loss, self, target, reduction, sigma, out);31+ npu_preparation::check_tensor({self, target}, out, out.scalar_type(), outputSize);
32- return out;32+ npu_preparation::check_memory({self, target}, {out});
33+ float sigma = static_cast<float>(beta);
34+ EXEC_NPU_CMD(aclnnSmoothL1Loss, self, target, reduction, sigma, out);
35+ return out;
33}36}
34 37 
35-at::Tensor smooth_l1_loss(const at::Tensor &self, const at::Tensor &target, int64_t reduction, double beta)38+at::Tensor smooth_l1_loss(const at::Tensor& self, const at::Tensor& target, int64_t reduction, double beta) {
36-{39+ DO_COMPATIBILITY(aclnnSmoothL1Loss, acl_op::smooth_l1_loss(self, target, reduction, beta));
37- DO_COMPATIBILITY(aclnnSmoothL1Loss, acl_op::smooth_l1_loss(self, target, reduction, beta));40+ auto outputSize = op_infer::smooth_l1_loss_npu_output_size(self, reduction);
38- auto outputSize = op_infer::smooth_l1_loss_npu_output_size(self, reduction);41+ at::ScalarType high_type = at::native::result_type(self, target);
39- at::ScalarType high_type = at::native::result_type(self, target);42+ at::Tensor result = npu_preparation::apply_tensor_without_format(outputSize, self.options().dtype(high_type));
40- at::Tensor result = npu_preparation::apply_tensor_without_format(outputSize, self.options().dtype(high_type));43+ float sigma = static_cast<float>(beta);
41- float sigma = static_cast<float>(beta);44+ EXEC_NPU_CMD(aclnnSmoothL1Loss, self, target, reduction, sigma, result);
42- EXEC_NPU_CMD(aclnnSmoothL1Loss, self, target, reduction, sigma, result);45+ return result;
43- return result;
44-}
45}46}
47+} // namespace op_api
Mop_plugin/ops/opapi/SortKernelNpuOpApi.cpp+71-65
@@ -20,90 +20,96 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-std::tuple<at::Tensor &, at::Tensor &> sort_output(const at::Tensor &self, bool stable, int64_t dim, bool descending,23+std::tuple<at::Tensor&, at::Tensor&> sort_output(
24- at::Tensor &values, at::Tensor &indices)24+ const at::Tensor& self,
25-{25+ bool stable,
26- EXEC_NPU_CMD(aclnnSort, self, stable, dim, descending, values, indices);26+ int64_t dim,
27- return std::tie(values, indices);27+ bool descending,
28+ at::Tensor& values,
29+ at::Tensor& indices) {
30+ EXEC_NPU_CMD(aclnnSort, self, stable, dim, descending, values, indices);
31+ return std::tie(values, indices);
28}32}
29 33 
30-std::tuple<at::Tensor, at::Tensor> sort(const at::Tensor &self, int64_t dim, bool descending)34+std::tuple<at::Tensor, at::Tensor> sort(const at::Tensor& self, int64_t dim, bool descending) {
31-{35+ DO_COMPATIBILITY(aclnnSort, acl_op::sort(self, dim, descending));
32- DO_COMPATIBILITY(aclnnSort, acl_op::sort(self, dim, descending));36+ at::Tensor values = npu_preparation::apply_tensor_without_format(self);
33- at::Tensor values = npu_preparation::apply_tensor_without_format(self);37+ at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));
34- at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));38+ bool stable = false;
35- bool stable = false;
36 39 
37- return sort_output(self, stable, dim, descending, values, indices);40+ return sort_output(self, stable, dim, descending, values, indices);
38}41}
39 42 
40#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)43#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)
41-std::tuple<at::Tensor, at::Tensor> sort(const at::Tensor &self, at::Dimname dim, bool descending)44+std::tuple<at::Tensor, at::Tensor> sort(const at::Tensor& self, at::Dimname dim, bool descending) {
42-{45+ DO_COMPATIBILITY(aclnnSort, acl_op::sort(self, dim, descending));
43- DO_COMPATIBILITY(aclnnSort, acl_op::sort(self, dim, descending));46+ at::Tensor values = npu_preparation::apply_tensor_without_format(self);
44- at::Tensor values = npu_preparation::apply_tensor_without_format(self);47+ at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));
45- at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));48+ bool stable = false;
46- bool stable = false;49+ int64_t argDim = dimname_to_position(self, dim);
47- int64_t argDim = dimname_to_position(self, dim);
48 50 
49- return sort_output(self, stable, argDim, descending, values, indices);51+ return sort_output(self, stable, argDim, descending, values, indices);
50}52}
51#endif53#endif
52 54 
53-std::tuple<at::Tensor &, at::Tensor &> sort_out(const at::Tensor &self, int64_t dim,55+std::tuple<at::Tensor&, at::Tensor&> sort_out(
54- bool descending, at::Tensor &values, at::Tensor &indices)56+ const at::Tensor& self,
55-{57+ int64_t dim,
56- DO_COMPATIBILITY(aclnnSort, acl_op::sort_out(self, dim, descending, values, indices));58+ bool descending,
57- npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());59+ at::Tensor& values,
58- npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());60+ at::Tensor& indices) {
59- bool stable = false;61+ DO_COMPATIBILITY(aclnnSort, acl_op::sort_out(self, dim, descending, values, indices));
62+ npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());
63+ npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());
64+ bool stable = false;
60 65 
61- return sort_output(self, stable, dim, descending, values, indices);66+ return sort_output(self, stable, dim, descending, values, indices);
62}67}
63 68 
64#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)69#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)
65-std::tuple<at::Tensor &, at::Tensor &> sort_out(const at::Tensor &self, at::Dimname dim,70+std::tuple<at::Tensor&, at::Tensor&> sort_out(
66- bool descending, at::Tensor &values, at::Tensor &indices)71+ const at::Tensor& self,
67-{72+ at::Dimname dim,
68- DO_COMPATIBILITY(aclnnSort, acl_op::sort_out(self, dim, descending, values, indices));73+ bool descending,
69- npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());74+ at::Tensor& values,
70- npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());75+ at::Tensor& indices) {
71- bool stable = false;76+ DO_COMPATIBILITY(aclnnSort, acl_op::sort_out(self, dim, descending, values, indices));
77+ npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());
78+ npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());
79+ bool stable = false;
72 80 
73- return sort_output(self, stable, dimname_to_position(self, dim), descending, values, indices);81+ return sort_output(self, stable, dimname_to_position(self, dim), descending, values, indices);
74}82}
75#endif83#endif
76 84 
77-std::tuple<at::Tensor, at::Tensor> sort(const at::Tensor &self,85+std::tuple<at::Tensor, at::Tensor> sort(
78- c10::optional<bool> stable,86+ const at::Tensor& self,
79- int64_t dim,87+ c10::optional<bool> stable,
80- bool descending)88+ int64_t dim,
81-{89+ bool descending) {
82- auto dtype = self.scalar_type();90+ auto dtype = self.scalar_type();
83- TORCH_CHECK(!(dtype == at::kDouble),91+ TORCH_CHECK(!(dtype == at::kDouble), "Input data type should not be float64 " + OPS_ERROR(ErrCode::TYPE));
84- "Input data type should not be float64 " + OPS_ERROR(ErrCode::TYPE));92+ at::Tensor values = npu_preparation::apply_tensor_without_format(self);
85- at::Tensor values = npu_preparation::apply_tensor_without_format(self);93+ at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));
86- at::Tensor indices = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kLong));94+ bool argStable = c10::value_or_else(stable, [] { return false; });
87- bool argStable = c10::value_or_else(stable, [] { return false; });95+ EXEC_NPU_CMD(aclnnSort, self, argStable, dim, descending, values, indices);
88- EXEC_NPU_CMD(aclnnSort, self, argStable, dim, descending, values, indices);96+ return std::tie(values, indices);
89- return std::tie(values, indices);
90}97}
91 98 
92-std::tuple<at::Tensor &, at::Tensor &> sort_out(const at::Tensor &self,99+std::tuple<at::Tensor&, at::Tensor&> sort_out(
93- c10::optional<bool> stable,100+ const at::Tensor& self,
94- int64_t dim,101+ c10::optional<bool> stable,
95- bool descending,102+ int64_t dim,
96- at::Tensor &values,103+ bool descending,
97- at::Tensor &indices)104+ at::Tensor& values,
98-{105+ at::Tensor& indices) {
99- auto dtype = self.scalar_type();106+ auto dtype = self.scalar_type();
100- TORCH_CHECK(!(dtype == at::kDouble),107+ TORCH_CHECK(!(dtype == at::kDouble), "Input data type should not be float64 " + OPS_ERROR(ErrCode::TYPE));
101- "Input data type should not be float64 " + OPS_ERROR(ErrCode::TYPE));108+ npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());
102- npu_preparation::check_tensor({self}, values, values.scalar_type(), self.sizes());109+ npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());
103- npu_preparation::check_tensor({self}, indices, indices.scalar_type(), self.sizes());110+ bool argStable = c10::value_or_else(stable, [] { return false; });
104- bool argStable = c10::value_or_else(stable, [] { return false; });111+ EXEC_NPU_CMD(aclnnSort, self, argStable, dim, descending, values, indices);
105- EXEC_NPU_CMD(aclnnSort, self, argStable, dim, descending, values, indices);112+ return std::tie(values, indices);
106- return std::tie(values, indices);
107}113}
108 114 
109-} // namespace op_api115+} // namespace op_api
Mop_plugin/ops/opapi/SparseFlashAttentionGradKernelNpuOpApi.cpp+100-65
@@ -26,79 +26,114 @@ const int DIMENSION_3D = 3;
26const int DIMENSION_4D = 4;26const int DIMENSION_4D = 4;
27 27 
28std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention_grad(28std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention_grad(
29- const at::Tensor &query,29+ const at::Tensor& query,
30- const at::Tensor &key,30+ const at::Tensor& key,
31- const at::Tensor &value,31+ const at::Tensor& value,
32- const at::Tensor &sparse_indices,32+ const at::Tensor& sparse_indices,
33- const at::Tensor &d_out,33+ const at::Tensor& d_out,
34- const at::Tensor &out,34+ const at::Tensor& out,
35- const at::Tensor &softmax_max,35+ const at::Tensor& softmax_max,
36- const at::Tensor &softmax_sum,36+ const at::Tensor& softmax_sum,
37 double scale_value,37 double scale_value,
38 int64_t sparse_block_size,38 int64_t sparse_block_size,
39- const c10::optional<at::Tensor> &query_rope,39+ const c10::optional<at::Tensor>& query_rope,
40- const c10::optional<at::Tensor> &key_rope,40+ const c10::optional<at::Tensor>& key_rope,
41- const c10::optional<at::Tensor> &actual_seq_qlen,41+ const c10::optional<at::Tensor>& actual_seq_qlen,
42- const c10::optional<at::Tensor> &actual_seq_kvlen,42+ const c10::optional<at::Tensor>& actual_seq_kvlen,
43 c10::optional<c10::string_view> layout,43 c10::optional<c10::string_view> layout,
44 c10::optional<int64_t> sparse_mode,44 c10::optional<int64_t> sparse_mode,
45 c10::optional<int64_t> pre_tokens,45 c10::optional<int64_t> pre_tokens,
46 c10::optional<int64_t> next_tokens,46 c10::optional<int64_t> next_tokens,
47- c10::optional<int64_t> attention_mode)47+ c10::optional<int64_t> attention_mode) {
48-{48+ const at::Tensor& query_rope_const = query_rope.value_or(at::Tensor());
49- 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- const at::Tensor &key_rope_const = key_rope.value_or(at::Tensor());50+ const at::Tensor& ac_seq_qlen = actual_seq_qlen.value_or(at::Tensor());
51- const at::Tensor &ac_seq_qlen = actual_seq_qlen.value_or(at::Tensor());51+ const at::Tensor& ac_seq_kvlen = actual_seq_kvlen.value_or(at::Tensor());
52- const at::Tensor &ac_seq_kvlen = actual_seq_kvlen.value_or(at::Tensor());52+ TORCH_CHECK(
53- TORCH_CHECK(query.dim() == DIMENSION_3D || query.dim() == DIMENSION_4D,53+ query.dim() == DIMENSION_3D || query.dim() == DIMENSION_4D,
54- "The shapes of the input query should be 3 or 4 dimensional, but got ",54+ "The shapes of the input query should be 3 or 4 dimensional, but got ",
55- query.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));55+ query.dim(),
56- if (query_rope_const.defined()) {56+ "-dimensional",
57- TORCH_CHECK(query_rope_const.dim() == DIMENSION_3D || query_rope_const.dim() == DIMENSION_4D,57+ OPS_ERROR(ErrCode::PARAM));
58- "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",58+ if (query_rope_const.defined()) {
59- query_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));59+ TORCH_CHECK(
60- }60+ query_rope_const.dim() == DIMENSION_3D || query_rope_const.dim() == DIMENSION_4D,
61- TORCH_CHECK(key.dim() == DIMENSION_3D || key.dim() == DIMENSION_4D,61+ "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",
62- "The shapes of the input key should be 3 or 4 dimensional, but got ", key.dim(),62+ query_rope_const.dim(),
63- "-dimensional", OPS_ERROR(ErrCode::PARAM));63+ "-dimensional",
64- if (key_rope_const.defined()) {64+ OPS_ERROR(ErrCode::PARAM));
65- TORCH_CHECK(key_rope_const.dim() == DIMENSION_3D || key_rope_const.dim() == DIMENSION_4D,65+ }
66- "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",66+ TORCH_CHECK(
67- key_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));67+ key.dim() == DIMENSION_3D || key.dim() == DIMENSION_4D,
68- }68+ "The shapes of the input key should be 3 or 4 dimensional, but got ",
69- TORCH_CHECK(value.dim() == DIMENSION_3D || value.dim() == DIMENSION_4D,69+ key.dim(),
70- "The shapes of the input value should be 3 or 4 dimensional, but got ",70+ "-dimensional",
71- value.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));71+ OPS_ERROR(ErrCode::PARAM));
72- at::Tensor d_query = OpPreparation::apply_tensor_without_format(query);72+ if (key_rope_const.defined()) {
73- at::Tensor d_key = OpPreparation::apply_tensor_without_format(key);73+ TORCH_CHECK(
74- at::Tensor d_value = OpPreparation::apply_tensor_without_format(value);74+ key_rope_const.dim() == DIMENSION_3D || key_rope_const.dim() == DIMENSION_4D,
75- at::Tensor d_query_rope;75+ "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",
76- at::Tensor d_key_rope;76+ key_rope_const.dim(),
77- if (query_rope_const.defined()) {77+ "-dimensional",
78- d_query_rope = OpPreparation::apply_tensor_without_format(query_rope_const);78+ OPS_ERROR(ErrCode::PARAM));
79- } else {79+ }
80- d_query_rope = at::empty({0}, query.options());80+ TORCH_CHECK(
81- }81+ value.dim() == DIMENSION_3D || value.dim() == DIMENSION_4D,
82- if (key_rope_const.defined()) {82+ "The shapes of the input value should be 3 or 4 dimensional, but got ",
83- d_key_rope = OpPreparation::apply_tensor_without_format(key_rope_const);83+ value.dim(),
84- } else {84+ "-dimensional",
85- d_key_rope = at::empty({0}, key.options());85+ OPS_ERROR(ErrCode::PARAM));
86- }86+ at::Tensor d_query = OpPreparation::apply_tensor_without_format(query);
87+ at::Tensor d_key = OpPreparation::apply_tensor_without_format(key);
88+ at::Tensor d_value = OpPreparation::apply_tensor_without_format(value);
89+ at::Tensor d_query_rope;
90+ at::Tensor d_key_rope;
91+ if (query_rope_const.defined()) {
92+ d_query_rope = OpPreparation::apply_tensor_without_format(query_rope_const);
93+ } else {
94+ d_query_rope = at::empty({0}, query.options());
95+ }
96+ if (key_rope_const.defined()) {
97+ d_key_rope = OpPreparation::apply_tensor_without_format(key_rope_const);
98+ } else {
99+ d_key_rope = at::empty({0}, key.options());
100+ }
87 101 
88- c10::string_view layout_str_view = layout.value_or("BSND");102+ c10::string_view layout_str_view = layout.value_or("BSND");
89- char *layout_ptr = const_cast<char *>(layout_str_view.data());103+ char* layout_ptr = const_cast<char*>(layout_str_view.data());
90 104 
91- const int64_t sparse_mode_const = sparse_mode.value_or(3);105+ const int64_t sparse_mode_const = sparse_mode.value_or(3);
92- const int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);106+ const int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
93- const int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);107+ const int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
94- const bool deterministic_const = true;108+ const bool deterministic_const = true;
95 109 
96- EXEC_NPU_NO_FORMAT_CHECK_CMD(110+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
97- aclnnSparseFlashAttentionGrad, query, key, value, sparse_indices, d_out,111+ aclnnSparseFlashAttentionGrad,
98- out, softmax_max, softmax_sum, ac_seq_qlen, ac_seq_kvlen, query_rope_const, key_rope_const,112+ query,
99- scale_value, sparse_block_size, layout_ptr, sparse_mode_const, pre_tokens_const, next_tokens_const,113+ key,
100- deterministic_const, d_query, d_key, d_value, d_query_rope, d_key_rope);114+ value,
115+ sparse_indices,
116+ d_out,
117+ out,
118+ softmax_max,
119+ softmax_sum,
120+ ac_seq_qlen,
121+ ac_seq_kvlen,
122+ query_rope_const,
123+ key_rope_const,
124+ scale_value,
125+ sparse_block_size,
126+ layout_ptr,
127+ sparse_mode_const,
128+ pre_tokens_const,
129+ next_tokens_const,
130+ deterministic_const,
131+ d_query,
132+ d_key,
133+ d_value,
134+ d_query_rope,
135+ d_key_rope);
101 136 
102- return std::make_tuple(d_query, d_key, d_value, d_query_rope, d_key_rope);137+ return std::make_tuple(d_query, d_key, d_value, d_query_rope, d_key_rope);
103-}
104}138}
139+} // namespace op_api
Mop_plugin/ops/opapi/SparseFlashAttentionKernelNpuOpApi.cpp+131-78
@@ -29,94 +29,147 @@ const static int64_t DIM_4 = 4;
29using npu_preparation = at_npu::native::OpPreparation;29using npu_preparation = at_npu::native::OpPreparation;
30 30 
31namespace {31namespace {
32-at::Tensor construct_sparse_flash_attention_output_tensor(32+at::Tensor construct_sparse_flash_attention_output_tensor(const at::Tensor& query, std::string layout) {
33- const at::Tensor& query, std::string layout)33+ TORCH_CHECK(
34-{34+ layout == "BSND" || layout == "TND",
35- TORCH_CHECK(layout == "BSND" || layout == "TND", "The layout of query only support BSND and TND, but got ",35+ "The layout of query only support BSND and TND, but got ",
36- layout, OPS_ERROR(ErrCode::PARAM));36+ layout,
37- at::SmallVector<int64_t, SIZE> output_size;37+ OPS_ERROR(ErrCode::PARAM));
38- for (size_t i = 0; i < query.sizes().size(); i++) {38+ at::SmallVector<int64_t, SIZE> output_size;
39- TORCH_CHECK(query.size(i) > 0, "All values within query's shape should be greater "39+ for (size_t i = 0; i < query.sizes().size(); i++) {
40- "than 0, but shape[", i, "] is ", query.size(i), OPS_ERROR(ErrCode::PARAM));40+ TORCH_CHECK(
41- }41+ query.size(i) > 0,
42- if (layout == "TND") {42+ "All values within query's shape should be greater "
43- TORCH_CHECK(query.dim() == DIM_3,43+ "than 0, but shape[",
44- "When the layout of query is TND, the query dimension must be 3, but got ",44+ i,
45- query.dim(), OPS_ERROR(ErrCode::PARAM));45+ "] is ",
46- output_size = {query.size(DIM_0), query.size(DIM_1), query.size(DIM_2)};46+ query.size(i),
47- } else {47+ OPS_ERROR(ErrCode::PARAM));
48- TORCH_CHECK(query.dim() == DIM_4,48+ }
49- "When the layout of query is BSND, the query dimension must be 4, but got ",49+ if (layout == "TND") {
50- query.dim(), OPS_ERROR(ErrCode::PARAM));50+ TORCH_CHECK(
51- output_size = {query.size(DIM_0), query.size(DIM_1), query.size(DIM_2), query.size(DIM_3)};51+ query.dim() == DIM_3,
52- }52+ "When the layout of query is TND, the query dimension must be 3, but got ",
53- at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, query.options().dtype(query.dtype()));53+ query.dim(),
54+ OPS_ERROR(ErrCode::PARAM));
55+ output_size = {query.size(DIM_0), query.size(DIM_1), query.size(DIM_2)};
56+ } else {
57+ TORCH_CHECK(
58+ query.dim() == DIM_4,
59+ "When the layout of query is BSND, the query dimension must be 4, but got ",
60+ query.dim(),
61+ OPS_ERROR(ErrCode::PARAM));
62+ output_size = {query.size(DIM_0), query.size(DIM_1), query.size(DIM_2), query.size(DIM_3)};
63+ }
64+ at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, query.options().dtype(query.dtype()));
54 65 
55- return output;66+ return output;
56}67}
57-}68+} // namespace
58- 
59 69 
60std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention(70std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_sparse_flash_attention(
61- const at::Tensor &query, const at::Tensor &key, const at::Tensor &value,71+ const at::Tensor& query,
62- const at::Tensor &sparse_indices, double scale_value,72+ const at::Tensor& key,
63- const c10::optional<at::Tensor> &block_table,73+ const at::Tensor& value,
64- const c10::optional<at::Tensor> &actual_seq_lengths_query,74+ const at::Tensor& sparse_indices,
65- const c10::optional<at::Tensor> &actual_seq_lengths_kv,75+ double scale_value,
66- const c10::optional<at::Tensor> &query_rope,76+ const c10::optional<at::Tensor>& block_table,
67- const c10::optional<at::Tensor> &key_rope, int64_t sparse_block_size,77+ const c10::optional<at::Tensor>& actual_seq_lengths_query,
68- c10::string_view layout_query, c10::string_view layout_kv,78+ const c10::optional<at::Tensor>& actual_seq_lengths_kv,
69- int64_t sparse_mode, int64_t pre_tokens, int64_t next_tokens,79+ const c10::optional<at::Tensor>& query_rope,
70- int64_t attention_mode, bool return_softmax_lse,80+ const c10::optional<at::Tensor>& key_rope,
71- const c10::optional<at::Tensor> &sinks)81+ int64_t sparse_block_size,
72-{82+ c10::string_view layout_query,
73- TORCH_CHECK(query.numel() > 0, "Tensor query is empty.", OPS_ERROR(ErrCode::PARAM));83+ c10::string_view layout_kv,
74- TORCH_CHECK(key.numel() > 0, "Tensor key is empty.", OPS_ERROR(ErrCode::PARAM));84+ int64_t sparse_mode,
75- TORCH_CHECK(sparse_indices.numel() > 0, "Tensor sparse_indices is empty.")85+ int64_t pre_tokens,
86+ int64_t next_tokens,
87+ int64_t attention_mode,
88+ bool return_softmax_lse,
89+ const c10::optional<at::Tensor>& sinks) {
90+ TORCH_CHECK(query.numel() > 0, "Tensor query is empty.", OPS_ERROR(ErrCode::PARAM));
91+ TORCH_CHECK(key.numel() > 0, "Tensor key is empty.", OPS_ERROR(ErrCode::PARAM));
92+ TORCH_CHECK(sparse_indices.numel() > 0, "Tensor sparse_indices is empty.")
76 93 
77- std::string layout_query_str = std::string(layout_query);94+ std::string layout_query_str = std::string(layout_query);
78- std::string layout_kv_str = std::string(layout_kv);95+ std::string layout_kv_str = std::string(layout_kv);
79 96 
80- // construct the output tensor97+ // construct the output tensor
81- at::Tensor sparse_flash_attention_output = construct_sparse_flash_attention_output_tensor(98+ at::Tensor sparse_flash_attention_output = construct_sparse_flash_attention_output_tensor(query, layout_query_str);
82- query, layout_query_str);99+ at::Tensor softmax_max;
83- at::Tensor softmax_max;100+ at::Tensor softmax_sum;
84- at::Tensor softmax_sum;101+ at::SmallVector<int64_t, SIZE> softmax_max_size;
85- at::SmallVector<int64_t, SIZE> softmax_max_size;102+ at::SmallVector<int64_t, SIZE> softmax_sum_size;
86- at::SmallVector<int64_t, SIZE> softmax_sum_size;103+ if (return_softmax_lse) {
87- if (return_softmax_lse) {104+ if (query.dim() == DIM_3) {
88- if (query.dim() == DIM_3) {105+ softmax_max_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
89- softmax_max_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};106+ softmax_sum_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
90- softmax_sum_size = {key.size(1), query.size(0), query.size(1) / key.size(1)};
91- } else {
92- softmax_max_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
93- softmax_sum_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
94- }
95 } else {107 } else {
96- softmax_max_size = {0};108+ softmax_max_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
97- softmax_sum_size = {0};109+ softmax_sum_size = {query.size(0), key.size(2), query.size(1), query.size(2) / key.size(2)};
98 }110 }
99- softmax_max = at::empty(softmax_max_size, query.options().dtype(at::kFloat));111+ } else {
100- softmax_sum = at::empty(softmax_sum_size, query.options().dtype(at::kFloat));112+ softmax_max_size = {0};
101- // convert str113+ softmax_sum_size = {0};
102- char *layout_query_ptr = const_cast<char *>(layout_query_str.c_str());114+ }
103- char *layout_kv_ptr = const_cast<char *>(layout_kv_str.c_str());115+ softmax_max = at::empty(softmax_max_size, query.options().dtype(at::kFloat));
116+ softmax_sum = at::empty(softmax_sum_size, query.options().dtype(at::kFloat));
117+ // convert str
118+ char* layout_query_ptr = const_cast<char*>(layout_query_str.c_str());
119+ char* layout_kv_ptr = const_cast<char*>(layout_kv_str.c_str());
104 120 
105- if (sinks.has_value() && sinks->defined()) {121+ if (sinks.has_value() && sinks->defined()) {
106- EXEC_NPU_NO_FORMAT_CHECK_CMD(aclnnSparseFlashAttentionV2, query,122+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
107- key, value, sparse_indices, block_table, actual_seq_lengths_query,123+ aclnnSparseFlashAttentionV2,
108- actual_seq_lengths_kv, query_rope, key_rope, sinks, scale_value, sparse_block_size,124+ query,
109- layout_query_ptr, layout_kv_ptr, sparse_mode, pre_tokens, next_tokens, attention_mode, return_softmax_lse,125+ key,
110- sparse_flash_attention_output, softmax_max, softmax_sum);126+ value,
111- } else {127+ sparse_indices,
112- EXEC_NPU_NO_FORMAT_CHECK_CMD(aclnnSparseFlashAttention, query,128+ block_table,
113- key, value, sparse_indices, block_table, actual_seq_lengths_query,129+ actual_seq_lengths_query,
114- actual_seq_lengths_kv, query_rope, key_rope, scale_value, sparse_block_size,130+ actual_seq_lengths_kv,
115- layout_query_ptr, layout_kv_ptr, sparse_mode, pre_tokens, next_tokens, attention_mode, return_softmax_lse,131+ query_rope,
116- sparse_flash_attention_output, softmax_max, softmax_sum);132+ key_rope,
117- }133+ sinks,
134+ scale_value,
135+ sparse_block_size,
136+ layout_query_ptr,
137+ layout_kv_ptr,
138+ sparse_mode,
139+ pre_tokens,
140+ next_tokens,
141+ attention_mode,
142+ return_softmax_lse,
143+ sparse_flash_attention_output,
144+ softmax_max,
145+ softmax_sum);
146+ } else {
147+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
148+ aclnnSparseFlashAttention,
149+ query,
150+ key,
151+ value,
152+ sparse_indices,
153+ block_table,
154+ actual_seq_lengths_query,
155+ actual_seq_lengths_kv,
156+ query_rope,
157+ key_rope,
158+ scale_value,
159+ sparse_block_size,
160+ layout_query_ptr,
161+ layout_kv_ptr,
162+ sparse_mode,
163+ pre_tokens,
164+ next_tokens,
165+ attention_mode,
166+ return_softmax_lse,
167+ sparse_flash_attention_output,
168+ softmax_max,
169+ softmax_sum);
170+ }
118 171 
119- return std::tuple<at::Tensor, at::Tensor, at::Tensor>(sparse_flash_attention_output, softmax_max, softmax_sum);172+ return std::tuple<at::Tensor, at::Tensor, at::Tensor>(sparse_flash_attention_output, softmax_max, softmax_sum);
120}173}
121 174 
122} // namespace op_api175} // namespace op_api
Mop_plugin/ops/opapi/SparseLightningIndexerGradKLLossKernelNpuOpApi.cpp+81-49
@@ -26,61 +26,93 @@ const int DIMENSION_3D = 3;
26const int DIMENSION_4D = 4;26const int DIMENSION_4D = 4;
27 27 
28std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_lightning_indexer_grad_kl_loss_symint(28std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_sparse_lightning_indexer_grad_kl_loss_symint(
29- const at::Tensor &query,29+ const at::Tensor& query,
30- const at::Tensor &key,30+ const at::Tensor& key,
31- const at::Tensor &query_index,31+ const at::Tensor& query_index,
32- const at::Tensor &key_index,32+ const at::Tensor& key_index,
33- const at::Tensor &weights,33+ const at::Tensor& weights,
34- const at::Tensor &sparse_indices,34+ const at::Tensor& sparse_indices,
35- const at::Tensor &softmax_max,35+ const at::Tensor& softmax_max,
36- const at::Tensor &softmax_sum,36+ const at::Tensor& softmax_sum,
37 double scale_value,37 double scale_value,
38- const c10::optional<at::Tensor> &query_rope,38+ const c10::optional<at::Tensor>& query_rope,
39- const c10::optional<at::Tensor> &key_rope,39+ const c10::optional<at::Tensor>& key_rope,
40 c10::OptionalArrayRef<c10::SymInt> actual_seq_qlen,40 c10::OptionalArrayRef<c10::SymInt> actual_seq_qlen,
41 c10::OptionalArrayRef<c10::SymInt> actual_seq_klen,41 c10::OptionalArrayRef<c10::SymInt> actual_seq_klen,
42 c10::optional<c10::string_view> layout,42 c10::optional<c10::string_view> layout,
43 c10::optional<int64_t> sparse_mode,43 c10::optional<int64_t> sparse_mode,
44 c10::optional<int64_t> pre_tokens,44 c10::optional<int64_t> pre_tokens,
45- c10::optional<int64_t> next_tokens)45+ c10::optional<int64_t> next_tokens) {
46-{46+ const at::Tensor& query_rope_const = query_rope.value_or(at::Tensor());
47- const at::Tensor &query_rope_const = query_rope.value_or(at::Tensor());47+ const at::Tensor& key_rope_const = key_rope.value_or(at::Tensor());
48- const at::Tensor &key_rope_const = key_rope.value_or(at::Tensor());48+ auto actual_seq_qlen_const =
49- auto actual_seq_qlen_const = actual_seq_qlen.has_value() ? c10::asIntArrayRefUnchecked(actual_seq_qlen.value()) : at::IntArrayRef{};49+ 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{};50+ auto actual_seq_klen_const =
51- c10::string_view layout_str = layout.value_or("BSND");51+ actual_seq_klen.has_value() ? c10::asIntArrayRefUnchecked(actual_seq_klen.value()) : at::IntArrayRef{};
52- char *layout_ptr = const_cast<char *>(layout_str.data());52+ c10::string_view layout_str = layout.value_or("BSND");
53- int64_t sparse_mode_const = sparse_mode.value_or(3);53+ char* layout_ptr = const_cast<char*>(layout_str.data());
54- int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);54+ int64_t sparse_mode_const = sparse_mode.value_or(3);
55- int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);55+ int64_t pre_tokens_const = pre_tokens.value_or(9223372036854775807);
56- bool deterministic_const = true;56+ int64_t next_tokens_const = next_tokens.value_or(9223372036854775807);
57- TORCH_CHECK(query.dim() == DIMENSION_3D || query.dim() == DIMENSION_4D,57+ bool deterministic_const = true;
58- "The shapes of the input query should be 3 or 4 dimensional, but got ",58+ TORCH_CHECK(
59- query.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));59+ query.dim() == DIMENSION_3D || query.dim() == DIMENSION_4D,
60- if (query_rope_const.defined()) {60+ "The shapes of the input query should be 3 or 4 dimensional, but got ",
61- TORCH_CHECK(query_rope_const.dim() == DIMENSION_3D || query_rope_const.dim() == DIMENSION_4D,61+ query.dim(),
62- "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",62+ "-dimensional",
63- query_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));63+ OPS_ERROR(ErrCode::PARAM));
64- }64+ if (query_rope_const.defined()) {
65- TORCH_CHECK(key.dim() == DIMENSION_3D || key.dim() == DIMENSION_4D,65+ TORCH_CHECK(
66- "The shapes of the input key should be 3 or 4 dimensional, but got ", key.dim(),66+ query_rope_const.dim() == DIMENSION_3D || query_rope_const.dim() == DIMENSION_4D,
67- "-dimensional", OPS_ERROR(ErrCode::PARAM));67+ "The shapes of the input query_rope should be 3 or 4 dimensional, but got ",
68- if (key_rope_const.defined()) {68+ query_rope_const.dim(),
69- TORCH_CHECK(key_rope_const.dim() == DIMENSION_3D || key_rope_const.dim() == DIMENSION_4D,69+ "-dimensional",
70- "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",70+ OPS_ERROR(ErrCode::PARAM));
71- key_rope_const.dim(), "-dimensional", OPS_ERROR(ErrCode::PARAM));71+ }
72- }72+ TORCH_CHECK(
73- at::Tensor d_query_index = OpPreparation::apply_tensor_without_format(query_index);73+ key.dim() == DIMENSION_3D || key.dim() == DIMENSION_4D,
74- at::Tensor d_key_index = OpPreparation::apply_tensor_without_format(key_index);74+ "The shapes of the input key should be 3 or 4 dimensional, but got ",
75- at::Tensor d_weights = OpPreparation::apply_tensor_without_format(weights);75+ key.dim(),
76- at::Tensor loss = OpPreparation::apply_tensor_without_format({1}, query.options().dtype(at::kFloat));76+ "-dimensional",
77+ OPS_ERROR(ErrCode::PARAM));
78+ if (key_rope_const.defined()) {
79+ TORCH_CHECK(
80+ key_rope_const.dim() == DIMENSION_3D || key_rope_const.dim() == DIMENSION_4D,
81+ "The shapes of the input key_rope should be 3 or 4 dimensional, but got ",
82+ key_rope_const.dim(),
83+ "-dimensional",
84+ OPS_ERROR(ErrCode::PARAM));
85+ }
86+ at::Tensor d_query_index = OpPreparation::apply_tensor_without_format(query_index);
87+ at::Tensor d_key_index = OpPreparation::apply_tensor_without_format(key_index);
88+ at::Tensor d_weights = OpPreparation::apply_tensor_without_format(weights);
89+ at::Tensor loss = OpPreparation::apply_tensor_without_format({1}, query.options().dtype(at::kFloat));
77 90 
78- EXEC_NPU_NO_FORMAT_CHECK_CMD(91+ EXEC_NPU_NO_FORMAT_CHECK_CMD(
79- aclnnSparseLightningIndexerGradKLLoss, query, key, query_index, key_index, weights,92+ aclnnSparseLightningIndexerGradKLLoss,
80- sparse_indices, softmax_max, softmax_sum, query_rope_const, key_rope_const, actual_seq_qlen_const,93+ query,
81- actual_seq_klen_const, scale_value, layout_ptr, sparse_mode_const, pre_tokens_const, next_tokens_const, deterministic_const,94+ key,
82- d_query_index, d_key_index, d_weights, loss);95+ query_index,
96+ key_index,
97+ weights,
98+ sparse_indices,
99+ softmax_max,
100+ softmax_sum,
101+ query_rope_const,
102+ key_rope_const,
103+ actual_seq_qlen_const,
104+ actual_seq_klen_const,
105+ scale_value,
106+ layout_ptr,
107+ sparse_mode_const,
108+ pre_tokens_const,
109+ next_tokens_const,
110+ deterministic_const,
111+ d_query_index,
112+ d_key_index,
113+ d_weights,
114+ loss);
83 115 
84- return std::make_tuple(d_query_index, d_key_index, d_weights, loss);116+ return std::make_tuple(d_query_index, d_key_index, d_weights, loss);
85-}
86}117}
118+} // namespace op_api
Mop_plugin/ops/opapi/StdKernelNpuOpApi.cpp+41-36
@@ -21,44 +21,49 @@
21namespace op_api {21namespace op_api {
22using npu_preparation = at_npu::native::OpPreparation;22using npu_preparation = at_npu::native::OpPreparation;
23 23 
24-at::Tensor& std_out(const at::Tensor& self, at::OptionalIntArrayRef dim,24+at::Tensor& std_out(
25- const c10::optional<c10::Scalar>& correction, bool keepdim, at::Tensor& out)25+ const at::Tensor& self,
26-{26+ at::OptionalIntArrayRef dim,
27- if (!correction_fits_aclnn_int64(correction)) {27+ const c10::optional<c10::Scalar>& correction,
28- at::Tensor cpu_out = out.cpu();28+ bool keepdim,
29- at::std_out(cpu_out, self.cpu(), dim, correction, keepdim);29+ at::Tensor& out) {
30- out.copy_(cpu_out);30+ if (!correction_fits_aclnn_int64(correction)) {
31- return out;31+ at::Tensor cpu_out = out.cpu();
32- }32+ at::std_out(cpu_out, self.cpu(), dim, correction, keepdim);
33- DO_COMPATIBILITY(aclnnStd, acl_op::std_out(self, dim, correction, keepdim, out));33+ out.copy_(cpu_out);
34- c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);
35- if (dim.has_value()) {
36- real_dim = op_infer::array_to_small_vector(dim.value());
37- }
38- auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);
39- int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;
40- auto real_dim_array = at::IntArrayRef(real_dim);
41- npu_preparation::check_tensor({self}, out, self, output_size);
42- EXEC_NPU_CMD(aclnnStd, self, real_dim_array, real_correction, keepdim, out);
43 return out;34 return out;
35+ }
36+ DO_COMPATIBILITY(aclnnStd, acl_op::std_out(self, dim, correction, keepdim, out));
37+ c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);
38+ if (dim.has_value()) {
39+ real_dim = op_infer::array_to_small_vector(dim.value());
40+ }
41+ auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);
42+ int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;
43+ auto real_dim_array = at::IntArrayRef(real_dim);
44+ npu_preparation::check_tensor({self}, out, self, output_size);
45+ EXEC_NPU_CMD(aclnnStd, self, real_dim_array, real_correction, keepdim, out);
46+ return out;
44}47}
45 48 
46-at::Tensor std(const at::Tensor& self, at::OptionalIntArrayRef dim,49+at::Tensor std(
47- const c10::optional<c10::Scalar>& correction, bool keepdim)50+ const at::Tensor& self,
48-{51+ at::OptionalIntArrayRef dim,
49- if (!correction_fits_aclnn_int64(correction)) {52+ const c10::optional<c10::Scalar>& correction,
50- return at::std(self.cpu(), dim, correction, keepdim).to(self.options());53+ bool keepdim) {
51- }54+ if (!correction_fits_aclnn_int64(correction)) {
52- DO_COMPATIBILITY(aclnnStd, acl_op::std(self, dim, correction, keepdim));55+ return at::std(self.cpu(), dim, correction, keepdim).to(self.options());
53- c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);56+ }
54- if (dim.has_value()) {57+ DO_COMPATIBILITY(aclnnStd, acl_op::std(self, dim, correction, keepdim));
55- real_dim = op_infer::array_to_small_vector(dim.value());58+ c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);
56- }59+ if (dim.has_value()) {
57- auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);60+ real_dim = op_infer::array_to_small_vector(dim.value());
58- auto result = npu_preparation::apply_tensor_without_format(output_size, self.options());61+ }
59- int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;62+ auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);
60- auto real_dim_array = at::IntArrayRef(real_dim);63+ auto result = npu_preparation::apply_tensor_without_format(output_size, self.options());
61- EXEC_NPU_CMD(aclnnStd, self, real_dim_array, real_correction, keepdim, result);64+ int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;
62- return result;65+ auto real_dim_array = at::IntArrayRef(real_dim);
66+ EXEC_NPU_CMD(aclnnStd, self, real_dim_array, real_correction, keepdim, result);
67+ return result;
63}68}
64} // namespace op_api69} // namespace op_api
Mop_plugin/ops/opapi/StdMeanKernelNpuOpApi.cpp+17-20
@@ -24,27 +24,24 @@ std::tuple<at::Tensor, at::Tensor> std_mean(
24 const at::Tensor& self,24 const at::Tensor& self,
25 at::OptionalIntArrayRef dim,25 at::OptionalIntArrayRef dim,
26 const c10::optional<at::Scalar>& correction,26 const c10::optional<at::Scalar>& correction,
27- bool keepdim)27+ bool keepdim) {
28-{28+ if (!correction_fits_aclnn_int64(correction)) {
29- if (!correction_fits_aclnn_int64(correction)) {29+ auto cpu_tup = at::std_mean(self.cpu(), dim, correction, keepdim);
30- auto cpu_tup = at::std_mean(self.cpu(), dim, correction, keepdim);30+ return std::make_tuple(std::get<0>(cpu_tup).to(self.options()), std::get<1>(cpu_tup).to(self.options()));
31- return std::make_tuple(31+ }
32- std::get<0>(cpu_tup).to(self.options()),32+ DO_COMPATIBILITY(aclnnStdMeanCorrection, acl_op::std_mean(self, dim, correction, keepdim));
33- std::get<1>(cpu_tup).to(self.options()));33+ c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);
34- }34+ if (dim.has_value()) {
35- DO_COMPATIBILITY(aclnnStdMeanCorrection, acl_op::std_mean(self, dim, correction, keepdim));35+ real_dim = op_infer::array_to_small_vector(dim.value());
36- c10::SmallVector<int64_t, SIZE> real_dim = op_plugin::utils::get_dimlist_for_tensor(self);36+ }
37- if (dim.has_value()) {37+ auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);
38- real_dim = op_infer::array_to_small_vector(dim.value());
39- }
40- auto output_size = op_infer::reduce_ops_npu_output_size(self, real_dim, keepdim);
41 38 
42- at::Tensor std_out = npu_preparation::apply_tensor_without_format(self, output_size);39+ at::Tensor std_out = npu_preparation::apply_tensor_without_format(self, output_size);
43- at::Tensor mean_out = npu_preparation::apply_tensor_without_format(self, output_size);40+ at::Tensor mean_out = npu_preparation::apply_tensor_without_format(self, output_size);
44 41 
45- int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;42+ int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1;
46- auto real_dim_array = at::IntArrayRef(real_dim);43+ auto real_dim_array = at::IntArrayRef(real_dim);
47- EXEC_NPU_CMD(aclnnStdMeanCorrection, self, real_dim_array, real_correction, keepdim, std_out, mean_out);44+ EXEC_NPU_CMD(aclnnStdMeanCorrection, self, real_dim_array, real_correction, keepdim, std_out, mean_out);
48- return std::tie(std_out, mean_out);45+ return std::tie(std_out, mean_out);
49}46}
50} // namespace op_api47} // namespace op_api
Mop_plugin/ops/opapi/SubKernelNpuOpApi.cpp+89-91
@@ -20,115 +20,113 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-inline void alpha_check_npu_tensor(const at::ScalarType self_dtype, const at::ScalarType other_dtype, at::Scalar alpha)23+inline void alpha_check_npu_tensor(
24-{24+ const at::ScalarType self_dtype,
25- TORCH_CHECK(isFloatingType(self_dtype) || isComplexType(self_dtype) ||25+ const at::ScalarType other_dtype,
26- isFloatingType(other_dtype) || isComplexType(other_dtype) || alpha.isIntegral(true),26+ at::Scalar alpha) {
27- "For integral input tensors, argument alpha must not be a floating point number.",27+ TORCH_CHECK(
28- OPS_ERROR(ErrCode::TYPE));28+ isFloatingType(self_dtype) || isComplexType(self_dtype) || isFloatingType(other_dtype) ||
29+ isComplexType(other_dtype) || alpha.isIntegral(true),
30+ "For integral input tensors, argument alpha must not be a floating point number.",
31+ OPS_ERROR(ErrCode::TYPE));
29}32}
30 33 
31-inline void alpha_check_npu_scalar(const at::ScalarType self_dtype, at::Scalar other, at::Scalar alpha)34+inline void alpha_check_npu_scalar(const at::ScalarType self_dtype, at::Scalar other, at::Scalar alpha) {
32-{35+ TORCH_CHECK(
33- TORCH_CHECK(isFloatingType(self_dtype) || isComplexType(self_dtype) ||36+ isFloatingType(self_dtype) || isComplexType(self_dtype) || other.isFloatingPoint() || alpha.isIntegral(true),
34- other.isFloatingPoint() || alpha.isIntegral(true),37+ "For integral input tensors, argument alpha must not be a floating point number.",
35- "For integral input tensors, argument alpha must not be a floating point number.",38+ OPS_ERROR(ErrCode::TYPE));
36- OPS_ERROR(ErrCode::TYPE));
37}39}
38 40 
39-static at::Tensor &sub_out_npu_nocheck(const at::Tensor &self, const at::Tensor &other, const at::Scalar alpha,41+static at::Tensor& sub_out_npu_nocheck(
40- at::Tensor &result)42+ const at::Tensor& self,
41-{43+ const at::Tensor& other,
42- if (npu_preparation::IsCPUScalar(other)) {44+ const at::Scalar alpha,
43- c10::Scalar other_scalar = other.item();45+ at::Tensor& result) {
44- EXEC_NPU_CMD(aclnnSubs, self, other_scalar, alpha, result);46+ if (npu_preparation::IsCPUScalar(other)) {
45- } else {47+ c10::Scalar other_scalar = other.item();
46- EXEC_NPU_CMD(aclnnSub, self, other, alpha, result);48+ EXEC_NPU_CMD(aclnnSubs, self, other_scalar, alpha, result);
47- }49+ } else {
48- return result;50+ EXEC_NPU_CMD(aclnnSub, self, other, alpha, result);
51+ }
52+ return result;
49}53}
50 54 
51-static at::Tensor& inplace_sub_out_npu_no_check(at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha)55+static at::Tensor& inplace_sub_out_npu_no_check(at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) {
52-{56+ if (npu_preparation::IsCPUScalar(other)) {
53- if (npu_preparation::IsCPUScalar(other)) {57+ c10::Scalar other_scalar = other.item();
54- c10::Scalar other_scalar = other.item();58+ EXEC_NPU_CMD(aclnnInplaceSubs, self, other_scalar, alpha);
55- EXEC_NPU_CMD(aclnnInplaceSubs, self, other_scalar, alpha);59+ } else {
56- } else {60+ EXEC_NPU_CMD(aclnnInplaceSub, self, other, alpha);
57- EXEC_NPU_CMD(aclnnInplaceSub, self, other, alpha);61+ }
58- }62+ return self;
59- return self;
60}63}
61 64 
62-static at::Tensor self_tensor_to_device(const at::Tensor &tensor, const at::ScalarType result_type,65+static at::Tensor self_tensor_to_device(
63- const c10::Device device)66+ const at::Tensor& tensor,
64-{67+ const at::ScalarType result_type,
65- if (npu_preparation::is_scalar_wrapped_to_tensor(tensor)) {68+ const c10::Device device) {
66- at::Scalar scalar = tensor.item();69+ if (npu_preparation::is_scalar_wrapped_to_tensor(tensor)) {
67- return npu_preparation::copy_scalar_to_device(scalar, result_type, device);70+ at::Scalar scalar = tensor.item();
68- }71+ return npu_preparation::copy_scalar_to_device(scalar, result_type, device);
69- return tensor;72+ }
73+ return tensor;
70}74}
71 75 
72-static at::Tensor sub_dest_output(const at::Tensor& self, const at::Tensor& other)76+static at::Tensor sub_dest_output(const at::Tensor& self, const at::Tensor& other) {
73-{77+ bool is_self_wrapped = npu_preparation::is_scalar_wrapped_to_tensor(self);
74- bool is_self_wrapped = npu_preparation::is_scalar_wrapped_to_tensor(self);78+ return is_self_wrapped ? other : self;
75- return is_self_wrapped ? other : self;
76}79}
77 80 
78-at::Tensor &sub_out(const at::Tensor &self, const at::Tensor &other, const at::Scalar &alpha, at::Tensor &result)81+at::Tensor& sub_out(const at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha, at::Tensor& result) {
79-{82+ DO_COMPATIBILITY(aclnnSub, acl_op::sub_out(self, other, alpha, result));
80- DO_COMPATIBILITY(aclnnSub, acl_op::sub_out(self, other, alpha, result));83+ DO_COMPATIBILITY(aclnnSubs, acl_op::sub_out(self, other, alpha, result));
81- DO_COMPATIBILITY(aclnnSubs, acl_op::sub_out(self, other, alpha, result));84+ alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);
82- alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);85+ auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
83- auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);86+ at::ScalarType result_type = at::native::result_type(self, other);
84- at::ScalarType result_type = at::native::result_type(self, other);87+ at::Tensor self_converted = self_tensor_to_device(self, result_type, result.device());
85- at::Tensor self_converted = self_tensor_to_device(self, result_type, result.device());88+ npu_preparation::check_tensor({self}, result, result, output_size);
86- npu_preparation::check_tensor({self}, result, result, output_size);89+ npu_preparation::check_memory({self, other}, {result});
87- npu_preparation::check_memory({self, other}, {result});90+ sub_out_npu_nocheck(self_converted, other, alpha, result);
88- sub_out_npu_nocheck(self_converted, other, alpha, result);91+ return result;
89- return result;
90}92}
91 93 
92-at::Tensor sub(const at::Tensor &self, const at::Tensor &other, const at::Scalar &alpha)94+at::Tensor sub(const at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) {
93-{95+ DO_COMPATIBILITY(aclnnSub, acl_op::sub(self, other, alpha));
94- DO_COMPATIBILITY(aclnnSub, acl_op::sub(self, other, alpha));96+ DO_COMPATIBILITY(aclnnSubs, acl_op::sub(self, other, alpha));
95- DO_COMPATIBILITY(aclnnSubs, acl_op::sub(self, other, alpha));97+ alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);
96- alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);98+ at::Tensor output_tensor = sub_dest_output(self, other);
97- at::Tensor output_tensor = sub_dest_output(self, other);99+ auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
98- auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);100+ at::ScalarType result_type = at::native::result_type(self, other);
99- at::ScalarType result_type = at::native::result_type(self, other);101+ at::Tensor self_converted = self_tensor_to_device(self, result_type, output_tensor.device());
100- at::Tensor self_converted = self_tensor_to_device(self, result_type, output_tensor.device());102+ auto result = npu_preparation::apply_tensor_without_format(output_size, output_tensor.options().dtype(result_type));
101- auto result = npu_preparation::apply_tensor_without_format(output_size, output_tensor.options().dtype(result_type));103+ sub_out_npu_nocheck(self_converted, other, alpha, result);
102- sub_out_npu_nocheck(self_converted, other, alpha, result);104+ return result;
103- return result;
104}105}
105 106 
106-at::Tensor sub(const at::Tensor &self, const at::Scalar &other, const at::Scalar &alpha)107+at::Tensor sub(const at::Tensor& self, const at::Scalar& other, const at::Scalar& alpha) {
107-{108+ DO_COMPATIBILITY(aclnnSubs, acl_op::sub(self, other, alpha));
108- DO_COMPATIBILITY(aclnnSubs, acl_op::sub(self, other, alpha));109+ alpha_check_npu_scalar(self.scalar_type(), other, alpha);
109- alpha_check_npu_scalar(self.scalar_type(), other, alpha);110+ auto output_size = op_infer::input_same_output_size(self);
110- auto output_size = op_infer::input_same_output_size(self);111+ at::ScalarType result_type = at::native::result_type(self, other);
111- at::ScalarType result_type = at::native::result_type(self, other);112+ auto result = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(result_type));
112- auto result = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(result_type));113+ EXEC_NPU_CMD(aclnnSubs, self, other, alpha, result);
113- EXEC_NPU_CMD(aclnnSubs, self, other, alpha, result);114+ return result;
114- return result;
115}115}
116 116 
117-at::Tensor &sub_(at::Tensor &self, const at::Tensor &other, const at::Scalar &alpha)117+at::Tensor& sub_(at::Tensor& self, const at::Tensor& other, const at::Scalar& alpha) {
118-{118+ DO_COMPATIBILITY(aclnnInplaceSub, acl_op::sub_(self, other, alpha));
119- DO_COMPATIBILITY(aclnnInplaceSub, acl_op::sub_(self, other, alpha));119+ DO_COMPATIBILITY(aclnnInplaceSubs, acl_op::sub_(self, other, alpha));
120- DO_COMPATIBILITY(aclnnInplaceSubs, acl_op::sub_(self, other, alpha));120+ alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);
121- alpha_check_npu_tensor(self.scalar_type(), other.scalar_type(), alpha);121+ npu_preparation::check_memory({self, other}, {self});
122- npu_preparation::check_memory({self, other}, {self});122+ inplace_sub_out_npu_no_check(self, other, alpha);
123- inplace_sub_out_npu_no_check(self, other, alpha);123+ return self;
124- return self;
125}124}
126 125 
127-at::Tensor &sub_(at::Tensor &self, const at::Scalar &other, const at::Scalar &alpha)126+at::Tensor& sub_(at::Tensor& self, const at::Scalar& other, const at::Scalar& alpha) {
128-{127+ DO_COMPATIBILITY(aclnnInplaceSubs, acl_op::sub_(self, other, alpha));
129- DO_COMPATIBILITY(aclnnInplaceSubs, acl_op::sub_(self, other, alpha));128+ alpha_check_npu_scalar(self.scalar_type(), other, alpha);
130- alpha_check_npu_scalar(self.scalar_type(), other, alpha);129+ EXEC_NPU_CMD(aclnnInplaceSubs, self, other, alpha);
131- EXEC_NPU_CMD(aclnnInplaceSubs, self, other, alpha);130+ return self;
132- return self;
133-}
134}131}
132+} // namespace op_api
Mop_plugin/ops/opapi/SumKernelNpuOpApi.cpp+28-33
@@ -21,46 +21,41 @@
21namespace op_api {21namespace op_api {
22 22 
23#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)23#if !VERSION_BETWEEN(V2R13, VERSION_NEWEST)
24-at::Tensor& sum_out(const at::Tensor &self,24+at::Tensor& sum_out(
25- at::DimnameList dim,25+ const at::Tensor& self,
26- bool keepdim,26+ at::DimnameList dim,
27- c10::optional<c10::ScalarType> dtype,27+ bool keepdim,
28- at::Tensor &result)28+ c10::optional<c10::ScalarType> dtype,
29-{29+ at::Tensor& result) {
30- DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum_out(self, dim, keepdim, dtype, result));30+ DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum_out(self, dim, keepdim, dtype, result));
31- return op_api::sum_out(self, dimnames_to_positions(self, dim), keepdim, dtype, result);31+ return op_api::sum_out(self, dimnames_to_positions(self, dim), keepdim, dtype, result);
32}32}
33 33 
34-at::Tensor sum(const at::Tensor &self,34+at::Tensor sum(const at::Tensor& self, at::DimnameList dim, bool keepdim, c10::optional<c10::ScalarType> dtype) {
35- at::DimnameList dim,35+ DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dim, keepdim, dtype));
36- bool keepdim,36+ return op_api::sum(self, dimnames_to_positions(self, dim), keepdim, dtype);
37- c10::optional<c10::ScalarType> dtype)
38-{
39- DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dim, keepdim, dtype));
40- return op_api::sum(self, dimnames_to_positions(self, dim), keepdim, dtype);
41}37}
42#endif38#endif
43 39 
44-at::Tensor sum(const at::Tensor &self, c10::optional<c10::ScalarType> dtype)40+at::Tensor sum(const at::Tensor& self, c10::optional<c10::ScalarType> dtype) {
45-{41+ DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dtype));
46- DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dtype));42+ return op_api::sum(self, c10::SmallVector<int64_t, N>{}, false, dtype);
47- return op_api::sum(self, c10::SmallVector<int64_t, N>{}, false, dtype);
48}43}
49 44 
50-at::Tensor& sum_out(const at::Tensor &self,45+at::Tensor& sum_out(
51- at::OptionalIntArrayRef dim,46+ const at::Tensor& self,
52- bool keepdim,47+ at::OptionalIntArrayRef dim,
53- c10::optional<c10::ScalarType> dtype,48+ bool keepdim,
54- at::Tensor &result)49+ c10::optional<c10::ScalarType> dtype,
55-{50+ at::Tensor& result) {
56- return op_api::sum_out_common_nocheck(self, dim.value_or(at::IntArrayRef{}), keepdim, dtype, result);51+ return op_api::sum_out_common_nocheck(self, dim.value_or(at::IntArrayRef{}), keepdim, dtype, result);
57}52}
58 53 
59-at::Tensor sum(const at::Tensor &self,54+at::Tensor sum(
60- at::OptionalIntArrayRef dim,55+ const at::Tensor& self,
61- bool keepdim,56+ at::OptionalIntArrayRef dim,
62- c10::optional<c10::ScalarType> dtype)57+ bool keepdim,
63-{58+ c10::optional<c10::ScalarType> dtype) {
64- return op_api::sum_common_nocheck(self, dim.value_or(at::IntArrayRef{}), keepdim, dtype);59+ return op_api::sum_common_nocheck(self, dim.value_or(at::IntArrayRef{}), keepdim, dtype);
65-}
66}60}
61+} // namespace op_api
Mop_plugin/ops/opapi/SwigluBackwardMxQuantWithDualAxisOpApi.cpp+78-59
@@ -33,76 +33,95 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> _npu_swiglu_backward_
33 c10::string_view round_mode,33 c10::string_view round_mode,
34 int64_t scale_alg,34 int64_t scale_alg,
35 int64_t dst_type,35 int64_t dst_type,
36- double dst_type_max)36+ double dst_type_max) {
37-{37+ TORCH_CHECK(x.dim() >= MIN_INPUT_DIM, "The input x should be at least 2D" + OPS_ERROR(ErrCode::PARAM));
38- TORCH_CHECK(x.dim() >= MIN_INPUT_DIM, "The input x should be at least 2D" + OPS_ERROR(ErrCode::PARAM));38+ TORCH_CHECK(
39- TORCH_CHECK(x.size(-1) % SPLIT_BLOCK_SIZE == 0 && x.size(-1) > 0, "The last dim of input must be divisible by 64 "39+ x.size(-1) % SPLIT_BLOCK_SIZE == 0 && x.size(-1) > 0,
40- "and more than 0." + OPS_ERROR(ErrCode::PARAM));40+ "The last dim of input must be divisible by 64 "
41+ "and more than 0." +
42+ OPS_ERROR(ErrCode::PARAM));
41 43 
42- static const bool is_available = check_aclnn_kernel_available("aclnnSwigluBackwardMxQuantWithDualAxis");44+ static const bool is_available = check_aclnn_kernel_available("aclnnSwigluBackwardMxQuantWithDualAxis");
43- TORCH_CHECK(is_available,45+ TORCH_CHECK(
44- "Current CANN version do not support this api: _npu_swiglu_backward_mx_quant_with_dual_axis. "46+ is_available,
45- "Please try to update the version of CANN." + OPS_ERROR(ErrCode::PARAM));47+ "Current CANN version do not support this api: _npu_swiglu_backward_mx_quant_with_dual_axis. "
48+ "Please try to update the version of CANN." +
49+ OPS_ERROR(ErrCode::PARAM));
46 50 
47- const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });51+ const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });
48- char *round_mode_ptr = const_cast<char *>(round_mode.data());52+ char* round_mode_ptr = const_cast<char*>(round_mode.data());
49- // y1_out and y2_out have the same shape as x53+ // y1_out and y2_out have the same shape as x
50- auto x_grad_shape = op_infer::array_to_small_vector(x.sizes());54+ auto x_grad_shape = op_infer::array_to_small_vector(x.sizes());
51 55 
52- // Infer mxscale1 shape: ceil(last_dim / 64) + append 256+ // Infer mxscale1 shape: ceil(last_dim / 64) + append 2
53- auto mxscale1_shape = op_infer::array_to_small_vector(x.sizes());57+ auto mxscale1_shape = op_infer::array_to_small_vector(x.sizes());
54- int64_t last_dim = x.size(-1);58+ int64_t last_dim = x.size(-1);
55- mxscale1_shape[mxscale1_shape.size() - 1] =59+ mxscale1_shape[mxscale1_shape.size() - 1] =
56- static_cast<int64_t>(std::ceil(static_cast<double>(last_dim) / SPLIT_BLOCK_SIZE));60+ static_cast<int64_t>(std::ceil(static_cast<double>(last_dim) / SPLIT_BLOCK_SIZE));
57- mxscale1_shape.emplace_back(NUM_TWO);61+ mxscale1_shape.emplace_back(NUM_TWO);
58 62 
59- // Infer mxscale2 shape63+ // Infer mxscale2 shape
60- auto mxscale2_shape = op_infer::array_to_small_vector(x.sizes());64+ auto mxscale2_shape = op_infer::array_to_small_vector(x.sizes());
61- int64_t second_to_last_dim = x.size(-2);65+ int64_t second_to_last_dim = x.size(-2);
62- int64_t quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));66+ int64_t quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));
63- if (group_index_opt.defined()) {67+ if (group_index_opt.defined()) {
64- quant_size = static_cast<int64_t>(std::floor(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE))68+ quant_size = static_cast<int64_t>(std::floor(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE)) +
65- + group_index_opt.size(0);69+ group_index_opt.size(0);
66- }70+ }
67- mxscale2_shape[mxscale2_shape.size() - 2] = quant_size;71+ mxscale2_shape[mxscale2_shape.size() - 2] = quant_size;
68- mxscale2_shape.emplace_back(NUM_TWO);72+ mxscale2_shape.emplace_back(NUM_TWO);
69 73 
70- aclDataType x_acltype = c10_npu::GetAclDataType(dst_type);74+ aclDataType x_acltype = c10_npu::GetAclDataType(dst_type);
71- ASCEND_LOGI("[npu_swiglu_backward_mx_quant_with_dual_axis]: "75+ ASCEND_LOGI(
72- "Getting aclTensor y1_out and y2_out dtype by Parameter(dst_type): %ld", dst_type);76+ "[npu_swiglu_backward_mx_quant_with_dual_axis]: "
77+ "Getting aclTensor y1_out and y2_out dtype by Parameter(dst_type): %ld",
78+ dst_type);
73 79 
74- bool special_output_type = (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||80+ bool special_output_type =
75- dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));81+ (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
82+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
76 83 
77- at::Tensor y1_out;84+ at::Tensor y1_out;
78- at::Tensor y2_out;85+ at::Tensor y2_out;
79- if (special_output_type) {86+ if (special_output_type) {
80- int64_t x_last_dim_val = x_grad_shape[x_grad_shape.size() - 1];87+ int64_t x_last_dim_val = x_grad_shape[x_grad_shape.size() - 1];
81- TORCH_CHECK(x_last_dim_val % NUM_TWO == 0,88+ TORCH_CHECK(
82- "The last dim of x_grad must be divisible by 2 if dtype is float4_e2m1 or float4_e1m2"89+ x_last_dim_val % NUM_TWO == 0,
83- + OPS_ERROR(ErrCode::PARAM));90+ "The last dim of x_grad must be divisible by 2 if dtype is float4_e2m1 or float4_e1m2" +
84- x_grad_shape[x_grad_shape.size() - 1] = x_last_dim_val / NUM_TWO;91+ OPS_ERROR(ErrCode::PARAM));
85- y1_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::ScalarType::Byte);92+ x_grad_shape[x_grad_shape.size() - 1] = x_last_dim_val / NUM_TWO;
86- y2_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::ScalarType::Byte);93+ y1_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::ScalarType::Byte);
87- } else {94+ y2_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::ScalarType::Byte);
88- at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(x_acltype);95+ } else {
89- y1_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::dtype(scalar_dtype));96+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(x_acltype);
90- y2_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::dtype(scalar_dtype));97+ y1_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::dtype(scalar_dtype));
91- }98+ y2_out = npu_preparation::apply_tensor_without_format(x_grad_shape, c10::dtype(scalar_dtype));
99+ }
92 100 
93- at::Tensor mxscale1 = npu_preparation::apply_tensor_without_format(mxscale1_shape, c10::dtype(at::ScalarType::Byte));101+ at::Tensor mxscale1 = npu_preparation::apply_tensor_without_format(mxscale1_shape, c10::dtype(at::ScalarType::Byte));
94- at::Tensor mxscale2 = npu_preparation::apply_tensor_without_format(mxscale2_shape, c10::dtype(at::ScalarType::Byte));102+ at::Tensor mxscale2 = npu_preparation::apply_tensor_without_format(mxscale2_shape, c10::dtype(at::ScalarType::Byte));
95 103 
96- TensorWrapper y1_out_wrapper = {y1_out, x_acltype};104+ TensorWrapper y1_out_wrapper = {y1_out, x_acltype};
97- TensorWrapper y2_out_wrapper = {y2_out, x_acltype};105+ TensorWrapper y2_out_wrapper = {y2_out, x_acltype};
98- TensorWrapper mxscale1_wrapper = {mxscale1, aclDataType::ACL_FLOAT8_E8M0};106+ TensorWrapper mxscale1_wrapper = {mxscale1, aclDataType::ACL_FLOAT8_E8M0};
99- TensorWrapper mxscale2_wrapper = {mxscale2, aclDataType::ACL_FLOAT8_E8M0};107+ TensorWrapper mxscale2_wrapper = {mxscale2, aclDataType::ACL_FLOAT8_E8M0};
100 108 
101- EXEC_NPU_CMD(aclnnSwigluBackwardMxQuantWithDualAxis, x, y_grad, group_index_opt, activate_left,109+ EXEC_NPU_CMD(
102- round_mode_ptr, scale_alg, x_acltype, dst_type_max,110+ aclnnSwigluBackwardMxQuantWithDualAxis,
103- y1_out_wrapper, mxscale1_wrapper, y2_out_wrapper, mxscale2_wrapper);111+ x,
112+ y_grad,
113+ group_index_opt,
114+ activate_left,
115+ round_mode_ptr,
116+ scale_alg,
117+ x_acltype,
118+ dst_type_max,
119+ y1_out_wrapper,
120+ mxscale1_wrapper,
121+ y2_out_wrapper,
122+ mxscale2_wrapper);
104 123 
105- return std::make_tuple(y1_out, mxscale1, y2_out, mxscale2);124+ return std::make_tuple(y1_out, mxscale1, y2_out, mxscale2);
106}125}
107 126 
108} // namespace op_api127} // namespace op_api
Mop_plugin/ops/opapi/SwigluGroupQuantBackwardKernelNpuOpApi.cpp+19-17
@@ -20,27 +20,29 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-std::tuple<at::Tensor, at::Tensor> npu_swiglu_group_quant_backward(const at::Tensor &grad_y, const at::Tensor &x,23+std::tuple<at::Tensor, at::Tensor> npu_swiglu_group_quant_backward(
24- const c10::optional<at::Tensor> &weight, const c10::optional<at::Tensor> &y_origin,24+ const at::Tensor& grad_y,
25- const c10::optional<at::Tensor> &group_index, double clamp_limit)25+ const at::Tensor& x,
26-{26+ const c10::optional<at::Tensor>& weight,
27+ const c10::optional<at::Tensor>& y_origin,
28+ const c10::optional<at::Tensor>& group_index,
29+ double clamp_limit) {
30+ // check x last dim
31+ int64_t x_last_dim = x.size(x.dim() - 1);
32+ TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));
27 33 
28- // check x last dim34+ at::Tensor grad_x = npu_preparation::apply_tensor_without_format(x.sizes(), x.options());
29- int64_t x_last_dim = x.size(x.dim() - 1);
30- TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));
31 35 
32- at::Tensor grad_x = npu_preparation::apply_tensor_without_format(x.sizes(), x.options());36+ at::Tensor grad_weight;
33 37 
34- at::Tensor grad_weight;38+ if (weight.has_value() && weight->defined()) {
39+ grad_weight = npu_preparation::apply_tensor_without_format(weight.value().sizes(), weight.value().options());
40+ } else {
41+ grad_weight = at::empty({0}, x.options().dtype(at::kFloat));
42+ }
35 43 
36- if (weight.has_value() && weight->defined()) {44+ EXEC_NPU_CMD(aclnnSwigluGroupQuantGrad, grad_y, x, weight, y_origin, group_index, clamp_limit, grad_x, grad_weight);
37- grad_weight = npu_preparation::apply_tensor_without_format(weight.value().sizes(), weight.value().options());
38- } else {
39- grad_weight = at::empty({0}, x.options().dtype(at::kFloat));
40- }
41 45 
42- EXEC_NPU_CMD(aclnnSwigluGroupQuantGrad, grad_y, x, weight, y_origin, group_index, clamp_limit, grad_x, grad_weight);46+ return std::make_tuple(grad_x, grad_weight);
43- 
44- return std::make_tuple(grad_x, grad_weight);
45}47}
46} // namespace op_api48} // namespace op_api
Mop_plugin/ops/opapi/SwigluGroupQuantKernelNpuOpApi.cpp+63-40
@@ -23,56 +23,79 @@ using npu_preparation = at_npu::native::OpPreparation;
23namespace {23namespace {
24constexpr int64_t QUANT_MODE_2 = 2;24constexpr int64_t QUANT_MODE_2 = 2;
25constexpr int64_t QUANT_MODE_3 = 3;25constexpr int64_t QUANT_MODE_3 = 3;
26-} // namespace26+} // namespace
27 27 
28-std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_swiglu_group_quant(const at::Tensor &x,28+std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_swiglu_group_quant(
29- const c10::optional<at::Tensor> &weight, const c10::optional<at::Tensor> &group_index,29+ const at::Tensor& x,
30- const c10::optional<at::Tensor> &scale, int64_t dst_type, int64_t quant_mode,30+ const c10::optional<at::Tensor>& weight,
31- int64_t block_size, bool round_scale, double clamp_limit, double dst_type_max, bool output_origin)31+ const c10::optional<at::Tensor>& group_index,
32-{32+ const c10::optional<at::Tensor>& scale,
33- at::Tensor y;33+ int64_t dst_type,
34- at::Tensor y_scale;34+ int64_t quant_mode,
35- at::Tensor y_origin = at::empty({0}, x.options());;35+ int64_t block_size,
36+ bool round_scale,
37+ double clamp_limit,
38+ double dst_type_max,
39+ bool output_origin) {
40+ at::Tensor y;
41+ at::Tensor y_scale;
42+ at::Tensor y_origin = at::empty({0}, x.options());
43+ ;
36 44 
37- int64_t x_last_dim = x.size(x.dim() - 1);45+ int64_t x_last_dim = x.size(x.dim() - 1);
38- TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));46+ TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));
39 47 
40- auto y_shape = op_infer::array_to_small_vector(x.sizes());48+ auto y_shape = op_infer::array_to_small_vector(x.sizes());
41- y_shape[x.dim() - 1] = x_last_dim / 2;49+ y_shape[x.dim() - 1] = x_last_dim / 2;
42 50 
43- aclDataType y_acltype;51+ aclDataType y_acltype;
44- if (quant_mode == QUANT_MODE_2 || quant_mode == QUANT_MODE_3) {52+ if (quant_mode == QUANT_MODE_2 || quant_mode == QUANT_MODE_3) {
45- y_acltype = aclDataType::ACL_HIFLOAT8;53+ y_acltype = aclDataType::ACL_HIFLOAT8;
46- y = npu_preparation::apply_tensor_without_format(y_shape, at::ScalarType::Byte);54+ y = npu_preparation::apply_tensor_without_format(y_shape, at::ScalarType::Byte);
55+ }
56+ 
57+ if (quant_mode == QUANT_MODE_2) {
58+ y_scale = at::empty({0}, x.options().dtype(at::kFloat));
59+ } else if (quant_mode == QUANT_MODE_3) {
60+ if (group_index.has_value() && group_index->defined()) {
61+ y_scale =
62+ npu_preparation::apply_tensor_without_format(group_index.value().sizes(), c10::dtype(at::ScalarType::Float));
63+ } else {
64+ y_scale = npu_preparation::apply_tensor_without_format({1}, c10::dtype(at::ScalarType::Float));
47 }65 }
66+ }
48 67 
49- if (quant_mode == QUANT_MODE_2) {68+ if (output_origin) {
50- y_scale = at::empty({0}, x.options().dtype(at::kFloat));69+ auto y_origin_shape = op_infer::array_to_small_vector(x.sizes());
51- } else if (quant_mode == QUANT_MODE_3) {70+ y_origin_shape[x.dim() - 1] = x_last_dim / 2;
52- if (group_index.has_value() && group_index->defined()) {71+ y_origin = npu_preparation::apply_tensor_without_format(y_origin_shape, x.options());
53- y_scale = npu_preparation::apply_tensor_without_format(group_index.value().sizes(), c10::dtype(at::ScalarType::Float));72+ }
54- } else {
55- y_scale = npu_preparation::apply_tensor_without_format({1}, c10::dtype(at::ScalarType::Float));
56- }
57- }
58 73 
59- if (output_origin) {74+ TensorWrapper y_wrapper = {y, y_acltype};
60- auto y_origin_shape = op_infer::array_to_small_vector(x.sizes());
61- y_origin_shape[x.dim() - 1] = x_last_dim / 2;
62- y_origin = npu_preparation::apply_tensor_without_format(y_origin_shape, x.options());
63- }
64 75 
65- TensorWrapper y_wrapper = {y, y_acltype};76+ aclDataType y_scale_acltype = npu_preparation::convert_to_acl_data_type(y_scale.scalar_type());
66 77 
67- aclDataType y_scale_acltype = npu_preparation::convert_to_acl_data_type(y_scale.scalar_type());78+ TensorWrapper y_scale_wrapper = {y_scale, y_scale_acltype};
68 79 
69- TensorWrapper y_scale_wrapper = {y_scale, y_scale_acltype};80+ TensorWrapper y_origin_wrapper = {y_origin, npu_preparation::convert_to_acl_data_type(y_origin.scalar_type())};
70 81 
71- TensorWrapper y_origin_wrapper = {y_origin, npu_preparation::convert_to_acl_data_type(y_origin.scalar_type())};82+ EXEC_NPU_CMD(
83+ aclnnSwigluGroupQuant,
84+ x,
85+ weight,
86+ group_index,
87+ scale,
88+ y_acltype,
89+ quant_mode,
90+ block_size,
91+ round_scale,
92+ clamp_limit,
93+ dst_type_max,
94+ output_origin,
95+ y_wrapper,
96+ y_scale_wrapper,
97+ y_origin_wrapper);
72 98 
73- EXEC_NPU_CMD(aclnnSwigluGroupQuant, x, weight, group_index, scale, y_acltype, quant_mode,99+ return std::make_tuple(y, y_scale, y_origin);
74- block_size, round_scale, clamp_limit, dst_type_max, output_origin, y_wrapper, y_scale_wrapper, y_origin_wrapper);
75- 
76- return std::make_tuple(y, y_scale, y_origin);
77}100}
78-} // namespace op_api101+} // namespace op_api
Mop_plugin/ops/opapi/SwigluMxQuantOpApi.cpp+129-89
@@ -21,111 +21,151 @@ const int64_t NUM_TWO = 2;
21constexpr float DEFAULT_BLOCKSIZE = 64.0;21constexpr float DEFAULT_BLOCKSIZE = 64.0;
22 22 
23std::tuple<at::Tensor, at::Tensor> npu_swiglu_mx_quant(23std::tuple<at::Tensor, at::Tensor> npu_swiglu_mx_quant(
24- const at::Tensor& x, const c10::optional<at::Tensor>& group_index,24+ const at::Tensor& x,
25- int64_t activate_dim, bool activate_left, int64_t swiglu_mode,25+ const c10::optional<at::Tensor>& group_index,
26- double clamp_limit, double glu_alpha, double glu_bias,26+ int64_t activate_dim,
27- int64_t group_mode, int64_t axis, int64_t dst_type,27+ bool activate_left,
28- c10::string_view round_mode, int64_t scale_alg, double max_dtype_value)28+ int64_t swiglu_mode,
29-{29+ double clamp_limit,
30- TORCH_CHECK(x.dim() > 1, "x dim should larger than 1", OPS_ERROR(ErrCode::PARAM));30+ double glu_alpha,
31- TORCH_CHECK(swiglu_mode == 0 || swiglu_mode == 1, "swiglu_mode only support 0 or 1, but got ", swiglu_mode,31+ double glu_bias,
32- OPS_ERROR(ErrCode::PARAM));32+ int64_t group_mode,
33- TORCH_CHECK(std::isfinite(clamp_limit) && clamp_limit > 0.0, "clamp_limit should be positive finite",33+ int64_t axis,
34- OPS_ERROR(ErrCode::PARAM));34+ int64_t dst_type,
35- TORCH_CHECK(std::isfinite(glu_alpha), "glu_alpha should be finite", OPS_ERROR(ErrCode::PARAM));35+ c10::string_view round_mode,
36- TORCH_CHECK(std::isfinite(glu_bias), "glu_bias should be finite", OPS_ERROR(ErrCode::PARAM));36+ int64_t scale_alg,
37+ double max_dtype_value) {
38+ TORCH_CHECK(x.dim() > 1, "x dim should larger than 1", OPS_ERROR(ErrCode::PARAM));
39+ TORCH_CHECK(
40+ swiglu_mode == 0 || swiglu_mode == 1,
41+ "swiglu_mode only support 0 or 1, but got ",
42+ swiglu_mode,
43+ OPS_ERROR(ErrCode::PARAM));
44+ TORCH_CHECK(
45+ std::isfinite(clamp_limit) && clamp_limit > 0.0,
46+ "clamp_limit should be positive finite",
47+ OPS_ERROR(ErrCode::PARAM));
48+ TORCH_CHECK(std::isfinite(glu_alpha), "glu_alpha should be finite", OPS_ERROR(ErrCode::PARAM));
49+ TORCH_CHECK(std::isfinite(glu_bias), "glu_bias should be finite", OPS_ERROR(ErrCode::PARAM));
37 50 
38- static const bool is_available = check_aclnn_kernel_available("aclnnSwigluMxQuant");51+ static const bool is_available = check_aclnn_kernel_available("aclnnSwigluMxQuant");
39- TORCH_CHECK(is_available,52+ TORCH_CHECK(
40- "Current CANN version do not support this api. Please try to update the version of CANN."53+ is_available,
41- + OPS_ERROR(ErrCode::PARAM));54+ "Current CANN version do not support this api. Please try to update the version of CANN." +
55+ OPS_ERROR(ErrCode::PARAM));
42 56 
43- const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });57+ const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });
44 58 
45- int64_t activate_dim_value = activate_dim;59+ int64_t activate_dim_value = activate_dim;
46- char *round_mode_ptr = const_cast<char *>(round_mode.data());60+ char* round_mode_ptr = const_cast<char*>(round_mode.data());
47 61 
48- // transform activate_dim62+ // transform activate_dim
49- if (activate_dim_value < 0) {63+ if (activate_dim_value < 0) {
50- activate_dim_value = activate_dim_value + x.dim();64+ activate_dim_value = activate_dim_value + x.dim();
51- }65+ }
52- TORCH_CHECK(activate_dim_value <= (x.dim() - 1) && activate_dim_value >= 0, "activate_dim should be in range [0, x.dim()-1]", OPS_ERROR(ErrCode::PARAM));66+ TORCH_CHECK(
67+ activate_dim_value <= (x.dim() - 1) && activate_dim_value >= 0,
68+ "activate_dim should be in range [0, x.dim()-1]",
69+ OPS_ERROR(ErrCode::PARAM));
53 70 
54- // Calculate quant_dim based on axis71+ // Calculate quant_dim based on axis
55- int64_t quant_dim_value = axis;72+ int64_t quant_dim_value = axis;
56- if (quant_dim_value < 0) {73+ if (quant_dim_value < 0) {
57- quant_dim_value = quant_dim_value + x.dim();74+ quant_dim_value = quant_dim_value + x.dim();
58- }75+ }
59 76 
60- TORCH_CHECK(quant_dim_value >= 0 && quant_dim_value <= (x.dim() - 1), "quant_dim should be in range [0, x.dim()-1]", OPS_ERROR(ErrCode::PARAM));77+ TORCH_CHECK(
78+ quant_dim_value >= 0 && quant_dim_value <= (x.dim() - 1),
79+ "quant_dim should be in range [0, x.dim()-1]",
80+ OPS_ERROR(ErrCode::PARAM));
61 81 
62- // select_dim is used to determine which dimension to divide by 2 for y and scale shapes82+ // select_dim is used to determine which dimension to divide by 2 for y and scale shapes
63- int64_t select_dim = activate_dim_value;83+ int64_t select_dim = activate_dim_value;
64 84 
65- at::SmallVector<int64_t, op_infer::SIZE> y_size;85+ at::SmallVector<int64_t, op_infer::SIZE> y_size;
66- at::SmallVector<int64_t, op_infer::SIZE> scale_size;86+ at::SmallVector<int64_t, op_infer::SIZE> scale_size;
67 87 
68- // Infer size of y, scale (divide by 2 at select_dim)88+ // Infer size of y, scale (divide by 2 at select_dim)
69- for (int i = 0; i < x.dim(); i++) {89+ for (int i = 0; i < x.dim(); i++) {
70- if (i == select_dim) {90+ if (i == select_dim) {
71- y_size.push_back(x.size(i) / NUM_TWO);91+ y_size.push_back(x.size(i) / NUM_TWO);
72- scale_size.push_back(x.size(i) / NUM_TWO);92+ scale_size.push_back(x.size(i) / NUM_TWO);
73- } else {
74- y_size.push_back(x.size(i));
75- scale_size.push_back(x.size(i));
76- }
77- }
78- 
79- // Calculate quant_size based on group_index and quant_dim (matching meta function logic)
80- int64_t quant_size = 0;
81- if (!group_index_opt.defined()) {
82- // group_index is None: quant_size = ceil(scale_size[quant_dim] / 64)
83- quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
84 } else {93 } else {
85- // group_index exists94+ y_size.push_back(x.size(i));
86- if (quant_dim_value == (x.dim() - 1)) {95+ scale_size.push_back(x.size(i));
87- // quant_dim is last dimension: quant_size = ceil(scale_size[quant_dim] / 64)
88- quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
89- } else {
90- // quant_dim is not last dimension: quant_size = ceil(scale_size[quant_dim] / 64) + group_index.shape[0]
91- quant_size = static_cast<int64_t>(std::floor(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
92- quant_size = quant_size + group_index_opt.sizes()[0];
93- }
94 }96 }
97+ }
95 98 
96- // Modify scale shape at quant_dim with quant_size, then append 299+ // Calculate quant_size based on group_index and quant_dim (matching meta function logic)
97- scale_size[quant_dim_value] = quant_size;100+ int64_t quant_size = 0;
98- scale_size.push_back(NUM_TWO);101+ if (!group_index_opt.defined()) {
99- 102+ // group_index is None: quant_size = ceil(scale_size[quant_dim] / 64)
100- at::Tensor y;103+ quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
101- aclDataType y_acltype;104+ } else {
102- 105+ // group_index exists
103- if (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||106+ if (quant_dim_value == (x.dim() - 1)) {
104- dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) {107+ // quant_dim is last dimension: quant_size = ceil(scale_size[quant_dim] / 64)
105- int64_t last_dim_val = y_size[x.dim() - 1];108+ quant_size =
106- TORCH_CHECK(last_dim_val % NUM_TWO == 0, "Y last dim should be even when type of y is float4_e1m2 or float4_e2m1", OPS_ERROR(ErrCode::PARAM));109+ static_cast<int64_t>(std::ceil(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
107- y_size[x.dim() - 1] = last_dim_val / NUM_TWO;
108- }
109- 
110- if (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
111- dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) {
112- y = npu_preparation::apply_tensor_without_format(y_size, c10::ScalarType::Byte);
113- y_acltype = c10_npu::GetAclDataType(dst_type);
114 } else {110 } else {
115- y_acltype = c10_npu::GetAclDataType(dst_type);111+ // quant_dim is not last dimension: quant_size = ceil(scale_size[quant_dim] / 64) + group_index.shape[0]
116- at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);112+ quant_size =
117- y = npu_preparation::apply_tensor_without_format(y_size, c10::dtype(scalar_dtype));113+ static_cast<int64_t>(std::floor(static_cast<double>(scale_size[quant_dim_value]) / DEFAULT_BLOCKSIZE));
114+ quant_size = quant_size + group_index_opt.sizes()[0];
118 }115 }
116+ }
119 117 
120- TensorWrapper y_wrapper = {y, y_acltype};118+ // Modify scale shape at quant_dim with quant_size, then append 2
119+ scale_size[quant_dim_value] = quant_size;
120+ scale_size.push_back(NUM_TWO);
121 121 
122- at::Tensor scale = npu_preparation::apply_tensor_without_format(scale_size, c10::dtype(c10::ScalarType::Byte));122+ at::Tensor y;
123- TensorWrapper mxscale_wrapper = {scale, aclDataType::ACL_FLOAT8_E8M0};123+ aclDataType y_acltype;
124 124 
125- EXEC_NPU_CMD(aclnnSwigluMxQuant, x, group_index_opt, activate_dim_value, activate_left,125+ if (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
126- swiglu_mode, clamp_limit, glu_alpha, glu_bias, group_mode, axis,126+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) {
127- y_acltype, round_mode_ptr, scale_alg, max_dtype_value, y_wrapper, mxscale_wrapper);127+ int64_t last_dim_val = y_size[x.dim() - 1];
128+ TORCH_CHECK(
129+ last_dim_val % NUM_TWO == 0,
130+ "Y last dim should be even when type of y is float4_e1m2 or float4_e2m1",
131+ OPS_ERROR(ErrCode::PARAM));
132+ y_size[x.dim() - 1] = last_dim_val / NUM_TWO;
133+ }
128 134 
129- return std::tie(y, scale);135+ if (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
136+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) {
137+ y = npu_preparation::apply_tensor_without_format(y_size, c10::ScalarType::Byte);
138+ y_acltype = c10_npu::GetAclDataType(dst_type);
139+ } else {
140+ y_acltype = c10_npu::GetAclDataType(dst_type);
141+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
142+ y = npu_preparation::apply_tensor_without_format(y_size, c10::dtype(scalar_dtype));
143+ }
144+ 
145+ TensorWrapper y_wrapper = {y, y_acltype};
146+ 
147+ at::Tensor scale = npu_preparation::apply_tensor_without_format(scale_size, c10::dtype(c10::ScalarType::Byte));
148+ TensorWrapper mxscale_wrapper = {scale, aclDataType::ACL_FLOAT8_E8M0};
149+ 
150+ EXEC_NPU_CMD(
151+ aclnnSwigluMxQuant,
152+ x,
153+ group_index_opt,
154+ activate_dim_value,
155+ activate_left,
156+ swiglu_mode,
157+ clamp_limit,
158+ glu_alpha,
159+ glu_bias,
160+ group_mode,
161+ axis,
162+ y_acltype,
163+ round_mode_ptr,
164+ scale_alg,
165+ max_dtype_value,
166+ y_wrapper,
167+ mxscale_wrapper);
168+ 
169+ return std::tie(y, scale);
130}170}
131-} // namespace op_api171+} // namespace op_api
Mop_plugin/ops/opapi/SwigluMxQuantWithDualAxisOpApi.cpp+75-61
@@ -32,81 +32,95 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_swiglu_mx_quant_w
32 c10::string_view round_mode,32 c10::string_view round_mode,
33 int64_t scale_alg,33 int64_t scale_alg,
34 int64_t dst_type,34 int64_t dst_type,
35- double dst_type_max)35+ double dst_type_max) {
36-{36+ at::Tensor y1;
37- at::Tensor y1;37+ at::Tensor mxscale1;
38- at::Tensor mxscale1;38+ at::Tensor y2;
39- at::Tensor y2;39+ at::Tensor mxscale2;
40- at::Tensor mxscale2;
41 40 
42- TORCH_CHECK(x.dim() >= MIN_INPUT_DIM, "The input x should be at least 2D" + OPS_ERROR(ErrCode::PARAM));41+ TORCH_CHECK(x.dim() >= MIN_INPUT_DIM, "The input x should be at least 2D" + OPS_ERROR(ErrCode::PARAM));
43- TORCH_CHECK(x.size(-1) % NUM_TWO == 0, "The last dim of input must be divisible by 2" + OPS_ERROR(ErrCode::PARAM));42+ TORCH_CHECK(x.size(-1) % NUM_TWO == 0, "The last dim of input must be divisible by 2" + OPS_ERROR(ErrCode::PARAM));
44 43 
45- static const bool is_available = check_aclnn_kernel_available("aclnnSwigluMxQuantWithDualAxis");44+ static const bool is_available = check_aclnn_kernel_available("aclnnSwigluMxQuantWithDualAxis");
46- TORCH_CHECK(is_available,45+ TORCH_CHECK(
47- "Current CANN version do not support this api: npu_swiglu_mx_quant_with_dual_axis. Please try to update the version of CANN."46+ is_available,
48- + OPS_ERROR(ErrCode::PARAM));47+ "Current CANN version do not support this api: npu_swiglu_mx_quant_with_dual_axis. Please try to update the version of CANN." +
48+ OPS_ERROR(ErrCode::PARAM));
49 49 
50- const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });50+ const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });
51 51 
52- char *round_mode_ptr = const_cast<char *>(round_mode.data());52+ char* round_mode_ptr = const_cast<char*>(round_mode.data());
53 53 
54- // Infer y shape: divide last dim by 254+ // Infer y shape: divide last dim by 2
55- auto y_shape = op_infer::array_to_small_vector(x.sizes());55+ auto y_shape = op_infer::array_to_small_vector(x.sizes());
56- y_shape[y_shape.size() - 1] = y_shape[y_shape.size() - 1] / NUM_TWO;56+ y_shape[y_shape.size() - 1] = y_shape[y_shape.size() - 1] / NUM_TWO;
57 57 
58- // Infer mxscale1 shape: ceil(last_dim / 64) + append 258+ // Infer mxscale1 shape: ceil(last_dim / 64) + append 2
59- auto mxscale1_shape = op_infer::array_to_small_vector(y_shape);59+ auto mxscale1_shape = op_infer::array_to_small_vector(y_shape);
60- int64_t last_dim = mxscale1_shape[mxscale1_shape.size() - 1];60+ int64_t last_dim = mxscale1_shape[mxscale1_shape.size() - 1];
61- mxscale1_shape[mxscale1_shape.size() - 1] = static_cast<int64_t>(std::ceil(static_cast<double>(last_dim) / SPLIT_BLOCK_SIZE));61+ mxscale1_shape[mxscale1_shape.size() - 1] =
62- mxscale1_shape.emplace_back(NUM_TWO);62+ static_cast<int64_t>(std::ceil(static_cast<double>(last_dim) / SPLIT_BLOCK_SIZE));
63+ mxscale1_shape.emplace_back(NUM_TWO);
63 64 
64- // Infer mxscale2 shape: floor(second_to_last_dim / 64) + group_num + append 265+ // Infer mxscale2 shape: floor(second_to_last_dim / 64) + group_num + append 2
65- auto mxscale2_shape = op_infer::array_to_small_vector(y_shape);66+ auto mxscale2_shape = op_infer::array_to_small_vector(y_shape);
66- int64_t second_to_last_dim = mxscale2_shape[mxscale2_shape.size() - 2];67+ int64_t second_to_last_dim = mxscale2_shape[mxscale2_shape.size() - 2];
67- int64_t quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));68+ int64_t quant_size = static_cast<int64_t>(std::ceil(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));
68- if (group_index_opt.defined()) {69+ if (group_index_opt.defined()) {
69- quant_size = static_cast<int64_t>(std::floor(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));70+ quant_size = static_cast<int64_t>(std::floor(static_cast<double>(second_to_last_dim) / SPLIT_BLOCK_SIZE));
70- quant_size = quant_size + group_index_opt.size(0);71+ quant_size = quant_size + group_index_opt.size(0);
71- }72+ }
72- mxscale2_shape[mxscale2_shape.size() - 2] = quant_size;73+ mxscale2_shape[mxscale2_shape.size() - 2] = quant_size;
73- mxscale2_shape.emplace_back(NUM_TWO);74+ mxscale2_shape.emplace_back(NUM_TWO);
74 75 
75- aclDataType y_acltype;76+ aclDataType y_acltype;
76- bool special_output_type = (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||77+ bool special_output_type =
77- dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));78+ (dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
79+ dst_type == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
78 80 
79- ASCEND_LOGI("[npu_swiglu_mx_quant_with_dual_axis]: Getting aclTensor y1 and y2 dtype by Parameter(dst_type): %ld", dst_type);81+ ASCEND_LOGI(
82+ "[npu_swiglu_mx_quant_with_dual_axis]: Getting aclTensor y1 and y2 dtype by Parameter(dst_type): %ld", dst_type);
80 83 
81- if (special_output_type) {84+ if (special_output_type) {
82- int64_t y_last_dim_val = y_shape[y_shape.size() - 1];85+ int64_t y_last_dim_val = y_shape[y_shape.size() - 1];
83- TORCH_CHECK(y_last_dim_val % NUM_TWO == 0,86+ TORCH_CHECK(
84- "The last dim of y must be divisible by 2 if y dtype is float4_e2m1 or float4_e1m2"87+ y_last_dim_val % NUM_TWO == 0,
85- + OPS_ERROR(ErrCode::PARAM));88+ "The last dim of y must be divisible by 2 if y dtype is float4_e2m1 or float4_e1m2" +
86- y_shape[y_shape.size() - 1] = y_last_dim_val / NUM_TWO;89+ OPS_ERROR(ErrCode::PARAM));
87- y1 = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);90+ y_shape[y_shape.size() - 1] = y_last_dim_val / NUM_TWO;
88- y2 = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);91+ y1 = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);
89- y_acltype = c10_npu::GetAclDataType(dst_type);92+ y2 = npu_preparation::apply_tensor_without_format(y_shape, c10::ScalarType::Byte);
90- } else {93+ y_acltype = c10_npu::GetAclDataType(dst_type);
91- y_acltype = c10_npu::GetAclDataType(dst_type);94+ } else {
92- at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);95+ y_acltype = c10_npu::GetAclDataType(dst_type);
93- y1 = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));96+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
94- y2 = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));97+ y1 = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));
95- }98+ y2 = npu_preparation::apply_tensor_without_format(y_shape, c10::dtype(scalar_dtype));
99+ }
96 100 
97- mxscale1 = npu_preparation::apply_tensor_without_format(mxscale1_shape, c10::dtype(at::ScalarType::Byte));101+ mxscale1 = npu_preparation::apply_tensor_without_format(mxscale1_shape, c10::dtype(at::ScalarType::Byte));
98- mxscale2 = npu_preparation::apply_tensor_without_format(mxscale2_shape, c10::dtype(at::ScalarType::Byte));102+ mxscale2 = npu_preparation::apply_tensor_without_format(mxscale2_shape, c10::dtype(at::ScalarType::Byte));
99 103 
100- TensorWrapper y1_wrapper = {y1, y_acltype};104+ TensorWrapper y1_wrapper = {y1, y_acltype};
101- TensorWrapper y2_wrapper = {y2, y_acltype};105+ TensorWrapper y2_wrapper = {y2, y_acltype};
102- TensorWrapper mxscale1_wrapper = {mxscale1, aclDataType::ACL_FLOAT8_E8M0};106+ TensorWrapper mxscale1_wrapper = {mxscale1, aclDataType::ACL_FLOAT8_E8M0};
103- TensorWrapper mxscale2_wrapper = {mxscale2, aclDataType::ACL_FLOAT8_E8M0};107+ TensorWrapper mxscale2_wrapper = {mxscale2, aclDataType::ACL_FLOAT8_E8M0};
104 108 
105- EXEC_NPU_CMD(aclnnSwigluMxQuantWithDualAxis, x, group_index_opt, activate_left,109+ EXEC_NPU_CMD(
106- round_mode_ptr, scale_alg, y_acltype, dst_type_max,110+ aclnnSwigluMxQuantWithDualAxis,
107- y1_wrapper, mxscale1_wrapper, y2_wrapper, mxscale2_wrapper);111+ x,
112+ group_index_opt,
113+ activate_left,
114+ round_mode_ptr,
115+ scale_alg,
116+ y_acltype,
117+ dst_type_max,
118+ y1_wrapper,
119+ mxscale1_wrapper,
120+ y2_wrapper,
121+ mxscale2_wrapper);
108 122 
109- return std::make_tuple(y1, mxscale1, y2, mxscale2);123+ return std::make_tuple(y1, mxscale1, y2, mxscale2);
110}124}
111 125 
112} // namespace op_api126} // namespace op_api
Mop_plugin/ops/opapi/SwigluquantKernelNpuOpApi.cpp+79-50
@@ -26,61 +26,90 @@ constexpr int64_t MAX_LAST_DIM = 8192;
26// the last dimension of x should be divisible by 426// the last dimension of x should be divisible by 4
27constexpr int64_t DIV_LAST_DIM = 4;27constexpr int64_t DIV_LAST_DIM = 4;
28constexpr int64_t INT4_IN_INT8_NUM = 2;28constexpr int64_t INT4_IN_INT8_NUM = 2;
29-} // namespace29+} // namespace
30 30 
31-std::tuple<at::Tensor, at::Tensor> npu_swiglu_quant(const at::Tensor& x, const c10::optional<at::Tensor>& smooth_scales,31+std::tuple<at::Tensor, at::Tensor> npu_swiglu_quant(
32- const c10::optional<at::Tensor>& offsets, const c10::optional<at::Tensor>& group_index, bool activate_left,32+ const at::Tensor& x,
33- int64_t quant_mode, int64_t group_list_type, c10::optional<at::ScalarType> dst_type)33+ const c10::optional<at::Tensor>& smooth_scales,
34-{34+ const c10::optional<at::Tensor>& offsets,
35- TORCH_CHECK(quant_mode == 0 || quant_mode == 1, "quant_mode only support 0(static) or 1(dynamic), but got ",35+ const c10::optional<at::Tensor>& group_index,
36- quant_mode, OPS_ERROR(ErrCode::PARAM));36+ bool activate_left,
37- TORCH_CHECK(group_list_type == 0 || group_list_type == 1,37+ int64_t quant_mode,
38- "group_list_type only support 0(cumsum) or 1(count), but got ",38+ int64_t group_list_type,
39- group_list_type, OPS_ERROR(ErrCode::PARAM));39+ c10::optional<at::ScalarType> dst_type) {
40- TORCH_CHECK(!dst_type.has_value() || dst_type == at::ScalarType::Char || dst_type == at::ScalarType::QUInt4x2,40+ TORCH_CHECK(
41- "dtype must be torch.int8 for int8 or torch.quint4x2 for int4" + OPS_ERROR(ErrCode::TYPE));41+ quant_mode == 0 || quant_mode == 1,
42- if (!dst_type.has_value()) {42+ "quant_mode only support 0(static) or 1(dynamic), but got ",
43- // dst_type default is torch.int843+ quant_mode,
44- dst_type = at::ScalarType::Char;44+ OPS_ERROR(ErrCode::PARAM));
45- }45+ TORCH_CHECK(
46+ group_list_type == 0 || group_list_type == 1,
47+ "group_list_type only support 0(cumsum) or 1(count), but got ",
48+ group_list_type,
49+ OPS_ERROR(ErrCode::PARAM));
50+ TORCH_CHECK(
51+ !dst_type.has_value() || dst_type == at::ScalarType::Char || dst_type == at::ScalarType::QUInt4x2,
52+ "dtype must be torch.int8 for int8 or torch.quint4x2 for int4" + OPS_ERROR(ErrCode::TYPE));
53+ if (!dst_type.has_value()) {
54+ // dst_type default is torch.int8
55+ dst_type = at::ScalarType::Char;
56+ }
46 57 
47- const at::Tensor& smooth_scales_opt = c10::value_or_else(smooth_scales, [] { return at::Tensor(); });58+ const at::Tensor& smooth_scales_opt = c10::value_or_else(smooth_scales, [] { return at::Tensor(); });
48- const at::Tensor& offsets_opt = c10::value_or_else(offsets, [] { return at::Tensor(); });59+ const at::Tensor& offsets_opt = c10::value_or_else(offsets, [] { return at::Tensor(); });
49- TORCH_CHECK(quant_mode != 0 || smooth_scales_opt.sizes() == offsets_opt.sizes(),60+ TORCH_CHECK(
50- "smooth_scales and offsets should have the same shape when quant_mode is 0",61+ quant_mode != 0 || smooth_scales_opt.sizes() == offsets_opt.sizes(),
51- OPS_ERROR(ErrCode::PARAM));62+ "smooth_scales and offsets should have the same shape when quant_mode is 0",
52- 63+ OPS_ERROR(ErrCode::PARAM));
53- // check x last dim
54- int64_t x_last_dim = x.size(x.dim() - 1);
55- TORCH_CHECK(x_last_dim <= MAX_LAST_DIM, "x last dim size should not be larger than ", MAX_LAST_DIM, ", but got ",
56- x_last_dim, OPS_ERROR(ErrCode::PARAM));
57- TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));
58- // to concatenate two INT4 into one INT8, the last dimension of x should be divisible by 4
59- TORCH_CHECK(dst_type != at::ScalarType::QUInt4x2 || x_last_dim % DIV_LAST_DIM == 0,
60- "x shape last dim must be divded by 4 when dst_type is torch.quint4x2, but got ",
61- x_last_dim, OPS_ERROR(ErrCode::PARAM));
62 64 
63- at::SmallVector<int64_t, op_infer::SIZE> y_size;65+ // check x last dim
64- at::SmallVector<int64_t, op_infer::SIZE> scale_size;66+ int64_t x_last_dim = x.size(x.dim() - 1);
65- for (int i = 0; i < x.dim() - 1; i++) {67+ TORCH_CHECK(
66- y_size.push_back(x.size(i));68+ x_last_dim <= MAX_LAST_DIM,
67- scale_size.push_back(x.size(i));69+ "x last dim size should not be larger than ",
68- }70+ MAX_LAST_DIM,
69- auto last_dim = dst_type == at::ScalarType::Char ? x_last_dim / 2 : x_last_dim / 2 / INT4_IN_INT8_NUM;71+ ", but got ",
70- y_size.push_back(last_dim);72+ x_last_dim,
73+ OPS_ERROR(ErrCode::PARAM));
74+ TORCH_CHECK(x_last_dim % 2 == 0, "x last dim size should be even", OPS_ERROR(ErrCode::PARAM));
75+ // to concatenate two INT4 into one INT8, the last dimension of x should be divisible by 4
76+ TORCH_CHECK(
77+ dst_type != at::ScalarType::QUInt4x2 || x_last_dim % DIV_LAST_DIM == 0,
78+ "x shape last dim must be divded by 4 when dst_type is torch.quint4x2, but got ",
79+ x_last_dim,
80+ OPS_ERROR(ErrCode::PARAM));
71 81 
72- // The dtype of y is INT8(char), change y_size for different quantization types82+ at::SmallVector<int64_t, op_infer::SIZE> y_size;
73- at::Tensor y = npu_preparation::apply_tensor_without_format(y_size, c10::dtype(c10::ScalarType::Char));83+ at::SmallVector<int64_t, op_infer::SIZE> scale_size;
74- at::Tensor scale = npu_preparation::apply_tensor_without_format(scale_size, c10::dtype(c10::ScalarType::Float));84+ for (int i = 0; i < x.dim() - 1; i++) {
85+ y_size.push_back(x.size(i));
86+ scale_size.push_back(x.size(i));
87+ }
88+ auto last_dim = dst_type == at::ScalarType::Char ? x_last_dim / 2 : x_last_dim / 2 / INT4_IN_INT8_NUM;
89+ y_size.push_back(last_dim);
75 90 
76- std::string quant_mode_str = quant_mode == 0 ? "static" : "dynamic";91+ // The dtype of y is INT8(char), change y_size for different quantization types
77- char* quant_mode_ptr = const_cast<char*>(quant_mode_str.c_str());92+ at::Tensor y = npu_preparation::apply_tensor_without_format(y_size, c10::dtype(c10::ScalarType::Char));
93+ at::Tensor scale = npu_preparation::apply_tensor_without_format(scale_size, c10::dtype(c10::ScalarType::Float));
78 94 
79- const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });95+ std::string quant_mode_str = quant_mode == 0 ? "static" : "dynamic";
80- int output_type = dst_type == at::ScalarType::Char ? INT8 : INT4;96+ char* quant_mode_ptr = const_cast<char*>(quant_mode_str.c_str());
81- EXEC_NPU_CMD(aclnnSwiGluQuantV2, x, smooth_scales_opt, offsets_opt, group_index_opt, activate_left, quant_mode_ptr,
82- group_list_type, output_type, y, scale);
83 97 
84- return std::tie(y, scale);98+ const at::Tensor& group_index_opt = c10::value_or_else(group_index, [] { return at::Tensor(); });
99+ int output_type = dst_type == at::ScalarType::Char ? INT8 : INT4;
100+ EXEC_NPU_CMD(
101+ aclnnSwiGluQuantV2,
102+ x,
103+ smooth_scales_opt,
104+ offsets_opt,
105+ group_index_opt,
106+ activate_left,
107+ quant_mode_ptr,
108+ group_list_type,
109+ output_type,
110+ y,
111+ scale);
112+ 
113+ return std::tie(y, scale);
85}114}
86-} // namespace op_api115+} // namespace op_api
Mop_plugin/ops/opapi/TanhKernelNpuOpApi.cpp+21-20
@@ -20,28 +20,29 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor& tanh_out(const at::Tensor& self, at::Tensor& result)23+at::Tensor& tanh_out(const at::Tensor& self, at::Tensor& result) {
24-{24+ DO_COMPATIBILITY(aclnnTanh, acl_op::tanh_out(self, result));
25- DO_COMPATIBILITY(aclnnTanh, acl_op::tanh_out(self, result));25+ TORCH_CHECK(
26- TORCH_CHECK(!isIntegralType(result.scalar_type(), true), "result dtype can't be cast to the desired output type.\n",26+ !isIntegralType(result.scalar_type(), true),
27- OPS_ERROR(ErrCode::TYPE));27+ "result dtype can't be cast to the desired output type.\n",
28- npu_preparation::check_tensor({self}, result, result, self.sizes());28+ OPS_ERROR(ErrCode::TYPE));
29- at_npu::native::OpPreparation::check_memory({self}, {result});29+ npu_preparation::check_tensor({self}, result, result, self.sizes());
30- EXEC_NPU_CMD(aclnnTanh, self, result);30+ at_npu::native::OpPreparation::check_memory({self}, {result});
31- at::namedinference::propagate_names(result, self);31+ EXEC_NPU_CMD(aclnnTanh, self, result);
32- return result;32+ at::namedinference::propagate_names(result, self);
33+ return result;
33}34}
34 35 
35at::Tensor tanh(const at::Tensor& self) {36at::Tensor tanh(const at::Tensor& self) {
36- DO_COMPATIBILITY(aclnnTanh, acl_op::tanh(self));37+ DO_COMPATIBILITY(aclnnTanh, acl_op::tanh(self));
37- auto output_dtype = self.dtype();38+ auto output_dtype = self.dtype();
38- if (isIntegralType(self.scalar_type(), true)) {39+ if (isIntegralType(self.scalar_type(), true)) {
39- output_dtype = at::kFloat;40+ output_dtype = at::kFloat;
40- }41+ }
41- at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(output_dtype));42+ at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(output_dtype));
42- EXEC_NPU_CMD(aclnnTanh, self, result);43+ EXEC_NPU_CMD(aclnnTanh, self, result);
43- at::namedinference::propagate_names(result, self);44+ at::namedinference::propagate_names(result, self);
44- return result;45+ return result;
45}46}
46 47 
47at::Tensor& tanh_(at::Tensor& self) {48at::Tensor& tanh_(at::Tensor& self) {
@@ -49,4 +50,4 @@ at::Tensor& tanh_(at::Tensor& self) {
49 EXEC_NPU_CMD(aclnnInplaceTanh, self);50 EXEC_NPU_CMD(aclnnInplaceTanh, self);
50 return self;51 return self;
51}52}
52-}53+} // namespace op_api
Mop_plugin/ops/opapi/sparse/SparseCoalesceKernelOpApi.cpp+41-59
@@ -30,67 +30,49 @@ namespace sparse {
30 30 
31using namespace at::sparse;31using namespace at::sparse;
32 32 
33-SparseTensor _coalesce_sparse(const SparseTensor& self)33+SparseTensor _coalesce_sparse(const SparseTensor& self) {
34-{34+ int64_t nnz = self._nnz();
35- int64_t nnz = self._nnz();35+ TORCH_CHECK(!self.is_coalesced(), OPS_ERROR(ErrCode::VALUE));
36- TORCH_CHECK(!self.is_coalesced(), OPS_ERROR(ErrCode::VALUE));36+ if (nnz < 2) {
37- if (nnz < 2) {37+ SparseTensor dst = self.clone();
38- SparseTensor dst = self.clone();38+ dst._coalesced_(true);
39- dst._coalesced_(true);39+ return dst;
40- return dst;40+ }
41- }
42 41 
43- at::Tensor values = self._values();42+ at::Tensor values = self._values();
44- at::Tensor indices = self._indices();43+ at::Tensor indices = self._indices();
45- at::Tensor indices_1d = at::sparse::flatten_indices(indices, self.sizes(), true);44+ at::Tensor indices_1d = at::sparse::flatten_indices(indices, self.sizes(), true);
46- auto unique_indices_info = at::_unique2(indices_1d, true, true);45+ auto unique_indices_info = at::_unique2(indices_1d, true, true);
47- at::Tensor unique_len = std::get<0>(unique_indices_info).to(at::kInt);46+ at::Tensor unique_len = std::get<0>(unique_indices_info).to(at::kInt);
48- auto new_nnz = unique_len.sizes()[0];47+ auto new_nnz = unique_len.sizes()[0];
49- auto new_values_size = values.sizes().vec();48+ auto new_values_size = values.sizes().vec();
50- new_values_size[0] = new_nnz;49+ new_values_size[0] = new_nnz;
51- at::Tensor new_indices_t = at::zeros(50+ at::Tensor new_indices_t = at::zeros({new_nnz, indices.sizes()[0]}, indices.options().dtype(at::kInt));
52- {new_nnz, indices.sizes()[0]},51+ at::Tensor indices_t = at_npu::native::NpuUtils::format_contiguous(indices.transpose(0, 1)).to(at::kInt);
53- indices.options().dtype(at::kInt));52+ at::Tensor unique_indices = std::get<1>(unique_indices_info).to(at::kInt);
54- at::Tensor indices_t = at_npu::native::NpuUtils::format_contiguous(indices.transpose(0, 1)).to(at::kInt);53+ if (values.scalar_type() == at::kHalf || (values.scalar_type() == at::kBFloat16)) {
55- at::Tensor unique_indices = std::get<1>(unique_indices_info).to(at::kInt);54+ at::Tensor values_f = values.to(at::kFloat);
56- if (values.scalar_type() == at::kHalf || (values.scalar_type() == at::kBFloat16)) {55+ at::Tensor new_values_f = at::zeros(new_values_size, values.options()).to(at::kFloat);
57- at::Tensor values_f = values.to(at::kFloat);56+ EXEC_NPU_CMD(aclnnCoalesceSparse, unique_len, unique_indices, indices_t, values_f, new_indices_t, new_values_f);
58- at::Tensor new_values_f = at::zeros(new_values_size, values.options()).to(at::kFloat);57+ at::Tensor new_indices = new_indices_t.transpose(0, 1).to(at::kLong);
59- EXEC_NPU_CMD(58+ if (values.scalar_type() == at::kHalf) {
60- aclnnCoalesceSparse,59+ at::Tensor new_values = new_values_f.to(at::kHalf);
61- unique_len,60+ SparseTensor dst =
62- unique_indices,61+ ::at::native::_sparse_coo_tensor_unsafe(new_indices, new_values, self.sizes())._coalesced_(true);
63- indices_t,62+ return dst;
64- values_f,
65- new_indices_t,
66- new_values_f);
67- at::Tensor new_indices = new_indices_t.transpose(0, 1).to(at::kLong);
68- if (values.scalar_type() == at::kHalf) {
69- at::Tensor new_values = new_values_f.to(at::kHalf);
70- SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(new_indices,
71- new_values, self.sizes())._coalesced_(true);
72- return dst;
73- } else {
74- at::Tensor new_values = new_values_f.to(at::kBFloat16);
75- SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(new_indices,
76- new_values, self.sizes())._coalesced_(true);
77- return dst;
78- }
79 } else {63 } else {
80- at::Tensor new_values = at::zeros(new_values_size, values.options());64+ at::Tensor new_values = new_values_f.to(at::kBFloat16);
81- EXEC_NPU_CMD(65+ SparseTensor dst =
82- aclnnCoalesceSparse,66+ ::at::native::_sparse_coo_tensor_unsafe(new_indices, new_values, self.sizes())._coalesced_(true);
83- unique_len,67+ return dst;
84- unique_indices,
85- indices_t,
86- values,
87- new_indices_t,
88- new_values);
89- at::Tensor new_indices = new_indices_t.transpose(0, 1).to(at::kLong);
90- SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(new_indices,
91- new_values, self.sizes())._coalesced_(true);
92- return dst;
93 }68 }
69+ } else {
70+ at::Tensor new_values = at::zeros(new_values_size, values.options());
71+ EXEC_NPU_CMD(aclnnCoalesceSparse, unique_len, unique_indices, indices_t, values, new_indices_t, new_values);
72+ at::Tensor new_indices = new_indices_t.transpose(0, 1).to(at::kLong);
73+ SparseTensor dst = ::at::native::_sparse_coo_tensor_unsafe(new_indices, new_values, self.sizes())._coalesced_(true);
74+ return dst;
75+ }
94}76}
95 77 
96-} // namespace at::native78+} // namespace sparse
Mop_plugin/ops/opapi/sparse/SparseMaxKernelNpuOpApi.cpp+6-8
@@ -18,16 +18,14 @@
18 18 
19namespace sparse {19namespace sparse {
20 20 
21-at::Tensor max_sparse(const at::Tensor &self)21+at::Tensor max_sparse(const at::Tensor& self) {
22-{22+ auto dense_tensor = self.to_dense();
23- auto dense_tensor = self.to_dense();23+ return at::max(dense_tensor);
24- return at::max(dense_tensor);
25}24}
26 25 
27-at::Tensor &max_out_sparse(const at::Tensor &self, const at::Tensor &other, at::Tensor &result)26+at::Tensor& max_out_sparse(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {
28-{27+ auto dense_tensor = self.to_dense();
29- auto dense_tensor = self.to_dense();28+ return at::max_out(dense_tensor, other, result);
30- return at::max_out(dense_tensor, other, result);
31}29}
32 30 
33} // namespace sparse31} // namespace sparse
Mop_plugin/ops/opapi/sparse/SparseTensorUtils.h+6-11
@@ -22,21 +22,16 @@
22namespace sparse {22namespace sparse {
23using namespace at::sparse;23using namespace at::sparse;
24 24 
25-inline at::SparseTensorImpl* get_sparse_impl(const SparseTensor& self)25+inline at::SparseTensorImpl* get_sparse_impl(const SparseTensor& self) {
26-{26+ TORCH_CHECK(self.is_sparse(), "_internal_get_SparseTensorImpl: not a sparse tensor", OPS_ERROR(ErrCode::VALUE));
27- TORCH_CHECK(self.is_sparse(), "_internal_get_SparseTensorImpl: not a sparse tensor", OPS_ERROR(ErrCode::VALUE));27+ return static_cast<at::SparseTensorImpl*>(self.unsafeGetTensorImpl());
28- return static_cast<at::SparseTensorImpl*>(self.unsafeGetTensorImpl());
29}28}
30 29 
31-inline void alias_into_sparse(30+inline void alias_into_sparse(const SparseTensor& self, const at::Tensor& indices, const at::Tensor& values) {
32- const SparseTensor& self,31+ get_sparse_impl(self)->set_indices_and_values_unsafe(indices, values);
33- const at::Tensor& indices,
34- const at::Tensor& values)
35-{
36- get_sparse_impl(self)->set_indices_and_values_unsafe(indices, values);
37}32}
38 33 
39SparseTensor& mul_out_sparse_scalar(SparseTensor& r, const SparseTensor& t, const at::Scalar& value);34SparseTensor& mul_out_sparse_scalar(SparseTensor& r, const SparseTensor& t, const at::Scalar& value);
40at::Tensor& mul_out_sparse(const at::Tensor& self, const at::Tensor& other, at::Tensor& out);35at::Tensor& mul_out_sparse(const at::Tensor& self, const at::Tensor& other, at::Tensor& out);
41 36 
42-}37+} // namespace sparse
Mop_plugin/third_party/atb/inc/svector.h+469-494
@@ -43,506 +43,481 @@ struct MaxSizeExceeded : public std::exception {};
43//!43//!
44//! 封装动态数组的顺序容器44//! 封装动态数组的顺序容器
45//!45//!
46-template <class T> class SVector {46+template <class T>
47-public:47+class SVector {
48- //! \brief 默认构造函数48+ public:
49- //!49+ //! \brief 默认构造函数
50- //! \note 容量为DEFAULT_SVECTOR_SIZE50+ //!
51- //!51+ //! \note 容量为DEFAULT_SVECTOR_SIZE
52- constexpr SVector() : size_(0)52+ //!
53- {53+ constexpr SVector() : size_(0) {
54- for (std::size_t i = 0; i < DEFAULT_SVECTOR_SIZE; ++i) {54+ for (std::size_t i = 0; i < DEFAULT_SVECTOR_SIZE; ++i) {
55- storage_[i] = {};55+ storage_[i] = {};
56- }
57 }56 }
58- //! \brief 初始化列表构造函数57+ }
59- //!58+ //! \brief 初始化列表构造函数
60- //! \param list59+ //!
61- //!60+ //! \param list
62- //! \note list长度需小于DEFAULT_SVECTOR_SIZE，否则会抛出异常61+ //!
63- //!62+ //! \note list长度需小于DEFAULT_SVECTOR_SIZE，否则会抛出异常
64- SVector(std::initializer_list<T> list)63+ //!
65- {64+ SVector(std::initializer_list<T> list) {
66- if (CHECK_BOUND && list.size() > DEFAULT_SVECTOR_SIZE) {65+ if (CHECK_BOUND && list.size() > DEFAULT_SVECTOR_SIZE) {
67- throw MaxSizeExceeded();66+ throw MaxSizeExceeded();
68- }
69- size_ = list.size();
70- size_t i = 0;
71- for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
72- storage_[i++] = *it;
73- }
74 }67 }
75- //! \brief 带参数的构造函数68+ size_ = list.size();
76- //!69+ size_t i = 0;
77- //! \param size70+ for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
78- //! \param value71+ storage_[i++] = *it;
79- //!
80- //! \note size大小需小于DEFAULT_SVECTOR_SIZE，否则会抛出异常
81- //!
82- explicit SVector(std::size_t size, const T &value = 0) : size_(0)
83- {
84- if (CHECK_BOUND && size > DEFAULT_SVECTOR_SIZE) {
85- throw MaxSizeExceeded();
86- }
87- size_ = size;
88- for (std::size_t i = 0; i < size_; ++i) {
89- storage_[i] = value;
90- }
91 }72 }
92- //! \brief 拷贝构造函数，创建一个新的SVector对象并将另一个SVector对象的值复制到新对象73+ }
93- //!74+ //! \brief 带参数的构造函数
94- //! \param other75+ //!
95- //!76+ //! \param size
96- SVector(const SVector<T> &other)77+ //! \param value
97- {78+ //!
98- if (other.heap_) {79+ //! \note size大小需小于DEFAULT_SVECTOR_SIZE，否则会抛出异常
99- heap_ = reinterpret_cast<T *>(malloc(other.size_ * sizeof(T)));80+ //!
100- if (!heap_) {81+ explicit SVector(std::size_t size, const T& value = 0) : size_(0) {
101- throw std::bad_alloc();82+ if (CHECK_BOUND && size > DEFAULT_SVECTOR_SIZE) {
102- }83+ throw MaxSizeExceeded();
103- size_ = other.size_;84+ }
104- for (std::size_t i = 0; i < other.size_; ++i) {85+ size_ = size;
105- heap_[i] = other.heap_[i];86+ for (std::size_t i = 0; i < size_; ++i) {
106- }87+ storage_[i] = value;
107- } else {88+ }
108- size_ = other.size_;89+ }
109- for (std::size_t i = 0; i < other.size_; ++i) {90+ //! \brief 拷贝构造函数，创建一个新的SVector对象并将另一个SVector对象的值复制到新对象
110- storage_[i] = other.storage_[i];91+ //!
111- }92+ //! \param other
112- }93+ //!
94+ SVector(const SVector<T>& other) {
95+ if (other.heap_) {
96+ heap_ = reinterpret_cast<T*>(malloc(other.size_ * sizeof(T)));
97+ if (!heap_) {
98+ throw std::bad_alloc();
99+ }
100+ size_ = other.size_;
101+ for (std::size_t i = 0; i < other.size_; ++i) {
102+ heap_[i] = other.heap_[i];
103+ }
104+ } else {
105+ size_ = other.size_;
106+ for (std::size_t i = 0; i < other.size_; ++i) {
107+ storage_[i] = other.storage_[i];
108+ }
109+ }
110+ }
111+ 
112+ ~SVector() {
113+ if (heap_) {
114+ free(heap_);
115+ }
116+ }
117+ 
118+ //! \brief 插入元素到指定容器
119+ //!
120+ //! \param val
121+ //!
122+ //! \note 待添加SVector内元素必须小于SVector容量，否则会抛出异常
123+ //!
124+ void push_back(const T& val) noexcept((!CHECK_BOUND) && std::is_nothrow_assignable<T, const T&>::value) {
125+ if (heap_) {
126+ if (CHECK_BOUND && size_ == capacity_) {
127+ throw MaxSizeExceeded();
128+ }
129+ heap_[size_++] = val;
130+ return;
131+ }
132+ if (CHECK_BOUND && size_ == DEFAULT_SVECTOR_SIZE) {
133+ throw MaxSizeExceeded();
134+ }
135+ storage_[size_++] = val;
136+ }
137+ 
138+ //! \brief 获取容器起始元素地址
139+ //!
140+ //! \return 指针
141+ //!
142+ T* begin() noexcept {
143+ if (heap_) {
144+ return &heap_[0];
145+ }
146+ return &storage_[0];
147+ }
148+ 
149+ //! \brief 获取容器起始元素地址
150+ //!
151+ //! \return 常量指针
152+ //!
153+ const T* begin() const noexcept {
154+ if (heap_) {
155+ return &heap_[0];
156+ }
157+ return &storage_[0];
158+ }
159+ 
160+ //! \brief 获取容器尾元素地址
161+ //!
162+ //! \return 指针
163+ //!
164+ T* end() noexcept {
165+ if (heap_) {
166+ return (&heap_[0]) + size_;
167+ }
168+ return (&storage_[0]) + size_;
169+ }
170+ 
171+ //! \brief 获取容器尾地址
172+ //!
173+ //! \return 常量指针
174+ //!
175+ const T* end() const noexcept {
176+ if (heap_) {
177+ return (&heap_[0]) + size_;
178+ }
179+ return (&storage_[0]) + size_;
180+ }
181+ 
182+ //! \brief 访问指定位置的元素
183+ //!
184+ //! \param i
185+ //!
186+ //! \return 引用
187+ //!
188+ T& operator[](std::size_t i) {
189+ if (heap_) {
190+ if (size_ == 0 || i >= size_) {
191+ throw std::out_of_range("out of range");
192+ }
193+ return heap_[i];
194+ }
195+ if (size_ == 0 || i >= size_) {
196+ throw std::out_of_range("out of range");
197+ }
198+ return storage_[i];
199+ }
200+ 
201+ //! \brief 访问指定位置的元素
202+ //!
203+ //! \param i
204+ //!
205+ //! \return 常量引用
206+ //!
207+ const T& operator[](std::size_t i) const {
208+ if (heap_) {
209+ if (size_ == 0 || i >= size_) {
210+ throw std::out_of_range("out of range");
211+ }
212+ return heap_[i];
213+ }
214+ if (size_ == 0 || i >= size_) {
215+ throw std::out_of_range("out of range");
216+ }
217+ return storage_[i];
218+ }
219+ 
220+ //! \brief 访问指定位置的元素
221+ //!
222+ //! \param i
223+ //!
224+ //! \return 引用
225+ //!
226+ T& at(std::size_t i) {
227+ if (heap_) {
228+ if (size_ == 0 || i >= size_) {
229+ throw std::out_of_range("out of range");
230+ }
231+ return heap_[i];
232+ } else {
233+ if (size_ == 0 || i >= size_ || i > DEFAULT_SVECTOR_SIZE) {
234+ throw std::out_of_range("out of range");
235+ }
236+ return storage_[i];
237+ }
238+ }
239+ 
240+ //! \brief 访问指定位置的元素
241+ //!
242+ //! \param i
243+ //!
244+ //! \return 引用
245+ //!
246+ const T& at(std::size_t i) const {
247+ if (heap_) {
248+ if (size_ == 0 || i >= size_) {
249+ throw std::out_of_range("heap out of range");
250+ }
251+ return heap_[i];
252+ }
253+ if (size_ == 0 || i >= size_ || i > DEFAULT_SVECTOR_SIZE) {
254+ throw std::out_of_range("stack out of range");
255+ }
256+ return storage_[i];
257+ }
258+ 
259+ //! \brief 获取容器的大小
260+ //!
261+ //! \return size
262+ //!
263+ std::size_t size() const noexcept {
264+ return size_;
265+ }
266+ 
267+ //! \brief 向容器内指定位置插入元素
268+ //!
269+ //! \param pos
270+ //! \param value
271+ //!
272+ //! \note pos必须小于SVector容量，否则会抛出异常
273+ //!
274+ void insert(const std::size_t pos, const T& value) noexcept(
275+ (!CHECK_BOUND) && std::is_nothrow_assignable<T, const T&>::value) {
276+ if (heap_) {
277+ if (pos > size_ || pos == capacity_) {
278+ throw MaxSizeExceeded();
279+ }
280+ for (auto it = size_; it != pos; it--) {
281+ heap_[it] = heap_[it - 1];
282+ }
283+ heap_[pos] = value;
284+ size_ += 1;
285+ return;
286+ }
287+ if (CHECK_BOUND && size_ == DEFAULT_SVECTOR_SIZE) {
288+ throw MaxSizeExceeded();
289+ }
290+ if (pos > size_) {
291+ throw MaxSizeExceeded();
113 }292 }
114 293 
115- ~SVector()294+ for (auto it = size_; it != pos; it--) {
116- {295+ storage_[it] = storage_[it - 1];
117- if (heap_) {296+ }
118- free(heap_);297+ storage_[pos] = value;
119- }298+ size_ += 1;
299+ return;
300+ }
301+ 
302+ //! \brief 判断容器是否为空
303+ //!
304+ //! \return bool值
305+ //!
306+ bool empty() const noexcept {
307+ return size_ == 0;
308+ }
309+ 
310+ //! \brief 清空容器
311+ void clear() noexcept {
312+ size_ = 0;
313+ }
314+ 
315+ //! \brief 获取容器起始元素地址
316+ //!
317+ //! \return 指针
318+ //!
319+ T* data() noexcept {
320+ if (heap_) {
321+ return &heap_[0];
322+ }
323+ return &storage_[0];
324+ }
325+ 
326+ //! \brief 获取容器起始元素地址
327+ //!
328+ //! \return 常量指针
329+ //!
330+ const T* data() const noexcept {
331+ if (heap_) {
332+ return &heap_[0];
333+ }
334+ return &storage_[0];
335+ }
336+ 
337+ //! \brief 改变SVector容器大小，不能改变SVector容量
338+ //!
339+ //! \param size
340+ //!
341+ //! \note 传入size参数不能超过SVector容量，反之，则会抛出异常。
342+ //!
343+ void resize(std::size_t size) {
344+ if (heap_ && size > capacity_) {
345+ throw MaxSizeExceeded();
346+ }
347+ size_ = size;
348+ }
349+ 
350+ //! \brief 改变SVector容量大小，清空内部数据，并将SVector容量大小定义为size大小
351+ //!
352+ //! \param size
353+ //!
354+ //! \note
355+ //! 用于预分配内存空间，SVector默认容量为DEFAULT_SVECTOR_SIZE，传入的size需大于DEFAULT_SVECTOR_SIZ且小于MAX_SVECTOR_SIZE，反之，则会抛出异常。
356+ //!
357+ void reserve(std::size_t size) {
358+ if (size > MAX_SVECTOR_SIZE) {
359+ throw MaxSizeExceeded();
120 }360 }
121 361 
122- //! \brief 插入元素到指定容器362+ if (size > DEFAULT_SVECTOR_SIZE) {
123- //!363+ if (heap_) {
124- //! \param val364+ free(heap_);
125- //!365+ }
126- //! \note 待添加SVector内元素必须小于SVector容量，否则会抛出异常366+ heap_ = reinterpret_cast<T*>(malloc(size * sizeof(T)));
127- //!367+ if (!heap_) {
128- void push_back(const T &val) noexcept((!CHECK_BOUND) && std::is_nothrow_assignable<T, const T &>::value)368+ throw std::bad_alloc();
129- {369+ }
130- if (heap_) {370+ for (std::size_t i = 0; i < size; ++i) {
131- if (CHECK_BOUND && size_ == capacity_) {371+ heap_[i] = {};
132- throw MaxSizeExceeded();372+ }
133- }373+ capacity_ = size;
134- heap_[size_++] = val;
135- return;
136- }
137- if (CHECK_BOUND && size_ == DEFAULT_SVECTOR_SIZE) {
138- throw MaxSizeExceeded();
139- }
140- storage_[size_++] = val;
141 }374 }
375+ }
142 376 
143- //! \brief 获取容器起始元素地址377+ //! \brief 判断两个容器中的元素是否全部相同
144- //!378+ //!
145- //! \return 指针379+ //! \param other
146- //!380+ //!
147- T *begin() noexcept381+ //! \return bool值
148- {382+ //!
149- if (heap_) {383+ bool operator==(const SVector<T>& other) const {
150- return &heap_[0];384+ if (heap_) {
385+ if (size_ != other.size_ || !other.heap_) {
386+ return false;
387+ }
388+ for (size_t i = 0; i < size_; ++i) {
389+ if (heap_[i] != other.heap_[i]) {
390+ return false;
151 }391 }
152- return &storage_[0];392+ }
393+ } else {
394+ if (size_ != other.size_) {
395+ return false;
396+ }
397+ for (size_t i = 0; i < size_; ++i) {
398+ if (storage_[i] != other.storage_[i]) {
399+ return false;
400+ }
401+ }
153 }402 }
403+ return true;
404+ }
154 405 
155- //! \brief 获取容器起始元素地址406+ //! \brief 判断两个容器中的元素是否存在不同
156- //!407+ //!
157- //! \return 常量指针408+ //! \param other
158- //!409+ //!
159- const T *begin() const noexcept410+ //! \return bool值
160- {411+ //!
161- if (heap_) {412+ bool operator!=(const SVector<T>& other) const {
162- return &heap_[0];413+ if (heap_) {
163- }414+ if (size_ != other.size_ || !other.heap_) {
164- return &storage_[0];
165- }
166- 
167- //! \brief 获取容器尾元素地址
168- //!
169- //! \return 指针
170- //!
171- T *end() noexcept
172- {
173- if (heap_) {
174- return (&heap_[0]) + size_;
175- }
176- return (&storage_[0]) + size_;
177- }
178- 
179- //! \brief 获取容器尾地址
180- //!
181- //! \return 常量指针
182- //!
183- const T *end() const noexcept
184- {
185- if (heap_) {
186- return (&heap_[0]) + size_;
187- }
188- return (&storage_[0]) + size_;
189- }
190- 
191- //! \brief 访问指定位置的元素
192- //!
193- //! \param i
194- //!
195- //! \return 引用
196- //!
197- T &operator[](std::size_t i)
198- {
199- if (heap_) {
200- if (size_ == 0 || i >= size_) {
201- throw std::out_of_range("out of range");
202- }
203- return heap_[i];
204- }
205- if (size_ == 0 || i >= size_) {
206- throw std::out_of_range("out of range");
207- }
208- return storage_[i];
209- }
210- 
211- //! \brief 访问指定位置的元素
212- //!
213- //! \param i
214- //!
215- //! \return 常量引用
216- //!
217- const T &operator[](std::size_t i) const
218- {
219- if (heap_) {
220- if (size_ == 0 || i >= size_) {
221- throw std::out_of_range("out of range");
222- }
223- return heap_[i];
224- }
225- if (size_ == 0 || i >= size_) {
226- throw std::out_of_range("out of range");
227- }
228- return storage_[i];
229- }
230- 
231- //! \brief 访问指定位置的元素
232- //!
233- //! \param i
234- //!
235- //! \return 引用
236- //!
237- T &at(std::size_t i)
238- {
239- if (heap_) {
240- if (size_ == 0 || i >= size_) {
241- throw std::out_of_range("out of range");
242- }
243- return heap_[i];
244- } else {
245- if (size_ == 0 || i >= size_ || i > DEFAULT_SVECTOR_SIZE) {
246- throw std::out_of_range("out of range");
247- }
248- return storage_[i];
249- }
250- }
251- 
252- //! \brief 访问指定位置的元素
253- //!
254- //! \param i
255- //!
256- //! \return 引用
257- //!
258- const T &at(std::size_t i) const
259- {
260- if (heap_) {
261- if (size_ == 0 || i >= size_) {
262- throw std::out_of_range("heap out of range");
263- }
264- return heap_[i];
265- }
266- if (size_ == 0 || i >= size_ || i > DEFAULT_SVECTOR_SIZE) {
267- throw std::out_of_range("stack out of range");
268- }
269- return storage_[i];
270- }
271- 
272- //! \brief 获取容器的大小
273- //!
274- //! \return size
275- //!
276- std::size_t size() const noexcept
277- {
278- return size_;
279- }
280- 
281- //! \brief 向容器内指定位置插入元素
282- //!
283- //! \param pos
284- //! \param value
285- //!
286- //! \note pos必须小于SVector容量，否则会抛出异常
287- //!
288- void insert(const std::size_t pos,
289- const T &value) noexcept((!CHECK_BOUND) && std::is_nothrow_assignable<T, const T &>::value)
290- {
291- if (heap_) {
292- if (pos > size_ || pos == capacity_) {
293- throw MaxSizeExceeded();
294- }
295- for (auto it = size_; it != pos; it--) {
296- heap_[it] = heap_[it - 1];
297- }
298- heap_[pos] = value;
299- size_ += 1;
300- return;
301- }
302- if (CHECK_BOUND && size_ == DEFAULT_SVECTOR_SIZE) {
303- throw MaxSizeExceeded();
304- }
305- if (pos > size_) {
306- throw MaxSizeExceeded();
307- }
308- 
309- for (auto it = size_; it != pos; it--) {
310- storage_[it] = storage_[it - 1];
311- }
312- storage_[pos] = value;
313- size_ += 1;
314- return;
315- }
316- 
317- //! \brief 判断容器是否为空
318- //!
319- //! \return bool值
320- //!
321- bool empty() const noexcept
322- {
323- return size_ == 0;
324- }
325- 
326- //! \brief 清空容器
327- void clear() noexcept
328- {
329- size_ = 0;
330- }
331- 
332- //! \brief 获取容器起始元素地址
333- //!
334- //! \return 指针
335- //!
336- T *data() noexcept
337- {
338- if (heap_) {
339- return &heap_[0];
340- }
341- return &storage_[0];
342- }
343- 
344- //! \brief 获取容器起始元素地址
345- //!
346- //! \return 常量指针
347- //!
348- const T *data() const noexcept
349- {
350- if (heap_) {
351- return &heap_[0];
352- }
353- return &storage_[0];
354- }
355- 
356- //! \brief 改变SVector容器大小，不能改变SVector容量
357- //!
358- //! \param size
359- //!
360- //! \note 传入size参数不能超过SVector容量，反之，则会抛出异常。
361- //!
362- void resize(std::size_t size)
363- {
364- if (heap_ && size > capacity_) {
365- throw MaxSizeExceeded();
366- }
367- size_ = size;
368- }
369- 
370- //! \brief 改变SVector容量大小，清空内部数据，并将SVector容量大小定义为size大小
371- //!
372- //! \param size
373- //!
374- //! \note 用于预分配内存空间，SVector默认容量为DEFAULT_SVECTOR_SIZE，传入的size需大于DEFAULT_SVECTOR_SIZ且小于MAX_SVECTOR_SIZE，反之，则会抛出异常。
375- //!
376- void reserve(std::size_t size)
377- {
378- if (size > MAX_SVECTOR_SIZE) {
379- throw MaxSizeExceeded();
380- }
381- 
382- if (size > DEFAULT_SVECTOR_SIZE) {
383- if (heap_) {
384- free(heap_);
385- }
386- heap_ = reinterpret_cast<T *>(malloc(size * sizeof(T)));
387- if (!heap_) {
388- throw std::bad_alloc();
389- }
390- for (std::size_t i = 0; i < size; ++i) {
391- heap_[i] = {};
392- }
393- capacity_ = size;
394- }
395- }
396- 
397- //! \brief 判断两个容器中的元素是否全部相同
398- //!
399- //! \param other
400- //!
401- //! \return bool值
402- //!
403- bool operator==(const SVector<T> &other) const
404- {
405- if (heap_) {
406- if (size_ != other.size_ || !other.heap_) {
407- return false;
408- }
409- for (size_t i = 0; i < size_; ++i) {
410- if (heap_[i] != other.heap_[i]) {
411- return false;
412- }
413- }
414- } else {
415- if (size_ != other.size_) {
416- return false;
417- }
418- for (size_t i = 0; i < size_; ++i) {
419- if (storage_[i] != other.storage_[i]) {
420- return false;
421- }
422- }
423- }
424 return true;415 return true;
425- }416+ }
426- 417+ for (size_t i = 0; i < size_; ++i) {
427- //! \brief 判断两个容器中的元素是否存在不同418+ if (heap_[i] != other.heap_[i]) {
428- //!419+ return true;
429- //! \param other
430- //!
431- //! \return bool值
432- //!
433- bool operator!=(const SVector<T> &other) const
434- {
435- if (heap_) {
436- if (size_ != other.size_ || !other.heap_) {
437- return true;
438- }
439- for (size_t i = 0; i < size_; ++i) {
440- if (heap_[i] != other.heap_[i]) {
441- return true;
442- }
443- }
444- } else {
445- if (size_ != other.size_) {
446- return true;
447- }
448- for (size_t i = 0; i < size_; ++i) {
449- if (storage_[i] != other.storage_[i]) {
450- return true;
451- }
452- }
453 }420 }
454- return false;421+ }
455- }422+ } else {
456- 423+ if (size_ != other.size_) {
457- //! \brief 判断一个容器中的元素是否比另一个容器小424+ return true;
458- //!425+ }
459- //! \param other426+ for (size_t i = 0; i < size_; ++i) {
460- //!427+ if (storage_[i] != other.storage_[i]) {
461- //! \return bool值428+ return true;
462- //!
463- bool operator<(const SVector<T> &other) const
464- {
465- if (heap_) {
466- if (size_ != other.size_ || !other.heap_) {
467- return size_ < other.size_;
468- }
469- for (size_t i = 0; i < size_; ++i) {
470- if (heap_[i] != other.heap_[i]) {
471- return heap_[i] < other.heap_[i];
472- }
473- }
474- } else {
475- if (size_ != other.size_) {
476- return size_ < other.size_;
477- }
478- for (size_t i = 0; i < size_; ++i) {
479- if (storage_[i] != other.storage_[i]) {
480- return storage_[i] < other.storage_[i];
481- }
482- }
483 }429 }
484- return false;430+ }
485 }431 }
432+ return false;
433+ }
486 434 
487- //! \brief 重载运算符函数，将初始化列表中的元素赋值给一个SVector对象435+ //! \brief 判断一个容器中的元素是否比另一个容器小
488- //!436+ //!
489- //! \param list437+ //! \param other
490- //!438+ //!
491- //! \return 容器引用439+ //! \return bool值
492- //!440+ //!
493- SVector &operator=(std::initializer_list<T> list)441+ bool operator<(const SVector<T>& other) const {
494- {442+ if (heap_) {
495- if (heap_) {443+ if (size_ != other.size_ || !other.heap_) {
496- if (CHECK_BOUND && list.size() > MAX_SVECTOR_SIZE) {444+ return size_ < other.size_;
497- throw MaxSizeExceeded();445+ }
498- }446+ for (size_t i = 0; i < size_; ++i) {
499- size_ = list.size();447+ if (heap_[i] != other.heap_[i]) {
500- size_t i = 0;448+ return heap_[i] < other.heap_[i];
501- for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
502- heap_[i++] = *it;
503- }
504- return *this;
505- } else {
506- if (CHECK_BOUND && list.size() > DEFAULT_SVECTOR_SIZE) {
507- throw MaxSizeExceeded();
508- }
509- size_ = list.size();
510- size_t i = 0;
511- for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
512- storage_[i++] = *it;
513- }
514- return *this;
515 }449 }
516- }450+ }
517- 451+ } else {
518- //! \brief 用一个容器给另一个容器赋值452+ if (size_ != other.size_) {
519- //!453+ return size_ < other.size_;
520- //! \param other454+ }
521- //!455+ for (size_t i = 0; i < size_; ++i) {
522- //! \return 容器引用456+ if (storage_[i] != other.storage_[i]) {
523- //!457+ return storage_[i] < other.storage_[i];
524- SVector &operator=(const SVector &other)
525- {
526- if (heap_) {
527- size_ = other.size_;
528- for (std::size_t i = 0; i < other.size_; ++i) {
529- heap_[i] = other.heap_[i];
530- }
531- return *this;
532- } else {
533- size_ = other.size_;
534- for (std::size_t i = 0; i < other.size_; ++i) {
535- storage_[i] = other.storage_[i];
536- }
537- return *this;
538 }458 }
459+ }
539 }460 }
461+ return false;
462+ }
540 463 
541-private:464+ //! \brief 重载运算符函数，将初始化列表中的元素赋值给一个SVector对象
542- std::size_t capacity_ = 0;465+ //!
543- std::size_t size_ = 0;466+ //! \param list
544- T storage_[DEFAULT_SVECTOR_SIZE + 1];467+ //!
545- T *heap_ = nullptr;468+ //! \return 容器引用
469+ //!
470+ SVector& operator=(std::initializer_list<T> list) {
471+ if (heap_) {
472+ if (CHECK_BOUND && list.size() > MAX_SVECTOR_SIZE) {
473+ throw MaxSizeExceeded();
474+ }
475+ size_ = list.size();
476+ size_t i = 0;
477+ for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
478+ heap_[i++] = *it;
479+ }
480+ return *this;
481+ } else {
482+ if (CHECK_BOUND && list.size() > DEFAULT_SVECTOR_SIZE) {
483+ throw MaxSizeExceeded();
484+ }
485+ size_ = list.size();
486+ size_t i = 0;
487+ for (auto it = list.begin(); it != list.end() && i < size_; ++it) {
488+ storage_[i++] = *it;
489+ }
490+ return *this;
491+ }
492+ }
493+ 
494+ //! \brief 用一个容器给另一个容器赋值
495+ //!
496+ //! \param other
497+ //!
498+ //! \return 容器引用
499+ //!
500+ SVector& operator=(const SVector& other) {
501+ if (heap_) {
502+ size_ = other.size_;
503+ for (std::size_t i = 0; i < other.size_; ++i) {
504+ heap_[i] = other.heap_[i];
505+ }
506+ return *this;
507+ } else {
508+ size_ = other.size_;
509+ for (std::size_t i = 0; i < other.size_; ++i) {
510+ storage_[i] = other.storage_[i];
511+ }
512+ return *this;
513+ }
514+ }
515+ 
516+ private:
517+ std::size_t capacity_ = 0;
518+ std::size_t size_ = 0;
519+ T storage_[DEFAULT_SVECTOR_SIZE + 1];
520+ T* heap_ = nullptr;
546};521};
547 522 
548//! \brief 输出容器中的元素523//! \brief 输出容器中的元素
@@ -552,21 +527,21 @@ private:
552//!527//!
553//! \return 输出流528//! \return 输出流
554//!529//!
555-template <class T> std::ostream &operator<<(std::ostream &os, const SVector<T> &svector)530+template <class T>
556-{531+std::ostream& operator<<(std::ostream& os, const SVector<T>& svector) {
557- if (svector.size() == 0) {532+ if (svector.size() == 0) {
558- return os;
559- }
560- 
561- std::string str = ",";
562- for (size_t i = 0; i < svector.size(); ++i) {
563- os << svector.at(i);
564- if (i != svector.size() - 1) {
565- os << str;
566- }
567- }
568- 
569 return os;533 return os;
534+ }
535+ 
536+ std::string str = ",";
537+ for (size_t i = 0; i < svector.size(); ++i) {
538+ os << svector.at(i);
539+ if (i != svector.size() - 1) {
540+ os << str;
541+ }
542+ }
543+ 
544+ return os;
570}545}
571} // namespace atb546} // namespace atb
572-#endif547+#endif
Mop_plugin/utils/SearchsortedSideUtil.h+5-5
@@ -17,11 +17,11 @@ namespace op_plugin {
17 17 
18/// When `side` is set it overrides `right` (see torch.searchsorted(side=...)).18/// When `side` is set it overrides `right` (see torch.searchsorted(side=...)).
19/// Invalid `side` and side/right conflicts are handled in searchsorted_pre_check_npu; call this after validate.19/// Invalid `side` and side/right conflicts are handled in searchsorted_pre_check_npu; call this after validate.
20-inline bool resolve_searchsorted_effective_right(bool right, const c10::optional<c10::string_view> &side_opt) {20+inline bool resolve_searchsorted_effective_right(bool right, const c10::optional<c10::string_view>& side_opt) {
21- if (!side_opt.has_value()) {21+ if (!side_opt.has_value()) {
22- return right;22+ return right;
23- }23+ }
24- return *side_opt == "right";24+ return *side_opt == "right";
25}25}
26 26 
27} // namespace op_plugin27} // namespace op_plugin
Mop_plugin/utils/SearchsortedValidateUtil.h+151-97
@@ -25,127 +25,181 @@
25 25 
26namespace op_plugin {26namespace op_plugin {
27 27 
28-inline bool searchsorted_dims_matched_before_last_dim(const at::Tensor &boundaries, const at::Tensor &input) {28+inline bool searchsorted_dims_matched_before_last_dim(const at::Tensor& boundaries, const at::Tensor& input) {
29- if (boundaries.dim() != input.dim()) {29+ if (boundaries.dim() != input.dim()) {
30- return false;30+ return false;
31+ }
32+ const auto& dims_bd = boundaries.sizes();
33+ const auto& dims_in = input.sizes();
34+ for (int64_t dim = 0; dim + 1 < boundaries.dim(); ++dim) {
35+ if (dims_bd[dim] != dims_in[dim]) {
36+ return false;
31 }37 }
32- const auto &dims_bd = boundaries.sizes();38+ }
33- const auto &dims_in = input.sizes();39+ return true;
34- for (int64_t dim = 0; dim + 1 < boundaries.dim(); ++dim) {
35- if (dims_bd[dim] != dims_in[dim]) {
36- return false;
37- }
38- }
39- return true;
40}40}
41 41 
42/// Full searchsorted pre_check for Tensor values; `output` may be null to skip out-dtype checks (scalar path).42/// Full searchsorted pre_check for Tensor values; `output` may be null to skip out-dtype checks (scalar path).
43-inline void searchsorted_pre_check_npu(const at::Tensor &boundaries, const at::Tensor &input, const at::Tensor *output,43+inline void searchsorted_pre_check_npu(
44- bool out_int32, bool right, const c10::optional<c10::string_view> &side_opt,44+ const at::Tensor& boundaries,
45- const c10::optional<at::Tensor> &sorter_opt) {45+ const at::Tensor& input,
46- if (side_opt.has_value()) {46+ const at::Tensor* output,
47- c10::string_view side = *side_opt;47+ bool out_int32,
48- TORCH_CHECK(side == "left" || side == "right",48+ bool right,
49- "torch.searchsorted(): side can only be 'left' or 'right' but got ", std::string(side.data(), side.size()));49+ const c10::optional<c10::string_view>& side_opt,
50+ const c10::optional<at::Tensor>& sorter_opt) {
51+ if (side_opt.has_value()) {
52+ c10::string_view side = *side_opt;
53+ TORCH_CHECK(
54+ side == "left" || side == "right",
55+ "torch.searchsorted(): side can only be 'left' or 'right' but got ",
56+ std::string(side.data(), side.size()));
50 57 
51- TORCH_CHECK(!right || side == "right",58+ TORCH_CHECK(
52- "torch.searchsorted(): side and right can't be set to opposites, got side of ",59+ !right || side == "right",
53- std::string(side.data(), side.size()), " while right was True");60+ "torch.searchsorted(): side and right can't be set to opposites, got side of ",
61+ std::string(side.data(), side.size()),
62+ " while right was True");
63+ }
64+ 
65+ TORCH_CHECK(
66+ boundaries.device() == input.device(),
67+ "torch.searchsorted(): boundaries and input value tensors should have same device type, but got "
68+ "boundaries tensor device type ",
69+ boundaries.device(),
70+ " and input value tensor device type ",
71+ input.device());
72+ 
73+ if (sorter_opt.has_value()) {
74+ const at::Tensor& sorter = *sorter_opt;
75+ TORCH_CHECK(sorter.defined(), "torch.searchsorted(): optional sorter was set but the sorter tensor is undefined");
76+ 
77+ TORCH_CHECK(
78+ sorter.device() == boundaries.device(),
79+ "torch.searchsorted(): sorter and boundary tensors should have same device type, but got sorter tensor "
80+ "device type ",
81+ sorter.device(),
82+ " and input value tensor device type ",
83+ boundaries.device());
84+ 
85+ TORCH_CHECK(
86+ sorter.sizes() == boundaries.sizes(),
87+ "torch.searchsorted(): boundary and sorter must have the same size, but got boundary tensor ",
88+ boundaries.sizes(),
89+ "and got sorter tensor ",
90+ sorter.sizes());
91+ 
92+ TORCH_CHECK(
93+ sorter.scalar_type() == at::ScalarType::Long,
94+ "torch.searchsorted(): sorter must be a tensor of long dtype but got dtype ",
95+ sorter.scalar_type());
96+ 
97+ if (sorter.numel() > 0) {
98+ auto minmax = sorter.aminmax();
99+ int64_t vmin = std::get<0>(minmax).item().toLong();
100+ int64_t vmax = std::get<1>(minmax).item().toLong();
101+ TORCH_CHECK(vmin >= 0 && vmax < sorter.sizes().back(), "torch.searchsorted(): sorter index out of range");
54 }102 }
103+ }
55 104 
56- TORCH_CHECK(boundaries.device() == input.device(),105+ TORCH_CHECK(
57- "torch.searchsorted(): boundaries and input value tensors should have same device type, but got "106+ input.dim() > 0 || (input.dim() == 0 && input.numel() == 1 && boundaries.dim() == 1),
58- "boundaries tensor device type ",107+ "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we "
59- boundaries.device(), " and input value tensor device type ", input.device());108+ "got boundaries tensor dim(",
109+ boundaries.dim(),
110+ ") and input value's dim(",
111+ input.dim(),
112+ ") numel(",
113+ input.numel(),
114+ ")");
60 115 
61- if (sorter_opt.has_value()) {116+ // Python `searchsorted(boundaries, 1)` may bind as a Tensor overload with shape [1] instead of a true
62- const at::Tensor &sorter = *sorter_opt;117+ // rank-0 scalar; CPU raises the scalar/boundaries-dim error, not "first N-1 dimensions...".
63- TORCH_CHECK(118+ if (boundaries.dim() != 1 && input.numel() == 1 && input.dim() != boundaries.dim()) {
64- sorter.defined(), "torch.searchsorted(): optional sorter was set but the sorter tensor is undefined");119+ TORCH_CHECK(
65- 120+ false,
66- TORCH_CHECK(sorter.device() == boundaries.device(),
67- "torch.searchsorted(): sorter and boundary tensors should have same device type, but got sorter tensor "
68- "device type ",
69- sorter.device(), " and input value tensor device type ", boundaries.device());
70- 
71- TORCH_CHECK(sorter.sizes() == boundaries.sizes(),
72- "torch.searchsorted(): boundary and sorter must have the same size, but got boundary tensor ",
73- boundaries.sizes(), "and got sorter tensor ", sorter.sizes());
74- 
75- TORCH_CHECK(sorter.scalar_type() == at::ScalarType::Long,
76- "torch.searchsorted(): sorter must be a tensor of long dtype but got dtype ", sorter.scalar_type());
77- 
78- if (sorter.numel() > 0) {
79- auto minmax = sorter.aminmax();
80- int64_t vmin = std::get<0>(minmax).item().toLong();
81- int64_t vmax = std::get<1>(minmax).item().toLong();
82- TORCH_CHECK(vmin >= 0 && vmax < sorter.sizes().back(), "torch.searchsorted(): sorter index out of range");
83- }
84- }
85- 
86- TORCH_CHECK(input.dim() > 0 || (input.dim() == 0 && input.numel() == 1 && boundaries.dim() == 1),
87 "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we "121 "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we "
88 "got boundaries tensor dim(",122 "got boundaries tensor dim(",
89- boundaries.dim(), ") and input value's dim(", input.dim(), ") numel(", input.numel(), ")");123+ boundaries.dim(),
124+ ") and input value's dim(",
125+ input.dim(),
126+ ") numel(",
127+ input.numel(),
128+ ")");
129+ }
90 130 
91- // Python `searchsorted(boundaries, 1)` may bind as a Tensor overload with shape [1] instead of a true131+ TORCH_CHECK(
92- // rank-0 scalar; CPU raises the scalar/boundaries-dim error, not "first N-1 dimensions...".132+ boundaries.dim() != 0,
93- if (boundaries.dim() != 1 && input.numel() == 1 && input.dim() != boundaries.dim()) {133+ "torch.searchsorted(): boundaries tensor should have positive dimension, but got 0 dimension");
94- TORCH_CHECK(false,
95- "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we "
96- "got boundaries tensor dim(",
97- boundaries.dim(), ") and input value's dim(", input.dim(), ") numel(", input.numel(), ")");
98- }
99 134 
100- TORCH_CHECK(boundaries.dim() != 0,135+ TORCH_CHECK(
101- "torch.searchsorted(): boundaries tensor should have positive dimension, but got 0 dimension");136+ boundaries.dim() == 1 || searchsorted_dims_matched_before_last_dim(boundaries, input),
137+ "torch.searchsorted(): boundaries tensor should be 1 dimension or the first N-1 dimensions of boundaries "
138+ "tensor and input value tensor must match, but we got boundaries tensor ",
139+ boundaries.sizes(),
140+ " and input value tensor ",
141+ input.sizes());
102 142 
103- TORCH_CHECK(boundaries.dim() == 1 || searchsorted_dims_matched_before_last_dim(boundaries, input),143+ if (output != nullptr) {
104- "torch.searchsorted(): boundaries tensor should be 1 dimension or the first N-1 dimensions of boundaries "144+ at::ScalarType output_dtype = output->scalar_type();
105- "tensor and input value tensor must match, but we got boundaries tensor ",145+ TORCH_CHECK(
106- boundaries.sizes(), " and input value tensor ", input.sizes());146+ (output_dtype == at::ScalarType::Long && !out_int32) || (output_dtype == at::ScalarType::Int && out_int32),
147+ "torch.searchsorted(): output tensor's dtype is wrong, it can only be Int(int32) or Long(int64) "
148+ "depending on whether out_int32 flag is True, but we got output tensor's dtype ",
149+ output_dtype,
150+ " and out_int32 flag is ",
151+ (out_int32 ? "True" : "False"));
152+ }
107 153 
108- if (output != nullptr) {154+ if (out_int32) {
109- at::ScalarType output_dtype = output->scalar_type();155+ TORCH_CHECK(
110- TORCH_CHECK(156+ boundaries.sizes().back() < INT_MAX,
111- (output_dtype == at::ScalarType::Long && !out_int32) || (output_dtype == at::ScalarType::Int && out_int32),157+ "torch.searchsorted(): the size of boundaries' last dimension should be less than ",
112- "torch.searchsorted(): output tensor's dtype is wrong, it can only be Int(int32) or Long(int64) "158+ INT_MAX,
113- "depending on whether out_int32 flag is True, but we got output tensor's dtype ",159+ ", but we got ",
114- output_dtype, " and out_int32 flag is ", (out_int32 ? "True" : "False"));160+ boundaries.sizes().back());
115- }161+ }
116- 
117- if (out_int32) {
118- TORCH_CHECK(boundaries.sizes().back() < INT_MAX,
119- "torch.searchsorted(): the size of boundaries' last dimension should be less than ", INT_MAX,
120- ", but we got ", boundaries.sizes().back());
121- }
122}162}
123 163 
124/// Functional Tensor overload: `out` is created after new_params in generated op_api, so validate everything164/// Functional Tensor overload: `out` is created after new_params in generated op_api, so validate everything
125/// except output dtype here (matches CPU before aclnn). Tensor_out uses searchsorted_validate_tensor_out_op.165/// except output dtype here (matches CPU before aclnn). Tensor_out uses searchsorted_validate_tensor_out_op.
126-inline int searchsorted_validate_core_no_output(const at::Tensor &sorted_sequence, const at::Tensor &self,166+inline int searchsorted_validate_core_no_output(
127- bool out_int32, bool right, const c10::optional<c10::string_view> &side_opt,167+ const at::Tensor& sorted_sequence,
128- const c10::optional<at::Tensor> &sorter_opt) {168+ const at::Tensor& self,
129- searchsorted_pre_check_npu(sorted_sequence, self, nullptr, out_int32, right, side_opt, sorter_opt);169+ bool out_int32,
130- return 0;170+ bool right,
171+ const c10::optional<c10::string_view>& side_opt,
172+ const c10::optional<at::Tensor>& sorter_opt) {
173+ searchsorted_pre_check_npu(sorted_sequence, self, nullptr, out_int32, right, side_opt, sorter_opt);
174+ return 0;
131}175}
132 176 
133/// Tensor_out / same shapes as aclnnSearchSorted. Returns 0 for use in yaml new_params.177/// Tensor_out / same shapes as aclnnSearchSorted. Returns 0 for use in yaml new_params.
134-inline int searchsorted_validate_tensor_out_op(const at::Tensor &sorted_sequence, const at::Tensor &self,178+inline int searchsorted_validate_tensor_out_op(
135- const at::Tensor &out, bool out_int32, bool right, const c10::optional<c10::string_view> &side_opt,179+ const at::Tensor& sorted_sequence,
136- const c10::optional<at::Tensor> &sorter_opt) {180+ const at::Tensor& self,
137- searchsorted_pre_check_npu(sorted_sequence, self, &out, out_int32, right, side_opt, sorter_opt);181+ const at::Tensor& out,
138- return 0;182+ bool out_int32,
183+ bool right,
184+ const c10::optional<c10::string_view>& side_opt,
185+ const c10::optional<at::Tensor>& sorter_opt) {
186+ searchsorted_pre_check_npu(sorted_sequence, self, &out, out_int32, right, side_opt, sorter_opt);
187+ return 0;
139}188}
140 189 
141/// Scalar value: materialize wrapped-number tensor like aten (no output-dtype check; out is created by kernel).190/// Scalar value: materialize wrapped-number tensor like aten (no output-dtype check; out is created by kernel).
142-inline int searchsorted_validate_scalar_op(const at::Tensor &sorted_sequence, const c10::Scalar &self, bool out_int32,191+inline int searchsorted_validate_scalar_op(
143- bool right, const c10::optional<c10::string_view> &side_opt, const c10::optional<at::Tensor> &sorter_opt) {192+ const at::Tensor& sorted_sequence,
144- at::Tensor input_t = at::empty({}, sorted_sequence.options());193+ const c10::Scalar& self,
145- input_t.fill_(self);194+ bool out_int32,
146- input_t.unsafeGetTensorImpl()->set_wrapped_number(true);195+ bool right,
147- searchsorted_pre_check_npu(sorted_sequence, input_t, nullptr, out_int32, right, side_opt, sorter_opt);196+ const c10::optional<c10::string_view>& side_opt,
148- return 0;197+ const c10::optional<at::Tensor>& sorter_opt) {
198+ at::Tensor input_t = at::empty({}, sorted_sequence.options());
199+ input_t.fill_(self);
200+ input_t.unsafeGetTensorImpl()->set_wrapped_number(true);
201+ searchsorted_pre_check_npu(sorted_sequence, input_t, nullptr, out_int32, right, side_opt, sorter_opt);
202+ return 0;
149}203}
150 204 
151} // namespace op_plugin205} // namespace op_plugin
Mop_plugin/utils/SearchsortedWarnUtil.h+65-62
@@ -21,79 +21,82 @@ namespace op_plugin {
21/// Row-major contiguous check aligned with dense strided layout (matches TensorImpl contiguous semantics).21/// Row-major contiguous check aligned with dense strided layout (matches TensorImpl contiguous semantics).
22/// Some backends may report is_contiguous() true while strides still require an explicit copy for aclnn; this22/// Some backends may report is_contiguous() true while strides still require an explicit copy for aclnn; this
23/// catches permute().to(...) cases that CPU flags in searchsorted_maybe_trim_input_tensors.23/// catches permute().to(...) cases that CPU flags in searchsorted_maybe_trim_input_tensors.
24-inline bool searchsorted_tensor_is_row_major_contiguous(const at::Tensor &t) {24+inline bool searchsorted_tensor_is_row_major_contiguous(const at::Tensor& t) {
25- if (!t.defined() || t.numel() == 0) {25+ if (!t.defined() || t.numel() == 0) {
26- return true;
27- }
28- if (t.layout() != c10::Layout::Strided) {
29- return t.is_contiguous();
30- }
31- const int64_t dim = t.dim();
32- if (dim == 0) {
33- return true;
34- }
35- int64_t z = 1;
36- for (int64_t d = dim - 1; d >= 0; --d) {
37- const int64_t size_d = t.size(d);
38- if (size_d != 1) {
39- if (t.stride(d) != z) {
40- return false;
41- }
42- if (size_d == 0) {
43- return true;
44- }
45- }
46- z *= size_d;
47- }
48 return true;26 return true;
27+ }
28+ if (t.layout() != c10::Layout::Strided) {
29+ return t.is_contiguous();
30+ }
31+ const int64_t dim = t.dim();
32+ if (dim == 0) {
33+ return true;
34+ }
35+ int64_t z = 1;
36+ for (int64_t d = dim - 1; d >= 0; --d) {
37+ const int64_t size_d = t.size(d);
38+ if (size_d != 1) {
39+ if (t.stride(d) != z) {
40+ return false;
41+ }
42+ if (size_d == 0) {
43+ return true;
44+ }
45+ }
46+ z *= size_d;
47+ }
48+ return true;
49}49}
50 50 
51/// Tensor, Tensor overload: warn once per process when inputs may need a contiguous copy (matches CPU/CUDA).51/// Tensor, Tensor overload: warn once per process when inputs may need a contiguous copy (matches CPU/CUDA).
52inline int warn_if_searchsorted_inputs_noncontiguous(52inline int warn_if_searchsorted_inputs_noncontiguous(
53- const at::Tensor &sorted_sequence, const at::Tensor &self, const c10::optional<at::Tensor> &sorter_opt) {53+ const at::Tensor& sorted_sequence,
54- if (!searchsorted_tensor_is_row_major_contiguous(self)) {54+ const at::Tensor& self,
55- TORCH_WARN_ONCE(55+ const c10::optional<at::Tensor>& sorter_opt) {
56- "torch.searchsorted(): input value tensor is non-contiguous, this will lower the performance due "56+ if (!searchsorted_tensor_is_row_major_contiguous(self)) {
57- "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous input value "57+ TORCH_WARN_ONCE(
58- "tensor if possible. This message will only appear once per program.");58+ "torch.searchsorted(): input value tensor is non-contiguous, this will lower the performance due "
59+ "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous input value "
60+ "tensor if possible. This message will only appear once per program.");
61+ }
62+ if (!searchsorted_tensor_is_row_major_contiguous(sorted_sequence)) {
63+ TORCH_WARN_ONCE(
64+ "torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "
65+ "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "
66+ "tensor if possible. This message will only appear once per program.");
67+ }
68+ if (sorter_opt.has_value()) {
69+ const at::Tensor& st = *sorter_opt;
70+ if (st.defined() && !searchsorted_tensor_is_row_major_contiguous(st)) {
71+ TORCH_WARN_ONCE(
72+ "torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
73+ "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
74+ "tensor if possible. This message will only appear once per program.");
59 }75 }
60- if (!searchsorted_tensor_is_row_major_contiguous(sorted_sequence)) {76+ }
61- TORCH_WARN_ONCE(77+ return 0;
62- "torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "
63- "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "
64- "tensor if possible. This message will only appear once per program.");
65- }
66- if (sorter_opt.has_value()) {
67- const at::Tensor &st = *sorter_opt;
68- if (st.defined() && !searchsorted_tensor_is_row_major_contiguous(st)) {
69- TORCH_WARN_ONCE(
70- "torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
71- "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
72- "tensor if possible. This message will only appear once per program.");
73- }
74- }
75- return 0;
76}78}
77 79 
78/// Tensor, Scalar overload: only boundaries / sorter apply (scalar value is materialized separately).80/// Tensor, Scalar overload: only boundaries / sorter apply (scalar value is materialized separately).
79inline int warn_if_searchsorted_scalar_inputs_noncontiguous(81inline int warn_if_searchsorted_scalar_inputs_noncontiguous(
80- const at::Tensor &sorted_sequence, const c10::optional<at::Tensor> &sorter_opt) {82+ const at::Tensor& sorted_sequence,
81- if (!searchsorted_tensor_is_row_major_contiguous(sorted_sequence)) {83+ const c10::optional<at::Tensor>& sorter_opt) {
82- TORCH_WARN_ONCE(84+ if (!searchsorted_tensor_is_row_major_contiguous(sorted_sequence)) {
83- "torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "85+ TORCH_WARN_ONCE(
84- "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "86+ "torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "
85- "tensor if possible. This message will only appear once per program.");87+ "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "
88+ "tensor if possible. This message will only appear once per program.");
89+ }
90+ if (sorter_opt.has_value()) {
91+ const at::Tensor& st = *sorter_opt;
92+ if (st.defined() && !searchsorted_tensor_is_row_major_contiguous(st)) {
93+ TORCH_WARN_ONCE(
94+ "torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
95+ "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
96+ "tensor if possible. This message will only appear once per program.");
86 }97 }
87- if (sorter_opt.has_value()) {98+ }
88- const at::Tensor &st = *sorter_opt;99+ return 0;
89- if (st.defined() && !searchsorted_tensor_is_row_major_contiguous(st)) {
90- TORCH_WARN_ONCE(
91- "torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
92- "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
93- "tensor if possible. This message will only appear once per program.");
94- }
95- }
96- return 0;
97}100}
98 101 
99} // namespace op_plugin102} // namespace op_plugin
Mop_plugin/utils/StdVarCorrectionUtils.h+23-24
@@ -16,31 +16,30 @@
16#include <c10/core/Scalar.h>16#include <c10/core/Scalar.h>
17#include <c10/util/Optional.h>17#include <c10/util/Optional.h>
18 18 
19-inline bool correction_fits_aclnn_int64(const c10::optional<c10::Scalar> &correction)19+inline bool correction_fits_aclnn_int64(const c10::optional<c10::Scalar>& correction) {
20-{20+ if (!correction.has_value()) {
21- if (!correction.has_value()) {
22- return true;
23- }
24- const c10::Scalar &s = correction.value();
25- if (s.isIntegral(true)) {
26- return true;
27- }
28- if (!s.isFloatingPoint()) {
29- return false;
30- }
31- const double v = s.toDouble();
32- if (!std::isfinite(v)) {
33- return false;
34- }
35- const double t = std::trunc(v);
36- if (v != t) {
37- return false;
38- }
39- if (v > static_cast<double>(std::numeric_limits<int64_t>::max()) ||
40- v < static_cast<double>(std::numeric_limits<int64_t>::min())) {
41- return false;
42- }
43 return true;21 return true;
22+ }
23+ const c10::Scalar& s = correction.value();
24+ if (s.isIntegral(true)) {
25+ return true;
26+ }
27+ if (!s.isFloatingPoint()) {
28+ return false;
29+ }
30+ const double v = s.toDouble();
31+ if (!std::isfinite(v)) {
32+ return false;
33+ }
34+ const double t = std::trunc(v);
35+ if (v != t) {
36+ return false;
37+ }
38+ if (v > static_cast<double>(std::numeric_limits<int64_t>::max()) ||
39+ v < static_cast<double>(std::numeric_limits<int64_t>::min())) {
40+ return false;
41+ }
42+ return true;
44}43}
45 44 
46#endif // OP_PLUGIN_UTILS_STD_VAR_CORRECTION_UTILS_H_45#endif // OP_PLUGIN_UTILS_STD_VAR_CORRECTION_UTILS_H_
Mop_plugin/utils/custom_functions/aclops/GeluBackwardKernelNpu.cpp+5-11
@@ -16,16 +16,10 @@
16#include "op_plugin/utils/OpAdapter.h"16#include "op_plugin/utils/OpAdapter.h"
17 17 
18namespace acl_op {18namespace acl_op {
19-at::Tensor gelu_backward_common_nocheck(const at::Tensor& grad, const at::Tensor& self)19+at::Tensor gelu_backward_common_nocheck(const at::Tensor& grad, const at::Tensor& self) {
20-{20+ at::Tensor grad_input = at_npu::native::OpPreparation::apply_tensor(self);
21- at::Tensor grad_input = at_npu::native::OpPreparation::apply_tensor(self);21+ at_npu::native::OpCommand cmd;
22- at_npu::native::OpCommand cmd;22+ cmd.Name("GeluGrad").Input(grad).Input(self).Input(grad).Output(grad_input).Run();
23- cmd.Name("GeluGrad")23+ return grad_input;
24- .Input(grad)
25- .Input(self)
26- .Input(grad)
27- .Output(grad_input)
28- .Run();
29- return grad_input;
30}24}
31} // namespace acl_op25} // namespace acl_op
Mop_plugin/utils/custom_functions/aclops/SoftplusBackwardKernelNpu.cpp+10-11
@@ -21,17 +21,16 @@ at::Tensor& softplus_backward_out_common_nocheck(
21 const at::Tensor& grad_output,21 const at::Tensor& grad_output,
22 const at::Tensor& self,22 const at::Tensor& self,
23 at::Scalar beta,23 at::Scalar beta,
24- at::Scalar threshold)24+ at::Scalar threshold) {
25-{25+ at_npu::native::OpCommand cmd;
26- at_npu::native::OpCommand cmd;26+ cmd.Name("SoftplusV2Grad")
27- cmd.Name("SoftplusV2Grad")27+ .Input(grad_output)
28- .Input(grad_output)28+ .Input(self)
29- .Input(self)29+ .Output(grad_input)
30- .Output(grad_input)30+ .Attr("beta", beta)
31- .Attr("beta", beta)31+ .Attr("threshold", threshold)
32- .Attr("threshold", threshold)32+ .Run();
33- .Run();
34 33 
35- return grad_input;34+ return grad_input;
36}35}
37} // namespace acl_op36} // namespace acl_op
Mop_plugin/utils/custom_functions/aclops/SumKernelNpu.cpp+58-81
@@ -23,35 +23,21 @@ using npu_preparation = at_npu::native::OpPreparation;
23using npu_utils = at_npu::native::NpuUtils;23using npu_utils = at_npu::native::NpuUtils;
24 24 
25namespace {25namespace {
26-at::Tensor& sum_out_npu_nocheck(26+at::Tensor& sum_out_npu_nocheck(at::Tensor& result, const at::Tensor& self, at::IntArrayRef dim, bool keepdim) {
27- at::Tensor& result,27+ at::dim_list_to_bitset(dim, self.dim());
28- const at::Tensor& self,28+ c10::SmallVector<int64_t, N> dim_list =
29- at::IntArrayRef dim,29+ dim.empty() ? op_plugin::utils::get_dimlist_for_tensor(self) : c10::SmallVector<int64_t, N>(dim);
30- bool keepdim)30+ at_npu::native::OpCommand cmd;
31-{31+ cmd.Name("ReduceSum").Input(self).Input(dim_list, at::kLong).Output(result).Attr("keep_dims", keepdim).Run();
32- at::dim_list_to_bitset(dim, self.dim());32+ return result;
33- c10::SmallVector<int64_t, N> dim_list = dim.empty() ? op_plugin::utils::get_dimlist_for_tensor(self) :
34- c10::SmallVector<int64_t, N>(dim);
35- at_npu::native::OpCommand cmd;
36- cmd.Name("ReduceSum")
37- .Input(self)
38- .Input(dim_list, at::kLong)
39- .Output(result)
40- .Attr("keep_dims", keepdim)
41- .Run();
42- return result;
43}33}
44 34 
45-at::Tensor check_dtype(35+at::Tensor check_dtype(const at::Tensor& self, c10::ScalarType out_type) {
46- const at::Tensor &self,36+ if (isIntegralType(out_type, true)) {
47- c10::ScalarType out_type)37+ out_type = at::kFloat;
48-{38+ }
49- if (isIntegralType(out_type, true)) {39+ at::Tensor self_cp = (self.scalar_type() == out_type) ? self : acl_op::_npu_dtype_cast(self, out_type);
50- out_type = at::kFloat;40+ return self_cp;
51- }
52- at::Tensor self_cp = (self.scalar_type() == out_type) ? self :
53- acl_op::_npu_dtype_cast(self, out_type);
54- return self_cp;
55}41}
56} // namespace42} // namespace
57 43 
@@ -60,72 +46,63 @@ at::Tensor& sum_out_common_nocheck(
60 const at::Tensor& self,46 const at::Tensor& self,
61 at::IntArrayRef dim,47 at::IntArrayRef dim,
62 bool keepdim,48 bool keepdim,
63- c10::optional<c10::ScalarType> dtype)49+ c10::optional<c10::ScalarType> dtype) {
64-{50+ auto output_size = op_infer::sum_npu_output_size(self, dim, keepdim);
65- auto output_size = op_infer::sum_npu_output_size(self, dim, keepdim);51+ auto res_type = dtype.has_value() ? dtype.value() : result.scalar_type();
66- auto res_type = dtype.has_value() ? dtype.value() : result.scalar_type();
67 52 
68- npu_preparation::CheckOut(53+ npu_preparation::CheckOut({self}, result, ACL_FORMAT_ND, res_type, output_size);
69- {self},
70- result,
71- ACL_FORMAT_ND,
72- res_type,
73- output_size);
74 54 
75- if (self.numel() == 0) {55+ if (self.numel() == 0) {
76- at::Tensor result_cast = at::empty(output_size, self.options().dtype(res_type));56+ at::Tensor result_cast = at::empty(output_size, self.options().dtype(res_type));
77- result.copy_(result_cast);57+ result.copy_(result_cast);
78- return result;
79- }
80- 
81- at::Tensor self_cp = check_dtype(self, res_type);
82- at::Tensor result_cp = result.scalar_type() == self_cp.scalar_type() ? result :
83- acl_op::_npu_dtype_cast(result, self_cp.scalar_type());
84- if (!npu_utils::check_match(&result_cp)) {
85- at::Tensor contiguous_result = npu_utils::format_contiguous(result_cp);
86- sum_out_npu_nocheck(contiguous_result, self_cp, dim, keepdim);
87- npu_utils::format_fresh_view(result, contiguous_result);
88- } else {
89- sum_out_npu_nocheck(result_cp, self_cp, dim, keepdim);
90- }
91- 
92- if (result_cp.scalar_type() != res_type) {
93- result_cp = acl_op::_npu_dtype_cast(result_cp, res_type);
94- result.copy_(result_cp);
95- } else {
96- result = result_cp;
97- }
98 return result;58 return result;
99-}59+ }
100 60 
61+ at::Tensor self_cp = check_dtype(self, res_type);
62+ at::Tensor result_cp =
63+ result.scalar_type() == self_cp.scalar_type() ? result : acl_op::_npu_dtype_cast(result, self_cp.scalar_type());
64+ if (!npu_utils::check_match(&result_cp)) {
65+ at::Tensor contiguous_result = npu_utils::format_contiguous(result_cp);
66+ sum_out_npu_nocheck(contiguous_result, self_cp, dim, keepdim);
67+ npu_utils::format_fresh_view(result, contiguous_result);
68+ } else {
69+ sum_out_npu_nocheck(result_cp, self_cp, dim, keepdim);
70+ }
71+ 
72+ if (result_cp.scalar_type() != res_type) {
73+ result_cp = acl_op::_npu_dtype_cast(result_cp, res_type);
74+ result.copy_(result_cp);
75+ } else {
76+ result = result_cp;
77+ }
78+ return result;
79+}
101 80 
102at::Tensor sum_common_nocheck(81at::Tensor sum_common_nocheck(
103 const at::Tensor& self,82 const at::Tensor& self,
104 at::IntArrayRef dim,83 at::IntArrayRef dim,
105 bool keepdim,84 bool keepdim,
106- c10::optional<c10::ScalarType> dtype)85+ c10::optional<c10::ScalarType> dtype) {
107-{86+ auto output_size = op_infer::reduce_ops_npu_output_size(self, dim, keepdim);
108- auto output_size = op_infer::reduce_ops_npu_output_size(self, dim, keepdim);87+ auto out_type = self.scalar_type();
109- auto out_type = self.scalar_type();
110 88 
111- if (dtype.has_value()) {89+ if (dtype.has_value()) {
112- out_type = dtype.value();90+ out_type = dtype.value();
113- } else if (isIntegralType(out_type, true)) {91+ } else if (isIntegralType(out_type, true)) {
114- out_type = at::kLong;92+ out_type = at::kLong;
115- }93+ }
116 94 
117- if (self.numel() == 0) {95+ if (self.numel() == 0) {
118- return at::zeros(output_size, self.options().dtype(out_type));96+ return at::zeros(output_size, self.options().dtype(out_type));
119- }97+ }
120 98 
121- at::Tensor self_cp = check_dtype(self, out_type);99+ at::Tensor self_cp = check_dtype(self, out_type);
122- at::Tensor result = npu_preparation::apply_tensor_with_format(100+ at::Tensor result = npu_preparation::apply_tensor_with_format(output_size, self_cp.options(), ACL_FORMAT_ND);
123- output_size, self_cp.options(), ACL_FORMAT_ND);101+ sum_out_npu_nocheck(result, self_cp, dim, keepdim);
124- sum_out_npu_nocheck(result, self_cp, dim, keepdim);
125 102 
126- if (result.scalar_type() != out_type) {103+ if (result.scalar_type() != out_type) {
127- result = acl_op::_npu_dtype_cast(result, out_type);104+ result = acl_op::_npu_dtype_cast(result, out_type);
128- }105+ }
129- return result;106+ return result;
130}107}
131} // namespace acl_op108} // namespace acl_op
Mop_plugin/utils/custom_functions/opapi/SumKernelNpuOpApi.cpp+34-34
@@ -20,45 +20,45 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-at::Tensor &sum_out_common_nocheck(const at::Tensor &self,23+at::Tensor& sum_out_common_nocheck(
24- at::IntArrayRef dim,24+ const at::Tensor& self,
25- bool keepdim,25+ at::IntArrayRef dim,
26- c10::optional<c10::ScalarType> dtype,26+ bool keepdim,
27- at::Tensor &result)27+ c10::optional<c10::ScalarType> dtype,
28-{28+ at::Tensor& result) {
29- DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum_out(self, dim, keepdim, dtype, result));29+ DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum_out(self, dim, keepdim, dtype, result));
30- auto output_size = op_infer::sum_npu_output_size(self, dim, keepdim);30+ auto output_size = op_infer::sum_npu_output_size(self, dim, keepdim);
31- auto res_type = dtype.has_value() ? dtype.value() : result.scalar_type();31+ auto res_type = dtype.has_value() ? dtype.value() : result.scalar_type();
32- npu_preparation::check_tensor({self}, result, res_type, output_size);32+ npu_preparation::check_tensor({self}, result, res_type, output_size);
33 33 
34- EXEC_NPU_CMD(aclnnReduceSum, self, dim, keepdim, res_type, result);34+ EXEC_NPU_CMD(aclnnReduceSum, self, dim, keepdim, res_type, result);
35- return result;35+ return result;
36}36}
37 37 
38-at::Tensor sum_common_nocheck(const at::Tensor &self,38+at::Tensor sum_common_nocheck(
39- at::IntArrayRef dim,39+ const at::Tensor& self,
40- bool keepdim,40+ at::IntArrayRef dim,
41- c10::optional<c10::ScalarType> dtype)41+ bool keepdim,
42-{42+ c10::optional<c10::ScalarType> dtype) {
43- DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dim, keepdim, dtype));43+ DO_COMPATIBILITY(aclnnReduceSum, acl_op::sum(self, dim, keepdim, dtype));
44- auto output_size = op_infer::reduce_ops_npu_output_size(self, dim, keepdim);44+ auto output_size = op_infer::reduce_ops_npu_output_size(self, dim, keepdim);
45- auto self_size = self.sizes();45+ auto self_size = self.sizes();
46- auto out_type = self.scalar_type();46+ auto out_type = self.scalar_type();
47 47 
48- if (dtype.has_value()) {48+ if (dtype.has_value()) {
49- out_type = dtype.value();49+ out_type = dtype.value();
50- } else if (isIntegralType(out_type, true)) {50+ } else if (isIntegralType(out_type, true)) {
51- out_type = at::kLong;51+ out_type = at::kLong;
52+ }
53+ 
54+ for (uint64_t i = 0; i < self_size.size(); i++) {
55+ if (self_size[i] == 0) {
56+ return at::zeros(output_size, self.options().dtype(out_type));
52 }57 }
58+ }
53 59 
54- for (uint64_t i = 0; i < self_size.size(); i++) {60+ at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(out_type));
55- if (self_size[i] == 0) {61+ EXEC_NPU_CMD(aclnnReduceSum, self, dim, keepdim, out_type, result);
56- return at::zeros(output_size, self.options().dtype(out_type));62+ return result;
57- }
58- }
59- 
60- at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(out_type));
61- EXEC_NPU_CMD(aclnnReduceSum, self, dim, keepdim, out_type, result);
62- return result;
63}63}
64} // namespace op_api64} // namespace op_api
Mtest/core_tests/test_compatibility.py+1-1
@@ -200,7 +200,7 @@ class TestOpApiCompatibility(TestCase):
200 allow_dict["torch_npu"].extend(allow_dict_info["torch_npu"][version_tag])200 allow_dict["torch_npu"].extend(allow_dict_info["torch_npu"][version_tag])
201 201 
202 # load torch_npu_OpApi_schema_all.json202 # load torch_npu_OpApi_schema_all.json
203- base_schema = {} 203+ base_schema = {}
204 with open(get_file_path_2(os.path.dirname(__file__), "torch_npu_OpApi_schema_all.json")) as fp:204 with open(get_file_path_2(os.path.dirname(__file__), "torch_npu_OpApi_schema_all.json")) as fp:
205 base_schema0 = json.load(fp)205 base_schema0 = json.load(fp)
206 for key, value in base_schema0.items():206 for key, value in base_schema0.items():
Mtest/test_base_ops/test_addmm_activation_out.py+7-7
@@ -57,13 +57,13 @@ class TestAddmmActivationOut(TestCase):
57 cpu_out = cpu_out.to(torch.float32)57 cpu_out = cpu_out.to(torch.float32)
58 58 
59 cpu_out_copy = copy.deepcopy(cpu_out)59 cpu_out_copy = copy.deepcopy(cpu_out)
60- cpu_ref = self.cpu_reference_out(cpu_input1, cpu_input2, cpu_input3, 60+ cpu_ref = self.cpu_reference_out(cpu_input1, cpu_input2, cpu_input3,
61 beta, alpha, use_gelu, cpu_out_copy)61 beta, alpha, use_gelu, cpu_out_copy)
62- 62+ 
63 if dtype == np.float16:63 if dtype == np.float16:
64 cpu_ref = cpu_ref.astype(np.float16)64 cpu_ref = cpu_ref.astype(np.float16)
65- 65+ 
66- npu_res = self.npu_op_exec_out(npu_input1, npu_input2, npu_input3, 66+ npu_res = self.npu_op_exec_out(npu_input1, npu_input2, npu_input3,
67 beta, alpha, use_gelu, npu_out)67 beta, alpha, use_gelu, npu_out)
68 self.assertRtolEqual(cpu_ref, npu_res, prec=1.e-3, prec16=1.e-3)68 self.assertRtolEqual(cpu_ref, npu_res, prec=1.e-3, prec16=1.e-3)
69 69 
@@ -71,7 +71,7 @@ class TestAddmmActivationOut(TestCase):
71 shape_a = (3, 3)71 shape_a = (3, 3)
72 shape_b = (3, 4)72 shape_b = (3, 4)
73 shape_c = (4, 3)73 shape_c = (4, 3)
74- 74+ 
75 # ReLU 测试75 # ReLU 测试
76 self._run_test(np.float32, shape_a, shape_b, shape_c, 1, (0, 2), False)76 self._run_test(np.float32, shape_a, shape_b, shape_c, 1, (0, 2), False)
77 # GeLU 测试77 # GeLU 测试
@@ -81,7 +81,7 @@ class TestAddmmActivationOut(TestCase):
81 shape_a = (3, 3)81 shape_a = (3, 3)
82 shape_b = (3, 4)82 shape_b = (3, 4)
83 shape_c = (4, 3)83 shape_c = (4, 3)
84- 84+ 
85 # ReLU 测试85 # ReLU 测试
86 self._run_test(np.float16, shape_a, shape_b, shape_c, 2, (0, 10), False)86 self._run_test(np.float16, shape_a, shape_b, shape_c, 2, (0, 10), False)
87 # GeLU 测试87 # GeLU 测试
@@ -89,4 +89,4 @@ class TestAddmmActivationOut(TestCase):
89 89 
90 90 
91if __name__ == '__main__':91if __name__ == '__main__':
92- run_tests()92+ run_tests()
Mtest/test_base_ops/test_npu_scaled_grouped_mm.py+1-1
@@ -528,4 +528,4 @@ class TestScaledGroupedMm(TestCase):
528 528 
529 529 
530if __name__ == "__main__":530if __name__ == "__main__":
531- run_tests()531+ run_tests()
Mtest/test_base_ops/test_npu_scaled_grouped_mm_v2.py+1-1
@@ -667,4 +667,4 @@ class TestScaledGroupedMmV2(TestCase):
667 self.assertTrue(output.device.type == 'npu')667 self.assertTrue(output.device.type == 'npu')
668 668 
669if __name__ == "__main__":669if __name__ == "__main__":
670- run_tests()670+ run_tests()
Mtest/test_custom_ops/test_npu_dequant_swiglu_quant.py+1-1
@@ -475,4 +475,4 @@ class TestNPUDequantSwigluQuant(TestCase):
475 475 
476 476 
477if __name__ == "__main__":477if __name__ == "__main__":
478- run_tests()478+ run_tests()
Mtest/test_custom_ops/test_npu_moe_finalize_routing_v2.py+8-8
@@ -20,12 +20,12 @@ class MoeFinalizeRoutingData:
20 20 
21 21 
22class TestMoeFinalizeRouting(TestCase):22class TestMoeFinalizeRouting(TestCase):
23- 23+ 
24 def moe_finalize_routing_np(self, data_struct):24 def moe_finalize_routing_np(self, data_struct):
25 NK = data_struct.expanded_src_to_dst_row.shape[0]25 NK = data_struct.expanded_src_to_dst_row.shape[0]
26 K = 126 K = 1
27 if data_struct.scales is not None:27 if data_struct.scales is not None:
28- K = data_struct.scales.shape[1] 28+ K = data_struct.scales.shape[1]
29 num_rows = NK // K29 num_rows = NK // K
30 H = data_struct.expanded_permuted_rows.shape[-1]30 H = data_struct.expanded_permuted_rows.shape[-1]
31 expanded_permuted_rows = data_struct.expanded_permuted_rows.reshape(-1, H)31 expanded_permuted_rows = data_struct.expanded_permuted_rows.reshape(-1, H)
@@ -45,7 +45,7 @@ class TestMoeFinalizeRouting(TestCase):
45 else:45 else:
46 dst_row = expanded_permuted_rows[value, :]46 dst_row = expanded_permuted_rows[value, :]
47 expert_id = data_struct.expert_for_source_row[i, k]47 expert_id = data_struct.expert_for_source_row[i, k]
48- 48+ 
49 scalesV = 1.049 scalesV = 1.0
50 if data_struct.scales is not None:50 if data_struct.scales is not None:
51 scalesV = data_struct.scales[i, k]51 scalesV = data_struct.scales[i, k]
@@ -54,7 +54,7 @@ class TestMoeFinalizeRouting(TestCase):
54 else:54 else:
55 out[i, :] += scalesV * dst_row55 out[i, :] += scalesV * dst_row
56 return out56 return out
57- 57+ 
58 def custom_op_exec(self, data_struct):58 def custom_op_exec(self, data_struct):
59 scales_npu = None59 scales_npu = None
60 if data_struct.scales is not None:60 if data_struct.scales is not None:
@@ -70,7 +70,7 @@ class TestMoeFinalizeRouting(TestCase):
70 torch.tensor(data_struct.expanded_src_to_dst_row).npu(),70 torch.tensor(data_struct.expanded_src_to_dst_row).npu(),
71 expert_for_source_row_npu,71 expert_for_source_row_npu,
72 data_struct.drop_pad_mode)72 data_struct.drop_pad_mode)
73- 73+ 
74 def generate_input_data(self, expert_num=16, token_len=10, top_k=4, num_rows=50):74 def generate_input_data(self, expert_num=16, token_len=10, top_k=4, num_rows=50):
75 expanded_permuted_rows = np.random.randn(num_rows * top_k, token_len).astype(np.float32)75 expanded_permuted_rows = np.random.randn(num_rows * top_k, token_len).astype(np.float32)
76 skip1 = np.random.randn(num_rows, token_len).astype(np.float32)76 skip1 = np.random.randn(num_rows, token_len).astype(np.float32)
@@ -83,7 +83,7 @@ class TestMoeFinalizeRouting(TestCase):
83 data_struct = MoeFinalizeRoutingData(expanded_permuted_rows, skip1, skip2_optional, bias, scales,83 data_struct = MoeFinalizeRoutingData(expanded_permuted_rows, skip1, skip2_optional, bias, scales,
84 expanded_src_to_dst_row, expert_for_source_row, 0)84 expanded_src_to_dst_row, expert_for_source_row, 0)
85 return data_struct85 return data_struct
86- 86+ 
87 def generate_input_data_drop_pad(self, expert_num=16, token_len=10, c=20, top_k=4, num_rows=50):87 def generate_input_data_drop_pad(self, expert_num=16, token_len=10, c=20, top_k=4, num_rows=50):
88 expanded_permuted_rows = np.random.randn(expert_num, c, token_len).astype(np.float32)88 expanded_permuted_rows = np.random.randn(expert_num, c, token_len).astype(np.float32)
89 skip1 = np.random.randn(num_rows, token_len).astype(np.float32)89 skip1 = np.random.randn(num_rows, token_len).astype(np.float32)
@@ -96,7 +96,7 @@ class TestMoeFinalizeRouting(TestCase):
96 data_struct = MoeFinalizeRoutingData(expanded_permuted_rows, skip1, skip2_optional, bias, scales,96 data_struct = MoeFinalizeRoutingData(expanded_permuted_rows, skip1, skip2_optional, bias, scales,
97 expanded_src_to_dst_row, expert_for_source_row, 1)97 expanded_src_to_dst_row, expert_for_source_row, 1)
98 return data_struct98 return data_struct
99- 99+ 
100 @SupportedDevices(['Ascend910B'])100 @SupportedDevices(['Ascend910B'])
101 def test_moe_finalize_routing(self, device="npu"):101 def test_moe_finalize_routing(self, device="npu"):
102 data_struct = self.generate_input_data(expert_num=16, token_len=5, top_k=4, num_rows=5)102 data_struct = self.generate_input_data(expert_num=16, token_len=5, top_k=4, num_rows=5)
@@ -114,4 +114,4 @@ class TestMoeFinalizeRouting(TestCase):
114 self.assertRtolEqual(expected_output, custom_output.cpu().numpy(), 0.0001)114 self.assertRtolEqual(expected_output, custom_output.cpu().numpy(), 0.0001)
115 115 
116if __name__ == "__main__":116if __name__ == "__main__":
117- run_tests() 117+ run_tests()
Mtorchnpugen/api/types/signatures.py+1-1
@@ -1,5 +1,5 @@
1# Copyright (c) 2023 Huawei Technologies Co., Ltd1# Copyright (c) 2023 Huawei Technologies Co., Ltd
2-# Copyright (c) 2019, Facebook CORPORATION. 2+# Copyright (c) 2019, Facebook CORPORATION.
3# All rights reserved.3# All rights reserved.
4#4#
5# Licensed under the BSD 3-Clause License (the "License");5# Licensed under the BSD 3-Clause License (the "License");