Star187
855
代码
代码
Issues181
Pull Requests333
项目讨论
Wiki
分析
Star187
855
已合并
[master][Fix] Fix static check errors detected by clang-format #5558
thickhair创建于 7月25日
[master][Fix] Fix static check errors detected by clang-format #5558
已合并
thickhair创建于 7月25日
从thickhair/op-pluginclang-format-5合入到Ascend/op-pluginmaster
共 49 个文件变更+7101-6842
M.pre-commit-config.yaml+15-9
@@ -79,15 +79,21 @@ repos:
79 args: ["--force-exclude", "--config", "pre-commit/typos.toml"]79 args: ["--force-exclude", "--config", "pre-commit/typos.toml"]
80 80 
81 #--------------- C++ 核心检查 ---------------------------------------------81 #--------------- C++ 核心检查 ---------------------------------------------
82- #- repo: https://gitcode.com/pre-commit-clang/mirrors-clang-format82+ # - repo: https://gitcode.com/pre-commit-clang/mirrors-clang-format
83- # rev: v18.1.883+ # rev: v18.1.8
84- # hooks:84+ # hooks:
85- # - id: clang-format85+ # - id: clang-format
86- # files: \.(c|h|cpp|hpp|cc|hh|cxx|hxx)$86+ # files: \.(c|h|cpp|hpp|cc|hh|cxx|hxx)$
87- # args:87+ # # 跳过 Python string.Template 模板目录，其中的 ${namespace} / ${declarations}
88- # - "--style=file" # 读取 .clang-format 文件88+ # # 占位符会被 clang-format 误改（在 $ 和 { 之间插入空格），导致代码生成失败
89- # - "--verbose"89+ # exclude: |
90- # - "-i"90+ # (?x)(
91+ # ^torchnpugen/templates/
92+ # )
93+ # args:
94+ # - "--style=file" # 读取 .clang-format 文件
95+ # - "--verbose"
96+ # - "-i"
91 97 
92 # -------------------------- Gitleaks 本地离线二进制扫描 --------------------------98 # -------------------------- Gitleaks 本地离线二进制扫描 --------------------------
93 - repo: local99 - repo: local
Mop_plugin/ops/opapi/RotaryMulKernelNpuOpApi.cpp+36-40
@@ -23,17 +23,12 @@ using npu_preparation = at_npu::native::OpPreparation;
23const static int64_t ROTATE_HALF = 0;23const static int64_t ROTATE_HALF = 0;
24const static int64_t ROTATE_INTERLEAVED = 1;24const static int64_t ROTATE_INTERLEAVED = 1;
25 25 
26-static bool isRotaryMulMixDtypeSupport(26+static bool isRotaryMulMixDtypeSupport(const at::Tensor& self, const at::Tensor& r1, const at::Tensor& r2) {
27- const at::Tensor& self,27+ return self.dtype() == r1.dtype() && self.dtype() == r2.dtype() ? false : true;
28- const at::Tensor& r1,
29- const at::Tensor& r2)
30-{
31- return self.dtype() == r1.dtype() && self.dtype() == r2.dtype() ? false : true;
32}28}
33 29 
34-static at::Tensor npu_dtype_cast_impl_op_api(const at::Tensor& self, at::ScalarType dtype)30+static at::Tensor npu_dtype_cast_impl_op_api(const at::Tensor& self, at::ScalarType dtype) {
35-{31+ return self.dtype() == dtype ? self : self.to(dtype);
36- return self.dtype() == dtype ? self : self.to(dtype);
37}32}
38 33 
39at::Tensor npu_rotary_mul(34at::Tensor npu_rotary_mul(
@@ -41,38 +36,39 @@ at::Tensor npu_rotary_mul(
41 const at::Tensor& r1,36 const at::Tensor& r1,
42 const at::Tensor& r2,37 const at::Tensor& r2,
43 c10::string_view rotary_mode,38 c10::string_view rotary_mode,
44- const c10::optional<at::Tensor>& rotate)39+ const c10::optional<at::Tensor>& rotate) {
45-{40+ static bool notNeedCheck = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
46- static bool notNeedCheck = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;41+ TORCH_CHECK(
47- TORCH_CHECK((notNeedCheck || (rotary_mode == "half" || rotary_mode == "interleave")),42+ (notNeedCheck || (rotary_mode == "half" || rotary_mode == "interleave")),
48- "The rotary_mode of npu_rotary_mul should be half or interleave, but got ", rotary_mode,43+ "The rotary_mode of npu_rotary_mul should be half or interleave, but got ",
49- OPS_ERROR(ErrCode::PARAM));44+ rotary_mode,
50- DO_COMPATIBILITY(aclnnRotaryPositionEmbedding, acl_op::npu_rotary_mul(self, r1, r2, rotary_mode));45+ OPS_ERROR(ErrCode::PARAM));
46+ DO_COMPATIBILITY(aclnnRotaryPositionEmbedding, acl_op::npu_rotary_mul(self, r1, r2, rotary_mode));
51 47 
52- int64_t mode = op_plugin::utils::get_rotary_mode(rotary_mode);48+ int64_t mode = op_plugin::utils::get_rotary_mode(rotary_mode);
53- if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend910B1) {49+ if (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend910B1) {
54- return acl_op::npu_rotary_mul(self, r1, r2, rotary_mode);50+ return acl_op::npu_rotary_mul(self, r1, r2, rotary_mode);
55- }51+ }
56- at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options());52+ at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options());
57- static const bool is_v2_available = check_aclnn_kernel_available("aclnnRotaryPositionEmbeddingV2");53+ static const bool is_v2_available = check_aclnn_kernel_available("aclnnRotaryPositionEmbeddingV2");
58- bool isMixDataType = isRotaryMulMixDtypeSupport(self, r1, r2);54+ bool isMixDataType = isRotaryMulMixDtypeSupport(self, r1, r2);
59- if (is_v2_available) {55+ if (is_v2_available) {
60- if (isMixDataType) {56+ if (isMixDataType) {
61- at::Tensor cosCast = npu_dtype_cast_impl_op_api(r1, self.scalar_type());57+ at::Tensor cosCast = npu_dtype_cast_impl_op_api(r1, self.scalar_type());
62- at::Tensor sinCast = npu_dtype_cast_impl_op_api(r2, self.scalar_type());58+ at::Tensor sinCast = npu_dtype_cast_impl_op_api(r2, self.scalar_type());
63- EXEC_NPU_CMD(aclnnRotaryPositionEmbeddingV2, self, cosCast, sinCast, mode, rotate, result);59+ EXEC_NPU_CMD(aclnnRotaryPositionEmbeddingV2, self, cosCast, sinCast, mode, rotate, result);
64- } else {
65- EXEC_NPU_CMD(aclnnRotaryPositionEmbeddingV2, self, r1, r2, mode, rotate, result);
66- }
67 } else {60 } else {
68- if (isMixDataType) {61+ EXEC_NPU_CMD(aclnnRotaryPositionEmbeddingV2, self, r1, r2, mode, rotate, result);
69- at::Tensor cosCast = npu_dtype_cast_impl_op_api(r1, self.scalar_type());
70- at::Tensor sinCast = npu_dtype_cast_impl_op_api(r2, self.scalar_type());
71- EXEC_NPU_CMD(aclnnRotaryPositionEmbedding, self, cosCast, sinCast, mode, result);
72- } else {
73- EXEC_NPU_CMD(aclnnRotaryPositionEmbedding, self, r1, r2, mode, result);
74- }
75 }62 }
76- return result;63+ } else {
77-}64+ if (isMixDataType) {
65+ at::Tensor cosCast = npu_dtype_cast_impl_op_api(r1, self.scalar_type());
66+ at::Tensor sinCast = npu_dtype_cast_impl_op_api(r2, self.scalar_type());
67+ EXEC_NPU_CMD(aclnnRotaryPositionEmbedding, self, cosCast, sinCast, mode, result);
68+ } else {
69+ EXEC_NPU_CMD(aclnnRotaryPositionEmbedding, self, r1, r2, mode, result);
70+ }
71+ }
72+ return result;
78}73}
74+} // namespace op_api
Mop_plugin/ops/opapi/RotateQuantKernelNpuOpApi.cpp+110-85
@@ -29,98 +29,123 @@ constexpr int64_t DEFAULT_SCALE_ALG = 0LL;
29constexpr int64_t DEFAULT_AXIS = -1LL;29constexpr int64_t DEFAULT_AXIS = -1LL;
30}; // namespace30}; // namespace
31 31 
32-std::tuple<at::Tensor, at::Tensor> npu_rotate_quant(const at::Tensor &x, const at::Tensor &rotation,32+std::tuple<at::Tensor, at::Tensor> npu_rotate_quant(
33- const c10::optional<at::Tensor> &alpha, c10::optional<int64_t> dst_dtype, c10::optional<int64_t> axis,33+ const at::Tensor& x,
34- c10::optional<c10::string_view> round_mode, c10::optional<int64_t> scale_alg, c10::optional<double> dst_type_max,34+ const at::Tensor& rotation,
35+ const c10::optional<at::Tensor>& alpha,
36+ c10::optional<int64_t> dst_dtype,
37+ c10::optional<int64_t> axis,
38+ c10::optional<c10::string_view> round_mode,
39+ c10::optional<int64_t> scale_alg,
40+ c10::optional<double> dst_type_max,
35 c10::optional<bool> transpose_y) {41 c10::optional<bool> transpose_y) {
36- TORCH_CHECK(x.defined(), "Input tensor(x) must be defined" + OPS_ERROR(ErrCode::PARAM));42+ TORCH_CHECK(x.defined(), "Input tensor(x) must be defined" + OPS_ERROR(ErrCode::PARAM));
37- TORCH_CHECK(rotation.defined(), "Input tensor(rotation) must be defined" + OPS_ERROR(ErrCode::PARAM));43+ TORCH_CHECK(rotation.defined(), "Input tensor(rotation) must be defined" + OPS_ERROR(ErrCode::PARAM));
38- if (alpha.has_value()) {44+ if (alpha.has_value()) {
39- TORCH_CHECK(alpha->defined(), "Input tensor(alpha) must be defined when provided" + OPS_ERROR(ErrCode::PARAM));45+ TORCH_CHECK(alpha->defined(), "Input tensor(alpha) must be defined when provided" + OPS_ERROR(ErrCode::PARAM));
40- }46+ }
41 47 
42- auto dim_num = x.dim();48+ auto dim_num = x.dim();
43- int64_t dst_dtype_val = dst_dtype.value_or(static_cast<int64_t>(c10_npu::DType::INT8));49+ int64_t dst_dtype_val = dst_dtype.value_or(static_cast<int64_t>(c10_npu::DType::INT8));
44- int64_t axis_val = axis.value_or(DEFAULT_AXIS);50+ int64_t axis_val = axis.value_or(DEFAULT_AXIS);
45- bool transpose_y_val = transpose_y.value_or(false);51+ bool transpose_y_val = transpose_y.value_or(false);
46 52 
47- TORCH_CHECK(!transpose_y_val,53+ TORCH_CHECK(
48- "In the current CANN version, for aclnnRotateQuant, the parameter transpose_y only supports False. "54+ !transpose_y_val,
49- "Please set transpose_y=False." +55+ "In the current CANN version, for aclnnRotateQuant, the parameter transpose_y only supports False. "
56+ "Please set transpose_y=False." +
57+ OPS_ERROR(ErrCode::PARAM));
58+ 
59+ TORCH_CHECK(
60+ axis_val >= -dim_num && axis_val < dim_num,
61+ "Param (axis) is out of input dimension range" + OPS_ERROR(ErrCode::PARAM));
62+ 
63+ bool is_int4_packed = (dst_dtype_val == DTYPE_NUM_FOR_QUINT4X2);
64+ aclDataType dst_acl_dtype = is_int4_packed ? aclDataType::ACL_DT_UNDEFINED : c10_npu::GetAclDataType(dst_dtype_val);
65+ bool is_fp4 = (dst_acl_dtype == aclDataType::ACL_FLOAT4_E2M1);
66+ bool is_mx_type =
67+ (dst_acl_dtype == aclDataType::ACL_FLOAT4_E2M1 || dst_acl_dtype == aclDataType::ACL_FLOAT8_E5M2 ||
68+ dst_acl_dtype == aclDataType::ACL_FLOAT8_E4M3FN);
69+ 
70+ ASCEND_LOGI("[npu_rotate_quant]: Getting aclTensor y dtype by Parameter(dst_dtype): %ld", dst_dtype_val);
71+ 
72+ auto output_size = op_infer::array_to_small_vector(x.sizes());
73+ aclDataType y_acltype;
74+ at::Tensor output_y;
75+ 
76+ if (is_int4_packed) {
77+ y_acltype = aclDataType::ACL_INT32;
78+ TORCH_CHECK(
79+ output_size[dim_num - 1] % INT4_IN_INT32_NUM == 0,
80+ "Input shape last dim must be divisible by 8 when int4 quantization" + OPS_ERROR(ErrCode::PARAM));
81+ output_size[dim_num - 1] /= INT4_IN_INT32_NUM;
82+ output_y = npu_preparation::apply_tensor_without_format(output_size, c10::ScalarType::Int);
83+ } else if (is_fp4) {
84+ y_acltype = aclDataType::ACL_FLOAT4_E2M1;
85+ TORCH_CHECK(
86+ output_size[dim_num - 1] % FP4_IN_UINT8_NUM == 0,
87+ "The last dim input shape must be divisible by 2 if "
88+ "output dtype is torch_npu.float4_e2m1" +
50 OPS_ERROR(ErrCode::PARAM));89 OPS_ERROR(ErrCode::PARAM));
90+ output_size[dim_num - 1] /= FP4_IN_UINT8_NUM;
91+ output_y = npu_preparation::apply_tensor_without_format(output_size, c10::ScalarType::Byte);
92+ } else {
93+ y_acltype = dst_acl_dtype;
94+ TORCH_CHECK(
95+ y_acltype != aclDataType::ACL_DT_UNDEFINED,
96+ "Unsupported dst_dtype value: ",
97+ dst_dtype_val,
98+ OPS_ERROR(ErrCode::PARAM));
99+ at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
100+ TORCH_CHECK(
101+ scalar_dtype != at::ScalarType::Undefined,
102+ "Cannot convert aclDataType to ScalarType for dst_dtype: ",
103+ dst_dtype_val,
104+ OPS_ERROR(ErrCode::PARAM));
105+ output_y = npu_preparation::apply_tensor_without_format(output_size, c10::dtype(scalar_dtype));
106+ }
107+ ASCEND_LOGI(
108+ "[npu_rotate_quant]: Setting aclTensor y dtype to: %s", at_npu::native::AclDataTypeToString(y_acltype).c_str());
51 109 
52- TORCH_CHECK(axis_val >= -dim_num && axis_val < dim_num,110+ TensorWrapper y_wrapper = {output_y, y_acltype};
53- "Param (axis) is out of input dimension range" + OPS_ERROR(ErrCode::PARAM));
54 111 
55- bool is_int4_packed = (dst_dtype_val == DTYPE_NUM_FOR_QUINT4X2);112+ at::Tensor output_scale;
56- aclDataType dst_acl_dtype = is_int4_packed ? aclDataType::ACL_DT_UNDEFINED113+ aclDataType scale_acltype;
57- : c10_npu::GetAclDataType(dst_dtype_val);114+ if (is_mx_type) {
58- bool is_fp4 = (dst_acl_dtype == aclDataType::ACL_FLOAT4_E2M1);115+ auto mxscale_shape = op_infer::array_to_small_vector(x.sizes());
59- bool is_mx_type = (dst_acl_dtype == aclDataType::ACL_FLOAT4_E2M1 ||116+ mxscale_shape.emplace_back(ALIGN_NUM);
60- dst_acl_dtype == aclDataType::ACL_FLOAT8_E5M2 ||117+ int64_t axis_change = axis_val < 0 ? axis_val + dim_num : axis_val;
61- dst_acl_dtype == aclDataType::ACL_FLOAT8_E4M3FN);118+ int64_t dim_size = op_infer::CeilDiv(mxscale_shape[axis_change], BLOCK_SIZE_BASE_NUM);
119+ dim_size = (dim_size + ALIGN_NUM - 1) / ALIGN_NUM;
120+ mxscale_shape[axis_change] = dim_size;
121+ at::ScalarType scale_scalar_type = npu_preparation::convert_to_scalar_type(aclDataType::ACL_FLOAT8_E8M0);
122+ output_scale = npu_preparation::apply_tensor_without_format(mxscale_shape, c10::dtype(scale_scalar_type));
123+ scale_acltype = aclDataType::ACL_FLOAT8_E8M0;
124+ } else {
125+ int64_t m = x.size(0);
126+ output_scale = npu_preparation::apply_tensor_without_format({m}, c10::dtype(c10::ScalarType::Float));
127+ scale_acltype = aclDataType::ACL_FLOAT;
128+ }
129+ TensorWrapper scale_wrapper = {output_scale, scale_acltype};
62 130 
63- ASCEND_LOGI("[npu_rotate_quant]: Getting aclTensor y dtype by Parameter(dst_dtype): %ld", dst_dtype_val);131+ const at::Tensor& alpha_real = alpha.value_or(at::Tensor());
132+ double dst_type_max_val = dst_type_max.value_or(0.0);
133+ std::string round_mode_str = std::string(round_mode.value_or("rint"));
134+ char* round_mode_ptr = const_cast<char*>(round_mode_str.data());
135+ int64_t scale_alg_val = scale_alg.value_or(DEFAULT_SCALE_ALG);
64 136 
65- auto output_size = op_infer::array_to_small_vector(x.sizes());137+ EXEC_NPU_CMD(
66- aclDataType y_acltype;138+ aclnnRotateQuant,
67- at::Tensor output_y;139+ x,
68- 140+ rotation,
69- if (is_int4_packed) {141+ alpha_real,
70- y_acltype = aclDataType::ACL_INT32;142+ axis_val,
71- TORCH_CHECK(output_size[dim_num - 1] % INT4_IN_INT32_NUM == 0,143+ round_mode_ptr,
72- "Input shape last dim must be divisible by 8 when int4 quantization" + OPS_ERROR(ErrCode::PARAM));144+ scale_alg_val,
73- output_size[dim_num - 1] /= INT4_IN_INT32_NUM;145+ dst_type_max_val,
74- output_y = npu_preparation::apply_tensor_without_format(output_size, c10::ScalarType::Int);146+ transpose_y_val,
75- } else if (is_fp4) {147+ y_wrapper,
76- y_acltype = aclDataType::ACL_FLOAT4_E2M1;148+ scale_wrapper);
77- TORCH_CHECK(output_size[dim_num - 1] % FP4_IN_UINT8_NUM == 0,149+ return std::tuple<at::Tensor, at::Tensor>(output_y, output_scale);
78- "The last dim input shape must be divisible by 2 if "
79- "output dtype is torch_npu.float4_e2m1" +
80- OPS_ERROR(ErrCode::PARAM));
81- output_size[dim_num - 1] /= FP4_IN_UINT8_NUM;
82- output_y = npu_preparation::apply_tensor_without_format(output_size, c10::ScalarType::Byte);
83- } else {
84- y_acltype = dst_acl_dtype;
85- TORCH_CHECK(y_acltype != aclDataType::ACL_DT_UNDEFINED, "Unsupported dst_dtype value: ", dst_dtype_val,
86- OPS_ERROR(ErrCode::PARAM));
87- at::ScalarType scalar_dtype = npu_preparation::convert_to_scalar_type(y_acltype);
88- TORCH_CHECK(scalar_dtype != at::ScalarType::Undefined,
89- "Cannot convert aclDataType to ScalarType for dst_dtype: ", dst_dtype_val, OPS_ERROR(ErrCode::PARAM));
90- output_y = npu_preparation::apply_tensor_without_format(output_size, c10::dtype(scalar_dtype));
91- }
92- ASCEND_LOGI(
93- "[npu_rotate_quant]: Setting aclTensor y dtype to: %s", at_npu::native::AclDataTypeToString(y_acltype).c_str());
94- 
95- TensorWrapper y_wrapper = {output_y, y_acltype};
96- 
97- at::Tensor output_scale;
98- aclDataType scale_acltype;
99- if (is_mx_type) {
100- auto mxscale_shape = op_infer::array_to_small_vector(x.sizes());
101- mxscale_shape.emplace_back(ALIGN_NUM);
102- int64_t axis_change = axis_val < 0 ? axis_val + dim_num : axis_val;
103- int64_t dim_size = op_infer::CeilDiv(mxscale_shape[axis_change], BLOCK_SIZE_BASE_NUM);
104- dim_size = (dim_size + ALIGN_NUM - 1) / ALIGN_NUM;
105- mxscale_shape[axis_change] = dim_size;
106- at::ScalarType scale_scalar_type = npu_preparation::convert_to_scalar_type(aclDataType::ACL_FLOAT8_E8M0);
107- output_scale = npu_preparation::apply_tensor_without_format(mxscale_shape, c10::dtype(scale_scalar_type));
108- scale_acltype = aclDataType::ACL_FLOAT8_E8M0;
109- } else {
110- int64_t m = x.size(0);
111- output_scale = npu_preparation::apply_tensor_without_format({m}, c10::dtype(c10::ScalarType::Float));
112- scale_acltype = aclDataType::ACL_FLOAT;
113- }
114- TensorWrapper scale_wrapper = {output_scale, scale_acltype};
115- 
116- const at::Tensor &alpha_real = alpha.value_or(at::Tensor());
117- double dst_type_max_val = dst_type_max.value_or(0.0);
118- std::string round_mode_str = std::string(round_mode.value_or("rint"));
119- char *round_mode_ptr = const_cast<char *>(round_mode_str.data());
120- int64_t scale_alg_val = scale_alg.value_or(DEFAULT_SCALE_ALG);
121- 
122- EXEC_NPU_CMD(aclnnRotateQuant, x, rotation, alpha_real, axis_val, round_mode_ptr, scale_alg_val, dst_type_max_val,
123- transpose_y_val, y_wrapper, scale_wrapper);
124- return std::tuple<at::Tensor, at::Tensor>(output_y, output_scale);
125}150}
126} // namespace op_api151} // namespace op_api
Mop_plugin/ops/opapi/RoundDecimalsKernelNpuOpApi.cpp+9-8
@@ -20,13 +20,14 @@
20namespace op_api {20namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23-static void round_decimals_check(const at::Tensor& self, int64_t decimals)23+static void round_decimals_check(const at::Tensor& self, int64_t decimals) {
24-{24+ TORCH_CHECK(
25- TORCH_CHECK(isFloatingType(self.scalar_type()) ||25+ isFloatingType(self.scalar_type()) || self.scalar_type() == at::ScalarType::Int ||
26- self.scalar_type() == at::ScalarType::Int ||26+ self.scalar_type() == at::ScalarType::Long,
27- self.scalar_type() == at::ScalarType::Long,27+ "\"round_npu\" not implemented for '",
28- "\"round_npu\" not implemented for '", toString(self.scalar_type()), "'",28+ toString(self.scalar_type()),
29- OPS_ERROR(ErrCode::TYPE));29+ "'",
30+ OPS_ERROR(ErrCode::TYPE));
30}31}
31 32 
32at::Tensor& round_out(const at::Tensor& self, int64_t decimals, at::Tensor& result) {33at::Tensor& round_out(const at::Tensor& self, int64_t decimals, at::Tensor& result) {
@@ -51,4 +52,4 @@ at::Tensor& round_(at::Tensor& self, int64_t decimals) {
51 EXEC_NPU_CMD(aclnnInplaceRoundDecimals, self, decimals);52 EXEC_NPU_CMD(aclnnInplaceRoundDecimals, self, decimals);
52 return self;53 return self;
53}54}
54-}55+} // namespace op_api
Mop_plugin/ops/opapi/ScaledGroupMmKernelNpuOpApi.cpp+546-457
@@ -29,449 +29,356 @@ const static int64_t K_SPLIT = 2;
29using npu_preparation = at_npu::native::OpPreparation;29using npu_preparation = at_npu::native::OpPreparation;
30 30 
31static void check_dims(int64_t split_item, size_t num_x, size_t num_weight, size_t num_group_list) {31static void check_dims(int64_t split_item, size_t num_x, size_t num_weight, size_t num_group_list) {
32- TORCH_CHECK(num_x > 0 && num_weight > 0,32+ TORCH_CHECK(
33- "Invalid inputs: neither x nor weight could be empty." + OPS_ERROR(ErrCode::PARAM));33+ num_x > 0 && num_weight > 0, "Invalid inputs: neither x nor weight could be empty." + OPS_ERROR(ErrCode::PARAM));
34- TORCH_CHECK(split_item == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item == IN_SPLIT_OUT_NOT_SPLIT ||34+ TORCH_CHECK(
35- split_item == IN_NOT_SPLIT_OUT_SPLIT || split_item == IN_SPLIT_OUT_SPLIT,35+ split_item == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item == IN_SPLIT_OUT_NOT_SPLIT ||
36- "Invalid value of split_item [", split_item,36+ split_item == IN_NOT_SPLIT_OUT_SPLIT || split_item == IN_SPLIT_OUT_SPLIT,
37- "], which should only be one of 0/1/2/3." + OPS_ERROR(ErrCode::PARAM));37+ "Invalid value of split_item [",
38- if (split_item == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item == IN_SPLIT_OUT_NOT_SPLIT) {38+ split_item,
39- if (num_group_list > 0) {39+ "], which should only be one of 0/1/2/3." + OPS_ERROR(ErrCode::PARAM));
40- TORCH_CHECK(num_x == 1 && num_weight == num_group_list,40+ if (split_item == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item == IN_SPLIT_OUT_NOT_SPLIT) {
41- "Invalid inputs. "41+ if (num_group_list > 0) {
42- "When split_item = 0 or 1 and input group_list is not None, "42+ TORCH_CHECK(
43- "the following two conditions are supposed to be satisfied: "43+ num_x == 1 && num_weight == num_group_list,
44- "(1) length of x equals 1; (2) length of weight equals that of group_list. "44+ "Invalid inputs. "
45- "Actual lengths: x [",45+ "When split_item = 0 or 1 and input group_list is not None, "
46- num_x, "], weight [", num_weight,46+ "the following two conditions are supposed to be satisfied: "
47- "], "47+ "(1) length of x equals 1; (2) length of weight equals that of group_list. "
48- "group_list [",48+ "Actual lengths: x [",
49- num_group_list, "]." + OPS_ERROR(ErrCode::PARAM));49+ num_x,
50- } else {50+ "], weight [",
51- TORCH_CHECK(num_x == num_weight,51+ num_weight,
52- "When split_item = 0 or 1 and input group_list is None, "52+ "], "
53- "the num of x tensors must equal the num of weight tensors."53+ "group_list [",
54- "Actual lengths: x [",54+ num_group_list,
55- num_x, "], weight [", num_weight, "]." + OPS_ERROR(ErrCode::PARAM));55+ "]." + OPS_ERROR(ErrCode::PARAM));
56- }56+ } else {
57+ TORCH_CHECK(
58+ num_x == num_weight,
59+ "When split_item = 0 or 1 and input group_list is None, "
60+ "the num of x tensors must equal the num of weight tensors."
61+ "Actual lengths: x [",
62+ num_x,
63+ "], weight [",
64+ num_weight,
65+ "]." + OPS_ERROR(ErrCode::PARAM));
57 }66 }
67+ }
58}68}
59 69 
60static void create_new_tensor_multi_dim(70static void create_new_tensor_multi_dim(
61- std::vector<at::Tensor> &y, const at::Tensor &x_i, size_t n, c10::TensorOptions options) {71+ std::vector<at::Tensor>& y,
62- auto x_sizes = x_i.sizes();72+ const at::Tensor& x_i,
63- std::vector<int64_t> y_sizes(x_sizes.begin(), x_sizes.end());73+ size_t n,
64- y_sizes.at(x_sizes.size() - 1) = static_cast<int64_t>(n);74+ c10::TensorOptions options) {
75+ auto x_sizes = x_i.sizes();
76+ std::vector<int64_t> y_sizes(x_sizes.begin(), x_sizes.end());
77+ y_sizes.at(x_sizes.size() - 1) = static_cast<int64_t>(n);
65 78 
66- auto output_size = op_infer::array_to_small_vector(y_sizes);79+ auto output_size = op_infer::array_to_small_vector(y_sizes);
67- y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));80+ y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));
68}81}
69 82 
70-static void create_new_tensor(std::vector<at::Tensor> &y, size_t dim_m, size_t dim_n, c10::TensorOptions options) {83+static void create_new_tensor(std::vector<at::Tensor>& y, size_t dim_m, size_t dim_n, c10::TensorOptions options) {
71- auto output_size = op_infer::array_to_small_vector({dim_m, dim_n});84+ auto output_size = op_infer::array_to_small_vector({dim_m, dim_n});
72- y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));85+ y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));
73}86}
74 87 
75static void create_new_tensor_batch(88static void create_new_tensor_batch(
76- std::vector<at::Tensor> &y, size_t batch, size_t dim_m, size_t dim_n, c10::TensorOptions options) {89+ std::vector<at::Tensor>& y,
77- auto output_size = op_infer::array_to_small_vector({batch, dim_m, dim_n});90+ size_t batch,
78- y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));91+ size_t dim_m,
92+ size_t dim_n,
93+ c10::TensorOptions options) {
94+ auto output_size = op_infer::array_to_small_vector({batch, dim_m, dim_n});
95+ y.emplace_back(npu_preparation::apply_tensor_without_format(output_size, options));
79}96}
80 97 
81-static void calculate_dim_m(size_t &dim_m, size_t num_x, const at::TensorList x) {98+static void calculate_dim_m(size_t& dim_m, size_t num_x, const at::TensorList x) {
82- for (size_t i = 0; i < num_x; i++) {99+ for (size_t i = 0; i < num_x; i++) {
83- dim_m += x[i].sizes()[0];100+ dim_m += x[i].sizes()[0];
84- }101+ }
85}102}
86 103 
87-static bool is_weight_trans(const at::Tensor &tensor) {104+static bool is_weight_trans(const at::Tensor& tensor) {
88- int64_t dim1 = tensor.dim() - 1;105+ int64_t dim1 = tensor.dim() - 1;
89- int64_t dim2 = tensor.dim() - 2;106+ int64_t dim2 = tensor.dim() - 2;
90- return tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2);107+ return tensor.stride(dim2) == 1 && tensor.stride(dim1) == tensor.size(dim2);
91}108}
92 109 
110+at::Tensor _scaled_grouped_mm(
111+ const at::Tensor& mat_a,
112+ const at::Tensor& mat_b,
113+ const at::Tensor& scale_a,
114+ const at::Tensor& scale_b,
115+ const c10::optional<at::Tensor>& offs, // group_list -tensor
116+ const c10::optional<at::Tensor>& bias, // torch not support
117+ const c10::optional<at::Tensor>& scale_result, // torch not support
118+ c10::optional<c10::ScalarType> out_dtype,
119+ bool use_fast_accum) {
120+ // check A5
121+ TORCH_CHECK(
122+ c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,
123+ "This interface is supported only on the Ascend950 platform and after.",
124+ OPS_ERROR(ErrCode::PARAM));
93 125 
94-at::Tensor _scaled_grouped_mm(const at::Tensor &mat_a, const at::Tensor &mat_b, const at::Tensor &scale_a,126+ // 1. 基本参数校验
95- const at::Tensor &scale_b,127+ const int32_t ndim_a = mat_a.dim();
96- const c10::optional<at::Tensor> &offs, // group_list -tensor128+ const int32_t ndim_b = mat_b.dim();
97- const c10::optional<at::Tensor> &bias, // torch not support129+ const int32_t ndim_sa = scale_a.dim();
98- const c10::optional<at::Tensor> &scale_result, // torch not support130+ const int32_t ndim_sb = scale_b.dim();
99- c10::optional<c10::ScalarType> out_dtype, bool use_fast_accum) {
100 131 
101- // check A5132+ TORCH_CHECK(ndim_a == 2 || ndim_a == 3, "mat_a dimension must be 2D or 3D, actual dimension: ", ndim_a);
102- TORCH_CHECK(c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950,133+ TORCH_CHECK(ndim_b == 2 || ndim_b == 3, "mat_b dimension must be 2D or 3D, actual dimension: ", ndim_b);
103- "This interface is supported only on the Ascend950 platform and after.", OPS_ERROR(ErrCode::PARAM));
104 134 
105- // 1. 基本参数校验135+ // Check scale dimension based on scale dtype
106- const int32_t ndim_a = mat_a.dim();136+ bool is_fp8_a = (scale_a.scalar_type() == at::kFloat);
107- const int32_t ndim_b = mat_b.dim();137+ bool is_fp8_b = (scale_b.scalar_type() == at::kFloat);
108- const int32_t ndim_sa = scale_a.dim();138+ bool is_mx_a = (scale_a.scalar_type() == at::kFloat8_e8m0fnu);
109- const int32_t ndim_sb = scale_b.dim();139+ bool is_mx_b = (scale_b.scalar_type() == at::kFloat8_e8m0fnu);
110 140 
141+ // scale_a dimension check
142+ if (is_fp8_a) {
111 TORCH_CHECK(143 TORCH_CHECK(
112- ndim_a == 2 || ndim_a == 3,144+ ndim_sa == 1 || ndim_sa == 2, "scale_a dimension must be 1D or 2D for fp8, actual dimension: ", ndim_sa);
113- "mat_a dimension must be 2D or 3D, actual dimension: ", ndim_a145+ } else if (is_mx_a) {
114- );146+ TORCH_CHECK(ndim_sa == 2 || ndim_sa == 3, "scale_a dimension must be 2D or 3D for mx, actual dimension: ", ndim_sa);
147+ } else {
148+ TORCH_CHECK(false, "scale_a must be float32 or float8_e8m0fnu, but got ", scale_a.dtype());
149+ }
150+ 
151+ // scale_b dimension check
152+ if (is_fp8_b) {
115 TORCH_CHECK(153 TORCH_CHECK(
116- ndim_b == 2 || ndim_b == 3,154+ ndim_sb == 1 || ndim_sb == 2, "scale_b dimension must be 1D or 2D for fp8, actual dimension: ", ndim_sb);
117- "mat_b dimension must be 2D or 3D, actual dimension: ", ndim_b155+ } else if (is_mx_b) {
118- );156+ TORCH_CHECK(
157+ ndim_sb == 2 || ndim_sb == 3 || ndim_sb == 4,
158+ "scale_b dimension must be 2D/3D/4D for mx, actual dimension: ",
159+ ndim_sb);
160+ } else {
161+ TORCH_CHECK(false, "scale_b must be float32 or float8_e8m0fnu, but got ", scale_b.dtype());
162+ }
119 163 
120- // Check scale dimension based on scale dtype164+ // Check fp8 scale size based on KC/GB quantization mode
121- bool is_fp8_a = (scale_a.scalar_type() == at::kFloat);165+ // KC mode: scale_a(perTokenScale) 1D(M,) or 2D(g,M), scale_b(scale) 2D(g,N)
122- bool is_fp8_b = (scale_b.scalar_type() == at::kFloat);166+ // GB mode: scale_a(perTokenScale) 2D(M,ceil(K/128)), scale_b(scale) 3D(g,ceil(K/128),ceil(N/128))
123- bool is_mx_a = (scale_a.scalar_type() == at::kFloat8_e8m0fnu);167+ if (is_fp8_a && is_fp8_b) {
124- bool is_mx_b = (scale_b.scalar_type() == at::kFloat8_e8m0fnu);168+ if (ndim_a == 2) {
125- 169+ // KC mode: perTokenScale shape (M,)
126- // scale_a dimension check170+ int scale_multiplier = 1;
127- if (is_fp8_a) {171+ if (ndim_b == 2) {
128- TORCH_CHECK(ndim_sa == 1 || ndim_sa == 2,172+ scale_multiplier = offs->size(0);
129- "scale_a dimension must be 1D or 2D for fp8, actual dimension: ", ndim_sa);173+ }
130- } else if (is_mx_a) {174+ TORCH_CHECK(ndim_sa == 1, "scale_a must be 1D for 2D mat_a (KC mode), but got ", ndim_sa, "D");
131- TORCH_CHECK(ndim_sa == 2 || ndim_sa == 3,175+ TORCH_CHECK(scale_a.is_contiguous(), "scale_a must be contiguous");
132- "scale_a dimension must be 2D or 3D for mx, actual dimension: ", ndim_sa);176+ TORCH_CHECK(
177+ scale_a.size(0) == mat_a.size(0) * scale_multiplier,
178+ "scale_a size[0] must equal ",
179+ mat_a.size(0) * scale_multiplier);
133 } else {180 } else {
134- TORCH_CHECK(false, "scale_a must be float32 or float8_e8m0fnu, but got ", scale_a.dtype());181+ // KC mode: perTokenScale shape (g, M) or GB mode: (M, ceil(K/128))
182+ TORCH_CHECK(ndim_sa == 2, "scale_a must be 2D for 3D mat_a, but got ", ndim_sa, "D");
183+ TORCH_CHECK(scale_a.stride(1) == 1, "scale_a must be contiguous in last dim");
184+ TORCH_CHECK(scale_a.size(0) == mat_a.size(0), "scale_a size[0] must equal mat_a batch dim (G)");
185+ // size[1] can be M (KC) or ceil(K/128) (GB) - both valid
135 }186 }
136- 187+ if (ndim_b == 2) {
137- // scale_b dimension check188+ // KC mode: scale shape (g, N) where g=1 for 2D
138- if (is_fp8_b) {189+ TORCH_CHECK(ndim_sb == 1, "scale_b must be 1D for 2D mat_b, but got ", ndim_sb, "D");
139- TORCH_CHECK(ndim_sb == 1 || ndim_sb == 2,190+ TORCH_CHECK(scale_b.is_contiguous(), "scale_b must be contiguous in last dim");
140- "scale_b dimension must be 1D or 2D for fp8, actual dimension: ", ndim_sb);191+ int scale_multiplier = 1;
141- } else if (is_mx_b) {192+ if (ndim_a == 2) {
142- TORCH_CHECK(ndim_sb == 2 || ndim_sb == 3 || ndim_sb == 4,193+ scale_multiplier = offs->size(0);
143- "scale_b dimension must be 2D/3D/4D for mx, actual dimension: ", ndim_sb);194+ }
195+ int64_t expected_n = mat_b.size(1);
196+ TORCH_CHECK(scale_b.size(0) == expected_n * scale_multiplier, "scale_b size mismatch");
144 } else {197 } else {
145- TORCH_CHECK(false, "scale_b must be float32 or float8_e8m0fnu, but got ", scale_b.dtype());198+ // KC mode: scale shape (g, N) or GB mode: (g, ceil(K/128), ceil(N/128))
199+ TORCH_CHECK(ndim_sb == 2, "scale_b must be 2 for 2D mat_b, but got ", ndim_sb, "D");
200+ if (ndim_sb == 2) {
201+ // KC mode: (g, N)
202+ TORCH_CHECK(scale_b.stride(1) == 1, "scale_b must be contiguous in last dim");
203+ TORCH_CHECK(scale_b.size(0) == mat_b.size(0), "scale_b size[0] must equal mat_b batch dim (G)");
204+ TORCH_CHECK(scale_b.size(1) == mat_b.size(2), "scale_b size[1] must equal mat_b N dim");
205+ }
146 }206 }
207+ }
147 208 
148- // Check fp8 scale size based on KC/GB quantization mode209+ // Batch dimension checks
149- // KC mode: scale_a(perTokenScale) 1D(M,) or 2D(g,M), scale_b(scale) 2D(g,N)210+ const bool use_a_2d = (ndim_a == 2);
150- // GB mode: scale_a(perTokenScale) 2D(M,ceil(K/128)), scale_b(scale) 3D(g,ceil(K/128),ceil(N/128))211+ const bool use_b_2d = (ndim_b == 2);
151- if (is_fp8_a && is_fp8_b) {212+ if (!use_a_2d || !use_b_2d) {
152- if (ndim_a == 2) {213+ TORCH_CHECK(mat_a.size(-1) == mat_b.size(-2), "contraction dimension mismatch between mat_a and mat_b");
153- // KC mode: perTokenScale shape (M,)214+ }
154- int scale_multiplier = 1;
155- if (ndim_b == 2) {
156- scale_multiplier = offs->size(0);
157- }
158- TORCH_CHECK(ndim_sa == 1, "scale_a must be 1D for 2D mat_a (KC mode), but got ", ndim_sa, "D");
159- TORCH_CHECK(scale_a.is_contiguous(), "scale_a must be contiguous");
160- TORCH_CHECK(scale_a.size(0) == mat_a.size(0) * scale_multiplier, "scale_a size[0] must equal ", mat_a.size(0) * scale_multiplier);
161- } else {
162- // KC mode: perTokenScale shape (g, M) or GB mode: (M, ceil(K/128))
163- TORCH_CHECK(ndim_sa == 2, "scale_a must be 2D for 3D mat_a, but got ", ndim_sa, "D");
164- TORCH_CHECK(scale_a.stride(1) == 1, "scale_a must be contiguous in last dim");
165- TORCH_CHECK(scale_a.size(0) == mat_a.size(0), "scale_a size[0] must equal mat_a batch dim (G)");
166- // size[1] can be M (KC) or ceil(K/128) (GB) - both valid
167- }
168- if (ndim_b == 2) {
169- // KC mode: scale shape (g, N) where g=1 for 2D
170- TORCH_CHECK(ndim_sb == 1, "scale_b must be 1D for 2D mat_b, but got ", ndim_sb, "D");
171- TORCH_CHECK(scale_b.is_contiguous(), "scale_b must be contiguous in last dim");
172- int scale_multiplier = 1;
173- if (ndim_a == 2) {
174- scale_multiplier = offs->size(0);
175- }
176- int64_t expected_n = mat_b.size(1);
177- TORCH_CHECK(scale_b.size(0) == expected_n * scale_multiplier, "scale_b size mismatch");
178- } else {
179- // KC mode: scale shape (g, N) or GB mode: (g, ceil(K/128), ceil(N/128))
180- TORCH_CHECK(ndim_sb == 2 , "scale_b must be 2 for 2D mat_b, but got ", ndim_sb, "D");
181- if (ndim_sb == 2) {
182- // KC mode: (g, N)
183- TORCH_CHECK(scale_b.stride(1) == 1, "scale_b must be contiguous in last dim");
184- TORCH_CHECK(scale_b.size(0) == mat_b.size(0), "scale_b size[0] must equal mat_b batch dim (G)");
185- TORCH_CHECK(scale_b.size(1) == mat_b.size(2), "scale_b size[1] must equal mat_b N dim");
186- }
187- }
188- }
189 215 
190- // Batch dimension checks216+ // Unsupported features
191- const bool use_a_2d = (ndim_a == 2);217+ TORCH_CHECK(!bias.has_value(), "NPU _scaled_grouped_mm does not support bias yet");
192- const bool use_b_2d = (ndim_b == 2);218+ TORCH_CHECK(!scale_result.has_value(), "NPU _scaled_grouped_mm does not support scale_result yet");
193- if (!use_a_2d || !use_b_2d) {
194- TORCH_CHECK(
195- mat_a.size(-1) == mat_b.size(-2),
196- "contraction dimension mismatch between mat_a and mat_b"
197- );
198- }
199 219 
200- // Unsupported features220+ // Offsets validation
201- TORCH_CHECK(221+ const bool req_offsets = (use_a_2d || use_b_2d);
202- !bias.has_value(),222+ TORCH_CHECK(offs.has_value() == req_offsets, "offsets required when using 2D input tensor");
203- "NPU _scaled_grouped_mm does not support bias yet"
204- );
205- TORCH_CHECK(
206- !scale_result.has_value(),
207- "NPU _scaled_grouped_mm does not support scale_result yet"
208- );
209 223 
210- // Offsets validation224+ if (offs.has_value()) {
211- const bool req_offsets = (use_a_2d || use_b_2d);225+ TORCH_CHECK(offs->dim() == 1, "offsets tensor must be 1D");
212- TORCH_CHECK(226+ TORCH_CHECK(offs->dtype() == at::kInt, "offsets data type must be int32");
213- offs.has_value() == req_offsets,227+ }
214- "offsets required when using 2D input tensor"
215- );
216 228 
217- if (offs.has_value()) {229+ // Output type restriction
218- TORCH_CHECK(offs->dim() == 1, "offsets tensor must be 1D");230+ auto out_type = out_dtype.value_or(at::kBFloat16);
219- TORCH_CHECK(offs->dtype() == at::kInt, "offsets data type must be int32");231+ TORCH_CHECK(out_type == at::kBFloat16, "_scaled_grouped_mm on NPU only supports BF16 output type");
220- }
221 232 
222- // Output type restriction233+ // 3. mat_a -> x (TensorList)
223- auto out_type = out_dtype.value_or(at::kBFloat16);234+ std::vector<at::Tensor> x_vec;
224- TORCH_CHECK(235+ if (use_a_2d) {
225- out_type == at::kBFloat16,236+ x_vec.push_back(mat_a);
226- "_scaled_grouped_mm on NPU only supports BF16 output type"237+ } else {
227- );238+ x_vec.push_back(mat_a.reshape({-1, mat_a.size(-1)}));
239+ }
228 240 
229- // 3. mat_a -> x (TensorList)241+ at::TensorList x = at::TensorList(x_vec);
230- std::vector<at::Tensor> x_vec;242+ 
231- if (use_a_2d) {243+ // 4. mat_b -> weight (TensorList)
232- x_vec.push_back(mat_a);244+ std::vector<at::Tensor> weight_vec;
245+ if (use_b_2d) {
246+ // mat_b 2D [K, N]: NPU singleWeight 需要 3D [G, K, N]
247+ auto b_expanded = mat_b.unsqueeze(0).expand({1, -1, -1}).contiguous();
248+ weight_vec.push_back(b_expanded);
249+ } else {
250+ // mat_b 3D [G, K, N]:
251+ weight_vec.push_back(mat_b);
252+ }
253+ 
254+ at::TensorList weight = at::TensorList(weight_vec);
255+ 
256+ // 5. scale_a -> per_token_scale (TensorList)
257+ std::vector<at::Tensor> per_token_scale_vec;
258+ per_token_scale_vec.push_back(scale_a);
259+ 
260+ c10::optional<at::TensorList> per_token_scale = c10::optional<at::TensorList>(at::TensorList(per_token_scale_vec));
261+ 
262+ // 6. scale_b -> scale (TensorList)
263+ std::vector<at::Tensor> scale_vec;
264+ if (scale_b.dim() == 2 || scale_b.dim() == 4) {
265+ scale_vec.push_back(scale_b);
266+ } else {
267+ scale_vec.push_back(scale_b.unsqueeze(0).expand({1, -1, -1, -1}).contiguous());
268+ }
269+ 
270+ c10::optional<at::TensorList> scale = c10::optional<at::TensorList>(at::TensorList(scale_vec));
271+ 
272+ // 7. offs -> group_list (c10::optional<at::Tensor>)
273+ c10::optional<at::Tensor> group_list = c10::nullopt;
274+ if (offs.has_value()) {
275+ group_list = offs->to(at::kLong);
276+ }
277+ 
278+ // 8. split_item / group_type / group_list_type / act_type
279+ int64_t split_item_val = IN_NOT_SPLIT_OUT_SPLIT; // 2
280+ c10::optional<int64_t> split_item = split_item_val;
281+ 
282+ // 根据入参自动推导 group_type
283+ // group_type 含义（矩阵乘 C[m,n]=A[m,k]×B[k,n]）：
284+ // -1 (DEFAULT_SPLIT): 不分组 - mat_a/mat_b 均为 3D 且 batch 维度一一对应
285+ // 0 (M_SPLIT): m 轴分组 - 单输入 mat_a(2D) 按 m 维度分割，对应多个 weight
286+ // 2 (K_SPLIT): k 轴分组 - 单 weight 被多个 group 共享
287+ c10::optional<int64_t> group_type = DEFAULT_SPLIT; // 默认不分组
288+ if (use_b_2d) {
289+ // mat_b 为 2D [K,N]：单 weight 共享模式，按 k 轴分组
290+ group_type = K_SPLIT;
291+ } else if (use_a_2d) {
292+ // mat_a 为 2D [M,K]，mat_b 为 3D [G,K,N]：单输入按 m 轴分割
293+ group_type = M_SPLIT;
294+ } else {
295+ // mat_a 为 3D [G_a,M,K]，mat_b 为 3D [G_b,K,N]
296+ if (mat_b.size(0) == 1) {
297+ // mat_b 只有 1 个 weight，被多个 group 共享，按 k 轴分组
298+ group_type = K_SPLIT;
299+ } else if (mat_a.size(0) == mat_b.size(0)) {
300+ // batch 维度一一对应，不分组
301+ group_type = DEFAULT_SPLIT;
233 } else {302 } else {
234- x_vec.push_back(mat_a.reshape({-1, mat_a.size(-1)}));303+ // mat_b 有多个 weight，按 m 轴分组
304+ group_type = M_SPLIT;
235 }305 }
306+ }
236 307 
237- at::TensorList x = at::TensorList(x_vec);308+ TORCH_CHECK(
309+ group_type != K_SPLIT,
310+ "K_SPLIT (group_type=2) is not supported yet. "
311+ "This occurs when mat_b is 2D or mat_b has only 1 weight shared by multiple groups. "
312+ "Current mat_a size: ",
313+ mat_a.sizes(),
314+ ", mat_b size: ",
315+ mat_b.sizes());
238 316 
239- // 4. mat_b -> weight (TensorList)317+ c10::optional<int64_t> group_list_type = 0;
240- std::vector<at::Tensor> weight_vec;318+ c10::optional<int64_t> act_type = 0;
241- if (use_b_2d) {
242- // mat_b 2D [K, N]: NPU singleWeight 需要 3D [G, K, N]
243- auto b_expanded = mat_b.unsqueeze(0).expand({1, -1, -1}).contiguous();
244- weight_vec.push_back(b_expanded);
245- } else {
246- // mat_b 3D [G, K, N]:
247- weight_vec.push_back(mat_b);
248- }
249 319 
250- at::TensorList weight = at::TensorList(weight_vec);320+ // 9. out_dtype -> output_dtype (ACL int64 格式)
321+ c10::optional<int64_t> output_dtype = static_cast<int64_t>(out_type);
251 322 
252- // 5. scale_a -> per_token_scale (TensorList)323+ // 10. 其他可选参数 (默认空)
253- std::vector<at::Tensor> per_token_scale_vec;324+ c10::optional<at::TensorList> bias_tl = c10::nullopt;
254- per_token_scale_vec.push_back(scale_a);325+ c10::optional<at::TensorList> offset_tl = c10::nullopt;
326+ c10::optional<at::TensorList> antiquant_scale_tl = c10::nullopt;
327+ c10::optional<at::TensorList> antiquant_offset_tl = c10::nullopt;
328+ c10::optional<at::TensorList> activation_input_tl = c10::nullopt;
329+ c10::optional<at::TensorList> activation_quant_scale_tl = c10::nullopt;
330+ c10::optional<at::TensorList> activation_quant_offset_tl = c10::nullopt;
331+ c10::OptionalIntArrayRef tuning_config = c10::OptionalIntArrayRef{};
332+ c10::optional<int64_t> x_dtype = c10::nullopt;
333+ c10::optional<int64_t> weight_dtype = c10::nullopt;
334+ c10::optional<int64_t> scale_dtype = c10::nullopt;
335+ c10::optional<int64_t> per_token_scale_dtype = c10::nullopt;
255 336 
256- c10::optional<at::TensorList> per_token_scale = c10::optional<at::TensorList>(at::TensorList(per_token_scale_vec));337+ /////////////////////////////////////////////////npu_group_matmul//////////////////////////////////////////////
257- 
258- // 6. scale_b -> scale (TensorList)
259- std::vector<at::Tensor> scale_vec;
260- if (scale_b.dim() == 2 || scale_b.dim() == 4) {
261- scale_vec.push_back(scale_b);
262- } else {
263- scale_vec.push_back(scale_b.unsqueeze(0).expand({1, -1, -1, -1}).contiguous());
264- }
265- 
266- c10::optional<at::TensorList> scale = c10::optional<at::TensorList>(at::TensorList(scale_vec));
267- 
268- // 7. offs -> group_list (c10::optional<at::Tensor>)
269- c10::optional<at::Tensor> group_list = c10::nullopt;
270- if (offs.has_value()) {
271- group_list = offs->to(at::kLong);
272- }
273- 
274- // 8. split_item / group_type / group_list_type / act_type
275- int64_t split_item_val = IN_NOT_SPLIT_OUT_SPLIT; // 2
276- c10::optional<int64_t> split_item = split_item_val;
277- 
278- 
279- // 根据入参自动推导 group_type
280- // group_type 含义（矩阵乘 C[m,n]=A[m,k]×B[k,n]）：
281- // -1 (DEFAULT_SPLIT): 不分组 - mat_a/mat_b 均为 3D 且 batch 维度一一对应
282- // 0 (M_SPLIT): m 轴分组 - 单输入 mat_a(2D) 按 m 维度分割，对应多个 weight
283- // 2 (K_SPLIT): k 轴分组 - 单 weight 被多个 group 共享
284- c10::optional<int64_t> group_type = DEFAULT_SPLIT; // 默认不分组
285- if (use_b_2d) {
286- // mat_b 为 2D [K,N]：单 weight 共享模式，按 k 轴分组
287- group_type = K_SPLIT;
288- } else if (use_a_2d) {
289- // mat_a 为 2D [M,K]，mat_b 为 3D [G,K,N]：单输入按 m 轴分割
290- group_type = M_SPLIT;
291- } else {
292- // mat_a 为 3D [G_a,M,K]，mat_b 为 3D [G_b,K,N]
293- if (mat_b.size(0) == 1) {
294- // mat_b 只有 1 个 weight，被多个 group 共享，按 k 轴分组
295- group_type = K_SPLIT;
296- } else if (mat_a.size(0) == mat_b.size(0)) {
297- // batch 维度一一对应，不分组
298- group_type = DEFAULT_SPLIT;
299- } else {
300- // mat_b 有多个 weight，按 m 轴分组
301- group_type = M_SPLIT;
302- }
303- }
304- 
305- TORCH_CHECK(group_type != K_SPLIT,
306- "K_SPLIT (group_type=2) is not supported yet. "
307- "This occurs when mat_b is 2D or mat_b has only 1 weight shared by multiple groups. "
308- "Current mat_a size: ", mat_a.sizes(), ", mat_b size: ", mat_b.sizes());
309- 
310- c10::optional<int64_t> group_list_type = 0;
311- c10::optional<int64_t> act_type = 0;
312- 
313- // 9. out_dtype -> output_dtype (ACL int64 格式)
314- c10::optional<int64_t> output_dtype = static_cast<int64_t>(out_type);
315- 
316- // 10. 其他可选参数 (默认空)
317- c10::optional<at::TensorList> bias_tl = c10::nullopt;
318- c10::optional<at::TensorList> offset_tl = c10::nullopt;
319- c10::optional<at::TensorList> antiquant_scale_tl = c10::nullopt;
320- c10::optional<at::TensorList> antiquant_offset_tl = c10::nullopt;
321- c10::optional<at::TensorList> activation_input_tl = c10::nullopt;
322- c10::optional<at::TensorList> activation_quant_scale_tl = c10::nullopt;
323- c10::optional<at::TensorList> activation_quant_offset_tl = c10::nullopt;
324- c10::OptionalIntArrayRef tuning_config = c10::OptionalIntArrayRef{};
325- c10::optional<int64_t> x_dtype = c10::nullopt;
326- c10::optional<int64_t> weight_dtype = c10::nullopt;
327- c10::optional<int64_t> scale_dtype = c10::nullopt;
328- c10::optional<int64_t> per_token_scale_dtype = c10::nullopt;
329- 
330- /////////////////////////////////////////////////npu_group_matmul//////////////////////////////////////////////
331 338 
339+ TORCH_CHECK(
340+ group_type.has_value(), "Requires manual passing group_type, current is None.", OPS_ERROR(ErrCode::VALUE));
341+ int64_t group_type_value = group_type.value();
342+ TORCH_CHECK(
343+ group_type_value == DEFAULT_SPLIT || group_type_value == M_SPLIT || group_type_value == K_SPLIT,
344+ "Use Tensor input with current cann version, "
345+ "The group type must be -1, 0 or 2, but now is [",
346+ group_type_value,
347+ "]",
348+ OPS_ERROR(ErrCode::VALUE));
349+ static const bool is_grouped_matmul_V4_available = check_aclnn_kernel_available("aclnnGroupedMatmulV4");
350+ if (C10_UNLIKELY(!is_grouped_matmul_V4_available)) {
332 TORCH_CHECK(351 TORCH_CHECK(
333- group_type.has_value(), "Requires manual passing group_type, current is None.", OPS_ERROR(ErrCode::VALUE));352+ !group_list.has_value(),
334- int64_t group_type_value = group_type.value();353+ "group_list don't support Tensor input with current cann version. "
335- TORCH_CHECK(group_type_value == DEFAULT_SPLIT || group_type_value == M_SPLIT || group_type_value == K_SPLIT,354+ "Please update cann version to 8.0.RC3 or higher, or use List[int] as input.",
336- "Use Tensor input with current cann version, "355+ OPS_ERROR(ErrCode::VALUE));
337- "The group type must be -1, 0 or 2, but now is [",
338- group_type_value, "]", OPS_ERROR(ErrCode::VALUE));
339- static const bool is_grouped_matmul_V4_available = check_aclnn_kernel_available("aclnnGroupedMatmulV4");
340- if (C10_UNLIKELY(!is_grouped_matmul_V4_available)) {
341- TORCH_CHECK(!group_list.has_value(),
342- "group_list don't support Tensor input with current cann version. "
343- "Please update cann version to 8.0.RC3 or higher, or use List[int] as input.",
344- OPS_ERROR(ErrCode::VALUE));
345- auto num_x = x.size();
346- auto num_weight = weight.size();
347- auto group_list_real = at::IntArrayRef{};
348- size_t num_group_list = 0;
349- int64_t split_item_value = split_item.value_or(0);
350- check_dims(split_item_value, num_x, num_weight, num_group_list);
351- 
352- std::vector<at::Tensor> y;
353- c10::TensorOptions options = x[0].options().dtype(output_dtype.has_value()
354- ? npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(output_dtype.value()))
355- : x[0].scalar_type());
356- 
357- if (split_item_value == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item_value == IN_SPLIT_OUT_NOT_SPLIT) {
358- y.reserve(num_x);
359- for (size_t i = 0; i < num_x; i++) {
360- create_new_tensor_multi_dim(y, x[i], weight[i].size(1), options);
361- }
362- } else if (split_item_value == IN_NOT_SPLIT_OUT_SPLIT || split_item_value == IN_SPLIT_OUT_SPLIT) {
363- if (num_x > 1) {
364- size_t dim_m = 0;
365- calculate_dim_m(dim_m, num_x, x);
366- create_new_tensor(y, dim_m, weight[0].sizes()[1], options);
367- } else if (num_x == 1) {
368- create_new_tensor(y, x[0].sizes()[0], weight[0].sizes()[1], options);
369- }
370- }
371- at::TensorList result = at::TensorList(y);
372- 
373- auto bias_real = bias_tl.value_or(at::TensorList());
374- auto scale_real = scale.value_or(at::TensorList());
375- auto offset_real = offset_tl.value_or(at::TensorList());
376- auto antiquant_scale_real = antiquant_scale_tl.value_or(at::TensorList());
377- auto antiquant_offset_real = antiquant_offset_tl.value_or(at::TensorList());
378- EXEC_NPU_CMD(aclnnGroupedMatmul, x, weight, bias_real, scale_real, offset_real, antiquant_scale_real,
379- antiquant_offset_real, group_list_real, split_item_value, result);
380- 
381- return y[0];
382- }
383- 
384 auto num_x = x.size();356 auto num_x = x.size();
385- bool singleWeight = weight.size() == 1 && weight[0].sizes().size() == 3;357+ auto num_weight = weight.size();
386- auto num_weight = singleWeight ? static_cast<size_t>(weight[0].size(0)) : static_cast<size_t>(weight.size());358+ auto group_list_real = at::IntArrayRef{};
387- auto group_list_real = group_list.value_or(at::Tensor());359+ size_t num_group_list = 0;
388- auto num_group_list = group_list_real.size(0);
389 int64_t split_item_value = split_item.value_or(0);360 int64_t split_item_value = split_item.value_or(0);
390 check_dims(split_item_value, num_x, num_weight, num_group_list);361 check_dims(split_item_value, num_x, num_weight, num_group_list);
391 362 
392 std::vector<at::Tensor> y;363 std::vector<at::Tensor> y;
393- c10::TensorOptions options = x[0].options().dtype(output_dtype.has_value()364+ c10::TensorOptions options = x[0].options().dtype(
365+ output_dtype.has_value()
394 ? npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(output_dtype.value()))366 ? npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(output_dtype.value()))
395 : x[0].scalar_type());367 : x[0].scalar_type());
396 368 
397- size_t dim_num_w = weight[0].sizes().size();
398- size_t n0 = static_cast<size_t>(weight[0].size(dim_num_w - 1));
399- // weight is trans or not
400- bool weight_trans = is_weight_trans(weight[0]);
401-#if VERSION_BETWEEN(V2R1, V2R7)
402- bool mxfp4_valid = x_dtype.has_value() && weight_dtype.has_value() &&
403- (x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
404- x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) &&
405- (weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2) ||
406- weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1));
407-#endif
408-#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)
409- bool mxfp4_valid = false;
410- if (x_dtype.has_value()) {
411- mxfp4_valid = (x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
412- x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
413- } else {
414- mxfp4_valid = x[0].scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
415- }
416- if (weight_dtype.has_value()) {
417- mxfp4_valid = mxfp4_valid &&
418- (weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
419- weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
420- } else {
421- mxfp4_valid = mxfp4_valid && weight[0].scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
422- }
423-#endif
424- size_t n_new = (mxfp4_valid && !weight_trans) ? (n0 * FP4_IN_INT8) : n0;
425- if (mxfp4_valid) {
426- TORCH_CHECK(x[0].size(1) != 1, "In mxfp4, dim K should not be 2.", OPS_ERROR(ErrCode::VALUE));
427- }
428 if (split_item_value == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item_value == IN_SPLIT_OUT_NOT_SPLIT) {369 if (split_item_value == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item_value == IN_SPLIT_OUT_NOT_SPLIT) {
429- if (num_group_list > 0) {370+ y.reserve(num_x);
430- y.reserve(num_group_list);371+ for (size_t i = 0; i < num_x; i++) {
431- int64_t glr_value_0 = group_list_real[0].item<int64_t>();372+ create_new_tensor_multi_dim(y, x[i], weight[i].size(1), options);
432- TORCH_CHECK(glr_value_0 >= 0, "group_list[0] should be larger than or equal to 0, but now is ", glr_value_0,373+ }
433- "." + OPS_ERROR(ErrCode::VALUE));
434- create_new_tensor(y, glr_value_0, n0, options);
435- int64_t glr_value_pre = glr_value_0;
436- for (int i = 1; i < num_group_list; i++) {
437- int64_t glr_value_cur = group_list_real[i].item<int64_t>();
438- TORCH_CHECK(glr_value_cur - glr_value_pre >= 0, "group_list[", i, "] - group_list[", i - 1,
439- "] should be larger than or equal to 0, but now is ", glr_value_cur - glr_value_pre,
440- "." + OPS_ERROR(ErrCode::VALUE));
441- size_t ni = singleWeight ? n0 : weight[i].size(dim_num_w - 1);
442- create_new_tensor(y, glr_value_cur - glr_value_pre, ni, options);
443- glr_value_pre = glr_value_cur;
444- }
445- } else {
446- y.reserve(num_x);
447- for (size_t i = 0; i < num_x; i++) {
448- size_t ni = singleWeight ? n0 : weight[i].size(dim_num_w - 1);
449- create_new_tensor_multi_dim(y, x[i], ni, options);
450- }
451- } // 校验NO_SPLIT时为特殊场景（groupList为空）或num_x > 1
452 } else if (split_item_value == IN_NOT_SPLIT_OUT_SPLIT || split_item_value == IN_SPLIT_OUT_SPLIT) {374 } else if (split_item_value == IN_NOT_SPLIT_OUT_SPLIT || split_item_value == IN_SPLIT_OUT_SPLIT) {
453- if (num_x > 1) {375+ if (num_x > 1) {
454- size_t dim_m = 0;376+ size_t dim_m = 0;
455- for (size_t i = 0; i < num_x; i++) {377+ calculate_dim_m(dim_m, num_x, x);
456- dim_m += static_cast<size_t>(x[i].size(0));378+ create_new_tensor(y, dim_m, weight[0].sizes()[1], options);
457- }379+ } else if (num_x == 1) {
458- weight[0].dtype() == at::ScalarType::Int ? create_new_tensor(y, dim_m, n0 * INT4_NUMS_IN_INT32, options)380+ create_new_tensor(y, x[0].sizes()[0], weight[0].sizes()[1], options);
459- : create_new_tensor(y, dim_m, n_new, options);381+ }
460- } else if (num_x == 1) {
461- if (group_type_value == K_SPLIT) {
462- TORCH_CHECK(num_weight == 1,
463- "When group_list is 2(K_SPLIT) and split_item is 2/3, the length of weight must equal x.");
464- weight[0].dtype() == at::ScalarType::Int
465- ? create_new_tensor_batch(y, num_group_list, x[0].size(0), n0 * INT4_NUMS_IN_INT32, options)
466- : create_new_tensor_batch(y, num_group_list, x[0].size(0), n_new, options);
467- } else {
468- (weight[0].dtype() == at::ScalarType::Int ||
469- (weight[0].dtype() == at::ScalarType::Float && weight[0].dtype() != x[0].dtype())) &&
470- (!weight_trans)
471- ? create_new_tensor(y, x[0].size(0), n0 * INT4_NUMS_IN_INT32, options)
472- : create_new_tensor(y, x[0].size(0), n_new, options);
473- }
474- }
475 }382 }
476 at::TensorList result = at::TensorList(y);383 at::TensorList result = at::TensorList(y);
477 384 
@@ -480,73 +387,255 @@ at::Tensor _scaled_grouped_mm(const at::Tensor &mat_a, const at::Tensor &mat_b,
480 auto offset_real = offset_tl.value_or(at::TensorList());387 auto offset_real = offset_tl.value_or(at::TensorList());
481 auto antiquant_scale_real = antiquant_scale_tl.value_or(at::TensorList());388 auto antiquant_scale_real = antiquant_scale_tl.value_or(at::TensorList());
482 auto antiquant_offset_real = antiquant_offset_tl.value_or(at::TensorList());389 auto antiquant_offset_real = antiquant_offset_tl.value_or(at::TensorList());
483- auto per_token_scale_real = per_token_scale.value_or(at::TensorList());390+ EXEC_NPU_CMD(
484- auto activation_input_real = activation_input_tl.value_or(at::TensorList());391+ aclnnGroupedMatmul,
485- auto activation_quant_scale_real = activation_quant_scale_tl.value_or(at::TensorList());392+ x,
486- auto activation_quant_offset_real = activation_quant_offset_tl.value_or(at::TensorList());393+ weight,
487- auto act_out = at::TensorList();394+ bias_real,
488- auto dynamic_quant_scale_out = at::TensorList();395+ scale_real,
489- int64_t group_list_type_value = group_list_type.value_or(0);396+ offset_real,
490- int64_t act_type_value = act_type.value_or(0);397+ antiquant_scale_real,
491- auto tuning_config_real = tuning_config.value_or(at::IntArrayRef{});398+ antiquant_offset_real,
399+ group_list_real,
400+ split_item_value,
401+ result);
492 402 
493- TensorListWrapper x_wrapper = {x,
494- x_dtype.has_value() ? c10_npu::GetAclDataType(x_dtype.value())
495- : npu_preparation::convert_to_acl_data_type(x[0].scalar_type())};
496- TensorListWrapper weight_wrapper = {weight,
497- weight_dtype.has_value() ? c10_npu::GetAclDataType(weight_dtype.value())
498- : npu_preparation::convert_to_acl_data_type(weight[0].scalar_type())};
499- TensorListWrapper scale_wrapper = {scale_real,
500- scale_dtype.has_value()
501- ? c10_npu::GetAclDataType(scale_dtype.value())
502- : (scale_real.empty() ? aclDataType::ACL_UINT64
503- : npu_preparation::convert_to_acl_data_type(scale_real[0].scalar_type()))};
504- TensorListWrapper per_token_scale_wrapper = {per_token_scale_real,
505- per_token_scale_dtype.has_value()
506- ? c10_npu::GetAclDataType(per_token_scale_dtype.value())
507- : (per_token_scale_real.empty()
508- ? aclDataType::ACL_FLOAT
509- : npu_preparation::convert_to_acl_data_type(per_token_scale_real[0].scalar_type()))};
510- TensorListWrapper antiquant_scale_wrapper = {antiquant_scale_real,
511- antiquant_scale_real.empty()
512- ? aclDataType::ACL_FLOAT16
513- : (antiquant_scale_real[0].scalar_type() == at::ScalarType::Byte
514- ? aclDataType::ACL_FLOAT8_E8M0
515- : npu_preparation::convert_to_acl_data_type(antiquant_scale_real[0].scalar_type()))};
516- 
517- int64_t weight_format = at_npu::native::custom_ops::get_npu_format(weight[0]);
518- const bool is_weight_nz = (weight_format == ACL_FORMAT_FRACTAL_NZ) ||
519- (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2) || (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_4) ||
520- (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_16);
521- if (is_weight_nz) {
522- static const bool is_weight_nz_available = check_aclnn_kernel_available("aclnnGroupedMatmulWeightNz");
523- TORCH_CHECK(is_weight_nz_available,
524- "Format of weight in npu_grouped_matmul is FRACTAL_NZ, current CANN version "
525- "do not support with this format. Please try to update the version of CANN." +
526- OPS_ERROR(ErrCode::PARAM));
527- int64_t quant_per_group_size = 0;
528- EXEC_NPU_CMD(aclnnGroupedMatmulWeightNz, x_wrapper, weight_wrapper, bias_real, scale_wrapper, offset_real,
529- antiquant_scale_wrapper, antiquant_offset_real, per_token_scale_wrapper, group_list_real,
530- activation_input_real, activation_quant_scale_real, activation_quant_offset_real, split_item_value,
531- group_type_value, group_list_type_value, act_type_value, tuning_config_real, quant_per_group_size, result,
532- act_out, dynamic_quant_scale_out);
533- return y[0];
534- }
535- static const bool is_grouped_matmul_V5_available = check_aclnn_kernel_available("aclnnGroupedMatmulV5");
536- static const bool dtypeValid = x[0].scalar_type() != at::ScalarType::Float8_e5m2 &&
537- x[0].scalar_type() != at::ScalarType::Float8_e4m3fn && !x_dtype.has_value() && !weight_dtype.has_value();
538- if (!is_grouped_matmul_V5_available || !dtypeValid || mxfp4_valid) {
539- EXEC_NPU_CMD(aclnnGroupedMatmulV4, x_wrapper, weight_wrapper, bias_real, scale_wrapper, offset_real,
540- antiquant_scale_real, antiquant_offset_real, per_token_scale_wrapper, group_list_real,
541- activation_input_real, activation_quant_scale_real, activation_quant_offset_real, split_item_value,
542- group_type_value, group_list_type_value, act_type_value, result, act_out, dynamic_quant_scale_out);
543- } else {
544- EXEC_NPU_CMD(aclnnGroupedMatmulV5, x_wrapper, weight_wrapper, bias_real, scale_wrapper, offset_real,
545- antiquant_scale_real, antiquant_offset_real, per_token_scale_wrapper, group_list_real,
546- activation_input_real, activation_quant_scale_real, activation_quant_offset_real, split_item_value,
547- group_type_value, group_list_type_value, act_type_value, tuning_config_real, result, act_out,
548- dynamic_quant_scale_out);
549- }
550 return y[0];403 return y[0];
404+ }
405+ 
406+ auto num_x = x.size();
407+ bool singleWeight = weight.size() == 1 && weight[0].sizes().size() == 3;
408+ auto num_weight = singleWeight ? static_cast<size_t>(weight[0].size(0)) : static_cast<size_t>(weight.size());
409+ auto group_list_real = group_list.value_or(at::Tensor());
410+ auto num_group_list = group_list_real.size(0);
411+ int64_t split_item_value = split_item.value_or(0);
412+ check_dims(split_item_value, num_x, num_weight, num_group_list);
413+ 
414+ std::vector<at::Tensor> y;
415+ c10::TensorOptions options = x[0].options().dtype(
416+ output_dtype.has_value() ? npu_preparation::convert_to_scalar_type(c10_npu::GetAclDataType(output_dtype.value()))
417+ : x[0].scalar_type());
418+ 
419+ size_t dim_num_w = weight[0].sizes().size();
420+ size_t n0 = static_cast<size_t>(weight[0].size(dim_num_w - 1));
421+ // weight is trans or not
422+ bool weight_trans = is_weight_trans(weight[0]);
423+#if VERSION_BETWEEN(V2R1, V2R7)
424+ bool mxfp4_valid = x_dtype.has_value() && weight_dtype.has_value() &&
425+ (x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
426+ x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2)) &&
427+ (weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2) ||
428+ weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1));
429+#endif
430+#if VERSION_BETWEEN(V2R8, VERSION_NEWEST)
431+ bool mxfp4_valid = false;
432+ if (x_dtype.has_value()) {
433+ mxfp4_valid =
434+ (x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
435+ x_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
436+ } else {
437+ mxfp4_valid = x[0].scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
438+ }
439+ if (weight_dtype.has_value()) {
440+ mxfp4_valid = mxfp4_valid &&
441+ (weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E2M1) ||
442+ weight_dtype.value() == static_cast<int64_t>(c10_npu::DType::FLOAT4_E1M2));
443+ } else {
444+ mxfp4_valid = mxfp4_valid && weight[0].scalar_type() == at::ScalarType::Float4_e2m1fn_x2;
445+ }
446+#endif
447+ size_t n_new = (mxfp4_valid && !weight_trans) ? (n0 * FP4_IN_INT8) : n0;
448+ if (mxfp4_valid) {
449+ TORCH_CHECK(x[0].size(1) != 1, "In mxfp4, dim K should not be 2.", OPS_ERROR(ErrCode::VALUE));
450+ }
451+ if (split_item_value == IN_NOT_SPLIT_OUT_NOT_SPLIT || split_item_value == IN_SPLIT_OUT_NOT_SPLIT) {
452+ if (num_group_list > 0) {
453+ y.reserve(num_group_list);
454+ int64_t glr_value_0 = group_list_real[0].item<int64_t>();
455+ TORCH_CHECK(
456+ glr_value_0 >= 0,
457+ "group_list[0] should be larger than or equal to 0, but now is ",
458+ glr_value_0,
459+ "." + OPS_ERROR(ErrCode::VALUE));
460+ create_new_tensor(y, glr_value_0, n0, options);
461+ int64_t glr_value_pre = glr_value_0;
462+ for (int i = 1; i < num_group_list; i++) {
463+ int64_t glr_value_cur = group_list_real[i].item<int64_t>();
464+ TORCH_CHECK(
465+ glr_value_cur - glr_value_pre >= 0,
466+ "group_list[",
467+ i,
468+ "] - group_list[",
469+ i - 1,
470+ "] should be larger than or equal to 0, but now is ",
471+ glr_value_cur - glr_value_pre,
472+ "." + OPS_ERROR(ErrCode::VALUE));
473+ size_t ni = singleWeight ? n0 : weight[i].size(dim_num_w - 1);
474+ create_new_tensor(y, glr_value_cur - glr_value_pre, ni, options);
475+ glr_value_pre = glr_value_cur;
476+ }
477+ } else {
478+ y.reserve(num_x);
479+ for (size_t i = 0; i < num_x; i++) {
480+ size_t ni = singleWeight ? n0 : weight[i].size(dim_num_w - 1);
481+ create_new_tensor_multi_dim(y, x[i], ni, options);
482+ }
483+ } // 校验NO_SPLIT时为特殊场景（groupList为空）或num_x > 1
484+ } else if (split_item_value == IN_NOT_SPLIT_OUT_SPLIT || split_item_value == IN_SPLIT_OUT_SPLIT) {
485+ if (num_x > 1) {
486+ size_t dim_m = 0;
487+ for (size_t i = 0; i < num_x; i++) {
488+ dim_m += static_cast<size_t>(x[i].size(0));
489+ }
490+ weight[0].dtype() == at::ScalarType::Int ? create_new_tensor(y, dim_m, n0 * INT4_NUMS_IN_INT32, options)
491+ : create_new_tensor(y, dim_m, n_new, options);
492+ } else if (num_x == 1) {
493+ if (group_type_value == K_SPLIT) {
494+ TORCH_CHECK(
495+ num_weight == 1, "When group_list is 2(K_SPLIT) and split_item is 2/3, the length of weight must equal x.");
496+ weight[0].dtype() == at::ScalarType::Int
497+ ? create_new_tensor_batch(y, num_group_list, x[0].size(0), n0 * INT4_NUMS_IN_INT32, options)
498+ : create_new_tensor_batch(y, num_group_list, x[0].size(0), n_new, options);
499+ } else {
500+ (weight[0].dtype() == at::ScalarType::Int ||
501+ (weight[0].dtype() == at::ScalarType::Float && weight[0].dtype() != x[0].dtype())) &&
502+ (!weight_trans)
503+ ? create_new_tensor(y, x[0].size(0), n0 * INT4_NUMS_IN_INT32, options)
504+ : create_new_tensor(y, x[0].size(0), n_new, options);
505+ }
506+ }
507+ }
508+ at::TensorList result = at::TensorList(y);
509+ 
510+ auto bias_real = bias_tl.value_or(at::TensorList());
511+ auto scale_real = scale.value_or(at::TensorList());
512+ auto offset_real = offset_tl.value_or(at::TensorList());
513+ auto antiquant_scale_real = antiquant_scale_tl.value_or(at::TensorList());
514+ auto antiquant_offset_real = antiquant_offset_tl.value_or(at::TensorList());
515+ auto per_token_scale_real = per_token_scale.value_or(at::TensorList());
516+ auto activation_input_real = activation_input_tl.value_or(at::TensorList());
517+ auto activation_quant_scale_real = activation_quant_scale_tl.value_or(at::TensorList());
518+ auto activation_quant_offset_real = activation_quant_offset_tl.value_or(at::TensorList());
519+ auto act_out = at::TensorList();
520+ auto dynamic_quant_scale_out = at::TensorList();
521+ int64_t group_list_type_value = group_list_type.value_or(0);
522+ int64_t act_type_value = act_type.value_or(0);
523+ auto tuning_config_real = tuning_config.value_or(at::IntArrayRef{});
524+ 
525+ TensorListWrapper x_wrapper = {
526+ x,
527+ x_dtype.has_value() ? c10_npu::GetAclDataType(x_dtype.value())
528+ : npu_preparation::convert_to_acl_data_type(x[0].scalar_type())};
529+ TensorListWrapper weight_wrapper = {
530+ weight,
531+ weight_dtype.has_value() ? c10_npu::GetAclDataType(weight_dtype.value())
532+ : npu_preparation::convert_to_acl_data_type(weight[0].scalar_type())};
533+ TensorListWrapper scale_wrapper = {
534+ scale_real,
535+ scale_dtype.has_value()
536+ ? c10_npu::GetAclDataType(scale_dtype.value())
537+ : (scale_real.empty() ? aclDataType::ACL_UINT64
538+ : npu_preparation::convert_to_acl_data_type(scale_real[0].scalar_type()))};
539+ TensorListWrapper per_token_scale_wrapper = {
540+ per_token_scale_real,
541+ per_token_scale_dtype.has_value()
542+ ? c10_npu::GetAclDataType(per_token_scale_dtype.value())
543+ : (per_token_scale_real.empty()
544+ ? aclDataType::ACL_FLOAT
545+ : npu_preparation::convert_to_acl_data_type(per_token_scale_real[0].scalar_type()))};
546+ TensorListWrapper antiquant_scale_wrapper = {
547+ antiquant_scale_real,
548+ antiquant_scale_real.empty()
549+ ? aclDataType::ACL_FLOAT16
550+ : (antiquant_scale_real[0].scalar_type() == at::ScalarType::Byte
551+ ? aclDataType::ACL_FLOAT8_E8M0
552+ : npu_preparation::convert_to_acl_data_type(antiquant_scale_real[0].scalar_type()))};
553+ 
554+ int64_t weight_format = at_npu::native::custom_ops::get_npu_format(weight[0]);
555+ const bool is_weight_nz = (weight_format == ACL_FORMAT_FRACTAL_NZ) || (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_2) ||
556+ (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_4) || (weight_format == ACL_FORMAT_FRACTAL_NZ_C0_16);
557+ if (is_weight_nz) {
558+ static const bool is_weight_nz_available = check_aclnn_kernel_available("aclnnGroupedMatmulWeightNz");
559+ TORCH_CHECK(
560+ is_weight_nz_available,
561+ "Format of weight in npu_grouped_matmul is FRACTAL_NZ, current CANN version "
562+ "do not support with this format. Please try to update the version of CANN." +
563+ OPS_ERROR(ErrCode::PARAM));
564+ int64_t quant_per_group_size = 0;
565+ EXEC_NPU_CMD(
566+ aclnnGroupedMatmulWeightNz,
567+ x_wrapper,
568+ weight_wrapper,
569+ bias_real,
570+ scale_wrapper,
571+ offset_real,
572+ antiquant_scale_wrapper,
573+ antiquant_offset_real,
574+ per_token_scale_wrapper,
575+ group_list_real,
576+ activation_input_real,
577+ activation_quant_scale_real,
578+ activation_quant_offset_real,
579+ split_item_value,
580+ group_type_value,
581+ group_list_type_value,
582+ act_type_value,
583+ tuning_config_real,
584+ quant_per_group_size,
585+ result,
586+ act_out,
587+ dynamic_quant_scale_out);
588+ return y[0];
589+ }
590+ static const bool is_grouped_matmul_V5_available = check_aclnn_kernel_available("aclnnGroupedMatmulV5");
591+ static const bool dtypeValid = x[0].scalar_type() != at::ScalarType::Float8_e5m2 &&
592+ x[0].scalar_type() != at::ScalarType::Float8_e4m3fn && !x_dtype.has_value() && !weight_dtype.has_value();
593+ if (!is_grouped_matmul_V5_available || !dtypeValid || mxfp4_valid) {
594+ EXEC_NPU_CMD(
595+ aclnnGroupedMatmulV4,
596+ x_wrapper,
597+ weight_wrapper,
598+ bias_real,
599+ scale_wrapper,
600+ offset_real,
601+ antiquant_scale_real,
602+ antiquant_offset_real,
603+ per_token_scale_wrapper,
604+ group_list_real,
605+ activation_input_real,
606+ activation_quant_scale_real,
607+ activation_quant_offset_real,
608+ split_item_value,
609+ group_type_value,
610+ group_list_type_value,
611+ act_type_value,
612+ result,
613+ act_out,
614+ dynamic_quant_scale_out);
615+ } else {
616+ EXEC_NPU_CMD(
617+ aclnnGroupedMatmulV5,
618+ x_wrapper,
619+ weight_wrapper,
620+ bias_real,
621+ scale_wrapper,
622+ offset_real,
623+ antiquant_scale_real,
624+ antiquant_offset_real,
625+ per_token_scale_wrapper,
626+ group_list_real,
627+ activation_input_real,
628+ activation_quant_scale_real,
629+ activation_quant_offset_real,
630+ split_item_value,
631+ group_type_value,
632+ group_list_type_value,
633+ act_type_value,
634+ tuning_config_real,
635+ result,
636+ act_out,
637+ dynamic_quant_scale_out);
638+ }
639+ return y[0];
551}640}
552-}641+} // namespace op_api
Mop_plugin/third_party/atb/inc/context.h+39-39
@@ -37,48 +37,48 @@ namespace atb {
37//! Context类会管理任务流队列比如Operation执行以及TilingCopy,管理tiling内存的申请与释放.37//! Context类会管理任务流队列比如Operation执行以及TilingCopy,管理tiling内存的申请与释放.
38//!38//!
39class Context {39class Context {
40- public:40+ public:
41- //! \brief 默认构造函数.41+ //! \brief 默认构造函数.
42- Context() = default;42+ Context() = default;
43 43 
44- //! \brief 默认析构函数.44+ //! \brief 默认析构函数.
45- virtual ~Context() = default;45+ virtual ~Context() = default;
46 46 
47- //!47+ //!
48- //! \brief 将传入stream队列设置为当前执行队列.48+ //! \brief 将传入stream队列设置为当前执行队列.
49- //!49+ //!
50- //! 将传入stream队列设置为当前执行队列,然后再去执行对应的Operation.50+ //! 将传入stream队列设置为当前执行队列,然后再去执行对应的Operation.
51- //!51+ //!
52- //! \param stream 传入的stream队列52+ //! \param stream 传入的stream队列
53- //!53+ //!
54- //! \return 状态值.如果设置成功，返回NO_ERROR.54+ //! \return 状态值.如果设置成功，返回NO_ERROR.
55- //!55+ //!
56- virtual Status SetExecuteStream(aclrtStream stream) = 0;56+ virtual Status SetExecuteStream(aclrtStream stream) = 0;
57 57 
58- //!58+ //!
59- //! \brief 获取当前执行stream队列.59+ //! \brief 获取当前执行stream队列.
60- //!60+ //!
61- //! \return 执行流队列61+ //! \return 执行流队列
62- //!62+ //!
63- virtual aclrtStream GetExecuteStream() const = 0;63+ virtual aclrtStream GetExecuteStream() const = 0;
64 64 
65- //!65+ //!
66- //! \brief 设置异步拷贝tiling信息功能.66+ //! \brief 设置异步拷贝tiling信息功能.
67- //!67+ //!
68- //! 设置异步拷贝tiling信息功能是否开启，如果是，则创建stream和event来进行tiling拷贝过程.68+ //! 设置异步拷贝tiling信息功能是否开启，如果是，则创建stream和event来进行tiling拷贝过程.
69- //!69+ //!
70- //! \param enable 传入的标志，bool类型70+ //! \param enable 传入的标志，bool类型
71- //!71+ //!
72- //! \return 状态值.如果设置成功，返回NO_ERROR.72+ //! \return 状态值.如果设置成功，返回NO_ERROR.
73- //!73+ //!
74- virtual Status SetAsyncTilingCopyStatus(bool enable) = 0;74+ virtual Status SetAsyncTilingCopyStatus(bool enable) = 0;
75 75 
76- //!76+ //!
77- //! \brief 获取tiling拷贝状态.77+ //! \brief 获取tiling拷贝状态.
78- //!78+ //!
79- //! \return 如果获取成功，返回True.79+ //! \return 如果获取成功，返回True.
80- //!80+ //!
81- virtual bool GetAsyncTilingCopyStatus() const = 0;81+ virtual bool GetAsyncTilingCopyStatus() const = 0;
82};82};
83 83 
84//!84//!
@@ -90,7 +90,7 @@ class Context {
90//!90//!
91//! \return 状态值.如果设置成功，返回NO_ERROR.91//! \return 状态值.如果设置成功，返回NO_ERROR.
92//!92//!
93-Status CreateContext(Context **context);93+Status CreateContext(Context** context);
94 94 
95//!95//!
96//! \brief 销毁上下文.96//! \brief 销毁上下文.
@@ -101,6 +101,6 @@ Status CreateContext(Context **context);
101//!101//!
102//! \return 状态值.如果设置成功，返回NO_ERROR.102//! \return 状态值.如果设置成功，返回NO_ERROR.
103//!103//!
104-Status DestroyContext(Context *context);104+Status DestroyContext(Context* context);
105} // namespace atb105} // namespace atb
106#endif106#endif
Mop_plugin/third_party/atb/inc/infer_op_params.h+579-578
@@ -41,8 +41,8 @@ namespace infer {
41//! \brief 数据排布类型41//! \brief 数据排布类型
42//!42//!
43enum InputLayout : int {43enum InputLayout : int {
44- TYPE_BSND = 0, //!< 默认值，表示数据排布为BSND44+ TYPE_BSND = 0, //!< 默认值，表示数据排布为BSND
45- TYPE_BNSD //!< 表示数据排布为BNSD45+ TYPE_BNSD //!< 表示数据排布为BNSD
46};46};
47 47 
48//!48//!
@@ -51,12 +51,12 @@ enum InputLayout : int {
51//! \brief 量化支持的类型51//! \brief 量化支持的类型
52//!52//!
53enum QuantType : int {53enum QuantType : int {
54- QUANT_UNDEFINED = 0, //!< 不量化54+ QUANT_UNDEFINED = 0, //!< 不量化
55- QUANT_INT4, //!< 当前不支持55+ QUANT_INT4, //!< 当前不支持
56- QUANT_INT8, //!< int8量化56+ QUANT_INT8, //!< int8量化
57- QUANT_INT16, //!< 当前不支持57+ QUANT_INT16, //!< 当前不支持
58- QUANT_FLOAT8, //!< 当前不支持58+ QUANT_FLOAT8, //!< 当前不支持
59- QUANT_FLOAT16, //!< 当前不支持59+ QUANT_FLOAT16, //!< 当前不支持
60};60};
61 61 
62//!62//!
@@ -65,9 +65,9 @@ enum QuantType : int {
65//! \brief 动态量化支持的类型65//! \brief 动态量化支持的类型
66//!66//!
67enum DynamicQuantType : int {67enum DynamicQuantType : int {
68- DYNAMIC_QUANT_UNDEFINED = 0, //!< 非动态量化68+ DYNAMIC_QUANT_UNDEFINED = 0, //!< 非动态量化
69- DYNAMIC_QUANT_SYMMETRIC, //!< 对称动态量化69+ DYNAMIC_QUANT_SYMMETRIC, //!< 对称动态量化
70- DYNAMIC_QUANT_ASYMMETRIC, //!< 非对称动态量化，暂不支持70+ DYNAMIC_QUANT_ASYMMETRIC, //!< 非对称动态量化，暂不支持
71};71};
72 72 
73//!73//!
@@ -80,17 +80,17 @@ enum DynamicQuantType : int {
80//! ACTIVATION_SWIGLU_BACKWARD: Swiglu正向激活函数的反向，求梯度时使用。只支持Atlas 800I A2推理产品。80//! ACTIVATION_SWIGLU_BACKWARD: Swiglu正向激活函数的反向，求梯度时使用。只支持Atlas 800I A2推理产品。
81//!81//!
82enum ActivationType : int {82enum ActivationType : int {
83- ACTIVATION_UNDEFINED = 0, //!< 未定义83+ ACTIVATION_UNDEFINED = 0, //!< 未定义
84- ACTIVATION_RELU, //!< RELU激活类型84+ ACTIVATION_RELU, //!< RELU激活类型
85- ACTIVATION_GELU, //!< GELU激活类型85+ ACTIVATION_GELU, //!< GELU激活类型
86- ACTIVATION_FAST_GELU, //!< FAST_GELU激活类型86+ ACTIVATION_FAST_GELU, //!< FAST_GELU激活类型
87- ACTIVATION_SWISH, //!< SWISH激活类型87+ ACTIVATION_SWISH, //!< SWISH激活类型
88- ACTIVATION_LOG, //!< LOG激活类型88+ ACTIVATION_LOG, //!< LOG激活类型
89- ACTIVATION_SWIGLU_FORWARD, //!< SWIGLU_FORWARD激活类型89+ ACTIVATION_SWIGLU_FORWARD, //!< SWIGLU_FORWARD激活类型
90- ACTIVATION_SWIGLU_BACKWARD, //!< SWIGLU_BACKWARD激活类型90+ ACTIVATION_SWIGLU_BACKWARD, //!< SWIGLU_BACKWARD激活类型
91- ACTIVATION_SIGMOID, //!< SIGMOID激活类型91+ ACTIVATION_SIGMOID, //!< SIGMOID激活类型
92- ACTIVATION_FASTER_GELU_FORWARD, //!< FASTER_GELU_FORWARD激活类型92+ ACTIVATION_FASTER_GELU_FORWARD, //!< FASTER_GELU_FORWARD激活类型
93- ACTIVATION_MAX, //!< 枚举最大值, 非激活类型93+ ACTIVATION_MAX, //!< 枚举最大值, 非激活类型
94};94};
95 95 
96//!96//!
@@ -99,9 +99,9 @@ enum ActivationType : int {
99//! \brief 通信算子支持的通信模式.99//! \brief 通信算子支持的通信模式.
100//!100//!
101enum CommMode : int {101enum CommMode : int {
102- COMM_UNDEFINED = -1, //!< 未定义102+ COMM_UNDEFINED = -1, //!< 未定义
103- COMM_MULTI_PROCESS, //!< 指定多进程通信103+ COMM_MULTI_PROCESS, //!< 指定多进程通信
104- COMM_MULTI_THREAD, //!< 指定多线程通信104+ COMM_MULTI_THREAD, //!< 指定多线程通信
105};105};
106 106 
107//!107//!
@@ -113,115 +113,116 @@ enum CommMode : int {
113//! Atlas 推理系列产品中不支持bf16类型数据。113//! Atlas 推理系列产品中不支持bf16类型数据。
114//!114//!
115struct RmsNormParam {115struct RmsNormParam {
116- //!116+ //!
117- //! \brief RmsNormType117+ //! \brief RmsNormType
118- //!118+ //!
119- enum RmsNormType : int {119+ enum RmsNormType : int {
120- RMS_NORM_UNDEFINED = 0, //!< 默认值，未定义120+ RMS_NORM_UNDEFINED = 0, //!< 默认值，未定义
121- RMS_NORM_NORM, //!< NORM参数。121+ RMS_NORM_NORM, //!< NORM参数。
122- RMS_NORM_PRENORM, //!< PRENORM参数。122+ RMS_NORM_PRENORM, //!< PRENORM参数。
123- RMS_NORM_POSTNORM, //!< POSTNORM参数123+ RMS_NORM_POSTNORM, //!< POSTNORM参数
124- };124+ };
125- //!125+ //!
126- //! \brief PrecisionMode126+ //! \brief PrecisionMode
127- //!127+ //!
128- enum PrecisionMode : int {128+ enum PrecisionMode : int {
129- HIGH_PRECISION_MODE = 0, //!< 中间计算使用float类型129+ HIGH_PRECISION_MODE = 0, //!< 中间计算使用float类型
130- HIGH_PERFORMANCE_MODE, //!< 中间计算使用float16类型130+ HIGH_PERFORMANCE_MODE, //!< 中间计算使用float16类型
131- };131+ };
132- //!132+ //!
133- //! \brief ModelType133+ //! \brief ModelType
134- //!134+ //!
135- enum ModelType : int {135+ enum ModelType : int {
136- LLAMA_MODEL = 0, //!< 默认值，使用Llama rmsnorm的公式136+ LLAMA_MODEL = 0, //!< 默认值，使用Llama rmsnorm的公式
137- GEMMA_MODEL, //!< 使用Gemma rmsnorm的公式137+ GEMMA_MODEL, //!< 使用Gemma rmsnorm的公式
138- };138+ };
139- //!139+ //!
140- //! \brief NormParam140+ //! \brief NormParam
141- //!141+ //!
142- struct NormParam {142+ struct NormParam {
143- //! \brief 量化类型。143+ //! \brief 量化类型。
144- //! 当前支持以下类型。144+ //! 当前支持以下类型。
145- //! QUANT_UNDEINFED, QUANT_INT8145+ //! QUANT_UNDEINFED, QUANT_INT8
146- QuantType quantType = QUANT_UNDEFINED;146+ QuantType quantType = QUANT_UNDEFINED;
147- //! \brief Epsilon，归一化时加在分母上防止除零。147+ //! \brief Epsilon，归一化时加在分母上防止除零。
148- float epsilon = 1e-5;148+ float epsilon = 1e-5;
149- //! \brief Epsilon，默认为1e-5，暂时不使用。149+ //! \brief Epsilon，默认为1e-5，暂时不使用。
150- double layerNormEps = 1e-5;150+ double layerNormEps = 1e-5;
151- //! \brief 默认为False，设置为true时会使用训练的rmsnormforward算子。仅在Atlas 800I A2推理产品上支持该设置。151+ //! \brief 默认为False，设置为true时会使用训练的rmsnormforward算子。仅在Atlas 800I A2推理产品上支持该设置。
152- //! 不支持和“precisionMode”，“modelType”同时设置。量化场景下不支持使用“rstd”。152+ //! 不支持和“precisionMode”，“modelType”同时设置。量化场景下不支持使用“rstd”。
153- bool rstd = false;153+ bool rstd = false;
154- //! \brief 默认为HIGH_PRECISION_MODE。154+ //! \brief 默认为HIGH_PRECISION_MODE。
155- //! 支持参数如下：155+ //! 支持参数如下：
156- //! HIGH_PRECISION_MODE：默认值，中间计算使用float类型156+ //! HIGH_PRECISION_MODE：默认值，中间计算使用float类型
157- //! HIGH_PERFORMANCE_MODE： 中间计算使用float16类型157+ //! HIGH_PERFORMANCE_MODE： 中间计算使用float16类型
158- //! 不支持和“rstd”，“modelType”同时设置。输入类型只支持float16。158+ //! 不支持和“rstd”，“modelType”同时设置。输入类型只支持float16。
159- //! 量化场景下不支持使用“precisionMode”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。159+ //! 量化场景下不支持使用“precisionMode”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
160- PrecisionMode precisionMode = HIGH_PRECISION_MODE;160+ PrecisionMode precisionMode = HIGH_PRECISION_MODE;
161- //! \brief 默认为LLAMA_MODEL，设置为GEMMA_MODEL时使用gemma模型的rmsnorm计算公式。161+ //! \brief 默认为LLAMA_MODEL，设置为GEMMA_MODEL时使用gemma模型的rmsnorm计算公式。
162- //! 支持参数如下：162+ //! 支持参数如下：
163- //! LLAMA_MODEL：默认值， Llama的rms norm计算公式。163+ //! LLAMA_MODEL：默认值， Llama的rms norm计算公式。
164- //! GEMMA_MODEL：Gemma的rms norm计算公式。164+ //! GEMMA_MODEL：Gemma的rms norm计算公式。
165- //! 不支持和“rstd”，“precisionMode”同时启用。165+ //! 不支持和“rstd”，“precisionMode”同时启用。
166- //! 量化场景下不支持使用“modelType”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。166+ //! 量化场景下不支持使用“modelType”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
167- ModelType modelType = LLAMA_MODEL;167+ ModelType modelType = LLAMA_MODEL;
168- //! \brief 动态量化类型。默认为DYNAMIC_QUANT_UNDEFINED非动态量化。当前版本暂不支持非对称动态量化。168+ //! \brief 动态量化类型。默认为DYNAMIC_QUANT_UNDEFINED非动态量化。当前版本暂不支持非对称动态量化。
169- DynamicQuantType dynamicQuantType = DYNAMIC_QUANT_UNDEFINED;169+ DynamicQuantType dynamicQuantType = DYNAMIC_QUANT_UNDEFINED;
170- //!
171- //! \brief 预留参数
172- //!
173- uint8_t rsv[32] = {0};
174- };
175- //!
176- //! \brief PreNormParam
177- //!
178- struct PreNormParam {
179- //! \brief 量化类型。
180- //! 当前支持以下类型。
181- //! QUANT_UNDEINFED
182- //! QUANT_INT8
183- QuantType quantType = QUANT_UNDEFINED;
184- //! \brief Epsilon，归一化时加在分母上防止除零。
185- float epsilon = 1e-5;
186- //! \brief 是否叠加偏置。默认为False，当需要输入beta时设置为True。量化场景下不支持使用“hasBias”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
187- bool hasBias = false;
188- //!
189- //! \brief 预留参数
190- //!
191- uint8_t rsv[23] = {0};
192- };
193- //!
194- //! \brief PostNormParam
195- //!
196- struct PostNormParam {
197- //! \brief 量化类型。
198- //! 当前仅支持QUANT_UNDEINFED。
199- QuantType quantType = QUANT_UNDEFINED;
200- //! \brief Epsilon，归一化时加在分母上防止除零。
201- float epsilon = 1e-5;
202- //! \brief 是否叠加偏置。默认为False，当需要输入beta时设置为True。
203- bool hasBias = false;
204- //!
205- //! \brief 预留参数
206- //!
207- uint8_t rsv[23] = {0};
208- };
209- //! \brief 归一化类型，参数如下：
210- //! RMS_NORM_UNDEFINED：默认值，未定义。
211- //! RMS_NORM_NORM：NORM参数。
212- //! RMS_NORM_PRENORM：PRENORM参数。
213- //! RMS_NORM_POSTNORM：POSTNORM参数。
214- RmsNormType layerType = RMS_NORM_UNDEFINED;
215- //! \brief NORM参数。
216- NormParam normParam;
217- //! \brief PRENORM参数。
218- PreNormParam preNormParam;
219- //! \brief POSTNORM参数。
220- PostNormParam postNormParam;
221 //!170 //!
222 //! \brief 预留参数171 //! \brief 预留参数
223 //!172 //!
224- uint8_t rsv[8] = {0};173+ uint8_t rsv[32] = {0};
174+ };
175+ //!
176+ //! \brief PreNormParam
177+ //!
178+ struct PreNormParam {
179+ //! \brief 量化类型。
180+ //! 当前支持以下类型。
181+ //! QUANT_UNDEINFED
182+ //! QUANT_INT8
183+ QuantType quantType = QUANT_UNDEFINED;
184+ //! \brief Epsilon，归一化时加在分母上防止除零。
185+ float epsilon = 1e-5;
186+ //! \brief
187+ //! 是否叠加偏置。默认为False，当需要输入beta时设置为True。量化场景下不支持使用“hasBias”，该场景下配置该参数将返回报错ERROR_INVALID_PARAM。
188+ bool hasBias = false;
189+ //!
190+ //! \brief 预留参数
191+ //!
192+ uint8_t rsv[23] = {0};
193+ };
194+ //!
195+ //! \brief PostNormParam
196+ //!
197+ struct PostNormParam {
198+ //! \brief 量化类型。
199+ //! 当前仅支持QUANT_UNDEINFED。
200+ QuantType quantType = QUANT_UNDEFINED;
201+ //! \brief Epsilon，归一化时加在分母上防止除零。
202+ float epsilon = 1e-5;
203+ //! \brief 是否叠加偏置。默认为False，当需要输入beta时设置为True。
204+ bool hasBias = false;
205+ //!
206+ //! \brief 预留参数
207+ //!
208+ uint8_t rsv[23] = {0};
209+ };
210+ //! \brief 归一化类型，参数如下：
211+ //! RMS_NORM_UNDEFINED：默认值，未定义。
212+ //! RMS_NORM_NORM：NORM参数。
213+ //! RMS_NORM_PRENORM：PRENORM参数。
214+ //! RMS_NORM_POSTNORM：POSTNORM参数。
215+ RmsNormType layerType = RMS_NORM_UNDEFINED;
216+ //! \brief NORM参数。
217+ NormParam normParam;
218+ //! \brief PRENORM参数。
219+ PreNormParam preNormParam;
220+ //! \brief POSTNORM参数。
221+ PostNormParam postNormParam;
222+ //!
223+ //! \brief 预留参数
224+ //!
225+ uint8_t rsv[8] = {0};
225};226};
226 227 
227//!228//!
@@ -229,115 +230,116 @@ struct RmsNormParam {
229//!230//!
230//! \brief 将A、B两个矩阵进行矩阵乘运算，同时可以选择对矩阵乘的运算结果进行叠加偏置、InplaceAdd融合或反量化操作。231//! \brief 将A、B两个矩阵进行矩阵乘运算，同时可以选择对矩阵乘的运算结果进行叠加偏置、InplaceAdd融合或反量化操作。
231//!232//!
232-//! \note 算子本质上是接收x和weight两个输入tensor作为A矩阵和B矩阵进行矩阵乘运算，可通过参数transposeA与transposeB控制做矩233+//! \note
234+//! 算子本质上是接收x和weight两个输入tensor作为A矩阵和B矩阵进行矩阵乘运算，可通过参数transposeA与transposeB控制做矩
233//! 阵乘前是否需要对A矩阵和B矩阵进行行列转置，根据参数转置后的A矩阵和B矩阵需满足矩阵乘维度关系。例如，当transposeA为false，235//! 阵乘前是否需要对A矩阵和B矩阵进行行列转置，根据参数转置后的A矩阵和B矩阵需满足矩阵乘维度关系。例如，当transposeA为false，
234//! transposeB为true时，x和weight的shape可以分别为[m, k]和[n, k]。236//! transposeB为true时，x和weight的shape可以分别为[m, k]和[n, k]。
235//!237//!
236//! \note 该算子支持浮点和量化场景，当参数outDataType值为ACL_DT_UNDEFINED时为浮点场景，否则为量化场景。238//! \note 该算子支持浮点和量化场景，当参数outDataType值为ACL_DT_UNDEFINED时为浮点场景，否则为量化场景。
237//!239//!
238struct LinearParam {240struct LinearParam {
239- //!241+ //!
240- //! \brief 是否转置A矩阵。242+ //! \brief 是否转置A矩阵。
241- //!243+ //!
242- //! \note 默认值为false，不转置。244+ //! \note 默认值为false，不转置。
243- //!245+ //!
244- //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为false。246+ //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为false。
245- //!247+ //!
246- bool transposeA = false;248+ bool transposeA = false;
247- //!249+ //!
248- //! \brief 是否转置B矩阵。250+ //! \brief 是否转置B矩阵。
249- //!251+ //!
250- //! \note 默认值为true，转置。252+ //! \note 默认值为true，转置。
251- //!253+ //!
252- //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为true。254+ //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为true。
253- //!255+ //!
254- bool transposeB = true;256+ bool transposeB = true;
255- //!257+ //!
256- //! \brief 是否叠加偏置。258+ //! \brief 是否叠加偏置。
257- //!259+ //!
258- //! \note 默认值为true，叠加偏置。260+ //! \note 默认值为true，叠加偏置。
259- //!261+ //!
260- //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为true。262+ //! \warning 在量化场景下，非Atlas 800I A2推理产品仅支持配置为true。
261- //!263+ //!
262- //! \warning enAccum为true时，仅支持配置为false。264+ //! \warning enAccum为true时，仅支持配置为false。
263- //!265+ //!
264- bool hasBias = true;266+ bool hasBias = true;
265- //!267+ //!
266- //! \brief 输出数据类型。268+ //! \brief 输出数据类型。
267- //!269+ //!
268- //! \note 默认值为ACL_DT_UNDEFINED。270+ //! \note 默认值为ACL_DT_UNDEFINED。
269- //!271+ //!
270- //! \warning 浮点场景下：支持配置为ACL_DT_UNDEFINED。272+ //! \warning 浮点场景下：支持配置为ACL_DT_UNDEFINED。
271- //!273+ //!
272- //! \warning 量化场景下：Atlas 800I A2推理产品支持配置为ACL_FLOAT16/ACL_BF16，否则，仅支持配置为ACL_FLOAT16。274+ //! \warning 量化场景下：Atlas 800I A2推理产品支持配置为ACL_FLOAT16/ACL_BF16，否则，仅支持配置为ACL_FLOAT16。
273- //!275+ //!
274- aclDataType outDataType = ACL_DT_UNDEFINED;276+ aclDataType outDataType = ACL_DT_UNDEFINED;
275- //!277+ //!
276- //! \brief 是否使能累加。278+ //! \brief 是否使能累加。
277- //!279+ //!
278- //! \note 默认值为false，不使能累加。280+ //! \note 默认值为false，不使能累加。
279- //!281+ //!
280- //! \warning 仅在Atlas 800I A2推理产品支持配置为true。282+ //! \warning 仅在Atlas 800I A2推理产品支持配置为true。
281- //!283+ //!
282- //! \warning hasBias为true时，仅支持配置为false。284+ //! \warning hasBias为true时，仅支持配置为false。
283- //!285+ //!
284- //! \warning 量化场景下，仅支持配置为false。286+ //! \warning 量化场景下，仅支持配置为false。
285- //!287+ //!
286- bool enAccum = false;288+ bool enAccum = false;
287- //!289+ //!
288- //! \brief 预留参数290+ //! \brief 预留参数
289- //!291+ //!
290- uint8_t rsv[23] = {0};292+ uint8_t rsv[23] = {0};
291};293};
292 294 
293struct GroupTopkParam {295struct GroupTopkParam {
294- //!296+ //!
295- //! \brief 每个token分组数量。注：“专家总数”为inTensor0Desc.shape.dims[1]的值。297+ //! \brief 每个token分组数量。注：“专家总数”为inTensor0Desc.shape.dims[1]的值。
296- //!298+ //!
297- //! \note 必传，默认值为1，取值范围为[1, 专家总数]。299+ //! \note 必传，默认值为1，取值范围为[1, 专家总数]。
298- //!300+ //!
299- //! \warning groupNum需要保证可以被inTensor0Desc.shape.dims[1]整除。301+ //! \warning groupNum需要保证可以被inTensor0Desc.shape.dims[1]整除。
300- //!302+ //!
301- int32_t groupNum = 1;303+ int32_t groupNum = 1;
302- //!304+ //!
303- //! \brief 选择top K组数量。305+ //! \brief 选择top K组数量。
304- //!306+ //!
305- //! \note 必传，默认值为0，取值范围为[1, groupNum]。307+ //! \note 必传，默认值为0，取值范围为[1, groupNum]。
306- //!308+ //!
307- //! \warning309+ //! \warning
308- //!310+ //!
309- int32_t k = 0;311+ int32_t k = 0;
310- //!312+ //!
311- //! \enum GroupMultiFlag313+ //! \enum GroupMultiFlag
312- //!314+ //!
313- //! \brief 指定GroupTopk每组中取值计算的方式。315+ //! \brief 指定GroupTopk每组中取值计算的方式。
314- //!316+ //!
315- //! \warning317+ //! \warning
316- //!318+ //!
317- enum GroupMultiFlag : uint16_t {319+ enum GroupMultiFlag : uint16_t {
318- UNDEFINED = 0, //!< 默认方式，每组内取最大值。320+ UNDEFINED = 0, //!< 默认方式，每组内取最大值。
319- SUM_MULTI_MAX //!< 每组内取n个最大值求和，需要设置参数n321+ SUM_MULTI_MAX //!< 每组内取n个最大值求和，需要设置参数n
320- };322+ };
321- //!323+ //!
322- //! \brief 指定GroupTopk每组中取值计算的方式。324+ //! \brief 指定GroupTopk每组中取值计算的方式。
323- //!325+ //!
324- //! \note 默认值为UNDEFINED。326+ //! \note 默认值为UNDEFINED。
325- //!327+ //!
326- //! \warning 取值为SUM_MULTI_MAX时需要传入参数n。328+ //! \warning 取值为SUM_MULTI_MAX时需要传入参数n。
327- //!329+ //!
328- GroupMultiFlag groupMultiFlag = UNDEFINED;330+ GroupMultiFlag groupMultiFlag = UNDEFINED;
329- //!331+ //!
330- //! \brief 每组内取值的个数。332+ //! \brief 每组内取值的个数。
331- //!333+ //!
332- //! \note 默认值为1，取值范围为[1,expert_num/groupNum]。334+ //! \note 默认值为1，取值范围为[1,expert_num/groupNum]。
333- //!335+ //!
334- //! \warning 只有当groupMultiFlag为SUM_MULTI_MAX时有效336+ //! \warning 只有当groupMultiFlag为SUM_MULTI_MAX时有效
335- //!337+ //!
336- uint16_t n = 1;338+ uint16_t n = 1;
337- //!339+ //!
338- //! \brief 预留参数340+ //! \brief 预留参数
339- //!341+ //!
340- uint8_t rsv[12] = {0};342+ uint8_t rsv[12] = {0};
341};343};
342 344 
343//!345//!
@@ -347,163 +349,163 @@ struct GroupTopkParam {
347//! 所以第二行和第四行为其目标block，context_lens则表示KV有多少个token，则代表仅有block_id为(3,4,5,9,10)是需要与Q进行计算的。349//! 所以第二行和第四行为其目标block，context_lens则表示KV有多少个token，则代表仅有block_id为(3,4,5,9,10)是需要与Q进行计算的。
348//!350//!
349struct PagedAttentionParam {351struct PagedAttentionParam {
350- //! query 头大小352+ //! query 头大小
351- int32_t headNum = 0;353+ int32_t headNum = 0;
352- //! 算子tor值, 在Q*K^T后乘354+ //! 算子tor值, 在Q*K^T后乘
353- float qkScale = 1.0;355+ float qkScale = 1.0;
354- //! kv头数量356+ //! kv头数量
355- int32_t kvHeadNum = 0;357+ int32_t kvHeadNum = 0;
356- //!358+ //!
357- //! \enum MaskType359+ //! \enum MaskType
358- //!360+ //!
359- //! \brief The type values of MaskType.361+ //! \brief The type values of MaskType.
360- //!362+ //!
361- enum MaskType : int {363+ enum MaskType : int {
362- UNDEFINED = 0, //!< 默认值，全0的mask364+ UNDEFINED = 0, //!< 默认值，全0的mask
363- MASK_TYPE_NORM, //!< 倒三角mask365+ MASK_TYPE_NORM, //!< 倒三角mask
364- MASK_TYPE_ALIBI, //!< alibi mask366+ MASK_TYPE_ALIBI, //!< alibi mask
365- MASK_TYPE_SPEC, //!< 并行解码mask367+ MASK_TYPE_SPEC, //!< 并行解码mask
366- MASK_TYPE_MASK_FREE, //! mask_free 只支持fp16368+ MASK_TYPE_MASK_FREE, //! mask_free 只支持fp16
367- MASK_TYPE_NORM_COMPRESS //!< 倒三角压缩mask369+ MASK_TYPE_NORM_COMPRESS //!< 倒三角压缩mask
368- };370+ };
369- //! mask类型371+ //! mask类型
370- MaskType maskType = UNDEFINED;372+ MaskType maskType = UNDEFINED;
371- //! 是否开启动态batch373+ //! 是否开启动态batch
372- bool batchRunStatusEnable = false;374+ bool batchRunStatusEnable = false;
373- //!375+ //!
374- //! \enum QuantType376+ //! \enum QuantType
375- //!377+ //!
376- //! \brief quant类型378+ //! \brief quant类型
377- //!379+ //!
378- enum QuantType : int {380+ enum QuantType : int {
379- TYPE_QUANT_UNDEFINED = 0, //!< 默认值，不与量化融合，此381+ TYPE_QUANT_UNDEFINED = 0, //!< 默认值，不与量化融合，此
380- TYPE_DEQUANT_FUSION, //!< 与反量化融合, 只支持Atlas 800I A2推理产品382+ TYPE_DEQUANT_FUSION, //!< 与反量化融合, 只支持Atlas 800I A2推理产品
381- TYPE_QUANT_QKV_OFFLINE, //!< 离线INT8量化, 只支持Atlas 800I A2推理产品383+ TYPE_QUANT_QKV_OFFLINE, //!< 离线INT8量化, 只支持Atlas 800I A2推理产品
382- TYPE_QUANT_QKV_ONLINE //!< 在线INT8量化, 只支持Atlas 800I A2推理产品384+ TYPE_QUANT_QKV_ONLINE //!< 在线INT8量化, 只支持Atlas 800I A2推理产品
383- };385+ };
384- //!386+ //!
385- //! 量化类型：387+ //! 量化类型：
386- //! 为TYPE_QUANT_UNDEFINED时q，keyCache，valueCache为bf16/float16。388+ //! 为TYPE_QUANT_UNDEFINED时q，keyCache，valueCache为bf16/float16。
387- //! 为TYPE_DEQUANT_FUSION时q为bf16/float16，keyCache，valueCache为int8。389+ //! 为TYPE_DEQUANT_FUSION时q为bf16/float16，keyCache，valueCache为int8。
388- //! 为TYPE_QUANT_QKV_OFFLINE或TYPE_QUANT_QKV_ONLINE时q，keyCache，valueCache为int8。390+ //! 为TYPE_QUANT_QKV_OFFLINE或TYPE_QUANT_QKV_ONLINE时q，keyCache，valueCache为int8。
389- //! keyCache,valueCache的headsize等长，范围为（0, 256]，且block_size * head_size ≤ 128 * 128。391+ //! keyCache,valueCache的headsize等长，范围为（0, 256]，且block_size * head_size ≤ 128 * 128。
390- //! outdatatype需要配置，只能是ACL_FLOAT16或ACL_BF16。inputLayout只支持TYPE_BSND。392+ //! outdatatype需要配置，只能是ACL_FLOAT16或ACL_BF16。inputLayout只支持TYPE_BSND。
391- QuantType quantType = TYPE_QUANT_UNDEFINED;393+ QuantType quantType = TYPE_QUANT_UNDEFINED;
392 394 
393- //! output数据类型（格式为aclDataType）395+ //! output数据类型（格式为aclDataType）
394- aclDataType outDataType = ACL_DT_UNDEFINED;396+ aclDataType outDataType = ACL_DT_UNDEFINED;
395 397 
396- //! 开启量化功能后是否使用offset398+ //! 开启量化功能后是否使用offset
397- bool hasQuantOffset = false;399+ bool hasQuantOffset = false;
398- //!400+ //!
399- //! \enum CompressType401+ //! \enum CompressType
400- //!402+ //!
401- //! \brief 压缩类型403+ //! \brief 压缩类型
402- //!404+ //!
403- enum CompressType : int {405+ enum CompressType : int {
404- COMPRESS_TYPE_UNDEFINED = 0, //!< 默认值，不压缩406+ COMPRESS_TYPE_UNDEFINED = 0, //!< 默认值，不压缩
405- COMPRESS_TYPE_KVHEAD, //!< 压缩key_cache, value_cache的kvHead维度, 只支持Atlas 800I A2推理产品。407+ COMPRESS_TYPE_KVHEAD, //!< 压缩key_cache, value_cache的kvHead维度, 只支持Atlas 800I A2推理产品。
406- COMPRESS_TYPE_KVHEAD_ROPE, //!< rope场景压缩key_cache, value_cache的kvHead维度, 只支持Atlas 800I A2推理产品。408+ COMPRESS_TYPE_KVHEAD_ROPE, //!< rope场景压缩key_cache, value_cache的kvHead维度, 只支持Atlas 800I A2推理产品。
407- COMPRESS_TYPE_MAX //!< 压缩类型边界值，仅用于判断是否出界，所有情况不能取该值。409+ COMPRESS_TYPE_MAX //!< 压缩类型边界值，仅用于判断是否出界，所有情况不能取该值。
408- };410+ };
409- //!411+ //!
410- //! 压缩方式412+ //! 压缩方式
411- //! 为COMPRESS_TYPE_KVHEAD时，不支持quanttype为2和3。413+ //! 为COMPRESS_TYPE_KVHEAD时，不支持quanttype为2和3。
412- //! 为COMPRESS_TYPE_KVHEAD_ROPE时, maskType需传0。不支持quanttype为2和3。414+ //! 为COMPRESS_TYPE_KVHEAD_ROPE时, maskType需传0。不支持quanttype为2和3。
413- CompressType compressType = COMPRESS_TYPE_UNDEFINED;415+ CompressType compressType = COMPRESS_TYPE_UNDEFINED;
414- //!416+ //!
415- //! \enum CalcType417+ //! \enum CalcType
416- //!418+ //!
417- //! \brief The type values of CalcType.419+ //! \brief The type values of CalcType.
418- //!420+ //!
419- enum CalcType : int {421+ enum CalcType : int {
420- CALC_TYPE_UNDEFINED = 0, //!< 默认值，不开启并行解码422+ CALC_TYPE_UNDEFINED = 0, //!< 默认值，不开启并行解码
421- CALC_TYPE_SPEC //!< 并行解码功能，此时只支持quantType = 0423+ CALC_TYPE_SPEC //!< 并行解码功能，此时只支持quantType = 0
422- };424+ };
423- //! 计算类型425+ //! 计算类型
424- CalcType calcType = CALC_TYPE_UNDEFINED;426+ CalcType calcType = CALC_TYPE_UNDEFINED;
425 427 
426- //!428+ //!
427- //! \enum ScaleType429+ //! \enum ScaleType
428- //!430+ //!
429- //! \brief The type values of ScaleType.431+ //! \brief The type values of ScaleType.
430- //!432+ //!
431- enum ScaleType : int {433+ enum ScaleType : int {
432- SCALE_TYPE_TOR = 0, //!< 默认值，不开启LogN缩放434+ SCALE_TYPE_TOR = 0, //!< 默认值，不开启LogN缩放
433- SCALE_TYPE_LOGN, //!< 注意力使用LogN缩放435+ SCALE_TYPE_LOGN, //!< 注意力使用LogN缩放
434- SCALE_TYPE_MAX //!< 边界值，仅用于判断是否出界436+ SCALE_TYPE_MAX //!< 边界值，仅用于判断是否出界
435- };437+ };
436- //! scale类型438+ //! scale类型
437- //! 为SCALE_TYPE_LOGN时，不支持quanttype为2和3。439+ //! 为SCALE_TYPE_LOGN时，不支持quanttype为2和3。
438- ScaleType scaleType = SCALE_TYPE_TOR;440+ ScaleType scaleType = SCALE_TYPE_TOR;
439 441 
440- //! 数据排布格式默认为BSND442+ //! 数据排布格式默认为BSND
441- InputLayout inputLayout = TYPE_BSND;443+ InputLayout inputLayout = TYPE_BSND;
442- //! \brief 大于0时开启MLA合并kvcache功能，表示kv合并传入时v的head_size444+ //! \brief 大于0时开启MLA合并kvcache功能，表示kv合并传入时v的head_size
443- //! \note 默认值为0445+ //! \note 默认值为0
444- //! \warning 取值范围为[0,576]446+ //! \warning 取值范围为[0,576]
445- uint32_t mlaVHeadSize = 0;447+ uint32_t mlaVHeadSize = 0;
446- //!448+ //!
447- //! \brief 预留参数449+ //! \brief 预留参数
448- //!450+ //!
449- uint8_t rsv[68] = {0};451+ uint8_t rsv[68] = {0};
450};452};
451 453 
452//!454//!
453//! \brief 遍历每个key和value，将key和value(num_heads, head_size)按照slotmapping填入key_cache/value_cache指定位置455//! \brief 遍历每个key和value，将key和value(num_heads, head_size)按照slotmapping填入key_cache/value_cache指定位置
454//!456//!
455struct ReshapeAndCacheParam {457struct ReshapeAndCacheParam {
456- //!458+ //!
457- //! \enum CompressType459+ //! \enum CompressType
458- //!460+ //!
459- //! \brief 压缩类型461+ //! \brief 压缩类型
460- //!462+ //!
461- //! \note 默认值为COMPRESS_TYPE_UNDEFINED(0)，不开启压缩功能。463+ //! \note 默认值为COMPRESS_TYPE_UNDEFINED(0)，不开启压缩功能。
462- //!464+ //!
463- //! \warning 仅在Atlas 800I A2推理产品上支持设置为非COMPRESS_TYPE_UNDEFINED(0)的值465+ //! \warning 仅在Atlas 800I A2推理产品上支持设置为非COMPRESS_TYPE_UNDEFINED(0)的值
464- //!466+ //!
465- enum CompressType : int {467+ enum CompressType : int {
466- COMPRESS_TYPE_UNDEFINED = 0, //!< 默认值，不压缩468+ COMPRESS_TYPE_UNDEFINED = 0, //!< 默认值，不压缩
467- COMPRESS_TYPE_KVHEAD, //!< alibi场景下压缩key_cache, value_cahe的kvHead维度469+ COMPRESS_TYPE_KVHEAD, //!< alibi场景下压缩key_cache, value_cahe的kvHead维度
468- COMPRESS_TYPE_KVHEAD_ROPE //!< rope场景下压缩key_cache, value_cahe的kvHead维度470+ COMPRESS_TYPE_KVHEAD_ROPE //!< rope场景下压缩key_cache, value_cahe的kvHead维度
469- };471+ };
470- //!472+ //!
471- //! \enum KvCacheCfg473+ //! \enum KvCacheCfg
472- //!474+ //!
473- //! \brief KvCache配置475+ //! \brief KvCache配置
474- //!476+ //!
475- //! \note 默认值为K_CACHE_V_CACHE(0)，传入key_cache和value_cache477+ //! \note 默认值为K_CACHE_V_CACHE(0)，传入key_cache和value_cache
476- //!478+ //!
477- //! \warning 仅在Atlas 800I A2推理产品上支持设置为K_CACHE_V_BYPASS(1)479+ //! \warning 仅在Atlas 800I A2推理产品上支持设置为K_CACHE_V_BYPASS(1)
478- //!480+ //!
479- enum KvCacheCfg : int {481+ enum KvCacheCfg : int {
480- K_CACHE_V_CACHE = 0, //!< 默认值,传入key_cache和value_cache482+ K_CACHE_V_CACHE = 0, //!< 默认值,传入key_cache和value_cache
481- K_CACHE_V_BYPASS, //!< 只传入key_cache483+ K_CACHE_V_BYPASS, //!< 只传入key_cache
482- K_CACHE_V_CACHE_NZ //!< 传入key_cache和value_cache,且为NZ格式484+ K_CACHE_V_CACHE_NZ //!< 传入key_cache和value_cache,且为NZ格式
483- };485+ };
484 486 
485- //! 压缩方式487+ //! 压缩方式
486- CompressType compressType = COMPRESS_TYPE_UNDEFINED;488+ CompressType compressType = COMPRESS_TYPE_UNDEFINED;
487- //! kvcache配置489+ //! kvcache配置
488- KvCacheCfg kvCacheCfg = K_CACHE_V_CACHE;490+ KvCacheCfg kvCacheCfg = K_CACHE_V_CACHE;
489- //!491+ //!
490- //! \brief 预留参数492+ //! \brief 预留参数
491- //!493+ //!
492- uint8_t rsv[16] = {0};494+ uint8_t rsv[16] = {0};
493};495};
494 496 
495//!497//!
496//! \brief 旋转位置编码。hiddenSizeQ必须是hiddenSizeK的整数倍且满足hiddenSizeQ = headDim * headNum。498//! \brief 旋转位置编码。hiddenSizeQ必须是hiddenSizeK的整数倍且满足hiddenSizeQ = headDim * headNum。
497//!499//!
498struct RopeParam {500struct RopeParam {
499- //! \brief rope，旋转系数，对半旋转是2，支持配置2、4或headDim / 2。501+ //! \brief rope，旋转系数，对半旋转是2，支持配置2、4或headDim / 2。
500- int32_t rotaryCoeff = 4;502+ int32_t rotaryCoeff = 4;
501- //! \brief 训练用参数，支持配置0或1503+ //! \brief 训练用参数，支持配置0或1
502- int32_t cosFormat = 0;504+ int32_t cosFormat = 0;
503- //!505+ //!
504- //! \brief 预留参数506+ //! \brief 预留参数
505- //!507+ //!
506- uint8_t rsv[8] = {0};508+ uint8_t rsv[8] = {0};
507};509};
508 510 
509//!511//!
@@ -513,160 +515,159 @@ struct RopeParam {
513//! \param right515//! \param right
514//! \return bool516//! \return bool
515//!517//!
516-inline bool operator==(const RopeParam &left, const RopeParam &right)518+inline bool operator==(const RopeParam& left, const RopeParam& right) {
517-{519+ return left.rotaryCoeff == right.rotaryCoeff && left.cosFormat == right.cosFormat;
518- return left.rotaryCoeff == right.rotaryCoeff && left.cosFormat == right.cosFormat;
519}520}
520 521 
521//!522//!
522//! \brief KVCache+KVCache+Muls+FlashAttention.523//! \brief KVCache+KVCache+Muls+FlashAttention.
523//!524//!
524struct SelfAttentionParam {525struct SelfAttentionParam {
525- //!526+ //!
526- //! \enum CalcType527+ //! \enum CalcType
527- //!528+ //!
528- //! \brief 计算类型529+ //! \brief 计算类型
529- //!530+ //!
530- enum CalcType : int {531+ enum CalcType : int {
531- UNDEFINED = 0, //!< decoder&encoder for flashAttention532+ UNDEFINED = 0, //!< decoder&encoder for flashAttention
532- ENCODER, //!< encoder for flashAttention533+ ENCODER, //!< encoder for flashAttention
533- DECODER, //!< decoder for flashAttention534+ DECODER, //!< decoder for flashAttention
534- PA_ENCODER, //!< encoder for pagedAttention535+ PA_ENCODER, //!< encoder for pagedAttention
535- PREFIX_ENCODER, //!< prefix encoder for flashAttention536+ PREFIX_ENCODER, //!< prefix encoder for flashAttention
536- };537+ };
537- //!538+ //!
538- //! \enum KernelType539+ //! \enum KernelType
539- //!540+ //!
540- //! \brief 算子内核精度类型541+ //! \brief 算子内核精度类型
541- //!542+ //!
542- enum KernelType : int {543+ enum KernelType : int {
543- KERNELTYPE_DEFAULT = 0, //!< i:float16, bmm:float16, o:float16544+ KERNELTYPE_DEFAULT = 0, //!< i:float16, bmm:float16, o:float16
544- KERNELTYPE_HIGH_PRECISION, //!< i:float16, bmm:float, o:float16545+ KERNELTYPE_HIGH_PRECISION, //!< i:float16, bmm:float, o:float16
545- KERNELTYPE_EXP_M8V2, //!< i:float16, bmm:float16, exp:m8v2, softmax:default546+ KERNELTYPE_EXP_M8V2, //!< i:float16, bmm:float16, exp:m8v2, softmax:default
546- };547+ };
547- //!548+ //!
548- //! \enum ClampType549+ //! \enum ClampType
549- //!550+ //!
550- //! \brief clamp类型551+ //! \brief clamp类型
551- //!552+ //!
552- enum ClampType : int {553+ enum ClampType : int {
553- CLAMP_TYPE_UNDEFINED = 0, //!< 不做clamp554+ CLAMP_TYPE_UNDEFINED = 0, //!< 不做clamp
554- CLAMP_TYPE_MIN_MAX //!< 做clamp，同时指定最大最小值555+ CLAMP_TYPE_MIN_MAX //!< 做clamp，同时指定最大最小值
555- };556+ };
556- //!557+ //!
557- //! \enum MaskType558+ //! \enum MaskType
558- //!559+ //!
559- //! \brief mask类型560+ //! \brief mask类型
560- //!561+ //!
561- enum MaskType : int {562+ enum MaskType : int {
562- MASK_TYPE_UNDEFINED = 0, //!< 默认值，全0mask563+ MASK_TYPE_UNDEFINED = 0, //!< 默认值，全0mask
563- MASK_TYPE_NORM, //!< 倒三角mask564+ MASK_TYPE_NORM, //!< 倒三角mask
564- MASK_TYPE_ALIBI, //!< alibi mask565+ MASK_TYPE_ALIBI, //!< alibi mask
565- MASK_TYPE_NORM_COMPRESS, //!< 倒三角压缩mask566+ MASK_TYPE_NORM_COMPRESS, //!< 倒三角压缩mask
566- MASK_TYPE_ALIBI_COMPRESS, //!< alibi压缩mask567+ MASK_TYPE_ALIBI_COMPRESS, //!< alibi压缩mask
567- MASK_TYPE_ALIBI_COMPRESS_SQRT, //!< alibi压缩开平方mask568+ MASK_TYPE_ALIBI_COMPRESS_SQRT, //!< alibi压缩开平方mask
568- MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN, //!< alibi压缩mask左对齐,只支持Atlas 800I A2推理产品569+ MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN, //!< alibi压缩mask左对齐,只支持Atlas 800I A2推理产品
569- MASK_TYPE_SLIDING_WINDOW_NORM, //!< sliding window attention mask570+ MASK_TYPE_SLIDING_WINDOW_NORM, //!< sliding window attention mask
570- MASK_TYPE_SLIDING_WINDOW_COMPRESS //!< sliding window attention压缩mask571+ MASK_TYPE_SLIDING_WINDOW_COMPRESS //!< sliding window attention压缩mask
571- };572+ };
572- //!573+ //!
573- //! \enum KvCacheCfg574+ //! \enum KvCacheCfg
574- //!575+ //!
575- //! \brief KvCache配置,不支持calcType为PA_ENCODER576+ //! \brief KvCache配置,不支持calcType为PA_ENCODER
576- //!577+ //!
577- enum KvCacheCfg : int {578+ enum KvCacheCfg : int {
578- K_CACHE_V_CACHE = 0, //!< 默认值,进行kvcache处理579+ K_CACHE_V_CACHE = 0, //!< 默认值,进行kvcache处理
579- K_BYPASS_V_BYPASS, //!< 直接传入kvcache580+ K_BYPASS_V_BYPASS, //!< 直接传入kvcache
580- };581+ };
581- //!582+ //!
582- //! \enum ScaleType583+ //! \enum ScaleType
583- //!584+ //!
584- //! \brief The type values of ScaleType.585+ //! \brief The type values of ScaleType.
585- //!586+ //!
586- enum ScaleType : int {587+ enum ScaleType : int {
587- SCALE_TYPE_TOR = 0, //!< 默认值，不开启LogN缩放588+ SCALE_TYPE_TOR = 0, //!< 默认值，不开启LogN缩放
588- SCALE_TYPE_LOGN, //!< 注意力使用LogN缩放，quantType只能是0589+ SCALE_TYPE_LOGN, //!< 注意力使用LogN缩放，quantType只能是0
589- SCALE_TYPE_MAX //!< 边界值，仅用于判断是否出界590+ SCALE_TYPE_MAX //!< 边界值，仅用于判断是否出界
590- };591+ };
591 592 
592- //! \enum QuantType593+ //! \enum QuantType
593- //!594+ //!
594- //! \brief quant类型595+ //! \brief quant类型
595- //!596+ //!
596- enum QuantType : int {597+ enum QuantType : int {
597- TYPE_QUANT_UNDEFINED = 0, //!< 默认值，不与量化融合，此时q，k，v为bf16/float16598+ TYPE_QUANT_UNDEFINED = 0, //!< 默认值，不与量化融合，此时q，k，v为bf16/float16
598- TYPE_QUANT_UNQUANT = 0, //!< 默认值，不与量化融合，此时q，k，v为bf16/float16599+ TYPE_QUANT_UNQUANT = 0, //!< 默认值，不与量化融合，此时q，k，v为bf16/float16
599- TYPE_DEQUANT_FUSION = 1, //!< 与反量化融合, 预留类型，当前不能够取此值。600+ TYPE_DEQUANT_FUSION = 1, //!< 与反量化融合, 预留类型，当前不能够取此值。
600- TYPE_QUANT_QKV_OFFLINE = 2, //!< 离线INT8量化, 只支持Atlas 800I A2推理产品601+ TYPE_QUANT_QKV_OFFLINE = 2, //!< 离线INT8量化, 只支持Atlas 800I A2推理产品
601- TYPE_QUANT_QKV_ONLINE = 3 //!< 在线INT8量化, 只支持Atlas 800I A2推理产品602+ TYPE_QUANT_QKV_ONLINE = 3 //!< 在线INT8量化, 只支持Atlas 800I A2推理产品
602- };603+ };
603- //!604+ //!
604- //! \enum CacheType605+ //! \enum CacheType
605- //!606+ //!
606- //! \brief cache内部排布类型, 为CACHE_TYPE_SWA开启SWA KVCache优化，只储存后windowSize个token的KVCache，607+ //! \brief cache内部排布类型, 为CACHE_TYPE_SWA开启SWA KVCache优化，只储存后windowSize个token的KVCache，
607- //! 控制KVCache的长度不超过windowSize, 以此减少显存占用608+ //! 控制KVCache的长度不超过windowSize, 以此减少显存占用
608- //!609+ //!
609- enum CacheType : int8_t {610+ enum CacheType : int8_t {
610- CACHE_TYPE_NORM = 0, //!< 正常cache611+ CACHE_TYPE_NORM = 0, //!< 正常cache
611- CACHE_TYPE_SWA = 1 //!< 固定长度cache612+ CACHE_TYPE_SWA = 1 //!< 固定长度cache
612- };613+ };
613- //!614+ //!
614- //! 量化类型(只支持PA_ENCODER)：615+ //! 量化类型(只支持PA_ENCODER)：
615- //! 当值为TYPE_QUANT_QKV_OFFLINE或TYPE_QUANT_QKV_ONLINE时q，k，v为int8。key,value的headsize等长，范围为（0, 256]，616+ //! 当值为TYPE_QUANT_QKV_OFFLINE或TYPE_QUANT_QKV_ONLINE时q，k，v为int8。key,value的headsize等长，范围为（0, 256]，
616- //! 且32对齐。outdatatype需要配置，只能是ACL_FLOAT16或ACL_BF16。inputLayout只支持TYPE_BSND，calcType只能为PA_ENCODER。617+ //! 且32对齐。outdatatype需要配置，只能是ACL_FLOAT16或ACL_BF16。inputLayout只支持TYPE_BSND，calcType只能为PA_ENCODER。
617- QuantType quantType = TYPE_QUANT_UNQUANT;618+ QuantType quantType = TYPE_QUANT_UNQUANT;
618 619 
619- //! output数据类型：只支持PA_ENCODER,且QuantType不为TYPE_QUANT_UNQUANT（格式为aclDataType）620+ //! output数据类型：只支持PA_ENCODER,且QuantType不为TYPE_QUANT_UNQUANT（格式为aclDataType）
620- aclDataType outDataType = ACL_DT_UNDEFINED;621+ aclDataType outDataType = ACL_DT_UNDEFINED;
621 622 
622- //! query头大小, 需大于0623+ //! query头大小, 需大于0
623- int32_t headNum = 0;624+ int32_t headNum = 0;
624- //! kv头数量, 该值需要用户根据使用的模型实际情况传入625+ //! kv头数量, 该值需要用户根据使用的模型实际情况传入
625- //! kvHeadNum = 0时，keyCache的k_head_num，valueCache的v_head_num与query的num_heads一致，均为num_heads的数值626+ //! kvHeadNum = 0时，keyCache的k_head_num，valueCache的v_head_num与query的num_heads一致，均为num_heads的数值
626- //! kvHeadNum != 0时，keyCache的k_head_num， valueCache的v_head_num与kvHeadNum值相同627+ //! kvHeadNum != 0时，keyCache的k_head_num， valueCache的v_head_num与kvHeadNum值相同
627- int32_t kvHeadNum = 0;628+ int32_t kvHeadNum = 0;
628- //! query缩放系数629+ //! query缩放系数
629- float qScale = 1;630+ float qScale = 1;
630- //! 算子tor值, 在Q*K^T后乘631+ //! 算子tor值, 在Q*K^T后乘
631- float qkScale = 1;632+ float qkScale = 1;
632- //! 是否开启动态batch633+ //! 是否开启动态batch
633- bool batchRunStatusEnable = false;634+ bool batchRunStatusEnable = false;
634- //! 是否开启倒三角优化, 只有mask为倒三角的时候才能开启优化635+ //! 是否开启倒三角优化, 只有mask为倒三角的时候才能开启优化
635- uint32_t isTriuMask = 0;636+ uint32_t isTriuMask = 0;
636- //! 计算类型637+ //! 计算类型
637- CalcType calcType = UNDEFINED;638+ CalcType calcType = UNDEFINED;
638- //! 内核精度类型639+ //! 内核精度类型
639- KernelType kernelType = KERNELTYPE_DEFAULT;640+ KernelType kernelType = KERNELTYPE_DEFAULT;
640- //! clamp类型641+ //! clamp类型
641- ClampType clampType = CLAMP_TYPE_UNDEFINED;642+ ClampType clampType = CLAMP_TYPE_UNDEFINED;
642- //! clamp功能最小值643+ //! clamp功能最小值
643- float clampMin = 0;644+ float clampMin = 0;
644- //! clamp功能最大值645+ //! clamp功能最大值
645- float clampMax = 0;646+ float clampMax = 0;
646- //! mask类型647+ //! mask类型
647- MaskType maskType = MASK_TYPE_UNDEFINED;648+ MaskType maskType = MASK_TYPE_UNDEFINED;
648- //! kvcache配置649+ //! kvcache配置
649- KvCacheCfg kvcacheCfg = K_CACHE_V_CACHE;650+ KvCacheCfg kvcacheCfg = K_CACHE_V_CACHE;
650- //! scale类型651+ //! scale类型
651- ScaleType scaleType = SCALE_TYPE_TOR;652+ ScaleType scaleType = SCALE_TYPE_TOR;
652- //! 数据排布格式默认为BSND653+ //! 数据排布格式默认为BSND
653- InputLayout inputLayout = TYPE_BSND;654+ InputLayout inputLayout = TYPE_BSND;
654- //! \brief 大于0时开启MLA合并kvcache功能，表示kv合并传入时v的head_size655+ //! \brief 大于0时开启MLA合并kvcache功能，表示kv合并传入时v的head_size
655- //! \note 默认值为0656+ //! \note 默认值为0
656- //! \warning 取值范围为[0,576]657+ //! \warning 取值范围为[0,576]
657- uint32_t mlaVHeadSize = 0;658+ uint32_t mlaVHeadSize = 0;
658- //! \brief cache内部排布，开启SWA特性并设置为CACHE_TYPE_SWA可以开启SWA cache优化659+ //! \brief cache内部排布，开启SWA特性并设置为CACHE_TYPE_SWA可以开启SWA cache优化
659- //! \note 默认值为CACHE_TYPE_NORM660+ //! \note 默认值为CACHE_TYPE_NORM
660- //! \warning 只有开启SWA特性后才可以是CACHE_TYPE_SWA661+ //! \warning 只有开启SWA特性后才可以是CACHE_TYPE_SWA
661- CacheType cacheType = CACHE_TYPE_NORM;662+ CacheType cacheType = CACHE_TYPE_NORM;
662- //! \brief windowSize大于0时开启SWA特性，开启SWA特性后表示sliding window 大小663+ //! \brief windowSize大于0时开启SWA特性，开启SWA特性后表示sliding window 大小
663- //! \note 默认值为0664+ //! \note 默认值为0
664- //! \warning windowSize大于0时需要将maskType设置为MASK_TYPE_SLIDING_WINDOW_NORM或MASK_TYPE_SLIDING_WINDOW_COMPRESS665+ //! \warning windowSize大于0时需要将maskType设置为MASK_TYPE_SLIDING_WINDOW_NORM或MASK_TYPE_SLIDING_WINDOW_COMPRESS
665- uint32_t windowSize = 0;666+ uint32_t windowSize = 0;
666- //!667+ //!
667- //! \brief 预留参数668+ //! \brief 预留参数
668- //!669+ //!
669- uint8_t rsv[64] = {0};670+ uint8_t rsv[64] = {0};
670};671};
671 672 
672//!673//!
@@ -678,72 +679,72 @@ struct SelfAttentionParam {
678//! 输入x、y对应维度的对应值要求相同或至少其中一个为1679//! 输入x、y对应维度的对应值要求相同或至少其中一个为1
679//!680//!
680struct ElewiseParam {681struct ElewiseParam {
681- //!682+ //!
682- //! \enum ElewiseType683+ //! \enum ElewiseType
683- //!684+ //!
684- //! \brief 计算类型685+ //! \brief 计算类型
685- //!686+ //!
686- enum ElewiseType : int {687+ enum ElewiseType : int {
687- ELEWISE_UNDEFINED = 0, //!< 默认值，未定义688+ ELEWISE_UNDEFINED = 0, //!< 默认值，未定义
688- ELEWISE_CAST, //!< 数据类型转换689+ ELEWISE_CAST, //!< 数据类型转换
689- ELEWISE_MULS, //!< 向量逐元素乘值690+ ELEWISE_MULS, //!< 向量逐元素乘值
690- ELEWISE_COS, //!< 逐元素计算余弦值691+ ELEWISE_COS, //!< 逐元素计算余弦值
691- ELEWISE_SIN, //!< 逐元素计算正弦值692+ ELEWISE_SIN, //!< 逐元素计算正弦值
692- ELEWISE_NEG, //!< 逐元素取相反数693+ ELEWISE_NEG, //!< 逐元素取相反数
693- ELEWISE_QUANT, //!< 量化, 仅在Atlas 800I A2推理产品上支持694+ ELEWISE_QUANT, //!< 量化, 仅在Atlas 800I A2推理产品上支持
694- ELEWISE_LOGICAL_NOT, //!< 逐元素逻辑非695+ ELEWISE_LOGICAL_NOT, //!< 逐元素逻辑非
695- ELEWISE_ADD, //!< 逐元素相加696+ ELEWISE_ADD, //!< 逐元素相加
696- ELEWISE_MUL, //!< 向量与向量逐元素相乘697+ ELEWISE_MUL, //!< 向量与向量逐元素相乘
697- ELEWISE_REALDIV, //!< 向量与向量逐元素相除698+ ELEWISE_REALDIV, //!< 向量与向量逐元素相除
698- ELEWISE_LOGICAL_AND, //!< 逐元素逻辑与699+ ELEWISE_LOGICAL_AND, //!< 逐元素逻辑与
699- ELEWISE_LOGICAL_OR, //!< 逐元素逻辑或700+ ELEWISE_LOGICAL_OR, //!< 逐元素逻辑或
700- ELEWISE_LESS, //!< 逐元素判断是否小于701+ ELEWISE_LESS, //!< 逐元素判断是否小于
701- ELEWISE_GREATER, //!< 逐元素判断是否大于702+ ELEWISE_GREATER, //!< 逐元素判断是否大于
702- ELEWISE_SUB, //!< 逐元素相减703+ ELEWISE_SUB, //!< 逐元素相减
703- ELEWISE_EQUAL, //!< 逐元素判断是否相等704+ ELEWISE_EQUAL, //!< 逐元素判断是否相等
704- ELEWISE_QUANT_PER_CHANNEL, //!< 每个通道量化705+ ELEWISE_QUANT_PER_CHANNEL, //!< 每个通道量化
705- ELEWISE_DEQUANT_PER_CHANNEL, //!< 每个通道反量化706+ ELEWISE_DEQUANT_PER_CHANNEL, //!< 每个通道反量化
706- ELEWISE_DYNAMIC_QUANT, //!< 逐行动态量化707+ ELEWISE_DYNAMIC_QUANT, //!< 逐行动态量化
707- ELEWISE_TANH, //!< 逐元素计算双曲正切值708+ ELEWISE_TANH, //!< 逐元素计算双曲正切值
708- ELEWISE_TYPE_MAX //!< 边界值，仅用于判断是否出界，所有情况不能取该值709+ ELEWISE_TYPE_MAX //!< 边界值，仅用于判断是否出界，所有情况不能取该值
709- };710+ };
710 711 
711- //! 量化（非每通道）所需参数712+ //! 量化（非每通道）所需参数
712- struct QuantParam {713+ struct QuantParam {
713- //! 量化的步长714+ //! 量化的步长
714- float inputScale = 1.0f;715+ float inputScale = 1.0f;
715- //! 动态量化的是否为非对称量化716+ //! 动态量化的是否为非对称量化
716- bool asymmetric = false; //!< false : symmetric，true : asymmetric717+ bool asymmetric = false; //!< false : symmetric，true : asymmetric
717- //! 量化的偏移度718+ //! 量化的偏移度
718- int inputOffset = 0;719+ int inputOffset = 0;
719- //!
720- //! \brief 预留参数
721- //!
722- uint8_t rsv[20] = {0};
723- };
724- 
725- //! 向量乘值所需参数
726- struct MulsParam {
727- //! 向量乘的值
728- float varAttr = 0.0f;
729- //!
730- //! \brief 预留参数
731- //!
732- uint8_t rsv[12] = {0};
733- };
734- 
735- //! 计算方式
736- ElewiseType elewiseType = ELEWISE_UNDEFINED;
737- //! 量化参数
738- QuantParam quantParam;
739- //! 乘值参数
740- MulsParam mulsParam;
741- //! 指定数据类型转换输出的数据类型
742- aclDataType outTensorType = ACL_DT_UNDEFINED;
743 //!720 //!
744 //! \brief 预留参数721 //! \brief 预留参数
745 //!722 //!
746- uint8_t rsv[8] = {0};723+ uint8_t rsv[20] = {0};
724+ };
725+ 
726+ //! 向量乘值所需参数
727+ struct MulsParam {
728+ //! 向量乘的值
729+ float varAttr = 0.0f;
730+ //!
731+ //! \brief 预留参数
732+ //!
733+ uint8_t rsv[12] = {0};
734+ };
735+ 
736+ //! 计算方式
737+ ElewiseType elewiseType = ELEWISE_UNDEFINED;
738+ //! 量化参数
739+ QuantParam quantParam;
740+ //! 乘值参数
741+ MulsParam mulsParam;
742+ //! 指定数据类型转换输出的数据类型
743+ aclDataType outTensorType = ACL_DT_UNDEFINED;
744+ //!
745+ //! \brief 预留参数
746+ //!
747+ uint8_t rsv[8] = {0};
747};748};
748 749 
749} // namespace infer750} // namespace infer
Mop_plugin/third_party/atb/inc/operation.h+70-63
@@ -36,70 +36,73 @@ namespace atb {
36//! 该接口类定义了算子准备与执行的需要的一系列的接口，通过创建Operation可以执行算子36//! 该接口类定义了算子准备与执行的需要的一系列的接口，通过创建Operation可以执行算子
37//!37//!
38class Operation {38class Operation {
39-public:39+ public:
40- //! \brief 默认构造函数.40+ //! \brief 默认构造函数.
41- Operation() = default;41+ Operation() = default;
42 42 
43- //! \brief 默认析构函数.43+ //! \brief 默认析构函数.
44- virtual ~Operation() = default;44+ virtual ~Operation() = default;
45- //!45+ //!
46- //! \brief 获取创建的Operation的名字46+ //! \brief 获取创建的Operation的名字
47- //!47+ //!
48- //! \return 返回字符串48+ //! \return 返回字符串
49- //!49+ //!
50- virtual std::string GetName() const = 0;50+ virtual std::string GetName() const = 0;
51 51 
52- //!52+ //!
53- //! \brief 根据输入Tensor描述信息推导出输出Tensor的描述信息。53+ //! \brief 根据输入Tensor描述信息推导出输出Tensor的描述信息。
54- //!54+ //!
55- //! \param inTensorDescs 存放所有输入tensor描述信息的SVector55+ //! \param inTensorDescs 存放所有输入tensor描述信息的SVector
56- //! \param outTensorDescs 存放所有输出tensor描述信息的SVector56+ //! \param outTensorDescs 存放所有输出tensor描述信息的SVector
57- //!57+ //!
58- //! \return 状态值，如果成功，返回NO_ERROR58+ //! \return 状态值，如果成功，返回NO_ERROR
59- //!59+ //!
60- virtual Status InferShape(const SVector<TensorDesc> &inTensorDescs, SVector<TensorDesc> &outTensorDescs) const = 0;60+ virtual Status InferShape(const SVector<TensorDesc>& inTensorDescs, SVector<TensorDesc>& outTensorDescs) const = 0;
61 61 
62- //!62+ //!
63- //! \brief 获取Op/GraphOp输入Tensor个数接口。63+ //! \brief 获取Op/GraphOp输入Tensor个数接口。
64- //!64+ //!
65- //! \return 整数值65+ //! \return 整数值
66- //!66+ //!
67- virtual uint32_t GetInputNum() const = 0;67+ virtual uint32_t GetInputNum() const = 0;
68 68 
69- //!69+ //!
70- //! \brief 获取Op/GraphOp输出Tensor个数接口。70+ //! \brief 获取Op/GraphOp输出Tensor个数接口。
71- //!71+ //!
72- //! \return 整数值72+ //! \return 整数值
73- //!73+ //!
74- virtual uint32_t GetOutputNum() const = 0;74+ virtual uint32_t GetOutputNum() const = 0;
75 75 
76- //!76+ //!
77- //! \brief Operation执行前的一系列准备工作77+ //! \brief Operation执行前的一系列准备工作
78- //!78+ //!
79- //! 主要是计算Operation执行过程需要分配的内存空间workspaceSize79+ //! 主要是计算Operation执行过程需要分配的内存空间workspaceSize
80- //!80+ //!
81- //! \param variantPack 输入与输出Tensor81+ //! \param variantPack 输入与输出Tensor
82- //! \param workspaceSize 获取Operation执行需要分配的内存空间82+ //! \param workspaceSize 获取Operation执行需要分配的内存空间
83- //! \param context Operation执行准备工作所在的上下文83+ //! \param context Operation执行准备工作所在的上下文
84- //!84+ //!
85- //! \return 状态值，如果成功，返回NO_ERROR85+ //! \return 状态值，如果成功，返回NO_ERROR
86- //!86+ //!
87- virtual Status Setup(const VariantPack &variantPack, uint64_t &workspaceSize, Context *context) = 0;87+ virtual Status Setup(const VariantPack& variantPack, uint64_t& workspaceSize, Context* context) = 0;
88 88 
89- //!89+ //!
90- //! \brief Operation执行的流程90+ //! \brief Operation执行的流程
91- //!91+ //!
92- //! 根据setup过程中得到的workspaceSize为Operation执行分配实际的内存，并执行Operation92+ //! 根据setup过程中得到的workspaceSize为Operation执行分配实际的内存，并执行Operation
93- //!93+ //!
94- //! \param variantPack 输入与输出Tensor94+ //! \param variantPack 输入与输出Tensor
95- //! \param workspace Operation执行分配的内存地址95+ //! \param workspace Operation执行分配的内存地址
96- //! \param workspaceSize Operation执行需要分配的内存空间96+ //! \param workspaceSize Operation执行需要分配的内存空间
97- //! \param context Operation执行所在的上下文97+ //! \param context Operation执行所在的上下文
98- //!98+ //!
99- //! \return 状态值，如果成功，返回NO_ERROR99+ //! \return 状态值，如果成功，返回NO_ERROR
100- //!100+ //!
101- virtual Status Execute(const VariantPack &variantPack, uint8_t *workspace, uint64_t workspaceSize,101+ virtual Status Execute(
102- Context *context) = 0;102+ const VariantPack& variantPack,
103+ uint8_t* workspace,
104+ uint64_t workspaceSize,
105+ Context* context) = 0;
103};106};
104 107 
105//!108//!
@@ -110,7 +113,8 @@ public:
110//!113//!
111//! \return 状态值，如果成功，返回NO_ERROR114//! \return 状态值，如果成功，返回NO_ERROR
112//!115//!
113-template <typename OpParam> Status CreateOperation(const OpParam &opParam, Operation **operation);116+template <typename OpParam>
117+Status CreateOperation(const OpParam& opParam, Operation** operation);
114 118 
115//!119//!
116//! \brief 销毁Operation120//! \brief 销毁Operation
@@ -119,9 +123,10 @@ template <typename OpParam> Status CreateOperation(const OpParam &opParam, Opera
119//!123//!
120//! \return 状态值，如果成功，返回NO_ERROR124//! \return 状态值，如果成功，返回NO_ERROR
121//!125//!
122-//! \note 调用CreateOperation接口创建Operation，执行完Operation后需要调用DestroyOperation接口进行销毁。否则将导致内存泄漏。126+//! \note
127+//! 调用CreateOperation接口创建Operation，执行完Operation后需要调用DestroyOperation接口进行销毁。否则将导致内存泄漏。
123//!128//!
124-Status DestroyOperation(Operation *operation);129+Status DestroyOperation(Operation* operation);
125 130 
126//!131//!
127//! \brief 拷贝Operation的Param参数132//! \brief 拷贝Operation的Param参数
@@ -131,7 +136,8 @@ Status DestroyOperation(Operation *operation);
131//!136//!
132//! \return 状态值，如果成功，返回NO_ERROR137//! \return 状态值，如果成功，返回NO_ERROR
133//!138//!
134-template <typename OpParam> Status CloneOperationParam(const Operation *operation, OpParam &opParam);139+template <typename OpParam>
140+Status CloneOperationParam(const Operation* operation, OpParam& opParam);
135 141 
136//!142//!
137//! \brief 更新Operation的Param参数143//! \brief 更新Operation的Param参数
@@ -141,7 +147,8 @@ template <typename OpParam> Status CloneOperationParam(const Operation *operatio
141//!147//!
142//! \return 状态值，如果成功，返回NO_ERROR148//! \return 状态值，如果成功，返回NO_ERROR
143//!149//!
144-template <typename OpParam> Status UpdateOperationParam(Operation *operation, const OpParam &opParam);150+template <typename OpParam>
151+Status UpdateOperationParam(Operation* operation, const OpParam& opParam);
145 152 
146} // namespace atb153} // namespace atb
147#endif154#endif
Mop_plugin/third_party/atb/libs/atb.cpp+9-8
@@ -3,19 +3,20 @@
3#include "operation.h"3#include "operation.h"
4 4 
5namespace atb {5namespace atb {
6-Status CreateContext(Context **context) {6+Status CreateContext(Context** context) {
7- return 0;7+ return 0;
8}8}
9 9 
10-Status DestroyContext(Context *context) {10+Status DestroyContext(Context* context) {
11- return 0;11+ return 0;
12}12}
13 13 
14-template <typename OpParam> Status CreateOperation(const OpParam &opParam, Operation **operation) {14+template <typename OpParam>
15- return 0;15+Status CreateOperation(const OpParam& opParam, Operation** operation) {
16+ return 0;
16}17}
17 18 
18-Status DestroyOperation(Operation *operation) {19+Status DestroyOperation(Operation* operation) {
19- return 0;20+ return 0;
20}21}
21} // namespace atb22} // namespace atb
Mop_plugin/utils/KernelNpuNewParams.cpp+21-20
@@ -14,34 +14,35 @@
14 14 
15namespace op_infer {15namespace op_infer {
16 16 
17-int64_t npu_gelu_approximate_mode(c10::string_view approximate)17+int64_t npu_gelu_approximate_mode(c10::string_view approximate) {
18-{18+ std::string approximate_str = std::string(approximate);
19- std::string approximate_str = std::string(approximate);19+ TORCH_CHECK(
20- TORCH_CHECK(approximate_str == "tanh" || approximate_str == "none",20+ approximate_str == "tanh" || approximate_str == "none",
21- "NPU error, approximate argument must be either none or tanh.", OPS_ERROR(ErrCode::PARAM));21+ "NPU error, approximate argument must be either none or tanh.",
22- int64_t approximate_mode = approximate_str == "tanh" ? 1 : 0;22+ OPS_ERROR(ErrCode::PARAM));
23- return approximate_mode;23+ int64_t approximate_mode = approximate_str == "tanh" ? 1 : 0;
24+ return approximate_mode;
24}25}
25 26 
26-std::string npu_gelu_approximate_str(c10::string_view approximate)27+std::string npu_gelu_approximate_str(c10::string_view approximate) {
27-{28+ std::string approximate_str = std::string(approximate);
28- std::string approximate_str = std::string(approximate);29+ TORCH_CHECK(
29- TORCH_CHECK(approximate_str == "tanh" || approximate_str == "none",30+ approximate_str == "tanh" || approximate_str == "none",
30- "NPU error, approximate argument must be either none or tanh.", OPS_ERROR(ErrCode::PARAM));31+ "NPU error, approximate argument must be either none or tanh.",
31- return approximate_str;32+ OPS_ERROR(ErrCode::PARAM));
33+ return approximate_str;
32}34}
33 35 
34bool npu_add_rms_norm_quant_param_check(36bool npu_add_rms_norm_quant_param_check(
35 c10::optional<at::Tensor> scales2,37 c10::optional<at::Tensor> scales2,
36 c10::optional<at::Tensor> zero_points2,38 c10::optional<at::Tensor> zero_points2,
37 int64_t axis,39 int64_t axis,
38- bool div_mode)40+ bool div_mode) {
39-{41+ TORCH_CHECK(!scales2.has_value(), "scales2 only support None.", OPS_ERROR(ErrCode::PARAM));
40- TORCH_CHECK(!scales2.has_value(), "scales2 only support None.", OPS_ERROR(ErrCode::PARAM));42+ TORCH_CHECK(!zero_points2.has_value(), "zero_points2 only support None.", OPS_ERROR(ErrCode::PARAM));
41- TORCH_CHECK(!zero_points2.has_value(), "zero_points2 only support None.", OPS_ERROR(ErrCode::PARAM));43+ TORCH_CHECK(axis == -1, "axis only support -1.", OPS_ERROR(ErrCode::PARAM));
42- TORCH_CHECK(axis == -1, "axis only support -1.", OPS_ERROR(ErrCode::PARAM));44+ TORCH_CHECK(div_mode == true, "div_mode only support True.", OPS_ERROR(ErrCode::PARAM));
43- TORCH_CHECK(div_mode == true, "div_mode only support True.", OPS_ERROR(ErrCode::PARAM));45+ return true;
44- return true;
45}46}
46 47 
47} // namespace op_infer48} // namespace op_infer
Mop_plugin/utils/KernelNpuNewParams.h+4-1
@@ -25,7 +25,10 @@ namespace op_infer {
25OP_PLUGIN_HIDDEN int64_t npu_gelu_approximate_mode(c10::string_view approximate);25OP_PLUGIN_HIDDEN int64_t npu_gelu_approximate_mode(c10::string_view approximate);
26OP_PLUGIN_HIDDEN std::string npu_gelu_approximate_str(c10::string_view approximate);26OP_PLUGIN_HIDDEN std::string npu_gelu_approximate_str(c10::string_view approximate);
27OP_PLUGIN_HIDDEN bool npu_add_rms_norm_quant_param_check(27OP_PLUGIN_HIDDEN bool npu_add_rms_norm_quant_param_check(
28- c10::optional<at::Tensor> scales2, c10::optional<at::Tensor> zero_points2, int64_t axis, bool div_mode);28+ c10::optional<at::Tensor> scales2,
29+ c10::optional<at::Tensor> zero_points2,
30+ int64_t axis,
31+ bool div_mode);
29 32 
30} // namespace op_infer33} // namespace op_infer
31 34 
Mop_plugin/utils/KernelNpuOutputDtype.cpp+101-98
@@ -21,129 +21,132 @@
21#include "op_plugin/OpApiInterface.h"21#include "op_plugin/OpApiInterface.h"
22namespace op_infer {22namespace op_infer {
23 23 
24-at::ScalarType angle_out_dtype(const at::Tensor& self)24+at::ScalarType angle_out_dtype(const at::Tensor& self) {
25-{25+ auto out_dtype = self.scalar_type();
26- auto out_dtype = self.scalar_type();26+ if (self.is_complex()) {
27- if (self.is_complex()) {27+ out_dtype = self.scalar_type() == at::kComplexFloat ? at::kFloat : at::kDouble;
28- out_dtype = self.scalar_type() == at::kComplexFloat ? at::kFloat : at::kDouble;28+ } else if (at::isIntegralType(out_dtype, true)) {
29- } else if (at::isIntegralType(out_dtype, true)) {29+ out_dtype = at::kFloat;
30- out_dtype = at::kFloat;30+ }
31- }31+ return out_dtype;
32- return out_dtype;
33}32}
34 33 
35-at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle)34+at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle) {
36-{35+ at::ScalarType high_type = at::native::result_type(abs, angle);
37- at::ScalarType high_type = at::native::result_type(abs, angle);36+ if (high_type == at::ScalarType::Float) {
38- if (high_type == at::ScalarType::Float) {37+ high_type = at::ScalarType::ComplexFloat;
39- high_type = at::ScalarType::ComplexFloat;38+ } else if (high_type == at::ScalarType::Double) {
40- } else if (high_type == at::ScalarType::Double) {39+ high_type = at::ScalarType::ComplexDouble;
41- high_type = at::ScalarType::ComplexDouble;40+ } else if (high_type == at::ScalarType::Half) {
42- } else if (high_type == at::ScalarType::Half) {41+ high_type = at::ScalarType::ComplexHalf;
43- high_type = at::ScalarType::ComplexHalf;42+ }
44- }43+ return high_type;
45- return high_type;
46}44}
47 45 
48-at::ScalarType npu_group_norm_silu_dst_type(const at::Tensor& input, const c10::optional<at::Tensor>& weight,46+at::ScalarType npu_group_norm_silu_dst_type(
49- const c10::optional<at::Tensor>& bias)47+ const at::Tensor& input,
50-{48+ const c10::optional<at::Tensor>& weight,
51- at::native::ResultTypeState state = {};49+ const c10::optional<at::Tensor>& bias) {
52- state = at::native::update_result_type_state(input, state);50+ at::native::ResultTypeState state = {};
51+ state = at::native::update_result_type_state(input, state);
53 52 
54- if (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950) {53+ if (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950) {
55- if (weight.has_value()) {54+ if (weight.has_value()) {
56- state = at::native::update_result_type_state(weight.value(), state);55+ state = at::native::update_result_type_state(weight.value(), state);
57- } else if (bias.has_value()) {56+ } else if (bias.has_value()) {
58- state = at::native::update_result_type_state(bias.value(), state);57+ state = at::native::update_result_type_state(bias.value(), state);
59- }
60 }58 }
59+ }
61 60 
62- return at::native::result_type(state);61+ return at::native::result_type(state);
63}62}
64 63 
65-at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype)64+at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype) {
66-{65+ at::ScalarType dst_type = c10::value_or_else(dst_dtype, [] { return at::ScalarType::Char; });
67- at::ScalarType dst_type = c10::value_or_else(dst_dtype, [] {return at::ScalarType::Char;});66+ if (dst_type == at::kQInt8) {
68- if (dst_type == at::kQInt8) {67+ dst_type = at::kChar;
69- dst_type = at::kChar;68+ }
70- }69+ TORCH_CHECK(
71- TORCH_CHECK(dst_type == at::ScalarType::Char || dst_type == at::ScalarType::QUInt4x2,70+ dst_type == at::ScalarType::Char || dst_type == at::ScalarType::QUInt4x2,
72- "dst_dtype must be Int8 or Int4" + OPS_ERROR(ErrCode::TYPE));71+ "dst_dtype must be Int8 or Int4" + OPS_ERROR(ErrCode::TYPE));
73- if (dst_type == at::ScalarType::QUInt4x2) {72+ if (dst_type == at::ScalarType::QUInt4x2) {
74- dst_type = at::ScalarType::Int;73+ dst_type = at::ScalarType::Int;
75- }74+ }
76- return dst_type;75+ return dst_type;
77}76}
78 77 
79-at::ScalarType npu_add_rms_norm_dynamic_quant_y_dtype(c10::optional<at::ScalarType> y_dtype)78+at::ScalarType npu_add_rms_norm_dynamic_quant_y_dtype(c10::optional<at::ScalarType> y_dtype) {
80-{79+ at::ScalarType dtype = c10::value_or_else(y_dtype, [] { return at::ScalarType::Char; });
81- at::ScalarType dtype = c10::value_or_else(y_dtype, [] { return at::ScalarType::Char; });80+ TORCH_CHECK(
82- TORCH_CHECK(dtype == at::ScalarType::Char || dtype == at::ScalarType::QUInt4x2,81+ dtype == at::ScalarType::Char || dtype == at::ScalarType::QUInt4x2,
83- "y_dtype must be torch.int8 or torch.quint4x2, but got ", dtype, OPS_ERROR(ErrCode::PARAM));82+ "y_dtype must be torch.int8 or torch.quint4x2, but got ",
84- // aclnn int4 output uses DT_INT32 (8 int4 packed per int32)83+ dtype,
85- if (dtype == at::ScalarType::QUInt4x2) {84+ OPS_ERROR(ErrCode::PARAM));
86- return at::ScalarType::Int;85+ // aclnn int4 output uses DT_INT32 (8 int4 packed per int32)
87- }86+ if (dtype == at::ScalarType::QUInt4x2) {
88- return dtype;87+ return at::ScalarType::Int;
88+ }
89+ return dtype;
89}90}
90 91 
91-at::ScalarType clamp_out_dtype(const at::Tensor& self, const c10::optional<at::Tensor>& min, const c10::optional<at::Tensor>& max)92+at::ScalarType clamp_out_dtype(
92-{93+ const at::Tensor& self,
93- TORCH_CHECK(min.has_value() || max.has_value(), "torch.clamp:At least one of 'min' or 'max' must be not None!");94+ const c10::optional<at::Tensor>& min,
95+ const c10::optional<at::Tensor>& max) {
96+ TORCH_CHECK(min.has_value() || max.has_value(), "torch.clamp:At least one of 'min' or 'max' must be not None!");
94 97 
95- at::native::ResultTypeState state = {};98+ at::native::ResultTypeState state = {};
96- state = at::native::update_result_type_state(self, state);99+ state = at::native::update_result_type_state(self, state);
97 100 
98- if (!min.has_value()) {101+ if (!min.has_value()) {
99- state = at::native::update_result_type_state(max.value(), state);102+ state = at::native::update_result_type_state(max.value(), state);
100- } else if (!max.has_value()) {103+ } else if (!max.has_value()) {
101- state = at::native::update_result_type_state(min.value(), state);104+ state = at::native::update_result_type_state(min.value(), state);
102- } else {105+ } else {
103- state = at::native::update_result_type_state(max.value(), state);106+ state = at::native::update_result_type_state(max.value(), state);
104- state = at::native::update_result_type_state(min.value(), state);107+ state = at::native::update_result_type_state(min.value(), state);
105- }108+ }
106 109 
107- return at::native::result_type(state);110+ return at::native::result_type(state);
108}111}
109 112 
110-at::ScalarType clamp_scalar_out_dtype(const at::Tensor& self, const c10::optional<at::Scalar>& min, const c10::optional<at::Scalar>& max)113+at::ScalarType clamp_scalar_out_dtype(
111-{114+ const at::Tensor& self,
112- TORCH_CHECK(min.has_value() || max.has_value(), "torch.clamp:At least one of 'min' or 'max' must be not None!");115+ const c10::optional<at::Scalar>& min,
116+ const c10::optional<at::Scalar>& max) {
117+ TORCH_CHECK(min.has_value() || max.has_value(), "torch.clamp:At least one of 'min' or 'max' must be not None!");
113 118 
114- at::native::ResultTypeState state = {};119+ at::native::ResultTypeState state = {};
115- state = at::native::update_result_type_state(self, state);120+ state = at::native::update_result_type_state(self, state);
116 121 
117- if (!min.has_value()) {122+ if (!min.has_value()) {
118- state = at::native::update_result_type_state(max.value(), state);123+ state = at::native::update_result_type_state(max.value(), state);
119- } else if (!max.has_value()) {124+ } else if (!max.has_value()) {
120- state = at::native::update_result_type_state(min.value(), state);125+ state = at::native::update_result_type_state(min.value(), state);
121- } else {126+ } else {
122- state = at::native::update_result_type_state(max.value(), state);127+ state = at::native::update_result_type_state(max.value(), state);
123- state = at::native::update_result_type_state(min.value(), state);128+ state = at::native::update_result_type_state(min.value(), state);
124- }129+ }
125 130 
126- return at::native::result_type(state);131+ return at::native::result_type(state);
127}132}
128 133 
129-at::ScalarType abs_out_dtype(const at::Tensor& self)134+at::ScalarType abs_out_dtype(const at::Tensor& self) {
130-{135+ at::ScalarType output = self.scalar_type();
131- at::ScalarType output = self.scalar_type();136+ if (output == at::ScalarType::ComplexFloat) {
132- if (output == at::ScalarType::ComplexFloat) {137+ output = at::ScalarType::Float;
133- output = at::ScalarType::Float;138+ } else if (output == at::ScalarType::ComplexDouble) {
134- } else if (output == at::ScalarType::ComplexDouble) {139+ output = at::ScalarType::Double;
135- output = at::ScalarType::Double;140+ } else if (output == at::ScalarType::ComplexHalf) {
136- } else if (output == at::ScalarType::ComplexHalf) {141+ output = at::ScalarType::Half;
137- output = at::ScalarType::Half;142+ }
138- }143+ return output;
139- return output;
140}144}
141 145 
142-at::ScalarType npu_moe_distribute_dispatch_setup_out_dtype(c10::optional<int64_t> y_dtype)146+at::ScalarType npu_moe_distribute_dispatch_setup_out_dtype(c10::optional<int64_t> y_dtype) {
143-{147+ auto output_dtype = at::kChar;
144- auto output_dtype = at::kChar;148+ output_dtype = at_npu::native::OpPreparation::convert_to_scalar_type(c10_npu::GetAclDataType(y_dtype.value()));
145- output_dtype = at_npu::native::OpPreparation::convert_to_scalar_type(c10_npu::GetAclDataType(y_dtype.value()));149+ return output_dtype;
146- return output_dtype;
147}150}
148 151 
149} // namespace op_infer152} // namespace op_infer
Mop_plugin/utils/KernelNpuOutputDtype.h+12-3
@@ -24,11 +24,20 @@ namespace op_infer {
24 24 
25OP_PLUGIN_HIDDEN at::ScalarType angle_out_dtype(const at::Tensor& self);25OP_PLUGIN_HIDDEN at::ScalarType angle_out_dtype(const at::Tensor& self);
26OP_PLUGIN_HIDDEN at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle);26OP_PLUGIN_HIDDEN at::ScalarType polar_out_dtype(const at::Tensor& abs, const at::Tensor& angle);
27-OP_PLUGIN_HIDDEN at::ScalarType npu_group_norm_silu_dst_type(const at::Tensor& input, const c10::optional<at::Tensor>& weight, const c10::optional<at::Tensor>& bias);27+OP_PLUGIN_HIDDEN at::ScalarType npu_group_norm_silu_dst_type(
28+ const at::Tensor& input,
29+ const c10::optional<at::Tensor>& weight,
30+ const c10::optional<at::Tensor>& bias);
28OP_PLUGIN_HIDDEN at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype);31OP_PLUGIN_HIDDEN at::ScalarType npu_group_quant_dst_type(c10::optional<at::ScalarType> dst_dtype);
29OP_PLUGIN_HIDDEN at::ScalarType npu_add_rms_norm_dynamic_quant_y_dtype(c10::optional<at::ScalarType> y_dtype);32OP_PLUGIN_HIDDEN at::ScalarType npu_add_rms_norm_dynamic_quant_y_dtype(c10::optional<at::ScalarType> y_dtype);
30-OP_PLUGIN_HIDDEN at::ScalarType clamp_out_dtype(const at::Tensor& self, const c10::optional<at::Tensor>& min, const c10::optional<at::Tensor>& max);33+OP_PLUGIN_HIDDEN at::ScalarType clamp_out_dtype(
31-OP_PLUGIN_HIDDEN at::ScalarType clamp_scalar_out_dtype(const at::Tensor& self, const c10::optional<at::Scalar>& min, const c10::optional<at::Scalar>& max);34+ const at::Tensor& self,
35+ const c10::optional<at::Tensor>& min,
36+ const c10::optional<at::Tensor>& max);
37+OP_PLUGIN_HIDDEN at::ScalarType clamp_scalar_out_dtype(
38+ const at::Tensor& self,
39+ const c10::optional<at::Scalar>& min,
40+ const c10::optional<at::Scalar>& max);
32OP_PLUGIN_HIDDEN at::ScalarType abs_out_dtype(const at::Tensor& self);41OP_PLUGIN_HIDDEN at::ScalarType abs_out_dtype(const at::Tensor& self);
33OP_PLUGIN_HIDDEN at::ScalarType npu_moe_distribute_dispatch_setup_out_dtype(c10::optional<int64_t> y_dtype);42OP_PLUGIN_HIDDEN at::ScalarType npu_moe_distribute_dispatch_setup_out_dtype(c10::optional<int64_t> y_dtype);
34 43