EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

408
csrc/attention/compressor/op_kernel/arch32/compressor_block_cube_perf.h Normal file
View File

@@ -0,0 +1,408 @@
/**
 * 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 compressor_block_cube_perf.h
* \brief
*/
#ifndef COMPRESSOR_BLOCK_CUBE_PERF_H
#define COMPRESSOR_BLOCK_CUBE_PERF_H
#include "compressor_comm.h"
#include "compressor_tools.h"
using namespace AscendC;
namespace Compressor {
template<typename COMP> class CompressorBlockCubePerf {
using MM1_OUT_T = float;
public:
__aicore__ inline CompressorBlockCubePerf(){};
__aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools);
__aicore__ inline void Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut);
__aicore__ inline void InitBuffers(TPipe *pipe);
__aicore__ inline void InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm);
__aicore__ inline void AllocEventID(TPipe *pipe);
__aicore__ inline void FreeEventID(TPipe *pipe);
__aicore__ inline void ComputeMm1(const RunInfo &info);
private:
using T = float;
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
__aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId);
__aicore__ inline void CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor, uint32_t hIdx, uint32_t kBase);
__aicore__ inline void CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
uint32_t hIdx, uint32_t kBase, uint32_t coffId);
__aicore__ inline void LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor, LocalTensor<X_T> xL1Tensor,
uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize);
__aicore__ inline void LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
uint32_t kStart, uint32_t kBase);
__aicore__ inline void MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C);
__aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
uint32_t coffId, uint32_t mStart, uint32_t mDealSize);
ConstInfo constInfo_ = {};
CompressorTools<COMP> tools_;
// GM
GlobalTensor<X_T> xGm_;
GlobalTensor<X_T> wkvGm_;
GlobalTensor<X_T> wgateGm_;
GlobalTensor<MM1_OUT_T>kvMm1ResGm;
GlobalTensor<MM1_OUT_T>scoreMm1ResGm;
GlobalTensor<int32_t> cuSeqlensGm_;
GlobalTensor<int32_t> sequsedGm_;
GlobalTensor<int32_t> startPosGm_;
bool isExistSeqUsed = false;
// =================================L1 Buffer=================================
static constexpr uint32_t L1_X_SIZE = 128 * 1024;
static constexpr uint32_t L1_W_SIZE = 64 * 1024;
// L1 Buffer
TBuf<TPosition::A1> xBufL1;
TBuf<TPosition::A1> wBufL1;
// =================================L0 Buffer=================================
// L0 buffer size
static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k
// L0_A
TBuf<TPosition::A2> tmpBufL0A;
// L0_B
TBuf<TPosition::B2> tmpBufL0B;
// L0_C
TBuf<TPosition::CO1> tmpBufL0C;
// =================================Event&Buffer ID===========================
// mte2 <> mte1 EventID
static constexpr uint32_t X_EVENT0 = EVENT_ID0;
static constexpr uint32_t X_EVENT1 = EVENT_ID1;
uint32_t xBufId = 0; // 用于DB计数
static constexpr uint32_t W_EVENT0 = EVENT_ID4;
static constexpr uint32_t W_EVENT1 = EVENT_ID5;
static constexpr uint32_t W_EVENT2 = EVENT_ID6;
static constexpr uint32_t W_EVENT3 = EVENT_ID7;
uint32_t wBufId = 0; // 用于DB计数
// mte1 <> mmad EventID
static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3;
static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4;
uint32_t l0abBufId = 0;
// mmad <> fixpipe EventID
static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID
static constexpr uint32_t L0C_EVENT1 = EVENT_ID1;
uint32_t l0cBufId = 0;
// =================================Loop======================================
uint32_t curBIdx_ = 0;
uint32_t curSIdx_ = 0;
};
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools)
{
this->constInfo_ = constInfo;
this->tools_ = tools;
}
template <typename COMP> __aicore__ inline void CompressorBlockCubePerf<COMP>::Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut)
{
xGm_.SetGlobalBuffer((__gm__ X_T *)x);
wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv);
wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate);
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
isExistSeqUsed = (seqUsed != nullptr);
if (isExistSeqUsed) {
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
}
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
}
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitBuffers(TPipe *pipe)
{
// L1
// 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K
// 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K
pipe->InitBuffer(xBufL1, L1_X_SIZE * 2);
// dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K
// cur和pre循环使用, 2份buffer就足够
pipe->InitBuffer(wBufL1, L1_W_SIZE * 4);
// L0
pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2);
pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2);
pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm)
{
this->kvMm1ResGm = kvMm1ResGm;
this->scoreMm1ResGm = scoreMm1ResGm;
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::AllocEventID(TPipe *pipe)
{
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
SetFlag<HardEvent::FIX_M>(L0C_EVENT0);
SetFlag<HardEvent::FIX_M>(L0C_EVENT1);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::FreeEventID(TPipe *pipe)
{
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0);
WaitFlag<HardEvent::FIX_M>(L0C_EVENT1);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor,
uint32_t hIdx, uint32_t kBase)
{
uint32_t tStart = tools_.GetTIdxByBatch(info.bStart) + info.sStart; // 此基本块在整个序列中的位置
uint32_t copySeqCnt = info.dealSeqCnt; // 此基本块处理的长度
uint32_t xL1Offset = 0 * (32 / sizeof(X_T));
uint64_t sIdx = tStart; // 起始s在整个T的起始点
uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx;
uint32_t nValue = copySeqCnt;
uint32_t dValue = kBase; // 拷贝的列数kBase
uint32_t srcDValue = constInfo_.hSize;
uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移，需要16对齐
CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
uint32_t hIdx, uint32_t kBase, uint32_t coffId)
{
// coffId=0, 搬运左矩阵的数据; coffId=1, 搬运右矩阵的数据
uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.dIdx * constInfo_.hSize + hIdx;
uint32_t wkvL1Offset = 0;
uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B
uint32_t nValue = constInfo_.dBaseSize;
uint32_t dValue = kBase;
uint32_t srcDValue = constInfo_.hSize;
uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor,
LocalTensor<X_T> xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize)
{
uint32_t mSize = info.dealSeqCnt;
uint32_t mSizeAlign = Align(mSize, 16U);
uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T));
uint32_t mLoop = Align(mDealSize, 16U) / 16;
for (uint32_t i = 0; i < mLoop; i++) {
LoadData2DParams loadData2DParams;
loadData2DParams.startIndex = i;
loadData2DParams.repeatTimes = kBase / (32 / sizeof(X_T));
loadData2DParams.srcStride = mSizeAlign / 16;
loadData2DParams.dstGap = 0;
loadData2DParams.ifTranspose = false;
LoadData(aL0Tensor[i * 16 * kBase], xL1Tensor[xTensorOffset], loadData2DParams); // 16: 一个分型的行数
}
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
uint32_t kStart, uint32_t kBase)
{
uint32_t rowCnt = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
uint64_t wTensorOffset = rowCnt * kStart;
LoadData2DParams loadData2DParams;
loadData2DParams.startIndex = 0;
loadData2DParams.repeatTimes = (rowCnt / 16) * (kBase / (32 / sizeof(X_T)));
loadData2DParams.srcStride = 1;
loadData2DParams.dstGap = 0;
loadData2DParams.ifTranspose = false;
LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParams);
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C)
{
MmadParams mmadParams;
mmadParams.m = (mActSize + 15) / 16 * 16;
mmadParams.n = nDealSize;
mmadParams.k = kActSize;
mmadParams.cmatrixInitVal = isInitL0C;
mmadParams.cmatrixSource = false;
Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
PipeBarrier<PIPE_M>();
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
uint32_t coffId, uint32_t mStart, uint32_t mDealSize)
{
// coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据
FixpipeParamsV220 fixParams;
fixParams.mSize = mDealSize;
fixParams.nSize = constInfo_.dBaseSize;
fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐
fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim;
fixParams.ndNum = 1;
uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize;
uint64_t gmOffset = coffId * constInfo_.headDim + constInfo_.dIdx + mStart * fixParams.dstStride + dbOffset;
uint32_t kvOffset = 0;
uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * constInfo_.dBaseSize;
Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams);
Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams);
}
template <typename COMP>
__aicore__ inline uint32_t CompressorBlockCubePerf<COMP>::GetMSize(const RunInfo &info, uint32_t coffId)
{
return info.dealSeqCnt;
}
template <typename COMP>
__aicore__ inline void CompressorBlockCubePerf<COMP>::ComputeMm1(const RunInfo &info)
{
static constexpr uint32_t K_SIZE = 512;
static constexpr uint32_t K_L1_BASE = 256;
static constexpr uint32_t M_L0_BASE = 128;
static constexpr uint32_t K_L0_BASE = 128;
uint32_t nCoff = (uint32_t)COMP::coff;
// hSize为K_SIZE=512的倍数
uint32_t hSize = constInfo_.hSize;
uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同
for (uint32_t h = 0; h < hSize; h += K_SIZE) {
for (uint32_t k = 0; k < K_SIZE; k += K_L1_BASE) {
bool isFirst = (h == 0 && k == 0);
bool isLast = ((h + K_SIZE >= hSize) && (k + K_L1_BASE >= K_SIZE));
uint32_t hIdx = (h + k + hIdxStart) % hSize; // h方向错位搬运
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
LocalTensor<X_T> xL1Tensor = xBufL1.GetWithOffset<X_T>(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE);
CopyXGmToL1(info, xL1Tensor, hIdx, K_L1_BASE);
SetFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
WaitFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
for (uint32_t i = nCoff; i > 0; i--) {
// coffId=0, 计算pre数据; coffId=1, 计算cur数据
uint32_t coffId = i - 1;
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
LocalTensor<X_T> wL1Tensor = wBufL1.GetWithOffset<X_T>(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE);
CopyWeightGmToL1(wL1Tensor, hIdx, K_L1_BASE, coffId);
SetFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
WaitFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
uint32_t mSize = GetMSize(info, coffId);
uint32_t actMDealSize = M_L0_BASE;
for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) {
if (mL0 + M_L0_BASE > mSize) {
actMDealSize = mSize - mL0;
}
l0cBufId = coffId + (mL0 / M_L0_BASE);
LocalTensor<T> cL0Tensor = tmpBufL0C.GetWithOffset<T>((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE);
if (isFirst) {
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
}
uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
for (uint32_t kL0 = 0; kL0 < K_L1_BASE; kL0 += K_L0_BASE) {
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
LocalTensor<X_T> aL0Tensor = tmpBufL0A.GetWithOffset<X_T>(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE);
LocalTensor<X_T> bL0Tensor = tmpBufL0B.GetWithOffset<X_T>(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE);
LoadAToL0(info, aL0Tensor, xL1Tensor, kL0, K_L0_BASE, mL0, actMDealSize);
LoadBToL0(bL0Tensor, wL1Tensor, kL0, K_L0_BASE);
SetFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
WaitFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
bool isInitL0C = isFirst && (kL0 == 0);
MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, nDealSize, K_L0_BASE, isInitL0C);
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
l0abBufId = (l0abBufId + 1) % 2;
}
if (isLast) {
SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize);
SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
}
}
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
wBufId = (wBufId + 1) % 4;
}
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
xBufId = (xBufId + 1) % 2;
}
}
}
} // namespace Compressor
#endif // COMPRESSOR_BLOCK_CUBE_PERF_H

1488
csrc/attention/compressor/op_kernel/arch32/compressor_block_vec_perf.h Normal file
View File

File diff suppressed because it is too large Load Diff

341
csrc/attention/compressor/op_kernel/arch32/compressor_comm.h Normal file
View File

@@ -0,0 +1,341 @@
/**
 * 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 compressor_comm.h
* \brief
*/
#ifndef COMPRESSOR_COMM_H
#define COMPRESSOR_COMM_H
#include "kernel_operator.h"
#include "kernel_operator_list_tensor_intf.h"
#include "kernel_tiling/kernel_tiling.h"
#include "lib/matmul_intf.h"
#include "lib/matrix/matmul/tiling.h"
using namespace AscendC;
namespace Compressor {
template <typename T>
__aicore__ inline T CeilDivT(T num1, T num2)
{
if (num2 == 0) {
return static_cast<T>(0);
}
return (num1 + num2 - 1) / num2;
}
template <typename T>
__aicore__ inline T Align(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd)));
}
template <typename T>
__aicore__ inline T Trunc(T num, T rnd)
{
return ((rnd) == 0) ? 0 : (((num) / (rnd) * (rnd)));
}
template <typename T>
__aicore__ inline T FloorPow2(T num)
{
if (num == 0) return 1;
for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
num |= (num >> i);
}
return num - (num >> 1);
}
template <typename T>
__aicore__ inline T CeilPow2(T num)
{
if (num <= 1) return 1;
num --;
for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
num |= (num >> i);
}
num ++;
return num;
}
enum class X_LAYOUT : std::uint8_t {
BSH = static_cast<std::uint8_t>(0),
TH = static_cast<std::uint8_t>(1)
};
enum class X_DTYPE : std::uint8_t {
BF16 = static_cast<std::uint8_t>(0),
FP16 = static_cast<std::uint8_t>(1)
};
enum class ROPE_DTYPE : std::uint8_t {
SAME_AS_X = static_cast<std::uint8_t>(0),
FP32 = static_cast<std::uint8_t>(1)
};
enum class COFF : std::uint8_t {
DISABLE = static_cast<std::uint8_t>(1),
OVERLAP = static_cast<std::uint8_t>(2)
};
enum class ROTARY_MODE : std::uint8_t {
HALF = static_cast<std::uint8_t>(1),
INTERLEAVE = static_cast<std::uint8_t>(2)
};
enum class CACHE_MODE : std::uint8_t {
CONTINUOUS = static_cast<std::uint8_t>(1),
CYCLE = static_cast<std::uint8_t>(2)
};
enum class TEMPLATE_ID : uint8_t {
NORMAL = 0,
EMPTY_X = 1,
PERF = 2
};
template <X_LAYOUT X_L, X_DTYPE X_T, ROPE_DTYPE R_T, COFF C, ROTARY_MODE Rotary_Mode, typename... Args>
struct COMPType {
static constexpr X_LAYOUT xLayout = X_L;
static constexpr X_DTYPE xDtype = X_T;
static constexpr ROPE_DTYPE ropeDtype = R_T;
static constexpr COFF coff = C;
static constexpr ROTARY_MODE rotaryMode = Rotary_Mode;
};
struct ConstInfo {
// 整个AICORE的任务信息, 左闭右开区间[ (bStart, s2Start), (bEnd, s2End) )
uint32_t bStart = 0U;
uint32_t sStart = 0U;
uint32_t bEnd = 0U;
uint32_t sEnd = 0U;
// 分核相关
uint32_t usedCoreNum = 0;
uint32_t dBaseSize = 0;
uint32_t mBaseSize = 0;
uint32_t tcSize = 0;
uint32_t tcBaseSize = 0;
uint32_t tcBasicBlockNum = 0;
uint32_t dBasicBlockNum = 0;
uint32_t coreGroupNum = 0;
uint32_t singleCoreDealTcBasicNum = 0;
uint32_t dIdx = 0;
uint32_t bIdxOfLastTc = 0;
uint32_t sIdxOfLastTc = 0;
// shape及参数
uint32_t batchSize = 0;
uint32_t hSize = 0;
uint32_t sSize = 0;
uint32_t headDim = 0;
uint32_t ropeHeadDim = 0;
uint32_t cmpRatio = 0;
float normEps = 1e-6;
float reciprocalD = 0;
uint32_t curGroupIdx = 0;
uint32_t tailGroupIdx = 0;
uint32_t tailBasicBlockNum = 0;
uint32_t realDealBasicBlockNum = 0;
// pageAttention
uint32_t blockNum = 0;
uint32_t blockSize = 0;
uint32_t maxBlockNumPerBatch = 0;
uint64_t stateCacheStrideDim0 = 0;
// workSpace
uint32_t dbWorkspaceRatio = 1;
uint32_t mm1KvResSize = 0;
uint32_t mm1ScoreResSize = 0;
uint32_t vec1TailCacheSize = 0;
uint32_t vec1ResSize = 0;
uint32_t mm1ResSize = 0; // 所有cube输出kv/score结果的总大小
uint32_t aiCoreIdx = 0;
uint32_t nSize = 0;
uint32_t dbSize = 0;
};
struct RunInfo {
bool isValid = false;
uint32_t cubeDbIdx = 0; // kernel主循环索引
// 增加字段
uint32_t dealTcNum = 0;
// 右边相关信息
uint32_t bStart = 0;
uint32_t sStart = 0;
uint32_t dealSeqCnt = 0;
// 左边相关信息
uint32_t preBStart = 0;
uint32_t preSStart = 0;
uint32_t preDealSeqCnt = 0; // 左边需要处理的s大小
uint32_t preFirstSeqCnt = 0; // 左边首块大小
uint32_t bEnd = 0;
uint32_t sEnd = 0;
uint32_t bStartSeqIdx = 0;
uint32_t bEndSeqIdx = 0;
// v2分核信息 sc是左闭右开
uint32_t scStart = 0;
uint32_t scEnd = 0;
uint32_t dealScSize = 0;
// vec1Res offset
uint64_t vec1ResOffset = 0;
};
struct Vec1RunInfo {
// vec相关信息，一次syncAll需处理数据的起始索引
bool resetResFlag = false; // v1积攒N轮 是否是N轮的起始轮
uint32_t c1v1DbIdx = 0; // vec1 doubleBuffer索引
uint32_t v1v2DbIdx = 0; // v1v2 doubleBuffer索引
uint32_t bStart = 0;
uint32_t sStart = 0;
uint32_t dealTcNum = 0;
uint32_t dealScSize = 0;
};
struct Vec2RunInfo {
// uint32_t bStart = 0;
uint32_t v2DbIdx = 0; // v2 doubleBuffer索引
uint32_t sStart = 0;
uint32_t bEnd = 0;
uint32_t sEnd = 0;
// v2分核信息 sc是左闭右开
uint32_t scStart = 0;
uint32_t scEnd = 0;
// uint32_t dealScSize = 0;
// 增加字段
uint32_t bStart = 0;
uint32_t compressedId = 0;
uint32_t bCompressedId = 0;
uint32_t dealScSize = 0;
};
struct MSplitInfo {
uint32_t vecStartB = 0U;
uint32_t vecStartS = 0U;
uint32_t vecEndB = 0U;
uint32_t vecEndS = 0U;
uint32_t dealTcNum = 0U;
// vec1Res offset
uint64_t vec1StartOffset = 0;
uint64_t vec1ResOffset = 0;
};
struct BlockInfo {
__aicore__ inline BlockInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealSeqSize) :
bIdx(bIdx), sIdx(sIdx), dealSeqSize(dealSeqSize) {};
uint32_t bIdx = 0U;
uint32_t sIdx = 0U;
uint32_t dealSeqSize = 0;
uint32_t isFirst = true;
uint32_t bSeqUsed = 0U;
uint32_t bStartPos = 0U;
uint32_t headHolderSeqCnt = 0U;
uint32_t validSeqCnt = 0U;
uint32_t tailHolderSeqCnt = 0U;
uint32_t dealTcSize = 0U;
uint32_t tailValidSeqCnt = 0U;
uint32_t compressTcSize = 0U;
};
// BUFFER的字节数
inline constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
inline constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
inline constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
inline constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
inline constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
inline constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
inline constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
inline constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
inline constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
inline constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
inline constexpr uint32_t BUFFER_SIZE_BYTE_64K = 65536;
// BLOCK和REPEAT的字节数
inline constexpr uint64_t BYTE_BLOCK = 32UL;
inline constexpr uint32_t REPEAT_BLOCK_BYTE = 256U;
// BLOCK和REPEAT的FP32元素数
inline constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); // 8
inline constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); // 64
inline constexpr uint32_t REPEAT_STRIDE_NUM = REPEAT_BLOCK_BYTE / BYTE_BLOCK; // 8
inline constexpr uint32_t REPEAT_MAX_NUM = 255;
inline constexpr uint32_t BRCB_NUM = 8;
inline constexpr uint32_t MAX_R = 256;
template <typename T>
__aicore__ inline void CopySingleMatrixNDToNZ(LocalTensor<T> l1Tensor, const GlobalTensor<T> gmTensor,
uint32_t nValue, uint32_t dValue, uint32_t srcDValue, uint32_t dstNzC0Stride)
{
Nd2NzParams nd2nzPara;
nd2nzPara.ndNum = 1;
nd2nzPara.nValue = nValue; // nd矩阵的行数
if constexpr (IsSameType<T, int4b_t>::value) {
constexpr uint32_t HALF_SIZE_DIVISOR = 2;
nd2nzPara.dValue = dValue / HALF_SIZE_DIVISOR;
nd2nzPara.srcDValue = srcDValue / HALF_SIZE_DIVISOR;
} else {
nd2nzPara.dValue = dValue; // nd矩阵的列数
nd2nzPara.srcDValue = srcDValue; // 同一nd矩阵相邻行起始地址间的偏移
}
nd2nzPara.dstNzC0Stride = dstNzC0Stride;
nd2nzPara.dstNzNStride = 1;
nd2nzPara.srcNdMatrixStride = 0;
nd2nzPara.dstNzMatrixStride = 0;
DataCopy(l1Tensor, gmTensor, nd2nzPara);
}
template <typename T>
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col)
{
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
AscendC::ShapeInfo shapeInfo(2, array2);
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
}
template <typename T>
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col)
{
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
AscendC::ShapeInfo shapeInfo(2, array2);
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
}
template <typename T>
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
{
uint32_t col = 32 / sizeof(T);
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
AscendC::ShapeInfo shapeInfo(2, array2);
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
}
template <typename T>
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
{
uint32_t col = 32 / sizeof(T);
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
AscendC::ShapeInfo shapeInfo(2, array2);
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
}
} // namespace Compressor
#endif

86
csrc/attention/compressor/op_kernel/arch32/compressor_kernel.h Normal file
View File

@@ -0,0 +1,86 @@
/**
 * 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 compressor_kernel.h
* \brief
*/
#ifndef COMPRESSOR_KERNEL
#define COMPRESSOR_KERNEL
#include "compressor_comm.h"
#include "compressor_template_tiling_key.h"
#include "compressor_kernel_perf.h"
#include "compressor_tiling_data.h"
#include "compressor_tools.h"
using namespace AscendC;
namespace Compressor {
template <typename COMP>
class CompressorKernel {
public:
__aicore__ inline CompressorKernel(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
: pipe_(pipe), tilingData_(tilingData) {}
__aicore__ inline void Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut,
__gm__ uint8_t *workspace);
__aicore__ inline void Process();
// ==============================TilingData&TPipe==============================
TPipe* pipe_;
const optiling::CompressorTilingData* __restrict tilingData_;
};
template <typename COMP>
__aicore__ inline void CompressorKernel<COMP>::Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut,
__gm__ uint8_t *workspace)
{
}
template <typename COMP>
__aicore__ inline void CompressorKernel<COMP>::Process()
{
}
} // namespace Compressor
#endif // COMPRESSOR_KERNEL

695
csrc/attention/compressor/op_kernel/arch32/compressor_kernel_perf.h Normal file
View File

@@ -0,0 +1,695 @@
/**
 * 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 compressor_kernel_perf.h
* \brief
*/
#ifndef COMPRESSOR_KERNEL_PERF_H
#define COMPRESSOR_KERNEL_PERF_H
#include "compressor_comm.h"
#include "compressor_template_tiling_key.h"
#include "compressor_tiling_data.h"
#include "compressor_tools.h"
#include "compressor_block_cube_perf.h"
#include "compressor_block_vec_perf.h"
using namespace AscendC;
namespace Compressor {
struct CmpBlockInfo {
__aicore__ inline CmpBlockInfo() {};
__aicore__ inline CmpBlockInfo(uint32_t bIdx, uint32_t sIdx, bool needReset = false) : bIdx(bIdx), sIdx(sIdx), needReset(needReset) {};
uint32_t bIdx = 0U;
uint32_t sIdx = 0U;
uint32_t bSeqUsed = 0U;
uint32_t bStartPos = 0U;
bool needReset = false;
bool isFirst = true;
uint32_t headSeqCnt = 0U;
uint32_t validSeqCnt = 0U;
uint32_t tailSeqCnt = 0U;
bool isCompress = 0U;
};
struct BasicBlockInfo {
uint32_t bIdx = 0;
uint32_t sIdx = 0;
uint32_t compressedTcNum = 0;
uint32_t dealSeqCnt = 0;
uint32_t dealTcNum = 0;
};
struct BatchInfo {
uint32_t tcNum = 0;
uint32_t compressedTcNum = 0;
uint32_t remSeqCnt = 0;
uint32_t seqCnt = 0;
uint32_t seqUsedCnt = 0;
uint32_t headHolderSeq = 0;
uint32_t bStartPos = 0;
uint32_t bIdx = 0;
uint32_t sIdx = 0;
};
template <typename COMP>
class CompressorKernelPerf {
public:
__aicore__ inline CompressorKernelPerf(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
: pipe_(pipe), tilingData_(tilingData) {}
__aicore__ inline void Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut,
__gm__ uint8_t *workspace);
__aicore__ inline void Process();
private:
// ================================Init functions==================================
__aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace);
// ================================Process functions================================
__aicore__ inline void InitTilingData();
__aicore__ inline void SetBaseSize();
// 获取基本块数量
__aicore__ inline uint32_t GetLoopTimes();
__aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo);
__aicore__ inline void UpdateCurGroup(BasicBlockInfo &basicBlockInfo, BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq);
__aicore__ inline BasicBlockInfo SkipOneLoop(BatchInfo &batchInfo);
// 计算分核基本信息
__aicore__ inline void CalcSplitCoreInfo();
__aicore__ inline void AllocEventID();
__aicore__ inline void FreeEventID();
__aicore__ inline void ComputeMm1(const RunInfo &info, bool isNeedExcute);
__aicore__ inline void ComputeVec1(const Vec1RunInfo &info);
__aicore__ inline void ComputeVec2(const Vec2RunInfo &info);
__aicore__ inline bool IsNeedExcuteC1(RunInfo info);
__aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx);
__aicore__ inline void CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx);
__aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info);
__aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info);
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
using T = float;
using MM1_OUT_T = T;
using VEC1_OUT_T = T;
// 常量
static constexpr uint64_t SYNC_MODE0 = 0;
static constexpr uint64_t SYNC_MODE2 = 2;
static constexpr uint32_t SYNC_C1_FLAG = 3;
static constexpr uint32_t SYNC_V1_FLAG = 4;
static constexpr uint32_t SYNC_V1_FLAG2 = 5;
static constexpr uint32_t SYNC_C1_V1_FLAG = 6;
static constexpr uint32_t SYNC_V1_C1_FLAG = 8;
// ==============================TilingData&TPipe==============================
TPipe* pipe_;
const optiling::CompressorTilingData* __restrict tilingData_;
// ===========================Workspace Global Tensor===========================
GlobalTensor<MM1_OUT_T> mm1KvResGm;
GlobalTensor<MM1_OUT_T> mm1ScoreResGm;
GlobalTensor<MM1_OUT_T> vec1KvCacheGm;
GlobalTensor<MM1_OUT_T> vec1ScoreCacheGm;
GlobalTensor<MM1_OUT_T> Vec1InputKvGm;
GlobalTensor<MM1_OUT_T> Vec1InputScoreGm;
GlobalTensor<VEC1_OUT_T> vec1ResGm;
GlobalTensor<VEC1_OUT_T> vec2InputGm;
// ================================Task Info====================================
CompressorTools<COMP> tools_;
ConstInfo constInfo{};
uint32_t aiCoreIdx = 0;
// ==============================Service Define==============================
CompressorBlockCubePerf<COMP> blockCube_;
CompressorBlockVectorPerf<COMP> blockVec_;
uint32_t allCompressedTcNum_ = 0;
uint32_t curCompressedTcNum_ = 0;
uint32_t accDealSize = 0;
uint32_t loopTimes = 0;
uint32_t cubeLoop = 0;
uint32_t vec1Loop = 0;
uint32_t vec2Loop = 0;
bool isFirstUpdateCurGroup = true;
};
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::Init(
__gm__ uint8_t *x,
__gm__ uint8_t *wKv,
__gm__ uint8_t *wGate,
__gm__ uint8_t *stateCache,
__gm__ uint8_t *ape,
__gm__ uint8_t *normWeight,
__gm__ uint8_t *ropeSin,
__gm__ uint8_t *ropeCos,
__gm__ uint8_t *stateBlockTable,
__gm__ uint8_t *cuSeqlens,
__gm__ uint8_t *seqUsed,
__gm__ uint8_t *startPos,
__gm__ uint8_t *cmpKvOut,
__gm__ uint8_t *workspace)
{
if ASCEND_IS_AIV {
constInfo.aiCoreIdx = GetBlockIdx() / 2;
} else {
constInfo.aiCoreIdx = GetBlockIdx();
}
InitTilingData();
// init tools
tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize;
tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio;
tools_.Init(startPos, seqUsed, cuSeqlens);
// 剔除尾部的无效batch
for (; constInfo.batchSize > 0; --constInfo.batchSize) {
uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1);
if (bSeqUsed > 0) {
break;
}
}
// 所有batch的有效序列都为0时, 直接退出
if (constInfo.batchSize == 0) {
return;
}
// 0. 计算最后一个Tc块的起始位置
constInfo.bIdxOfLastTc = constInfo.batchSize - 1;
// 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息
SetBaseSize(); // 设置基本块大小
CalcSplitCoreInfo();
// 2. 计算循环次数
loopTimes = GetLoopTimes();
// 3. 初始化workspace
InitWorkspace(workspace);
// 4. 初始化block层
if ASCEND_IS_AIC {
#if __CCE_AICORE__ == 310
blockCube_.InitParams(constInfo, tools_);
#else
blockCube_.InitParams(constInfo, tools_);
#endif
blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos,
stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);
blockCube_.InitBuffers(pipe_);
#if __CCE_AICORE__ == 310
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
#else
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
#endif
} else {
blockVec_.InitParams(constInfo, tools_);
blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable,
cuSeqlens, seqUsed, startPos, cmpKvOut);
blockVec_.InitBuffers(pipe_);
#if __CCE_AICORE__ == 310
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
#else
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
#endif
}
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::InitTilingData() {
constInfo.cmpRatio = tilingData_->baseParams.cmpRatio;
constInfo.batchSize = tilingData_->baseParams.batchSize;
constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize;
constInfo.headDim = tilingData_->baseParams.headDim;
constInfo.hSize = tilingData_->baseParams.hiddenSize;
constInfo.sSize = tilingData_->baseParams.seqSize;
constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim;
constInfo.normEps = tilingData_->baseParams.normEps;
constInfo.reciprocalD = tilingData_->baseParams.reciprocalD;
constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum;
constInfo.blockNum = tilingData_->pageAttentionParams.blockNum;
constInfo.blockSize = tilingData_->pageAttentionParams.blockSize;
constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch;
constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0;
constInfo.nSize = tilingData_->baseParams.nSize;
constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize;
constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::SetBaseSize()
{
uint32_t mSize = 0;
uint32_t minMBaseSize = 0;
bool sameSeqUsed = true;
uint32_t firstBatchSeqUsed = tools_.GetSeqLength(0);
for (uint32_t i = 0; i < constInfo.batchSize; i++) {
uint32_t bSeqUsed = tools_.GetSeqLength(i);
uint32_t bStartPos = tools_.GetStartPos(i);
// 获取m大小
mSize += bSeqUsed;
// 获取是否等长
if (sameSeqUsed && (bSeqUsed != firstBatchSeqUsed)) {
sameSeqUsed = false;
}
// 获取m轴最小切分大小
if (minMBaseSize != constInfo.cmpRatio) {
uint32_t startCmpIdx = bStartPos / constInfo.cmpRatio;
uint32_t endCmpIdx = (bStartPos + bSeqUsed) / constInfo.cmpRatio;
if (startCmpIdx == endCmpIdx) {
if (bSeqUsed > minMBaseSize) {
minMBaseSize = bSeqUsed;
}
} else if (startCmpIdx + 1 == endCmpIdx) {
uint32_t startCmpValidSeqCnt = constInfo.cmpRatio - (bStartPos % constInfo.cmpRatio);
uint32_t endCmpValidSeqCnt = (bStartPos + bSeqUsed) % constInfo.cmpRatio;
if (startCmpValidSeqCnt > minMBaseSize) {
minMBaseSize = startCmpValidSeqCnt;
}
if (endCmpValidSeqCnt > minMBaseSize) {
minMBaseSize = endCmpValidSeqCnt;
}
} else {
minMBaseSize = constInfo.cmpRatio;
}
}
}
uint32_t aiCoreNum = constInfo.usedCoreNum;
constInfo.dBaseSize = 64;
uint32_t dBaseBlockNum = constInfo.headDim / constInfo.dBaseSize;
if (sameSeqUsed && mSize <= (constInfo.mBaseSize * (aiCoreNum / dBaseBlockNum))) {
if constexpr (COMP::coff == COFF::OVERLAP) {
if (constInfo.headDim == 128) {
dBaseBlockNum = 8;
} else if (constInfo.headDim == 512) {
dBaseBlockNum = 16;
}
} else {
if (constInfo.headDim == 128) {
dBaseBlockNum = 8;
} else if (constInfo.headDim == 512) {
dBaseBlockNum = 16;
}
}
// 核数足够时, 修改才生效
if (aiCoreNum >= dBaseBlockNum) {
constInfo.dBaseSize = constInfo.headDim / dBaseBlockNum;
// 开启全核
uint32_t coreGroupNum = aiCoreNum / dBaseBlockNum;
uint32_t newMBaseSize = (constInfo.batchSize + coreGroupNum - 1) / coreGroupNum * firstBatchSeqUsed;
if (newMBaseSize > minMBaseSize && newMBaseSize < constInfo.mBaseSize) {
constInfo.mBaseSize = newMBaseSize;
}
}
}
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::SkipInvalidBatch(BatchInfo &batchInfo)
{
for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) {
batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx);
if (batchInfo.seqCnt > 0) {
break;
}
}
batchInfo.remSeqCnt = batchInfo.seqCnt;
if (tools_.isExistSeqUsed_) {
batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx);
} else {
batchInfo.seqUsedCnt = batchInfo.seqCnt;
}
if (batchInfo.bIdx < constInfo.batchSize) {
batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx);
batchInfo.sIdx = 0;
batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1);
batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
}
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::UpdateCurGroup(BasicBlockInfo &basicBlockInfo,
BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq)
{
// 更新当前组的信息
if (curGroupQuota == 0 && !isFirstUpdateCurGroup) {
return;
}
isFirstUpdateCurGroup = false;
basicBlockInfo.bIdx = batchInfo.bIdx;
uint32_t curGroupDealSeq = curGroupQuota < curDealSeq ? curGroupQuota : curDealSeq;
basicBlockInfo.sIdx = batchInfo.sIdx + curGroupDealSeq;
basicBlockInfo.dealSeqCnt += curGroupDealSeq;
curGroupQuota -= curGroupDealSeq;
// 结尾需要跳batch，需要考虑在当前组起始为末尾，或者当前组起始大于整个M轴
if ((curGroupQuota == 0 || basicBlockInfo.bIdx == constInfo.batchSize - 1) && basicBlockInfo.sIdx == batchInfo.seqCnt) {
basicBlockInfo.sIdx = 0;
for (basicBlockInfo.bIdx++; basicBlockInfo.bIdx < constInfo.batchSize; ++basicBlockInfo.bIdx) {
uint32_t seqCnt = tools_.GetSeqLength(basicBlockInfo.bIdx);
if (seqCnt > 0) {
break;
}
}
}
}
template <typename COMP>
__aicore__ inline BasicBlockInfo CompressorKernelPerf<COMP>::SkipOneLoop(BatchInfo &batchInfo)
{
BasicBlockInfo basicBlockInfo{};
isFirstUpdateCurGroup = true;
uint32_t curGroupQuota = constInfo.mBaseSize * constInfo.curGroupIdx; // m轴当前组起始
bool curGroupStartFlag = false;
uint32_t quota = constInfo.coreGroupNum * constInfo.mBaseSize;
for (; batchInfo.bIdx < constInfo.batchSize;) {
uint32_t curDealSeq = 0;
uint32_t curDealTcNum = 0;
uint32_t curDealCompressedTcNum = 0;
// 无法处理完当前整个batch
if (quota < batchInfo.remSeqCnt) {
// 向下对齐r，
if (quota > constInfo.cmpRatio - batchInfo.headHolderSeq) {
uint32_t delta = (batchInfo.bStartPos + batchInfo.sIdx + quota) & (constInfo.cmpRatio - 1); // 超出对齐的部分
curDealSeq = quota - delta;
quota -= curDealSeq;
curDealTcNum = (curDealSeq + constInfo.cmpRatio - 1) / constInfo.cmpRatio;
curDealCompressedTcNum = min(curDealTcNum, batchInfo.compressedTcNum);
// 更新当前组所需信息
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
// 更新batch信息
batchInfo.remSeqCnt = batchInfo.remSeqCnt - curDealSeq;
batchInfo.sIdx = batchInfo.sIdx + curDealSeq;
batchInfo.compressedTcNum -= curDealCompressedTcNum;
batchInfo.tcNum -= curDealTcNum;
// 更新loop信息
basicBlockInfo.dealTcNum += curDealTcNum;
basicBlockInfo.compressedTcNum += curDealCompressedTcNum;
}
break;
} else {
// 处理整个batch
quota -= batchInfo.remSeqCnt;
curDealSeq = batchInfo.remSeqCnt;
curDealTcNum = batchInfo.tcNum;
// 更新当前组所需信息
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
// 更新batch和loop信息
batchInfo.remSeqCnt = 0;
basicBlockInfo.dealTcNum += batchInfo.tcNum;
basicBlockInfo.compressedTcNum += batchInfo.compressedTcNum;
batchInfo.bIdx++;
SkipInvalidBatch(batchInfo);
}
}
uint32_t totalDataSize = constInfo.coreGroupNum * constInfo.mBaseSize - quota;
// 2. 当前组的起始偏移
uint32_t currentGroupStart = constInfo.curGroupIdx * constInfo.mBaseSize;
// 3. 安全判断
if (currentGroupStart >= totalDataSize) {
// 超出尾块
basicBlockInfo.dealSeqCnt = 0;
} else {
// 还在有效范围内，计算剩余量
uint32_t remaining = totalDataSize - currentGroupStart;
basicBlockInfo.dealSeqCnt = (remaining < constInfo.mBaseSize) ? remaining : constInfo.mBaseSize;
}
return basicBlockInfo;
}
template <typename COMP>
__aicore__ inline uint32_t CompressorKernelPerf<COMP>::GetLoopTimes()
{
// 计算主循环次数
uint32_t loopTimes = 0;
BatchInfo batchInfo{};
SkipInvalidBatch(batchInfo);
for (;batchInfo.bIdx < constInfo.batchSize; ++loopTimes) {
SkipOneLoop(batchInfo);
}
return loopTimes;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::CalcSplitCoreInfo()
{
// D方向的基本块数量
constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize;
// 核的组数
constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum;
// 每个核处理的d方向的索引
constInfo.dIdx = (constInfo.aiCoreIdx % constInfo.dBasicBlockNum) * constInfo.dBaseSize;
// 当前组id
constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum;
constInfo.mm1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.coreGroupNum;
uint32_t coff = (uint32_t)COMP::coff;
constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff;
constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff;
constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize;
constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::InitWorkspace(__gm__ uint8_t *workspace) {
uint64_t offset = 0;
uint64_t mm1KvResStartOffset = offset;
// mm1KvResGm
mm1KvResGm.SetGlobalBuffer(
(__gm__ MM1_OUT_T *)(workspace + offset +
constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T)));
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T);
uint64_t mm1ScoreResStartOffset = offset;
// mm1ScoreResGm
mm1ScoreResGm.SetGlobalBuffer(
(__gm__ MM1_OUT_T *)(workspace + offset +
constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T)));
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T);
Vec1InputKvGm.SetGlobalBuffer(
(__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset));
Vec1InputScoreGm.SetGlobalBuffer(
(__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset));
vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
uint64_t beforeVecOffset = offset;
// vec1Res
vec1ResGm.SetGlobalBuffer(
(__gm__ VEC1_OUT_T *)(workspace + offset));
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T);
// vec2Input
vec2InputGm.SetGlobalBuffer(
(__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset));
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeMm1(const RunInfo &info, bool isNeedExcute) {
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + info.cubeDbIdx);
if (isNeedExcute) {
blockCube_.ComputeMm1(info);
}
CrossCoreSetFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
CrossCoreWaitFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_C1_V1_FLAG + info.cubeDbIdx);
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeVec1(const Vec1RunInfo &info) {
CrossCoreWaitFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_C1_V1_FLAG + info.c1v1DbIdx);
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + info.c1v1DbIdx);
blockVec_.ComputeVec1(info);
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + info.c1v1DbIdx);
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio);
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeVec2(const Vec2RunInfo &info) {
blockVec_.ComputeVec2(info);
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::AllocEventID()
{
if ASCEND_IS_AIC {
blockCube_.AllocEventID(pipe_);
} else {
blockVec_.AllocEventID();
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + i);
}
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2);
}
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::FreeEventID()
{
if ASCEND_IS_AIC {
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + i);
}
blockCube_.FreeEventID(pipe_);
} else {
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio);
blockVec_.FreeEventID();
}
}
template <typename COMP>
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedExcuteC1(RunInfo info)
{
// B超出范围则cube不执行
return info.bStart < constInfo.batchSize;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx)
{
vec1Info.bStart = batchInfo.bIdx;
vec1Info.sStart = batchInfo.sIdx;
vec1Info.resetResFlag = (loopIdx & (constInfo.nSize - 1)) == 0;
vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1));
vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
BasicBlockInfo basicBlockInfo = SkipOneLoop(batchInfo);
info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1));
info.dealSeqCnt = basicBlockInfo.dealSeqCnt;
info.dealTcNum = basicBlockInfo.dealTcNum;
info.bStart = basicBlockInfo.bIdx;
info.sStart = basicBlockInfo.sIdx;
vec1Info.dealTcNum = basicBlockInfo.dealTcNum;
vec1Info.dealScSize = basicBlockInfo.compressedTcNum;
allCompressedTcNum_ += basicBlockInfo.compressedTcNum;
}
template <typename COMP>
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedExcuteV2(Vec2RunInfo &vec2Info)
{
return (vec2Info.dealScSize > 0);
}
template <typename COMP>
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedSyncAll(uint32_t curBasicBlockIdx)
{
if (allCompressedTcNum_ == 0) {
return false;
}
uint32_t cnt = curBasicBlockIdx + 1;
if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) {
return true;
}
return false;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::UpdateVec2Info(
Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info)
{
// nSize轮起始先重置v2Info信息
if (curBasicBlockIdx % constInfo.nSize == 0) {
vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
vec2Info.bStart = info.bStart;
vec2Info.sStart = info.sStart;
// 将sStart转成bCompressedId
uint32_t startPos = tools_.GetStartPos(info.bStart);
if (tools_.isExistSeqUsed_) {
uint32_t seqUsed = tools_.GetSeqUsed(info.bStart);
if (vec2Info.sStart >= seqUsed) {
vec2Info.bStart++;
vec2Info.sStart = 0;
}
}
vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio;
vec2Info.dealScSize = 0;
} else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) {
vec2Loop++;
}
vec2Info.dealScSize += info.dealScSize;
vec2Info.compressedId += info.dealScSize;
}
template <typename COMP>
__aicore__ inline void CompressorKernelPerf<COMP>::Process()
{
// 所有batch的有效序列都为0时, 直接退出
if (constInfo.batchSize == 0) {
return;
}
AllocEventID();
BatchInfo batchInfo{};
RunInfo extraInfo[1];
Vec1RunInfo vec1Info{};
Vec2RunInfo vec2Info{};
SkipInvalidBatch(batchInfo);
for (uint32_t i = 0; i < loopTimes; ++i) {
RunInfo &extraInfo0 = extraInfo[0];
CalcC1V1Params(extraInfo0, vec1Info, batchInfo, i);
bool isNeedExcuteC1 = IsNeedExcuteC1(extraInfo0);
if ASCEND_IS_AIC {
ComputeMm1(extraInfo0, isNeedExcuteC1);
} else {
ComputeVec1(vec1Info);
UpdateVec2Info(vec2Info, i, vec1Info);
if (IsNeedSyncAll(i)) {
SyncAll();
if (IsNeedExcuteV2(vec2Info)) {
ComputeVec2(vec2Info);
}
}
}
}
FreeEventID();
}
} // namespace Compressor
#endif // COMPRESSOR_KERNEL_PERF_H

56
csrc/attention/compressor/op_kernel/arch32/compressor_template_tiling_key.h Normal file
View File

@@ -0,0 +1,56 @@
/**
 * 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 COMPRESSOR_template_tiling_key.h
* \brief
*/
#ifndef COMPRESSOR_TEMPLATE_TILING_KEY_H
#define COMPRESSOR_TEMPLATE_TILING_KEY_H
#include "ascendc/host_api/tiling/template_argument.h"
#define ASCENDC_TPL_1_BW 1 // 每个参数占用1个bit位
#define ASCENDC_TPL_2_BW 2 // 每个参数占用2个bit位
#define ASCENDC_TPL_4_BW 4 // 每个参数占用4个bit位
// 可表示的tilingkey范围为64bit，注意不可超过限制
ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识，与opType保持一致
// 可能需要切分之后的headdim
// bit:0 LAYOUT 0:BSH 1:TH
ASCENDC_TPL_UINT_DECL(X_LAYOUT, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1),
// bit:1-4 x的dtype 0:BF16 1:FP16
ASCENDC_TPL_UINT_DECL(X_DTYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, 0, 1),
// bit:5-6 coff 1:无需overlap 2:需要overlap
ASCENDC_TPL_UINT_DECL(COFF, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
// bit:7-8 rotary_mode 1:half 2:interleave
ASCENDC_TPL_UINT_DECL(ROTARY_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
// bit:9-10 cache_mode 1:CONTINUOUS 2:cycle
ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
// bit:11-12 template_id 0:empty_tensor 1:normal 2:full load
ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),
// bit:13 rope dtype 0:same as x 1:fp32
ASCENDC_TPL_UINT_DECL(ROPE_DTYPE, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1),
);
ASCENDC_TPL_SEL(
ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(X_LAYOUT, ASCENDC_TPL_UI_LIST, 0, 1),
ASCENDC_TPL_UINT_SEL(X_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1),
ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),
ASCENDC_TPL_UINT_SEL(ROTARY_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),
ASCENDC_TPL_UINT_SEL(ROPE_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1),
ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)),
);
#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H

69
csrc/attention/compressor/op_kernel/arch32/compressor_tiling_data.h Normal file
View File

@@ -0,0 +1,69 @@
/**
 * 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 COMPRESSOR_tiling_datay.h
* \brief
*/
#ifndef COMPRESSOR_TILING_DATA_H
#define COMPRESSOR_TILING_DATA_H
#include <cstdint>
#include "kernel_tiling/kernel_tiling.h"
const uint32_t CMP_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐
namespace optiling {
// 1. 基础参数结构体
struct CompressorBaseParams {
uint32_t batchSize = 0; // bastch size（批大小）
uint32_t seqSize = 0; // sequence size（kvs大小）
uint32_t hiddenSize = 0; // hidden size（隐藏层大小）
uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数：批大小x序列1长度)
uint32_t headDim = 0; // head size of kv
uint32_t ropeHeadDim = 64; // dim size per rope head 64（单个带RoPE头的维度）
uint32_t csSize = 0; // Compress sequence len
uint32_t cmpRatio = 4; // Compress ratio
uint32_t cgSize = 0; // Compress group size
float normEps = 1e-6; // RMSNorm eps
float reciprocalD = 0; // 1分之D
uint32_t usedCoreNum = 0; // 使用核数
uint32_t nSize = 0; // 控制v2积攒的轮数
uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride
};
struct CompressorPageAttentionParams {
uint32_t blockNum = 0;
uint32_t blockSize = 1;
uint32_t maxBlockNumPerBatch = 1;
};
struct CompressorInnerSplitParams {
uint32_t mBaseSize;
uint32_t dBaseSize;
};
struct CompressorWorkspaceParams {
uint32_t mm1KvResSize;
uint32_t mm1ScoreResSize;
uint32_t vec1ResSize;
uint32_t vec1TailCacheSize;
uint32_t dbWorkspaceRatio = 1;
};
struct CompressorTilingData {
CompressorBaseParams baseParams;
CompressorPageAttentionParams pageAttentionParams;
CompressorInnerSplitParams innerSplitParams;
CompressorWorkspaceParams workspaceParams;
};
} // optiling
#endif // COMPRESSOR_TILING_DATA_H

761
csrc/attention/compressor/op_kernel/arch32/compressor_tools.h Normal file
View File

@@ -0,0 +1,761 @@
/**
 * 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 compressor_tools.h
* \brief 放算子都需要、与算子联系紧密、但是又不方便单独独立出来的公共工具
*/
#ifndef COMPRESSOR_TOOLS_H
#define COMPRESSOR_TOOLS_H
#include "compressor_comm.h"
using namespace AscendC;
namespace Compressor {
struct ToolsParams {
uint32_t seqSize = 0U;
uint32_t cmpRatio = 0U;
};
template <typename COMP>
class CompressorTools {
public:
__aicore__ inline CompressorTools()
{
}
__aicore__ inline void Init(__gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos);
__aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx);
__aicore__ inline uint32_t GetStartPos(uint32_t bIdx);
__aicore__ inline uint32_t GetSeqLength(uint32_t bIdx);
__aicore__ inline uint32_t GetTIdxByBatch(uint32_t bIdx);
public:
ToolsParams toolParams_{};
bool isExistSeqUsed_ = false;
private:
bool isExistStartPos_ = false;
GlobalTensor<int32_t> cuSeqlensGm_;
GlobalTensor<int32_t> sequsedGm_;
GlobalTensor<int32_t> startPosGm_;
};
template <typename COMP>
__aicore__ inline void CompressorTools<COMP>::Init(__gm__ uint8_t *startPos, __gm__ uint8_t *seqUsed,
__gm__ uint8_t *cuSeqlens)
{
isExistStartPos_ = (startPos != nullptr);
if (isExistStartPos_) {
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
}
isExistSeqUsed_ = (seqUsed != nullptr);
if (isExistSeqUsed_) {
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
}
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
}
}
template <typename COMP>
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqUsed(uint32_t bIdx)
{
if (isExistSeqUsed_) {
return (uint32_t)sequsedGm_.GetValue(bIdx);
} else {
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx));
} else {
return toolParams_.seqSize;
}
}
}
template <typename COMP>
__aicore__ inline uint32_t CompressorTools<COMP>::GetStartPos(uint32_t bIdx)
{
if (isExistStartPos_) {
return (uint32_t)startPosGm_.GetValue(bIdx);
} else {
return 0;
}
}
template <typename COMP>
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqLength(uint32_t bIdx)
{
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx);
} else {
return toolParams_.seqSize;
}
}
template <typename COMP>
__aicore__ inline uint32_t CompressorTools<COMP>::GetTIdxByBatch(uint32_t bIdx)
{
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx));
} else {
return toolParams_.seqSize * bIdx;
}
}
// iterator
struct SliceInfo {
__aicore__ inline SliceInfo(){};
__aicore__ inline SliceInfo(uint32_t bIdx, uint32_t sIdx) : bIdx(bIdx), sIdx(sIdx){};
uint32_t bIdx = 0U;
uint32_t sIdx = 0U;
uint32_t bSeqUsed = 0U;
uint32_t bStartPos = 0U;
uint32_t headHolderSeqCnt = 0U;
uint32_t validSeqCnt = 0U;
uint32_t tailHolderSeqCnt = 0U;
uint32_t dealSeqCnt = 0;
uint32_t dealTcSize = 0U;
uint32_t compressTcSize = 0U;
};
template <typename COMP>
class CompressorSliceIterator {
public:
__aicore__ inline CompressorSliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
{
}
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
__aicore__ inline bool IsEnd();
__aicore__ inline void IteratorSlice();
__aicore__ inline SliceInfo &GetSlice();
__aicore__ inline SliceInfo &GetSliceByCmp();
bool isFirst_ = true;
SliceInfo sliceInfo_{};
private:
CompressorTools<COMP> &tools_;
// iterator
uint32_t maxDealSeqCnt_ = 0;
uint32_t batch_size_ = 0;
};
template <typename COMP>
__aicore__ inline void CompressorSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
{
sliceInfo_.bIdx = bIdx;
sliceInfo_.sIdx = sIdx;
isFirst_ = true;
}
template <typename COMP>
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
{
this->batch_size_ = batch_size;
}
template <typename COMP>
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
{
this->maxDealSeqCnt_ = maxDealSeqCnt;
}
template <typename COMP>
__aicore__ inline bool CompressorSliceIterator<COMP>::IsEnd()
{
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
}
template <typename COMP>
__aicore__ inline void CompressorSliceIterator<COMP>::IteratorSlice()
{
bool isUpdateBatchInfo = false;
if (!isFirst_) {
// 更新剩余未处理的行数
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
sliceInfo_.sIdx = 0;
sliceInfo_.bIdx++;
isUpdateBatchInfo = true;
}
} else {
isUpdateBatchInfo = true;
isFirst_ = false;
}
// 更新与bIdx相关的bStartPos和bSeqUsed
if (isUpdateBatchInfo) {
// SkipInvalidBatch
while (sliceInfo_.bIdx < batch_size_) {
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
if (sliceInfo_.bSeqUsed > 0) {
break;
}
sliceInfo_.bIdx++;
}
if (sliceInfo_.bIdx < batch_size_) {
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
}
}
}
template <typename COMP>
__aicore__ inline SliceInfo &CompressorSliceIterator<COMP>::GetSliceByCmp()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if (isFirst_) {
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
isFirst_ = false;
}
// 计算头部占位行数、有效数据行数、尾部占位行数
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
}
sliceInfo_.tailHolderSeqCnt =
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
sliceInfo_.tailHolderSeqCnt = 0;
}
// 头和尾处理，否则需要处理的seq等于cmpRatio
if (sliceInfo_.validSeqCnt < cmpRatio) {
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx == 0) {
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
}
} else {
sliceInfo_.dealSeqCnt = cmpRatio;
}
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
// 计算本次可以处理的Tc个数
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
if (sliceInfo_.tailHolderSeqCnt > 0) {
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩
}
return sliceInfo_;
}
template <typename COMP>
__aicore__ inline SliceInfo &CompressorSliceIterator<COMP>::GetSlice()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if (isFirst_) {
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
isFirst_ = false;
}
// 计算头部占位行数、有效数据行数、尾部占位行数
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
}
sliceInfo_.tailHolderSeqCnt =
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
sliceInfo_.tailHolderSeqCnt = 0;
}
sliceInfo_.dealSeqCnt = sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt;
// 计算本次可以处理的Tc个数
sliceInfo_.dealTcSize = sliceInfo_.dealSeqCnt / cmpRatio;
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
if (sliceInfo_.tailHolderSeqCnt > 0) {
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩
}
return sliceInfo_;
}
struct SplitCoreSliceInfo : public SliceInfo {
__aicore__ inline SplitCoreSliceInfo(){};
__aicore__ inline SplitCoreSliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){};
uint32_t preFirstSeqCnt = 0U; // 左边每次迭代基本块的第一个seqCnt大小
};
template <typename COMP>
class CompressorSplitCoreSliceIterator {
public:
__aicore__ inline CompressorSplitCoreSliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
{
}
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
__aicore__ inline bool IsEnd();
__aicore__ inline void IteratorSlice();
__aicore__ inline SplitCoreSliceInfo &GetSlice();
__aicore__ inline SplitCoreSliceInfo &GetSliceByCmp();
__aicore__ inline uint32_t GetBIdx();
__aicore__ inline SplitCoreSliceInfo &GetLeftNextCmpSeqCnt();
__aicore__ inline SplitCoreSliceInfo &GetRightNextCmpSeqCnt();
bool isFirst_ = true;
bool isLeftFirstBath = false;
bool isMaxDealSeqCntFirst = false;
SplitCoreSliceInfo sliceInfo_{};
private:
CompressorTools<COMP> &tools_;
// iterator
uint32_t maxDealSeqCnt_ = 0;
uint32_t batch_size_ = 0;
};
template <typename COMP>
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
{
sliceInfo_.bIdx = bIdx;
sliceInfo_.sIdx = sIdx;
isFirst_ = true;
}
template <typename COMP>
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
{
this->batch_size_ = batch_size;
isMaxDealSeqCntFirst = true;
}
template <typename COMP>
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
{
this->maxDealSeqCnt_ = maxDealSeqCnt;
}
template <typename COMP>
__aicore__ inline bool CompressorSplitCoreSliceIterator<COMP>::IsEnd()
{
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
}
template <typename COMP>
__aicore__ inline uint32_t CompressorSplitCoreSliceIterator<COMP>::GetBIdx()
{
return sliceInfo_.bIdx;
}
template <typename COMP>
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::IteratorSlice()
{
bool isUpdateBatchInfo = false;
if (isMaxDealSeqCntFirst) {
isMaxDealSeqCntFirst = false;
}
if (!isFirst_) {
// 更新剩余未处理的行数
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
sliceInfo_.sIdx = 0;
// 左边最后一块跳到b=0 s=0处理
if (isLeftFirstBath) {
isLeftFirstBath = false;
} else {
sliceInfo_.bIdx++;
}
isUpdateBatchInfo = true;
}
} else {
isUpdateBatchInfo = true;
isFirst_ = false;
}
// 更新与bIdx相关的bStartPos和bSeqUsed
if (isUpdateBatchInfo) {
// SkipInvalidBatch
while (sliceInfo_.bIdx < batch_size_) {
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
if (sliceInfo_.bSeqUsed > 0) {
break;
}
sliceInfo_.bIdx++;
}
if (sliceInfo_.bIdx < batch_size_) {
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
}
}
}
template <typename COMP>
__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator<COMP>::GetLeftNextCmpSeqCnt()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if (isFirst_) {
// 左边 T轴首次减去T轴最后一块
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(batch_size_ - 1);
sliceInfo_.bStartPos = tools_.GetStartPos(batch_size_ - 1);
// 处理最后一块是中间整块或者尾块的情况
uint32_t lastSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio == 0 ?
cmpRatio :
(sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio;
// 处理最后一块是头块的情况
if (sliceInfo_.bSeqUsed < cmpRatio) {
lastSeqCnt = sliceInfo_.bSeqUsed;
}
sliceInfo_.sIdx = sliceInfo_.bSeqUsed - lastSeqCnt;
isLeftFirstBath = true;
isFirst_ = false;
}
// 计算头部占位行数、有效数据行数、尾部占位行数
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
}
sliceInfo_.tailHolderSeqCnt =
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
sliceInfo_.tailHolderSeqCnt = 0;
}
// 头和尾处理，否则需要处理的seq等于cmpRatio
if (sliceInfo_.validSeqCnt < cmpRatio) {
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx == 0) {
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
}
} else {
sliceInfo_.dealSeqCnt = cmpRatio;
}
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
// 计算本次可以处理的Tc个数
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
if (sliceInfo_.tailHolderSeqCnt > 0) {
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩
}
// 记录左边第一个块
if (isMaxDealSeqCntFirst) {
sliceInfo_.preFirstSeqCnt = sliceInfo_.dealSeqCnt;
}
return sliceInfo_;
}
template <typename COMP>
__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator<COMP>::GetRightNextCmpSeqCnt()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if (isFirst_) {
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
isFirst_ = false;
}
// 计算头部占位行数、有效数据行数、尾部占位行数
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
}
sliceInfo_.tailHolderSeqCnt =
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
sliceInfo_.tailHolderSeqCnt = 0;
}
// 头和尾处理，否则需要处理的seq等于cmpRatio
if (sliceInfo_.validSeqCnt < cmpRatio) {
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx == 0) {
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
}
} else {
sliceInfo_.dealSeqCnt = cmpRatio;
}
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
// 计算本次可以处理的Tc个数
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
if (sliceInfo_.tailHolderSeqCnt > 0) {
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时，其不需要压缩
}
return sliceInfo_;
}
struct Vec1SliceInfo : public SliceInfo {
__aicore__ inline Vec1SliceInfo(){};
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){};
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt)
: SliceInfo(bIdx, sIdx), dealedSeqCnt(dealedSeqCnt){};
uint32_t dealedSeqCnt = 0U;
uint32_t dealedTcCnt = 0U;
uint32_t bSeqLength = 0U;
uint32_t compressoredScCnt = 0U;
bool isFirst = false;
bool isLast = false;
};
struct StatisticInfo {
__aicore__ inline StatisticInfo(){};
__aicore__ inline StatisticInfo(uint32_t actualTcCnt, uint32_t dealSeqCnt, uint32_t compressorScCnt)
: actualTcCnt(actualTcCnt), dealSeqCnt(dealSeqCnt), compressorScCnt(compressorScCnt){};
uint32_t actualTcCnt = 0U;
uint32_t dealSeqCnt = 0U;
uint32_t compressorScCnt = 0U;
};
template <typename COMP>
class CompressorVec1SliceIterator {
public:
__aicore__ inline CompressorVec1SliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
{
}
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, uint32_t compressoredScCnt);
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
__aicore__ inline void SetDealedSeqCnt(uint32_t dealedSeqCnt);
__aicore__ inline void SetDealedTcCnt(uint32_t dealedTcCnt);
__aicore__ inline void SetCompressoredScCnt(uint32_t compressoredScCnt);
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize);
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize, uint32_t canDealTcSize);
__aicore__ inline uint32_t GetNeedDealTcSize();
__aicore__ inline bool IsEnd();
template <bool IS_STATISTIC = false>
__aicore__ inline void IteratorSlice();
__aicore__ inline Vec1SliceInfo &GetSlice();
template <bool IS_STATISTIC = false>
__aicore__ inline StatisticInfo &FullIteratorSlice();
private:
CompressorTools<COMP> &tools_;
bool isFirst_ = true;
Vec1SliceInfo sliceInfo_{};
StatisticInfo statisticInfo_{};
uint32_t needDealTcSize_ = 0U;
uint32_t batch_size_ = 0U;
};
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
{
sliceInfo_.bIdx = bIdx;
sliceInfo_.sIdx = sIdx;
while (tools_.GetSeqLength(sliceInfo_.bIdx) == 0) {
sliceInfo_.bIdx++;
if (sliceInfo_.bIdx == batch_size_) {
sliceInfo_.bIdx = 0;
}
}
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
isFirst_ = true;
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt,
uint32_t compressoredScCnt)
{
Reset(bIdx, sIdx);
SetDealedSeqCnt(dealedSeqCnt);
SetCompressoredScCnt(compressoredScCnt);
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
{
this->batch_size_ = batch_size;
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedSeqCnt(uint32_t dealedSeqCnt)
{
this->sliceInfo_.dealedSeqCnt = dealedSeqCnt;
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetCompressoredScCnt(uint32_t compressoredScCnt)
{
this->sliceInfo_.compressoredScCnt = compressoredScCnt;
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedTcCnt(uint32_t dealedTcCnt)
{
this->sliceInfo_.dealedTcCnt = dealedTcCnt;
}
template <typename COMP>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetNeedDealTcSize(uint32_t needDealTcSize)
{
this->needDealTcSize_ = needDealTcSize;
}
template <typename COMP>
template <bool IS_STATISTIC>
__aicore__ inline void CompressorVec1SliceIterator<COMP>::IteratorSlice()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if constexpr (IS_STATISTIC) {
statisticInfo_.actualTcCnt += sliceInfo_.dealTcSize;
statisticInfo_.compressorScCnt += sliceInfo_.compressTcSize;
}
needDealTcSize_ -= sliceInfo_.dealTcSize;
sliceInfo_.dealedSeqCnt += sliceInfo_.validSeqCnt;
sliceInfo_.compressoredScCnt += sliceInfo_.compressTcSize;
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) {
do {
uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
if (sliceInfo_.bSeqUsed < seqLength) {
uint32_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, cmpRatio) - sliceInfo_.bStartPos;
sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx;
uint32_t tcGap = CeilDivT(static_cast<int32_t>(seqLength - nextAlignSIdx),
static_cast<int32_t>(cmpRatio));
if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) {
// 此时bseqused所在压缩块未被纳入计算
tcGap++;
}
sliceInfo_.sIdx = nextAlignSIdx;
if (needDealTcSize_ < tcGap) {
sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio;
sliceInfo_.sIdx += needDealTcSize_ * cmpRatio;
needDealTcSize_ = 0;
break;
}
sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx;
needDealTcSize_ -= tcGap;
}
sliceInfo_.bIdx++;
if (sliceInfo_.bIdx == batch_size_) {
sliceInfo_.bIdx = 0;
}
sliceInfo_.sIdx = 0;
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
} while (sliceInfo_.bSeqUsed == 0);
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
}
if (isFirst_) {
isFirst_ = false;
}
}
template <typename COMP>
__aicore__ inline uint32_t CompressorVec1SliceIterator<COMP>::GetNeedDealTcSize()
{
return needDealTcSize_;
}
template <typename COMP>
__aicore__ inline bool CompressorVec1SliceIterator<COMP>::IsEnd()
{
return (needDealTcSize_ == 0);
}
template <typename COMP>
__aicore__ inline Vec1SliceInfo &CompressorVec1SliceIterator<COMP>::GetSlice()
{
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) {
sliceInfo_.headHolderSeqCnt = 0;
sliceInfo_.validSeqCnt = 0;
sliceInfo_.tailHolderSeqCnt = 0;
sliceInfo_.dealTcSize = 0;
sliceInfo_.compressTcSize = 0;
} else {
// 计算头部占位行数、有效数据行数、尾部占位行数
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
if (CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt, cmpRatio) > needDealTcSize_) {
sliceInfo_.validSeqCnt = needDealTcSize_ * cmpRatio - sliceInfo_.headHolderSeqCnt;
}
uint32_t globalTotalSeqCnt = sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt;
sliceInfo_.tailHolderSeqCnt = Align(globalTotalSeqCnt, cmpRatio) - globalTotalSeqCnt;
// 计算本次可以处理的Tc个数
sliceInfo_.dealTcSize =
(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt) / cmpRatio;
sliceInfo_.compressTcSize =
(sliceInfo_.headHolderSeqCnt + min(sliceInfo_.validSeqCnt, sliceInfo_.bSeqUsed - sliceInfo_.sIdx)) /
cmpRatio;
}
sliceInfo_.isFirst = isFirst_;
sliceInfo_.isLast =
sliceInfo_.bSeqUsed > sliceInfo_.sIdx &&
CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.bSeqUsed - sliceInfo_.sIdx, cmpRatio) >= needDealTcSize_;
return sliceInfo_;
}
template <typename COMP>
template <bool IS_STATISTIC>
__aicore__ inline StatisticInfo &CompressorVec1SliceIterator<COMP>::FullIteratorSlice()
{
if constexpr (IS_STATISTIC) {
statisticInfo_ = {0U, 0U, 0U};
Vec1SliceInfo tempSliceInfo = GetSlice();
while (!IsEnd()) {
GetSlice();
IteratorSlice<IS_STATISTIC>();
}
Vec1SliceInfo sliceInfo = GetSlice();
statisticInfo_.dealSeqCnt = sliceInfo.dealedSeqCnt - tempSliceInfo.dealedSeqCnt;
} else {
while (!IsEnd()) {
GetSlice();
IteratorSlice<IS_STATISTIC>();
}
}
return statisticInfo_;
}
} // namespace Compressor
#endif

347
csrc/attention/compressor/op_kernel/arch32/compressor_vector_comm.h Normal file
View File

@@ -0,0 +1,347 @@
/**
 * 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 compressor_vector_comm.h
* \brief 存放各种vector的公共组件
*/
#ifndef COMPRESSOR_VECTOR_COMM_H
#define COMPRESSOR_VECTOR_COMM_H
#include "compressor_comm.h"
namespace Compressor {
struct MatRpeatParam {
uint32_t row;
uint32_t col;
uint32_t dtypeMask;
uint32_t loopTimes;
uint32_t colRemain;
uint8_t repeatStride;
};
struct RmsNormParam {
float reciprocal;
float epsilon;
uint32_t row;
uint32_t col;
};
/**
* @brief ColumnSum 对矩阵按列进行求和
* @param dstLocal 输出tensor [1, col]，支持和shareTmpUb是同一块空间
* @param srcLocal 输入tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)]
* @param row 行数
* @param col 列数
*/
__aicore__ inline void ColumnSum(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
{
// 行数为1时，直接将srcLocal复制到dstLocal
if (unlikely(row == 1)) {
DataCopy(dstLocal, srcLocal, row * col);
PipeBarrier<PIPE_V>();
return;
}
for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) {
if (row & mask) {
// 将输入对半求和后放进临时空间
Add(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算
PipeBarrier<PIPE_V>();
// 将余量加到前一半上
if (unlikely(row > mask)) {
if ((row - mask) > (mask >> 1)) {
Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算
PipeBarrier<PIPE_V>();
Add(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col);
PipeBarrier<PIPE_V>();
} else {
Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col);
PipeBarrier<PIPE_V>();
}
}
// 每次将后一半行加到前一半上
for (uint32_t i = mask >> 2; i > 1; i >>= 1) {
Add(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col);
PipeBarrier<PIPE_V>();
}
if (mask == 2) { // 2:最后一次矩阵运算处理
DataCopy(dstLocal, shareTmpUb, col);
} else {
Add(dstLocal, shareTmpUb, shareTmpUb[col], col);
}
PipeBarrier<PIPE_V>();
break;
}
}
}
/**
* @brief ColumnMax 对矩阵按列进行求最大值
* @param dstLocal 输出tensor [1, col]，支持和shareTmpUb是同一块空间
* @param srcLocal 输入tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)]
* @param row 行数
* @param col 列数
*/
__aicore__ inline void ColumnMax(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
{
// 行数为1时，直接将srcLocal复制到dstLocal
if (unlikely(row == 1)) {
DataCopy(dstLocal, srcLocal, row * col);
PipeBarrier<PIPE_V>();
return;
}
for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) {
if (row & mask) {
// 将输入对半求最大值后放进临时空间
Max(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算
PipeBarrier<PIPE_V>();
// 将余量和前一半求最大值后加到前一半上
if (unlikely(row > mask)) {
if ((row - mask) > (mask >> 1)) {
Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算
PipeBarrier<PIPE_V>();
Max(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col);
PipeBarrier<PIPE_V>();
} else {
Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col);
PipeBarrier<PIPE_V>();
}
}
// 每次将后一半行和前一半最大值后加到前一半上
for (uint32_t i = mask >> 2; i > 1; i >>= 1) {
Max(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col);
PipeBarrier<PIPE_V>();
}
if (mask == 2) { // 2:最后一次矩阵运算处理
DataCopy(dstLocal, shareTmpUb, col);
} else {
Max(dstLocal, shareTmpUb, shareTmpUb[col], col);
}
PipeBarrier<PIPE_V>();
break;
}
}
}
/**
* @brief MatSubVec 矩阵逐行减向量
* @param dstLocal 输出tensor [row, col]
* @param src0Local 输入tensor [row, col]
* @param src1Local 输入tensor [1, col]
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void MatSubVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
{
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
uint32_t offset = 0;
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.dtypeMask, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
offset += repeatParam.dtypeMask;
}
if (repeatParam.colRemain > 0) {
Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.colRemain, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
}
}
}
/**
* @brief MatDivVec 矩阵逐行除以向量
* @param dstLocal 输出tensor [row, col]
* @param src0Local 输入tensor [row, col]
* @param src1Local 输入tensor [1, col]
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void MatDivVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
{
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
uint32_t offset = 0;
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.dtypeMask, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
offset += repeatParam.dtypeMask;
}
if (repeatParam.colRemain > 0) {
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.colRemain, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
}
}
}
/**
* @brief MatMulVec 矩阵逐行乘以向量
* @param dstLocal 输出tensor [row, col]
* @param src0Local 输入tensor [row, col]
* @param src1Local 输入tensor [1, col]
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void MatMulVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
{
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
uint32_t offset = 0;
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.dtypeMask, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
offset += repeatParam.dtypeMask;
}
if (repeatParam.colRemain > 0) {
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
repeatParam.colRemain, repeatRowTimes,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
}
}
}
/**
* @brief RowSum 矩阵对每行求和
* @param dstLocal 输出tensor [1, row]
* @param srcLocal 输入tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [row, col]，支持和srcLocal是同一块空间
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void RowSum(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &shareTmpUb, const MatRpeatParam &repeatParam)
{
uint32_t blockCount = repeatParam.loopTimes;
if (blockCount > 0 && repeatParam.colRemain > 0) {
Add(shareTmpUb, srcLocal, srcLocal[blockCount * repeatParam.dtypeMask], repeatParam.colRemain,
repeatParam.row,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride});
AscendC::PipeBarrier<PIPE_V>();
}
for (uint32_t loopCount = blockCount >> 1; loopCount > 0; loopCount = blockCount >> 1) {
blockCount = (blockCount + 1) >> 1;
for (uint32_t i = 0; i < loopCount; i++) {
Add(shareTmpUb[i * repeatParam.dtypeMask], srcLocal[i * repeatParam.dtypeMask],
srcLocal[(i + blockCount) * repeatParam.dtypeMask], repeatParam.dtypeMask, repeatParam.row,
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride});
}
AscendC::PipeBarrier<PIPE_V>();
}
WholeReduceSum(dstLocal, shareTmpUb,
(repeatParam.col < repeatParam.dtypeMask) ? repeatParam.col :
repeatParam.dtypeMask,
repeatParam.row, 1, 1, repeatParam.repeatStride);
}
/**
* @brief RowDivs 矩阵每行除以对应元素
* @param dstLocal 输出tensor [row, col]
* @param src0Local 输入tensor [row, col]
* @param src1Local 输入tensor [row, 1]，需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM])
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void RowDivs(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
{
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
uint32_t offset = 0;
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
repeatParam.dtypeMask, repeatRowTimes,
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
offset += repeatParam.dtypeMask;
}
if (repeatParam.colRemain > 0) {
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
repeatParam.colRemain, repeatRowTimes,
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
}
}
}
/**
* @brief RowMuls 矩阵每行乘以相同元素
* @param dstLocal 输出tensor [row, col]
* @param src0Local 输入tensor [row, col]
* @param src1Local 输入tensor [row, 1]，需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM])
* @param repeatParam 描述待处理数据的排布，包括
row 行数
col 列数
dtypeMask 一次迭代参与计算元素数
loopTimes 循环次数
colRemain 剩余列数
repeatStride 循环步长（内存中实际列长度）
*/
__aicore__ inline void RowMuls(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
{
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
uint32_t offset = 0;
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
repeatParam.dtypeMask, repeatRowTimes,
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
offset += repeatParam.dtypeMask;
}
if (repeatParam.colRemain > 0) {
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
repeatParam.colRemain, repeatRowTimes,
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
}
}
}
} // namespace Compressor
#endif // COMPRESSOR_VECTOR_COMM_H

87
csrc/attention/compressor/op_kernel/arch32/rms_norm.h Normal file
View File

@@ -0,0 +1,87 @@
/**
 * 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 rms_norm.h
* \brief
*/
#ifndef RMS_NORM_H
#define RMS_NORM_H
#include "compressor_comm.h"
#include "compressor_vector_comm.h"
namespace Compressor {
/**
* @brief RmsNorm 对矩阵进行rmsnorm
* @param dstLocal 输出tensor [row, col]，支持和srcLocal是同一块空间
* @param srcLocal 输入tensor [row, col]
* @param gammaLocal 系数gamma [1, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [(row * col + row) * sizeof(float)]
* @param rmsNormParams rms所需系数，包括
reciprocal rmsnorm系数reciprocal
epsilon rmsnorm系数epsilon
row 处理的行数
col 列数
*/
template <typename GammaType>
__aicore__ inline void RmsNorm(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<GammaType> &gammaLocal, const LocalTensor<float> &shareTmpUb,
const RmsNormParam &rmsNormParams)
{
uint64_t cnt = rmsNormParams.row * rmsNormParams.col;
LocalTensor<float> temp1Local = shareTmpUb.ReinterpretCast<float>();
LocalTensor<float> temp2Local = temp1Local[cnt];
// temp1Local = srcLocal ^ 2
Mul(temp1Local, srcLocal, srcLocal, cnt);
PipeBarrier<PIPE_V>();
MatRpeatParam repeatParams = {
rmsNormParams.row, // row
rmsNormParams.col, // col
FP32_REPEAT_ELEMENT_NUM, // dtypeMask
rmsNormParams.col / FP32_REPEAT_ELEMENT_NUM, // loopTimes
rmsNormParams.col % FP32_REPEAT_ELEMENT_NUM, // colsRemain
static_cast<uint8_t>(rmsNormParams.col / FP32_BLOCK_ELEMENT_NUM), // repeatStride
};
// temp2Local[row] = Sum(temp1Local)
RowSum(temp2Local, temp1Local, temp1Local, repeatParams);
PipeBarrier<PIPE_V>();
// temp2Local[row] = temp2Local[row] * reciprocal(1/N)
Muls(temp2Local, temp2Local, rmsNormParams.reciprocal, rmsNormParams.row);
PipeBarrier<PIPE_V>();
// temp2Local[row] = temp2Local[row] + epsilon
Adds(temp2Local, temp2Local, rmsNormParams.epsilon, rmsNormParams.row);
PipeBarrier<PIPE_V>();
// temp2Local[row] = Sqrt(temp2Local[row])
Sqrt(temp2Local, temp2Local, rmsNormParams.row);
PipeBarrier<PIPE_V>();
// temp1Local[row, 8] = brc(temp2Local[row, 1])
Brcb(temp1Local, temp2Local, CeilDivT(rmsNormParams.row, BRCB_NUM), {1, 8});
PipeBarrier<PIPE_V>();
// dstLocal = srcLocal / temp1Local(sum)
RowDivs(dstLocal, srcLocal, temp1Local, repeatParams);
PipeBarrier<PIPE_V>();
// dstLocal = dstLocal * gammaLocal
MatMulVec(dstLocal, dstLocal, gammaLocal, repeatParams);
}
} // namespace Compressor
#endif // MLA_PROLOG_RMS_NORM_H

130
csrc/attention/compressor/op_kernel/arch32/rope.h Normal file
View File

@@ -0,0 +1,130 @@
/**
 * 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 rope.h
* \brief
*/
#ifndef ROPE_H
#define ROPE_H
#include "compressor_comm.h"
#include "compressor_vector_comm.h"
namespace Compressor {
/**
* @brief SetGatherSrcOffset 计算用于interleave模式的offset
* @param gatherOffsetLocal 输出tensor [count]，数据类型需要为int64_t，使用时要转换
* @param count offset的元素个数，一般为列数
*/
template <typename T>
__aicore__ inline void SetGatherSrcOffset(const LocalTensor<int32_t> &gatherOffsetLocal, uint32_t count)
{
for (uint32_t i = 0; i < 8; i++) {
gatherOffsetLocal.SetValue(i, i ^ 1);
}
event_t eventId_S_V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
SetFlag<HardEvent::S_V>(eventId_S_V);
WaitFlag<HardEvent::S_V>(eventId_S_V);
int32_t scalarValue = 8;
while (scalarValue < count) {
int32_t nextValue = scalarValue * 2;
PipeBarrier<PIPE_V>();
if (nextValue < count) {
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, scalarValue);
} else {
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, count - scalarValue);
break;
}
scalarValue = nextValue;
}
PipeBarrier<PIPE_V>();
Muls(gatherOffsetLocal, gatherOffsetLocal, static_cast<int32_t>(sizeof(T)), count);
}
/**
* @brief RotaryPosEmb 同时做row行的RotaryPosEmb，每一行的元素为col
* @param dstLocal 输出tensor [row, actualCol]，支持和srcLocal是同一块空间
* @param srcLocal 输入tensor [row, actualCol]
* @param cosLocal cos系数tensor [row, col]
* @param sinLocal sin系数tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [row * col * sizeof(float)]
* @param gatherOffsetcastLocal 用于interleave模式的offset，数据类型需要为uint64_t
* @param row 待处理的行数
* @param col 待处理的列数
* @param actualCol 实际列数
* @param baseAddr 计算基地址
*/
template <ROTARY_MODE MODE>
__aicore__ inline void RotaryPosEmb(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &cosLocal, const LocalTensor<float> &sinLocal,
const LocalTensor<float> &shareTmpUb,
const LocalTensor<uint32_t> &gatherOffsetcastLocal, uint32_t row, uint32_t col,
uint32_t actualCol, uint64_t baseAddr)
{
uint64_t cnt = row * col;
uint32_t half_col = col >> 1;
uint64_t rsvdCnt = 0;
LocalTensor<float> reArrLocal = shareTmpUb.ReinterpretCast<float>();
if constexpr (MODE == ROTARY_MODE::HALF) {
DataCopy(reArrLocal, srcLocal[baseAddr + half_col],
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
DataCopy(reArrLocal[half_col], srcLocal[baseAddr],
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
PipeBarrier<PIPE_V>();
Muls(reArrLocal, reArrLocal, float(-1), half_col, row,
{1, 1, static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM))});
} else if constexpr (MODE == ROTARY_MODE::INTERLEAVE) {
for (uint32_t i = 0; i < row; i++) {
Gather(reArrLocal[i * col], srcLocal[i * actualCol + baseAddr], gatherOffsetcastLocal, 0, col);
}
PipeBarrier<PIPE_V>();
uint32_t repeatTimes = cnt / FP32_REPEAT_ELEMENT_NUM;
uint32_t remainder = cnt % FP32_REPEAT_ELEMENT_NUM;
uint64_t fullMask = 0x5555555555555555;
uint64_t partialMask = 0x55;
SetVectorMask<float, MaskMode::NORMAL>(0, fullMask);
Muls<float, false>(reArrLocal, reArrLocal, float(-1), MASK_PLACEHOLDER, repeatTimes,
{1, 1, FP32_BLOCK_ELEMENT_NUM, FP32_BLOCK_ELEMENT_NUM});
if (unlikely(remainder > 0)) {
SetVectorMask<float, MaskMode::NORMAL>(0, partialMask);
Muls<float, false>(reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM],
reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM], float(-1), MASK_PLACEHOLDER,
remainder / FP32_BLOCK_ELEMENT_NUM, {1, 1, 1, 1});
}
ResetMask();
}
PipeBarrier<PIPE_V>();
BinaryRepeatParams computeParams{1,
1,
1,
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
static_cast<uint8_t>(CeilDivT(col, FP32_BLOCK_ELEMENT_NUM))};
Mul(dstLocal[baseAddr], srcLocal[baseAddr], cosLocal, col, row, computeParams);
Mul(reArrLocal, reArrLocal, sinLocal, cnt);
PipeBarrier<PIPE_V>();
Add(dstLocal[baseAddr], dstLocal[baseAddr], reArrLocal, col, row, computeParams);
}
}
#endif

52
csrc/attention/compressor/op_kernel/arch32/soft_max.h Normal file
View File

@@ -0,0 +1,52 @@
/**
 * 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 soft_max.h
* \brief
*/
#ifndef SOFT_MAX_H
#define SOFT_MAX_H
#include "compressor_comm.h"
#include "compressor_vector_comm.h"
namespace Compressor {
/**
* @brief ColumnSoftMax 对矩阵按列进行SoftMax
* @param dstLocal 输出tensor [row, col]，支持和srcLocal是同一块空间
* @param srcLocal 输入tensor [row, col]
* @param shareTmpUb 临时buffer 内部需要的空间为 [floor(row / 2) * col * sizeof(float)]
* @param row 行数
* @param col 列数
*/
__aicore__ inline void ColumnSoftMax(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
{
uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
uint32_t dLoop = col / dtypeMask;
uint32_t dRemain = col % dtypeMask;
uint8_t repeatStride = col / FP32_BLOCK_ELEMENT_NUM;
ColumnMax(shareTmpUb, srcLocal, shareTmpUb, row, col);
PipeBarrier<PIPE_V>();
MatSubVec(dstLocal, srcLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride});
PipeBarrier<PIPE_V>();
Exp(dstLocal, dstLocal, row * col);
PipeBarrier<PIPE_V>();
ColumnSum(shareTmpUb, dstLocal, shareTmpUb, row, col);
PipeBarrier<PIPE_V>();
MatDivVec(dstLocal, dstLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride});
}
} // namespace Compressor
#endif