已合并
add quantmatmulactivationquant #7272
zhuoyunhang创建于 7月9日
add quantmatmulactivationquant #7272
已合并
zhuoyunhang创建于 7月9日
54 个文件变更+6306-33
@@ -7,7 +7,7 @@
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
10-set(OPTENSOR_TAG_ID d9b85b07061630f67f4c27d30c949e4fa5133193)10+set(OPTENSOR_TAG_ID 7c1bca75345a0efa3f37eb3c7cf6860751af2749)
11 11 
12if(EXISTS "${PROJECT_SOURCE_DIR}/../ops-tensor")12if(EXISTS "${PROJECT_SOURCE_DIR}/../ops-tensor")
13 get_filename_component(OPTENSOR_SOURCE_PATH13 get_filename_component(OPTENSOR_SOURCE_PATH
@@ -115,4 +115,4 @@ else()
115 get_filename_component(TENSOR_API115 get_filename_component(TENSOR_API
116 ${OPTENSOR_SOURCE_PATH}/include/tensor_api REALPATH)116 ${OPTENSOR_SOURCE_PATH}/include/tensor_api REALPATH)
117 endif()117 endif()
118-endif()118+endif()
@@ -326,6 +326,7 @@
326- [aclnnQuantConvolution](../../conv/convolution_forward/docs/aclnnQuantConvolution.md)326- [aclnnQuantConvolution](../../conv/convolution_forward/docs/aclnnQuantConvolution.md)
327- [aclnnQuantConvolutionWeightNz](../../conv/convolution_forward/docs/aclnnQuantConvolutionWeightNz.md)327- [aclnnQuantConvolutionWeightNz](../../conv/convolution_forward/docs/aclnnQuantConvolutionWeightNz.md)
328- [aclnnQuantMatmul](../../matmul/quant_matmul/docs/aclnnQuantMatmul.md)328- [aclnnQuantMatmul](../../matmul/quant_matmul/docs/aclnnQuantMatmul.md)
329+- [aclnnQuantMatmulActivationQuantWeightNz](../../matmul/quant_matmul_activation_quant/docs/aclnnQuantMatmulActivationQuantWeightNz.md)
329- [aclnnQuantMatmulDequant](../../matmul/quant_matmul_dequant/docs/aclnnQuantMatmulDequant.md)330- [aclnnQuantMatmulDequant](../../matmul/quant_matmul_dequant/docs/aclnnQuantMatmulDequant.md)
330- [aclnnQuantMatmulReduceSumWeightNz](../../matmul/quant_matmul_reduce_sum/docs/aclnnQuantMatmulReduceSumWeightNz.md)331- [aclnnQuantMatmulReduceSumWeightNz](../../matmul/quant_matmul_reduce_sum/docs/aclnnQuantMatmulReduceSumWeightNz.md)
331- [aclnnQuantMatmulV2](../../matmul/quant_matmul/docs/aclnnQuantMatmulV2.md)332- [aclnnQuantMatmulV2](../../matmul/quant_matmul/docs/aclnnQuantMatmulV2.md)
@@ -4,3 +4,4 @@
4- [swiglu_group](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group/swiglu_group.md)4- [swiglu_group](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group/swiglu_group.md)
5- [swiglu_group_quant](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group_quant/swiglu_group_quant.md)5- [swiglu_group_quant](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group_quant/swiglu_group_quant.md)
6- [rms_norm_dynamic_quant](../../torch_extension/cann_ops_nn/ops/norm/rms_norm_dynamic_quant/rms_norm_dynamic_quant.md)6- [rms_norm_dynamic_quant](../../torch_extension/cann_ops_nn/ops/norm/rms_norm_dynamic_quant/rms_norm_dynamic_quant.md)
7+- [quant_matmul_activation_quant](../../matmul/quant_matmul_activation_quant/docs/torchapi_quant_matmul_activation_quant.md)
@@ -360,6 +360,7 @@
360| [aclnnQuantBatchMatmulInplaceAdd](../../matmul/quant_batch_matmul_inplace_add/docs/aclnnQuantBatchMatmulInplaceAdd.md) | 实现量化矩阵乘计算和原地累加加法计算,基本功能为矩阵乘和加法的组合。 | - | 默认确定性实现 |360| [aclnnQuantBatchMatmulInplaceAdd](../../matmul/quant_batch_matmul_inplace_add/docs/aclnnQuantBatchMatmulInplaceAdd.md) | 实现量化矩阵乘计算和原地累加加法计算,基本功能为矩阵乘和加法的组合。 | - | 默认确定性实现 |
361| [aclnnQuantMatmulReduceSumWeightNz](../../matmul/quant_matmul_reduce_sum/docs/aclnnQuantMatmulReduceSumWeightNz.md) | 完成量化的分组矩阵计算,然后所有组的矩阵计算结果相加后输出。 | 默认非确定性实现,支持配置开启。 | - |361| [aclnnQuantMatmulReduceSumWeightNz](../../matmul/quant_matmul_reduce_sum/docs/aclnnQuantMatmulReduceSumWeightNz.md) | 完成量化的分组矩阵计算,然后所有组的矩阵计算结果相加后输出。 | 默认非确定性实现,支持配置开启。 | - |
362| [aclnnQuantMatmulWeightNz](../../matmul/quant_batch_matmul_v3/docs/aclnnQuantMatmulWeightNz.md) | 完成量化的矩阵乘计算。 | 默认确定性实现 | 默认确定性实现 |362| [aclnnQuantMatmulWeightNz](../../matmul/quant_batch_matmul_v3/docs/aclnnQuantMatmulWeightNz.md) | 完成量化的矩阵乘计算。 | 默认确定性实现 | 默认确定性实现 |
363+| [aclnnQuantMatmulActivationQuantWeightNz](../../matmul/quant_matmul_activation_quant/docs/aclnnQuantMatmulActivationQuantWeightNz.md) | 融合量化的矩阵乘、激活以及动态量化计算,weight仅支持NZ格式。 | - | 默认确定性实现 |
363| [aclnnQuantMatmulDequant](../../matmul/quant_matmul_dequant/docs/aclnnQuantMatmulDequant.md) | 对输入x进行量化,矩阵乘以及反量化。 | - | - |364| [aclnnQuantMatmulDequant](../../matmul/quant_matmul_dequant/docs/aclnnQuantMatmulDequant.md) | 对输入x进行量化,矩阵乘以及反量化。 | - | - |
364| [aclnnQuantMax](../../quant/quant_max/docs/aclnnQuantMax.md) | 完成输入的量化计算并得到输入绝对值的最大值。 | - | 默认确定性实现 |365| [aclnnQuantMax](../../quant/quant_max/docs/aclnnQuantMax.md) | 完成输入的量化计算并得到输入绝对值的最大值。 | - | 默认确定性实现 |
365| [aclnnGroupedQuantMax](../../quant/grouped_quant_max/docs/aclnnGroupedQuantMax.md) | 完成输入分组的量化计算并得到输入分组绝对值的最大值。 | - | 默认确定性实现 |366| [aclnnGroupedQuantMax](../../quant/grouped_quant_max/docs/aclnnGroupedQuantMax.md) | 完成输入分组的量化计算并得到输入分组绝对值的最大值。 | - | 默认确定性实现 |
@@ -2767,6 +2767,16 @@
2767 <td>AI Core</td>2767 <td>AI Core</td>
2768 <td>完成量化的矩阵乘计算后,进行反量化。</td>2768 <td>完成量化的矩阵乘计算后,进行反量化。</td>
2769 </tr>2769 </tr>
2770+ <tr>
2771+ <td>matmul</td>
2772+ <td><a href="../../matmul/quant_matmul_activation_quant/README.md">quant_matmul_activation_quant</a></td>
2773+ <td>✓</td>
2774+ <td>✓</td>
2775+ <td>✓</td>
2776+ <td>✓</td>
2777+ <td>AI Core</td>
2778+ <td>完成量化的矩阵乘计算后计算Gelu激活再执行动态MX量化。</td>
2779+ </tr>
2770 <tr>2780 <tr>
2771 <td>matmul</td>2781 <td>matmul</td>
2772 <td><a href="../../matmul/quant_matmul_reduce_sum/README.md">quant_matmul_reduce_sum</a></td>2782 <td><a href="../../matmul/quant_matmul_reduce_sum/README.md">quant_matmul_reduce_sum</a></td>
@@ -27,3 +27,4 @@
27|[swiglu_group](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group/swiglu_group.md)|SwiGLU分组激活算子,对输入张量按最后一维拆分为两部分,分别进行clamp和sigmoid操作后相乘,支持可选的权重和分组索引。|-|-|27|[swiglu_group](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group/swiglu_group.md)|SwiGLU分组激活算子,对输入张量按最后一维拆分为两部分,分别进行clamp和sigmoid操作后相乘,支持可选的权重和分组索引。|-|-|
28|[swiglu_group_quant](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group_quant/swiglu_group_quant.md)|融合SwiGLU分组激活与量化的算子,在SwiGLU计算基础上支持FP8/MXFP4等多种量化模式输出。|-|-|28|[swiglu_group_quant](../../torch_extension/cann_ops_nn/ops/activation/swiglu_group_quant/swiglu_group_quant.md)|融合SwiGLU分组激活与量化的算子,在SwiGLU计算基础上支持FP8/MXFP4等多种量化模式输出。|-|-|
29|[rms_norm_dynamic_quant](../../torch_extension/cann_ops_nn/ops/norm/rms_norm_dynamic_quant/rms_norm_dynamic_quant.md)|融合RMS Normalization与INT8动态量化,输出量化后的张量及缩放因子。|-|-|29|[rms_norm_dynamic_quant](../../torch_extension/cann_ops_nn/ops/norm/rms_norm_dynamic_quant/rms_norm_dynamic_quant.md)|融合RMS Normalization与INT8动态量化,输出量化后的张量及缩放因子。|-|-|
30+|[quant_matmul_activation_quant](../../matmul/quant_matmul_activation_quant/docs/torchapi_quant_matmul_activation_quant.md)| 融合量化的矩阵乘、激活以及动态量化计算,weight仅支持NZ格式。 |-|默认支持确定性计算|
@@ -426,9 +426,7 @@ void AdaptiveSlidingWindowMXBasicAPITiling::NormalizeSingleRoundTailSplitBasicBl
426 const uint64_t baseMAlignSize = inputParams_.transA ?426 const uint64_t baseMAlignSize = inputParams_.transA ?
427 GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.aDtype) :427 GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.aDtype) :
428 qmmv3_tiling_const::CUBE_BLOCK;428 qmmv3_tiling_const::CUBE_BLOCK;
429- const uint64_t baseNAlignSize = inputParams_.transB ?429+ const uint64_t baseNAlignSize = GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE);
430- qmmv3_tiling_const::CUBE_BLOCK :
431- GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.bDtype);
432 const uint64_t newBaseM = GetSingleRoundTailSplitBase(inputParams_.mSize, adaptiveWin_.baseM,430 const uint64_t newBaseM = GetSingleRoundTailSplitBase(inputParams_.mSize, adaptiveWin_.baseM,
433 adaptiveWin_.mTailTile, baseMAlignSize);431 adaptiveWin_.mTailTile, baseMAlignSize);
434 const uint64_t newBaseN = GetSingleRoundTailSplitBase(inputParams_.nSize, adaptiveWin_.baseN,432 const uint64_t newBaseN = GetSingleRoundTailSplitBase(inputParams_.nSize, adaptiveWin_.baseN,
@@ -429,6 +429,12 @@ bool AdaptiveSlidingWindowTiling::IsInValidWeighNzTailSplit(uint64_t splitCnt, b
429 return tailN % GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.bDtype) != 0UL;429 return tailN % GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.bDtype) != 0UL;
430}430}
431 431 
432+uint64_t AdaptiveSlidingWindowTiling::GetBaseNAlignSize(uint64_t innerAlignSize) const
433+{
434+ return inputParams_.transB ? qmmv3_tiling_const::CUBE_BLOCK :
435+ GetShapeWithDataType(innerAlignSize, inputParams_.bDtype);
436+}
437+ 
432uint64_t AdaptiveSlidingWindowTiling::GetTailBasicBlockSplitMax(bool isMSplit, uint64_t tileMax,438uint64_t AdaptiveSlidingWindowTiling::GetTailBasicBlockSplitMax(bool isMSplit, uint64_t tileMax,
433 uint64_t splitSize) const439 uint64_t splitSize) const
434{440{
@@ -442,9 +448,7 @@ uint64_t AdaptiveSlidingWindowTiling::GetTailBasicBlockSplitMax(bool isMSplit, u
442 const uint64_t splitAlignNum = isMSplit ? (inputParams_.transA ?448 const uint64_t splitAlignNum = isMSplit ? (inputParams_.transA ?
443 GetShapeWithDataType(baseMAlignSize, inputParams_.aDtype) :449 GetShapeWithDataType(baseMAlignSize, inputParams_.aDtype) :
444 qmmv3_tiling_const::CUBE_BLOCK) :450 qmmv3_tiling_const::CUBE_BLOCK) :
445- (!inputParams_.transB ?451+ GetBaseNAlignSize(baseNAlignSize);
446- GetShapeWithDataType(baseNAlignSize, inputParams_.bDtype) :
447- qmmv3_tiling_const::CUBE_BLOCK);
448 return std::min(tileMax, MathUtil::CeilDivision(splitSize, splitAlignNum));452 return std::min(tileMax, MathUtil::CeilDivision(splitSize, splitAlignNum));
449}453}
450 454 
@@ -88,8 +88,8 @@ protected:
88 virtual void CalcTailBasicBlockBfullLoad();88 virtual void CalcTailBasicBlockBfullLoad();
89 virtual void CalcTailBasicBlock4MmadS8S4();89 virtual void CalcTailBasicBlock4MmadS8S4();
90 uint64_t GetTailBasicBlockSplitMax(bool isMSplit, uint64_t tileMax, uint64_t splitSize) const;90 uint64_t GetTailBasicBlockSplitMax(bool isMSplit, uint64_t tileMax, uint64_t splitSize) const;
91- bool CanIncreaseTailSplit(bool isPreSplitM, bool isPreSplit, uint64_t preSplit, uint64_t secSplit,91+ virtual bool CanIncreaseTailSplit(bool isPreSplitM, bool isPreSplit, uint64_t preSplit, uint64_t secSplit,
92- uint64_t splitMax);92+ uint64_t splitMax);
93 uint64_t GetTailSplitState(bool isPreSplitM, bool isPreSplit, uint64_t split, uint64_t splitSize) const;93 uint64_t GetTailSplitState(bool isPreSplitM, bool isPreSplit, uint64_t split, uint64_t splitSize) const;
94 void CalcTailBasicBlockSplit(bool isPreSplitM, uint64_t preSplitMax, uint64_t secSplitMax, uint64_t preSplitSize,94 void CalcTailBasicBlockSplit(bool isPreSplitM, uint64_t preSplitMax, uint64_t secSplitMax, uint64_t preSplitSize,
95 uint64_t secSplitSize);95 uint64_t secSplitSize);
@@ -111,6 +111,7 @@ protected:
111 void GetOuterMAxisTailCnt(uint64_t& baseTailSplitCnt, uint64_t& tailMain);111 void GetOuterMAxisTailCnt(uint64_t& baseTailSplitCnt, uint64_t& tailMain);
112 void GetOuterNAxisTailCnt(uint64_t& baseTailSplitCnt, uint64_t& tailMain);112 void GetOuterNAxisTailCnt(uint64_t& baseTailSplitCnt, uint64_t& tailMain);
113 virtual bool CheckCoreNum() const;113 virtual bool CheckCoreNum() const;
114+ virtual uint64_t GetBaseNAlignSize(uint64_t innerAlignSize) const;
114 115 
115 DequantBmm::QuantBatchMatmulV3TilingDataParams tilingDataSelf_;116 DequantBmm::QuantBatchMatmulV3TilingDataParams tilingDataSelf_;
116 DequantBmm::QuantBatchMatmulV3TilingDataParams& tilingData_;117 DequantBmm::QuantBatchMatmulV3TilingDataParams& tilingData_;
@@ -111,7 +111,7 @@ bool BaseBlockCalculator::ValidateInput() const
111 static_cast<uint64_t>(compileInfo_.aicNum), batchCoreCnt_),111 static_cast<uint64_t>(compileInfo_.aicNum), batchCoreCnt_),
112 return false);112 return false);
113 uint64_t baseMAlignSize = GetBaseMAlignSize();113 uint64_t baseMAlignSize = GetBaseMAlignSize();
114- uint64_t baseNAlignSize = GetBaseNAlignSize();114+ uint64_t baseNAlignSize = GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE);
115 uint64_t baseKAlignSize = GetBaseKAlignSize();115 uint64_t baseKAlignSize = GetBaseKAlignSize();
116 OP_TILING_CHECK(116 OP_TILING_CHECK(
117 baseMAlignSize == 0UL || baseNAlignSize == 0UL || baseKAlignSize == 0UL,117 baseMAlignSize == 0UL || baseNAlignSize == 0UL || baseKAlignSize == 0UL,
@@ -140,7 +140,7 @@ void BaseBlockCalculator::ComputeBaseBlockDefault()
140 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),140 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
141 GetBaseMAlignSize());141 GetBaseMAlignSize());
142 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),142 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
143- GetBaseNAlignSize());143+ GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE));
144 144 
145 uint64_t baseKDefaultSize = GetShapeWithDataType(qmmv3_tiling_const::BASIC_BLOCK_SIZE_128, inputParams_.aDtype);145 uint64_t baseKDefaultSize = GetShapeWithDataType(qmmv3_tiling_const::BASIC_BLOCK_SIZE_128, inputParams_.aDtype);
146 baseBlockRes_.baseK = ops::CeilAlign(std::min(baseKDefaultSize, inputParams_.kSize), GetBaseKAlignSize());146 baseBlockRes_.baseK = ops::CeilAlign(std::min(baseKDefaultSize, inputParams_.kSize), GetBaseKAlignSize());
@@ -169,7 +169,7 @@ void BaseBlockCalculator::ComputeBaseBlockMmadS8S4()
169 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),169 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
170 GetBaseMAlignSize());170 GetBaseMAlignSize());
171 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),171 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
172- GetBaseNAlignSize());172+ GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE));
173 173 
174 uint64_t basicBlockSizeA = qmmv3_tiling_const::BASIC_BLOCK_SIZE_128;174 uint64_t basicBlockSizeA = qmmv3_tiling_const::BASIC_BLOCK_SIZE_128;
175 uint64_t basicBlockSizeB = qmmv3_tiling_const::BASIC_BLOCK_SIZE_128;175 uint64_t basicBlockSizeB = qmmv3_tiling_const::BASIC_BLOCK_SIZE_128;
@@ -196,7 +196,7 @@ bool BaseBlockCalculator::InitStreamKBaseBlock()
196 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),196 baseBlockRes_.baseM = ops::CeilAlign(std::min(inputParams_.mSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
197 GetBaseMAlignSize());197 GetBaseMAlignSize());
198 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),198 baseBlockRes_.baseN = ops::CeilAlign(std::min(inputParams_.nSize, qmmv3_tiling_const::BASIC_BLOCK_SIZE_256),
199- GetBaseNAlignSize());199+ GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE));
200 OP_TILING_CHECK(baseBlockRes_.baseM == 0UL || baseBlockRes_.baseN == 0UL,200 OP_TILING_CHECK(baseBlockRes_.baseM == 0UL || baseBlockRes_.baseN == 0UL,
201 CUBE_INNER_ERR_REPORT(inputParams_.opName, "Invalid StreamK base divisor: baseM(%lu), baseN(%lu).",201 CUBE_INNER_ERR_REPORT(inputParams_.opName, "Invalid StreamK base divisor: baseM(%lu), baseN(%lu).",
202 baseBlockRes_.baseM, baseBlockRes_.baseN),202 baseBlockRes_.baseM, baseBlockRes_.baseN),
@@ -221,7 +221,7 @@ bool BaseBlockCalculator::UpdateSmallMnStreamKBase()
221 baseBlockRes_.baseM = ops::CeilAlign(MathUtil::CeilDivision(inputParams_.mSize, baseBlockRes_.mCnt),221 baseBlockRes_.baseM = ops::CeilAlign(MathUtil::CeilDivision(inputParams_.mSize, baseBlockRes_.mCnt),
222 GetBaseMAlignSize());222 GetBaseMAlignSize());
223 baseBlockRes_.baseN = ops::CeilAlign(MathUtil::CeilDivision(inputParams_.nSize, baseBlockRes_.nCnt),223 baseBlockRes_.baseN = ops::CeilAlign(MathUtil::CeilDivision(inputParams_.nSize, baseBlockRes_.nCnt),
224- GetBaseNAlignSize());224+ GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE));
225 baseBlockRes_.preSplitKBlockCnt = batchCoreCnt_ * baseBlockRes_.mCnt * baseBlockRes_.nCnt;225 baseBlockRes_.preSplitKBlockCnt = batchCoreCnt_ * baseBlockRes_.mCnt * baseBlockRes_.nCnt;
226 OP_TILING_CHECK(226 OP_TILING_CHECK(
227 baseBlockRes_.preSplitKBlockCnt == 0UL,227 baseBlockRes_.preSplitKBlockCnt == 0UL,
@@ -289,10 +289,10 @@ uint64_t BaseBlockCalculator::GetBaseMAlignSize() const
289 qmmv3_tiling_const::CUBE_BLOCK;289 qmmv3_tiling_const::CUBE_BLOCK;
290}290}
291 291 
292-uint64_t BaseBlockCalculator::GetBaseNAlignSize() const292+uint64_t BaseBlockCalculator::GetBaseNAlignSize(uint64_t innerAlignSize) const
293{293{
294 return inputParams_.transB ? qmmv3_tiling_const::CUBE_BLOCK :294 return inputParams_.transB ? qmmv3_tiling_const::CUBE_BLOCK :
295- GetShapeWithDataType(qmmv3_tiling_const::L1_ALIGN_SIZE, inputParams_.bDtype);295+ GetShapeWithDataType(innerAlignSize, inputParams_.bDtype);
296}296}
297 297 
298uint64_t BaseBlockCalculator::GetBaseKAlignSize() const298uint64_t BaseBlockCalculator::GetBaseKAlignSize() const
@@ -352,9 +352,7 @@ void BaseBlockCalculator::SearchLoadBalanceBaseBlock(uint64_t roundLimit, uint64
352 uint64_t baseMAlignNum = inputParams_.transA ?352 uint64_t baseMAlignNum = inputParams_.transA ?
353 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :353 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :
354 qmmv3_tiling_const::CUBE_BLOCK;354 qmmv3_tiling_const::CUBE_BLOCK;
355- uint64_t baseNAlignNum = inputParams_.transB ?355+ uint64_t baseNAlignNum = GetBaseNAlignSize(qmmv3_tiling_const::L2_ALIGN_SIZE);
356- qmmv3_tiling_const::CUBE_BLOCK :
357- GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.bDtype);
358 uint64_t searchBaseM = ops::CeilAlign(baseBlockRes_.baseM, baseMAlignNum);356 uint64_t searchBaseM = ops::CeilAlign(baseBlockRes_.baseM, baseMAlignNum);
359 uint64_t searchBaseN = ops::CeilAlign(baseBlockRes_.baseN, baseNAlignNum);357 uint64_t searchBaseN = ops::CeilAlign(baseBlockRes_.baseN, baseNAlignNum);
360 bool hasCandidate = false;358 bool hasCandidate = false;
@@ -415,7 +413,7 @@ void BaseBlockCalculator::TryApplyLoadBalanceBase(uint64_t bestBaseM, uint64_t b
415 baseBlockRes_.baseM = ops::CeilAlign(inputParams_.mSize, GetBaseMAlignSize());413 baseBlockRes_.baseM = ops::CeilAlign(inputParams_.mSize, GetBaseMAlignSize());
416 }414 }
417 if (baseBlockRes_.baseN > inputParams_.nSize) {415 if (baseBlockRes_.baseN > inputParams_.nSize) {
418- baseBlockRes_.baseN = ops::CeilAlign(inputParams_.nSize, GetBaseNAlignSize());416+ baseBlockRes_.baseN = ops::CeilAlign(inputParams_.nSize, GetBaseNAlignSize(qmmv3_tiling_const::L1_ALIGN_SIZE));
419 }417 }
420}418}
421 419 
@@ -447,9 +445,7 @@ bool BaseBlockCalculator::AdjustBaseBlockDefault()
447 uint64_t baseMAlignNum = inputParams_.transA ?445 uint64_t baseMAlignNum = inputParams_.transA ?
448 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :446 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :
449 qmmv3_tiling_const::CUBE_BLOCK;447 qmmv3_tiling_const::CUBE_BLOCK;
450- uint64_t baseNAlignNum = inputParams_.transB ?448+ uint64_t baseNAlignNum = GetBaseNAlignSize(qmmv3_tiling_const::L2_ALIGN_SIZE);
451- qmmv3_tiling_const::CUBE_BLOCK :
452- GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.bDtype);
453 uint64_t baseKAlignNum = (inputParams_.transA && !inputParams_.transB) ?449 uint64_t baseKAlignNum = (inputParams_.transA && !inputParams_.transB) ?
454 GetShapeWithDataType(qmmv3_tiling_const::BASIC_BLOCK_SIZE_32, inputParams_.aDtype) :450 GetShapeWithDataType(qmmv3_tiling_const::BASIC_BLOCK_SIZE_32, inputParams_.aDtype) :
455 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype);451 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype);
@@ -552,6 +548,7 @@ void BaseBlockCalculator::TrySwapBaseMNForMxFalseTrue(uint64_t& baseM, uint64_t&
552 if (swapUsedCore >= curUsedCore && (swapCoreShapeBetter || swapBaseNIsFriendly)) {548 if (swapUsedCore >= curUsedCore && (swapCoreShapeBetter || swapBaseNIsFriendly)) {
553 baseM = swapBaseM;549 baseM = swapBaseM;
554 baseN = swapBaseN;550 baseN = swapBaseN;
551+ baseN = ops::CeilAlign(baseN, GetBaseNAlignSize(qmmv3_tiling_const::L2_ALIGN_SIZE));
555 }552 }
556}553}
557 554 
@@ -579,9 +576,7 @@ bool BaseBlockCalculator::AdjustBaseBlockPertile(uint64_t coreNumMN)
579 uint64_t baseMAlignNum = inputParams_.transA ?576 uint64_t baseMAlignNum = inputParams_.transA ?
580 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :577 GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.aDtype) :
581 qmmv3_tiling_const::CUBE_BLOCK;578 qmmv3_tiling_const::CUBE_BLOCK;
582- uint64_t baseNAlignNum = !inputParams_.transB ?579+ uint64_t baseNAlignNum = GetBaseNAlignSize(qmmv3_tiling_const::L2_ALIGN_SIZE);
583- GetShapeWithDataType(qmmv3_tiling_const::L2_ALIGN_SIZE, inputParams_.bDtype) :
584- qmmv3_tiling_const::CUBE_BLOCK;
585 uint64_t adjustBaseM = baseBlockRes_.baseM;580 uint64_t adjustBaseM = baseBlockRes_.baseM;
586 uint64_t adjustBaseN = baseBlockRes_.baseN;581 uint64_t adjustBaseN = baseBlockRes_.baseN;
587 uint64_t adjustMCore = MathUtil::CeilDivision(inputParams_.mSize, adjustBaseM);582 uint64_t adjustMCore = MathUtil::CeilDivision(inputParams_.mSize, adjustBaseM);
@@ -43,6 +43,14 @@ public:
43 bool Compute(BaseBlockMode mode);43 bool Compute(BaseBlockMode mode);
44 const BaseBlockRes& GetOutput() const;44 const BaseBlockRes& GetOutput() const;
45 45 
46+protected:
47+ // QBMMActivationQuant需要修改为32对齐
48+ virtual uint64_t GetBaseNAlignSize(uint64_t innerAlignSize) const;
49+ const QuantBatchMatmulInfo& inputParams_;
50+ const QuantBatchMatmulV3CompileInfo& compileInfo_;
51+ uint64_t batchCoreCnt_ = 1UL;
52+ BaseBlockRes baseBlockRes_;
53+ 
46private:54private:
47 bool ValidateInput() const;55 bool ValidateInput() const;
48 bool ValidateBaseBlock() const;56 bool ValidateBaseBlock() const;
@@ -55,7 +63,6 @@ private:
55 void UpdateTailStreamKBase();63 void UpdateTailStreamKBase();
56 bool FinalizeStreamKBaseK();64 bool FinalizeStreamKBaseK();
57 uint64_t GetBaseMAlignSize() const;65 uint64_t GetBaseMAlignSize() const;
58- uint64_t GetBaseNAlignSize() const;
59 uint64_t GetBaseKAlignSize() const;66 uint64_t GetBaseKAlignSize() const;
60 bool OptimizeBaseBlockForCoreUtilization(BaseBlockMode mode);67 bool OptimizeBaseBlockForCoreUtilization(BaseBlockMode mode);
61 void OptimizeBaseBlockForLoadBalance();68 void OptimizeBaseBlockForLoadBalance();
@@ -75,11 +82,6 @@ private:
75 bool CalculateOptimalSplit(uint64_t& baseM, uint64_t& baseN, uint64_t baseMAlignNum, uint64_t baseNAlignNum,82 bool CalculateOptimalSplit(uint64_t& baseM, uint64_t& baseN, uint64_t baseMAlignNum, uint64_t baseNAlignNum,
76 uint64_t baseKAlignNum) const;83 uint64_t baseKAlignNum) const;
77 bool IsMxBackwardTrans() const;84 bool IsMxBackwardTrans() const;
78- 
79- const QuantBatchMatmulInfo& inputParams_;
80- const QuantBatchMatmulV3CompileInfo& compileInfo_;
81- uint64_t batchCoreCnt_ = 1UL;
82- BaseBlockRes baseBlockRes_;
83};85};
84 86 
85} // namespace optiling87} // namespace optiling
@@ -0,0 +1,10 @@
1+# This program is free software, you can redistribute it and/or modify.
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This file is a part of the CANN Open Software.
4+# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
7+# See LICENSE in the root of the software repository for the full text of the License.
8+#/
9+ 
10+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE quant_matmul_activation_quant ACLNNTYPE aclnn_exclude DEPENDENCIES quant_batch_matmul_v3)
@@ -0,0 +1,177 @@
1+# QuantMatmulActivationQuant
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| ---- | :----:|
7+|<term>Ascend 950PR/Ascend 950DT</term>|√|
8+|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>|×|
9+|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>|×|
10+|<term>Atlas 200I/500 A2 推理产品</term>|×|
11+|<term>Atlas 推理系列产品</term>|×|
12+|<term>Atlas 训练系列产品</term>|×|
13+ 
14+## 功能说明
15+ 
16+- 算子功能:融合量化的矩阵乘、激活以及动态量化,当前支持激活为gelu、MX [量化模式](../../docs/zh/context/quant_mode_introduction.md)。
17+ 
18+- 计算公式:
19+ 
20+ - <term>Ascend 950PR/Ascend 950DT</term>
21+ 
22+ - QuantMatmul MX量化模式:
23+ 
24+ $$
25+ matmulOut[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice))* (x1Scale[m/gsM, j] * x2Scale[j, n/gsN]))+bias[n]
26+ $$
27+ 
28+ 其中,gsM,gsN和gsK分别代表groupSizeM,groupSizeN和groupSizeK;x1Slice代表x1第m行长度为groupSizeK的向量,x2Slice代表x2第n列长度为groupSizeK的向量;K轴均从j*groupSizeK起始切片,j的取值范围为[0, kLoops),kLoops = ceil(K / groupSizeK),K为K轴长度,支持最后的切片长度不足groupSizeK。
29+ 
30+ - 激活计算公式:
31+ 
32+ - gelu_tanh(高性能近似):
33+ $$
34+ activationOut=GELU(matmulOut)=matmulOut × Φ(matmulOut)=0.5 * matmulOut * (1 + tanh( \sqrt{2 / \pi} * (matmulOut + 0.044715 * matmulOut^{3})))
35+ $$
36+ 
37+ - gelu_erf:
38+ $$
39+ activationOut=GELU(matmulOut)=0.5 * matmulOut * (1 + erf(matmulOut / \sqrt{2}))
40+ $$
41+ 
42+ - 动态量化计算公式:
43+ 
44+ - 场景1,当scaleAlg为0时:
45+ - 将输入activationOut在尾轴上按$k = 32$个数分组,一组k个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\},\space k = 32$
46+ 
47+ $$
48+ shared\_exp = floor(log_2(max_i(|V_i|))) - emax \\
49+ mxscale = 2^{shared\_exp}\\
50+ P_i = cast\_to\_dst\_type(V_i/mxscale, round\_mode), \space i\space from\space 1\space to\space k\\
51+ $$
52+ 
53+ - 量化后的 $P_{i}$ 按对应的 $V_{i}$ 的位置组成输出yOut,mxscale按尾轴上的分组输出yScaleOut。
54+ 
55+ - emax: 对应数据类型的最大正则数的指数位。
56+ 
57+ | DataType | emax |
58+ | :-----------: | :--: |
59+ | FLOAT4_E2M1 | 2 |
60+ | FLOAT4_E1M2 | 0 |
jayshu
jayshujayshu7月9日

E1M2的这个emax写错了,应该为1

likedislike
陈勿静
陈勿静
7月22日 评论:
61+ | FLOAT8_E4M3FN | 8 |
62+ | FLOAT8_E5M2 | 15 |
63+ 
64+ - 场景2,当scaleAlg为1时,只涉及FP8类型:
65+ - 将输入activationOut在尾轴上按$k = 32$个数分块,对每块单独计算一个块缩放因子$S_{fp32}^b$,再把块内所有元素用同一个$S_{fp32}^b$映射到目标低精度类型FP8。如果最后一块不足$k = 32$个元素,把缺失值视为0,按照完整块处理。
66+ - 找到该块中数值的最大绝对值:
67+ 
68+ $$
69+ Amax(D_{fp32}^b)=max(\{|d_{i}|\}_{i=1}^{k})
70+ $$
71+ 
72+ - 将FP32映射到目标数据类型FP8可表示的范围内,其中$Amax(DType)$是目标精度能表示的最大值:
73+ 
74+ $$
75+ S_{fp32}^b = \frac{Amax(D_{fp32}^b)}{Amax(DType)}
76+ $$
77+ 
78+ - 将块缩放因子$S_{fp32}^b$转换为FP8格式下可表示的缩放值$S_{ue8m0}^b$
79+ - 从块的浮点缩放因子$S_{fp32}^b$中提取无偏指数$E_{int}^b$和尾数$M_{fixp}^b$
80+ - 为保证量化时不溢出,对指数进行向上取整,且在FP8可表示的范围内:
81+ 
82+ $$
83+ E_{int}^b = \begin{cases} E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为正规数,且} E_{int}^b < 254 \text{且} M_{fixp}^b > 0 \\ E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为非正规数,且} M_{fixp}^b > 0.5 \\ E_{int}^b, & \text{否则} \end{cases}
84+ $$
85+ 
86+ - 计算块缩放因子:$S_{ue8m0}^b=2^{E_{int}^b}$
87+ - 计算块转换因子:$R_{fp32}^b=\frac{1}{fp32(S_{ue8m0}^b)}$
88+ - 应用到量化的最终步骤,对于每个块内元素,$d^i = DType(d_{fp32}^i \cdot R_{fp32}^n)$,最终输出的量化结果是$\left(S^b, [d^i]_{i=1}^k\right)$,其中$S^b$代表块的缩放因子,这里指$S_{ue8m0}^b$,$[d^i]_{i=1}^k$代表块内量化后的数据。
89+ 
90+## 参数说明
91+ 
92+<table class="tg" style="undefined;table-layout: fixed; width: 1166px"><colgroup>
93+<col style="width: 81px">
94+<col style="width: 121px">
95+<col style="width: 430px">
96+<col style="width: 390px">
97+<col style="width: 144px">
98+</colgroup>
99+<thead>
100+ <tr>
101+ <th class="tg-xbcz"><span style="font-weight:700;color:var(--theme-text);background-color:var(--theme-table-header-bg)">参数名</span></th>
102+ <th class="tg-xbcz"><span style="font-weight:700;color:var(--theme-text);background-color:var(--theme-table-header-bg)">输入/输出/属性</span></th>
103+ <th class="tg-xbcz"><span style="font-weight:700;color:var(--theme-text);background-color:var(--theme-table-header-bg)">描述</span></th>
104+ <th class="tg-xbcz"><span style="font-weight:700;color:var(--theme-text);background-color:var(--theme-table-header-bg)">数据类型</span></th>
105+ <th class="tg-xbcz"><span style="font-weight:700;color:var(--theme-text);background-color:var(--theme-table-header-bg)">数据格式</span></th>
106+ </tr></thead>
107+<tbody>
108+ <tr>
109+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">x1</span></td>
110+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">输入</span></td>
111+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">矩阵乘运算中的左矩阵。</span></td>
112+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">FLOAT8_E4M3FN, FLOAT8_E5M2</span></td>
113+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">ND</span></td>
114+ </tr>
115+ <tr>
116+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">x2</span></td>
117+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">输入</span></td>
118+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">矩阵乘运算中的右矩阵。</span></td>
119+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">FLOAT8_E4M3FN</span></td>
120+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">FRACTAL_NZ</span></td>
121+ </tr>
122+ <tr>
123+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">x1_scale_optional</span></td>
124+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">可选输入</span></td>
125+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">矩阵乘计算时,量化参数的缩放因子,对应公式的x1Scale。</span></td>
126+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">FLOAT8_E8M0</span></td>
127+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">ND</span></td>
128+ </tr>
129+ <tr>
130+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">x2_scale</span></td>
131+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">输入</span></td>
132+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">矩阵乘计算时,量化参数的缩放因子,对应公式的x2Scale。</span></td>
133+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">FLOAT8_E8M0</span></td>
134+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">ND</span></td>
135+ </tr>
136+ <tr>
137+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">bias_optional</span></td>
138+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">可选输入</span></td>
139+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">矩阵乘运算后累加的偏置,对应公式中的bias。</span></td>
L
Lliubo757月16日

评审应该说过,可选输入,后面都有optional,命名要规范

likedislike
陈勿静
陈勿静
7月22日 评论:
140+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">FLOAT32</span></td>
141+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">ND</span></td>
142+ </tr>
143+ <tr>
144+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">y</span></td>
145+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">输出</span></td>
146+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">动态量化后的矩阵乘及激活计算结果。</span></td>
147+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">FLOAT8_E4M3FN, FLOAT8_E5M2</span></td>
148+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--devui-base-bg, #ffffff)">ND</span></td>
149+ </tr>
150+ <tr>
151+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">y_scale</span></td>
152+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">输出</span></td>
153+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">动态量化后每个分组对应的量化尺度。</span></td>
154+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">FLOAT8_E8M0</span></td>
155+ <td class="tg-zgfj"><span style="color:var(--theme-aide-text);background-color:var(--theme-table-header-bg)">ND</span></td>
156+ </tr>
157+</tbody></table>
158+ 
159+## 约束说明
160+ 
161+- 不支持空tensor。
162+- 支持连续tensor,[非连续tensor](../../docs/zh/context/non_contiguous_tensor.md)仅支持最后两根轴转置场景。
163+- 输入和输出支持以下数据类型组合:
164+ 
165+ - <term>Ascend 950PR/Ascend 950DT</term>
166+ 
167+ | x1 | x2 | x1_scale | x2_scale | bias | y | y_scale |
168+ |---------------|---------------|-------------|-------------|--------------|---------------------------|-------------|
169+ | FLOAT8_E4M3FN | FLOAT8_E4M3FN | FLOAT8_E8M0 | FLOAT8_E8M0 | null/FLOAT32 | FLOAT8_E4M3FN/FLOAT8_E5M2 | FLOAT8_E8M0 |
170+ | FLOAT8_E5M2 | FLOAT8_E4M3FN | FLOAT8_E8M0 | FLOAT8_E8M0 | null/FLOAT32 | FLOAT8_E4M3FN/FLOAT8_E5M2 | FLOAT8_E8M0 |
171+ 
172+## 调用说明
173+ 
174+ | 调用方式 | 样例代码 | 说明 |
175+ | ---------------- | --------------------------- | --------------------------------------------------- |
176+ | aclnn接口 | [test_aclnn_quant_matmul_activation_quant](examples/arch35/test_aclnn_quant_matmul_activation_quant.cpp) | 通过<br>[aclnnQuantMatmulActivationQuantWeightNz](docs/aclnnQuantMatmulActivationQuantWeightNz.md)<br>调用QuantMatmulActivationQuant算子。 |
177+ | PyTorch API | - | 通过<br>[quant_matmul_activation_quant](docs/torchapi_quant_matmul_activation_quant.md)<br>调用QuantMatmulActivationQuant算子。 |
Zzhuoyunhang7月20日

y、yscale 放最后

likedislike
陈勿静
陈勿静
7月22日 评论:
@@ -0,0 +1,839 @@
1+# aclnnQuantMatmulActivationQuantWeightNz
2+ 
3+## 产品支持情况
4+ 
5+| 产品 | 是否支持 |
6+| :----------------------------------------------------------- |:-------:|
7+| <term>Ascend 950PR/Ascend 950DT</term> | ✓ |
8+| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | × |
9+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | × |
10+| <term>Atlas 200I/500 A2 推理产品</term> | × |
11+| <term>Atlas 推理系列产品</term> | × |
12+| <term>Atlas 训练系列产品</term> | × |
13+ 
14+## 功能说明
15+ 
16+- 接口功能:融合量化的矩阵乘、激活以及动态量化计算。当前支持激活为gelu(包括gelu_tanh和gelu_erf),量化模式为MX [量化模式](../../../docs/zh/context/quant_mode_introduction.md)。最小支持输入维度为2维,最大支持输入维度为6维。相似接口有aclnnMm(仅支持2维Tensor作为输入的矩阵乘)和aclnnBatchMatMul(仅支持三维的矩阵乘,其中第一维是Batch维度),本接口在其他接口的基础上融合了gelu激活及动态MX量化,提高硬件计算效率。
17+ 
18+- 计算公式:
19+ 
20+ - <term>Ascend 950PR/Ascend 950DT</term>
21+ 
22+ 支持MX[量化模式](../../../docs/zh/context/quant_mode_introduction.md),不同量化模式对应的输入输出数据类型组合参见[约束说明](#约束说明)。
23+ 
24+ <details>
25+ 
26+ <summary><strong>QuantMatmul MX量化模式</strong></summary>
27+ 
28+ $$
29+ matmulOut[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice))* (x1Scale[m/gsM, j] * x2Scale[j, n/gsN]))+bias[n]
30+ $$
31+ 
32+ 其中,gsM,gsN和gsK分别代表groupSizeM,groupSizeN和groupSizeK;x1Slice代表x1第m行长度为groupSizeK的向量,x2Slice代表x2第n列长度为groupSizeK的向量;K轴均从j*groupSizeK起始切片,j的取值范围为[0, kLoops),kLoops = ceil(K / groupSizeK),K为K轴长度,支持最后的切片长度不足groupSizeK。仅MX量化模式下包含bias。
33+ 
34+ </details>
35+ 
36+ <details>
37+ 
38+ <summary><strong>激活计算公式</strong></summary>
39+ 
40+ - gelu_tanh(高性能近似):
41+ 
42+ $$
43+ activationOut=GELU(matmulOut)=matmulOut × Φ(matmulOut)=0.5 * matmulOut * (1 + tanh( \sqrt{2 / \pi} * (matmulOut + 0.044715 * matmulOut^{3})))
44+ $$
45+ 
46+ - gelu_erf:
47+ 
48+ $$
49+ activationOut=GELU(matmulOut)=0.5 * matmulOut * (1 + erf(matmulOut / \sqrt{2}))
50+ $$
51+ 
52+ </details>
53+ 
54+ <details>
55+ 
56+ <summary><strong>动态量化计算公式</strong></summary>
57+ 
58+ - 场景1,当scaleAlg为0时:
59+ - 将输入activationOut在尾轴上按$k=32$个数分组,一组k个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\},\space k = 32$
60+ 
61+ $$
62+ shared\_exp = floor(log_2(max_i(|V_i|))) - emax \\
63+ mxscale = 2^{shared\_exp}\\
64+ P_i = cast\_to\_dst\_type(V_i/mxscale, round\_mode), \space i\space from\space 1\space to\space k\\
65+ $$
66+ 
67+ - 量化后的 $P_{i}$ 按对应的 $V_{i}$ 的位置组成输出yOut,mxscale按尾轴上的分组输出yScaleOut。
68+ 
69+ - emax: 对应数据类型的最大正则数的指数位。
70+ 
71+ | DataType | emax |
72+ | :-----------: | :--: |
73+ | FLOAT4_E2M1 | 2 |
74+ | FLOAT4_E1M2 | 0 |
75+ | FLOAT8_E4M3FN | 8 |
76+ | FLOAT8_E5M2 | 15 |
77+ 
78+ - 场景2,当scaleAlg为1时,只涉及FP8类型:
79+ - 将输入activationOut在尾轴上按$k = 32$个数分块,对每块单独计算一个块缩放因子$S_{fp32}^b$,再把块内所有元素用同一个$S_{fp32}^b$映射到目标低精度类型FP8。如果最后一块不足$k = 32$个元素,把缺失值视为0,按照完整块处理。
80+ - 找到该块中数值的最大绝对值:
81+ 
82+ $$
83+ Amax(D_{fp32}^b)=max(\{|d_{i}|\}_{i=1}^{k})
84+ $$
85+ 
86+ - 将FP32映射到目标数据类型FP8可表示的范围内,其中$Amax(DType)$是目标精度能表示的最大值:
87+ 
88+ $$
89+ S_{fp32}^b = \frac{Amax(D_{fp32}^b)}{Amax(DType)}
90+ $$
91+ 
92+ - 将块缩放因子$S_{fp32}^b$转换为FP8格式下可表示的缩放值$S_{ue8m0}^b$
93+ - 从块的浮点缩放因子$S_{fp32}^b$中提取无偏指数$E_{int}^b$和尾数$M_{fixp}^b$
94+ - 为保证量化时不溢出,对指数进行向上取整,且在FP8可表示的范围内:
95+ 
96+ $$
97+ E_{int}^b = \begin{cases} E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为正规数,且} E_{int}^b < 254 \text{且} M_{fixp}^b > 0 \\ E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为非正规数,且} M_{fixp}^b > 0.5 \\ E_{int}^b, & \text{否则} \end{cases}
98+ $$
99+ 
100+ - 计算块缩放因子:$S_{ue8m0}^b=2^{E_{int}^b}$
101+ - 计算块转换因子:$R_{fp32}^b=\frac{1}{fp32(S_{ue8m0}^b)}$
102+ - 应用到量化的最终步骤,对于每个块内元素,$d^i = DType(d_{fp32}^i \cdot R_{fp32}^n)$,最终输出的量化结果是$\left(S^b, [d^i]_{i=1}^k\right)$,其中$S^b$代表块的缩放因子,这里指$S_{ue8m0}^b$,$[d^i]_{i=1}^k$代表块内量化后的数据。
103+ 
104+ </details>
105+ 
106+## 函数原型
107+ 
108+每个算子分为[两段式接口](../../../docs/zh/context/two_phase_api.md),必须先调用"aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize"接口获取计算所需workspace大小以及包含了算子计算流程的执行器,再调用"aclnnQuantMatmulActivationQuantWeightNz"接口执行计算。
109+ 
110+```cpp
111+aclnnStatus aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize(
112+ const aclTensor *x1,
113+ const aclTensor *x2,
114+ const aclTensor *x1ScaleOptional,
115+ const aclTensor *x2Scale,
116+ const aclTensor *biasOptional,
117+ bool transposeX1,
118+ bool transposeX2,
119+ int64_t groupSize,
120+ const char *activationType,
121+ const char *quantMode,
122+ const char *roundMode,
123+ int64_t scaleAlg,
124+ double dstTypeMax,
125+ aclTensor *yOut,
126+ aclTensor *yScaleOut,
127+ uint64_t *workspaceSize,
128+ aclOpExecutor **executor)
129+```
130+ 
131+```cpp
132+aclnnStatus aclnnQuantMatmulActivationQuantWeightNz(
133+ void *workspace,
134+ uint64_t workspaceSize,
135+ aclOpExecutor *executor,
136+ aclrtStream stream)
137+```
138+ 
139+## aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize
140+ 
141+- **参数说明:**
142+ 
143+ <table style="table-layout: fixed; width: 1554px"><colgroup>
144+ <col style="width: 248px">
145+ <col style="width: 121px">
146+ <col style="width: 210px">
147+ <col style="width: 327px">
148+ <col style="width: 250px">
149+ <col style="width: 115px">
150+ <col style="width: 138px">
151+ <col style="width: 145px">
152+ </colgroup>
153+ <thead>
154+ <tr>
155+ <th>参数名</th>
156+ <th>输入/输出</th>
157+ <th>描述</th>
158+ <th>使用说明</th>
159+ <th>数据类型</th>
160+ <th>数据格式</th>
161+ <th>维度(shape)</th>
162+ <th>非连续Tensor</th>
163+ </tr></thead>
164+ <tbody>
165+ <tr>
166+ <td>x1</td>
167+ <td>输入</td>
168+ <td>公式中的输入x1。</td>
169+ <td>
170+ <ul>
171+ <li>不支持空Tensor。</li>
172+ <li>支持最后两根轴转置情况下的<a href="../../../docs/zh/context/non_contiguous_tensor.md">非连续的Tensor</a>,其他轴方向不支持非连续的Tensor。</li>
173+ <li>transposeX1为false情况下x1各个维度表示:(batch, m, k),batch可不存在。</li>
174+ <li>transposeX1为true情况下x1各个维度表示:(batch, k, m),batch可不存在。</li>
175+ </ul>
176+ </td>
177+ <td>FLOAT8_E4M3FN、FLOAT8_E5M2</td>
178+ <td>ND</td>
179+ <td>2-6</td>
180+ <td>✓</td>
181+ </tr>
182+ <tr>
183+ <td>x2</td>
184+ <td>输入</td>
185+ <td>公式中的输入x2。</td>
186+ <td>
187+ <ul>
188+ <li>不支持空Tensor。</li>
189+ <li>NZ格式下,shape支持4-8维。</li>
190+ <li>transposeX2为true时x2各个维度表示:(batch, k1, n1, n0, k0),batch可不存在,k0 = 32,n0 = 16。</li>
191+ <li>transposeX2为false时x2各个维度表示:(batch, n1, k1, k0, n0),batch可不存在,k0 = 16,n0 = 32。</li>
192+ <li>x1 shape中的k和x2 shape中的k1需要满足ceil(k / k0) = k1,x2 shape中的n1与y的n需要满足ceil(n / n0) = n1。</li>
193+ <li>可使用aclnnCalculateMatmulWeightSizeV2接口以及aclnnTransMatmulWeight接口完成输入Format从ND到NZ格式的转换。</li>
194+ </ul>
195+ </td>
196+ <td>FLOAT8_E4M3FN</td>
197+ <td>NZ</td>
198+ <td>4-8</td>
199+ <td>✓</td>
200+ </tr>
201+ <tr>
202+ <td>x1ScaleOptional</td>
203+ <td>可选输入</td>
204+ <td>公式中的输入x1Scale。</td>
205+ <td>
206+ <ul>
207+ <li>不支持空Tensor。</li>
208+ <li>当x1Scale为FLOAT8_E8M0时,x1Scale为3维,各个维度表示:transposeX1为false时为(m, ceil(k / 64), 2),transposeX1为true时为(ceil(k / 64), m, 2)。</li>
209+ <li>如果不使用该参数,传入nullptr。</li>
210+ </ul>
211+ </td>
212+ <td>FLOAT8_E8M0</td>
213+ <td>ND</td>
214+ <td>1、3</td>
215+ <td>✓</td>
216+ </tr>
217+ <tr>
218+ <td>x2Scale</td>
219+ <td>输入</td>
220+ <td>表示量化参数,公式中的输入x2Scale。</td>
221+ <td>
222+ <ul>
223+ <li>不支持空Tensor。</li>
224+ <li>当x2Scale为FLOAT8_E8M0时,x2Scale为3维,各个维度表示:transposeX2为false时为(ceil(k / 64), n, 2),transposeX2为true时为(n, ceil(k / 64), 2)。</li>
225+ </ul>
226+ </td>
227+ <td>FLOAT8_E8M0</td>
228+ <td>ND</td>
229+ <td>1、3</td>
230+ <td>✓</td>
231+ </tr>
232+ <tr>
233+ <td>biasOptional</td>
234+ <td>可选输入</td>
235+ <td>公式中的输入bias。</td>
236+ <td>
237+ <ul>
238+ <li>不支持空Tensor。</li>
239+ <li>shape支持1维(n,)或3维(batch, 1, n),n与x2的n一致。</li>
240+ <li>当y的shape为2、4、5、6维时,bias的shape只支持1维(n,)。</li>
241+ <li>如果不使用该参数,传入nullptr。</li>
242+ </ul>
243+ </td>
244+ <td>FLOAT32</td>
245+ <td>ND</td>
246+ <td>1、3</td>
247+ <td>×</td>
248+ </tr>
249+ <tr>
250+ <td>transposeX1</td>
251+ <td>输入</td>
252+ <td>表示x1的输入shape是否转置。</td>
253+ <td>-</td>
254+ <td>BOOL</td>
255+ <td>-</td>
256+ <td>-</td>
257+ <td>-</td>
258+ </tr>
259+ <tr>
260+ <td>transposeX2</td>
261+ <td>输入</td>
262+ <td>表示x2的输入shape是否转置。</td>
263+ <td>-</td>
264+ <td>BOOL</td>
265+ <td>-</td>
266+ <td>-</td>
267+ <td>-</td>
268+ </tr>
269+ <tr>
270+ <td>groupSize</td>
271+ <td>可选输入</td>
272+ <td>用于输入m、n、k方向上的量化分组大小。</td>
273+ <td>由3个方向的groupSizeM,groupSizeN,groupSizeK三个值拼接组成,每个值占16位,共占用int64_t类型groupSize的低48位(groupSize中的高16位的数值无效),计算公式见表格下方公式一。不支持groupSize的场景,传入0。</td>
274+ <td>INT64</td>
275+ <td>-</td>
276+ <td>-</td>
277+ <td>-</td>
278+ </tr>
279+ <tr>
280+ <td>activationType</td>
281+ <td>输入</td>
282+ <td>激活的类型。</td>
283+ <td>支持{"gelu_tanh", "gelu_erf"}。</td>
284+ <td>STRING</td>
285+ <td>-</td>
286+ <td>-</td>
287+ <td>-</td>
288+ </tr>
289+ <tr>
290+ <td>quantMode</td>
291+ <td>输入</td>
292+ <td>量化的模式。</td>
293+ <td>支持{"mx"}。</td>
294+ <td>STRING</td>
295+ <td>-</td>
296+ <td>-</td>
297+ <td>-</td>
298+ </tr>
299+ <tr>
300+ <td>roundMode</td>
301+ <td>输入</td>
302+ <td>表示数据转换的模式,对应公式中的round_mode。</td>
303+ <td>
304+ <ul>
305+ <li>当yDtype为FLOAT4_E2M1/FLOAT4_E1M2时,支持{"rint", "floor", "round"}。</li>
306+ <li>当yDtype为FLOAT8_E4M3FN/FLOAT8_E5M2时,仅支持{"rint"}。</li>
307+ <li>传入空指针时,采用"rint"模式。</li>
308+ </ul>
309+ </td>
310+ <td>STRING</td>
311+ <td>-</td>
312+ <td>-</td>
313+ <td>-</td>
314+ </tr>
315+ <tr>
316+ <td>scaleAlg</td>
317+ <td>输入</td>
318+ <td>表示mxscaleOut的计算方法,对应公式中的scaleAlg。</td>
319+ <td>
320+ <ul>
321+ <li>支持取值0、1,取值为0代表场景1,为1代表场景2。</li>
322+ <li>当yDtype为FLOAT4_E1M2/FLOAT4_E2M1时仅支持取值为0。</li>
323+ <li>当yDtype为FLOAT8_E4M3FN/FLOAT8_E5M2时支持取值为0和1。</li>
324+ </ul>
325+ </td>
326+ <td>INT64</td>
327+ <td>-</td>
328+ <td>-</td>
329+ <td>-</td>
330+ </tr>
331+ <tr>
332+ <td>dstTypeMax</td>
333+ <td>输入</td>
334+ <td>表示maxType的取值,对应公式中的Amax(DType)。</td>
335+ <td>
336+ <ul>
337+ <li>支持取值0.0和6.0-12.0。</li>
338+ </ul>
339+ </td>
340+ <td>DOUBLE</td>
341+ <td>-</td>
342+ <td>-</td>
343+ <td>-</td>
344+ </tr>
345+ <tr>
346+ <td>yOut</td>
347+ <td>输出</td>
348+ <td>公式中的输出y,动态量化后的矩阵乘及激活计算结果。</td>
349+ <td>
350+ <ul>
351+ <li>不支持空Tensor。</li>
352+ <li>shape和矩阵乘计算结果一致,(batch, m, n),batch可不存在。</li>
353+ </ul>
354+ </td>
355+ <td>FLOAT8_E4M3FN、FLOAT8_E5M2</td>
356+ <td>ND</td>
357+ <td>2-6</td>
358+ <td>-</td>
359+ </tr>
360+ <tr>
361+ <td>yScaleOut</td>
362+ <td>输出</td>
363+ <td>动态量化后每个分组对应的量化尺度。</td>
364+ <td>
365+ <ul>
366+ <li>不支持空Tensor。</li>
367+ <li>shape在尾轴轴上为y对应轴的值除以32向上取整,并对其进行偶数pad,pad填充值为0。</li>
368+ <li>yScale输出需要对每两行数据进行交织处理。</li>
369+ <li>shape和矩阵乘计算结果一致,(batch, m, n/32, 2),batch可不存在。</li>
370+ </ul>
371+ </td>
372+ <td>FLOAT8_E8M0</td>
373+ <td>ND</td>
374+ <td>1-8</td>
375+ <td>-</td>
376+ </tr>
377+ <tr>
378+ <td>workspaceSize</td>
379+ <td>输出</td>
380+ <td>返回需要在Device侧申请的workspace大小。</td>
381+ <td>-</td>
382+ <td>-</td>
383+ <td>-</td>
384+ <td>-</td>
385+ <td>-</td>
386+ </tr>
387+ <tr>
388+ <td style="white-space: nowrap">executor</td>
389+ <td>输出</td>
390+ <td>返回op执行器,包含了算子计算流程。</td>
391+ <td>-</td>
392+ <td>-</td>
393+ <td>-</td>
394+ <td>-</td>
395+ <td>-</td>
396+ </tr>
397+ </tbody></table>
398+ 
399+ - 注:可选输入是指可选的量化参数,支持传入空指针。
400+ 
401+ - 公式一:<a name='f1'></a>
402+ 
403+ $$
404+ groupSize = groupSizeK | groupSizeN << 16 | groupSizeM << 32
405+ $$
406+ 
407+ <details>
408+ 
409+ <summary><term>Ascend 950PR/Ascend 950DT</term></summary>
410+ 
411+ - x2支持最后两根轴转置情况下的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md),其他场景的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md)不支持。
412+ - 支持groupSize传非0。
413+ </details>
414+ 
415+- **返回值:**
416+ 
417+ aclnnStatus:返回状态码,具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
418+ 
419+ 第一段接口完成入参校验,出现以下场景时报错:
420+ 
421+ <table style="table-layout: fixed; width: 1149px"><colgroup>
422+ <col style="width: 281px">
423+ <col style="width: 119px">
424+ <col style="width: 749px">
425+ </colgroup>
426+ <thead>
427+ <tr>
428+ <th>返回值</th>
429+ <th>错误码</th>
430+ <th>描述</th>
431+ </tr></thead>
432+ <tbody>
433+ <tr>
434+ <td>ACLNN_ERR_PARAM_NULLPTR</td>
435+ <td>161001</td>
436+ <td>传入的x1、x2、x1Scale或x2Scale是空指针。</td>
437+ </tr>
438+ <tr>
439+ <td rowspan="4">ACLNN_ERR_PARAM_INVALID</td>
440+ <td rowspan="4">161002</td>
441+ <td>x1、x2、biasOptional、x1ScaleOptional、x2Scale、yOut或yScaleOut的数据类型和数据格式不在支持的范围之内。</td>
442+ </tr>
443+ <tr>
444+ <td>x1、x2、biasOptional、x1ScaleOptional、x2Scale、yOut或yScaleOut的shape不满足校验条件。</td>
445+ </tr>
446+ <tr>
447+ <td>x1、x2、biasOptional、x1ScaleOptional、x2Scale、yOut或yScaleOut是空tensor。</td>
448+ </tr>
449+ </tbody></table>
450+ 
451+## aclnnQuantMatmulActivationQuantWeightNz
452+ 
453+- **参数说明:**
454+ 
455+ <table style="table-layout: fixed; width: 1150px"><colgroup>
456+ <col style="width: 168px">
457+ <col style="width: 128px">
458+ <col style="width: 854px">
459+ </colgroup>
460+ <thead>
461+ <tr>
462+ <th>参数名</th>
463+ <th>输入/输出</th>
464+ <th>描述</th>
465+ </tr></thead>
466+ <tbody>
467+ <tr>
468+ <td>workspace</td>
469+ <td>输入</td>
470+ <td>在Device侧申请的workspace内存地址。</td>
471+ </tr>
472+ <tr>
473+ <td>workspaceSize</td>
474+ <td>输入</td>
475+ <td>在Device侧申请的workspace大小,由第一段接口aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize获取。</td>
476+ </tr>
477+ <tr>
478+ <td>executor</td>
479+ <td>输入</td>
480+ <td>op执行器,包含了算子计算流程。</td>
481+ </tr>
482+ <tr>
483+ <td>stream</td>
484+ <td>输入</td>
485+ <td>指定执行任务的Stream。</td>
486+ </tr>
487+ </tbody></table>
488+ 
489+- **返回值:**
490+ 
491+ aclnnStatus:返回状态码,具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
492+ 
493+## 约束说明
494+ 
495+- 确定性计算:
496+ - aclnnQuantMatmulActivationQuantWeightNz默认确定性实现。
497+ 
498+<details>
499+ 
500+<summary><term>Ascend 950PR/Ascend 950DT</term></summary>
501+ 
502+- **公共约束:**
503+ <a id="公共约束"></a>
504+ - x1仅支持ND格式,x2仅支持NZ格式。
505+ - 当k或n为1时,无法使用weightNz特性,本接口不支持此种场景。
506+ - 支持调用本接口前,通过[aclnnTransMatmulWeight](https://gitcode.com/cann/ops-math/blob/master/conversion/trans_data/docs/aclnnTransMatmulWeight.md)或[aclnnNpuFormatCast](https://gitcode.com/cann/ops-math/blob/master/conversion/npu_format_cast/docs/aclnnNpuFormatCast.md)对format为ND的x2处理得到NZ格式,在使用时必须使用0来填充以防引入脏数据。
507+ - transposeX1为false时x1的shape:(batch, m, k)。transposeX1为true时x1的shape:(batch, k, m)。其中batch代表前0~4维,0维表示batch不存在。
508+ - transposeX2为false时x2的shape:(batch, n1, k1, k0, n0)。transposeX2为true时x2的shape:(batch, k1, n1, n0, k0)。其中batch代表前0~4维,0维表示batch不存在。k与x1的shape中的k一致。
509+ - x1支持最后两根轴转置情况下的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md),其他场景的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md)不支持。
510+ - x2支持最后两根轴转置情况下的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md),其他场景的[非连续的Tensor](../../../docs/zh/context/non_contiguous_tensor.md)不支持。
511+ - 当x1数据类型为FLOAT8_E5M2时,x2数据类型必须为FLOAT8_E4M3FN。
512+ - bias相关约束:
513+ - 可选参数,支持传入空指针, 不支持空tensor。
514+ - shape支持1维(n,)或3维(batch, 1, n),n与x2的n一致。
515+ - 当y的shape为2、4、5、6维时,bias的shape只支持1维(n,)。
516+ - groupSize相关约束:
517+ - 仅在MX[量化模式](../../../docs/zh/context/quant_mode_introduction.md)中生效。
518+ - 传入的groupSize内部会按如下公式分解得到groupSizeM、groupSizeN、groupSizeK, 当前支持[0, 0, 0]或[1, 1, 32]的组合, [0, 0, 0]会自动推导为[1, 1, 32]。
519+ 
520+ $$
521+ groupSize = groupSizeK | groupSizeN << 16 | groupSizeM << 32
522+ $$
523+ 
524+ - y的shape支持2~6维,(batch, m, n),batch可不存在,m与x1的m一致,n与x2的n一致。
525+ 
526+ <details>
527+ 
528+ <summary><strong>MX量化场景约束:</strong></summary>
529+ <a id="MX量化"></a>
530+ 
531+ - 输入和输出支持以下数据类型组合:
532+ <a id="输入和输出支持以下数据类型组合MX"></a>
533+ 
534+ | x1 | x2 | x1Scale | x2Scale | bias | y | yScale |
535+ |---------------|---------------|-------------|-------------|--------------|---------------------------|-------------|
536+ | FLOAT8_E4M3FN | FLOAT8_E4M3FN | FLOAT8_E8M0 | FLOAT8_E8M0 | null/FLOAT32 | FLOAT8_E4M3FN/FLOAT8_E5M2 | FLOAT8_E8M0 |
537+ | FLOAT8_E5M2 | FLOAT8_E4M3FN | FLOAT8_E8M0 | FLOAT8_E8M0 | null/FLOAT32 | FLOAT8_E4M3FN/FLOAT8_E5M2 | FLOAT8_E8M0 |
538+ 
539+ - x1数据类型、x2数据类型、x1、x2、x1Scale、x2Scale和groupSize的取值关系:
540+ 
541+ |量化类型|x1数据类型|x2数据类型|x1 shape|x2 shape|x1Scale shape|x2Scale shape|bias shape|yScale shape|[gsM, gsN, gsK]|groupSize|
542+ |-------|--------|--------|--------|--------|-------------|-------------|------------|---------------------------------------|--|--|
543+ |MX量化|FLOAT8_E4M3FN|FLOAT8_E4M3FN|<li>非转置:(batch, m, k)</li>|<li>非转置:(batch, k, n)</li><li>转置:(batch, n, k)</li>|<li>非转置:(batch, m, ceil(k / 64), 2)</li>|<li>非转置:(batch, ceil(k / 64), n, 2)</li><li>转置:(batch, n, ceil(k / 64), 2)</li>|(n,)或(batch, 1, n)|null|[1, 1, 32]|4295032864|
544+ |MX量化|FLOAT8_E5M2|FLOAT8_E4M3FN|<li>非转置:(batch, m, k)</li>|<li>非转置:(batch, k, n)</li><li>转置:(batch, n, k)</li>|<li>非转置:(batch, m, ceil(k / 64), 2)</li>|<li>非转置:(batch, ceil(k / 64), n, 2)</li><li>转置:(batch, n, ceil(k / 64), 2)</li>|(n,)或(batch, 1, n)|null|[1, 1, 32]|4295032864|
545+ 
546+ - 注:上表中gsM、gsK和gsN分别表示groupSizeM、groupSizeK和groupSizeN。
547+ - MX量化场景下,x1和x1Scale的转置属性需要保持一致,x2和x2Scale的转置属性需要保持一致。
548+ 
549+ </details>
550+ 
551+</details>
552+ 
553+## 调用示例
554+ 
555+示例代码如下,仅供参考,具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/compile_and_run_sample.md)。
556+ 
557+- <term>Ascend 950PR/Ascend 950DT</term>
558+x1为FLOAT8_E4M3FN,x2为FLOAT8_E4M3FN(NZ格式),x1Scale为FLOAT8_E8M0,x2Scale为FLOAT8_E8M0,激活为gelu_tanh,scaleAlg为0(OCP)。
559+ 
560+ ```cpp
561+ #include <iostream>
562+ #include <memory>
563+ #include <cmath>
564+ #include <vector>
565+ #include "acl/acl.h"
566+ #include "aclnnop/aclnn_cast.h"
567+ #include "aclnnop/aclnn_npu_format_cast.h"
568+ #include "aclnnop/aclnn_quant_matmul_activation_quant_weight_nz.h"
569+ #define CHECK_RET(cond, return_expr) \
570+ do { \
571+ if (!(cond)) { \
572+ return_expr; \
573+ } \
574+ } while (0)
575+ #define CHECK_FREE_RET(cond, return_expr) \
576+ do { \
577+ if (!(cond)) { \
578+ Finalize(deviceId, stream); \
579+ return_expr; \
580+ } \
581+ } while (0)
582+ #define LOG_PRINT(message, ...) \
583+ do { \
584+ printf(message, ##__VA_ARGS__); \
585+ } while (0)
586+ int64_t GetShapeSize(const std::vector<int64_t>& shape)
587+ {
588+ int64_t shapeSize = 1;
589+ for (auto i : shape) {
590+ shapeSize *= i;
591+ }
592+ return shapeSize;
593+ }
594+ int Init(int32_t deviceId, aclrtStream* stream)
595+ {
596+ // 固定写法,资源初始化
597+ auto ret = aclInit(nullptr);
598+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
599+ ret = aclrtSetDevice(deviceId);
600+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
601+ ret = aclrtCreateStream(stream);
602+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
603+ return 0;
604+ }
605+ template <typename T>
606+ int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
607+ aclDataType dataType, aclTensor** tensor)
608+ {
609+ auto size = GetShapeSize(shape) * sizeof(T);
610+ // 调用aclrtMalloc申请device侧内存
611+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
612+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
613+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
614+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
615+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
616+ // 计算连续tensor的strides
617+ std::vector<int64_t> strides(shape.size(), 1);
618+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
619+ strides[i] = shape[i + 1] * strides[i + 1];
620+ }
621+ // 调用aclCreateTensor接口创建aclTensor
622+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
623+ shape.data(), shape.size(), *deviceAddr);
624+ return 0;
625+ }
626+ void Finalize(int32_t deviceId, aclrtStream stream)
627+ {
628+ aclrtDestroyStream(stream);
629+ aclrtResetDevice(deviceId);
630+ aclFinalize();
631+ }
632+ // 将float8_e4m3的uint8_t表示转换为float表示
633+ float Fp4E4M3ToFloat(uint8_t h)
634+ {
635+ 
636+ int sign = (h >> 7) & 0x1;
637+ int exponent = (h >> 3) & 0xF;
638+ int mantissa = h & 0x7U;
639+ float value = 0.0f;
640+ if (exponent ==0) {
641+ if (mantissa == 0) {
642+ return sign ? -0.0f : 0.0f;
643+ } else {
644+ value = static_cast<float>(mantissa) / 8.0f;
645+ value = ldexp(value, -6);
646+ }
647+ } else {
648+ value = static_cast<float>(mantissa) / 8.0f + 1.0f;
649+ value = ldexp(value, exponent - 7);
650+ }
651+ 
652+ return sign ? -value : value;
653+ }
654+ 
655+ float Fp4E8M0ToFloat(uint8_t h)
656+ {
657+ uint32_t exponent = h & 0x00FFU; // exponent bits
658+ // mantissa 左移 23 - 7
659+ uint32_t fBits = exponent << 23;
660+ // 强转float
661+ return *reinterpret_cast<float*>(&fBits);
662+ }
663+ 
664+ template <typename T>
665+ int CreateAclTensorWithFormat(const std::vector<T>& hostData, const std::vector<int64_t>& shape, int64_t** storageShape,
666+ uint64_t* storageShapeSize, void** deviceAddr, aclDataType dataType, aclTensor** tensor,
667+ aclFormat format)
668+ {
669+ auto size = hostData.size() * sizeof(T);
670+ // 调用aclrtMalloc申请device侧内存
671+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
672+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
673+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
674+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
675+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
676+ // 计算连续tensor的strides
677+ std::vector<int64_t> strides(shape.size(), 1);
678+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
679+ strides[i] = shape[i + 1] * strides[i + 1];
680+ }
681+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, format, *storageShape,
682+ *storageShapeSize, *deviceAddr);
683+ return 0;
684+ }
685+ int AclnnQuantMatmulWeightNzActivationQuantTest(int32_t deviceId, aclrtStream& stream)
686+ {
687+ auto ret = Init(deviceId, &stream);
688+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
689+ // 2. 构造输入与输出,需要根据API的接口自定义构造
690+ int64_t m = 5;
691+ int64_t k = 64;
692+ int64_t n = 128;
693+ bool transposeX1 = false;
694+ bool transposeX2 = true;
695+ int64_t groupSize = 32;
696+ std::vector<int64_t> x1Shape = {m, k};
697+ std::vector<int64_t> x2Shape = {n, k};
698+ std::vector<int64_t> x1ScaleShape = {m, k / groupSize / 2, 2};
699+ std::vector<int64_t> x2ScaleShape = {n, k / groupSize / 2, 2};
700+ std::vector<int64_t> outShape = {m, n};
701+ std::vector<int64_t> outScaleShape = {m, n / groupSize / 2, 2};
702+ void* x1DeviceAddr = nullptr;
703+ void* x2DeviceAddr = nullptr;
704+ void* x2NzDeviceAddr = nullptr;
705+ void* x1ScaleDeviceAddr = nullptr;
706+ void* x2ScaleDeviceAddr = nullptr;
707+ void* outDeviceAddr = nullptr;
708+ void* outScaleDeviceAddr = nullptr;
709+ aclTensor* x1 = nullptr;
710+ aclTensor* x2 = nullptr;
711+ aclTensor* x1Scale = nullptr;
712+ aclTensor* x2Scale = nullptr;
713+ aclTensor* bias = nullptr;
714+ aclTensor* out = nullptr;
715+ aclTensor* outScale = nullptr;
716+ std::vector<uint8_t> x1HostData(m * k, 0b00111000); // float8_e4m3的1.0
717+ std::vector<uint8_t> x2HostData(n * k, 0b00111000); // float8_e4m3的1.0
718+ std::vector<uint8_t> x1ScaleHostData(m * k / groupSize, 0b01111111); // float8_e8m0的1.0
719+ std::vector<uint8_t> x2ScaleHostData(n * k / groupSize, 0b01111111); // float8_e8m0的1.0
720+ std::vector<uint8_t> outHostData(m * n, 0);
721+ std::vector<uint8_t> outScaleHostData(m * n / groupSize, 0);
722+ std::vector<int32_t> x2NzHostData(n * m, 0);
723+ int64_t* dstShape = nullptr;
724+ uint64_t dstShapeSize = 0;
725+ aclTensor* x2Nz = nullptr;
726+ int actualFormat;
727+ // 创建x1 aclTensor
728+ ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x1);
729+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1TensorPtr(x1, aclDestroyTensor);
730+ std::unique_ptr<void, aclError (*)(void*)> x1DeviceAddrPtr(x1DeviceAddr, aclrtFree);
731+ CHECK_RET(ret == ACL_SUCCESS, return ret);
732+ // 创建x2 aclTensor
733+ ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x2);
734+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2TensorPtr(x2, aclDestroyTensor);
735+ std::unique_ptr<void, aclError (*)(void*)> x2DeviceAddrPtr(x2DeviceAddr, aclrtFree);
736+ CHECK_RET(ret == ACL_SUCCESS, return ret);
737+ // 创建x1Scale aclTensor
738+ ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &x1Scale);
739+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1ScaleTensorPtr(x1Scale, aclDestroyTensor);
740+ std::unique_ptr<void, aclError (*)(void*)> x1ScaleDeviceAddrPtr(x1ScaleDeviceAddr, aclrtFree);
741+ CHECK_RET(ret == ACL_SUCCESS, return ret);
742+ // 创建x2Scale aclTensor
743+ ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &x2Scale);
744+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2ScaleTensorPtr(x2Scale, aclDestroyTensor);
745+ std::unique_ptr<void, aclError (*)(void*)> x2ScaleDeviceAddrPtr(x2ScaleDeviceAddr, aclrtFree);
746+ CHECK_RET(ret == ACL_SUCCESS, return ret);
747+ // 创建out aclTensor
748+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &out);
749+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outTensorPtr(out, aclDestroyTensor);
750+ std::unique_ptr<void, aclError (*)(void*)> outDeviceAddrPtr(outDeviceAddr, aclrtFree);
751+ CHECK_RET(ret == ACL_SUCCESS, return ret);
752+ // 创建mxscale aclTensor
753+ ret = CreateAclTensor(outScaleHostData, outScaleShape, &outScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &outScale);
754+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outScaleTensorPtr(outScale, aclDestroyTensor);
755+ std::unique_ptr<void, aclError (*)(void*)> outScaleDeviceAddrPtr(outScaleDeviceAddr, aclrtFree);
756+ CHECK_RET(ret == ACL_SUCCESS, return ret);
757+ // 3. 调用CANN算子库API,需要修改为具体的Api名称
758+ uint64_t workspaceSize = 0;
759+ aclOpExecutor* executor = nullptr;
760+ // x2转Nz
761+ // 计算目标tensor的shape和format
762+ aclDataType srcDtype = aclDataType::ACL_FLOAT8_E4M3FN;
763+ 
764+ ret = aclnnNpuFormatCastCalculateSizeAndFormat(x2, 29, aclDataType::ACL_FLOAT8_E4M3FN, &dstShape, &dstShapeSize, &actualFormat);
765+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCastCalculateSizeAndFormat failed. ERROR: %d\n", ret);
766+ return ret);
767+ ret = CreateAclTensorWithFormat(x2NzHostData, x2Shape, &dstShape, &dstShapeSize, &x2NzDeviceAddr, srcDtype, &x2Nz,
768+ static_cast<aclFormat>(actualFormat));
769+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2NzTensorPtr(x2Nz, aclDestroyTensor);
770+ std::unique_ptr<void, aclError (*)(void*)> x2NzDeviceAddrPtr(x2NzDeviceAddr, aclrtFree);
771+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("CreateAclTensorWithFormat failed. ERROR: %d\n", ret); return ret);
772+ // 调用aclnnNpuFormatCastGetWorkspaceSize第一段接口
773+ ret = aclnnNpuFormatCastGetWorkspaceSize(x2, x2Nz, &workspaceSize, &executor);
774+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCastGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
775+ // 根据第一段接口计算出的workspaceSize申请device内存
776+ void* workspaceNzAddr = nullptr;
777+ std::unique_ptr<void, aclError (*)(void*)> workspaceNzAddrPtr(nullptr, aclrtFree);
778+ if (workspaceSize > 0) {
779+ ret = aclrtMalloc(&workspaceNzAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
780+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
781+ workspaceNzAddrPtr.reset(workspaceNzAddr);
782+ }
783+ // 调用aclnnNpuFormatCastGetWorkspaceSize第二段接口
784+ ret = aclnnNpuFormatCast(workspaceNzAddr, workspaceSize, executor, stream);
785+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCast failed. ERROR: %d\n", ret); return ret);
786+ ret = aclrtSynchronizeStream(stream);
787+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
788+ workspaceSize = 0;
789+ executor = nullptr;
790+ ret = aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize(x1, x2Nz, x1Scale, x2Scale, nullptr,
791+ transposeX1, transposeX2, groupSize, "gelu_tanh", "mx", "rint", 0, 0.0, out, outScale, &workspaceSize, &executor);
792+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize failed. ERROR: %d\n", ret);
793+ return ret);
794+ // 根据第一段接口计算出的workspaceSize申请device内存
795+ void* workspaceAddr = nullptr;
796+ std::unique_ptr<void, aclError (*)(void*)> workspaceAddrPtr(nullptr, aclrtFree);
797+ if (workspaceSize > 0) {
798+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
799+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
800+ workspaceAddrPtr.reset(workspaceAddr);
801+ }
802+ // 调用aclnnQuantMatmulActivationQuantWeightNz第二段接口
803+ ret = aclnnQuantMatmulActivationQuantWeightNz(workspaceAddr, workspaceSize, executor, stream);
804+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulActivationQuantWeightNz failed. ERROR: %d\n", ret); return ret);
805+ // 4. (固定写法)同步等待任务执行结束
806+ ret = aclrtSynchronizeStream(stream);
807+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
808+ // 5. 获取输出的值,将device侧内存上的结果拷贝至host侧,需要根据具体API的接口定义修改
809+ auto size = GetShapeSize(outShape);
810+ std::vector<uint8_t> resultData(size, 0);
811+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
812+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
813+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
814+ for (int64_t i = 0; i < size; i++) {
815+ LOG_PRINT("result[%ld] is: %f\n", i, Fp4E4M3ToFloat(resultData[i]));
816+ }
817+ size = GetShapeSize(outScaleShape);
818+ std::vector<uint8_t> scaleData(size, 0);
819+ ret = aclrtMemcpy(scaleData.data(), scaleData.size() * sizeof(scaleData[0]), outScaleDeviceAddr,
820+ size * sizeof(scaleData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
821+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy scale result from device to host failed. ERROR: %d\n", ret); return ret);
822+ for (int64_t i = 0; i < size; i++) {
823+ LOG_PRINT("scale[%ld] is: %f\n", i, Fp4E8M0ToFloat(scaleData[i]));
824+ }
825+ return ACL_SUCCESS;
826+ }
827+ 
828+ int main()
829+ {
830+ // 1. (固定写法)device/stream初始化,参考acl API手册
831+ // 根据自己的实际device填写deviceId
832+ int32_t deviceId = 0;
833+ aclrtStream stream;
834+ auto ret = AclnnQuantMatmulWeightNzActivationQuantTest(deviceId, stream);
835+ CHECK_FREE_RET(ret == ACL_SUCCESS, LOG_PRINT("AclnnQuantMatmulWeightNzActivationQuantTest failed. ERROR: %d\n", ret); return ret);
836+ Finalize(deviceId, stream);
837+ return 0;
838+ }
839+ ```
@@ -0,0 +1,175 @@
1+# cann_ops_nn.quant_matmul_activation_quant
2+ 
3+## 产品支持情况
4+ 
5+- <term>Ascend 950PR/Ascend 950DT</term>:支持
6+- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>:不支持
7+- <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>:不支持
8+- <term>Atlas 200I/500 A2 推理产品</term>:不支持
9+- <term>Atlas 推理系列产品</term>:不支持
10+- <term>Atlas 训练系列产品</term>:不支持
11+ 
12+## 功能说明
13+ 
14+- 接口功能:
15+ 
16+ 融合量化的矩阵乘、激活以及动态量化,封装 `aclnnQuantMatmulActivationQuantWeightNz`。当前支持激活为 gelu、MX 量化模式。输入 `x1``x2` 为 FP8 量化矩阵,必选输入 `x2_scale`、可选输入 `x1_scale` 为 MX 量化缩放因子,`bias` 为偏置项;矩阵乘结果经激活函数后做动态量化,输出量化结果 `y` 和量化尺度 `y_scale`。M/N/K 维度及转置标志由 `x1``x2` 最后两维自动匹配推导,无需显式传入。
17+ 
18+- 计算公式:
19+ 
20+ - QuantMatmul MX 量化模式:
21+ 
22+ $$
23+ matmulOut[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice))* (x1Scale[m/gsM, j] * x2Scale[j, n/gsN]))+bias[n]
24+ $$
25+ 
26+ 其中,gsM、gsN 和 gsK 分别代表 groupSizeM、groupSizeN 和 groupSizeK;x1Slice 代表 x1 第 m 行长度为 groupSizeK 的向量,x2Slice 代表 x2 第 n 列长度为 groupSizeK 的向量;K 轴均从 j*groupSizeK 起始切片,j 的取值范围为 [0, kLoops),kLoops = ceil(K / groupSizeK),K 为 K 轴长度,支持最后的切片长度不足 groupSizeK。
27+ 
28+ - 激活计算公式:
29+ 
30+ - gelu_tanh(高性能近似):
31+ $$
32+ activationOut=GELU(matmulOut)=matmulOut × Φ(matmulOut)=0.5 * matmulOut * (1 + tanh( \sqrt{2 / \pi} * (matmulOut + 0.044715 * matmulOut^{3})))
33+ $$
34+ 
35+ - gelu_erf:
36+ $$
37+ activationOut=GELU(matmulOut)=0.5 * matmulOut * (1 + erf(matmulOut / \sqrt{2}))
38+ $$
39+ 
40+ - 动态量化计算公式:
41+ 
42+ - 场景 1,当 scale_alg 为 0 时:
43+ - 将输入 activationOut 在尾轴上按 $k = 32$ 个数分组,一组 k 个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\}, k = 32$
44+ 
45+ $$
46+ shared\_exp = floor(log_2(max_i(|V_i|))) - emax \\
47+ mxscale = 2^{shared\_exp}\\
48+ P_i = cast\_to\_dst\_type(V_i/mxscale, round\_mode), \space i\space from\space 1\space to\space k\\
49+ $$
50+ 
51+ - 量化后的 $P_{i}$ 按对应的 $V_{i}$ 的位置组成输出yOut,mxscale按尾轴上的分组输出yScaleOut。
52+ 
53+ - emax:对应数据类型的最大正则数的指数位。
54+ 
55+ | DataType | emax |
56+ | :-----------: | :--: |
57+ | FLOAT4_E2M1 | 2 |
58+ | FLOAT4_E1M2 | 0 |
59+ | FLOAT8_E4M3FN | 8 |
60+ | FLOAT8_E5M2 | 15 |
61+ 
62+ - 场景 2,当 scale_alg 为 1 时,只涉及 FP8 类型:
63+ - 将输入activationOut在尾轴上按$k = 32$个数分块,对每块单独计算一个块缩放因子$S_{fp32}^b$,再把块内所有元素用同一个$S_{fp32}^b$映射到目标低精度类型FP8。如果最后一块不足$k = 32$个元素,把缺失值视为0,按照完整块处理。
64+ - 找到该块中数值的最大绝对值:
65+ 
66+ $$
67+ Amax(D_{fp32}^b)=max(\{|d_{i}|\}_{i=1}^{k})
68+ $$
69+ 
70+ - 将 FP32 映射到目标数据类型 FP8 可表示的范围内,其中 $Amax(DType)$ 是目标精度能表示的最大值:
71+ 
72+ $$
73+ S_{fp32}^b = \frac{Amax(D_{fp32}^b)}{Amax(DType)}
74+ $$
75+ 
76+ - 将块缩放因子 $S_{fp32}^b$ 转换为 FP8 格式下可表示的缩放值 $S_{ue8m0}^b$
77+ - 从块的浮点缩放因子 $S_{fp32}^b$ 中提取无偏指数 $E_{int}^b$ 和尾数 $M_{fixp}^b$
78+ - 为保证量化时不溢出,对指数进行向上取整,且在 FP8 可表示的范围内:
79+ 
80+ $$
81+ E_{int}^b = \begin{cases} E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为正规数,且} E_{int}^b < 254 \text{且} M_{fixp}^b > 0 \\ E_{int}^b + 1, & \text{如果} S_{fp32}^b \text{为非正规数,且} M_{fixp}^b > 0.5 \\ E_{int}^b, & \text{否则} \end{cases}
82+ $$
83+ 
84+ - 计算块缩放因子:$S_{ue8m0}^b=2^{E_{int}^b}$
85+ - 计算块转换因子:$R_{fp32}^b=\frac{1}{fp32(S_{ue8m0}^b)}$
86+ - 应用到量化的最终步骤,对于每个块内元素,$d^i = DType(d_{fp32}^i \cdot R_{fp32}^n)$,最终输出的量化结果是 $\left(S^b, [d^i]_{i=1}^k\right)$,其中 $S^b$ 代表块的缩放因子,这里指 $S_{ue8m0}^b$,$[d^i]_{i=1}^k$ 代表块内量化后的数据。
87+ 
88+## 函数原型
89+ 
90+```python
91+torch.ops.cann_ops_nn.quant_matmul_activation_quant(x1, x2, x2_scale, *, x1_scale=None, bias=None,
92+ output_dtype=None, x1_dtype=None, x2_dtype=None, x1scale_dtype=None, x2scale_dtype=None,
93+ group_sizes=None, activation_type="gelu_tanh", quant_mode="mx", round_mode="rint",
94+ scale_alg=0, dst_type_max=0.0) -> (Tensor y, Tensor y_scale)
95+```
96+ 
97+## 参数说明
98+ 
99+| 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) |
100+| --- | --- | --- | --- | --- | --- |
101+| `x1` | Tensor | 必选 | 矩阵乘运算中的左矩阵。数据格式为 ND。最后两维为 `(M, K)``(K, M)`,由与 `x2` 的维度匹配自动推导。 | torch.float8_e4m3fn、torch.float8_e5m2 | 2-6 维,`(..., M, K)``(..., K, M)` |
102+| `x2` | Tensor | 必选 | 矩阵乘运算中的右矩阵。数据格式为 FRACTAL_NZ。最后两维为 `(K, N)` 或 `(N, K)`,由与 `x1` 的维度匹配自动推导。 | torch.float8_e4m3fn | 2-6 维,`(..., K, N)``(..., N, K)` |
103+| `x2_scale` | Tensor | 必选 | 矩阵乘计算时 x2 的 MX 量化缩放因子。数据格式为 ND。batch 维须与 `x2` 一致。 | torch.float8_e8m0fnu | `(..., K//64, N, 2)` 或 `(..., N, K//64, 2)`(随 `x2` 方向) |
104+| `x1_scale` | Tensor | 可选 | 矩阵乘计算时 x1 的 MX 量化缩放因子。数据格式为 ND。batch 维须与 `x1` 一致。 | torch.float8_e8m0fnu | `(..., M, K//64, 2)` 或 `(..., K//64, M, 2)`(随 `x1` 方向) |
105+| `bias` | Tensor | 可选 | 矩阵乘运算后累加的偏置。数据格式为 ND。 | float32 | `(N,)` |
106+| `output_dtype` | int | 可选 | 输出 `y` 的数据类型枚举值。支持 torch.float8_e4m3fn、torch.float8_e5m2 等。默认值 None(等价于 0,表示与 `x1` 同类型)。 | int | - |
107+| `x1_dtype` | int | 可选 | `x1` 的数据类型枚举值。不传入时根据 `x1` 的 scalar_type 自动推导。 | int | - |
108+| `x2_dtype` | int | 可选 | `x2` 的数据类型枚举值。不传入时根据 `x2` 的 scalar_type 自动推导。 | int | - |
109+| `x1scale_dtype` | int | 可选 | `x1_scale` 的数据类型枚举值。不传入时根据 `x1_scale` 的 scalar_type 自动推导。 | int | - |
110+| `x2scale_dtype` | int | 可选 | `x2_scale` 的数据类型枚举值。不传入时根据 `x2_scale` 的 scalar_type 自动推导。 | int | - |
111+| `group_sizes` | List[int] | 可选 | 分组量化大小 `[groupSizeM, groupSizeN, groupSizeK]`,每个元素取值范围为 [0, 65535]。 | list | `(3,)` |
112+| `activation_type` | str | 可选 | 激活函数类型,支持 `"gelu_tanh"``"gelu_erf"`,默认值 `"gelu_tanh"`。 | string | - |
113+| `quant_mode` | str | 可选 | 量化模式,当前支持 `"mx"`,默认值 `"mx"`。 | string | - |
114+| `round_mode` | str | 可选 | 舍入模式,支持 `"rint"``"floor"``"round"`,当前仅支持 `"rint"`。 | string | - |
115+| `scale_alg` | int | 可选 | 缩放算法。0 表示 MX 动态量化(场景 1),1 表示 FP8 块缩放量化(场景 2),默认值 0。 | int | - |
116+| `dst_type_max` | float | 可选 | 目标数据类型最大值,用于量化范围控制,默认值 0.0(表示使用目标类型的默认最大值)。 | float32 | - |
117+ 
118+## 返回值说明
119+ 
120+| 输出名 | 输出类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) |
121+| --- | --- | --- | --- | --- | --- |
122+| `y` | Tensor | 必选 | 动态量化后的矩阵乘及激活计算结果。 | torch.float8_e4m3fn、torch.float8_e5m2 | `(..., M, N)`; |
123+| `y_scale` | Tensor | 必选 | 动态量化后每个分组对应的量化尺度,最后一维固定为 2。 | torch.float8_e8m0fnu | `(..., M, CeilDiv(N, 64), 2)` |
124+ 
125+## 约束说明
126+ 
127+- 该接口支持训练、推理场景下使用。
128+- 该接口支持单算子模式调用。
129+- 不支持空 Tensor。
130+- 支持连续 Tensor,非连续 Tensor 仅支持最后两根轴转置场景。
131+- `x1` 支持 2-6 维,`x2`仅支持NZ 4-8 维。
132+-`K``N` 为 1 时,无法使用weightNz特性,本接口不支持此种场景。
133+- M/N/K 维度及 `transpose_x1``transpose_x2``x1``x2` 最后两维自动匹配推导:取 `x1` 最后两维和 `x2` 最后两维共四个值中相等的一对作为 K,`x1` 中剩余的为 M,`x2` 中剩余的为 N。若四组组合中无相等维度则报错。
134+- `x1``x2` 的 batch 维度(除最后两维外的维度)支持广播(右对齐),如 `x1=(1,M,K)``x2=(8,K,N)` 输出 `(8,M,N)`
135+- `x1_scale``x2_scale` 若传入,其 batch 维度(除最后三维外的维度)的数量和每一维的值必须与对应的 `x1``x2` 完全一致;若 `x1` 无 batch 维度(2D),则 `x1_scale``x2_scale` 须为 3D。
136+- `x1_scale``x2_scale` 最后一维必须为 2。
137+- `group_sizes` 若传入,必须包含三个元素 `[groupSizeM, groupSizeN, groupSizeK]`,每个元素取值范围为 [0, 65535],当前mx场景仅支持[1, 1, 32]。
138+- 输入和输出支持以下数据类型组合:
139+ 
140+ - <term>Ascend 950PR/Ascend 950DT</term>
141+ 
142+ | x1 | x2 | x1_scale | x2_scale | bias | y | y_scale |
143+ |---------------|---------------|-------------|-------------|------------------|---------------------------|-------------|
144+ | torch.float8_e4m3fn | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e4m3fn2 | torch.float8_e8m0fnu |
145+ | torch.float8_e5m2 | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e5m2 | torch.float8_e8m0fnu |
146+ 
147+## 确定性计算
148+ 
149+默认支持确定性计算。
150+ 
151+## 调用说明
152+ 
153+- 单算子模式调用
154+ 
155+ ```python
156+ import torch
157+ import torch_npu
158+ import cann_ops_nn
159+ 
160+ m, k, n = 5, 64, 128
161+ group_size = 32
162+ # x1 物理形状 (M, K);x2 物理形状 (K, N)
163+ x1 = torch.randn(m, k, dtype=torch.float32).to(torch.float8_e4m3fn).npu()
164+ x2 = torch.randn(k, n, dtype=torch.float32).to(torch.float8_e4m3fn).npu()
165+ x2_nz = torch_npu.npu_format_cast(x2, 29) # 29为NZ格式
166+ x1_scale = torch.ones(m, k // group_size // 2, 2, dtype=torch.float8_e8m0fnu).npu()
167+ x2_scale = torch.ones(k // group_size // 2, n, 2, dtype=torch.float8_e8m0fnu).npu()
168+ 
169+ y, y_scale = torch.ops.cann_ops_nn.quant_matmul_activation_quant(
170+ x1, x2_nz, x2_scale, x1_scale=x1_scale, bias=None,
171+ activation_type="gelu_tanh", quant_mode="mx", round_mode="rint",
172+ scale_alg=0, dst_type_max=0.0)
173+ print("y: ", y)
174+ print("y_scale: ", y_scale)
175+ ```
@@ -0,0 +1,295 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include <memory>
13+#include <cmath>
14+#include <vector>
15+#include "acl/acl.h"
16+#include "aclnnop/aclnn_cast.h"
17+#include "aclnnop/aclnn_npu_format_cast.h"
18+#include "aclnnop/aclnn_quant_matmul_activation_quant_weight_nz.h"
19+#define CHECK_RET(cond, return_expr) \
20+ do { \
21+ if (!(cond)) { \
22+ return_expr; \
23+ } \
24+ } while (0)
25+#define CHECK_FREE_RET(cond, return_expr) \
26+ do { \
27+ if (!(cond)) { \
28+ Finalize(deviceId, stream); \
29+ return_expr; \
30+ } \
31+ } while (0)
32+#define LOG_PRINT(message, ...) \
33+ do { \
34+ printf(message, ##__VA_ARGS__); \
35+ } while (0)
36+int64_t GetShapeSize(const std::vector<int64_t>& shape)
37+{
38+ int64_t shapeSize = 1;
39+ for (auto i : shape) {
40+ shapeSize *= i;
41+ }
42+ return shapeSize;
43+}
44+int Init(int32_t deviceId, aclrtStream* stream)
45+{
46+ // 固定写法,资源初始化
47+ auto ret = aclInit(nullptr);
48+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
49+ ret = aclrtSetDevice(deviceId);
50+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
51+ ret = aclrtCreateStream(stream);
52+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
53+ return 0;
54+}
55+template <typename T>
56+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
57+ aclDataType dataType, aclTensor** tensor)
58+{
59+ auto size = GetShapeSize(shape) * sizeof(T);
60+ // 调用aclrtMalloc申请device侧内存
61+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
62+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
63+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
64+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
65+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
66+ // 计算连续tensor的strides
67+ std::vector<int64_t> strides(shape.size(), 1);
68+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
69+ strides[i] = shape[i + 1] * strides[i + 1];
70+ }
71+ // 调用aclCreateTensor接口创建aclTensor
72+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
73+ shape.data(), shape.size(), *deviceAddr);
74+ return 0;
75+}
76+void Finalize(int32_t deviceId, aclrtStream stream)
77+{
78+ aclrtDestroyStream(stream);
79+ aclrtResetDevice(deviceId);
80+ aclFinalize();
81+}
82+// 将float8_e4m3的uint8_t表示转换为float表示
83+float Fp4E4M3ToFloat(uint8_t h)
84+{
85+ int sign = (h >> 7) & 0x1;
86+ int exponent = (h >> 3) & 0xF;
87+ int mantissa = h & 0x7U;
88+ float value = 0.0f;
89+ if (exponent == 0) {
90+ if (mantissa == 0) {
91+ return sign ? -0.0f : 0.0f;
92+ } else {
93+ value = static_cast<float>(mantissa) / 8.0f;
94+ value = ldexp(value, -6);
95+ }
96+ } else {
97+ value = static_cast<float>(mantissa) / 8.0f + 1.0f;
98+ value = ldexp(value, exponent - 7);
99+ }
100+ 
101+ return sign ? -value : value;
102+}
103+ 
104+float Fp4E8M0ToFloat(uint8_t h)
105+{
106+ uint32_t exponent = h & 0x00FFU; // exponent bits
107+ // mantissa 左移 23 - 7
108+ uint32_t fBits = exponent << 23;
109+ // 强转float
110+ return *reinterpret_cast<float*>(&fBits);
111+}
112+ 
113+template <typename T>
114+int CreateAclTensorWithFormat(const std::vector<T>& hostData, const std::vector<int64_t>& shape, int64_t** storageShape,
115+ uint64_t* storageShapeSize, void** deviceAddr, aclDataType dataType, aclTensor** tensor,
116+ aclFormat format)
117+{
118+ auto size = hostData.size() * sizeof(T);
119+ // 调用aclrtMalloc申请device侧内存
120+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
121+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
122+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
123+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
124+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
125+ // 计算连续tensor的strides
126+ std::vector<int64_t> strides(shape.size(), 1);
127+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
128+ strides[i] = shape[i + 1] * strides[i + 1];
129+ }
130+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, format, *storageShape,
131+ *storageShapeSize, *deviceAddr);
132+ return 0;
133+}
134+int AclnnQuantMatmulWeightNzActivationQuantTest(int32_t deviceId, aclrtStream& stream)
135+{
136+ auto ret = Init(deviceId, &stream);
137+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
138+ // 2. 构造输入与输出,需要根据API的接口自定义构造
139+ int64_t m = 5;
140+ int64_t k = 64;
141+ int64_t n = 128;
142+ bool transposeX1 = false;
143+ bool transposeX2 = true;
144+ int64_t groupSize = 32;
145+ std::vector<int64_t> x1Shape = {m, k};
146+ std::vector<int64_t> x2Shape = {n, k};
147+ std::vector<int64_t> x1ScaleShape = {m, k / groupSize / 2, 2};
148+ std::vector<int64_t> x2ScaleShape = {n, k / groupSize / 2, 2};
149+ std::vector<int64_t> outShape = {m, n};
150+ std::vector<int64_t> outScaleShape = {m, n / groupSize / 2, 2};
151+ void* x1DeviceAddr = nullptr;
152+ void* x2DeviceAddr = nullptr;
153+ void* x2NzDeviceAddr = nullptr;
154+ void* x1ScaleDeviceAddr = nullptr;
155+ void* x2ScaleDeviceAddr = nullptr;
156+ void* outDeviceAddr = nullptr;
157+ void* outScaleDeviceAddr = nullptr;
158+ aclTensor* x1 = nullptr;
159+ aclTensor* x2 = nullptr;
160+ aclTensor* x1Scale = nullptr;
161+ aclTensor* x2Scale = nullptr;
162+ aclTensor* bias = nullptr;
163+ aclTensor* out = nullptr;
164+ aclTensor* outScale = nullptr;
165+ std::vector<uint8_t> x1HostData(m * k, 0b00111000); // float8_e4m3的1.0
166+ std::vector<uint8_t> x2HostData(n * k, 0b00111000); // float8_e4m3的1.0
167+ std::vector<uint8_t> x1ScaleHostData(m * k / groupSize, 0b01111111); // float8_e8m0的1.0
168+ std::vector<uint8_t> x2ScaleHostData(n * k / groupSize, 0b01111111); // float8_e8m0的1.0
169+ std::vector<uint8_t> outHostData(m * n, 0);
170+ std::vector<uint8_t> outScaleHostData(m * n / groupSize, 0);
171+ std::vector<int32_t> x2NzHostData(n * m, 0);
172+ int64_t* dstShape = nullptr;
173+ uint64_t dstShapeSize = 0;
174+ aclTensor* x2Nz = nullptr;
175+ int actualFormat;
176+ // 创建x1 aclTensor
177+ ret = CreateAclTensor(x1HostData, x1Shape, &x1DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x1);
178+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1TensorPtr(x1, aclDestroyTensor);
179+ std::unique_ptr<void, aclError (*)(void*)> x1DeviceAddrPtr(x1DeviceAddr, aclrtFree);
180+ CHECK_RET(ret == ACL_SUCCESS, return ret);
181+ // 创建x2 aclTensor
182+ ret = CreateAclTensor(x2HostData, x2Shape, &x2DeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &x2);
183+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2TensorPtr(x2, aclDestroyTensor);
184+ std::unique_ptr<void, aclError (*)(void*)> x2DeviceAddrPtr(x2DeviceAddr, aclrtFree);
185+ CHECK_RET(ret == ACL_SUCCESS, return ret);
186+ // 创建x1Scale aclTensor
187+ ret = CreateAclTensor(x1ScaleHostData, x1ScaleShape, &x1ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &x1Scale);
188+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x1ScaleTensorPtr(x1Scale, aclDestroyTensor);
189+ std::unique_ptr<void, aclError (*)(void*)> x1ScaleDeviceAddrPtr(x1ScaleDeviceAddr, aclrtFree);
190+ CHECK_RET(ret == ACL_SUCCESS, return ret);
191+ // 创建x2Scale aclTensor
192+ ret = CreateAclTensor(x2ScaleHostData, x2ScaleShape, &x2ScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0, &x2Scale);
193+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2ScaleTensorPtr(x2Scale, aclDestroyTensor);
194+ std::unique_ptr<void, aclError (*)(void*)> x2ScaleDeviceAddrPtr(x2ScaleDeviceAddr, aclrtFree);
195+ CHECK_RET(ret == ACL_SUCCESS, return ret);
196+ // 创建out aclTensor
197+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT8_E4M3FN, &out);
198+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outTensorPtr(out, aclDestroyTensor);
199+ std::unique_ptr<void, aclError (*)(void*)> outDeviceAddrPtr(outDeviceAddr, aclrtFree);
200+ CHECK_RET(ret == ACL_SUCCESS, return ret);
201+ // 创建mxscale aclTensor
202+ ret = CreateAclTensor(outScaleHostData, outScaleShape, &outScaleDeviceAddr, aclDataType::ACL_FLOAT8_E8M0,
203+ &outScale);
204+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> outScaleTensorPtr(outScale, aclDestroyTensor);
205+ std::unique_ptr<void, aclError (*)(void*)> outScaleDeviceAddrPtr(outScaleDeviceAddr, aclrtFree);
206+ CHECK_RET(ret == ACL_SUCCESS, return ret);
207+ // 3. 调用CANN算子库API,需要修改为具体的Api名称
208+ uint64_t workspaceSize = 0;
209+ aclOpExecutor* executor = nullptr;
210+ // x2转Nz
211+ // 计算目标tensor的shape和format
212+ aclDataType srcDtype = aclDataType::ACL_FLOAT8_E4M3FN;
213+ 
214+ ret = aclnnNpuFormatCastCalculateSizeAndFormat(x2, 29, aclDataType::ACL_FLOAT8_E4M3FN, &dstShape, &dstShapeSize,
215+ &actualFormat);
216+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCastCalculateSizeAndFormat failed. ERROR: %d\n", ret);
217+ return ret);
218+ ret = CreateAclTensorWithFormat(x2NzHostData, x2Shape, &dstShape, &dstShapeSize, &x2NzDeviceAddr, srcDtype, &x2Nz,
219+ static_cast<aclFormat>(actualFormat));
220+ std::unique_ptr<aclTensor, aclnnStatus (*)(const aclTensor*)> x2NzTensorPtr(x2Nz, aclDestroyTensor);
221+ std::unique_ptr<void, aclError (*)(void*)> x2NzDeviceAddrPtr(x2NzDeviceAddr, aclrtFree);
222+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("CreateAclTensorWithFormat failed. ERROR: %d\n", ret); return ret);
223+ // 调用aclnnNpuFormatCastGetWorkspaceSize第一段接口
224+ ret = aclnnNpuFormatCastGetWorkspaceSize(x2, x2Nz, &workspaceSize, &executor);
225+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCastGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
226+ // 根据第一段接口计算出的workspaceSize申请device内存
227+ void* workspaceNzAddr = nullptr;
228+ std::unique_ptr<void, aclError (*)(void*)> workspaceNzAddrPtr(nullptr, aclrtFree);
229+ if (workspaceSize > 0) {
230+ ret = aclrtMalloc(&workspaceNzAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
231+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
232+ workspaceNzAddrPtr.reset(workspaceNzAddr);
233+ }
234+ // 调用aclnnNpuFormatCastGetWorkspaceSize第二段接口
235+ ret = aclnnNpuFormatCast(workspaceNzAddr, workspaceSize, executor, stream);
236+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNpuFormatCast failed. ERROR: %d\n", ret); return ret);
237+ ret = aclrtSynchronizeStream(stream);
238+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
239+ workspaceSize = 0;
240+ executor = nullptr;
241+ ret = aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize(x1, x2Nz, x1Scale, x2Scale, nullptr, transposeX1,
242+ transposeX2, groupSize, "gelu_tanh", "mx", "rint", 0,
243+ 0.0, out, outScale, &workspaceSize, &executor);
244+ CHECK_RET(ret == ACL_SUCCESS,
245+ LOG_PRINT("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize failed. ERROR: %d\n", ret);
246+ return ret);
247+ // 根据第一段接口计算出的workspaceSize申请device内存
248+ void* workspaceAddr = nullptr;
249+ std::unique_ptr<void, aclError (*)(void*)> workspaceAddrPtr(nullptr, aclrtFree);
250+ if (workspaceSize > 0) {
251+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
252+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
253+ workspaceAddrPtr.reset(workspaceAddr);
254+ }
255+ // 调用aclnnQuantMatmulActivationQuantWeightNz第二段接口
256+ ret = aclnnQuantMatmulActivationQuantWeightNz(workspaceAddr, workspaceSize, executor, stream);
257+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnQuantMatmulActivationQuantWeightNz failed. ERROR: %d\n", ret);
258+ return ret);
259+ // 4. (固定写法)同步等待任务执行结束
260+ ret = aclrtSynchronizeStream(stream);
261+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
262+ // 5. 获取输出的值,将device侧内存上的结果拷贝至host侧,需要根据具体API的接口定义修改
263+ auto size = GetShapeSize(outShape);
264+ std::vector<uint8_t> resultData(size, 0);
265+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
266+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
267+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
268+ for (int64_t i = 0; i < size; i++) {
269+ LOG_PRINT("result[%ld] is: %f\n", i, Fp4E4M3ToFloat(resultData[i]));
270+ }
271+ size = GetShapeSize(outScaleShape);
272+ std::vector<uint8_t> scaleData(size, 0);
273+ ret = aclrtMemcpy(scaleData.data(), scaleData.size() * sizeof(scaleData[0]), outScaleDeviceAddr,
274+ size * sizeof(scaleData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
275+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy scale result from device to host failed. ERROR: %d\n", ret);
276+ return ret);
277+ for (int64_t i = 0; i < size; i++) {
278+ LOG_PRINT("scale[%ld] is: %f\n", i, Fp4E8M0ToFloat(scaleData[i]));
279+ }
280+ return ACL_SUCCESS;
281+}
282+ 
283+int main()
284+{
285+ // 1. (固定写法)device/stream初始化,参考acl API手册
286+ // 根据自己的实际device填写deviceId
287+ int32_t deviceId = 0;
288+ aclrtStream stream;
289+ auto ret = AclnnQuantMatmulWeightNzActivationQuantTest(deviceId, stream);
290+ CHECK_FREE_RET(ret == ACL_SUCCESS,
291+ LOG_PRINT("AclnnQuantMatmulWeightNzActivationQuantTest failed. ERROR: %d\n", ret);
292+ return ret);
293+ Finalize(deviceId, stream);
294+ return 0;
295+}
@@ -0,0 +1,891 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "aclnn_kernels/common/op_error_check.h"
12+#include "aclnn_kernels/transdata.h"
13+#include "aclnn_kernels/transpose.h"
14+#include "aclnn_kernels/contiguous.h"
15+#include "aclnn_kernels/reshape.h"
16+#include "aclnn_quant_matmul_activation_quant_weight_nz.h"
17+#include "quant_matmul_activation_quant_util.h"
18+#include "matmul/common/op_host/op_api/matmul_util.h"
19+#include <dlfcn.h>
20+#include "securec.h"
21+#include "opdev/common_types.h"
22+#include "opdev/op_dfx.h"
23+#include "opdev/op_executor.h"
24+#include "opdev/op_log.h"
25+#include "opdev/platform.h"
26+#include "log/log.h"
27+#include "matmul/common/op_host/log_format_util.h"
28+#include "quant_matmul_activation_quant.h"
29+#include "quant_matmul_activation_quant_check.h"
30+#include "util/math_util.h"
31+ 
32+using namespace op;
33+using namespace QBMMActivationQuant;
34+using Ops::NN::FormatString;
35+using Ops::NN::StripEnclosingSquareBrackets;
36+using Ops::NN::SwapLastTwoDimValue;
37+ 
38+namespace {
39+struct MatmulShapeInfo {
40+ int64_t mDim;
41+ int64_t kDim;
42+ int64_t nDim;
43+};
44+ 
45+constexpr int IDX_0 = 0;
46+constexpr int IDX_1 = 1;
47+ 
48+static aclnnStatus CheckNotNull(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
49+{
50+ OP_CHECK_NULL(params.x1, return ACLNN_ERR_PARAM_NULLPTR);
51+ OP_CHECK_NULL(params.x2, return ACLNN_ERR_PARAM_NULLPTR);
52+ OP_CHECK_NULL(params.x1Scale, return ACLNN_ERR_PARAM_NULLPTR);
53+ OP_CHECK_NULL(params.x2Scale, return ACLNN_ERR_PARAM_NULLPTR);
ddssz
ddsszddssz7月20日

[安全性][必须修复] CheckNotNull 只检查了四个输入,没有检查必选输出 params.yparams.yScale。后续 CheckOptioanlAlg 会解引用 params.y,并且两个指针都会传给 ViewCopy;调用者传空输出时无法按 ACLNN 约定返回 ACLNN_ERR_PARAM_NULLPTR,而可能直接崩溃。请把两个输出加入入口判空,并保证判空发生在首次访问之前。

likedislike
十方君主
十方君主
7月21日 评论:
54+ OP_CHECK_NULL(params.y, return ACLNN_ERR_PARAM_NULLPTR);
55+ OP_CHECK_NULL(params.yScale, return ACLNN_ERR_PARAM_NULLPTR);
56+ return ACLNN_SUCCESS;
57+}
58+ 
59+static aclnnStatus CheckFormat(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
60+{
61+ if (params.x1->GetStorageFormat() != Format::FORMAT_ND) {
62+ OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1",
63+ op::ToString(params.x1->GetStorageFormat()).GetString(),
64+ "the format of x1 must be ND");
65+ return ACLNN_ERR_PARAM_INVALID;
66+ }
67+ if (params.x2->GetStorageFormat() != Format::FORMAT_FRACTAL_NZ) {
68+ OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2",
69+ op::ToString(params.x2->GetStorageFormat()).GetString(),
70+ "the format of x2 must be FORMAT_FRACTAL_NZ");
71+ return ACLNN_ERR_PARAM_INVALID;
72+ }
wangzitao
wangzitaowangzitao7月9日

【🔴 必须修改】x2 格式校验错误消息与实际检查矛盾

代码检查 x2 必须为 FORMAT_FRACTAL_NZ,但错误消息说 "the format of x2 must be ND",会误导用户。实际上 x2 必须是 FRACTAL_NZ 格式。

建议:将错误消息改为 "the format of x2 must be FRACTAL_NZ"

likedislike
zhuoyunhang
7月14日 评论:
73+ if (params.x1Scale->GetStorageFormat() != Format::FORMAT_ND) {
74+ OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1Scale",
75+ op::ToString(params.x1Scale->GetStorageFormat()).GetString(),
76+ "the format of x1Scale must be ND");
77+ return ACLNN_ERR_PARAM_INVALID;
78+ }
79+ if (params.x2Scale->GetStorageFormat() != Format::FORMAT_ND) {
80+ OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2Scale",
81+ op::ToString(params.x2Scale->GetStorageFormat()).GetString(),
82+ "the format of x2Scale must be ND");
83+ return ACLNN_ERR_PARAM_INVALID;
84+ }
85+ if (params.bias != nullptr && params.bias->GetStorageFormat() != Format::FORMAT_ND) {
86+ OP_LOGE_FOR_INVALID_FORMATS_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "bias",
87+ op::ToString(params.bias->GetStorageFormat()).GetString(),
88+ "the format of bias must be ND");
89+ return ACLNN_ERR_PARAM_INVALID;
90+ }
91+ return ACLNN_SUCCESS;
92+}
93+ 
94+static aclnnStatus IsMxQuantDim(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
95+{
96+ // scale 由固定 3 维(M/K/2 或 K/N/2)加上与 x1 一致的 batch 维度组成
97+ int64_t x1DimNum = static_cast<int64_t>(params.x1->GetViewShape().GetDimNum());
98+ int64_t x1BatchDimNum = std::max<int64_t>(x1DimNum - static_cast<int64_t>(MX_X1_DIM), 0);
99+ int64_t expectedScaleDimNum = static_cast<int64_t>(MX_X1_SCALE_DIM) + x1BatchDimNum;
100+ 
101+ auto x1ScaleDimNum = params.x1Scale->GetViewShape().GetDimNum();
102+ auto x2ScaleDimNum = params.x2Scale->GetViewShape().GetDimNum();
103+ if (static_cast<int64_t>(x2ScaleDimNum) != expectedScaleDimNum) {
104+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(
105+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2Scale",
106+ FormatString("%zuD", x2ScaleDimNum).c_str(),
107+ FormatString("when the quantization mode is mx, the shape dim of x2Scale must be %ld "
108+ "(batch dim of x1 %ld + fixed dim %zu)",
109+ expectedScaleDimNum, x1BatchDimNum, MX_X2_SCALE_DIM)
110+ .c_str());
111+ return ACLNN_ERR_PARAM_INVALID;
112+ }
113+ if (static_cast<int64_t>(x1ScaleDimNum) != expectedScaleDimNum) {
114+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(
115+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1Scale",
116+ FormatString("%zuD", x1ScaleDimNum).c_str(),
117+ FormatString("when the quantization mode is mx, the shape dim of x1Scale must be %ld "
118+ "(batch dim of x1 %ld + fixed dim %zu)",
119+ expectedScaleDimNum, x1BatchDimNum, MX_X1_SCALE_DIM)
120+ .c_str());
121+ return ACLNN_ERR_PARAM_INVALID;
122+ }
123+ 
124+ return ACLNN_SUCCESS;
125+}
126+ 
127+static aclnnStatus CheckInputDtypeValid(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
128+{
129+ if (!CheckType(params.x1->GetDataType(), X1_DTYPE_SUPPORT_LIST)) {
130+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1",
131+ op::ToString(params.x1->GetDataType()).GetString(),
132+ FormatString("the dtype of x1 must be in dtype support list %s",
133+ op::ToString(X1_DTYPE_SUPPORT_LIST).GetString())
134+ .c_str());
135+ return ACLNN_ERR_PARAM_INVALID;
136+ }
137+ if (!CheckType(params.x2->GetDataType(), X2_DTYPE_SUPPORT_LIST)) {
138+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2",
139+ op::ToString(params.x2->GetDataType()).GetString(),
140+ FormatString("the dtype of x2 must be in dtype support list %s",
141+ op::ToString(X2_DTYPE_SUPPORT_LIST).GetString())
142+ .c_str());
143+ return ACLNN_ERR_PARAM_INVALID;
144+ }
145+ return ACLNN_SUCCESS;
146+}
147+ 
148+static aclnnStatus CheckMxfp8DtypeValid(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
149+{
150+ if (CheckInputDtypeValid(params) != ACLNN_SUCCESS) {
151+ return ACLNN_ERR_PARAM_INVALID;
152+ }
153+ if (params.x1Scale->GetDataType() != op::DataType::DT_FLOAT8_E8M0) {
154+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
155+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1Scale",
156+ op::ToString(params.x1Scale->GetDataType()).GetString(),
157+ "when the quantization mode is mx, the dtype of x1Scale must be FLOAT8_E8M0");
158+ return ACLNN_ERR_PARAM_INVALID;
159+ }
160+ if (params.x2Scale->GetDataType() != op::DataType::DT_FLOAT8_E8M0) {
161+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
162+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2Scale",
163+ op::ToString(params.x2Scale->GetDataType()).GetString(),
164+ "when the quantization mode is mx, the dtype of x2Scale must be FLOAT8_E8M0");
165+ return ACLNN_ERR_PARAM_INVALID;
166+ }
167+ OP_LOGD("QuantMatmulActivationQuant CheckMxfp8DtypeValid success.");
168+ return ACLNN_SUCCESS;
ddssz
ddsszddssz7月20日

[正确性][必须修复] 当前 dtype 校验只覆盖输入和两个输入 scale,完全没有校验输出 yScale。调用者传入 FLOAT16 等类型时仍会继续到 ViewCopy,而内部实际产生的是 FLOAT8_E8M0;本 PR 的 exception_dtype_yScale_invalid 用例也无法由这里稳定返回参数错误。请显式要求 yScaleFLOAT8_E8M0,并在格式校验中同步要求 ND。

likedislike
十方君主
十方君主
7月21日 评论:
169+}
170+ 
171+static aclnnStatus CheckDtype(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
172+{
173+ auto x1Dtype = params.x1->GetDataType();
174+ auto x2Dtype = params.x2->GetDataType();
175+ auto x1ScaleDtype = params.x1Scale->GetDataType();
176+ auto x2ScaleDtype = params.x2Scale->GetDataType();
177+ auto yDtype = params.y->GetDataType();
178+ auto yScaleDtype = params.yScale->GetDataType();
179+ if ((x1Dtype == DataType::DT_FLOAT8_E4M3FN || x1Dtype == DataType::DT_FLOAT8_E5M2) &&
180+ x2Dtype == DataType::DT_FLOAT8_E4M3FN &&
181+ (yDtype == DataType::DT_FLOAT8_E4M3FN ||
182+ yDtype == DataType::DT_FLOAT8_E5M2 && yScaleDtype == DataType::DT_FLOAT8_E8M0)) {
183+ CHECK_COND(IsMxQuantDim(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "Check IsMxQuantDim failed.");
184+ if (params.bias != nullptr && params.bias->GetDataType() != op::DataType::DT_FLOAT) {
185+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "bias",
186+ op::ToString(params.bias->GetDataType()).GetString(),
187+ "the dtype of bias must be FLOAT");
188+ return ACLNN_ERR_PARAM_INVALID;
189+ }
190+ return CheckMxfp8DtypeValid(params);
191+ } else {
192+ OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(
193+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1, x2, x1Scale, x2Scale, y, yScale",
194+ FormatString("%s, %s, %s, %s", op::ToString(x1Dtype).GetString(), op::ToString(x2Dtype).GetString(),
195+ op::ToString(x1ScaleDtype).GetString(), op::ToString(x2ScaleDtype).GetString())
196+ .c_str(),
197+ FormatString("when the dtypes of x1 and x2 are %s and %s, and the dtypes of x1Scale and x2Scale are %s "
198+ "and %s, this dtype combination can not be supported",
199+ op::ToString(x1Dtype).GetString(), op::ToString(x2Dtype).GetString(),
200+ op::ToString(x1ScaleDtype).GetString(), op::ToString(x2ScaleDtype).GetString())
201+ .c_str());
202+ return ACLNN_ERR_PARAM_INVALID;
203+ }
204+}
205+ 
206+static aclnnStatus CheckOptioanlAlg(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
207+{
208+ CHECK_RET(params.activationType != nullptr, ACLNN_ERR_PARAM_NULLPTR);
209+ const std::string activationType(params.activationType);
ddssz
ddsszddssz7月20日

[安全性][必须修复] 这是公开 C 接口传入的字符指针,但在任何判空之前就用它构造 std::stringactivationType == nullptr 时行为未定义,quantModeroundMode 也有同样问题。请在 API 入口统一校验三个字符串指针(或明确并安全应用默认值),再进行 DFX/字符串构造和取值校验。

likedislike
十方君主
十方君主
7月21日 评论:
210+ if (activationType != "gelu_tanh" && activationType != "gelu_erf") {
211+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize",
212+ "activationType", activationType,
213+ "The activationType must be gelu_tanh or gelu_erf");
214+ return ACLNN_ERR_PARAM_INVALID;
215+ }
216+ CHECK_RET(params.quantMode != nullptr, ACLNN_ERR_PARAM_NULLPTR);
217+ const std::string quantMode(params.quantMode);
218+ if (quantMode != "mx") {
219+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "quantMode",
220+ quantMode, "The quantMode must be mx");
221+ return ACLNN_ERR_PARAM_INVALID;
222+ }
223+ CHECK_RET(params.roundMode != nullptr, ACLNN_ERR_PARAM_NULLPTR);
224+ std::string roundMode(params.roundMode);
225+ if (roundMode != "rint" && roundMode != "floor" && roundMode != "round") {
226+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "roundMode",
227+ roundMode,
228+ "roundMode optional values are rint/floor/round, it's enabled when "
229+ "dynamic mx quant, fp8 only support rint, fp4 support all");
230+ return ACLNN_ERR_PARAM_INVALID;
231+ }
232+ if (params.scaleAlg != 0 && params.scaleAlg != 1) {
233+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "scaleAlg",
234+ std::to_string(params.scaleAlg), "The scaleAlg optional values are 0/1");
235+ return ACLNN_ERR_PARAM_INVALID;
236+ }
237+ if (CheckType(params.y->GetDataType(), Y_DTYPE_SUPPORT_LIST) && roundMode != "rint") {
238+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "roundMode",
239+ roundMode,
240+ FormatString("roundMode must be rint when the dtype of y in %s",
241+ op::ToString(Y_DTYPE_SUPPORT_LIST).GetString())
242+ .c_str());
243+ return ACLNN_ERR_PARAM_INVALID;
244+ }
245+ if (!CheckType(params.y->GetDataType(), Y_DTYPE_SUPPORT_LIST) && params.scaleAlg == 1) {
246+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "scaleAlg",
247+ std::to_string(params.scaleAlg),
248+ FormatString("scaleAlg can't be 1 when the dtype of y not in %s",
249+ op::ToString(Y_DTYPE_SUPPORT_LIST).GetString())
250+ .c_str());
251+ return ACLNN_ERR_PARAM_INVALID;
252+ }
253+ return ACLNN_SUCCESS;
254+}
255+ 
256+static MatmulShapeInfo GetMatmulShapeInfo(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
257+{
258+ int64_t x1DimNum = params.x1->GetViewShape().GetDimNum();
259+ int64_t x2DimNum = params.x2->GetViewShape().GetDimNum();
260+ return {
261+ params.transposeX1 ?
262+ params.x1->GetViewShape().GetDim(x1DimNum - 1) :
263+ params.x1->GetViewShape().GetDim(x1DimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM),
264+ params.transposeX1 ?
265+ params.x1->GetViewShape().GetDim(x1DimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM) :
266+ params.x1->GetViewShape().GetDim(x1DimNum - 1),
267+ params.transposeX2 ?
268+ params.x2->GetViewShape().GetDim(x2DimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM) :
269+ params.x2->GetViewShape().GetDim(x2DimNum - 1),
270+ };
271+}
272+ 
273+static aclnnStatus CheckInputOutDims(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
274+{
275+ auto x1DimNum = params.x1->GetViewShape().GetDimNum();
276+ auto x2DimNum = params.x2->GetStorageShape().GetDimNum();
277+ if (x1DimNum < MX_X1_DIM_MIN || x1DimNum > MX_X1_DIM_MAX) {
278+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1",
279+ FormatString("%zuD", x1DimNum).c_str(),
280+ FormatString("the shape dim of x1 must be in the range of 2 to 6"));
281+ return ACLNN_ERR_PARAM_INVALID;
282+ }
283+ if (x2DimNum < MX_X2_DIM_MIN || x2DimNum > MX_X2_DIM_MAX) {
284+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2",
285+ FormatString("%zuD", x2DimNum).c_str(),
286+ FormatString("the shape dim of x2 must be in the range of 4 to 8"));
287+ return ACLNN_ERR_PARAM_INVALID;
288+ }
289+ 
290+ return ACLNN_SUCCESS;
291+}
292+ 
293+static aclnnStatus CheckShapeInfoMatch(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params,
294+ const MatmulShapeInfo& shapeInfo)
295+{
296+ int64_t x2DimNum = params.x2->GetViewShape().GetDimNum();
297+ int64_t x2KDim = params.transposeX2 ? params.x2->GetViewShape().GetDim(x2DimNum - 1) :
298+ params.x2->GetViewShape().GetDim(
299+ x2DimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM);
300+ if (shapeInfo.kDim != x2KDim) {
301+ OP_LOGE_FOR_INVALID_VALUES_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1 K, x2 K",
302+ FormatString("%ld, %ld", shapeInfo.kDim, x2KDim).c_str(),
303+ "the K dimension of x1 and x2 must be equal");
304+ return ACLNN_ERR_PARAM_INVALID;
305+ }
306+ return ACLNN_SUCCESS;
307+}
308+ 
309+static bool CheckMKN(int64_t m, int64_t k, int64_t n)
310+{
311+ if (m <= 0) {
312+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1 M",
313+ std::to_string(m).c_str(), "the M dimension of x1 must be positive");
314+ return false;
315+ }
316+ if (k <= 0) {
317+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "K",
318+ std::to_string(k).c_str(),
319+ "the K dimension of x1 and x2 must be positive");
320+ return false;
321+ }
322+ if (n <= 0) {
323+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2 N",
324+ std::to_string(n).c_str(), "the N dimension of x2 must be positive");
325+ return false;
326+ }
327+ return true;
328+}
329+ 
330+static aclnnStatus CheckWeightNzParamsDAV3510(const aclTensor* x1, const aclTensor* x2)
331+{
332+ if (op::GetCurrentPlatformInfo().GetCurNpuArch() != NpuArch::DAV_3510) {
333+ return ACLNN_SUCCESS;
334+ }
335+ 
336+ if (x1 == nullptr) {
337+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1", "null",
338+ "x1 can not be null");
339+ return ACLNN_ERR_PARAM_NULLPTR;
340+ }
341+ if (x2 == nullptr) {
342+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2", "null",
343+ "x2 can not be null");
344+ return ACLNN_ERR_PARAM_NULLPTR;
345+ }
346+ 
347+ if (static_cast<ge::Format>(ge::GetPrimaryFormat(x2->GetStorageFormat())) != Format::FORMAT_FRACTAL_NZ) {
348+ OP_LOGE_FOR_INVALID_FORMAT_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2",
349+ op::ToString(x2->GetStorageFormat()).GetString(),
350+ "the format of x2 must be FRACTAL_NZ");
351+ return ACLNN_ERR_PARAM_INVALID;
352+ }
353+ 
354+ // NZ情况下,x2的k和n不能为1
355+ int64_t dim1 = x2->GetViewShape().GetDimNum() - 1;
356+ int64_t dim2 = x2->GetViewShape().GetDimNum() - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM;
357+ if (x2->GetViewShape().GetDim(dim2) == 1 || x2->GetViewShape().GetDim(dim1) == 1) {
358+ OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(
359+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2 K, x2 N",
360+ FormatString("%ld, %ld", x2->GetViewShape().GetDim(dim2), x2->GetViewShape().GetDim(dim1)).c_str(),
361+ "when the format of x2 is FRACTAL_NZ, the k dimension and n dimension of x2 can not be 1");
362+ return ACLNN_ERR_PARAM_INVALID;
363+ }
364+ 
365+ OP_LOGD("QuantMatmulWeightNz check params success.");
366+ return ACLNN_SUCCESS;
367+}
368+ 
369+static inline bool IsMicroScaling(const aclTensor* x1Scale, const aclTensor* x2Scale)
370+{
371+ if (x1Scale == nullptr || x2Scale == nullptr) {
372+ return false;
373+ }
374+ return x1Scale->GetDataType() == op::DataType::DT_FLOAT8_E8M0 &&
375+ x2Scale->GetDataType() == op::DataType::DT_FLOAT8_E8M0;
376+}
377+ 
378+static aclnnStatus CheckMxScaleLastDim(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
379+{
380+ if (!IsMicroScaling(params.x1Scale, params.x2Scale)) {
381+ return ACLNN_SUCCESS;
382+ }
383+ 
384+ auto scale1LastDimValue = params.x1Scale->GetViewShape().GetDim(params.x1Scale->GetViewShape().GetDimNum() - 1);
385+ auto scale2LastDimValue = params.x2Scale->GetViewShape().GetDim(params.x2Scale->GetViewShape().GetDimNum() - 1);
386+ if (scale1LastDimValue != MXFP_MULTI_BASE_SIZE) {
387+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
388+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1Scale",
389+ StripEnclosingSquareBrackets(op::ToString(params.x1Scale->GetViewShape()).GetString()).c_str(),
390+ FormatString("when the quantization mode is mx, the last dimension of x1Scale must be %d",
391+ MXFP_MULTI_BASE_SIZE)
392+ .c_str());
393+ return ACLNN_ERR_PARAM_INVALID;
394+ }
395+ if (scale2LastDimValue != MXFP_MULTI_BASE_SIZE) {
396+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
397+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2Scale",
398+ StripEnclosingSquareBrackets(op::ToString(params.x2Scale->GetViewShape()).GetString()).c_str(),
399+ FormatString("when the quantization mode is mx, the last dimension of x2Scale must be %d",
400+ MXFP_MULTI_BASE_SIZE)
401+ .c_str());
402+ return ACLNN_ERR_PARAM_INVALID;
403+ }
404+ return ACLNN_SUCCESS;
405+}
406+ 
407+static void GetExpectedScaleShape(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params,
408+ const MatmulShapeInfo& shapeInfo, op::Shape& x1ScaleExpectShape,
409+ op::Shape& x2ScaleExpectShape)
410+{
411+ if (!IsMicroScaling(params.x1Scale, params.x2Scale)) {
412+ x1ScaleExpectShape = {1};
413+ x2ScaleExpectShape = {1};
414+ return;
415+ }
416+ 
417+ // batch 维度与 x1 保持一致
418+ const auto& x1View = params.x1->GetViewShape();
419+ int64_t x1DimNum = static_cast<int64_t>(x1View.GetDimNum());
420+ int64_t x1BatchDimNum = std::max<int64_t>(x1DimNum - static_cast<int64_t>(MX_X1_DIM), 0);
421+ 
422+ x1ScaleExpectShape = op::Shape();
423+ x2ScaleExpectShape = op::Shape();
424+ for (int64_t i = 0; i < x1BatchDimNum; ++i) {
425+ x1ScaleExpectShape.AppendDim(x1View.GetDim(i));
426+ x2ScaleExpectShape.AppendDim(x1View.GetDim(i));
427+ }
428+ if (params.transposeX1) {
429+ x1ScaleExpectShape.AppendDim(Ops::Base::CeilDiv(shapeInfo.kDim, SPLIT_SIZE));
430+ x1ScaleExpectShape.AppendDim(shapeInfo.mDim);
431+ } else {
432+ x1ScaleExpectShape.AppendDim(shapeInfo.mDim);
433+ x1ScaleExpectShape.AppendDim(Ops::Base::CeilDiv(shapeInfo.kDim, SPLIT_SIZE));
434+ }
435+ x1ScaleExpectShape.AppendDim(MXFP_MULTI_BASE_SIZE);
436+ 
437+ if (params.transposeX2) {
438+ x2ScaleExpectShape.AppendDim(shapeInfo.nDim);
439+ x2ScaleExpectShape.AppendDim(Ops::Base::CeilDiv(shapeInfo.kDim, SPLIT_SIZE));
440+ } else {
441+ x2ScaleExpectShape.AppendDim(Ops::Base::CeilDiv(shapeInfo.kDim, SPLIT_SIZE));
442+ x2ScaleExpectShape.AppendDim(shapeInfo.nDim);
443+ }
444+ x2ScaleExpectShape.AppendDim(MXFP_MULTI_BASE_SIZE);
445+}
446+ 
447+static aclnnStatus CheckExpectedShapes(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params,
448+ const MatmulShapeInfo& shapeInfo)
449+{
450+ auto& x1View = params.x1->GetViewShape();
451+ auto& x2View = params.x2->GetViewShape();
452+ int64_t x1DimNum = x1View.GetDimNum();
453+ int64_t x2DimNum = x2View.GetDimNum();
454+ 
455+ // 维度下限校验:x1/x2 此处均为 viewShape(ND 逻辑维度),下限与 x1 一致
456+ if (x1DimNum < MX_X1_DIM || x2DimNum < MX_X1_DIM) {
457+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "unsupported dim combination: x1DimNum=%ld, x2DimNum=%ld", x1DimNum, x2DimNum);
458+ return ACLNN_ERR_PARAM_INVALID;
459+ }
460+ 
461+ // 1. batch 维 broadcast 合法性校验(2D 时 batchDimNum=0,循环不执行)
462+ int64_t x1BatchCount = x1DimNum - 2;
463+ int64_t x2BatchCount = x2DimNum - 2;
464+ int64_t batchDimNum = std::max(x1BatchCount, x2BatchCount);
465+ for (int64_t i = 0; i < batchDimNum; ++i) {
466+ // 右对齐:batch 维从左向右编号,不足的补 1
467+ int64_t x1Idx = i - (batchDimNum - x1BatchCount);
468+ int64_t x2Idx = i - (batchDimNum - x2BatchCount);
469+ int64_t x1BatchDim = (x1Idx >= 0) ? x1View.GetDim(x1Idx) : 1;
470+ int64_t x2BatchDim = (x2Idx >= 0) ? x2View.GetDim(x2Idx) : 1;
471+ if (x1BatchDim != x2BatchDim && x1BatchDim != 1 && x2BatchDim != 1) {
472+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "batch dim %ld mismatch: x1=%ld, x2=%ld", i, x1BatchDim, x2BatchDim);
473+ return ACLNN_ERR_PARAM_INVALID;
474+ }
475+ }
476+ 
477+ // 2. 校验 x1 最后两维
478+ int64_t x1M = params.transposeX1 ? x1View.GetDim(x1DimNum - 1) : x1View.GetDim(x1DimNum - 2);
479+ int64_t x1K = params.transposeX1 ? x1View.GetDim(x1DimNum - 2) : x1View.GetDim(x1DimNum - 1);
480+ if (x1M != shapeInfo.mDim || x1K != shapeInfo.kDim) {
481+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1",
482+ StripEnclosingSquareBrackets(op::ToString(x1View).GetString()).c_str(),
483+ FormatString("x1 last two dims must be [%ld, %ld], but got [%ld, %ld]",
484+ params.transposeX1 ? shapeInfo.kDim : shapeInfo.mDim,
485+ params.transposeX1 ? shapeInfo.mDim : shapeInfo.kDim,
486+ x1View.GetDim(x1DimNum - 2), x1View.GetDim(x1DimNum - 1))
487+ .c_str());
488+ return ACLNN_ERR_PARAM_INVALID;
489+ }
490+ 
491+ // 3. 校验 x2 最后两维
492+ int64_t x2K = params.transposeX2 ? x2View.GetDim(x2DimNum - 1) : x2View.GetDim(x2DimNum - 2);
493+ int64_t x2N = params.transposeX2 ? x2View.GetDim(x2DimNum - 2) : x2View.GetDim(x2DimNum - 1);
494+ if (x2K != shapeInfo.kDim || x2N != shapeInfo.nDim) {
495+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2",
496+ StripEnclosingSquareBrackets(op::ToString(x2View).GetString()).c_str(),
497+ FormatString("x2 last two dims must be [%ld, %ld], but got [%ld, %ld]",
498+ params.transposeX2 ? shapeInfo.nDim : shapeInfo.kDim,
499+ params.transposeX2 ? shapeInfo.kDim : shapeInfo.nDim,
500+ x2View.GetDim(x2DimNum - 2), x2View.GetDim(x2DimNum - 1))
501+ .c_str());
502+ return ACLNN_ERR_PARAM_INVALID;
503+ }
504+ 
505+ op::Shape x1ScaleExpectShape;
506+ op::Shape x2ScaleExpectShape;
507+ GetExpectedScaleShape(params, shapeInfo, x1ScaleExpectShape, x2ScaleExpectShape);
508+ 
509+ if (params.x1Scale->GetViewShape() != x1ScaleExpectShape) {
510+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
511+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1Scale",
512+ StripEnclosingSquareBrackets(op::ToString(params.x1Scale->GetViewShape()).GetString()).c_str(),
513+ FormatString("the shape of x1Scale must be %s", op::ToString(x1ScaleExpectShape).GetString()).c_str());
514+ return ACLNN_ERR_PARAM_INVALID;
515+ }
516+ if (params.x2Scale->GetViewShape() != x2ScaleExpectShape) {
517+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
518+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x2Scale",
519+ StripEnclosingSquareBrackets(op::ToString(params.x2Scale->GetViewShape()).GetString()).c_str(),
520+ FormatString("the shape of x2Scale must be %s", op::ToString(x2ScaleExpectShape).GetString()).c_str());
521+ return ACLNN_ERR_PARAM_INVALID;
522+ }
523+ return ACLNN_SUCCESS;
524+}
525+ 
526+static int64_t InferOutputShape(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
527+{
528+ int64_t inferedOutbatchValue = 1;
529+ auto x1DimNum = params.x1->GetViewShape().GetDimNum();
530+ auto x2DimNum = params.x2->GetViewShape().GetDimNum();
531+ auto outDimNum = std::max(x1DimNum, x2DimNum);
532+ auto& longShapeTensor = x1DimNum > x2DimNum ? params.x1 : params.x2;
533+ auto& shortShapeTensor = x1DimNum > x2DimNum ? params.x2 : params.x1;
534+ size_t validOffset = outDimNum - std::min(x1DimNum, x2DimNum);
535+ for (size_t i = 0; i + QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM < outDimNum; i++) {
536+ auto shortDimValue = i < validOffset ? 1 : shortShapeTensor->GetViewShape().GetDim(i - validOffset);
537+ auto longDimValue = longShapeTensor->GetViewShape().GetDim(i);
538+ if (shortDimValue > 1 && longDimValue > 1 && shortDimValue != longDimValue) {
539+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
540+ "Current short dim value %ld and long dim value %ld are not supported for broadcasting.",
541+ shortDimValue, longDimValue);
542+ return OUTPUT_INFER_FAIL;
543+ }
544+ int64_t curBatchValue = static_cast<int64_t>(std::max(shortDimValue, longDimValue));
545+ inferedOutbatchValue = inferedOutbatchValue * curBatchValue;
546+ }
547+ return inferedOutbatchValue;
548+}
549+ 
550+static aclnnStatus CheckShape(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
551+{
552+ CHECK_COND(CheckInputOutDims(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "Check CheckInputOutDims failed.");
553+ MatmulShapeInfo shapeInfo = GetMatmulShapeInfo(params);
554+ CHECK_COND(CheckShapeInfoMatch(params, shapeInfo) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
555+ "CheckShapeInfoMatch failed.");
556+ 
557+ if (!CheckMKN(shapeInfo.mDim, shapeInfo.kDim, shapeInfo.nDim)) {
558+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "CheckMKN failed.");
559+ return ACLNN_ERR_PARAM_INVALID;
560+ }
561+ CHECK_COND(CheckMxScaleLastDim(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "CheckMxScaleLastDim failed.");
562+ 
563+ if (params.bias != nullptr) {
564+ auto biasDimNum = params.bias->GetViewShape().GetDimNum();
565+ auto outDimNum = params.y->GetViewShape().GetDimNum();
566+ auto nDim = shapeInfo.nDim;
567+ if (biasDimNum != 1 && biasDimNum != 3) {
568+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "bias",
569+ FormatString("%zuD", biasDimNum).c_str(),
570+ "the shape dim of bias must be 1 or 3");
571+ return ACLNN_ERR_PARAM_INVALID;
572+ }
573+ if (biasDimNum == 1) {
574+ CHECK_COND(params.bias->GetViewShape().GetDim(0) == nDim, ACLNN_ERR_PARAM_INVALID,
575+ "bias dim should be equal to N dim %ld, but is %ld", nDim,
576+ params.bias->GetViewShape().GetDim(0));
577+ } else {
578+ if (outDimNum == 2 || outDimNum == 4 || outDimNum == 5 || outDimNum == 6) {
579+ OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(
580+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "bias",
581+ FormatString("%zuD", biasDimNum).c_str(),
582+ FormatString("when out dim-num is %zu, bias only support 1D, but is 3D", outDimNum).c_str());
583+ return ACLNN_ERR_PARAM_INVALID;
584+ }
585+ CHECK_COND(params.bias->GetViewShape().GetDim(1) == 1, ACLNN_ERR_PARAM_INVALID,
586+ "bias 2nd dim should be 1, but is %ld", params.bias->GetViewShape().GetDim(1));
587+ CHECK_COND(params.bias->GetViewShape().GetDim(2) == nDim, ACLNN_ERR_PARAM_INVALID,
588+ "bias 3rd dim should be equal to N dim %ld, but is %ld", nDim,
589+ params.bias->GetViewShape().GetDim(2));
590+ int64_t inferedOutbatchValue = InferOutputShape(params);
591+ if (inferedOutbatchValue == OUTPUT_INFER_FAIL) {
592+ return ACLNN_ERR_PARAM_INVALID;
593+ }
594+ CHECK_COND(params.bias->GetViewShape().GetDim(0) == inferedOutbatchValue, ACLNN_ERR_PARAM_INVALID,
595+ "bias 1st dim should be batch, but is %ld", params.bias->GetViewShape().GetDim(0));
596+ }
597+ }
598+ 
599+ return CheckExpectedShapes(params, shapeInfo);
600+}
wangzitao
wangzitaowangzitao7月9日

【🔴 必须修改】CheckExpectedShapes 被注释掉,缺少形状校验

CheckExpectedShapes 函数已实现但被注释掉,导致 x1/x2/x1Scale/x2Scale 的形状匹配检查被完全跳过。非法形状的输入将不会被拦截,可能导致后续计算错误或崩溃。

建议:取消注释 return CheckExpectedShapes(params, shapeInfo); 或说明跳过校验的原因。

likedislike
十方君主
十方君主
7月21日 评论:
601+ 
602+static aclnnStatus CheckParams(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
603+{
604+ OP_LOGD("QuantMatmulActivationQuant check params.");
605+ CHECK_RET(CheckNotNull(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
606+ CHECK_RET(CheckDtype(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
607+ CHECK_RET(CheckShape(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
608+ CHECK_RET(CheckFormat(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
609+ CHECK_RET(CheckOptioanlAlg(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
ddssz
ddsszddssz7月20日

[正确性][必须修复] CheckParams 到这里已经结束,但可选 bias 从未做过 dtype、format 或 shape 校验。非空的非法 bias 会直接进入 L0/kernel,而 kernel 将其固定按 float 读取,错误类型或不满足 N/batch 约束的形状可能导致错误结果或越界。请按接口支持范围补齐 bias 的 FLOAT32、ND 以及 (N,)/受支持 batch 形状校验。

likedislike
十方君主
十方君主
7月21日 评论:
610+ OP_LOGD("QuantMatmulActivationQuant check params success.");
611+ 
612+ return ACLNN_SUCCESS;
613+}
614+static aclnnStatus PreProcessOriginalShape(const aclTensor* x1, const aclTensor* x1Scale, const aclTensor* x2Scale)
615+{
616+ // original shape must be set before contiguous
617+ if (x1 != nullptr) {
618+ x1->SetOriginalShape(x1->GetViewShape());
619+ OP_LOGD("x1 original shape set to view shape.");
620+ }
621+ 
622+ if (x1Scale != nullptr) {
623+ x1Scale->SetOriginalShape(x1Scale->GetViewShape());
624+ OP_LOGD("x1Scale original shape set to view shape.");
625+ }
626+ 
627+ if (x2Scale != nullptr) {
628+ x2Scale->SetOriginalShape(x2Scale->GetViewShape());
629+ OP_LOGD("x2Scale original shape set to view shape.");
630+ }
631+ 
632+ return ACLNN_SUCCESS;
633+}
634+ 
635+static inline bool MxScaleContiguousProcess(const aclTensor*& mxScaleTensor, bool transpose, aclOpExecutor* executor)
636+{
637+ if (mxScaleTensor == nullptr || mxScaleTensor->GetViewShape().GetDimNum() < MX_SCALE_MAX_DIM) {
638+ OP_LOGD("MX scale no need to do contiguous process.");
639+ return true;
640+ }
641+ auto transposeFlag = false;
642+ int64_t dimNum = mxScaleTensor->GetViewShape().GetDimNum();
643+ int64_t lastDim = mxScaleTensor->GetViewShape().GetDim(dimNum - 1);
644+ int64_t lastSecondDim = mxScaleTensor->GetViewShape().GetDim(dimNum -
645+ QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM);
646+ // 3: 倒数第3维
647+ int64_t lastThirdDim = mxScaleTensor->GetViewShape().GetDim(dimNum - 3);
648+ if (mxScaleTensor->GetViewStrides()[dimNum - 3] == lastDim &&
649+ mxScaleTensor->GetViewStrides()[dimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM] ==
650+ lastDim * lastThirdDim) {
651+ int64_t tmpNxD = lastDim * lastSecondDim * lastThirdDim;
652+ transposeFlag = true;
653+ // 4: batch维度从倒数第4维起
654+ for (int64_t batchDim = dimNum - 4; batchDim >= 0; batchDim--) {
655+ if (mxScaleTensor->GetViewStrides()[batchDim] != tmpNxD) {
656+ transposeFlag = false;
657+ break;
658+ }
659+ tmpNxD *= mxScaleTensor->GetViewShape().GetDim(batchDim);
660+ }
661+ if (transpose) {
662+ if (lastSecondDim == 1 && lastThirdDim == 1) {
663+ transposeFlag = false;
664+ }
665+ } else {
666+ if (lastSecondDim == 1 || lastThirdDim == 1) {
667+ transposeFlag = false;
668+ }
669+ }
670+ }
671+ 
672+ if (transposeFlag) {
673+ op::Shape swapedShape = mxScaleTensor->GetViewShape();
674+ swapedShape.SetDim(dimNum - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM, lastThirdDim);
675+ // 3: 倒数第3维
676+ swapedShape.SetDim(dimNum - 3, lastSecondDim);
677+ mxScaleTensor = executor->CreateView(mxScaleTensor, swapedShape, mxScaleTensor->GetViewOffset());
678+ } else {
679+ mxScaleTensor = l0op::Contiguous(mxScaleTensor, executor);
680+ }
681+ CHECK_RET(mxScaleTensor != nullptr, false);
682+ return true;
683+}
684+ 
685+static bool CheckSpecialCase(const aclTensor* tensor, int64_t firstLastDim, int64_t secondLastDim)
686+{
687+ if ((tensor->GetViewShape().GetDim(firstLastDim) == tensor->GetViewShape().GetDim(secondLastDim)) &&
688+ (tensor->GetViewShape().GetDim(secondLastDim) == 1)) {
689+ OP_LOGD("QuantMatmulActivationQuant special case, no need to set transpose attr value.");
690+ return true;
691+ }
692+ return false;
693+}
694+ 
695+static bool GetTransposeAttrValue(const aclTensor* tensor, bool transpose, bool checkSpecialCase = true)
696+{
697+ int64_t dim1 = tensor->GetViewShape().GetDimNum() - 1;
698+ int64_t dim2 = tensor->GetViewShape().GetDimNum() - QuantMatmulActivationQuantAclnnCheck::PENULTIMATE_DIM;
699+ // check if tensor is contiguous layout
700+ if (tensor->GetViewStrides()[dim2] == 1 &&
701+ (tensor->GetViewStrides()[dim1] == tensor->GetViewShape().GetDim(dim2))) {
702+ OP_LOGD("QuantMatmulActivationQuant GetTransposeAttrValue, find tensor is not contiguous.");
703+ const_cast<aclTensor*>(tensor)->SetViewShape(SwapLastTwoDimValue(tensor->GetViewShape()));
704+ // 如果不需要校验特殊case,则直接返回
705+ if (!checkSpecialCase) {
706+ return !transpose;
707+ }
708+ if (!CheckSpecialCase(tensor, dim1, dim2)) {
709+ return !transpose;
710+ }
711+ }
712+ return transpose;
713+}
714+ 
715+static void GetTranspose(const QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params, bool& transposeX1,
716+ bool& transposeX2)
717+{
718+ transposeX1 = GetTransposeAttrValue(params.x1, transposeX1, true);
719+ transposeX2 = GetTransposeAttrValue(params.x2, transposeX2, true);
720+ OP_LOGD("QuantMatmulActivationQuant attr transposeX1 is %d, transposeX2 is %d.", transposeX1, transposeX2);
721+}
722+ 
723+static bool CheckGroupSize(QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params)
724+{
725+ auto groupSize = params.groupSize;
726+ if (groupSize < 0) {
727+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "groupSize",
728+ std::to_string(groupSize).c_str(), "groupSize can not be negative");
729+ return false;
730+ }
731+ uint64_t groupSizeM = (static_cast<uint64_t>(groupSize) >> GROUP_M_OFFSET) & GROUP_MNK_BIT_SIZE;
732+ uint64_t groupSizeN = (static_cast<uint64_t>(groupSize) >> GROUP_N_OFFSET) & GROUP_MNK_BIT_SIZE;
733+ uint64_t groupSizeK = static_cast<uint64_t>(groupSize) & GROUP_MNK_BIT_SIZE;
734+ 
735+ if (groupSizeK == 0 && groupSizeM == 0 && groupSizeN == 0) {
736+ params.groupSize = (1UL << GROUP_M_OFFSET) | (1UL << GROUP_N_OFFSET) |
737+ static_cast<uint64_t>(PERGROUP_GROUP_SIZE);
738+ } else if (groupSizeK != static_cast<uint64_t>(PERGROUP_GROUP_SIZE) || groupSizeM != 1UL || groupSizeN != 1UL) {
739+ OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(
740+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "groupSize, groupSizeM, groupSizeN, groupSizeK",
741+ FormatString("%ld, %lu, %lu, %lu", groupSize, groupSizeM, groupSizeN, groupSizeK).c_str(),
742+ "when the quantization mode is mx, groupSize must be 4295032864 and Torch API group_sizes must be [1, "
743+ "1, 32]");
744+ return false;
745+ }
746+ 
747+ OP_LOGD("QuantMatmulActivationQuant check group_size success.");
748+ return true;
749+}
750+ 
751+static aclnnStatus aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSizeCommon(
752+ QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams& params, aclOpExecutor* executor)
753+{
754+ params.x2Scale = QuantMatmulActivationQuantAclnnCheck::SetTensorToNDFormat(params.x2Scale);
755+ params.x1Scale = QuantMatmulActivationQuantAclnnCheck::SetTensorToNDFormat(params.x1Scale);
756+ 
757+ if (params.bias != nullptr) {
758+ params.bias = QuantMatmulActivationQuantAclnnCheck::SetTensorToNDFormat(params.bias);
759+ }
760+ 
761+ auto reformatedX1 = QuantMatmulActivationQuantAclnnCheck::SetTensorToNDFormat(params.x1);
762+ params.x1 = reformatedX1;
763+ QuantMatmulActivationQuantAclnnCheck::TensorContiguousProcess(params.x1, params.transposeX1, executor);
764+ MxScaleContiguousProcess(params.x1Scale, params.transposeX1, executor);
765+ MxScaleContiguousProcess(params.x2Scale, params.transposeX2, executor);
766+ 
767+ // 设置x2的OriginalShape为它的ViewShape
768+ auto retNZProcess = QuantMatmulActivationQuantAclnnCheck::WeightNZCaseProcess(params.x2, params.transposeX2,
769+ executor);
770+ CHECK_RET(retNZProcess == ACLNN_SUCCESS, retNZProcess);
771+ 
772+ GetTranspose(params, params.transposeX1, params.transposeX2);
773+ 
774+ CHECK_COND(CheckGroupSize(params), ACLNN_ERR_PARAM_INVALID, "CheckGroupSize failed.");
775+ 
776+ // 固定写法,参数检查
777+ auto ret = CheckParams(params);
778+ CHECK_RET(ret == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
779+ // Invoke l0 operator QuantMatmulActivationQuant for calculation.
780+ auto quantMatmulActivationQuantResults = l0op::QuantMatmulActivationQuant(
781+ params.x1, params.x2, params.bias, params.x1Scale, params.x2Scale, params.transposeX1, params.transposeX2,
782+ params.groupSize, params.activationType, params.y_dtype, params.quantMode, params.roundMode, params.scaleAlg,
783+ params.dstTypeMax, executor);
784+ 
785+ auto yComputeOut = std::get<IDX_0>(quantMatmulActivationQuantResults);
786+ auto yScaleComputeOut = std::get<IDX_1>(quantMatmulActivationQuantResults);
787+ 
788+ // 校验输出不为空
789+ CHECK_RET(yComputeOut != nullptr, ACLNN_ERR_INNER_NULLPTR);
790+ CHECK_RET(yScaleComputeOut != nullptr, ACLNN_ERR_INNER_NULLPTR);
791+ 
792+ // 将结果拷贝到输出tensor
793+ auto viewCopyYResult = l0op::ViewCopy(yComputeOut, params.y, executor);
794+ CHECK_RET(viewCopyYResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
795+ 
796+ auto viewCopyYScaleResult = l0op::ViewCopy(yScaleComputeOut, params.yScale, executor);
797+ CHECK_RET(viewCopyYScaleResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
798+ 
799+ return ACLNN_SUCCESS;
800+}
801+ 
802+} // namespace
803+ 
804+#ifdef __cplusplus
805+extern "C" {
806+#endif
807+aclnnStatus aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize(
808+ const aclTensor* x1, const aclTensor* x2, const aclTensor* x1ScaleOptional, const aclTensor* x2Scale,
809+ const aclTensor* biasOptional, bool transposeX1, bool transposeX2, int64_t groupSize, const char* activationType,
810+ const char* quantMode, const char* roundMode, int64_t scaleAlg, double dstTypeMax, aclTensor* y, aclTensor* yScale,
811+ uint64_t* workspaceSize, aclOpExecutor** executor)
812+{
813+ L2_DFX_PHASE_1(aclnnQuantMatmulActivationQuantWeightNz,
814+ DFX_IN(x1, x2, x1ScaleOptional, x2Scale, biasOptional, transposeX1, transposeX2, groupSize,
815+ activationType, quantMode, roundMode, scaleAlg, dstTypeMax),
816+ DFX_OUT(y, yScale));
817+ 
818+ auto ret = CheckWeightNzParamsDAV3510(x1, x2);
819+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
820+ 
821+ auto y_dtype = x1->GetDataType();
822+ QBMMActivationQuant::QuantMatmulActivationQuantWeightNzParams params{
823+ x1, x2, x1ScaleOptional, x2Scale, biasOptional, y, yScale, transposeX1,
824+ transposeX2, groupSize, activationType, y_dtype, quantMode, roundMode, scaleAlg, dstTypeMax};
825+ 
826+ // 空tensor 处理
827+ if (params.x1->IsEmpty() || params.x2->IsEmpty() || (params.x1Scale != nullptr && params.x1Scale->IsEmpty()) ||
828+ (params.x2Scale != nullptr && params.x2Scale->IsEmpty()) ||
829+ (params.bias != nullptr && params.bias->IsEmpty())) {
830+ OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
831+ "aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize", "x1, x2, x1Scale, x2Scale, bias",
832+ Ops::NN::FormatString(
833+ "%s, %s, %s, %s, %s", op::ToString(x1->GetViewShape()).GetString(),
834+ op::ToString(x2->GetViewShape()).GetString(), op::ToString(x1ScaleOptional->GetViewShape()).GetString(),
835+ op::ToString(x2Scale->GetViewShape()).GetString(),
836+ params.bias != nullptr ? op::ToString(params.bias->GetViewShape()).GetString() : "null")
837+ .c_str(),
838+ Ops::NN::FormatString("The shapes of %s cannot be %s", "x1, x2, x1Scale, x2Scale, bias", "empty").c_str());
839+ return ACLNN_ERR_PARAM_INVALID;
840+ }
841+ 
842+ // Step 1: 设置original_shape(必须在Contiguous之前)
843+ ret = PreProcessOriginalShape(params.x1, params.x1Scale, params.x2Scale);
844+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
845+ 
846+ CHECK_COND(CheckInputOutDims(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "Check CheckInputOutDims failed.");
847+ 
848+ CHECK_RET(x2 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
849+ params.transposeX2 = GetTransposeAttrValue(x2, transposeX2, false);
ddssz
ddsszddssz7月9日

CheckParams/CheckNotNull 在 common 流程后面才执行,但这里已经调用 GetTransposeAttrValue(x2, ...),随后还会访问 x2->GetStorageShape()。当 x2 为空时会先解引用崩溃,无法返回 ACLNN_ERR_PARAM_NULLPTR。建议在任何 x2 使用前执行必选输入判空,或把这段移动到参数校验之后。

likedislike
十方君主
十方君主
7月21日 评论:
850+ 
851+ op::Shape weightNzShape = QuantMatmulActivationQuantAclnnCheck::GetWeightNzShape(x2, transposeX2);
852+ if (!QuantMatmulActivationQuantAclnnCheck::CheckWeightNzStorageShape(weightNzShape, x2->GetStorageShape())) {
853+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
854+ "x2'format only support NZ, but now x2's format is not NZ(Ascend affinity format). \
855+ aclnnCalculateMatmulWeightSizeV2 and aclnnTransMatmulWeight can be used to convert the input format from ND to Ascend \
856+ affinity format.");
857+ return ACLNN_ERR_PARAM_INVALID;
858+ }
859+ 
860+ // 固定写法,创建OpExecutor
861+ auto uniqueExecutor = CREATE_EXECUTOR();
862+ auto executorPtr = uniqueExecutor.get();
863+ CHECK_RET(executorPtr != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
864+ x2 = QuantMatmulActivationQuantAclnnCheck::SetTensorToNZFormat(x2, weightNzShape, executorPtr);
ddssz
ddsszddssz7月20日

[健壮性][必须修复] SetTensorToNZFormat 已经会在 CreateView 失败时返回 nullptr,但调用方没有检查就写入 params.x2。随后 common 流程会在 WeightNZCaseProcess 中直接访问 x2->GetStorageFormat(),把一次可返回的资源分配失败变成空指针崩溃。请在这里立即校验返回值并返回合适的 ACLNN_ERR_INNER_NULLPTR

likedislike
十方君主
十方君主
7月21日 评论:
865+ CHECK_RET(x2 != nullptr, ACLNN_ERR_PARAM_NULLPTR);
866+ params.x2 = x2;
867+ 
868+ ret = aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSizeCommon(params, executorPtr);
869+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
870+ 
871+ // Standard syntax, get the size of workspace needed during computation.
872+ CHECK_RET(workspaceSize != nullptr, ACLNN_ERR_PARAM_NULLPTR);
873+ CHECK_RET(executor != nullptr, ACLNN_ERR_PARAM_NULLPTR);
874+ *workspaceSize = uniqueExecutor->GetWorkspaceSize();
ddssz
ddsszddssz7月9日

【空指针】workspaceSizeexecutor 是 ACLNN API 的输出参数,当前入口没有校验,最后直接写 *workspaceSizeReleaseTo(executor)。用户传空指针时会直接崩溃;需要按 ACLNN 规范在前置参数检查阶段返回 ACLNN_ERR_PARAM_NULLPTR

likedislike
十方君主
十方君主
7月21日 评论:
875+ uniqueExecutor.ReleaseTo(executor);
876+ 
877+ return ACLNN_SUCCESS;
878+}
879+ 
880+aclnnStatus aclnnQuantMatmulActivationQuantWeightNz(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
881+ aclrtStream stream)
882+{
883+ L2_DFX_PHASE_2(aclnnQuantMatmulActivationQuantWeightNz);
884+ CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER,
885+ "This is an error in QuantMatmulActivationQuantWeightNz launch aicore.");
886+ return ACLNN_SUCCESS;
887+}
888+ 
889+#ifdef __cplusplus
890+}
891+#endif
@@ -0,0 +1,53 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef OP_API_INC_QUANT_MATMUL_ACTIVATION_QUANT_NZ_H
12+#define OP_API_INC_QUANT_MATMUL_ACTIVATION_QUANT_NZ_H
13+ 
14+#include "aclnn/aclnn_base.h"
15+#include "aclnn_util.h"
16+ 
17+#ifdef __cplusplus
18+extern "C" {
19+#endif
20+ 
21+/**
22+ * @brief aclnnQuantMatmulActivationQuantWeightNz的第一段接口,根据具体的计算流程,计算workspace大小。
23+ * @domain aclnn_ops_infer
24+ * 算子功能:实现QuantMatmulActivationQuantWeightNz计算
25+ * @param [in] x1: matmul左矩阵,数据类型支持:FLOAT8_E4M3FN、FLOAT8_E5M2数据类型,数据格式支持ND。
26+ * @param [in] x2: matmul右矩阵,数据类型支持:FLOAT8_E4M3FN数据类型,数据格式支持NZ。
27+ * @param [out] workspaceSize: 返回用户需要在npu device侧申请的workspace大小。
28+ * @param [out] executor: 返回op执行器,包含算子计算流程。
29+ * @return aclnnStatus: 返回状态码。
30+ */
31+ACLNN_API aclnnStatus aclnnQuantMatmulActivationQuantWeightNzGetWorkspaceSize(
32+ const aclTensor* x1, const aclTensor* x2, const aclTensor* x1ScaleOptional, const aclTensor* x2Scale,
33+ const aclTensor* biasOptional, bool transposeX1, bool transposeX2, int64_t groupSize, const char* activationType,
34+ const char* quantMode, const char* roundMode, int64_t scaleAlg, double dstTypeMax, aclTensor* yOut,
35+ aclTensor* yScaleOut, uint64_t* workspaceSize, aclOpExecutor** executor);
36+ 
37+/**
38+ * @brief aclnnQuantMatmulActivatiionQuant的第二段接口,用于执行计算。
39+ * @param [in] workspace: 在npu device侧申请的workspace内存起址。
40+ * @param [in] workspace_size: 在npu
41+ * device侧申请的workspace大小,由第一段接口aclnnQuantMatmulActivatiionQuantGetWorkspaceSize获取。
42+ * @param [in] exector: op执行器,包含了算子计算流程。
43+ * @param [in] stream: acl stream流。
44+ * @return aclnnStatus: 返回状态码
45+ */
46+ACLNN_API aclnnStatus aclnnQuantMatmulActivationQuantWeightNz(void* workspace, uint64_t workspaceSize,
47+ aclOpExecutor* executor, aclrtStream stream);
48+ 
49+#ifdef __cplusplus
50+}
51+#endif
52+ 
53+#endif // QuantMatmulActivationQuantWeightNz
@@ -0,0 +1,90 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "aclnn_kernels/common/op_error_check.h"
12+#include "opdev/make_op_executor.h"
13+#include "opdev/op_def.h"
14+#include "opdev/op_dfx.h"
15+#include "opdev/op_executor.h"
16+#include "opdev/op_log.h"
17+#include "opdev/shape_utils.h"
18+#include "quant_matmul_activation_quant.h"
19+#include "util/math_util.h"
20+#include "quant_matmul_activation_quant_util.h"
21+ 
22+using namespace op;
23+using namespace QBMMActivationQuant;
24+ 
25+namespace l0op {
26+ 
27+OP_TYPE_REGISTER(QuantMatmulActivationQuant);
28+constexpr int64_t DIM_TWO = 2L;
29+constexpr int64_t BLOCKSIZE = 32L;
30+ 
31+const std::array<aclTensor*, QUANT_MATMUL_ACTIVATION_QUANT_OUT_NUM> QuantMatmulActivationQuant(
32+ const aclTensor* x1, const aclTensor* x2, const aclTensor* bias, const aclTensor* x1Scale, const aclTensor* x2Scale,
33+ bool transposeX1, bool transposeX2, int64_t groupSize, const char* activationType, int64_t y_dtype,
34+ const char* quantMode, const char* roundMode, int64_t scaleAlg, double dstTypeMax, aclOpExecutor* executor)
35+{
36+ L0_DFX(QuantMatmulActivationQuant, x1, x2, bias, x1Scale, x2Scale, transposeX1, transposeX2, groupSize,
37+ activationType, y_dtype, quantMode, roundMode, scaleAlg, dstTypeMax);
38+ 
39+ Format format = Format::FORMAT_ND;
40+ op::Shape x1Shape = x1->GetViewShape();
41+ op::Shape x2Shape = x2->GetViewShape();
42+ op::Shape yOutShape = x1Shape;
43+ auto x1DimNum = x1Shape.GetDimNum();
44+ auto x2DimNum = x2Shape.GetDimNum();
45+ 
46+ auto yOutDim0 = transposeX1 ? x1->GetViewShape().GetDim(x1DimNum - LAST_FIRST_DIM_INDEX) :
47+ x1->GetViewShape().GetDim(x1DimNum - LAST_SECOND_DIM_INDEX);
48+ auto yOutDim1 = transposeX2 ? x2->GetViewShape().GetDim(x2DimNum - LAST_SECOND_DIM_INDEX) :
49+ x2->GetViewShape().GetDim(x2DimNum - LAST_FIRST_DIM_INDEX);
50+ yOutShape.SetDim(x1DimNum - LAST_SECOND_DIM_INDEX, yOutDim0);
51+ yOutShape.SetDim(x1DimNum - LAST_FIRST_DIM_INDEX, yOutDim1);
52+ 
53+ op::Shape yScaleOutShape = yOutShape;
54+ auto yScaleOutDim1 = (Ops::Base::CeilDiv(yOutDim1, BLOCKSIZE) + MXFP_MULTI_BASE_SIZE - 1) / MXFP_MULTI_BASE_SIZE;
55+ yScaleOutShape.SetDim(x1DimNum - LAST_FIRST_DIM_INDEX, yScaleOutDim1);
56+ yScaleOutShape.AppendDim(DIM_TWO);
wangzitao
wangzitaowangzitao7月9日

【🟡 建议修改】使用字面量 2 代替常量,可读性差

(Ops::Base::CeilDiv(yOutDim1, BLOCKSIZE) + 2 - 1) / 2+ 2 - 1 等价于 + 1,且使用字面量 2 不如使用 MXFP_MULTI_BASE_SIZE 常量清晰。

建议:改写为 Ops::Base::CeilDiv(Ops::Base::CeilDiv(yOutDim1, BLOCKSIZE), MXFP_MULTI_BASE_SIZE)

likedislike
十方君主
十方君主
7月21日 评论:
57+ 
58+ auto yOut = executor->AllocTensor(yOutShape, x1->GetDataType(), format);
ddssz
ddsszddssz7月9日

AllocTensor 的返回值没有检查,yOut/yScaleOut 可能为 nullptr,但后面马上传给 INFER_SHAPE 和 L0 算子。建议分配后立即 CHECK_RET,否则内存分配失败会变成更深层的空指针问题。

likedislike
十方君主
十方君主
7月21日 评论:
59+ auto yScaleOut = executor->AllocTensor(yScaleOutShape, x1Scale->GetDataType(), format);
60+ 
61+ if (yOut == nullptr || yScaleOut == nullptr) {
62+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "alloc tensor failed.");
63+ return {yOut, yScaleOut};
64+ }
65+ 
66+ // 运行QuantMatmulActivationQuant算子的InferShape函数,推导输出shape
67+ auto ret = INFER_SHAPE(QuantMatmulActivationQuant, OP_INPUT(x1, x2, bias, x1Scale, x2Scale),
68+ OP_OUTPUT(yOut, yScaleOut),
69+ OP_ATTR(transposeX1, transposeX2, groupSize, activationType, y_dtype, quantMode, roundMode,
70+ scaleAlg, dstTypeMax));
71+ 
72+ if (ret != ACLNN_SUCCESS) {
73+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed.");
74+ return {nullptr, nullptr};
75+ }
76+ 
77+ OP_LOGD("l0 transposeX1 = %s, transposeX2 = %s", transposeX1 ? "true" : "false", transposeX2 ? "true" : "false");
78+ 
79+ ret = ADD_TO_LAUNCHER_LIST_AICORE(QuantMatmulActivationQuant, OP_INPUT(x1, x2, bias, x1Scale, x2Scale),
80+ OP_OUTPUT(yOut, yScaleOut),
81+ OP_ATTR(transposeX1, transposeX2, groupSize, activationType, y_dtype, quantMode,
82+ roundMode, scaleAlg, dstTypeMax));
83+ if (ret != ACLNN_SUCCESS) {
84+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE failed.");
85+ return {nullptr, nullptr};
86+ }
87+ return {yOut, yScaleOut};
88+}
89+ 
90+} // namespace l0op
@@ -0,0 +1,26 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef OP_HOST_OP_API_QUANT_MATMUL_ACTIVATION_QUANT_H
12+#define OP_HOST_OP_API_QUANT_MATMUL_ACTIVATION_QUANT_H
13+ 
14+#include "opdev/op_executor.h"
15+#include "opdev/make_op_executor.h"
16+ 
17+namespace l0op {
18+constexpr size_t QUANT_MATMUL_ACTIVATION_QUANT_OUT_NUM = 2; // output y and yScale
19+ 
20+const std::array<aclTensor*, QUANT_MATMUL_ACTIVATION_QUANT_OUT_NUM> QuantMatmulActivationQuant(
21+ const aclTensor* x1, const aclTensor* x2, const aclTensor* bias, const aclTensor* x1Scale, const aclTensor* x2Scale,
22+ bool transposeX1, bool transposeX2, int64_t groupSize, const char* activationType, int64_t y_dtype,
23+ const char* quantMode, const char* roundMode, int64_t scaleAlg, double dstTypeMax, aclOpExecutor* executor);
24+} // namespace l0op
25+ 
26+#endif // OP_HOST_OP_API_QUANT_MATMUL_ACTIVATION_QUANT_H
@@ -0,0 +1,259 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file quant_matmul_activation_quant_check.cpp
13+ * \brief
14+ */
15+#include "quant_matmul_activation_quant_check.h"
16+ 
17+namespace QuantMatmulActivationQuantAclnnCheck {
18+ 
19+using namespace op;
20+using namespace ge;
21+using Ops::Base::CeilDiv;
22+using Ops::NN::IsTransposeLastTwoDims;
23+using Ops::NN::SwapLastTwoDimValue;
24+ 
25+bool CheckSpecialCase(const aclTensor* tensor, int64_t firstLastDim, int64_t secondLastDim)
26+{
27+ if (tensor->GetViewShape().GetDim(firstLastDim) == tensor->GetViewShape().GetDim(secondLastDim)) {
28+ OP_LOGD("QuantMatmul special case, no need to set transpose attr value.");
29+ return true;
30+ }
31+ return false;
32+}
33+ 
34+bool GetTransposeAttrValue(const aclTensor* tensor, bool transpose)
35+{
36+ int64_t dim1 = tensor->GetViewShape().GetDimNum() - 1;
37+ int64_t dim2 = tensor->GetViewShape().GetDimNum() - PENULTIMATE_DIM;
38+ // check if tensor is contiguous layout
39+ if (tensor->GetViewStrides()[dim2] == 1 && tensor->GetViewStrides()[dim1] == tensor->GetViewShape().GetDim(dim2)) {
40+ OP_LOGD("QuantMatmul GetTransposeAttrValue, find tensor is not contiguous.");
41+ const_cast<aclTensor*>(tensor)->SetViewShape(SwapLastTwoDimValue(tensor->GetViewShape()));
42+ if (!CheckSpecialCase(tensor, dim1, dim2)) {
43+ return !transpose;
44+ }
45+ }
46+ return transpose;
47+}
48+ 
49+op::Shape GetWeightNzShape(const aclTensor* input, bool transpose)
50+{
51+ size_t viewDimNum = input->GetViewShape().GetDimNum();
52+ int64_t k = transpose ? input->GetViewShape().GetDim(viewDimNum - 1) :
53+ input->GetViewShape().GetDim(viewDimNum - LAST_SECOND_DIM_INDEX);
54+ int64_t n = transpose ? input->GetViewShape().GetDim(viewDimNum - LAST_SECOND_DIM_INDEX) :
55+ input->GetViewShape().GetDim(viewDimNum - 1);
56+ 
57+ int64_t nz_k0_value_trans = NZ_K0_VALUE_INT8_TRANS;
58+ int64_t k1 = transpose ? CeilDiv(k, nz_k0_value_trans) : CeilDiv(k, NZ_K0_VALUE_BMM_BLOCK_NUM);
59+ int64_t n1 = transpose ? CeilDiv(n, NZ_K0_VALUE_BMM_BLOCK_NUM) : CeilDiv(n, nz_k0_value_trans);
60+ 
61+ op::Shape weightNzShape;
62+ for (size_t i = 0; i < viewDimNum - LAST_SECOND_DIM_INDEX; i++) {
63+ weightNzShape.AppendDim(input->GetViewShape().GetDim(i));
64+ }
65+ if (transpose) {
66+ weightNzShape.AppendDim(k1);
67+ weightNzShape.AppendDim(n1);
68+ } else {
69+ weightNzShape.AppendDim(n1);
70+ weightNzShape.AppendDim(k1);
71+ }
72+ weightNzShape.AppendDim(NZ_STORAGE_PENULTIMATE_DIM);
73+ weightNzShape.AppendDim(nz_k0_value_trans);
74+ return weightNzShape;
75+}
76+ 
77+bool CheckWeightNzStorageShape(const op::Shape& nzShape, const op::Shape& storageShape)
78+{
79+ uint64_t nzDimMultiply = 1;
80+ uint64_t nzDimNum = nzShape.GetDimNum();
81+ for (uint64_t i = 0; i < nzDimNum; i++) {
82+ nzDimMultiply *= nzShape[i];
83+ }
84+ 
85+ uint64_t storageDimMultiply = 1;
86+ uint64_t storageDimNum = storageShape.GetDimNum();
87+ for (uint64_t i = 0; i < storageDimNum; i++) {
88+ storageDimMultiply *= storageShape[i];
89+ }
90+ 
91+ return nzDimMultiply == storageDimMultiply;
92+}
93+ 
94+const aclTensor* SetTensorToNZFormat(const aclTensor* input, op::Shape& shape, aclOpExecutor* executor)
95+{
96+ if (executor == nullptr) {
97+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "QuantMatmul SetTensorToNZFormat, executor is null");
98+ return nullptr;
99+ }
100+ auto formatTensor = executor->CreateView(input, shape, input->GetViewOffset());
ddssz
ddsszddssz7月9日

【空指针】executorCreateView 返回值都没有保护,下一行直接 formatTensor->SetStorageFormatCreateView 失败时会空指针崩溃;建议在函数入口检查 executor,并对 formatTensorCHECK_RET(..., ACLNN_ERR_INNER_NULLPTR)

likedislike
十方君主
十方君主
7月21日 评论:
101+ if (formatTensor == nullptr) {
102+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "QuantMatmul SetTensorToNZFormat, formatTensor is null");
103+ return nullptr;
104+ }
105+ formatTensor->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ);
106+ formatTensor->SetOriginalFormat(op::Format::FORMAT_ND);
107+ formatTensor->SetViewShape(input->GetViewShape());
108+ return formatTensor;
109+}
110+ 
111+bool TensorContiguousProcess(const aclTensor*& contiguousTensor, bool& transpose, aclOpExecutor* executor)
112+{
113+ if (contiguousTensor == nullptr) {
114+ OP_LOGD("QuantMatmul no need to do contiguous process.");
115+ return true;
116+ }
117+ bool isNZTensor = static_cast<ge::Format>(ge::GetPrimaryFormat(contiguousTensor->GetStorageFormat())) ==
118+ op::Format::FORMAT_FRACTAL_NZ;
119+ auto storageShape = contiguousTensor->GetStorageShape();
120+ auto transposeFlag = IsTransposeLastTwoDims(contiguousTensor);
121+ // swap tensor if its viewshape not satisfy request shape without adding a transpose node
122+ if (transposeFlag) {
123+ contiguousTensor = executor->CreateView(contiguousTensor, SwapLastTwoDimValue(contiguousTensor->GetViewShape()),
124+ contiguousTensor->GetViewOffset());
125+ transpose = !transpose;
126+ } else {
127+ contiguousTensor = l0op::Contiguous(contiguousTensor, executor);
128+ }
129+ CHECK_RET(contiguousTensor != nullptr, false);
130+ if (isNZTensor) {
131+ contiguousTensor->SetStorageShape(storageShape); // 对NZ的场景需要用原NZshape刷新
132+ contiguousTensor->SetOriginalShape(storageShape);
133+ }
134+ return true;
135+}
136+ 
137+aclnnStatus WeightNZCaseProcess(const aclTensor*& x2, bool& transposeX2, aclOpExecutor* executor)
138+{
139+ // if weight is already in nz format, no need to set contiguous
140+ if (ge::GetPrimaryFormat(x2->GetStorageFormat()) == op::Format::FORMAT_FRACTAL_NZ ||
141+ ge::GetPrimaryFormat(x2->GetStorageFormat()) == op::Format::FORMAT_FRACTAL_NZ_C0_32) {
142+ x2->SetOriginalShape(x2->GetViewShape());
143+ if (ge::GetPrimaryFormat(x2->GetStorageFormat()) == op::Format::FORMAT_FRACTAL_NZ_C0_32) {
144+ CHECK_RET(SetSpecilNZTensorToNormalNZFormat(x2, executor) == ACLNN_SUCCESS, ACLNN_ERR_INNER_NULLPTR);
145+ }
146+ } else {
147+ CHECK_RET(TensorContiguousProcess(x2, transposeX2, executor), ACLNN_ERR_INNER_NULLPTR);
148+ }
149+ return ACLNN_SUCCESS;
150+}
151+ 
152+aclTensor* ConvertTensorToInt4(const aclTensor* input, aclOpExecutor* executor)
153+{
154+ // 将int32的输入dtype修改为int4, 同时ViewShape和ViewStrides也从int32修改为int4所对应的。
155+ auto viewShape = input->GetViewShape();
156+ viewShape[viewShape.GetDimNum() - 1] = viewShape[viewShape.GetDimNum() - 1] * INT4_NUMS_IN_INT32;
157+ auto inputTemp = executor->CreateView(input, viewShape, input->GetViewOffset());
158+ inputTemp->SetDataType(DataType::DT_INT4);
159+ OP_LOGD("The conversion from int32 to int4 is completed.");
160+ return inputTemp;
161+}
162+ 
163+void InputPreProcessA4W4(const aclTensor*& x1, const aclTensor*& x2, aclOpExecutor* executor)
164+{
165+ if (x2->GetDataType() == DataType::DT_INT32) {
166+ x2 = ConvertTensorToInt4(x2, executor);
167+ }
168+ if (x1->GetDataType() == DataType::DT_INT32) {
169+ x1 = ConvertTensorToInt4(x1, executor);
170+ }
171+}
172+ 
173+aclnnStatus A4W4CaseProcess(const aclTensor*& x1, const aclTensor*& x2, aclOpExecutor* executor)
174+{
175+ InputPreProcessA4W4(x1, x2, executor);
176+ return ACLNN_SUCCESS;
177+}
178+ 
179+const aclTensor* SetTensorToNDFormat(const aclTensor* input)
180+{
181+ OP_LOGD("QuantMatmul set tensor to ND format.");
182+ const aclTensor* output = nullptr;
183+ if (input == nullptr) {
184+ return output;
185+ }
186+ if (input->GetStorageFormat() != Format::FORMAT_FRACTAL_NZ) {
187+ output = l0op::ReFormat(input, op::Format::FORMAT_ND);
188+ } else {
189+ output = input;
190+ }
191+ return output;
192+}
193+ 
194+const aclTensor* GetNDFormat(const aclTensor* input)
195+{
196+ const aclTensor* reformatedInput = input;
197+ if (input != nullptr) {
198+ reformatedInput = SetTensorToNDFormat(input);
199+ }
200+ return reformatedInput;
201+}
202+ 
203+void GetDtypeAndTranspose(TupleTensor mandatoryTensors, int64_t& dtype, bool& transposeX1, bool& transposeX2)
204+{
205+ auto x1 = std::get<0>(mandatoryTensors);
206+ auto x2 = std::get<1>(mandatoryTensors);
207+ auto out = std::get<INDEX_OUT_IN_TUPLE>(mandatoryTensors);
208+ dtype = static_cast<int64_t>(out->GetDataType());
209+ transposeX1 = GetTransposeAttrValue(x1, transposeX1);
210+ transposeX2 = GetTransposeAttrValue(x2, transposeX2);
211+ OP_LOGD("QuantMatmul attr transposeX1 is %d, transposeX2 is %d.", transposeX1, transposeX2);
212+}
213+ 
214+aclnnStatus SetSpecilNZTensorToNormalNZFormat(const aclTensor*& input, aclOpExecutor* executor)
215+{
216+ OP_LOGD("QuantMatmulV4 set special NZ format to normal NZ format.");
217+ auto nzTensorTmp = executor->CreateView(input, input->GetViewShape(), input->GetViewOffset());
218+ CHECK_RET(nzTensorTmp != nullptr, ACLNN_ERR_INNER_NULLPTR);
219+ nzTensorTmp->SetViewFormat(op::Format::FORMAT_ND);
220+ nzTensorTmp->SetOriginalFormat(op::Format::FORMAT_ND);
221+ nzTensorTmp->SetStorageFormat(op::Format::FORMAT_FRACTAL_NZ);
222+ nzTensorTmp->SetStorageShape(input->GetStorageShape());
223+ nzTensorTmp->SetOriginalShape(input->GetOriginalShape());
224+ input = nzTensorTmp;
225+ return ACLNN_SUCCESS;
226+}
227+ 
228+aclnnStatus SpecialOutputProcess(const aclTensor* x1, const aclTensor* x2, const aclTensor* out,
229+ const aclTensor*& matmulRet, aclOpExecutor* executor)
230+{
231+ // we have to reshape for case which x1 and x2 are 2 dims and out is 3 dims, otherwise, viewcopy will fail
232+ OP_LOGD("QuantMatmul enter SpecialOutputProcess func.");
233+ auto outShape = out->GetViewShape();
234+ auto outDimNum = outShape.GetDimNum();
235+ int64_t outMDim = outShape.GetDim(outDimNum - 2);
236+ auto x1DimNum = x1->GetViewShape().GetDimNum();
237+ auto x2DimNum = x2->GetViewShape().GetDimNum();
238+ // speical case : x1 and x2 are 2 dim, output is 3 dim, have to reshape matmul result, otherwise viewcopy will fail.
239+ if (x1DimNum == 2 && x2DimNum == 2 && outDimNum == 3 && outMDim == 1) {
240+ matmulRet = l0op::Reshape(matmulRet, outShape, executor);
241+ }
242+ CHECK_RET(matmulRet != nullptr, ACLNN_ERR_INNER_NULLPTR);
243+ return ACLNN_SUCCESS;
244+}
245+ 
246+aclnnStatus PostMatmulCalcProcess(const aclTensor* matmulRet, const aclTensor* x1, const aclTensor* x2,
247+ const aclTensor* out, aclOpExecutor* executor)
248+{
249+ CHECK_RET(matmulRet != nullptr, ACLNN_ERR_INNER_NULLPTR);
250+ CHECK_RET(SpecialOutputProcess(x1, x2, out, matmulRet, executor) == ACLNN_SUCCESS, ACLNN_ERR_INNER_NULLPTR);
251+ 
252+ // 如果出参out是非连续Tensor,需要把计算完的连续Tensor转非连续
253+ auto viewCopyResult = l0op::ViewCopy(matmulRet, out, executor);
254+ CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
255+ 
256+ return ACLNN_SUCCESS;
257+}
258+ 
259+} // namespace QuantMatmulActivationQuantAclnnCheck
@@ -0,0 +1,96 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#ifndef OP_API_INC_QUANT_MATMUL_ACTIVATION_QUANT_CHECK_H
12+#define OP_API_INC_QUANT_MATMUL_ACTIVATION_QUANT_CHECK_H
13+#include <map>
14+#include "aclnn/aclnn_base.h"
15+#include "aclnn_kernels/common/op_error_check.h"
16+#include "opdev/common_types.h"
17+#include "opdev/op_dfx.h"
18+#include "opdev/op_log.h"
19+#include "opdev/platform.h"
20+#include "util/math_util.h"
21+#include "matmul/common/op_host/op_api/matmul_util.h"
22+#include "aclnn_kernels/common/op_error_check.h"
23+#include "aclnn_kernels/transdata.h"
24+#include "aclnn_kernels/transpose.h"
25+#include "aclnn_kernels/contiguous.h"
26+#include "aclnn_kernels/reshape.h"
27+#include "opdev/op_executor.h"
28+ 
29+namespace QuantMatmulActivationQuantAclnnCheck {
30+ 
31+using TupleInput = std::tuple<const aclTensor*, const aclTensor*>;
32+using TupleQuant = std::tuple<const aclTensor*, const aclTensor*, const aclTensor*, const aclTensor*, const aclTensor*,
33+ const aclTensor*, const aclTensor*, const int64_t&, const int64_t&>;
34+using TupleFused = std::tuple<const aclTensor*, const char*>;
35+using TupleAttr = std::tuple<bool, bool>;
36+using TupleTensor = std::tuple<const aclTensor*, const aclTensor*, const aclTensor*>;
37+ 
38+static constexpr int INDEX_X1_IN_INPUT_TUPLE = 0;
39+static constexpr int INDEX_X2_IN_INPUT_TUPLE = 1;
40+static constexpr int INDEX_X1_SCALE_IN_QUANT_TUPLE = 0;
41+static constexpr int INDEX_X2_SCALE_IN_QUANT_TUPLE = 1;
42+static constexpr int INDEX_Y_SCALE_IN_QUANT_TUPLE = 2;
43+static constexpr int INDEX_X1_OFFSET_IN_QUANT_TUPLE = 3;
44+static constexpr int INDEX_X2_OFFSET_IN_QUANT_TUPLE = 4;
45+static constexpr int INDEX_Y_OFFSET_IN_QUANT_TUPLE = 5;
46+static constexpr int INDEX_BIAS_IN_QUANT_TUPLE = 6;
47+static constexpr int INDEX_GROUP_SIZE_IN_QUANT_TUPLE = 7;
48+static constexpr int INDEX_INTERFACE_TYPE_IN_QUANT_TUPLE = 8;
49+static constexpr int INDEX_X3_IN_FUSED_TUPLE = 0;
50+static constexpr int INDEX_FUSEDOPTYPE_IN_FUSED_TUPLE = 1;
51+static constexpr int INDEX_OUT_IN_TUPLE = 2;
52+ 
53+static constexpr int64_t NZ_K0_VALUE_BMM_BLOCK_NUM = 16;
54+ 
55+static constexpr size_t LAST_SECOND_DIM_INDEX = 2;
56+static const int64_t NZ_K0_VALUE_INT8_INT4 = 16;
57+static const int64_t NZ_K0_VALUE_INT8_TRANS = 32;
58+static const int64_t NZ_K0_VALUE_INT4_TRANS = 64;
59+static const int NZ_STORAGE_PENULTIMATE_DIM = 16;
60+static const int NZ_STORAGE_LAST_DIM = 32;
61+ 
62+static const int64_t INT4_NUMS_IN_INT32 = 8;
63+static const size_t MIN_DIM_NUM_ND = 2;
64+static const size_t MAX_DIM_NUM_ND = 6;
65+static const size_t MIN_DIM_NUM_NZ = 4;
66+static const size_t MAX_DIM_NUM_NZ = 8;
67+ 
68+static const size_t PENULTIMATE_DIM = 2;
69+ 
70+bool CheckSpecialCase(const aclTensor* tensor, int64_t firstLastDim, int64_t secondLastDim);
71+bool GetTransposeAttrValue(const aclTensor* tensor, bool transpose);
72+op::Shape GetWeightNzShape(const aclTensor* input, bool transpose);
73+bool CheckWeightNzStorageShape(const op::Shape& nzShape, const op::Shape& storageShape);
74+const aclTensor* SetTensorToNZFormat(const aclTensor* input, op::Shape& shape, aclOpExecutor* executor);
75+ 
76+bool TensorContiguousProcess(const aclTensor*& contiguousTensor, bool& transpose, aclOpExecutor* executor);
77+ 
78+aclnnStatus WeightNZCaseProcess(const aclTensor*& x2, bool& transposeX2, aclOpExecutor* executor);
79+aclnnStatus SetSpecilNZTensorToNormalNZFormat(const aclTensor*& input, aclOpExecutor* executor);
80+ 
81+aclTensor* ConvertTensorToInt4(const aclTensor* input, aclOpExecutor* executor);
82+void InputPreProcessA4W4(const aclTensor*& x1, const aclTensor*& x2, aclOpExecutor* executor);
83+aclnnStatus A4W4CaseProcess(const aclTensor*& x1, const aclTensor*& x2, aclOpExecutor* executor);
84+ 
85+const aclTensor* SetTensorToNDFormat(const aclTensor* input);
86+const aclTensor* GetNDFormat(const aclTensor* input);
87+ 
88+void GetDtypeAndTranspose(TupleTensor mandatoryTensors, int64_t& dtype, bool& transposeX1, bool& transposeX2);
89+ 
90+aclnnStatus SpecialOutputProcess(const aclTensor* x1, const aclTensor* x2, const aclTensor* out,
91+ const aclTensor*& matmulRet, aclOpExecutor* executor);
92+aclnnStatus PostMatmulCalcProcess(const aclTensor* matmulRet, const aclTensor* x1, const aclTensor* x2,
93+ const aclTensor* out, aclOpExecutor* executor);
94+ 
95+} // namespace QuantMatmulActivationQuantAclnnCheck
96+#endif // OP_API_INC_QUANT_MATMUL_ACTIVATION_QUANT_CHECK_H