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

71
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog.cpp Normal file
View File

@@ -0,0 +1,71 @@
/**
* 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 swiglu_clip_quant.cpp
* \brief
*/
#include "indexer_compress_epilog_multi_row.h"
#include "indexer_compress_epilog_single_row.h"
#include "indexer_compress_epilog_multi_row_mx_fp8.h"
#include "indexer_compress_epilog_single_row_mx_fp8.h"
#define SINGLE_ROW_NORMAL_QUANT 10001
#define SINGLE_ROW_MXFP8_QUANT 10000
#define MULTI_ROW_NORMAL_QUANT 10011
#define MULTI_ROW_MXFP8_QUANT 10010
using namespace AscendC;
extern "C" __global__ __aicore__ void indexer_compress_epilog(
GM_ADDR indexer_compress_cache,
GM_ADDR indexer_compress_cache_scale,
GM_ADDR x,
GM_ADDR slot_mapping,
GM_ADDR indexer_compress_cache_out,
GM_ADDR indexer_compress_cache_scale_out,
GM_ADDR workspace,
GM_ADDR tiling)
{
if (workspace == nullptr) {
return;
}
GM_ADDR userWs = GetUserWorkspace(workspace);
if (userWs == nullptr) {
return;
}
GET_TILING_DATA(tilingData, tiling);
TPipe pipe;
int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();
if (TILING_KEY_IS(MULTI_ROW_NORMAL_QUANT)) {
IndexerCompressEpilog::IndexerCompressEpilogMultiRow<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE, DTYPE_INDEXER_COMPRESS_CACHE_SCALE> op;
op.Init(x, slot_mapping, indexer_compress_cache, indexer_compress_cache_scale, userWs, &tilingData, &pipe);
op.Process();
return;
} else if (TILING_KEY_IS(SINGLE_ROW_NORMAL_QUANT)) {
IndexerCompressEpilog::IndexerCompressEpilogSingleRow<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE, DTYPE_INDEXER_COMPRESS_CACHE_SCALE> op;
op.Init(x, slot_mapping, indexer_compress_cache, indexer_compress_cache_scale, userWs, &tilingData, &pipe);
op.Process();
return;
} else if (TILING_KEY_IS(MULTI_ROW_MXFP8_QUANT)){
IndexerCompressEpilog::IndexerCompressEpilogMultiRowMxFp8<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE, DTYPE_INDEXER_COMPRESS_CACHE_SCALE> op;
op.Init(x, slot_mapping, indexer_compress_cache, indexer_compress_cache_scale, userWs, &tilingData, &pipe);
op.Process();
return;
} else if (TILING_KEY_IS(SINGLE_ROW_MXFP8_QUANT)){
IndexerCompressEpilog::IndexerCompressEpilogSingleRowMxFp8<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE, DTYPE_INDEXER_COMPRESS_CACHE_SCALE> op;
op.Init(x, slot_mapping, indexer_compress_cache, indexer_compress_cache_scale, userWs, &tilingData, &pipe);
op.Process();
return;
}
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(oriOverflowMode);
}

400
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog_base.h Normal file
View File

@@ -0,0 +1,400 @@
/**
* 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 swiglu_block_quant_base.h
* \brief
*/
#ifndef SWIGLU_BLOCK_QUANT_BASE_H
#define SWIGLU_BLOCK_QUANT_BASE_H
#include "kernel_operator.h"
namespace IndexerCompressEpilog {
using namespace AscendC;
using namespace AscendC::MicroAPI;
using AscendC::MicroAPI::MaskReg;
using AscendC::MicroAPI::RegTensor;
using AscendC::MicroAPI::UnalignReg;
constexpr int32_t BLOCK_SIZE = 32;
constexpr int32_t VL_FP32 = 64;
constexpr int32_t PER_BLOCK_FP16 = 128;
constexpr float FP8_E5M2_MAX_VALUE = 57344.0f;
constexpr float FP8_E4M3FN_MAX_VALUE = 448.0f;
constexpr float FP8_E5M2_MIN_VALUE = -57344.0f;
constexpr float FP8_E4M3FN_MIN_VALUE = -448.0f;
constexpr uint32_t FAST_LOG_SHIFT_BITS = 23U;
constexpr uint32_t FAST_LOG_AND_VALUE1 = 0xFF;
constexpr uint32_t FAST_LOG_AND_VALUE2 = (((uint32_t)1 << (uint32_t)23) - (uint32_t)1);
constexpr uint32_t INV_FP8_E5M2_MAX_VALUE = 0x37924925;
constexpr uint32_t INV_FP8_E4M3_MAX_VALUE = 0x3b124925;
#define FLOAT_OVERFLOW_MODE_CTRL 60
#ifndef INFINITY
#define INFINITY (__builtin_inff())
#endif
constexpr float POS_INFINITY = INFINITY;
constexpr float NEG_INFINITY = -INFINITY;
__aicore__ inline int32_t CeilDiv(int32_t a, int b)
{
if (b == 0) {
return a;
}
return (a + b - 1) / b;
}
__aicore__ inline int32_t CeilAlign(int32_t a, int b)
{
return CeilDiv(a, b) * b;
}
template <typename T>
__aicore__ inline int32_t RoundUp(int32_t num)
{
int32_t elemNum = BLOCK_SIZE / sizeof(T);
return CeilAlign(num, elemNum);
}
template <typename T>
__aicore__ inline int32_t RoundUp(int32_t num, int32_t elemNum)
{
return CeilAlign(num, elemNum);
}
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::UNKNOWN,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::UNKNOWN,
};
constexpr AscendC::MicroAPI::CastTrait castTraitB322B16Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_RINT,
};
constexpr static AscendC::MicroAPI::CastTrait castTraitF32toFp8Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_RINT,
};
constexpr static AscendC::MicroAPI::CastTrait castTraitU32toU8Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::NO_SAT,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::CAST_NONE,
};
template <typename T>
__aicore__ inline void LoadInputData(RegTensor<float>& dst, __local_mem__ T* src, MaskReg pregLoop, uint32_t srcOffset)
{
if constexpr (IsSameType<T, float>::value) {
DataCopy(dst, src + srcOffset);
} else if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
RegTensor<T> tmp;
DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(tmp, src + srcOffset);
Cast<float, T, castTraitB162B32Even>(dst, tmp, pregLoop);
}
}
template <typename T>
__aicore__ inline void StoreOutputData(
__local_mem__ T* dst, RegTensor<float>& src, MaskReg pregLoop, uint32_t dstOffset)
{
if constexpr (IsSameType<T, float>::value) {
DataCopy(dst + dstOffset, src, pregLoop);
} else if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
RegTensor<T> tmp;
Cast<T, float, castTraitB322B16Even>(tmp, src, pregLoop);
DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + dstOffset, tmp, pregLoop);
} else if constexpr (IsSameType<T, fp8_e4m3fn_t>::value || IsSameType<T, fp8_e5m2_t>::value) {
RegTensor<T> tmp;
Cast<T, float, castTraitF32toFp8Even>(tmp, src, pregLoop);
DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(dst + dstOffset, tmp, pregLoop);
}
}
template <typename T>
__aicore__ inline void StoreMxFp8Scale(
__local_mem__ T* dst, RegTensor<float>& src, MaskReg pregLoop, uint32_t dstOffset)
{
RegTensor<uint32_t> tmp;
RegTensor<uint8_t> tmp1;
ShiftRights(tmp, (RegTensor<uint32_t> &)src, static_cast<int16_t>(FAST_LOG_SHIFT_BITS), pregLoop);
Cast<uint8_t, uint32_t, castTraitU32toU8Even>(tmp1, tmp, pregLoop);
DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B8>(dst + dstOffset, (RegTensor<T> &)tmp1, pregLoop);
}
template <typename T0, typename T1, typename T2, bool roundScale = true>
__aicore__ inline void VFProcessDynamicMxFp8Quant(
const LocalTensor<T0>& yLocal, const LocalTensor<T1>& scaleLocal, const LocalTensor<T2>& xLocal,
float coeff, float fp8Min, float fp8Max, const uint16_t curRowNum, const uint32_t curColNum)
{
__local_mem__ T0* yLocalAddr = (__local_mem__ T0*)yLocal.GetPhyAddr();
__local_mem__ T1* scaleLocalAddr = (__local_mem__ T1*)scaleLocal.GetPhyAddr();
__local_mem__ T2* xLocalAddr = (__local_mem__ T2*)xLocal.GetPhyAddr();
uint16_t loopCount = CeilDiv(curColNum, VL_FP32);
uint32_t curColNumAlign = RoundUp<T2>(curColNum);
uint32_t dstCurColNumAlign = RoundUp<T0>(curColNum);
uint16_t loopCountFoldTwo = loopCount / 2;
uint16_t loopCountReminder = loopCount % 2;
uint32_t tailReminder = curColNum - (loopCount - 1) * VL_FP32;
uint32_t scaleColNumAlign = RoundUp<T1>((curColNum + 128 - 1) / 128);
uint32_t sregNum = loopCountReminder == 0 ? curColNum - loopCountFoldTwo * VL_FP32 : loopCountFoldTwo * VL_FP32;
__VEC_SCOPE__
{
RegTensor<float> x0;
RegTensor<float> x0Abs;
RegTensor<float> x1;
RegTensor<float> x1Abs;
RegTensor<float> max0;
RegTensor<float> max1;
RegTensor<float> max2;
RegTensor<uint32_t> tmp0;
RegTensor<uint32_t> tmp1;
RegTensor<uint32_t> vreg0;
RegTensor<uint32_t> vreg1;
RegTensor<uint32_t> vreg2;
RegTensor<uint32_t> vreg3;
RegTensor<uint32_t> vreg4;
RegTensor<int32_t> vreg5;
RegTensor<uint32_t> zero;
RegTensor<uint32_t> one;
RegTensor<uint32_t> tmp3;
RegTensor<float> dupScale;
MaskReg pregLoop;
MaskReg preg1 = CreateMask<T1, AscendC::MicroAPI::MaskPattern::VL1>();
MaskReg pregMerge = CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
MaskReg pregMain = CreateMask<float>();
MaskReg cmpMask;
Duplicate(tmp0, FAST_LOG_AND_VALUE1, pregMerge);
Duplicate(tmp1, FAST_LOG_AND_VALUE2, pregMerge);
Duplicate(zero, static_cast<uint32_t>(0), pregMerge);
Duplicate(one, static_cast<uint32_t>(1), pregMerge);
Duplicate(tmp3, static_cast<uint32_t>(127), pregMerge);
for (uint16_t i = 0; i < curRowNum; i++) {
uint32_t sreg = sregNum;
for (uint16_t j = 0; j < loopCountFoldTwo; j++) {
pregLoop = UpdateMask<float>(sreg);
LoadInputData<T2>(x0, xLocalAddr, pregMain, 2 * j * VL_FP32 + i * curColNumAlign);
LoadInputData<T2>(x1, xLocalAddr, pregLoop, (2 * j + 1) * VL_FP32 + i * curColNumAlign);
Abs(x0Abs, x0, pregMain);
Abs(x1Abs, x1, pregLoop);
ReduceMax(max0, x0Abs, pregMain);
ReduceMax(max1, x1Abs, pregLoop);
Max(max2, max0, max1, pregMerge);
Maxs(max2, max2, static_cast<float>(1e-4), pregMerge);
Muls(max2, max2, coeff, pregMerge);
if constexpr (roundScale) {
ShiftRights(vreg0, (RegTensor<uint32_t> &)max2, static_cast<int16_t>(FAST_LOG_SHIFT_BITS), pregMerge);
And(vreg1, vreg0, tmp0, pregMerge);
And(vreg2, vreg1, tmp1, pregMerge);
Compare<uint32_t, AscendC::CMPMODE::NE>(cmpMask, vreg2, zero, pregMerge);
Select(vreg4, one, zero, cmpMask);
Sub(vreg1, vreg1, tmp3, pregMerge);
Add(vreg1, vreg1, vreg4, pregMerge);
Adds(vreg5, (RegTensor<int32_t> &)vreg1, static_cast<int32_t>(127), pregMerge);
ShiftLefts((RegTensor<int32_t> &)max2, vreg5, static_cast<int16_t>(23), pregMerge);
}
Duplicate(dupScale, max2, pregMain);
Div(x0, x0, dupScale, pregMain);
Div(x1, x1, dupScale, pregLoop);
Maxs(x0, x0, fp8Min, pregMain);
Mins(x0, x0, fp8Max, pregMain);
Maxs(x1, x1, fp8Min, pregLoop);
Mins(x1, x1, fp8Max, pregLoop);
StoreOutputData<T0>(yLocalAddr, x0, pregMain, 2 * j * VL_FP32 + i * dstCurColNumAlign);
StoreOutputData<T0>(yLocalAddr, x1, pregLoop, (2 * j + 1) * VL_FP32 + i * dstCurColNumAlign);
StoreMxFp8Scale<T1>(scaleLocalAddr, max2, preg1, j + i * scaleColNumAlign);
}
// 处理尾块, 这里只有一个for循环
pregLoop = UpdateMask<float>(tailReminder);
for (uint16_t j = 0; j < loopCountReminder; j++) {
LoadInputData<T2>(x0, xLocalAddr, pregLoop, 2 * loopCountFoldTwo * VL_FP32 + i * curColNumAlign);
Abs(x0Abs, x0, pregLoop);
ReduceMax(max0, x0Abs, pregLoop);
Maxs(max2, max0, static_cast<float>(1e-4), pregMerge);
Muls(max2, max2, coeff, pregMerge);
if constexpr (roundScale) {
ShiftRights(vreg0, (RegTensor<uint32_t> &)max2, static_cast<int16_t>(FAST_LOG_SHIFT_BITS), pregMerge);
And(vreg1, vreg0, tmp0, pregMerge);
And(vreg2, vreg1, tmp1, pregMerge);
Compare<uint32_t, AscendC::CMPMODE::NE>(cmpMask, vreg2, zero, pregMerge);
Select(vreg4, one, zero, cmpMask);
Sub(vreg1, vreg1, tmp3, pregMerge);
Add(vreg1, vreg1, vreg4, pregMerge);
Adds(vreg5, (RegTensor<int32_t> &)vreg1, static_cast<int32_t>(127), pregMerge);
ShiftLefts((RegTensor<int32_t> &)max2, vreg5, static_cast<int16_t>(23), pregMerge);
}
Duplicate(dupScale, max2, pregLoop);
Div(x0, x0, dupScale, pregLoop);
Maxs(x0, x0, fp8Min, pregLoop);
Mins(x0, x0, fp8Max, pregLoop);
StoreOutputData<T0>(yLocalAddr, x0, pregLoop, 2 * loopCountFoldTwo * VL_FP32 + i * dstCurColNumAlign);
StoreMxFp8Scale<T1>(scaleLocalAddr, max2, preg1, loopCountFoldTwo + i * scaleColNumAlign);
}
}
}
}
template <typename T0, typename T1>
__aicore__ inline void VFProcessDynamicBlockQuant(
const LocalTensor<T0>& yLocal, const LocalTensor<float>& scaleLocal, const LocalTensor<T1>& xLocal,
float coeff, const uint16_t curRowNum, const uint32_t curColNum)
{
__local_mem__ T0* yLocalAddr = (__local_mem__ T0*)yLocal.GetPhyAddr();
__local_mem__ float* scaleLocalAddr = (__local_mem__ float*)scaleLocal.GetPhyAddr();
__local_mem__ T1* xLocalAddr = (__local_mem__ T1*)xLocal.GetPhyAddr();
uint16_t loopCount = CeilDiv(curColNum, VL_FP32);
uint32_t curColNumAlign = RoundUp<T1>(curColNum);
uint32_t dstCurColNumAlign = RoundUp<T0>(curColNum);
uint16_t loopCountFoldTwo = loopCount / 2;
uint16_t loopCountReminder = loopCount % 2;
uint32_t tailReminder = curColNum - (loopCount - 1) * VL_FP32;
uint32_t scaleColNumAlign = RoundUp<float>((curColNum + 128 - 1) / 128);
uint32_t sregNum = loopCountReminder == 0 ? curColNum - loopCountFoldTwo * VL_FP32 : loopCountFoldTwo * VL_FP32;
static constexpr AscendC::MicroAPI::DivSpecificMode mode = {AscendC::MicroAPI::MaskMergeMode::ZEROING, false};
uint32_t maxValueInt = 0;
if constexpr (IsSameType<T0, fp8_e5m2_t>::value) {
maxValueInt = INV_FP8_E5M2_MAX_VALUE;
} else if constexpr (IsSameType<T0, fp8_e4m3fn_t>::value) {
maxValueInt = INV_FP8_E4M3_MAX_VALUE;
}
__VEC_SCOPE__
{
RegTensor<float> xLeft;
RegTensor<float> xRight;
RegTensor<float> x0Left;
RegTensor<float> x0Right;
RegTensor<float> x1Left;
RegTensor<float> x1Right;
RegTensor<float> xAbsLeft;
RegTensor<float> xAbsRight;
RegTensor<float> xMax;
RegTensor<float> tmp;
RegTensor<float> dupScale;
RegTensor<float> scale;
RegTensor<float> scale0;
RegTensor<float> scale1;
RegTensor<float> inf;
RegTensor<float> one;
RegTensor<float> zero;
RegTensor<uint32_t> coeffReg;
MaskReg pregLoop = CreateMask<float>();
Duplicate(one, static_cast<float>(1.0f), pregLoop);
Duplicate(coeffReg, maxValueInt, pregLoop);
Duplicate(zero, 0.0f);
Duplicate(inf, 1.0f);
Div<float, &mode>(inf, inf, zero, pregLoop);
MaskReg pregMain = CreateMask<float>();
MaskReg preg1 = CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
MaskReg compareLeft;
MaskReg compareRight;
MaskReg compareScalar;
for (uint16_t i = 0; i < curRowNum; i++) {
uint32_t sreg = sregNum;
for (uint16_t j = 0; j < loopCountFoldTwo; j++) {
pregLoop = UpdateMask<float>(sreg);
LoadInputData<T1>(xLeft, xLocalAddr, pregMain, 2 * j * VL_FP32 + i * curColNumAlign);
LoadInputData<T1>(xRight, xLocalAddr, pregLoop, (2 * j + 1) * VL_FP32 + i * curColNumAlign);
Muls(xAbsLeft, xLeft, 0.0f, pregMain);
Compare<float, CMPMODE::NE>(compareLeft, xAbsLeft, xAbsLeft, pregMain);
MaskNot(compareLeft, compareLeft, pregMain);
Abs(xAbsLeft, xLeft, compareLeft);
ReduceMax(scale0, xAbsLeft, pregMain);
Muls(xAbsRight, xRight, 0.0f, pregLoop);
Compare<float, CMPMODE::NE>(compareRight, xAbsRight, xAbsRight, pregLoop);
MaskNot(compareRight, compareRight, pregLoop);
Abs(xAbsRight, xRight, compareRight);
ReduceMax(scale1, xAbsRight, pregLoop);
Max(scale, scale0, scale1, preg1);
CompareScalar<float, CMPMODE::NE>(compareScalar, scale, (float)0.0, preg1);
Mul(scale, scale, (RegTensor<float>&)coeffReg, compareScalar);
Min(scale, scale, inf, preg1);
Duplicate(dupScale, scale, pregMain);
DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(scaleLocalAddr + j + i * scaleColNumAlign, scale, preg1);
Div<float, &mode>(x0Left, xLeft, dupScale, pregMain);
Muls(x1Left, x0Left, 0.0f, pregMain);
Compare<float, CMPMODE::NE>(compareLeft, x1Left, x1Left, pregMain);
Select(xLeft, xLeft, x0Left, compareLeft);
Div<float, &mode>(x0Right, xRight, dupScale, pregLoop);
Muls(x1Right, x0Right, 0.0f, pregLoop);
Compare<float, CMPMODE::NE>(compareRight, x1Right, x1Right, pregLoop);
Select(xRight, xRight, x0Right, compareRight);
StoreOutputData<T0>(yLocalAddr, xLeft, pregMain, 2 * j * VL_FP32 + i * dstCurColNumAlign);
StoreOutputData<T0>(yLocalAddr, xRight, pregLoop, (2 * j + 1) * VL_FP32 + i * dstCurColNumAlign);
}
// 处理尾块, 这里只有一个for循环
pregLoop = UpdateMask<float>(tailReminder);
for (uint16_t j = 0; j < loopCountReminder; j++) {
LoadInputData<T1>(xLeft, xLocalAddr, pregLoop, loopCountFoldTwo * 2 * VL_FP32 + i * curColNumAlign);
Abs(xAbsLeft, xLeft, pregLoop);
ReduceMax(scale, xAbsLeft, pregLoop);
CompareScalar<float, CMPMODE::NE>(compareScalar, scale, (float)0.0, preg1);
Mul(scale, scale, (RegTensor<float>&)coeffReg, compareScalar);
Min(scale, scale, inf, preg1);
Duplicate(dupScale, scale, pregLoop);
DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(scaleLocalAddr + loopCountFoldTwo + i * scaleColNumAlign, scale, preg1);
Div<float, &mode>(x0Left, xLeft, dupScale, pregLoop);
Muls(x1Left, x0Left, 0.0f, pregLoop);
Compare<float, CMPMODE::NE>(compareLeft, x1Left, x1Left, pregLoop);
Select(xLeft, xLeft, x0Left, compareLeft);
StoreOutputData(yLocalAddr, xLeft, pregLoop, loopCountFoldTwo * 2 * VL_FP32 + i * dstCurColNumAlign);
}
}
}
}
template <typename T>
__aicore__ inline void CopyIn(
const GlobalTensor<T>& inputGm, const LocalTensor<T>& inputTensor, const uint16_t nBurst, const uint32_t copyLen,
uint32_t srcStride = 0)
{
DataCopyPadExtParams<T> dataCopyPadExtParams;
dataCopyPadExtParams.isPad = false;
dataCopyPadExtParams.leftPadding = 0;
dataCopyPadExtParams.rightPadding = 0;
dataCopyPadExtParams.paddingValue = 0;
DataCopyExtParams dataCoptExtParams;
dataCoptExtParams.blockCount = nBurst;
dataCoptExtParams.blockLen = copyLen * sizeof(T);
dataCoptExtParams.srcStride = srcStride * sizeof(T);
dataCoptExtParams.dstStride = 0;
DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams);
}
template <typename T, AscendC::PaddingMode mode = AscendC::PaddingMode::Normal>
__aicore__ inline void CopyOut(
const LocalTensor<T>& outputTensor, const GlobalTensor<T>& outputGm, const uint16_t nBurst, const uint32_t copyLen,
uint32_t dstStride = 0)
{
DataCopyExtParams dataCopyParams;
dataCopyParams.blockCount = nBurst;
dataCopyParams.blockLen = copyLen * sizeof(T);
dataCopyParams.srcStride = 0;
dataCopyParams.dstStride = dstStride * sizeof(T);
DataCopyPad<T, mode>(outputGm, outputTensor, dataCopyParams);
}
} // namespace SwigluBlockQuant
#endif

149
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog_multi_row.h Normal file
View File

@@ -0,0 +1,149 @@
/**
* 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 indexer_compress_epilog_multi_row.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_MULTI_ROW_H
#define INDEXER_COMPRESS_EPILOG_MULTI_ROW_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_base.h"
namespace IndexerCompressEpilog {
using namespace AscendC;
template <typename T0, typename T1, typename T2>
class IndexerCompressEpilogMultiRow {
public:
__aicore__ inline IndexerCompressEpilogMultiRow()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR indexerCompressCacheScale,
GM_ADDR workspace, const IndexerCompressEpilogTilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
indexerCompressCacheScaleGm.SetGlobalBuffer((__gm__ float*)indexerCompressCacheScale);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<T1>(tilingData->d) * sizeof(T1));
pipe->InitBuffer(
indexerCompressCacheScaleQue, 2,
tilingData->rowFactor * RoundUp<float>(tilingData->scaleCol) * sizeof(float));
pipe->InitBuffer(indexBuf, RoundUp<int32_t>(tilingData->rowFactor) * sizeof(int32_t));
indexLocal = indexBuf.Get<int32_t>();
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
}
__aicore__ inline void Process()
{
SetMaxValue();
int64_t curBlockIdx = GetBlockIdx();
int64_t rowOuterLoop =
(curBlockIdx == GetBlockNum() - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock;
int64_t tailRowFactor = (curBlockIdx == GetBlockNum() - 1) ? tilingData->tailRowFactorOfTailBlock :
tilingData->tailRowFactorOfFormerBlock;
int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d;
for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->rowFactor;
xLocal = xQue.template AllocTensor<T0>();
validIdx = 0;
for (int64_t rowInnerIdx = 0; rowInnerIdx < curRowFactor; rowInnerIdx++) {
int64_t curSlotIdx = curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor + rowInnerIdx;
int64_t slot = slotMappingGm.GetValue(curSlotIdx);
if (slot == -1) {
continue;
}
CopyIn(
xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->d +
rowInnerIdx * tilingData->d],
xLocal[validIdx * RoundUp<T0>(tilingData->d)], 1, tilingData->d);
indexLocal.SetValue(validIdx, slot);
validIdx++;
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
SetFlag<HardEvent::S_MTE3>(eventId);
WaitFlag<HardEvent::S_MTE3>(eventId);
}
xQue.template EnQue(xLocal);
xLocal = xQue.template DeQue<T0>();
indexerCompressCacheLocal = indexerCompressCacheQue.template AllocTensor<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.AllocTensor<float>();
VFProcessDynamicBlockQuant(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, validIdx, tilingData->d);
xQue.template FreeTensor(xLocal);
indexerCompressCacheQue.template EnQue(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template EnQue(indexerCompressCacheScaleLocal);
indexerCompressCacheLocal = indexerCompressCacheQue.template DeQue<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<float>();
for (int64_t curValidIdx = 0; curValidIdx < validIdx; curValidIdx++) {
int64_t curSlotIdx = indexLocal.GetValue(curValidIdx);
CopyOut(
indexerCompressCacheLocal[curValidIdx * RoundUp<T1>(tilingData->d)],
indexerCompressCacheGm[curSlotIdx * tilingData->d], 1, tilingData->d);
CopyOut(
indexerCompressCacheScaleLocal[curValidIdx * RoundUp<float>(tilingData->scaleCol)],
indexerCompressCacheScaleGm[curSlotIdx * CeilDiv(tilingData->d, PER_BLOCK_FP16)], 1,
tilingData->scaleCol);
}
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
if constexpr (IsSameType<T1, fp8_e5m2_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E5M2_MAX_VALUE;
fp8Max = FP8_E5M2_MAX_VALUE;
fp8Min = FP8_E5M2_MIN_VALUE;
} else if constexpr (IsSameType<T1, fp8_e4m3fn_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
fp8Max = FP8_E4M3FN_MAX_VALUE;
fp8Min = FP8_E4M3FN_MIN_VALUE;
}
}
private:
TPipe* pipe;
const IndexerCompressEpilogTilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
GlobalTensor<float> indexerCompressCacheScaleGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
TBuf<QuePosition::VECCALC> indexBuf;
LocalTensor<T0> xLocal;
LocalTensor<T1> indexerCompressCacheLocal;
LocalTensor<float> indexerCompressCacheScaleLocal;
LocalTensor<int32_t> indexLocal;
int64_t validIdx = 0;
float maxValue = 0.0f;
float fp8Min = 0.0f;
float fp8Max = 0.0f;
};
} // namespace IndexerCompressEpilog
#endif

156
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog_multi_row_mx_fp8.h Normal file
View File

@@ -0,0 +1,156 @@
/**
* 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 indexer_compress_epilog_multi_row.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_MULTI_ROW_MX_FP8_H
#define INDEXER_COMPRESS_EPILOG_MULTI_ROW_MX_FP8_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_base.h"
namespace IndexerCompressEpilog {
using namespace AscendC;
template <typename T0, typename T1, typename T2>
class IndexerCompressEpilogMultiRowMxFp8 {
public:
__aicore__ inline IndexerCompressEpilogMultiRowMxFp8()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR indexerCompressCacheScale,
GM_ADDR workspace, const IndexerCompressEpilogTilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
indexerCompressCacheScaleGm.SetGlobalBuffer((__gm__ T2*)indexerCompressCacheScale);
roundScale = (tilingData->roundScale == 1);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<T1>(tilingData->d) * sizeof(T1));
pipe->InitBuffer(
indexerCompressCacheScaleQue, 2,
tilingData->rowFactor * RoundUp<T2>(tilingData->scaleCol) * sizeof(T2));
pipe->InitBuffer(indexBuf, RoundUp<int32_t>(tilingData->rowFactor) * sizeof(int32_t));
indexLocal = indexBuf.Get<int32_t>();
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
}
__aicore__ inline void Process()
{
SetMaxValue();
int64_t curBlockIdx = GetBlockIdx();
int64_t rowOuterLoop =
(curBlockIdx == GetBlockNum() - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock;
int64_t tailRowFactor = (curBlockIdx == GetBlockNum() - 1) ? tilingData->tailRowFactorOfTailBlock :
tilingData->tailRowFactorOfFormerBlock;
int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d;
for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
int64_t curRowFactor = (rowOuterIdx == rowOuterLoop - 1) ? tailRowFactor : tilingData->rowFactor;
xLocal = xQue.template AllocTensor<T0>();
validIdx = 0;
for (int64_t rowInnerIdx = 0; rowInnerIdx < curRowFactor; rowInnerIdx++) {
int64_t curSlotIdx = curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor + rowInnerIdx;
int64_t slot = slotMappingGm.GetValue(curSlotIdx);
if (slot == -1) {
continue;
}
CopyIn(
xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->d +
rowInnerIdx * tilingData->d],
xLocal[validIdx * RoundUp<T0>(tilingData->d)], 1, tilingData->d);
indexLocal.SetValue(validIdx, slot);
validIdx++;
event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
SetFlag<HardEvent::S_MTE3>(eventId);
WaitFlag<HardEvent::S_MTE3>(eventId);
}
xQue.template EnQue(xLocal);
xLocal = xQue.template DeQue<T0>();
indexerCompressCacheLocal = indexerCompressCacheQue.template AllocTensor<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.AllocTensor<T2>();
if (roundScale) {
VFProcessDynamicMxFp8Quant<T1, T2, T0, true>(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, fp8Min, fp8Max, validIdx, tilingData->d);
} else {
VFProcessDynamicMxFp8Quant<T1, T2, T0, false>(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, fp8Min, fp8Max, validIdx, tilingData->d);
}
xQue.template FreeTensor(xLocal);
indexerCompressCacheQue.template EnQue(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template EnQue(indexerCompressCacheScaleLocal);
indexerCompressCacheLocal = indexerCompressCacheQue.template DeQue<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<T2>();
for (int64_t curValidIdx = 0; curValidIdx < validIdx; curValidIdx++) {
int64_t curSlotIdx = indexLocal.GetValue(curValidIdx);
CopyOut(
indexerCompressCacheLocal[curValidIdx * RoundUp<T1>(tilingData->d)],
indexerCompressCacheGm[curSlotIdx * tilingData->d], 1, tilingData->d);
CopyOut(
indexerCompressCacheScaleLocal[curValidIdx * RoundUp<T2>(tilingData->scaleCol)],
indexerCompressCacheScaleGm[curSlotIdx * CeilDiv(tilingData->d, PER_BLOCK_FP16)], 1,
tilingData->scaleCol);
}
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
if constexpr (IsSameType<T1, fp8_e5m2_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E5M2_MAX_VALUE;
fp8Max = FP8_E5M2_MAX_VALUE;
fp8Min = FP8_E5M2_MIN_VALUE;
} else if constexpr (IsSameType<T1, fp8_e4m3fn_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
fp8Max = FP8_E4M3FN_MAX_VALUE;
fp8Min = FP8_E4M3FN_MIN_VALUE;
}
}
private:
TPipe* pipe;
const IndexerCompressEpilogTilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
GlobalTensor<T2> indexerCompressCacheScaleGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
TBuf<QuePosition::VECCALC> indexBuf;
LocalTensor<T0> xLocal;
LocalTensor<T1> indexerCompressCacheLocal;
LocalTensor<T2> indexerCompressCacheScaleLocal;
LocalTensor<int32_t> indexLocal;
int64_t validIdx = 0;
float maxValue = 0.0f;
float fp8Min = 0.0f;
float fp8Max = 0.0f;
bool roundScale = true;
};
} // namespace IndexerCompressEpilog
#endif

126
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog_single_row.h Normal file
View File

@@ -0,0 +1,126 @@
/**
* 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 indexer_compress_epilog_d_full_load.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_SINGLE_ROW_H
#define INDEXER_COMPRESS_EPILOG_SINGLE_ROW_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_base.h"
namespace IndexerCompressEpilog {
using namespace AscendC;
template <typename T0, typename T1, typename T2>
class IndexerCompressEpilogSingleRow {
public:
__aicore__ inline IndexerCompressEpilogSingleRow()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR indexerCompressCacheScale,
GM_ADDR workspace, const IndexerCompressEpilogTilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
indexerCompressCacheScaleGm.SetGlobalBuffer((__gm__ float*)indexerCompressCacheScale);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<T1>(tilingData->d) * sizeof(T1));
pipe->InitBuffer(
indexerCompressCacheScaleQue, 2,
tilingData->rowFactor * RoundUp<float>(tilingData->scaleCol) * sizeof(float));
pipe->InitBuffer(indexBuf, RoundUp<int32_t>(tilingData->rowFactor) * sizeof(int32_t));
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
}
__aicore__ inline void Process()
{
SetMaxValue();
int64_t curBlockIdx = GetBlockIdx();
int64_t rowOuterLoop =
(curBlockIdx == GetBlockNum() - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock;
int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d;
for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
xLocal = xQue.template AllocTensor<T0>();
int64_t curSlotIdx = curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor;
int64_t slot = slotMappingGm.GetValue(curSlotIdx);
if (slot == -1) {
continue;
}
CopyIn(xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->d],
xLocal, 1, tilingData->d);
xQue.template EnQue(xLocal);
xLocal = xQue.template DeQue<T0>();
indexerCompressCacheLocal = indexerCompressCacheQue.template AllocTensor<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.AllocTensor<float>();
VFProcessDynamicBlockQuant(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, 1, tilingData->d);
xQue.template FreeTensor(xLocal);
indexerCompressCacheQue.template EnQue(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template EnQue(indexerCompressCacheScaleLocal);
indexerCompressCacheLocal = indexerCompressCacheQue.template DeQue<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<float>();
CopyOut(
indexerCompressCacheLocal, indexerCompressCacheGm[slot * tilingData->d], 1, tilingData->d);
CopyOut(
indexerCompressCacheScaleLocal, indexerCompressCacheScaleGm[slot * CeilDiv(tilingData->d, PER_BLOCK_FP16)], 1,
tilingData->scaleCol);
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
if constexpr (IsSameType<T1, fp8_e5m2_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E5M2_MAX_VALUE;
} else if constexpr (IsSameType<T1, fp8_e4m3fn_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
}
}
private:
TPipe* pipe;
const IndexerCompressEpilogTilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
GlobalTensor<float> indexerCompressCacheScaleGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
TBuf<QuePosition::VECCALC> indexBuf;
LocalTensor<T0> xLocal;
LocalTensor<T1> indexerCompressCacheLocal;
LocalTensor<float> indexerCompressCacheScaleLocal;
float maxValue = 0.0f;
};
} // namespace IndexerCompressEpilog
#endif

140
csrc/attention/indexer_compress_epilog/op_kernel/indexer_compress_epilog_single_row_mx_fp8.h Normal file
View File

@@ -0,0 +1,140 @@
/**
* 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 indexer_compress_epilog_d_full_load.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_SINGLE_ROW_MX_FP8_H
#define INDEXER_COMPRESS_EPILOG_SINGLE_ROW_MX_FP8_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_base.h"
namespace IndexerCompressEpilog {
using namespace AscendC;
template <typename T0, typename T1, typename T2>
class IndexerCompressEpilogSingleRowMxFp8 {
public:
__aicore__ inline IndexerCompressEpilogSingleRowMxFp8()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR indexerCompressCacheScale,
GM_ADDR workspace, const IndexerCompressEpilogTilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
indexerCompressCacheScaleGm.SetGlobalBuffer((__gm__ T2*)indexerCompressCacheScale);
roundScale = (tilingData->roundScale == 1);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<T1>(tilingData->d) * sizeof(T1));
pipe->InitBuffer(
indexerCompressCacheScaleQue, 2,
tilingData->rowFactor * RoundUp<T2>(tilingData->scaleCol) * sizeof(T2));
pipe->InitBuffer(indexBuf, RoundUp<int32_t>(tilingData->rowFactor) * sizeof(int32_t));
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
}
__aicore__ inline void Process()
{
SetMaxValue();
int64_t curBlockIdx = GetBlockIdx();
int64_t rowOuterLoop =
(curBlockIdx == GetBlockNum() - 1) ? tilingData->rowLoopOfTailBlock : tilingData->rowLoopOfFormerBlock;
int64_t xGmBaseOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->d;
for (int64_t rowOuterIdx = 0; rowOuterIdx < rowOuterLoop; rowOuterIdx++) {
xLocal = xQue.template AllocTensor<T0>();
int64_t curSlotIdx = curBlockIdx * tilingData->rowOfFormerBlock + rowOuterIdx * tilingData->rowFactor;
int64_t slot = slotMappingGm.GetValue(curSlotIdx);
if (slot == -1) {
continue;
}
CopyIn(xGm[xGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->d],
xLocal, 1, tilingData->d);
xQue.template EnQue(xLocal);
xLocal = xQue.template DeQue<T0>();
indexerCompressCacheLocal = indexerCompressCacheQue.template AllocTensor<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.AllocTensor<T2>();
if (roundScale) {
VFProcessDynamicMxFp8Quant<T1, T2, T0, true>(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, fp8Min, fp8Max, 1, tilingData->d);
} else {
VFProcessDynamicMxFp8Quant<T1, T2, T0, false>(
indexerCompressCacheLocal, indexerCompressCacheScaleLocal, xLocal, maxValue, fp8Min, fp8Max, 1, tilingData->d);
}
xQue.template FreeTensor(xLocal);
indexerCompressCacheQue.template EnQue(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template EnQue(indexerCompressCacheScaleLocal);
indexerCompressCacheLocal = indexerCompressCacheQue.template DeQue<T1>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<T2>();
CopyOut(
indexerCompressCacheLocal, indexerCompressCacheGm[slot * tilingData->d], 1, tilingData->d);
CopyOut(
indexerCompressCacheScaleLocal, indexerCompressCacheScaleGm[slot * CeilDiv(tilingData->d, PER_BLOCK_FP16)], 1,
tilingData->scaleCol);
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
if constexpr (IsSameType<T1, fp8_e5m2_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E5M2_MAX_VALUE;
fp8Max = FP8_E5M2_MAX_VALUE;
fp8Min = FP8_E5M2_MIN_VALUE;
} else if constexpr (IsSameType<T1, fp8_e4m3fn_t>::value) {
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
fp8Max = FP8_E4M3FN_MAX_VALUE;
fp8Min = FP8_E4M3FN_MIN_VALUE;
}
}
private:
TPipe* pipe;
const IndexerCompressEpilogTilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
GlobalTensor<T2> indexerCompressCacheScaleGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
TBuf<QuePosition::VECCALC> indexBuf;
LocalTensor<T0> xLocal;
LocalTensor<T1> indexerCompressCacheLocal;
LocalTensor<T2> indexerCompressCacheScaleLocal;
float maxValue = 0.0f;
float fp8Min = 0.0f;
float fp8Max = 0.0f;
bool roundScale = true;
};
} // namespace IndexerCompressEpilog
#endif