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

148
csrc/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_common.h Normal file
View File

@@ -0,0 +1,148 @@
/**
 * 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 quant_lightning_indexer_common.h
* \brief
*/
#ifndef QUANT_LIGHTNING_INDEXER_COMMON_H
#define QUANT_LIGHTNING_INDEXER_COMMON_H
namespace QLICommon {
// 与tiling的layout保持一致
enum class LI_LAYOUT : uint32_t {
BSND = 0,
TND = 1,
PA_BSND = 2
};
template <typename Q_T, typename K_T, typename OUT_T, const bool PAGE_ATTENTION = false,
LI_LAYOUT Q_LAYOUT_T = LI_LAYOUT::BSND, LI_LAYOUT K_LAYOUT_T = LI_LAYOUT::PA_BSND, typename... Args>
struct QLIType {
using queryType = Q_T;
using keyType = K_T;
using outputType = OUT_T;
static constexpr bool pageAttention = PAGE_ATTENTION;
static constexpr LI_LAYOUT layout = Q_LAYOUT_T;
static constexpr LI_LAYOUT keyLayout = K_LAYOUT_T;
};
struct RunInfo {
uint32_t loop;
uint32_t bN2Idx;
uint32_t bIdx;
uint32_t n2Idx = 0;
uint32_t gS1Idx;
uint32_t s2Idx;
uint32_t actS1Size = 1;
uint32_t actS2Size = 1;
uint32_t actS2SizeOrig = 1;
uint32_t actMBaseSize;
uint32_t actualSingleProcessSInnerSize;
uint32_t actualSingleProcessSInnerSizeAlign;
uint64_t tensorQueryOffset;
uint64_t tensorKeyOffset;
uint64_t tensorKeyScaleOffset;
uint64_t tensorWeightsOffset;
uint64_t indiceOutOffset;
bool isFirstS2InnerLoop;
bool isLastS2InnerLoop;
bool isValid = false;
};
struct ConstInfo {
// CUBE与VEC核间同步的模式
static constexpr uint32_t FIA_SYNC_MODE2 = 2;
// BUFFER的字节数
static constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
static constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
static constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
static constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
static constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
static constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
static constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
static constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
static constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
static constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
// 无效索引
static constexpr int INVALID_IDX = -1;
// CUBE和VEC的核间同步EventID
uint32_t syncC1V1 = 0U;
uint32_t syncC1V0 = 2U;
uint32_t syncV1C1 = 0U;
uint32_t syncV0C1 = 1U;
// 基本块大小
uint32_t mBaseSize = 1ULL;
uint32_t s1BaseSize = 1ULL;
uint32_t s2BaseSize = 1ULL;
uint64_t batchSize = 0ULL;
uint64_t gSize = 0ULL;
uint64_t qHeadNum = 0ULL;
uint64_t kHeadNum;
uint64_t headDim;
uint64_t sparseCount; // topK选取大小
uint64_t kSeqSize = 0ULL; // kv最大S长度
uint64_t qSeqSize = 1ULL; // q最大S长度
uint32_t kCacheBlockSize = 0; // PA场景的block size
uint32_t maxBlockNumPerBatch = 0; // PA场景的最大单batch block number
LI_LAYOUT outputLayout; // 输出的格式
bool attenMaskFlag = false;
uint32_t cmpRatio = 1; // 压缩率
bool batchSupperFlag = false; // Qactual_se长度是否为B+1
int64_t stride = 1;
int64_t scaleStride = 1;
uint32_t actualLenQDims = 0U; // query的actualSeqLength 的维度
uint32_t actualLenDims = 0U; // KV 的actualSeqLength 的维度
bool isAccumSeqS1 = false; // 是否累加模式
bool isAccumSeqS2 = false; // 是否累加模式
uint32_t s2Start = 0U;
uint32_t s2End = 0U;
uint32_t bN2Start = 0U;
uint32_t bN2End = 0U;
uint32_t gS1Start = 0U;
uint32_t gS1End = 0U;
uint32_t coreEnable = 0U;
};
template <typename T1, typename T2>
__aicore__ inline T1 Align(T1 num, T2 rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd)));
}
template <typename T1, typename T2>
__aicore__ inline T1 Min(T1 a, T2 b)
{
return (a > b) ? (b) : (a);
}
template <typename T1, typename T2>
__aicore__ inline T1 Max(T1 a, T2 b)
{
return (a > b) ? (a) : (b);
}
template <typename T>
__aicore__ inline T CeilDiv(T num, T rnd)
{
return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd)));
}
} // namespace QLICommon
#endif // QUANT_LIGHTNING_INDEXER_COMMON_H

667
csrc/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_kernel.h Normal file
View File

@@ -0,0 +1,667 @@
/**
 * 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 quant_lightning_indexer_kernel.h
* \brief
*/
#ifndef QUANT_LIGHTNING_INDEXER_KERNEL_H
#define QUANT_LIGHTNING_INDEXER_KERNEL_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"
#include "quant_lightning_indexer_common.h"
#include "quant_lightning_indexer_service_vector.h"
#include "quant_lightning_indexer_service_cube.h"
#include "../vllm_quant_lightning_indexer_metadata.h"
namespace QLIKernel {
using namespace QLICommon;
using namespace QLIServiceVec;
using namespace matmul;
using namespace optiling::detail;
using namespace optiling;
using AscendC::CacheMode;
using AscendC::CrossCoreSetFlag;
using AscendC::CrossCoreWaitFlag;
// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算
struct TempLoopInfo {
uint32_t bN2Idx = 0;
uint32_t bIdx = 0U;
uint32_t n2Idx = 0U;
uint32_t gS1Idx = 0U;
uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx
uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx
uint32_t actS1Size = 1ULL; // 当前Batch循环处理的S1轴的实际大小
uint32_t actS2Size = 0ULL;
uint32_t actS2SizeOrig = 0ULL;
bool curActSeqLenIsZero = false;
bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出
uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小
uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小
uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小
uint32_t validS2Len = 0U;
};
template <typename QLIT>
class QLIPreload {
public:
__aicore__ inline QLIPreload(){};
__aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights,
__gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale, __gm__ uint8_t *actualSeqLengthsQ,
__gm__ uint8_t *actualSeqLengthsK, __gm__ uint8_t *blockTable, __gm__ uint8_t *metadata,
__gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace,
const QLITilingData *__restrict tiling, TPipe *tPipe);
__aicore__ inline void Process();
// =================================类型定义区=================================
using Q_T = typename QLIT::queryType;
using K_T = typename QLIT::keyType;
using OUT_T = typename QLIT::outputType;
static constexpr bool PAGE_ATTENTION = QLIT::pageAttention;
static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout;
static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout;
using MM1_OUT_T = float;
QLIMatmul<QLIT> matmulService;
QLIVector<QLIT> vectorService;
// =================================常量区=================================
static constexpr uint32_t SYNC_C1_V1_FLAG = 4;
static constexpr uint32_t SYNC_V1_C1_FLAG = 5;
static constexpr uint32_t M_BASE_SIZE = 256;
static constexpr uint32_t S2_BASE_SIZE = 2048;
static constexpr uint32_t HEAD_DIM = 128;
static constexpr uint32_t K_HEAD_NUM = 1;
static constexpr uint32_t GM_ALIGN_BYTES = 512;
static constexpr uint32_t LI_QUANT_PRELOAD_TASK_CACHE_SIZE = 2;
// for workspace double
static constexpr uint32_t WS_DOBULE = 2;
static constexpr uint32_t ELE_NUM_PER_BLOCK = 16;
protected:
TPipe *pipe = nullptr;
// offset
uint64_t queryCoreOffset = 0ULL;
uint64_t keyCoreOffset = 0ULL;
uint64_t keyScaleCoreOffset = 0ULL;
uint64_t weightsCoreOffset = 0ULL;
uint64_t indiceOutCoreOffset = 0ULL;
uint32_t coreZeroEnable = 1U;
// ================================Global Buffer区=================================
GlobalTensor<Q_T> queryGm;
GlobalTensor<K_T> keyGm;
GlobalTensor<half> weightsGm;
GlobalTensor<uint32_t> metadataGm;
GlobalTensor<int32_t> indiceOutGm;
GlobalTensor<int32_t> blockTableGm;
GlobalTensor<uint32_t> actualSeqLengthsGmQ;
GlobalTensor<uint32_t> actualSeqLengthsGm;
// ================================类成员变量====================================
// aic、aiv核信息
uint32_t tmpBlockIdx = 0U;
uint32_t aiCoreIdx = 0U;
QLICommon::ConstInfo constInfo{};
TempLoopInfo tempLoopInfo{};
// ================================Init functions==================================
__aicore__ inline void InitTilingData(const QLITilingData *__restrict tilingData);
__aicore__ inline void InitBuffers();
__aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK);
// ================================Split Core================================
__aicore__ inline void SplitCore();
__aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig,
uint32_t &validS2Len);
__aicore__ inline uint32_t GetTotalBaseBlockNum();
// ================================Process functions================================
__aicore__ inline void ProcessMain();
__aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx,
QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE]);
__aicore__ inline void ProcessInvalid();
// ================================Params Calc=====================================
__aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx);
__aicore__ inline void GetBN2Idx(uint32_t bN2Idx);
__aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq,
GlobalTensor<uint32_t> &actualSeqLengthsGm, uint32_t defaultSeqLen);
__aicore__ inline uint32_t GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq,
GlobalTensor<uint32_t> &actualSeqLengthsGm, uint32_t defaultSeqLen, uint32_t cmpRatio);
__aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig);
__aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx);
__aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo);
__aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start);
};
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::InitTilingData(const QLITilingData *__restrict tilingData)
{
constInfo.batchSize = tilingData->bSize;
constInfo.qHeadNum = constInfo.gSize = tilingData->gSize;
constInfo.kSeqSize = tilingData->s2Size;
constInfo.qSeqSize = tilingData->s1Size;
constInfo.attenMaskFlag = (tilingData->sparseMode == 3);
constInfo.kCacheBlockSize = tilingData->blockSize;
constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch;
constInfo.sparseCount = tilingData->sparseCount;
constInfo.cmpRatio = tilingData->cmpRatio;
constInfo.batchSupperFlag = tilingData->batchSupperFlag;
constInfo.stride = tilingData->stride;
constInfo.scaleStride = tilingData->scaleStride;
constInfo.outputLayout = Q_LAYOUT_T; // 输出和输入形状一致
if (Q_LAYOUT_T == LI_LAYOUT::TND) {
constInfo.isAccumSeqS1 = true;
}
if (K_LAYOUT_T == LI_LAYOUT::TND) {
constInfo.isAccumSeqS2 = true;
}
constInfo.kHeadNum = K_HEAD_NUM;
constInfo.headDim = HEAD_DIM;
constInfo.mBaseSize = M_BASE_SIZE;
constInfo.s2BaseSize = S2_BASE_SIZE;
constInfo.s1BaseSize = (constInfo.mBaseSize + constInfo.gSize - 1) / constInfo.gSize;
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::InitBuffers()
{
if ASCEND_IS_AIV {
vectorService.InitBuffers(pipe);
} else {
matmulService.InitBuffers(pipe);
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ,
__gm__ uint8_t *actualSeqLengthsK)
{
if (actualSeqLengthsQ == nullptr) {
constInfo.actualLenQDims = 0;
} else {
constInfo.actualLenQDims = (constInfo.batchSupperFlag) ? constInfo.batchSize + 1 : constInfo.batchSize;
actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims);
}
if (actualSeqLengthsK == nullptr) {
constInfo.actualLenDims = 0;
} else {
constInfo.actualLenDims = constInfo.batchSize;
actualSeqLengthsGm.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims);
}
}
template <typename QLIT>
__aicore__ inline uint32_t QLIPreload<QLIT>::GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq,
GlobalTensor<uint32_t> &actualSeqLengthsGm,
uint32_t defaultSeqLen)
{
if (actualLenDims == 0) {
return defaultSeqLen;
} else if (constInfo.batchSupperFlag) {
return actualSeqLengthsGm.GetValue(bIdx + 1) - actualSeqLengthsGm.GetValue(bIdx);
} else if (isAccumSeq && bIdx > 0) {
return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1);
} else {
return actualSeqLengthsGm.GetValue(bIdx);
}
}
template <typename QLIT>
__aicore__ inline uint32_t QLIPreload<QLIT>::GetActualSeqLenKey(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq,
GlobalTensor<uint32_t> &actualSeqLengthsGm,
uint32_t defaultSeqLen, uint32_t cmpRatio)
{
if (actualLenDims == 0) {
return defaultSeqLen * cmpRatio;
} else if (isAccumSeq && bIdx > 0) {
return actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1);
} else {
return actualSeqLengthsGm.GetValue(bIdx);
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size, uint32_t &actS2Size, uint32_t &actS2SizeOrig)
{
actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ,
constInfo.qSeqSize);
actS2SizeOrig =
GetActualSeqLenKey(bIdx, constInfo.actualLenDims, constInfo.isAccumSeqS2, actualSeqLengthsGm, constInfo.kSeqSize, constInfo.cmpRatio); // 压缩前的actS2Size
actS2Size = actS2SizeOrig / constInfo.cmpRatio; // 真实使用的压缩后S2长度
}
template <typename QLIT>
__aicore__ inline uint32_t QLIPreload<QLIT>::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size,
uint32_t actS2SizeOrig, uint32_t &validS2Len)
{
if (actS2SizeOrig / constInfo.cmpRatio == 0) {
validS2Len = 0;
return 0;
}
uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx;
int32_t validS2LenBase = static_cast<int32_t>(actS2SizeOrig) - static_cast<int32_t>(actS1Size); // 压缩前的validS2LenBase
validS2Len = (static_cast<int32_t>(s1Offset) + validS2LenBase + static_cast<int32_t>(constInfo.s1BaseSize)) / static_cast<int32_t>(constInfo.cmpRatio);
validS2Len = Min(validS2Len, static_cast<int32_t>(actS2SizeOrig) / constInfo.cmpRatio);
validS2Len = Max(validS2Len, 1);
return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
}
template <typename QLIT>
__aicore__ inline uint32_t QLIPreload<QLIT>::GetTotalBaseBlockNum()
{
uint32_t totalBlockNum = 0;
uint32_t actS1Size, actS2Size, actS2SizeOrig;
uint32_t s1GBaseNum, s2BaseNum;
uint32_t validS2Len = 0;
for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) {
GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig);
s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize);
if (!constInfo.attenMaskFlag) {
s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize);
totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum;
continue;
}
for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) {
s2BaseNum = GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2SizeOrig, validS2Len);
totalBlockNum += s2BaseNum * constInfo.kHeadNum;
}
}
return totalBlockNum;
}
// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块
template <typename QLIT>
__aicore__ void inline QLIPreload<QLIT>::SplitCore()
{
constInfo.coreEnable = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_CORE_ENABLE_INDEX, false));
if (aiCoreIdx != 0) {
constInfo.bN2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_BN2_START_INDEX, false));
constInfo.gS1Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_M_START_INDEX, false));
constInfo.s2Start = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_S2_START_INDEX, false));
}
constInfo.bN2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_BN2_END_INDEX, false));
constInfo.gS1End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_M_END_INDEX, false));
constInfo.s2End = metadataGm.GetValue(GetAttrAbsIndex(aiCoreIdx, LI_S2_END_INDEX, false));
// 如果0核都没有启动，说明所有核都没启动
coreZeroEnable = metadataGm.GetValue(GetAttrAbsIndex(0, LI_CORE_ENABLE_INDEX, false));
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start)
{
if ASCEND_IS_AIV {
if (constInfo.outputLayout == LI_LAYOUT::TND) {
uint32_t tSizeIdx = (constInfo.batchSupperFlag) ? constInfo.batchSize : constInfo.batchSize - 1;
uint32_t tBaseIdx = (constInfo.batchSupperFlag) ? bIdx : bIdx - 1;
uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1);
uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(tBaseIdx);
uint32_t s1Count = tempLoopInfo.actS1Size;
for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) {
uint64_t indiceOutOffset =
(tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移
n2Idx * constInfo.sparseCount; // N2轴偏移
vectorService.CleanInvalidOutput(indiceOutOffset);
}
} else if (constInfo.outputLayout == LI_LAYOUT::BSND) {
for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) {
// B,S1,N2,K
uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount +
s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移
n2Idx * constInfo.sparseCount; // N2轴偏移
vectorService.CleanInvalidOutput(indiceOutOffset);
}
}
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights,
__gm__ uint8_t *queryScale, __gm__ uint8_t *keyScale,
__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK,
__gm__ uint8_t *blockTable, __gm__ uint8_t *metadata,
__gm__ uint8_t *sparseIndices, __gm__ uint8_t *workspace,
const QLITilingData *__restrict tiling, TPipe *tPipe)
{
if ASCEND_IS_AIV {
tmpBlockIdx = GetBlockIdx(); // vec:0-47
aiCoreIdx = tmpBlockIdx / 2;
} else {
tmpBlockIdx = GetBlockIdx(); // cube:0-23
aiCoreIdx = tmpBlockIdx;
}
InitTilingData(tiling);
InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK);
if (metadata != nullptr) {
metadataGm.SetGlobalBuffer((__gm__ uint32_t *)metadata);
// 计算分核
SplitCore();
}
pipe = tPipe;
// workspace 内存排布
// |mm1ResGm(存S)
uint64_t offset = 0;
// mm1开DoubleBuffer
GlobalTensor<MM1_OUT_T> mm1ResGm; // 存放S
uint64_t singleCoreMm1ResSize = WS_DOBULE * constInfo.s1BaseSize * constInfo.s2BaseSize * sizeof(MM1_OUT_T);
mm1ResGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + aiCoreIdx * singleCoreMm1ResSize));
offset += GetBlockNum() * singleCoreMm1ResSize;
GlobalTensor<half> weightWorkspaceGm; // v1阶段处理w*scale后的结果
uint64_t weightMemSize = BLOCK_CUBE * constInfo.mBaseSize * WS_DOBULE * sizeof(half);
weightWorkspaceGm.SetGlobalBuffer((__gm__ half *)(workspace + offset + aiCoreIdx * weightMemSize));
offset += GetBlockNum() * weightMemSize;
GlobalTensor<half> qScaleGm;
GlobalTensor<half> kScaleGm;
if ASCEND_IS_AIV {
vectorService.InitParams(constInfo, tiling);
indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices);
weightsGm.SetGlobalBuffer((__gm__ half *)weights);
qScaleGm.SetGlobalBuffer((__gm__ half *)queryScale);
kScaleGm.SetGlobalBuffer((__gm__ half *)keyScale);
blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
vectorService.InitVecInputTensor(weightsGm, qScaleGm, kScaleGm, indiceOutGm, blockTableGm);
vectorService.InitVecWorkspaceTensor(weightWorkspaceGm, mm1ResGm);
} else {
matmulService.InitParams(constInfo);
queryGm.SetGlobalBuffer((__gm__ Q_T *)query);
if constexpr (PAGE_ATTENTION) {
blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
}
keyGm.SetGlobalBuffer((__gm__ K_T *)key);
matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm, mm1ResGm, weightWorkspaceGm);
}
InitBuffers();
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::GetBN2Idx(uint32_t bN2Idx)
{
tempLoopInfo.bN2Idx = bN2Idx;
tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum;
tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum;
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx)
{
tempLoopInfo.gS1Idx = gS1LoopIdx;
tempLoopInfo.actMBaseSize = constInfo.mBaseSize;
uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize;
if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) {
tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail;
}
bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd);
uint32_t s2BlockNum;
uint32_t validS2Len = 0;
if (constInfo.attenMaskFlag) {
s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig,
tempLoopInfo.validS2Len);
} else {
s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
tempLoopInfo.validS2Len = tempLoopInfo.actS2Size;
}
tempLoopInfo.s2LoopEnd = (isEnd && constInfo.s2End != 0) ? constInfo.s2End : s2BlockNum;
tempLoopInfo.s2BasicSizeTail = tempLoopInfo.validS2Len % constInfo.s2BaseSize;
tempLoopInfo.s2BasicSizeTail = (tempLoopInfo.s2BasicSizeTail == 0) ?
constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail;
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::CalcGS1LoopParams(uint32_t bN2LoopIdx)
{
GetBN2Idx(bN2LoopIdx);
GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig);
if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) {
tempLoopInfo.curActSeqLenIsZero = true;
return;
}
tempLoopInfo.curActSeqLenIsZero = false;
tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize;
tempLoopInfo.mBasicSizeTail =
(tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail;
uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
tempLoopInfo.gS1LoopEnd = (bN2LoopIdx + 1 == constInfo.bN2End && constInfo.gS1End != 0) ? constInfo.gS1End : gS1SplitNum;
if constexpr (Q_LAYOUT_T == LI_LAYOUT::BSND) {
if (tempLoopInfo.gS1LoopEnd == gS1SplitNum && constInfo.qSeqSize > tempLoopInfo.actS1Size) {
tempLoopInfo.needDealActS1LessThanS1 = true;
}
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, QLICommon::RunInfo &runInfo)
{
runInfo.loop = loop;
runInfo.bIdx = tempLoopInfo.bIdx;
runInfo.gS1Idx = tempLoopInfo.gS1Idx;
runInfo.s2Idx = s2LoopIdx;
runInfo.bN2Idx = tempLoopInfo.bN2Idx;
runInfo.isValid = s2LoopIdx < tempLoopInfo.s2LoopEnd;
if (!runInfo.isValid) {
return; // 需要验证， v1 时候需要runInfo
}
runInfo.actS1Size = tempLoopInfo.actS1Size;
runInfo.actS2Size = tempLoopInfo.actS2Size;
runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig;
// 计算实际基本块size
runInfo.actMBaseSize = tempLoopInfo.actMBaseSize;
runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize;
uint32_t s2SplitNum = (tempLoopInfo.validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
if (runInfo.s2Idx == s2SplitNum - 1) {
runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail;
}
runInfo.actualSingleProcessSInnerSizeAlign =
QLICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, QLICommon::ConstInfo::BUFFER_SIZE_BYTE_32B);
runInfo.isFirstS2InnerLoop = s2LoopIdx == constInfo.s2Start;
runInfo.isLastS2InnerLoop = (s2LoopIdx + 1 == tempLoopInfo.s2LoopEnd);
if (runInfo.isFirstS2InnerLoop) {
uint64_t actualSeqQPrefixSum;
if constexpr (Q_LAYOUT_T == LI_LAYOUT::TND) {
uint32_t actualSeqLengthsGmQIdx = (constInfo.batchSupperFlag) ? runInfo.bIdx : runInfo.bIdx - 1;
actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(actualSeqLengthsGmQIdx);
} else { // BSND
actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize;
}
uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim;
// B,S1,N1(N2,G),D
queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim;
// B,S1,N1(N2,G)/T,N1(N2,G)
weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize;
// B,S1,N2,k/T,N2,k
indiceOutCoreOffset =
actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount;
}
uint64_t actualSeqKPrefixSum;
if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D
actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGm.GetValue(runInfo.bIdx - 1);
} else {
actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize;
}
uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim;
keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim;
keyScaleCoreOffset = (actualSeqKPrefixSum + runInfo.s2Idx * constInfo.s2BaseSize) * constInfo.kHeadNum;
runInfo.tensorQueryOffset = queryCoreOffset;
runInfo.tensorKeyOffset = keyCoreOffset;
runInfo.tensorKeyScaleOffset = keyScaleCoreOffset;
runInfo.tensorWeightsOffset = weightsCoreOffset;
runInfo.indiceOutOffset = indiceOutCoreOffset;
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::Process()
{
// 没有计算任务，直接清理输出
if (coreZeroEnable == 0) {
ProcessInvalid();
return;
}
ProcessMain();
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::ProcessInvalid()
{
if ASCEND_IS_AIV {
uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2
uint64_t totalOutputSize =
constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount;
uint64_t singleCoreSize =
QLICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T));
uint64_t baseSize = tmpBlockIdx * singleCoreSize;
if (baseSize < totalOutputSize) {
uint64_t dealSize =
(baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize;
GlobalTensor<OUT_T> output = indiceOutGm[baseSize];
AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX);
}
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::ProcessMain()
{
// 无任务核直接返回
if (constInfo.coreEnable == 0) {
return;
}
if ASCEND_IS_AIV {
vectorService.AllocEventID();
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE2>(constInfo.syncV1C1);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE2>(constInfo.syncV1C1);
} else {
matmulService.AllocEventID();
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V0);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V0);
}
QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE];
// 适配左闭右开
if (constInfo.bN2Start == constInfo.bN2End) {
if (constInfo.gS1Start != constInfo.gS1End || constInfo.s2Start != constInfo.s2End) {
constInfo.bN2End += 1;
}
} else if ((constInfo.gS1End != 0) || (constInfo.s2End != 0)){
constInfo.bN2End += 1;
}
uint32_t gloop = 0;
for (uint32_t bN2LoopIdx = constInfo.bN2Start; bN2LoopIdx < constInfo.bN2End; bN2LoopIdx++) {
CalcGS1LoopParams(bN2LoopIdx);
if (tempLoopInfo.curActSeqLenIsZero) {
DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U);
if ASCEND_IS_AIV {
if (bN2LoopIdx + 1 == constInfo.bN2End && gloop > 0) {
CrossCoreWaitFlag(constInfo.syncC1V1);
vectorService.ProcessVec1(runInfo[1 - gloop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE]);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE3>(
constInfo.syncV1C1); // 反向同步 1
}
}
continue;
}
for (uint32_t gS1LoopIdx = constInfo.gS1Start; gS1LoopIdx < tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) {
CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx);
bool isEnd = (bN2LoopIdx + 1 == constInfo.bN2End) && (gS1LoopIdx + 1 == tempLoopInfo.gS1LoopEnd);
uint32_t extraLoop = isEnd ? LI_QUANT_PRELOAD_TASK_CACHE_SIZE - 1 : 0; // 只preload一轮
for (uint32_t s2LoopIdx = constInfo.s2Start; s2LoopIdx < (tempLoopInfo.s2LoopEnd + extraLoop); s2LoopIdx++) {
ProcessBaseBlock(gloop, s2LoopIdx, runInfo);
++gloop;
}
constInfo.s2Start = 0;
}
if (tempLoopInfo.needDealActS1LessThanS1) {
DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size);
}
constInfo.gS1Start = 0;
}
if ASCEND_IS_AIV {
vectorService.FreeEventID();
CrossCoreWaitFlag(constInfo.syncC1V0);
CrossCoreWaitFlag(constInfo.syncC1V0);
} else {
matmulService.FreeEventID();
CrossCoreWaitFlag(constInfo.syncV1C1);
CrossCoreWaitFlag(constInfo.syncV1C1);
}
}
template <typename QLIT>
__aicore__ inline void QLIPreload<QLIT>::ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx,
QLICommon::RunInfo runInfo[LI_QUANT_PRELOAD_TASK_CACHE_SIZE])
{
int32_t curTaskId = loop % LI_QUANT_PRELOAD_TASK_CACHE_SIZE;
QLICommon::RunInfo &curRunInfo = runInfo[curTaskId];
QLICommon::RunInfo &lastRunInfo = runInfo[1 - curTaskId];
CalcRunInfo(loop, s2LoopIdx, curRunInfo);
if (curRunInfo.isValid) {
if ASCEND_IS_AIC {
if (curRunInfo.isFirstS2InnerLoop) {
CrossCoreWaitFlag(constInfo.syncV0C1);
}
CrossCoreWaitFlag(constInfo.syncV1C1); // 反向同步 1
matmulService.ComputeMm1(curRunInfo);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V1);
if (curRunInfo.isLastS2InnerLoop) {
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_FIX>(constInfo.syncC1V0); // 反向同步 0
}
} else {
if (curRunInfo.isFirstS2InnerLoop) {
CrossCoreWaitFlag(constInfo.syncC1V0); // 反向同步 0
vectorService.ProcessVec0(curRunInfo);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV0C1);
}
}
}
if (lastRunInfo.isValid) {
if ASCEND_IS_AIV {
CrossCoreWaitFlag(constInfo.syncC1V1);
vectorService.ProcessVec1(lastRunInfo);
CrossCoreSetFlag<QLICommon::ConstInfo::FIA_SYNC_MODE2, PIPE_MTE3>(constInfo.syncV1C1); // 反向同步 1
}
lastRunInfo.isValid = false;
}
}
} // namespace QLIKernel
#endif // QUANT_LIGHTNING_INDEXER_KERNEL_H

613
csrc/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_cube.h Normal file
View File

@@ -0,0 +1,613 @@
/**
 * 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 quant_lightning_indexer_service_cube.h
* \brief use 5 buffer for matmul l1, better pipeline
*/
#ifndef QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_H
#define QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_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"
#include "quant_lightning_indexer_common.h"
namespace QLIKernel {
using namespace QLICommon;
struct MmInfo {
int64_t s2L0LoopId;
int64_t s1gL0LoopId;
int64_t s2L0RealSize;
int64_t s2GmOffset;
};
template <typename QLIT>
class QLIMatmul {
public:
using Q_T = typename QLIT::queryType;
using K_T = typename QLIT::keyType;
__aicore__ inline QLIMatmul(){};
__aicore__ inline void InitBuffers(TPipe *pipe);
__aicore__ inline void InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<K_T> &keyGm,
const GlobalTensor<Q_T> &queryGm, const GlobalTensor<float> &mm1ResGm,
const GlobalTensor<half> &weightWorkspaceGm);
__aicore__ inline void InitParams(const ConstInfo &constInfo);
__aicore__ inline void AllocEventID();
__aicore__ inline void FreeEventID();
__aicore__ inline void ComputeMm1(const QLICommon::RunInfo &runInfo);
static constexpr IsResetLoad3dConfig LOAD3DV2_CONFIG = {true, true}; // isSetFMatrix isSetPadding;
static constexpr uint64_t DOUBLE_BUF_NUM = 2;
static constexpr uint64_t L0AB_BUF_NUM = 4;
static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2;
static constexpr uint32_t QW_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + DOUBLE_BUF_NUM;
static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3;
static constexpr uint32_t M_FIX_EVENT = EVENT_ID0;
static constexpr uint32_t FIX_M_EVENT = EVENT_ID2;
static constexpr uint32_t FIX_MTE1_EVENT = EVENT_ID4;
static constexpr uint64_t S8_BLOCK_CUBE = 32;
static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2;
static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2;
static constexpr uint64_t D_BASIC_BLOCK = 128;
static constexpr uint64_t S1G_BASIC_BLOCK_L1 = 256;
static constexpr uint64_t S1G_BASIC_BLOCK_L0 = 128;
static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128;
static constexpr uint64_t QUERY_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK;
static constexpr uint64_t SL1_BUFFER_OFFSET = S1G_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0;
static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK;
static constexpr uint64_t WEIGHT_BUFFER_OFFSET = S1G_BASIC_BLOCK_L1 * BLOCK_CUBE;
static constexpr uint64_t L0AB_BUFFER_OFFSET_S8_16K = 16 * 1024;
static constexpr uint64_t L0AB_BUFFER_OFFSET_FP16_16K = 16 * 512;
static constexpr uint64_t L0C_BUFFER_OFFSET = 64 * 256;
private:
__aicore__ inline void WeightDmaCopy(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo);
__aicore__ inline void LoadKeyToL0b(uint64_t s2L0RealSize);
__aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize, uint64_t s1gL0RealSize);
__aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo);
__aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo);
__aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo);
__aicore__ inline void FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize);
__aicore__ inline void LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx,
int64_t mStartPt);
__aicore__ inline void LoadWeightToL0a(uint64_t s1gL1Offset);
__aicore__ inline void ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset);
__aicore__ inline void FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset,
uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo);
__aicore__ inline void ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize);
__aicore__ inline void ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx,
const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo);
__aicore__ inline void ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt,
const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo);
__aicore__ inline void CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt, const MmInfo &lastMmInfo,
const QLICommon::RunInfo &runInfo);
static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout;
static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout;
GlobalTensor<int32_t> blkTableGm_;
GlobalTensor<K_T> keyGm_;
GlobalTensor<Q_T> queryGm_;
GlobalTensor<half> weightGm_;
GlobalTensor<float> mm1ResGm_;
TBuf<TPosition::A1> bufQL1_;
LocalTensor<Q_T> queryL1_;
TBuf<TPosition::B1> bufKeyL1_;
LocalTensor<K_T> keyL1_;
TBuf<TPosition::A1> bufWeightL1_;
LocalTensor<half> weightL1_;
TBuf<TPosition::B1> bufSL1_;
LocalTensor<half> sL1_;
TBuf<TPosition::A2> bufL0A_;
LocalTensor<Q_T> l0a_;
TBuf<TPosition::B2> bufL0B_;
LocalTensor<K_T> l0b_;
TBuf<TPosition::CO1> bufL0C_;
LocalTensor<int32_t> cL0_;
uint64_t keyL1BufIdx_ = 0;
uint64_t qwL1Mte2BufIdx_ = 0;
uint64_t sL1BufIdx_ = 0;
uint64_t l0BufIdx_ = 0;
uint64_t l0cBufIdx_ = 0;
ConstInfo constInfo_;
};
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::InitParams(const ConstInfo &constInfo)
{
constInfo_ = constInfo;
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::InitBuffers(TPipe *pipe)
{
pipe->InitBuffer(bufQL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * D_BASIC_BLOCK * sizeof(Q_T));
queryL1_ = bufQL1_.Get<Q_T>();
pipe->InitBuffer(bufKeyL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * D_BASIC_BLOCK * sizeof(K_T));
keyL1_ = bufKeyL1_.Get<K_T>();
pipe->InitBuffer(bufWeightL1_, DOUBLE_BUF_NUM * S1G_BASIC_BLOCK_L1 * BLOCK_CUBE * sizeof(half));
weightL1_ = bufWeightL1_.Get<half>();
pipe->InitBuffer(bufSL1_, DOUBLE_BUF_NUM * S2_BASIC_BLOCK_L0 * S1G_BASIC_BLOCK_L0 * sizeof(half));
sL1_ = bufSL1_.Get<half>();
pipe->InitBuffer(bufL0A_, 64 * 1024);
l0a_ = bufL0A_.Get<Q_T>();
pipe->InitBuffer(bufL0B_, 64 * 1024);
l0b_ = bufL0B_.Get<K_T>();
pipe->InitBuffer(bufL0C_, 128 * 1024);
cL0_ = bufL0C_.Get<int32_t>();
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm,
const GlobalTensor<K_T> &keyGm,
const GlobalTensor<Q_T> &queryGm,
const GlobalTensor<float> &mm1ResGm,
const GlobalTensor<half> &weightWorkspaceGm)
{
blkTableGm_ = blkTableGm;
keyGm_ = keyGm;
queryGm_ = queryGm;
mm1ResGm_ = mm1ResGm;
weightGm_ = weightWorkspaceGm;
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::ProcessWs(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t sL1BufIdx,
const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo)
{
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM);
for (int64_t s1gOffset = 0; s1gOffset < s1gL0RealSize; s1gOffset += constInfo_.gSize) {
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM);
LoadSToL0b(s1gL0RealSize, mmInfo.s2L0RealSize, sL1BufIdx, s1gOffset);
LoadWeightToL0a(s1gOffset + s1gL1Offset);
ComputeWs(s1gL0RealSize, mmInfo.s2L0RealSize, s1gOffset);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM);
l0BufIdx_++;
}
FixpResToGm(s1gL0RealSize / constInfo_.gSize, mmInfo.s2L0RealSize, s1gL1Offset / constInfo_.gSize,
mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0, runInfo);
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM);
l0cBufIdx_++;
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::ProcessQk(uint64_t s1gL0RealSize, uint64_t s1gL1Offset, uint64_t s1L0LoopCnt,
const MmInfo &mmInfo, const QLICommon::RunInfo &runInfo)
{
if (mmInfo.s1gL0LoopId == 0) {
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM);
if constexpr (K_LAYOUT_T == LI_LAYOUT::PA_BSND) {
KeyNd2NzForPA(mmInfo.s2L0RealSize, runInfo.s2Idx * constInfo_.s2BaseSize + mmInfo.s2GmOffset, runInfo);
} else {
KeyNd2Nz(mmInfo.s2L0RealSize, mmInfo, runInfo);
}
SetFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
WaitFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
}
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM);
LoadQueryToL0a(s1gL1Offset, runInfo.actMBaseSize, s1gL0RealSize);
LoadKeyToL0b(mmInfo.s2L0RealSize);
if (mmInfo.s1gL0LoopId + 1 >= s1L0LoopCnt) {
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % DOUBLE_BUF_NUM);
keyL1BufIdx_++;
}
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM);
ComputeQk(s1gL0RealSize, mmInfo.s2L0RealSize);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0AB_BUF_NUM);
FixpSToL1(s1gL0RealSize, mmInfo.s2L0RealSize);
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0cBufIdx_ % DOUBLE_BUF_NUM);
l0BufIdx_++;
l0cBufIdx_++;
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::CalcMmInfo(MmInfo &mmInfo, uint64_t loopIdx, uint64_t s1L0LoopCnt,
const MmInfo &lastMmInfo, const QLICommon::RunInfo &runInfo)
{
mmInfo.s2L0LoopId = loopIdx / s1L0LoopCnt;
mmInfo.s1gL0LoopId = loopIdx % s1L0LoopCnt;
if (mmInfo.s1gL0LoopId == 0) {
mmInfo.s2GmOffset = mmInfo.s2L0LoopId * S2_BASIC_BLOCK_L0;
mmInfo.s2L0RealSize = mmInfo.s2GmOffset + S2_BASIC_BLOCK_L0 > runInfo.actualSingleProcessSInnerSize
? runInfo.actualSingleProcessSInnerSize - mmInfo.s2GmOffset
: S2_BASIC_BLOCK_L0;
} else {
mmInfo.s2L0RealSize = lastMmInfo.s2L0RealSize;
}
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::ComputeMm1(const QLICommon::RunInfo &runInfo)
{
if (runInfo.isFirstS2InnerLoop) {
WaitFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM);
QueryNd2Nz(runInfo.actMBaseSize, runInfo); // 256 * 128 // L1BasicBlock
WeightDmaCopy(runInfo.actMBaseSize, runInfo);
}
int64_t loopIdx = 0;
int64_t s2L0LoopCnt = CeilDiv(runInfo.actualSingleProcessSInnerSize, S2_BASIC_BLOCK_L0); // 2048取128
int64_t s1L0LoopCnt = CeilDiv(runInfo.actMBaseSize, S1G_BASIC_BLOCK_L0); // 256取128
int64_t s1gL1Offset[2] = {0, static_cast<int64_t>(S1G_BASIC_BLOCK_L0)};
int64_t s1gL0RealSize[2] = {s1L0LoopCnt > 1 ? static_cast<int64_t>(S1G_BASIC_BLOCK_L0) : runInfo.actMBaseSize,
runInfo.actMBaseSize - s1gL1Offset[1]};
MmInfo mmInfo[2];
CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo);
ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt],
s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1],
runInfo);
SetFlag<HardEvent::FIX_MTE1>(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM);
sL1BufIdx_++;
loopIdx++;
while (loopIdx < s2L0LoopCnt * s1L0LoopCnt) {
CalcMmInfo(mmInfo[loopIdx & 1], loopIdx, s1L0LoopCnt, mmInfo[(loopIdx + 1) & 1], runInfo);
ProcessQk(s1gL0RealSize[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt],
s1gL1Offset[mmInfo[loopIdx & 1].s1gL0LoopId % s1L0LoopCnt], s1L0LoopCnt, mmInfo[loopIdx & 1],
runInfo);
SetFlag<HardEvent::FIX_MTE1>(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM);
sL1BufIdx_++;
WaitFlag<HardEvent::FIX_MTE1>(FIX_MTE1_EVENT + sL1BufIdx_ % DOUBLE_BUF_NUM);
ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt],
s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_,
mmInfo[(loopIdx + 1) & 1], runInfo);
loopIdx++;
}
WaitFlag<HardEvent::FIX_MTE1>(FIX_MTE1_EVENT + (sL1BufIdx_ + 1) % DOUBLE_BUF_NUM);
ProcessWs(s1gL0RealSize[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt],
s1gL1Offset[mmInfo[(loopIdx + 1) & 1].s1gL0LoopId % s1L0LoopCnt], sL1BufIdx_ - 1,
mmInfo[(loopIdx + 1) & 1], runInfo);
if (runInfo.isLastS2InnerLoop) {
SetFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM);
qwL1Mte2BufIdx_++;
}
}
// blkNum, blkSize, N2, D
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset,
const QLICommon::RunInfo &runInfo)
{
uint64_t s2L1Offset = 0;
while (s2L1Offset < s2L1RealSize) {
uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize;
uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize;
uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) *
constInfo_.stride +
s2BlkOffset * constInfo_.headDim;
uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset;
s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize ? constInfo_.kCacheBlockSize - s2BlkOffset
: s2Mte2Size;
Nd2NzParams nd2nzPara;
nd2nzPara.ndNum = 1;
nd2nzPara.nValue = s2Mte2Size; // 行数
nd2nzPara.dValue = constInfo_.headDim;
nd2nzPara.srcDValue = constInfo_.headDim;
nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
nd2nzPara.dstNzNStride = 1;
nd2nzPara.srcNdMatrixStride = 0;
nd2nzPara.dstNzMatrixStride = 0;
DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * S8_BLOCK_CUBE],
keyGm_[keyGmOffset], nd2nzPara);
s2L1Offset += s2Mte2Size;
}
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::KeyNd2Nz(uint64_t s2L1RealSize, const MmInfo &mmInfo,
const QLICommon::RunInfo &runInfo)
{
uint64_t dStride = constInfo_.headDim;
if constexpr (K_LAYOUT_T == LI_LAYOUT::BSND || K_LAYOUT_T == LI_LAYOUT::TND) {
dStride = constInfo_.headDim * constInfo_.kHeadNum; // constInfo_.kHeadNum
}
Nd2NzParams nd2nzPara;
nd2nzPara.ndNum = 1;
nd2nzPara.nValue = s2L1RealSize; // 行数
nd2nzPara.dValue = constInfo_.headDim;
nd2nzPara.srcDValue = dStride;
nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
nd2nzPara.dstNzNStride = 1;
nd2nzPara.srcNdMatrixStride = 0;
nd2nzPara.dstNzMatrixStride = 0;
// 默认一块buf最多放两份
DataCopy(keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET],
keyGm_[runInfo.tensorKeyOffset + mmInfo.s2GmOffset * constInfo_.headDim], nd2nzPara);
}
// batch, s1, g, 1
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::WeightDmaCopy(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo)
{
DataCopyParams copyInParams;
copyInParams.blockCount = 1;
copyInParams.blockLen = s1gL1RealSize;
copyInParams.srcStride = 0;
copyInParams.dstStride = 0;
DataCopy(weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET],
weightGm_[runInfo.loop % DOUBLE_BUF_NUM * BLOCK_CUBE * constInfo_.mBaseSize], copyInParams);
}
// batch, s1, n2, g, d
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::QueryNd2Nz(uint64_t s1gL1RealSize, const QLICommon::RunInfo &runInfo)
{
Nd2NzParams nd2nzPara;
nd2nzPara.ndNum = 1;
nd2nzPara.nValue = s1gL1RealSize; // 行数
nd2nzPara.dValue = constInfo_.headDim;
nd2nzPara.srcDValue = constInfo_.headDim;
nd2nzPara.dstNzC0Stride = CeilAlign(s1gL1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
nd2nzPara.dstNzNStride = 1;
nd2nzPara.srcNdMatrixStride = 0;
nd2nzPara.dstNzMatrixStride = 0;
// 默认一块buf最多放两份
DataCopy(queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET], queryGm_[runInfo.tensorQueryOffset],
nd2nzPara);
}
// s1g, d
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL1RealSize,
uint64_t s1gL0RealSize)
{
LoadData3DParamsV2<Q_T> loadData3DParams;
// SetFmatrixParams
loadData3DParams.l1H = CeilDiv(s1gL1RealSize, BLOCK_CUBE); // Hin=M1=8
loadData3DParams.l1W = BLOCK_CUBE; // Win=M0
loadData3DParams.channelSize = constInfo_.headDim; // Cin=K
loadData3DParams.padList[0] = 0;
loadData3DParams.padList[1] = 0;
loadData3DParams.padList[2] = 0;
loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果
// SetLoadToA0Params
loadData3DParams.mExtension = CeilAlign(s1gL0RealSize, BLOCK_CUBE); // M height维度目的
loadData3DParams.kExtension = constInfo_.headDim; // K width维度目的
loadData3DParams.mStartPt = s1gL1Offset;
loadData3DParams.kStartPt = 0;
loadData3DParams.strideW = 1;
loadData3DParams.strideH = 1;
loadData3DParams.filterW = 1;
loadData3DParams.filterSizeW = (1 >> 8) & 255;
loadData3DParams.filterH = 1;
loadData3DParams.filterSizeH = (1 >> 8) & 255;
loadData3DParams.dilationFilterW = 1;
loadData3DParams.dilationFilterH = 1;
loadData3DParams.enTranspose = 0;
loadData3DParams.fMatrixCtrl = 0;
LoadData<Q_T, LOAD3DV2_CONFIG>(l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K],
queryL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * QUERY_BUFFER_OFFSET],
loadData3DParams);
}
// s1, g, s2 --> 2 * 64* 128
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::LoadSToL0b(uint64_t s1gL1RealSize, uint64_t s2L0RealSize, uint64_t sL1BufIdx,
int64_t mStartPt)
{
LoadData3DParamsV2<half> loadData3DParams;
// SetFmatrixParams
loadData3DParams.l1H = S1G_BASIC_BLOCK_L0 / BLOCK_CUBE; // Hin=M1=8
loadData3DParams.l1W = BLOCK_CUBE; // Win=M0
loadData3DParams.channelSize = CeilAlign(s2L0RealSize, BLOCK_CUBE); // Cin=K
loadData3DParams.padList[0] = 0;
loadData3DParams.padList[1] = 0;
loadData3DParams.padList[2] = 0;
loadData3DParams.padList[3] = 255; // 尾部数据不影响滑窗的结果
// SetLoadToA0Params
loadData3DParams.mExtension = constInfo_.gSize; // M height维度目的
loadData3DParams.kExtension = CeilAlign(s2L0RealSize, BLOCK_CUBE); // K width维度目的
loadData3DParams.kStartPt = 0;
loadData3DParams.strideW = 1;
loadData3DParams.strideH = 1;
loadData3DParams.filterW = 1;
loadData3DParams.filterSizeW = (1 >> 8) & 255;
loadData3DParams.filterH = 1;
loadData3DParams.filterSizeH = (1 >> 8) & 255;
loadData3DParams.dilationFilterW = 1;
loadData3DParams.dilationFilterH = 1;
loadData3DParams.enTranspose = 1;
loadData3DParams.fMatrixCtrl = 0;
loadData3DParams.mStartPt = mStartPt;
LoadData<half, LOAD3DV2_CONFIG>(
l0b_.template ReinterpretCast<half>()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K],
sL1_[(sL1BufIdx % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET], loadData3DParams);
}
// s1,g,1(16), 2,64,16
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::LoadWeightToL0a(uint64_t s1gL1Offset)
{
LoadData2DParams loadData2DParams;
loadData2DParams.startIndex = 0;
loadData2DParams.repeatTimes = CeilDiv(constInfo_.gSize, BLOCK_CUBE);
loadData2DParams.srcStride = 1;
loadData2DParams.dstGap = 0;
loadData2DParams.ifTranspose = true;
LoadData(l0a_.template ReinterpretCast<half>()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K],
weightL1_[(qwL1Mte2BufIdx_ % DOUBLE_BUF_NUM) * WEIGHT_BUFFER_OFFSET + s1gL1Offset* BLOCK_CUBE],
loadData2DParams);
}
// s2, d -> 128,128
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::LoadKeyToL0b(uint64_t s2L0RealSize)
{
LoadData2DParams loadData2DParams;
loadData2DParams.startIndex = 0;
loadData2DParams.repeatTimes = CeilDiv(s2L0RealSize, BLOCK_CUBE) * CeilDiv(constInfo_.headDim, S8_BLOCK_CUBE);
loadData2DParams.srcStride = 1;
loadData2DParams.dstGap = 0;
loadData2DParams.ifTranspose = false;
LoadData(l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K],
keyL1_[(keyL1BufIdx_ % DOUBLE_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParams);
}
// A: s1,g,1(16) B: s1,g,s2 C: s1, 1(16), s2
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::ComputeWs(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, int64_t s1gOffset)
{
SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
WaitFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
MmadParams mmadParams;
mmadParams.m = BLOCK_CUBE;
mmadParams.n = s2L0RealSize;
mmadParams.k = constInfo_.gSize;
mmadParams.cmatrixInitVal = true;
mmadParams.cmatrixSource = false;
Mmad(cL0_.template ReinterpretCast<float>()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET +
s1gOffset * S2_BASIC_BLOCK_L0],
l0a_.template ReinterpretCast<half>()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K],
l0b_.template ReinterpretCast<half>()[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_FP16_16K],
mmadParams);
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::ComputeQk(uint64_t s1gL0RealSize, uint64_t s2L0RealSize)
{
SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
WaitFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
MmadParams mmadParams;
mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE);
mmadParams.n = s2L0RealSize;
mmadParams.k = constInfo_.headDim;
mmadParams.cmatrixInitVal = true;
mmadParams.cmatrixSource = false;
Mmad(cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET],
l0a_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K],
l0b_[(l0BufIdx_ % L0AB_BUF_NUM) * L0AB_BUFFER_OFFSET_S8_16K], mmadParams);
if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) {
PipeBarrier<PIPE_M>();
}
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::FixpSToL1(uint64_t s1gL0RealSize, uint64_t s2L0RealSize)
{
SetFlag<HardEvent::M_FIX>(M_FIX_EVENT);
WaitFlag<HardEvent::M_FIX>(M_FIX_EVENT);
DataCopyCO12DstParams params;
params.mSize = CeilAlign(s1gL0RealSize, BLOCK_CUBE);
params.nSize = CeilAlign(s2L0RealSize, BLOCK_CUBE);
params.dstStride = S1G_BASIC_BLOCK_L0;
params.srcStride = params.mSize;
params.quantPre = QuantMode_t::DEQF16;
params.reluPre = 1;
params.channelSplit = 0;
params.nz2ndEn = 0;
SetFixpipePreQuantFlag(0x3a800000);
DataCopy(sL1_[(sL1BufIdx_ % DOUBLE_BUF_NUM) * SL1_BUFFER_OFFSET],
cL0_[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET], params);
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::FixpResToGm(uint64_t s1L0RealCount, uint64_t s2L0RealSize, uint64_t s1GmOffset,
uint64_t s2GmOffset, const QLICommon::RunInfo &runInfo)
{
SetFlag<HardEvent::M_FIX>(M_FIX_EVENT);
WaitFlag<HardEvent::M_FIX>(M_FIX_EVENT);
AscendC::DataCopyCO12DstParams intriParams;
intriParams.mSize = 1;
intriParams.nSize = s2L0RealSize;
intriParams.dstStride = constInfo_.s2BaseSize;
intriParams.srcStride = 16;
// set mode according to dtype
intriParams.quantPre = QuantMode_t::NoQuant;
intriParams.nz2ndEn = true;
intriParams.reluPre = 0;
AscendC::SetFixpipeNz2ndFlag(s1L0RealCount, CeilDiv(constInfo_.gSize, BLOCK_CUBE) * S2_BASIC_BLOCK_L0 / BLOCK_CUBE,
2048);
AscendC::DataCopy(mm1ResGm_[(runInfo.loop % 2) * constInfo_.mBaseSize / constInfo_.gSize * constInfo_.s2BaseSize +
s1GmOffset * intriParams.dstStride + s2GmOffset],
cL0_.template ReinterpretCast<float>()[(l0cBufIdx_ % DOUBLE_BUF_NUM) * L0C_BUFFER_OFFSET],
intriParams);
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::AllocEventID()
{
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
SetFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + 0);
SetFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + 1);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 2);
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 3);
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + 0);
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + 1);
}
template <typename QLIT>
__aicore__ inline void QLIMatmul<QLIT>::FreeEventID()
{
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
WaitFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + 0);
WaitFlag<HardEvent::MTE1_MTE2>(QW_MTE1_MTE2_EVENT + 1);
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 2);
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 3);
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + 0);
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + 1);
}
} // namespace QLIKernel
#endif // QUANT_LIGHTNING_INDEXER_SERVICE_CUBE_H

437
csrc/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_service_vector.h Normal file
View File

@@ -0,0 +1,437 @@
/**
 * 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 quant_lightning_indexer_service_vector.h
* \brief
*/
#ifndef QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_H
#define QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_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"
#include "quant_lightning_indexer_common.h"
#include "quant_lightning_indexer_vector.h"
namespace QLIKernel {
using namespace QLICommon;
using namespace QLIServiceVec;
constexpr uint32_t BASE_TOPK = 2048;
constexpr uint32_t BASE_TOPK_VALUE_IDX_SIZE = 4096;
constexpr uint32_t ELE_NUM_32 = 32;
constexpr uint32_t ELE_NUM_128 = 128;
constexpr uint32_t ELE_NUM_512 = 512;
template <typename QLIT>
class QLIVector {
public:
// =================================类型定义区=================================
static constexpr LI_LAYOUT Q_LAYOUT_T = QLIT::layout;
static constexpr LI_LAYOUT K_LAYOUT_T = QLIT::keyLayout;
static constexpr bool PAGE_ATTENTION = QLIT::pageAttention;
// MM输出数据类型, 当前只支持float
using MM1_OUT_T = float;
__aicore__ inline QLIVector(){};
__aicore__ inline void ProcessVec0(const QLICommon::RunInfo &info);
__aicore__ inline void ProcessVec1(const QLICommon::RunInfo &info);
__aicore__ inline void InitBuffers(TPipe *pipe);
__aicore__ inline void InitParams(const struct QLICommon::ConstInfo &constInfo,
const QLITilingData *__restrict tilingData);
__aicore__ inline void InitVecWorkspaceTensor(GlobalTensor<half> vec0OutGm, GlobalTensor<MM1_OUT_T> mm1ResGm);
__aicore__ inline void InitVecInputTensor(GlobalTensor<half> weightsGm, GlobalTensor<half> qScaleGm,
GlobalTensor<half> kScaleGm, GlobalTensor<int32_t> indiceOutGm,
GlobalTensor<int32_t> blockTableGm);
__aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset);
__aicore__ inline int32_t AlignS2(int32_t cuS2Len);
__aicore__ inline void AllocEventID();
__aicore__ inline void FreeEventID();
protected:
GlobalTensor<MM1_OUT_T> mm1ResGm;
GlobalTensor<half> weightsGm;
GlobalTensor<half> qScaleGm;
GlobalTensor<half> kScaleGm;
GlobalTensor<half> vec0OutGm;
GlobalTensor<int32_t> indiceOutGm;
GlobalTensor<int32_t> blockTableGm;
// =================================常量区=================================
private:
__aicore__ inline void GetKeyScale(const QLICommon::RunInfo &runInfo, const LocalTensor<half> &resUb,
int64_t batchId, int64_t startS2, int64_t getLen);
// ================================Local Buffer区====================================
// queue
TQue<QuePosition::VECIN, 1> inQueue_;
TQue<QuePosition::VECOUT, 1> outQueue_;
// tmp buff for vector
TBuf<TPosition::VECCALC> sortOutBuf_;
TBuf<TPosition::VECCALC> indexBuf_;
TBuf<TPosition::VECCALC> tmpBuf_;
LocalTensor<int32_t> globalTopkIndice_;
LocalTensor<float> globalTopkUb_;
int32_t blockId_ = -1;
// para for vector
int32_t groupInner_ = 0;
int32_t globalTopkNum_ = 0;
int64_t blockS2StartIdx_ = 0;
int32_t gSize_ = 0;
int32_t kSeqSize_ = 0;
int32_t kHeadNum_ = 0;
int32_t qHeadNum_ = 0;
int32_t s1BaseSize_ = 0;
int32_t s2BaseSize_ = 0;
int32_t kCacheBlockSize_ = 0;
int32_t maxBlockNumPerBatch_ = 0;
struct QLICommon::ConstInfo constInfo_;
};
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::GetKeyScale(const QLICommon::RunInfo &runInfo, const LocalTensor<half> &resUb,
int64_t batchId, int64_t startS2, int64_t getLen)
{
// startS2一定能整除kCacheBlockSize_
AscendC::DataCopyPadExtParams<half> padParams{false, 0, 0, 0};
AscendC::DataCopyExtParams copyInParams;
if constexpr (PAGE_ATTENTION) {
int32_t startBlockTableIdx = startS2 / kCacheBlockSize_;
int32_t startBlockTableOffset = startS2 % kCacheBlockSize_;
int32_t blockTableBatchOffset = batchId * maxBlockNumPerBatch_;
copyInParams.blockCount = 1;
copyInParams.srcStride = 0;
copyInParams.dstStride = 0;
copyInParams.rsv = 0;
int32_t resUbBaseOffset = 0;
if (startBlockTableOffset > 0) {
int32_t firstPartLen =
kCacheBlockSize_ - startBlockTableOffset > getLen ? getLen : kCacheBlockSize_ - startBlockTableOffset;
copyInParams.blockLen = firstPartLen * sizeof(half);
int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx);
SetWaitFlag<HardEvent::S_MTE2>(HardEvent::S_MTE2);
AscendC::DataCopyPad(resUb, kScaleGm[blockId * constInfo_.scaleStride + startBlockTableOffset],
copyInParams, padParams);
startBlockTableIdx++;
getLen = getLen - firstPartLen;
resUbBaseOffset = firstPartLen;
}
int32_t getLoopNum = CeilDiv(getLen, kCacheBlockSize_);
copyInParams.blockLen = kCacheBlockSize_ * sizeof(half);
for (int32_t i = 0; i < getLoopNum; i++) {
if (i == getLoopNum - 1) {
copyInParams.blockLen = (getLen - i * kCacheBlockSize_) * sizeof(half);
}
int32_t blockId = blockTableGm.GetValue(blockTableBatchOffset + startBlockTableIdx + i);
SetWaitFlag<HardEvent::S_MTE2>(HardEvent::S_MTE2);
AscendC::DataCopyPad(resUb[resUbBaseOffset + i * kCacheBlockSize_], kScaleGm[blockId * constInfo_.scaleStride],
copyInParams, padParams);
}
} else {
copyInParams.blockCount = 1;
copyInParams.blockLen = getLen * sizeof(half);
copyInParams.srcStride = 0;
copyInParams.dstStride = 0;
copyInParams.rsv = 0;
AscendC::DataCopyPad(resUb, kScaleGm[runInfo.tensorKeyScaleOffset], copyInParams, padParams);
}
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::InitBuffers(TPipe *pipe)
{
pipe->InitBuffer(inQueue_, 2, s2BaseSize_ * sizeof(float) * 2); // 32KB
pipe->InitBuffer(outQueue_, 1, BASE_TOPK * sizeof(float)); // 8 KB
pipe->InitBuffer(indexBuf_, s2BaseSize_ * sizeof(int32_t)); // 8 KB
pipe->InitBuffer(tmpBuf_, 64 * 1024); // 64KB
pipe->InitBuffer(sortOutBuf_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE * sizeof(float)); // 32KB
globalTopkIndice_ = indexBuf_.Get<int32_t>();
globalTopkUb_ = sortOutBuf_.Get<float>();
globalTopkNum_ = 0;
// 基本块执行前初始化UB和GM
// step1. 初始化一个有序索引 0 - s2BaseSize_
ArithProgression<int32_t>(globalTopkIndice_, 0, 1, s2BaseSize_);
// step2. globalTopkUb_ [CeilDiv(s1BaseSize_, 2), BASE_TOPK, 2] -inf,-1
InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE);
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::InitParams(const struct QLICommon::ConstInfo &constInfo,
const QLITilingData *__restrict tilingData)
{
this->constInfo_ = constInfo;
blockS2StartIdx_ = 0;
gSize_ = constInfo.gSize;
kSeqSize_ = constInfo.kSeqSize;
// define N2 para
kHeadNum_ = constInfo.kHeadNum;
qHeadNum_ = constInfo.qHeadNum;
// define MMBase para
s1BaseSize_ = constInfo.s1BaseSize; // 4
s2BaseSize_ = constInfo.s2BaseSize; // 2048
kCacheBlockSize_ = constInfo.kCacheBlockSize;
maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch;
blockId_ = GetBlockIdx();
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::InitVecInputTensor(GlobalTensor<half> weightsGm, GlobalTensor<half> qScaleGm,
GlobalTensor<half> kScaleGm,
GlobalTensor<int32_t> indiceOutGm,
GlobalTensor<int32_t> blockTableGm)
{
this->weightsGm = weightsGm;
this->qScaleGm = qScaleGm;
this->kScaleGm = kScaleGm;
this->indiceOutGm = indiceOutGm;
this->blockTableGm = blockTableGm;
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::InitVecWorkspaceTensor(GlobalTensor<half> vec0OutGm,
GlobalTensor<MM1_OUT_T> mm1ResGm)
{
this->mm1ResGm = mm1ResGm;
this->vec0OutGm = vec0OutGm;
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::AllocEventID()
{
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::FreeEventID()
{
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::CleanInvalidOutput(int64_t invalidS1offset)
{
// init -1 and copy to output
LocalTensor<float> valueULocal = outQueue_.AllocTensor<float>();
LocalTensor<int32_t> idxULocal1 = valueULocal.template ReinterpretCast<int32_t>();
Duplicate(idxULocal1, constInfo_.INVALID_IDX, constInfo_.sparseCount);
outQueue_.EnQue<float>(valueULocal);
valueULocal = outQueue_.DeQue<float>();
QLIServiceVec::CopyOut(indiceOutGm[invalidS1offset], idxULocal1, constInfo_.sparseCount);
outQueue_.FreeTensor(valueULocal);
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::ProcessVec0(const QLICommon::RunInfo &info)
{
// 只需要一个v核做
if (blockId_ % 2 != 0) {
return;
}
int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_;
// 计算输出w基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*64 + aic_offset
int64_t vec0OutGmOffset = (info.loop % 2) * ((s1BaseSize_ * gSize_ * BLOCK_CUBE));
// 计算输入weight的地址偏移，qScale的地址偏移与weight相同
int64_t weightGmOffset = info.tensorWeightsOffset + cuBaseS1Idx * qHeadNum_;
// 当前需要计算的S1行数，处理尾块场景
int32_t cuS1ProcNum = cuBaseS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_;
int32_t cuProcEleNum = cuS1ProcNum * gSize_;
LocalTensor<half> inWeightsUb = inQueue_.AllocTensor<half>();
LocalTensor<half> inQScaleUb = inWeightsUb[cuProcEleNum];
AscendC::DataCopyPadExtParams<half> padParams{false, 0, 0, 0};
AscendC::DataCopyExtParams copyInParams;
copyInParams.blockCount = 1;
copyInParams.blockLen = cuProcEleNum * sizeof(half);
copyInParams.srcStride = 0;
copyInParams.dstStride = 0;
copyInParams.rsv = 0;
AscendC::DataCopyPad(inWeightsUb, weightsGm[weightGmOffset], copyInParams, padParams);
AscendC::DataCopyPad(inQScaleUb, qScaleGm[weightGmOffset], copyInParams, padParams);
inQueue_.EnQue<half>(inWeightsUb);
inWeightsUb = inQueue_.DeQue<half>();
AscendC::Mul(inWeightsUb, inWeightsUb, inQScaleUb, cuProcEleNum);
PipeBarrier<PIPE_V>();
LocalTensor<half> resUb = outQueue_.AllocTensor<half>();
AscendC::Brcb(resUb, inWeightsUb, static_cast<uint8_t>(cuProcEleNum / 8), {1, 8});
inQueue_.FreeTensor(inWeightsUb);
outQueue_.EnQue<half>(resUb);
resUb = outQueue_.DeQue<half>();
AscendC::DataCopyParams copyOutParams;
copyOutParams.blockCount = 1;
copyOutParams.blockLen = cuProcEleNum * BLOCK_CUBE * sizeof(half);
copyOutParams.srcStride = 0;
copyOutParams.dstStride = 0;
AscendC::DataCopyPad(vec0OutGm[vec0OutGmOffset], resUb, copyOutParams);
outQueue_.FreeTensor(resUb);
}
template <typename QLIT>
__aicore__ inline int32_t QLIVector<QLIT>::AlignS2(int32_t cuS2Len)
{
// 限制：当前cuS2Len最大为2048，暂不考虑更长
// 该函数目的是将cuS2Len对齐到形如 32*(4^n)*m 的形式 (m ∈ [1, 3])，方便后续sort/merge
if (cuS2Len <= ELE_NUM_128) {
return Align(cuS2Len, ELE_NUM_32);
} else if (cuS2Len <= ELE_NUM_512) {
return Align(cuS2Len, ELE_NUM_128);
} else {
return Align(cuS2Len, ELE_NUM_512);
}
}
template <typename QLIT>
__aicore__ inline void QLIVector<QLIT>::ProcessVec1(const QLICommon::RunInfo &info)
{
int32_t cuBaseS1Idx = info.gS1Idx * s1BaseSize_;
int32_t cuBaseS2Idx = info.s2Idx * s2BaseSize_;
// 计算基本块基地址偏移 偶数循环 -> 0 + aic_offset 奇数循环 -> 4*2048 + aic_offset
int64_t mmGmOffset = (info.loop % 2) * (s1BaseSize_ * s2BaseSize_);
// cuS1BeginIdxPerAiv: 每个AIV的S1起始偏移
int32_t cuS1BeginIdxPerAiv = cuBaseS1Idx;
int32_t cuS1ProcNum =
cuS1BeginIdxPerAiv + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_;
// cuS1ProcNumPerAiv: 每个AIv的S1计算量
int32_t cuS1ProcNumPerAiv = blockId_ % 2 == 0 ? CeilDiv(cuS1ProcNum, 2) : (cuS1ProcNum / 2);
cuS1BeginIdxPerAiv += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2);
// 基本块基地址偏移奇数核加一个S1地址偏移
mmGmOffset += (blockId_ % 2) * CeilDiv(cuS1ProcNum, 2) * s2BaseSize_;
// 非首个基本块, M(S1)轴发生切换需要初始化
if (info.loop != 0 && info.s2Idx == 0) {
// globalTopkUb_ value,index=-inf,-1
InitSortOutBuf(globalTopkUb_, CeilDiv(s1BaseSize_, 2) * BASE_TOPK_VALUE_IDX_SIZE);
blockS2StartIdx_ = 0;
} else if (info.loop == 0) {
blockS2StartIdx_ = info.s2Idx;
}
// cuRealAcSeq: 当前基本块S1对应的AcSeq
int32_t cuRealAcSeq = info.actS2Size;
int32_t cuRealAcSeqCount = 0;
if (constInfo_.attenMaskFlag) {
// attenMask true场景
cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + cuS1BeginIdxPerAiv;
}
int32_t cuRealAcSeqIni = cuRealAcSeq;
// LD输出S1方向偏移，保证2个Vector输出的内容连续
uint32_t ldS1Offset = (blockId_ % 2 == 0) ? s1BaseSize_ / 2 - cuS1ProcNumPerAiv : 0;
for (int innerS1Idx = 0; innerS1Idx < cuS1ProcNumPerAiv; innerS1Idx++) {
if (constInfo_.attenMaskFlag) {
cuRealAcSeqCount += 1;
cuRealAcSeq = (cuRealAcSeqCount + cuRealAcSeqIni) / static_cast<int32_t>(constInfo_.cmpRatio);
}
int32_t cuS2Len = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq ? cuRealAcSeq - cuBaseS2Idx : s2BaseSize_;
int32_t cuS1Idx = cuS1BeginIdxPerAiv + innerS1Idx;
if (cuRealAcSeq > 0 && cuS2Len > 0) {
int32_t cuS2LenVecAlign = AlignS2(cuS2Len);
LocalTensor<float> mmInUb = inQueue_.AllocTensor<float>();
LocalTensor<float> kScaleUb = mmInUb[cuS2LenVecAlign];
LocalTensor<half> kScaleTUb = kScaleUb.template ReinterpretCast<half>()[cuS2LenVecAlign];
AscendC::DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
AscendC::DataCopyPadExtParams<half> padTParams{false, 0, 0, 0};
AscendC::DataCopyExtParams copyInParams;
copyInParams.blockCount = 1;
copyInParams.blockLen = cuS2Len * sizeof(float);
copyInParams.srcStride = 0;
copyInParams.dstStride = 0;
copyInParams.rsv = 0;
AscendC::DataCopyPad(mmInUb, mm1ResGm[mmGmOffset + innerS1Idx * s2BaseSize_], copyInParams, padParams);
GetKeyScale(info, kScaleTUb, info.bIdx, cuBaseS2Idx, cuS2Len);
inQueue_.EnQue<float>(mmInUb);
mmInUb = inQueue_.DeQue<float>();
AscendC::Cast(kScaleUb, kScaleTUb, RoundMode::CAST_NONE, cuS2Len);
PipeBarrier<PIPE_V>();
AscendC::Mul(mmInUb, mmInUb, kScaleUb, cuS2Len);
PipeBarrier<PIPE_V>();
LocalTensor<float> sortBuff = tmpBuf_.Get<float>();
LocalTensor<float> sortScoreUb = sortBuff;
LocalTensor<float> sortIndiceUb = sortBuff[cuS2LenVecAlign];
PipeBarrier<PIPE_V>();
Duplicate(sortScoreUb.template ReinterpretCast<int32_t>(), QLIServiceVec::NEG_INF, cuS2LenVecAlign);
PipeBarrier<PIPE_V>();
Adds(sortScoreUb, mmInUb, 0.0f, cuS2Len);
PipeBarrier<PIPE_V>();
inQueue_.FreeTensor(mmInUb);
LocalTensor<int32_t> sortIndiceUbInt = sortIndiceUb.template ReinterpretCast<int32_t>();
// 无效数据索引填充为-1
if (cuS2LenVecAlign != cuS2Len) {
Duplicate(sortIndiceUbInt, -1, cuS2LenVecAlign);
PipeBarrier<PIPE_V>();
}
Adds(sortIndiceUbInt, globalTopkIndice_, static_cast<int32_t>(cuBaseS2Idx), cuS2Len);
PipeBarrier<PIPE_V>();
LocalTensor<float> tmpSortBuf = sortBuff[2 * cuS2LenVecAlign];
QLIServiceVec::SortAll(sortBuff, tmpSortBuf, cuS2LenVecAlign);
PipeBarrier<PIPE_V>();
QLIServiceVec::MergeSort(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK, sortBuff,
cuS2LenVecAlign, tmpSortBuf);
PipeBarrier<PIPE_V>();
bool isS2End = cuBaseS2Idx + s2BaseSize_ >= cuRealAcSeq;
bool needCopyOutGm = blockS2StartIdx_ == 0 && isS2End;
if (needCopyOutGm) {
LocalTensor<uint32_t> idxULocal = outQueue_.AllocTensor<uint32_t>();
ExtractIndex(idxULocal,
globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE].template ReinterpretCast<uint32_t>(),
BASE_TOPK);
PipeBarrier<PIPE_V>();
InitSortOutBuf(globalTopkUb_[innerS1Idx * BASE_TOPK_VALUE_IDX_SIZE], BASE_TOPK_VALUE_IDX_SIZE);
outQueue_.EnQue<uint32_t>(idxULocal);
idxULocal = outQueue_.DeQue<uint32_t>();
QLIServiceVec::CopyOut(indiceOutGm[info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount],
idxULocal.template ReinterpretCast<int32_t>(), constInfo_.sparseCount);
outQueue_.FreeTensor(idxULocal);
}
} else if (cuRealAcSeq <= 0) {
CleanInvalidOutput(info.indiceOutOffset + cuS1Idx * constInfo_.sparseCount);
}
}
// BNSD场景无效S1 输出-1
if (Q_LAYOUT_T == LI_LAYOUT::BSND) {
// 最后一个S1的基本块, 需要 >= info.actS1Size
bool isS1LoopEnd = (cuBaseS1Idx + s1BaseSize_) >= info.actS1Size;
int32_t invalidS1Num = constInfo_.qSeqSize - info.actS1Size;
// blockS2StartIdx_ == 0 控制S2从开始的核去做冗余清理
if (invalidS1Num > 0 && isS1LoopEnd && blockS2StartIdx_ == 0) {
int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num, 2) : (invalidS1Num / 2);
int32_t s1OffsetPerAiv = info.actS1Size + (blockId_ % 2) * CeilDiv(invalidS1Num, 2);
for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) {
CleanInvalidOutput(info.indiceOutOffset + (s1OffsetPerAiv + innerS1Idx) * constInfo_.sparseCount);
}
}
int32_t invalidS1Num2 = info.actS1Size - info.actS2SizeOrig;
if (invalidS1Num2 > 0 && isS1LoopEnd && blockS2StartIdx_ == 0 && constInfo_.attenMaskFlag) {
int32_t s1NumPerAiv = blockId_ % 2 == 0 ? CeilDiv(invalidS1Num2, 2) : (invalidS1Num2 / 2);
int32_t s1OffsetPerAiv = (blockId_ % 2) * CeilDiv(invalidS1Num2, 2);
for (int innerS1Idx = 0; innerS1Idx < s1NumPerAiv; innerS1Idx++) {
CleanInvalidOutput((info.bN2Idx * constInfo_.qSeqSize + s1OffsetPerAiv + innerS1Idx) *
constInfo_.sparseCount);
}
}
}
if (info.isLastS2InnerLoop) {
// S2最后一个Loop后, 下一个基本块初始从0开始
blockS2StartIdx_ = 0;
}
}
} // namespace QLIKernel
#endif // QUANT_LIGHTNING_INDEXER_SERVICE_VECTOR_H

193
csrc/attention/vllm_quant_lightning_indexer/op_kernel/arch32/quant_lightning_indexer_vector.h Normal file
View File

@@ -0,0 +1,193 @@
/**
 * 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 quant_lightning_indexer_vector.h
* \brief
*/
#ifndef QUANT_LIGHTNING_INDEXER_VECTOR_H
#define QUANT_LIGHTNING_INDEXER_VECTOR_H
#include "kernel_operator.h"
#include "quant_lightning_indexer_vector.h"
namespace QLIServiceVec {
using namespace AscendC;
constexpr int32_t NEG_INF = 0xFF800000;
constexpr int32_t INVALID_INDEX = -1;
constexpr uint8_t VEC_REPEAT_MAX = 255;
constexpr uint8_t B32_VEC_ELM_NUM = 64;
constexpr uint8_t B32_BLOCK_ALIGN_NUM = 8;
constexpr uint8_t B32_VEC_REPEAT_STRIDE = 8;
constexpr uint64_t VEC_REPEAT_BYTES = 256;
constexpr int32_t CONST_TWO = 2;
constexpr int64_t VALUE_AND_INDEX_NUM = 2;
constexpr int64_t BLOCK_BYTES = 32;
constexpr int64_t MRG_QUE_0 = 0;
constexpr int64_t MRG_QUE_1 = 1;
constexpr int64_t MRG_QUE_2 = 2;
constexpr int64_t MRG_QUE_3 = 3;
constexpr int64_t MRG_BLOCK_2 = 2;
constexpr int64_t MRG_BLOCK_3 = 3;
constexpr int64_t MRG_BLOCK_4 = 4;
template <typename T>
__aicore__ inline void CopyOut(const GlobalTensor<T> &dstGm, const LocalTensor<T> &srcUb, int64_t copyCount)
{
AscendC::DataCopyParams dataCopyOutyParams;
dataCopyOutyParams.blockCount = 1;
dataCopyOutyParams.blockLen = copyCount * sizeof(T);
dataCopyOutyParams.srcStride = 0;
dataCopyOutyParams.dstStride = 0;
AscendC::DataCopyPad(dstGm, srcUb, dataCopyOutyParams);
}
/**
src: 传入的初始化空间
eleNum: 需要初始化的元素个数需为64整数倍，元素将被初始化为交错排布的-inf，-1
*/
__aicore__ inline void InitSortOutBuf(const LocalTensor<float> &src, int64_t eleNum)
{
uint64_t mask1[2] = {0x5555555555555555, 0};
uint64_t mask0[2] = {0xaaaaaaaaaaaaaaaa, 0};
int64_t repeatNum = eleNum / B32_VEC_ELM_NUM;
int64_t forLoop = repeatNum / VEC_REPEAT_MAX;
int64_t forRemain = repeatNum % VEC_REPEAT_MAX;
for (int i = 0; i < forLoop; i++) {
AscendC::Duplicate(src.template ReinterpretCast<int32_t>(), NEG_INF, mask1, VEC_REPEAT_MAX, 1,
B32_VEC_REPEAT_STRIDE);
AscendC::Duplicate(src.template ReinterpretCast<int32_t>(), INVALID_INDEX, mask0, VEC_REPEAT_MAX, 1,
B32_VEC_REPEAT_STRIDE);
}
if (forRemain > 0) {
AscendC::Duplicate(src.template ReinterpretCast<int32_t>()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM], NEG_INF,
mask1, forRemain, 1, B32_VEC_REPEAT_STRIDE);
AscendC::Duplicate(src.template ReinterpretCast<int32_t>()[forLoop * VEC_REPEAT_MAX * B32_VEC_ELM_NUM],
INVALID_INDEX, mask0, forRemain, 1, B32_VEC_REPEAT_STRIDE);
}
AscendC::PipeBarrier<PIPE_V>();
}
/**
src: logits和索引，前logitsNum为logits，后logitsNum为索引
tmp: 计算使用到的临时空间，大小与src一致
logitsNum: 排序的元素个数, 暂只支持[128,256,384,512,1024,2048]
*/
__aicore__ inline void SortAll(LocalTensor<float> &src, LocalTensor<float> &tmp, int64_t logitsNum)
{
int64_t sort32Repeats = logitsNum / BLOCK_BYTES;
AscendC::Sort32(tmp, src, src[logitsNum].ReinterpretCast<uint32_t>(), sort32Repeats);
AscendC::PipeBarrier<PIPE_V>();
int64_t mrgGroups = sort32Repeats;
int64_t mrgElements = BLOCK_BYTES;
int64_t i = 0;
AscendC::LocalTensor<float> srcTensor;
AscendC::LocalTensor<float> dstTensor;
while (true) {
if (i % CONST_TWO == 0) {
srcTensor = tmp;
dstTensor = src;
} else {
srcTensor = src;
dstTensor = tmp;
}
AscendC::MrgSort4Info params;
params.elementLengths[0] = mrgElements;
params.elementLengths[MRG_QUE_1] = mrgElements;
params.elementLengths[MRG_QUE_2] = mrgElements;
params.elementLengths[MRG_QUE_3] = mrgElements;
params.ifExhaustedSuspension = false;
params.validBit = 0b1111;
AscendC::MrgSortSrcList<float> srcList;
srcList.src1 = srcTensor[0];
srcList.src2 = srcTensor[MRG_QUE_1 * VALUE_AND_INDEX_NUM * mrgElements];
srcList.src3 = srcTensor[MRG_QUE_2 * VALUE_AND_INDEX_NUM * mrgElements];
srcList.src4 = srcTensor[MRG_QUE_3 * VALUE_AND_INDEX_NUM * mrgElements];
if (mrgGroups <= MRG_BLOCK_4) {
params.repeatTimes = 1;
if (mrgGroups == 1) {
break;
} else if (mrgGroups == MRG_BLOCK_2) {
params.validBit = 0b0011;
} else if (mrgGroups == MRG_BLOCK_3) {
params.validBit = 0b0111;
} else if (mrgGroups == MRG_BLOCK_4) {
params.validBit = 0b1111;
}
AscendC::MrgSort<float>(dstTensor, srcList, params);
i += 1;
break;
} else {
params.repeatTimes = mrgGroups / MRG_BLOCK_4;
AscendC::MrgSort<float>(dstTensor, srcList, params);
i += 1;
mrgElements = mrgElements * MRG_BLOCK_4;
mrgGroups = mrgGroups / MRG_BLOCK_4;
}
AscendC::PipeBarrier<PIPE_V>();
}
if (i % CONST_TWO == 0) {
AscendC::DataCopy(src, tmp, logitsNum * VALUE_AND_INDEX_NUM);
AscendC::PipeBarrier<PIPE_V>();
}
}
/**
mrgDst: 合并进的Tensor
mrgSrc: 待合并的Tensor
tmpTensor：空间为mrgDst+mrgSrc
*/
__aicore__ inline void MergeSort(const LocalTensor<float> &mrgDst, int32_t mrgDstNum, LocalTensor<float> &mrgSrc,
int32_t mrgSrcNum, LocalTensor<float> &tmpTensor)
{
AscendC::MrgSort4Info params;
params.elementLengths[0] = mrgSrcNum;
params.elementLengths[1] = mrgDstNum;
params.ifExhaustedSuspension = false;
params.validBit = 0b0011;
params.repeatTimes = 1;
AscendC::MrgSortSrcList<float> srcList;
srcList.src1 = mrgSrc;
srcList.src2 = mrgDst;
AscendC::MrgSort<float>(tmpTensor, srcList, params);
AscendC::PipeBarrier<PIPE_V>();
AscendC::DataCopy(mrgDst, tmpTensor, mrgDstNum * VALUE_AND_INDEX_NUM);
AscendC::PipeBarrier<PIPE_V>();
}
__aicore__ inline void ExtractIndex(const LocalTensor<uint32_t> &idxULocal, const LocalTensor<uint32_t> &sortLocal,
int64_t extractNum)
{
AscendC::GatherMaskParams gatherMaskParams;
gatherMaskParams.repeatTimes = Ceil(extractNum * sizeof(float) * VALUE_AND_INDEX_NUM, VEC_REPEAT_BYTES);
gatherMaskParams.src0BlockStride = 1;
gatherMaskParams.src0RepeatStride = B32_VEC_REPEAT_STRIDE;
gatherMaskParams.src1RepeatStride = 0;
uint64_t rsvdCnt = 0; // 用于保存筛选后保留下来的元素个数
uint8_t src1Pattern = 2; // 固定模式2,表示筛选出奇数索引的数
AscendC::GatherMask(idxULocal, sortLocal, src1Pattern, false, static_cast<uint32_t>(0), gatherMaskParams, rsvdCnt);
AscendC::PipeBarrier<PIPE_V>();
}
template <HardEvent event>
__aicore__ inline void SetWaitFlag(HardEvent evt)
{
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(evt));
AscendC::SetFlag<event>(eventId);
AscendC::WaitFlag<event>(eventId);
}
} // namespace QLIServiceVec
#endif // QUANT_LIGHTNING_INDEXER_VECTOR_H