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

19
csrc/attention/indexer_compress_epilog_v2/CMakeLists.txt Normal file
View File

@@ -0,0 +1,19 @@
# -----------------------------------------------------------------------------------------------------------
# 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(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
if(NOT ENABLE_TEST AND NOT BENCHMARK)
list(REMOVE_ITEM CURRENT_DIRS tests)
endif()
foreach(SUB_DIR ${CURRENT_DIRS})
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
add_subdirectory(${SUB_DIR})
endif()
endforeach()

63
csrc/attention/indexer_compress_epilog_v2/op_host/CMakeLists.txt Normal file
View File

@@ -0,0 +1,63 @@
# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
# add_ops_compile_options(
# OP_NAME IndexerCompressEpilogV2
# OPTIONS --cce-auto-sync=off
# -Wno-deprecated-declarations
# -Werror
# -mllvm -cce-aicore-hoist-movemask=false
# --op_relocatable_kernel_binary=true
# )
# set(indexer_compress_epilog_v2_depends transformer/attention/indexer_compress_epilog_v2 PARENT_SCOPE)
# target_sources(op_host_aclnn PRIVATE
# op_host/indexer_compress_epilog_v2_def.cpp
# )
# target_sources(optiling PRIVATE
# op_host/indexer_compress_epilog_v2_tiling.cpp
# )
# if (NOT BUILD_OPEN_PROJECT)
# target_sources(opmaster_ct PRIVATE
# op_host/indexer_compress_epilog_v2_tiling.cpp
# )
# endif ()
# target_include_directories(optiling PRIVATE
# ${CMAKE_CURRENT_SOURCE_DIR}/op_host
# )
# target_sources(opsproto PRIVATE
# op_host/indexer_compress_epilog_v2_proto.cpp
# )
add_op_to_compiled_list()
if (BUILD_OPEN_PROJECT)
target_sources(op_host_aclnn PRIVATE
indexer_compress_epilog_v2_def.cpp
)
endif()
add_ops_compile_options(
OP_NAME IndexerCompressEpilogV2
OPTIONS --cce-auto-sync=off
-Wno-deprecated-declarations
-mllvm -cce-aicore-hoist-movemask=false
--op_relocatable_kernel_binary=true
)
if (NOT BUILD_OPS_RTY_KERNEL)
add_modules_sources(OPTYPE indexer_compress_epilog_v2 ACLNNTYPE aclnn)
endif()

52
csrc/attention/indexer_compress_epilog_v2/op_host/indexer_compress_epilog_v2_def.cpp Normal file
View File

@@ -0,0 +1,52 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file indexer_compress_epilog_v2_def.cpp
* \brief
*/
#include "register/op_def_registry.h"
namespace ops {
class IndexerCompressEpilogV2 : public OpDef {
public:
explicit IndexerCompressEpilogV2(const char* name) : OpDef(name)
{
this->Input("indexer_compress_cache")
.ParamType(REQUIRED)
.DataType({ge::DT_UINT8, ge::DT_UINT8})
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
.IgnoreContiguous();
this->Input("x")
.ParamType(REQUIRED)
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
this->Input("slot_mapping")
.ParamType(REQUIRED)
.DataType({ge::DT_INT32, ge::DT_INT32})
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
this->Output("indexer_compress_cache")
.ParamType(REQUIRED)
.DataType({ge::DT_UINT8, ge::DT_UINT8})
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
.IgnoreContiguous();
this->Attr("layout").AttrType(OPTIONAL).Int(2);
this->Attr("block_stride").AttrType(OPTIONAL).Int(0);
this->AICore().AddConfig("ascend950");
}
};
OP_ADD(IndexerCompressEpilogV2);
} // namespace ops

36
csrc/attention/indexer_compress_epilog_v2/op_host/indexer_compress_epilog_v2_proto.cpp Normal file
View File

@@ -0,0 +1,36 @@
/**
* 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_v2_proto.cpp
* \brief
*/
#include <graph/utils/type_utils.h>
#include <register/op_impl_registry.h>
#include "error/ops_error.h"
using namespace ge;
namespace ops {
graphStatus InferShape4IndexerCompressEpilogV2(gert::InferShapeContext* context)
{
return ge::GRAPH_SUCCESS;
}
graphStatus InferDtype4IndexerCompressEpilogV2(gert::InferDataTypeContext* context)
{
return GRAPH_SUCCESS;
}
IMPL_OP_INFERSHAPE(IndexerCompressEpilogV2)
.InferShape(InferShape4IndexerCompressEpilogV2)
.InferDataType(InferDtype4IndexerCompressEpilogV2);
} // namespace ops

273
csrc/attention/indexer_compress_epilog_v2/op_host/indexer_compress_epilog_v2_tiling.cpp Normal file
View File

@@ -0,0 +1,273 @@
/**
* 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_v2_tiling.cpp
* \brief
*/
#include <sstream>
#include "indexer_compress_epilog_v2_tiling.h"
using namespace ge;
namespace optiling {
namespace {
constexpr uint64_t WORKSPACE_SIZE = 32;
int64_t CeilDiv(int64_t x, int64_t y)
{
if (y != 0) {
return (x + y - 1) / y;
}
return x;
}
int64_t DownAlign(int64_t x, int64_t y) {
if (y == 0) {
return x;
}
return (x / y) * y;
}
int64_t RoundUp(int64_t x, int64_t y) {
return CeilDiv(x, y) * y;
}
constexpr int64_t INPUT_CACHE_IDX = 0;
constexpr int64_t INPUT_X_IDX = 1;
constexpr int64_t INPUT_SLOT_MAPPING_IDX = 2;
constexpr int64_t ATTR_LAYOUT_INDEX = 0;
constexpr int64_t ATTR_BLOCK_STRIDE_INDEX = 1;
constexpr int64_t BLOCK_SIZE = 32;
constexpr int64_t REPEAT_SIZE = 256;
constexpr int64_t DOUBLE_BUFFER = 2;
// per_block量化,每128个f16需要量化出一个scale, 因此切分尾轴时，以128为factor进行切分
constexpr int64_t PER_BLOCK_FP16 = 128;
constexpr int64_t SINGLE_ROW = 1;
constexpr int64_t DIM_0 = 0;
constexpr int64_t DIM_1 = 1;
constexpr int64_t DIM_2 = 2;
constexpr int64_t DIM_3 = 3;
constexpr int64_t SINGLE_ROW_TILING_KEY = 20020;
constexpr int64_t MULTI_ROW_TILING_KEY = 20021;
constexpr int64_t INPUTS_DIM_LIMIT = 4;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::GetPlatformInfo()
{
auto platformInfo = context_->GetPlatformInfo();
if (platformInfo == nullptr) {
auto compileInfoPtr = context_->GetCompileInfo<IndexerCompressEpilogV2CompileInfo>();
OPS_ERR_IF(compileInfoPtr == nullptr, OPS_LOG_E(context_, "compile info is null"),
return ge::GRAPH_FAILED);
coreNum_ = compileInfoPtr->coreNum;
ubSize_ = compileInfoPtr->ubSize;
} else {
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
coreNum_ = ascendcPlatform.GetCoreNumAiv();
uint64_t ubSizePlatForm;
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
ubSize_ = ubSizePlatForm;
socVersion_ = ascendcPlatform.GetSocVersion();
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::GetAttr()
{
auto* attrs = context_->GetAttrs();
OPS_LOG_E_IF_NULL(context_, attrs, return ge::GRAPH_FAILED);
auto layout = attrs->GetAttrPointer<int64_t>(ATTR_LAYOUT_INDEX);
layout_ = layout == nullptr ? 2 : *layout;
auto blockStride = attrs->GetAttrPointer<int64_t>(ATTR_BLOCK_STRIDE_INDEX);
blockStride_ = blockStride == nullptr ? 0 : *blockStride;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::GetShapeAttrsInfoInner()
{
auto shapeCache = context_->GetInputShape(INPUT_CACHE_IDX);
OPS_LOG_E_IF_NULL(context_, shapeCache, return ge::GRAPH_FAILED);
auto cacheStorageShape = shapeCache->GetStorageShape();
OPS_ERR_IF(
(cacheStorageShape.GetDimNum() != INPUTS_DIM_LIMIT),
OPS_LOG_E(context_, "the dim of indexer_compress_cache only support %d, please check.", INPUTS_DIM_LIMIT),
return ge::GRAPH_FAILED);
cacheBs_ = cacheStorageShape.GetDim(DIM_1);
OPS_ERR_IF(
(cacheStorageShape.GetDim(DIM_2) != 1),
OPS_LOG_E(context_, "the third dim of indexer_compress_cache must be 1, please check."),
return ge::GRAPH_FAILED);
cacheD_ = cacheStorageShape.GetDim(DIM_3);
auto shapeX = context_->GetInputShape(INPUT_X_IDX);
OPS_LOG_E_IF_NULL(context_, shapeX, return ge::GRAPH_FAILED);
auto xStorageShape = shapeX->GetStorageShape();
d_ = xStorageShape.GetDim(xStorageShape.GetDimNum() - 1);
scaleCol_ = CeilDiv(d_, PER_BLOCK_FP16);
auto shapeSlotMapping = context_->GetInputShape(INPUT_SLOT_MAPPING_IDX);
OPS_LOG_E_IF_NULL(context_, shapeSlotMapping, return ge::GRAPH_FAILED);
auto slotMappingStorageShape = shapeSlotMapping->GetStorageShape();
tnd_ = 1;
for (int i = 0; i < slotMappingStorageShape.GetDimNum(); i++) {
tnd_ = tnd_ * slotMappingStorageShape.GetDim(i);
}
OPS_ERR_IF(GetAttr() != ge::GRAPH_SUCCESS,
OPS_LOG_E(context_->GetNodeName(), "get attr failed."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::CheckShapesAndAttrs()
{
int64_t minBlockSize = cacheBs_ * cacheD_;
OPS_ERR_IF((blockStride_ != 0 && blockStride_ < minBlockSize),
OPS_LOG_E(context_, "block_stride must be greater than min block size %ld.", minBlockSize),
return ge::GRAPH_FAILED);
OPS_ERR_IF((layout_ != 2),
OPS_LOG_E(context_, "layout only support 2, please check."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::CalcOpTiling()
{
rowOfFormerBlock_ = CeilDiv(tnd_, static_cast<int64_t>(coreNum_));
usedCoreNums_ = std::min(CeilDiv(tnd_, rowOfFormerBlock_), static_cast<int64_t>(coreNum_));
rowOfTailBlock_ = tnd_ - (usedCoreNums_ - 1) * rowOfFormerBlock_;
int64_t minRowPerCore = 1;
int64_t rowOnceLoop = std::min(rowOfFormerBlock_, minRowPerCore);
rowFactor_ = rowOnceLoop;
int64_t scaleByteSize = 4;
int64_t perBlockScaleElemNum = BLOCK_SIZE / scaleByteSize;
// d全载,尝试搬入更多的bs
while (rowFactor_ <= rowOfFormerBlock_) {
int64_t xSize = rowFactor_ * RoundUp(d_, 16) * 2 * DOUBLE_BUFFER;
int64_t ySize = rowFactor_ * RoundUp(d_, 32) * 1 * DOUBLE_BUFFER;
int64_t scaleSize = rowFactor_ * RoundUp(scaleCol_, perBlockScaleElemNum) * scaleByteSize * DOUBLE_BUFFER;
int64_t tmpBufferSize = RoundUp(rowFactor_, 8) * 4;
int64_t totalSize = xSize + ySize + scaleSize + tmpBufferSize;
if (totalSize > ubSize_) {
rowFactor_ = rowFactor_ - 1;
break;
}
rowFactor_ = rowFactor_ + 1;
}
if (rowFactor_ > rowOfFormerBlock_) {
rowFactor_--;
}
rowLoopOfFormerBlock_ = CeilDiv(rowOfFormerBlock_, rowFactor_);
rowLoopOfTailBlock_ = CeilDiv(rowOfTailBlock_, rowFactor_);
tailRowFactorOfFormerBlock_ = rowOfFormerBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfFormerBlock_ % rowFactor_;
tailRowFactorOfTailBlock_ = rowOfTailBlock_ % rowFactor_ == 0 ? rowFactor_ : rowOfTailBlock_ % rowFactor_;
// 如果未指定blockStride，设置默认值为单个block大小（连续排列）
if (blockStride_ == 0) {
blockStride_ = cacheBs_ * cacheD_;
}
tilingData_.set_d(d_);
tilingData_.set_cacheBs(cacheBs_);
tilingData_.set_scaleCol(scaleCol_);
tilingData_.set_rowOfFormerBlock(rowOfFormerBlock_);
tilingData_.set_rowOfTailBlock(rowOfTailBlock_);
tilingData_.set_rowLoopOfFormerBlock(rowLoopOfFormerBlock_);
tilingData_.set_rowLoopOfTailBlock(rowLoopOfTailBlock_);
tilingData_.set_rowFactor(rowFactor_);
tilingData_.set_tailRowFactorOfFormerBlock(tailRowFactorOfFormerBlock_);
tilingData_.set_tailRowFactorOfTailBlock(tailRowFactorOfTailBlock_);
tilingData_.set_blockStride(blockStride_);
// SINGLE_ROW TILING_KEY : 20020
// MULTI_ROW TILING_KEY : 20021
if (rowFactor_ == SINGLE_ROW) {
tilingKey_ = SINGLE_ROW_TILING_KEY;
} else {
tilingKey_ = MULTI_ROW_TILING_KEY;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::DoOpTiling()
{
if (GetPlatformInfo() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
if (GetShapeAttrsInfoInner() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
if (CheckShapesAndAttrs() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
if (CalcOpTiling() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
if (GetWorkspaceSize() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
if (PostTiling() == ge::GRAPH_FAILED) {
return ge::GRAPH_FAILED;
}
context_->SetTilingKey(tilingKey_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::GetWorkspaceSize()
{
workspaceSize_ = WORKSPACE_SIZE;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus IndexerCompressEpilogV2Tiling::PostTiling()
{
context_->SetBlockDim(usedCoreNums_);
size_t* workspaces = context_->GetWorkspaceSizes(1);
workspaces[0] = workspaceSize_;
tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
return ge::GRAPH_SUCCESS;
}
ge::graphStatus TilingPrepareForIndexerCompressEpilogV2(gert::TilingParseContext *context)
{
(void)context;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus TilingForIndexerCompressEpilogV2(gert::TilingContext *context)
{
OPS_ERR_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("IndexerCompressEpilogV2", "Tiling context is null"),
return ge::GRAPH_FAILED);
IndexerCompressEpilogV2Tiling IndexerCompressEpilogV2Tiling(context);
return IndexerCompressEpilogV2Tiling.DoOpTiling();
}
IMPL_OP_OPTILING(IndexerCompressEpilogV2)
.Tiling(TilingForIndexerCompressEpilogV2)
.TilingParse<IndexerCompressEpilogV2CompileInfo>(TilingPrepareForIndexerCompressEpilogV2);
} // namespace optiling

109
csrc/attention/indexer_compress_epilog_v2/op_host/indexer_compress_epilog_v2_tiling.h Normal file
View File

@@ -0,0 +1,109 @@
/**
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
/*!
* \file indexer_compress_epilog_v2_tiling.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_V2_TILING_H
#define INDEXER_COMPRESS_EPILOG_V2_TILING_H
#include <vector>
#include <iostream>
#include "register/op_impl_registry.h"
#include "platform/platform_infos_def.h"
#include "exe_graph/runtime/tiling_context.h"
#include "tiling/platform/platform_ascendc.h"
#include "register/op_def_registry.h"
#include "register/tilingdata_base.h"
#include "tiling/tiling_api.h"
#include "error/ops_error.h"
#include "platform/platform_info.h"
namespace optiling {
// ----------公共定义----------
struct TilingRequiredParaInfo {
const gert::CompileTimeTensorDesc *desc;
const gert::StorageShape *shape;
};
struct TilingOptionalParaInfo {
const gert::CompileTimeTensorDesc *desc;
const gert::Tensor *tensor;
};
// ----------算子TilingData定义----------
BEGIN_TILING_DATA_DEF(IndexerCompressEpilogV2TilingData)
TILING_DATA_FIELD_DEF(int64_t, d); // head_dim of x
TILING_DATA_FIELD_DEF(int64_t, cacheBs); // second dim of cache (cacheBn, cacheBs, 1, cacheD)
TILING_DATA_FIELD_DEF(int64_t, scaleCol); // 一行多少个scale
TILING_DATA_FIELD_DEF(int64_t, rowOfFormerBlock); // 头核共需要处理多少行
TILING_DATA_FIELD_DEF(int64_t, rowOfTailBlock); // 尾核共需要处理多少行
TILING_DATA_FIELD_DEF(int64_t, rowLoopOfFormerBlock); // 头核需要几次ub搬入
TILING_DATA_FIELD_DEF(int64_t, rowLoopOfTailBlock); // 尾核需要几次ub搬入
TILING_DATA_FIELD_DEF(int64_t, rowFactor); // ub一次标准处理行数
TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfFormerBlock); // 头核最后一次ub处理行数
TILING_DATA_FIELD_DEF(int64_t, tailRowFactorOfTailBlock); // 尾核最后一次ub处理行数
TILING_DATA_FIELD_DEF(int64_t, blockStride);
END_TILING_DATA_DEF;
REGISTER_TILING_DATA_CLASS(IndexerCompressEpilogV2, IndexerCompressEpilogV2TilingData)
// ----------算子CompileInfo定义----------
struct IndexerCompressEpilogV2CompileInfo {
uint64_t coreNum = 0;
uint64_t ubSize = 0;
};
// ----------算子Tiling入参信息解析及check类----------
class IndexerCompressEpilogV2Tiling {
public:
explicit IndexerCompressEpilogV2Tiling(gert::TilingContext* tilingContext) : context_(tilingContext)
{
}
~IndexerCompressEpilogV2Tiling() = default;
ge::graphStatus GetPlatformInfo();
ge::graphStatus DoOpTiling();
ge::graphStatus GetWorkspaceSize();
ge::graphStatus PostTiling();
ge::graphStatus GetAttr();
ge::graphStatus GetShapeAttrsInfoInner();
ge::graphStatus CheckShapesAndAttrs();
ge::graphStatus CalcOpTiling();
private:
gert::TilingContext *context_ = nullptr;
IndexerCompressEpilogV2TilingData tilingData_;
uint64_t coreNum_ = 0;
uint64_t workspaceSize_ = 0;
uint64_t usedCoreNums_ = 0;
uint64_t ubSize_ = 0;
int64_t tnd_ = 0; // x shape (tnd, d)
int64_t cacheBs_ = 0; // cache second dim : block_size
int64_t d_ = 0;
int64_t cacheD_ = 0;
int64_t scaleCol_ = 0;
int64_t rowOfFormerBlock_ = 0;
int64_t rowOfTailBlock_ = 0;
int64_t rowLoopOfFormerBlock_ = 0;
int64_t rowLoopOfTailBlock_ = 0;
int64_t rowFactor_ = 0;
int64_t tailRowFactorOfFormerBlock_ = 0;
int64_t tailRowFactorOfTailBlock_= 0;
int64_t layout_ = 2;
int64_t blockStride_ = 0;
platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
int64_t tilingKey_ = 0;
};
} // namespace optiling
#endif // INDEXER_COMPRESS_EPILOG_V2_TILING_H

55
csrc/attention/indexer_compress_epilog_v2/op_kernel/indexer_compress_epilog_v2.cpp Normal file
View File

@@ -0,0 +1,55 @@
/**
* 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_v2.cpp
* \brief
*/
#include "indexer_compress_epilog_v2_multi_row.h"
#include "indexer_compress_epilog_v2_single_row.h"
#define SINGLE_ROW_TILING_KEY 20020
#define MULTI_ROW_TILING_KEY 20021
using namespace AscendC;
extern "C" __global__ __aicore__ void indexer_compress_epilog_v2(
GM_ADDR indexer_compress_cache,
GM_ADDR x,
GM_ADDR slot_mapping,
GM_ADDR indexer_compress_cache_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_TILING_KEY)) {
IndexerCompressEpilogV2::IndexerCompressEpilogV2MultiRow<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE> op;
op.Init(x, slot_mapping, indexer_compress_cache, userWs, &tilingData, &pipe);
op.Process();
return;
} else if (TILING_KEY_IS(SINGLE_ROW_TILING_KEY)) {
IndexerCompressEpilogV2::IndexerCompressEpilogV2SingleRow<DTYPE_X, DTYPE_INDEXER_COMPRESS_CACHE> op;
op.Init(x, slot_mapping, indexer_compress_cache, userWs, &tilingData, &pipe);
op.Process();
return;
}
AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(oriOverflowMode);
}

281
csrc/attention/indexer_compress_epilog_v2/op_kernel/indexer_compress_epilog_v2_base.h Normal file
View File

@@ -0,0 +1,281 @@
/**
* 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_v2_base.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_V2_BASE_H
#define INDEXER_COMPRESS_EPILOG_V2_BASE_H
#include "kernel_operator.h"
namespace IndexerCompressEpilogV2 {
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;
constexpr int64_t B32_INTERPRE_TO_B8_RATIO = 4;
#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 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循环
uint32_t sregTail = tailReminder;
pregLoop = UpdateMask<float>(sregTail);
for (uint16_t j = 0; j < loopCountReminder; j++) {
LoadInputData<T1>(xLeft, xLocalAddr, pregLoop, loopCountFoldTwo * 2 * VL_FP32 + i * curColNumAlign);
Muls(xAbsLeft, xLeft, 0.0f, pregLoop);
Compare<float, CMPMODE::NE>(compareLeft, xAbsLeft, xAbsLeft, pregLoop);
MaskNot(compareLeft, compareLeft, pregLoop);
Abs(xAbsLeft, xLeft, compareLeft);
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 IndexerCompressEpilogV2
#endif

147
csrc/attention/indexer_compress_epilog_v2/op_kernel/indexer_compress_epilog_v2_multi_row.h Normal file
View File

@@ -0,0 +1,147 @@
/**
* 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_v2_multi_row.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_V2_MULTI_ROW_H
#define INDEXER_COMPRESS_EPILOG_V2_MULTI_ROW_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_v2_base.h"
namespace IndexerCompressEpilogV2 {
using namespace AscendC;
template <typename T0, typename T1>
class IndexerCompressEpilogV2MultiRow {
public:
__aicore__ inline IndexerCompressEpilogV2MultiRow()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR workspace,
const IndexerCompressEpilogV2TilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<fp8_e4m3fn_t>(tilingData->d) * sizeof(fp8_e4m3fn_t));
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<fp8_e4m3fn_t>();
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<fp8_e4m3fn_t>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<float>();
for (int64_t curValidIdx = 0; curValidIdx < validIdx; curValidIdx++) {
int64_t curSlotIdx = indexLocal.GetValue(curValidIdx);
int64_t blkNumIdx = curSlotIdx / tilingData->cacheBs;
int64_t blkSizeIdx = curSlotIdx % tilingData->cacheBs;
int64_t valueOffset = blkNumIdx * tilingData->blockStride + blkSizeIdx * tilingData->d;
int64_t scaleOffset = blkNumIdx * tilingData->blockStride + tilingData->cacheBs * tilingData->d
+ blkSizeIdx * tilingData->scaleCol * B32_INTERPRE_TO_B8_RATIO;
LocalTensor<T1> valueInterpreLocal = indexerCompressCacheLocal.ReinterpretCast<T1>();
LocalTensor<T1> scaleInterpreLocal = indexerCompressCacheScaleLocal.ReinterpretCast<T1>();
CopyOut(
valueInterpreLocal[curValidIdx * RoundUp<fp8_e4m3fn_t>(tilingData->d)],
indexerCompressCacheGm[valueOffset], 1, tilingData->d);
CopyOut(
scaleInterpreLocal[curValidIdx * RoundUp<float>(tilingData->scaleCol) * B32_INTERPRE_TO_B8_RATIO],
indexerCompressCacheGm[scaleOffset], 1, tilingData->scaleCol * B32_INTERPRE_TO_B8_RATIO);
}
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
fp8Max = FP8_E4M3FN_MAX_VALUE;
fp8Min = FP8_E4M3FN_MIN_VALUE;
}
private:
TPipe* pipe;
const IndexerCompressEpilogV2TilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
TBuf<QuePosition::VECCALC> indexBuf;
LocalTensor<T0> xLocal;
LocalTensor<fp8_e4m3fn_t> 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 IndexerCompressEpilogV2
#endif

124
csrc/attention/indexer_compress_epilog_v2/op_kernel/indexer_compress_epilog_v2_single_row.h Normal file
View File

@@ -0,0 +1,124 @@
/**
* 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_v2_single_row.h
* \brief
*/
#ifndef INDEXER_COMPRESS_EPILOG_V2_SINGLE_ROW_H
#define INDEXER_COMPRESS_EPILOG_V2_SINGLE_ROW_H
#include "kernel_operator.h"
#include "indexer_compress_epilog_v2_base.h"
namespace IndexerCompressEpilogV2 {
using namespace AscendC;
template <typename T0, typename T1>
class IndexerCompressEpilogV2SingleRow {
public:
__aicore__ inline IndexerCompressEpilogV2SingleRow()
{}
__aicore__ inline void Init(
GM_ADDR x, GM_ADDR slotMapping, GM_ADDR indexerCompressCache, GM_ADDR workspace,
const IndexerCompressEpilogV2TilingData* tilingDataPtr, TPipe* pipePtr)
{
pipe = pipePtr;
tilingData = tilingDataPtr;
xGm.SetGlobalBuffer((__gm__ T0*)x);
slotMappingGm.SetGlobalBuffer((__gm__ int32_t*)slotMapping);
indexerCompressCacheGm.SetGlobalBuffer((__gm__ T1*)indexerCompressCache);
pipe->InitBuffer(xQue, 2, tilingData->rowFactor * RoundUp<T0>(tilingData->d) * sizeof(T0));
pipe->InitBuffer(indexerCompressCacheQue, 2, tilingData->rowFactor * RoundUp<fp8_e4m3fn_t>(tilingData->d) * sizeof(fp8_e4m3fn_t));
pipe->InitBuffer(
indexerCompressCacheScaleQue, 2,
tilingData->rowFactor * RoundUp<float>(tilingData->scaleCol) * sizeof(float));
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<fp8_e4m3fn_t>();
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<fp8_e4m3fn_t>();
indexerCompressCacheScaleLocal = indexerCompressCacheScaleQue.template DeQue<float>();
int64_t blkNumIdx = slot / tilingData->cacheBs;
int64_t blkSizeIdx = slot % tilingData->cacheBs;
int64_t valueOffset = blkNumIdx * tilingData->blockStride + blkSizeIdx * tilingData->d;
int64_t scaleOffset = blkNumIdx * tilingData->blockStride + tilingData->cacheBs * tilingData->d
+ blkSizeIdx * tilingData->scaleCol * B32_INTERPRE_TO_B8_RATIO;
LocalTensor<T1> valueInterpreLocal = indexerCompressCacheLocal.ReinterpretCast<T1>();
LocalTensor<T1> scaleInterpreLocal = indexerCompressCacheScaleLocal.ReinterpretCast<T1>();
CopyOut(
valueInterpreLocal, indexerCompressCacheGm[valueOffset], 1, tilingData->d);
CopyOut(
scaleInterpreLocal, indexerCompressCacheGm[scaleOffset], 1, tilingData->scaleCol * B32_INTERPRE_TO_B8_RATIO);
indexerCompressCacheQue.template FreeTensor(indexerCompressCacheLocal);
indexerCompressCacheScaleQue.template FreeTensor(indexerCompressCacheScaleLocal);
}
}
__aicore__ inline void SetMaxValue()
{
maxValue = static_cast<float>(1.0) / FP8_E4M3FN_MAX_VALUE;
}
private:
TPipe* pipe;
const IndexerCompressEpilogV2TilingData* tilingData;
GlobalTensor<T0> xGm;
GlobalTensor<int32_t> slotMappingGm;
GlobalTensor<T1> indexerCompressCacheGm;
TQue<QuePosition::VECIN, 1> xQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheQue;
TQue<QuePosition::VECOUT, 1> indexerCompressCacheScaleQue;
LocalTensor<T0> xLocal;
LocalTensor<fp8_e4m3fn_t> indexerCompressCacheLocal;
LocalTensor<float> indexerCompressCacheScaleLocal;
float maxValue = 0.0f;
};
} // namespace IndexerCompressEpilogV2
#endif