Star1143
2757
代码
代码
Issues242
Pull Requests406
Actions
项目讨论
Wiki
分析
Star1143
2757
已合并
GMM support groupListType=2 #4844
lijiamin2025创建于 4月29日
GMM support groupListType=2 #4844
已合并
lijiamin2025创建于 4月29日
从lijiamin2025/ops-transformergmm2合入到cann/ops-transformermaster
共 15 个文件变更+717-273
Mgmm/grouped_matmul/op_api/aclnn_grouped_matmul.cpp+49-42
@@ -1120,12 +1120,18 @@ static aclnnStatus CheckA8W4SymmQuantParams(const gmm::GroupedMatmulParams &gmmP
1120}1120}
1121 1121 
1122static aclnnStatus CheckA8W4QuantParams(const gmm::GroupedMatmulParams &gmmParams) {1122static aclnnStatus CheckA8W4QuantParams(const gmm::GroupedMatmulParams &gmmParams) {
1123- if (!isA8W8AsymmetricQuant(gmmParams)) {1123+ if (gmmParams.groupListType != 1) {
1124- CheckA8W4SymmQuantParams(gmmParams);1124+ OP_LOGW("GMM A8W4: groupListType only support 1(count), but now is %ld.",
1125- return ACLNN_SUCCESS;1125+ gmmParams.groupListType);
1126- }1126+ }
1127- CHECK_COND(CheckA8W4AsymQuantParams(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "CheckA8W4AsymQuantParams failed.");1127+ if (!isA8W8AsymmetricQuant(gmmParams)) {
1128- return ACLNN_SUCCESS;1128+ CheckA8W4SymmQuantParams(gmmParams);
1129+ return ACLNN_SUCCESS;
1130+ }
1131+ CHECK_COND(CheckA8W4AsymQuantParams(gmmParams) == ACLNN_SUCCESS,
1132+ ACLNN_ERR_PARAM_INVALID,
1133+ "CheckA8W4AsymQuantParams failed.");
1134+ return ACLNN_SUCCESS;
1129}1135}
1130 1136 
1131static aclnnStatus CheckA4W4ParamsShape(const gmm::GroupedMatmulParams &gmmParams) {1137static aclnnStatus CheckA4W4ParamsShape(const gmm::GroupedMatmulParams &gmmParams) {
@@ -1154,41 +1160,44 @@ static aclnnStatus CheckA4W4ParamsShape(const gmm::GroupedMatmulParams &gmmParam
1154}1160}
1155 1161 
1156static aclnnStatus CheckA4W4QuantParams(const gmm::GroupedMatmulParams &gmmParams) {1162static aclnnStatus CheckA4W4QuantParams(const gmm::GroupedMatmulParams &gmmParams) {
1157- CHECK_COND(gmmParams.groupListType == 0 || gmmParams.groupListType == 1, ACLNN_ERR_PARAM_INVALID,1163+ // 0: cumsum, 1: count, 2: sparse.
1158- "GMM A4W4: groupListType only support 0 or 1, but now is %ld.", gmmParams.groupListType);1164+ CHECK_COND(gmmParams.groupListType == 0 || gmmParams.groupListType == 1 || gmmParams.groupListType == 2,
1159- DataType yDtype = (*gmmParams.y)[0]->GetDataType();1165+ ACLNN_ERR_PARAM_INVALID,
1160- CHECK_COND(yDtype == DataType::DT_FLOAT16 || yDtype == DataType::DT_BF16, ACLNN_ERR_PARAM_INVALID,1166+ "GMM A4W4: groupListType only support 0(cumsum) or 1(count) or 2(sparsem), but now is %ld.",
1161- "GMM A4W4: output y dtype should be float16 or bfloat16, current dtype is %s.",1167+ gmmParams.groupListType);
1162- gmm::dTypeToString(yDtype).c_str());1168+ DataType yDtype = (*gmmParams.y)[0]->GetDataType();
1163- CHECK_COND(gmmParams.offsetOptional == nullptr, ACLNN_ERR_PARAM_INVALID,1169+ CHECK_COND(yDtype == DataType::DT_FLOAT16 || yDtype == DataType::DT_BF16, ACLNN_ERR_PARAM_INVALID,
1164- "GMM A4W4: offset must be null.");1170+ "GMM A4W4: output y dtype should be float16 or bfloat16, current dtype is %s.",
1165- CHECK_COND(gmmParams.biasOptional == nullptr, ACLNN_ERR_PARAM_INVALID,1171+ gmm::dTypeToString(yDtype).c_str());
1166- "GMM A4W4: bias must be null.");1172+ CHECK_COND(gmmParams.offsetOptional == nullptr, ACLNN_ERR_PARAM_INVALID,
1173+ "GMM A4W4: offset must be null.");
1174+ CHECK_COND(gmmParams.biasOptional == nullptr, ACLNN_ERR_PARAM_INVALID,
1175+ "GMM A4W4: bias must be null.");
1176+ CHECK_COND(gmmParams.scaleOptional != nullptr, ACLNN_ERR_PARAM_INVALID,
1177+ "GMM A4W4: scale must not be null.");
1178+ DataType scaleDtype = (*gmmParams.scaleOptional)[0]->GetDataType();
1179+ CHECK_COND(scaleDtype == DataType::DT_UINT64, ACLNN_ERR_PARAM_INVALID,
1180+ "GMM A4W4: scale dtype does not match with required dtype uint64, current dtype is %s.",
1181+ gmm::dTypeToString(scaleDtype).c_str());
1167 1182 
1168- CHECK_COND(gmmParams.scaleOptional != nullptr, ACLNN_ERR_PARAM_INVALID,1183+ bool isPerTokenQuant = gmmParams.perTokenScaleOptional != nullptr;
1169- "GMM A4W4: scale must not be null.");1184+ if (isPerTokenQuant) {
1170- DataType scaleDtype = (*gmmParams.scaleOptional)[0]->GetDataType();1185+ DataType perTokenScaleDtype = (*gmmParams.perTokenScaleOptional)[0]->GetDataType();
1171- CHECK_COND(scaleDtype == DataType::DT_UINT64, ACLNN_ERR_PARAM_INVALID,1186+ CHECK_COND(perTokenScaleDtype == DataType::DT_FLOAT, ACLNN_ERR_PARAM_INVALID,
1172- "GMM A4W4: scale dtype does not match with required dtype uint64, current dtype is %s.",1187+ "GMM A4W4: perTokenScale dtype does not match with required dtype float32, current dtype is %s.",
1173- gmm::dTypeToString(scaleDtype).c_str());1188+ gmm::dTypeToString(perTokenScaleDtype).c_str());
1174- 1189+ CHECK_COND(CheckPerTokenScale(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
1175- bool isPerTokenQuant = gmmParams.perTokenScaleOptional != nullptr;1190+ "GMM A4W4: Check perTokenScale failed!");
1176- if (isPerTokenQuant) {1191+ }
1177- DataType perTokenScaleDtype = (*gmmParams.perTokenScaleOptional)[0]->GetDataType();1192+ CHECK_COND(CheckA4W4ParamsShape(gmmParams) == ACLNN_SUCCESS,
1178- CHECK_COND(perTokenScaleDtype == DataType::DT_FLOAT, ACLNN_ERR_PARAM_INVALID,1193+ ACLNN_ERR_PARAM_INVALID,
1179- "GMM A4W4: perTokenScale dtype does not match with required dtype float32, current dtype is %s.",1194+ "CheckA4W4ParamsShape failed.");
1180- gmm::dTypeToString(perTokenScaleDtype).c_str());1195+ CHECK_COND(IsGmmAntiQuantEmpty(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
1181- 1196+ "GMM A4W4: Detected quant, but antiquant inputs are not empty!");
1182- CHECK_COND(CheckPerTokenScale(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,1197+ CHECK_COND(gmmParams.groupType == gmm::SPLIT_M && gmmParams.x->Size() == 1 && gmmParams.weight->Size() == 1
1183- "GMM A4W4: Check perTokenScale failed!");1198+ && gmmParams.y->Size() == 1, ACLNN_ERR_PARAM_INVALID,
1184- }1199+ "A4W4 only support split m, single x, single weight, single y.");
1185- CHECK_COND(CheckA4W4ParamsShape(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "CheckA4W4ParamsShape failed.");1200+ return ACLNN_SUCCESS;
1186- CHECK_COND(IsGmmAntiQuantEmpty(gmmParams) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID,
1187- "GMM A4W4: Detected quant, but antiquant inputs are not empty!");
1188- CHECK_COND(gmmParams.groupType == gmm::SPLIT_M && gmmParams.x->Size() == 1 && gmmParams.weight->Size() == 1
1189- && gmmParams.y->Size() == 1, ACLNN_ERR_PARAM_INVALID,
1190- "A4W4 only support split m, single x, single weight, single y.");
1191- return ACLNN_SUCCESS;
1192}1201}
1193 1202 
1194bool isActivationAllowed(int64_t act_type) {1203bool isActivationAllowed(int64_t act_type) {
@@ -2489,8 +2498,6 @@ aclnnStatus CheckCommonParam(const aclTensorList *x , const aclTensorList *weigh
2489 if (groupListType == gmm::GROUP_LIST_SPARSE_M) {2498 if (groupListType == gmm::GROUP_LIST_SPARSE_M) {
2490 CHECK_COND(npuArch == NpuArch::DAV_2201 || npuArch == NpuArch::DAV_3510, ACLNN_ERR_PARAM_INVALID,2499 CHECK_COND(npuArch == NpuArch::DAV_2201 || npuArch == NpuArch::DAV_3510, ACLNN_ERR_PARAM_INVALID,
2491 "This platform not support groupListType is 2.");2500 "This platform not support groupListType is 2.");
2492- CHECK_COND(groupType == gmm::SPLIT_M, ACLNN_ERR_PARAM_INVALID,
2493- "When groupListType is 2 only support groupType 0, but get groupType %ld.", groupType);
2494 } else {2501 } else {
2495 CHECK_COND(groupListType == 0 || groupListType == 1, ACLNN_ERR_PARAM_INVALID, "groupListType shoule be 0 or 1.");2502 CHECK_COND(groupListType == 0 || groupListType == 1, ACLNN_ERR_PARAM_INVALID, "groupListType shoule be 0 or 1.");
2496 }2503 }
Mgmm/grouped_matmul/op_host/grouped_matmul_infershape.cpp+10-5
@@ -142,8 +142,12 @@ static ge::graphStatus UpdateMultipleShapeY(gert::InferShapeContext* context, co
142 OP_CHECK_NULL_WITH_CONTEXT(context, x0Shape);142 OP_CHECK_NULL_WITH_CONTEXT(context, x0Shape);
143 const gert::Shape* weight0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0);143 const gert::Shape* weight0Shape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, 0);
144 OP_CHECK_NULL_WITH_CONTEXT(context, weight0Shape);144 OP_CHECK_NULL_WITH_CONTEXT(context, weight0Shape);
145+ // For SPARSEM (groupListType=2), groupList shape is [E, 2], so loop count should be E (first dim)
146+ // For CUMSUM/COUNT (groupListType=0/1), groupList is 1D, loop count is shape size
147+ int64_t loopCount = (*groupListTypePtr == GROUP_LIST_SPARSE) ?
148+ groupListTensor->GetStorageShape().GetDim(0) : groupListTensor->GetShapeSize();
145 int64_t preOffset = 0;149 int64_t preOffset = 0;
146- for (int idx = 0; idx < groupListTensor->GetShapeSize(); ++idx) {150+ for (int idx = 0; idx < loopCount; ++idx) {
147 const gert::Shape* weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idx);151 const gert::Shape* weightShape = context->GetDynamicInputShape(GMM_INDEX_IN_WEIGHT, idx);
148 if (weightShape == nullptr) {152 if (weightShape == nullptr) {
149 weightShape = weight0Shape;153 weightShape = weight0Shape;
@@ -160,15 +164,16 @@ static ge::graphStatus UpdateMultipleShapeY(gert::InferShapeContext* context, co
160 std::vector<int64_t> yDims;164 std::vector<int64_t> yDims;
161 if (*groupListTypePtr == 0) {165 if (*groupListTypePtr == 0) {
162 yDims = {groupListData[idx] - preOffset, weightShape->GetDim(weightDimN)};166 yDims = {groupListData[idx] - preOffset, weightShape->GetDim(weightDimN)};
167+ preOffset = groupListData[idx];
163 } else if (*groupListTypePtr == 1) {168 } else if (*groupListTypePtr == 1) {
164 yDims = {groupListData[idx], weightShape->GetDim(weightDimN)};169 yDims = {groupListData[idx], weightShape->GetDim(weightDimN)};
165- } else {170+ } else if (*groupListTypePtr == GROUP_LIST_SPARSE) {
166- OP_LOGE(context->GetNodeName(), "Invalid groupListType = %ld", *groupListTypePtr);171+ // SPARSEM: groupList shape is [E, 2], second column (idx*2+1) is token count per group
167- return GRAPH_FAILED;172+ yDims = {groupListData[idx * GROUP_LIST_SPARSE + GROUP_LIST_SPARSE_OFFSET],
173+ weightShape->GetDim(weightDimN)};
168 }174 }
169 OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y + idx, yDims) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(),175 OP_CHECK_IF(UpdateShapeY(context, GMM_INDEX_OUT_Y + idx, yDims) != GRAPH_SUCCESS, OP_LOGE(context->GetNodeName(),
170 "Failed to update shape of y."), return GRAPH_FAILED);176 "Failed to update shape of y."), return GRAPH_FAILED);
171- preOffset = groupListData[idx];
172 }177 }
173 }178 }
174 179 
Mgmm/grouped_matmul/op_host/grouped_matmul_infershape_common_util.h+2-0
@@ -70,6 +70,8 @@ constexpr int64_t MXFP_DIVISOR_SIZE = 64;
70constexpr int64_t MXFP_MULTI_BASE_SIZE = 2;70constexpr int64_t MXFP_MULTI_BASE_SIZE = 2;
71constexpr int64_t MXFP_TYPEM_SCALE_DIM_NUM = 4;71constexpr int64_t MXFP_TYPEM_SCALE_DIM_NUM = 4;
72constexpr int64_t MXFP_TYPEK_SCALE_DIM_NUM = 3;72constexpr int64_t MXFP_TYPEK_SCALE_DIM_NUM = 3;
73+constexpr int64_t GROUP_LIST_SPARSE = 2;
74+constexpr int64_t GROUP_LIST_SPARSE_OFFSET = 1;
73struct GMMAttrs {75struct GMMAttrs {
74 int64_t splitItem;76 int64_t splitItem;
75 int64_t outputDtype;77 int64_t outputDtype;
Mgmm/grouped_matmul/op_host/op_tiling/grouped_matmul_tiling.cpp+1-10
@@ -1197,18 +1197,9 @@ ge::graphStatus GMMTiling::GMMGetAttrs(const gert::TilingContext* context) {
1197 OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null1197 OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr); // check compileInfoPtr is not null
1198 if (groupListType_ == GROUP_LIST_SPARSE_M) {1198 if (groupListType_ == GROUP_LIST_SPARSE_M) {
1199 OP_CHECK_IF((!(compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910B ||1199 OP_CHECK_IF((!(compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910B ||
1200- compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910_93)),1200+ compileInfoPtr->socVersion == platform_ascendc::SocVersion::ASCEND910_93)),
1201 OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "This platform not support groupListType is 2"),1201 OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(), "This platform not support groupListType is 2"),
1202 return ge::GRAPH_FAILED);1202 return ge::GRAPH_FAILED);
1203- OP_CHECK_IF(!isA8W8_,
1204- OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(),
1205- "Only both x and weight dtype are int8 support groupListType is 2"),
1206- return ge::GRAPH_FAILED);
1207- OP_CHECK_IF(groupType_ != SPLIT_M,
1208- OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(),
1209- "When groupListType is 2 only support groupType 0, but get groupType %ld",
1210- groupType_),
1211- return ge::GRAPH_FAILED);
1212 }1203 }
1213 1204 
1214 auto perTokenScalePtr = context->GetOptionalInputTensor(PER_TOKEN_SCALE_INDEX);1205 auto perTokenScalePtr = context->GetOptionalInputTensor(PER_TOKEN_SCALE_INDEX);
Mgmm/grouped_matmul/op_kernel/a16w4_msd/grouped_matmul_weight_quant_a16w4_msd_controller.h+137-7
@@ -27,6 +27,10 @@ namespace GROUPED_MATMUL::A16W4Msd {
27 template <typename xType, typename wType, typename biasType, int groupListType>27 template <typename xType, typename wType, typename biasType, int groupListType>
28 28 
29#define GMM_WQ_A16W4_MSD_CONTROLLER_CLASS GMMWeightQuantA16W4MsdController<xType, wType, biasType, groupListType>29#define GMM_WQ_A16W4_MSD_CONTROLLER_CLASS GMMWeightQuantA16W4MsdController<xType, wType, biasType, groupListType>
30+#define GMM_WQ_A16W4_MSD_CONTROLLER_CLASS_M_SPARSE \
31+ GMMWeightQuantA16W4MsdControllerMSparse<xType, wType, biasType, groupListType>
32+constexpr uint32_t SPARSE_GROUP_LIST_SPLIT = 2;
33+constexpr uint32_t SPARSE_GROUP_LIST_SPLIT_VALUE_OFFSET = 1;
30 34 
31GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM35GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
32class GMMWeightQuantA16W4MsdController {36class GMMWeightQuantA16W4MsdController {
@@ -36,7 +40,7 @@ public:
36 GM_ADDR y, const GMMBaseParams *__restrict baseTiling);40 GM_ADDR y, const GMMBaseParams *__restrict baseTiling);
37 __aicore__ inline void Process(GM_ADDR workspace, TPipe *tPipe);41 __aicore__ inline void Process(GM_ADDR workspace, TPipe *tPipe);
38 42 
39-private:43+protected:
40 __aicore__ inline void PreProcess(const A16W4MsdConstParam &constParams,44 __aicore__ inline void PreProcess(const A16W4MsdConstParam &constParams,
41 WeightQuantA16W4MsdBasicBlock<xType, wType, biasType> &basicBlock);45 WeightQuantA16W4MsdBasicBlock<xType, wType, biasType> &basicBlock);
42 __aicore__ inline void InitConstParam(A16W4MsdConstParam &constParams);46 __aicore__ inline void InitConstParam(A16W4MsdConstParam &constParams);
@@ -45,7 +49,8 @@ private:
45 __aicore__ inline void InitWorkspaceSize(uint64_t cubeBlockIdx, GM_ADDR workspace);49 __aicore__ inline void InitWorkspaceSize(uint64_t cubeBlockIdx, GM_ADDR workspace);
46 __aicore__ inline void SetOffsetParam(uint64_t mSize, uint64_t mOffset, uint64_t nOffset, uint64_t kOffset,50 __aicore__ inline void SetOffsetParam(uint64_t mSize, uint64_t mOffset, uint64_t nOffset, uint64_t kOffset,
47 const A16W4MsdConstParam &constParams,51 const A16W4MsdConstParam &constParams,
48- A16W4MsdBasicBlockOffsetParam &offsetParam);52+ A16W4MsdBasicBlockOffsetParam &offsetParam,
53+ __gm__ xType *antiquantScaleGm);
49 54 
50 const GMMBaseParams *gmmBaseTiling_;55 const GMMBaseParams *gmmBaseTiling_;
51 56 
@@ -123,7 +128,7 @@ __aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::Process(GM_ADDR worksp
123 for (uint64_t kOffset = 0; kOffset < constParams.kSize; kOffset += BASE_KUB_SIZE) {128 for (uint64_t kOffset = 0; kOffset < constParams.kSize; kOffset += BASE_KUB_SIZE) {
124 SetOffsetParam(mSize, ((curCoreStartBlockId - startBasicBlockId) / nBlockNum) * BASE_ML1_SIZE,129 SetOffsetParam(mSize, ((curCoreStartBlockId - startBasicBlockId) / nBlockNum) * BASE_ML1_SIZE,
125 ((curCoreStartBlockId - startBasicBlockId) % nBlockNum) * constParams.nL1BaseSize,130 ((curCoreStartBlockId - startBasicBlockId) % nBlockNum) * constParams.nL1BaseSize,
126- kOffset, constParams, offsetParam[processId]);131+ kOffset, constParams, offsetParam[processId], antiquantScaleGm_);
127 basicBlock.ComputeBasicBlock(constParams, offsetParam[processId], offsetParam[1 - processId]);132 basicBlock.ComputeBasicBlock(constParams, offsetParam[processId], offsetParam[1 - processId]);
128 processId = 1 - processId;133 processId = 1 - processId;
129 }134 }
@@ -159,7 +164,8 @@ GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
159__aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::SetOffsetParam(uint64_t mSize, uint64_t mOffset,164__aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::SetOffsetParam(uint64_t mSize, uint64_t mOffset,
160 uint64_t nOffset, uint64_t kOffset,165 uint64_t nOffset, uint64_t kOffset,
161 const A16W4MsdConstParam &constParams,166 const A16W4MsdConstParam &constParams,
162- A16W4MsdBasicBlockOffsetParam &offsetParam)167+ A16W4MsdBasicBlockOffsetParam &offsetParam,
168+ __gm__ xType *antiquantScaleGm)
163{169{
164 offsetParam.kOffset = kOffset;170 offsetParam.kOffset = kOffset;
165 offsetParam.kUbSize = kOffset + BASE_KUB_SIZE > constParams.kSize ? constParams.kSize - kOffset : BASE_KUB_SIZE;171 offsetParam.kUbSize = kOffset + BASE_KUB_SIZE > constParams.kSize ? constParams.kSize - kOffset : BASE_KUB_SIZE;
@@ -173,7 +179,7 @@ __aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::SetOffsetParam(uint64_
173 ? constParams.nSize - offsetParam.nOffset179 ? constParams.nSize - offsetParam.nOffset
174 : constParams.nL1BaseSize;180 : constParams.nL1BaseSize;
175 offsetParam.yGmAddr = reinterpret_cast<GM_ADDR>(yGm_);181 offsetParam.yGmAddr = reinterpret_cast<GM_ADDR>(yGm_);
176- offsetParam.antiquantScaleGm = reinterpret_cast<GM_ADDR>(antiquantScaleGm_);182+ offsetParam.antiquantScaleGm = reinterpret_cast<GM_ADDR>(antiquantScaleGm);
177 offsetParam.aMaxGmAddr = reinterpret_cast<GM_ADDR>(aMaxWs_);183 offsetParam.aMaxGmAddr = reinterpret_cast<GM_ADDR>(aMaxWs_);
178}184}
179 185 
@@ -190,8 +196,14 @@ __aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::PreProcess(
190 GlobalTensor<int8_t> aUnfoldGlobal;196 GlobalTensor<int8_t> aUnfoldGlobal;
191 aUnfoldGlobal.SetGlobalBuffer(aUnfoldS8WsAddr_);197 aUnfoldGlobal.SetGlobalBuffer(aUnfoldS8WsAddr_);
192 uint64_t preSum = 0;198 uint64_t preSum = 0;
193- for (uint32_t groupIdx = 0, startBasicBlockId = 0; groupIdx < gmmBaseTiling_->groupNum; ++groupIdx) {199+ // 2: groupList shape: [e, 2]; 1: groupList shape: [e]
200+ uint32_t groupListInnerShape = groupListType == SPARSE_GROUP_LIST_SPLIT ? 2 : 1;
201+ uint32_t groupListShapeSize = gmmBaseTiling_->groupNum * groupListInnerShape;
202+ for (uint32_t groupIdx = 0, startBasicBlockId = 0; groupIdx < groupListShapeSize; groupIdx += groupListInnerShape) {
194 uint64_t mSize = GetSplitValueFromGroupList(groupIdx);203 uint64_t mSize = GetSplitValueFromGroupList(groupIdx);
204+ if (groupListType == SPARSE_GROUP_LIST_SPLIT && mSize == 0) {
205+ break;
206+ }
195 if (mSize > 0) {207 if (mSize > 0) {
196 xGlobal.SetGlobalBuffer(xGm_);208 xGlobal.SetGlobalBuffer(xGm_);
197 uint32_t totalBasicBlockCount = startBasicBlockId + mSize;209 uint32_t totalBasicBlockCount = startBasicBlockId + mSize;
@@ -243,14 +255,132 @@ __aicore__ inline uint64_t GMM_WQ_A16W4_MSD_CONTROLLER_CLASS::GetSplitValueFromG
243{255{
244 uint64_t splitValue = 0;256 uint64_t splitValue = 0;
245 if constexpr (groupListType == 0) {257 if constexpr (groupListType == 0) {
258+ // 0: cumsum mode
246 uint64_t offset = static_cast<uint64_t>(groupListGm_.GetValue(groupIdx));259 uint64_t offset = static_cast<uint64_t>(groupListGm_.GetValue(groupIdx));
247 splitValue = offset - preOffset_;260 splitValue = offset - preOffset_;
248 preOffset_ = offset;261 preOffset_ = offset;
249- } else {262+ } else if (groupListType == 1) {
263+ // 1: count mode
250 splitValue = static_cast<uint64_t>(groupListGm_.GetValue(groupIdx));264 splitValue = static_cast<uint64_t>(groupListGm_.GetValue(groupIdx));
265+ } else {
266+ // 2： sparse mode
267+ splitValue = static_cast<int32_t>(groupListGm_.GetValue(groupIdx + SPARSE_GROUP_LIST_SPLIT_VALUE_OFFSET));
251 }268 }
252 return splitValue;269 return splitValue;
253}270}
271+ 
272+GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
273+class GMMWeightQuantA16W4MsdControllerMSparse : public GMM_WQ_A16W4_MSD_CONTROLLER_CLASS {
274+public:
275+ __aicore__ inline GMMWeightQuantA16W4MsdControllerMSparse() : GMM_WQ_A16W4_MSD_CONTROLLER_CLASS() {};
276+ __aicore__ inline void Process(GM_ADDR workspace, TPipe *tPipe);
277+private:
278+ __aicore__ inline void PreUpdateGmAddrForSparse(GlobalTensor<wType> &weightGmUpdate,
279+ GlobalTensor<xType> &antiquantScaleGmUpdate,
280+ GlobalTensor<biasType> &biasGmUpdate,
281+ uint64_t kSize, uint64_t nSize, uint32_t expertIdx);
282+ __aicore__ inline void PostUpdateGmAddrForSparse(uint64_t mSize, uint64_t kSize, uint64_t nSize);
283+ 
284+ static constexpr uint64_t BASIC_BLOCK_PROCESS_NUM = 2;
285+ static constexpr uint64_t BASE_ML1_SIZE = 32;
286+ static constexpr uint64_t GROUP_SIZE = 32;
287+};
288+ 
289+GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
290+__aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS_M_SPARSE::Process(GM_ADDR workspace, TPipe *tPipe)
291+{
292+ uint32_t cubeBlockIdx = GetBlockIdx();
293+ if ASCEND_IS_AIV {
294+ cubeBlockIdx = GetBlockIdx() >> 1;
295+ }
296+ A16W4MsdConstParam constParams;
297+ this->InitConstParam(constParams);
298+ 
299+ WeightQuantA16W4MsdBasicBlock<xType, wType, biasType> basicBlock;
300+ basicBlock.InitPreProcess(tPipe);
301+ 
302+ this->InitWorkspaceSize(cubeBlockIdx, workspace);
303+ this->PreProcess(constParams, basicBlock);
304+ SyncAll();
305+ 
306+ basicBlock.InitMsd(this->weightS8WsAddr_, this->aUnfoldS8WsAddr_, this->cF32WsAddr_, GROUP_SIZE, tPipe);
307+ 
308+ A16W4MsdBasicBlockOffsetParam offsetParam[BASIC_BLOCK_PROCESS_NUM];
309+ GlobalTensor<wType> weightGmUpdate;
310+ GlobalTensor<xType> antiquantScaleGmUpdate;
311+ GlobalTensor<biasType> biasGmUpdate;
312+ uint64_t processId = 0;
313+ 
314+ uint32_t groupListSplitValueOffset = 1;
315+ uint32_t groupListInnerShape = 2u; // groupList shape: [e, 2]
316+ uint32_t groupListShapeSize = this->gmmBaseTiling_->groupNum * groupListInnerShape;
317+ for (uint32_t loop = 0, startBasicBlockId = 0; loop < groupListShapeSize; loop += groupListInnerShape) {
318+ uint64_t mSize = this->GetSplitValueFromGroupList(loop);
319+ if (mSize == 0) {
320+ // all tokens have be processed
321+ break;
322+ }
323+ uint32_t expertIdx = static_cast<uint32_t>(this->groupListGm_.GetValue(loop));
324+ if (mSize > 0) {
325+ PreUpdateGmAddrForSparse(weightGmUpdate, antiquantScaleGmUpdate, biasGmUpdate,
326+ constParams.kSize, constParams.nSize, expertIdx);
327+ basicBlock.UpdateGlobalAddr(this->aUnfoldS8WsAddr_, (__gm__ wType *)weightGmUpdate.GetPhyAddr(),
328+ (__gm__ xType *)antiquantScaleGmUpdate.GetPhyAddr(),
329+ (__gm__ biasType *)biasGmUpdate.GetPhyAddr(),
330+ this->yGm_, this->aMaxWs_, this->gmmBaseTiling_->withOffset);
331+ 
332+ uint32_t mBlockNum = A16W4Msd::CeilDiv(mSize, BASE_ML1_SIZE);
333+ uint32_t nBlockNum = A16W4Msd::CeilDiv(constParams.nSize, constParams.nL1BaseSize);
334+ 
335+ uint32_t totalBasicBlockCount = startBasicBlockId + mBlockNum * nBlockNum;
336+ uint32_t curCoreStartBlockId =
337+ cubeBlockIdx >= startBasicBlockId ? cubeBlockIdx : cubeBlockIdx + this->gmmBaseTiling_->coreNum;
338+ 
339+ while (curCoreStartBlockId < totalBasicBlockCount) {
340+ for (uint64_t kOffset = 0; kOffset < constParams.kSize; kOffset += BASE_KUB_SIZE) {
341+ this->SetOffsetParam(mSize,
342+ ((curCoreStartBlockId - startBasicBlockId) / nBlockNum) * BASE_ML1_SIZE,
343+ ((curCoreStartBlockId - startBasicBlockId) % nBlockNum) * constParams.nL1BaseSize,
344+ kOffset, constParams, offsetParam[processId],
345+ (__gm__ xType *)antiquantScaleGmUpdate.GetPhyAddr());
346+ basicBlock.ComputeBasicBlock(constParams, offsetParam[processId], offsetParam[1 - processId]);
347+ processId = 1 - processId;
348+ }
349+ curCoreStartBlockId += this->gmmBaseTiling_->coreNum;
350+ }
351+ startBasicBlockId = totalBasicBlockCount % this->gmmBaseTiling_->coreNum;
352+ }
353+ PostUpdateGmAddrForSparse(mSize, constParams.kSize, constParams.nSize);
354+ }
355+ 
356+ basicBlock.EndMsd(constParams, offsetParam[1 - processId]);
357+}
358+ 
359+GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
360+__aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS_M_SPARSE::PreUpdateGmAddrForSparse(
361+ GlobalTensor<wType> &weightGmUpdate, GlobalTensor<xType> &antiquantScaleGmUpdate,
362+ GlobalTensor<biasType> &biasGmUpdate, uint64_t kSize, uint64_t nSize, uint32_t expertIdx)
363+{
364+ uint32_t wBaseOffset = 0;
365+ if constexpr (IsSameType<wType, int4b_t>::value) {
366+ wBaseOffset = (expertIdx * nSize * kSize) >> 1;
367+ } else {
368+ wBaseOffset = expertIdx * nSize * kSize;
369+ }
370+ weightGmUpdate.SetGlobalBuffer(this->weightGm_ + wBaseOffset);
371+ antiquantScaleGmUpdate.SetGlobalBuffer(
372+ this->antiquantScaleGm_ + expertIdx * nSize * this->gmmBaseTiling_->quantGroupNum);
373+ biasGmUpdate.SetGlobalBuffer(this->biasGm_ + expertIdx * nSize);
374+}
375+ 
376+GMM_WQ_A16W4_MSD_CONTROLLER_TEMPLATE_PARAM
377+__aicore__ inline void GMM_WQ_A16W4_MSD_CONTROLLER_CLASS_M_SPARSE::PostUpdateGmAddrForSparse(
378+ uint64_t mSize, uint64_t kSize, uint64_t nSize)
379+{
380+ this->aUnfoldS8WsAddr_ += 2 * mSize * kSize;
381+ this->aMaxWs_ += mSize * 8;
382+ this->yGm_ += mSize * nSize;
383+}
L
Llijiamin20255月18日

增加tiling UT

likedislike
254} // namespace GROUPED_MATMUL::A16W4Msd384} // namespace GROUPED_MATMUL::A16W4Msd
255 385 
256#endif // GROUPED_MATMUL_WEIGHT_QUANT_A16W4_MSD_CONTROLLER_H386#endif // GROUPED_MATMUL_WEIGHT_QUANT_A16W4_MSD_CONTROLLER_H
Mgmm/grouped_matmul/op_kernel/grouped_matmul.cpp+86-28
@@ -397,16 +397,32 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
397 // ANTIQUANT397 // ANTIQUANT
398 if constexpr ((D_T_A == GMM_TPL_BF16) &&398 if constexpr ((D_T_A == GMM_TPL_BF16) &&
399 A16W8_KERNEL_TEMPLATE == GROUPED_MATMUL_A16W4_KERNEL_TEMPLATE_MSD_ANTIQUANT_GS32) {399 A16W8_KERNEL_TEMPLATE == GROUPED_MATMUL_A16W4_KERNEL_TEMPLATE_MSD_ANTIQUANT_GS32) {
400- GMM_CV_SPLIT_IMP_A16W4_MSD(A16W4Msd::GMMWeightQuantA16W4MsdController, false);400+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
401+ GMM_CV_SPLIT_IMP_A16W4_MSD(A16W4Msd::GMMWeightQuantA16W4MsdControllerMSparse, false);
402+ } else {
403+ GMM_CV_SPLIT_IMP_A16W4_MSD(A16W4Msd::GMMWeightQuantA16W4MsdController, false);
404+ }
401 } else if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) &&405 } else if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) &&
402 A16W8_KERNEL_TEMPLATE != GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) {406 A16W8_KERNEL_TEMPLATE != GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) {
403 // ANTIQUANT_A16W4 & ANTIQUANT_A16W8_NOT_MSD407 // ANTIQUANT_A16W4 & ANTIQUANT_A16W8_NOT_MSD
404 if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {408 if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {
405- GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, false, false, matmulCFG);409+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
410+ GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantSparseProcess, false, false, false, matmulCFG);
411+ } else {
412+ GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, false, false, matmulCFG);
413+ }
406 } else if constexpr (TRANS_B == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {414 } else if constexpr (TRANS_B == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {
407- GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, true, false, matmulCFG);415+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
416+ GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantSparseProcess, false, true, false, matmulCFG);
417+ } else {
418+ GMM_IMP(GMMAntiquantComputeNorm, GMMAntiquantProcess, false, true, false, matmulCFG);
419+ }
408 } else if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_2) {420 } else if constexpr (TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_2) {
409- GMM_IMP(GMMAntiquantComputePerformance, GMMAntiquantProcess, false, false, false, matmulCFG);421+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
422+ GMM_IMP(GMMAntiquantComputePerformance, GMMAntiquantSparseProcess, false, false, false, matmulCFG);
423+ } else {
424+ GMM_IMP(GMMAntiquantComputePerformance, GMMAntiquantProcess, false, false, false, matmulCFG);
425+ }
410 }426 }
411 }427 }
412#if defined(ORIG_DTYPE_WEIGHT) && defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8428#if defined(ORIG_DTYPE_WEIGHT) && defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8
@@ -414,11 +430,21 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
414 if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) && D_T_B == GMM_TPL_INT8 &&430 if constexpr ((D_T_A == GMM_TPL_FLOAT16 || D_T_A == GMM_TPL_BF16) && D_T_B == GMM_TPL_INT8 &&
415 A16W8_KERNEL_TEMPLATE == GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) {431 A16W8_KERNEL_TEMPLATE == GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD) {
416 if constexpr (TRANS_B == 0) {432 if constexpr (TRANS_B == 0) {
417- GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, false, false, matmulCFG,433+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
418- xTypeMSD, weightTypeMSD, yTypeMSD);434+ GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDMSparseProcess, false, false, false,
435+ matmulCFG, xTypeMSD, weightTypeMSD, yTypeMSD);
436+ } else {
437+ GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, false, false,
438+ matmulCFG, xTypeMSD, weightTypeMSD, yTypeMSD);
439+ }
419 } else if constexpr (TRANS_B == 1) {440 } else if constexpr (TRANS_B == 1) {
420- GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, true, false, matmulCFG,441+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
421- xTypeMSD, weightTypeMSD, yTypeMSD);442+ GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDMSparseProcess, false, true, false,
443+ matmulCFG, xTypeMSD, weightTypeMSD, yTypeMSD);
444+ } else {
445+ GMM_CV_SPLIT_IMP(GMMA16W8MSDCompute, GMMA16W8MSDProcess, false, true, false,
446+ matmulCFG, xTypeMSD, weightTypeMSD, yTypeMSD);
447+ }
422 }448 }
423 }449 }
424#endif450#endif
@@ -484,7 +510,7 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
484#elif defined(GMM_A4W4)510#elif defined(GMM_A4W4)
485 // QUANT_A4W4511 // QUANT_A4W4
486 if constexpr (D_T_A == GMM_TPL_INT4 && D_T_B == GMM_TPL_INT4) {512 if constexpr (D_T_A == GMM_TPL_INT4 && D_T_B == GMM_TPL_INT4) {
487- if constexpr (TRANS_B == 0){513+ if constexpr (TRANS_B == 0) {
488 GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling);514 GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling);
489 if (gmmBaseParams_.isA4W4Optimize) {515 if (gmmBaseParams_.isA4W4Optimize) {
490 tPipe.Destroy();516 tPipe.Destroy();
@@ -496,10 +522,18 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
496#endif522#endif
497 AscendC::SetMMLayoutTransform(false);523 AscendC::SetMMLayoutTransform(false);
498 } else {524 } else {
499- GMM_A4W4_IMP(GMMA4W4Compute, false, false, matmulCFG, xType, weightType, yType);525+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
526+ GMM_A4W4_IMP(GMMA4W4SparseCompute, false, false, matmulCFG, xType, weightType, yType);
527+ } else {
528+ GMM_A4W4_IMP(GMMA4W4Compute, false, false, matmulCFG, xType, weightType, yType);
529+ }
530+ }
531+ } else {
532+ if constexpr (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM) {
533+ GMM_A4W4_IMP(GMMA4W4SparseCompute, false, true, matmulCFG, xType, weightType, yType);
534+ } else {
535+ GMM_A4W4_IMP(GMMA4W4Compute, false, true, matmulCFG, xType, weightType, yType);
500 }536 }
501- }else{
502- GMM_A4W4_IMP(GMMA4W4Compute, false, true, matmulCFG, xType, weightType, yType);
503 }537 }
504 }538 }
505#elif defined(GMM_QUANT_INT8) || defined(GMM_QUANT_INT32)539#elif defined(GMM_QUANT_INT8) || defined(GMM_QUANT_INT32)
@@ -531,22 +565,33 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
531 }565 }
532#elif defined(GMM_FLOAT)566#elif defined(GMM_FLOAT)
533 // NO_QUANT567 // NO_QUANT
534- if (GROUP_LIST_TYPE != GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM &&568+ if (IS_STATIC_TILING_API == 0 && A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE) {
535- IS_STATIC_TILING_API == 0 &&569+ GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling);
536- A8W4_KERNEL_TEMPLATE == GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE) {570+ if constexpr (TRANS_A == 0 && TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) {
537- if constexpr (TRANS_A == 0 && TRANS_B == 0 && AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) {571+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
538- GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFGUnitFlag);572+ GMM_CUBE_IMP(GMMGroupMSparseProcess, false, false, false, matmulCFGUnitFlag);
539- } else if constexpr (TRANS_A == 0 && TRANS_B == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) {573+ } else {
540- GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFGUnitFlag);574+ GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFGUnitFlag);
541- } else if constexpr (TRANS_A == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {575+ }
542- if ASCEND_IS_AIV {576+ } else if constexpr (TRANS_A == 0 && TRANS_B == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_CUBE_ONLY) {
543- GET_TILING_DATA(tilingData, tiling);577+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
544- EmptyTensorCompute<DTYPE_Y>(groupList, y, &tilingData);578+ GMM_CUBE_IMP(GMMGroupMSparseProcess, false, true, false, matmulCFGUnitFlag);
545- }579+ } else {
546- if ASCEND_IS_AIC {580+ GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFGUnitFlag);
581+ }
582+ } else if constexpr (TRANS_A == 1 && AIV_AIC_RATIO == GROUPED_MATMUL_AIV_AIC_RATIO_1) {
583+ if ASCEND_IS_AIV {
584+ GET_TILING_DATA(tilingData, tiling);
585+ EmptyTensorCompute<DTYPE_Y>(groupList, y, &tilingData);
586+ }
587+ if ASCEND_IS_AIC {
588+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
589+ GMM_CUBE_IMP(GMMGroupMSparseProcess, true, false, false, matmulCFG);
590+ } else {
547 GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG);591 GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG);
548 }592 }
549 }593 }
594+ }
550 }595 }
551 596 
552#endif597#endif
@@ -554,17 +599,30 @@ __global__ __aicore__ void grouped_matmul(GM_ADDR x, GM_ADDR weight, GM_ADDR bia
554 599 
555#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200600#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 200
556#if defined(GMM_FLOAT)601#if defined(GMM_FLOAT)
602+ GET_TILING_DATA_MEMBER(GMMTilingData, gmmBaseParams, gmmBaseParams_, tiling);
557 if constexpr (TRANS_A == 0 && TRANS_B == 0) {603 if constexpr (TRANS_A == 0 && TRANS_B == 0) {
558- GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFG);604+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
605+ GMM_CUBE_IMP(GMMGroupMSparseProcess, false, false, false, matmulCFG);
606+ } else {
607+ GMM_CUBE_IMP(GMMProcess, false, false, false, matmulCFG);
608+ }
559 } else if constexpr (TRANS_A == 0 && TRANS_B == 1) {609 } else if constexpr (TRANS_A == 0 && TRANS_B == 1) {
560- GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFG);610+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
611+ GMM_CUBE_IMP(GMMGroupMSparseProcess, false, true, false, matmulCFG);
612+ } else {
613+ GMM_CUBE_IMP(GMMProcess, false, true, false, matmulCFG);
614+ }
561 } else if constexpr (TRANS_A == 1 && TRANS_B == 0) {615 } else if constexpr (TRANS_A == 1 && TRANS_B == 0) {
562 if ASCEND_IS_AIV {616 if ASCEND_IS_AIV {
563 GET_TILING_DATA(tilingData, tiling);617 GET_TILING_DATA(tilingData, tiling);
564 EmptyTensorCompute<DTYPE_Y>(groupList, y, &tilingData);618 EmptyTensorCompute<DTYPE_Y>(groupList, y, &tilingData);
565 }619 }
566 if ASCEND_IS_AIC {620 if ASCEND_IS_AIC {
567- GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG);621+ if (GROUP_LIST_TYPE == GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM && gmmBaseParams_.groupType == 0) {
622+ GMM_CUBE_IMP(GMMGroupMSparseProcess, true, false, false, matmulCFG);
623+ } else {
624+ GMM_CUBE_IMP(GMMProcess, true, false, false, matmulCFG);
625+ }
568 }626 }
569 }627 }
570#endif628#endif
Mgmm/grouped_matmul/op_kernel/grouped_matmul.h+55-20
@@ -150,6 +150,9 @@ class GMMProcess {
150 __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);150 __aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
151 151 
152 __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);152 __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
153+ 
154+ __aicore__ inline bool UpdateMnConfigForGroupListMSparse(
155+ MNConfig &mnConfig, uint32_t splitValue, uint32_t groupIdx);
153};156};
154 157 
155template <typename ComputeType>158template <typename ComputeType>
@@ -247,6 +250,28 @@ __aicore__ inline void GMMProcess<ComputeType>::UpdateMnConfig(MNConfig &mnConfi
247 mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;250 mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
248}251}
249 252 
253+template <typename ComputeType>
254+__aicore__ inline bool GMMProcess<ComputeType>::UpdateMnConfigForGroupListMSparse(
255+ MNConfig &mnConfig, uint32_t splitValue, uint32_t groupIdx)
256+{
257+ mnConfig.mAxisBaseOffset += mnConfig.m;
258+ mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
259+ mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
260+ mnConfig.scaleIndex++;
261+ SetMNConfig(splitValue, groupIdx, mnConfig);
262+ if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
263+ return true; // skip
264+ }
265+ 
266+ mnConfig.nAxisBaseOffset = groupIdx * mnConfig.n;
267+ if constexpr (GMMProcess<ComputeType>::B::format == CubeFormat::NZ) {
268+ mnConfig.wBaseOffset = AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.nAxisBaseOffset);
269+ } else {
270+ mnConfig.wBaseOffset = mnConfig.k * mnConfig.nAxisBaseOffset;
271+ }
272+ return false; // no skip
273+}
274+ 
250template <typename ComputeType>275template <typename ComputeType>
251__aicore__ inline void GMMProcess<ComputeType>::Process() {276__aicore__ inline void GMMProcess<ComputeType>::Process() {
252 MNConfig mnConfig;277 MNConfig mnConfig;
@@ -257,9 +282,16 @@ __aicore__ inline void GMMProcess<ComputeType>::Process() {
257 preOffset = 0;282 preOffset = 0;
258 }283 }
259 AscendC::WaitPreTaskEnd();284 AscendC::WaitPreTaskEnd();
260- for (uint32_t groupIdx = 0, count = 0; groupIdx < groupNum; ++groupIdx) {285+ // 2: groupList shape: [e, 2]; 1: groupList shape: [e]
286+ uint32_t groupListType = gmmBaseParams->groupListType;
287+ uint32_t groupListInnerShape = groupListType == GROUP_LIST_TYPE_SPARSE ? 2 : 1;
288+ uint32_t groupListShapeSize = groupNum * groupListInnerShape;
289+ for (uint32_t groupIdx = 0, count = 0; groupIdx < groupListShapeSize; groupIdx += groupListInnerShape) {
261 UpdateMnConfig(mnConfig);290 UpdateMnConfig(mnConfig);
262 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm);291 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm);
292+ if (groupListType == GROUP_LIST_TYPE_SPARSE && splitValue <= 0) {
293+ break;
294+ }
263 SetMNConfig(splitValue, groupIdx, mnConfig);295 SetMNConfig(splitValue, groupIdx, mnConfig);
264 if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {296 if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
265 continue;297 continue;
@@ -303,27 +335,18 @@ public:
303 uint32_t groupListInnerShape = 2u; // shape: [e, 2]335 uint32_t groupListInnerShape = 2u; // shape: [e, 2]
304 uint32_t groupListShapeSize = this->groupNum * groupListInnerShape;336 uint32_t groupListShapeSize = this->groupNum * groupListInnerShape;
305 AscendC::WaitPreTaskEnd();337 AscendC::WaitPreTaskEnd();
306- for (uint32_t loop = 0, count = 0; loop < groupListShapeSize; loop += groupListInnerShape) {338+ for (uint32_t loop = 0, listIndex = 0, count = 0;
339+ loop < groupListShapeSize; loop += groupListInnerShape, listIndex++) {
307 int32_t splitValue = static_cast<int32_t>(this->groupListGm.GetValue(loop + 1));340 int32_t splitValue = static_cast<int32_t>(this->groupListGm.GetValue(loop + 1));
308 if (splitValue <= 0) {341 if (splitValue <= 0) {
309 break;342 break;
310 }343 }
311 344 
312 uint32_t groupIdx = static_cast<int32_t>(this->groupListGm.GetValue(loop));345 uint32_t groupIdx = static_cast<int32_t>(this->groupListGm.GetValue(loop));
313- mnConfig.mAxisBaseOffset += mnConfig.m;346+ bool skip = this->UpdateMnConfigForGroupListMSparse(mnConfig, splitValue, groupIdx);
314- mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;347+ if (skip) {
315- mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
316- this->SetMNConfig(splitValue, groupIdx, mnConfig);
317- if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
318 continue;348 continue;
319 }349 }
320- mnConfig.nAxisBaseOffset = groupIdx * mnConfig.n;
321- if constexpr (GMMProcess<ComputeType>::B::format == CubeFormat::NZ) {
322- // 16: nz format last two dim size
323- mnConfig.wBaseOffset = AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.nAxisBaseOffset);
324- } else {
325- mnConfig.wBaseOffset = mnConfig.k * mnConfig.nAxisBaseOffset;
326- }
327 350 
328 mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);351 mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
329 mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);352 mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
@@ -334,7 +357,7 @@ public:
334 357 
335 while (curBlock < curCount) {358 while (curBlock < curCount) {
336 MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);359 MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
337- this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx);360+ this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx, listIndex);
338 this->computeOp.VectorCompute(mnConfig);361 this->computeOp.VectorCompute(mnConfig);
339 curBlock += this->gmmBaseParams->coreNum;362 curBlock += this->gmmBaseParams->coreNum;
340 }363 }
@@ -346,7 +369,6 @@ public:
346 }369 }
347};370};
348 371 
349- 
350/** @brief intenal computation class372/** @brief intenal computation class
351*/373*/
352template <class mmType, bool sync = false>374template <class mmType, bool sync = false>
@@ -371,7 +393,7 @@ class GMMCompute {
371 const GMMBaseParams* __restrict gmmBaseParams,393 const GMMBaseParams* __restrict gmmBaseParams,
372 const TCubeTiling* __restrict mmTilingData, TPipe* tPipe);394 const TCubeTiling* __restrict mmTilingData, TPipe* tPipe);
373 395 
374- __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx);396+ __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx, uint32_t listIndex = 0);
375 397 
376 __aicore__ inline void VectorCompute(MNConfig& mnConfig) {}398 __aicore__ inline void VectorCompute(MNConfig& mnConfig) {}
377 399 
@@ -410,6 +432,8 @@ class GMMCompute {
410 uint32_t subBlockIdx;432 uint32_t subBlockIdx;
411 bool mmWaitStatus;433 bool mmWaitStatus;
412 uint32_t activeType;434 uint32_t activeType;
435+ uint32_t groupListType;
436+ uint32_t groupType;
413};437};
414 438 
415template <typename mmType, bool sync>439template <typename mmType, bool sync>
@@ -437,6 +461,8 @@ __aicore__ inline void GMMCompute<mmType, sync>::Init(GM_ADDR x, GM_ADDR weight,
437 }461 }
438 activeType = gmmBaseParams->activeType;462 activeType = gmmBaseParams->activeType;
439 mmWaitStatus = false;463 mmWaitStatus = false;
464+ groupListType = gmmBaseParams->groupListType;
465+ groupType = gmmBaseParams->groupType;
440#if defined(GMM_QUANT_INT8)466#if defined(GMM_QUANT_INT8)
441 scaleTensorPtr = scale;467 scaleTensorPtr = scale;
442#endif468#endif
@@ -515,7 +541,8 @@ __aicore__ inline GlobalTensor<typename mmType::BT::T> GMMCompute<mmType, sync>:
515}541}
516 542 
517template <typename mmType, bool sync>543template <typename mmType, bool sync>
518-__aicore__ inline void GMMCompute<mmType, sync>::MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx) {544+__aicore__ inline void GMMCompute<mmType, sync>::MMCompute(
545+ uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx, uint32_t listIndex) {
519 if (subBlockIdx != 0) {546 if (subBlockIdx != 0) {
520 return;547 return;
521 }548 }
@@ -530,7 +557,11 @@ __aicore__ inline void GMMCompute<mmType, sync>::MMCompute(uint32_t groupIdx, MN
530 uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;557 uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
531 // init global buffer558 // init global buffer
532 if (singleX == 0) {559 if (singleX == 0) {
533- xGm.SetGlobalBuffer(GetTensorAddr<AT>(groupIdx, xTensorPtr));560+ if (groupListType == GROUP_LIST_TYPE_SPARSE && groupType == 0) { // 0: split M
561+ xGm.SetGlobalBuffer(GetTensorAddr<AT>(listIndex, xTensorPtr));
562+ } else {
563+ xGm.SetGlobalBuffer(GetTensorAddr<AT>(groupIdx, xTensorPtr));
564+ }
534 } else {565 } else {
535 xGm.SetGlobalBuffer(GetTensorAddr<AT>(0, xTensorPtr) + mnConfig.xBaseOffset);566 xGm.SetGlobalBuffer(GetTensorAddr<AT>(0, xTensorPtr) + mnConfig.xBaseOffset);
536 }567 }
@@ -549,7 +580,11 @@ __aicore__ inline void GMMCompute<mmType, sync>::MMCompute(uint32_t groupIdx, MN
549#endif580#endif
550 SetGlobalBufferBias(groupIdx, tailN, mnConfig);581 SetGlobalBufferBias(groupIdx, tailN, mnConfig);
551 if (singleY == 0) {582 if (singleY == 0) {
552- yGm.SetGlobalBuffer(GetTensorAddr<CT>(groupIdx, yTensorPtr));583+ if (groupListType == GROUP_LIST_TYPE_SPARSE && groupType == 0) {
584+ yGm.SetGlobalBuffer(GetTensorAddr<CT>(listIndex, yTensorPtr));
585+ } else {
586+ yGm.SetGlobalBuffer(GetTensorAddr<CT>(groupIdx, yTensorPtr));
587+ }
553 } else {588 } else {
554 yGm.SetGlobalBuffer(GetTensorAddr<CT>(0, yTensorPtr) + mnConfig.yBaseOffset);589 yGm.SetGlobalBuffer(GetTensorAddr<CT>(0, yTensorPtr) + mnConfig.yBaseOffset);
555 }590 }
Mgmm/grouped_matmul/op_kernel/grouped_matmul_a4w4.h+120-73
@@ -56,22 +56,24 @@ __aicore__ inline void DataCopyPad2DA4W4(const LocalTensor<T> dst, const GlobalT
56 return;56 return;
57}57}
58 58
59- template <class mmType>59+template <class mmType>
60- class GMMA4W4Compute {60+class GMMA4W4Compute {
61- public:61+public:
62- using aT = MatmulType<TPosition::GM, CubeFormat::ND, DTYPE_X_A4W4>;62+ using aT = MatmulType<TPosition::GM, CubeFormat::ND, DTYPE_X_A4W4>;
63- using bT = typename mmType::BT;63+ using bT = typename mmType::BT;
64- using biasT = MatmulType<TPosition::GM, CubeFormat::ND, int32_t>;64+ using biasT = MatmulType<TPosition::GM, CubeFormat::ND, int32_t>;
65- using cT = MatmulType<TPosition::GM, CubeFormat::ND, half>;65+ using cT = MatmulType<TPosition::GM, CubeFormat::ND, half>;
66- using DTYPE_OUT = DTYPE_Y_A4W4;66+ using DTYPE_OUT = DTYPE_Y_A4W4;
67 67
68- public:68+public:
69 __aicore__ inline GMMA4W4Compute(typename mmType::MT &matmul) : mm(matmul) {}69 __aicore__ inline GMMA4W4Compute(typename mmType::MT &matmul) : mm(matmul) {}
70 __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList,70 __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList,
71 GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams,71 GM_ADDR perTokenScale, GM_ADDR y, GM_ADDR workspace, const GMMBaseParams* __restrict gmmBaseParams,
72 const TCubeTiling* __restrict mmTilingData, TPipe* tPipe);72 const TCubeTiling* __restrict mmTilingData, TPipe* tPipe);
73 __aicore__ inline void Process();73 __aicore__ inline void Process();
74- private:74+protected:
75+ __aicore__ inline void ProcessCommon(
76+ MNConfig &mnConfig, uint32_t splitValue, uint32_t groupIdx, uint32_t &preCount);
75 __aicore__ inline void InitUbBuffer();77 __aicore__ inline void InitUbBuffer();
76 __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig);78 __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig);
77 __aicore__ inline void MMComputePerGroup(uint32_t groupIdx, MNConfig& mnConfig, uint32_t curSingleM, uint32_t curSingleN,79 __aicore__ inline void MMComputePerGroup(uint32_t groupIdx, MNConfig& mnConfig, uint32_t curSingleM, uint32_t curSingleN,
@@ -86,36 +88,36 @@ __aicore__ inline void DataCopyPad2DA4W4(const LocalTensor<T> dst, const GlobalT
86 __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN,88 __aicore__ inline void DataCopyPerTokenScaleAndBrcb(MNConfig& mnConfig, uint32_t curBaseM, uint32_t alignBaseN,
87 uint32_t offsetM);89 uint32_t offsetM);
88 90 
89- private:91+protected:
90- typename mmType::MT& mm;92+ typename mmType::MT& mm;
91- const uint32_t HALF_ALIGN = 16;93+ const uint32_t HALF_ALIGN = 16;
92- GlobalTensor<DTYPE_X_A4W4> xGm;94+ GlobalTensor<DTYPE_X_A4W4> xGm;
93- GlobalTensor<DTYPE_WEIGHT_A4W4> weightGm;95+ GlobalTensor<DTYPE_WEIGHT_A4W4> weightGm;
94- GlobalTensor<cT::T> mmOutGm;96+ GlobalTensor<cT::T> mmOutGm;
95- GlobalTensor<DTYPE_SCALE_A4W4> scaleGm;97+ GlobalTensor<DTYPE_SCALE_A4W4> scaleGm;
96- GlobalTensor<DTYPE_PERTOKEN_SCALE_A4W4> perTokenScaleGm;98+ GlobalTensor<DTYPE_PERTOKEN_SCALE_A4W4> perTokenScaleGm;
97- GlobalTensor<int64_t> groupListGm;99+ GlobalTensor<int64_t> groupListGm;
98- GlobalTensor<DTYPE_OUT> yGm;100+ GlobalTensor<DTYPE_OUT> yGm;
99- // define the que101+ // define the que
100- TQue<QuePosition::VECIN, 1> vecInQueue;102+ TQue<QuePosition::VECIN, 1> vecInQueue;
101- TQue<QuePosition::VECOUT, 1> vecOutQueue;103+ TQue<QuePosition::VECOUT, 1> vecOutQueue;
102- TQue<QuePosition::VECIN, 1> scaleInQueue;104+ TQue<QuePosition::VECIN, 1> scaleInQueue;
103- TQue<QuePosition::VECIN, 1> perTokenScaleInQueue;105+ TQue<QuePosition::VECIN, 1> perTokenScaleInQueue;
104- TBuf<TPosition::VECCALC> tmpBuff;106+ TBuf<TPosition::VECCALC> tmpBuff;
105- LocalTensor<float> mmOutFp32Buf;107+ LocalTensor<float> mmOutFp32Buf;
106- LocalTensor<float> pertokenBrcbLocal;108+ LocalTensor<float> pertokenBrcbLocal;
107- LocalTensor<float> perTokenResBuf;109+ LocalTensor<float> perTokenResBuf;
108- LocalTensor<uint8_t> calcTmpBuf;110+ LocalTensor<uint8_t> calcTmpBuf;
109- uint32_t subBlockIdx;111+ uint32_t subBlockIdx;
110- uint32_t coreIdx;112+ uint32_t coreIdx;
111- uint32_t quantGroupSize_;113+ uint32_t quantGroupSize_;
112- uint32_t cubeCount = 0;114+ uint32_t cubeCount = 0;
113- uint32_t vecCount_ = 0;115+ uint32_t vecCount_ = 0;
114- uint32_t mmBaseBlockOffset_ = 0;116+ uint32_t mmBaseBlockOffset_ = 0;
115- TPipe *pipe;117+ TPipe *pipe;
116- const GMMBaseParams *tiling;118+ const GMMBaseParams *tiling;
117- const TCubeTiling* mmTilingData;119+ const TCubeTiling* mmTilingData;
118- };120+};
119 121
120 template <typename mmType>122 template <typename mmType>
121 __aicore__ inline void GMMA4W4Compute<mmType>::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList,123 __aicore__ inline void GMMA4W4Compute<mmType>::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR scale, GM_ADDR groupList,
@@ -173,42 +175,49 @@ __aicore__ inline void DataCopyPad2DA4W4(const LocalTensor<T> dst, const GlobalT
173 template <typename mmType>175 template <typename mmType>
174 __aicore__ inline void GMMA4W4Compute<mmType>::Process()176 __aicore__ inline void GMMA4W4Compute<mmType>::Process()
175 {177 {
176- MNConfig mnConfig;178+ MNConfig mnConfig;
177- mnConfig.baseM = mmTilingData->baseM;179+ mnConfig.baseM = mmTilingData->baseM;
178- mnConfig.baseN = mmTilingData->baseN;180+ mnConfig.baseN = mmTilingData->baseN;
179- mnConfig.singleM = mnConfig.baseM;181+ mnConfig.singleM = mnConfig.baseM;
180- //仅Perchannel模式使能singleN动态调整，其他情况singleN等于baseN182+ // 仅Perchannel模式使能singleN动态调整，其他情况singleN等于baseN
181- if(tiling->quantGroupNum == 1 && tiling->singleN != 0){183+ if (tiling->quantGroupNum == 1 && tiling->singleN != 0) {
182 mnConfig.singleN = tiling->singleN;184 mnConfig.singleN = tiling->singleN;
183- } else {185+ } else {
184 mnConfig.singleN = mnConfig.baseN;186 mnConfig.singleN = mnConfig.baseN;
185- }187+ }
186- mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN);188+ mnConfig.blockDimN = Ceil(tiling->n, mnConfig.singleN);
187- int32_t preOffset = 0;189+ int32_t preOffset = 0;
188- for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) {190+ for (uint32_t groupIdx = 0, preCount = 0; groupIdx < tiling->groupNum; ++groupIdx) {
189- int32_t m = GetSplitValueFromGroupList(groupIdx, preOffset, tiling, groupListGm);191+ int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, tiling, groupListGm);
190- if (m <= 0) {192+ if (splitValue <= 0) {
191- continue;193+ continue;
192- }194+ }
193- mnConfig.m = static_cast<uint32_t>(m);195+ ProcessCommon(mnConfig, splitValue, groupIdx, preCount);
194- mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);196+ }
195- mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k);197+}
196- uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM;
197- uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + tiling->coreNum;
198- uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN;
199 198 
200- while (curBlock < curCount) {199+template <typename mmType>
201- MNBlockIdxCompute(mnConfig, curBlock, preCount, thresholdM_dimN);200+__aicore__ inline void GMMA4W4Compute<mmType>::ProcessCommon(
202- MMCompute(groupIdx, mnConfig);201+ MNConfig &mnConfig, uint32_t splitValue, uint32_t groupIdx, uint32_t &preCount)
203- if ASCEND_IS_AIV {202+{
204- VectorCompute(groupIdx, mnConfig);203+ uint32_t coreNum = tiling->coreNum;
205- }204+ mnConfig.m = static_cast<uint32_t>(splitValue);
206- curBlock += tiling->coreNum;205+ mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
207- }206+ mm.SetOrgShape(mnConfig.m, tiling->n, tiling->k);
208- preCount = curCount % tiling->coreNum;207+ uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM;
209- mnConfig.offsetM += mnConfig.m;208+ uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + coreNum;
210- }209+ uint32_t thresholdM_dimN = thresholdBlockNum * mnConfig.blockDimN;
211- }210+ while (curBlock < curCount) {
211+ MNBlockIdxCompute(mnConfig, curBlock, preCount, thresholdM_dimN);
212+ MMCompute(groupIdx, mnConfig);
213+ if ASCEND_IS_AIV {
214+ VectorCompute(groupIdx, mnConfig);
215+ }
216+ curBlock += coreNum;
217+ }
218+ preCount = curCount % coreNum;
219+ mnConfig.offsetM += mnConfig.m;
220+}
212 221 
213 template <typename mmType>222 template <typename mmType>
214 __aicore__ inline void GMMA4W4Compute<mmType>::MMComputePerChannel(uint32_t groupIdx, MNConfig& mnConfig, uint32_t curSingleM, uint32_t curSingleN,223 __aicore__ inline void GMMA4W4Compute<mmType>::MMComputePerChannel(uint32_t groupIdx, MNConfig& mnConfig, uint32_t curSingleM, uint32_t curSingleN,
@@ -290,6 +299,7 @@ __aicore__ inline void GMMA4W4Compute<mmType>::MMCompute(uint32_t groupIdx, MNCo
290 } else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) {299 } else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) {
291 weightOffset = groupIdx * tiling->n * tiling->k + tailN * tiling->k; 300 weightOffset = groupIdx * tiling->n * tiling->k + tailN * tiling->k;
292 } else {301 } else {
302+ // B trans and nd do not support groupListType = 2, refuse in tiling process.
293 weightOffset = groupIdx * tiling->n * tiling->k + tailN; 303 weightOffset = groupIdx * tiling->n * tiling->k + tailN;
294 }304 }
295 if(tiling->quantGroupNum == 1) {305 if(tiling->quantGroupNum == 1) {
@@ -400,6 +410,43 @@ __aicore__ inline void GMMA4W4Compute<mmType>::DataCopyPerTokenScaleAndBrcb(MNCo
400 perTokenScaleInQueue.FreeTensor(perTokenScaleLocal);410 perTokenScaleInQueue.FreeTensor(perTokenScaleLocal);
401}411}
402 412 
413+// GroupMatmul A4W4 Sparse operator Class (for groupListType=2)
414+template <class mmType>
415+class GMMA4W4SparseCompute : public GMMA4W4Compute<mmType> {
416+public:
417+ __aicore__ inline GMMA4W4SparseCompute(typename mmType::MT &matmul) : GMMA4W4Compute<mmType>(matmul) {}
418+ __aicore__ inline void Process();
419+};
420+ 
421+template <typename mmType>
422+__aicore__ inline void GMMA4W4SparseCompute<mmType>::Process()
423+{
424+ MNConfig mnConfig;
425+ mnConfig.baseM = this->mmTilingData->baseM;
426+ mnConfig.baseN = this->mmTilingData->baseN;
427+ mnConfig.singleM = mnConfig.baseM;
428+ // 仅Perchannel模式使能singleN动态调整，其他情况singleN等于baseN
429+ if (this->tiling->quantGroupNum == 1 && this->tiling->singleN != 0) {
430+ mnConfig.singleN = this->tiling->singleN;
431+ } else {
432+ mnConfig.singleN = mnConfig.baseN;
433+ }
434+ mnConfig.blockDimN = Ceil(this->tiling->n, mnConfig.singleN);
435+ int32_t preOffset = 0;
436+ 
437+ uint32_t groupListSplitValueOffset = 1;
438+ uint32_t groupListInnerShape = 2u; // groupList shape: [e, 2]
439+ uint32_t groupListShapeSize = this->tiling->groupNum * groupListInnerShape;
440+ for (uint32_t loop = 0, preCount = 0; loop < groupListShapeSize; loop += groupListInnerShape) {
441+ int32_t splitValue = static_cast<int32_t>(this->groupListGm.GetValue(loop + groupListSplitValueOffset));
442+ if (splitValue <= 0) {
443+ break;
444+ }
445+ uint32_t groupIdx = static_cast<uint32_t>(this->groupListGm.GetValue(loop));
446+ this->ProcessCommon(mnConfig, splitValue, groupIdx, preCount);
447+ }
448+}
449+ 
403 } // namespace GROUPED_MATMUL450 } // namespace GROUPED_MATMUL
404 #endif451 #endif
405 #endif452 #endif
Mgmm/grouped_matmul/op_kernel/grouped_matmul_antiquant.h+96-43
@@ -45,21 +45,22 @@ struct CastWeightConfig {
45*/45*/
46template <typename ComputeType>46template <typename ComputeType>
47class GMMAntiquantProcess : public GMMProcess<ComputeType>{47class GMMAntiquantProcess : public GMMProcess<ComputeType>{
48- protected:48+public:
49- constexpr static bool antiquantPerformance = ComputeType::antiquantPerformanceFlag;
50- public:
51 /** @brief constructor */49 /** @brief constructor */
52 __aicore__ inline GMMAntiquantProcess(ComputeType& computeOp_) : GMMProcess<ComputeType>(computeOp_) {}50 __aicore__ inline GMMAntiquantProcess(ComputeType& computeOp_) : GMMProcess<ComputeType>(computeOp_) {}
53 51 
54 __aicore__ inline void Process();52 __aicore__ inline void Process();
55 53 
56- private:54+protected:
55+ constexpr static bool antiquantPerformance = ComputeType::antiquantPerformanceFlag;
56+ 
57+ __aicore__ inline void ProcessCommon(MNConfig &mnConfig, CastWeightConfig &castConfig, uint32_t groupIdx,
58+ uint32_t &count, uint32_t coreNum, uint32_t listIndex = 0);
57 __aicore__ inline void SetAntiquantMNConfig(const uint64_t singleWorkSpaceSize, const uint32_t curBlock, bool& validCore,59 __aicore__ inline void SetAntiquantMNConfig(const uint64_t singleWorkSpaceSize, const uint32_t curBlock, bool& validCore,
58 CastWeightConfig& castConfig, MNConfig &mnConfig);60 CastWeightConfig& castConfig, MNConfig &mnConfig);
59- 
60 __aicore__ inline void SetAntiquantCastConfig(uint32_t& curCount, MNConfig mnConfig,61 __aicore__ inline void SetAntiquantCastConfig(uint32_t& curCount, MNConfig mnConfig,
61 CastWeightConfig& castConfig);62 CastWeightConfig& castConfig);
62- __aicore__ inline void AntiquantUpdateSingleM(MNConfig& mnConfig, uint32_t& dimM, uint32_t dimN);63+ __aicore__ inline void AntiquantUpdateSingleM(MNConfig& mnConfig, uint32_t& dimM, uint32_t dimN);
63};64};
64 65 
65template <typename ComputeType>66template <typename ComputeType>
@@ -113,9 +114,9 @@ template <typename ComputeType>
113__aicore__ inline void GMMAntiquantProcess<ComputeType>::Process() {114__aicore__ inline void GMMAntiquantProcess<ComputeType>::Process() {
114 MNConfig mnConfig;115 MNConfig mnConfig;
115 CastWeightConfig castConfig;116 CastWeightConfig castConfig;
116- castConfig.coreNum = this->gmmBaseParams->coreNum;117+ uint32_t coreNum = this->gmmBaseParams->coreNum;
117- bool validCore = true;118+ castConfig.coreNum = coreNum;
118- uint64_t singleWorkSpaceSize = this->gmmBaseParams->workspaceSize / 2; // 2: antiQuantNormal use 2 block workspace119+ 
119 if (this->gmmBaseParams->groupType != -1) { // -1: no need to split120 if (this->gmmBaseParams->groupType != -1) { // -1: no need to split
120 this->preOffset = 0;121 this->preOffset = 0;
121 if (unlikely(this->groupListPtr == nullptr)) {this->groupNum = 0;} // not continue Process122 if (unlikely(this->groupListPtr == nullptr)) {this->groupNum = 0;} // not continue Process
@@ -123,44 +124,55 @@ __aicore__ inline void GMMAntiquantProcess<ComputeType>::Process() {
123 for (uint32_t groupIdx = 0, count = 0; groupIdx < this->groupNum; ++groupIdx) {124 for (uint32_t groupIdx = 0, count = 0; groupIdx < this->groupNum; ++groupIdx) {
124 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, this->preOffset, this->gmmBaseParams, this->groupListGm);125 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, this->preOffset, this->gmmBaseParams, this->groupListGm);
125 this->SetMNConfig(splitValue, groupIdx, mnConfig);126 this->SetMNConfig(splitValue, groupIdx, mnConfig);
126- uint32_t dimM = Ceil(mnConfig.m, mnConfig.singleM);127+ ProcessCommon(mnConfig, castConfig, groupIdx, count, coreNum);
127- uint32_t dimN = Ceil(mnConfig.n, mnConfig.singleN);
128- if constexpr (!antiquantPerformance) {
129- AntiquantUpdateSingleM(mnConfig, dimM, dimN);
130- }
131- mnConfig.blockDimM = dimM;
132- mnConfig.blockDimN = dimN;
133- uint32_t curCount = count + dimM * dimN;
134- uint32_t curBlock = this->coreIdx >= count ? this->coreIdx : this->coreIdx + this->gmmBaseParams->coreNum;
135- uint32_t thresholdM_dimN = thresholdBlockNum * dimN;
136- 
137- if constexpr (antiquantPerformance) {
138- SetAntiquantCastConfig(curCount, mnConfig, castConfig);
139- }
140- 
141- while (curBlock < curCount) {
142- if constexpr (antiquantPerformance) { // performance verison, will split dimN
143- SetAntiquantMNConfig(singleWorkSpaceSize, curBlock, validCore, castConfig, mnConfig);
144- } else {
145- mnConfig.workSpaceOffset = mnConfig.wBaseOffset;
146- MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
147- }
148- this->computeOp.PreCompute(groupIdx, this->coreIdx, mnConfig, castConfig);
149- this->computeOp.MMSync();
150- if (validCore) {
151- mnConfig.workSpaceOffset += mnConfig.nIdx * mnConfig.singleN;
152- if constexpr (antiquantPerformance) {
153- mnConfig.nIdx += castConfig.castRoundIdx * castConfig.nUsedCore;
154- }
155- this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx);
156- }
157- curBlock += this->gmmBaseParams->coreNum;
158- }
159 this->UpdateMnConfig(mnConfig);128 this->UpdateMnConfig(mnConfig);
160- count = curCount % this->gmmBaseParams->coreNum;
161 }129 }
162}130}
163 131 
132+// for split core
133+template <typename ComputeType>
134+__aicore__ inline void GMMAntiquantProcess<ComputeType>::ProcessCommon(MNConfig &mnConfig, CastWeightConfig &castConfig,
135+ uint32_t groupIdx, uint32_t &count, uint32_t coreNum, uint32_t listIndex)
136+{
137+ bool validCore = true;
138+ // 2: antiQuantNormal use 2 block workspace
139+ uint64_t singleWorkSpaceSize = this->gmmBaseParams->workspaceSize / 2;
140+ 
141+ uint32_t dimM = Ceil(mnConfig.m, mnConfig.singleM);
142+ uint32_t dimN = Ceil(mnConfig.n, mnConfig.singleN);
143+ if constexpr (!antiquantPerformance) {
144+ AntiquantUpdateSingleM(mnConfig, dimM, dimN);
145+ }
146+ 
147+ mnConfig.blockDimM = dimM;
148+ mnConfig.blockDimN = dimN;
149+ uint32_t curCount = count + dimM * dimN;
150+ uint32_t curBlock = this->coreIdx >= count ? this->coreIdx : this->coreIdx + coreNum;
151+ uint32_t thresholdM_dimN = thresholdBlockNum * dimN;
152+ if constexpr (antiquantPerformance) {
153+ SetAntiquantCastConfig(curCount, mnConfig, castConfig);
154+ }
155+ 
156+ while (curBlock < curCount) {
157+ if constexpr (antiquantPerformance) { // performance verison, will split dimN
158+ SetAntiquantMNConfig(singleWorkSpaceSize, curBlock, validCore, castConfig, mnConfig);
159+ } else {
160+ mnConfig.workSpaceOffset = mnConfig.wBaseOffset;
161+ MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
162+ }
163+ this->computeOp.PreCompute(groupIdx, this->coreIdx, mnConfig, castConfig);
164+ this->computeOp.MMSync();
165+ if (validCore) {
166+ mnConfig.workSpaceOffset += mnConfig.nIdx * mnConfig.singleN;
167+ if constexpr (antiquantPerformance) {
168+ mnConfig.nIdx += castConfig.castRoundIdx * castConfig.nUsedCore;
169+ }
170+ this->computeOp.MMCompute(groupIdx, mnConfig, this->coreIdx, listIndex);
171+ }
172+ curBlock += coreNum;
173+ }
174+ count = curCount % coreNum;
175+}
164 176 
165/** @brief intenal computation class177/** @brief intenal computation class
166*/178*/
@@ -541,6 +553,47 @@ GMMAntiquantCompute<mmType, sync, antiquantPerformance>::DataCopyScaleAndOffset(
541 offsetInUb.SetSize(alignBaseN);553 offsetInUb.SetSize(alignBaseN);
542}554}
543 555 
556+/** @brief GroupMatmul Antiquant Sparse operator Class (for groupListType=2)
557+ */
558+template <typename ComputeType>
559+class GMMAntiquantSparseProcess : public GMMAntiquantProcess<ComputeType> {
560+public:
561+ /** @brief constructor */
562+ __aicore__ inline GMMAntiquantSparseProcess(ComputeType& computeOp_)
563+ : GMMAntiquantProcess<ComputeType>(computeOp_) {};
564+ 
565+ __aicore__ inline void Process();
566+};
567+ 
568+template <typename ComputeType>
569+__aicore__ inline void GMMAntiquantSparseProcess<ComputeType>::Process()
570+{
571+ MNConfig mnConfig;
572+ CastWeightConfig castConfig;
573+ uint32_t coreNum = this->gmmBaseParams->coreNum;
574+ castConfig.coreNum = coreNum;
575+ 
576+ if (this->gmmBaseParams->groupType != -1) {
577+ if (unlikely(this->groupListPtr == nullptr)) { this->groupNum = 0; }
578+ }
579+ 
580+ uint32_t groupListSplitValueOffset = 1;
581+ uint32_t groupListInnerShape = 2u; // groupList shape: [e, 2]
582+ uint32_t groupListShapeSize = this->groupNum * groupListInnerShape;
583+ for (uint32_t loop = 0, listIndex = 0, count = 0;
584+ loop < groupListShapeSize; loop += groupListInnerShape, listIndex++) {
585+ int32_t splitValue = static_cast<int32_t>(this->groupListGm.GetValue(loop + groupListSplitValueOffset));
586+ if (splitValue <= 0) { break; }
587+ uint32_t groupIdx = static_cast<uint32_t>(this->groupListGm.GetValue(loop));
588+ bool skip = this->UpdateMnConfigForGroupListMSparse(mnConfig, splitValue, groupIdx);
589+ if (skip) {
590+ continue;
591+ }
592+ 
593+ this->ProcessCommon(mnConfig, castConfig, groupIdx, count, coreNum, listIndex);
594+ }
595+}
596+ 
544template <class mmType, bool sync = false>597template <class mmType, bool sync = false>
545using GMMAntiquantComputePerformance = GMMAntiquantCompute<mmType, sync, true>;598using GMMAntiquantComputePerformance = GMMAntiquantCompute<mmType, sync, true>;
546 599 
Mgmm/grouped_matmul/op_kernel/grouped_matmul_antiquant_a16w8_msd.h+110-27
@@ -52,7 +52,17 @@ struct PreBaseMNConfig {
52*/52*/
53template <typename ComputeType>53template <typename ComputeType>
54class GMMA16W8MSDProcess{54class GMMA16W8MSDProcess{
55- protected:55+public:
56+ /** @brief constructor */
57+ __aicore__ inline GMMA16W8MSDProcess(ComputeType& computeOp_) : computeOp(computeOp_) {}
58+ 
59+ __aicore__ inline void Init(const GMMBaseParams* __restrict gmmBaseParamsIn,
60+ const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn,
61+ GM_ADDR groupList, GM_ADDR tiling);
62+ 
63+ __aicore__ inline void Process();
64+ 
65+protected:
56 using B = typename ComputeType::B;66 using B = typename ComputeType::B;
57 ComputeType& computeOp; // internal computation operator67 ComputeType& computeOp; // internal computation operator
58 const GMMBaseParams* __restrict gmmBaseParams;68 const GMMBaseParams* __restrict gmmBaseParams;
@@ -70,17 +80,6 @@ class GMMA16W8MSDProcess{
70 TILING_TYPE* kListGm;80 TILING_TYPE* kListGm;
71 TILING_TYPE* nListGm;81 TILING_TYPE* nListGm;
72 82 
73- public:
74- /** @brief constructor */
75- __aicore__ inline GMMA16W8MSDProcess(ComputeType& computeOp_) : computeOp(computeOp_) {}
76- 
77- __aicore__ inline void Init(const GMMBaseParams* __restrict gmmBaseParamsIn,
78- const TCubeTiling* __restrict mmTilingDataIn, TILING_TYPE* gmmArrayAddrIn,
79- GM_ADDR groupList, GM_ADDR tiling);
80- 
81- __aicore__ inline void Process();
82- 
83- private:
84 __aicore__ inline void PreProcess(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig,83 __aicore__ inline void PreProcess(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig,
85 uint32_t &preGroupIdx, uint32_t &preCoreCount, bool &isPreRequired);84 uint32_t &preGroupIdx, uint32_t &preCoreCount, bool &isPreRequired);
86 85 
@@ -89,6 +88,8 @@ class GMMA16W8MSDProcess{
89 __aicore__ inline void SetMNConfigs(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig);88 __aicore__ inline void SetMNConfigs(PreBaseMNConfig &preBaseMNConfig, MNConfig &mnConfig);
90 89 
91 __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);90 __aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
91+ __aicore__ inline void ProcessCommon(
92+ MNConfig &mnConfig, uint32_t &count, uint32_t &curCount, uint32_t &curBlock, uint32_t &secondHalfIterCount);
92};93};
93 94 
94template <typename ComputeType>95template <typename ComputeType>
@@ -153,10 +154,14 @@ __aicore__ inline void GMMA16W8MSDProcess<ComputeType>::PreProcess(
153 bool &isPreRequired) {154 bool &isPreRequired) {
154 PreBaseMNConfig preBaseMNConfigs[A16W8_MSD_PREPROCESS_MAX_GROUP];155 PreBaseMNConfig preBaseMNConfigs[A16W8_MSD_PREPROCESS_MAX_GROUP];
155 uint32_t preValidGroupCount = 0;156 uint32_t preValidGroupCount = 0;
156- while (preCoreCount < coreNum && preValidGroupCount < A16W8_MSD_PREPROCESS_MAX_GROUP && preGroupIdx < groupNum) {157+ // 2: groupList shape: [e, 2]; 1: groupList shape: [e]
158+ uint32_t groupListInnerShape = gmmBaseParams->groupListType == GROUP_LIST_TYPE_SPARSE ? 2 : 1;
159+ uint32_t groupListShapeSize = groupNum * groupListInnerShape;
160+ while (preCoreCount < coreNum &&
161+ preValidGroupCount < A16W8_MSD_PREPROCESS_MAX_GROUP && preGroupIdx < groupListShapeSize) {
157 preBaseMNConfig.mAxisBaseOffset += preBaseMNConfig.m;162 preBaseMNConfig.mAxisBaseOffset += preBaseMNConfig.m;
158 preBaseMNConfig.m = GetSplitValueFromGroupList(preGroupIdx, preOffsetPre, gmmBaseParams, groupListGm);163 preBaseMNConfig.m = GetSplitValueFromGroupList(preGroupIdx, preOffsetPre, gmmBaseParams, groupListGm);
159- preGroupIdx++;164+ preGroupIdx += groupListInnerShape;
160 if (preBaseMNConfig.m <= 0) {165 if (preBaseMNConfig.m <= 0) {
161 continue;166 continue;
162 }167 }
@@ -216,23 +221,30 @@ __aicore__ inline void GMMA16W8MSDProcess<ComputeType>::Process() {
216 if (mnConfig.m <= 0) {221 if (mnConfig.m <= 0) {
217 continue;222 continue;
218 }223 }
219- mnConfig.blockDimM = Ceil(A16W8_MSD_STEP * mnConfig.m, mnConfig.singleM);224+ ProcessCommon(mnConfig, count, curCount, curBlock, secondHalfIterCount);
220- mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
221- curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
222- curBlock = coreIdx >= count ? coreIdx : coreIdx + coreNum;
223- while (curBlock < curCount) {
224- mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
225- mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
226- computeOp.MMCompute(mnConfig);
227- computeOp.PostProcess(mnConfig, false, secondHalfIterCount);
228- secondHalfIterCount++;
229- curBlock += coreNum;
230- }
231- count = curCount % coreNum;
232 }225 }
233 TailProcess(mnConfig, secondHalfIterCount);226 TailProcess(mnConfig, secondHalfIterCount);
234}227}
235 228 
229+template <typename ComputeType>
230+__aicore__ inline void GMMA16W8MSDProcess<ComputeType>::ProcessCommon(
231+ MNConfig &mnConfig, uint32_t &count, uint32_t &curCount, uint32_t &curBlock, uint32_t &secondHalfIterCount)
232+{
233+ mnConfig.blockDimM = Ceil(A16W8_MSD_STEP * mnConfig.m, mnConfig.singleM);
234+ mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
235+ curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
236+ curBlock = coreIdx >= count ? coreIdx : coreIdx + coreNum;
237+ while (curBlock < curCount) {
238+ mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
239+ mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
240+ computeOp.MMCompute(mnConfig);
241+ computeOp.PostProcess(mnConfig, false, secondHalfIterCount);
242+ secondHalfIterCount++;
243+ curBlock += coreNum;
244+ }
245+ count = curCount % coreNum;
246+}
247+ 
236/** @brief intenal computation class248/** @brief intenal computation class
237*/249*/
238template <class mmType, bool sync = false>250template <class mmType, bool sync = false>
@@ -944,6 +956,77 @@ __aicore__ inline void GMMA16W8MSDCompute<mmType, sync>::CopyOutFinalResult(
944 vecOutQueue.FreeTensor(output);956 vecOutQueue.FreeTensor(output);
945}957}
946 958 
959+ 
960+template <typename ComputeType>
961+class GMMA16W8MSDMSparseProcess : public GMMA16W8MSDProcess<ComputeType> {
962+public:
963+ /** @brief constructor */
964+ __aicore__ inline GMMA16W8MSDMSparseProcess(ComputeType& computeOp_)
965+ : GMMA16W8MSDProcess<ComputeType>(computeOp_) {}
966+ __aicore__ inline void Process();
967+private:
968+ __aicore__ inline void UpdateMnConfigForGroupListMSparse(
969+ MNConfig &mnConfig, uint32_t splitValue, uint32_t expertIdx, uint32_t groupIdx);
970+};
971+ 
972+template <typename ComputeType>
973+__aicore__ inline void GMMA16W8MSDMSparseProcess<ComputeType>::Process()
974+{
975+ PreBaseMNConfig preBaseMNConfig;
976+ MNConfig mnConfig;
977+ uint32_t preValidGroupCount = 0;
978+ uint32_t preGroupIdx = 0;
979+ bool isPreRequired = false;
980+ uint32_t secondHalfIterCount = 0;
981+ this->SetMNConfigs(preBaseMNConfig, mnConfig);
982+ if (mnConfig.k <= 0 || mnConfig.n <= 0) {
983+ return;
984+ }
985+ 
986+ uint32_t groupListSplitValueOffset = 1;
987+ uint32_t groupListInnerShape = 2u; // groupList shape: [e, 2]
988+ uint32_t groupListShapeSize = this->groupNum * groupListInnerShape;
989+ for (uint32_t loop(0), count(0), curBlock(0), curCount(0), preCoreCount(0);
990+ loop < groupListShapeSize; loop += groupListInnerShape) {
991+ isPreRequired = preGroupIdx == loop; // loop groupNum
992+ if (isPreRequired) {
993+ this->PreProcess(preBaseMNConfig, mnConfig, preGroupIdx, preCoreCount, isPreRequired);
994+ }
995+ 
996+ uint32_t splitValue = static_cast<int32_t>(this->groupListGm.GetValue(loop + groupListSplitValueOffset));
997+ if (splitValue <= 0) { break; }
998+ uint32_t expertIdx = static_cast<uint32_t>(this->groupListGm.GetValue(loop));
999+ UpdateMnConfigForGroupListMSparse(mnConfig, splitValue, expertIdx, loop);
1000+ 
1001+ if ASCEND_IS_AIC {
1002+ if (isPreRequired) {
1003+ CrossCoreWaitFlag(SYNC_AIV_AIC_FLAG);
1004+ }
1005+ }
1006+ this->ProcessCommon(mnConfig, count, curCount, curBlock, secondHalfIterCount);
1007+ }
1008+ this->TailProcess(mnConfig, secondHalfIterCount);
1009+}
1010+ 
1011+template <typename ComputeType>
1012+__aicore__ inline void GMMA16W8MSDMSparseProcess<ComputeType>::UpdateMnConfigForGroupListMSparse(
1013+ MNConfig &mnConfig, uint32_t splitValue, uint32_t expertIdx, uint32_t groupIdx)
1014+{
1015+ if (groupIdx > 0) {
1016+ mnConfig.mAxisBaseOffset += mnConfig.m;
1017+ mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
1018+ mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
1019+ }
1020+ 
1021+ mnConfig.nAxisBaseOffset = expertIdx * mnConfig.n;
1022+ if constexpr (GMMA16W8MSDProcess<ComputeType>::B::format == CubeFormat::NZ) {
1023+ // 16: nz format last two dim size
1024+ mnConfig.wBaseOffset = AlignUp<16>(mnConfig.k) * AlignUp<16>(mnConfig.nAxisBaseOffset);
1025+ } else {
1026+ mnConfig.wBaseOffset = mnConfig.k * mnConfig.nAxisBaseOffset;
1027+ }
1028+ mnConfig.m = splitValue;
1029+}
947} // namespace GROUPED_MATMUL1030} // namespace GROUPED_MATMUL
948 1031 
949#endif1032#endif
Mgmm/grouped_matmul/op_kernel/grouped_matmul_quant_mixcore.h+2-2
@@ -53,7 +53,7 @@ class GMMQuantMixCoreCompute : public GMMCompute<mmType, sync> {
53 __aicore__ inline void InitStaticTiling(const GMMBaseParams* __restrict gmmBaseParams, GM_ADDR workspace,53 __aicore__ inline void InitStaticTiling(const GMMBaseParams* __restrict gmmBaseParams, GM_ADDR workspace,
54 int32_t baseM, int32_t baseN);54 int32_t baseM, int32_t baseN);
55 55 
56- __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx);56+ __aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig& mnConfig, uint32_t coreIdx, uint32_t listIndex = 0);
57 57 
58 __aicore__ inline void VectorCompute(MNConfig& mnConfig);58 __aicore__ inline void VectorCompute(MNConfig& mnConfig);
59 59 
@@ -174,7 +174,7 @@ __aicore__ inline void GMMQuantMixCoreCompute<mmType, sync>::PostCompute() {
174 174 
175template <typename mmType, bool sync>175template <typename mmType, bool sync>
176__aicore__ inline void GMMQuantMixCoreCompute<mmType, sync>::MMCompute(uint32_t groupIdx, MNConfig& mnConfig,176__aicore__ inline void GMMQuantMixCoreCompute<mmType, sync>::MMCompute(uint32_t groupIdx, MNConfig& mnConfig,
177- uint32_t coreIdx) {177+ uint32_t coreIdx, uint32_t listIndex) {
178 uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;178 uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
179 uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;179 uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;
180 uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM180 uint32_t curSingleM = mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM
Mgmm/grouped_matmul/op_kernel/grouped_matmul_tiling_key.h+20-11
@@ -68,7 +68,8 @@
68 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \68 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \
69 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \69 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \
70 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \70 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \
71- ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\71+ ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, \
72+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), \
72 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \73 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \
73 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \74 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \
74 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \75 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \
@@ -82,10 +83,12 @@
82 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \83 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \
83 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \84 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \
84 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), \85 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1), \
85- ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\86+ ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, \
87+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), \
86 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \88 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \
87 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \89 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \
88- ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), \90+ ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_MSD, \
91+ GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT), \
89 ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), \92 ASCENDC_TPL_UINT_SEL(AIV_AIC_RATIO, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_AIV_AIC_RATIO_1), \
90 ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0)93 ASCENDC_TPL_BOOL_SEL(IS_ENABLE_FIXED_AXIS, 0)
91 94 
@@ -96,7 +99,8 @@
96 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \99 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \
97 ASCENDC_TPL_BOOL_SEL(TRANS_A, trans_a), \100 ASCENDC_TPL_BOOL_SEL(TRANS_A, trans_a), \
98 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \101 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \
99- ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\102+ ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, \
103+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), \
100 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \104 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \
101 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \105 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \
102 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \106 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE), \
@@ -110,7 +114,8 @@
110 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \114 ASCENDC_TPL_DTYPE_SEL(D_T_Y, dtype), \
111 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \115 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0), \
112 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \116 ASCENDC_TPL_BOOL_SEL(TRANS_B, trans_b), \
113- ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),\117+ ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM, \
118+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT, GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM), \
114 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \119 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0), \
115 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \120 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE), \
116 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W4_KERNEL_TEMPLATE_MSD_ANTIQUANT_GS32), \121 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W4_KERNEL_TEMPLATE_MSD_ANTIQUANT_GS32), \
@@ -206,8 +211,9 @@ ASCENDC_TPL_SEL(
206 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0),211 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0),
207 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0),212 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0),
208 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,213 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,
209- GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,214+ GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,
210- GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),215+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT,
216+ GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM),
211 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),217 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),
212 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),218 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),
213 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT),219 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT),
@@ -222,8 +228,9 @@ ASCENDC_TPL_SEL(
222 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0),228 ASCENDC_TPL_BOOL_SEL(TRANS_A, 0),
223 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0),229 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0),
224 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,230 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,
225- GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,231+ GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,
226- GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),232+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT,
233+ GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM),
227 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),234 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),
228 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),235 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),
229 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT),236 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_ANTIQUANT),
@@ -293,7 +300,8 @@ ASCENDC_TPL_SEL(
293 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1),300 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1),
294 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,301 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,
295 GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,302 GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,
296- GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),303+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT,
304+ GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM),
297 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),305 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),
298 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),306 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),
299 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE),307 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE),
@@ -478,7 +486,8 @@ ASCENDC_TPL_SEL(
478 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1),486 ASCENDC_TPL_BOOL_SEL(TRANS_B, 0, 1),
479 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,487 ASCENDC_TPL_UINT_SEL(GROUP_LIST_TYPE, ASCENDC_TPL_UI_LIST,
480 GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,488 GROUPED_MATMUL_GROUP_LIST_TYPE_CUMSUM,
481- GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT),489+ GROUPED_MATMUL_GROUP_LIST_TYPE_COUNT,
490+ GROUPED_MATMUL_GROUP_LIST_TYPE_SPARSEM),
482 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),491 ASCENDC_TPL_BOOL_SEL(IS_STATIC_TILING_API, 0),
483 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),492 ASCENDC_TPL_UINT_SEL(A8W4_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A8W4_KERNEL_TEMPLATE_NONE),
484 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE),493 ASCENDC_TPL_UINT_SEL(A16W8_KERNEL_TEMPLATE, ASCENDC_TPL_UI_LIST, GROUPED_MATMUL_A16W8_KERNEL_TEMPLATE_NONE),
Mgmm/grouped_matmul/op_kernel/grouped_matmul_utils.h+4-0
@@ -179,6 +179,8 @@ constexpr int32_t STATIC_TILING_STEP_KA_KB = 4;
179constexpr uint64_t DOUBLE_BUFFER_L0A_L0B = 2;179constexpr uint64_t DOUBLE_BUFFER_L0A_L0B = 2;
180constexpr uint32_t STATIC_TILING_MAX_K = 8192;180constexpr uint32_t STATIC_TILING_MAX_K = 8192;
181constexpr uint32_t STATIC_TILING_MAX_SINGLE_N = 1024;181constexpr uint32_t STATIC_TILING_MAX_SINGLE_N = 1024;
182+constexpr uint32_t GROUP_LIST_TYPE_SPARSE = 2;
183+constexpr uint32_t SPARSE_GROUP_LIST_SPLIT_VALUE_OFFSET = 1;
182 184 
183template<class AT_, class BT_, class CT_, class BiasT_, const auto& MM_CFG = CFG_MDL>185template<class AT_, class BT_, class CT_, class BiasT_, const auto& MM_CFG = CFG_MDL>
184struct MMType {186struct MMType {
@@ -307,6 +309,8 @@ __aicore__ inline int32_t GetSplitValueFromGroupList(uint32_t groupIdx, int32_t
307 int32_t offset = static_cast<int32_t>(groupListGm.GetValue(groupIdx));309 int32_t offset = static_cast<int32_t>(groupListGm.GetValue(groupIdx));
308 splitValue = offset - preOffset;310 splitValue = offset - preOffset;
309 preOffset = offset;311 preOffset = offset;
312+ } else if (gmmBaseParams->groupListType == GROUP_LIST_TYPE_SPARSE) {
313+ splitValue = static_cast<int32_t>(groupListGm.GetValue(groupIdx + SPARSE_GROUP_LIST_SPLIT_VALUE_OFFSET));
310 } else {314 } else {
311 splitValue = static_cast<int32_t>(groupListGm.GetValue(groupIdx));315 splitValue = static_cast<int32_t>(groupListGm.GetValue(groupIdx));
312 }316 }
Mgmm/grouped_matmul/op_kernel/grouped_matmul_vector.h+8-5
@@ -49,12 +49,15 @@ __aicore__ inline void EmptyTensorCompute(GM_ADDR groupListPtr, GM_ADDR y, const
49 if (coreRation > 1) {49 if (coreRation > 1) {
50 coreIdx /= coreRation;50 coreIdx /= coreRation;
51 }51 }
52- 52+ // 2: groupList shape: [e, 2]; 1: groupList shape: [e]
53- for (uint32_t groupIdx = 0; groupIdx < gmmBaseParams->groupNum; ++groupIdx) {53+ uint32_t groupListInnerShape = gmmBaseParams->groupListType == GROUP_LIST_TYPE_SPARSE ? 2 : 1;
54+ uint32_t groupListShapeSize = gmmBaseParams->groupNum * groupListInnerShape;
55+ for (uint32_t groupIdx = 0; groupIdx < groupListShapeSize; groupIdx += groupListInnerShape) {
54 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm);56 int32_t splitValue = GetSplitValueFromGroupList(groupIdx, preOffset, gmmBaseParams, groupListGm);
55- uint32_t m = isAllSingleTensor && gmmBaseParams->groupType == 2 ? *ubM : *(ubM + groupIdx);57+ uint32_t m = isAllSingleTensor && gmmBaseParams->groupType == 2 ?
56- uint32_t k = *ubK < 0 && gmmBaseParams->groupType == 2 ? splitValue : *(ubK + groupIdx);58+ *ubM : *(ubM + groupIdx / groupListInnerShape);
57- uint32_t n = isAllSingleTensor ? *ubN : *(ubN + groupIdx);59+ uint32_t k = *ubK < 0 && gmmBaseParams->groupType == 2 ? splitValue : *(ubK + groupIdx / groupListInnerShape);
60+ uint32_t n = isAllSingleTensor ? *ubN : *(ubN + groupIdx / groupListInnerShape);
58 61 
59 if (k == 0) {62 if (k == 0) {
60 uint32_t singleM = Ceil(m, gmmBaseParams->coreNum);63 uint32_t singleM = Ceil(m, gmmBaseParams->coreNum);
Mgmm/grouped_matmul/tests/ut/op_host/test_grouped_matmul_tiling.csv+17-0
@@ -148,3 +148,20 @@ Ascend910B,test_tiling_a8w4_k_too_large_invalid,true,test_tiling_a8w4_k_too_larg
148Ascend910B,test_tiling_a8w4_offset_invalid_shape,true,test_tiling_a8w4_offset_invalid_shape,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,0,0,NONE,512:1028,512:1028,INT8,ND,2:1028:1024,2:1028:1024,INT4,ND,NONE,NONE,FLOAT,ND,2:4:1024,2:4:1024,UINT64,ND,2:2:1024,2:2:1024,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,2,2,INT64,ND,512,512,FLOAT,ND,512:1024,512:1024,FLOAT16,ND148Ascend910B,test_tiling_a8w4_offset_invalid_shape,true,test_tiling_a8w4_offset_invalid_shape,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,0,0,NONE,512:1028,512:1028,INT8,ND,2:1028:1024,2:1028:1024,INT4,ND,NONE,NONE,FLOAT,ND,2:4:1024,2:4:1024,UINT64,ND,2:2:1024,2:2:1024,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,2,2,INT64,ND,512,512,FLOAT,ND,512:1024,512:1024,FLOAT16,ND
149Ascend950,test_tiling_no_quant_sparse_grouplist_invalid_950,true,test_tiling_no_quant_sparse_grouplist_invalid_950,false,0,32,64,262144,524288,196608,262144,65536,65536,3,0,false,false,0,2,0,0,256:512,256:512,FLOAT16,ND,2:512:256,2:512:256,FLOAT16,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,2,2,INT64,ND,NONE,NONE,FLOAT,ND,256:256,256:256,FLOAT16,ND149Ascend950,test_tiling_no_quant_sparse_grouplist_invalid_950,true,test_tiling_no_quant_sparse_grouplist_invalid_950,false,0,32,64,262144,524288,196608,262144,65536,65536,3,0,false,false,0,2,0,0,256:512,256:512,FLOAT16,ND,2:512:256,2:512:256,FLOAT16,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,2,2,INT64,ND,NONE,NONE,FLOAT,ND,256:256,256:256,FLOAT16,ND
150Ascend950,test_tiling_no_quant_groupnum_large_invalid_950,true,test_tiling_no_quant_groupnum_large_invalid_950,false,0,32,64,262144,524288,196608,262144,65536,65536,3,0,false,false,0,0,0,0,256:512,256:512,FLOAT16,ND,1205:512:256,1205:512:256,FLOAT16,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,1205,1205,INT64,ND,NONE,NONE,FLOAT,ND,256:256,256:256,FLOAT16,ND150Ascend950,test_tiling_no_quant_groupnum_large_invalid_950,true,test_tiling_no_quant_groupnum_large_invalid_950,false,0,32,64,262144,524288,196608,262144,65536,65536,3,0,false,false,0,0,0,0,256:512,256:512,FLOAT16,ND,1205:512:256,1205:512:256,FLOAT16,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT16,ND,NONE,NONE,FLOAT16,ND,1205,1205,INT64,ND,NONE,NONE,FLOAT,ND,256:256,256:256,FLOAT16,ND
151+Ascend910B,test_tiling_a16w4ofp16_pergroup_sparsem_1aic2aiv,true,test_tiling_a16w4ofp16_pergroup_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,4096:2048,4096:2048,FLOAT16,ND,4:2048:7168,4:2048:7168,INT4,ND,4096:7168,4096:7168,FLOAT16,ND,4:7168,4:7168,FLOAT,ND,NONE,NONE,FLOAT,ND,4:7168,4:7168,FLOAT16,ND,4:7168,4:7168,FLOAT16,ND,4:2,4:2,INT64,ND,4096,4096,FLOAT,ND,4096,7168,FLOAT16,ND
152+Ascend910B,test_tiling_a16w4obf16_pergroup_sparsem_1aic2aiv,true,test_tiling_a16w4obf16_pergroup_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,4096:2048,4096:2048,BF16,ND,4:2048:7168,4:2048:7168,INT4,ND,4096:7168,4096:7168,FLOAT,ND,4:7168,4:7168,FLOAT,ND,NONE,NONE,FLOAT,ND,4:256:7168,4:256:7168,BF16,ND,4:256:7168,4:256:7168,BF16,ND,4:2,4:2,INT64,ND,4096,4096,FLOAT,ND,4096,7168,BF16,ND
153+Ascend910B,test_tiling_a16w4ofp16_msd_transw_sparsem_1aic2aiv,true,test_tiling_a16w4ofp16_msd_transw_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,true,0,2,0,NONE,512:2048,512:2048,FLOAT16,ND,4:7168:2048,4:7168:2048,INT4,ND,512:2048,512:2048,FLOAT16,ND,4:2048,4:2048,FLOAT,ND,NONE,NONE,FLOAT,ND,4:2048,4:2048,FLOAT16,ND,4:2048,4:2048,FLOAT16,ND,4:2,4:2,INT64,ND,512,512,FLOAT,ND,512,2048,FLOAT16,ND
154+Ascend910B,test_tiling_a16w4obf16_msd_transw_sparsem_1aic2aiv,true,test_tiling_a16w4obf16_msd_transw_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,true,0,2,0,NONE,512:2048,512:2048,BF16,ND,4:7168:2048,4:7168:2048,INT4,ND,512:2048,512:2048,FLOAT,ND,4:2048,4:2048,FLOAT,ND,NONE,NONE,FLOAT,ND,4:256:2048,4:256:2048,BF16,ND,4:256:2048,4:256:2048,BF16,ND,4:2,4:2,INT64,ND,512,512,FLOAT,ND,512,2048,BF16,ND
155+Ascend910B,test_tiling_a16w8ofp16_msd_sparsem_1aic1aiv,true,test_tiling_a16w8ofp16_msd_sparsem_1aic1aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,512:2048,512:2048,FLOAT16,ND,8:2048:1024,8:2048:1024,INT8,ND,512:1024,512:1024,FLOAT16,ND,8:1024,8:1024,FLOAT,ND,NONE,NONE,FLOAT,ND,8:1024,8:1024,FLOAT16,ND,8:1024,8:1024,FLOAT16,ND,8:2,8:2,INT64,ND,512,512,FLOAT,ND,512,1024,FLOAT16,ND
156+Ascend910B,test_tiling_a16w8obf16_msd_sparsem_1aic1aiv,true,test_tiling_a16w8obf16_msd_sparsem_1aic1aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,512:2048,512:2048,BF16,ND,8:2048:1024,8:2048:1024,INT8,ND,512:1024,512:1024,FLOAT,ND,8:1024,8:1024,FLOAT,ND,NONE,NONE,FLOAT,ND,8:1024,8:1024,BF16,ND,8:1024,8:1024,BF16,ND,8:2,8:2,INT64,ND,512,512,FLOAT,ND,512,1024,BF16,ND
157+Ascend910B,test_tiling_a16w8ofp16_antiquant_sparsem_1aic2aiv,true,test_tiling_a16w8ofp16_antiquant_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,4096:2048,4096:2048,FLOAT16,ND,4:2048:7168,4:2048:7168,INT8,ND,4096:7168,4096:7168,FLOAT16,ND,4:7168,4:7168,FLOAT,ND,NONE,NONE,FLOAT,ND,4:7168,4:7168,FLOAT16,ND,4:7168,4:7168,FLOAT16,ND,4:2,4:2,INT64,ND,4096,4096,FLOAT,ND,4096,7168,FLOAT16,ND
158+Ascend910B,test_tiling_a16w8obf16_antiquant_sparsem_1aic2aiv,true,test_tiling_a16w8obf16_antiquant_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,4096:2048,4096:2048,BF16,ND,4:2048:7168,4:2048:7168,INT8,ND,4096:7168,4096:7168,FLOAT,ND,4:7168,4:7168,FLOAT,ND,NONE,NONE,FLOAT,ND,4:7168,4:7168,BF16,ND,4:7168,4:7168,BF16,ND,4:2,4:2,INT64,ND,4096,4096,FLOAT,ND,4096,7168,BF16,ND
159+Ascend910B,test_tiling_a4w4ofp16_sparsem_1aic2aiv,true,test_tiling_a4w4ofp16_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,350:1280,350:1280,INT4,ND,2:1280:580,2:1280:580,INT4,ND,NONE,NONE,INT32,ND,2:580,2:580,UINT64,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,350,580,FLOAT16,ND
160+Ascend910B,test_tiling_a4w4obf16_sparsem_1aic2aiv,true,test_tiling_a4w4obf16_sparsem_1aic2aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,350:1280,350:1280,INT4,ND,2:1280:580,2:1280:580,INT4,ND,NONE,NONE,INT32,ND,2:580,2:580,UINT64,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,350,580,BF16,ND
161+Ascend910B,test_tiling_fp16_sparsem_1aic,true,test_tiling_fp16_sparsem_1aic,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,345:1280,345:1280,FLOAT16,ND,2:1280:567,2:1280:567,FLOAT16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,FLOAT16,ND
162+Ascend910B,test_tiling_bf16_sparsem_1aic,true,test_tiling_bf16_sparsem_1aic,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,345:1280,345:1280,BF16,ND,2:1280:567,2:1280:567,BF16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,BF16,ND
163+Ascend910B,test_tiling_fp32_sparsem_1aic,true,test_tiling_fp32_sparsem_1aic,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,false,0,2,0,NONE,345:1280,345:1280,FLOAT,ND,2:1280:567,2:1280:567,FLOAT,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,FLOAT,ND
164+Ascend910B,test_tiling_fp16_transw_sparsem_1aic,true,test_tiling_fp16_transw_sparsem_1aic,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,true,0,2,0,NONE,345:1280,345:1280,FLOAT16,ND,2:1280:567,2:1280:567,FLOAT16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,FLOAT16,ND
165+Ascend910B,test_tiling_bf16_transw_sparsem_1aic,true,test_tiling_bf16_transw_sparsem_1aic,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,false,true,0,2,0,NONE,345:1280,345:1280,BF16,ND,2:1280:567,2:1280:567,BF16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,BF16,ND
166+Ascend910B,test_tiling_fp16_transx_sparsem_1aic1aiv,true,test_tiling_fp16_transx_sparsem_1aic1aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,true,false,0,2,0,NONE,345:1280,345:1280,FLOAT16,ND,2:1280:567,2:1280:567,FLOAT16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,FLOAT16,ND
167+Ascend910B,test_tiling_bf16_transx_sparsem_1aic1aiv,true,test_tiling_bf16_transx_sparsem_1aic1aiv,false,0,24,48,196608,524288,196608,131072,65536,65536,0,0,true,false,0,2,0,NONE,345:1280,345:1280,BF16,ND,2:1280:567,2:1280:567,BF16,ND,345:567,345:567,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,NONE,NONE,FLOAT,ND,2:2,2:2,INT64,ND,NONE,NONE,FLOAT,ND,345,567,BF16,ND