@@ -0,0 +1,76 @@
|
||||
/**
|
||||
* 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_v2_mxquant.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef GROUPED_MATMUL_SWIGLU_QUANT_V2_MXQUANT_H
|
||||
#define GROUPED_MATMUL_SWIGLU_QUANT_V2_MXQUANT_H
|
||||
|
||||
#include "cgmct/kernel/kernel_gmm_swiglu_mxquant.h"
|
||||
#include "cgmct/block/block_mx_mm_aic_to_aiv_builder.h"
|
||||
#include "cgmct/block/block_scheduler_gmm_aswt_with_tail_split.h"
|
||||
|
||||
using namespace Cgmct::Gemm;
|
||||
using namespace Cgmct::Gemm::Kernel;
|
||||
|
||||
template <typename layoutA, typename layoutB>
|
||||
__aicore__ inline void GmmSwigluAswt(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale, GM_ADDR groupList,
|
||||
GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling); \
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, mmTilingData, mmTilingData_, tiling); \
|
||||
// 定义L1和L0的TileShape
|
||||
using L1TileShape = AscendC::Shape<_0, _0, _0>;
|
||||
using L0TileShape = AscendC::Shape<_0, _0, _0>;
|
||||
// 定义矩阵的类型和布局
|
||||
using AType = DTYPE_X;
|
||||
using BType = DTYPE_WEIGHT;
|
||||
using CType = DTYPE_Y;
|
||||
using LayoutA = layoutA;
|
||||
using LayoutB = layoutB;
|
||||
using LayoutC = layout::RowMajorAlign;
|
||||
using weightscaleType = AscendC::fp8_e8m0_t;
|
||||
using BiasType = float;
|
||||
// 定义scheduler类型
|
||||
using BlockScheduler = GroupedMatmulAswtWithTailSplitScheduler;
|
||||
// 定义MMAD类型
|
||||
using C1Type = float;
|
||||
// 定义BlockEpilogue类型
|
||||
using BlockEpilogue = Block::BlockEpilogueSwigluQuant<L0TileShape, CType, C1Type, weightscaleType, weightscaleType,
|
||||
true>;
|
||||
// 定义shape的形状，tuple保存 m n k batch
|
||||
using ProblemShape = MatmulShape;
|
||||
using BlockMmad = Block::BlockMxMmAicToAivBuilder<AType, LayoutA, BType, LayoutB, BiasType, C1Type, LayoutC, L1TileShape,
|
||||
L0TileShape, BlockScheduler, QuantMatmulWithTileMultiBlock<>,
|
||||
Tile::TileCopy<Arch::DAV3510, Tile::CopyInAndCopyOutSplitMWithParams>>;
|
||||
using QGmmKernel =
|
||||
Kernel::KernelGmmSwiGluMixOnlineDynamic<ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler>;
|
||||
using Params = typename QGmmKernel::Params;
|
||||
using GMMTiling = typename QGmmKernel::GMMTiling;
|
||||
GMMTiling gmmParams{gmmSwigluQuantParams_.groupNum, gmmSwigluQuantParams_.groupListType, mmTilingData_.baseM,
|
||||
mmTilingData_.baseN, mmTilingData_.baseK};
|
||||
gmmParams.matmulTiling = &mmTilingData_;
|
||||
Params params = {// template shape, gmm shape can not get now
|
||||
{1, 1, 1, 1},
|
||||
// mmad args
|
||||
{x, weight, weightScale, xScale, y, groupList},
|
||||
{y, yScale, nullptr, nullptr, nullptr, static_cast<uint32_t>(mmTilingData_.baseM),
|
||||
static_cast<uint32_t>(mmTilingData_.baseN)},
|
||||
// gmm tiling data
|
||||
gmmParams};
|
||||
QGmmKernel op;
|
||||
op(params);
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,115 @@
|
||||
/**
|
||||
* Copyright (c) 2026 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_v2_pertoken_quant.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef GROUPED_MATMUL_SWIGLU_QUANT_V2_PERTOKEN_QUANT_H
|
||||
#define GROUPED_MATMUL_SWIGLU_QUANT_V2_PERTOKEN_QUANT_H
|
||||
|
||||
#include "cgmct/kernel/kernel_gmm_swiglu_pertoken_quant.h"
|
||||
#include "cgmct/block/block_mmad_builder.h"
|
||||
#include "cgmct/block/block_scheduler_gmm_aswt_with_tail_split.h"
|
||||
|
||||
using namespace Cgmct::Gemm;
|
||||
using namespace Cgmct::Gemm::Kernel;
|
||||
|
||||
static constexpr uint8_t BF16_VALUE = 27;
|
||||
|
||||
template <uint8_t dequantDtype, typename layoutA, typename layoutB>
|
||||
__aicore__ inline void GmmSwigluAswtPertokenKernel(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale,
|
||||
GM_ADDR groupList, GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace,
|
||||
GM_ADDR tiling, TPipe *pipe)
|
||||
{
|
||||
/* 1. 取 tiling 数据 */
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, mmTilingData, mmTilingData_, tiling);
|
||||
|
||||
/* 2. 编译期常量决定 DequantType / C1Type */
|
||||
using DequantType =
|
||||
std::conditional_t<dequantDtype == 1, half, std::conditional_t<dequantDtype == BF16_VALUE, bfloat16_t, float>>;
|
||||
using AType = DTYPE_X;
|
||||
using BType = DTYPE_WEIGHT;
|
||||
using CType = DTYPE_Y; // y dtype
|
||||
using C1Type = std::conditional_t<std::is_same_v<AType, int8_t>, int32_t, float>; // matmul output dtype
|
||||
|
||||
/* 3. 其余别名 */
|
||||
using L0TileShape = AscendC::Shape<_0, _0, _0>;
|
||||
using L1TileShape = AscendC::Shape<_0, _0, _0>;
|
||||
using LayoutA = layoutA;
|
||||
using LayoutB = layoutB;
|
||||
using LayoutC = layout::RowMajorAlign;
|
||||
using weightscaleType = DTYPE_WEIGHT_SCALE;
|
||||
using xscaleType = float;
|
||||
using BiasType = float;
|
||||
using BlockScheduler = GroupedMatmulAswtWithTailSplitScheduler;
|
||||
using BlockEpilogueDequantAndSwiglu =
|
||||
Block::BlockEpilogueDequantSwiglu<L0TileShape, DequantType, C1Type, weightscaleType, xscaleType, true>;
|
||||
using BlockEpiloguePertokenQuant = Block::BlockEpiloguePertokenQuant<DequantType, CType>;
|
||||
using ProblemShape = MatmulShape;
|
||||
using BlockMmad =
|
||||
Block::BlockMmadBuilder<AType, LayoutA, BType, LayoutB, C1Type, LayoutC, BiasType, layout::RowMajor,
|
||||
L1TileShape, L0TileShape, BlockScheduler, MatmulMultiBlock<>,
|
||||
Tile::TileCopy<Arch::DAV3510, Tile::CopyInAndCopyOutSplitMWithParams>>;
|
||||
using QGmmKernel =
|
||||
Kernel::KernelGmmSwiGluPertokenQuant<ProblemShape, BlockMmad, BlockEpilogueDequantAndSwiglu,
|
||||
BlockEpiloguePertokenQuant, BlockScheduler, weightscaleType, xscaleType>;
|
||||
|
||||
/* 4. 拼参数、launch */
|
||||
using Params = typename QGmmKernel::Params;
|
||||
using GMMTiling = typename QGmmKernel::GMMTiling;
|
||||
GMMTiling gmmParams{gmmSwigluQuantParams_.groupNum, gmmSwigluQuantParams_.groupListType, mmTilingData_.baseM,
|
||||
mmTilingData_.baseN, mmTilingData_.baseK};
|
||||
gmmParams.matmulTiling = &mmTilingData_;
|
||||
Params params = {
|
||||
{1, 1, 1, 1},
|
||||
// mmad args
|
||||
{x, weight, y, nullptr, groupList},
|
||||
{workspace, weightScale, xScale, static_cast<uint32_t>(mmTilingData_.baseM),
|
||||
static_cast<uint32_t>(mmTilingData_.baseN)},
|
||||
{workspace, smoothScale, y, yScale, gmmSwigluQuantParams_.rowLen, gmmSwigluQuantParams_.ubAvail, false},
|
||||
// gmm tiling data
|
||||
gmmParams};
|
||||
QGmmKernel op(pipe);
|
||||
op(params);
|
||||
}
|
||||
|
||||
/* ----------------------------------------------------------
|
||||
* 5. 最外层入口：只做 switch，把运行期值 → 编译期常量
|
||||
* ---------------------------------------------------------- */
|
||||
template <typename layoutA, typename layoutB>
|
||||
__aicore__ inline void GmmSwigluAswtPertoken(GM_ADDR x, GM_ADDR weight, GM_ADDR weightScale, GM_ADDR xScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR smoothScale, GM_ADDR groupList,
|
||||
GM_ADDR y, GM_ADDR yScale, GM_ADDR workspace, GM_ADDR tiling, TPipe *pipe)
|
||||
{
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantTilingDataParams, gmmSwigluQuantParams, gmmSwigluQuantParams_, tiling);
|
||||
|
||||
switch (gmmSwigluQuantParams_.dequantDtype) {
|
||||
case 1:
|
||||
GmmSwigluAswtPertokenKernel<1, layoutA, layoutB>(x, weight, weightScale, xScale, weightAssistanceMatrix,
|
||||
smoothScale, groupList, y, yScale, workspace, tiling,
|
||||
pipe);
|
||||
break;
|
||||
case BF16_VALUE:
|
||||
GmmSwigluAswtPertokenKernel<BF16_VALUE, layoutA, layoutB>(x, weight, weightScale, xScale,
|
||||
weightAssistanceMatrix, smoothScale, groupList, y,
|
||||
yScale, workspace, tiling, pipe);
|
||||
break;
|
||||
default:
|
||||
GmmSwigluAswtPertokenKernel<0, layoutA, layoutB>(x, weight, weightScale, xScale, weightAssistanceMatrix,
|
||||
smoothScale, groupList, y, yScale, workspace, tiling,
|
||||
pipe);
|
||||
break;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,40 @@
|
||||
/**
|
||||
* Copyright (c) 2025-2026 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_v2_tiling_key.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef __OP_KERNEL_GMM_SWIGLU_QUANT_V2_TILING_KEY_H__
|
||||
#define __OP_KERNEL_GMM_SWIGLU_QUANT_V2_TILING_KEY_H__
|
||||
|
||||
#include "ascendc/host_api/tiling/template_argument.h"
|
||||
|
||||
#define GMM_SWIGLU_QUANT_NO_TRANS 0
|
||||
#define GMM_SWIGLU_QUANT_TRANS 1
|
||||
|
||||
// 模板参数
|
||||
ASCENDC_TPL_ARGS_DECL(GroupedMatmulSwigluQuantV2, // 算子OpType
|
||||
ASCENDC_TPL_UINT_DECL(QUANT_B_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST,
|
||||
GMM_SWIGLU_QUANT_NO_TRANS, GMM_SWIGLU_QUANT_TRANS),
|
||||
ASCENDC_TPL_UINT_DECL(QUANT_A_TRANS, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST,
|
||||
GMM_SWIGLU_QUANT_NO_TRANS, GMM_SWIGLU_QUANT_TRANS));
|
||||
|
||||
// 模板参数组合
|
||||
// 用于调用GET_TPL_TILING_KEY获取TilingKey时，接口内部校验TilingKey是否合法
|
||||
ASCENDC_TPL_SEL(
|
||||
ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
|
||||
ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS),
|
||||
ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS)),
|
||||
ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
|
||||
ASCENDC_TPL_UINT_SEL(QUANT_B_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_TRANS),
|
||||
ASCENDC_TPL_UINT_SEL(QUANT_A_TRANS, ASCENDC_TPL_UI_LIST, GMM_SWIGLU_QUANT_NO_TRANS)));
|
||||
#endif
|
||||
@@ -0,0 +1,591 @@
|
||||
/**
|
||||
* 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 grrouped_matmul_swiglu_quant_spilit_fusion.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_FUSION_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_SPLIT_FUSION_H
|
||||
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
using namespace AscendC;
|
||||
constexpr int64_t BLOCK_SIZE = 32;
|
||||
constexpr int64_t BLOCK_ELEM = BLOCK_SIZE / sizeof(float);
|
||||
constexpr int64_t SWI_FACTOR = 2;
|
||||
constexpr float DYNAMIC_QUANT_FACTOR = 1.0 / static_cast<float>(127.0);
|
||||
constexpr uint64_t MAX_CALC_NUM = 64;
|
||||
constexpr uint64_t REDUCEMAX_CALC_NUM = 64;
|
||||
constexpr uint64_t SPILI_NUM = 2;
|
||||
constexpr uint64_t VC_SYNC_MAX_TIMES = 14;
|
||||
constexpr uint64_t RESRERVE_MEM_SIZE = 192;
|
||||
|
||||
class GroupedMatmulDequantSwigluQuantFusion {
|
||||
public:
|
||||
using aType = MatmulType<TPosition::GM, CubeFormat::ND, int8_t>;
|
||||
using bType = MatmulType<TPosition::GM, CubeFormat::NZ, int8_t>;
|
||||
using cType = MatmulType<TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
using biasType = MatmulType<TPosition::GM, CubeFormat::ND, int32_t>;
|
||||
using matmulType = MMImplType<aType, bType, cType, biasType, matmulCFGUnitFlag>;
|
||||
matmulType::MT mm;
|
||||
|
||||
__aicore__ inline GroupedMatmulDequantSwigluQuantFusion(
|
||||
TPipe* pipe, const GMMSwigluQuantV2TilingFusionData* __restrict tiling,
|
||||
const TCubeTiling* __restrict matmulTilingData)
|
||||
: pipe_(pipe), tilingData_(tiling), matmulTilingData_(matmulTilingData) {
|
||||
}
|
||||
|
||||
__aicore__ inline int CeilDiv(int a, int b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR weight_scale, GM_ADDR activation_scale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR group_list,
|
||||
GM_ADDR y, GM_ADDR scale, GM_ADDR workspace) {
|
||||
xGm_.SetGlobalBuffer((__gm__ int8_t*)x);
|
||||
groupListGm_.SetGlobalBuffer((__gm__ int64_t*)group_list);
|
||||
weightGm_.SetGlobalBuffer(GetTensorAddr<int8_t>(0, weight));
|
||||
weightScaleGm_.SetGlobalBuffer(GetTensorAddr<float>(0, weight_scale));
|
||||
workspaceGm_.SetGlobalBuffer((__gm__ int32_t*)workspace);
|
||||
activateScaleGm_.SetGlobalBuffer((__gm__ float*)activation_scale);
|
||||
scaleGm_.SetGlobalBuffer((__gm__ float*)scale);
|
||||
yGm_.SetGlobalBuffer((__gm__ int8_t*)y);
|
||||
weightScaleTensorPtr_ = weight_scale;
|
||||
weightTensorPtr_ = weight;
|
||||
|
||||
nBasicsBlocks = CeilDiv(tilingData_->N, matmulTilingData_->baseN);
|
||||
totalBasicBlocks = 0;
|
||||
for (int groupId = 0; groupId < tilingData_->groupNum; groupId++) {
|
||||
int tokens = groupListGm_.GetValue(groupId);
|
||||
if (tilingData_->groupListType == 0 && groupId > 0) {
|
||||
tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1);
|
||||
}
|
||||
int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM);
|
||||
totalBasicBlocks += mBasicBlocks * nBasicsBlocks;
|
||||
}
|
||||
|
||||
totalSyncTimes = CeilDiv(totalBasicBlocks, tilingData_->cubeBlockDim);
|
||||
if ASCEND_IS_AIV {
|
||||
pipe_->InitBuffer(xActQueue_, 1, (tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * SWI_FACTOR + tilingData_->ubFactorDimx * BLOCK_ELEM) * sizeof(int32_t));
|
||||
pipe_->InitBuffer(inScaleQueue_, 1, ((tilingData_->N / SPILI_NUM) * SWI_FACTOR + (tilingData_->N / SPILI_NUM)) * sizeof(float));
|
||||
pipe_->InitBuffer(outQueue_, 1, tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t) + tilingData_->ubFactorDimx * sizeof(float) + RESRERVE_MEM_SIZE);
|
||||
pipe_->InitBuffer(tmpBuf1_, tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * SWI_FACTOR * sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void FindCurrentGroup(uint32_t basicBlockIdxInGlobal, uint32_t& currentGroupId,
|
||||
uint32_t& globalMOffset, uint32_t& processedBasicBlock) {
|
||||
for (int groupId = currentGroupId; groupId < tilingData_->groupNum; groupId++) {
|
||||
int tokens = groupListGm_.GetValue(groupId);
|
||||
if (tilingData_->groupListType == 0 && groupId > 0) {
|
||||
tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1);
|
||||
}
|
||||
int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM);
|
||||
if (processedBasicBlock + mBasicBlocks * nBasicsBlocks > basicBlockIdxInGlobal) {
|
||||
currentGroupId = groupId;
|
||||
break;
|
||||
} else {
|
||||
globalMOffset += tokens;
|
||||
processedBasicBlock += mBasicBlocks * nBasicsBlocks;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void CalculateBlockSizes(int tokens, int currentBasicBlockMId, int currentBasicBlockNId,
|
||||
int& realMSize, int& realNSize) {
|
||||
realMSize = matmulTilingData_->baseM;
|
||||
if (currentBasicBlockMId * matmulTilingData_->baseM + realMSize > tokens) {
|
||||
realMSize = tokens - currentBasicBlockMId * matmulTilingData_->baseM;
|
||||
}
|
||||
realNSize = matmulTilingData_->baseN;
|
||||
if (currentBasicBlockNId * matmulTilingData_->baseN + realNSize > tilingData_->N) {
|
||||
realNSize = tilingData_->N - currentBasicBlockNId * matmulTilingData_->baseN;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void SetupMatmulShape(int tokens, int realMSize, int realNSize) {
|
||||
mm.SetOrgShape(tokens, tilingData_->N, tilingData_->K);
|
||||
mm.SetSingleShape(realMSize, realNSize, tilingData_->K);
|
||||
}
|
||||
|
||||
__aicore__ inline void SetupMatmulWeight(int currentGroupId, int currentBasicBlockNId) {
|
||||
if (tilingData_->isSingleTensor == 0) {
|
||||
weightGm_.SetGlobalBuffer(GetTensorAddr<int8_t>(currentGroupId, weightTensorPtr_));
|
||||
mm.SetTensorB(weightGm_[0x8 * currentBasicBlockNId * tilingData_->K * 0x20]);
|
||||
} else {
|
||||
int64_t tensorBOffset = currentGroupId * tilingData_->K * tilingData_->N + 0x8 * currentBasicBlockNId * tilingData_->K * 0x20;
|
||||
mm.SetTensorB(weightGm_[tensorBOffset]);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessCubeBlock(uint32_t basicBlockIdxInGlobal, uint32_t& currentGroupId,
|
||||
uint32_t& globalMOffset, uint32_t& processedBasicBlock) {
|
||||
FindCurrentGroup(basicBlockIdxInGlobal, currentGroupId, globalMOffset, processedBasicBlock);
|
||||
int tokens = groupListGm_.GetValue(currentGroupId);
|
||||
if (tilingData_->groupListType == 0 && currentGroupId > 0) {
|
||||
tokens = groupListGm_.GetValue(currentGroupId) - groupListGm_.GetValue(currentGroupId - 1);
|
||||
}
|
||||
int basicBlockIdxInCurrentGroup = basicBlockIdxInGlobal - processedBasicBlock;
|
||||
int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM);
|
||||
int currentBasicBlockMId = basicBlockIdxInCurrentGroup / nBasicsBlocks;
|
||||
int currentBasicBlockNId = basicBlockIdxInCurrentGroup % nBasicsBlocks;
|
||||
int realMSize = 0;
|
||||
int realNSize = 0;
|
||||
CalculateBlockSizes(tokens, currentBasicBlockMId, currentBasicBlockNId, realMSize, realNSize);
|
||||
SetupMatmulShape(tokens, realMSize, realNSize);
|
||||
int64_t tensorAOffset = currentBasicBlockMId * matmulTilingData_->baseM * tilingData_->K + globalMOffset * tilingData_->K;
|
||||
mm.SetTensorA(xGm_[tensorAOffset]);
|
||||
SetupMatmulWeight(currentGroupId, currentBasicBlockNId);
|
||||
int64_t workspaceOffset = globalMOffset * tilingData_->N + currentBasicBlockMId * matmulTilingData_->baseM * tilingData_->N
|
||||
+ currentBasicBlockNId * matmulTilingData_->baseN;
|
||||
mm.template IterateAll<false>(workspaceGm_[workspaceOffset]);
|
||||
}
|
||||
|
||||
__aicore__ inline void FinalizeCubeSync(uint32_t& syncId) {
|
||||
while (syncId < totalSyncTimes) {
|
||||
AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(0x8);
|
||||
syncId += 1;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void CubeProcess() {
|
||||
if ASCEND_IS_AIC {
|
||||
uint32_t currentBlockId = GetBlockIdx();
|
||||
uint32_t rsvBlockNum = 0;
|
||||
uint32_t calcBlockNum = 0;
|
||||
uint32_t cvTimes = 0;
|
||||
uint32_t syncId = 0;
|
||||
uint32_t globalMOffset = 0;
|
||||
uint32_t processedBasicBlock = 0;
|
||||
uint32_t currentGroupId = 0;
|
||||
uint32_t realSyncId = 0;
|
||||
while (currentBlockId < totalBasicBlocks) {
|
||||
cvTimes = CeilDiv(nBasicsBlocks - rsvBlockNum, tilingData_->cubeBlockDim);
|
||||
calcBlockNum += cvTimes * tilingData_->cubeBlockDim;
|
||||
rsvBlockNum = calcBlockNum % nBasicsBlocks;
|
||||
|
||||
for (uint32_t cvId = 0; cvId < cvTimes; cvId++) {
|
||||
uint32_t basicBlockIdxInGlobal = currentBlockId;
|
||||
if (basicBlockIdxInGlobal >= totalBasicBlocks) {
|
||||
break;
|
||||
}
|
||||
ProcessCubeBlock(basicBlockIdxInGlobal, currentGroupId, globalMOffset, processedBasicBlock);
|
||||
currentBlockId += tilingData_->cubeBlockDim;
|
||||
syncId += 1;
|
||||
}
|
||||
AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(0x8);
|
||||
|
||||
realSyncId += 1;
|
||||
if (realSyncId > 0 && realSyncId % VC_SYNC_MAX_TIMES == 0) {
|
||||
AscendC::CrossCoreWaitFlag(0x9);
|
||||
}
|
||||
}
|
||||
FinalizeCubeSync(syncId);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void CalculateEndGroupInfo(int endBasicBlockId, int endGroupId, int& endGroupMOffset,
|
||||
int& basicBlockCountBeforeEndGroup) {
|
||||
endGroupMOffset = 0;
|
||||
basicBlockCountBeforeEndGroup = 0;
|
||||
for (int gId = 0; gId < endGroupId; gId++) {
|
||||
int tokens = groupListGm_.GetValue(gId);
|
||||
if (tilingData_->groupListType == 0 && gId > 0) {
|
||||
tokens = groupListGm_.GetValue(gId) - groupListGm_.GetValue(gId - 1);
|
||||
}
|
||||
int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM);
|
||||
basicBlockCountBeforeEndGroup += mBasicBlocks * nBasicsBlocks;
|
||||
endGroupMOffset += tokens;
|
||||
}
|
||||
int basicBlockIdxInCurrentGroup = endBasicBlockId - basicBlockCountBeforeEndGroup;
|
||||
int currentBasicBlockMId = basicBlockIdxInCurrentGroup / nBasicsBlocks;
|
||||
endGroupMOffset += currentBasicBlockMId * matmulTilingData_->baseM;
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessGroupRange(int startGroupId, int endGroupId, int endGroupMOffset,
|
||||
uint32_t& globalMOffset, bool &isSyncAll) {
|
||||
int currentGroupMOffset = 0;
|
||||
for (int gId = 0; gId < startGroupId; gId++) {
|
||||
currentGroupMOffset += groupListGm_.GetValue(gId);
|
||||
}
|
||||
for (int groupId = startGroupId; groupId <= endGroupId; groupId++) {
|
||||
if (tilingData_->groupListType == 0 && groupId > 0) {
|
||||
currentGroupMOffset = groupListGm_.GetValue(groupId);
|
||||
} else {
|
||||
currentGroupMOffset += groupListGm_.GetValue(groupId);
|
||||
}
|
||||
int calcCount = 0;
|
||||
if (currentGroupMOffset <= endGroupMOffset) {
|
||||
calcCount = currentGroupMOffset - globalMOffset;
|
||||
} else {
|
||||
calcCount = endGroupMOffset - globalMOffset;
|
||||
}
|
||||
ProcessDSQ(groupId, globalMOffset, calcCount, isSyncAll);
|
||||
globalMOffset += calcCount;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessVectorBlock(uint32_t syncId, bool& isSyncAll, uint32_t& globalMOffset) {
|
||||
int startBasicBlockId = syncId * tilingData_->cubeBlockDim;
|
||||
int endBasicBlockId = startBasicBlockId + tilingData_->cubeBlockDim;
|
||||
if (totalBasicBlocks < endBasicBlockId) {
|
||||
endBasicBlockId = totalBasicBlocks;
|
||||
}
|
||||
int startGroupId = GetGroupId(startBasicBlockId);
|
||||
int endGroupId = GetGroupId(endBasicBlockId);
|
||||
int endGroupMOffset = 0;
|
||||
int basicBlockCountBeforeEndGroup = 0;
|
||||
CalculateEndGroupInfo(endBasicBlockId, endGroupId, endGroupMOffset, basicBlockCountBeforeEndGroup);
|
||||
ProcessGroupRange(startGroupId, endGroupId, endGroupMOffset, globalMOffset, isSyncAll);
|
||||
}
|
||||
|
||||
__aicore__ inline void VectorProcess() {
|
||||
if ASCEND_IS_AIV {
|
||||
weightCacheGroupId_ = -1;
|
||||
uint32_t currentBlockId = GetBlockIdx() / 2;
|
||||
uint32_t rsvBlockNum = 0;
|
||||
uint32_t calcBlockNum = 0;
|
||||
uint32_t cvTimes = 0;
|
||||
uint32_t syncId = 0;
|
||||
uint32_t globalMOffset = 0;
|
||||
uint32_t processedBasicBlock = 0;
|
||||
uint32_t currentGroupId = 0;
|
||||
uint32_t realSyncId = 0;
|
||||
bool isSyncAll = false;
|
||||
while (syncId < totalSyncTimes) {
|
||||
cvTimes = CeilDiv(nBasicsBlocks - rsvBlockNum, tilingData_->cubeBlockDim);
|
||||
calcBlockNum += cvTimes * tilingData_->cubeBlockDim;
|
||||
rsvBlockNum = calcBlockNum % nBasicsBlocks;
|
||||
isSyncAll = true;
|
||||
|
||||
for (uint32_t cvId = 0; cvId < cvTimes; cvId++) {
|
||||
ProcessVectorBlock(syncId, isSyncAll, globalMOffset);
|
||||
currentBlockId += tilingData_->cubeBlockDim;
|
||||
syncId += 1;
|
||||
}
|
||||
|
||||
realSyncId += 1;
|
||||
if (realSyncId > 0 && (realSyncId % VC_SYNC_MAX_TIMES == 0)) {
|
||||
AscendC::CrossCoreSetFlag<0x2, PIPE_MTE2>(0x9);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void Process() {
|
||||
CubeProcess();
|
||||
VectorProcess();
|
||||
}
|
||||
|
||||
__aicore__ inline int GetGroupId(int basicBlockId) {
|
||||
int processedBasicBlock = 0;
|
||||
int currentGroupId = 0;
|
||||
int globalMOffset = 0;
|
||||
for (int groupId = 0; groupId < tilingData_->groupNum; groupId++) {
|
||||
int tokens = groupListGm_.GetValue(groupId);
|
||||
if (tilingData_->groupListType == 0 && groupId > 0) {
|
||||
tokens = groupListGm_.GetValue(groupId) - groupListGm_.GetValue(groupId - 1);
|
||||
}
|
||||
int mBasicBlocks = CeilDiv(tokens, matmulTilingData_->baseM);
|
||||
if (processedBasicBlock + mBasicBlocks * nBasicsBlocks >= basicBlockId) {
|
||||
return groupId;
|
||||
} else {
|
||||
processedBasicBlock += mBasicBlocks * nBasicsBlocks;
|
||||
}
|
||||
}
|
||||
return tilingData_->groupNum - 1;
|
||||
}
|
||||
|
||||
__aicore__ inline void ComputeReduceMax(const LocalTensor<float>& tempRes, int32_t calcCount) {
|
||||
uint32_t vectorCycles = calcCount / MAX_CALC_NUM;
|
||||
uint32_t remainElements = calcCount % MAX_CALC_NUM;
|
||||
|
||||
BinaryRepeatParams repeatParams;
|
||||
repeatParams.dstBlkStride = 1;
|
||||
repeatParams.src0BlkStride = 1;
|
||||
repeatParams.src1BlkStride = 1;
|
||||
repeatParams.dstRepStride = 0;
|
||||
repeatParams.src0RepStride = 0x8;
|
||||
repeatParams.src1RepStride = 0;
|
||||
|
||||
if (vectorCycles > 0 && remainElements > 0) {
|
||||
Max(tempRes, tempRes, tempRes[vectorCycles * MAX_CALC_NUM], remainElements, 1, repeatParams);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
if (vectorCycles > 1) {
|
||||
Max(tempRes, tempRes[MAX_CALC_NUM], tempRes, MAX_CALC_NUM, vectorCycles - 1, repeatParams);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessDSQ(int groupId, int globalOffset, int calcCount, bool &isSyncAll) {
|
||||
int32_t blockDimxFactor = (calcCount + tilingData_->vectorBlockDim - 1) / tilingData_->vectorBlockDim;
|
||||
int32_t realCoreDim = calcCount == 0 ? 0 : (calcCount + blockDimxFactor - 1) / blockDimxFactor;
|
||||
|
||||
if (GetBlockIdx() >= realCoreDim) {
|
||||
if (isSyncAll) {
|
||||
AscendC::CrossCoreWaitFlag(0x8);
|
||||
SyncAll<true>();
|
||||
isSyncAll = false;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
DataCopyPadParams padParams{false, 0, 0, 0};
|
||||
LocalTensor<float> inScaleLocal = inScaleQueue_.AllocTensor<float>();
|
||||
|
||||
if (weightCacheGroupId_ != groupId) {
|
||||
DataCopyParams dataCopyWeightScaleParams;
|
||||
dataCopyWeightScaleParams.blockCount = 1;
|
||||
dataCopyWeightScaleParams.blockLen = tilingData_->N * sizeof(float);
|
||||
dataCopyWeightScaleParams.srcStride = 0;
|
||||
dataCopyWeightScaleParams.dstStride = 0;
|
||||
if (tilingData_->isSingleTensor == 0) {
|
||||
weightScaleGm_.SetGlobalBuffer(GetTensorAddr<float>(groupId, weightScaleTensorPtr_));
|
||||
DataCopyPad(inScaleLocal, weightScaleGm_, dataCopyWeightScaleParams, padParams);
|
||||
} else {
|
||||
DataCopyPad(inScaleLocal, weightScaleGm_[groupId * tilingData_->N], dataCopyWeightScaleParams, padParams);
|
||||
}
|
||||
DataCopyParams dataCopyQuantScaleParams;
|
||||
dataCopyQuantScaleParams.blockCount = 1;
|
||||
dataCopyQuantScaleParams.blockLen = (tilingData_->N / SPILI_NUM) * sizeof(float);
|
||||
dataCopyQuantScaleParams.srcStride = 0;
|
||||
dataCopyQuantScaleParams.dstStride = 0;
|
||||
weightCacheGroupId_ = groupId;
|
||||
}
|
||||
|
||||
inScaleQueue_.EnQue(inScaleLocal);
|
||||
inScaleLocal = inScaleQueue_.DeQue<float>();
|
||||
|
||||
int32_t blockDimxTailFactor = calcCount - blockDimxFactor * (realCoreDim - 1);
|
||||
int32_t DimxCore = GetBlockIdx() == (realCoreDim - 1) ? blockDimxTailFactor : blockDimxFactor;
|
||||
|
||||
int32_t ubDimxLoop = (DimxCore + tilingData_->ubFactorDimx - 1) / tilingData_->ubFactorDimx;
|
||||
int32_t ubDimxTailFactor = DimxCore - tilingData_->ubFactorDimx * (ubDimxLoop - 1);
|
||||
|
||||
int64_t coreDimxOffset = blockDimxFactor * GetBlockIdx();
|
||||
int32_t actOffset = tilingData_->actRight * tilingData_->ubFactorDimy;
|
||||
int32_t gateOffset = tilingData_->ubFactorDimy - actOffset;
|
||||
|
||||
LocalTensor<float> weightScaleLocal = inScaleLocal;
|
||||
LocalTensor<float> quantScaleLocal = inScaleLocal[tilingData_->N];
|
||||
|
||||
for (uint32_t loopIdx = 0; loopIdx < ubDimxLoop; loopIdx++) {
|
||||
int64_t xDimxOffset = (coreDimxOffset + loopIdx * tilingData_->ubFactorDimx) + globalOffset;
|
||||
int32_t proDimsx = loopIdx == (ubDimxLoop - 1) ? ubDimxTailFactor : tilingData_->ubFactorDimx;
|
||||
LocalTensor<float> tmpUbF32 = tmpBuf1_.AllocTensor<float>();
|
||||
SetMaskCount();
|
||||
SetVectorMask<float, MaskMode::COUNTER>(tilingData_->ubFactorDimy * SWI_FACTOR);
|
||||
Copy<float, false>(tmpUbF32, weightScaleLocal, MASK_PLACEHOLDER, proDimsx,
|
||||
{1, 1, static_cast<uint16_t>((tilingData_->ubFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 0});
|
||||
SetMaskNorm();
|
||||
ResetMask();
|
||||
|
||||
LocalTensor<int32_t> xActLocal = xActQueue_.AllocTensor<int32_t>();
|
||||
DataCopyParams dataCopyActScaleParams;
|
||||
dataCopyActScaleParams.blockCount = proDimsx;
|
||||
dataCopyActScaleParams.blockLen = sizeof(float);
|
||||
dataCopyActScaleParams.srcStride = 0;
|
||||
dataCopyActScaleParams.dstStride = 0;
|
||||
LocalTensor<float> xActLocalF32 = xActLocal.template ReinterpretCast<float>();
|
||||
DataCopyPad(xActLocalF32[tilingData_->ubFactorDimx * tilingData_->N], activateScaleGm_[xDimxOffset],
|
||||
dataCopyActScaleParams, padParams);
|
||||
|
||||
if (isSyncAll) {
|
||||
AscendC::CrossCoreWaitFlag(0x8);
|
||||
SyncAll<true>();
|
||||
isSyncAll = false;
|
||||
}
|
||||
|
||||
DataCopyParams dataCopyXParams;
|
||||
dataCopyXParams.blockCount = proDimsx;
|
||||
dataCopyXParams.blockLen = tilingData_->N * sizeof(int32_t);
|
||||
dataCopyXParams.srcStride = 0;
|
||||
dataCopyXParams.dstStride = 0;
|
||||
DataCopyPad(xActLocal, workspaceGm_[xDimxOffset * tilingData_->N], dataCopyXParams, padParams);
|
||||
xActQueue_.EnQue(xActLocal);
|
||||
xActLocal = xActQueue_.DeQue<int32_t>();
|
||||
|
||||
LocalTensor<int32_t> xLocal = xActLocal;
|
||||
xActLocalF32 = xActLocal.template ReinterpretCast<float>();
|
||||
LocalTensor<float> xLocalF32 = xActLocalF32;
|
||||
LocalTensor<float> activationScaleLocal = xActLocalF32[tilingData_->ubFactorDimx * tilingData_->N];
|
||||
|
||||
Cast(xLocalF32, xLocal, RoundMode::CAST_NONE, SWI_FACTOR * proDimsx * tilingData_->ubFactorDimy);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Mul(xLocalF32, tmpUbF32, xLocalF32, tilingData_->ubFactorDimy * SWI_FACTOR * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
SetMaskCount();
|
||||
SetVectorMask<float, MaskMode::COUNTER>(tilingData_->ubFactorDimy * SWI_FACTOR);
|
||||
Copy<float, false>(tmpUbF32, activationScaleLocal, AscendC::MASK_PLACEHOLDER, proDimsx,
|
||||
{1, 0, static_cast<uint16_t>((tilingData_->ubFactorDimy * SWI_FACTOR) / BLOCK_ELEM), 1});
|
||||
SetMaskNorm();
|
||||
ResetMask();
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Mul(xLocalF32, tmpUbF32, xLocalF32, tilingData_->ubFactorDimy * SWI_FACTOR * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
LocalTensor<float> tmpUbF32Act = tmpUbF32;
|
||||
LocalTensor<float> tmpUbF32Gate = tmpUbF32[tilingData_->ubFactorDimy * proDimsx];
|
||||
SetMaskCount();
|
||||
SetVectorMask<float, MaskMode::COUNTER>(tilingData_->ubFactorDimy);
|
||||
Copy<float, false>(tmpUbF32Act, xLocalF32[actOffset], AscendC::MASK_PLACEHOLDER, proDimsx,
|
||||
{1, 1, static_cast<uint16_t>(tilingData_->ubFactorDimy / BLOCK_ELEM),
|
||||
static_cast<uint16_t>(tilingData_->ubFactorDimy / BLOCK_ELEM * SWI_FACTOR)});
|
||||
Copy<float, false>(tmpUbF32Gate, xLocalF32[gateOffset], AscendC::MASK_PLACEHOLDER, proDimsx,
|
||||
{1, 1, static_cast<uint16_t>(tilingData_->ubFactorDimy / BLOCK_ELEM),
|
||||
static_cast<uint16_t>(tilingData_->ubFactorDimy / BLOCK_ELEM * SWI_FACTOR)});
|
||||
SetMaskNorm();
|
||||
ResetMask();
|
||||
PipeBarrier<PIPE_V>();
|
||||
limited=tilingData_->swigluLimit;
|
||||
|
||||
if (limited > 0.0f) {
|
||||
Mins(tmpUbF32Gate, tmpUbF32Gate, limited, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(tmpUbF32Gate, tmpUbF32Gate, (-1.0f * limited), tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(tmpUbF32Act, tmpUbF32Act, limited, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
Muls(xLocalF32, tmpUbF32Act, static_cast<float>(-1.0), tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Exp(xLocalF32, xLocalF32, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Adds(xLocalF32, xLocalF32, static_cast<float>(1.0), tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Div(tmpUbF32Act, tmpUbF32Act, xLocalF32, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
xActQueue_.FreeTensor(xActLocal);
|
||||
Mul(tmpUbF32Act, tmpUbF32Gate, tmpUbF32Act, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Abs(tmpUbF32Gate, tmpUbF32Act, tilingData_->ubFactorDimy * proDimsx);
|
||||
|
||||
LocalTensor<float> outLocal = outQueue_.AllocTensor<float>();
|
||||
|
||||
uint64_t scaleOutOffset = tilingData_->ubFactorDimx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t) / sizeof(float);
|
||||
uint64_t alignScaleOutOffset = Ceil(scaleOutOffset, uint32_t(8)) * 8; // 8: num int32_t in 32B ub block
|
||||
LocalTensor<float> scaleOut = outLocal[alignScaleOutOffset];
|
||||
LocalTensor<int8_t> yOut = outLocal.template ReinterpretCast<int8_t>();
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
for (uint32_t i = 0; i < proDimsx; i++) {
|
||||
ComputeReduceMax(tmpUbF32Gate[i * tilingData_->ubFactorDimy], tilingData_->ubFactorDimy);
|
||||
}
|
||||
|
||||
uint64_t realReduceMaxCalcNum = REDUCEMAX_CALC_NUM;
|
||||
if (tilingData_->ubFactorDimy < REDUCEMAX_CALC_NUM) {
|
||||
realReduceMaxCalcNum = tilingData_->ubFactorDimy;
|
||||
}
|
||||
|
||||
WholeReduceMax(tmpUbF32Gate, tmpUbF32Gate, realReduceMaxCalcNum, proDimsx, 1, 1,
|
||||
tilingData_->ubFactorDimy / BLOCK_ELEM, ReduceOrder::ORDER_ONLY_VALUE);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Muls(scaleOut, tmpUbF32Gate, DYNAMIC_QUANT_FACTOR, proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
int64_t blockCount = (proDimsx + BLOCK_ELEM - 1) / BLOCK_ELEM;
|
||||
Brcb(outLocal, scaleOut, blockCount, {1, 8});
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
SetMaskCount();
|
||||
SetVectorMask<float, MaskMode::COUNTER>(tilingData_->ubFactorDimy);
|
||||
Copy<float, false>(tmpUbF32Gate, outLocal, AscendC::MASK_PLACEHOLDER, proDimsx,
|
||||
{1, 0, static_cast<uint16_t>(tilingData_->ubFactorDimy / BLOCK_ELEM), 1});
|
||||
SetMaskNorm();
|
||||
ResetMask();
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Div(tmpUbF32Act, tmpUbF32Act, tmpUbF32Gate, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
LocalTensor<int32_t> tmpUbF32ActI32 = tmpUbF32Act.ReinterpretCast<int32_t>();
|
||||
Cast(tmpUbF32ActI32, tmpUbF32Act, RoundMode::CAST_RINT, tilingData_->ubFactorDimy * proDimsx);
|
||||
SetDeqScale((half)1.000000e+00f);
|
||||
|
||||
LocalTensor<half> tmpUbF32Gate16 = tmpUbF32Gate.template ReinterpretCast<half>();
|
||||
Cast(tmpUbF32Gate16, tmpUbF32ActI32, RoundMode::CAST_ROUND, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
Cast(yOut, tmpUbF32Gate16, RoundMode::CAST_TRUNC, tilingData_->ubFactorDimy * proDimsx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
tmpBuf1_.FreeTensor(tmpUbF32);
|
||||
outQueue_.EnQue<float>(outLocal);
|
||||
outLocal = outQueue_.DeQue<float>();
|
||||
scaleOut = outLocal[alignScaleOutOffset];
|
||||
yOut = outLocal.template ReinterpretCast<int8_t>();
|
||||
|
||||
DataCopyParams dataCopyOutScaleParams;
|
||||
dataCopyOutScaleParams.blockCount = 1;
|
||||
dataCopyOutScaleParams.blockLen = proDimsx * sizeof(float);
|
||||
dataCopyOutScaleParams.srcStride = 0;
|
||||
dataCopyOutScaleParams.dstStride = 0;
|
||||
DataCopyPad(scaleGm_[xDimxOffset], scaleOut, dataCopyOutScaleParams);
|
||||
|
||||
DataCopyParams dataCopyOutyParams;
|
||||
dataCopyOutyParams.blockCount = 1;
|
||||
dataCopyOutyParams.blockLen = proDimsx * (tilingData_->N / SPILI_NUM) * sizeof(int8_t);
|
||||
dataCopyOutyParams.srcStride = 0;
|
||||
dataCopyOutyParams.dstStride = 0;
|
||||
DataCopyPad(yGm_[xDimxOffset * (tilingData_->N / SPILI_NUM)], yOut, dataCopyOutyParams);
|
||||
outQueue_.FreeTensor(outLocal);
|
||||
}
|
||||
inScaleQueue_.FreeTensor(inScaleLocal);
|
||||
}
|
||||
|
||||
private:
|
||||
TPipe *pipe_ = nullptr;
|
||||
const GMMSwigluQuantV2TilingFusionData* __restrict tilingData_;
|
||||
const TCubeTiling* __restrict matmulTilingData_;
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
GlobalTensor<int8_t> xGm_;
|
||||
GlobalTensor<int8_t> weightGm_;
|
||||
GlobalTensor<int8_t> yGm_;
|
||||
GlobalTensor<int32_t> workspaceGm_;
|
||||
GlobalTensor<float> weightScaleGm_;
|
||||
GlobalTensor<float> activateScaleGm_;
|
||||
GlobalTensor<float> scaleGm_;
|
||||
GlobalTensor<int64_t> groupListGm_;
|
||||
int nBasicsBlocks = 0;
|
||||
int totalBasicBlocks = 0;
|
||||
int totalSyncTimes = 0;
|
||||
int32_t weightCacheGroupId_ = -1;
|
||||
float limited = FLOAT_INF;
|
||||
|
||||
TQue<TPosition::VECIN, 1> inQue_;
|
||||
TQue<TPosition::VECIN, 1> xQue_;
|
||||
TBuf<TPosition::VECCALC> tmpBuf_;
|
||||
TQue<TPosition::VECOUT, 1> scaleOutQue_;
|
||||
TQue<TPosition::VECOUT, 1> yOutQue_;
|
||||
TQue<QuePosition::VECIN, 1> xActQueue_;
|
||||
TQue<QuePosition::VECOUT, 1> outQueue_;
|
||||
TQue<QuePosition::VECIN, 1> inScaleQueue_;
|
||||
TBuf<TPosition::VECCALC> tmpBuf1_;
|
||||
|
||||
GM_ADDR weightTensorPtr_;
|
||||
GM_ADDR weightScaleTensorPtr_;
|
||||
};
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,109 @@
|
||||
/**
|
||||
* Copyright (c) 2026 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_v2.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "kernel_operator.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "grouped_matmul_swiglu_quant_spilit_fusion.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_pipeline.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_a4w4_pipeline.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
using namespace GroupedMatmulDequantSwigluQuant;
|
||||
extern "C" __global__ __aicore__ void grouped_matmul_swiglu_quant_v2(GM_ADDR x, GM_ADDR xScale, GM_ADDR groupList,
|
||||
GM_ADDR weight, GM_ADDR weightScale,
|
||||
GM_ADDR weightAssistanceMatrix, GM_ADDR bias,
|
||||
GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
TPipe tPipe;
|
||||
GM_ADDR userWorkspace = GetUserWorkspace(workspace);
|
||||
|
||||
#if defined(GMM_SWIGLU_QUANT_V2_A8W4_MSD)
|
||||
if (TILING_KEY_IS(2)) {
|
||||
KERNEL_TASK_TYPE(2, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_,
|
||||
tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, 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 = MMImplTypeCustom<xType, weightType, yType>;
|
||||
matmulType::MT mm;
|
||||
if ASCEND_IS_AIC {
|
||||
mm.SetSubBlockIdx(0);
|
||||
mm.Init(&mmTilingData_);
|
||||
}
|
||||
GMMSwigluQuantPipelineSchedule<matmulType> op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe);
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, userWorkspace);
|
||||
op.Process();
|
||||
}
|
||||
#endif
|
||||
#if defined(GMM_SWIGLU_QUANT_V2_A4W4)
|
||||
if (TILING_KEY_IS(4)) {
|
||||
KERNEL_TASK_TYPE(4, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_,
|
||||
tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, 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 = MMImplTypeCustom<xType, weightType, yType>;
|
||||
matmulType::MT mm;
|
||||
if ASCEND_IS_AIC {
|
||||
mm.SetSubBlockIdx(0);
|
||||
mm.Init(&mmTilingData_);
|
||||
}
|
||||
GMMSwigluQuantPipelineSchedule<matmulType> op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe);
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, smoothScale, y, yScale, userWorkspace);
|
||||
op.Process();
|
||||
} else if (TILING_KEY_IS(5)) {
|
||||
KERNEL_TASK_TYPE(5, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2BaseParams_,
|
||||
tiling);
|
||||
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, gmmSwigluQuantV2, gmmSwiglu_, tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingData, mmTilingData, mmTilingData_, tiling);
|
||||
using xType = MatmulType<TPosition::GM, CubeFormat::ND, int4b_t, false>;
|
||||
using weightType = MatmulType<TPosition::GM, wFormat, int4b_t, true>;
|
||||
using yType = MatmulType<TPosition::GM, CubeFormat::ND, half, false>;
|
||||
using matmulType = MMImplTypeCustom<xType, weightType, yType>;
|
||||
matmulType::MT mm;
|
||||
if ASCEND_IS_AIC {
|
||||
mm.SetSubBlockIdx(0);
|
||||
mm.Init(&mmTilingData_);
|
||||
}
|
||||
GMMSwigluQuantPipelineSchedule<matmulType> op(mm, &gmmSwigluQuantV2BaseParams_, &gmmSwiglu_, &tPipe);
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, smoothScale, y, yScale, userWorkspace);
|
||||
op.Process();
|
||||
}
|
||||
#endif
|
||||
if (TILING_KEY_IS(3)) {
|
||||
KERNEL_TASK_TYPE(3, KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GET_TILING_DATA_WITH_STRUCT(GMMSwigluQuantV2TilingFusionData, tilingData, tiling);
|
||||
GET_TILING_DATA_MEMBER(GMMSwigluQuantV2TilingFusionData, matmulTiling, matmulTilingData, tiling);
|
||||
GroupedMatmulDequantSwigluQuantFusion op(&tPipe, &tilingData, &matmulTilingData);
|
||||
if ASCEND_IS_AIC {
|
||||
op.mm.SetSubBlockIdx(0);
|
||||
op.mm.Init(&matmulTilingData, &tPipe);
|
||||
}
|
||||
|
||||
op.Init(x, weight, weightScale, xScale, weightAssistanceMatrix, groupList, y, yScale, userWorkspace);
|
||||
op.Process();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
/**
|
||||
* 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_v2_a4w4_mid.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A4W4
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
using namespace matmul;
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr uint32_t BUFFER_NUM = 1;
|
||||
|
||||
template <class mmType>
|
||||
class GMMA4W4MidProcess {
|
||||
public:
|
||||
using bT = typename mmType::BT;
|
||||
|
||||
public:
|
||||
__aicore__ inline GMMA4W4MidProcess(typename mmType::MT &matmul) : mm(matmul)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN);
|
||||
__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, bool resetOutputOffset);
|
||||
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
const uint32_t HALF_ALIGN = 16;
|
||||
GlobalTensor<int4b_t> xGM;
|
||||
GlobalTensor<int4b_t> weightGM;
|
||||
|
||||
GlobalTensor<half> mmOutGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<uint64_t> weightScaleGM;
|
||||
|
||||
GM_ADDR weightTensorPtr;
|
||||
GM_ADDR weightScaleTensorPtr;
|
||||
|
||||
MNConfig mnConfig;
|
||||
|
||||
// define the que
|
||||
uint32_t subBlockIdx = 0;
|
||||
uint32_t coreIdx = 0;
|
||||
uint32_t quantGroupSize = 0;
|
||||
uint32_t vecCount = 0;
|
||||
uint32_t xRowSumCount = 0;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams = nullptr;
|
||||
};
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void
|
||||
GMMA4W4MidProcess<mmType>::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN;
|
||||
xRowSumCount = gmmSwigluQuantV2BaseParams->M;
|
||||
xGM.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.xGM);
|
||||
weightGM.SetGlobalBuffer(GetTensorAddr<int4b_t>(0, gmAddrParams.weightGM));
|
||||
weightScaleGM.SetGlobalBuffer(GetTensorAddr<uint64_t>(0, gmAddrParams.weightScaleGM));
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
mmOutGM1.SetGlobalBuffer((__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1));
|
||||
quantGroupSize = gmmSwigluQuantV2BaseParams->K / gmmSwigluQuantV2BaseParams->quantGroupNum; // 约束为整除关系
|
||||
subBlockIdx = GetSubBlockIdx();
|
||||
coreIdx = GetBlockIdx();
|
||||
weightTensorPtr = gmAddrParams.weightGM;
|
||||
weightScaleTensorPtr = gmAddrParams.weightScaleGM;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA4W4MidProcess<mmType>::UpdateMnConfig(MNConfig &mnConfig, bool resetOutputOffset)
|
||||
{
|
||||
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;
|
||||
if (resetOutputOffset) {
|
||||
mnConfig.yBaseOffset = 0;
|
||||
} else {
|
||||
mnConfig.yBaseOffset += mnConfig.m * mnConfig.n;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA4W4MidProcess<mmType>::SetMNConfig(const int32_t splitValue, MNConfig &mnConfig)
|
||||
{
|
||||
mnConfig.m = static_cast<int64_t>(splitValue);
|
||||
mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.singleN = gmmSwigluQuantV2BaseParams->singleN != 0 && gmmSwigluQuantV2BaseParams->quantGroupNum == 1?
|
||||
gmmSwigluQuantV2BaseParams->singleN : gmmSwigluQuantV2BaseParams->baseN;
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA4W4MidProcess<mmType>::Process(WorkSpaceSplitConfig &workspaceSplitConfig,
|
||||
int64_t workspaceSplitLoopIdx)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
if (workspaceSplitLoopIdx >= workspaceSplitConfig.loopCount || workspaceSplitLoopIdx < 0) {
|
||||
return;
|
||||
}
|
||||
mmOutGM = (workspaceSplitLoopIdx % NUM_2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
mnConfig.baseM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.singleN = gmmSwigluQuantV2BaseParams->singleN != 0 && gmmSwigluQuantV2BaseParams->quantGroupNum == 1?
|
||||
gmmSwigluQuantV2BaseParams->singleN : gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.k = gmmSwigluQuantV2BaseParams->K; // tilingData
|
||||
mnConfig.n = gmmSwigluQuantV2BaseParams->N; // tilingData
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
int32_t totalTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
// 当workspace切换时,需要将输出的地址偏移初始化为0,使用resetOutputOffset控制
|
||||
bool resetOutputOffset = true;
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig, resetOutputOffset);
|
||||
resetOutputOffset = false;
|
||||
int32_t currSplitValue = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
} else {
|
||||
totalTmp += static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = totalTmp;
|
||||
}
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit :
|
||||
currSplitValue;
|
||||
|
||||
int32_t splitValue = (currSplitValue - prevSplitValue);
|
||||
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 + gmmSwigluQuantV2BaseParams->coreNum;
|
||||
while (curBlock < curCount) {
|
||||
mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN;
|
||||
MMCompute(groupIdx, mnConfig, workspaceSplitConfig);
|
||||
curBlock += gmmSwigluQuantV2BaseParams->coreNum;
|
||||
}
|
||||
preCount = curCount % gmmSwigluQuantV2BaseParams->coreNum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void GMMA4W4MidProcess<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 = gmmSwigluQuantV2BaseParams->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;
|
||||
mm.SetSingleShape(curSingleM, curSingleN, quantGroupSize);
|
||||
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 < gmmSwigluQuantV2BaseParams->quantGroupNum; loopK++) {
|
||||
mm.SetTensorA(
|
||||
xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]);
|
||||
if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) {
|
||||
weightGM.SetGlobalBuffer(GetTensorAddr<int4b_t>(groupIdx, weightTensorPtr));
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == true) {
|
||||
weightOffset = tailN * 64;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
} else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) {
|
||||
weightOffset = tailN * gmmSwigluQuantV2BaseParams->K;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightOffset = tailN;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
}
|
||||
} else {
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == true) {
|
||||
weightOffset = static_cast<uint64_t>(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K +
|
||||
tailN * 64;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
} else if constexpr (mmType::BT::format == CubeFormat::NZ && mmType::BT::isTrans == false) {
|
||||
weightOffset =
|
||||
static_cast<uint64_t>(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K +
|
||||
tailN * gmmSwigluQuantV2BaseParams->K;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightOffset =
|
||||
static_cast<uint64_t>(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K +
|
||||
tailN;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
}
|
||||
}
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mm.SetTensorB(weightSlice, mmType::BT::isTrans);
|
||||
if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) {
|
||||
weightScaleGM.SetGlobalBuffer(GetTensorAddr<uint64_t>(groupIdx, weightScaleTensorPtr));
|
||||
mm.SetQuantVector(weightScaleGM[loopK * gmmSwigluQuantV2BaseParams->N + tailN]);
|
||||
} else {
|
||||
mm.SetQuantVector(
|
||||
weightScaleGM[groupIdx * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->quantGroupNum +
|
||||
loopK * gmmSwigluQuantV2BaseParams->N + tailN]);
|
||||
}
|
||||
mm.IterateAll(mmOutGM[mnConfig.workspaceOffset], loopK == 0 ? 0 : 1);
|
||||
}
|
||||
}
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A4W4
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_MID_H
|
||||
@@ -0,0 +1,229 @@
|
||||
/**
|
||||
* 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_v2_a4w4_pipeline.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H
|
||||
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_v2_a4w4_mid.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_a4w4_post.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A4W4
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
|
||||
template <class mmType>
|
||||
class GMMSwigluQuantPipelineSchedule {
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
TPipe *pipe;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2;
|
||||
// WorkSpaceSplitConfig控制Workspace切割方式的结构体;
|
||||
WorkSpaceSplitConfig workspaceSplitConfig;
|
||||
WorkSpaceSplitConfig tempWorkspaceSplitConfig;
|
||||
// 记录GM_ADDR的结构体
|
||||
GMAddrParams gmAddrParams;
|
||||
// 中间处理GMMA4W4MidProcess类
|
||||
GMMA4W4MidProcess<mmType> midProcess;
|
||||
// 后处理GMMA4W4PostProcess类
|
||||
GMMA4W4PostProcess 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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
: mm(mm_), midProcess(mm), gmmSwigluQuantV2BaseParams(gmmSwigluQuantV2BaseParamsIN),
|
||||
gmmSwigluQuantV2(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 smoothScale, 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 smoothScale, 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.smoothScaleGM = smoothScale;
|
||||
gmAddrParams.yGM = y;
|
||||
gmAddrParams.yScaleGM = yScale;
|
||||
gmAddrParams.workSpaceGM = workspace;
|
||||
gmAddrParams.workSpaceOffset1 = gmmSwigluQuantV2BaseParams->workSpaceOffset1;
|
||||
gmAddrParams.workSpaceOffset2 = 0;
|
||||
gmAddrParams.workSpaceOffset3 = 0;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM);
|
||||
InitWorkSpaceSplitConfig(workspaceSplitConfig);
|
||||
}
|
||||
|
||||
template <class mmType>
|
||||
__aicore__ inline void GMMSwigluQuantPipelineSchedule<mmType>::Process()
|
||||
{
|
||||
// 1.对每次workspace切分做大循环。
|
||||
midProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams);
|
||||
postProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2);
|
||||
|
||||
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次后处理 并行
|
||||
midProcess.Process(workspaceSplitConfig, workspaceSplitLoopIdx);
|
||||
|
||||
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)
|
||||
{
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwigluQuantV2->groupListLen - 1);
|
||||
} else {
|
||||
int64_t totalTmp = 0;
|
||||
for (uint32_t i = 0; i < gmmSwigluQuantV2->groupListLen; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
workspaceSplitConfig.M = totalTmp;
|
||||
}
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmSwigluQuantV2BaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmSwigluQuantV2BaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmSwigluQuantV2BaseParams->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 * gmmSwigluQuantV2BaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
int32_t curTaskNumTmp = 0;
|
||||
int32_t nextTaskNumTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) {
|
||||
curTaskNumTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
if (workspaceSplitConfig.rightMatrixExpertEndIndex == 0) {
|
||||
nextTaskNumTmp = groupListGM.GetValue(0);
|
||||
} else {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) {
|
||||
nextTaskNumTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen) {
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
} else {
|
||||
curTaskNumTmp += groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex);
|
||||
curTaskNum = curTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
}
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen - 1 ?
|
||||
gmmSwigluQuantV2->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
} else {
|
||||
if (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen - 1) {
|
||||
nextTaskNumTmp += groupListGM.GetValue(nextTaskIdx);
|
||||
}
|
||||
nextTaskNum = nextTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
}
|
||||
if (curTaskNum > gmmSwigluQuantV2BaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmSwigluQuantV2BaseParams->mLimit &&
|
||||
nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmSwigluQuantV2BaseParams->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 >= gmmSwigluQuantV2->groupListLen ?
|
||||
gmmSwigluQuantV2->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A4W4
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_PIPELINE_H
|
||||
@@ -0,0 +1,392 @@
|
||||
/**
|
||||
* 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_v2_a4w4_post.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A4W4
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
using namespace AscendC;
|
||||
#define DOUBLE_BUFFER 2
|
||||
constexpr float DEFAULT_MUL_SCALE = 16.0f;
|
||||
class GMMA4W4PostProcess {
|
||||
public:
|
||||
__aicore__ inline GMMA4W4PostProcess(){};
|
||||
__aicore__ inline void Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN);
|
||||
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
|
||||
private:
|
||||
__aicore__ inline void UpdateVecConfig(uint32_t blockIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
|
||||
__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 Quant(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void Swiglu(uint32_t loopIdx, VecConfig &vecConfig);
|
||||
|
||||
__aicore__ inline void MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
__aicore__ inline void ApplySmoothScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig);
|
||||
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
GlobalTensor<float> perTokenScaleGM;
|
||||
GlobalTensor<int64_t> groupListGM;
|
||||
GlobalTensor<float> smoothScaleGM;
|
||||
GlobalTensor<int8_t> quantOutputGM;
|
||||
GlobalTensor<float> quantScaleOutputGM;
|
||||
GlobalTensor<half> mmOutGM1;
|
||||
GlobalTensor<half> mmOutGM2;
|
||||
GlobalTensor<half> mmOutGM;
|
||||
LocalTensor<float> mmLocal_fp32;
|
||||
LocalTensor<half> mmLocal_fp16;
|
||||
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;
|
||||
float limited = FLOAT_INF;
|
||||
};
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * NUM_2;
|
||||
blockIdx = GetBlockIdx();
|
||||
gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN;
|
||||
gmmSwigluQuantV2 = gmmSwigluIN;
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwigluQuantV2->groupListLen);
|
||||
mmOutGM1.SetGlobalBuffer((__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM));
|
||||
mmOutGM2.SetGlobalBuffer(
|
||||
(__gm__ half *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1));
|
||||
perTokenScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.xScaleGM, gmmSwigluQuantV2BaseParams->M);
|
||||
smoothScaleGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.smoothScaleGM);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM, gmmSwigluQuantV2BaseParams->M *
|
||||
gmmSwigluQuantV2->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmSwigluQuantV2BaseParams->M);
|
||||
limited = gmmSwigluQuantV2BaseParams->swigluLimit;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::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 = vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen;
|
||||
DataCopyExtParams copyParams_0{1, static_cast<uint32_t>(processNum * SIZE_OF_HALF_2), 0, 0, 0};
|
||||
DataCopyPadExtParams<half> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(mmLocal_fp16[processNum], mmOutGM[vecConfig.curOffset], 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>();
|
||||
|
||||
vecConfig.curIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwigluQuantV2->tokenLen;
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::VectorCompute(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
// 1.perToken反量化
|
||||
MulPertokenScale(loopIdx, vecConfig, workspaceSplitConfig);
|
||||
// 2.Swiglu
|
||||
Swiglu(loopIdx, vecConfig);
|
||||
// 3.ApplySmoothScale（smoothScaleDimNum为0时跳过，表示smoothScale为空指针）
|
||||
if (gmmSwigluQuantV2BaseParams->smoothScaleDimNum != 0) {
|
||||
ApplySmoothScale(loopIdx, vecConfig, workspaceSplitConfig);
|
||||
}
|
||||
// 4.Quant
|
||||
Quant(loopIdx, vecConfig);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
if (loopIdx != 0) {
|
||||
mmLocal_fp32 = mmOutQueue.DeQue<float>();
|
||||
}
|
||||
float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], scale,
|
||||
gmmSwigluQuantV2->tokenLen);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::Swiglu(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 高阶API swiglu
|
||||
float beta = 1.0f;
|
||||
LocalTensor<float> workspaceLocal = reduceWorkspace.Get<float>();
|
||||
LocalTensor<float> src0Local =
|
||||
mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen];
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{
|
||||
1, static_cast<uint16_t>((gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, 0};
|
||||
DataCopy(mmLocal_fp32[loopIdx * gmmSwigluQuantV2->tokenLen], workspaceLocal, repeatParams);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::ApplySmoothScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
int64_t smoothScaleDimNum = gmmSwigluQuantV2BaseParams->smoothScaleDimNum;
|
||||
int64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR;
|
||||
int64_t currentTokenIdx = workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx + loopIdx;
|
||||
|
||||
// 找到当前token所属的group
|
||||
uint32_t groupIdx = 0;
|
||||
int64_t prevM = 0;
|
||||
int64_t totalTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
for (uint32_t i = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
i <= workspaceSplitConfig.rightMatrixExpertEndIndex; i++) {
|
||||
int64_t currM = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
currM = groupListGM.GetValue(i);
|
||||
} else {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
currM = totalTmp;
|
||||
}
|
||||
if (currentTokenIdx < currM) {
|
||||
groupIdx = i;
|
||||
break;
|
||||
}
|
||||
prevM = currM;
|
||||
}
|
||||
|
||||
uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen;
|
||||
|
||||
if (smoothScaleDimNum == NUM_2) {
|
||||
// smoothScale形状为 (E, N/2)，只需要当前group的那一行
|
||||
for (uint32_t j = 0; j < halfTokenLen; j++) {
|
||||
float scale = smoothScaleGM.GetValue(groupIdx * halfTokenLen + j);
|
||||
float val = mmLocal_fp32.GetValue(preOffset + j);
|
||||
mmLocal_fp32.SetValue(preOffset + j, val * scale);
|
||||
}
|
||||
} else if (smoothScaleDimNum == 1) {
|
||||
// smoothScale形状为 (E,)，需要广播到 (N/2)
|
||||
float scale = smoothScaleGM.GetValue(groupIdx);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(mmLocal_fp32[preOffset], mmLocal_fp32[preOffset], scale, halfTokenLen);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / BISECT;
|
||||
PipeBarrier<PIPE_V>();
|
||||
Abs(mmLocal_fp32[preOffset + gmmSwigluQuantV2->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 + gmmSwigluQuantV2->tokenLen / BISECT],
|
||||
reduceTmpLocal, static_cast<uint32_t>(halfTokenLen));
|
||||
float quantScale = reduceResLocal.GetValue(0) / QUANT_SCALE_INT8;
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantScaleLocal.SetValue(loopIdx, quantScale);
|
||||
quantScale = QUANT_SCALE_INT8 / reduceResLocal.GetValue(0);
|
||||
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);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::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 * gmmSwigluQuantV2->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
int64_t totalTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
currM = groupListGM.GetValue(groupIdx);
|
||||
} else {
|
||||
totalTmp += groupListGM.GetValue(groupIdx);
|
||||
currM = totalTmp;
|
||||
}
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum =
|
||||
(vecConfig.taskNum + gmmSwigluQuantV2->maxProcessRowNum - 1) / gmmSwigluQuantV2->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum :
|
||||
gmmSwigluQuantV2->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,
|
||||
gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, 1,
|
||||
gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR *
|
||||
sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, 1,
|
||||
AlignUp<int32_t>(gmmSwigluQuantV2->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(float) +
|
||||
UB_BLOCK_UNIT_SIZE + UB_BLOCK_UNIT_SIZE);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::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 % NUM_2 == 0 ? mmOutGM1 : mmOutGM2);
|
||||
LocalTensor<half> mmLocal = mmOutQueue.AllocTensor<half>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
for (uint32_t outLoopIdx = 0; outLoopIdx < vecConfig.outLoopNum; outLoopIdx++) {
|
||||
vecConfig.innerLoopNum = outLoopIdx == (vecConfig.outLoopNum - 1) ? vecConfig.tailLoopNum :
|
||||
gmmSwigluQuantV2->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. 四步vector计算（perToken反量化、Swiglu、SmoothScale、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);
|
||||
}
|
||||
mmLocal = mmOutQueue.DeQue<half>();
|
||||
quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
quantScaleOutQueue.FreeTensor(quantScaleLocal);
|
||||
quantOutQueue.FreeTensor(quantLocal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA4W4PostProcess::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 * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0,
|
||||
0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) *
|
||||
gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A4W4
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A4W4_POST_H
|
||||
@@ -0,0 +1,288 @@
|
||||
/**
|
||||
* 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_v2_a8w4_msd_mid.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN);
|
||||
__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;
|
||||
|
||||
GM_ADDR weightTensorPtr;
|
||||
GM_ADDR weightScaleTensorPtr;
|
||||
|
||||
// define the que
|
||||
uint32_t subBlockIdx = 0;
|
||||
uint32_t coreIdx = 0;
|
||||
uint32_t quantGroupSize = 0;
|
||||
uint32_t vecCount = 0;
|
||||
uint32_t xRowSumCount = 0;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
};
|
||||
|
||||
template <typename mmType>
|
||||
__aicore__ inline void
|
||||
GMMA8W4MidProcess<mmType>::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN)
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN;
|
||||
xRowSumCount = gmmSwigluQuantV2BaseParams->M;
|
||||
xGM1.SetGlobalBuffer((__gm__ int4b_t *)gmAddrParams.workSpaceGM); // 从前处理中获得的结果
|
||||
xGM2.SetGlobalBuffer(
|
||||
(__gm__ int4b_t *)((__gm__ int8_t *)gmAddrParams.workSpaceGM + gmAddrParams.workSpaceOffset1));
|
||||
weightGM.SetGlobalBuffer(GetTensorAddr<int4b_t>(0, gmAddrParams.weightGM));
|
||||
weightScaleGM.SetGlobalBuffer(GetTensorAddr<uint64_t>(0, 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 = gmmSwigluQuantV2BaseParams->K / gmmSwigluQuantV2BaseParams->quantGroupNum; // 约束为整除关系
|
||||
subBlockIdx = GetSubBlockIdx();
|
||||
coreIdx = GetBlockIdx();
|
||||
weightTensorPtr = gmAddrParams.weightGM;
|
||||
weightScaleTensorPtr = gmAddrParams.weightScaleGM;
|
||||
}
|
||||
}
|
||||
|
||||
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 = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.singleN = gmmSwigluQuantV2BaseParams->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 = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.baseN = gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.singleM = gmmSwigluQuantV2BaseParams->baseM;
|
||||
mnConfig.singleN = gmmSwigluQuantV2BaseParams->baseN;
|
||||
mnConfig.k = gmmSwigluQuantV2BaseParams->K; // tilingData
|
||||
mnConfig.n = gmmSwigluQuantV2BaseParams->N; // tilingData
|
||||
mnConfig.blockDimN = Ceil(mnConfig.n, mnConfig.singleN);
|
||||
int32_t prevSplitValue = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
int32_t totalTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex, preCount = 0;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; ++groupIdx) {
|
||||
UpdateMnConfig(mnConfig);
|
||||
int32_t currSplitValue = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
currSplitValue = static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
} else {
|
||||
totalTmp += static_cast<int32_t>(groupListGM.GetValue(groupIdx));
|
||||
currSplitValue = totalTmp;
|
||||
}
|
||||
currSplitValue = currSplitValue > (workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->mLimit ?
|
||||
(workspaceSplitLoopIdx + 1) * gmmSwigluQuantV2BaseParams->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 + gmmSwigluQuantV2BaseParams->coreNum;
|
||||
while (curBlock < curCount) {
|
||||
mnConfig.mIdx = (curBlock - preCount) / mnConfig.blockDimN;
|
||||
mnConfig.nIdx = (curBlock - preCount) % mnConfig.blockDimN;
|
||||
MMCompute(groupIdx, mnConfig, workspaceSplitConfig);
|
||||
curBlock += gmmSwigluQuantV2BaseParams->coreNum;
|
||||
}
|
||||
preCount = curCount % gmmSwigluQuantV2BaseParams->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 = gmmSwigluQuantV2BaseParams->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;
|
||||
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 < gmmSwigluQuantV2BaseParams->quantGroupNum; loopK++) {
|
||||
mm.SetTensorA(
|
||||
xGM[mnConfig.xBaseOffset + mnConfig.mIdx * mnConfig.k * mnConfig.singleM + loopK * quantGroupSize]);
|
||||
if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) {
|
||||
weightGM.SetGlobalBuffer(GetTensorAddr<int4b_t>(groupIdx, weightTensorPtr));
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightOffset = tailN * gmmSwigluQuantV2BaseParams->K;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightOffset = tailN;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
}
|
||||
} else {
|
||||
if constexpr (mmType::BT::format == CubeFormat::NZ) {
|
||||
weightOffset =
|
||||
static_cast<uint64_t>(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K +
|
||||
tailN * gmmSwigluQuantV2BaseParams->K;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * 64];
|
||||
} else {
|
||||
weightOffset =
|
||||
static_cast<uint64_t>(groupIdx) * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->K +
|
||||
tailN;
|
||||
weightSlice = weightGM[weightOffset + loopK * quantGroupSize * gmmSwigluQuantV2BaseParams->N];
|
||||
}
|
||||
}
|
||||
if (mnConfig.blockDimM == 1) {
|
||||
weightSlice.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
|
||||
}
|
||||
mm.SetTensorB(weightSlice);
|
||||
if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) {
|
||||
weightScaleGM.SetGlobalBuffer(GetTensorAddr<uint64_t>(groupIdx, weightScaleTensorPtr));
|
||||
mm.SetQuantVector(weightScaleGM[loopK * gmmSwigluQuantV2BaseParams->N + tailN]);
|
||||
} else {
|
||||
mm.SetQuantVector(
|
||||
weightScaleGM[groupIdx * gmmSwigluQuantV2BaseParams->N * gmmSwigluQuantV2BaseParams->quantGroupNum +
|
||||
loopK * gmmSwigluQuantV2BaseParams->N + tailN]);
|
||||
}
|
||||
mm.Iterate();
|
||||
mm.GetTensorC(mmOutGM[mnConfig.workspaceOffset], loopK == 0 ? 0 : 1);
|
||||
}
|
||||
}
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_MID_H
|
||||
@@ -0,0 +1,236 @@
|
||||
/**
|
||||
* 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_v2_a8w4_msd_pipeline.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H
|
||||
|
||||
#include <typeinfo>
|
||||
#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_pre.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_mid.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_a8w4_msd_post.h"
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
|
||||
template <class mmType>
|
||||
class GMMSwigluQuantPipelineSchedule {
|
||||
private:
|
||||
typename mmType::MT &mm;
|
||||
TPipe *pipe;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2;
|
||||
// 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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN, TPipe *tPipeIN)
|
||||
: mm(mm_), midProcess(mm), gmmSwigluQuantV2BaseParams(gmmSwigluQuantV2BaseParamsIN),
|
||||
gmmSwigluQuantV2(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 = gmmSwigluQuantV2BaseParams->workSpaceOffset1 / 2;
|
||||
gmAddrParams.workSpaceOffset2 = gmmSwigluQuantV2BaseParams->workSpaceOffset1;
|
||||
gmAddrParams.workSpaceOffset3 =
|
||||
gmmSwigluQuantV2BaseParams->workSpaceOffset1 + gmmSwigluQuantV2BaseParams->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, gmmSwigluQuantV2BaseParams);
|
||||
midProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams);
|
||||
postProcess.Init(gmAddrParams, gmmSwigluQuantV2BaseParams, gmmSwigluQuantV2);
|
||||
|
||||
// 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)
|
||||
{
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
workspaceSplitConfig.M = groupListGM.GetValue(gmmSwigluQuantV2->groupListLen - 1);
|
||||
} else {
|
||||
int64_t totalTmp = 0;
|
||||
for (uint32_t i = 0; i < gmmSwigluQuantV2->groupListLen; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
workspaceSplitConfig.M = totalTmp;
|
||||
}
|
||||
workspaceSplitConfig.loopCount = Ceil(workspaceSplitConfig.M, gmmSwigluQuantV2BaseParams->mLimit);
|
||||
workspaceSplitConfig.notLastTaskSize = gmmSwigluQuantV2BaseParams->mLimit;
|
||||
workspaceSplitConfig.lastLoopTaskSize =
|
||||
workspaceSplitConfig.M - (workspaceSplitConfig.loopCount - 1) * gmmSwigluQuantV2BaseParams->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 * gmmSwigluQuantV2BaseParams->mLimit;
|
||||
workspaceSplitConfig.rightMatrixExpertStartIndex = workspaceSplitConfig.rightMatrixExpertNextStartIndex;
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
// 计算右专家矩阵的终止索引(rightMatrixExpertEndIndex) 和下一次的起始索引(rightMatrixExpertNextStartIndex)
|
||||
int32_t curTaskNum = 0;
|
||||
int32_t nextTaskNum = 0;
|
||||
int32_t curTaskNumTmp = 0;
|
||||
int32_t nextTaskNumTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) {
|
||||
curTaskNumTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
if (workspaceSplitConfig.rightMatrixExpertEndIndex == 0) {
|
||||
nextTaskNumTmp = groupListGM.GetValue(0);
|
||||
} else {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertEndIndex; i++) {
|
||||
nextTaskNumTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
while (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen) {
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
curTaskNum = groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex) -
|
||||
workspaceSplitConfig.leftMatrixStartIndex;
|
||||
} else {
|
||||
curTaskNumTmp += groupListGM.GetValue(workspaceSplitConfig.rightMatrixExpertEndIndex);
|
||||
curTaskNum = curTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
}
|
||||
int32_t nextTaskIdx = workspaceSplitConfig.rightMatrixExpertEndIndex >= gmmSwigluQuantV2->groupListLen - 1 ?
|
||||
gmmSwigluQuantV2->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
nextTaskNum = groupListGM.GetValue(nextTaskIdx) - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
} else {
|
||||
if (workspaceSplitConfig.rightMatrixExpertEndIndex < gmmSwigluQuantV2->groupListLen - 1) {
|
||||
nextTaskNumTmp += groupListGM.GetValue(nextTaskIdx);
|
||||
}
|
||||
nextTaskNum = nextTaskNumTmp - workspaceSplitConfig.leftMatrixStartIndex;
|
||||
}
|
||||
if (curTaskNum > gmmSwigluQuantV2BaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
break;
|
||||
} else if (curTaskNum == gmmSwigluQuantV2BaseParams->mLimit &&
|
||||
nextTaskNum > gmmSwigluQuantV2BaseParams->mLimit) {
|
||||
workspaceSplitConfig.rightMatrixExpertNextStartIndex = workspaceSplitConfig.rightMatrixExpertEndIndex + 1;
|
||||
break;
|
||||
} else if (nextTaskNum > gmmSwigluQuantV2BaseParams->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 >= gmmSwigluQuantV2->groupListLen ?
|
||||
gmmSwigluQuantV2->groupListLen - 1 :
|
||||
workspaceSplitConfig.rightMatrixExpertEndIndex;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PIPELINE_H
|
||||
@@ -0,0 +1,444 @@
|
||||
/**
|
||||
* 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_v2_a8w4_msd_post.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN);
|
||||
|
||||
__aicore__ inline void Process(WorkSpaceSplitConfig &workspaceSplitConfig, int64_t workspaceSplitLoopIdx,
|
||||
TPipe *pipe);
|
||||
static constexpr float FLOAT_INF = 3e+99;
|
||||
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 GMMSwigluQuantV2 *__restrict gmmSwigluQuantV2;
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
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;
|
||||
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;
|
||||
GM_ADDR weightAuxiliaryMatrixTensorPtr;
|
||||
float limited = FLOAT_INF;
|
||||
};
|
||||
|
||||
__aicore__ inline void
|
||||
GMMA8W4PostProcess::Init(const GMAddrParams gmAddrParams,
|
||||
const GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN,
|
||||
const GMMSwigluQuantV2 *__restrict gmmSwigluIN)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
aicCoreNum = GetBlockNum();
|
||||
aivCoreNum = aicCoreNum * 2;
|
||||
blockIdx = GetBlockIdx();
|
||||
gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN;
|
||||
gmmSwigluQuantV2 = gmmSwigluIN;
|
||||
weightAuxiliaryMatrixGM.SetGlobalBuffer(GetTensorAddr<float>(0, gmAddrParams.weightAuxiliaryMatrixGM));
|
||||
groupListGM.SetGlobalBuffer((__gm__ int64_t *)gmAddrParams.groupListGM, gmmSwigluQuantV2->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, gmmSwigluQuantV2BaseParams->M);
|
||||
quantOutputGM.SetGlobalBuffer((__gm__ int8_t *)gmAddrParams.yGM, gmmSwigluQuantV2BaseParams->M *
|
||||
gmmSwigluQuantV2->tokenLen /
|
||||
SWIGLU_REDUCE_FACTOR);
|
||||
quantScaleOutputGM.SetGlobalBuffer((__gm__ float *)gmAddrParams.yScaleGM, gmmSwigluQuantV2BaseParams->M);
|
||||
weightAuxiliaryMatrixTensorPtr = gmAddrParams.weightAuxiliaryMatrixGM;
|
||||
limited = gmmSwigluQuantV2BaseParams->swigluLimit;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::customDataCopyIn(uint32_t outLoopIdx, GlobalTensor<half> &mmOutGM,
|
||||
VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
LocalTensor<half> _inMMLocal_0 = mmOutQueue.DeQue<half>();
|
||||
const int64_t processNum = 2 * vecConfig.innerLoopNum * gmmSwigluQuantV2->tokenLen;
|
||||
DataCopyExtParams copyParams_0{1, static_cast<uint32_t>(processNum * SIZE_OF_HALF_2), 0, 0, 0};
|
||||
DataCopyPadExtParams<half> padParams_0{false, 0, 0, 0};
|
||||
DataCopyPad(_inMMLocal_0[processNum], mmOutGM[vecConfig.curOffset * DOUBLE_ROW], copyParams_0, padParams_0);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_0);
|
||||
|
||||
LocalTensor<half> _inMMLocal_1 = mmOutQueue.DeQue<half>();
|
||||
// 1. fp16 -> fp32
|
||||
Cast(_inMMLocal_1.ReinterpretCast<float>(), _inMMLocal_1[processNum], RoundMode::CAST_NONE, processNum);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal_1);
|
||||
LocalTensor<float> _inMMLocal_2 = mmOutQueue.DeQue<float>();
|
||||
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(_inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen],
|
||||
_inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen], DEFAULT_MUL_SCALE,
|
||||
gmmSwigluQuantV2->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Add(_inMMLocal_2[i * gmmSwigluQuantV2->tokenLen], _inMMLocal_2[(DOUBLE_ROW * i) * gmmSwigluQuantV2->tokenLen],
|
||||
_inMMLocal_2[(DOUBLE_ROW * i + 1) * gmmSwigluQuantV2->tokenLen], gmmSwigluQuantV2->tokenLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
vecConfig.curIdx++;
|
||||
}
|
||||
vecConfig.curOffset = vecConfig.curIdx * gmmSwigluQuantV2->tokenLen;
|
||||
mmOutQueue.EnQue(_inMMLocal_2);
|
||||
}
|
||||
|
||||
__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
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
Add(mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen],
|
||||
weightAuxiliaryMatrixLocal, gmmSwigluQuantV2->tokenLen);
|
||||
vecConfig.nextUpdateInterVal--;
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::MulPertokenScale(uint32_t loopIdx, VecConfig &vecConfig,
|
||||
WorkSpaceSplitConfig &workspaceSplitConfig)
|
||||
{
|
||||
LocalTensor<float> mmLocal = mmOutQueue.DeQue<float>();
|
||||
int32_t eventIdSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
float scale = perTokenScaleGM.GetValue(loopIdx + workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
Muls(mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], mmLocal[loopIdx * gmmSwigluQuantV2->tokenLen], scale,
|
||||
gmmSwigluQuantV2->tokenLen);
|
||||
SetFlag<HardEvent::S_V>(eventIdSToV);
|
||||
WaitFlag<HardEvent::S_V>(eventIdSToV);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::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 * gmmSwigluQuantV2->tokenLen + gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR];
|
||||
LocalTensor<float> src1Local = _inMMLocal[loopIdx * gmmSwigluQuantV2->tokenLen];
|
||||
|
||||
if (limited > 0.0f) {
|
||||
Mins(src0Local, src0Local, limited, gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Maxs(src0Local, src0Local, (-1.0f * limited), gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Mins(src1Local, src1Local, limited, gmmSwigluQuantV2->tokenLen / 2);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
SwiGLU<float, false>(workspaceLocal, src0Local, src1Local, beta, gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR);
|
||||
PipeBarrier<PIPE_V>();
|
||||
DataCopyParams repeatParams{
|
||||
1, static_cast<uint16_t>((gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR) / ALIGN_8_ELE), 0, 0};
|
||||
DataCopy(_inMMLocal[loopIdx * gmmSwigluQuantV2->tokenLen], workspaceLocal, repeatParams);
|
||||
|
||||
mmOutQueue.EnQue(_inMMLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::Quant(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
LocalTensor<float> _inMMLocal = mmOutQueue.DeQue<float>();
|
||||
uint64_t preOffset = loopIdx * gmmSwigluQuantV2->tokenLen;
|
||||
uint64_t halfTokenLen = gmmSwigluQuantV2->tokenLen / BISECT;
|
||||
Abs(_inMMLocal[preOffset + gmmSwigluQuantV2->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 + gmmSwigluQuantV2->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);
|
||||
}
|
||||
|
||||
__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 * gmmSwigluQuantV2->tokenLen;
|
||||
vecConfig.curOffset = vecConfig.startOffset;
|
||||
int64_t curStartIdx = vecConfig.startIdx;
|
||||
int64_t prevM = workspaceSplitLoopIdx * workspaceSplitConfig.notLastTaskSize;
|
||||
int64_t totalTmp = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 1) {
|
||||
for (uint32_t i = 0; i < workspaceSplitConfig.rightMatrixExpertStartIndex; i++) {
|
||||
totalTmp += groupListGM.GetValue(i);
|
||||
}
|
||||
}
|
||||
for (uint32_t groupIdx = workspaceSplitConfig.rightMatrixExpertStartIndex;
|
||||
groupIdx <= workspaceSplitConfig.rightMatrixExpertEndIndex; groupIdx++) {
|
||||
int64_t currM = 0;
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
currM = groupListGM.GetValue(groupIdx);
|
||||
} else {
|
||||
totalTmp += groupListGM.GetValue(groupIdx);
|
||||
currM = totalTmp;
|
||||
}
|
||||
int64_t tempM = currM - prevM;
|
||||
prevM = currM;
|
||||
if (curStartIdx >= 0 && curStartIdx - tempM < 0) {
|
||||
vecConfig.curGroupIdx = groupIdx;
|
||||
vecConfig.nextUpdateInterVal = tempM - curStartIdx;
|
||||
}
|
||||
curStartIdx -= tempM;
|
||||
}
|
||||
// 第三步 计算总数据量
|
||||
vecConfig.outLoopNum =
|
||||
(vecConfig.taskNum + gmmSwigluQuantV2->maxProcessRowNum - 1) / gmmSwigluQuantV2->maxProcessRowNum;
|
||||
vecConfig.tailLoopNum = vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum ?
|
||||
vecConfig.taskNum % gmmSwigluQuantV2->maxProcessRowNum :
|
||||
gmmSwigluQuantV2->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 * gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen * sizeof(float));
|
||||
pipe->InitBuffer(quantOutQueue, 1,
|
||||
gmmSwigluQuantV2->maxProcessRowNum * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR *
|
||||
sizeof(int8_t));
|
||||
pipe->InitBuffer(quantScaleOutQueue, 1,
|
||||
AlignUp<int32_t>(gmmSwigluQuantV2->maxProcessRowNum, ALIGN_8_ELE) * sizeof(float));
|
||||
pipe->InitBuffer(weightAuxiliaryMatrixInQueue, 1, gmmSwigluQuantV2->tokenLen * sizeof(float));
|
||||
// two 32 byte buffer for reduceMax calculation in Quant.
|
||||
pipe->InitBuffer(reduceWorkspace, gmmSwigluQuantV2->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>(gmmSwigluQuantV2->tokenLen * sizeof(float)), 0,
|
||||
0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
if (gmmSwigluQuantV2BaseParams->isSingleTensor == 0) {
|
||||
weightAuxiliaryMatrixGM.SetGlobalBuffer(
|
||||
GetTensorAddr<float>(vecConfig.curGroupIdx, weightAuxiliaryMatrixTensorPtr));
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal, weightAuxiliaryMatrixGM, copyAuxiliaryMatrixParams, padParams);
|
||||
} else {
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal,
|
||||
weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwigluQuantV2->tokenLen],
|
||||
copyAuxiliaryMatrixParams, padParams);
|
||||
}
|
||||
weightAuxiliaryMatrixInQueue.EnQue(weightAuxiliaryMatrixLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void GMMA8W4PostProcess::UpdateAuxiliaryMatrix(uint32_t loopIdx, VecConfig &vecConfig)
|
||||
{
|
||||
// 更新weightAuxiliaryMatrix
|
||||
if (unlikely(vecConfig.nextUpdateInterVal == 0)) {
|
||||
int64_t loop = gmmSwigluQuantV2->groupListLen - vecConfig.curGroupIdx;
|
||||
while (loop--) {
|
||||
if (gmmSwigluQuantV2BaseParams->groupListType == 0) {
|
||||
int64_t curTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextTemp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextTemp != curTemp) {
|
||||
vecConfig.nextUpdateInterVal = nextTemp - curTemp;
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
vecConfig.curGroupIdx++;
|
||||
int64_t nextUpdateInterValTmp = groupListGM.GetValue(vecConfig.curGroupIdx);
|
||||
if (nextUpdateInterValTmp != 0) {
|
||||
vecConfig.nextUpdateInterVal = nextUpdateInterValTmp;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
LocalTensor<float> weightAuxiliaryMatrixLocal = weightAuxiliaryMatrixInQueue.DeQue<float>();
|
||||
DataCopyExtParams copyParams{1, static_cast<uint32_t>(gmmSwigluQuantV2->tokenLen * sizeof(float)), 0, 0, 0};
|
||||
DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
|
||||
DataCopyPad(weightAuxiliaryMatrixLocal,
|
||||
weightAuxiliaryMatrixGM[vecConfig.curGroupIdx * gmmSwigluQuantV2->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 = mmOutQueue.AllocTensor<half>();
|
||||
LocalTensor<float> quantScaleLocal = quantScaleOutQueue.AllocTensor<float>();
|
||||
LocalTensor<int8_t> quantLocal = quantOutQueue.AllocTensor<int8_t>();
|
||||
|
||||
mmOutQueue.EnQue(mmLocal);
|
||||
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 :
|
||||
gmmSwigluQuantV2->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 = mmOutQueue.DeQue<half>();
|
||||
quantScaleLocal = quantScaleOutQueue.DeQue<float>();
|
||||
quantLocal = quantOutQueue.DeQue<int8_t>();
|
||||
|
||||
weightAuxiliaryMatrixInQueue.FreeTensor(weightAuxiliaryMatrixLocal);
|
||||
mmOutQueue.FreeTensor(mmLocal);
|
||||
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 * gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR * sizeof(int8_t)), 0,
|
||||
0};
|
||||
DataCopyPad(quantOutputGM[(workspaceSplitConfig.leftMatrixStartIndex + vecConfig.startIdx) *
|
||||
gmmSwigluQuantV2->tokenLen / SWIGLU_REDUCE_FACTOR],
|
||||
quantLocal, copyParams_1);
|
||||
|
||||
vecConfig.startIdx += vecConfig.innerLoopNum;
|
||||
vecConfig.startOffset = vecConfig.startIdx * gmmSwigluQuantV2->tokenLen;
|
||||
quantOutQueue.EnQue(quantLocal);
|
||||
quantScaleOutQueue.EnQue(quantScaleLocal);
|
||||
}
|
||||
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_POST_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_v2_a8w4_msd_pre.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H
|
||||
|
||||
#include "grouped_matmul_swiglu_quant_v2_utils.h"
|
||||
#include "kernel_operator.h"
|
||||
|
||||
#ifdef GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN);
|
||||
__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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParams;
|
||||
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<float> xRowSumTensor;
|
||||
|
||||
GlobalTensor<int8_t> xGM;
|
||||
GlobalTensor<int8_t> yGm;
|
||||
GlobalTensor<int8_t> yGm1;
|
||||
GlobalTensor<int8_t> yGm2;
|
||||
|
||||
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 GMMSwigluQuantV2BaseParams *__restrict gmmSwigluQuantV2BaseParamsIN)
|
||||
{
|
||||
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);
|
||||
gmmSwigluQuantV2BaseParams = gmmSwigluQuantV2BaseParamsIN;
|
||||
vK = gmmSwigluQuantV2BaseParams->K;
|
||||
groupNum = static_cast<uint32_t>(gmmSwigluQuantV2BaseParams->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 * SIZE_OF_HALF_2); // 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 * SIZE_OF_HALF_2;
|
||||
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 GroupedMatmulDequantSwigluQuant
|
||||
#endif // GMM_SWIGLU_QUANT_V2_A8W4_MSD
|
||||
#endif // OP_KERNEL_GROUPED_MATMUL_SWIGLU_QUANT_V2_A8W4_MSD_PRE_H
|
||||
@@ -0,0 +1,66 @@
|
||||
/**
|
||||
* 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_v2_apt.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#if ORIG_DTYPE_X_SCALE == DT_FLOAT8_E8M0
|
||||
#include "arch35/grouped_matmul_swiglu_quant_v2_mxquant.h"
|
||||
#elif ORIG_DTYPE_X_SCALE == DT_FLOAT
|
||||
#include "arch35/grouped_matmul_swiglu_quant_v2_pertoken_quant.h"
|
||||
#endif
|
||||
#include "arch35/grouped_matmul_swiglu_quant_v2_tiling_key.h"
|
||||
|
||||
#define FLOAT_OVERFLOW_MODE_CTRL 60
|
||||
|
||||
using namespace AscendC;
|
||||
using namespace matmul;
|
||||
|
||||
template <int8_t QUANT_B_TRANS, int8_t QUANT_A_TRANS>
|
||||
__global__ __aicore__ void grouped_matmul_swiglu_quant_v2(GM_ADDR x, GM_ADDR xScale, GM_ADDR groupList, GM_ADDR weight,
|
||||
GM_ADDR weightScale, GM_ADDR weightAssistanceMatrix,
|
||||
GM_ADDR bias, GM_ADDR smoothScale, GM_ADDR y, GM_ADDR yScale,
|
||||
GM_ADDR workspace, GM_ADDR tiling)
|
||||
{
|
||||
TPipe tPipe;
|
||||
GM_ADDR userWorkspace = GetUserWorkspace(workspace);
|
||||
int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();
|
||||
// enable overflow mode to avoid nan/inf value
|
||||
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
|
||||
#if ORIG_DTYPE_X_SCALE == DT_FLOAT8_E8M0
|
||||
if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_NO_TRANS && QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = false
|
||||
GmmSwigluAswt<Cgmct::Gemm::layout::RowMajor, Cgmct::Gemm::layout::RowMajor>(
|
||||
x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace,
|
||||
tiling);
|
||||
} else if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_TRANS && QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = true
|
||||
GmmSwigluAswt<Cgmct::Gemm::layout::RowMajor, Cgmct::Gemm::layout::ColumnMajor>(
|
||||
x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace,
|
||||
tiling);
|
||||
}
|
||||
#elif ORIG_DTYPE_X_SCALE == DT_FLOAT
|
||||
if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_NO_TRANS &&
|
||||
QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = false
|
||||
GmmSwigluAswtPertoken<Cgmct::Gemm::layout::RowMajor, Cgmct::Gemm::layout::RowMajor>(
|
||||
x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace,
|
||||
tiling, &tPipe);
|
||||
} else if (QUANT_B_TRANS == GMM_SWIGLU_QUANT_TRANS &&
|
||||
QUANT_A_TRANS == GMM_SWIGLU_QUANT_NO_TRANS) { // transX = false, transW = true
|
||||
GmmSwigluAswtPertoken<Cgmct::Gemm::layout::RowMajor, Cgmct::Gemm::layout::ColumnMajor>(
|
||||
x, weight, weightScale, xScale, weightAssistanceMatrix, smoothScale, groupList, y, yScale, workspace,
|
||||
tiling, &tPipe);
|
||||
}
|
||||
#endif
|
||||
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(oriOverflowMode);
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
/**
|
||||
* 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_v2_utils.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef OP_KERNEL_GROUPED_MATMUL_DEQUANT_SWIGLU_QUANT_V2_UTILS_H
|
||||
#define OP_KERNEL_GROUPED_MATMUL_DEQUANT_SWIGLU_QUANT_V2_UTILS_H
|
||||
|
||||
// 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_V2_A8W4_MSD
|
||||
using DTYPE_X_A8W4_MSD = AscendC::int4b_t;
|
||||
// A4W4 场景
|
||||
#elif defined(ORIG_DTYPE_X) && defined(DT_INT4) && ORIG_DTYPE_X == DT_INT4 && defined(ORIG_DTYPE_WEIGHT) && \
|
||||
defined(DT_INT4) && ORIG_DTYPE_WEIGHT == DT_INT4
|
||||
#define GMM_SWIGLU_QUANT_V2_A4W4
|
||||
// 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_V2_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格式分类
|
||||
|
||||
|
||||
namespace GroupedMatmulDequantSwigluQuant {
|
||||
using namespace AscendC;
|
||||
|
||||
|
||||
constexpr uint32_t UB_BLOCK_UNIT_SIZE = 32;
|
||||
constexpr uint32_t FLOAT_UB_BLOCK_UNIT_SIZE = 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 uint32_t NUM_2 = 2;
|
||||
constexpr int64_t SWIGLU_REDUCE_FACTOR = 2;
|
||||
constexpr int64_t DOUBLE_BUFFER = 2;
|
||||
constexpr int64_t DOUBLE_ROW = 2;
|
||||
constexpr int64_t SIZE_OF_HALF_2 = 2;
|
||||
constexpr uint8_t NUM_8 = 8;
|
||||
constexpr float QUANT_SCALE_INT8 = 127.0f;
|
||||
|
||||
constexpr MatmulConfig matmulCFGUnitFlag{false, false, true, 0, 0, 0, false, false, false, false, false, 0, 0, 0,
|
||||
0, 0, 0, 0, false};
|
||||
constexpr MatmulConfig NZ_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, false);
|
||||
constexpr MatmulConfig CUSTOM_CFG_MDL = GetMDLConfig(false, false, 0, true, false, false, true);
|
||||
|
||||
template <class AT_, class BT_, class CT_, class BiasT_, const MatmulConfig& MMCFG_>
|
||||
struct MMImplType {
|
||||
using AT = AT_;
|
||||
using BT = BT_;
|
||||
using CT = CT_;
|
||||
using BiasT = BiasT_;
|
||||
using MT = matmul::MatmulImpl<AT, BT, CT, BiasT, MMCFG_>;
|
||||
};
|
||||
|
||||
template <class AT_, class BT_, class CT_>
|
||||
struct MMImplTypeCustom {
|
||||
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 {
|
||||
uint32_t m = 0;
|
||||
uint32_t k = 0;
|
||||
uint32_t n = 0;
|
||||
uint32_t baseM = 0;
|
||||
uint32_t baseN = 0;
|
||||
uint32_t baseK = 0;
|
||||
uint32_t mIdx = 0;
|
||||
uint32_t nIdx = 0;
|
||||
uint32_t blockDimM = 0;
|
||||
uint32_t blockDimN = 0;
|
||||
uint32_t singleM = 0;
|
||||
uint32_t singleN = 0;
|
||||
uint64_t wBaseOffset = 0;
|
||||
uint64_t mAxisBaseOffset = 0;
|
||||
uint64_t nAxisBaseOffset = 0;
|
||||
uint64_t xBaseOffset = 0;
|
||||
uint64_t yBaseOffset = 0;
|
||||
uint64_t wOutOffset = 0;
|
||||
uint64_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 nextUpdateInterVal = 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_ADDR smoothScaleGM; // 平滑缩放因子
|
||||
// 输出 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 <>
|
||||
__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>();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline __gm__ T* GetTensorAddr(uint16_t index, GM_ADDR tensorPtr)
|
||||
{
|
||||
__gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorPtr);
|
||||
uint64_t tensorPtrOffset = *dataAddr;
|
||||
|
||||
__gm__ uint64_t* retPtr = dataAddr + (tensorPtrOffset >> 3);
|
||||
return reinterpret_cast<__gm__ T*>(*(retPtr + index));
|
||||
}
|
||||
} // namespace GroupedMatmulDequantSwigluQuant
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user