30
csrc/attention/sparse_flash_attention/op_host/CMakeLists.txt
Normal file
30
csrc/attention/sparse_flash_attention/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,30 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# 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_op_to_compiled_list()
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
set(sparse_flash_attention_depends attention/common CACHE INTERNAL "Dependencies for sparse_flash_attention")
|
||||
target_sources(op_host_aclnnInner PRIVATE
|
||||
sparse_flash_attention_def.cpp
|
||||
)
|
||||
endif()
|
||||
|
||||
add_ops_compile_options(
|
||||
OP_NAME SparseFlashAttention
|
||||
OPTIONS --cce-auto-sync=off
|
||||
-Wno-deprecated-declarations
|
||||
-mllvm -cce-vf-remove-membar=false
|
||||
-mllvm -cce-aicore-hoist-movemask=false
|
||||
)
|
||||
|
||||
if (NOT BUILD_OPS_RTY_KERNEL)
|
||||
add_modules_sources(OPTYPE sparse_flash_attention ACLNNTYPE aclnn_inner)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
#include <cstring>
|
||||
#include "graph/types.h"
|
||||
#include "aclnn_sparse_flash_attention.h"
|
||||
|
||||
#include "opdev/make_op_executor.h"
|
||||
#include "opdev/op_dfx.h"
|
||||
#include "opdev/op_executor.h"
|
||||
#include "opdev/tensor_view_utils.h"
|
||||
#include "opdev/op_def.h"
|
||||
#include "opdev/op_log.h"
|
||||
#include "opdev/shape_utils.h"
|
||||
#include "opdev/common_types.h"
|
||||
#include "opdev/data_type_utils.h"
|
||||
#include "opdev/format_utils.h"
|
||||
|
||||
using namespace op;
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
extern aclnnStatus aclnnInnerSparseFlashAttentionGetWorkspaceSize(
|
||||
const aclTensor *query, const aclTensor *key, const aclTensor *value, const aclTensor *sparse_indices,
|
||||
const aclTensor *blockTableOptional, const aclTensor *actualSeqLengthsQueryOptional, const aclTensor *actualSeqLengthsKvOptional,
|
||||
const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, double scaleValue,
|
||||
int64_t sparseBlockSizeOptional, char *layoutQueryOptional, char *layoutKvOptional,
|
||||
int64_t sparseMode, int64_t preTokens, int64_t nextTokens, int64_t attentionMode,
|
||||
bool returnSoftmaxLse, const aclTensor *attentionOut, const aclTensor *softmaxMax,
|
||||
const aclTensor *softmaxSum, uint64_t *workspaceSize, aclOpExecutor **executor);
|
||||
|
||||
extern aclnnStatus aclnnInnerSparseFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
|
||||
const aclrtStream stream);
|
||||
|
||||
class TensorHolder {
|
||||
public:
|
||||
TensorHolder(const aclTensor *&output, aclDataType dataType, std::string varName) {
|
||||
inner_ = nullptr;
|
||||
name_ = varName;
|
||||
if (output == nullptr) {
|
||||
std::vector<int64_t> shape = {0};
|
||||
int64_t addr = 0xff;
|
||||
inner_ = aclCreateTensor(shape.data(), shape.size(),
|
||||
dataType, shape.data(), 0, ACL_FORMAT_ND,
|
||||
shape.data(), shape.size(), static_cast<void *>(&addr));
|
||||
output = inner_;
|
||||
}
|
||||
}
|
||||
|
||||
~TensorHolder() {
|
||||
if (inner_) {
|
||||
aclDestroyTensor(inner_);
|
||||
inner_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void CheckTensorConditionalNotNull(bool conditional) const {
|
||||
if (inner_ && conditional) {
|
||||
OP_LOGW("Check %s != nullptr failed!", name_.c_str());
|
||||
} else if (!inner_ && !conditional) {
|
||||
OP_LOGW("Check %s == nullptr failed!", name_.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
bool IsTensorNotNull() const {
|
||||
return inner_ == nullptr;
|
||||
}
|
||||
|
||||
private:
|
||||
const aclTensor *inner_;
|
||||
std::string name_;
|
||||
};
|
||||
|
||||
aclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize(
|
||||
const aclTensor *query,
|
||||
const aclTensor *key,
|
||||
const aclTensor *value,
|
||||
const aclTensor *sparseIndices,
|
||||
const aclTensor *blockTableOptional,
|
||||
const aclTensor *actualSeqLengthsQueryOptional,
|
||||
const aclTensor *actualSeqLengthsKvOptional,
|
||||
const aclTensor *queryRopeOptional,
|
||||
const aclTensor *keyRopeOptional,
|
||||
double scaleValue,
|
||||
int64_t sparseBlockSizeOptional,
|
||||
char *layoutQueryOptional,
|
||||
char *layoutKvOptional,
|
||||
int64_t sparseMode,
|
||||
int64_t preTokens,
|
||||
int64_t nextTokens,
|
||||
int64_t attentionMode,
|
||||
bool returnSoftmaxLse,
|
||||
const aclTensor *attentionOut,
|
||||
const aclTensor *softmaxMax,
|
||||
const aclTensor *softmaxSum,
|
||||
uint64_t *workspaceSize,
|
||||
aclOpExecutor **executor)
|
||||
{
|
||||
if (returnSoftmaxLse) {
|
||||
if (softmaxMax == nullptr || softmaxSum == nullptr) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "when returnSoftmaxLse is true, softmaxMax and softmaxSum cannot be nullptr.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
} else {
|
||||
if (softmaxMax == nullptr && softmaxSum == nullptr) {
|
||||
auto softmaxMaxHolder = TensorHolder(softmaxMax, aclDataType::ACL_FLOAT, std::string("softmaxMax"));
|
||||
auto softmaxSumHolder = TensorHolder(softmaxSum, aclDataType::ACL_FLOAT, std::string("softmaxSum"));
|
||||
if (softmaxMax == nullptr) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Failed to create the holder of tensor softmaxMax!");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (softmaxSum == nullptr) {
|
||||
OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Failed to create the holder of tensor softmaxSum!");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
}
|
||||
}
|
||||
return aclnnInnerSparseFlashAttentionGetWorkspaceSize(
|
||||
query, key, value, sparseIndices, blockTableOptional, actualSeqLengthsQueryOptional, actualSeqLengthsKvOptional, queryRopeOptional, keyRopeOptional,
|
||||
scaleValue, sparseBlockSizeOptional, layoutQueryOptional, layoutKvOptional, sparseMode, preTokens,
|
||||
nextTokens, attentionMode, returnSoftmaxLse, attentionOut,
|
||||
softmaxMax, softmaxSum, workspaceSize, executor);
|
||||
}
|
||||
|
||||
aclnnStatus aclnnSparseFlashAttention(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
|
||||
const aclrtStream stream)
|
||||
{
|
||||
return aclnnInnerSparseFlashAttention(workspace, workspaceSize, executor, stream);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
@@ -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.
|
||||
*/
|
||||
|
||||
#ifndef ACLNN_SPARSE_FLASH_ATTENTION_H
|
||||
#define ACLNN_SPARSE_FLASH_ATTENTION_H
|
||||
|
||||
#include "aclnn/acl_meta.h"
|
||||
#include "aclnn/aclnn_base.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief The first interface of aclnnMlaPrologV2WeightNz calculates the workspace size based on the specific calculation process.
|
||||
* @domain aclnn_ops_infer
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus aclnnSparseFlashAttentionGetWorkspaceSize(
|
||||
const aclTensor *query,
|
||||
const aclTensor *key,
|
||||
const aclTensor *value,
|
||||
const aclTensor *sparseIndices,
|
||||
const aclTensor *blockTableOptional,
|
||||
const aclTensor *actualSeqLengthsQueryOptional,
|
||||
const aclTensor *actualSeqLengthsKvOptional,
|
||||
const aclTensor *queryRopeOptional,
|
||||
const aclTensor *keyRopeOptional,
|
||||
double scaleValue,
|
||||
int64_t sparseBlockSizeOptional,
|
||||
char *layoutQueryOptional,
|
||||
char *layoutKvOptional,
|
||||
int64_t sparseMode,
|
||||
int64_t preTokens,
|
||||
int64_t nextTokens,
|
||||
int64_t attentionMode,
|
||||
bool returnSoftmaxLse,
|
||||
const aclTensor *attentionOut,
|
||||
const aclTensor *softmaxMax,
|
||||
const aclTensor *softmaxSum,
|
||||
uint64_t *workspaceSize,
|
||||
aclOpExecutor **executor);
|
||||
|
||||
/**
|
||||
* @brief The second interface of ACLNN_SPARSE_FLASH_ATTENTION_H is used to perform calculations.
|
||||
*/
|
||||
__attribute__((visibility("default"))) aclnnStatus aclnnSparseFlashAttention(void *workspace,
|
||||
uint64_t workspaceSize,
|
||||
aclOpExecutor *executor,
|
||||
const aclrtStream stream);
|
||||
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif // ACLNN_SPARSE_FLASH_ATTENTION_H
|
||||
@@ -0,0 +1,102 @@
|
||||
/**
|
||||
* 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 sparse_flash_attention_def.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include "register/op_def_registry.h"
|
||||
|
||||
namespace ops {
|
||||
class SparseFlashAttention : public OpDef {
|
||||
public:
|
||||
explicit SparseFlashAttention(const char *name) : OpDef(name)
|
||||
{
|
||||
this->Input("query")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("key")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("value")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("sparse_indices")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("block_table")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("actual_seq_lengths_query")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("actual_seq_lengths_kv")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("query_rope")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("key_rope")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Output("attention_out")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("softmax_max")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("softmax_sum")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Attr("scale_value").AttrType(REQUIRED).Float(1.0);
|
||||
this->Attr("sparse_block_size").AttrType(OPTIONAL).Int(1);
|
||||
this->Attr("layout_query").AttrType(OPTIONAL).String("BSND");
|
||||
this->Attr("layout_kv").AttrType(OPTIONAL).String("BSND");
|
||||
this->Attr("sparse_mode").AttrType(OPTIONAL).Int(3); // 3:默认值，只计算下三角
|
||||
this->Attr("pre_tokens").AttrType(OPTIONAL).Int(INT64_MAX);
|
||||
this->Attr("next_tokens").AttrType(OPTIONAL).Int(INT64_MAX);
|
||||
this->Attr("attention_mode").AttrType(OPTIONAL).Int(2);
|
||||
this->Attr("return_softmax_lse").AttrType(OPTIONAL).Bool(false);
|
||||
OpAICoreConfig aicore_config;
|
||||
aicore_config.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true);
|
||||
this->AICore().AddConfig("ascend910b", aicore_config);
|
||||
this->AICore().AddConfig("ascend910_93", aicore_config);
|
||||
this->AICore().AddConfig("ascend950", aicore_config);
|
||||
}
|
||||
};
|
||||
OP_ADD(SparseFlashAttention);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,130 @@
|
||||
/**
|
||||
* 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 sparse_flash_attention_proto.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#include <graph/utils/type_utils.h>
|
||||
#include <register/op_impl_registry.h>
|
||||
#include "err/ops_err.h"
|
||||
|
||||
using namespace ge;
|
||||
|
||||
namespace ops {
|
||||
constexpr size_t QUERY_INPUT_INDEX = 0;
|
||||
constexpr size_t KEY_INPUT_INDEX = 1;
|
||||
|
||||
constexpr uint32_t DIM_NUM_1 = 1;
|
||||
constexpr uint32_t DIM_NUM_3 = 3;
|
||||
constexpr uint32_t DIM_NUM_4 = 4;
|
||||
constexpr uint32_t DIM_INDEX_0 = 0;
|
||||
constexpr uint32_t DIM_INDEX_1 = 1;
|
||||
constexpr uint32_t DIM_INDEX_2 = 2;
|
||||
constexpr uint32_t DIM_INDEX_3 = 3;
|
||||
constexpr uint32_t LAYOUT_KEY_ATTR_INDEX = 3;
|
||||
constexpr uint32_t RETURN_SOFTMAX_LSE_INDEX = 8;
|
||||
|
||||
constexpr uint32_t OUTPUT_INDEX_0 = 0;
|
||||
constexpr uint32_t OUTPUT_INDEX_1 = 1;
|
||||
constexpr uint32_t OUTPUT_INDEX_2 = 2;
|
||||
|
||||
ge::graphStatus InferShapeSparseFlashAttention(gert::InferShapeContext *context)
|
||||
{
|
||||
OP_CHECK_IF(context == nullptr, OP_LOGE("SparseFlashAttention", "InferShapeContext is nullptr"),
|
||||
return ge::GRAPH_FAILED);
|
||||
const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, queryShape);
|
||||
|
||||
const gert::Shape *keyShape = context->GetInputShape(KEY_INPUT_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, keyShape);
|
||||
|
||||
gert::Shape *attentionOutShape = context->GetOutputShape(0);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, attentionOutShape);
|
||||
*attentionOutShape = *queryShape;
|
||||
|
||||
gert::Shape *softmaxMaxShape = context->GetOutputShape(1);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, softmaxMaxShape);
|
||||
|
||||
gert::Shape *softmaxSumShape = context->GetOutputShape(2);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, softmaxSumShape);
|
||||
|
||||
auto attrs = context->GetAttrs();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
|
||||
const char *inputLayoutKeyPtr = attrs->GetAttrPointer<char>(LAYOUT_KEY_ATTR_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, inputLayoutKeyPtr);
|
||||
std::string inputLayoutKeyPtrStr = std::string(inputLayoutKeyPtr);
|
||||
const bool *lse_flag = attrs->GetAttrPointer<bool>(RETURN_SOFTMAX_LSE_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, lse_flag);
|
||||
bool return_softmax_lse = (lse_flag != nullptr)? *lse_flag : false;
|
||||
|
||||
if(return_softmax_lse){
|
||||
if(queryShape->GetDimNum() == DIM_NUM_3){
|
||||
if (inputLayoutKeyPtrStr == "PA_BSND") {
|
||||
softmaxMaxShape->SetDimNum(DIM_NUM_3);
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2));
|
||||
|
||||
softmaxSumShape->SetDimNum(DIM_NUM_3);
|
||||
softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_2));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_2));
|
||||
} else {
|
||||
softmaxMaxShape->SetDimNum(DIM_NUM_3);
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1));
|
||||
|
||||
softmaxSumShape->SetDimNum(DIM_NUM_3);
|
||||
softmaxSumShape->SetDim(DIM_INDEX_0, keyShape->GetDim(DIM_INDEX_1));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_1, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1) / keyShape->GetDim(DIM_INDEX_1));
|
||||
}
|
||||
} else {
|
||||
softmaxMaxShape->SetDimNum(DIM_NUM_4);
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1));
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2));
|
||||
|
||||
softmaxSumShape->SetDimNum(DIM_NUM_4);
|
||||
softmaxSumShape->SetDim(DIM_INDEX_0, queryShape->GetDim(DIM_INDEX_0));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_1, keyShape->GetDim(DIM_INDEX_2));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_2, queryShape->GetDim(DIM_INDEX_1));
|
||||
softmaxSumShape->SetDim(DIM_INDEX_3, queryShape->GetDim(DIM_INDEX_2) / keyShape->GetDim(DIM_INDEX_2));
|
||||
}
|
||||
} else {
|
||||
softmaxMaxShape->SetDimNum(DIM_NUM_1);
|
||||
softmaxMaxShape->SetDim(DIM_INDEX_0, 0);
|
||||
softmaxSumShape->SetDimNum(DIM_NUM_1);
|
||||
softmaxSumShape->SetDim(DIM_INDEX_0, 0);
|
||||
}
|
||||
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus InferDataTypeSparseFlashAttention(gert::InferDataTypeContext *context)
|
||||
{
|
||||
OP_CHECK_IF(context == nullptr, OP_LOGE("SparseFlashAttention", "InferShapeContext is nullptr"),
|
||||
return ge::GRAPH_FAILED);
|
||||
const auto inputDataType = context->GetInputDataType(QUERY_INPUT_INDEX);
|
||||
context->SetOutputDataType(OUTPUT_INDEX_0, inputDataType);
|
||||
context->SetOutputDataType(OUTPUT_INDEX_1, ge::DT_FLOAT);
|
||||
context->SetOutputDataType(OUTPUT_INDEX_2, ge::DT_FLOAT);
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_INFERSHAPE(SparseFlashAttention)
|
||||
.InferShape(InferShapeSparseFlashAttention)
|
||||
.InferDataType(InferDataTypeSparseFlashAttention);
|
||||
} // namespace ops
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,639 @@
|
||||
/**
|
||||
* 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 sparse_flash_attention_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef SPARSE_FLASH_ATTENTION_TILING_H
|
||||
#define SPARSE_FLASH_ATTENTION_TILING_H
|
||||
|
||||
#include <sstream>
|
||||
#include <graph/utils/type_utils.h>
|
||||
#include <exe_graph/runtime/tiling_context.h>
|
||||
#include <tiling/platform/platform_ascendc.h>
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "exe_graph/runtime/tiling_context.h"
|
||||
#include "platform/soc_spec.h"
|
||||
|
||||
namespace optiling {
|
||||
// ------------------算子原型索引常量定义----------------
|
||||
// Inputs Index
|
||||
constexpr uint32_t QUERY_INPUT_INDEX = 0;
|
||||
constexpr uint32_t KEY_INPUT_INDEX = 1;
|
||||
constexpr uint32_t VALUE_INPUT_INDEX = 2;
|
||||
constexpr uint32_t SPARSE_INDICES_INPUT_INDEX = 3;
|
||||
constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 4;
|
||||
constexpr uint32_t ACT_SEQ_LEN_Q_INPUT_INDEX = 5;
|
||||
constexpr uint32_t ACT_SEQ_LEN_KV_INPUT_INDEX = 6;
|
||||
constexpr uint32_t QUERY_ROPE_INPUT_INDEX = 7;
|
||||
constexpr uint32_t KEY_ROPE_INPUT_INDEX = 8;
|
||||
// Outputs Index
|
||||
constexpr uint32_t OUTPUT_INDEX = 0;
|
||||
constexpr uint32_t SOFTMAXMAX_INDEX = 1;
|
||||
constexpr uint32_t SOFTMAXSUM_INDEX = 2;
|
||||
|
||||
// Attributes Index
|
||||
constexpr uint32_t SCALE_VALUE_ATTR_INDEX = 0;
|
||||
constexpr uint32_t SPARSE_BLOCK_SIZE_ATTR_INDEX = 1;
|
||||
constexpr uint32_t LAYOUT_QUERY_ATTR_INDEX = 2;
|
||||
constexpr uint32_t LAYOUT_KV_ATTR_INDEX = 3;
|
||||
constexpr uint32_t SPARSE_MODE_ATTR_INDEX = 4;
|
||||
constexpr uint32_t PRE_TOKENS_ATTR_INDEX = 5;
|
||||
constexpr uint32_t NEXT_TOKENS_ATTR_INDEX = 6;
|
||||
constexpr uint32_t ATTENTION_MODE_ATTR_INDEX = 7;
|
||||
constexpr uint32_t RETURN_SOFTMAX_LSE_ATTR_INDEX = 8;
|
||||
// Dim Num
|
||||
constexpr size_t DIM_NUM_TWO = 2;
|
||||
constexpr size_t DIM_NUM_THREE = 3;
|
||||
constexpr size_t DIM_NUM_FOUR = 4;
|
||||
// 常量
|
||||
constexpr uint32_t MAX_BLOCK_SIZE = 1024;
|
||||
constexpr uint32_t COPYND2NZ_SRC_STRIDE_LIMITATION = 65535;
|
||||
constexpr uint32_t NUM_BYTES_FLOAT = 4;
|
||||
constexpr uint32_t NUM_BYTES_FLOAT16 = 2;
|
||||
constexpr uint32_t NUM_BYTES_BF16 = 2;
|
||||
constexpr uint32_t BYTE_BLOCK = 32;
|
||||
const uint32_t SFA_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐
|
||||
|
||||
// ------------------公共定义--------------------------
|
||||
enum class SFALayout : uint32_t {
|
||||
BSND = 0,
|
||||
TND = 1,
|
||||
PA_BSND = 2,
|
||||
BNSG = 3,
|
||||
NTG = 4
|
||||
};
|
||||
|
||||
struct SFATilingShapeCompareParam {
|
||||
int64_t B = 1;
|
||||
int64_t S = 1;
|
||||
int64_t N = 1;
|
||||
int64_t D = 1;
|
||||
int64_t T = 1;
|
||||
int64_t G = 1;
|
||||
// PA
|
||||
int64_t Bs = 1;
|
||||
int64_t Bn = 1;
|
||||
};
|
||||
|
||||
enum class KvStorageMode : uint32_t {
|
||||
BATCH_CONTINUOUS = 0,
|
||||
PAGE_ATTENTION = 1
|
||||
};
|
||||
|
||||
enum class SFAPerfMode : uint32_t {
|
||||
C_TEMPLATE_MODE = 0,
|
||||
V_TEMPLATE_MODE
|
||||
};
|
||||
|
||||
enum class SFAAxis : uint32_t {
|
||||
B = 0,
|
||||
S = 1,
|
||||
N = 2,
|
||||
D = 3,
|
||||
K = 3, // sparse_indices的K和key的D枚举值相同，表达相同位置, 最后一维
|
||||
T = 5,
|
||||
Bn = 6, // block number
|
||||
Bs = 7, // block size
|
||||
G = 8,
|
||||
};
|
||||
|
||||
struct SFARequiredParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::StorageShape *shape;
|
||||
};
|
||||
|
||||
struct SFAOptionalParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::Tensor *tensor;
|
||||
};
|
||||
|
||||
// -----------算子Tiling入参结构体定义---------------
|
||||
struct SFAParaInfo {
|
||||
SFARequiredParaInfo query = {nullptr, nullptr};
|
||||
SFARequiredParaInfo key = {nullptr, nullptr};
|
||||
SFARequiredParaInfo value = {nullptr, nullptr};
|
||||
SFARequiredParaInfo sparseIndices = {nullptr, nullptr};
|
||||
SFAOptionalParaInfo blockTable = {nullptr, nullptr};
|
||||
SFAOptionalParaInfo actualSeqLengthsQ = {nullptr, nullptr};
|
||||
SFAOptionalParaInfo actualSeqLengths = {nullptr, nullptr};
|
||||
SFAOptionalParaInfo queryRope = {nullptr, nullptr};
|
||||
SFAOptionalParaInfo keyRope = {nullptr, nullptr};
|
||||
SFARequiredParaInfo attenOut = {nullptr, nullptr};
|
||||
SFARequiredParaInfo softmaxMax = {nullptr, nullptr};
|
||||
SFARequiredParaInfo softmaxSum = {nullptr, nullptr};
|
||||
|
||||
const char *layoutQuery = nullptr;
|
||||
const char *layoutKV = nullptr;
|
||||
const int64_t *sparseBlockSize = nullptr;
|
||||
const float *scaleValue = nullptr;
|
||||
const int64_t *sparseMode = nullptr;
|
||||
const int64_t *preTokens = nullptr;
|
||||
const int64_t *nextTokens = nullptr;
|
||||
const int64_t *attentionMode = nullptr;
|
||||
const bool *returnSoftmaxLse = nullptr;
|
||||
};
|
||||
|
||||
struct InnerSplitParams {
|
||||
uint32_t s1GBaseSize = 1;
|
||||
uint32_t s2BaseSize = 1;
|
||||
};
|
||||
|
||||
// -----------算子TilingData定义---------------
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionBaseParamsMla)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, batchSize)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, seqSize)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, qSeqSize)
|
||||
TILING_DATA_FIELD_DEF(int64_t, blockSize)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch)
|
||||
TILING_DATA_FIELD_DEF(float, scaleValue)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, nNumOfQInOneGroup)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsQ)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, actualLenDimsKV)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, outputLayout)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, sparseMode)
|
||||
TILING_DATA_FIELD_DEF(int64_t, preTokens)
|
||||
TILING_DATA_FIELD_DEF(int64_t, nextTokens)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, attentionMode)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, returnSoftmaxLse)
|
||||
TILING_DATA_FIELD_DEF(int64_t, sparseBlockSize)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, sparseBlockCount)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsNull)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, isActualLenDimsKVNull)
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttentionBaseParamsMlaOp, SparseFlashAttentionBaseParamsMla)
|
||||
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionSingleCoreParamsMla)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSingleCoreParamsMlaOp, SparseFlashAttentionSingleCoreParamsMla)
|
||||
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionSingleCoreTensorSizeMla)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, mmResUbSize);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, bmm2ResUbSize);
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSingleCoreTensorSizeMlaOp, SparseFlashAttentionSingleCoreTensorSizeMla)
|
||||
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionSplitKVParamsMla)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, s2) // S2切分份数
|
||||
TILING_DATA_FIELD_DEF(uint32_t, accumOutSize) // FD workspace
|
||||
TILING_DATA_FIELD_DEF(uint32_t, logSumExpSize) // FD workspace
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttentionSplitKVParamsMlaOp, SparseFlashAttentionSplitKVParamsMla)
|
||||
|
||||
// 内切基本块参数
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionInnerSplitParams)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, mBaseSize)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, s2BaseSize)
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttentionInnerSplitParamsOp, SparseFlashAttentionInnerSplitParams)
|
||||
|
||||
BEGIN_TILING_DATA_DEF(SparseFlashAttentionTilingDataMla)
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionBaseParamsMla, baseParams);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSplitKVParamsMla, splitKVParams);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSingleCoreParamsMla, singleCoreParams);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionSingleCoreTensorSizeMla, singleCoreTensorSize);
|
||||
TILING_DATA_FIELD_DEF_STRUCT(SparseFlashAttentionInnerSplitParams, innerSplitParams);
|
||||
END_TILING_DATA_DEF
|
||||
REGISTER_TILING_DATA_CLASS(SparseFlashAttention, SparseFlashAttentionTilingDataMla)
|
||||
|
||||
template <typename T> inline T Align(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd) - 1) / (rnd) * (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::string SFAShape2String(const T &shape)
|
||||
{
|
||||
std::ostringstream oss;
|
||||
oss << "[";
|
||||
if (shape.GetDimNum() > 0) {
|
||||
for (size_t i = 0; i < shape.GetDimNum() - 1; ++i) {
|
||||
oss << shape.GetDim(i) << ", ";
|
||||
}
|
||||
oss << shape.GetDim(shape.GetDimNum() - 1);
|
||||
}
|
||||
oss << "]";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
static std::string GetShapeStr(gert::Shape shape);
|
||||
static std::string SFADataTypeToSerialString(ge::DataType type);
|
||||
std::string SFATensorDesc2String(const gert::StorageShape *shape, const gert::CompileTimeTensorDesc *tensor);
|
||||
std::string SFADebugTilingContext(const gert::TilingContext *context);
|
||||
std::string SFALayoutToSerialString(SFALayout layout);
|
||||
|
||||
// -----------算子Tiling入参信息类---------------
|
||||
struct SFATilingInfo {
|
||||
const char *opName = nullptr;
|
||||
fe::PlatFormInfos *platformInfo = nullptr;
|
||||
SFAParaInfo opParamInfo;
|
||||
|
||||
// Base Param
|
||||
NpuArch npuArch = NpuArch::DAV_2201;
|
||||
bool isA5 = false;
|
||||
uint32_t bSize = 0;
|
||||
uint32_t n1Size = 0;
|
||||
uint32_t n2Size = 0;
|
||||
uint32_t s1Size = 0;
|
||||
int64_t s2Size = 0;
|
||||
uint32_t qkHeadDim = 0;
|
||||
uint32_t vHeadDim = 0;
|
||||
uint32_t gSize = 0;
|
||||
uint32_t ropeHeadDim = 0;
|
||||
uint32_t qTSize = 0; // 仅TND时生效
|
||||
uint32_t kvTSize = 0; // 仅TND时生效
|
||||
float scaleValue = 0;
|
||||
uint32_t innerPrecise = 0;
|
||||
uint32_t l2CacheOffFlag = 0;
|
||||
int64_t sparseBlockSize = 0;
|
||||
int64_t sparseBlockCount = 0;
|
||||
|
||||
bool pageAttentionFlag = false;
|
||||
int64_t blockSize = 0;
|
||||
uint32_t blockTypeSize = 0;
|
||||
uint32_t maxBlockNumPerBatch = 0;
|
||||
uint32_t totalBlockNum = 0;
|
||||
|
||||
uint32_t actualLenDimsQ = 0;
|
||||
uint32_t maxActualseq = 0;
|
||||
|
||||
bool actualQSeqLenFlag = false;
|
||||
bool actualSeqLenFlag = false;
|
||||
bool isSameSeqAllKVTensor = true;
|
||||
bool isSameActualseq = true;
|
||||
uint32_t actualLenDimsKV = 0;
|
||||
std::vector<int64_t> kvListSeqLens {};
|
||||
|
||||
uint32_t sparseMode = 0;
|
||||
int64_t preTokens = INT64_MAX;
|
||||
int64_t nextTokens = INT64_MAX;
|
||||
uint32_t attentionMode = 2;
|
||||
bool returnSoftmaxLse = false;
|
||||
|
||||
ge::DataType inputQType = ge::DT_FLOAT16;
|
||||
ge::DataType inputKvType = ge::DT_FLOAT16;
|
||||
ge::DataType outputType = ge::DT_FLOAT16;
|
||||
|
||||
KvStorageMode kvStorageMode = KvStorageMode::BATCH_CONTINUOUS;
|
||||
|
||||
SFALayout qLayout = SFALayout::BSND;
|
||||
SFALayout topkLayout = SFALayout::BSND;
|
||||
SFALayout outLayout = SFALayout::BSND;
|
||||
SFALayout kvLayout = SFALayout::BSND;
|
||||
SFALayout softmaxMaxLayout = SFALayout::BNSG;
|
||||
SFALayout softmaxSumLayout = SFALayout::BNSG;
|
||||
|
||||
ge::DataType inputQRopeType = ge::DT_FLOAT16;
|
||||
ge::DataType inputKRopeType = ge::DT_FLOAT16;
|
||||
|
||||
uint64_t l2CacheSize = 0;
|
||||
};
|
||||
|
||||
// ---------------算子Tiling类---------------
|
||||
class SFAMlaTiling {
|
||||
public:
|
||||
explicit SFAMlaTiling(gert::TilingContext *context) : context_(context) {}
|
||||
ge::graphStatus DoOpTiling(SFATilingInfo *sfaInfo);
|
||||
|
||||
private:
|
||||
ge::graphStatus SetBlockDim(uint32_t blockDim) const;
|
||||
ge::graphStatus SetTilingKey(uint64_t tilingKey) const;
|
||||
ge::graphStatus SetWorkspaceSize(uint64_t workspaceSize) const;
|
||||
ge::graphStatus SetTilingData(TilingDef &tilingData) const;
|
||||
gert::TilingContext *context_ = nullptr;
|
||||
ge::graphStatus GetPlatformInfo();
|
||||
void GenTilingKey();
|
||||
bool DealSameSeqEachBatch();
|
||||
|
||||
void ZeroTensorProcess() const;
|
||||
void InitParams();
|
||||
|
||||
void Split();
|
||||
bool IsBalanceSplitCore();
|
||||
|
||||
void SplitBalanced();
|
||||
void CalcInnerSize(uint32_t s2Size);
|
||||
|
||||
bool IsFlashDecode(uint32_t coreNum);
|
||||
|
||||
void FillTilingBaseParamsMla();
|
||||
void FillTilingSplitKVMla();
|
||||
|
||||
void FillTilingSingleCoreParamsMla();
|
||||
void FillTilingSingleCoreTensorSizeMla();
|
||||
void FillTiling();
|
||||
|
||||
void CalcUbBmm();
|
||||
void CheckUbSpace();
|
||||
void NormalCalcFDWorkSpace(const uint32_t actCoreNum);
|
||||
void CalcFDWorkSpace(const uint32_t actCoreNum);
|
||||
void GetWorkspaceSize();
|
||||
|
||||
uint32_t CalcBalanceFDParamNums(const uint32_t actCoreNum) const;
|
||||
|
||||
void CalcBlockDim();
|
||||
|
||||
uint32_t GetTypeSize(ge::DataType dtype) const;
|
||||
|
||||
bool balanceModeFlag_ = false;
|
||||
bool splitKVFlag_ = false;
|
||||
|
||||
uint32_t coreNum_ = 0;
|
||||
SFAPerfMode perfMode_ = SFAPerfMode::V_TEMPLATE_MODE;
|
||||
uint32_t kvSplitPart_ = 1;
|
||||
size_t mmResUbSize_ = 0;
|
||||
size_t bmm2ResUbSize_ = 0;
|
||||
size_t qPreSizeMla_= 0;
|
||||
uint32_t sInnerLoopTimes_ = 0;
|
||||
uint32_t sInnerSize_ = 0;
|
||||
uint32_t sInnerSizeTail_ = 0;
|
||||
uint32_t sInnerSizeAlign_ = 0;
|
||||
uint32_t kvSplit_ = 0;
|
||||
uint32_t usedCoreNum_ = 0;
|
||||
uint32_t formerCoreNum_ = 0;
|
||||
uint32_t blockSplitBn2Range_ = 0;
|
||||
uint32_t tailSplitedBatchRange_ = 0;
|
||||
|
||||
uint32_t aicNum_ = 0;
|
||||
uint32_t aivNum_ = 0;
|
||||
size_t libapiSize_ = 0;
|
||||
|
||||
SparseFlashAttentionTilingDataMla tilingData_;
|
||||
uint32_t blockDim_{0};
|
||||
uint64_t workspaceSize_{0};
|
||||
uint64_t tilingKey_{0};
|
||||
|
||||
uint32_t headDimAlign_ = 0;
|
||||
uint32_t mBaseSize_ = 128;
|
||||
uint32_t mFdBaseSize_ = 8;
|
||||
|
||||
SFATilingInfo *sfaInfo_ = nullptr;
|
||||
};
|
||||
|
||||
// -----------算子Tiling入参信息解析及Check类---------------
|
||||
class SFATilingCheck {
|
||||
public:
|
||||
explicit SFATilingCheck(const SFATilingInfo &sfaInfo) : sfaInfo_(sfaInfo) {};
|
||||
~SFATilingCheck() = default;
|
||||
virtual ge::graphStatus Process();
|
||||
private:
|
||||
void Init();
|
||||
void LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList,
|
||||
const ge::DataType &actualDtype, const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc,
|
||||
const std::string &name) const;
|
||||
template <typename T> void LogErrorNumberSupport(const std::vector<T> &expectNumberList,
|
||||
const T &actualValue, const std::string &name, const std::string subName) const;
|
||||
template <typename T> void LogErrorDimNumSupport(const std::vector<T> &expectNumberList,
|
||||
const T &actualValue, const std::string &name) const;
|
||||
ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape,
|
||||
const std::vector<size_t> &expectDimNumList, const std::string &name) const;
|
||||
ge::graphStatus CheckDimNumInLayoutSupport(const SFALayout &layout,
|
||||
const gert::StorageShape *shape, const std::string &name) const;
|
||||
void LogErrorLayoutSupport(const std::vector<SFALayout> &expectLayoutList,
|
||||
const SFALayout &actualLayout, const std::string &name) const;
|
||||
ge::graphStatus GetExpectedShape(gert::Shape &shapeExpected,
|
||||
const SFATilingShapeCompareParam ¶m, const SFALayout &layout) const;
|
||||
ge::graphStatus CompareShape(SFATilingShapeCompareParam ¶m,
|
||||
const gert::Shape &shape, const SFALayout &layout, const std::string &name) const;
|
||||
ge::graphStatus CheckLayoutSupport(const SFALayout &actualLayout, const std::string &name) const;
|
||||
ge::graphStatus CheckSingleParaQuery() const;
|
||||
ge::graphStatus CheckSingleParaKey() const;
|
||||
ge::graphStatus CheckSingleParaValue() const;
|
||||
ge::graphStatus CheckSingleParaQueryRope() const;
|
||||
ge::graphStatus CheckSingleParaKeyRope() const;
|
||||
ge::graphStatus CheckSingleParaAttenOut() const;
|
||||
ge::graphStatus CheckSingleParaNumHeads() const;
|
||||
ge::graphStatus CheckSingleParaKvHeadNums() const;
|
||||
ge::graphStatus CheckSingleParaLayout() const;
|
||||
ge::graphStatus CheckSingleParaSparseMode() const;
|
||||
ge::graphStatus CheckSingleParaSparseBlockSize() const;
|
||||
ge::graphStatus CheckSingleParaSparseIndices() const;
|
||||
ge::graphStatus CheckSinglePara() const;
|
||||
ge::graphStatus CheckMultiParaConsistency() const;
|
||||
ge::graphStatus CheckRopeExistence();
|
||||
ge::graphStatus CheckExists(const void *pointer, const std::string &name) const;
|
||||
ge::graphStatus CheckNotExists(const void *pointer, const std::string &name) const;
|
||||
ge::graphStatus CheckExistsByMap(const std::map<std::string, const void *> ¶mMap) const;
|
||||
ge::graphStatus CheckNotExistsByMap(const std::map<std::string, const void *> ¶mMap) const;
|
||||
ge::graphStatus CheckExistenceByMap(std::map<std::string, const void *> &existMap,
|
||||
std::map<std::string, const void *> ¬ExistMap) const;
|
||||
template <typename T> ge::graphStatus CheckAttrValueByMap(
|
||||
std::map<std::string, std::pair<const T *, T>> &attrMap) const;
|
||||
ge::graphStatus CheckParaExistenceMlaNoquant() const;
|
||||
ge::graphStatus CheckParaExistenceGqaNoquant() const;
|
||||
ge::graphStatus CheckParaExistenceMla() const;
|
||||
ge::graphStatus CheckParaExistence();
|
||||
ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
|
||||
const SFALayout &layout, const std::string &name) const;
|
||||
void SetSFAShapeCompare();
|
||||
ge::graphStatus CheckQRope();
|
||||
ge::graphStatus CheckQRopeShape();
|
||||
ge::graphStatus CheckVAndKRopeShapeForBatchContinuous();
|
||||
ge::graphStatus CheckVAndKRopeShapeForPageAttention();
|
||||
ge::graphStatus CheckVAndKRopeShape();
|
||||
ge::graphStatus CheckVAndKRope();
|
||||
ge::graphStatus CheckTopK();
|
||||
ge::graphStatus CheckTopkShape();
|
||||
ge::graphStatus CheckBlockTable() const;
|
||||
ge::graphStatus CheckDTypeConsistency(const ge::DataType &actualDtype,
|
||||
const ge::DataType &expectDtype, const std::string &name) const;
|
||||
|
||||
ge::graphStatus CheckAttenOut();
|
||||
ge::graphStatus CheckAttenOutShape();
|
||||
ge::graphStatus CheckSoftmaxMax();
|
||||
ge::graphStatus CheckSoftmaxMaxShape();
|
||||
ge::graphStatus CheckSoftmaxSum();
|
||||
ge::graphStatus CheckSoftmaxSumShape();
|
||||
ge::graphStatus CheckActualSeqLensQ();
|
||||
ge::graphStatus CheckActualSeqLensQShape();
|
||||
ge::graphStatus CheckActualSeqLensQDType();
|
||||
ge::graphStatus CheckActualSeqLens();
|
||||
ge::graphStatus CheckActualSeqLensDType();
|
||||
ge::graphStatus CheckActualSeqLensShape();
|
||||
ge::graphStatus CheckMultiParaConsistency();
|
||||
|
||||
ge::graphStatus CheckFeatureMlaNoQuantShape() const;
|
||||
ge::graphStatus CheckFeatureMlaNoQuantLayout() const;
|
||||
ge::graphStatus CheckFeatureMlaNoQuantDtype() const;
|
||||
ge::graphStatus CheckFeatureMlaNoquantPa() const;
|
||||
ge::graphStatus CheckFeatureMlaNoquant() const;
|
||||
ge::graphStatus CheckFeatureMla() const;
|
||||
ge::graphStatus CheckFeature() const;
|
||||
|
||||
ge::graphStatus CheckSingleParaPreTokens() const;
|
||||
ge::graphStatus CheckSingleParaNextTokens() const;
|
||||
|
||||
private:
|
||||
const char *opName_;
|
||||
fe::PlatFormInfos *platformInfo_;
|
||||
SFAParaInfo opParamInfo_;
|
||||
const SFATilingInfo &sfaInfo_;
|
||||
|
||||
uint32_t bSize_ = 0;
|
||||
uint32_t n1Size_ = 0;
|
||||
uint32_t n2Size_ = 0;
|
||||
uint32_t gSize_ = 0;
|
||||
uint32_t s1Size_ = 0;
|
||||
int64_t s2Size_ = 0;
|
||||
uint32_t qkHeadDim_ = 0;
|
||||
uint32_t vHeadDim_ = 0;
|
||||
uint32_t ropeHeadDim_ = 0;
|
||||
uint32_t qTSize_ = 0; // 仅TND时生效
|
||||
uint32_t kvTSize_ = 0; // 仅TND时生效
|
||||
KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS;
|
||||
uint32_t sparseBlockCount_ = 0;
|
||||
int64_t sparseBlockSize_ = 0;
|
||||
|
||||
SFALayout qLayout_ = SFALayout::BSND;
|
||||
SFALayout topkLayout_ = SFALayout::BSND;
|
||||
SFALayout outLayout_ = SFALayout::BSND;
|
||||
SFALayout kvLayout_ = SFALayout::BSND;
|
||||
SFALayout softmaxMaxLayout_ = SFALayout::BNSG;
|
||||
SFALayout softmaxSumLayout_ = SFALayout::BNSG;
|
||||
|
||||
uint32_t maxBlockNumPerBatch_ = 0;
|
||||
int64_t blockSize_ = 0;
|
||||
|
||||
uint32_t aicNum_ = 0;
|
||||
uint32_t aivNum_ = 0;
|
||||
NpuArch npuArch_ = NpuArch::DAV_2201;
|
||||
bool isA5_ = false;
|
||||
uint64_t l2CacheSize_ = 0;
|
||||
|
||||
ge::DataType inputQType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputKvType_ = ge::DT_FLOAT16;
|
||||
ge::DataType outputType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputQRopeType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputKRopeType_ = ge::DT_FLOAT16;
|
||||
|
||||
gert::Shape queryShapeCmp_{};
|
||||
gert::Shape keyShapeCmp_{};
|
||||
gert::Shape valueShapeCmp_{};
|
||||
gert::Shape topkShapeCmp_{};
|
||||
gert::Shape queryRopeShapeCmp_{};
|
||||
gert::Shape keyRopeShapeCmp_{};
|
||||
gert::Shape attenOutShapeCmp_{};
|
||||
gert::Shape softmaxMaxShapeCmp_{};
|
||||
gert::Shape softmaxSumShapeCmp_{};
|
||||
};
|
||||
|
||||
class SFAInfoParser {
|
||||
public:
|
||||
explicit SFAInfoParser(const gert::TilingContext *context) : context_(context) {}
|
||||
~SFAInfoParser() = default;
|
||||
|
||||
ge::graphStatus CheckRequiredInOutExistence() const;
|
||||
ge::graphStatus CheckRequiredAttrExistence() const;
|
||||
ge::graphStatus CheckRequiredParaExistence() const;
|
||||
|
||||
ge::graphStatus GetActualSeqLenSize(uint32_t &size, const gert::Tensor *tensor,
|
||||
SFALayout &layout, const std::string &name) const;
|
||||
ge::graphStatus GetActualSeqLenQSize(uint32_t &size);
|
||||
ge::graphStatus GetOpName();
|
||||
ge::graphStatus GetNpuInfo();
|
||||
void GetOptionalInputParaInfo();
|
||||
void GetInputParaInfo();
|
||||
void GetOutputParaInfo();
|
||||
ge::graphStatus GetAttrParaInfo();
|
||||
ge::graphStatus GetKvCache();
|
||||
ge::graphStatus GetOpParaInfo();
|
||||
|
||||
ge::graphStatus GetInOutDataType();
|
||||
ge::graphStatus GetBatchSize();
|
||||
ge::graphStatus GetQTSize();
|
||||
ge::graphStatus GetKVTSize();
|
||||
ge::graphStatus GetQkHeadDim();
|
||||
ge::graphStatus GetS1Size();
|
||||
ge::graphStatus GetKvStorageMode();
|
||||
ge::graphStatus GetKvLayout();
|
||||
void SetSFAShape();
|
||||
ge::graphStatus GetS2SizeForBatchContinuous();
|
||||
ge::graphStatus GetMaxBlockNumPerBatch();
|
||||
ge::graphStatus GetBlockSize();
|
||||
ge::graphStatus GetS2SizeForPageAttention();
|
||||
ge::graphStatus GetS2Size();
|
||||
ge::graphStatus GetValueHeadDim();
|
||||
ge::graphStatus GetRopeHeadDim();
|
||||
ge::graphStatus GetQueryAndOutLayout();
|
||||
ge::graphStatus GetTopkLayout();
|
||||
ge::graphStatus GetSoftmaxMaxAndSumLayout();
|
||||
ge::graphStatus GetN1Size();
|
||||
ge::graphStatus GetN2Size();
|
||||
ge::graphStatus GetGSize();
|
||||
ge::graphStatus GetSparseBlockCount();
|
||||
ge::graphStatus GetActualseqInfo();
|
||||
void GenerateInfo(SFATilingInfo &sfaInfo);
|
||||
ge::graphStatus Parse(SFATilingInfo &sfaInfo);
|
||||
|
||||
public:
|
||||
bool HasAxis(const SFAAxis &axis, const SFALayout &layout, const gert::Shape &shape) const;
|
||||
size_t GetAxisIdx(const SFAAxis &axis, const SFALayout &layout) const;
|
||||
uint32_t GetAxisNum(const gert::Shape &shape, const SFAAxis &axis,const SFALayout &layout) const;
|
||||
|
||||
const gert::TilingContext *context_ = nullptr;
|
||||
|
||||
const char *opName_;
|
||||
fe::PlatFormInfos *platformInfo_;
|
||||
SFAParaInfo opParamInfo_;
|
||||
static constexpr int64_t invalidDimValue_ = std::numeric_limits<int64_t>::min();
|
||||
|
||||
uint32_t bSize_ = 0;
|
||||
uint32_t n1Size_ = 0;
|
||||
uint32_t n2Size_ = 0;
|
||||
uint32_t gSize_ = 0;
|
||||
uint32_t s1Size_ = 0;
|
||||
int64_t s2Size_ = 0;
|
||||
uint32_t qkHeadDim_ = 0;
|
||||
uint32_t vHeadDim_ = 0;
|
||||
uint32_t ropeHeadDim_ = 0;
|
||||
uint32_t qTSize_ = 0; // 仅TND时生效
|
||||
uint32_t kvTSize_ = 0; // 仅TND时生效
|
||||
KvStorageMode kvStorageMode_ = KvStorageMode::BATCH_CONTINUOUS;
|
||||
uint32_t sparseBlockCount_ = 0;
|
||||
|
||||
SFALayout qLayout_ = SFALayout::BSND;
|
||||
SFALayout topkLayout_ = SFALayout::BSND;
|
||||
SFALayout outLayout_ = SFALayout::BSND;
|
||||
SFALayout kvLayout_ = SFALayout::BSND;
|
||||
SFALayout softmaxMaxLayout_ = SFALayout::BNSG;
|
||||
SFALayout softmaxSumLayout_ = SFALayout::BNSG;
|
||||
uint32_t maxBlockNumPerBatch_ = 0;
|
||||
uint32_t blockSize_ = 0;
|
||||
|
||||
NpuArch npuArch_ = NpuArch::DAV_2201;
|
||||
bool isA5_ = false;
|
||||
|
||||
ge::DataType inputQType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputKvType_ = ge::DT_FLOAT16;
|
||||
ge::DataType outputType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputQRopeType_ = ge::DT_FLOAT16;
|
||||
ge::DataType inputKRopeType_ = ge::DT_FLOAT16;
|
||||
|
||||
uint64_t l2CacheSize_ = 0;
|
||||
|
||||
bool isSameSeqAllKVTensor_ = true;
|
||||
bool isSameActualseq_ = true;
|
||||
uint32_t maxActualseq_ = 0;
|
||||
|
||||
uint32_t actualLenDimsQ_ = 0;
|
||||
uint32_t actualLenDimsKV_ = 0;
|
||||
|
||||
gert::Shape queryShape_{};
|
||||
gert::Shape keyShape_{};
|
||||
gert::Shape valueShape_{};
|
||||
gert::Shape sparseIndicesShape_{};
|
||||
gert::Shape queryRopeShape_{};
|
||||
gert::Shape keyRopeShape_{};
|
||||
};
|
||||
} // namespace optiling
|
||||
#endif // SPARSE_FLASH_ATTENTION_TILING_H
|
||||
Reference in New Issue
Block a user