Star872
1766
代码
代码
Issues242
Pull Requests326
Actions
项目讨论
Wiki
分析
Star872
1766
已合并
addRmsNormDynamicMxQuant support split-r template #3573
huanghaihong创建于 4月7日
addRmsNormDynamicMxQuant support split-r template #3573
已合并
huanghaihong创建于 4月7日
从huanghaihong/ops-nnarnd_split_r合入到cann/ops-nnmaster
共 18 个文件变更+3368-1271
Mnorm/add_rms_norm_dynamic_mx_quant/op_graph/add_rms_norm_dynamic_mx_quant_proto.h+4-0
@@ -32,9 +32,13 @@ namespace ge {
32 32 
33* @par Inputs:33* @par Inputs:
34* @li x1: A tensor for add compute. Support dtype: float16, bfloat16, support format: ND.34* @li x1: A tensor for add compute. Support dtype: float16, bfloat16, support format: ND.
35+* The shape supports at least 1 dimension, and at most 7 dimensions.
35* @li x2: A tensor for add compute. Support dtype: float16, bfloat16, support format: ND.36* @li x2: A tensor for add compute. Support dtype: float16, bfloat16, support format: ND.
37+* The shape must be the same as x1.
36* @li gamma: A tensor for rms norm weight params. Support dtype: float32, float16, bfloat16, support format: ND.38* @li gamma: A tensor for rms norm weight params. Support dtype: float32, float16, bfloat16, support format: ND.
39+* The shape must match the normalized dimension of x1.
37* @li beta: An optional tensor for rms norm weight params. Support dtype: float32, float16, bfloat16, support format: ND.40* @li beta: An optional tensor for rms norm weight params. Support dtype: float32, float16, bfloat16, support format: ND.
41+* The shape must match the normalized dimension of x1.
38 42 
39* @par Attributes:43* @par Attributes:
40* @li epsilon: An optional attribute for numerical stability in rms norm, the type is float32. Defaults to 1e-6.44* @li epsilon: An optional attribute for numerical stability in rms norm, the type is float32. Defaults to 1e-6.
Mnorm/add_rms_norm_dynamic_mx_quant/op_host/add_rms_norm_dynamic_mx_quant_tiling.h+78-4
@@ -29,10 +29,53 @@
29#include "log/log.h"29#include "log/log.h"
30#include "util/math_util.h"30#include "util/math_util.h"
31#include "../op_kernel/arch35/add_rms_norm_dynamic_mx_quant_tiling_data.h"31#include "../op_kernel/arch35/add_rms_norm_dynamic_mx_quant_tiling_data.h"
32+#include "../op_kernel/arch35/add_rms_norm_dynamic_mx_quant_tiling_key.h"
32 33 
33using namespace Ops::NN::Optiling;34using namespace Ops::NN::Optiling;
34 35 
35namespace optiling {36namespace optiling {
37+ 
38+namespace add_rms_norm_dynamic_mx_quant {
39+ 
40+enum class ComputeMode : uint64_t
41+{
42+ FULL_LOAD = 0,
43+ SPLIT_R = 1,
44+ REDUCE_EMPTY = 2,
45+};
46+ 
47+enum class YDataType : uint64_t
48+{
49+ FP8 = 0,
50+ FP4 = 1,
51+};
52+ 
53+class AddRmsNormDynamicMxQuantTilingKey {
54+public:
55+ AddRmsNormDynamicMxQuantTilingKey& SetComputeMode(ComputeMode mode)
56+ {
57+ computeMode_ = mode;
58+ return *this;
59+ }
60+ 
61+ AddRmsNormDynamicMxQuantTilingKey& SetYDataType(YDataType type)
62+ {
63+ yDataType_ = type;
64+ return *this;
65+ }
66+ 
67+ uint64_t GetTilingKey() const
68+ {
69+ return GET_TPL_TILING_KEY(static_cast<uint64_t>(computeMode_), static_cast<uint64_t>(yDataType_));
70+ }
71+ 
72+private:
73+ ComputeMode computeMode_ = ComputeMode::FULL_LOAD;
74+ YDataType yDataType_ = YDataType::FP8;
75+};
76+ 
77+} // namespace add_rms_norm_dynamic_mx_quant
78+ 
36// Input indices79// Input indices
37constexpr uint64_t X1_INDEX = 0;80constexpr uint64_t X1_INDEX = 0;
38constexpr uint64_t X2_INDEX = 1;81constexpr uint64_t X2_INDEX = 1;
@@ -62,6 +105,7 @@ constexpr uint32_t FP32_SIZE = 4;
62constexpr uint32_t FP8_SIZE = 1;105constexpr uint32_t FP8_SIZE = 1;
63constexpr uint32_t B16_SIZE = 2;106constexpr uint32_t B16_SIZE = 2;
64constexpr uint32_t NUM_TWO = 2;107constexpr uint32_t NUM_TWO = 2;
108+constexpr uint32_t NUM_FOUR = 4;
65constexpr uint32_t B32_BLOCK_NUM = 8;109constexpr uint32_t B32_BLOCK_NUM = 8;
66constexpr uint64_t ALIGN_FACTOR_512 = 512;110constexpr uint64_t ALIGN_FACTOR_512 = 512;
67constexpr uint64_t COL_ALIGN_NUM = 64;111constexpr uint64_t COL_ALIGN_NUM = 64;
@@ -69,6 +113,7 @@ constexpr uint32_t UB_RESERVE_FOR_RSTD_ALIGN = 1024;
69constexpr uint32_t UB_RESERVE_FOR_OUTPUT_Y_ALIGN = 1536;113constexpr uint32_t UB_RESERVE_FOR_OUTPUT_Y_ALIGN = 1536;
70constexpr uint64_t ARND_REDUCE_EMPTY_PRIORITY = 500;114constexpr uint64_t ARND_REDUCE_EMPTY_PRIORITY = 500;
71constexpr uint64_t ARND_R_FULL_LOAD_PRIORITY = 1000;115constexpr uint64_t ARND_R_FULL_LOAD_PRIORITY = 1000;
116+constexpr uint64_t ARND_SPLIT_R_PRIORITY = 2000;
72constexpr uint64_t FULL_LOAD_R_MAX = 16384;117constexpr uint64_t FULL_LOAD_R_MAX = 16384;
73 118 
74// DstType enum values119// DstType enum values
@@ -77,10 +122,6 @@ constexpr int64_t DST_TYPE_E4M3FN = 36;
77constexpr int64_t DST_TYPE_E2M1 = 40;122constexpr int64_t DST_TYPE_E2M1 = 40;
78constexpr int64_t DST_TYPE_E1M2 = 41;123constexpr int64_t DST_TYPE_E1M2 = 41;
79 124 
80-// Tiling key values
81-constexpr int64_t TILING_KEY_FP8_R_FULL_LOAD = 100;
82-constexpr int64_t TILING_KEY_FP4_R_FULL_LOAD = 101;
83- 
84const std::set<ge::DataType> Y_SUPPORT_DTYPE_FP4_SET = {ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2};125const std::set<ge::DataType> Y_SUPPORT_DTYPE_FP4_SET = {ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2};
85const std::set<ge::DataType> Y_SUPPORT_DTYPE_FP8_SET = {ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2};126const std::set<ge::DataType> Y_SUPPORT_DTYPE_FP8_SET = {ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2};
86const std::set<ge::DataType> Y_SUPPORT_DTYPE_SET = {ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2};127const std::set<ge::DataType> Y_SUPPORT_DTYPE_SET = {ge::DT_FLOAT4_E2M1, ge::DT_FLOAT4_E1M2, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E5M2};
@@ -225,6 +266,39 @@ private:
225 AddRmsNormDynamicMxQuantReduceEmptyTilingData td_;266 AddRmsNormDynamicMxQuantReduceEmptyTilingData td_;
226};267};
227 268 
269+class AddRmsNormDynamicMxQuantSplitRTiling : public AddRmsNormDynamicMxQuantRegbaseTilingBase {
270+public:
271+ explicit AddRmsNormDynamicMxQuantSplitRTiling(gert::TilingContext* context)
272+ : AddRmsNormDynamicMxQuantRegbaseTilingBase(context)
273+ {}
274+ ~AddRmsNormDynamicMxQuantSplitRTiling() override = default;
275+ void Reset(gert::TilingContext* context) override
276+ {
277+ AddRmsNormDynamicMxQuantRegbaseTilingBase::Reset(context);
278+ }
279+protected:
280+ bool IsCapable() override;
281+ ge::graphStatus DoOpTiling() override;
282+ ge::graphStatus DoLibApiTiling() override;
283+ ge::graphStatus PostTiling() override;
284+ uint64_t GetTilingKey() const override;
285+ int64_t GetPowerSplit(uint64_t numN);
286+ int64_t GetCacheId(int64_t idx);
287+ uint64_t GetMaxBaseN(uint64_t initialN);
288+ void SetTilingData();
289+ void PrintTilingData();
290+ 
291+private:
292+ uint64_t baseN_{64}; // initial baseN, power of 2
293+ uint64_t baseM_{128}; // 32-byte aligned, for rstd batch alignment
294+ uint64_t nUbLoops_{0};
295+ uint64_t powerSplit_{0};
296+ uint64_t mainFoldCount_{0};
297+ uint64_t foldTail_{0};
298+ uint64_t baseNBlockSize_{64};
299+ AddRmsNormDynamicMxQuantSplitRTilingData tilingData;
300+};
301+ 
228} // namespace optiling302} // namespace optiling
229 303 
230#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_TILING_H304#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_TILING_H
Mnorm/add_rms_norm_dynamic_mx_quant/op_host/add_rms_norm_dynamic_mx_quant_tiling_base_arch35.cpp+168-123
@@ -67,12 +67,18 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckInputShapeDim()
67 size_t x2DimNum = x2Shape->GetStorageShape().GetDimNum();67 size_t x2DimNum = x2Shape->GetStorageShape().GetDimNum();
68 OP_CHECK_IF(68 OP_CHECK_IF(
69 (x1DimNum > MAX_DIM_CNT) || (x2DimNum > MAX_DIM_CNT) || (x1DimNum < 1) || (x2DimNum < 1),69 (x1DimNum > MAX_DIM_CNT) || (x2DimNum > MAX_DIM_CNT) || (x1DimNum < 1) || (x2DimNum < 1),
70- OP_LOGE(context_->GetNodeName(), "Input x1/x2 dim should be greater than 0,"70+ OP_LOGE_FOR_INVALID_SHAPEDIM(
71- "and not bigger than %u.", MAX_DIM_CNT), return ge::GRAPH_FAILED);71+ context_->GetNodeName(), "x1/x2",
72+ (std::to_string(x1DimNum) + "/" + std::to_string(x2DimNum)).c_str(), "1D to 7D"),
73+ return ge::GRAPH_FAILED);
72 OP_CHECK_IF(!CheckDimBiggerZero(x1Shape, x1DimNum, nodeName, "x1"),74 OP_CHECK_IF(!CheckDimBiggerZero(x1Shape, x1DimNum, nodeName, "x1"),
73- OP_LOGE(context_->GetNodeName(), "Input x1 shape is invalid, please check."), return ge::GRAPH_FAILED);75+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
76+ context_->GetNodeName(), "x1", "", "all dims should be greater than 0"),
77+ return ge::GRAPH_FAILED);
74 OP_CHECK_IF(!CheckDimBiggerZero(x2Shape, x2DimNum, nodeName, "x2"),78 OP_CHECK_IF(!CheckDimBiggerZero(x2Shape, x2DimNum, nodeName, "x2"),
75- OP_LOGE(context_->GetNodeName(), "Input x2 shape is invalid, please check."),return ge::GRAPH_FAILED);79+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
80+ context_->GetNodeName(), "x2", "", "all dims should be greater than 0"),
81+ return ge::GRAPH_FAILED);
76 return ge::GRAPH_SUCCESS;82 return ge::GRAPH_SUCCESS;
77}83}
78 84 
@@ -87,31 +93,40 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckInputShapeValue(
87 OP_CHECK_NULL_WITH_CONTEXT(context_, gammaShape);93 OP_CHECK_NULL_WITH_CONTEXT(context_, gammaShape);
88 94 
89 // x1 and x2 shapes must be equal95 // x1 and x2 shapes must be equal
90- if (!NormCheck::CheckShapeSame(x1Shape, x2Shape, nodeName, "x1", "x2")) {96+ OP_CHECK_IF(
91- OP_LOGE(context_->GetNodeName(), "Input x1 shape is not same with x2 shape.");97+ !NormCheck::CheckShapeSame(x1Shape, x2Shape, nodeName, "x1", "x2"),
92- return ge::GRAPH_FAILED;98+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
93- }99+ context_->GetNodeName(), "x2",
100+ Ops::Base::ToString(x2Shape->GetStorageShape()).c_str(), "same as x1"),
101+ return ge::GRAPH_FAILED);
94 102 
95 // gamma dim num should be 1103 // gamma dim num should be 1
96- if (1 != gammaShape->GetStorageShape().GetDimNum()) {104+ OP_CHECK_IF(
97- OP_LOGE(context_->GetNodeName(), "The shape dim of gamma/beta only support 1, please check.");105+ 1 != gammaShape->GetStorageShape().GetDimNum(),
98- return ge::GRAPH_FAILED;106+ OP_LOGE_FOR_INVALID_SHAPEDIM(
99- }107+ context_->GetNodeName(), "gamma",
108+ std::to_string(gammaShape->GetStorageShape().GetDimNum()).c_str(), "1D"),
109+ return ge::GRAPH_FAILED);
100 110 
101 // gamma should match last dim of x111 // gamma should match last dim of x
102- if (!NormCheck::CheckShapeBC(x1Shape, gammaShape, nodeName, "x1", "gamma", true)) {112+ OP_CHECK_IF(
103- OP_LOGE(context_->GetNodeName(), "Input gamma shape value is not valid with x shape value.");113+ !NormCheck::CheckShapeBC(x1Shape, gammaShape, nodeName, "x1", "gamma", true),
104- return ge::GRAPH_FAILED;114+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
105- }115+ context_->GetNodeName(), "gamma",
116+ Ops::Base::ToString(gammaShape->GetStorageShape()).c_str(),
117+ "should match last dim of x1"),
118+ return ge::GRAPH_FAILED);
106 119 
107 // If beta exists, it should match gamma shape120 // If beta exists, it should match gamma shape
108 if (betaFlag_) {121 if (betaFlag_) {
109 const gert::StorageShape* betaShape = context_->GetOptionalInputShape(BETA_INDEX);122 const gert::StorageShape* betaShape = context_->GetOptionalInputShape(BETA_INDEX);
110 OP_CHECK_NULL_WITH_CONTEXT(context_, betaShape);123 OP_CHECK_NULL_WITH_CONTEXT(context_, betaShape);
111- if (!NormCheck::CheckShapeSame(gammaShape, betaShape, nodeName, "gamma", "beta")) {124+ OP_CHECK_IF(
112- OP_LOGE(context_->GetNodeName(), "Input beta shape is not same with gamma shape.");125+ !NormCheck::CheckShapeSame(gammaShape, betaShape, nodeName, "gamma", "beta"),
113- return ge::GRAPH_FAILED;126+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
114- }127+ context_->GetNodeName(), "beta",
128+ Ops::Base::ToString(betaShape->GetStorageShape()).c_str(), "same as gamma"),
129+ return ge::GRAPH_FAILED);
115 }130 }
116 return ge::GRAPH_SUCCESS;131 return ge::GRAPH_SUCCESS;
117}132}
@@ -130,26 +145,36 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckInputDtype()
130 ge::DataType gammaDtype = context_->GetInputTensor(GAMMA_INDEX)->GetDataType();145 ge::DataType gammaDtype = context_->GetInputTensor(GAMMA_INDEX)->GetDataType();
131 146 
132 // x1 and x2 must have same dtype147 // x1 and x2 must have same dtype
133- if (x1Dtype != x2Dtype) {148+ std::string x2ReasonStr = "same as x1 (" + Ops::Base::ToString(static_cast<ge::DataType>(x1Dtype)) + ")";
134- OP_LOGE(context_->GetNodeName(), "Input x1/x2 dtype should be equal.");149+ OP_CHECK_IF(
135- return ge::GRAPH_FAILED;150+ x1Dtype != x2Dtype,
136- }151+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
152+ context_->GetNodeName(), "x2", Ops::Base::ToString(static_cast<ge::DataType>(x2Dtype)).c_str(),
153+ x2ReasonStr.c_str()),
154+ return ge::GRAPH_FAILED);
137 // x dtype must be FP16 or BF16155 // x dtype must be FP16 or BF16
138- if (supportedXDtypes.count(x1Dtype) == 0) {156+ OP_CHECK_IF(
139- OP_LOGE(context_->GetNodeName(), "Input x1/x2 dtype should be float16 or bfloat16.");157+ supportedXDtypes.count(x1Dtype) == 0,
140- return ge::GRAPH_FAILED;158+ OP_LOGE_FOR_INVALID_DTYPE(
141- }159+ context_->GetNodeName(), "x1", Ops::Base::ToString(static_cast<ge::DataType>(x1Dtype)).c_str(),
160+ "float16 or bfloat16"),
161+ return ge::GRAPH_FAILED);
142 // gamma must be FP16/BF16/FP32162 // gamma must be FP16/BF16/FP32
143- if (supportedGammaDtypes.count(gammaDtype) == 0) {163+ OP_CHECK_IF(
144- OP_LOGE(context_->GetNodeName(), "Input gamma dtype should be float16, bfloat16 or float32.");164+ supportedGammaDtypes.count(gammaDtype) == 0,
145- return ge::GRAPH_FAILED;165+ OP_LOGE_FOR_INVALID_DTYPE(
146- }166+ context_->GetNodeName(), "gamma", Ops::Base::ToString(static_cast<ge::DataType>(gammaDtype)).c_str(),
167+ "float16, bfloat16 or float32"),
168+ return ge::GRAPH_FAILED);
147 if (betaFlag_) {169 if (betaFlag_) {
148 ge::DataType betaDtype = context_->GetInputTensor(BETA_INDEX)->GetDataType();170 ge::DataType betaDtype = context_->GetInputTensor(BETA_INDEX)->GetDataType();
149- if (gammaDtype != betaDtype) {171+ std::string betaReasonStr = "same as gamma (" + Ops::Base::ToString(static_cast<ge::DataType>(gammaDtype)) + ")";
150- OP_LOGE(context_->GetNodeName(), "Input gamma/beta dtype should be equal.");172+ OP_CHECK_IF(
151- return ge::GRAPH_FAILED;173+ gammaDtype != betaDtype,
152- }174+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
175+ context_->GetNodeName(), "beta", Ops::Base::ToString(static_cast<ge::DataType>(betaDtype)).c_str(),
176+ betaReasonStr.c_str()),
177+ return ge::GRAPH_FAILED);
153 }178 }
154 179 
155 xDtype_ = x1Dtype;180 xDtype_ = x1Dtype;
@@ -166,10 +191,12 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckOutputDtype()
166 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(Y_INDEX));191 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(Y_INDEX));
167 ge::DataType yDtype = context_->GetOutputDesc(Y_INDEX)->GetDataType();192 ge::DataType yDtype = context_->GetOutputDesc(Y_INDEX)->GetDataType();
168 193 
169- if (Y_SUPPORT_DTYPE_SET.count(yDtype) == 0) {194+ OP_CHECK_IF(
170- OP_LOGE(context_->GetNodeName(), "Output y dtype should be FP4_E2M1/E1M2 or FP8_E4M3FN/E5M2.");195+ Y_SUPPORT_DTYPE_SET.count(yDtype) == 0,
171- return ge::GRAPH_FAILED;196+ OP_LOGE_FOR_INVALID_DTYPE(
172- }197+ context_->GetNodeName(), "y", Ops::Base::ToString(static_cast<ge::DataType>(yDtype)).c_str(),
198+ "float4_e2m1, float4_e1m2, float8_e4m3fn or float8_e5m2"),
199+ return ge::GRAPH_FAILED);
173 200 
174 yDtype_ = yDtype;201 yDtype_ = yDtype;
175 202 
@@ -177,14 +204,18 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckOutputDtype()
177 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(MXSCALE_INDEX));204 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(MXSCALE_INDEX));
178 ge::DataType outputXDtype = context_->GetOutputDesc(X_INDEX)->GetDataType();205 ge::DataType outputXDtype = context_->GetOutputDesc(X_INDEX)->GetDataType();
179 ge::DataType mxScaleDtype = context_->GetOutputDesc(MXSCALE_INDEX)->GetDataType();206 ge::DataType mxScaleDtype = context_->GetOutputDesc(MXSCALE_INDEX)->GetDataType();
180- if (outputXDtype != xDtype_){207+ OP_CHECK_IF(
181- OP_LOGE(context_->GetNodeName(), "Output X dtype should be equal to input X1/X2 dtype.");208+ outputXDtype != xDtype_,
182- return ge::GRAPH_FAILED;209+ OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
183- }210+ context_->GetNodeName(), "output_x", Ops::Base::ToString(static_cast<ge::DataType>(outputXDtype)).c_str(),
184- if (mxScaleDtype != ge::DT_FLOAT8_E8M0){211+ "same as input x1/x2"),
185- OP_LOGE(context_->GetNodeName(), "Output mxScale dtype should be FLOAT8_E8M0.");212+ return ge::GRAPH_FAILED);
186- return ge::GRAPH_FAILED;213+ OP_CHECK_IF(
187- }214+ mxScaleDtype != ge::DT_FLOAT8_E8M0,
215+ OP_LOGE_FOR_INVALID_DTYPE(
216+ context_->GetNodeName(), "mxscale", Ops::Base::ToString(static_cast<ge::DataType>(mxScaleDtype)).c_str(),
217+ "float8_e8m0"),
218+ return ge::GRAPH_FAILED);
188 219 
189 // output_rstd ATTR校验：设置rstd_flag220 // output_rstd ATTR校验：设置rstd_flag
190 rstdFlag_ = 0;221 rstdFlag_ = 0;
@@ -201,10 +232,12 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckOutputDtype()
201 if (rstdFlag_) {232 if (rstdFlag_) {
202 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(RSTD_INDEX));233 OP_CHECK_NULL_WITH_CONTEXT(context_, context_->GetOutputDesc(RSTD_INDEX));
203 ge::DataType rstdDtype = context_->GetOutputDesc(RSTD_INDEX)->GetDataType();234 ge::DataType rstdDtype = context_->GetOutputDesc(RSTD_INDEX)->GetDataType();
204- if (rstdDtype != ge::DT_FLOAT){235+ OP_CHECK_IF(
205- OP_LOGE(context_->GetNodeName(), "Output rstd dtype should be FLOAT32.");236+ rstdDtype != ge::DT_FLOAT,
206- return ge::GRAPH_FAILED;237+ OP_LOGE_FOR_INVALID_DTYPE(
207- }238+ context_->GetNodeName(), "rstd", Ops::Base::ToString(static_cast<ge::DataType>(rstdDtype)).c_str(),
239+ "float"),
240+ return ge::GRAPH_FAILED);
208 }241 }
209 return ge::GRAPH_SUCCESS;242 return ge::GRAPH_SUCCESS;
210}243}
@@ -223,60 +256,55 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckMxQuantParams()
223 256 
224 OP_CHECK_IF(257 OP_CHECK_IF(
225 (rm == MxRoundMode::UNDEFINED),258 (rm == MxRoundMode::UNDEFINED),
226- OP_LOGE(context_->GetNodeName(),259+ OP_LOGE_WITH_INVALID_ATTR(context_->GetNodeName(), "round_mode", roundModeStr, "rint, round or floor"),
227- "invalid round_mode:%s; round_mode should be one of {rint, floor, round}", roundModeStr),
228 return ge::GRAPH_FAILED);260 return ge::GRAPH_FAILED);
229 261 
230 // FP8输出类型仅支持rint262 // FP8输出类型仅支持rint
231 OP_CHECK_IF(263 OP_CHECK_IF(
232 (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0 && rm != MxRoundMode::RINT),264 (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0 && rm != MxRoundMode::RINT),
233- OP_LOGE(context_->GetNodeName(),265+ OP_LOGE_WITH_INVALID_ATTR(context_->GetNodeName(), "round_mode", roundModeStr, "rint"),
234- "When output y's data type is FLOAT8_E4M3FN/FLOAT8_E5M2, round_mode:[%s] only support rint, "
235- "please check.",
236- roundModeStr),
237 return ge::GRAPH_FAILED);266 return ge::GRAPH_FAILED);
238 267 
239 // 2. dst_type校验：必须与y的dtype对应268 // 2. dst_type校验：必须与y的dtype对应
240 const int64_t* dstTypePtr = attrs->GetAttrPointer<int64_t>(DST_TYPE_ATTR_INDEX);269 const int64_t* dstTypePtr = attrs->GetAttrPointer<int64_t>(DST_TYPE_ATTR_INDEX);
241 if (dstTypePtr != nullptr) {270 if (dstTypePtr != nullptr) {
242 int64_t dstType = *dstTypePtr;271 int64_t dstType = *dstTypePtr;
243- if ((yDtype_ == ge::DT_FLOAT4_E2M1 && dstType != DST_TYPE_E2M1) ||272+ OP_CHECK_IF(
244- (yDtype_ == ge::DT_FLOAT4_E1M2 && dstType != DST_TYPE_E1M2) ||273+ (yDtype_ == ge::DT_FLOAT4_E2M1 && dstType != DST_TYPE_E2M1) ||
245- (yDtype_ == ge::DT_FLOAT8_E4M3FN && dstType != DST_TYPE_E4M3FN) ||274+ (yDtype_ == ge::DT_FLOAT4_E1M2 && dstType != DST_TYPE_E1M2) ||
246- (yDtype_ == ge::DT_FLOAT8_E5M2 && dstType != DST_TYPE_E5M2)) {275+ (yDtype_ == ge::DT_FLOAT8_E4M3FN && dstType != DST_TYPE_E4M3FN) ||
247- OP_LOGE(context_->GetNodeName(),276+ (yDtype_ == ge::DT_FLOAT8_E5M2 && dstType != DST_TYPE_E5M2),
248- "y's data type and dst_type is not corresponded. "277+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
249- "FLOAT4_E2M1/FLOAT4_E1M2/FLOAT8_E4M3FN/FLOAT8_E5M2 correspond to dst_type: 40/41/36/35.");278+ context_->GetNodeName(), "dst_type", std::to_string(dstType).c_str(),
250- return ge::GRAPH_FAILED;279+ "40/41/36/35 for float4_e2m1/float4_e1m2/float8_e4m3fn/float8_e5m2"),
251- }280+ return ge::GRAPH_FAILED);
252 }281 }
253 282 
254 // 3. quant_alg校验：必须为0或1，FP4仅支持0283 // 3. quant_alg校验：必须为0或1，FP4仅支持0
255 const int64_t* quantAlgPtr = attrs->GetAttrPointer<int64_t>(QUANT_ALG_ATTR_INDEX);284 const int64_t* quantAlgPtr = attrs->GetAttrPointer<int64_t>(QUANT_ALG_ATTR_INDEX);
256 if (quantAlgPtr != nullptr) {285 if (quantAlgPtr != nullptr) {
257 int64_t quantAlg = *quantAlgPtr;286 int64_t quantAlg = *quantAlgPtr;
258- if (quantAlg < 0 || quantAlg > 1) {287+ OP_CHECK_IF(
259- OP_LOGE(context_->GetNodeName(), "The quant_alg[%ld] should be 0 or 1.", quantAlg);288+ quantAlg < 0 || quantAlg > 1,
260- return ge::GRAPH_FAILED;289+ OP_LOGE_WITH_INVALID_ATTR(
261- }290+ context_->GetNodeName(), "quant_alg", std::to_string(quantAlg).c_str(), "0 or 1"),
262- if (quantAlg == 1 && Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {291+ return ge::GRAPH_FAILED);
263- OP_LOGE(context_->GetNodeName(),292+ OP_CHECK_IF(
264- "When y's data type is FLOAT4_E2M1/FLOAT4_E1M2, quant_alg must be set to 0.");293+ quantAlg == 1 && Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0,
265- return ge::GRAPH_FAILED;294+ OP_LOGE_WITH_INVALID_ATTR(context_->GetNodeName(), "quant_alg", std::to_string(quantAlg).c_str(), "0"),
266- }295+ return ge::GRAPH_FAILED);
267 }296 }
268 297 
269 // 4. FP4输出时，x的最后一维必须为偶数298 // 4. FP4输出时，x的最后一维必须为偶数
270 if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {299 if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {
271 const gert::Shape x1Shape = context_->GetInputShape(X1_INDEX)->GetStorageShape();300 const gert::Shape x1Shape = context_->GetInputShape(X1_INDEX)->GetStorageShape();
272 size_t lastDim = x1Shape.GetDimNum() - 1;301 size_t lastDim = x1Shape.GetDimNum() - 1;
273- if (x1Shape.GetDim(lastDim) % NUM_TWO != 0) {302+ OP_CHECK_IF(
274- OP_LOGE(context_->GetNodeName(),303+ x1Shape.GetDim(lastDim) % NUM_TWO != 0,
275- "When output y's data type is FLOAT4_E2M1/FLOAT4_E1M2, "304+ OP_LOGE_FOR_INVALID_SHAPESIZE(
276- "the last axis of x should be even, but got %ld.",305+ context_->GetNodeName(), "x1", std::to_string(x1Shape.GetDim(lastDim)).c_str(),
277- x1Shape.GetDim(lastDim));306+ "even number when y dtype is fp4"),
278- return ge::GRAPH_FAILED;307+ return ge::GRAPH_FAILED);
279- }
280 }308 }
281 return ge::GRAPH_SUCCESS;309 return ge::GRAPH_SUCCESS;
282}310}
@@ -294,16 +322,20 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckOutputShapeValue
294 OP_CHECK_NULL_WITH_CONTEXT(context_, outputYShape);322 OP_CHECK_NULL_WITH_CONTEXT(context_, outputYShape);
295 323 
296 // outputX and inputX1 shapes must be equal324 // outputX and inputX1 shapes must be equal
297- if (!NormCheck::CheckShapeSame(x1Shape, outputXShape, nodeName, "inputX1", "outputX")) {325+ OP_CHECK_IF(
298- OP_LOGE(context_->GetNodeName(), "Output X shape is not same with inputX1 shape.");326+ !NormCheck::CheckShapeSame(x1Shape, outputXShape, nodeName, "inputX1", "outputX"),
299- return ge::GRAPH_FAILED;327+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
300- }328+ context_->GetNodeName(), "output_x",
329+ Ops::Base::ToString(outputXShape->GetStorageShape()).c_str(), "same as input x1"),
330+ return ge::GRAPH_FAILED);
301 331 
302 // outputY and inputX1 shapes must be equal332 // outputY and inputX1 shapes must be equal
303- if (!NormCheck::CheckShapeSame(x1Shape, outputYShape, nodeName, "inputX1", "outputY")) {333+ OP_CHECK_IF(
304- OP_LOGE(context_->GetNodeName(), "Output Y shape is not same with inputX1 shape.");334+ !NormCheck::CheckShapeSame(x1Shape, outputYShape, nodeName, "inputX1", "outputY"),
305- return ge::GRAPH_FAILED;335+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
306- }336+ context_->GetNodeName(), "output_y",
337+ Ops::Base::ToString(outputYShape->GetStorageShape()).c_str(), "same as input x1"),
338+ return ge::GRAPH_FAILED);
307 339
308 return ge::GRAPH_SUCCESS;340 return ge::GRAPH_SUCCESS;
309}341}
@@ -323,34 +355,41 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckMxScaleRstdShape
323 size_t gammaRank = gammaShape.GetDimNum();355 size_t gammaRank = gammaShape.GetDimNum();
324 356 
325 // mxscale rank必须等于xRank + 1357 // mxscale rank必须等于xRank + 1
326- if (mxscaleRank != xRank + 1) {358+ OP_CHECK_IF(
327- OP_LOGE(context_->GetNodeName(), "Output mxscale rank [%zu] should be equal to xRank + 1 [%zu].", 359+ mxscaleRank != xRank + 1,
328- mxscaleRank, xRank + 1);360+ OP_LOGE_FOR_INVALID_SHAPEDIM(
329- return ge::GRAPH_FAILED;361+ context_->GetNodeName(), "mxscale", std::to_string(mxscaleRank).c_str(),
330- }362+ std::to_string(xRank + 1).c_str()),
363+ return ge::GRAPH_FAILED);
331 364 
332 // A维度的轴必须一致365 // A维度的轴必须一致
333 for (size_t i = 0; i < xRank - gammaRank; i++) {366 for (size_t i = 0; i < xRank - gammaRank; i++) {
334- if (mxscaleShape.GetDim(i) != x1Shape.GetDim(i)) {367+ OP_CHECK_IF(
335- OP_LOGE(context_->GetNodeName(),368+ mxscaleShape.GetDim(i) != x1Shape.GetDim(i),
336- "mxscale shape dim[%zu]=%ld should match x1 shape dim[%zu]=%ld.",369+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
337- i, mxscaleShape.GetDim(i), i, x1Shape.GetDim(i));370+ context_->GetNodeName(), "mxscale",
338- return ge::GRAPH_FAILED;371+ Ops::Base::ToString(mxscaleShape).c_str(),
339- }372+ "batch dims should match x1"),
373+ return ge::GRAPH_FAILED);
340 }374 }
341 375 
342 // MxScale输出的 (-2轴, -1轴) shape value: (CeilDiv(CeilDiv(R, 32), 2), 2)376 // MxScale输出的 (-2轴, -1轴) shape value: (CeilDiv(CeilDiv(R, 32), 2), 2)
343 uint64_t expectedLastDim = Ops::Base::CeilDiv(377 uint64_t expectedLastDim = Ops::Base::CeilDiv(
344 Ops::Base::CeilDiv(numCol_, static_cast<uint64_t>(MX_BLOCK_SIZE_32)), static_cast<uint64_t>(NUM_TWO));378 Ops::Base::CeilDiv(numCol_, static_cast<uint64_t>(MX_BLOCK_SIZE_32)), static_cast<uint64_t>(NUM_TWO));
345- if (mxscaleShape.GetDim(mxscaleRank - NUM_TWO) != static_cast<int64_t>(expectedLastDim)) {379+ std::string mxscaleReasonStr = "mxscale -2nd dim should be CeilDiv(CeilDiv(R=" +
346- OP_LOGE(context_->GetNodeName(), "mxscale -2last dim[%ld] should be CeilDiv(CeilDiv(R=%ld, 32), 2) = %ld.",380+ std::to_string(numCol_) + ", 32), 2) = " + std::to_string(expectedLastDim);
347- mxscaleShape.GetDim(mxscaleRank - NUM_TWO), static_cast<int64_t>(numCol_), static_cast<int64_t>(expectedLastDim));381+ OP_CHECK_IF(
348- return ge::GRAPH_FAILED;382+ mxscaleShape.GetDim(mxscaleRank - NUM_TWO) != static_cast<int64_t>(expectedLastDim),
349- }383+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
350- if (mxscaleShape.GetDim(mxscaleRank - 1) != static_cast<int64_t>(NUM_TWO)) {384+ context_->GetNodeName(), "mxscale",
351- OP_LOGE(context_->GetNodeName(), "mxscale last dim[%ld] should be 2.", mxscaleShape.GetDim(mxscaleRank - 1));385+ Ops::Base::ToString(mxscaleShape).c_str(), mxscaleReasonStr.c_str()),
352- return ge::GRAPH_FAILED;386+ return ge::GRAPH_FAILED);
353- }387+ OP_CHECK_IF(
388+ mxscaleShape.GetDim(mxscaleRank - 1) != static_cast<int64_t>(NUM_TWO),
389+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
390+ context_->GetNodeName(), "mxscale",
391+ Ops::Base::ToString(mxscaleShape).c_str(), "last dim should be 2"),
392+ return ge::GRAPH_FAILED);
354 393 
355 if (rstdFlag_) {394 if (rstdFlag_) {
356 const gert::StorageShape* rstdShapePtr = context_->GetOutputShape(RSTD_INDEX);395 const gert::StorageShape* rstdShapePtr = context_->GetOutputShape(RSTD_INDEX);
@@ -358,22 +397,28 @@ ge::graphStatus AddRmsNormDynamicMxQuantRegbaseTilingBase::CheckMxScaleRstdShape
358 const gert::Shape rstdShape = rstdShapePtr->GetStorageShape();397 const gert::Shape rstdShape = rstdShapePtr->GetStorageShape();
359 OP_CHECK_IF(398 OP_CHECK_IF(
360 (rstdShape.GetDimNum() != xRank),399 (rstdShape.GetDimNum() != xRank),
361- OP_LOGE(context_->GetNodeName(), "Invalid rstd shape dim num (must same with input x1)."),400+ OP_LOGE_FOR_INVALID_SHAPEDIM(
362- return ge::GRAPH_FAILED);401+ context_->GetNodeName(), "rstd",
402+ std::to_string(rstdShape.GetDimNum()).c_str(), std::to_string(xRank).c_str()),
403+ return ge::GRAPH_FAILED);
363 // A维度的轴必须一致，其他维度为1404 // A维度的轴必须一致，其他维度为1
364 for (size_t i = 0; i < xRank; i++) {405 for (size_t i = 0; i < xRank; i++) {
365 if (i >= xRank - gammaRank) {406 if (i >= xRank - gammaRank) {
366- if (rstdShape.GetDim(i) != 1) {407+ OP_CHECK_IF(
367- OP_LOGE(context_->GetNodeName(), "rstd shape dim[%zu]=%ld should be 1.", i, rstdShape.GetDim(i));408+ rstdShape.GetDim(i) != 1,
368- return ge::GRAPH_FAILED;409+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
369- }410+ context_->GetNodeName(), "rstd",
411+ Ops::Base::ToString(rstdShape).c_str(),
412+ "norm dims should be 1"),
413+ return ge::GRAPH_FAILED);
370 } else {414 } else {
371- if (rstdShape.GetDim(i) != x1Shape.GetDim(i)) {415+ OP_CHECK_IF(
372- OP_LOGE(context_->GetNodeName(),416+ rstdShape.GetDim(i) != x1Shape.GetDim(i),
373- "rstd shape dim[%zu]=%ld should match x1 shape dim[%zu]=%ld.",417+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
374- i, rstdShape.GetDim(i), i, x1Shape.GetDim(i));418+ context_->GetNodeName(), "rstd",
375- return ge::GRAPH_FAILED;419+ Ops::Base::ToString(rstdShape).c_str(),
376- }420+ "batch dims should match x1"),
421+ return ge::GRAPH_FAILED);
377 }422 }
378 }423 }
379 }424 }
Mnorm/add_rms_norm_dynamic_mx_quant/op_host/add_rms_norm_dynamic_mx_quant_tiling_r_full_load_arch35.cpp+13-8
@@ -15,6 +15,8 @@
15#include "add_rms_norm_dynamic_mx_quant_tiling.h"15#include "add_rms_norm_dynamic_mx_quant_tiling.h"
16#include "norm/norm_common/op_host/norm_tiling_check_common.h"16#include "norm/norm_common/op_host/norm_tiling_check_common.h"
17 17 
18+using namespace optiling::add_rms_norm_dynamic_mx_quant;
19+ 
18namespace optiling {20namespace optiling {
19using namespace NormCheck;21using namespace NormCheck;
20 22 
@@ -93,7 +95,7 @@ ge::graphStatus AddRmsNormDynamicMxQuantRFullLoadTiling::SetTilingParams()
93 }95 }
94 96 
95 if (rowFactor < 1) {97 if (rowFactor < 1) {
96- OP_LOGE(context_->GetNodeName(), "Cannot fit even 1 row in UB for R-full-load. R=%lu.", numCol_);98+ OP_LOGI(context_->GetNodeName(), "Cannot fit even 1 row in UB for R-full-load. R=%lu.", numCol_);
97 return ge::GRAPH_PARAM_INVALID; // R轴不能全载，继续调下个模板99 return ge::GRAPH_PARAM_INVALID; // R轴不能全载，继续调下个模板
98 }100 }
99 101 
@@ -110,6 +112,11 @@ bool AddRmsNormDynamicMxQuantRFullLoadTiling::IsCapable()
110 return false;112 return false;
111 }113 }
112 if (numCol_ > FULL_LOAD_R_MAX) {114 if (numCol_ > FULL_LOAD_R_MAX) {
115+ OP_LOGD(
116+ context_->GetNodeName(),
117+ "FullLoad IsCapable false: numCol=%ld >= fullLoadRMax=%ld, "
118+ "binary add rounds increase, recommend SplitR mode.",
119+ numCol_, FULL_LOAD_R_MAX);
113 return false;120 return false;
114 }121 }
115 return true;122 return true;
@@ -211,16 +218,14 @@ ge::graphStatus AddRmsNormDynamicMxQuantRFullLoadTiling::PostTiling()
211 218 
212uint64_t AddRmsNormDynamicMxQuantRFullLoadTiling::GetTilingKey() const219uint64_t AddRmsNormDynamicMxQuantRFullLoadTiling::GetTilingKey() const
213{220{
214- // Tiling key221+ AddRmsNormDynamicMxQuantTilingKey tilingKey;
215- uint64_t tilingKey = 0;222+ tilingKey.SetComputeMode(ComputeMode::FULL_LOAD);
216 if (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0) {223 if (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0) {
217- tilingKey = TILING_KEY_FP8_R_FULL_LOAD;224+ tilingKey.SetYDataType(YDataType::FP8);
218- OP_LOGD(context_->GetNodeName(), "TilingKey is %lu.", TILING_KEY_FP8_R_FULL_LOAD);
219 } else if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {225 } else if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {
220- tilingKey = TILING_KEY_FP4_R_FULL_LOAD;226+ tilingKey.SetYDataType(YDataType::FP4);
221- OP_LOGD(context_->GetNodeName(), "TilingKey is %lu.", TILING_KEY_FP4_R_FULL_LOAD);
222 }227 }
223- return tilingKey;228+ return tilingKey.GetTilingKey();
224}229}
225 230 
226REGISTER_OPS_TILING_TEMPLATE(AddRmsNormDynamicMxQuant, AddRmsNormDynamicMxQuantRFullLoadTiling, ARND_R_FULL_LOAD_PRIORITY);231REGISTER_OPS_TILING_TEMPLATE(AddRmsNormDynamicMxQuant, AddRmsNormDynamicMxQuantRFullLoadTiling, ARND_R_FULL_LOAD_PRIORITY);
Mnorm/add_rms_norm_dynamic_mx_quant/op_host/add_rms_norm_dynamic_mx_quant_tiling_reduce_empty_arch35.cpp+151-147
@@ -1,147 +1,151 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.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 of3+ * 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").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.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,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.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- 10+ 
11-/*!11+/*!
12- * \file add_rms_norm_dynamic_mx_quant_tiling_reduce_empty_arch35.cpp12+ * \file add_rms_norm_dynamic_mx_quant_tiling_reduce_empty_arch35.cpp
13- * \brief13+ * \brief
14- */14+ */
15-#include "add_rms_norm_dynamic_mx_quant_tiling.h"15+#include "add_rms_norm_dynamic_mx_quant_tiling.h"
16- 16+ 
17-namespace optiling {17+using namespace optiling::add_rms_norm_dynamic_mx_quant;
18- 18+ 
19-constexpr uint64_t TILING_KEY_REDUCE_EMPTY = 300;19+namespace optiling {
20-constexpr static int64_t BUFFER_NUM = 1;20+ 
21-constexpr static int64_t SINGLE_AIV_CORE_THRESHOLD_BYTES = 32L * 1024L;21+constexpr static int64_t BUFFER_NUM = 1;
22-constexpr static uint64_t DEFAULT_NUM = 0;22+constexpr static int64_t SINGLE_AIV_CORE_THRESHOLD_BYTES = 32L * 1024L;
23- 23+constexpr static uint64_t DEFAULT_NUM = 0;
24-bool AddRmsNormDynamicMxQuantReduceEmptyTiling::IsCapable()24+ 
25-{25+bool AddRmsNormDynamicMxQuantReduceEmptyTiling::IsCapable()
26- if (numCol_ != 0 && numRow_ != 0) {26+{
27- OP_LOGI(context_->GetNodeName(), "ReduceEmptyTiling not applicable: numCol=%lu != 0 or numRow=%lu != 0.", numCol_, numRow_);27+ if (numCol_ != 0 && numRow_ != 0) {
28- return false;28+ OP_LOGI(context_->GetNodeName(), "ReduceEmptyTiling not applicable: numCol=%lu != 0 or numRow=%lu != 0.", numCol_, numRow_);
29- }29+ return false;
30- OP_LOGI(context_->GetNodeName(), "ReduceEmptyTiling IsCapable: true (numCol=%lu, numRow=%lu).", numCol_, numRow_);30+ }
31- return true;31+ OP_LOGI(context_->GetNodeName(), "ReduceEmptyTiling IsCapable: true (numCol=%lu, numRow=%lu).", numCol_, numRow_);
32-}32+ return true;
33- 33+}
34-ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::DoOpTiling()34+ 
35-{35+ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::DoOpTiling()
36- OP_LOGD(context_->GetNodeName(), "Enter ReduceEmpty DoOpTiling.");36+{
37- 37+ OP_LOGD(context_->GetNodeName(), "Enter ReduceEmpty DoOpTiling.");
38- td_.rstdFlag = rstdFlag_;38+ 
39- td_.numRow = numRow_;39+ td_.rstdFlag = rstdFlag_;
40- 40+ td_.numRow = numRow_;
41- if (rstdFlag_ == 0 || numRow_ == 0) {41+ 
42- usedCoreNum_ = 1;42+ if (rstdFlag_ == 0 || numRow_ == 0) {
43- td_.perCoreElements = DEFAULT_NUM;43+ usedCoreNum_ = 1;
44- td_.lastCoreElements = DEFAULT_NUM;44+ td_.perCoreElements = DEFAULT_NUM;
45- td_.perCoreLoops = DEFAULT_NUM;45+ td_.lastCoreElements = DEFAULT_NUM;
46- td_.perCorePerLoopElements = DEFAULT_NUM;46+ td_.perCoreLoops = DEFAULT_NUM;
47- td_.perCoreLastLoopElements = DEFAULT_NUM;47+ td_.perCorePerLoopElements = DEFAULT_NUM;
48- td_.lastCoreLoops = DEFAULT_NUM;48+ td_.perCoreLastLoopElements = DEFAULT_NUM;
49- td_.lastCorePerLoopElements = DEFAULT_NUM;49+ td_.lastCoreLoops = DEFAULT_NUM;
50- td_.lastCoreLastLoopElements = DEFAULT_NUM;50+ td_.lastCorePerLoopElements = DEFAULT_NUM;
51- return ge::GRAPH_SUCCESS;51+ td_.lastCoreLastLoopElements = DEFAULT_NUM;
52- }52+ return ge::GRAPH_SUCCESS;
53- 53+ }
54- // rstd output: numRow elements of FP3254+ 
55- int64_t totalLength = static_cast<int64_t>(numRow_);55+ // rstd output: numRow elements of FP32
56- int64_t elemSize = FP32_SIZE;56+ int64_t totalLength = static_cast<int64_t>(numRow_);
57- int64_t aivNum = static_cast<int64_t>(totalCoreNum_);57+ int64_t elemSize = FP32_SIZE;
58- int64_t perLoopMaxElements =58+ int64_t aivNum = static_cast<int64_t>(totalCoreNum_);
59- static_cast<int64_t>(maxUbSize_) / elemSize / BUFFER_NUM;59+ int64_t perLoopMaxElements =
60- 60+ static_cast<int64_t>(maxUbSize_) / elemSize / BUFFER_NUM;
61- // Core splitting61+ 
62- int64_t blockNum = Ops::Base::CeilDiv(totalLength * elemSize, SINGLE_AIV_CORE_THRESHOLD_BYTES);62+ // Core splitting
63- if (blockNum > aivNum) {63+ int64_t blockNum = Ops::Base::CeilDiv(totalLength * elemSize, SINGLE_AIV_CORE_THRESHOLD_BYTES);
64- blockNum = aivNum;64+ if (blockNum > aivNum) {
65- }65+ blockNum = aivNum;
66- 66+ }
67- usedCoreNum_ = static_cast<uint64_t>(blockNum);67+ 
68- 68+ usedCoreNum_ = static_cast<uint64_t>(blockNum);
69- int64_t perCoreElements = Ops::Base::CeilDiv(totalLength, blockNum);69+ 
70- int64_t lastCoreElements = totalLength - (blockNum - 1) * perCoreElements;70+ int64_t perCoreElements = Ops::Base::CeilDiv(totalLength, blockNum);
71- 71+ int64_t lastCoreElements = totalLength - (blockNum - 1) * perCoreElements;
72- td_.perCoreElements = static_cast<uint64_t>(perCoreElements);72+ 
73- td_.lastCoreElements = static_cast<uint64_t>(lastCoreElements);73+ td_.perCoreElements = static_cast<uint64_t>(perCoreElements);
74- 74+ td_.lastCoreElements = static_cast<uint64_t>(lastCoreElements);
75- // Intra-core loop splitting (non-last core)75+ 
76- int64_t ubAlignElems = static_cast<int64_t>(ubBlockSize_) / elemSize;76+ // Intra-core loop splitting (non-last core)
77- int64_t perCorePerLoopElements = Ops::Base::FloorAlign(77+ int64_t ubAlignElems = static_cast<int64_t>(ubBlockSize_) / elemSize;
78- std::min(perLoopMaxElements, perCoreElements), ubAlignElems);78+ int64_t perCorePerLoopElements = Ops::Base::FloorAlign(
79- 79+ std::min(perLoopMaxElements, perCoreElements), ubAlignElems);
80- int64_t perCoreLoops = Ops::Base::CeilDiv(perCoreElements, perCorePerLoopElements);80+ 
81- int64_t perCoreLastLoopElements = perCoreElements - (perCoreLoops - 1) * perCorePerLoopElements;81+ int64_t perCoreLoops = Ops::Base::CeilDiv(perCoreElements, perCorePerLoopElements);
82- 82+ int64_t perCoreLastLoopElements = perCoreElements - (perCoreLoops - 1) * perCorePerLoopElements;
83- td_.perCoreLoops = static_cast<uint64_t>(perCoreLoops);83+ 
84- td_.perCorePerLoopElements = static_cast<uint64_t>(perCorePerLoopElements);84+ td_.perCoreLoops = static_cast<uint64_t>(perCoreLoops);
85- td_.perCoreLastLoopElements = static_cast<uint64_t>(perCoreLastLoopElements);85+ td_.perCorePerLoopElements = static_cast<uint64_t>(perCorePerLoopElements);
86- 86+ td_.perCoreLastLoopElements = static_cast<uint64_t>(perCoreLastLoopElements);
87- // Intra-core loop splitting (last core)87+ 
88- int64_t lastCorePerLoopElements = Ops::Base::FloorAlign(88+ // Intra-core loop splitting (last core)
89- std::min(perLoopMaxElements, lastCoreElements), ubAlignElems);89+ int64_t lastCorePerLoopElements = Ops::Base::FloorAlign(
90- if (lastCorePerLoopElements < 1) {90+ std::min(perLoopMaxElements, lastCoreElements), ubAlignElems);
91- lastCorePerLoopElements = lastCoreElements;91+ if (lastCorePerLoopElements < 1) {
92- }92+ lastCorePerLoopElements = lastCoreElements;
93- int64_t lastCoreLoops = Ops::Base::CeilDiv(lastCoreElements, lastCorePerLoopElements);93+ }
94- int64_t lastCoreLastLoopElements = lastCoreElements - (lastCoreLoops - 1) * lastCorePerLoopElements;94+ int64_t lastCoreLoops = Ops::Base::CeilDiv(lastCoreElements, lastCorePerLoopElements);
95- 95+ int64_t lastCoreLastLoopElements = lastCoreElements - (lastCoreLoops - 1) * lastCorePerLoopElements;
96- td_.lastCoreLoops = static_cast<uint64_t>(lastCoreLoops);96+ 
97- td_.lastCorePerLoopElements = static_cast<uint64_t>(lastCorePerLoopElements);97+ td_.lastCoreLoops = static_cast<uint64_t>(lastCoreLoops);
98- td_.lastCoreLastLoopElements = static_cast<uint64_t>(lastCoreLastLoopElements);98+ td_.lastCorePerLoopElements = static_cast<uint64_t>(lastCorePerLoopElements);
99- 99+ td_.lastCoreLastLoopElements = static_cast<uint64_t>(lastCoreLastLoopElements);
100- OP_LOGI(context_->GetNodeName(),100+ 
101- "ReduceEmpty DoOpTiling: blockNum=%lu, perCoreElements=%lu, lastCoreElements=%lu, "101+ OP_LOGI(context_->GetNodeName(),
102- "perCoreLoops=%lu, perCorePerLoopElements=%lu, perCoreLastLoopElements=%lu, "102+ "ReduceEmpty DoOpTiling: blockNum=%lu, perCoreElements=%lu, lastCoreElements=%lu, "
103- "lastCoreLoops=%lu, lastCorePerLoopElements=%lu, lastCoreLastLoopElements=%lu, rstdFlag=%u.",103+ "perCoreLoops=%lu, perCorePerLoopElements=%lu, perCoreLastLoopElements=%lu, "
104- usedCoreNum_, td_.perCoreElements, td_.lastCoreElements,104+ "lastCoreLoops=%lu, lastCorePerLoopElements=%lu, lastCoreLastLoopElements=%lu, rstdFlag=%u.",
105- td_.perCoreLoops, td_.perCorePerLoopElements, td_.perCoreLastLoopElements,105+ usedCoreNum_, td_.perCoreElements, td_.lastCoreElements,
106- td_.lastCoreLoops, td_.lastCorePerLoopElements, td_.lastCoreLastLoopElements, td_.rstdFlag);106+ td_.perCoreLoops, td_.perCorePerLoopElements, td_.perCoreLastLoopElements,
107- 107+ td_.lastCoreLoops, td_.lastCorePerLoopElements, td_.lastCoreLastLoopElements, td_.rstdFlag);
108- return ge::GRAPH_SUCCESS;108+ 
109-}109+ return ge::GRAPH_SUCCESS;
110- 110+}
111-ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::DoLibApiTiling()111+ 
112-{112+ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::DoLibApiTiling()
113- return ge::GRAPH_SUCCESS;113+{
114-}114+ return ge::GRAPH_SUCCESS;
115- 115+}
116-uint64_t AddRmsNormDynamicMxQuantReduceEmptyTiling::GetTilingKey() const116+ 
117-{117+uint64_t AddRmsNormDynamicMxQuantReduceEmptyTiling::GetTilingKey() const
118- return TILING_KEY_REDUCE_EMPTY;118+{
119-}119+ AddRmsNormDynamicMxQuantTilingKey tilingKey;
120- 120+ tilingKey.SetComputeMode(ComputeMode::REDUCE_EMPTY);
121-ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::PostTiling()121+ tilingKey.SetYDataType(YDataType::FP8);
122-{122+ return tilingKey.GetTilingKey();
123- OP_LOGD(context_->GetNodeName(), "ReduceEmpty PostTiling: usedCoreNum=%lu.", usedCoreNum_);123+}
124- context_->SetBlockDim(usedCoreNum_);124+ 
125- 125+ge::graphStatus AddRmsNormDynamicMxQuantReduceEmptyTiling::PostTiling()
126- auto rawTilingData = context_->GetRawTilingData();126+{
127- OP_CHECK_IF(sizeof(td_) > rawTilingData->GetCapacity(),127+ OP_LOGD(context_->GetNodeName(), "ReduceEmpty PostTiling: usedCoreNum=%lu.", usedCoreNum_);
128- OP_LOGE(context_->GetNodeName(), "actual tiling data size %zu > context tiling data size %zu",128+ context_->SetBlockDim(usedCoreNum_);
129- sizeof(td_), rawTilingData->GetCapacity()),129+ 
130- return ge::GRAPH_FAILED);130+ auto rawTilingData = context_->GetRawTilingData();
131- auto capSize = rawTilingData->GetCapacity();131+ OP_CHECK_IF(sizeof(td_) > rawTilingData->GetCapacity(),
132- void* ptrData = rawTilingData->GetData();132+ OP_LOGE(context_->GetNodeName(), "actual tiling data size %zu > context tiling data size %zu",
133- OP_CHECK_NULL_WITH_CONTEXT(context_, ptrData);133+ sizeof(td_), rawTilingData->GetCapacity()),
134- void* ptrStruct = static_cast<void*>(&td_);134+ return ge::GRAPH_FAILED);
135- OP_CHECK_NULL_WITH_CONTEXT(context_, ptrStruct);135+ auto capSize = rawTilingData->GetCapacity();
136- OP_CHECK_IF(memcpy_s(ptrData, capSize, ptrStruct, sizeof(td_)) != 0,136+ void* ptrData = rawTilingData->GetData();
137- OP_LOGE(context_->GetNodeName(), "Set tiling data is failed!"), return ge::GRAPH_FAILED);137+ OP_CHECK_NULL_WITH_CONTEXT(context_, ptrData);
138- rawTilingData->SetDataSize(sizeof(td_));138+ void* ptrStruct = static_cast<void*>(&td_);
139- 139+ OP_CHECK_NULL_WITH_CONTEXT(context_, ptrStruct);
140- size_t* currentWorkspace = context_->GetWorkspaceSizes(1);140+ OP_CHECK_IF(memcpy_s(ptrData, capSize, ptrStruct, sizeof(td_)) != 0,
141- OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);141+ OP_LOGE(context_->GetNodeName(), "Set tiling data is failed!"), return ge::GRAPH_FAILED);
142- currentWorkspace[0] = workspaceSize_;142+ rawTilingData->SetDataSize(sizeof(td_));
143- return ge::GRAPH_SUCCESS;143+ 
144-}144+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
145- 145+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
146-REGISTER_OPS_TILING_TEMPLATE(AddRmsNormDynamicMxQuant, AddRmsNormDynamicMxQuantReduceEmptyTiling, ARND_REDUCE_EMPTY_PRIORITY);146+ currentWorkspace[0] = workspaceSize_;
147-} // namespace optiling147+ return ge::GRAPH_SUCCESS;
148+}
149+ 
150+REGISTER_OPS_TILING_TEMPLATE(AddRmsNormDynamicMxQuant, AddRmsNormDynamicMxQuantReduceEmptyTiling, ARND_REDUCE_EMPTY_PRIORITY);
151+} // namespace optiling
Anorm/add_rms_norm_dynamic_mx_quant/op_host/add_rms_norm_dynamic_mx_quant_tiling_split_r_arch35.cpp+223-0
@@ -0,0 +1,223 @@
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 add_rms_norm_dynamic_mx_quant_tiling_split_r_arch35.cpp
13+ * \brief
14+ */
15+#include "add_rms_norm_dynamic_mx_quant_tiling.h"
16+#include "norm/norm_common/op_host/norm_tiling_check_common.h"
17+ 
18+using namespace optiling::add_rms_norm_dynamic_mx_quant;
19+ 
20+namespace optiling {
21+using namespace NormCheck;
22+ 
23+constexpr uint32_t RETAINED_SIZE_256 = 256;
24+ 
25+int64_t AddRmsNormDynamicMxQuantSplitRTiling::GetPowerSplit(uint64_t numN)
26+{
27+ uint64_t ubLoops = Ops::Base::CeilDiv(numCol_, numN);
28+ int64_t powerSplit = ubLoops == 0 ? 1 : (1L << (ULONG_BIT_LEN - 1 - __builtin_clzl(ubLoops)));
29+ powerSplit = (static_cast<uint64_t>(powerSplit) == ubLoops) ? powerSplit / NUM_TWO : powerSplit;
30+ if (ubLoops == 1) {
31+ powerSplit = 1;
32+ }
33+ return powerSplit;
34+}
35+ 
36+int64_t AddRmsNormDynamicMxQuantSplitRTiling::GetCacheId(int64_t idx)
37+{
38+ return __builtin_popcountll(idx ^ (idx + 1)) - 1;
39+}
40+ 
41+uint64_t AddRmsNormDynamicMxQuantSplitRTiling::GetMaxBaseN(uint64_t initialN)
42+{
43+ uint64_t vlfp32 = vecLengthFP32_;
44+ uint64_t binaryAddElemtMaxLen = vlfp32 * vlfp32 * NUM_TWO * NUM_TWO;
45+ 
46+ // Fixed buffers (not dependent on baseN)
47+ uint64_t rstdBufSize = Ops::Base::CeilAlign(baseM_ * FP32_SIZE, ubBlockSize_) * DOUBLE_BUFFER;
48+ uint64_t binaryAddBufSize = Ops::Base::CeilAlign(vlfp32 * FP32_SIZE * NUM_TWO, ubBlockSize_);
49+ 
50+ int64_t powerSplit = GetPowerSplit(initialN);
51+ uint64_t cacheBufSize = static_cast<uint64_t>(GetCacheId(powerSplit - 1) + 1) * ubBlockSize_;
52+ 
53+ while (powerSplit > 1 && NUM_TWO * initialN <= binaryAddElemtMaxLen) {
54+ uint64_t candidateN = NUM_TWO * initialN;
55+ uint64_t xInputBuf = NUM_TWO * candidateN * xDtypeSize_ * DOUBLE_BUFFER;
56+ uint64_t gammaBetaBuf = (1 + betaFlag_) * candidateN * gammaDtypeSize_ * DOUBLE_BUFFER;
57+ uint64_t yBuf = 0;
58+ if (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0) {
59+ yBuf = Ops::Base::CeilAlign(candidateN * FP8_SIZE, ubBlockSize_) * DOUBLE_BUFFER;
60+ } else if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {
61+ yBuf = Ops::Base::CeilAlign(candidateN / NUM_TWO, ubBlockSize_) * DOUBLE_BUFFER;
62+ }
63+ uint64_t xOutBuf = candidateN * xDtypeSize_ * DOUBLE_BUFFER;
64+ uint64_t mxScaleBuf = Ops::Base::CeilAlign(
65+ Ops::Base::CeilDiv(candidateN, static_cast<uint64_t>(MX_BLOCK_SIZE_32)) * sizeof(uint8_t), ubBlockSize_);
66+ 
67+ uint64_t xFp32Buf = Ops::Base::CeilAlign(candidateN * FP32_SIZE, ubBlockSize_);
68+ uint64_t yTmpBuf = Ops::Base::CeilAlign(candidateN * xDtypeSize_, ubBlockSize_);
69+ uint64_t maxExpBuf = Ops::Base::CeilAlign(
70+ Ops::Base::CeilDiv(candidateN, static_cast<uint64_t>(MX_BLOCK_SIZE_32)) * sizeof(uint16_t), ubBlockSize_);
71+ uint64_t halfScaleBuf = maxExpBuf;
72+ 
73+ uint64_t totalNDependent =
74+ xInputBuf + gammaBetaBuf + yBuf + xOutBuf + mxScaleBuf + xFp32Buf + yTmpBuf + maxExpBuf + halfScaleBuf;
75+ 
76+ if (totalNDependent > maxUbSize_ - RETAINED_SIZE_256 - (rstdBufSize + cacheBufSize + binaryAddBufSize)) {
77+ break;
78+ }
79+ 
80+ initialN = candidateN;
81+ powerSplit = GetPowerSplit(initialN);
82+ cacheBufSize = static_cast<uint64_t>(GetCacheId(powerSplit - 1) + 1) * ubBlockSize_;
83+ }
84+ return initialN;
85+}
86+ 
87+bool AddRmsNormDynamicMxQuantSplitRTiling::IsCapable()
88+{
89+ if (Y_SUPPORT_DTYPE_SET.count(yDtype_) == 0) {
90+ return false;
91+ }
92+ return numCol_ >= baseN_;
93+}
94+ 
95+ge::graphStatus AddRmsNormDynamicMxQuantSplitRTiling::DoOpTiling()
96+{
97+ OP_LOGD(context_->GetNodeName(), "Enter DoOpTiling for SplitR.");
98+ 
99+ // Multi-core split on A axis (same as full-load)
100+ mPerCore_ = Ops::Base::CeilDiv(numRow_, totalCoreNum_);
101+ usedCoreNum_ = Ops::Base::CeilDiv(numRow_, mPerCore_);
102+ mLastCore_ = numRow_ - (usedCoreNum_ - 1) * mPerCore_;
103+ blockFactor_ = mPerCore_;
104+ 
105+ // Maximize baseN within UB constraints , baseN_ is a multiple of 64
106+ baseN_ = GetMaxBaseN(baseN_);
107+ baseNBlockSize_ = baseN_ / MX_BLOCK_SIZE_32;
108+ 
109+ // R-axis split
110+ nUbLoops_ = Ops::Base::CeilDiv(numCol_, baseN_);
111+ 
112+ binAddQuotient_ = baseN_ == 0 ? 1 : (1UL << (ULONG_BIT_LEN - 1 - __builtin_clzl(baseN_)));
113+ binAddQuotient_ = (binAddQuotient_ == baseN_) ? binAddQuotient_ / NUM_TWO : binAddQuotient_;
114+ powerSplit_ = GetPowerSplit(baseN_);
115+ mainFoldCount_ = powerSplit_ * baseN_ > numCol_ ? 0 : (numCol_ - powerSplit_ * baseN_) / baseN_;
116+ 
117+ foldTail_ = numCol_ % baseN_;
118+ 
119+ OP_LOGI(
120+ context_->GetNodeName(),
121+ "SplitR: baseN=%lu, baseM=%lu, nUbLoops=%lu, powerSplit=%lu, "
122+ "mainFoldCount=%lu, foldTail=%lu, blockFactor=%lu.",
123+ baseN_, baseM_, nUbLoops_, powerSplit_, mainFoldCount_, foldTail_, blockFactor_);
124+ 
125+ SetTilingData();
126+ PrintTilingData();
127+ return ge::GRAPH_SUCCESS;
128+}
129+ 
130+void AddRmsNormDynamicMxQuantSplitRTiling::SetTilingData()
131+{
132+ // AddRmsNorm fields
133+ tilingData.numCol = numCol_;
134+ tilingData.numColAlign = numColAlign_;
135+ tilingData.blockFactor = blockFactor_;
136+ tilingData.mLastCore = mLastCore_;
137+ // SplitR specific fields
138+ tilingData.baseN = baseN_;
139+ tilingData.baseM = baseM_;
140+ tilingData.baseNBlockSize = baseNBlockSize_;
141+ tilingData.nUbLoops = nUbLoops_;
142+ tilingData.binAddQuotient = binAddQuotient_;
143+ tilingData.powerSplit = powerSplit_;
144+ tilingData.mainFoldCount = mainFoldCount_;
145+ tilingData.foldTail = foldTail_;
146+ tilingData.epsilon = epsilon_;
147+ tilingData.avgFactor = avgFactor_;
148+ // DynamicMxQuant fields
149+ tilingData.roundMode = roundMode_;
150+ tilingData.mxBlockSize = mxBlockSize_;
151+ tilingData.scaleAlg = scaleAlg_;
152+ tilingData.mxScaleSize = mxScaleSize_;
153+ // Flags
154+ tilingData.betaFlag = betaFlag_;
155+ tilingData.rstdFlag = rstdFlag_;
156+}
157+ 
158+void AddRmsNormDynamicMxQuantSplitRTiling::PrintTilingData()
159+{
160+ OP_LOGI(
161+ context_->GetNodeName(),
162+ "TilingData numCol: %lu, numColAlign: %lu, "
163+ "blockFactor: %lu, mLastCore: %lu, baseN: %lu, baseM: %lu, "
164+ "nUbLoops: %lu, binAddQuotient: %lu, powerSplit: %lu, "
165+ "mainFoldCount: %lu, foldTail: %lu, epsilon: %f, avgFactor: %f.",
166+ tilingData.numCol, tilingData.numColAlign, tilingData.blockFactor, tilingData.mLastCore,
167+ tilingData.baseN, tilingData.baseM, tilingData.nUbLoops, tilingData.binAddQuotient, tilingData.powerSplit,
168+ tilingData.mainFoldCount, tilingData.foldTail, tilingData.epsilon, tilingData.avgFactor);
169+ OP_LOGI(
170+ context_->GetNodeName(),
171+ "TilingData roundMode: %lu, mxBlockSize: %lu, scaleAlg: %ld, "
172+ "mxScaleSize: %lu, betaFlag: %u, rstdFlag: %u.",
173+ tilingData.roundMode, tilingData.mxBlockSize, tilingData.scaleAlg,
174+ tilingData.mxScaleSize, tilingData.betaFlag, tilingData.rstdFlag);
175+}
176+ 
177+ge::graphStatus AddRmsNormDynamicMxQuantSplitRTiling::DoLibApiTiling()
178+{
179+ return ge::GRAPH_SUCCESS;
180+}
181+ 
182+ge::graphStatus AddRmsNormDynamicMxQuantSplitRTiling::PostTiling()
183+{
184+ OP_LOGD(context_->GetNodeName(), "Tiling usedCoreNum is %lu.", usedCoreNum_);
185+ context_->SetBlockDim(usedCoreNum_);
186+ 
187+ auto rawTilingData = context_->GetRawTilingData();
188+ OP_CHECK_IF(
189+ sizeof(tilingData) > rawTilingData->GetCapacity(),
190+ OP_LOGE(
191+ context_->GetNodeName(), "actual tiling data size %zu > context tiling data size %zu", sizeof(tilingData),
192+ rawTilingData->GetCapacity()),
193+ return ge::GRAPH_FAILED);
194+ auto capSize = rawTilingData->GetCapacity();
195+ void* ptrData = rawTilingData->GetData();
196+ OP_CHECK_NULL_WITH_CONTEXT(context_, ptrData);
197+ void* ptrStruct = static_cast<void*>(&tilingData);
198+ OP_CHECK_NULL_WITH_CONTEXT(context_, ptrStruct);
199+ OP_CHECK_IF(
200+ memcpy_s(ptrData, capSize, ptrStruct, sizeof(tilingData)) != 0,
201+ OP_LOGE(context_->GetNodeName(), "Set tiling data is failed!"), return ge::GRAPH_FAILED);
202+ rawTilingData->SetDataSize(sizeof(tilingData));
203+ 
204+ size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
205+ OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
206+ currentWorkspace[0] = workspaceSize_;
207+ return ge::GRAPH_SUCCESS;
208+}
209+ 
210+uint64_t AddRmsNormDynamicMxQuantSplitRTiling::GetTilingKey() const
211+{
212+ AddRmsNormDynamicMxQuantTilingKey tilingKey;
213+ tilingKey.SetComputeMode(ComputeMode::SPLIT_R);
214+ if (Y_SUPPORT_DTYPE_FP8_SET.count(yDtype_) != 0) {
215+ tilingKey.SetYDataType(YDataType::FP8);
216+ } else if (Y_SUPPORT_DTYPE_FP4_SET.count(yDtype_) != 0) {
217+ tilingKey.SetYDataType(YDataType::FP4);
218+ }
219+ return tilingKey.GetTilingKey();
220+}
221+ 
222+REGISTER_OPS_TILING_TEMPLATE(AddRmsNormDynamicMxQuant, AddRmsNormDynamicMxQuantSplitRTiling, ARND_SPLIT_R_PRIORITY);
223+} // namespace optiling
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/add_rms_norm_dynamic_mx_quant_apt.cpp+21-9
@@ -15,16 +15,20 @@
15#include "arch35/add_rms_norm_dynamic_mx_quant_fp8_r_full_load.h"15#include "arch35/add_rms_norm_dynamic_mx_quant_fp8_r_full_load.h"
16#include "arch35/add_rms_norm_dynamic_mx_quant_fp4_r_full_load.h"16#include "arch35/add_rms_norm_dynamic_mx_quant_fp4_r_full_load.h"
17#include "arch35/add_rms_norm_dynamic_mx_quant_reduce_empty.h"17#include "arch35/add_rms_norm_dynamic_mx_quant_reduce_empty.h"
18+#include "arch35/add_rms_norm_dynamic_mx_quant_fp8_split_r.h"
19+#include "arch35/add_rms_norm_dynamic_mx_quant_fp4_split_r.h"
18#include "arch35/add_rms_norm_dynamic_mx_quant_tiling_data.h"20#include "arch35/add_rms_norm_dynamic_mx_quant_tiling_data.h"
21+#include "arch35/add_rms_norm_dynamic_mx_quant_tiling_key.h"
19 22 
20using namespace AscendC;23using namespace AscendC;
21using namespace AddRmsNormDynamicMxQuant;24using namespace AddRmsNormDynamicMxQuant;
22 25 
23-#define TILING_KEY_FP8_R_FULL_LOAD 10026+#define FLOAT_OVERFLOW_MODE_CTRL 60
Z
Zzzz20484月20日

FLOAT_OVERFLOW_MODE_CTRL 在多个被 include 的头文件（common.h:17, fp4_split_r.h:17, fp8_split_r.h:17）中已定义。本文件不需要再次定义，可直接使用被引入的定义。建议删除此行避免编译警告。

likedislike
24-#define TILING_KEY_FP4_R_FULL_LOAD 101
25-#define TILING_KEY_REDUCE_EMPTY 300
26 27 
27-extern "C" __global__ __aicore__ void add_rms_norm_dynamic_mx_quant(28+REGISTER_TILING_DEFAULT(AddRmsNormDynamicMxQuantTilingData);
29+ 
30+template <int8_t COMPUTE_MODE, int8_t Y_DATA_TYPE>
31+__global__ __aicore__ void add_rms_norm_dynamic_mx_quant(
28 GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta,32 GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta,
29 GM_ADDR y, GM_ADDR x, GM_ADDR mxscale, GM_ADDR rstd,33 GM_ADDR y, GM_ADDR x, GM_ADDR mxscale, GM_ADDR rstd,
30 GM_ADDR workspace, GM_ADDR tiling)34 GM_ADDR workspace, GM_ADDR tiling)
@@ -34,29 +38,37 @@ extern "C" __global__ __aicore__ void add_rms_norm_dynamic_mx_quant(
34 }38 }
35 39
36 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);40 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
37- REGISTER_TILING_DEFAULT(AddRmsNormDynamicMxQuantTilingData);
38- REGISTER_TILING_FOR_TILINGKEY("TILING_KEY_VAR == 300", AddRmsNormDynamicMxQuantReduceEmptyTilingData);
39 TPipe pipe;41 TPipe pipe;
40 42 
41#if (__NPU_ARCH__ == 3510)43#if (__NPU_ARCH__ == 3510)
42 int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();44 int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();
43#endif45#endif
44 46 
45- if (TILING_KEY_IS(TILING_KEY_REDUCE_EMPTY)) {47+ if constexpr (COMPUTE_MODE == COMPUTE_MODE_REDUCE_EMPTY) {
46 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantReduceEmptyTilingData, tilingDataIn, tiling);48 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantReduceEmptyTilingData, tilingDataIn, tiling);
47 AddRmsNormDynamicMxQuantReduceEmpty op(&tilingDataIn);49 AddRmsNormDynamicMxQuantReduceEmpty op(&tilingDataIn);
48 op.Init(rstd);50 op.Init(rstd);
49 op.Process();51 op.Process();
50- } else if (TILING_KEY_IS(TILING_KEY_FP8_R_FULL_LOAD)) {52+ } else if constexpr (COMPUTE_MODE == COMPUTE_MODE_FULL_LOAD && Y_DATA_TYPE == Y_DATA_TYPE_FP8) {
51 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantTilingData, tilingDataIn, tiling);53 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantTilingData, tilingDataIn, tiling);
52 AddRmsNormDynamicMxQuantFP8RFullLoad<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);54 AddRmsNormDynamicMxQuantFP8RFullLoad<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);
53 op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);55 op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);
54 op.Process();56 op.Process();
55- } else if (TILING_KEY_IS(TILING_KEY_FP4_R_FULL_LOAD)) {57+ } else if constexpr (COMPUTE_MODE == COMPUTE_MODE_FULL_LOAD && Y_DATA_TYPE == Y_DATA_TYPE_FP4) {
56 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantTilingData, tilingDataIn, tiling);58 GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantTilingData, tilingDataIn, tiling);
57 AddRmsNormDynamicMxQuantFP4RFullLoad<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);59 AddRmsNormDynamicMxQuantFP4RFullLoad<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);
58 op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);60 op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);
59 op.Process();61 op.Process();
62+ } else if constexpr (COMPUTE_MODE == COMPUTE_MODE_SPLIT_R && Y_DATA_TYPE == Y_DATA_TYPE_FP8) {
63+ GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantSplitRTilingData, tilingDataIn, tiling);
64+ AddRmsNormDynamicMxQuantFP8SplitR<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);
65+ op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);
66+ op.Process();
67+ } else if constexpr (COMPUTE_MODE == COMPUTE_MODE_SPLIT_R && Y_DATA_TYPE == Y_DATA_TYPE_FP4) {
68+ GET_TILING_DATA_WITH_STRUCT(AddRmsNormDynamicMxQuantSplitRTilingData, tilingDataIn, tiling);
69+ AddRmsNormDynamicMxQuantFP4SplitR<DTYPE_X1, DTYPE_GAMMA, DTYPE_Y> op(&pipe);
70+ op.Init(x1, x2, gamma, beta, y, x, mxscale, mxscale, rstd, &tilingDataIn);
71+ op.Process();
60 }72 }
61 73 
62#if (__NPU_ARCH__ == 3510)74#if (__NPU_ARCH__ == 3510)
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_common.h+743-0
@@ -14,6 +14,7 @@
14 */14 */
15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H
16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H
17+#define FLOAT_OVERFLOW_MODE_CTRL 60
17 18 
18#include "kernel_operator.h"19#include "kernel_operator.h"
19#include "kernel_tiling/kernel_tiling.h"20#include "kernel_tiling/kernel_tiling.h"
@@ -78,6 +79,8 @@ constexpr int32_t NEG_ZERO = 0x80000000;
78constexpr float ONE_FOURTH = 0.25;79constexpr float ONE_FOURTH = 0.25;
79constexpr int32_t FP32_BIAS_NEG = -127;80constexpr int32_t FP32_BIAS_NEG = -127;
80constexpr int32_t FP32_BIAS = 127;81constexpr int32_t FP32_BIAS = 127;
82+constexpr int32_t NEG_ONE = -1;
83+constexpr float FOUR = 4.0;
81constexpr uint16_t FP4_E2M1_BF16_MAX_EXP = 0x0100;84constexpr uint16_t FP4_E2M1_BF16_MAX_EXP = 0x0100;
82constexpr uint16_t FP4_E1M2_MAX_EXP = 0x0000;85constexpr uint16_t FP4_E1M2_MAX_EXP = 0x0000;
83constexpr int64_t MODE_ROUND = 0;86constexpr int64_t MODE_ROUND = 0;
@@ -85,9 +88,14 @@ constexpr int64_t MODE_FLOOR = 1;
85constexpr int64_t MODE_RINT = 4;88constexpr int64_t MODE_RINT = 4;
86 89 
87constexpr static uint32_t VL_F32 = platform::GetVRegSize() / sizeof(float); // 6490constexpr static uint32_t VL_F32 = platform::GetVRegSize() / sizeof(float); // 64
91+constexpr static uint32_t VL_B16 = platform::GetVRegSize() / sizeof(half); // 128
92+constexpr static uint16_t ELEMENT_AFTER_REDUCE = platform::GetVRegSize() / Ops::Base::GetUbBlockSize(); // 8
88constexpr static uint32_t BLOCK_F32_ALIGN_NUM = Ops::Base::GetUbBlockSize() / sizeof(float); // 893constexpr static uint32_t BLOCK_F32_ALIGN_NUM = Ops::Base::GetUbBlockSize() / sizeof(float); // 8
89constexpr static uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize();94constexpr static uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize();
90 95 
96+constexpr int64_t AR_RECOMPUTE_SUM_BUFFER_BYTES = 32;
97+constexpr int64_t AR_RECOMPUTE_SUM_LEN = AR_RECOMPUTE_SUM_BUFFER_BYTES / sizeof(float);
98+ 
91constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {99constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {
92 AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,100 AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
93 AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN,101 AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN,
@@ -114,6 +122,11 @@ __aicore__ inline T Min(T a, T b)
114 return a > b ? b : a;122 return a > b ? b : a;
115}123}
116 124 
125+__aicore__ inline int64_t GetCacheId(const int64_t idx)
126+{
127+ return AscendC::ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1;
128+}
129+ 
117template <typename T_IN>130template <typename T_IN>
118__aicore__ inline void LoadTensorForDtypeTIn(131__aicore__ inline void LoadTensorForDtypeTIn(
119 __local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)132 __local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
@@ -140,5 +153,735 @@ __aicore__ inline void StoreTensorForDtypeTOut(
140 }153 }
141}154}
142 155 
156+// ===================== Rstd Computation =====================
157+ 
158+__aicore__ inline void CalculateRstd(
159+ LocalTensor<float>& xReduceLocal, LocalTensor<float>& rstdLocal,
160+ uint32_t curRows, float avgFactor, float epsilon)
161+{
162+ static constexpr float POS_INF = 3.40282366920938E+38;
163+ static constexpr float SCALAR1 = -0.5;
164+ static constexpr float SCALAR2 = 1.5;
165+ static constexpr float SCALAR3 = 0.5;
166+ static constexpr float SCALAR0 = -99.99;
167+ 
168+ __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr();
169+ __local_mem__ float* xReduceUb = (__local_mem__ float*)xReduceLocal.GetPhyAddr();
170+ uint16_t loopRows = static_cast<uint16_t>((curRows + VL_F32 - 1) / VL_F32);
171+ __VEC_SCOPE__
172+ {
173+ RegTensor<float> var, rstd, r, y, s, t, one, scalar1;
174+ RegTensor<float> t1, t2, t3, t4, scalarInf, scalarZero;
175+ MaskReg cmpRegZero, cmpRegInf;
176+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
177+ MaskReg pregLoop;
178+ 
179+ uint32_t sreg = static_cast<uint32_t>(curRows);
180+ for (uint16_t i = 0; i < loopRows; ++i) {
181+ pregLoop = UpdateMask<float>(sreg);
182+ AscendC::MicroAPI::Duplicate(scalarInf, POS_INF, pregLoop);
183+ AscendC::MicroAPI::Duplicate(scalarZero, float(0.0), pregLoop);
184+ AscendC::MicroAPI::Duplicate(one, float(1.0), pregLoop);
185+ AscendC::MicroAPI::Duplicate(scalar1, SCALAR3, pregLoop);
186+ AscendC::MicroAPI::Duplicate(t1, SCALAR2, pregLoop);
187+ AscendC::MicroAPI::Duplicate(s, float(1.0), pregLoop);
188+ // rstd = 1/sqrt(mean(x^2) + eps)
189+ AscendC::MicroAPI::DataCopy(var, xReduceUb + i * VL_F32);
190+ AscendC::MicroAPI::Muls(var, var, avgFactor, pregLoop);
191+ AscendC::MicroAPI::Adds(var, var, epsilon, pregLoop);
192+ AscendC::MicroAPI::Maxs(var, var, SCALAR0, pregLoop);
193+ AscendC::MicroAPI::Div(r, one, var, pregLoop);
194+ AscendC::MicroAPI::Sqrt(y, r, pregLoop);
195+ AscendC::MicroAPI::Muls(t, var, SCALAR1, pregLoop);
196+ AscendC::MicroAPI::Mul(t, t, y, pregLoop);
197+ AscendC::MicroAPI::Mula(t1, t, y, pregLoop);
198+ AscendC::MicroAPI::Mul(rstd, y, t1, pregLoop);
199+ AscendC::MicroAPI::Muls(t3, var, float(-1.0), pregLoop);
200+ AscendC::MicroAPI::Mula(s, t3, r, pregLoop);
201+ AscendC::MicroAPI::Muls(t4, rstd, float(-1.0), pregLoop);
202+ AscendC::MicroAPI::Mula(r, t4, rstd, pregLoop);
203+ AscendC::MicroAPI::Mula(s, var, r, pregLoop);
204+ AscendC::MicroAPI::Mul(s, s, rstd, pregLoop);
205+ AscendC::MicroAPI::Mula(rstd, s, scalar1, pregLoop);
206+ AscendC::MicroAPI::CompareScalar(cmpRegZero, var, POS_INF, pregLoop);
207+ AscendC::MicroAPI::Select(rstd, scalarZero, rstd, cmpRegZero);
208+ AscendC::MicroAPI::CompareScalar(cmpRegInf, var, float(0.0), pregLoop);
209+ AscendC::MicroAPI::Select(rstd, scalarInf, rstd, cmpRegInf);
210+ AscendC::MicroAPI::DataCopy(rstdInUb + i * VL_F32, rstd, pregLoop);
211+ }
212+ }
213+}
214+ 
215+// ===================== Split-R Common Free Functions =====================
216+ 
217+// (x1 + x2)² → xFp32Tmp
218+template <typename T_X>
219+__aicore__ inline void MainBlockSquareVF(
220+ LocalTensor<T_X>& x1Local, LocalTensor<T_X>& x2Local,
221+ LocalTensor<float>& xFp32Tmp, uint32_t count)
222+{
223+ __local_mem__ T_X* x1InUb = (__local_mem__ T_X*)x1Local.GetPhyAddr();
224+ __local_mem__ T_X* x2InUb = (__local_mem__ T_X*)x2Local.GetPhyAddr();
225+ __local_mem__ float* xFp32TmpBuf = (__local_mem__ float*)xFp32Tmp.GetPhyAddr();
226+ 
227+ uint16_t loops = (count + VL_F32 - 1) / VL_F32;
228+ uint32_t sreg = count;
229+ __VEC_SCOPE__
230+ {
231+ RegTensor<float> x1Reg, x2Reg, xSum;
232+ MaskReg pregLoop;
233+ for (uint16_t vi = 0; vi < loops; ++vi) {
234+ uint32_t offset = vi * VL_F32;
235+ pregLoop = UpdateMask<float>(sreg);
236+ LoadTensorForDtypeTIn<T_X>(x1InUb, x1Reg, pregLoop, offset);
237+ LoadTensorForDtypeTIn<T_X>(x2InUb, x2Reg, pregLoop, offset);
238+ AscendC::MicroAPI::Add(xSum, x1Reg, x2Reg, pregLoop);
239+ AscendC::MicroAPI::Mul(xSum, xSum, xSum, pregLoop);
240+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_NORM_B32>(
241+ xFp32TmpBuf + offset, xSum, pregLoop);
242+ }
243+ }
244+}
245+ 
246+// xFp32Tmp += (x1Fold + x2Fold)²
247+template <typename T_X>
248+__aicore__ inline void FoldBlockSquareAddVF(
249+ LocalTensor<T_X>& x1FoldLocal, LocalTensor<T_X>& x2FoldLocal,
250+ LocalTensor<float>& xFp32Tmp, uint32_t tailCount)
251+{
252+ __local_mem__ T_X* x1FoldInUb = (__local_mem__ T_X*)x1FoldLocal.GetPhyAddr();
253+ __local_mem__ T_X* x2FoldInUb = (__local_mem__ T_X*)x2FoldLocal.GetPhyAddr();
254+ __local_mem__ float* xFp32TmpBuf = (__local_mem__ float*)xFp32Tmp.GetPhyAddr();
255+ uint16_t tailLoops = (tailCount + VL_F32 - 1) / VL_F32;
256+ uint32_t sregTail = tailCount;
257+ __VEC_SCOPE__
258+ {
259+ RegTensor<float> x1FoldReg, x2FoldReg, foldSquare, mainReg, sum;
260+ MaskReg pregLoop;
261+ for (uint16_t i = 0; i < tailLoops; ++i) {
262+ pregLoop = UpdateMask<float>(sregTail);
263+ uint32_t offset = i * VL_F32;
264+ // Fold tile: (x1Fold + x2Fold)²
265+ LoadTensorForDtypeTIn<T_X>(x1FoldInUb, x1FoldReg, pregLoop, offset);
266+ LoadTensorForDtypeTIn<T_X>(x2FoldInUb, x2FoldReg, pregLoop, offset);
267+ AscendC::MicroAPI::Add(x1FoldReg, x1FoldReg, x2FoldReg, pregLoop);
268+ AscendC::MicroAPI::Mul(foldSquare, x1FoldReg, x1FoldReg, pregLoop);
269+ // Read back main block result and accumulate
270+ AscendC::MicroAPI::DataCopy(mainReg, xFp32TmpBuf + offset);
271+ AscendC::MicroAPI::Add(sum, mainReg, foldSquare, pregLoop);
272+ AscendC::MicroAPI::Select(sum, sum, mainReg, pregLoop);
273+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_NORM_B32>(
274+ xFp32TmpBuf + offset, sum, pregLoop);
275+ }
276+ }
277+}
278+ 
279+__aicore__ inline void UpdateCache(
280+ const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,
281+ const int64_t cacheId, const int64_t stride)
282+{
283+ uint16_t innerLoopTimes = cacheId;
284+ uint32_t innerLoopStride = stride;
285+ __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();
286+ __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride;
287+ __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();
288+ 
289+ __VEC_SCOPE__
290+ {
291+ RegTensor<float> aReg, bReg;
292+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
293+ 
294+ DataCopy(aReg, (__local_mem__ float*)src);
295+ for (uint16_t j = 0; j < innerLoopTimes; ++j) {
296+ DataCopy(bReg, dst + j * innerLoopStride);
297+ AscendC::MicroAPI::Add(aReg, aReg, bReg, pregOne);
298+ }
299+ DataCopy((__local_mem__ float*)cache, aReg, pregOne);
300+ }
301+}
302+ 
303+template <typename T_X>
304+__aicore__ inline void CalculateXAdd(
305+ LocalTensor<T_X>& xLocal1, LocalTensor<T_X>& xLocal2,
306+ LocalTensor<T_X>& xOutLocal, LocalTensor<float>& xFp32Local, uint32_t count)
307+{
308+ __local_mem__ T_X* x1InUb = (__local_mem__ T_X*)xLocal1.GetPhyAddr();
309+ __local_mem__ T_X* x2InUb = (__local_mem__ T_X*)xLocal2.GetPhyAddr();
310+ __local_mem__ T_X* xOutInUb = (__local_mem__ T_X*)xOutLocal.GetPhyAddr();
311+ __local_mem__ float* xFp32Tmp = (__local_mem__ float*)xFp32Local.GetPhyAddr();
312+ 
313+ uint32_t sreg = count;
314+ uint16_t loopCount = (sreg + VL_F32 - 1) / VL_F32;
315+ 
316+ __VEC_SCOPE__
317+ {
318+ RegTensor<float> x1, x2, xSum;
319+ MaskReg pregLoop;
320+ for (uint16_t i = 0; i < loopCount; ++i) {
321+ uint32_t offset = i * VL_F32;
322+ pregLoop = UpdateMask<float>(sreg);
323+ LoadTensorForDtypeTIn<T_X>(x1InUb, x1, pregLoop, offset);
324+ LoadTensorForDtypeTIn<T_X>(x2InUb, x2, pregLoop, offset);
325+ AscendC::MicroAPI::Add(xSum, x1, x2, pregLoop);
326+ StoreTensorForDtypeTOut<T_X>(xOutInUb, xSum, pregLoop, offset);
327+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_NORM_B32>(xFp32Tmp + offset, xSum, pregLoop);
328+ }
329+ }
330+}
331+template <typename T_X>
332+__aicore__ inline void MxQuantComputeMaxExpOCP(
333+ __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint16_t loopNum)
334+{
335+ __VEC_SCOPE__
336+ {
337+ AscendC::MicroAPI::RegTensor<T_X> vdExp0;
338+ AscendC::MicroAPI::RegTensor<T_X> vdExp1;
339+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
340+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
341+ AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect0;
342+ AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect1;
343+ AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract0;
344+ AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract1;
345+ AscendC::MicroAPI::RegTensor<uint16_t> expMaskBF16;
346+ AscendC::MicroAPI::Duplicate(expMaskBF16, MAX_EXP_FOR_BF16);
347+ AscendC::MicroAPI::RegTensor<uint16_t> invalidMaskFP16;
348+ AscendC::MicroAPI::Duplicate(invalidMaskFP16, INVALID_FLOAT16);
349+ AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
350+ AscendC::MicroAPI::MaskReg Mask = AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
351+ AscendC::MicroAPI::MaskReg invalidDataMask0;
352+ AscendC::MicroAPI::MaskReg invalidDataMask1;
353+ AscendC::MicroAPI::UnalignReg u1;
354+ static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
355+ AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
356+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
357+ for (uint16_t i = 0; i < loopNum; i++) {
358+ AscendC::MicroAPI::LoadAlign<
359+ T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
360+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(
361+ vdExp0, vdExp1, srcAddr, VL_B16 * DIGIT_TWO);
362+ if constexpr (IsSame<T_X, half>::value) {
363+ AscendC::MicroAPI::And(
364+ vdExpSelect0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, invalidMaskFP16, Mask);
365+ AscendC::MicroAPI::And(
366+ vdExpSelect1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, invalidMaskFP16, Mask);
367+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
368+ invalidDataMask0, vdExpSelect0, invalidMaskFP16, Mask);
369+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
370+ invalidDataMask1, vdExpSelect1, invalidMaskFP16, Mask);
371+ AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp0BF16, vdExp0, Mask);
372+ AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp1BF16, vdExp1, Mask);
373+ AscendC::MicroAPI::And(
374+ vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0BF16, expMaskBF16, Mask);
375+ AscendC::MicroAPI::And(
376+ vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1BF16, expMaskBF16, Mask);
377+ AscendC::MicroAPI::Select<uint16_t>(vdExpExtract0, vdExpExtract0, expMaskBF16, invalidDataMask0);
378+ AscendC::MicroAPI::Select<uint16_t>(vdExpExtract1, vdExpExtract1, expMaskBF16, invalidDataMask1);
379+ } else {
380+ AscendC::MicroAPI::And(
381+ vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, expMaskBF16, Mask);
382+ AscendC::MicroAPI::And(
383+ vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, expMaskBF16, Mask);
384+ }
385+ AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, Mask);
386+ AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, Mask);
387+ AscendC::MicroAPI::StoreUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
388+ maxExpAddr, vdMaxExp, u1, ELEMENT_AFTER_REDUCE);
389+ }
390+ AscendC::MicroAPI::StoreUnAlignPost(maxExpAddr, u1, 0);
391+ }
392+}
393+ 
394+template <typename T_Y>
395+__aicore__ inline void MxQuantComputeScaleOCP(
396+ __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr,
397+ __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB,
398+ uint16_t loopNumScale)
399+{
400+ uint16_t emax;
401+ if constexpr (IsSame<T_Y, fp8_e4m3fn_t>::value) {
402+ emax = MX_FP8_E4M3_MAX_EXP;
403+ } else if constexpr (IsSame<T_Y, fp8_e5m2_t>::value) {
404+ emax = MX_FP8_E5M2_MAX_EXP;
405+ } else if constexpr (IsSame<T_Y, fp4x2_e2m1_t>::value) {
406+ emax = FP4_E2M1_BF16_MAX_EXP;
407+ } else {
408+ emax = FP4_E1M2_MAX_EXP;
409+ }
410+ __VEC_SCOPE__
411+ {
412+ AscendC::MicroAPI::RegTensor<uint16_t> expMask;
413+ AscendC::MicroAPI::Duplicate(expMask, MAX_EXP_FOR_BF16);
414+ AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
415+ AscendC::MicroAPI::MaskReg cmpResult;
416+ AscendC::MicroAPI::MaskReg zeroMask;
417+ AscendC::MicroAPI::MaskReg preMaskScale;
418+ AscendC::MicroAPI::RegTensor<uint16_t> maxExpValue;
419+ AscendC::MicroAPI::Duplicate(maxExpValue, emax);
420+ AscendC::MicroAPI::RegTensor<uint16_t> sharedExp;
421+ AscendC::MicroAPI::RegTensor<uint16_t> scaleValue;
422+ AscendC::MicroAPI::RegTensor<uint16_t> scaleBias;
423+ AscendC::MicroAPI::Duplicate(scaleBias, BF16_EXP_BIAS);
424+ AscendC::MicroAPI::RegTensor<uint16_t> halfScale;
425+ AscendC::MicroAPI::RegTensor<uint16_t> fp8NanRegTensor;
426+ AscendC::MicroAPI::Duplicate(fp8NanRegTensor, MAX_EXP_FOR_FP8);
427+ AscendC::MicroAPI::RegTensor<uint16_t> zeroRegTensor;
428+ AscendC::MicroAPI::Duplicate(zeroRegTensor, 0);
429+ AscendC::MicroAPI::RegTensor<uint16_t> nanRegTensor;
430+ AscendC::MicroAPI::Duplicate(nanRegTensor, NAN_CUSTOMIZATION);
431+ AscendC::MicroAPI::MaskReg invalidDataMask;
432+ AscendC::MicroAPI::MaskReg specialDataMask;
433+ AscendC::MicroAPI::RegTensor<uint16_t> specialExpRegTensor;
434+ AscendC::MicroAPI::Duplicate(specialExpRegTensor, SPECIAL_EXP_THRESHOLD);
435+ for (uint16_t i = 0; i < loopNumScale; i++) {
436+ preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);
437+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
438+ vdMaxExp, maxExpAddr, VL_B16);
439+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(cmpResult, vdMaxExp, expMask, preMaskScale);
440+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(zeroMask, vdMaxExp, zeroRegTensor, preMaskScale);
441+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::LE>(invalidDataMask, vdMaxExp, maxExpValue, preMaskScale);
442+ AscendC::MicroAPI::Select<uint16_t>(vdMaxExp, maxExpValue, vdMaxExp, invalidDataMask);
443+ AscendC::MicroAPI::Sub(sharedExp, vdMaxExp, maxExpValue, preMaskScale);
444+ AscendC::MicroAPI::ShiftRights(scaleValue, sharedExp, SHR_NUM_FOR_BF16, preMaskScale);
445+ AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, fp8NanRegTensor, cmpResult);
446+ AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, zeroRegTensor, zeroMask);
447+ AscendC::MicroAPI::StoreAlign<
448+ uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
449+ AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
450+ mxScaleLocalAddr, scaleValue, VL_F32, preMaskScale);
451+ AscendC::MicroAPI::Compare<uint16_t, CMPMODE::EQ>(specialDataMask, sharedExp, scaleBias, preMaskScale);
452+ AscendC::MicroAPI::Sub(halfScale, scaleBias, sharedExp, preMaskScale);
453+ AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, nanRegTensor, cmpResult);
454+ AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, zeroRegTensor, zeroMask);
455+ AscendC::MicroAPI::Select<uint16_t>(halfScale, specialExpRegTensor, halfScale, specialDataMask);
456+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
457+ halfScaleLocalAddr, halfScale, VL_B16, preMaskScale);
458+ }
459+ }
460+}
461+ 
462+constexpr uint32_t CUBLAS_ZERO_FOR_ALL = 0x00000000;
463+constexpr uint32_t CUBLAS_EXP254 = 0x000000fe;
464+constexpr uint32_t CUBLAS_HALF_FOR_MAN = 0x00400000;
465+ 
466+template <typename T_X>
467+__aicore__ inline void MxQuantComputeMaxExpcuBLAS(
468+ __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint16_t loopNum)
469+{
470+ __VEC_SCOPE__
471+ {
472+ AscendC::MicroAPI::RegTensor<T_X> vdExp0;
473+ AscendC::MicroAPI::RegTensor<T_X> vdExp1;
474+ AscendC::MicroAPI::RegTensor<uint16_t> absMask16Bit;
475+ AscendC::MicroAPI::Duplicate(absMask16Bit, ABS_MASK_FOR_16BIT);
476+ AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
477+ AscendC::MicroAPI::MaskReg Mask = AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
478+ AscendC::MicroAPI::UnalignReg u1;
479+ for (uint16_t i = 0; i < loopNum; i++) {
480+ AscendC::MicroAPI::LoadAlign<
481+ T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
482+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, VL_B16 * DIGIT_TWO);
483+ AscendC::MicroAPI::And(
484+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
485+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, absMask16Bit, Mask);
486+ AscendC::MicroAPI::And(
487+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1,
488+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, absMask16Bit, Mask);
489+ AscendC::MicroAPI::Max(
490+ vdMaxExp, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
491+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, Mask);
492+ AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, Mask);
493+ AscendC::MicroAPI::StoreUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
494+ maxExpAddr, vdMaxExp, u1, ELEMENT_AFTER_REDUCE);
495+ }
496+ AscendC::MicroAPI::StoreUnAlignPost(maxExpAddr, u1, 0);
497+ }
498+}
499+ 
500+template <typename T_X, typename T_Y>
501+__aicore__ inline void MxQuantComputeScalecuBLAS(
502+ __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr,
503+ __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB,
504+ uint16_t loopNumScale4NV)
505+{
506+ uint32_t dtypeMax;
507+ if constexpr (IsSame<T_Y, fp8_e4m3fn_t>::value) {
508+ dtypeMax = MX_FP8_E4M3_MAX;
509+ } else {
510+ dtypeMax = MX_FP8_E5M2_MAX;
511+ }
512+ __VEC_SCOPE__
513+ {
514+ AscendC::MicroAPI::RegTensor<uint16_t> max16;
515+ AscendC::MicroAPI::RegTensor<uint32_t> max32;
516+ AscendC::MicroAPI::RegTensor<uint32_t> exp32;
517+ AscendC::MicroAPI::RegTensor<uint32_t> man32;
518+ AscendC::MicroAPI::RegTensor<uint32_t> normalExp32;
519+ AscendC::MicroAPI::RegTensor<uint32_t> expAddOne32;
520+ AscendC::MicroAPI::RegTensor<uint32_t> extractExp;
521+ AscendC::MicroAPI::RegTensor<uint16_t> expOut;
522+ AscendC::MicroAPI::RegTensor<uint32_t> halfScale;
523+ AscendC::MicroAPI::RegTensor<uint16_t> recExpOut;
524+ AscendC::MicroAPI::RegTensor<uint32_t> invMax;
525+ AscendC::MicroAPI::Duplicate(invMax, dtypeMax);
526+ AscendC::MicroAPI::RegTensor<uint32_t> manMaskFP32;
527+ AscendC::MicroAPI::Duplicate(manMaskFP32, MAN_MASK_FLOAT);
528+ AscendC::MicroAPI::RegTensor<uint32_t> expMask;
529+ AscendC::MicroAPI::Duplicate(expMask, MAX_EXP_FOR_FP32);
530+ AscendC::MicroAPI::RegTensor<uint32_t> zeroRegTensor32;
531+ AscendC::MicroAPI::Duplicate(zeroRegTensor32, 0);
532+ AscendC::MicroAPI::RegTensor<uint32_t> scaleBias;
533+ AscendC::MicroAPI::Duplicate(scaleBias, FP32_EXP_BIAS_CUBLAS);
534+ AscendC::MicroAPI::RegTensor<uint32_t> nanRegTensor;
535+ AscendC::MicroAPI::Duplicate(nanRegTensor, NAN_CUSTOMIZATION_PACK);
536+ AscendC::MicroAPI::RegTensor<uint32_t> fp8NanRegTensor;
537+ AscendC::MicroAPI::Duplicate(fp8NanRegTensor, MAX_EXP_FOR_FP8_IN_FP32);
538+ AscendC::MicroAPI::MaskReg cmpResult;
539+ AscendC::MicroAPI::MaskReg zeroMask;
540+ AscendC::MicroAPI::MaskReg p0;
541+ AscendC::MicroAPI::MaskReg p1;
542+ AscendC::MicroAPI::MaskReg p2;
543+ AscendC::MicroAPI::MaskReg preMaskScale;
544+ AscendC::MicroAPI::MaskReg maskHalf;
545+ uint32_t SixtyFour = 64;
546+ uint32_t ThirtyTwo = 32;
547+ preMaskScale = AscendC::MicroAPI::CreateMask<uint32_t>();
548+ maskHalf = AscendC::MicroAPI::UpdateMask<uint16_t>(SixtyFour);
549+ static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Float = {
550+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
551+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
552+ for (uint16_t i = 0; i < loopNumScale4NV; i++) {
553+ AscendC::MicroAPI::LoadAlign<
554+ uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
555+ AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(max16, maxExpAddr, VL_F32);
556+ AscendC::MicroAPI::Cast<float, T_X, castTraitHalf2Float>(
557+ (AscendC::MicroAPI::RegTensor<float>&)max32,
558+ (AscendC::MicroAPI::RegTensor<T_X>&)max16, preMaskScale);
559+ AscendC::MicroAPI::Compare<uint32_t, CMPMODE::LT>(cmpResult, max32, expMask, preMaskScale);
560+ AscendC::MicroAPI::Compare<uint32_t, CMPMODE::NE>(zeroMask, max32, zeroRegTensor32, preMaskScale);
561+ AscendC::MicroAPI::Mul(
562+ (AscendC::MicroAPI::RegTensor<float>&)max32,
563+ (AscendC::MicroAPI::RegTensor<float>&)max32,
564+ (AscendC::MicroAPI::RegTensor<float>&)invMax, preMaskScale);
565+ AscendC::MicroAPI::ShiftRights(exp32, max32, SHR_NUM_FOR_FP32, preMaskScale);
566+ AscendC::MicroAPI::And(man32, max32, manMaskFP32, preMaskScale);
567+ AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p0, exp32, CUBLAS_ZERO_FOR_ALL, preMaskScale);
568+ AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::LT>(p1, exp32, CUBLAS_EXP254, preMaskScale);
569+ AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, CUBLAS_ZERO_FOR_ALL, preMaskScale);
570+ AscendC::MicroAPI::MaskAnd(p0, p0, p1, preMaskScale);
571+ AscendC::MicroAPI::MaskAnd(p0, p0, p2, preMaskScale);
572+ AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::EQ>(p1, exp32, CUBLAS_ZERO_FOR_ALL, preMaskScale);
573+ AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, CUBLAS_HALF_FOR_MAN, preMaskScale);
574+ AscendC::MicroAPI::MaskAnd(p1, p1, p2, preMaskScale);
575+ AscendC::MicroAPI::MaskOr(p0, p0, p1, preMaskScale);
576+ AscendC::MicroAPI::Adds(expAddOne32, exp32, 1, preMaskScale);
577+ AscendC::MicroAPI::Select(extractExp, expAddOne32, exp32, p0);
578+ AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, fp8NanRegTensor, cmpResult);
579+ AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, zeroRegTensor32, zeroMask);
580+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(expOut, extractExp);
581+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
582+ mxScaleLocalAddr + i * ThirtyTwo, expOut, maskHalf);
583+ AscendC::MicroAPI::ShiftLefts(extractExp, extractExp, SHR_NUM_FOR_BF16, preMaskScale);
584+ AscendC::MicroAPI::Sub(halfScale, scaleBias, extractExp, preMaskScale);
585+ AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, nanRegTensor, cmpResult);
586+ AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, zeroRegTensor32, zeroMask);
587+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(recExpOut, halfScale);
588+ AscendC::MicroAPI::StoreAlign<uint16_t>(halfScaleLocalAddr + i * VL_F32, recExpOut, maskHalf);
589+ }
590+ }
591+}
592+ 
593+template <AscendC::RoundMode roundMode, typename T_X, typename T_Y>
594+__aicore__ inline void MxQuantComputeData(
595+ __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr,
596+ __ubuf__ int8_t* outLocalAddr, uint16_t loopNum)
597+{
598+ __VEC_SCOPE__
599+ {
600+ AscendC::MicroAPI::MaskReg dataMask1 = AscendC::MicroAPI::CreateMask<T_X>();
601+ AscendC::MicroAPI::MaskReg dataMask2 = AscendC::MicroAPI::CreateMask<T_X>();
602+ AscendC::MicroAPI::MaskReg dataMask3 = AscendC::MicroAPI::CreateMask<T_X>();
603+ AscendC::MicroAPI::MaskReg dataMask4 = AscendC::MicroAPI::CreateMask<T_X>();
604+ AscendC::MicroAPI::MaskReg dataMask5 = AscendC::MicroAPI::CreateMask<T_Y>();
605+ AscendC::MicroAPI::MaskReg maskAll =
606+ AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
607+ AscendC::MicroAPI::RegTensor<uint16_t> halfScaleForMul;
608+ AscendC::MicroAPI::RegTensor<float> floatScaleForMul;
609+ AscendC::MicroAPI::RegTensor<T_X> vdExp0;
610+ AscendC::MicroAPI::RegTensor<T_X> vdExp1;
611+ AscendC::MicroAPI::RegTensor<T_X> vdExp0Convert;
612+ AscendC::MicroAPI::RegTensor<T_X> vdExp1Convert;
613+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
614+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
615+ AscendC::MicroAPI::RegTensor<float> vdExp0FP32Zero;
616+ AscendC::MicroAPI::RegTensor<float> vdExp0FP32One;
617+ AscendC::MicroAPI::RegTensor<float> vdExp1FP32Zero;
618+ AscendC::MicroAPI::RegTensor<float> vdExp1FP32One;
619+ AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP8Zero;
620+ AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP8One;
621+ AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP8Zero;
622+ AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP8One;
623+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdBF16Exp0FP4;
624+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdBF16Exp1FP4;
625+ static constexpr AscendC::MicroAPI::CastTrait castTrait = {
626+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
627+ AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
628+ static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
629+ AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
630+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
631+ static constexpr AscendC::MicroAPI::CastTrait castTraitZero = {
632+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
633+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
634+ static constexpr AscendC::MicroAPI::CastTrait castTraitOne = {
635+ AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN,
636+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
637+ static constexpr AscendC::MicroAPI::CastTrait castTrait32to80 = {
638+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::SAT,
639+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
640+ static constexpr AscendC::MicroAPI::CastTrait castTrait32to81 = {
641+ AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::SAT,
642+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
643+ static constexpr AscendC::MicroAPI::CastTrait castTrait32to82 = {
644+ AscendC::MicroAPI::RegLayout::TWO, AscendC::MicroAPI::SatMode::SAT,
645+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
646+ static constexpr AscendC::MicroAPI::CastTrait castTrait32to83 = {
647+ AscendC::MicroAPI::RegLayout::THREE, AscendC::MicroAPI::SatMode::SAT,
648+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
649+ for (uint16_t i = 0; i < loopNum; i++) {
650+ AscendC::MicroAPI::LoadAlign<
651+ T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
652+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(
653+ vdExp0, vdExp1, srcAddr, VL_B16 * DIGIT_TWO);
654+ AscendC::MicroAPI::LoadAlign<
655+ uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
656+ AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(
657+ halfScaleForMul, halfScaleLocalAddr, ELEMENT_AFTER_REDUCE);
658+ if constexpr (IsSame<T_X, half>::value) {
659+ AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);
660+ AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);
661+ AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitZero>(
662+ floatScaleForMul, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul, maskAll);
663+ AscendC::MicroAPI::Mul(vdExp0FP32Zero, vdExp0FP32Zero, floatScaleForMul, dataMask3);
664+ AscendC::MicroAPI::Mul(vdExp0FP32One, vdExp0FP32One, floatScaleForMul, dataMask4);
665+ AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp1FP32Zero, vdExp1, dataMask1);
666+ AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp1FP32One, vdExp1, dataMask1);
667+ AscendC::MicroAPI::Mul(vdExp1FP32Zero, vdExp1FP32Zero, floatScaleForMul, dataMask3);
668+ AscendC::MicroAPI::Mul(vdExp1FP32One, vdExp1FP32One, floatScaleForMul, dataMask4);
669+ } else {
670+ AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
671+ AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
672+ AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);
673+ AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);
674+ AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp1FP32Zero, vdExp1, dataMask2);
675+ AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp1FP32One, vdExp1, dataMask2);
676+ }
677+ AscendC::MicroAPI::Cast<T_Y, float, castTrait32to80>(vdExp0FP8Zero, vdExp0FP32Zero, dataMask3);
678+ AscendC::MicroAPI::Cast<T_Y, float, castTrait32to82>(vdExp0FP8One, vdExp0FP32One, dataMask3);
679+ AscendC::MicroAPI::Cast<T_Y, float, castTrait32to81>(vdExp1FP8Zero, vdExp1FP32Zero, dataMask4);
680+ AscendC::MicroAPI::Cast<T_Y, float, castTrait32to83>(vdExp1FP8One, vdExp1FP32One, dataMask4);
681+ 
682+ AscendC::MicroAPI::Add(
683+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
684+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
685+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8One, dataMask5);
686+ AscendC::MicroAPI::Add(
687+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
688+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
689+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8Zero, dataMask5);
690+ AscendC::MicroAPI::Add(
691+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
692+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
693+ (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8One, dataMask5);
694+ 
695+ AscendC::MicroAPI::StoreAlign<
696+ int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
697+ AscendC::MicroAPI::StoreDist::DIST_NORM_B8>(
698+ outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP8Zero, OUT_ALL, dataMask5);
699+ }
700+ }
701+}
702+ 
703+template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode, typename T_Y>
704+__aicore__ inline void ComputeFP4FromHalf(AscendC::MicroAPI::RegTensor<float>& Reg)
705+{
706+ AscendC::MicroAPI::MaskReg pregAll32 =
707+ AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
708+ AscendC::MicroAPI::MaskReg zeroMask;
709+ AscendC::MicroAPI::MaskReg specialMask;
710+ AscendC::MicroAPI::MaskReg negInfMask;
711+ AscendC::MicroAPI::RegTensor<int32_t> negZero;
712+ AscendC::MicroAPI::RegTensor<int32_t> maxExpFP32;
713+ AscendC::MicroAPI::RegTensor<int32_t> exp0FP32;
714+ AscendC::MicroAPI::RegTensor<int32_t> exp1FP32;
715+ 
716+ AscendC::MicroAPI::Duplicate(negZero, NEG_ZERO);
717+ AscendC::MicroAPI::Compare<int32_t, CMPMODE::EQ>(
718+ negInfMask, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, negZero, pregAll32);
719+ 
720+ if constexpr (IsSame<T_Y, fp4x2_e1m2_t>::value) {
721+ AscendC::MicroAPI::Muls(Reg, Reg, FOUR, pregAll32);
722+ AscendC::MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
723+ AscendC::MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
724+ AscendC::MicroAPI::Muls(Reg, Reg, ONE_FOURTH, pregAll32);
725+ } else {
726+ AscendC::MicroAPI::Duplicate(maxExpFP32, MAX_EXP_FOR_FP32);
727+ AscendC::MicroAPI::And(exp0FP32, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, maxExpFP32, pregAll32);
728+ AscendC::MicroAPI::ShiftRights(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);
729+ AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS_NEG, pregAll32);
730+ AscendC::MicroAPI::Maxs(exp0FP32, exp0FP32, 0, pregAll32);
731+ AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, NEG_ONE, pregAll32);
732+ AscendC::MicroAPI::Muls(exp1FP32, exp0FP32, NEG_ONE, pregAll32);
733+ AscendC::MicroAPI::Adds(exp1FP32, exp1FP32, FP32_BIAS, pregAll32);
734+ AscendC::MicroAPI::ShiftLefts(exp1FP32, exp1FP32, SHR_NUM_FOR_FP32, pregAll32);
735+ 
736+ AscendC::MicroAPI::Mul(Reg, Reg, (AscendC::MicroAPI::RegTensor<float>&)exp1FP32, pregAll32);
737+ AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS, pregAll32);
738+ AscendC::MicroAPI::ShiftLefts(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);
739+ AscendC::MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
740+ AscendC::MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
741+ AscendC::MicroAPI::Mul(Reg, Reg, (AscendC::MicroAPI::RegTensor<float>&)exp0FP32, pregAll32);
742+ }
743+ 
744+ AscendC::MicroAPI::CompareScalar<float, CMPMODE::EQ>(zeroMask, Reg, 0, pregAll32);
745+ AscendC::MicroAPI::MaskAnd(zeroMask, specialMask, zeroMask, pregAll32);
746+ AscendC::MicroAPI::MaskOr(zeroMask, negInfMask, zeroMask, pregAll32);
747+ AscendC::MicroAPI::Select<int32_t>(
748+ (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, negZero, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg,
749+ zeroMask);
750+}
751+ 
752+template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode, typename T_X, typename T_Y>
753+__aicore__ inline void MxQuantComputeDataFP4(
754+ __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr,
755+ __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB, uint16_t loopNum)
756+{
757+ __VEC_SCOPE__
758+ {
759+ AscendC::MicroAPI::MaskReg dataMask1;
760+ AscendC::MicroAPI::RegTensor<uint16_t> halfScaleForMul;
761+ AscendC::MicroAPI::RegTensor<T_X> vdExp0;
762+ AscendC::MicroAPI::RegTensor<T_X> vdExp1;
763+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
764+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
765+ AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP4;
766+ AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP4;
767+ AscendC::MicroAPI::RegTensor<float> halfScaleForMulFP32;
768+ AscendC::MicroAPI::RegTensor<float> vdExp0ZeroFP32;
769+ AscendC::MicroAPI::RegTensor<float> vdExp0OneFP32;
770+ AscendC::MicroAPI::RegTensor<float> vdExp1ZeroFP32;
771+ AscendC::MicroAPI::RegTensor<float> vdExp1OneFP32;
772+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0ZeroBF16;
773+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0OneBF16;
774+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1ZeroBF16;
775+ AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1OneBF16;
776+ AscendC::MicroAPI::MaskReg dataMaskB16 = AscendC::MicroAPI::CreateMask<half>();
777+ AscendC::MicroAPI::MaskReg dataMaskB32 = AscendC::MicroAPI::CreateMask<float>();
778+ 
779+ static constexpr AscendC::MicroAPI::CastTrait castTrait = {
780+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
781+ AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
782+ static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
783+ AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
784+ AscendC::MicroAPI::MaskMergeMode::ZEROING, toBf16RoundMode};
785+ static constexpr AscendC::MicroAPI::CastTrait castTraitF16toFp32Zero = {
786+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
787+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
788+ static constexpr AscendC::MicroAPI::CastTrait castTraitF16toFp32One = {
789+ AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN,
790+ AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
791+ static constexpr AscendC::MicroAPI::CastTrait castTraitFp32toBF16 = {
792+ AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::NO_SAT,
793+ AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
794+ 
795+ for (uint16_t i = 0; i < loopNum; i++) {
796+ dataMask1 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
797+ AscendC::MicroAPI::LoadAlign<
798+ T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
799+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, VL_B16 * DIGIT_TWO);
800+ AscendC::MicroAPI::LoadAlign<
801+ uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
802+ AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr, ELEMENT_AFTER_REDUCE);
803+ 
804+ if constexpr (IsSame<T_X, half>::value) {
805+ if constexpr (roundMode == RoundMode::CAST_RINT || roundMode == RoundMode::CAST_ROUND) {
806+ // tail_axis_optimize_fp16 -> fp4 (rint, round)
807+ AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitF16toFp32Zero>(
808+ halfScaleForMulFP32, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
809+ dataMaskB16);
810+ AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32Zero>(
811+ vdExp0ZeroFP32, vdExp0, dataMaskB16);
812+ AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32One>(vdExp0OneFP32, vdExp0, dataMaskB16);
813+ AscendC::MicroAPI::Mul(vdExp0ZeroFP32, vdExp0ZeroFP32, halfScaleForMulFP32, dataMaskB32);
814+ AscendC::MicroAPI::Mul(vdExp0OneFP32, vdExp0OneFP32, halfScaleForMulFP32, dataMaskB32);
815+ ComputeFP4FromHalf<toBf16RoundMode, roundMode, T_Y>(vdExp0ZeroFP32);
816+ ComputeFP4FromHalf<toBf16RoundMode, roundMode, T_Y>(vdExp0OneFP32);
817+ AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
818+ vdExp0ZeroBF16, vdExp0ZeroFP32, dataMaskB32);
819+ AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
820+ vdExp0OneBF16, vdExp0OneFP32, dataMaskB32);
821+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
822+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0ZeroBF16,
823+ (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp0ZeroBF16);
824+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
825+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0OneBF16,
826+ (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp0OneBF16);
827+ AscendC::MicroAPI::Interleave(vdExp0ZeroBF16, vdExp0OneBF16, vdExp0ZeroBF16, vdExp0OneBF16);
828+ AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32Zero>(
829+ vdExp1ZeroFP32, vdExp1, dataMaskB16);
830+ AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32One>(vdExp1OneFP32, vdExp1, dataMaskB16);
831+ AscendC::MicroAPI::Mul(vdExp1ZeroFP32, vdExp1ZeroFP32, halfScaleForMulFP32, dataMaskB32);
832+ AscendC::MicroAPI::Mul(vdExp1OneFP32, vdExp1OneFP32, halfScaleForMulFP32, dataMaskB32);
833+ ComputeFP4FromHalf<toBf16RoundMode, roundMode, T_Y>(vdExp1ZeroFP32);
834+ ComputeFP4FromHalf<toBf16RoundMode, roundMode, T_Y>(vdExp1OneFP32);
835+ AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
836+ vdExp1ZeroBF16, vdExp1ZeroFP32, dataMaskB32);
837+ AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
838+ vdExp1OneBF16, vdExp1OneFP32, dataMaskB32);
839+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
840+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1ZeroBF16,
841+ (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp1ZeroBF16);
842+ AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
843+ (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1OneBF16,
844+ (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp1OneBF16);
845+ AscendC::MicroAPI::Interleave(vdExp1ZeroBF16, vdExp1OneBF16, vdExp1ZeroBF16, vdExp1OneBF16);
846+ AscendC::MicroAPI::Interleave(vdExp0ZeroBF16, vdExp1ZeroBF16, vdExp0ZeroBF16, vdExp1ZeroBF16);
847+ AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp0FP4, vdExp0ZeroBF16, dataMask1);
848+ AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp1FP4, vdExp1ZeroBF16, dataMask1);
849+ } else {
850+ // for fp16 -> fp4 (floor)
851+ AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp0BF16, vdExp0, dataMask1);
852+ AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp1BF16, vdExp1, dataMask1);
853+ AscendC::MicroAPI::Mul(
854+ vdExp0BF16, vdExp0BF16, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
855+ dataMask1);
856+ AscendC::MicroAPI::Mul(
857+ vdExp1BF16, vdExp1BF16, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
858+ dataMask1);
859+ AscendC::MicroAPI::Interleave(vdExp0BF16, vdExp1BF16, vdExp0BF16, vdExp1BF16);
860+ AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp0FP4, vdExp0BF16, dataMask1);
861+ AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp1FP4, vdExp1BF16, dataMask1);
862+ }
863+ } else {
864+ // for bf16
865+ AscendC::MicroAPI::Mul(
866+ vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
867+ AscendC::MicroAPI::Mul(
868+ vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
869+ AscendC::MicroAPI::Interleave(vdExp0, vdExp1, vdExp0, vdExp1);
870+ AscendC::MicroAPI::Cast<T_Y, T_X, castTrait>(vdExp0FP4, vdExp0, dataMask1);
871+ AscendC::MicroAPI::Cast<T_Y, T_X, castTrait>(vdExp1FP4, vdExp1, dataMask1);
872+ }
873+ 
874+ AscendC::MicroAPI::StoreAlign<
875+ int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
876+ AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
877+ outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
878+ AscendC::MicroAPI::StoreAlign<
879+ int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
880+ AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
881+ outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
882+ }
883+ }
884+}
885+ 
143} // namespace AddRmsNormDynamicMxQuant886} // namespace AddRmsNormDynamicMxQuant
144#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H887#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_COMMON_H
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_fp4_r_full_load.h+6-381
@@ -14,7 +14,6 @@
14 */14 */
15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H
16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H
17-#define FLOAT_OVERFLOW_MODE_CTRL 60
18 17 
19#include "add_rms_norm_dynamic_mx_quant_common.h"18#include "add_rms_norm_dynamic_mx_quant_common.h"
20 19 
@@ -55,15 +54,6 @@ public:
55 betaFlag_ = tiling->betaFlag;54 betaFlag_ = tiling->betaFlag;
56 rstdFlag_ = tiling->rstdFlag;55 rstdFlag_ = tiling->rstdFlag;
57 56 
58- vlForHalfNumber = platform::GetVRegSize() / sizeof(T_X);
59- elementAfterReduce = platform::GetVRegSize() / UB_BLOCK_SIZE;
60- 
61- if constexpr (IsSame<T_Y, fp4x2_e2m1_t>::value) {
62- f4Emax_ = FP4_E2M1_BF16_MAX_EXP;
63- } else {
64- f4Emax_ = FP4_E1M2_MAX_EXP;
65- }
66- 
67 uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;57 uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;
68 x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, rowWork * numCol_);58 x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, rowWork * numCol_);
69 x2Gm.SetGlobalBuffer((__gm__ T_X*)x2 + blockOffset, rowWork * numCol_);59 x2Gm.SetGlobalBuffer((__gm__ T_X*)x2 + blockOffset, rowWork * numCol_);
@@ -152,7 +142,7 @@ private:
152 LocalTensor<float> xReduceLocal = xReduceBuff.Get<float>();142 LocalTensor<float> xReduceLocal = xReduceBuff.Get<float>();
153 CalculateSquareReduceSum(xFp32Local, xReduceLocal, curRows);143 CalculateSquareReduceSum(xFp32Local, xReduceLocal, curRows);
154 144 
155- CalculateRstd(xReduceLocal, rstdLocal, curRows);145+ CalculateRstd(xReduceLocal, rstdLocal, curRows, avgFactor_, epsilon_);
156 outQueueRstd.EnQue<float>(rstdLocal);146 outQueueRstd.EnQue<float>(rstdLocal);
157 147 
158 rstdLocal = outQueueRstd.DeQue<float>();148 rstdLocal = outQueueRstd.DeQue<float>();
@@ -454,60 +444,6 @@ private:
454 }444 }
455 }445 }
456 446 
457- __aicore__ inline void CalculateRstd(
458- LocalTensor<float>& xReduceLocal, LocalTensor<float>& rstdLocal, uint32_t curRows)
459- {
460- static constexpr float POS_INF = 3.40282366920938E+38;
461- static constexpr float SCALAR1 = -0.5;
462- static constexpr float SCALAR2 = 1.5;
463- static constexpr float SCALAR3 = 0.5;
464- static constexpr float SCALAR0 = -99.99;
465- 
466- __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr();
467- __local_mem__ float* xReduceUb = (__local_mem__ float*)xReduceLocal.GetPhyAddr();
468- uint16_t loopRows = static_cast<uint16_t>((curRows + VL_F32 - 1) / VL_F32);
469- __VEC_SCOPE__
470- {
471- RegTensor<float> var, rstd, r, y, s, t, one, scalar1, t1, t2, t3, t4, scalarInf, scalarZero;
472- MaskReg cmpRegZero, cmpRegInf;
473- MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
474- MaskReg pregLoop;
475- 
476- uint32_t sreg = static_cast<uint32_t>(curRows);
477- for (uint16_t i = 0; i < loopRows; ++i) {
478- pregLoop = UpdateMask<float>(sreg);
479- AscendC::MicroAPI::Duplicate(scalarInf, POS_INF, pregLoop);
480- AscendC::MicroAPI::Duplicate(scalarZero, float(0.0), pregLoop);
481- AscendC::MicroAPI::Duplicate(one, float(1.0), pregLoop);
482- AscendC::MicroAPI::Duplicate(scalar1, SCALAR3, pregLoop);
483- AscendC::MicroAPI::Duplicate(t1, SCALAR2, pregLoop);
484- AscendC::MicroAPI::Duplicate(s, float(1.0), pregLoop);
485- AscendC::MicroAPI::DataCopy(var, xReduceUb + i * VL_F32);
486- AscendC::MicroAPI::Muls(var, var, avgFactor_, pregLoop);
487- AscendC::MicroAPI::Adds(var, var, epsilon_, pregLoop);
488- AscendC::MicroAPI::Maxs(var, var, SCALAR0, pregLoop);
489- AscendC::MicroAPI::Div(r, one, var, pregLoop);
490- AscendC::MicroAPI::Sqrt(y, r, pregLoop);
491- AscendC::MicroAPI::Muls(t, var, SCALAR1, pregLoop);
492- AscendC::MicroAPI::Mul(t, t, y, pregLoop);
493- AscendC::MicroAPI::Mula(t1, t, y, pregLoop);
494- AscendC::MicroAPI::Mul(rstd, y, t1, pregLoop);
495- AscendC::MicroAPI::Muls(t3, var, float(-1.0), pregLoop);
496- AscendC::MicroAPI::Mula(s, t3, r, pregLoop);
497- AscendC::MicroAPI::Muls(t4, rstd, float(-1.0), pregLoop);
498- AscendC::MicroAPI::Mula(r, t4, rstd, pregLoop);
499- AscendC::MicroAPI::Mula(s, var, r, pregLoop);
500- AscendC::MicroAPI::Mul(s, s, rstd, pregLoop);
501- AscendC::MicroAPI::Mula(rstd, s, scalar1, pregLoop);
502- AscendC::MicroAPI::CompareScalar(cmpRegZero, var, POS_INF, pregLoop);
503- AscendC::MicroAPI::Select(rstd, scalarZero, rstd, cmpRegZero);
504- AscendC::MicroAPI::CompareScalar(cmpRegInf, var, float(0.0), pregLoop);
505- AscendC::MicroAPI::Select(rstd, scalarInf, rstd, cmpRegInf);
506- AscendC::MicroAPI::DataCopy(rstdInUb + i * VL_F32, rstd, pregLoop);
507- }
508- }
509- }
510- 
511 template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode>447 template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode>
512 __aicore__ inline void DynamicMxQuantPhaseFP4(LocalTensor<T_X>& yLocal, uint32_t curRows)448 __aicore__ inline void DynamicMxQuantPhaseFP4(LocalTensor<T_X>& yLocal, uint32_t curRows)
513 {449 {
@@ -516,8 +452,8 @@ private:
516 uint32_t totalScaleInUB = curRows * blockNumInColAxis_;452 uint32_t totalScaleInUB = curRows * blockNumInColAxis_;
517 uint32_t totalCountInUB = curRows * blockNumInColAxis_ * mxBlockSize_;453 uint32_t totalCountInUB = curRows * blockNumInColAxis_ * mxBlockSize_;
518 454 
519- uint16_t loopNum = (totalCountInUB + vlForHalfNumber * DIGIT_TWO - 1) / (vlForHalfNumber * DIGIT_TWO);455+ uint16_t loopNum = (totalCountInUB + VL_B16 * DIGIT_TWO - 1) / (VL_B16 * DIGIT_TWO);
520- uint16_t loopNumScale = (totalScaleInUB + vlForHalfNumber - 1) / vlForHalfNumber;456+ uint16_t loopNumScale = (totalScaleInUB + VL_B16 - 1) / VL_B16;
521 457 
522 auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());458 auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
523 auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());459 auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
@@ -532,327 +468,19 @@ private:
532 auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());468 auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());
533 469 
534 maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());470 maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
535- MxQuantComputeMaxExpOCP(srcAddr, maxExpAddr, totalCountInUB, loopNum);471+ MxQuantComputeMaxExpOCP<T_X>(srcAddr, maxExpAddr, loopNum);
536- MxQuantComputeScaleOCP(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);472+ MxQuantComputeScaleOCP<T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);
537 473 
538 srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());474 srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
539 halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());475 halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
540 476 
541- MxQuantComputeDataFP4<toBf16RoundMode, roundMode>(477+ MxQuantComputeDataFP4<toBf16RoundMode, roundMode, T_X, T_Y>(
542 srcAddr, halfScaleLocalAddr, outLocalAddr, totalCountInUB, loopNum);478 srcAddr, halfScaleLocalAddr, outLocalAddr, totalCountInUB, loopNum);
543 479 
544 outQueueQuantY.EnQue(outLocal);480 outQueueQuantY.EnQue(outLocal);
545 mxScaleQueue.EnQue(mxScaleLocal);481 mxScaleQueue.EnQue(mxScaleLocal);
546 }482 }
547 483 
548- __aicore__ inline void MxQuantComputeMaxExpOCP(
549- __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint32_t totalCountInUB, uint16_t loopNum)
550- {
551- __VEC_SCOPE__
552- {
553- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
554- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
555- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
556- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
557- AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect0;
558- AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect1;
559- AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract0;
560- AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract1;
561- AscendC::MicroAPI::RegTensor<uint16_t> expMaskBF16;
562- AscendC::MicroAPI::Duplicate(expMaskBF16, MAX_EXP_FOR_BF16);
563- AscendC::MicroAPI::RegTensor<uint16_t> invalidMaskFP16;
564- AscendC::MicroAPI::Duplicate(invalidMaskFP16, INVALID_FLOAT16);
565- AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
566- AscendC::MicroAPI::MaskReg scaleMask1;
567- AscendC::MicroAPI::MaskReg scaleMask2;
568- AscendC::MicroAPI::MaskReg invalidDataMask0;
569- AscendC::MicroAPI::MaskReg invalidDataMask1;
570- AscendC::MicroAPI::UnalignReg u1;
571- static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
572- AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
573- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
574- for (uint16_t i = 0; i < loopNum; i++) {
575- scaleMask1 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
576- scaleMask2 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
577- AscendC::MicroAPI::DataCopy<
578- T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
579- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, vlForHalfNumber * DIGIT_TWO);
580- if constexpr (IsSame<T_X, half>::value) {
581- AscendC::MicroAPI::And(
582- vdExpSelect0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, invalidMaskFP16, scaleMask1);
583- AscendC::MicroAPI::And(
584- vdExpSelect1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, invalidMaskFP16, scaleMask1);
585- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
586- invalidDataMask0, vdExpSelect0, invalidMaskFP16, scaleMask1);
587- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
588- invalidDataMask1, vdExpSelect1, invalidMaskFP16, scaleMask1);
589- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp0BF16, vdExp0, scaleMask1);
590- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp1BF16, vdExp1, scaleMask1);
591- AscendC::MicroAPI::And(
592- vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0BF16, expMaskBF16, scaleMask1);
593- AscendC::MicroAPI::And(
594- vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1BF16, expMaskBF16, scaleMask1);
595- AscendC::MicroAPI::Select<uint16_t>(vdExpExtract0, vdExpExtract0, expMaskBF16, invalidDataMask0);
596- AscendC::MicroAPI::Select<uint16_t>(vdExpExtract1, vdExpExtract1, expMaskBF16, invalidDataMask1);
597- } else {
598- AscendC::MicroAPI::And(
599- vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, expMaskBF16, scaleMask1);
600- AscendC::MicroAPI::And(
601- vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, expMaskBF16, scaleMask1);
602- }
603- AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, scaleMask1);
604- AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, scaleMask1);
605- AscendC::MicroAPI::DataCopyUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
606- maxExpAddr, vdMaxExp, u1, elementAfterReduce);
607- }
608- AscendC::MicroAPI::DataCopyUnAlignPost(maxExpAddr, u1, 0);
609- }
610- }
611- 
612- __aicore__ inline void MxQuantComputeScaleOCP(
613- __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, __ubuf__ uint16_t* halfScaleLocalAddr,
614- uint32_t totalScaleInUB, uint16_t loopNumScale)
615- {
616- __VEC_SCOPE__
617- {
618- AscendC::MicroAPI::RegTensor<uint16_t> expMask;
619- AscendC::MicroAPI::Duplicate(expMask, MAX_EXP_FOR_BF16);
620- AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
621- AscendC::MicroAPI::MaskReg cmpResult;
622- AscendC::MicroAPI::MaskReg zeroMask;
623- AscendC::MicroAPI::MaskReg preMaskScale;
624- AscendC::MicroAPI::RegTensor<uint16_t> maxExpValue;
625- AscendC::MicroAPI::Duplicate(maxExpValue, f4Emax_);
626- AscendC::MicroAPI::RegTensor<uint16_t> sharedExp;
627- AscendC::MicroAPI::RegTensor<uint16_t> scaleValue;
628- AscendC::MicroAPI::RegTensor<uint16_t> scaleBias;
629- AscendC::MicroAPI::Duplicate(scaleBias, BF16_EXP_BIAS);
630- AscendC::MicroAPI::RegTensor<uint16_t> halfScale;
631- AscendC::MicroAPI::RegTensor<uint16_t> fp8NanRegTensor;
632- AscendC::MicroAPI::Duplicate(fp8NanRegTensor, MAX_EXP_FOR_FP8);
633- AscendC::MicroAPI::RegTensor<uint16_t> zeroRegTensor;
634- AscendC::MicroAPI::Duplicate(zeroRegTensor, 0);
635- AscendC::MicroAPI::RegTensor<uint16_t> nanRegTensor;
636- AscendC::MicroAPI::Duplicate(nanRegTensor, NAN_CUSTOMIZATION);
637- AscendC::MicroAPI::MaskReg invalidDataMask;
638- AscendC::MicroAPI::MaskReg specialDataMask;
639- AscendC::MicroAPI::RegTensor<uint16_t> specialExpRegTensor;
640- AscendC::MicroAPI::Duplicate(specialExpRegTensor, SPECIAL_EXP_THRESHOLD);
641- for (uint16_t i = 0; i < loopNumScale; i++) {
642- preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);
643- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
644- vdMaxExp, maxExpAddr, vlForHalfNumber);
645- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(cmpResult, vdMaxExp, expMask, preMaskScale);
646- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(zeroMask, vdMaxExp, zeroRegTensor, preMaskScale);
647- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::LE>(invalidDataMask, vdMaxExp, maxExpValue, preMaskScale);
648- 
649- AscendC::MicroAPI::Select<uint16_t>(vdMaxExp, maxExpValue, vdMaxExp, invalidDataMask);
650- 
651- AscendC::MicroAPI::Sub(sharedExp, vdMaxExp, maxExpValue, preMaskScale);
652- AscendC::MicroAPI::ShiftRights(scaleValue, sharedExp, SHR_NUM_FOR_BF16, preMaskScale);
653- 
654- AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, fp8NanRegTensor, cmpResult);
655- AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, zeroRegTensor, zeroMask);
656- 
657- AscendC::MicroAPI::DataCopy<
658- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
659- AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
660- mxScaleLocalAddr, scaleValue, vlForHalfNumber / DIGIT_TWO, preMaskScale);
661- 
662- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::EQ>(specialDataMask, sharedExp, scaleBias, preMaskScale);
663- AscendC::MicroAPI::Sub(halfScale, scaleBias, sharedExp, preMaskScale);
664- AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, nanRegTensor, cmpResult);
665- AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, zeroRegTensor, zeroMask);
666- AscendC::MicroAPI::Select<uint16_t>(halfScale, specialExpRegTensor, halfScale, specialDataMask);
667- 
668- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
669- halfScaleLocalAddr, halfScale, vlForHalfNumber, preMaskScale);
670- }
671- }
672- }
673- 
674- template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode>
675- __aicore__ inline void ComputeFP4FromHalf(AscendC::MicroAPI::RegTensor<float>& Reg)
676- {
677- AscendC::MicroAPI::MaskReg pregAll32 =
678- AscendC::MicroAPI::CreateMask<uint32_t, AscendC::MicroAPI::MaskPattern::ALL>();
679- AscendC::MicroAPI::MaskReg zeroMask;
680- AscendC::MicroAPI::MaskReg specialMask;
681- AscendC::MicroAPI::MaskReg negInfMask;
682- AscendC::MicroAPI::RegTensor<int32_t> negZero;
683- AscendC::MicroAPI::RegTensor<int32_t> maxExpFP32;
684- AscendC::MicroAPI::RegTensor<int32_t> exp0FP32;
685- AscendC::MicroAPI::RegTensor<int32_t> exp1FP32;
686- 
687- AscendC::MicroAPI::Duplicate(negZero, NEG_ZERO);
688- AscendC::MicroAPI::Compare<int32_t, CMPMODE::EQ>(
689- negInfMask, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, negZero, pregAll32);
690- 
691- if constexpr (IsSame<T_Y, fp4x2_e1m2_t>::value) {
692- AscendC::MicroAPI::Muls(Reg, Reg, FOUR, pregAll32);
693- AscendC::MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
694- AscendC::MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
695- AscendC::MicroAPI::Muls(Reg, Reg, ONE_FOURTH, pregAll32);
696- } else {
697- AscendC::MicroAPI::Duplicate(maxExpFP32, MAX_EXP_FOR_FP32);
698- AscendC::MicroAPI::And(exp0FP32, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, maxExpFP32, pregAll32);
699- AscendC::MicroAPI::ShiftRights(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);
700- AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS_NEG, pregAll32);
701- AscendC::MicroAPI::Maxs(exp0FP32, exp0FP32, 0, pregAll32);
702- AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, NEG_ONE, pregAll32);
703- AscendC::MicroAPI::Muls(exp1FP32, exp0FP32, NEG_ONE, pregAll32);
704- AscendC::MicroAPI::Adds(exp1FP32, exp1FP32, FP32_BIAS, pregAll32);
705- AscendC::MicroAPI::ShiftLefts(exp1FP32, exp1FP32, SHR_NUM_FOR_FP32, pregAll32);
706- 
707- AscendC::MicroAPI::Mul(Reg, Reg, (AscendC::MicroAPI::RegTensor<float>&)exp1FP32, pregAll32);
708- AscendC::MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS, pregAll32);
709- AscendC::MicroAPI::ShiftLefts(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);
710- AscendC::MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
711- AscendC::MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
712- AscendC::MicroAPI::Mul(Reg, Reg, (AscendC::MicroAPI::RegTensor<float>&)exp0FP32, pregAll32);
713- }
714- 
715- AscendC::MicroAPI::CompareScalar<float, CMPMODE::EQ>(zeroMask, Reg, 0, pregAll32);
716- AscendC::MicroAPI::MaskAnd(zeroMask, specialMask, zeroMask, pregAll32);
717- AscendC::MicroAPI::MaskOr(zeroMask, negInfMask, zeroMask, pregAll32);
718- AscendC::MicroAPI::Select<int32_t>(
719- (AscendC::MicroAPI::RegTensor<int32_t>&)Reg, negZero, (AscendC::MicroAPI::RegTensor<int32_t>&)Reg,
720- zeroMask);
721- }
722- 
723- template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode>
724- __aicore__ inline void MxQuantComputeDataFP4(
725- __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr, __ubuf__ int8_t* outLocalAddr,
726- uint32_t totalCountInUB, uint16_t loopNum)
727- {
728- __VEC_SCOPE__
729- {
730- AscendC::MicroAPI::MaskReg dataMask1;
731- AscendC::MicroAPI::RegTensor<uint16_t> halfScaleForMul;
732- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
733- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
734- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
735- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
736- AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP4;
737- AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP4;
738- AscendC::MicroAPI::RegTensor<float> halfScaleForMulFP32;
739- AscendC::MicroAPI::RegTensor<float> vdExp0ZeroFP32;
740- AscendC::MicroAPI::RegTensor<float> vdExp0OneFP32;
741- AscendC::MicroAPI::RegTensor<float> vdExp1ZeroFP32;
742- AscendC::MicroAPI::RegTensor<float> vdExp1OneFP32;
743- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0ZeroBF16;
744- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0OneBF16;
745- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1ZeroBF16;
746- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1OneBF16;
747- AscendC::MicroAPI::MaskReg dataMaskB16 = AscendC::MicroAPI::CreateMask<half>();
748- AscendC::MicroAPI::MaskReg dataMaskB32 = AscendC::MicroAPI::CreateMask<float>();
749- 
750- static constexpr AscendC::MicroAPI::CastTrait castTrait = {
751- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
752- AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
753- static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
754- AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
755- AscendC::MicroAPI::MaskMergeMode::ZEROING, toBf16RoundMode};
756- static constexpr AscendC::MicroAPI::CastTrait castTraitF16toFp32Zero = {
757- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
758- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
759- static constexpr AscendC::MicroAPI::CastTrait castTraitF16toFp32One = {
760- AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN,
761- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
762- static constexpr AscendC::MicroAPI::CastTrait castTraitFp32toBF16 = {
763- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::NO_SAT,
764- AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
765- 
766- for (uint16_t i = 0; i < loopNum; i++) {
767- dataMask1 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
768- AscendC::MicroAPI::DataCopy<
769- T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
770- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, vlForHalfNumber * DIGIT_TWO);
771- AscendC::MicroAPI::DataCopy<
772- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
773- AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr, elementAfterReduce);
774- 
775- if constexpr (IsSame<T_X, half>::value) {
776- if constexpr (roundMode == RoundMode::CAST_RINT || roundMode == RoundMode::CAST_ROUND) {
777- // tail_axis_optimize_fp16 -> fp4 (rint、round)
778- AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitF16toFp32Zero>(
779- halfScaleForMulFP32, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
780- dataMaskB16);
781- AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32Zero>(
782- vdExp0ZeroFP32, vdExp0, dataMaskB16);
783- AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32One>(vdExp0OneFP32, vdExp0, dataMaskB16);
784- AscendC::MicroAPI::Mul(vdExp0ZeroFP32, vdExp0ZeroFP32, halfScaleForMulFP32, dataMaskB32);
785- AscendC::MicroAPI::Mul(vdExp0OneFP32, vdExp0OneFP32, halfScaleForMulFP32, dataMaskB32);
786- ComputeFP4FromHalf<toBf16RoundMode, roundMode>(vdExp0ZeroFP32);
787- ComputeFP4FromHalf<toBf16RoundMode, roundMode>(vdExp0OneFP32);
788- AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
789- vdExp0ZeroBF16, vdExp0ZeroFP32, dataMaskB32);
790- AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
791- vdExp0OneBF16, vdExp0OneFP32, dataMaskB32);
792- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
793- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0ZeroBF16,
794- (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp0ZeroBF16);
795- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
796- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0OneBF16,
797- (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp0OneBF16);
798- AscendC::MicroAPI::Interleave(vdExp0ZeroBF16, vdExp0OneBF16, vdExp0ZeroBF16, vdExp0OneBF16);
799- AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32Zero>(
800- vdExp1ZeroFP32, vdExp1, dataMaskB16);
801- AscendC::MicroAPI::Cast<float, T_X, castTraitF16toFp32One>(vdExp1OneFP32, vdExp1, dataMaskB16);
802- AscendC::MicroAPI::Mul(vdExp1ZeroFP32, vdExp1ZeroFP32, halfScaleForMulFP32, dataMaskB32);
803- AscendC::MicroAPI::Mul(vdExp1OneFP32, vdExp1OneFP32, halfScaleForMulFP32, dataMaskB32);
804- ComputeFP4FromHalf<toBf16RoundMode, roundMode>(vdExp1ZeroFP32);
805- ComputeFP4FromHalf<toBf16RoundMode, roundMode>(vdExp1OneFP32);
806- AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
807- vdExp1ZeroBF16, vdExp1ZeroFP32, dataMaskB32);
808- AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitFp32toBF16>(
809- vdExp1OneBF16, vdExp1OneFP32, dataMaskB32);
810- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
811- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1ZeroBF16,
812- (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp1ZeroBF16);
813- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(
814- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1OneBF16,
815- (AscendC::MicroAPI::RegTensor<uint32_t>&)vdExp1OneBF16);
816- AscendC::MicroAPI::Interleave(vdExp1ZeroBF16, vdExp1OneBF16, vdExp1ZeroBF16, vdExp1OneBF16);
817- AscendC::MicroAPI::Interleave(vdExp0ZeroBF16, vdExp1ZeroBF16, vdExp0ZeroBF16, vdExp1ZeroBF16);
818- AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp0FP4, vdExp0ZeroBF16, dataMask1);
819- AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp1FP4, vdExp1ZeroBF16, dataMask1);
820- } else {
821- // for fp16 -> fp4 (floor)
822- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp0BF16, vdExp0, dataMask1);
823- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp1BF16, vdExp1, dataMask1);
824- AscendC::MicroAPI::Mul(
825- vdExp0BF16, vdExp0BF16, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
826- dataMask1);
827- AscendC::MicroAPI::Mul(
828- vdExp1BF16, vdExp1BF16, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul,
829- dataMask1);
830- AscendC::MicroAPI::Interleave(vdExp0BF16, vdExp1BF16, vdExp0BF16, vdExp1BF16);
831- AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp0FP4, vdExp0BF16, dataMask1);
832- AscendC::MicroAPI::Cast<T_Y, bfloat16_t, castTrait>(vdExp1FP4, vdExp1BF16, dataMask1);
833- }
834- } else {
835- // for bf16
836- AscendC::MicroAPI::Mul(
837- vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
838- AscendC::MicroAPI::Mul(
839- vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
840- AscendC::MicroAPI::Interleave(vdExp0, vdExp1, vdExp0, vdExp1);
841- AscendC::MicroAPI::Cast<T_Y, T_X, castTrait>(vdExp0FP4, vdExp0, dataMask1);
842- AscendC::MicroAPI::Cast<T_Y, T_X, castTrait>(vdExp1FP4, vdExp1, dataMask1);
843- }
844- 
845- AscendC::MicroAPI::DataCopy<
846- int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
847- AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
848- outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
849- AscendC::MicroAPI::DataCopy<
850- int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
851- AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
852- outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
853- }
854- }
855- }
856 484 
857 __aicore__ inline void CopyInXMultiMoveAlign(uint64_t offset, uint32_t curRows)485 __aicore__ inline void CopyInXMultiMoveAlign(uint64_t offset, uint32_t curRows)
858 {486 {
@@ -972,9 +600,6 @@ private:
972 uint32_t betaFlag_;600 uint32_t betaFlag_;
973 uint32_t rstdFlag_;601 uint32_t rstdFlag_;
974 602 
975- uint16_t f4Emax_;
976- uint32_t vlForHalfNumber;
977- uint16_t elementAfterReduce;
978};603};
979} // namespace AddRmsNormDynamicMxQuant604} // namespace AddRmsNormDynamicMxQuant
980#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H605#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_R_FULL_LOAD_H
Anorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_fp4_split_r.h+646-0
@@ -0,0 +1,646 @@
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 add_rms_norm_dynamic_mx_quant_fp4_split_r.h
13+ * \brief
14+ */
15+#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_SPLIT_R_H
16+#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_SPLIT_R_H
17+ 
18+#include "add_rms_norm_dynamic_mx_quant_common.h"
19+ 
20+namespace AddRmsNormDynamicMxQuant {
21+ 
22+template <typename T_X, typename T_GAMMA, typename T_Y>
23+class AddRmsNormDynamicMxQuantFP4SplitR {
24+public:
25+ __aicore__ inline AddRmsNormDynamicMxQuantFP4SplitR(TPipe* pipe)
26+ {
27+ pPipe = pipe;
28+ }
29+ 
30+ __aicore__ inline void Init(
31+ GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y,
32+ GM_ADDR x, GM_ADDR mxscale, GM_ADDR workspace, GM_ADDR rstd,
33+ const AddRmsNormDynamicMxQuantSplitRTilingData* tiling)
34+ {
35+ #if (__NPU_ARCH__ == 3510)
36+ AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
37+ #endif
38+ ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
39+ 
40+ numCol_ = tiling->numCol;
41+ numColAlign_ = tiling->numColAlign;
42+ blockFactor_ = tiling->blockFactor;
43+ mLastCore_ = tiling->mLastCore;
44+ baseN_ = tiling->baseN;
45+ baseNBlockSize_ = tiling->baseNBlockSize;
46+ baseM_ = tiling->baseM;
47+ nUbLoops_ = tiling->nUbLoops;
48+ binAddQuotient_ = tiling->binAddQuotient;
49+ powerSplit_ = tiling->powerSplit;
50+ mainFoldCount_ = tiling->mainFoldCount;
51+ foldTail_ = tiling->foldTail;
52+ epsilon_ = tiling->epsilon;
53+ avgFactor_ = tiling->avgFactor;
54+ roundMode_ = tiling->roundMode;
55+ mxBlockSize_ = tiling->mxBlockSize;
56+ scaleAlg_ = tiling->scaleAlg;
57+ mxScaleSize_ = tiling->mxScaleSize;
58+ betaFlag_ = tiling->betaFlag;
59+ rstdFlag_ = tiling->rstdFlag;
60+ 
61+ resultCacheID_ = GetCacheId(powerSplit_ - 1);
62+ mCurCore_ = (GetBlockIdx() == GetBlockNum() - 1) ? mLastCore_ : blockFactor_;
63+ 
64+ // === Setup GM tensors ===
65+ uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;
66+ x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, mCurCore_ * numCol_);
67+ x2Gm.SetGlobalBuffer((__gm__ T_X*)x2 + blockOffset, mCurCore_ * numCol_);
68+ gammaGm.SetGlobalBuffer((__gm__ T_GAMMA*)gamma, numCol_);
69+ if (betaFlag_ != 0) {
70+ betaGm.SetGlobalBuffer((__gm__ T_GAMMA*)beta, numCol_);
71+ }
72+ xOutGm.SetGlobalBuffer((__gm__ T_X*)x + blockOffset, mCurCore_ * numCol_);
73+ if (rstdFlag_ != 0) {
74+ rstdGm.SetGlobalBuffer((__gm__ float*)rstd + GetBlockIdx() * blockFactor_, blockFactor_);
75+ }
76+ yFp4Gm.SetGlobalBuffer((__gm__ uint8_t*)y + blockOffset / DIGIT_TWO, mCurCore_ * numCol_ / DIGIT_TWO);
77+ mxScaleGm.SetGlobalBuffer(
78+ (__gm__ uint8_t*)mxscale + GetBlockIdx() * blockFactor_ * mxScaleSize_, mCurCore_ * mxScaleSize_);
79+ 
80+ // === Compute buffer sizes ===
81+ uint64_t xBufSize = CeilAlign(baseN_ * sizeof(T_X), UB_BLOCK_SIZE);
82+ uint64_t xFp32BufSize = CeilAlign(baseN_ * sizeof(float), UB_BLOCK_SIZE);
83+ uint64_t yTmpBufSize = CeilAlign(baseN_ * sizeof(T_X), UB_BLOCK_SIZE);
84+ uint64_t rstdBufSize = CeilAlign(baseM_ * sizeof(float), UB_BLOCK_SIZE);
85+ uint64_t cacheBufSize = CeilAlign(
86+ static_cast<uint64_t>((resultCacheID_ + 1) * sizeof(float)) * AR_RECOMPUTE_SUM_LEN, UB_BLOCK_SIZE);
87+ uint64_t binaryAddBufSize = CeilAlign(VL_F32 * DIGIT_TWO * sizeof(float), UB_BLOCK_SIZE);
88+ // FP4: baseN_ elements -> baseN_/2 bytes (fp4x2 packing)
89+ uint64_t quantYBufSize = CeilAlign(baseN_ / DIGIT_TWO, UB_BLOCK_SIZE);
90+ uint64_t maxExpBufSize = CeilAlign(baseNBlockSize_ * sizeof(uint16_t), UB_BLOCK_SIZE);
91+ uint64_t halfScaleBufSize = maxExpBufSize;
92+ uint64_t scaleBufSize = CeilAlign(baseNBlockSize_ * sizeof(uint8_t), UB_BLOCK_SIZE);
93+ 
94+ // === Init buffers ===
95+ pPipe->InitBuffer(inQueueX1, DOUBLE_BUFFER_NUM, xBufSize);
96+ pPipe->InitBuffer(inQueueX2, DOUBLE_BUFFER_NUM, xBufSize);
97+ if (betaFlag_ != 0) {
98+ pPipe->InitBuffer(
99+ inQueueGammabeta, DOUBLE_BUFFER_NUM, DIGIT_TWO * CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE));
100+ } else {
101+ pPipe->InitBuffer(inQueueGammabeta, DOUBLE_BUFFER_NUM, CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE));
102+ }
103+ 
104+ pPipe->InitBuffer(outQueueX, DOUBLE_BUFFER_NUM, xBufSize);
105+ pPipe->InitBuffer(outQueueRstd, DOUBLE_BUFFER_NUM, rstdBufSize);
106+ pPipe->InitBuffer(outQueueQuantY, DOUBLE_BUFFER_NUM, quantYBufSize);
107+ pPipe->InitBuffer(mxScaleQueue, DOUBLE_BUFFER_NUM, scaleBufSize);
108+ 
109+ pPipe->InitBuffer(xFp32Buf, xFp32BufSize);
110+ pPipe->InitBuffer(yTmpBuf, yTmpBufSize);
111+ pPipe->InitBuffer(cacheBuf, cacheBufSize);
112+ pPipe->InitBuffer(binaryAddBuf, binaryAddBufSize);
113+ pPipe->InitBuffer(maxExpBuff, maxExpBufSize);
114+ pPipe->InitBuffer(halfScaleBuff, halfScaleBufSize);
115+ }
116+ 
117+ __aicore__ inline void Process()
118+ {
119+ uint32_t mCnt = CeilDiv(mCurCore_, baseM_);
120+ for (uint64_t i = 0; i < mCnt; ++i) {
121+ uint32_t curM = (i == mCnt - 1) ? static_cast<uint32_t>(mCurCore_ - (mCnt - 1) * baseM_)
122+ : static_cast<uint32_t>(baseM_);
123+ 
124+ // Phase1: compute rstd
125+ LocalTensor<float> rstdLocal = outQueueRstd.AllocTensor<float>();
126+ for (uint32_t j = 0; j < curM; ++j) {
127+ int64_t gmRowOffset = (i * baseM_ + j) * numCol_;
128+ ComputeOneLineXSquareSum(rstdLocal, gmRowOffset, j);
129+ }
130+ CalculateRstd(rstdLocal, rstdLocal, curM, avgFactor_, epsilon_);
131+ outQueueRstd.EnQue<float>(rstdLocal);
132+ rstdLocal = outQueueRstd.DeQue<float>();
133+ 
134+ // Phase2: compute y、MxQuant、x_out
135+ for (uint64_t j = 0; j < nUbLoops_; ++j) {
136+ uint32_t curN = (j == nUbLoops_ - 1)
137+ ? static_cast<uint32_t>(numCol_ - (nUbLoops_ - 1) * baseN_)
138+ : static_cast<uint32_t>(baseN_);
139+ 
140+ // Load gamma (and beta) per-tile
141+ LocalTensor<T_GAMMA> gammabetaLocal = inQueueGammabeta.AllocTensor<T_GAMMA>();
142+ CopyInGammabeta(gammabetaLocal, j * baseN_, curN);
143+ inQueueGammabeta.EnQue(gammabetaLocal);
144+ inQueueGammabeta.DeQue<T_GAMMA>();
145+ 
146+ for (uint32_t k = 0; k < curM; ++k) {
147+ int64_t gmOffset = (i * baseM_ + k) * numCol_ + j * baseN_;
148+ 
149+ // Re-load x1, x2
150+ CopyInX(gmOffset, curN, j);
151+ LocalTensor<T_X> xLocal1 = inQueueX1.DeQue<T_X>();
152+ LocalTensor<T_X> xLocal2 = inQueueX2.DeQue<T_X>();
153+ LocalTensor<T_X> xOutLocal = outQueueX.AllocTensor<T_X>();
154+ LocalTensor<float> xFp32Local = xFp32Buf.Get<float>();
155+ CalculateXAdd(xLocal1, xLocal2, xOutLocal, xFp32Local, curN);
156+ inQueueX1.FreeTensor(xLocal1);
157+ inQueueX2.FreeTensor(xLocal2);
158+ outQueueX.EnQue<T_X>(xOutLocal);
159+ 
160+ // CopyOut x_out
161+ CopyOutX(gmOffset, curN, j);
162+ 
163+ // Compute y_local
164+ LocalTensor<T_X> yLocal = yTmpBuf.Get<T_X>();
165+ if ((j == nUbLoops_ - 1) && (numCol_ != numColAlign_)) {
166+ Duplicate<T_X>(yLocal, static_cast<T_X>(0), baseN_);
167+ PipeBarrier<PIPE_V>();
168+ }
169+ if (betaFlag_ != 0) {
170+ CalculateY<true>(xFp32Local, yLocal, rstdLocal, curN, k);
171+ } else {
172+ CalculateY<false>(xFp32Local, yLocal, rstdLocal, curN, k);
173+ }
174+ 
175+ // MxQuant FP4 with roundMode dispatch
176+ if (roundMode_ == MODE_RINT) {
177+ DynamicMxQuantPhaseFP4<RoundMode::CAST_TRUNC, RoundMode::CAST_RINT>(yLocal, j);
178+ } else if (roundMode_ == MODE_ROUND) {
179+ DynamicMxQuantPhaseFP4<RoundMode::CAST_TRUNC, RoundMode::CAST_ROUND>(yLocal, j);
180+ } else if (roundMode_ == MODE_FLOOR) {
181+ DynamicMxQuantPhaseFP4<RoundMode::CAST_FLOOR, RoundMode::CAST_FLOOR>(yLocal, j);
182+ }
183+ 
184+ // CopyOut y_quant and mxscale
185+ CopyOutQuantYFP4(gmOffset, curN, j);
186+ CopyOutMxScale(i * baseM_ + k, j);
187+ }
188+ 
189+ inQueueGammabeta.FreeTensor(gammabetaLocal);
190+ }
191+ // CopyOut rstd
192+ if (rstdFlag_ != 0) {
193+ DataCopyExtParams copyParams{
194+ static_cast<uint16_t>(1), static_cast<uint32_t>(curM * sizeof(float)), static_cast<uint32_t>(0),
195+ static_cast<uint32_t>(0), 0};
196+ DataCopyPad(rstdGm[i * baseM_], rstdLocal, copyParams);
197+ }
198+ outQueueRstd.FreeTensor(rstdLocal);
199+ }
200+ }
201+ 
202+private:
203+ __aicore__ inline void ComputeOneLineXSquareSum(
204+ LocalTensor<float>& rstdLocal, int64_t gmRowOffset, uint32_t rowIndex)
205+ {
206+ DataCopyPadParams padParams{false, 0, 0, 0};
207+ DataCopyParams xDataCopyParams;
208+ xDataCopyParams.blockCount = 1;
209+ xDataCopyParams.srcStride = 0;
210+ xDataCopyParams.dstStride = 0;
211+ DataCopyParams xFoldDataCopyParams;
212+ xFoldDataCopyParams.blockCount = 1;
213+ xFoldDataCopyParams.srcStride = 0;
214+ xFoldDataCopyParams.dstStride = 0;
215+ 
216+ LocalTensor<float> cacheLocal = cacheBuf.Get<float>();
217+ LocalTensor<float> xFp32Tmp = xFp32Buf.Get<float>();
218+ 
219+ for (int64_t r = 0; r < powerSplit_; ++r) {
220+ int64_t xGmOffset1 = gmRowOffset + baseN_ * r;
221+ int64_t xGmOffset2 = gmRowOffset + baseN_ * (r + powerSplit_);
222+ 
223+ // Step1: Load main tile, compute (x1+x2)² → xFp32Tmp, then Free
224+ xDataCopyParams.blockLen = baseN_ * sizeof(T_X);
225+ LocalTensor<T_X> x1Local = inQueueX1.AllocTensor<T_X>();
226+ DataCopyPad(x1Local, x1Gm[xGmOffset1], xDataCopyParams, padParams);
227+ inQueueX1.EnQue<T_X>(x1Local);
228+ x1Local = inQueueX1.DeQue<T_X>();
229+ 
230+ LocalTensor<T_X> x2Local = inQueueX2.AllocTensor<T_X>();
231+ DataCopyPad(x2Local, x2Gm[xGmOffset1], xDataCopyParams, padParams);
232+ inQueueX2.EnQue<T_X>(x2Local);
233+ x2Local = inQueueX2.DeQue<T_X>();
234+ 
235+ MainBlockSquareVF(x1Local, x2Local, xFp32Tmp, baseN_);
236+ inQueueX1.FreeTensor(x1Local);
237+ inQueueX2.FreeTensor(x2Local);
238+ 
239+ // Step2: Load fold tile, compute (x1Fold+x2Fold)² and accumulate to xFp32Tmp, then Free
240+ if (r < mainFoldCount_) {
241+ xFoldDataCopyParams.blockLen = baseN_ * sizeof(T_X);
242+ LocalTensor<T_X> x1FoldLocal = inQueueX1.AllocTensor<T_X>();
243+ DataCopyPad(x1FoldLocal, x1Gm[xGmOffset2], xFoldDataCopyParams, padParams);
244+ inQueueX1.EnQue<T_X>(x1FoldLocal);
245+ x1FoldLocal = inQueueX1.DeQue<T_X>();
246+ 
247+ LocalTensor<T_X> x2FoldLocal = inQueueX2.AllocTensor<T_X>();
248+ DataCopyPad(x2FoldLocal, x2Gm[xGmOffset2], xFoldDataCopyParams, padParams);
249+ inQueueX2.EnQue<T_X>(x2FoldLocal);
250+ x2FoldLocal = inQueueX2.DeQue<T_X>();
251+ 
252+ FoldBlockSquareAddVF(x1FoldLocal, x2FoldLocal, xFp32Tmp, baseN_);
253+ inQueueX1.FreeTensor(x1FoldLocal);
254+ inQueueX2.FreeTensor(x2FoldLocal);
255+ } else if (r == mainFoldCount_ && foldTail_ > 0) {
256+ xFoldDataCopyParams.blockLen = foldTail_ * sizeof(T_X);
257+ LocalTensor<T_X> x1FoldLocal = inQueueX1.AllocTensor<T_X>();
258+ DataCopyPad(x1FoldLocal, x1Gm[xGmOffset2], xFoldDataCopyParams, padParams);
259+ inQueueX1.EnQue<T_X>(x1FoldLocal);
260+ x1FoldLocal = inQueueX1.DeQue<T_X>();
261+ 
262+ LocalTensor<T_X> x2FoldLocal = inQueueX2.AllocTensor<T_X>();
263+ DataCopyPad(x2FoldLocal, x2Gm[xGmOffset2], xFoldDataCopyParams, padParams);
264+ inQueueX2.EnQue<T_X>(x2FoldLocal);
265+ x2FoldLocal = inQueueX2.DeQue<T_X>();
266+ 
267+ FoldBlockSquareAddVF(x1FoldLocal, x2FoldLocal, xFp32Tmp, foldTail_);
268+ inQueueX1.FreeTensor(x1FoldLocal);
269+ inQueueX2.FreeTensor(x2FoldLocal);
270+ }
271+ // reduce sum
272+ CalculateSquareReduceSum(xFp32Tmp, xFp32Tmp, baseN_);
273+ int64_t cacheId = GetCacheId(r);
274+ UpdateCache(cacheLocal, xFp32Tmp, cacheId, AR_RECOMPUTE_SUM_LEN);
275+ }
276+ 
277+ // final accumulated result to rstdLocal
278+ __local_mem__ float* dstPtr = (__local_mem__ float*)rstdLocal.GetPhyAddr();
279+ __local_mem__ float* cachePtr =
280+ (__local_mem__ float*)cacheLocal.GetPhyAddr() + resultCacheID_ * AR_RECOMPUTE_SUM_LEN;
281+ __VEC_SCOPE__
282+ {
283+ RegTensor<float> a;
284+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
285+ DataCopy<float, LoadDist::DIST_NORM>(a, cachePtr);
286+ DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);
287+ }
288+ }
289+ 
290+ __aicore__ inline void CalculateSquareReduceSum(
291+ LocalTensor<float>& xFp32Local, LocalTensor<float>& xReduceLocal, uint32_t reduceNum)
292+ {
293+ LocalTensor<float> binaryAddBuffTmp = binaryAddBuf.Get<float>();
294+ __local_mem__ float* xReduceUb = (__local_mem__ float*)xReduceLocal.GetPhyAddr();
295+ __local_mem__ float* tmpUb = (__local_mem__ float*)binaryAddBuffTmp.GetPhyAddr();
296+ __local_mem__ float* xFp32Tmp = (__local_mem__ float*)xFp32Local.GetPhyAddr();
297+ 
298+ if (reduceNum <= VL_F32) {
299+ CalculateSquareReduceSumLessThanVL(xFp32Tmp, xReduceUb, reduceNum);
300+ } else if (reduceNum <= VL_F32 + VL_F32) {
301+ CalculateSquareReduceSumLessThanTwoVL(xFp32Tmp, xReduceUb, reduceNum);
302+ } else if (reduceNum <= VL_F32 * VL_F32 * DIGIT_TWO) {
303+ CalculateSquareReduceSumCommon<DIGIT_ONE>(xFp32Tmp, xReduceUb, tmpUb, reduceNum);
304+ } else {
305+ CalculateSquareReduceSumCommon<DIGIT_TWO>(xFp32Tmp, xReduceUb, tmpUb, reduceNum);
306+ }
307+ }
308+ 
309+ __aicore__ inline void CalculateSquareReduceSumLessThanVL(
310+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, uint32_t reduceNum)
311+ {
312+ __VEC_SCOPE__
313+ {
314+ RegTensor<float> x, vMean, onesReg;
315+ uint32_t sreg0 = reduceNum;
316+ MaskReg pregLoop = UpdateMask<float>(sreg0);
317+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
318+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
319+ 
320+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregLoop, 0);
321+ AscendC::MicroAPI::ReduceSum(vMean, x, pregLoop);
322+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
323+ }
324+ }
325+ 
326+ __aicore__ inline void CalculateSquareReduceSumLessThanTwoVL(
327+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, uint32_t reduceNum)
328+ {
329+ uint32_t tailLen = reduceNum - VL_F32;
330+ __VEC_SCOPE__
331+ {
332+ RegTensor<float> x, xFold, sumReg, vMean, onesReg;
333+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
334+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
335+ MaskReg pregTail = UpdateMask<float>(tailLen);
336+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
337+ 
338+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregFull, 0);
339+ LoadTensorForDtypeTIn<float>(xFp32Tmp + VL_F32, xFold, pregTail, 0);
340+ AscendC::MicroAPI::ShiftLefts(
341+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregTail);
342+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
343+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
344+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
345+ }
346+ }
347+ 
348+ template <int32_t LAST_LOOP_NUMS>
349+ __aicore__ inline void CalculateSquareReduceSumCommon(
350+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, __local_mem__ float* tmpUb,
351+ uint32_t reduceNum)
352+ {
353+ uint32_t binaryAddQuotient = binAddQuotient_;
354+ uint16_t binaryAddQuotientLoop = (binaryAddQuotient + VL_F32 - 1) / VL_F32;
355+ uint32_t lastBinaryAddNum = binaryAddQuotient / VL_F32;
356+ 
357+ uint32_t binaryAddRemainder = reduceNum - binaryAddQuotient;
358+ uint16_t binaryAddRemainderCeilLoop = (binaryAddRemainder + VL_F32 - 1) / VL_F32;
359+ uint16_t binaryAddRemainderFloorLoop = binaryAddRemainder / VL_F32;
360+ __VEC_SCOPE__
361+ {
362+ RegTensor<float> x, xFold, sumReg, vMean, onesReg;
363+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
364+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
365+ MaskReg pregLoop;
366+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
367+ 
368+ for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; ++r) {
369+ uint32_t offset = r * VL_F32;
370+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregFull, offset);
371+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddQuotient, xFold, pregFull, offset);
372+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
373+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
374+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
375+ tmpUb + static_cast<uint32_t>(r), vMean, pregOne);
376+ }
377+ uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_F32;
378+ for (uint16_t r = 0;
379+ r < static_cast<uint16_t>(binaryAddRemainderCeilLoop - binaryAddRemainderFloorLoop); ++r) {
380+ pregLoop = UpdateMask<float>(sregRemainder);
381+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddRemainderFloorLoop * VL_F32, x, pregFull, 0);
382+ LoadTensorForDtypeTIn<float>(
383+ xFp32Tmp + binaryAddRemainderFloorLoop * VL_F32 + binaryAddQuotient, xFold, pregLoop, 0);
384+ AscendC::MicroAPI::ShiftLefts(
385+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregLoop);
386+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
387+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
388+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
389+ tmpUb + static_cast<uint32_t>(binaryAddRemainderFloorLoop), vMean, pregOne);
390+ }
391+ for (uint16_t r = 0;
392+ r < static_cast<uint16_t>(binaryAddQuotientLoop - binaryAddRemainderCeilLoop); ++r) {
393+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddRemainderCeilLoop * VL_F32, x, pregFull, 0);
394+ AscendC::MicroAPI::ReduceSum(vMean, x, pregFull);
395+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
396+ tmpUb + static_cast<uint32_t>(binaryAddRemainderCeilLoop + r), vMean, pregOne);
397+ }
398+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
399+ if constexpr (LAST_LOOP_NUMS == 1) {
400+ MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum);
401+ DataCopy(x, tmpUb);
402+ AscendC::MicroAPI::ReduceSum(vMean, x, pregLast);
403+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
404+ } else if constexpr (LAST_LOOP_NUMS == 2) {
405+ lastBinaryAddNum -= VL_F32;
406+ MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum);
407+ DataCopy(x, tmpUb);
408+ DataCopy(xFold, tmpUb + VL_F32);
409+ AscendC::MicroAPI::ShiftLefts(
410+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregLast);
411+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
412+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
413+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
414+ }
415+ }
416+ }
417+ 
418+ 
419+ template <bool hasBeta>
420+ __aicore__ inline void CalculateY(
421+ LocalTensor<float>& xFp32Local, LocalTensor<T_X>& yLocal, LocalTensor<float>& rstdLocal,
422+ uint32_t curN, uint32_t rowIdx)
423+ {
424+ __local_mem__ float* xFp32Tmp = (__local_mem__ float*)xFp32Local.GetPhyAddr();
425+ __local_mem__ T_GAMMA* gammaInUb = (__local_mem__ T_GAMMA*)gammaLocal_.GetPhyAddr();
426+ __local_mem__ T_X* yInUb = (__local_mem__ T_X*)yLocal.GetPhyAddr();
427+ __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr();
428+ __local_mem__ T_GAMMA* betaInUb;
429+ if constexpr (hasBeta) {
430+ betaInUb = (__local_mem__ T_GAMMA*)betaLocal_.GetPhyAddr();
431+ }
432+ 
433+ uint16_t loopCols = static_cast<uint16_t>((curN + VL_F32 - 1) / VL_F32);
434+ 
435+ __VEC_SCOPE__
436+ {
437+ RegTensor<float> xRegFp32, gammaRegFp32, rstdReg, betaRegFp32;
438+ MaskReg maskReg;
439+ 
440+ AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdInUb + rowIdx);
441+ uint32_t sregCount = curN;
442+ for (uint16_t r = 0; r < loopCols; ++r) {
443+ uint32_t offset = r * VL_F32;
444+ maskReg = UpdateMask<float>(sregCount);
445+ LoadTensorForDtypeTIn<float>(xFp32Tmp, xRegFp32, maskReg, offset);
446+ LoadTensorForDtypeTIn<T_GAMMA>(gammaInUb, gammaRegFp32, maskReg, offset);
447+ AscendC::MicroAPI::Mul(xRegFp32, xRegFp32, rstdReg, maskReg);
448+ AscendC::MicroAPI::Mul(xRegFp32, xRegFp32, gammaRegFp32, maskReg);
449+ if constexpr (hasBeta) {
450+ LoadTensorForDtypeTIn<T_GAMMA>(betaInUb, betaRegFp32, maskReg, offset);
451+ AscendC::MicroAPI::Add(xRegFp32, xRegFp32, betaRegFp32, maskReg);
452+ }
453+ StoreTensorForDtypeTOut<T_X>(yInUb, xRegFp32, maskReg, offset);
454+ }
455+ }
456+ }
457+ 
458+ template <AscendC::RoundMode toBf16RoundMode, AscendC::RoundMode roundMode>
459+ __aicore__ inline void DynamicMxQuantPhaseFP4(LocalTensor<T_X>& yLocal, uint64_t ubLoopIdx)
460+ {
461+ // each call processes a single row's single tile
462+ uint32_t curBlockNumInColAxis;
463+ uint32_t curN;
464+ if (ubLoopIdx == nUbLoops_ - 1) {
465+ curN = static_cast<uint32_t>(numColAlign_ - (nUbLoops_ - 1) * baseN_);
466+ curBlockNumInColAxis = CeilDiv(static_cast<uint64_t>(curN), static_cast<uint64_t>(mxBlockSize_));
467+ } else {
468+ curN = static_cast<uint32_t>(baseN_);
469+ curBlockNumInColAxis = CeilDiv(baseN_, mxBlockSize_);
470+ }
471+ 
472+ uint32_t totalScaleInUB = curBlockNumInColAxis;
473+ uint32_t totalCountInUB = curBlockNumInColAxis * mxBlockSize_;
474+ 
475+ uint16_t loopNum = (totalCountInUB + VL_B16 * DIGIT_TWO - 1) / (VL_B16 * DIGIT_TWO);
476+ uint16_t loopNumScale = (totalScaleInUB + VL_B16 - 1) / VL_B16;
477+ 
478+ LocalTensor<uint16_t> maxExpLocal = maxExpBuff.Get<uint16_t>();
479+ auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
480+ auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
481+ 
482+ LocalTensor<uint16_t> mxScaleLocal = mxScaleQueue.AllocTensor<uint16_t>();
483+ auto mxScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(mxScaleLocal.GetPhyAddr());
484+ 
485+ LocalTensor<uint16_t> halfScaleLocal = halfScaleBuff.Get<uint16_t>();
486+ auto halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
487+ 
488+ LocalTensor<int8_t> outLocal = outQueueQuantY.AllocTensor<int8_t>();
489+ auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());
490+ 
491+ MxQuantComputeMaxExpOCP<T_X>(srcAddr, maxExpAddr, loopNum);
492+ MxQuantComputeScaleOCP<T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);
493+ 
494+ srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
495+ halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
496+ 
497+ MxQuantComputeDataFP4<toBf16RoundMode, roundMode, T_X, T_Y>(
498+ srcAddr, halfScaleLocalAddr, outLocalAddr, totalCountInUB, loopNum);
499+ 
500+ outQueueQuantY.EnQue(outLocal);
501+ mxScaleQueue.EnQue(mxScaleLocal);
502+ }
503+ 
504+ __aicore__ inline void CopyInGammabeta(LocalTensor<T_GAMMA>& gammabetaLocal, int64_t offset, uint32_t len)
505+ {
506+ DataCopyExtParams copyParams{
507+ static_cast<uint16_t>(1),
508+ static_cast<uint32_t>(len * sizeof(T_GAMMA)),
509+ static_cast<uint32_t>(0),
510+ static_cast<uint32_t>(0),
511+ 0
512+ };
513+ DataCopyPadExtParams<T_GAMMA> padParams{false, static_cast<uint8_t>(0), static_cast<uint8_t>(0), static_cast<T_GAMMA>(0.0)};
514+ gammaLocal_ = gammabetaLocal;
515+ DataCopyPad<T_GAMMA>(gammaLocal_, gammaGm[offset], copyParams, padParams);
516+ if (betaFlag_ != 0) {
Z
Zzzz20484月20日

gammaLocal_ 和 betaLocal_ 是成员变量，在 CopyInGammabeta 函数中赋值后供后续 CalculateY 使用。建议在类定义中添加注释说明这两个成员变量的用途和生命周期，方便后续维护理解数据流。

likedislike
517+ betaLocal_ = gammabetaLocal[CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE) / sizeof(T_GAMMA)];
518+ DataCopyPad<T_GAMMA>(betaLocal_, betaGm[offset], copyParams, padParams);
519+ }
520+ }
521+ 
522+ __aicore__ inline void CopyInX(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
523+ {
524+ LocalTensor<T_X> xLocal1 = inQueueX1.AllocTensor<T_X>();
525+ LocalTensor<T_X> xLocal2 = inQueueX2.AllocTensor<T_X>();
526+ 
527+ DataCopyExtParams extParams{
528+ static_cast<uint16_t>(1),
529+ static_cast<uint32_t>(curN * sizeof(T_X)),
530+ static_cast<uint32_t>(0),
531+ static_cast<uint32_t>(0),
532+ 0
533+ };
534+ DataCopyPadExtParams<T_X> padParams{
535+ false, static_cast<uint8_t>(0), static_cast<uint8_t>(0), static_cast<T_X>(0.0)};
536+ 
537+ DataCopyPad(xLocal1, x1Gm[gmOffset], extParams, padParams);
538+ DataCopyPad(xLocal2, x2Gm[gmOffset], extParams, padParams);
539+ inQueueX1.EnQue(xLocal1);
540+ inQueueX2.EnQue(xLocal2);
541+ }
542+ 
543+ __aicore__ inline void CopyOutX(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
544+ {
545+ LocalTensor<T_X> xLocal = outQueueX.DeQue<T_X>();
546+ 
547+ DataCopyExtParams copyParams{
548+ static_cast<uint16_t>(1),
549+ static_cast<uint32_t>(curN * sizeof(T_X)),
550+ static_cast<uint32_t>(0),
551+ static_cast<uint32_t>(0),
552+ 0
553+ };
554+ DataCopyPad(xOutGm[gmOffset], xLocal, copyParams);
555+ outQueueX.FreeTensor(xLocal);
556+ }
557+ 
558+ __aicore__ inline void CopyOutQuantYFP4(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
559+ {
560+ LocalTensor<uint8_t> quantYLocal = outQueueQuantY.DeQue<uint8_t>();
561+ uint32_t fp4ByteLen = curN / DIGIT_TWO;
562+ uint32_t srcStride = 0;
563+ if ((ubLoopIdx == nUbLoops_ - 1) && (numCol_ != numColAlign_)) {
564+ srcStride = (numColAlign_ - numCol_) / DIGIT_TWO / UB_BLOCK_SIZE;
565+ }
566+ DataCopyExtParams copyParams{
567+ static_cast<uint16_t>(1),
568+ static_cast<uint32_t>(fp4ByteLen),
569+ static_cast<uint32_t>(srcStride),
570+ static_cast<uint32_t>(0),
571+ 0
572+ };
573+ DataCopyPad<uint8_t>(yFp4Gm[gmOffset / DIGIT_TWO], quantYLocal, copyParams);
574+ outQueueQuantY.FreeTensor(quantYLocal);
575+ }
576+ 
577+ __aicore__ inline void CopyOutMxScale(uint64_t rowIdx, uint64_t tileIdx)
578+ {
579+ LocalTensor<uint8_t> mxScaleLocal = mxScaleQueue.DeQue<uint8_t>();
580+ uint32_t curScaleSize;
581+ if (tileIdx == nUbLoops_ - 1) {
582+ uint32_t curN = static_cast<uint32_t>(numColAlign_ - (nUbLoops_ - 1) * baseN_);
583+ curScaleSize = CeilDiv(static_cast<uint64_t>(curN), mxBlockSize_);
584+ } else {
585+ curScaleSize = CeilDiv(baseN_, mxBlockSize_);
586+ }
587+ uint64_t scaleGmOffset = rowIdx * mxScaleSize_ + tileIdx * CeilDiv(baseN_, mxBlockSize_);
588+ DataCopyExtParams copyParams{
589+ static_cast<uint16_t>(1),
590+ static_cast<uint32_t>(curScaleSize),
591+ static_cast<uint32_t>(0),
592+ static_cast<uint32_t>(0),
593+ 0
594+ };
595+ DataCopyPad<uint8_t, PaddingMode::Compact>(mxScaleGm[scaleGmOffset], mxScaleLocal, copyParams);
596+ mxScaleQueue.FreeTensor(mxScaleLocal);
597+ }
598+ 
599+private:
600+ TPipe* pPipe = nullptr;
601+ 
602+ // Input Queues
603+ TQue<QuePosition::VECIN, 1> inQueueX1;
604+ TQue<QuePosition::VECIN, 1> inQueueX2;
605+ TQue<QuePosition::VECIN, 1> inQueueGammabeta;
606+ 
607+ LocalTensor<T_GAMMA> gammaLocal_;
608+ LocalTensor<T_GAMMA> betaLocal_;
609+ 
610+ // Output Queues
611+ TQue<QuePosition::VECOUT, 1> outQueueX;
612+ TQue<QuePosition::VECOUT, 1> outQueueRstd;
613+ TQue<QuePosition::VECOUT, 1> outQueueQuantY;
614+ TQue<QuePosition::VECOUT, 1> mxScaleQueue;
615+ 
616+ // TBuf
617+ TBuf<TPosition::VECCALC> xFp32Buf;
618+ TBuf<TPosition::VECCALC> cacheBuf;
619+ TBuf<TPosition::VECCALC> binaryAddBuf;
620+ TBuf<TPosition::VECCALC> yTmpBuf;
621+ TBuf<TPosition::VECCALC> maxExpBuff;
622+ TBuf<TPosition::VECCALC> halfScaleBuff;
623+ 
624+ // GM Tensors
625+ GlobalTensor<T_X> x1Gm, x2Gm, xOutGm;
626+ GlobalTensor<T_GAMMA> gammaGm, betaGm;
627+ GlobalTensor<float> rstdGm;
628+ GlobalTensor<uint8_t> yFp4Gm, mxScaleGm;
629+ 
630+ // Tiling Parameters
631+ uint64_t numCol_, numColAlign_;
632+ uint64_t blockFactor_, mLastCore_, mCurCore_;
633+ uint64_t baseN_, baseM_, baseNBlockSize_;
634+ uint64_t nUbLoops_;
635+ uint64_t binAddQuotient_, powerSplit_;
636+ uint64_t mainFoldCount_, foldTail_;
637+ int64_t resultCacheID_;
638+ float epsilon_, avgFactor_;
639+ uint64_t roundMode_, mxBlockSize_;
640+ int64_t scaleAlg_;
641+ uint64_t mxScaleSize_;
642+ uint32_t betaFlag_, rstdFlag_;
643+};
644+ 
645+} // namespace AddRmsNormDynamicMxQuant
646+#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP4_SPLIT_R_H
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_fp8_r_full_load.h+9-463
@@ -14,7 +14,6 @@
14 */14 */
15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H15#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H
16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H16#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H
17-#define FLOAT_OVERFLOW_MODE_CTRL 60
18 17 
19#include "add_rms_norm_dynamic_mx_quant_common.h"18#include "add_rms_norm_dynamic_mx_quant_common.h"
20 19 
@@ -57,18 +56,6 @@ public:
57 betaFlag_ = tiling->betaFlag;56 betaFlag_ = tiling->betaFlag;
58 rstdFlag_ = tiling->rstdFlag;57 rstdFlag_ = tiling->rstdFlag;
59 58 
60- vlForHalfNumber = platform::GetVRegSize() / sizeof(T_X); // VL in T_X elements
61- vlForFloat32Number = platform::GetVRegSize() / sizeof(float); // VL in float elements = 64
62- elementAfterReduce = platform::GetVRegSize() / UB_BLOCK_SIZE; // 8
63- 
64- if constexpr (IsSame<T_Y, fp8_e4m3fn_t>::value) {
65- f8Emax = MX_FP8_E4M3_MAX_EXP;
66- dtypeMax = MX_FP8_E4M3_MAX;
67- } else {
68- f8Emax = MX_FP8_E5M2_MAX_EXP;
69- dtypeMax = MX_FP8_E5M2_MAX;
70- }
71- 
72 // === Setup GM tensors ===59 // === Setup GM tensors ===
73 uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;60 uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;
74 x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, rowWork * numCol_);61 x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, rowWork * numCol_);
@@ -166,7 +153,7 @@ private:
166 CalculateSquareReduceSum(xFp32Local, xReduceLocal, curRows);153 CalculateSquareReduceSum(xFp32Local, xReduceLocal, curRows);
167 154 
168 // --- Rstd = 1/sqrt(mean + epsilon_) ---155 // --- Rstd = 1/sqrt(mean + epsilon_) ---
169- CalculateRstd(xReduceLocal, rstdLocal, curRows);156+ CalculateRstd(xReduceLocal, rstdLocal, curRows, avgFactor_, epsilon_);
170 outQueueRstd.EnQue<float>(rstdLocal);157 outQueueRstd.EnQue<float>(rstdLocal);
171 158 
172 // --- CopyOut rstd ---159 // --- CopyOut rstd ---
@@ -451,64 +438,6 @@ private:
451 }438 }
452 }439 }
453 }440 }
454- 
455- __aicore__ inline void CalculateRstd(
456- LocalTensor<float>& xReduceLocal, LocalTensor<float>& rstdLocal, uint32_t curRows)
457- {
458- static constexpr float POS_INF = 3.40282366920938E+38;
459- static constexpr float SCALAR1 = -0.5;
460- static constexpr float SCALAR2 = 1.5;
461- static constexpr float SCALAR3 = 0.5;
462- static constexpr float SCALAR0 = -99.99;
463- 
464- __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr();
465- __local_mem__ float* xReduceUb = (__local_mem__ float*)xReduceLocal.GetPhyAddr();
466- uint16_t loopRows = static_cast<uint16_t>((curRows + VL_F32 - 1) / VL_F32);
467- __VEC_SCOPE__
468- {
469- RegTensor<float> var, rstd, r, y, s, t, one, scalar1;
470- RegTensor<float> t1, t2, t3, t4, scalarInf, scalarZero;
471- MaskReg cmpRegZero;
472- MaskReg cmpRegInf;
473- MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
474- MaskReg pregLoop;
475- 
476- uint32_t sreg = static_cast<uint32_t>(curRows);
477- for (uint16_t i = 0; i < loopRows; ++i) {
478- pregLoop = UpdateMask<float>(sreg);
479- AscendC::MicroAPI::Duplicate(scalarInf, POS_INF, pregLoop);
480- AscendC::MicroAPI::Duplicate(scalarZero, float(0.0), pregLoop);
481- AscendC::MicroAPI::Duplicate(one, float(1.0), pregLoop);
482- AscendC::MicroAPI::Duplicate(scalar1, SCALAR3, pregLoop);
483- AscendC::MicroAPI::Duplicate(t1, SCALAR2, pregLoop);
484- AscendC::MicroAPI::Duplicate(s, float(1.0), pregLoop);
485- // rstd computation
486- AscendC::MicroAPI::DataCopy(var, xReduceUb + i * VL_F32);
487- AscendC::MicroAPI::Muls(var, var, avgFactor_, pregLoop);
488- AscendC::MicroAPI::Adds(var, var, epsilon_, pregLoop);
489- AscendC::MicroAPI::Maxs(var, var, SCALAR0, pregLoop);
490- AscendC::MicroAPI::Div(r, one, var, pregLoop);
491- AscendC::MicroAPI::Sqrt(y, r, pregLoop);
492- AscendC::MicroAPI::Muls(t, var, SCALAR1, pregLoop);
493- AscendC::MicroAPI::Mul(t, t, y, pregLoop);
494- AscendC::MicroAPI::Mula(t1, t, y, pregLoop);
495- AscendC::MicroAPI::Mul(rstd, y, t1, pregLoop);
496- AscendC::MicroAPI::Muls(t3, var, float(-1.0), pregLoop);
497- AscendC::MicroAPI::Mula(s, t3, r, pregLoop);
498- AscendC::MicroAPI::Muls(t4, rstd, float(-1.0), pregLoop);
499- AscendC::MicroAPI::Mula(r, t4, rstd, pregLoop);
500- AscendC::MicroAPI::Mula(s, var, r, pregLoop);
501- AscendC::MicroAPI::Mul(s, s, rstd, pregLoop);
502- AscendC::MicroAPI::Mula(rstd, s, scalar1, pregLoop);
503- AscendC::MicroAPI::CompareScalar(cmpRegZero, var, POS_INF, pregLoop);
504- AscendC::MicroAPI::Select(rstd, scalarZero, rstd, cmpRegZero);
505- AscendC::MicroAPI::CompareScalar(cmpRegInf, var, float(0.0), pregLoop);
506- AscendC::MicroAPI::Select(rstd, scalarInf, rstd, cmpRegInf);
507- AscendC::MicroAPI::DataCopy(rstdInUb + i * VL_F32, rstd, pregLoop);
508- }
509- }
510- }
511- 
512 template <AscendC::RoundMode roundMode>441 template <AscendC::RoundMode roundMode>
513 __aicore__ inline void DynamicMxQuantPhase(LocalTensor<T_X>& yLocal, uint32_t curRows)442 __aicore__ inline void DynamicMxQuantPhase(LocalTensor<T_X>& yLocal, uint32_t curRows)
514 {443 {
@@ -517,9 +446,9 @@ private:
517 uint32_t totalScaleInUB = curRows * blockNumInColAxis_;446 uint32_t totalScaleInUB = curRows * blockNumInColAxis_;
518 uint32_t totalCountInUB = curRows * blockNumInColAxis_ * mxBlockSize_;447 uint32_t totalCountInUB = curRows * blockNumInColAxis_ * mxBlockSize_;
519 448 
520- uint16_t loopNum = (totalCountInUB + vlForHalfNumber * DIGIT_TWO - 1) / (vlForHalfNumber * DIGIT_TWO);449+ uint16_t loopNum = (totalCountInUB + VL_B16 * DIGIT_TWO - 1) / (VL_B16 * DIGIT_TWO);
521- uint16_t loopNumScale = (totalScaleInUB + vlForHalfNumber - 1) / vlForHalfNumber;450+ uint16_t loopNumScale = (totalScaleInUB + VL_B16 - 1) / VL_B16;
522- uint16_t loopNumScale4NV = (totalScaleInUB + vlForFloat32Number - 1) / vlForFloat32Number;451+ uint16_t loopNumScale4NV = (totalScaleInUB + VL_F32 - 1) / VL_F32;
523 452 
524 auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());453 auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
525 auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());454 auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
@@ -534,397 +463,23 @@ private:
534 auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());463 auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());
535 maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());464 maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
536 if (scaleAlg_ == 0) {465 if (scaleAlg_ == 0) {
537- MxQuantComputeMaxExpOCP(srcAddr, maxExpAddr, totalCountInUB, loopNum);466+ MxQuantComputeMaxExpOCP<T_X>(srcAddr, maxExpAddr, loopNum);
538- MxQuantComputeScaleOCP(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);467+ MxQuantComputeScaleOCP<T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);
539 } else {468 } else {
540- MxQuantComputeMaxExpcuBLAS(srcAddr, maxExpAddr, totalCountInUB, loopNum);469+ MxQuantComputeMaxExpcuBLAS<T_X>(srcAddr, maxExpAddr, loopNum);
541- MxQuantComputeScalecuBLAS(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale4NV);470+ MxQuantComputeScalecuBLAS<T_X, T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale4NV);
542 }471 }
543 472 
544 srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());473 srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
545 halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());474 halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
546 475 
547- MxQuantComputeData<roundMode>(srcAddr, halfScaleLocalAddr, outLocalAddr, totalCountInUB, loopNum);476+ MxQuantComputeData<roundMode, T_X, T_Y>(srcAddr, halfScaleLocalAddr, outLocalAddr, loopNum);
548 477 
549 outQueueQuantY.EnQue(outLocal);478 outQueueQuantY.EnQue(outLocal);
550 mxScaleQueue.EnQue(mxScaleLocal);479 mxScaleQueue.EnQue(mxScaleLocal);
551 return;480 return;
552 }481 }
553 482 
554- __aicore__ inline void MxQuantComputeMaxExpOCP(
555- __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint32_t totalCountInUB, uint16_t loopNum)
556- {
557- __VEC_SCOPE__
558- {
559- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
560- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
561- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
562- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
563- AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect0;
564- AscendC::MicroAPI::RegTensor<uint16_t> vdExpSelect1;
565- AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract0;
566- AscendC::MicroAPI::RegTensor<uint16_t> vdExpExtract1;
567- AscendC::MicroAPI::RegTensor<uint16_t> expMaskBF16;
568- AscendC::MicroAPI::Duplicate(expMaskBF16, MAX_EXP_FOR_BF16);
569- AscendC::MicroAPI::RegTensor<uint16_t> invalidMaskFP16;
570- AscendC::MicroAPI::Duplicate(invalidMaskFP16, INVALID_FLOAT16);
571- AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
572- AscendC::MicroAPI::MaskReg scaleMask1;
573- AscendC::MicroAPI::MaskReg scaleMask2;
574- AscendC::MicroAPI::MaskReg invalidDataMask0;
575- AscendC::MicroAPI::MaskReg invalidDataMask1;
576- AscendC::MicroAPI::UnalignReg u1;
577- static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
578- AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
579- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
580- for (uint16_t i = 0; i < loopNum; i++) {
581- scaleMask1 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
582- scaleMask2 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
583- AscendC::MicroAPI::DataCopy<
584- T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE, AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(
585- vdExp0, vdExp1, srcAddr, vlForHalfNumber * DIGIT_TWO);
586- if constexpr (IsSame<T_X, half>::value) {
587- AscendC::MicroAPI::And(
588- vdExpSelect0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, invalidMaskFP16, scaleMask1);
589- AscendC::MicroAPI::And(
590- vdExpSelect1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, invalidMaskFP16, scaleMask1);
591- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
592- invalidDataMask0, vdExpSelect0, invalidMaskFP16, scaleMask1);
593- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(
594- invalidDataMask1, vdExpSelect1, invalidMaskFP16, scaleMask1);
595- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp0BF16, vdExp0, scaleMask1);
596- AscendC::MicroAPI::Cast<bfloat16_t, T_X, castTraitHalf2Bf16>(vdExp1BF16, vdExp1, scaleMask1);
597- AscendC::MicroAPI::And(
598- vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0BF16, expMaskBF16, scaleMask1);
599- AscendC::MicroAPI::And(
600- vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1BF16, expMaskBF16, scaleMask1);
601- AscendC::MicroAPI::Select<uint16_t>(vdExpExtract0, vdExpExtract0, expMaskBF16, invalidDataMask0);
602- AscendC::MicroAPI::Select<uint16_t>(vdExpExtract1, vdExpExtract1, expMaskBF16, invalidDataMask1);
603- } else {
604- AscendC::MicroAPI::And(
605- vdExpExtract0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, expMaskBF16, scaleMask1);
606- AscendC::MicroAPI::And(
607- vdExpExtract1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, expMaskBF16, scaleMask1);
608- }
609- 
610- AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, scaleMask1);
611- AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, scaleMask1);
612- 
613- AscendC::MicroAPI::DataCopyUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
614- maxExpAddr, vdMaxExp, u1, elementAfterReduce);
615- }
616- AscendC::MicroAPI::DataCopyUnAlignPost(maxExpAddr, u1, 0);
617- }
618- return;
619- }
620-
621- __aicore__ inline void MxQuantComputeScaleOCP(
622- __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, __ubuf__ uint16_t* halfScaleLocalAddr,
623- uint32_t totalScaleInUB, uint16_t loopNumScale)
624- {
625- __VEC_SCOPE__
626- {
627- AscendC::MicroAPI::RegTensor<uint16_t> expMask;
628- AscendC::MicroAPI::Duplicate(expMask, MAX_EXP_FOR_BF16);
629- AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
630- 
631- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
632- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
633- 
634- AscendC::MicroAPI::MaskReg cmpResult;
635- AscendC::MicroAPI::MaskReg zeroMask;
636- AscendC::MicroAPI::MaskReg cmpResultSub;
637- AscendC::MicroAPI::MaskReg preMaskScale;
638- AscendC::MicroAPI::RegTensor<uint16_t> maxExpValue;
639- AscendC::MicroAPI::Duplicate(maxExpValue, f8Emax);
640- AscendC::MicroAPI::RegTensor<uint16_t> sharedExp;
641- AscendC::MicroAPI::RegTensor<uint16_t> scaleValue;
642- AscendC::MicroAPI::RegTensor<uint16_t> scaleBias;
643- AscendC::MicroAPI::Duplicate(scaleBias, BF16_EXP_BIAS);
644- AscendC::MicroAPI::RegTensor<uint16_t> halfScale;
645- AscendC::MicroAPI::RegTensor<uint16_t> fp8NanRegTensor;
646- AscendC::MicroAPI::Duplicate(fp8NanRegTensor, MAX_EXP_FOR_FP8);
647- AscendC::MicroAPI::RegTensor<uint16_t> zeroRegTensor;
648- AscendC::MicroAPI::Duplicate(zeroRegTensor, 0);
649- AscendC::MicroAPI::RegTensor<uint16_t> nanRegTensor;
650- AscendC::MicroAPI::Duplicate(nanRegTensor, NAN_CUSTOMIZATION);
651- AscendC::MicroAPI::MaskReg invalidDataMask;
652- AscendC::MicroAPI::MaskReg specialDataMask;
653- AscendC::MicroAPI::RegTensor<uint16_t> specialExpRegTensor;
654- AscendC::MicroAPI::Duplicate(specialExpRegTensor, SPECIAL_EXP_THRESHOLD);
655- for (uint16_t i = 0; i < loopNumScale; i++) {
656- preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);
657- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
658- vdMaxExp, maxExpAddr, vlForHalfNumber);
659- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(cmpResult, vdMaxExp, expMask, preMaskScale); // INF/NAN
660- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(zeroMask, vdMaxExp, zeroRegTensor, preMaskScale);
661- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::LE>(invalidDataMask, vdMaxExp, maxExpValue, preMaskScale);
662- 
663- AscendC::MicroAPI::Select<uint16_t>(vdMaxExp, maxExpValue, vdMaxExp, invalidDataMask);
664- 
665- AscendC::MicroAPI::Sub(sharedExp, vdMaxExp, maxExpValue, preMaskScale);
666- AscendC::MicroAPI::ShiftRights(scaleValue, sharedExp, SHR_NUM_FOR_BF16, preMaskScale);
667- 
668- AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, fp8NanRegTensor, cmpResult);
669- AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, zeroRegTensor, zeroMask);
670- 
671- AscendC::MicroAPI::DataCopy<
672- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
673- AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
674- mxScaleLocalAddr, scaleValue, vlForHalfNumber / DIGIT_TWO, preMaskScale);
675- 
676- AscendC::MicroAPI::Compare<uint16_t, CMPMODE::EQ>(specialDataMask, sharedExp, scaleBias, preMaskScale);
677- AscendC::MicroAPI::Sub(halfScale, scaleBias, sharedExp, preMaskScale);
678- AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, nanRegTensor, cmpResult);
679- AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, zeroRegTensor, zeroMask);
680- AscendC::MicroAPI::Select<uint16_t>(halfScale, specialExpRegTensor, halfScale, specialDataMask);
681- 
682- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
683- halfScaleLocalAddr, halfScale, vlForHalfNumber, preMaskScale);
684- }
685- }
686- return;
687- }
688- 
689- __aicore__ inline void MxQuantComputeMaxExpcuBLAS(
690- __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint32_t totalCountInUB, uint16_t loopNum)
691- {
692- __VEC_SCOPE__
693- {
694- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
695- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
696- AscendC::MicroAPI::RegTensor<uint16_t> absMask16Bit;
697- AscendC::MicroAPI::Duplicate(absMask16Bit, ABS_MASK_FOR_16BIT);
698- AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
699- AscendC::MicroAPI::MaskReg scaleMask1;
700- AscendC::MicroAPI::UnalignReg u1;
701- for (uint16_t i = 0; i < loopNum; i++) {
702- scaleMask1 = AscendC::MicroAPI::UpdateMask<T_X>(totalCountInUB);
703- AscendC::MicroAPI::DataCopy<
704- T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE, AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(
705- vdExp0, vdExp1, srcAddr, vlForHalfNumber * DIGIT_TWO);
706- AscendC::MicroAPI::And(
707- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
708- absMask16Bit, scaleMask1);
709- AscendC::MicroAPI::And(
710- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1,
711- absMask16Bit, scaleMask1);
712- AscendC::MicroAPI::Max(
713- vdMaxExp, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
714- (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, scaleMask1);
715- AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, scaleMask1);
716- AscendC::MicroAPI::DataCopyUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
717- maxExpAddr, vdMaxExp, u1, elementAfterReduce);
718- }
719- AscendC::MicroAPI::DataCopyUnAlignPost(maxExpAddr, u1, 0);
720- }
721- return;
722- }
723- 
724- __aicore__ inline void MxQuantComputeScalecuBLAS(
725- __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, __ubuf__ uint16_t* halfScaleLocalAddr,
726- uint32_t totalScaleInUB, uint16_t loopNumScale4NV)
727- {
728- __VEC_SCOPE__
729- {
730- AscendC::MicroAPI::RegTensor<uint16_t> max16;
731- AscendC::MicroAPI::RegTensor<uint32_t> max32;
732- AscendC::MicroAPI::RegTensor<uint32_t> exp32;
733- AscendC::MicroAPI::RegTensor<uint32_t> man32;
734- AscendC::MicroAPI::RegTensor<uint32_t> normalExp32;
735- AscendC::MicroAPI::RegTensor<uint32_t> expAddOne32;
736- AscendC::MicroAPI::RegTensor<uint32_t> extractExp;
737- AscendC::MicroAPI::RegTensor<uint16_t> expOut;
738- AscendC::MicroAPI::RegTensor<uint32_t> halfScale;
739- AscendC::MicroAPI::RegTensor<uint16_t> recExpOut;
740- 
741- AscendC::MicroAPI::RegTensor<uint32_t> invMax;
742- AscendC::MicroAPI::Duplicate(invMax, dtypeMax);
743- AscendC::MicroAPI::RegTensor<uint32_t> manMaskFP32;
744- AscendC::MicroAPI::Duplicate(manMaskFP32, MAN_MASK_FLOAT);
745- AscendC::MicroAPI::RegTensor<uint32_t> expMask;
746- AscendC::MicroAPI::Duplicate(expMask, MAX_EXP_FOR_FP32);
747- AscendC::MicroAPI::RegTensor<uint32_t> zeroRegTensor32;
748- AscendC::MicroAPI::Duplicate(zeroRegTensor32, 0);
749- AscendC::MicroAPI::RegTensor<uint32_t> scaleBias;
750- AscendC::MicroAPI::Duplicate(scaleBias, FP32_EXP_BIAS_CUBLAS);
751- AscendC::MicroAPI::RegTensor<uint32_t> nanRegTensor;
752- AscendC::MicroAPI::Duplicate(nanRegTensor, NAN_CUSTOMIZATION_PACK);
753- AscendC::MicroAPI::RegTensor<uint32_t> fp8NanRegTensor;
754- AscendC::MicroAPI::Duplicate(fp8NanRegTensor, MAX_EXP_FOR_FP8_IN_FP32);
755- 
756- AscendC::MicroAPI::MaskReg cmpResult;
757- AscendC::MicroAPI::MaskReg zeroMask;
758- AscendC::MicroAPI::MaskReg p0;
759- AscendC::MicroAPI::MaskReg p1;
760- AscendC::MicroAPI::MaskReg p2;
761- AscendC::MicroAPI::MaskReg preMaskScale;
762- AscendC::MicroAPI::MaskReg maskHalf;
763- preMaskScale = AscendC::MicroAPI::CreateMask<uint32_t>();
764- maskHalf = AscendC::MicroAPI::CreateMask<uint16_t>();
765- static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Float = {
766- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
767- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
768- for (uint16_t i = 0; i < loopNumScale4NV; i++) {
769- // preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);
770- AscendC::MicroAPI::DataCopy<
771- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
772- AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(max16, maxExpAddr, vlForFloat32Number);
773- 
774- AscendC::MicroAPI::Cast<float, T_X, castTraitHalf2Float>(
775- (AscendC::MicroAPI::RegTensor<float>&)max32, (AscendC::MicroAPI::RegTensor<T_X>&)max16, preMaskScale);
776- AscendC::MicroAPI::Compare<uint32_t, CMPMODE::LT>(cmpResult, max32, expMask, preMaskScale);
777- AscendC::MicroAPI::Compare<uint32_t, CMPMODE::NE>(zeroMask, max32, zeroRegTensor32, preMaskScale);
778- 
779- AscendC::MicroAPI::Mul(
780- (AscendC::MicroAPI::RegTensor<float>&)max32, (AscendC::MicroAPI::RegTensor<float>&)max32,
781- (AscendC::MicroAPI::RegTensor<float>&)invMax, preMaskScale);
782- AscendC::MicroAPI::ShiftRights(exp32, max32, SHR_NUM_FOR_FP32, preMaskScale);
783- AscendC::MicroAPI::And(man32, max32, manMaskFP32, preMaskScale);
784- 
785- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p0, exp32, zeroForAll, preMaskScale);
786- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::LT>(p1, exp32, Exp254, preMaskScale);
787- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, zeroForAll, preMaskScale);
788- AscendC::MicroAPI::MaskAnd(p0, p0, p1, preMaskScale);
789- AscendC::MicroAPI::MaskAnd(p0, p0, p2, preMaskScale);
790- 
791- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::EQ>(p1, exp32, zeroForAll, preMaskScale);
792- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, halfForMan, preMaskScale);
793- AscendC::MicroAPI::MaskAnd(p1, p1, p2, preMaskScale);
794- AscendC::MicroAPI::MaskOr(p0, p0, p1, preMaskScale);
795- 
796- AscendC::MicroAPI::Adds(expAddOne32, exp32, 1, preMaskScale);
797- AscendC::MicroAPI::Select(extractExp, expAddOne32, exp32, p0);
798- AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, fp8NanRegTensor, cmpResult);
799- AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, zeroRegTensor32, zeroMask);
800- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(expOut, extractExp);
801- 
802- AscendC::MicroAPI::DataCopy<
803- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
804- AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(
805- mxScaleLocalAddr, expOut, vlForFloat32Number / DIGIT_TWO, maskHalf);
806- 
807- AscendC::MicroAPI::ShiftLefts(extractExp, extractExp, SHR_NUM_FOR_BF16, preMaskScale);
808- AscendC::MicroAPI::Sub(halfScale, scaleBias, extractExp, preMaskScale);
809- AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, nanRegTensor, cmpResult);
810- AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, zeroRegTensor32, zeroMask);
811- AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(recExpOut, halfScale);
812- 
813- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
814- halfScaleLocalAddr, recExpOut, vlForFloat32Number, maskHalf);
815- }
816- }
817- return;
818- }
819- 
820- template <AscendC::RoundMode roundMode>
821- __aicore__ inline void MxQuantComputeData(
822- __ubuf__ T_X* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr, __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB,
823- uint16_t loopNum)
824- {
825- __VEC_SCOPE__
826- {
827- AscendC::MicroAPI::MaskReg dataMask1;
828- AscendC::MicroAPI::MaskReg dataMask2;
829- AscendC::MicroAPI::MaskReg dataMask3;
830- AscendC::MicroAPI::MaskReg dataMask4;
831- AscendC::MicroAPI::MaskReg dataMask5;
832- AscendC::MicroAPI::MaskReg maskAll =
833- AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
834- AscendC::MicroAPI::RegTensor<uint16_t> halfScaleForMul;
835- AscendC::MicroAPI::RegTensor<float> floatScaleForMul;
836- AscendC::MicroAPI::RegTensor<T_X> vdExp0;
837- AscendC::MicroAPI::RegTensor<T_X> vdExp1;
838- AscendC::MicroAPI::RegTensor<T_X> vdExp0Convert;
839- AscendC::MicroAPI::RegTensor<T_X> vdExp1Convert;
840- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp0BF16;
841- AscendC::MicroAPI::RegTensor<bfloat16_t> vdExp1BF16;
842- AscendC::MicroAPI::RegTensor<float> vdExp0FP32Zero;
843- AscendC::MicroAPI::RegTensor<float> vdExp0FP32One;
844- AscendC::MicroAPI::RegTensor<float> vdExp1FP32Zero;
845- AscendC::MicroAPI::RegTensor<float> vdExp1FP32One;
846- AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP8Zero;
847- AscendC::MicroAPI::RegTensor<T_Y> vdExp0FP8One;
848- AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP8Zero;
849- AscendC::MicroAPI::RegTensor<T_Y> vdExp1FP8One;
850- AscendC::MicroAPI::RegTensor<bfloat16_t> vdBF16Exp0FP4;
851- AscendC::MicroAPI::RegTensor<bfloat16_t> vdBF16Exp1FP4;
852- static constexpr AscendC::MicroAPI::CastTrait castTrait = {
853- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
854- AscendC::MicroAPI::MaskMergeMode::ZEROING, roundMode};
855- static constexpr AscendC::MicroAPI::CastTrait castTraitHalf2Bf16 = {
856- AscendC::MicroAPI::RegLayout::UNKNOWN, AscendC::MicroAPI::SatMode::UNKNOWN,
857- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_TRUNC};
858- static constexpr AscendC::MicroAPI::CastTrait castTraitZero = {
859- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN,
860- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
861- static constexpr AscendC::MicroAPI::CastTrait castTraitOne = {
862- AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::UNKNOWN,
863- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
864- static constexpr AscendC::MicroAPI::CastTrait castTrait32to8 = {
865- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::SAT,
866- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
867- static constexpr AscendC::MicroAPI::CastTrait castTrait32to80 = {
868- AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::SAT,
869- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
870- static constexpr AscendC::MicroAPI::CastTrait castTrait32to81 = {
871- AscendC::MicroAPI::RegLayout::ONE, AscendC::MicroAPI::SatMode::SAT,
872- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
873- static constexpr AscendC::MicroAPI::CastTrait castTrait32to82 = {
874- AscendC::MicroAPI::RegLayout::TWO, AscendC::MicroAPI::SatMode::SAT,
875- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
876- static constexpr AscendC::MicroAPI::CastTrait castTrait32to83 = {
877- AscendC::MicroAPI::RegLayout::THREE, AscendC::MicroAPI::SatMode::SAT,
878- AscendC::MicroAPI::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
879- dataMask1 = AscendC::MicroAPI::CreateMask<T_X>();
880- dataMask2 = AscendC::MicroAPI::CreateMask<T_X>();
881- dataMask3 = AscendC::MicroAPI::CreateMask<T_X>();
882- dataMask4 = AscendC::MicroAPI::CreateMask<T_X>();
883- dataMask5 = AscendC::MicroAPI::CreateMask<T_Y>();
884- for (uint16_t i = 0; i < loopNum; i++) {
885- AscendC::MicroAPI::DataCopy<
886- T_X, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE, AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(
887- vdExp0, vdExp1, srcAddr, vlForHalfNumber * DIGIT_TWO);
888- AscendC::MicroAPI::DataCopy<
889- uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE, AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(
890- halfScaleForMul, halfScaleLocalAddr, elementAfterReduce);
891- if constexpr (IsSame<T_X, half>::value) {
892- AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);
893- AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);
894- AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitZero>(
895- floatScaleForMul, (AscendC::MicroAPI::RegTensor<bfloat16_t>&)halfScaleForMul, maskAll);
896- AscendC::MicroAPI::Mul(vdExp0FP32Zero, vdExp0FP32Zero, floatScaleForMul, dataMask3);
897- AscendC::MicroAPI::Mul(vdExp0FP32One, vdExp0FP32One, floatScaleForMul, dataMask4);
898- 
899- AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp1FP32Zero, vdExp1, dataMask1);
900- AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp1FP32One, vdExp1, dataMask1);
901- AscendC::MicroAPI::Mul(vdExp1FP32Zero, vdExp1FP32Zero, floatScaleForMul, dataMask3);
902- AscendC::MicroAPI::Mul(vdExp1FP32One, vdExp1FP32One, floatScaleForMul, dataMask4);
903- } else {
904- AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
905- AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor<T_X>&)halfScaleForMul, dataMask1);
906- 
907- AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);
908- AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);
909- AscendC::MicroAPI::Cast<float, T_X, castTraitZero>(vdExp1FP32Zero, vdExp1, dataMask2);
910- AscendC::MicroAPI::Cast<float, T_X, castTraitOne>(vdExp1FP32One, vdExp1, dataMask2);
911- }
912- AscendC::MicroAPI::Cast<T_Y, float, castTrait32to80>(vdExp0FP8Zero, vdExp0FP32Zero, dataMask3);
913- AscendC::MicroAPI::Cast<T_Y, float, castTrait32to82>(vdExp0FP8One, vdExp0FP32One, dataMask3);
914- AscendC::MicroAPI::Cast<T_Y, float, castTrait32to81>(vdExp1FP8Zero, vdExp1FP32Zero, dataMask4);
915- AscendC::MicroAPI::Cast<T_Y, float, castTrait32to83>(vdExp1FP8One, vdExp1FP32One, dataMask4);
916-
917- AscendC::MicroAPI::Add((AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8One, dataMask5);
918- AscendC::MicroAPI::Add((AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8Zero, dataMask5);
919- AscendC::MicroAPI::Add((AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero, (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8One, dataMask5);
920- 
921- AscendC::MicroAPI::DataCopy<
922- int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE, AscendC::MicroAPI::StoreDist::DIST_NORM_B8>(
923- outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP8Zero, OUT_ALL, dataMask5);
924- }
925- }
926- return;
927- }
928 483 
929 __aicore__ inline void MxQuantDeletePadData(484 __aicore__ inline void MxQuantDeletePadData(
930 __ubuf__ int8_t* outLocalAddr, __ubuf__ int8_t* outBufferLocalAddr, uint16_t loopNum, uint32_t inputUpdateStride,485 __ubuf__ int8_t* outLocalAddr, __ubuf__ int8_t* outBufferLocalAddr, uint16_t loopNum, uint32_t inputUpdateStride,
@@ -1079,15 +634,6 @@ private:
1079 uint64_t mxScaleSize_;634 uint64_t mxScaleSize_;
1080 uint32_t betaFlag_;635 uint32_t betaFlag_;
1081 uint32_t rstdFlag_;636 uint32_t rstdFlag_;
1082- 
1083- uint16_t f8Emax;
1084- uint32_t vlForHalfNumber;
1085- uint32_t vlForFloat32Number;
1086- uint32_t dtypeMax;
1087- uint16_t elementAfterReduce;
1088- uint32_t zeroForAll = 0x00000000;
1089- uint32_t Exp254 = 0x000000fe;
1090- uint32_t halfForMan = 0x00400000;
1091};637};
1092} // namespace AddRmsNormDynamicMxQuant638} // namespace AddRmsNormDynamicMxQuant
1093#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H639#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_R_FULL_LOAD_H
Anorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_fp8_split_r.h+642-0
@@ -0,0 +1,642 @@
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 add_rms_norm_dynamic_mx_quant_fp8_split_r.h
13+ * \brief
14+ */
15+#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_SPLIT_R_H
16+#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_SPLIT_R_H
17+ 
18+#include "add_rms_norm_dynamic_mx_quant_common.h"
19+ 
20+namespace AddRmsNormDynamicMxQuant {
21+ 
22+template <typename T_X, typename T_GAMMA, typename T_Y>
23+class AddRmsNormDynamicMxQuantFP8SplitR {
24+public:
25+ __aicore__ inline AddRmsNormDynamicMxQuantFP8SplitR(TPipe* pipe)
26+ {
27+ pPipe = pipe;
28+ }
29+ 
30+ __aicore__ inline void Init(
31+ GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y,
32+ GM_ADDR x, GM_ADDR mxscale, GM_ADDR workspace, GM_ADDR rstd,
33+ const AddRmsNormDynamicMxQuantSplitRTilingData* tiling)
34+ {
35+ #if (__NPU_ARCH__ == 3510)
36+ AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
37+ #endif
38+ ASSERT(GetBlockNum() != 0 && "Block dim can not be zero!");
39+ 
40+ numCol_ = tiling->numCol;
41+ numColAlign_ = tiling->numColAlign;
42+ blockFactor_ = tiling->blockFactor;
43+ mLastCore_ = tiling->mLastCore;
44+ baseN_ = tiling->baseN;
45+ baseNBlockSize_ = tiling->baseNBlockSize;
46+ baseM_ = tiling->baseM;
47+ nUbLoops_ = tiling->nUbLoops;
48+ binAddQuotient_ = tiling->binAddQuotient;
49+ powerSplit_ = tiling->powerSplit;
50+ mainFoldCount_ = tiling->mainFoldCount;
51+ foldTail_ = tiling->foldTail;
52+ epsilon_ = tiling->epsilon;
53+ avgFactor_ = tiling->avgFactor;
54+ roundMode_ = tiling->roundMode;
55+ mxBlockSize_ = tiling->mxBlockSize;
56+ scaleAlg_ = tiling->scaleAlg;
57+ mxScaleSize_ = tiling->mxScaleSize;
58+ betaFlag_ = tiling->betaFlag;
59+ rstdFlag_ = tiling->rstdFlag;
60+ 
61+ resultCacheID_ = GetCacheId(powerSplit_ - 1);
62+ mCurCore_ = (GetBlockIdx() == GetBlockNum() - 1) ? mLastCore_ : blockFactor_;
63+ 
64+ // === Setup GM tensors ===
65+ uint64_t blockOffset = GetBlockIdx() * blockFactor_ * numCol_;
66+ x1Gm.SetGlobalBuffer((__gm__ T_X*)x1 + blockOffset, mCurCore_ * numCol_);
67+ x2Gm.SetGlobalBuffer((__gm__ T_X*)x2 + blockOffset, mCurCore_ * numCol_);
68+ gammaGm.SetGlobalBuffer((__gm__ T_GAMMA*)gamma, numCol_);
69+ if (betaFlag_ != 0) {
70+ betaGm.SetGlobalBuffer((__gm__ T_GAMMA*)beta, numCol_);
71+ }
72+ xOutGm.SetGlobalBuffer((__gm__ T_X*)x + blockOffset, mCurCore_ * numCol_);
73+ if (rstdFlag_ != 0) {
74+ rstdGm.SetGlobalBuffer((__gm__ float*)rstd + GetBlockIdx() * blockFactor_, blockFactor_);
75+ }
76+ yFp8Gm.SetGlobalBuffer((__gm__ uint8_t*)y + blockOffset, mCurCore_ * numCol_);
77+ mxScaleGm.SetGlobalBuffer(
78+ (__gm__ uint8_t*)mxscale + GetBlockIdx() * blockFactor_ * mxScaleSize_, mCurCore_ * mxScaleSize_);
79+ 
80+ // === Compute buffer sizes ===
81+ uint64_t xBufSize = CeilAlign(baseN_ * sizeof(T_X), UB_BLOCK_SIZE);
82+ uint64_t xFp32BufSize = CeilAlign(baseN_ * sizeof(float), UB_BLOCK_SIZE);
83+ uint64_t yTmpBufSize = CeilAlign(baseN_ * sizeof(T_X), UB_BLOCK_SIZE);
84+ uint64_t rstdBufSize = CeilAlign(baseM_ * sizeof(float), UB_BLOCK_SIZE);
85+ uint64_t cacheBufSize = CeilAlign(
86+ static_cast<uint64_t>((resultCacheID_ + 1) * sizeof(float)) * AR_RECOMPUTE_SUM_LEN, UB_BLOCK_SIZE);
87+ uint64_t binaryAddBufSize = CeilAlign(VL_F32 * DIGIT_TWO * sizeof(float), UB_BLOCK_SIZE);
88+ uint64_t quantYBufSize = CeilAlign(baseN_ * sizeof(T_Y), UB_BLOCK_SIZE);
89+ uint64_t maxExpBufSize = CeilAlign(baseNBlockSize_ * sizeof(uint16_t), UB_BLOCK_SIZE);
90+ uint64_t halfScaleBufSize = maxExpBufSize;
91+ uint64_t scaleBufSize = CeilAlign(baseNBlockSize_ * sizeof(uint8_t), UB_BLOCK_SIZE);
92+ 
93+ // === Init buffers ===
94+ pPipe->InitBuffer(inQueueX1, DOUBLE_BUFFER_NUM, xBufSize);
95+ pPipe->InitBuffer(inQueueX2, DOUBLE_BUFFER_NUM, xBufSize);
96+ if (betaFlag_ != 0) {
97+ pPipe->InitBuffer(
98+ inQueueGammabeta, DOUBLE_BUFFER_NUM, DIGIT_TWO * CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE));
99+ } else {
100+ pPipe->InitBuffer(inQueueGammabeta, DOUBLE_BUFFER_NUM, CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE));
101+ }
102+ 
103+ pPipe->InitBuffer(outQueueX, DOUBLE_BUFFER_NUM, xBufSize);
104+ pPipe->InitBuffer(outQueueRstd, DOUBLE_BUFFER_NUM, rstdBufSize);
105+ pPipe->InitBuffer(outQueueQuantY, DOUBLE_BUFFER_NUM, quantYBufSize);
106+ pPipe->InitBuffer(mxScaleQueue, DOUBLE_BUFFER_NUM, scaleBufSize);
107+ 
108+ pPipe->InitBuffer(xFp32Buf, xFp32BufSize);
109+ pPipe->InitBuffer(yTmpBuf, yTmpBufSize);
110+ pPipe->InitBuffer(cacheBuf, cacheBufSize);
111+ pPipe->InitBuffer(binaryAddBuf, binaryAddBufSize);
112+ pPipe->InitBuffer(maxExpBuff, maxExpBufSize);
113+ pPipe->InitBuffer(halfScaleBuff, halfScaleBufSize);
114+ }
115+ 
116+ __aicore__ inline void Process()
117+ {
118+ uint32_t mCnt = CeilDiv(mCurCore_, baseM_);
119+ for (uint64_t i = 0; i < mCnt; ++i) {
120+ uint32_t curM = (i == mCnt - 1) ? static_cast<uint32_t>(mCurCore_ - (mCnt - 1) * baseM_)
121+ : static_cast<uint32_t>(baseM_);
122+ 
123+ // Phase1: compute rstd
124+ LocalTensor<float> rstdLocal = outQueueRstd.AllocTensor<float>();
125+ for (uint32_t j = 0; j < curM; ++j) {
126+ int64_t gmRowOffset = (i * baseM_ + j) * numCol_;
127+ ComputeOneLineXSquareSum(rstdLocal, gmRowOffset, j);
128+ }
129+ CalculateRstd(rstdLocal, rstdLocal, curM, avgFactor_, epsilon_);
130+ outQueueRstd.EnQue<float>(rstdLocal);
131+ rstdLocal = outQueueRstd.DeQue<float>();
132+ 
133+ // Phase2: compute y、MxQuant、x_out
134+ for (uint64_t j = 0; j < nUbLoops_; ++j) {
135+ uint32_t curN = (j == nUbLoops_ - 1)
136+ ? static_cast<uint32_t>(numCol_ - (nUbLoops_ - 1) * baseN_)
137+ : static_cast<uint32_t>(baseN_);
138+ 
139+ // Load gamma (and beta) per-tile
140+ LocalTensor<T_GAMMA> gammabetaLocal = inQueueGammabeta.AllocTensor<T_GAMMA>();
141+ CopyInGammabeta(gammabetaLocal, j * baseN_, curN);
142+ inQueueGammabeta.EnQue(gammabetaLocal);
143+ inQueueGammabeta.DeQue<T_GAMMA>();
144+ 
145+ for (uint32_t k = 0; k < curM; ++k) {
146+ int64_t gmOffset = (i * baseM_ + k) * numCol_ + j * baseN_;
147+ 
148+ // Re-load x1, x2
149+ CopyInX(gmOffset, curN, j);
150+ LocalTensor<T_X> xLocal1 = inQueueX1.DeQue<T_X>();
151+ LocalTensor<T_X> xLocal2 = inQueueX2.DeQue<T_X>();
152+ LocalTensor<T_X> xOutLocal = outQueueX.AllocTensor<T_X>();
153+ LocalTensor<float> xFp32Local = xFp32Buf.Get<float>();
154+ CalculateXAdd(xLocal1, xLocal2, xOutLocal, xFp32Local, curN);
155+ inQueueX1.FreeTensor(xLocal1);
156+ inQueueX2.FreeTensor(xLocal2);
157+ outQueueX.EnQue<T_X>(xOutLocal);
158+ 
159+ // CopyOut x_out
160+ CopyOutX(gmOffset, curN, j);
161+ 
162+ // Compute y_local
163+ LocalTensor<T_X> yLocal = yTmpBuf.Get<T_X>();
164+ if ((j == nUbLoops_ - 1) && (numCol_ != numColAlign_)) {
165+ Duplicate<T_X>(yLocal, static_cast<T_X>(0), baseN_);
166+ PipeBarrier<PIPE_V>();
167+ }
168+ if (betaFlag_ != 0) {
169+ CalculateY<true>(xFp32Local, yLocal, rstdLocal, curN, k);
170+ } else {
171+ CalculateY<false>(xFp32Local, yLocal, rstdLocal, curN, k);
172+ }
173+ 
174+ // MxQuant
175+ DynamicMxQuantPhase<RoundMode::CAST_RINT>(yLocal, j);
176+ 
177+ // CopyOut y_quant and mxscale
178+ CopyOutQuantY(gmOffset, curN, j);
179+ CopyOutMxScale(i * baseM_ + k, j);
180+ }
181+ 
182+ inQueueGammabeta.FreeTensor(gammabetaLocal);
183+ }
184+ // CopyOut rstd
185+ if (rstdFlag_ != 0) {
186+ DataCopyExtParams copyParams{
187+ static_cast<uint16_t>(1), static_cast<uint32_t>(curM * sizeof(float)), static_cast<uint32_t>(0),
188+ static_cast<uint32_t>(0), 0};
189+ DataCopyPad(rstdGm[i * baseM_], rstdLocal, copyParams);
190+ }
191+ outQueueRstd.FreeTensor(rstdLocal);
192+ }
193+ }
194+ 
195+private:
196+ __aicore__ inline void ComputeOneLineXSquareSum(
197+ LocalTensor<float>& rstdLocal, int64_t gmRowOffset, uint32_t rowIndex)
198+ {
199+ DataCopyPadParams padParams{false, 0, 0, 0};
200+ DataCopyParams xDataCopyParams;
201+ xDataCopyParams.blockCount = 1;
202+ xDataCopyParams.srcStride = 0;
203+ xDataCopyParams.dstStride = 0;
204+ DataCopyParams xFoldDataCopyParams;
205+ xFoldDataCopyParams.blockCount = 1;
206+ xFoldDataCopyParams.srcStride = 0;
207+ xFoldDataCopyParams.dstStride = 0;
208+ 
209+ LocalTensor<float> cacheLocal = cacheBuf.Get<float>();
210+ LocalTensor<float> xFp32Tmp = xFp32Buf.Get<float>();
211+ 
212+ for (int64_t r = 0; r < powerSplit_; ++r) {
213+ int64_t xGmOffset1 = gmRowOffset + baseN_ * r;
214+ int64_t xGmOffset2 = gmRowOffset + baseN_ * (r + powerSplit_);
215+ 
216+ // Step1: Load main tile, compute (x1+x2)² → xFp32Tmp, then Free
217+ xDataCopyParams.blockLen = baseN_ * sizeof(T_X);
218+ LocalTensor<T_X> x1Local = inQueueX1.AllocTensor<T_X>();
219+ DataCopyPad(x1Local, x1Gm[xGmOffset1], xDataCopyParams, padParams);
220+ inQueueX1.EnQue<T_X>(x1Local);
221+ x1Local = inQueueX1.DeQue<T_X>();
222+ 
223+ LocalTensor<T_X> x2Local = inQueueX2.AllocTensor<T_X>();
224+ DataCopyPad(x2Local, x2Gm[xGmOffset1], xDataCopyParams, padParams);
225+ inQueueX2.EnQue<T_X>(x2Local);
226+ x2Local = inQueueX2.DeQue<T_X>();
227+ 
228+ MainBlockSquareVF(x1Local, x2Local, xFp32Tmp, baseN_);
229+ inQueueX1.FreeTensor(x1Local);
230+ inQueueX2.FreeTensor(x2Local);
231+ 
232+ // Step2: Load fold tile, compute (x1Fold+x2Fold)² and accumulate to xFp32Tmp, then Free
233+ if (r < mainFoldCount_) {
234+ xFoldDataCopyParams.blockLen = baseN_ * sizeof(T_X);
235+ LocalTensor<T_X> x1FoldLocal = inQueueX1.AllocTensor<T_X>();
236+ DataCopyPad(x1FoldLocal, x1Gm[xGmOffset2], xFoldDataCopyParams, padParams);
237+ inQueueX1.EnQue<T_X>(x1FoldLocal);
238+ x1FoldLocal = inQueueX1.DeQue<T_X>();
239+ 
240+ LocalTensor<T_X> x2FoldLocal = inQueueX2.AllocTensor<T_X>();
241+ DataCopyPad(x2FoldLocal, x2Gm[xGmOffset2], xFoldDataCopyParams, padParams);
242+ inQueueX2.EnQue<T_X>(x2FoldLocal);
243+ x2FoldLocal = inQueueX2.DeQue<T_X>();
244+ 
245+ FoldBlockSquareAddVF(x1FoldLocal, x2FoldLocal, xFp32Tmp, baseN_);
246+ inQueueX1.FreeTensor(x1FoldLocal);
247+ inQueueX2.FreeTensor(x2FoldLocal);
248+ } else if (r == mainFoldCount_ && foldTail_ > 0) {
249+ xFoldDataCopyParams.blockLen = foldTail_ * sizeof(T_X);
250+ LocalTensor<T_X> x1FoldLocal = inQueueX1.AllocTensor<T_X>();
251+ DataCopyPad(x1FoldLocal, x1Gm[xGmOffset2], xFoldDataCopyParams, padParams);
252+ inQueueX1.EnQue<T_X>(x1FoldLocal);
253+ x1FoldLocal = inQueueX1.DeQue<T_X>();
254+ 
255+ LocalTensor<T_X> x2FoldLocal = inQueueX2.AllocTensor<T_X>();
256+ DataCopyPad(x2FoldLocal, x2Gm[xGmOffset2], xFoldDataCopyParams, padParams);
257+ inQueueX2.EnQue<T_X>(x2FoldLocal);
258+ x2FoldLocal = inQueueX2.DeQue<T_X>();
259+ 
260+ FoldBlockSquareAddVF(x1FoldLocal, x2FoldLocal, xFp32Tmp, foldTail_);
261+ inQueueX1.FreeTensor(x1FoldLocal);
262+ inQueueX2.FreeTensor(x2FoldLocal);
263+ }
264+ // reduce sum
265+ CalculateSquareReduceSum(xFp32Tmp, xFp32Tmp, baseN_);
266+ int64_t cacheId = GetCacheId(r);
267+ UpdateCache(cacheLocal, xFp32Tmp, cacheId, AR_RECOMPUTE_SUM_LEN);
268+ }
269+ 
270+ // final accumulated result to rstdLocal
271+ __local_mem__ float* dstPtr = (__local_mem__ float*)rstdLocal.GetPhyAddr();
272+ __local_mem__ float* cachePtr =
273+ (__local_mem__ float*)cacheLocal.GetPhyAddr() + resultCacheID_ * AR_RECOMPUTE_SUM_LEN;
274+ __VEC_SCOPE__
275+ {
276+ RegTensor<float> a;
277+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
278+ DataCopy<float, LoadDist::DIST_NORM>(a, cachePtr);
279+ DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);
280+ }
281+ }
282+ 
283+ __aicore__ inline void CalculateSquareReduceSum(
284+ LocalTensor<float>& xFp32Local, LocalTensor<float>& xReduceLocal, uint32_t reduceNum)
285+ {
286+ LocalTensor<float> binaryAddBuffTmp = binaryAddBuf.Get<float>();
287+ __local_mem__ float* xReduceUb = (__local_mem__ float*)xReduceLocal.GetPhyAddr();
288+ __local_mem__ float* tmpUb = (__local_mem__ float*)binaryAddBuffTmp.GetPhyAddr();
289+ __local_mem__ float* xFp32Tmp = (__local_mem__ float*)xFp32Local.GetPhyAddr();
290+ 
291+ if (reduceNum <= VL_F32) {
292+ CalculateSquareReduceSumLessThanVL(xFp32Tmp, xReduceUb, reduceNum);
293+ } else if (reduceNum <= VL_F32 + VL_F32) {
294+ CalculateSquareReduceSumLessThanTwoVL(xFp32Tmp, xReduceUb, reduceNum);
295+ } else if (reduceNum <= VL_F32 * VL_F32 * DIGIT_TWO) {
296+ CalculateSquareReduceSumCommon<1>(xFp32Tmp, xReduceUb, tmpUb, reduceNum);
297+ } else {
298+ CalculateSquareReduceSumCommon<2>(xFp32Tmp, xReduceUb, tmpUb, reduceNum);
299+ }
300+ }
301+ 
302+ __aicore__ inline void CalculateSquareReduceSumLessThanVL(
303+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, uint32_t reduceNum)
304+ {
305+ __VEC_SCOPE__
306+ {
307+ RegTensor<float> x, vMean, onesReg;
308+ uint32_t sreg0 = reduceNum;
309+ MaskReg pregLoop = UpdateMask<float>(sreg0);
310+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
311+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
312+ 
313+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregLoop, 0);
314+ AscendC::MicroAPI::ReduceSum(vMean, x, pregLoop);
315+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
316+ }
317+ }
318+ 
319+ __aicore__ inline void CalculateSquareReduceSumLessThanTwoVL(
320+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, uint32_t reduceNum)
321+ {
322+ uint32_t tailLen = reduceNum - VL_F32;
323+ __VEC_SCOPE__
324+ {
325+ RegTensor<float> x, xFold, sumReg, vMean, onesReg;
326+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
327+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
328+ MaskReg pregTail = UpdateMask<float>(tailLen);
329+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
330+ 
331+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregFull, 0);
332+ LoadTensorForDtypeTIn<float>(xFp32Tmp + VL_F32, xFold, pregTail, 0);
333+ AscendC::MicroAPI::ShiftLefts(
334+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregTail);
335+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
336+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
337+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
338+ }
339+ }
340+ 
341+ template <int32_t LAST_LOOP_NUMS>
342+ __aicore__ inline void CalculateSquareReduceSumCommon(
343+ __local_mem__ float* xFp32Tmp, __local_mem__ float* xReduceUb, __local_mem__ float* tmpUb,
344+ uint32_t reduceNum)
345+ {
346+ uint32_t binaryAddQuotient = binAddQuotient_;
347+ uint16_t binaryAddQuotientLoop = (binaryAddQuotient + VL_F32 - 1) / VL_F32;
348+ uint32_t lastBinaryAddNum = binaryAddQuotient / VL_F32;
349+ 
350+ uint32_t binaryAddRemainder = reduceNum - binaryAddQuotient;
351+ uint16_t binaryAddRemainderCeilLoop = (binaryAddRemainder + VL_F32 - 1) / VL_F32;
352+ uint16_t binaryAddRemainderFloorLoop = binaryAddRemainder / VL_F32;
353+ __VEC_SCOPE__
354+ {
355+ RegTensor<float> x, xFold, sumReg, vMean, onesReg;
356+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
357+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
358+ MaskReg pregLoop;
359+ AscendC::MicroAPI::Duplicate(onesReg, float(1.0), pregOne);
360+ 
361+ for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; ++r) {
362+ uint32_t offset = r * VL_F32;
363+ LoadTensorForDtypeTIn<float>(xFp32Tmp, x, pregFull, offset);
364+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddQuotient, xFold, pregFull, offset);
365+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
366+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
367+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
368+ tmpUb + static_cast<uint32_t>(r), vMean, pregOne);
369+ }
370+ uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_F32;
371+ for (uint16_t r = 0;
372+ r < static_cast<uint16_t>(binaryAddRemainderCeilLoop - binaryAddRemainderFloorLoop); ++r) {
373+ pregLoop = UpdateMask<float>(sregRemainder);
374+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddRemainderFloorLoop * VL_F32, x, pregFull, 0);
375+ LoadTensorForDtypeTIn<float>(
376+ xFp32Tmp + binaryAddRemainderFloorLoop * VL_F32 + binaryAddQuotient, xFold, pregLoop, 0);
377+ AscendC::MicroAPI::ShiftLefts(
378+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregLoop);
379+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
380+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
381+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
382+ tmpUb + static_cast<uint32_t>(binaryAddRemainderFloorLoop), vMean, pregOne);
383+ }
384+ for (uint16_t r = 0;
385+ r < static_cast<uint16_t>(binaryAddQuotientLoop - binaryAddRemainderCeilLoop); ++r) {
386+ LoadTensorForDtypeTIn<float>(xFp32Tmp + binaryAddRemainderCeilLoop * VL_F32, x, pregFull, 0);
387+ AscendC::MicroAPI::ReduceSum(vMean, x, pregFull);
388+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
389+ tmpUb + static_cast<uint32_t>(binaryAddRemainderCeilLoop + r), vMean, pregOne);
390+ }
391+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
392+ if constexpr (LAST_LOOP_NUMS == 1) {
393+ MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum);
394+ DataCopy(x, tmpUb);
395+ AscendC::MicroAPI::ReduceSum(vMean, x, pregLast);
396+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
397+ } else if constexpr (LAST_LOOP_NUMS == 2) {
398+ lastBinaryAddNum -= VL_F32;
399+ MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum);
400+ DataCopy(x, tmpUb);
401+ DataCopy(xFold, tmpUb + VL_F32);
402+ AscendC::MicroAPI::ShiftLefts(
403+ (RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregLast);
404+ AscendC::MicroAPI::Add(sumReg, x, xFold, pregFull);
405+ AscendC::MicroAPI::ReduceSum(vMean, sumReg, pregFull);
406+ AscendC::MicroAPI::DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(xReduceUb, vMean, pregOne);
407+ }
408+ }
409+ }
410+ 
411+ template <bool hasBeta>
412+ __aicore__ inline void CalculateY(
413+ LocalTensor<float>& xFp32Local, LocalTensor<T_X>& yLocal, LocalTensor<float>& rstdLocal,
414+ uint32_t curN, uint32_t rowIdx)
415+ {
416+ __local_mem__ float* xFp32Tmp = (__local_mem__ float*)xFp32Local.GetPhyAddr();
417+ __local_mem__ T_GAMMA* gammaInUb = (__local_mem__ T_GAMMA*)gammaLocal_.GetPhyAddr();
418+ __local_mem__ T_X* yInUb = (__local_mem__ T_X*)yLocal.GetPhyAddr();
419+ __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr();
420+ __local_mem__ T_GAMMA* betaInUb;
421+ if constexpr (hasBeta) {
422+ betaInUb = (__local_mem__ T_GAMMA*)betaLocal_.GetPhyAddr();
423+ }
424+ 
425+ uint16_t loopCols = static_cast<uint16_t>((curN + VL_F32 - 1) / VL_F32);
426+ 
427+ __VEC_SCOPE__
428+ {
429+ RegTensor<float> xRegFp32, gammaRegFp32, rstdReg, betaRegFp32;
430+ MaskReg maskReg;
431+ 
432+ AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdInUb + rowIdx);
433+ uint32_t sregCount = curN;
434+ for (uint16_t r = 0; r < loopCols; ++r) {
435+ uint32_t offset = r * VL_F32;
436+ maskReg = UpdateMask<float>(sregCount);
437+ LoadTensorForDtypeTIn<float>(xFp32Tmp, xRegFp32, maskReg, offset);
438+ LoadTensorForDtypeTIn<T_GAMMA>(gammaInUb, gammaRegFp32, maskReg, offset);
439+ AscendC::MicroAPI::Mul(xRegFp32, xRegFp32, rstdReg, maskReg);
440+ AscendC::MicroAPI::Mul(xRegFp32, xRegFp32, gammaRegFp32, maskReg);
441+ if constexpr (hasBeta) {
442+ LoadTensorForDtypeTIn<T_GAMMA>(betaInUb, betaRegFp32, maskReg, offset);
443+ AscendC::MicroAPI::Add(xRegFp32, xRegFp32, betaRegFp32, maskReg);
444+ }
445+ StoreTensorForDtypeTOut<T_X>(yInUb, xRegFp32, maskReg, offset);
446+ }
447+ }
448+ }
449+ 
450+ template <AscendC::RoundMode roundMode>
451+ __aicore__ inline void DynamicMxQuantPhase(LocalTensor<T_X>& yLocal, uint64_t ubLoopIdx)
452+ {
453+ // each call processes a single row's single tile
454+ uint32_t curBlockNumInColAxis;
455+ uint32_t curN;
456+ if (ubLoopIdx == nUbLoops_ - 1) {
457+ curN = static_cast<uint32_t>(numColAlign_ - (nUbLoops_ - 1) * baseN_);
458+ curBlockNumInColAxis = CeilDiv(static_cast<uint64_t>(curN), static_cast<uint64_t>(mxBlockSize_));
459+ } else {
460+ curN = static_cast<uint32_t>(baseN_);
461+ curBlockNumInColAxis = CeilDiv(baseN_, mxBlockSize_);
462+ }
463+ 
464+ uint32_t totalScaleInUB = curBlockNumInColAxis;
465+ uint32_t totalCountInUB = curBlockNumInColAxis * mxBlockSize_;
466+ 
467+ uint16_t loopNum = (totalCountInUB + VL_B16 * DIGIT_TWO - 1) / (VL_B16 * DIGIT_TWO);
468+ uint16_t loopNumScale = (totalScaleInUB + VL_B16 - 1) / VL_B16;
469+ uint16_t loopNumScale4NV = (totalScaleInUB + VL_F32 - 1) / VL_F32;
470+ 
471+ LocalTensor<uint16_t> maxExpLocal = maxExpBuff.Get<uint16_t>();
472+ auto srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
473+ auto maxExpAddr = reinterpret_cast<__ubuf__ uint16_t*>(maxExpLocal.GetPhyAddr());
474+ 
475+ LocalTensor<uint16_t> mxScaleLocal = mxScaleQueue.AllocTensor<uint16_t>();
476+ auto mxScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(mxScaleLocal.GetPhyAddr());
477+ 
478+ LocalTensor<uint16_t> halfScaleLocal = halfScaleBuff.Get<uint16_t>();
479+ auto halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
480+ 
481+ LocalTensor<int8_t> outLocal = outQueueQuantY.AllocTensor<int8_t>();
482+ auto outLocalAddr = reinterpret_cast<__ubuf__ int8_t*>(outLocal.GetPhyAddr());
483+ 
484+ if (scaleAlg_ == 0) {
485+ MxQuantComputeMaxExpOCP<T_X>(srcAddr, maxExpAddr, loopNum);
486+ MxQuantComputeScaleOCP<T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale);
487+ } else {
488+ MxQuantComputeMaxExpcuBLAS<T_X>(srcAddr, maxExpAddr, loopNum);
489+ MxQuantComputeScalecuBLAS<T_X, T_Y>(maxExpAddr, mxScaleLocalAddr, halfScaleLocalAddr, totalScaleInUB, loopNumScale4NV);
490+ }
491+ 
492+ srcAddr = reinterpret_cast<__ubuf__ T_X*>(yLocal.GetPhyAddr());
493+ halfScaleLocalAddr = reinterpret_cast<__ubuf__ uint16_t*>(halfScaleLocal.GetPhyAddr());
494+ 
495+ MxQuantComputeData<roundMode, T_X, T_Y>(srcAddr, halfScaleLocalAddr, outLocalAddr, loopNum);
496+ 
497+ outQueueQuantY.EnQue(outLocal);
498+ mxScaleQueue.EnQue(mxScaleLocal);
499+ }
500+ 
501+ __aicore__ inline void CopyInGammabeta(LocalTensor<T_GAMMA>& gammabetaLocal, int64_t offset, uint32_t len)
502+ {
503+ DataCopyExtParams copyParams{
504+ static_cast<uint16_t>(1),
505+ static_cast<uint32_t>(len * sizeof(T_GAMMA)),
506+ static_cast<uint32_t>(0),
507+ static_cast<uint32_t>(0),
508+ 0
509+ };
510+ DataCopyPadExtParams<T_GAMMA> padParams{false, static_cast<uint8_t>(0), static_cast<uint8_t>(0), static_cast<T_GAMMA>(0.0)};
511+ gammaLocal_ = gammabetaLocal;
512+ DataCopyPad<T_GAMMA>(gammaLocal_, gammaGm[offset], copyParams, padParams);
513+ if (betaFlag_ != 0) {
514+ betaLocal_ = gammabetaLocal[CeilAlign(baseN_ * sizeof(T_GAMMA), UB_BLOCK_SIZE) / sizeof(T_GAMMA)];
515+ DataCopyPad<T_GAMMA>(betaLocal_, betaGm[offset], copyParams, padParams);
516+ }
517+ }
518+ 
519+ __aicore__ inline void CopyInX(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
520+ {
521+ LocalTensor<T_X> xLocal1 = inQueueX1.AllocTensor<T_X>();
522+ LocalTensor<T_X> xLocal2 = inQueueX2.AllocTensor<T_X>();
523+ 
524+ DataCopyExtParams extParams{
525+ static_cast<uint16_t>(1),
526+ static_cast<uint32_t>(curN * sizeof(T_X)),
527+ static_cast<uint32_t>(0),
528+ static_cast<uint32_t>(0),
529+ 0
530+ };
531+ DataCopyPadExtParams<T_X> padParams{
532+ false, static_cast<uint8_t>(0), static_cast<uint8_t>(0), static_cast<T_X>(0.0)};
533+ 
534+ DataCopyPad(xLocal1, x1Gm[gmOffset], extParams, padParams);
535+ DataCopyPad(xLocal2, x2Gm[gmOffset], extParams, padParams);
536+ inQueueX1.EnQue(xLocal1);
537+ inQueueX2.EnQue(xLocal2);
538+ }
539+ 
540+ __aicore__ inline void CopyOutX(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
541+ {
542+ LocalTensor<T_X> xLocal = outQueueX.DeQue<T_X>();
543+ 
544+ DataCopyExtParams copyParams{
545+ static_cast<uint16_t>(1),
546+ static_cast<uint32_t>(curN * sizeof(T_X)),
547+ static_cast<uint32_t>(0),
548+ static_cast<uint32_t>(0),
549+ 0
550+ };
551+ DataCopyPad(xOutGm[gmOffset], xLocal, copyParams);
552+ outQueueX.FreeTensor(xLocal);
553+ }
554+ 
555+ __aicore__ inline void CopyOutQuantY(int64_t gmOffset, uint32_t curN, uint64_t ubLoopIdx)
556+ {
557+ LocalTensor<uint8_t> quantYLocal = outQueueQuantY.DeQue<uint8_t>();
558+ uint32_t srcStride = 0;
559+ if ((ubLoopIdx == nUbLoops_ - 1) && (numCol_ != numColAlign_)) {
560+ srcStride = (numColAlign_ - numCol_) * sizeof(uint8_t) / UB_BLOCK_SIZE;
561+ }
562+ DataCopyExtParams copyParams{
563+ static_cast<uint16_t>(1),
564+ static_cast<uint32_t>(curN),
565+ static_cast<uint32_t>(srcStride),
566+ static_cast<uint32_t>(0),
567+ 0
568+ };
569+ DataCopyPad<uint8_t>(yFp8Gm[gmOffset], quantYLocal, copyParams);
570+ outQueueQuantY.FreeTensor(quantYLocal);
571+ }
572+ 
573+ __aicore__ inline void CopyOutMxScale(uint64_t rowIdx, uint64_t tileIdx)
574+ {
575+ LocalTensor<uint8_t> mxScaleLocal = mxScaleQueue.DeQue<uint8_t>();
576+ uint32_t curScaleSize;
577+ if (tileIdx == nUbLoops_ - 1) {
578+ uint32_t curN = static_cast<uint32_t>(numColAlign_ - (nUbLoops_ - 1) * baseN_);
579+ curScaleSize = CeilDiv(static_cast<uint64_t>(curN), mxBlockSize_);
580+ } else {
581+ curScaleSize = CeilDiv(baseN_, mxBlockSize_);
582+ }
583+ uint64_t scaleGmOffset = rowIdx * mxScaleSize_ + tileIdx * CeilDiv(baseN_, mxBlockSize_);
584+ DataCopyExtParams copyParams{
585+ static_cast<uint16_t>(1),
586+ static_cast<uint32_t>(curScaleSize),
587+ static_cast<uint32_t>(0),
588+ static_cast<uint32_t>(0),
589+ 0
590+ };
591+ DataCopyPad<uint8_t, PaddingMode::Compact>(mxScaleGm[scaleGmOffset], mxScaleLocal, copyParams);
592+ mxScaleQueue.FreeTensor(mxScaleLocal);
593+ }
594+ 
595+private:
596+ TPipe* pPipe = nullptr;
597+ 
598+ // Input Queues
599+ TQue<QuePosition::VECIN, 1> inQueueX1;
600+ TQue<QuePosition::VECIN, 1> inQueueX2;
601+ TQue<QuePosition::VECIN, 1> inQueueGammabeta;
602+ 
603+ LocalTensor<T_GAMMA> gammaLocal_;
604+ LocalTensor<T_GAMMA> betaLocal_;
605+ 
606+ // Output Queues
607+ TQue<QuePosition::VECOUT, 1> outQueueX;
608+ TQue<QuePosition::VECOUT, 1> outQueueRstd;
609+ TQue<QuePosition::VECOUT, 1> outQueueQuantY;
610+ TQue<QuePosition::VECOUT, 1> mxScaleQueue;
611+ 
612+ // TBuf
613+ TBuf<TPosition::VECCALC> xFp32Buf;
614+ TBuf<TPosition::VECCALC> cacheBuf;
615+ TBuf<TPosition::VECCALC> binaryAddBuf;
616+ TBuf<TPosition::VECCALC> yTmpBuf;
617+ TBuf<TPosition::VECCALC> maxExpBuff;
618+ TBuf<TPosition::VECCALC> halfScaleBuff;
619+ 
620+ // GM Tensors
621+ GlobalTensor<T_X> x1Gm, x2Gm, xOutGm;
622+ GlobalTensor<T_GAMMA> gammaGm, betaGm;
623+ GlobalTensor<float> rstdGm;
624+ GlobalTensor<uint8_t> yFp8Gm, mxScaleGm;
625+ 
626+ // Tiling Parameters
627+ uint64_t numCol_, numColAlign_;
628+ uint64_t blockFactor_, mLastCore_, mCurCore_;
629+ uint64_t baseN_, baseM_, baseNBlockSize_;
630+ uint64_t nUbLoops_;
631+ uint64_t binAddQuotient_, powerSplit_;
632+ uint64_t mainFoldCount_, foldTail_;
633+ int64_t resultCacheID_;
634+ float epsilon_, avgFactor_;
635+ uint64_t roundMode_, mxBlockSize_;
636+ int64_t scaleAlg_;
637+ uint64_t mxScaleSize_;
638+ uint32_t betaFlag_, rstdFlag_;
639+};
640+ 
641+} // namespace AddRmsNormDynamicMxQuant
642+#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_FP8_SPLIT_R_H
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_reduce_empty.h+113-113
@@ -1,113 +1,113 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.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 of3+ * 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").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.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,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.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- 10+ 
11-/*!11+/*!
12- * \file add_rms_norm_dynamic_mx_quant_reduce_empty.h12+ * \file add_rms_norm_dynamic_mx_quant_reduce_empty.h
13- * \brief When numCol=0 and numRow!=0, fill rstd output with NaN.13+ * \brief
14- */14+ */
15-#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_15+#ifndef ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_
16-#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_16+#define ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_
17- 17+ 
18-#include "add_rms_norm_dynamic_mx_quant_common.h"18+#include "add_rms_norm_dynamic_mx_quant_common.h"
19- 19+ 
20-namespace AddRmsNormDynamicMxQuant {20+namespace AddRmsNormDynamicMxQuant {
21-using namespace AscendC;21+using namespace AscendC;
22- 22+ 
23-class AddRmsNormDynamicMxQuantReduceEmpty {23+class AddRmsNormDynamicMxQuantReduceEmpty {
24-public:24+public:
25- __aicore__ inline explicit AddRmsNormDynamicMxQuantReduceEmpty(25+ __aicore__ inline explicit AddRmsNormDynamicMxQuantReduceEmpty(
26- const AddRmsNormDynamicMxQuantReduceEmptyTilingData* tilingDataIn)26+ const AddRmsNormDynamicMxQuantReduceEmptyTilingData* tilingDataIn)
27- {27+ {
28- tilingData_ = tilingDataIn;28+ tilingData_ = tilingDataIn;
29- }29+ }
30- 30+ 
31- __aicore__ inline void Init(GM_ADDR rstd)31+ __aicore__ inline void Init(GM_ADDR rstd)
32- {32+ {
33- blockIdx_ = GetBlockIdx();33+ blockIdx_ = GetBlockIdx();
34- usedCoreNum_ = GetBlockNum();34+ usedCoreNum_ = GetBlockNum();
35- ASSERT(usedCoreNum_ != 0 && "block dim can not be zero!");35+ ASSERT(usedCoreNum_ != 0 && "block dim can not be zero!");
36- if (blockIdx_ >= usedCoreNum_) {36+ if (blockIdx_ >= usedCoreNum_) {
37- return;37+ return;
38- }38+ }
39- 39+ 
40- rstdFlag_ = tilingData_->rstdFlag;40+ rstdFlag_ = tilingData_->rstdFlag;
41- numRow_ = tilingData_->numRow;41+ numRow_ = tilingData_->numRow;
42- if (rstdFlag_ == 0 || numRow_ == 0) {42+ if (rstdFlag_ == 0 || numRow_ == 0) {
43- return;43+ return;
44- }44+ }
45- 45+ 
46- perCoreElements_ = tilingData_->perCoreElements;46+ perCoreElements_ = tilingData_->perCoreElements;
47- if (blockIdx_ < usedCoreNum_ - 1) {47+ if (blockIdx_ < usedCoreNum_ - 1) {
48- curCoreElements_ = tilingData_->perCoreElements;48+ curCoreElements_ = tilingData_->perCoreElements;
49- coreLoopsNum_ = tilingData_->perCoreLoops;49+ coreLoopsNum_ = tilingData_->perCoreLoops;
50- perLoopElements_ = tilingData_->perCorePerLoopElements;50+ perLoopElements_ = tilingData_->perCorePerLoopElements;
51- lastLoopElements_ = tilingData_->perCoreLastLoopElements;51+ lastLoopElements_ = tilingData_->perCoreLastLoopElements;
52- } else {52+ } else {
53- curCoreElements_ = tilingData_->lastCoreElements;53+ curCoreElements_ = tilingData_->lastCoreElements;
54- coreLoopsNum_ = tilingData_->lastCoreLoops;54+ coreLoopsNum_ = tilingData_->lastCoreLoops;
55- perLoopElements_ = tilingData_->lastCorePerLoopElements;55+ perLoopElements_ = tilingData_->lastCorePerLoopElements;
56- lastLoopElements_ = tilingData_->lastCoreLastLoopElements;56+ lastLoopElements_ = tilingData_->lastCoreLastLoopElements;
57- }57+ }
58- 58+ 
59- rstdGm_.SetGlobalBuffer(59+ rstdGm_.SetGlobalBuffer(
60- (__gm__ float*)rstd + perCoreElements_ * blockIdx_, curCoreElements_);60+ (__gm__ float*)rstd + perCoreElements_ * blockIdx_, curCoreElements_);
61- 61+ 
62- pipe_.InitBuffer(outNanQueue_, BUFFER_NUM, perLoopElements_ * sizeof(float));62+ pipe_.InitBuffer(outNanQueue_, BUFFER_NUM, perLoopElements_ * sizeof(float));
63- }63+ }
64- 64+ 
65- __aicore__ inline void Process()65+ __aicore__ inline void Process()
66- {66+ {
67- if (blockIdx_ >= usedCoreNum_ || rstdFlag_ == 0 || numRow_ == 0) {67+ if (blockIdx_ >= usedCoreNum_ || rstdFlag_ == 0 || numRow_ == 0) {
68- return;68+ return;
69- }69+ }
70- 70+ 
71- LocalTensor<float> nanLocal = outNanQueue_.AllocTensor<float>();71+ LocalTensor<float> nanLocal = outNanQueue_.AllocTensor<float>();
72- 72+ 
73- float nanVal = AscendC::NumericLimits<float>::QuietNaN();73+ float nanVal = AscendC::NumericLimits<float>::QuietNaN();
74- Duplicate(nanLocal, nanVal, perLoopElements_);74+ Duplicate(nanLocal, nanVal, perLoopElements_);
75- 75+ 
76- outNanQueue_.EnQue(nanLocal);76+ outNanQueue_.EnQue(nanLocal);
77- nanLocal = outNanQueue_.DeQue<float>();77+ nanLocal = outNanQueue_.DeQue<float>();
78- 78+ 
79- for (uint64_t i = 0; i < coreLoopsNum_; i++) {79+ for (uint64_t i = 0; i < coreLoopsNum_; i++) {
80- uint64_t curElements = (i == coreLoopsNum_ - 1) ? lastLoopElements_ : perLoopElements_;80+ uint64_t curElements = (i == coreLoopsNum_ - 1) ? lastLoopElements_ : perLoopElements_;
81- DataCopyExtParams copyParams;81+ DataCopyExtParams copyParams;
82- copyParams.blockCount = 1;82+ copyParams.blockCount = 1;
83- copyParams.blockLen = static_cast<uint32_t>(curElements * sizeof(float));83+ copyParams.blockLen = static_cast<uint32_t>(curElements * sizeof(float));
84- copyParams.srcStride = 0;84+ copyParams.srcStride = 0;
85- copyParams.dstStride = 0;85+ copyParams.dstStride = 0;
86- DataCopyPad(rstdGm_[i * perLoopElements_], nanLocal, copyParams);86+ DataCopyPad(rstdGm_[i * perLoopElements_], nanLocal, copyParams);
87- }87+ }
88- 88+ 
89- outNanQueue_.FreeTensor(nanLocal);89+ outNanQueue_.FreeTensor(nanLocal);
90- }90+ }
91- 91+ 
92-private:92+private:
93- constexpr static int64_t BUFFER_NUM = 1;93+ constexpr static int64_t BUFFER_NUM = 1;
94- 94+ 
95- TPipe pipe_;95+ TPipe pipe_;
96- const AddRmsNormDynamicMxQuantReduceEmptyTilingData* tilingData_;96+ const AddRmsNormDynamicMxQuantReduceEmptyTilingData* tilingData_;
97- 97+ 
98- GlobalTensor<float> rstdGm_;98+ GlobalTensor<float> rstdGm_;
99- TQue<QuePosition::VECOUT, 1> outNanQueue_;99+ TQue<QuePosition::VECOUT, 1> outNanQueue_;
100- 100+ 
101- uint64_t usedCoreNum_{0};101+ uint64_t usedCoreNum_{0};
102- uint64_t blockIdx_{0};102+ uint64_t blockIdx_{0};
103- uint32_t rstdFlag_{0};103+ uint32_t rstdFlag_{0};
104- uint64_t numRow_{0};104+ uint64_t numRow_{0};
105- uint64_t perCoreElements_{0};105+ uint64_t perCoreElements_{0};
106- uint64_t curCoreElements_{0};106+ uint64_t curCoreElements_{0};
107- uint64_t coreLoopsNum_{0};107+ uint64_t coreLoopsNum_{0};
108- uint64_t perLoopElements_{0};108+ uint64_t perLoopElements_{0};
109- uint64_t lastLoopElements_{0};109+ uint64_t lastLoopElements_{0};
110-};110+};
111- 111+ 
112-} // namespace AddRmsNormDynamicMxQuant112+} // namespace AddRmsNormDynamicMxQuant
113-#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_113+#endif // ADD_RMS_NORM_DYNAMIC_MX_QUANT_REDUCE_EMPTY_H_
Mnorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_tiling_data.h+27-4
@@ -22,14 +22,14 @@ struct AddRmsNormDynamicMxQuantTilingData {
22 uint64_t blockFactor; // rows per core22 uint64_t blockFactor; // rows per core
23 uint64_t rowFactor; // rows per UB iteration23 uint64_t rowFactor; // rows per UB iteration
24 uint64_t binAddQuotient; // binary add quotient point24 uint64_t binAddQuotient; // binary add quotient point
25- float epsilon;
26- float avgFactor; // 1.0 / R
27 uint64_t roundMode; // rounding mode (0=round, 1=floor, 4=rint)25 uint64_t roundMode; // rounding mode (0=round, 1=floor, 4=rint)
28 uint64_t mxBlockSize; // MX block size (32)26 uint64_t mxBlockSize; // MX block size (32)
29- int64_t scaleAlg; // scale algorithm (0=standard, 1=cublas)
30 uint64_t blockNumInColAxis; // CeilDiv(R, 32)27 uint64_t blockNumInColAxis; // CeilDiv(R, 32)
31 uint64_t dstStrideUbBlocks; // R axis needs dstStrideUbBlocks to align numColAlign28 uint64_t dstStrideUbBlocks; // R axis needs dstStrideUbBlocks to align numColAlign
32 uint64_t mxScaleSize; // mxscale output size per row29 uint64_t mxScaleSize; // mxscale output size per row
30+ int64_t scaleAlg; // scale algorithm (0=standard, 1=cublas)
31+ float epsilon;
32+ float avgFactor; // 1.0 / R
33 uint32_t betaFlag; // whether beta input exists33 uint32_t betaFlag; // whether beta input exists
34 uint32_t rstdFlag; // whether rstd output is needed34 uint32_t rstdFlag; // whether rstd output is needed
35};35};
@@ -43,8 +43,31 @@ struct AddRmsNormDynamicMxQuantReduceEmptyTilingData {
43 uint64_t lastCoreLoops; // loops for last core43 uint64_t lastCoreLoops; // loops for last core
44 uint64_t lastCorePerLoopElements; // elements per loop (last core, non-last loop)44 uint64_t lastCorePerLoopElements; // elements per loop (last core, non-last loop)
45 uint64_t lastCoreLastLoopElements; // elements for last loop (last core)45 uint64_t lastCoreLastLoopElements; // elements for last loop (last core)
46- uint32_t rstdFlag; // whether rstd output is needed
47 uint64_t numRow;46 uint64_t numRow;
47+ uint32_t rstdFlag; // whether rstd output is needed
48+};
49+ 
50+struct AddRmsNormDynamicMxQuantSplitRTilingData {
51+ uint64_t numCol; // R: norm dimension size
52+ uint64_t numColAlign; // R aligned to block boundary
53+ uint64_t blockFactor; // mPerCore: rows per core
Z
Zzzz20484月20日

TilingData 结构体字段排列建议优化：uint32_t 类型字段（betaFlag, rstdFlag）与 uint64_t 类型字段交叉排列会产生填充字节。建议将同宽度字段集中排列（如将 betaFlag/rstdFlag 移到结构体末尾或 uint32_t 区域），可减少内存占用并提升数据拷贝效率。

likedislike
54+ uint64_t mLastCore; // last core row count
55+ uint64_t baseN; // tile size along R axis (power of 2)
56+ uint64_t baseNBlockSize; // CeilDiv(baseN, 32)
57+ uint64_t baseM; // batch size for rstd output alignment
58+ uint64_t nUbLoops; // number of UB iterations along R axis
59+ uint64_t binAddQuotient; // intra-tile binary fold point
60+ uint64_t powerSplit; // inter-tile binary fold power
61+ uint64_t mainFoldCount; // main fold tile count after powerSplit
62+ uint64_t foldTail; // tail length after folding
63+ uint64_t roundMode; // rounding mode (0=round, 1=floor, 4=rint)
64+ uint64_t mxBlockSize; // MX block size (32)
65+ uint64_t mxScaleSize; // mxscale output size per row
66+ int64_t scaleAlg; // scale algorithm (0=standard, 1=cublas)
67+ float epsilon;
68+ float avgFactor; // 1.0 / R
69+ uint32_t betaFlag; // whether beta input exists
70+ uint32_t rstdFlag; // whether rstd output is needed
48};71};
49 72 
50#endif73#endif
Anorm/add_rms_norm_dynamic_mx_quant/op_kernel/arch35/add_rms_norm_dynamic_mx_quant_tiling_key.h+62-0
@@ -0,0 +1,62 @@
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 add_rms_norm_dynamic_mx_quant_tiling_key.h
13+ * \brief
14+ */
15+ 
16+#ifndef OP_KERNEL_ADD_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H
17+#define OP_KERNEL_ADD_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H
18+ 
19+#include "ascendc/host_api/tiling/template_argument.h"
20+ 
21+#define COMPUTE_MODE_FULL_LOAD 0
22+#define COMPUTE_MODE_SPLIT_R 1
23+#define COMPUTE_MODE_REDUCE_EMPTY 2
24+ 
25+#define Y_DATA_TYPE_FP8 0
26+#define Y_DATA_TYPE_FP4 1
27+ 
28+ASCENDC_TPL_ARGS_DECL(
29+ AddRmsNormDynamicMxQuant,
30+ ASCENDC_TPL_UINT_DECL(
31+ COMPUTE_MODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_FULL_LOAD, COMPUTE_MODE_SPLIT_R,
32+ COMPUTE_MODE_REDUCE_EMPTY),
33+ ASCENDC_TPL_UINT_DECL(
34+ Y_DATA_TYPE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP8, Y_DATA_TYPE_FP4));
35+ 
36+ASCENDC_TPL_SEL(
37+ ASCENDC_TPL_ARGS_SEL(
38+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
39+ ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_FULL_LOAD),
40+ ASCENDC_TPL_UINT_SEL(Y_DATA_TYPE, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP8),
41+ ASCENDC_TPL_TILING_STRUCT_SEL(AddRmsNormDynamicMxQuantTilingData)),
42+ ASCENDC_TPL_ARGS_SEL(
43+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
44+ ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_FULL_LOAD),
45+ ASCENDC_TPL_UINT_SEL(Y_DATA_TYPE, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP4),
46+ ASCENDC_TPL_TILING_STRUCT_SEL(AddRmsNormDynamicMxQuantTilingData)),
47+ ASCENDC_TPL_ARGS_SEL(
48+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
49+ ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_SPLIT_R),
50+ ASCENDC_TPL_UINT_SEL(Y_DATA_TYPE, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP8),
51+ ASCENDC_TPL_TILING_STRUCT_SEL(AddRmsNormDynamicMxQuantSplitRTilingData)),
52+ ASCENDC_TPL_ARGS_SEL(
53+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
54+ ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_SPLIT_R),
55+ ASCENDC_TPL_UINT_SEL(Y_DATA_TYPE, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP4),
56+ ASCENDC_TPL_TILING_STRUCT_SEL(AddRmsNormDynamicMxQuantSplitRTilingData)),
57+ ASCENDC_TPL_ARGS_SEL(
58+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
59+ ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_REDUCE_EMPTY),
60+ ASCENDC_TPL_UINT_SEL(Y_DATA_TYPE, ASCENDC_TPL_UI_LIST, Y_DATA_TYPE_FP8, Y_DATA_TYPE_FP4),
61+ ASCENDC_TPL_TILING_STRUCT_SEL(AddRmsNormDynamicMxQuantReduceEmptyTilingData)));
62+#endif
Mnorm/add_rms_norm_dynamic_mx_quant/tests/ut/op_host/test_add_rms_norm_dynamic_mx_quant_tiling.cpp+413-2
@@ -51,7 +51,7 @@ static string to_string(void* buf, size_t size)
51 return result;51 return result;
52}52}
53 53 
54-TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_100)54+TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_fp8_full_load)
55{55{
56 // dlog_setlevel(0, 0, 0);56 // dlog_setlevel(0, 0, 0);
57 gert::StorageShape x1_shape = {{8, 64}, {8, 64}};57 gert::StorageShape x1_shape = {{8, 64}, {8, 64}};
@@ -149,7 +149,418 @@ TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_100)
149 // workspaces nullptr return failed149 // workspaces nullptr return failed
150 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);150 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
151 auto tiling_key = tiling_context->GetTilingKey();151 auto tiling_key = tiling_context->GetTilingKey();
152- ASSERT_EQ(tiling_key, 100);152+ ASSERT_EQ(tiling_key, 0);
153+ auto tilingData = tiling_context->GetRawTilingData();
154+ ASSERT_NE(tilingData, nullptr);
155+}
156+ 
157+TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_fp8_split_r)
158+{
159+ // dlog_setlevel(0, 0, 0);
160+ gert::StorageShape x1_shape = {{10, 200000}, {10, 200000}};
161+ gert::StorageShape x2_shape = {{10, 200000}, {10, 200000}};
162+ gert::StorageShape gamma_shape = {{200000}, {200000}};
163+ gert::StorageShape beta_shape = {{200000}, {200000}};
164+ 
165+ gert::StorageShape out_y_shape = {{10, 200000}, {10, 200000}};
166+ gert::StorageShape out_x_shape = {{10, 200000}, {10, 200000}};
167+ gert::StorageShape out_mxScale_shape = {{10, 3125, 2}, {10, 3125, 2}};
168+ gert::StorageShape out_rstd_shape = {{10, 1}, {10, 1}};
169+ 
170+ std::map<std::string, std::string> soc_version_infos = {{"NpuArch", "3510"}};
171+ string compile_info_string = R"({
172+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
173+ "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
174+ "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
175+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
176+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
177+ "CORE_NUM": 64, "socVersion": "Ascend950"}
178+ })";
179+ map<string, string> soc_infos;
180+ map<string, string> aicore_spec;
181+ map<string, string> intrinsics;
182+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
183+ 
184+ // platform info
185+ fe::PlatFormInfos platform_info;
186+ platform_info.Init();
187+ // compile info
188+ optiling::AddRmsNormDynamicMxQuantCompileInfo compile_info;
189+ 
190+ std::string op_type("AddRmsNormDynamicMxQuant");
191+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
192+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
193+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
194+ 
195+ // tilingParseFunc simulate
196+ auto kernel_holder =
197+ gert::KernelRunContextFaker()
198+ .KernelIONum(4, 4)
199+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
200+ .Outputs({&compile_info})
201+ .Build();
202+ 
203+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
204+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
205+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
206+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
207+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
208+ "AICoreintrinsicDtypeMap", intrinsics);
209+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
210+ "version", soc_version_infos);
211+ 
212+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
213+ 
214+ // tilingFunc simulate
215+ auto param = gert::TilingData::CreateCap(4096);
216+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
217+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
218+ ASSERT_NE(param, nullptr);
219+ auto holder = gert::TilingContextFaker()
220+ .SetOpType(op_type)
221+ .NodeIoNum(4, 4)
222+ .IrInstanceNum({1, 1, 1, 1})
223+ .InputShapes({&x1_shape, &x2_shape, &gamma_shape, &beta_shape})
224+ .OutputShapes({&out_y_shape, &out_x_shape, &out_mxScale_shape, &out_rstd_shape})
225+ .CompileInfo(&compile_info)
226+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
227+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
228+ .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
229+ .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
230+ .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
231+ .NodeOutputTd(0, ge::DT_FLOAT8_E4M3FN, ge::FORMAT_ND, ge::FORMAT_ND)
232+ .NodeOutputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
233+ .NodeOutputTd(2, ge::DT_FLOAT8_E8M0, ge::FORMAT_ND, ge::FORMAT_ND)
234+ .NodeOutputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
235+ .NodeAttrs(
236+ {{"epsilon", Ops::NN::AnyValue::CreateFrom<float>(1e-06)},
237+ {"scale_alg", Ops::NN::AnyValue::CreateFrom<int64_t>(0)},
238+ {"round_mode", Ops::NN::AnyValue::CreateFrom<std::string>("rint")},
239+ {"dst_type", Ops::NN::AnyValue::CreateFrom<int64_t>(36)},
240+ {"output_rstd", Ops::NN::AnyValue::CreateFrom<bool>(true)}})
241+ .TilingData(param.get())
242+ .Workspace(ws_size)
243+ .Build();
244+ 
245+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
246+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
247+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
248+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
249+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
250+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
251+ 
252+ // workspaces nullptr return failed
253+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
254+ auto tiling_key = tiling_context->GetTilingKey();
255+ ASSERT_EQ(tiling_key, 1);
256+ auto tilingData = tiling_context->GetRawTilingData();
257+ ASSERT_NE(tilingData, nullptr);
258+}
259+ 
260+TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_fp4_full_load)
261+{
262+ // dlog_setlevel(0, 0, 0);
263+ gert::StorageShape x1_shape = {{8, 64}, {8, 64}};
264+ gert::StorageShape x2_shape = {{8, 64}, {8, 64}};
265+ gert::StorageShape gamma_shape = {{64}, {64}};
266+ gert::StorageShape beta_shape = {{64}, {64}};
267+ 
268+ gert::StorageShape out_y_shape = {{8, 64}, {8, 64}};
269+ gert::StorageShape out_x_shape = {{8, 64}, {8, 64}};
270+ gert::StorageShape out_mxScale_shape = {{8, 1, 2}, {8, 1, 2}};
271+ gert::StorageShape out_rstd_shape = {{8, 1}, {8, 1}};
272+ 
273+ std::map<std::string, std::string> soc_version_infos = {{"NpuArch", "3510"}};
274+ string compile_info_string = R"({
275+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
276+ "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
277+ "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
278+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
279+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
280+ "CORE_NUM": 64, "socVersion": "Ascend950"}
281+ })";
282+ map<string, string> soc_infos;
283+ map<string, string> aicore_spec;
284+ map<string, string> intrinsics;
285+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
286+ 
287+ // platform info
288+ fe::PlatFormInfos platform_info;
289+ platform_info.Init();
290+ // compile info
291+ optiling::AddRmsNormDynamicMxQuantCompileInfo compile_info;
292+ 
293+ std::string op_type("AddRmsNormDynamicMxQuant");
294+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
295+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
296+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
297+ 
298+ // tilingParseFunc simulate
299+ auto kernel_holder =
300+ gert::KernelRunContextFaker()
301+ .KernelIONum(4, 4)
302+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
303+ .Outputs({&compile_info})
304+ .Build();
305+ 
306+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
307+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
308+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
309+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
310+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
311+ "AICoreintrinsicDtypeMap", intrinsics);
312+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
313+ "version", soc_version_infos);
314+ 
315+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
316+ 
317+ // tilingFunc simulate
318+ auto param = gert::TilingData::CreateCap(4096);
319+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
320+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
321+ ASSERT_NE(param, nullptr);
322+ auto holder = gert::TilingContextFaker()
323+ .SetOpType(op_type)
324+ .NodeIoNum(4, 4)
325+ .IrInstanceNum({1, 1, 1, 1})
326+ .InputShapes({&x1_shape, &x2_shape, &gamma_shape, &beta_shape})
327+ .OutputShapes({&out_y_shape, &out_x_shape, &out_mxScale_shape, &out_rstd_shape})
328+ .CompileInfo(&compile_info)
329+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
330+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
331+ .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
332+ .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
333+ .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
334+ .NodeOutputTd(0, ge::DT_FLOAT4_E1M2, ge::FORMAT_ND, ge::FORMAT_ND)
335+ .NodeOutputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
336+ .NodeOutputTd(2, ge::DT_FLOAT8_E8M0, ge::FORMAT_ND, ge::FORMAT_ND)
337+ .NodeAttrs(
338+ {{"epsilon", Ops::NN::AnyValue::CreateFrom<float>(1e-06)},
339+ {"scale_alg", Ops::NN::AnyValue::CreateFrom<int64_t>(0)},
340+ {"round_mode", Ops::NN::AnyValue::CreateFrom<std::string>("floor")},
341+ {"dst_type", Ops::NN::AnyValue::CreateFrom<int64_t>(41)},
342+ {"output_rstd", Ops::NN::AnyValue::CreateFrom<bool>(false)}})
343+ .TilingData(param.get())
344+ .Workspace(ws_size)
345+ .Build();
346+ 
347+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
348+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
349+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
350+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
351+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
352+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
353+ 
354+ // workspaces nullptr return failed
355+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
356+ auto tiling_key = tiling_context->GetTilingKey();
357+ ASSERT_EQ(tiling_key, 16);
358+ auto tilingData = tiling_context->GetRawTilingData();
359+ ASSERT_NE(tilingData, nullptr);
360+}
361+ 
362+TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_fp4_split_r)
363+{
364+ // dlog_setlevel(0, 0, 0);
365+ gert::StorageShape x1_shape = {{10, 200000}, {10, 200000}};
366+ gert::StorageShape x2_shape = {{10, 200000}, {10, 200000}};
367+ gert::StorageShape gamma_shape = {{200000}, {200000}};
368+ gert::StorageShape beta_shape = {{200000}, {200000}};
369+ 
370+ gert::StorageShape out_y_shape = {{10, 200000}, {10, 200000}};
371+ gert::StorageShape out_x_shape = {{10, 200000}, {10, 200000}};
372+ gert::StorageShape out_mxScale_shape = {{10, 3125, 2}, {10, 3125, 2}};
373+ gert::StorageShape out_rstd_shape = {{10, 1}, {10, 1}};
374+ 
375+ std::map<std::string, std::string> soc_version_infos = {{"NpuArch", "3510"}};
376+ string compile_info_string = R"({
377+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
378+ "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
379+ "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
380+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
381+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
382+ "CORE_NUM": 64, "socVersion": "Ascend950"}
383+ })";
384+ map<string, string> soc_infos;
385+ map<string, string> aicore_spec;
386+ map<string, string> intrinsics;
387+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
388+ 
389+ // platform info
390+ fe::PlatFormInfos platform_info;
391+ platform_info.Init();
392+ // compile info
393+ optiling::AddRmsNormDynamicMxQuantCompileInfo compile_info;
394+ 
395+ std::string op_type("AddRmsNormDynamicMxQuant");
396+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
397+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
398+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
399+ 
400+ // tilingParseFunc simulate
401+ auto kernel_holder =
402+ gert::KernelRunContextFaker()
403+ .KernelIONum(4, 4)
404+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
405+ .Outputs({&compile_info})
406+ .Build();
407+ 
408+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
409+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
410+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
411+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
412+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
413+ "AICoreintrinsicDtypeMap", intrinsics);
414+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
415+ "version", soc_version_infos);
416+ 
417+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
418+ 
419+ // tilingFunc simulate
420+ auto param = gert::TilingData::CreateCap(4096);
421+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
422+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
423+ ASSERT_NE(param, nullptr);
424+ auto holder = gert::TilingContextFaker()
425+ .SetOpType(op_type)
426+ .NodeIoNum(4, 4)
427+ .IrInstanceNum({1, 1, 1, 1})
428+ .InputShapes({&x1_shape, &x2_shape, &gamma_shape, &beta_shape})
429+ .OutputShapes({&out_y_shape, &out_x_shape, &out_mxScale_shape, &out_rstd_shape})
430+ .CompileInfo(&compile_info)
431+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
432+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
433+ .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
434+ .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
435+ .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
436+ .NodeOutputTd(0, ge::DT_FLOAT4_E2M1, ge::FORMAT_ND, ge::FORMAT_ND)
437+ .NodeOutputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
438+ .NodeOutputTd(2, ge::DT_FLOAT8_E8M0, ge::FORMAT_ND, ge::FORMAT_ND)
439+ .NodeOutputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
440+ .NodeAttrs(
441+ {{"epsilon", Ops::NN::AnyValue::CreateFrom<float>(1e-06)},
442+ {"scale_alg", Ops::NN::AnyValue::CreateFrom<int64_t>(0)},
443+ {"round_mode", Ops::NN::AnyValue::CreateFrom<std::string>("rint")},
444+ {"dst_type", Ops::NN::AnyValue::CreateFrom<int64_t>(40)},
445+ {"output_rstd", Ops::NN::AnyValue::CreateFrom<bool>(true)}})
446+ .TilingData(param.get())
447+ .Workspace(ws_size)
448+ .Build();
449+ 
450+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
451+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
452+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
453+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
454+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
455+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
456+ 
457+ // workspaces nullptr return failed
458+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
459+ auto tiling_key = tiling_context->GetTilingKey();
460+ ASSERT_EQ(tiling_key, 17);
461+ auto tilingData = tiling_context->GetRawTilingData();
462+ ASSERT_NE(tilingData, nullptr);
463+}
464+ 
465+TEST_F(AddRmsNormDynamicMxQuantTilingTest, add_rms_norm_dynamic_mx_quant_reduce_empty)
466+{
467+ // dlog_setlevel(0, 0, 0);
468+ gert::StorageShape x1_shape = {{8, 0}, {8, 0}};
469+ gert::StorageShape x2_shape = {{8, 0}, {8, 0}};
470+ gert::StorageShape gamma_shape = {{0}, {0}};
471+ gert::StorageShape beta_shape = {{0}, {0}};
472+ 
473+ gert::StorageShape out_y_shape = {{8, 0}, {8, 0}};
474+ gert::StorageShape out_x_shape = {{8, 0}, {8, 0}};
475+ gert::StorageShape out_mxScale_shape = {{8, 0, 2}, {8, 0, 2}};
476+ gert::StorageShape out_rstd_shape = {{8, 1}, {8, 1}};
477+ 
478+ std::map<std::string, std::string> soc_version_infos = {{"NpuArch", "3510"}};
479+ string compile_info_string = R"({
480+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
481+ "Intrinsic_fix_pipe_l0c2out": false, "Intrinsic_data_move_l12ub": true,
482+ "Intrinsic_data_move_l0c2ub": true, "Intrinsic_data_move_out2l1_nd2nz": false,
483+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
484+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
485+ "CORE_NUM": 64, "socVersion": "Ascend950"}
486+ })";
487+ map<string, string> soc_infos;
488+ map<string, string> aicore_spec;
489+ map<string, string> intrinsics;
490+ GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
491+ 
492+ // platform info
493+ fe::PlatFormInfos platform_info;
494+ platform_info.Init();
495+ // compile info
496+ optiling::AddRmsNormDynamicMxQuantCompileInfo compile_info;
497+ 
498+ std::string op_type("AddRmsNormDynamicMxQuant");
499+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr);
500+ auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling;
501+ auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse;
502+ 
503+ // tilingParseFunc simulate
504+ auto kernel_holder =
505+ gert::KernelRunContextFaker()
506+ .KernelIONum(4, 4)
507+ .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)})
508+ .Outputs({&compile_info})
509+ .Build();
510+ 
511+ ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init());
512+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
513+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
514+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
515+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
516+ "AICoreintrinsicDtypeMap", intrinsics);
517+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
518+ "version", soc_version_infos);
519+ 
520+ ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
521+ 
522+ // tilingFunc simulate
523+ auto param = gert::TilingData::CreateCap(4096);
524+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
525+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
526+ ASSERT_NE(param, nullptr);
527+ auto holder = gert::TilingContextFaker()
528+ .SetOpType(op_type)
529+ .NodeIoNum(4, 4)
530+ .IrInstanceNum({1, 1, 1, 1})
531+ .InputShapes({&x1_shape, &x2_shape, &gamma_shape, &beta_shape})
532+ .OutputShapes({&out_y_shape, &out_x_shape, &out_mxScale_shape, &out_rstd_shape})
533+ .CompileInfo(&compile_info)
534+ .PlatformInfo(reinterpret_cast<char*>(&platform_info))
535+ .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
536+ .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
537+ .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
538+ .NodeInputTd(3, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
539+ .NodeOutputTd(0, ge::DT_FLOAT8_E4M3FN, ge::FORMAT_ND, ge::FORMAT_ND)
540+ .NodeOutputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND)
541+ .NodeOutputTd(2, ge::DT_FLOAT8_E8M0, ge::FORMAT_ND, ge::FORMAT_ND)
542+ .NodeOutputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
543+ .NodeAttrs(
544+ {{"epsilon", Ops::NN::AnyValue::CreateFrom<float>(1e-06)},
545+ {"scale_alg", Ops::NN::AnyValue::CreateFrom<int64_t>(0)},
546+ {"round_mode", Ops::NN::AnyValue::CreateFrom<std::string>("rint")},
547+ {"dst_type", Ops::NN::AnyValue::CreateFrom<int64_t>(36)},
548+ {"output_rstd", Ops::NN::AnyValue::CreateFrom<bool>(true)}})
549+ .TilingData(param.get())
550+ .Workspace(ws_size)
551+ .Build();
552+ 
553+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
554+ ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr);
555+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos);
556+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec);
557+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
558+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
559+ 
560+ // workspaces nullptr return failed
561+ EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
562+ auto tiling_key = tiling_context->GetTilingKey();
563+ ASSERT_EQ(tiling_key, 2);
153 auto tilingData = tiling_context->GetRawTilingData();564 auto tilingData = tiling_context->GetRawTilingData();
154 ASSERT_NE(tilingData, nullptr);565 ASSERT_NE(tilingData, nullptr);
155}566}
Mnorm/add_rms_norm_dynamic_mx_quant/tests/ut/op_kernel/test_add_rms_norm_dynamic_mx_quant_apt.cpp+47-15
@@ -21,15 +21,20 @@
21#include "gtest/gtest.h"21#include "gtest/gtest.h"
22#include "tikicpulib.h"22#include "tikicpulib.h"
23#include "data_utils.h"23#include "data_utils.h"
24+#include "add_rms_norm_dynamic_mx_quant_tiling_def.h"
24 25 
25-#include <cstdint>26+#ifdef __CCE_KT_TEST__
27+#include "add_rms_norm_dynamic_mx_quant_apt.cpp"
28+#endif
26 29 
27using namespace std;30using namespace std;
28 31 
29-extern "C" __global__ __aicore__ void add_rms_norm_dynamic_mx_quant(32+#define COMPUTE_MODE_FULL_LOAD 0
30- GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta,33+#define COMPUTE_MODE_SPLIT_R 1
31- GM_ADDR y, GM_ADDR x, GM_ADDR mxscale, GM_ADDR rstd,34+#define COMPUTE_MODE_REDUCE_EMPTY 2
32- GM_ADDR workspace, GM_ADDR tiling);35+ 
36+#define Y_DATA_TYPE_FP8 0
37+#define Y_DATA_TYPE_FP4 1
33 38 
34class add_rms_norm_dynamic_mx_quant_test : public testing::Test {39class add_rms_norm_dynamic_mx_quant_test : public testing::Test {
35protected:40protected:
@@ -43,7 +48,7 @@ protected:
43 }48 }
44};49};
45 50 
46-TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_r_full_load_fp8_tilingkey100_case1)51+TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_r_full_load_fp8_case1)
47{52{
48 int64_t numA = 1;53 int64_t numA = 1;
49 int64_t numR = 64;54 int64_t numR = 64;
@@ -88,9 +93,16 @@ TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_r_full_load_fp8_tilingkey10
88 tilingDatafromBin->rstdFlag = 1;93 tilingDatafromBin->rstdFlag = 1;
89 94 
90 AscendC::SetKernelMode(KernelMode::AIV_MODE);95 AscendC::SetKernelMode(KernelMode::AIV_MODE);
91- ICPU_SET_TILING_KEY(100);96+ ICPU_SET_TILING_KEY(2);
97+ 
98+ auto add_rms_norm_dynamic_mx_quant_wrapper = [](GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR x,
99+ GM_ADDR mxscale, GM_ADDR rstd, GM_ADDR workspace, GM_ADDR tiling) {
100+ ::add_rms_norm_dynamic_mx_quant<COMPUTE_MODE_FULL_LOAD, Y_DATA_TYPE_FP8>(
101+ x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, tiling);
102+ };
92 ICPU_RUN_KF(103 ICPU_RUN_KF(
93- add_rms_norm_dynamic_mx_quant, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, (uint8_t*)(tilingDatafromBin));104+ add_rms_norm_dynamic_mx_quant_wrapper, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace,
105+ (uint8_t*)(tilingDatafromBin));
94 106 
95 AscendC::GmFree(x1);107 AscendC::GmFree(x1);
96 AscendC::GmFree(x2);108 AscendC::GmFree(x2);
@@ -104,7 +116,6 @@ TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_r_full_load_fp8_tilingkey10
104 AscendC::GmFree(tiling);116 AscendC::GmFree(tiling);
105}117}
106 118 
107-// empty_tensor, tilingKey(3000), numA=176, numR=0, putput_rstd=true
108TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_m_not0_n_0_outputRstd_true)119TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_m_not0_n_0_outputRstd_true)
109{120{
110 int64_t numA = 176;121 int64_t numA = 176;
@@ -145,9 +156,16 @@ TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_m_not0_n_0_outputRstd_true)
145 tilingDatafromBin->numRow = 176;156 tilingDatafromBin->numRow = 176;
146 157 
147 AscendC::SetKernelMode(KernelMode::AIV_MODE);158 AscendC::SetKernelMode(KernelMode::AIV_MODE);
148- ICPU_SET_TILING_KEY(300);159+ ICPU_SET_TILING_KEY(2);
160+ 
161+ auto add_rms_norm_dynamic_mx_quant_wrapper = [](GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR x,
162+ GM_ADDR mxscale, GM_ADDR rstd, GM_ADDR workspace, GM_ADDR tiling) {
163+ ::add_rms_norm_dynamic_mx_quant<COMPUTE_MODE_REDUCE_EMPTY, Y_DATA_TYPE_FP8>(
164+ x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, tiling);
165+ };
149 ICPU_RUN_KF(166 ICPU_RUN_KF(
150- add_rms_norm_dynamic_mx_quant, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, (uint8_t*)(tilingDatafromBin));167+ add_rms_norm_dynamic_mx_quant_wrapper, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace,
168+ (uint8_t*)(tilingDatafromBin));
151 169 
152 AscendC::GmFree(x1);170 AscendC::GmFree(x1);
153 AscendC::GmFree(x2);171 AscendC::GmFree(x2);
@@ -202,9 +220,16 @@ TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_m_0_n_not_0_outputRstd_true
202 tilingDatafromBin->numRow = 0;220 tilingDatafromBin->numRow = 0;
203 221 
204 AscendC::SetKernelMode(KernelMode::AIV_MODE);222 AscendC::SetKernelMode(KernelMode::AIV_MODE);
205- ICPU_SET_TILING_KEY(300);223+ ICPU_SET_TILING_KEY(2);
224+ 
225+ auto add_rms_norm_dynamic_mx_quant_wrapper = [](GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR x,
226+ GM_ADDR mxscale, GM_ADDR rstd, GM_ADDR workspace, GM_ADDR tiling) {
227+ ::add_rms_norm_dynamic_mx_quant<COMPUTE_MODE_REDUCE_EMPTY, Y_DATA_TYPE_FP8>(
228+ x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, tiling);
229+ };
206 ICPU_RUN_KF(230 ICPU_RUN_KF(
207- add_rms_norm_dynamic_mx_quant, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, (uint8_t*)(tilingDatafromBin));231+ add_rms_norm_dynamic_mx_quant_wrapper, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace,
232+ (uint8_t*)(tilingDatafromBin));
208 233 
209 AscendC::GmFree(x1);234 AscendC::GmFree(x1);
210 AscendC::GmFree(x2);235 AscendC::GmFree(x2);
@@ -259,9 +284,16 @@ TEST_F(add_rms_norm_dynamic_mx_quant_test, test_case_m_0_n_0_outputRstd_true)
259 tilingDatafromBin->numRow = 0;284 tilingDatafromBin->numRow = 0;
260 285 
261 AscendC::SetKernelMode(KernelMode::AIV_MODE);286 AscendC::SetKernelMode(KernelMode::AIV_MODE);
262- ICPU_SET_TILING_KEY(300);287+ ICPU_SET_TILING_KEY(2);
288+ 
289+ auto add_rms_norm_dynamic_mx_quant_wrapper = [](GM_ADDR x1, GM_ADDR x2, GM_ADDR gamma, GM_ADDR beta, GM_ADDR y, GM_ADDR x,
290+ GM_ADDR mxscale, GM_ADDR rstd, GM_ADDR workspace, GM_ADDR tiling) {
291+ ::add_rms_norm_dynamic_mx_quant<COMPUTE_MODE_REDUCE_EMPTY, Y_DATA_TYPE_FP8>(
292+ x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, tiling);
293+ };
263 ICPU_RUN_KF(294 ICPU_RUN_KF(
264- add_rms_norm_dynamic_mx_quant, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace, (uint8_t*)(tilingDatafromBin));295+ add_rms_norm_dynamic_mx_quant_wrapper, blockDim, x1, x2, gamma, beta, y, x, mxscale, rstd, workspace,
296+ (uint8_t*)(tilingDatafromBin));
265 297 
266 AscendC::GmFree(x1);298 AscendC::GmFree(x1);
267 AscendC::GmFree(x2);299 AscendC::GmFree(x2);
Mnorm/rms_norm_dynamic_mx_quant/op_kernel/arch35/rms_norm_dynamic_mx_quant_tiling_key.h+2-2
@@ -15,7 +15,6 @@
15 15 
Z
Zzzz20484月20日

修复正确。此前 #endif 在第15行导致 #ifndef 保护失效，后续代码未被保护，可能引发重复包含问题。本次将 #endif 移至文件末尾是正确的修复。建议确认此修复有对应的测试覆盖。

likedislike
16#ifndef OP_KERNEL_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H16#ifndef OP_KERNEL_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H
17#define OP_KERNEL_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H17#define OP_KERNEL_RMS_NORM_DYNAMIC_MX_QUANT_TILING_KEY_H
18-#endif
19 18 
20#include "ascendc/host_api/tiling/template_argument.h"19#include "ascendc/host_api/tiling/template_argument.h"
21 20 
@@ -59,4 +58,5 @@ ASCENDC_TPL_SEL(
59 ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),58 ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
60 ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_REDUCE_EMPTY),59 ASCENDC_TPL_UINT_SEL(COMPUTE_MODE, ASCENDC_TPL_UI_LIST, COMPUTE_MODE_REDUCE_EMPTY),
61 ASCENDC_TPL_UINT_SEL(OPTIMIZE_MODE, ASCENDC_TPL_UI_LIST, OPTIMIZE_MODE_NORMAL, OPTIMIZE_MODE_OPTIMIZE),60 ASCENDC_TPL_UINT_SEL(OPTIMIZE_MODE, ASCENDC_TPL_UI_LIST, OPTIMIZE_MODE_NORMAL, OPTIMIZE_MODE_OPTIMIZE),
62- ASCENDC_TPL_TILING_STRUCT_SEL(RmsNormDynamicMxQuantReduceEmptyTilingData)));61+ ASCENDC_TPL_TILING_STRUCT_SEL(RmsNormDynamicMxQuantReduceEmptyTilingData)));
62+#endif