@@ -0,0 +1,197 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.cpp
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_pre.h"
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_mid.h"
|
||||
#include "grouped_matmul_swiglu_quant_a8w4_msd_post.h"
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
|
||||
template <class mmType>
|
||||
class GMMSwigluQuantPipelineSchedule {
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
TPipe *pipe;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
// WorkSpaceSplitConfig控制Workspace切割方式的结构体;
|
||||
WorkSpaceSplitConfig workspaceSplitConfig;
|
||||
WorkSpaceSplitConfig tempWorkspaceSplitConfig;
|
||||
// 记录GM_ADDR的结构体
|
||||
GMAddrParams gmAddrParams;
|
||||
// 前处理GMMA8W4PreProcess类
|
||||
GMMA8W4PreProcess preProcess;
|
||||
// 中间处理GMMA8W4MidProcess类
|
||||
GMMA8W4MidProcess<mmType> midProcess;
|
||||
// 后处理GMMA8W4PostProcess类
|
||||
GMMA8W4PostProcess postProcess;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
__aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx);
|
||||
|
||||
public:
|
||||
__aicore__ inline GMMSwigluQuantPipelineSchedule(typename mmType::MT &mm_,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
: mm(mm_), midProcess(mm), gmmBaseParams(gmmBaseParamsIN), gmmSwiglu(gmmSwigluIN), pipe(tPipeIN)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace);
|
||||
__aicore__ inline void Process();
|
||||
};
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void GMMSwigluQuantPipelineSchedule<mmType>::Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale,
|
||||
GM_ADDR xScale, GM_ADDR weightAssistanceMatrix,
|
||||
GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace)
|
||||
{
|
||||
gmAddrParams.xGM = x;
|
||||
gmAddrParams.weightGM = weight;
|
||||
gmAddrParams.weightScaleGM = weightScale;
|
||||
gmAddrParams.xScaleGM = xScale;
|
||||
gmAddrParams.weightAuxiliaryMatrixGM = weightAssistanceMatrix;
|
||||
gmAddrParams.groupListGM = groupList;
|
||||
gmAddrParams.yGM = y;
|
||||
gmAddrParams.yScaleGM = yScale;
|
||||
gmAddrParams.workSpaceGM = workspace;
|
||||
gmAddrParams.workSpaceOffset1 = gmmBaseParams->workSpaceOffset1 / 2;
|
||||
gmAddrParams.workSpaceOffset2 = gmmBaseParams->workSpaceOffset1;
|
||||
gmAddrParams.workSpaceOffset3 = gmmBaseParams->workSpaceOffset1 + gmmBaseParams->workSpaceOffset2 / 2;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
InitWorkSpaceSplitConfig(workspaceSplitConfig);
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void GMMSwigluQuantPipelineSchedule<mmType>::Process()
|
||||
{
|
||||
// 1.对每次workspace切分做大循环。
|
||||
preProcess.Init(gmAddrParams, gmmBaseParams);
|
||||
midProcess.Init(gmAddrParams, gmmBaseParams);
|
||||
postProcess.Init(gmAddrParams, gmmBaseParams, gmmSwiglu);
|
||||
|
||||
// 1.前处理提前下发一次
|
||||
preProcess.Process(workspaceSplitConfig, 0, pipe);
|
||||
for (int64_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount;
|
||||
workspaceSplitLoopIdx++) {
|
||||
// 更新workspaceSplitConfig
|
||||
UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
|
||||
SyncAll<false>();
|
||||
// 2.第n次中处理 && 第n+1次前处理 && 第n-1次后处理 并行
|
||||
midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
|
||||
preProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx + 1, pipe);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
SyncAll<true>();
|
||||
}
|
||||
postProcess.Process(tempWorkspaceSplitConfig, workspaceSplitLoopIdx - 1, pipe);
|
||||
// 3.第n-1次后处理需要保留第n次的切分数据
|
||||
tempWorkspaceSplitConfig = workspaceSplitConfig;
|
||||
// reset
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
SyncAll<false>();
|
||||
// 3.前一次后处理 && 后一次MM 并行
|
||||
}
|
||||
// reset
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Reset();
|
||||
}
|
||||
SyncAll<false>();
|
||||
// // 4.最后一次后处理
|
||||
postProcess.Process(workspaceSplitConfig, workspaceSplitConfig.loopCount - 1, pipe);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe->Destroy();
|
||||
}
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void
|
||||
GMMSwigluQuantPipelineSchedule<mmType>::InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1);
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0;
|
||||
workspaceSplitConfig.isLastLoop = false;
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void
|
||||
GMMSwigluQuantPipelineSchedule<mmType>::UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx)
|
||||
{
|
||||
if (workspaceSplitLoopIdx < 0)
|
||||
return;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
if (curTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
}
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
}
|
||||
workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false;
|
||||
|
||||
if (workspaceSplitConfig.isLastLoop) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex =
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_PIPELINE_H
|
||||
@@ -0,0 +1,102 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.cpp
|
||||
* \brief
|
||||
*/
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
#include "grouped_matmul_swiglu_pipeline.h"
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_split_ws.h"
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
using namespace GROUPED_MATMUL_SWIGLU_QUANT;
|
||||
|
||||
#define GMM_CV_SPLIT_IMP(computeClass, dtypeWeightScale, transA, transB, sync) \
|
||||
do { \
|
||||
using xType = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, DTYPE_X, false>; \
|
||||
using weightType = MatmulType<AscendC::TPosition::GM, CubeFormat::NZ, DTYPE_WEIGHT, false>; \
|
||||
using yType = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>; \
|
||||
using matmulType = MMImplTypeStatic<xType, weightType, yType>; \
|
||||
matmulType::MT mm; \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \
|
||||
if ASCEND_IS_AIC { \
|
||||
mm.SetSubBlockIdx(0); \
|
||||
mm.Init(&mmTilingData_, &tPipe); \
|
||||
} \
|
||||
computeClass<matmulType, sync, dtypeWeightScale> computeOp(mm); \
|
||||
computeOp.Init(x, weight, weightScale, xScale, groupList, y, yScale, user1, &gmmSwigluBaseParams_, \
|
||||
&mmTilingData_, &gmmSwiglu_, &tPipe); \
|
||||
computeOp.Process(); \
|
||||
} while (0)
|
||||
|
||||
#define GMM_CV_SPLIT_IMP_A8W4_MSD(computeClass, dtypeWeightScale, transA, transB, sync) \
|
||||
do { \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwigluBaseParams, gmmSwigluBaseParams_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, mmTilingData, mmTilingData_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingData, gmmSwiglu, gmmSwiglu_, tiling); \
|
||||
using xType = MatmulType<TPosition::GM, CubeFormat::ND, int4b_t, false>; \
|
||||
using weightType = MatmulType<TPosition::GM, wFormat, int4b_t, false>; \
|
||||
using yType = MatmulType<TPosition::GM, CubeFormat::ND, half, false>; \
|
||||
using matmulType = MMImplType<xType, weightType, yType>; \
|
||||
matmulType::MT mm; \
|
||||
if ASCEND_IS_AIC { \
|
||||
mm.SetSubBlockIdx(0); \
|
||||
mm.Init(&mmTilingData_); \
|
||||
} \
|
||||
computeClass<matmulType> op(mm, &gmmSwigluBaseParams_, &gmmSwiglu_, &tPipe); \
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, user1); \
|
||||
\
|
||||
op.Process(); \
|
||||
} while (0)
|
||||
|
||||
extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale,
|
||||
GM_ADDR xScale, GM_ADDR weightAssistanceMatrix,
|
||||
GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
TPipe tPipe;
|
||||
AscendCUtils::SetOverflow(1);
|
||||
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GM_ADDR user1 = GetUserWorkspace(workspace);
|
||||
#if defined(GMM_SWIGLU_QUANT_A8W8)
|
||||
if (TILING_KEY_IS(0)) { // antiquant msd
|
||||
KERNEL_TASK_TYPE(0, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP(GMMSwigluCompute, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
} else if (TILING_KEY_IS(1)) {
|
||||
KERNEL_TASK_TYPE(1, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP(GMMSwigluSplitWorkSpaceCompute, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
}
|
||||
#elif defined(GMM_SWIGLU_QUANT_A8W4_MSD)
|
||||
if (TILING_KEY_IS(2)) {
|
||||
KERNEL_TASK_TYPE(2, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GMM_CV_SPLIT_IMP_A8W4_MSD(GMMSwigluQuantPipelineSchedule, // computeClass
|
||||
DTYPE_WEIGHT_SCALE,
|
||||
false, // transA
|
||||
false, // transB
|
||||
false // sync
|
||||
);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,547 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
/** @brief internal computation class
|
||||
*/
|
||||
template <class mmType, bool sync = false, typename CHANNELDTYPE = float>
|
||||
class GMMSwigluCompute {
|
||||
public:
|
||||
using AT = typename mmType::AT::T;
|
||||
using BT = typename mmType::BT::T;
|
||||
using B = typename mmType::BT;
|
||||
using CT = typename mmType::CT::T;
|
||||
using BiasT = typename mmType::BiasT::T;
|
||||
using WT = int8_t;
|
||||
constexpr static bool transposeX = mmType::AT::isTrans;
|
||||
constexpr static bool transposeW = mmType::BT::isTrans;
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
/** @brief constructor */
|
||||
__aicore__ inline GMMSwigluCompute(typename mmType::MT &mm_) : mm(mm_)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale,
|
||||
GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN,
|
||||
TPipe *tPipeIN);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx);
|
||||
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k);
|
||||
|
||||
__aicore__ inline void CubeProcess(MNConfig &mnConfig);
|
||||
__aicore__ inline void VecProcess(VecConfig &vecConfig);
|
||||
__aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count,
|
||||
const uint32_t thresholdM_dimN);
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig);
|
||||
__aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig);
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig);
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig);
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, VecConfig &vecConfig);
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig);
|
||||
__aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
__aicore__ inline void Quant(uint32_t loopidx);
|
||||
__aicore__ inline void Swiglu(uint32_t loopidx);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const TCubeTiling *__restrict mmTilingData;
|
||||
TPipe *pipe;
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> weightGM;
|
||||
GlobalTensor<CHANNELDTYPE> perChannelScaleGM;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<int32_t> mmOutGM;
|
||||
// define the que
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECIN, 1> perChannelScaleInQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int32_t preOffset = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
float limited = FLOAT_INF;
|
||||
GM_ADDR xTensorPtr;
|
||||
GM_ADDR weightTensorPtr;
|
||||
};
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Init(
|
||||
GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput,
|
||||
GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
{
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
mmTilingData = mmTilingDataIN;
|
||||
gmmBaseParams = gmmSwigluBaseParamsIn;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
pipe = tPipeIN;
|
||||
xTensorPtr = x;
|
||||
limited = gmmBaseParams->limited;
|
||||
weightTensorPtr = weight;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen);
|
||||
mmOutGM.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->M * gmmSwiglu->tokenLen);
|
||||
if ASCEND_IS_AIV {
|
||||
perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale,
|
||||
gmmSwiglu->groupListLen * gmmSwiglu->tokenLen);
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Process()
|
||||
{
|
||||
MNConfig mnConfig;
|
||||
VecConfig vecConfig;
|
||||
CubeProcess(mnConfig);
|
||||
VecProcess(vecConfig);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void
|
||||
GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
DataCopyExtParams copyChannelParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = channelScaleLocal.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(channelScaleLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
uint32_t coreIdx)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;
|
||||
uint32_t curSingleM =
|
||||
mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k;
|
||||
if constexpr (transposeX) {
|
||||
xOffset = mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset);
|
||||
weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + mnConfig.wBaseOffset + GetWOffset(tailN, mnConfig.k));
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k);
|
||||
mm.SetTensorA(xGM[xOffset], transposeX);
|
||||
mm.SetTensorB(weightGM, transposeW);
|
||||
mm.template IterateAll<sync>(mmOutGM[mnConfig.workSpaceOffset], 0);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (B::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::SetMNConfig(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
SetMKN(splitValue, groupIdx, mnConfig);
|
||||
mnConfig.baseM = BASIC_M;
|
||||
mnConfig.baseN = BASIC_N;
|
||||
mnConfig.singleM = SINGLE_CORE_M;
|
||||
mnConfig.singleN = SINGLE_CORE_N;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::SetMKN(const int32_t splitValue,
|
||||
const uint32_t groupIdx, MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<uint32_t>(splitValue);
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline uint64_t GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::GetWOffset(uint32_t tailN, uint32_t k)
|
||||
{
|
||||
uint64_t wOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size
|
||||
} else {
|
||||
wOffset = tailN;
|
||||
}
|
||||
return wOffset;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::CubeProcess(MNConfig &mnConfig)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
preOffset = 0;
|
||||
int32_t prevSplitValue = 0;
|
||||
for (uint32_t groupIdx = 0, count = 0; groupIdx < gmmSwiglu->groupListLen; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
int32_t splitValue = currSplitValue - prevSplitValue;
|
||||
prevSplitValue = currSplitValue;
|
||||
SetMNConfig(splitValue, groupIdx, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
|
||||
uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
|
||||
uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum;
|
||||
uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN;
|
||||
|
||||
while (curBlock < curCount) {
|
||||
MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
|
||||
MMCompute(groupIdx, mnConfig, blockIdx);
|
||||
curBlock += aicCoreNum;
|
||||
}
|
||||
count = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
SyncAll<false>();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::VecProcess(VecConfig &vecConfig)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
UpdateVecConfig(blockIdx, vecConfig);
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.AllocTensor<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.AllocTensor<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
PreLoadTokenAndChannel<CHANNELDTYPE>(channelScaleLocal, vecConfig);
|
||||
}
|
||||
SyncAll<false>();
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
customDataCopyIn(outLoopIdx, vecConfig);
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
UpdateChannelScale<CHANNELDTYPE>(innerLoopIdx, vecConfig);
|
||||
VectorCompute(innerLoopIdx, vecConfig);
|
||||
}
|
||||
customDataCopyOut(vecConfig);
|
||||
}
|
||||
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.DeQue<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
perChannelScaleInQueue.FreeTensor(channelScaleLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock,
|
||||
const uint32_t count, const uint32_t thresholdM_dimN)
|
||||
{
|
||||
mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateVecConfig(uint32_t blockIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
int64_t prevM = 0;
|
||||
for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
vecConfig.M += tempM;
|
||||
}
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
prevM = 0;
|
||||
for (uint32_t groupIdx = 0; groupIdx < gmmSwiglu->groupListLen; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
pipe->Reset();
|
||||
// 第四步 申请空间
|
||||
pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t));
|
||||
pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER,
|
||||
AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::customDataCopyIn(uint32_t outLoopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<int32_t> _inMMLocal_0 = mmOutQueue.DeQue<int32_t>();
|
||||
DataCopyExtParams copyParams_0{
|
||||
1, static_cast<uint32_t>(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_0);
|
||||
|
||||
LocalTensor<int32_t> _inMMLocal_1 = mmOutQueue.DeQue<int32_t>();
|
||||
|
||||
Cast(_inMMLocal_1.ReinterpretCast<float>(), _inMMLocal_1, RoundMode::CAST_NONE,
|
||||
vecConfig.innerLoopNum * gmmSwiglu->tokenLen);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_1);
|
||||
LocalTensor<float> _inMMLocal_2 = mmOutQueue.DeQue<float>();
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
float scale = perTokenScaleGM.GetValue(vecConfig.curIdx);
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
SetFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
WaitFlag<HardEvent::S_V>(EVENT_ID0);
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
mmOutQueue.EnQue(_inMMLocal_2);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::UpdateChannelScale(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 更新perChannel
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> _inChannel = perChannelScaleInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = _inChannel.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams,
|
||||
padParams);
|
||||
}
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
perChannelScaleInQueue.EnQue(_inChannel);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::VectorCompute(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
Dequant(loopIdx, vecConfig);
|
||||
Swiglu(loopIdx);
|
||||
Quant(loopIdx);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Dequant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> perChannelLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
perChannelScaleInQueue.EnQue(perChannelLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Swiglu(uint32_t loopIdx)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
_inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen];
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::Quant(uint32_t loopIdx)
|
||||
{
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluCompute<mmType, sync, CHANNELDTYPE>::customDataCopyOut(VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyPad(quantScaleOutputGM[vecConfig.startIdx], quantScaleLocal, copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
DataCopyPad(quantOutputGM[vecConfig.startIdx * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR], quantLocal,
|
||||
copyParams_1);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // ASCENDC_GROUPED_MATMUL_QUANT_MIXCORE_H
|
||||
@@ -0,0 +1,251 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_antiquant_a8w4_msd.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "grouped_matmul_swiglu_quant.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace matmul;
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr uint32_t BUFFER_NUM = 1;
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4(const LocalTensor<T> dst, const GlobalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T);
|
||||
// 32: int32 -> float16, 为防止跨行数据进入同一32B block，提前每行按偶数block对齐
|
||||
params.dstStride = Ceil(dim0 * sizeof(T), 32) % 2;
|
||||
|
||||
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
|
||||
DataCopyPad(dst, src, params, padParams);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4ND(const LocalTensor<T> dst, const GlobalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T);
|
||||
params.dstStride = 0;
|
||||
|
||||
DataCopyPadExtParams<T> padParams{true, 0, 0, 0};
|
||||
DataCopyPad(dst, src, params, padParams);
|
||||
return;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DataCopyPad2DA8W4(const GlobalTensor<T> dst, const LocalTensor<T> src, uint32_t dim1,
|
||||
uint32_t dim0, uint32_t srcDim0, uint32_t dstDim0)
|
||||
{
|
||||
DataCopyExtParams params;
|
||||
params.blockCount = dim1;
|
||||
params.blockLen = dim0 * sizeof(T);
|
||||
// 32: ub访问粒度为32B
|
||||
params.srcStride = (srcDim0 - dim0) * sizeof(T) / 32;
|
||||
params.dstStride = (dstDim0 - dim0) * sizeof(T);
|
||||
DataCopyPad(dst, src, params);
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
class GMMA8W4MidProcess {
|
||||
public:
|
||||
using bT = typename mmType::BT;
|
||||
|
||||
public:
|
||||
__aicore__ inline GMMA8W4MidProcess(typename mmType::MT &matmul) : mm(matmul)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN);
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx);
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, MNConfig &mnConfig);
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const uint32_t HALF_ALIGN = 16;
|
||||
GlobalTensor<int4b_t> xGM;
|
||||
GlobalTensor<int4b_t> xGM1;
|
||||
GlobalTensor<int4b_t> xGM2;
|
||||
GlobalTensor<int4b_t> weightGM;
|
||||
|
||||
GlobalTensor<half> mmOutGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<uint64_t> weightScaleGM;
|
||||
// define the que
|
||||
uint32_t subBlockIdx = 0;
|
||||
uint32_t coreIdx = 0;
|
||||
uint32_t quantGroupSize = 0;
|
||||
uint32_t vecCount = 0;
|
||||
uint32_t xRowSumCount = 0;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
};
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
gmmBaseParams = gmmSwigluBaseParamsIN;
|
||||
xRowSumCount = gmmBaseParams->M;
|
||||
xGM1.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.workSpaceGM); // 从前处理中获得的结果
|
||||
xGM2.SetGlobalBuffer(
|
||||
(__gm__ int4b_t *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1));
|
||||
weightGM.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.weightGM);
|
||||
weightScaleGM.SetGlobalBuffer((__gm__ uint64_t *)gmAddrParams.weightScaleGM);
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
mmOutGM1.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3));
|
||||
quantGroupSize = gmmBaseParams->K / gmmBaseParams->quantGroupNum; // 约束为整除关系
|
||||
subBlockIdx = GetSubBlockIdx();
|
||||
coreIdx = GetBlockIdx();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (bT::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<int64_t>(splitValue);
|
||||
mnConfig.baseM = gmmBaseParams->baseM;
|
||||
mnConfig.baseN = gmmBaseParams->baseN;
|
||||
mnConfig.singleM = gmmBaseParams->baseM;
|
||||
mnConfig.singleN = gmmBaseParams->baseN;
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) {
|
||||
return;
|
||||
}
|
||||
xGM = (workspaceSplitLoopIdx % 2 == 0 ? xGM1 : xGM2);
|
||||
mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
MNConfig mnConfig;
|
||||
mnConfig.baseM = gmmBaseParams->baseM;
|
||||
mnConfig.baseN = gmmBaseParams->baseN;
|
||||
mnConfig.singleM = gmmBaseParams->baseM;
|
||||
mnConfig.singleN = gmmBaseParams->baseN;
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit :
|
||||
currSplitValue;
|
||||
|
||||
int32_t splitValue = (currSplitValue - prevSplitValue) * 2; // 2: int8 has been split in 2 int4
|
||||
prevSplitValue = currSplitValue;
|
||||
|
||||
SetMNConfig(splitValue, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
uint32_t curCount = preCount + mnConfig.blockDimN * mnConfig.blockDimM;
|
||||
uint32_t curBlock = coreIdx >= preCount ? coreIdx : coreIdx + gmmBaseParams->coreNum;
|
||||
while (curBlock < curCount) {
|
||||
mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN;
|
||||
MMCompute(groupIdx, mnConfig, workspaceSplitConfig);
|
||||
curBlock += gmmBaseParams->coreNum;
|
||||
}
|
||||
preCount = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA8W4MidProcess<mmType>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.singleN;
|
||||
if (unlikely(mnConfig.nIdx == mnConfig.blockDimN - 1)) {
|
||||
curSingleN = gmmBaseParams->N - tailN;
|
||||
}
|
||||
uint32_t curSingleM = mnConfig.singleM;
|
||||
if (unlikely(mnConfig.mIdx == mnConfig.blockDimM - 1)) {
|
||||
curSingleM = mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t weightOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightOffset = static_cast<uint64_t>(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN * gmmBaseParams->K;
|
||||
} else {
|
||||
weightOffset = static_cast<uint64_t>(groupIdx) * gmmBaseParams->N * gmmBaseParams->K + tailN;
|
||||
}
|
||||
mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize); // 8, 256, 512 --> 514us
|
||||
GlobalTensor<int4b_t> weightSlice;
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
for (uint32_t loopK = 0; loopK < gmmBaseParams->quantGroupNum; loopK++) {
|
||||
mm.SetTensorA(
|
||||
xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]);
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmBaseParams->N];
|
||||
}
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mm.SetTensorB(weightSlice);
|
||||
mm.SetQuantVector(weightScaleGM[groupIdx * gmmBaseParams->N * gmmBaseParams->quantGroupNum +
|
||||
loopK * gmmBaseParams->N + tailN]);
|
||||
mm.Iterate();
|
||||
mm.GetTensorC(mmOutGM[mnConfig.workSpaceOffset], loopK == 0 ? 0 : 1);
|
||||
}
|
||||
}
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_MID_H
|
||||
@@ -0,0 +1,385 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_a8w4_msd_post.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_POST_H
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
#define DOUBLE_BUFFER 2
|
||||
constexpr float DEFAULT_MUL_SCALE = 16.0f;
|
||||
class GMMA8W4PostProcess {
|
||||
public:
|
||||
__aicore__ inline GMMA8W4PostProcess(){};
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN);
|
||||
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
|
||||
private:
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
|
||||
__aicore__ inline void UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void VectorCompute(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<half> &mmOutGM, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig, WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void PreLoadAuxiliaryMatrix(VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Quant(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> weightAuxiliaryMatrixGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<half> mmOutGM;
|
||||
LocalTensor<float> mmLocal_fp32;
|
||||
LocalTensor<half> mmLocal_fp16;
|
||||
TQue<QuePosition::VECIN, 1> weightAuxiliaryMatrixInQueue;
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
};
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN,
|
||||
const GMMSwiglu *__restrict gmmSwigluIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
gmmBaseParams = gmmSwigluBaseParamsIN;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
weightAuxiliaryMatrixGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.weightAuxiliaryMatrixGM); // E, N
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwiglu->groupListLen);
|
||||
mmOutGM1.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset2));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset3));
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<half> &mmOutGM,
|
||||
VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
mmLocal_fp16 = mmOutQueue.DeQue<half>();
|
||||
mmLocal_fp32 = mmLocal_fp16.ReinterpretCast<float>();
|
||||
const int64_t processNum = 2 * vecConfig.innerLoopNum * gmmSwiglu->tokenLen;
|
||||
DataCopyExtParams copyParams_0{1, static_cast<uint32_t>(processNum * sizeof(half)), 0, 0, 0};
|
||||
DataCopyPadExtParams<half> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(mmLocal_fp16[processNum], mmOutGM[vecConfig.curOffset * DOUBLE_ROW], copyParams_0, padParams_0);
|
||||
|
||||
mmOutQueue.EnQue(mmLocal_fp16);
|
||||
mmLocal_fp16 = mmOutQueue.DeQue<half>();
|
||||
// 1. fp16 -> fp32
|
||||
Cast(mmLocal_fp32, mmLocal_fp16[processNum], RoundMode::CAST_NONE, processNum);
|
||||
PipeBarrier<PIPE_V>();
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
// 2. high_4bit * 16 + low_4bit
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
Muls(mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen],
|
||||
DEFAULT_MUL_SCALE, gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Add(mmLocal_fp32[i * gmmSwiglu->tokenLen], mmLocal_fp32[(DOUBLE_ROW * i) * gmmSwiglu->tokenLen],
|
||||
mmLocal_fp32[(DOUBLE_ROW * i + 1) * gmmSwiglu->tokenLen], gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
// 1.辅助矩阵加回
|
||||
MergeAuxiliaryMatrix(loopIdx, vecConfig);
|
||||
// 2.perToken反量化
|
||||
MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig);
|
||||
// 3.Swiglu
|
||||
Swiglu(loopIdx, vecConfig);
|
||||
// 4.Quant
|
||||
Quant(loopIdx, vecConfig);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::MergeAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
mmLocal_fp32 = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
Add(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], weightAuxiliaryMatrixLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
PipeBarrier<PIPE_V>();
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx);
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen];
|
||||
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(mmLocal_fp32[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], mmLocal_fp32[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, mmLocal_fp32[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, mmLocal_fp32, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(mmLocal_fp32);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
vecConfig.M = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize :
|
||||
workspaceSplitConfig.lastLoopTaskSize;
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
|
||||
// 第四步 申请空间
|
||||
// 2 * row * n * sizeof(float) + row * n / 2 * sizeof(int8) + alignUp<row, 8> * sizeof(float) + n * sizeof(float) +
|
||||
// n / 2 *sizeof(float) + 64 < 191 * 1024
|
||||
pipe->InitBuffer(mmOutQueue, 1, 2 * gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, 1,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, 1, AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
pipe->InitBuffer(weightAuxiliaryMatrixInQueue, 1, gmmSwiglu->tokenLen * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::PreLoadAuxiliaryMatrix(VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyAuxiliaryMatrixParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyAuxiliaryMatrixParams, padParams);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 更新weightAuxiliaryMatrix
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) {
|
||||
return;
|
||||
}
|
||||
VecConfig vecConfig;
|
||||
UpdateVecConfig(blockIdx, vecConfig, workspaceSplitConfig, workspaceSplitLoopIdx, pipe);
|
||||
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
mmOutGM = (workspaceSplitLoopIdx % 2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.AllocTensor<float>();
|
||||
LocalTensor<half> mmLocal_fp32 = mmOutQueue.AllocTensor<half>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
|
||||
mmOutQueue.EnQue(mmLocal_fp32);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
PreLoadAuxiliaryMatrix(vecConfig);
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
int32_t eventIdMTE3ToMTE2 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
|
||||
SetFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
WaitFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
// 1.matmul中间结果搬入 + 高四位与低四位合并
|
||||
customDataCopyIn(outLoopIdx, mmOutGM, vecConfig, workspaceSplitConfig);
|
||||
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
// 2.如果涉及group切换，更新辅助矩阵
|
||||
UpdateAuxiliaryMatrix(innerLoopIdx, vecConfig);
|
||||
// 3. 四步vector计算（辅助矩阵加回、perToken反量化、Swiglu、Quant）
|
||||
VectorCompute(innerLoopIdx, vecConfig, workspaceSplitConfig);
|
||||
}
|
||||
int32_t eventIdVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
|
||||
SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
WaitFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
customDataCopyOut(vecConfig, workspaceSplitConfig);
|
||||
}
|
||||
weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
mmLocal_fp32 = mmOutQueue.DeQue<half>();
|
||||
quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
|
||||
weightAuxiliaryMatrixInQueue.FreeTensor(weightAuxiliaryMatrixLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal_fp32);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::customDataCopyOut(VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal,
|
||||
copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_AFTER_H
|
||||
@@ -0,0 +1,219 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_a8w4_msd_pre.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
#ifdef GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
#define BUFFER_NUM_A8W4_PRE 1
|
||||
constexpr int TWO = 2;
|
||||
constexpr int EIGHT = 8;
|
||||
constexpr size_t LEN_128 = 128; // 16bit operator
|
||||
constexpr int DATA_BLOCK_SIZE_32 = 32;
|
||||
class GMMA8W4PreProcess {
|
||||
public:
|
||||
__aicore__ inline GMMA8W4PreProcess(){};
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN);
|
||||
__aicore__ inline void CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_, uint32_t &curCoreStartOffset_);
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
__aicore__ inline void CustomInitBuffer(TPipe *pipe);
|
||||
|
||||
private:
|
||||
TQue<QuePosition::VECIN, BUFFER_NUM_A8W4_PRE> vecInQueueX, vecInQueueXBak;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA1;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA2;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueA3;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueue0F;
|
||||
TQue<QuePosition::VECOUT, BUFFER_NUM_A8W4_PRE> vecOutQueueRowSum;
|
||||
TBuf<TPosition::VECCALC> tempBuff;
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParams;
|
||||
LocalTensor<int8_t> xTensor;
|
||||
LocalTensor<half> xHighHalfTensor;
|
||||
LocalTensor<float> xHighFloatTensor;
|
||||
LocalTensor<half> xLowHalfTensor;
|
||||
LocalTensor<half> xLowHalfTensor2;
|
||||
LocalTensor<int4b_t> xHighI4Tensor;
|
||||
LocalTensor<int4b_t> xLowI4Tensor;
|
||||
LocalTensor<int16_t> xLowI16Tensor;
|
||||
LocalTensor<int64_t> groupListTensor;
|
||||
LocalTensor<float> groupListFTensor;
|
||||
LocalTensor<float> xRowSumTensor;
|
||||
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> yGm;
|
||||
GlobalTensor<int8_t> yGm1;
|
||||
GlobalTensor<int8_t> yGm2;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
|
||||
uint32_t vK{0};
|
||||
uint32_t vKAlign{0};
|
||||
uint32_t totalM{0};
|
||||
uint32_t blockDim{0};
|
||||
uint32_t curCoreId{0};
|
||||
uint32_t curCoreTaskNum{0};
|
||||
uint32_t curCoreStartOffset{0};
|
||||
uint32_t curCoreOuterLoopNum{0};
|
||||
uint32_t curCoreInnerTailLoopNum{0};
|
||||
uint32_t groupNum{0};
|
||||
};
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.xGM);
|
||||
yGm1.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM);
|
||||
yGm2.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1);
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
gmmSwigluBaseParams = gmmSwigluBaseParamsIN;
|
||||
vK = gmmSwigluBaseParams->K;
|
||||
groupNum = static_cast<uint32_t>(gmmSwigluBaseParams->groupNum);
|
||||
// M * K * 7B (1B + 0.5B + 0.5B + 2B + 4B) <= UBsize - 256B
|
||||
blockDim = GetBlockNum() * GetTaskRation();
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::CustomInitBuffer(TPipe *pipe)
|
||||
{
|
||||
pipe->InitBuffer(vecInQueueX, BUFFER_NUM_A8W4_PRE, vK * sizeof(int8_t)); // K * 1B
|
||||
pipe->InitBuffer(vecOutQueueA1, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B
|
||||
pipe->InitBuffer(vecOutQueueA2, BUFFER_NUM_A8W4_PRE, vK * sizeof(int4b_t)); // K * 0.5B
|
||||
pipe->InitBuffer(vecOutQueueA3, BUFFER_NUM_A8W4_PRE, vK * sizeof(half)); // K * 2B
|
||||
// xLowHalfTensor, xLowHalfTensor2 and xHighFloatTensor share the same buffer
|
||||
pipe->InitBuffer(tempBuff, vK * sizeof(float)); // K * 4B
|
||||
constexpr int BUFFER_SIZE_256B = 128 * sizeof(int16_t);
|
||||
pipe->InitBuffer(vecOutQueue0F, BUFFER_NUM_A8W4_PRE, BUFFER_SIZE_256B); // 256B
|
||||
}
|
||||
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::CalculateTaskInfoEachCore(uint32_t &curCoreTaskNum_,
|
||||
uint32_t &curCoreStartOffset_)
|
||||
{
|
||||
// 均分任务数
|
||||
int64_t eachCoreTaskNum = (totalM + blockDim - 1) / blockDim; // 每个核处理的数据量
|
||||
// 尾核任务数
|
||||
int64_t taskNumPertailCore = eachCoreTaskNum - 1;
|
||||
// 实际使用核数
|
||||
int64_t usedCoreNum = totalM >= blockDim ? blockDim : totalM;
|
||||
// 尾核起始索引
|
||||
uint32_t tailCoreIdx = totalM - (eachCoreTaskNum - 1) * usedCoreNum;
|
||||
curCoreId = GetBlockIdx();
|
||||
// 每个核处理的任务数量 = 是否为尾核 ？均分任务数 ：(均分任务数 - 1)
|
||||
curCoreTaskNum_ = curCoreId < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
// 每个核处理的起始偏移地址 = 是否为尾核 ？均分任务数 * blockId : (均分任务数 - 1) * blockId + 尾核起始索引
|
||||
curCoreStartOffset_ =
|
||||
curCoreId < tailCoreIdx ? eachCoreTaskNum * curCoreId : ((eachCoreTaskNum - 1) * curCoreId + tailCoreIdx);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PreProcess::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx, TPipe *pipe)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount) {
|
||||
return;
|
||||
}
|
||||
yGm = (workspaceSplitLoopIdx % 2 == 0 ? yGm1 : yGm2);
|
||||
CustomInitBuffer(pipe);
|
||||
constexpr int32_t MASK = 128;
|
||||
xTensor = vecInQueueX.AllocTensor<int8_t>();
|
||||
xHighI4Tensor = vecOutQueueA1.AllocTensor<int4b_t>();
|
||||
xLowI4Tensor = vecOutQueueA2.AllocTensor<int4b_t>();
|
||||
xHighHalfTensor = vecOutQueueA3.AllocTensor<half>();
|
||||
const uint32_t xLowHalfOffset = vK * sizeof(half);
|
||||
xLowHalfTensor = tempBuff.GetWithOffset<half>(xLowHalfOffset, 0);
|
||||
xLowHalfTensor2 = tempBuff.GetWithOffset<half>(xLowHalfOffset, xLowHalfOffset);
|
||||
xLowI16Tensor = vecOutQueue0F.AllocTensor<int16_t>();
|
||||
|
||||
Duplicate(xLowI16Tensor, static_cast<int16_t>(0x0F0F), MASK); // get rid of high 4 bits in every int8
|
||||
PipeBarrier<PIPE_V>();
|
||||
const size_t LEN_VK = (vK / 2) / 128;
|
||||
const size_t LAST_LEN_VK = (vK % 256) / 2;
|
||||
const half ONE_SIXTEENTH = static_cast<half>(0.0625f);
|
||||
// groupList仅支持count
|
||||
SetFlag<HardEvent::MTE2_S>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::MTE2_S>(EVENT_ID0);
|
||||
totalM = workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - 1 ? workspaceSplitConfig.notLastTaskSize :
|
||||
workspaceSplitConfig.lastLoopTaskSize;
|
||||
SetFlag<HardEvent::S_MTE2>(EVENT_ID0);
|
||||
WaitFlag<HardEvent::S_MTE2>(EVENT_ID0);
|
||||
CalculateTaskInfoEachCore(curCoreTaskNum, curCoreStartOffset);
|
||||
SetFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
|
||||
for (uint32_t xloop = 0; xloop < curCoreTaskNum; xloop++) {
|
||||
uint64_t relStartAddr = (xloop + curCoreStartOffset) * vK;
|
||||
uint64_t absStartAddr = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize * vK + relStartAddr;
|
||||
// 高四位处理开始
|
||||
WaitFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
DataCopy(xTensor, xGM[absStartAddr], vK);
|
||||
SetFlag<HardEvent::MTE2_V>(EVENT_ID0); // 3
|
||||
WaitFlag<HardEvent::MTE2_V>(EVENT_ID0); // 3
|
||||
Cast(xHighHalfTensor, xTensor, AscendC::RoundMode::CAST_NONE, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(xHighHalfTensor, xHighHalfTensor, ONE_SIXTEENTH, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
Cast(xHighI4Tensor, xHighHalfTensor, AscendC::RoundMode::CAST_FLOOR, vK);
|
||||
SetFlag<HardEvent::V_MTE3>(EVENT_ID0); // 4
|
||||
WaitFlag<HardEvent::V_MTE3>(EVENT_ID0); // 4
|
||||
DataCopy(yGm[relStartAddr], xHighI4Tensor.ReinterpretCast<int8_t>(), vK / 2);
|
||||
// 高四位处理结束
|
||||
|
||||
// 低四位处理开始
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
And(xLowHalfTensor.ReinterpretCast<int16_t>(), xTensor.ReinterpretCast<int16_t>(), xLowI16Tensor, LEN_128,
|
||||
LEN_VK, {1, 1, 1, 8, 8, 0});
|
||||
if (LAST_LEN_VK > 0) {
|
||||
And(xLowHalfTensor[LEN_VK * LEN_128].ReinterpretCast<int16_t>(),
|
||||
xTensor[LEN_VK * LEN_128 * TWO].ReinterpretCast<int16_t>(), xLowI16Tensor, LAST_LEN_VK, 1,
|
||||
{1, 1, 1, 8, 8, 0});
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
SetFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
Cast(xLowHalfTensor2.ReinterpretCast<half>(), xLowHalfTensor.ReinterpretCast<int8_t>(),
|
||||
AscendC::RoundMode::CAST_NONE, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
const half MINUS_EIGHT = static_cast<half>(-8);
|
||||
Adds(xHighHalfTensor, xLowHalfTensor2, MINUS_EIGHT, vK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
Cast(xLowI4Tensor, xHighHalfTensor.ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE, vK);
|
||||
SetFlag<HardEvent::V_MTE3>(EVENT_ID1); // 5
|
||||
WaitFlag<HardEvent::V_MTE3>(EVENT_ID1); // 5
|
||||
DataCopy(yGm[relStartAddr + vK / TWO], xLowI4Tensor.ReinterpretCast<int8_t>(), vK / TWO);
|
||||
SetFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
// 低四位处理结束
|
||||
}
|
||||
|
||||
WaitFlag<HardEvent::V_MTE2>(EVENT_ID0); // 0
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID0); // 1
|
||||
WaitFlag<HardEvent::MTE3_V>(EVENT_ID1); // 2
|
||||
vecInQueueX.FreeTensor(xTensor);
|
||||
vecOutQueueA1.FreeTensor(xHighI4Tensor);
|
||||
vecOutQueueA2.FreeTensor(xLowI4Tensor);
|
||||
vecOutQueueA3.FreeTensor(xHighHalfTensor);
|
||||
vecOutQueue0F.FreeTensor(xLowI16Tensor);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_A8W4_MSD_PRE_H
|
||||
@@ -0,0 +1,635 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_split_ws.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_utils.h"
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
/** @brief internal computation class
|
||||
*/
|
||||
|
||||
template <class mmType, bool sync = false, typename CHANNELDTYPE = float>
|
||||
class GMMSwigluSplitWorkSpaceCompute {
|
||||
public:
|
||||
using AT = typename mmType::AT::T;
|
||||
using BT = typename mmType::BT::T;
|
||||
using B = typename mmType::BT;
|
||||
using CT = typename mmType::CT::T;
|
||||
using BiasT = typename mmType::BiasT::T;
|
||||
using WT = int8_t;
|
||||
constexpr static bool transposeX = mmType::AT::isTrans;
|
||||
constexpr static bool transposeW = mmType::BT::isTrans;
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
/** @brief constructor */
|
||||
__aicore__ inline GMMSwigluSplitWorkSpaceCompute(typename mmType::MT &mm_) : mm(mm_)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale,
|
||||
GM_ADDR groupList, GM_ADDR quantOutput, GM_ADDR quantScaleOutput, GM_ADDR workspace,
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParamsIN,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN,
|
||||
TPipe *tPipeIN);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
__aicore__ inline void MMCompute(uint32_t groupIdx, MNConfig &mnConfig, uint32_t coreIdx,
|
||||
GlobalTensor<int32_t> &mmOutGM);
|
||||
|
||||
__aicore__ inline void UpdateMnConfig(MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMNConfig(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline void SetMKN(const int32_t splitValue, const uint32_t groupIdx, MNConfig &mnConfig);
|
||||
|
||||
__aicore__ inline uint64_t GetWOffset(uint32_t tailN, uint32_t k);
|
||||
|
||||
__aicore__ inline void MNBlockIdxCompute(MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count,
|
||||
const uint32_t thresholdM_dimN);
|
||||
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void UpdateChannelScale(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void VectorCompute(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void PreLoadTokenAndChannel(LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void UpdateWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int32_t workspaceSplitLoopIdx);
|
||||
|
||||
__aicore__ inline void InitWorkSpaceSplitConfig(WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<int32_t> &mmOutGM, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void customDataCopyOut(VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Dequant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Quant(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Swiglu(uint32_t loopidx, VecConfig &vecConfig);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const GMMSwigluBaseParams *__restrict gmmBaseParams;
|
||||
const GMMSwiglu *__restrict gmmSwiglu;
|
||||
const TCubeTiling *__restrict mmTilingData;
|
||||
WorkSpaceSplitConfig workspaceSplitConfig;
|
||||
TPipe *pipe;
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> weightGM;
|
||||
GlobalTensor<CHANNELDTYPE> perChannelScaleGM;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<int32_t> mmOutGM1;
|
||||
GlobalTensor<int32_t> mmOutGM2;
|
||||
// define the que
|
||||
TQue<QuePosition::VECIN, 1> mmOutQueue;
|
||||
TQue<QuePosition::VECIN, 1> perChannelScaleInQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantOutQueue;
|
||||
TQue<QuePosition::VECOUT, 1> quantScaleOutQueue;
|
||||
TBuf<TPosition::VECCALC> reduceWorkspace;
|
||||
uint32_t blockIdx = 0;
|
||||
int64_t aicCoreNum = 0;
|
||||
int64_t aivCoreNum = 0;
|
||||
GM_ADDR xTensorPtr;
|
||||
GM_ADDR weightTensorPtr;
|
||||
float limited = FLOAT_INF;
|
||||
};
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Init(
|
||||
GM_ADDR x, GM_ADDR weight, GM_ADDR perChannelScale, GM_ADDR perTokenScale, GM_ADDR groupList, GM_ADDR quantOutput,
|
||||
GM_ADDR quantScaleOutput, GM_ADDR workspace, const GMMSwigluBaseParams *__restrict gmmSwigluBaseParamsIn,
|
||||
const TCubeTiling *__restrict mmTilingDataIN, const GMMSwiglu *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
{
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
pipe = tPipeIN;
|
||||
xTensorPtr = x;
|
||||
weightTensorPtr = weight;
|
||||
mmTilingData = mmTilingDataIN;
|
||||
gmmBaseParams = gmmSwigluBaseParamsIn;
|
||||
gmmSwiglu = gmmSwigluIN;
|
||||
limited = gmmBaseParams->limited;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)groupList, gmmSwiglu->groupListLen);
|
||||
mmOutGM1.SetGlobalBuffer((__gm__ int32_t *)workspace, gmmBaseParams->mLimit * gmmSwiglu->tokenLen);
|
||||
mmOutGM2.SetGlobalBuffer((__gm__ int32_t *)workspace + gmmBaseParams->mLimit * gmmSwiglu->tokenLen,
|
||||
gmmBaseParams->mLimit * gmmSwiglu->tokenLen);
|
||||
if ASCEND_IS_AIV {
|
||||
perChannelScaleGM.SetGlobalBuffer((__gm__ CHANNELDTYPE *)perChannelScale,
|
||||
gmmSwiglu->groupListLen * gmmSwiglu->tokenLen);
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)perTokenScale, gmmBaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)quantOutput,
|
||||
gmmBaseParams->M * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)quantScaleOutput, gmmBaseParams->M);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::InitWorkSpaceSplitConfig(
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwiglu->groupListLen - 1);
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmBaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.leftMatrixStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = 0;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = 0;
|
||||
workspaceSplitConfig.isLastLoop = false;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateWorkSpaceSplitConfig(
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig, int32_t workspaceSplitLoopIdx)
|
||||
{
|
||||
workspaceSplitConfig.leftMatrixStartIndex = workspaceSplitLoopIdx * gmmBaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwiglu->groupListLen) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen - 1 ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
if (curTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmBaseParams->mLimit && nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmBaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
}
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex++;
|
||||
}
|
||||
workspaceSplitConfig.isLastLoop = workspaceSplitLoopIdx == workspaceSplitConfig.loopCount - 1 ? true : false;
|
||||
|
||||
if (workspaceSplitConfig.isLastLoop) {
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex =
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwiglu->groupListLen ?
|
||||
gmmSwiglu->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Process()
|
||||
{
|
||||
InitWorkSpaceSplitConfig(workspaceSplitConfig);
|
||||
int32_t parallelNum = 2; // 2: double workspace buffer
|
||||
for (int32_t workspaceSplitLoopIdx = 0; workspaceSplitLoopIdx < workspaceSplitConfig.loopCount;
|
||||
workspaceSplitLoopIdx++) {
|
||||
UpdateWorkSpaceSplitConfig(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
GlobalTensor<int32_t> mmOutGM = (workspaceSplitLoopIdx % parallelNum == 0) ? mmOutGM1 : mmOutGM2;
|
||||
|
||||
if ASCEND_IS_AIC {
|
||||
if (workspaceSplitLoopIdx >= parallelNum) { // first parallelNum core no need to wait
|
||||
SyncAll<false>();
|
||||
}
|
||||
MNConfig mnConfig;
|
||||
int32_t prevSplitValue = workspaceSplitConfig.leftMatrixStartIndex;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, count = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmBaseParams->mLimit :
|
||||
currSplitValue;
|
||||
int32_t splitValue = currSplitValue - prevSplitValue;
|
||||
prevSplitValue = currSplitValue;
|
||||
SetMNConfig(splitValue, groupIdx, mnConfig);
|
||||
if (mnConfig.m <= 0 || mnConfig.k <= 0 || mnConfig.n <= 0) {
|
||||
continue;
|
||||
}
|
||||
mnConfig.blockDimM = Ceil(mnConfig.m, mnConfig.singleM);
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
|
||||
uint32_t curCount = count + mnConfig.blockDimM * mnConfig.blockDimN;
|
||||
uint32_t curBlock = blockIdx >= count ? blockIdx : blockIdx + gmmBaseParams->coreNum;
|
||||
uint32_t thresholdM_dimN = THRESHOLD_BLOCK_NUM * mnConfig.blockDimN;
|
||||
|
||||
while (curBlock < curCount) {
|
||||
MNBlockIdxCompute(mnConfig, curBlock, count, thresholdM_dimN);
|
||||
MMCompute(groupIdx, mnConfig, blockIdx, mmOutGM);
|
||||
curBlock += aicCoreNum;
|
||||
}
|
||||
count = curCount % gmmBaseParams->coreNum;
|
||||
}
|
||||
SyncAll<false>();
|
||||
}
|
||||
|
||||
if ASCEND_IS_AIV {
|
||||
VecConfig vecConfig;
|
||||
UpdateVecConfig(blockIdx, vecConfig);
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.AllocTensor<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.AllocTensor<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
PreLoadTokenAndChannel<CHANNELDTYPE>(channelScaleLocal, vecConfig);
|
||||
}
|
||||
SyncAll<false>();
|
||||
if (blockIdx < vecConfig.usedCoreNum) {
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum =
|
||||
outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum : gmmSwiglu->maxProcessRowNum;
|
||||
int32_t eventIdMTE3ToMTE2 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
|
||||
SetFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
WaitFlag<HardEvent::MTE3_MTE2>(eventIdMTE3ToMTE2);
|
||||
customDataCopyIn(outLoopIdx, mmOutGM, vecConfig);
|
||||
for (uint32_t innerLoopIdx = 0; innerLoopIdx < vecConfig.innerLoopNum; innerLoopIdx++) {
|
||||
UpdateChannelScale<CHANNELDTYPE>(innerLoopIdx, vecConfig);
|
||||
VectorCompute(innerLoopIdx, vecConfig);
|
||||
}
|
||||
int32_t eventIdVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
|
||||
SetFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
WaitFlag<HardEvent::V_MTE3>(eventIdVToMTE3);
|
||||
customDataCopyOut(vecConfig);
|
||||
}
|
||||
|
||||
LocalTensor<float> channelScaleLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
LocalTensor<int32_t> mmLocal = mmOutQueue.DeQue<int32_t>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
perChannelScaleInQueue.FreeTensor(channelScaleLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
}
|
||||
if (workspaceSplitLoopIdx < workspaceSplitConfig.loopCount - parallelNum) {
|
||||
SyncAll<false>();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::PreLoadTokenAndChannel(
|
||||
LocalTensor<float> &channelScaleLocal, VecConfig &vecConfig)
|
||||
{
|
||||
DataCopyExtParams copyChannelParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = channelScaleLocal.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
int32_t eventIdMTE2ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
Cast(channelScaleLocal, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(channelScaleLocal, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyChannelParams, padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(channelScaleLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::MMCompute(uint32_t groupIdx, MNConfig &mnConfig,
|
||||
uint32_t coreIdx, GlobalTensor<int32_t> &mmOutGM)
|
||||
{
|
||||
uint32_t tailN = mnConfig.nIdx * mnConfig.singleN;
|
||||
uint32_t curSingleN = mnConfig.nIdx < mnConfig.blockDimN - 1 ? mnConfig.singleN : mnConfig.n - tailN;
|
||||
uint32_t curSingleM =
|
||||
mnConfig.mIdx < mnConfig.blockDimM - 1 ? mnConfig.singleM : mnConfig.m - mnConfig.mIdx * mnConfig.singleM;
|
||||
uint64_t xOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.k;
|
||||
if constexpr (transposeX) {
|
||||
xOffset = mnConfig.mIdx * mnConfig.singleM;
|
||||
}
|
||||
uint64_t outOffset = mnConfig.mIdx * mnConfig.singleM * mnConfig.n + tailN;
|
||||
xGM.SetGlobalBuffer((__gm__ int8_t *)xTensorPtr + mnConfig.xBaseOffset +
|
||||
workspaceSplitConfig.leftMatrixStartIndex * mnConfig.k);
|
||||
weightGM.SetGlobalBuffer((__gm__ int8_t *)weightTensorPtr + groupIdx * mnConfig.k * mnConfig.n +
|
||||
GetWOffset(tailN, mnConfig.k));
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
} else {
|
||||
weightGM.SetL2CacheHint(CacheMode::CACHE_MODE_NORMAL);
|
||||
}
|
||||
mnConfig.workSpaceOffset = outOffset + mnConfig.yBaseOffset;
|
||||
mm.SetOrgShape(mnConfig.m, mnConfig.n, mnConfig.k);
|
||||
mm.SetSingleShape(curSingleM, curSingleN, mnConfig.k);
|
||||
mm.SetTensorA(xGM[xOffset], transposeX);
|
||||
mm.SetTensorB(weightGM, transposeW);
|
||||
mm.template IterateAll<sync>(mmOutGM[mnConfig.workSpaceOffset], 0);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateMnConfig(MNConfig &mnConfig)
|
||||
{
|
||||
if constexpr (B::format == CubeFormat::NZ) {
|
||||
mnConfig.wBaseOffset += AlignUp<16>(mnConfig.k) * AlignUp<32>(mnConfig.n); // 16: nz format last two dim size
|
||||
} else {
|
||||
mnConfig.wBaseOffset += mnConfig.k * mnConfig.n;
|
||||
}
|
||||
mnConfig.nAxisBaseOffset += mnConfig.n;
|
||||
mnConfig.mAxisBaseOffset += mnConfig.m;
|
||||
mnConfig.xBaseOffset += mnConfig.m * mnConfig.k;
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::SetMNConfig(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
SetMKN(splitValue, groupIdx, mnConfig);
|
||||
mnConfig.baseM = BASIC_M;
|
||||
mnConfig.baseN = BASIC_N;
|
||||
mnConfig.singleM = SINGLE_CORE_M;
|
||||
mnConfig.singleN = SINGLE_CORE_N;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::SetMKN(const int32_t splitValue,
|
||||
const uint32_t groupIdx,
|
||||
MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<int64_t>(splitValue);
|
||||
mnConfig.k = gmmBaseParams->K; // tilingData
|
||||
mnConfig.n = gmmBaseParams->N; // tilingData
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline uint64_t GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::GetWOffset(uint32_t tailN,
|
||||
uint32_t k)
|
||||
{
|
||||
uint64_t wOffset = 0;
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
wOffset = tailN * AlignUp<16>(k); // 16: nz format last two dim size
|
||||
} else {
|
||||
wOffset = tailN;
|
||||
}
|
||||
return wOffset;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::MNBlockIdxCompute(
|
||||
MNConfig &mnConfig, const uint32_t curBlock, const uint32_t count, const uint32_t thresholdM_dimN)
|
||||
{
|
||||
mnConfig.mIdx = (curBlock - count) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - count) % mnConfig.blockDimN;
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateVecConfig(uint32_t blockIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 第一步 读取grouplist reduceSum 计算总数据个数
|
||||
vecConfig.M =
|
||||
workspaceSplitConfig.isLastLoop ? workspaceSplitConfig.lastLoopTaskSize : workspaceSplitConfig.notLastTaskSize;
|
||||
// 第二步 计算分核
|
||||
uint32_t eachCoreTaskNum = (vecConfig.M + aivCoreNum - 1) / aivCoreNum;
|
||||
vecConfig.usedCoreNum = vecConfig.M >= aivCoreNum ? aivCoreNum : vecConfig.M;
|
||||
uint32_t tailCoreIdx = vecConfig.M - (eachCoreTaskNum - 1) * vecConfig.usedCoreNum;
|
||||
vecConfig.taskNum = blockIdx < tailCoreIdx ? eachCoreTaskNum : eachCoreTaskNum - 1;
|
||||
vecConfig.startIdx =
|
||||
blockIdx < tailCoreIdx ? eachCoreTaskNum * blockIdx : ((eachCoreTaskNum - 1) * blockIdx + tailCoreIdx);
|
||||
vecConfig.curIdx = vecConfig.startIdx;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitConfig.leftMatrixStartIndex;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = groupListGM.GetValue(groupIdx);
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpadteInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum = (vecConfig.taskNum + gmmSwiglu->maxProcessRowNum - 1) / gmmSwiglu->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwiglu->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwiglu->maxProcessRowNum :
|
||||
gmmSwiglu->maxProcessRowNum;
|
||||
pipe->Reset();
|
||||
// 第四步 申请空间
|
||||
pipe->InitBuffer(mmOutQueue, DOUBLE_BUFFER, gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen * sizeof(int32_t));
|
||||
pipe->InitBuffer(perChannelScaleInQueue, DOUBLE_BUFFER, gmmSwiglu->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, DOUBLE_BUFFER,
|
||||
gmmSwiglu->maxProcessRowNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, DOUBLE_BUFFER,
|
||||
AlignUp<int32_t>(gmmSwiglu->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) + UB_BLOCK_UNIT_SIZE +
|
||||
UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::customDataCopyIn(
|
||||
uint32_t outLoopIdx, GlobalTensor<int32_t> &mmOutGM, VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<int32_t> _inMMLocal_0 = mmOutQueue.DeQue<int32_t>();
|
||||
DataCopyExtParams copyParams_0{
|
||||
1, static_cast<uint32_t>(vecConfig.innerLoopNum * gmmSwiglu->tokenLen * sizeof(int32_t)), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(_inMMLocal_0, mmOutGM[vecConfig.curOffset], copyParams_0, padParams_0);
|
||||
mmOutQueue.EnQue(_inMMLocal_0);
|
||||
|
||||
LocalTensor<int32_t> _inMMLocal_1 = mmOutQueue.DeQue<int32_t>();
|
||||
|
||||
Cast(_inMMLocal_1.ReinterpretCast<float>(), _inMMLocal_1, RoundMode::CAST_NONE,
|
||||
vecConfig.innerLoopNum * gmmSwiglu->tokenLen);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_1);
|
||||
LocalTensor<float> _inMMLocal_2 = mmOutQueue.DeQue<float>();
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
for (uint32_t i = 0; i < vecConfig.innerLoopNum; i++) {
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
float scale = perTokenScaleGM.GetValue(vecConfig.curIdx + workspaceSplitConfig.leftMatrixStartIndex);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal_2[i * gmmSwiglu->tokenLen], _inMMLocal_2[i * gmmSwiglu->tokenLen], scale, gmmSwiglu->tokenLen);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwiglu->tokenLen;
|
||||
mmOutQueue.EnQue(_inMMLocal_2);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
template <typename DTYPE_CS>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::UpdateChannelScale(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 更新perChannel
|
||||
if (unlikely(vecConfig.nextUpadteInterVal == 0)) {
|
||||
int64_t loop = gmmSwiglu->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpadteInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
LocalTensor<float> _inChannel = perChannelScaleInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwiglu->tokenLen * sizeof(DTYPE_CS)), 0, 0, 0};
|
||||
DataCopyPadExtParams<DTYPE_CS> padParams{false, 0, 0, 0};
|
||||
if constexpr (!IsSameType<DTYPE_CS, float>::value) {
|
||||
LocalTensor<DTYPE_CS> dstLocalT = _inChannel.template ReinterpretCast<DTYPE_CS>();
|
||||
DataCopyPad(dstLocalT[gmmSwiglu->tokenLen], perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen],
|
||||
copyParams, padParams);
|
||||
int32_t eventIdMTE2ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
|
||||
SetFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
WaitFlag<HardEvent::MTE2_V>(eventIdMTE2ToV);
|
||||
Cast(_inChannel, dstLocalT[gmmSwiglu->tokenLen], RoundMode::CAST_NONE, gmmSwiglu->tokenLen);
|
||||
} else {
|
||||
DataCopyPad(_inChannel, perChannelScaleGM[vecConfig.curGroupIdx * gmmSwiglu->tokenLen], copyParams,
|
||||
padParams);
|
||||
}
|
||||
perChannelScaleInQueue.EnQue(_inChannel);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::VectorCompute(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
Dequant(loopIdx, vecConfig);
|
||||
Swiglu(loopIdx, vecConfig);
|
||||
Quant(loopIdx, vecConfig);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Dequant(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// perChanelScale * perTokenScale
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> perChannelLocal = perChannelScaleInQueue.DeQue<float>();
|
||||
Mul(mmLocal[loopIdx * gmmSwiglu->tokenLen], mmLocal[loopIdx * gmmSwiglu->tokenLen], perChannelLocal,
|
||||
gmmSwiglu->tokenLen);
|
||||
vecConfig.nextUpadteInterVal--;
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
perChannelScaleInQueue.EnQue(perChannelLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Swiglu(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
_inMMLocal[loopIdx * gmmSwiglu->tokenLen + gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = _inMMLocal[loopIdx * gmmSwiglu->tokenLen];
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwiglu->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{1, static_cast<uint16_t>((gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0,
|
||||
0};
|
||||
DataCopy(_inMMLocal[loopIdx * gmmSwiglu->tokenLen], workspaceLocal, repeatParams);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::Quant(uint32_t loopIdx,
|
||||
VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
uint64_t preOffset = loopIdx * gmmSwiglu->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwiglu->tokenLen / BISECT;
|
||||
Abs(_inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], _inMMLocal[preOffset], halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// reduceMax
|
||||
LocalTensor<float> workLocal = reduceWorkspace.Get<float>(halfTokenLen);
|
||||
LocalTensor<float> reduceResLocal =
|
||||
reduceWorkspace.GetWithOffset<float>(FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float));
|
||||
LocalTensor<float> reduceTmpLocal = reduceWorkspace.GetWithOffset<float>(
|
||||
FLOAT_UB_BLOCK_UNIT_SIZE, halfTokenLen * sizeof(float) + UB_BLOCK_UNIT_SIZE);
|
||||
ReduceMaxTemplate(reduceResLocal, workLocal, _inMMLocal[preOffset + gmmSwiglu->tokenLen / BISECT], reduceTmpLocal,
|
||||
static_cast<uint32_t>(halfTokenLen));
|
||||
|
||||
int32_t eventIdVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
|
||||
SetFlag<HardEvent::V_S>(eventIdVToS);
|
||||
WaitFlag<HardEvent::V_S>(eventIdVToS);
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = 1 / quantScale;
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(_inMMLocal[preOffset], _inMMLocal[preOffset], quantScale, halfTokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
int32_t dstTempOffset = static_cast<int32_t>(preOffset / BISECT);
|
||||
int32_t srcTempOffset = static_cast<int32_t>(preOffset);
|
||||
int32_t tempCount = static_cast<int32_t>(halfTokenLen);
|
||||
LocalTensor<int8_t> castSpace = reduceWorkspace.Get<int8_t>(UB_BLOCK_UNIT_SIZE);
|
||||
CastFp32ToInt8Template(quantLocal, _inMMLocal, castSpace, dstTempOffset, srcTempOffset, tempCount);
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
}
|
||||
|
||||
template <typename mmType, bool sync, typename CHANNELDTYPE>
|
||||
__aicore__ inline void
|
||||
GMMSwigluSplitWorkSpaceCompute<mmType, sync, CHANNELDTYPE>::customDataCopyOut(VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
DataCopyParams copyParams_0{1, (uint16_t)(vecConfig.innerLoopNum * sizeof(float)), 0, 0};
|
||||
DataCopyPad(quantScaleOutputGM[workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx], quantScaleLocal,
|
||||
copyParams_0);
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
DataCopyParams copyParams_1{
|
||||
1, (uint16_t)(vecConfig.innerLoopNum * gmmSwiglu->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0, 0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) * gmmSwiglu->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwiglu->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
#endif // ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_WS_H
|
||||
@@ -0,0 +1,332 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file grouped_matmul_swiglu_quant_utils.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H
|
||||
#define ASCENDC_GROUPED_MATMUL_SWIGLU_QUANT_UTILS_H
|
||||
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "kernel_operator.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
|
||||
#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220
|
||||
// A8W4 MSD场景
|
||||
#if defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \
|
||||
defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4
|
||||
#define GMM_SWIGLU_QUANT_A8W4_MSD
|
||||
using DTYPE_X_A8W4_MSD = AscendC::int4b_t;
|
||||
// A8W8 场景
|
||||
#elif defined(ORIG_DTYPE_X) && defined(DT_INT8) && ORIG_DTYPE_X == DT_INT8 && defined(ORIG_DTYPE_WEIGHT) && \
|
||||
defined(DT_INT8) && ORIG_DTYPE_WEIGHT == DT_INT8
|
||||
#define GMM_SWIGLU_QUANT_A8W8
|
||||
#endif // 场景分类
|
||||
|
||||
#if defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_FRACTAL_NZ
|
||||
constexpr CubeFormat wFormat = CubeFormat::NZ;
|
||||
#elif defined(FORMAT_WEIGHT) && FORMAT_WEIGHT == FORMAT_ND
|
||||
constexpr CubeFormat wFormat = CubeFormat::ND;
|
||||
#endif // weight格式分类
|
||||
|
||||
#endif // 芯片型号分类
|
||||
|
||||
namespace GROUPED_MATMUL_SWIGLU_QUANT {
|
||||
using namespace AscendC;
|
||||
constexpr uint32_t INT8_BITS = 8; // a int8 number has 8 bits
|
||||
constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32; // 32: a block has 32 bytes data
|
||||
constexpr uint32_t THRESHOLD_BLOCK_NUM = 8;
|
||||
constexpr uint32_t UB_BLOCK_DOUBLE_UNIT_SIZE = 64; // 64: a block has 64 bytes data
|
||||
constexpr uint32_t HALF_UB_BLOCK_UNIT_SIZE = UB_BLOCK_UNIT_SIZE / 2; // 2: a float16 data has two bytes
|
||||
constexpr uint32_t FLOAT_UB_BLOCK_UNIT_SIZE = 8; // 2: a float16 data has two bytes
|
||||
constexpr uint32_t SINGLE_CORE_M = 128;
|
||||
constexpr uint32_t SINGLE_CORE_N = 256;
|
||||
constexpr uint32_t SINGLE_CORE_K = 7168;
|
||||
constexpr uint32_t BASIC_M = 128;
|
||||
constexpr uint32_t BASIC_N = 256;
|
||||
constexpr uint32_t BASIC_K = 128;
|
||||
constexpr uint32_t STEP_M = 1;
|
||||
constexpr uint32_t STEP_N = 1;
|
||||
constexpr uint32_t STEP_Ka = 4;
|
||||
constexpr uint32_t STEP_Kb = 4;
|
||||
constexpr uint32_t DEPTH_A1 = 8;
|
||||
constexpr uint32_t DEPTH_B1 = 8;
|
||||
constexpr uint32_t VEC_LEN_ONCE_REPEAT_ELE = 64;
|
||||
constexpr uint32_t VEC_LEN_ONCE_REPEAT_BLOCK = 8;
|
||||
constexpr uint32_t FP32_LEN_64_REPEAT = 4096;
|
||||
constexpr uint32_t REPEAT_64 = 64;
|
||||
constexpr uint32_t REPEAT_8 = 8;
|
||||
constexpr uint32_t BISECT = 2;
|
||||
constexpr uint32_t MOD_32_MASK = 0x1F;
|
||||
constexpr uint32_t MOD_16_MASK = 0x0F;
|
||||
constexpr uint32_t ALIGN_8_ELE = 8;
|
||||
constexpr uint32_t ALIGN_16_ELE = 16;
|
||||
constexpr float QUANT_SCALE_INT8 = 127.0f;
|
||||
constexpr int64_t SWIGLU_REDUCE_FACTOR = 2;
|
||||
constexpr int64_t DOUBLE_BUFFER = 2;
|
||||
constexpr uint8_t NUM_8 = 8;
|
||||
constexpr bool NO_BIAS = false;
|
||||
constexpr int64_t DOUBLE_ROW = 2;
|
||||
constexpr MatmulConfig CUSTOM_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true);
|
||||
constexpr MatmulConfig GetMMStaticCFG()
|
||||
{
|
||||
MatmulConfig MM_CFG = CUSTOM_CFG_MDL;
|
||||
MM_CFG.singleCoreM = SINGLE_CORE_M;
|
||||
MM_CFG.singleCoreN = SINGLE_CORE_N;
|
||||
MM_CFG.singleCoreK = SINGLE_CORE_K;
|
||||
MM_CFG.basicM = BASIC_M;
|
||||
MM_CFG.basicN = BASIC_N;
|
||||
MM_CFG.basicK = BASIC_K;
|
||||
return MM_CFG;
|
||||
}
|
||||
|
||||
constexpr static MatmulApiStaticTiling GetMMTiling(const MatmulApiStaticTiling &mmTiling)
|
||||
{
|
||||
MatmulApiStaticTiling tiling = mmTiling;
|
||||
tiling.stepM = STEP_M;
|
||||
tiling.stepN = STEP_N;
|
||||
tiling.stepKa = STEP_Ka;
|
||||
tiling.stepKb = STEP_Kb;
|
||||
tiling.depthA1 = DEPTH_A1;
|
||||
tiling.depthB1 = DEPTH_B1;
|
||||
tiling.isBias = NO_BIAS;
|
||||
return tiling;
|
||||
}
|
||||
|
||||
template <class AT_, class BT_, class CT_>
|
||||
struct MMImplTypeStatic {
|
||||
using AT = AT_;
|
||||
using BT = BT_;
|
||||
using CT = CT_;
|
||||
// bias未被使用但高阶模板参数需要传入
|
||||
using BiasT = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
static constexpr MatmulConfig cfg = GetMMStaticCFG();
|
||||
static constexpr MatmulApiStaticTiling mdl = GetMMTiling(GetMatmulApiTiling<AT, BT, CT, BiasT>(cfg));
|
||||
using MT = matmul::MatmulImpl<AT, BT, CT, BiasT, mdl>;
|
||||
};
|
||||
|
||||
template <class AT_, class BT_, class CT_>
|
||||
struct MMImplType {
|
||||
using AT = AT_;
|
||||
using BT = BT_;
|
||||
using CT = CT_;
|
||||
// bias未被使用但高阶模板参数需要传入
|
||||
using BiasT = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
using MT = matmul::MatmulImpl<AT, BT, CT, BiasT, CUSTOM_CFG_MDL>;
|
||||
};
|
||||
|
||||
struct MNConfig {
|
||||
int64_t m = 0;
|
||||
int64_t k = 0;
|
||||
int64_t n = 0;
|
||||
int64_t baseM = 0;
|
||||
int64_t baseN = 0;
|
||||
int64_t mIdx = 0;
|
||||
int64_t nIdx = 0;
|
||||
int64_t blockDimM = 0;
|
||||
int64_t blockDimN = 0;
|
||||
int64_t singleM = 0;
|
||||
int64_t singleN = 0;
|
||||
int64_t wBaseOffset = 0;
|
||||
int64_t nAxisBaseOffset = 0;
|
||||
int64_t mAxisBaseOffset = 0;
|
||||
int64_t xBaseOffset = 0;
|
||||
int64_t yBaseOffset = 0;
|
||||
int64_t wOutOffset = 0;
|
||||
int64_t workSpaceOffset = 0;
|
||||
};
|
||||
|
||||
struct VecConfig {
|
||||
int64_t M = 0;
|
||||
int64_t usedCoreNum = 0;
|
||||
int64_t startOffset = 0;
|
||||
int64_t curOffset = 0;
|
||||
int64_t startIdx = 0;
|
||||
int64_t curIdx = 0;
|
||||
int64_t taskNum = 0;
|
||||
int64_t curGroupIdx = 0;
|
||||
int64_t outLoopNum = 0;
|
||||
int64_t innerLoopNum = 0;
|
||||
int64_t tailLoopNum = 0;
|
||||
int64_t nextUpadteInterVal = 0;
|
||||
};
|
||||
|
||||
struct WorkSpaceSplitConfig {
|
||||
int64_t M = 0;
|
||||
int64_t loopCount = 0;
|
||||
int64_t leftMatrixStartIndex = 0;
|
||||
int64_t rightMatrixExpertStartIndex = 0;
|
||||
int64_t rightMatrixExpertNextStartIndex = 0;
|
||||
int64_t rightMatrixExpertEndIndex = 0;
|
||||
int64_t notLastTaskSize = 0;
|
||||
int64_t lastLoopTaskSize = 0;
|
||||
bool isLastLoop = false;
|
||||
};
|
||||
|
||||
struct GMAddrParams {
|
||||
// 输入 GM Tensor
|
||||
GM_ADDR xGM; // 左矩阵
|
||||
GM_ADDR weightGM; // 右矩阵
|
||||
GM_ADDR weightScaleGM; // 权重scale
|
||||
GM_ADDR xScaleGM; // 激活scale
|
||||
GM_ADDR weightAuxiliaryMatrixGM; // 权重辅助矩阵
|
||||
GM_ADDR groupListGM; // 分组矩阵
|
||||
// 输出 GM Tensor
|
||||
GM_ADDR yGM; // 输出量化矩阵
|
||||
GM_ADDR yScaleGM; // 输出scale矩阵
|
||||
// workspace GM Tensor
|
||||
GM_ADDR workSpaceGM; // 左矩阵前处理结果矩阵 (double workspace) + 中间处理结果矩阵 (double workspace)
|
||||
int64_t workSpaceOffset1;
|
||||
int64_t workSpaceOffset2;
|
||||
int64_t workSpaceOffset3;
|
||||
};
|
||||
|
||||
template <uint32_t base, typename T = uint32_t>
|
||||
__aicore__ inline auto AlignUp(T a) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return (a + base - 1) / base * base;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline auto AlignUp(T a, T base) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return (a + base - 1) / base * base;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline auto AlignDown(T a, T base) -> T
|
||||
{
|
||||
if (unlikely(base == 0)) {
|
||||
return a;
|
||||
}
|
||||
return a / base * base;
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<4, uint32_t>(uint32_t a)
|
||||
{
|
||||
// to be Multiple of 4, result should be in a format of b(xxxx,x100).
|
||||
// This means last two bits should be zero, requiring that
|
||||
// result = num & b(1111,1100) = num & (~3).
|
||||
// &(~3) operator may reduces num into the range [num, num - 3].
|
||||
// As the result should be no less than a (result >= a), it means num - 3 >= a in the worst case.
|
||||
// In this case, num >= a+3. On the other hand, num should also be less then a+4, otherwise,
|
||||
// the result will not be least multiple of 4 for 3. In other cases like [num, num - 2],
|
||||
// num = a + 3 also satisfies the goal condition.
|
||||
return (a + 3) & ~3; // & ~3: set last two bits of (a+3) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<8, uint32_t>(uint32_t a)
|
||||
{
|
||||
// In general, if we want to get the least multiple of b (b is the power of 2) for a,
|
||||
// it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b)
|
||||
return (a + 7) & ~7; // & ~7: set last four bits of (a+7) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<16, uint32_t>(uint32_t a)
|
||||
{
|
||||
// In general, if we want to get the least multiple of b (b is the power of 2) for a,
|
||||
// it comes to a conclusion from the above comment: result = (a + (b - 1)) & (~b)
|
||||
return (a + 15) & ~15; // & ~15: set last four bits of (a+15) to be zero
|
||||
}
|
||||
|
||||
template <>
|
||||
__aicore__ inline uint32_t AlignUp<32, uint32_t>(uint32_t a)
|
||||
{
|
||||
// refer to the above comments.
|
||||
return (a + 31) & ~31; // & ~31: set last five bits of (a+31) to be zero}
|
||||
}
|
||||
|
||||
__aicore__ inline void ReduceMaxSmall(const LocalTensor<float> &dstLocal, const LocalTensor<float> &workLocal,
|
||||
const LocalTensor<float> &srcLocal, uint32_t count)
|
||||
{
|
||||
/**
|
||||
* @brief ReduceMaxSmall 此函数仅支持入参count小于4096。
|
||||
*/
|
||||
uint32_t repeat = count / VEC_LEN_ONCE_REPEAT_ELE;
|
||||
uint32_t tailNum = count % VEC_LEN_ONCE_REPEAT_ELE;
|
||||
if (likely(repeat > 0)) {
|
||||
WholeReduceMax(workLocal, srcLocal, VEC_LEN_ONCE_REPEAT_ELE, repeat, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (unlikely(tailNum != 0)) {
|
||||
WholeReduceMax(workLocal[repeat], srcLocal[count - tailNum], tailNum, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
repeat += 1;
|
||||
}
|
||||
WholeReduceMax(dstLocal, workLocal, repeat, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK, ReduceOrder::ORDER_ONLY_VALUE);
|
||||
}
|
||||
|
||||
__aicore__ inline void ReduceMaxTemplate(const LocalTensor<float> &dstLocal, const LocalTensor<float> &workLocal,
|
||||
const LocalTensor<float> &srcLocal, const LocalTensor<float> &resTmpLocal,
|
||||
uint32_t count)
|
||||
{
|
||||
/**
|
||||
* @brief 当前算子仅支持[32, 10240]长度的词向量维度N，对应此函数count入参范围在[16, 5120]。
|
||||
* @param [in] count: 本函数支持count范围为[1,8192]。
|
||||
*/
|
||||
if (count <= FP32_LEN_64_REPEAT) {
|
||||
ReduceMaxSmall(dstLocal, workLocal, srcLocal, count);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else {
|
||||
BlockReduceMax(workLocal, srcLocal, REPEAT_64, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
BlockReduceMax(workLocal, workLocal, REPEAT_8, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
WholeReduceMax(resTmpLocal, workLocal, VEC_LEN_ONCE_REPEAT_ELE, 1, 1, 1, VEC_LEN_ONCE_REPEAT_BLOCK,
|
||||
ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
ReduceMaxSmall(dstLocal, workLocal, srcLocal[FP32_LEN_64_REPEAT], count - FP32_LEN_64_REPEAT);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
const BinaryRepeatParams repeatParams = {1, 1, 1, NUM_8, NUM_8, NUM_8};
|
||||
Max(dstLocal, dstLocal, resTmpLocal, 1, 1, repeatParams);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void CastFp32ToInt8Template(LocalTensor<int8_t> &dstLocal, LocalTensor<float> &srcLocal,
|
||||
LocalTensor<int8_t> &oneBlockWorkspace, int32_t dstOffset,
|
||||
int32_t srcOffset, int32_t count)
|
||||
{
|
||||
Cast(srcLocal[srcOffset].ReinterpretCast<half>(), srcLocal[srcOffset], RoundMode::CAST_RINT, count);
|
||||
PipeBarrier<PIPE_V>();
|
||||
if ((dstOffset & MOD_32_MASK) == 0) {
|
||||
Cast(dstLocal[dstOffset], srcLocal[srcOffset].ReinterpretCast<half>(), RoundMode::CAST_RINT, count);
|
||||
} else if ((dstOffset & MOD_16_MASK) == 0) {
|
||||
Cast(dstLocal[dstOffset + ALIGN_16_ELE], srcLocal[srcOffset + ALIGN_8_ELE].ReinterpretCast<half>(),
|
||||
RoundMode::CAST_RINT, count - ALIGN_16_ELE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(oneBlockWorkspace, srcLocal[srcOffset].ReinterpretCast<half>(), RoundMode::CAST_RINT, ALIGN_16_ELE);
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
for (int32_t i = 0; i < ALIGN_16_ELE; i++) {
|
||||
int8_t temp = oneBlockWorkspace.GetValue(i);
|
||||
dstLocal.SetValue(dstOffset + i, temp);
|
||||
}
|
||||
PipeBarrier<PIPE_ALL>();
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GROUPED_MATMUL_SWIGLU_QUANT
|
||||
|
||||
#endif // ASCENDC_GROUPED_MATMUL_UTILS_H
|
||||
Reference in New Issue
Block a user