* Copyright (c) 2025 Huawei Technologies Co., Ltd.
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
*/
* \file lightning_indexer_tiling.cpp
* \brief
*/
#include "lightning_indexer_tiling.h"
#include "../op_kernel/lightning_indexer_template_tiling_key.h"
using namespace ge;
using namespace AscendC;
using std::map;
using std::string;
namespace optiling {
ge::graphStatus LIInfoParser::CheckRequiredInOutExistence() const
{
OP_CHECK_IF(opParamInfo_.query.shape == nullptr, OP_LOGE(opName_, "Shape of tensor query is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.query.desc == nullptr, OP_LOGE(opName_, "Desc of tensor query is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.key.shape == nullptr, OP_LOGE(opName_, "Shape of tensor k is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.key.desc == nullptr, OP_LOGE(opName_, "Desc of tensor k is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.weights.shape == nullptr, OP_LOGE(opName_, "Shape of tensor value is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.weights.desc == nullptr, OP_LOGE(opName_, "Desc of tensor value is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.attenOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.attenOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.valuesOut.shape == nullptr, OP_LOGE(opName_, "Shape of tensor output values is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.valuesOut.desc == nullptr, OP_LOGE(opName_, "Desc of tensor output values is nullptr"),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::CheckRequiredAttrExistence() const
{
OP_CHECK_IF(opParamInfo_.layOut == nullptr, OP_LOGE(opName_, "attr layout_query is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.layOutKey == nullptr, OP_LOGE(opName_, "attr layout_key is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.sparseCount == nullptr, OP_LOGE(opName_, "attr sparse_count is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.sparseMode == nullptr, OP_LOGE(opName_, "attr sparse_mode is nullptr"),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::CheckRequiredParaExistence() const
{
if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetOpName()
{
if (context_->GetNodeName() == nullptr) {
OP_LOGE("LightningIndexer", "opName got from TilingContext is nullptr");
return ge::GRAPH_FAILED;
}
opName_ = context_->GetNodeName();
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetNpuInfo()
{
platformInfo_ = context_->GetPlatformInfo();
OP_CHECK_IF(platformInfo_ == nullptr, OP_LOGE(opName_, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo_);
uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
OP_CHECK_IF(aicNum == 0 || aivNum == 0, OP_LOGE(opName_, "num of core obtained is 0."), return GRAPH_FAILED);
socVersion_ = ascendcPlatform.GetSocVersion();
if ((socVersion_ != platform_ascendc::SocVersion::ASCEND910B) &&
(socVersion_ != platform_ascendc::SocVersion::ASCEND910_93) &&
(socVersion_ != platform_ascendc::SocVersion::ASCEND950)) {
OP_LOGE(opName_, "SOC Version[%d] is not support.", static_cast<int32_t>(socVersion_));
return GRAPH_FAILED;
}
OP_CHECK_IF(context_->GetWorkspaceSizes(1) == nullptr, OP_LOGE(opName_, "workSpaceSize got from ge is nullptr"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(context_->GetRawTilingData() == nullptr,
OP_LOGE(context_->GetNodeName(), "RawTilingData got from GE context is nullptr."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
void LIInfoParser::GetOptionalInputParaInfo()
{
opParamInfo_.actualSeqLengthsQ.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_Q_INDEX);
opParamInfo_.actualSeqLengthsQ.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_Q_INDEX);
opParamInfo_.actualSeqLengths.tensor = context_->GetOptionalInputTensor(ACTUAL_SEQ_K_INDEX);
opParamInfo_.actualSeqLengths.desc = context_->GetOptionalInputDesc(ACTUAL_SEQ_K_INDEX);
opParamInfo_.blockTable.tensor = context_->GetOptionalInputTensor(BLOCK_TABLE_INDEX);
opParamInfo_.blockTable.desc = context_->GetOptionalInputDesc(BLOCK_TABLE_INDEX);
}
void LIInfoParser::GetInputParaInfo()
{
opParamInfo_.query.desc = context_->GetInputDesc(QUERY_INDEX);
opParamInfo_.query.shape = context_->GetInputShape(QUERY_INDEX);
opParamInfo_.key.desc = context_->GetInputDesc(KEY_INDEX);
opParamInfo_.key.shape = context_->GetInputShape(KEY_INDEX);
opParamInfo_.weights.desc = context_->GetInputDesc(WEIGTHS_INDEX);
opParamInfo_.weights.shape = context_->GetInputShape(WEIGTHS_INDEX);
GetOptionalInputParaInfo();
}
void LIInfoParser::GetOutputParaInfo()
{
opParamInfo_.attenOut.desc = context_->GetOutputDesc(LIGHTNING_INDEXER);
opParamInfo_.attenOut.shape = context_->GetOutputShape(LIGHTNING_INDEXER);
opParamInfo_.valuesOut.desc = context_->GetOutputDesc(LIGHTNING_VALUES);
opParamInfo_.valuesOut.shape = context_->GetOutputShape(LIGHTNING_VALUES);
}
ge::graphStatus LIInfoParser::GetAndCheckAttrParaInfo()
{
auto attrs = context_->GetAttrs();
OP_CHECK_IF(attrs == nullptr, OPS_REPORT_VECTOR_INNER_ERR(context_->GetNodeName(), "attrs got from ge is nullptr"),
return ge::GRAPH_FAILED);
OP_LOGI(context_->GetNodeName(), "GetAndCheckAttrParaInfo start");
opParamInfo_.layOut = attrs->GetStr(ATTR_QUERY_LAYOUT_INDEX);
opParamInfo_.layOutKey = attrs->GetStr(ATTR_KEY_LAYOUT_INDEX);
opParamInfo_.sparseCount = attrs->GetAttrPointer<int32_t>(ATTR_SPARSE_COUNT_INDEX);
opParamInfo_.sparseMode = attrs->GetAttrPointer<int32_t>(ATTR_SPARSE_MODE_INDEX);
opParamInfo_.preTokens = attrs->GetAttrPointer<int64_t>(ATTR_PRE_TOKENS_INDEX);
opParamInfo_.nextTokens = attrs->GetAttrPointer<int64_t>(ATTR_NEXT_TOKENS_INDEX);
opParamInfo_.returnValue = attrs->GetAttrPointer<bool>(ATTR_RETURN_VALUE_INDEX);
if (opParamInfo_.layOut != nullptr) {
OP_LOGI(context_->GetNodeName(), "layout_query is:%s", opParamInfo_.layOut);
}
if (opParamInfo_.layOutKey != nullptr) {
OP_LOGI(context_->GetNodeName(), "layout_key is:%s", opParamInfo_.layOutKey);
}
if (opParamInfo_.sparseCount != nullptr) {
OP_LOGI(context_->GetNodeName(), "selscted count is:%d", *opParamInfo_.sparseCount);
}
if (opParamInfo_.sparseMode != nullptr) {
OP_LOGI(context_->GetNodeName(), "sparse mode is:%d", *opParamInfo_.sparseMode);
}
if (opParamInfo_.preTokens != nullptr) {
OP_LOGI(context_->GetNodeName(), "pre tokens is:%d", *opParamInfo_.preTokens);
}
if (opParamInfo_.nextTokens != nullptr) {
OP_LOGI(context_->GetNodeName(), "next tokens is:%d", *opParamInfo_.nextTokens);
}
if (opParamInfo_.returnValue != nullptr) {
OP_LOGI(context_->GetNodeName(), "return value is:%d", *opParamInfo_.returnValue);
}
OP_LOGI(context_->GetNodeName(), "GetAndCheckAttrParaInfo end");
OP_CHECK_IF(
((std::string(opParamInfo_.layOutKey) != "PA_BSND")
&& (std::string(opParamInfo_.layOut) != std::string(opParamInfo_.layOutKey))),
OP_LOGE(opName_, "under non-PA conditions, layout_query and layout_key should be equal."),
return ge::GRAPH_FAILED);
OP_CHECK_IF(
((std::string(opParamInfo_.layOutKey) != "PA_BSND") && (std::string(opParamInfo_.layOutKey) != "BSND")
&& (std::string(opParamInfo_.layOutKey) != "TND")),
OP_LOGE(opName_, "input attr layout_key only supported PA_BSND, BSND or TND"), return ge::GRAPH_FAILED);
OP_CHECK_IF(((std::string(opParamInfo_.layOut) != "BSND") && (std::string(opParamInfo_.layOut) != "TND")),
OP_LOGE(opName_, "input attr layout_query only supported BSND or TND."), return ge::GRAPH_FAILED);
OP_CHECK_IF((!((*opParamInfo_.sparseCount > 0) && (*opParamInfo_.sparseCount <= SPARSE_LIMIT)) &&
*opParamInfo_.sparseCount % 1024 != 0),
OP_LOGE(opName_, "input attr sparse_count must > 0 and <= 8192. And when sparse_count > 2048, sparse_count must be an interger multiple of 1024."),
return ge::GRAPH_FAILED);
OP_CHECK_IF(!((*opParamInfo_.sparseMode == 0) || (*opParamInfo_.sparseMode == SPARSE_MODE_LOWER)),
OP_LOGE(opName_, "input attr sparse_mode only supported 0 or 3."), return ge::GRAPH_FAILED);
OP_CHECK_IF(*opParamInfo_.preTokens != INT64_MAX,
OP_LOGE(opName_, "input attr pre_tokens only supported INT64_MAX."), return ge::GRAPH_FAILED);
OP_CHECK_IF(*opParamInfo_.nextTokens != INT64_MAX,
OP_LOGE(opName_, "input attr nextTokens only supported INT64_MAX."), return ge::GRAPH_FAILED);
OP_CHECK_IF(*opParamInfo_.returnValue && std::string(opParamInfo_.layOutKey) == "PA_BSND",
OP_LOGE(opName_, "when return_value is true, key layout do not support PA_BSND."), return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetOpParaInfo()
{
GetInputParaInfo();
GetOutputParaInfo();
if (ge::GRAPH_SUCCESS != GetAndCheckAttrParaInfo()) {
return ge::GRAPH_FAILED;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetAndCheckInOutDataType()
{
inputQType_ = opParamInfo_.query.desc->GetDataType();
inputKType_ = opParamInfo_.key.desc->GetDataType();
weightsType_ = opParamInfo_.weights.desc->GetDataType();
outputType_ = opParamInfo_.attenOut.desc->GetDataType();
valuesOutType_ = opParamInfo_.valuesOut.desc->GetDataType();
bool inDTypeAllEqual = (inputQType_ == inputKType_);
OP_CHECK_IF(!inDTypeAllEqual,
OP_LOGE(opName_, "The data types of the input query and key must be the same."),
return ge::GRAPH_FAILED);
OP_CHECK_IF(((inputQType_ != ge::DT_FLOAT16) && (inputQType_ != ge::DT_BF16)),
OP_LOGE(opName_, "The data types of the input query, key must be float16 or bfloat16."),
return ge::GRAPH_FAILED);
if (socVersion_ == platform_ascendc::SocVersion::ASCEND950) {
OP_CHECK_IF((inputQType_ != weightsType_),
OP_LOGE(opName_, "The data types of the input query, key, and weights must be the same."),
return ge::GRAPH_FAILED);
} else {
if (weightsType_ != ge::DT_FLOAT) {
OP_CHECK_IF((inputQType_ != weightsType_),
OP_LOGE(opName_, "The data types of the input query, key, and weights must be the same."),
return ge::GRAPH_FAILED);
} else {
OP_CHECK_IF((weightsType_ != ge::DT_FLOAT),
OP_LOGE(opName_, "The data types of the input weights must be float32."),
return ge::GRAPH_FAILED);
}
}
OP_CHECK_IF(outputType_ != ge::DT_INT32,
OP_LOGE(opName_, "The data types of the output sparse_indices must be int32."),
return ge::GRAPH_FAILED);
OP_CHECK_IF(valuesOutType_ != inputQType_,
OP_LOGE(opName_, "The data types of the output sparse_values must be same as inputQType."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetQueryKeyAndOutLayout()
{
const map<string, DataLayout> layoutMap = {
{"BSND", DataLayout::BSND},
{"TND", DataLayout::TND},
{"PA_BSND", DataLayout::BnBsND}
};
std::string layout(opParamInfo_.layOut);
auto it = layoutMap.find(layout);
if (it != layoutMap.end()) {
qLayout_ = it->second;
}
std::string layoutKey(opParamInfo_.layOutKey);
auto itKey = layoutMap.find(layoutKey);
if (itKey != layoutMap.end()) {
kLayout_ = itKey->second;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetAndCheckOptionalInput()
{
if (kLayout_ == DataLayout::BnBsND) {
OP_CHECK_IF(opParamInfo_.blockTable.tensor == nullptr,
OP_LOGE(opName_, "when layout_key is PA_BSND, input block_table must not be null"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(
opParamInfo_.actualSeqLengths.tensor == nullptr,
OP_LOGE(opName_, "when layout_key is PA_BSND, input actual_seq_lengths_key must not be null"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.blockTable.desc->GetDataType() != ge::DT_INT32,
OP_LOGE(opName_, "input block_table data type only support int32"), return ge::GRAPH_FAILED);
} else if (kLayout_ == DataLayout::TND) {
OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor == nullptr,
OP_LOGE(opName_, "when layout_key is TND, input actual_seq_lengths_key must not be null"),
return ge::GRAPH_FAILED);
}
OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor != nullptr &&
opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32,
OP_LOGE(opName_, "input actual_seq_lengths_key data type only support int32"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(opParamInfo_.actualSeqLengths.tensor != nullptr &&
opParamInfo_.actualSeqLengths.desc->GetDataType() != ge::DT_INT32,
OP_LOGE(opName_, "input actual_seq_lengths_key data type only support int32"),
return ge::GRAPH_FAILED);
if (qLayout_ == DataLayout::TND) {
OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor == nullptr,
OP_LOGE(opName_, "when layout_query is TND, input actual_seq_lengths_query must not be null"),
return ge::GRAPH_FAILED);
}
OP_CHECK_IF(opParamInfo_.actualSeqLengthsQ.tensor != nullptr &&
opParamInfo_.actualSeqLengthsQ.desc->GetDataType() != ge::DT_INT32,
OP_LOGE(opName_, "input actual_seq_lengths_query data type only support int32"),
return ge::GRAPH_FAILED);
OP_CHECK_IF(kLayout_ != DataLayout::BnBsND && opParamInfo_.blockTable.tensor != nullptr,
OP_LOGE(opName_, "when key layout is not PA_BSND, input block_table must be null"),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::CheckShapeDim()
{
OP_CHECK_IF((opParamInfo_.blockTable.tensor != nullptr) &&
(opParamInfo_.blockTable.tensor->GetStorageShape().GetDimNum() != DIM_NUM_TWO),
OP_LOGE(opName_, "the dim num of block_table's shape should be 2"), return ge::GRAPH_FAILED);
uint32_t kShapeDim = opParamInfo_.key.shape->GetStorageShape().GetDimNum();
uint32_t qShapeDim = opParamInfo_.query.shape->GetStorageShape().GetDimNum();
uint32_t weightsShapeDim = opParamInfo_.weights.shape->GetStorageShape().GetDimNum();
uint32_t outShapeDim = opParamInfo_.attenOut.shape->GetStorageShape().GetDimNum();
uint32_t valuesOutShapeDim = opParamInfo_.valuesOut.shape->GetStorageShape().GetDimNum();
uint32_t qExpectShapeDim = DIM_NUM_FOUR;
uint32_t kExpectShapeDim = DIM_NUM_FOUR;
if (qLayout_ == DataLayout::TND) {
qExpectShapeDim = DIM_NUM_THREE;
}
if (kLayout_ == DataLayout::TND) {
kExpectShapeDim = DIM_NUM_THREE;
}
OP_CHECK_IF(kShapeDim != kExpectShapeDim,
OP_LOGE(opName_, "the dim num of key's shape should be %u, but now is %u", kExpectShapeDim, kShapeDim),
return ge::GRAPH_FAILED);
OP_CHECK_IF(qShapeDim != qExpectShapeDim,
OP_LOGE(opName_, "the dim num of query's shape should be %u, but now is %u",
qExpectShapeDim, qShapeDim),
return ge::GRAPH_FAILED);
OP_CHECK_IF(outShapeDim != qExpectShapeDim,
OP_LOGE(opName_, "the dim num of sparse_indices's shape should be %u, but now is %u",
qExpectShapeDim, outShapeDim),
return ge::GRAPH_FAILED);
OP_CHECK_IF(valuesOutShapeDim != qExpectShapeDim && (*opParamInfo_.returnValue),
OP_LOGE(opName_, "the dim num of sparse_values's shape should be %u, but now is %u",
qExpectShapeDim, valuesOutShapeDim),
return ge::GRAPH_FAILED);
OP_CHECK_IF(!(weightsShapeDim == qExpectShapeDim - 1),
OP_LOGE(opName_, "the dim num of weights's shape should be %u, but now is %u", qExpectShapeDim - 1,
weightsShapeDim),
return ge::GRAPH_FAILED);
if (opParamInfo_.valuesOut.shape->GetStorageShape().GetShapeSize() != 0 && !(*opParamInfo_.returnValue)) {
OP_LOGW(opName_, "when returnValue is false, valuesOut must be null.");
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetN1Size()
{
if (qLayout_ == DataLayout::BSND) {
n1Size_ = static_cast<uint32_t>(opParamInfo_.query.shape->GetStorageShape().GetDim(DIM_IDX_TWO));
} else {
n1Size_ = static_cast<uint32_t>(opParamInfo_.query.shape->GetStorageShape().GetDim(1));
}
OP_LOGI(context_->GetNodeName(), "n1Size is %d", n1Size_);
OP_CHECK_IF(n1Size_ > QUERY_HEAD_NUM_LIMIT, OP_LOGE(opName_, "N1 is %u, but N1 must be no greater than %u.",
n1Size_, QUERY_HEAD_NUM_LIMIT), return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
const std::string &actualSeqLenName) const
{
size = static_cast<uint32_t>(tensor->GetShapeSize());
if (size <= 0) {
OP_LOGE(opName_, "%s's shape size is %u, it should be greater than 0.", actualSeqLenName.c_str(), size);
return ge::GRAPH_FAILED;
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetAndCheckN2Size()
{
uint32_t n2Index = (kLayout_ == DataLayout::TND) ? DIM_IDX_ONE : DIM_IDX_TWO;
n2Size_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(n2Index));
OP_LOGI(context_->GetNodeName(), "n2Size_ is %d", n2Size_);
OP_CHECK_IF(n2Size_ != 1, OP_LOGE(opName_, "key shape[%u] is numhead, only support 1.", n2Index),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetGSize()
{
if (n1Size_ % n2Size_ != 0) {
OP_LOGE(opName_, "input query's head_num %u can not be a multiple of key's head_num %u.", n1Size_, n2Size_);
return ge::GRAPH_FAILED;
}
gSize_ = n1Size_ / n2Size_;
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetBatchSize()
{
if ((qLayout_ == DataLayout::TND)) {
return GetActualSeqLenSize(bSize_, opParamInfo_.actualSeqLengthsQ.tensor, "input actual_seq_lengths_query");
} else {
bSize_ = opParamInfo_.query.shape->GetStorageShape().GetDim(0);
return ge::GRAPH_SUCCESS;
}
}
ge::graphStatus LIInfoParser::GetHeadDim()
{
uint32_t dIndex = DIM_IDX_TWO;
switch (qLayout_) {
case DataLayout::TND:
dIndex = DIM_IDX_TWO;
break;
case DataLayout::BSND:
dIndex = DIM_IDX_THREE;
break;
default:
OP_LOGE(opName_, "unsupported layout for getting head dim.");
return ge::GRAPH_FAILED;
}
headDim_ = opParamInfo_.query.shape->GetStorageShape().GetDim(dIndex);
OP_CHECK_IF(headDim_ != HEAD_DIM_LIMIT, OP_LOGE(opName_, "input query's last dim head_dim only support 128."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetS1Size()
{
if (qLayout_ == DataLayout::BSND) {
s1Size_ = opParamInfo_.query.shape->GetStorageShape().GetDim(1);
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetAndCheckBlockSize()
{
blockSize_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(1));
OP_LOGI(context_->GetNodeName(), "blockSize_ is %d", blockSize_);
OP_CHECK_IF(((blockSize_ % 16 != 0) || (blockSize_ == 0) || (blockSize_ > 1024)),
OP_LOGE(opName_, "input key's block_size must be a multiple of 16 and belong to (0, 1024]."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::CheckBlockCount()
{
int32_t blockCount_ = static_cast<uint32_t>(opParamInfo_.key.shape->GetStorageShape().GetDim(0));
OP_CHECK_IF((blockCount_ == 0),
OP_LOGE(opName_, "input key's block_count cannot be 0."),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetS2SizeForPageAttention()
{
if (GetAndCheckBlockSize() != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
if (CheckBlockCount() != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
maxBlockNumPerBatch_ = opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(1);
s2Size_ = maxBlockNumPerBatch_ * blockSize_;
OP_LOGI(context_->GetNodeName(), "maxBlockNumPerBatch_ is %d, blockSize_ is %d, s2Size_ is %d",
maxBlockNumPerBatch_, blockSize_, s2Size_);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::GetS2Size()
{
if (kLayout_ == DataLayout::BnBsND) {
return GetS2SizeForPageAttention();
} else if (kLayout_ == DataLayout::TND) {
s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(0);
} else if (kLayout_ == DataLayout::BSND) {
s2Size_ = opParamInfo_.key.shape->GetStorageShape().GetDim(1);
}
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::ValidateInputShapesMatchQtnd()
{
if (kLayout_ == DataLayout::TND) {
OP_CHECK_IF(
(opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_),
OP_LOGE(opName_,
"TND case input actual_seq_lengths_query, actual_seq_lengths_key are %u, %ld respectively, they must be same.",
bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()),
return ge::GRAPH_FAILED);
} else {
OP_CHECK_IF(
(opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_) ||
(opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_),
OP_LOGE(
opName_,
"TND case input actual_seq_lengths_query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.",
bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(),
opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
}
uint32_t qTsize = opParamInfo_.query.shape->GetStorageShape().GetDim(0);
OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != qTsize) ||
(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != qTsize),
OP_LOGE(opName_, "TND case input query, weights and sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.",
qTsize, opParamInfo_.weights.shape->GetStorageShape().GetDim(0),
opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0) != qTsize &&
(*opParamInfo_.returnValue)),
OP_LOGE(opName_, "TND case input query and sparse_values dim 0 are %u, %ld respectively, they must be same.",
qTsize, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::ValidateInputShapesMatchQbsnd()
{
if (kLayout_ == DataLayout::BnBsND) {
OP_CHECK_IF((opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0) != bSize_) ||
(opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_),
OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_key, block_table dim 0 are %u, %ld, %ld respectively, they must be same.",
bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize(),
opParamInfo_.blockTable.tensor->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
} else if (kLayout_ == DataLayout::BSND) {
OP_CHECK_IF(opParamInfo_.key.shape->GetStorageShape().GetDim(0) != bSize_,
OP_LOGE(opName_, "BSND case input query, key dim 0 are %u, %ld respectively, they must be same.",
bSize_, opParamInfo_.key.shape->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.actualSeqLengths.tensor != nullptr) &&
(opParamInfo_.actualSeqLengths.tensor->GetShapeSize() != bSize_),
OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_key dim 0 are %u, %ld respectively, they must be same.",
bSize_, opParamInfo_.actualSeqLengths.tensor->GetShapeSize()),
return ge::GRAPH_FAILED);
}
OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(0) != bSize_) ||
(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0) != bSize_),
OP_LOGE(opName_, "BSND case input query, weight and sparse_indices dim 0 are %u, %ld, %ld respectively, they must be same.",
bSize_, opParamInfo_.weights.shape->GetStorageShape().GetDim(0),
opParamInfo_.attenOut.shape->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0) != bSize_ &&
(*opParamInfo_.returnValue)),
OP_LOGE(opName_, "BSND case input query, sparse_values dim 0 are %u, %ld respectively, they must be same.",
bSize_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(0)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.actualSeqLengthsQ.tensor != nullptr) &&
(opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize() != bSize_),
OP_LOGE(opName_, "BSND case input query, actual_seq_lengths_query dim 0 are %u, %ld respectively, they must be same",
bSize_, opParamInfo_.actualSeqLengthsQ.tensor->GetShapeSize()),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(1) != s1Size_) ||
(opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1) != s1Size_),
OP_LOGE(opName_, "BSND case input query, weight and sparse_indices dim 1 are %u, %ld, %ld, they must be same.",
s1Size_, opParamInfo_.weights.shape->GetStorageShape().GetDim(1),
opParamInfo_.attenOut.shape->GetStorageShape().GetDim(1)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(1) != s1Size_ &&
(*opParamInfo_.returnValue)),
OP_LOGE(opName_, "BSND case input query and sparse_values dim 1 are %u, %ld, they must be same.",
s1Size_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(1)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LIInfoParser::ValidateInputShapesMatch()
{
TND:
query [T,N1,D],
key [BlockNum,BlockSize,N2,D],
weight [T,N1],
block_table [BatchSize, BatchMaxBlockNum],
act_seq_k [BatchSize]
act_seq_q [BatchSize],
out [T,N2,topk]
----------------------
BSND:
query [BatchSize,S1,N1,D],
key [BlockNum,BlockSize,N2,D],
weight [BatchSize,S1,N1],
block_table [BatchSize, BatchMaxBlockNum],
act_seq_k [BatchSize]
act_seq_q [BatchSize] 可选
out [BatchSize,S1,N2,topk]
*/
uint32_t queryWeightsN1Dim = 1;
uint32_t outN2Dim = 1;
if (qLayout_ == DataLayout::TND) {
if (ValidateInputShapesMatchQtnd() != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
} else {
if (ValidateInputShapesMatchQbsnd() != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
queryWeightsN1Dim = DIM_IDX_TWO;
outN2Dim = DIM_IDX_TWO;
}
OP_CHECK_IF((opParamInfo_.weights.shape->GetStorageShape().GetDim(queryWeightsN1Dim) != n1Size_),
OP_LOGE(opName_, "input query, weight shape dim N1 must be same."), return ge::GRAPH_FAILED);
uint32_t keyDDim = kLayout_ == DataLayout::TND ? DIM_IDX_TWO : DIM_IDX_THREE;
OP_CHECK_IF((opParamInfo_.key.shape->GetStorageShape().GetDim(keyDDim) != headDim_),
OP_LOGE(opName_, "input query, key shape last dim must be same."), return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_),
OP_LOGE(opName_, "input query and output sparse_indices shape n2 dim must be same,"
"but now they are %u, %ld respectively.",
n2Size_, opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim) != n2Size_ &&
(*opParamInfo_.returnValue)),
OP_LOGE(opName_, "input query and sparse_values shape n2 dim must be same,"
"but now they are %u, %ld respectively.",
n2Size_, opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount),
OP_LOGE(opName_, "output sparse_indices shape last dim must be same as attr sparse_count,"
"but now they are %u, %ld respectively.", *opParamInfo_.sparseCount,
opParamInfo_.attenOut.shape->GetStorageShape().GetDim(outN2Dim + 1)),
return ge::GRAPH_FAILED);
OP_CHECK_IF((opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim + 1) != *opParamInfo_.sparseCount &&
(*opParamInfo_.returnValue)),
OP_LOGE(opName_, "output sparse_values shape last dim must be same as attr sparse_count,"
"but now they are %u, %ld respectively.", *opParamInfo_.sparseCount,
opParamInfo_.valuesOut.shape->GetStorageShape().GetDim(outN2Dim + 1)),
return ge::GRAPH_FAILED);
return ge::GRAPH_SUCCESS;
}
void LIInfoParser::GenerateInfo(LITilingInfo &liInfo)
{
liInfo.opName = opName_;
liInfo.platformInfo = platformInfo_;
liInfo.opParamInfo = opParamInfo_;
liInfo.socVersion = socVersion_;
liInfo.bSize = bSize_;
liInfo.n1Size = n1Size_;
liInfo.n2Size = n2Size_;
liInfo.s1Size = s1Size_;
liInfo.s2Size = s2Size_;
liInfo.gSize = gSize_;
liInfo.inputQType = inputQType_;
liInfo.inputKType = inputKType_;
liInfo.weightsType = weightsType_;
liInfo.outputType = outputType_;
liInfo.blockSize = blockSize_;
liInfo.maxBlockNumPerBatch = maxBlockNumPerBatch_;
std::string layOutKeyStr(opParamInfo_.layOutKey);
liInfo.pageAttentionFlag = layOutKeyStr == "PA_BSND" ? true : false;
liInfo.sparseMode = *opParamInfo_.sparseMode;
liInfo.sparseCount = *opParamInfo_.sparseCount;
liInfo.preTokens = *opParamInfo_.preTokens;
liInfo.nextTokens = *opParamInfo_.nextTokens;
liInfo.returnValue = *opParamInfo_.returnValue;
liInfo.inputQLayout = qLayout_;
liInfo.inputKLayout = kLayout_;
}
ge::graphStatus LIInfoParser::ParseAndCheck(LITilingInfo &liInfo)
{
if (ge::GRAPH_SUCCESS != GetOpName() || ge::GRAPH_SUCCESS != GetNpuInfo() || ge::GRAPH_SUCCESS != GetOpParaInfo() ||
ge::GRAPH_SUCCESS != CheckRequiredParaExistence()) {
return ge::GRAPH_FAILED;
}
if (ge::GRAPH_SUCCESS != GetAndCheckInOutDataType() || ge::GRAPH_SUCCESS != GetQueryKeyAndOutLayout() ||
ge::GRAPH_SUCCESS != GetAndCheckOptionalInput()) {
return ge::GRAPH_FAILED;
}
if (ge::GRAPH_SUCCESS != CheckShapeDim() || ge::GRAPH_SUCCESS != GetN1Size() ||
ge::GRAPH_SUCCESS != GetAndCheckN2Size() || ge::GRAPH_SUCCESS != GetGSize()) {
return ge::GRAPH_FAILED;
}
if (ge::GRAPH_SUCCESS != GetBatchSize() || ge::GRAPH_SUCCESS != GetS1Size() || ge::GRAPH_SUCCESS != GetHeadDim() ||
ge::GRAPH_SUCCESS != GetS2Size()) {
return ge::GRAPH_FAILED;
}
if (ge::GRAPH_SUCCESS != ValidateInputShapesMatch()) {
return ge::GRAPH_FAILED;
}
GenerateInfo(liInfo);
return ge::GRAPH_SUCCESS;
}
static ge::graphStatus TilingPrepareForLightningIndexer(gert::TilingParseContext * )
{
return ge::GRAPH_SUCCESS;
}
ge::graphStatus LightningIndexerTiling::DoTiling(LITilingInfo *tilingInfo)
{
auto ascendcPlatform = platform_ascendc::PlatformAscendC(tilingInfo->platformInfo);
uint32_t aivNum = ascendcPlatform.GetCoreNumAiv();
uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
uint32_t blockDim = ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum);
context_->SetBlockDim(blockDim);
constexpr uint32_t MM1_RES_ELEM_SIZE = 4;
constexpr uint32_t DOUBLE_BUFFER = 2;
constexpr uint32_t M_BASE_SIZE = 512;
constexpr uint32_t S2_BASE_SIZE = 512;
constexpr uint32_t V1_RES_ELEM_SIZE = 4;
constexpr uint32_t V1_RES_ELEM_TYPE = 2;
constexpr uint32_t V1_DECODE_PARAM_ELEM_SIZE = 8;
constexpr uint32_t V1_DECODE_PARAM_NUM = 16;
constexpr uint32_t V1_DECODE_DATA_NUM = 2;
constexpr uint32_t S1_BASE_SIZE = 8;
constexpr uint32_t TOPK_MAX_SIZE = 2048;
uint32_t workspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510) {
constexpr uint32_t s1BaseSize = 4;
constexpr uint32_t s2BaseSize = 128;
workspaceSize +=
s1BaseSize * ((tilingInfo->s2Size + s2BaseSize - 1) / s2BaseSize) * s2BaseSize * sizeof(uint16_t) * aicNum;
} else {
constexpr uint32_t mm1ResSize = M_BASE_SIZE * S2_BASE_SIZE;
workspaceSize += mm1ResSize * MM1_RES_ELEM_SIZE * DOUBLE_BUFFER * aicNum;
workspaceSize +=
V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_RES_ELEM_TYPE * TOPK_MAX_SIZE * V1_RES_ELEM_SIZE * aicNum;
workspaceSize +=
V1_DECODE_DATA_NUM * S1_BASE_SIZE * V1_DECODE_PARAM_NUM * V1_DECODE_PARAM_ELEM_SIZE * aicNum;
}
size_t *workSpaces = context_->GetWorkspaceSizes(1);
workSpaces[0] = workspaceSize;
tilingData_.set_bSize(tilingInfo->bSize);
tilingData_.set_s2Size(tilingInfo->s2Size);
tilingData_.set_s1Size(tilingInfo->s1Size);
tilingData_.set_sparseCount(tilingInfo->sparseCount);
tilingData_.set_gSize(tilingInfo->gSize);
tilingData_.set_blockSize(tilingInfo->blockSize);
tilingData_.set_maxBlockNumPerBatch(tilingInfo->maxBlockNumPerBatch);
tilingData_.set_sparseMode(tilingInfo->sparseMode);
tilingData_.set_preTokens(tilingInfo->preTokens);
tilingData_.set_nextTokens(tilingInfo->nextTokens);
tilingData_.set_returnValue(tilingInfo->returnValue);
tilingData_.set_usedCoreNum(blockDim);
tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
uint32_t inputQType = static_cast<uint32_t>(tilingInfo->inputQType);
uint32_t inputKType = static_cast<uint32_t>(tilingInfo->inputKType);
uint32_t weightsType = static_cast<uint32_t>(tilingInfo->weightsType);
uint32_t outputType = static_cast<uint32_t>(tilingInfo->outputType);
uint32_t pageAttentionFlag = static_cast<uint32_t>(tilingInfo->pageAttentionFlag);
uint32_t inputQLayout = static_cast<uint32_t>(tilingInfo->inputQLayout);
uint32_t inputKLayout = static_cast<uint32_t>(tilingInfo->inputKLayout);
uint32_t weightTypeFlag = (weightsType == ge::DT_FLOAT) ? 1 : 0;
uint64_t tilingKey =
GET_TPL_TILING_KEY(inputQType, inputKType, outputType, pageAttentionFlag, inputQLayout, inputKLayout, weightTypeFlag);
context_->SetTilingKey(tilingKey);
context_->SetScheduleMode(1);
return ge::GRAPH_SUCCESS;
}
ge::graphStatus TilingForLightningIndexer(gert::TilingContext *context)
{
OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("LightningIndexer", "Tiling context is null."),
return ge::GRAPH_FAILED);
LITilingInfo liInfo;
LIInfoParser LIInfoParser(context);
if (LIInfoParser.ParseAndCheck(liInfo) != ge::GRAPH_SUCCESS) {
return ge::GRAPH_FAILED;
}
LightningIndexerTiling liTiling(context);
return liTiling.DoTiling(&liInfo);
}
IMPL_OP_OPTILING(LightningIndexer)
.Tiling(TilingForLightningIndexer)
.TilingParse<LICompileInfo>(TilingPrepareForLightningIndexer);
}