* Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
* MindIE is licensed under Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
* http://license.coscl.org.cn/MulanPSL2
* 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 FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
*/
#ifndef OP_API_INC_LEVEL0_FLASH_ATTN_H_
#define OP_API_INC_LEVEL0_FLASH_ATTN_H_
#include <array>
#include "opdev/op_executor.h"
namespace l0op {
* @brief FlashAttn level-0 operator。
* 封装FlashAttn算子的底层调度,完成InferShape与Kernel Launch注册。
* 该接口为内部接口,仅供aclnn层调用。
*
* @param q query tensor
* @param k key tensor
* @param v value tensor
* @param blockTableOptional 分页KV缓存块映射表(可选,INT32)
* @param cuSeqlensQOptional query累积序列长度tensor(可选,INT32)
* @param cuSeqlensKvOptional kv累积序列长度tensor(可选,INT32)
* @param sequsedQOptional query各batch实际序列长度tensor(可选,INT32)
* @param sequsedKvOptional kv各batch实际序列长度tensor(可选,INT32)
* @param sinksOptional 可学习sink权重(可选,FLOAT32)
* @param attnMaskOptional attnMask参数(可选,INT8)
* @param metadataOptional 预计算tiling元数据(可选,INT32)
* @param softmaxScale softmax缩放系数(float)
* @param maskMode 掩码模式(int64_t)
* @param winLeft 左窗口大小(int64_t)
* @param winRight 右窗口大小(int64_t)
* @param layoutQ query布局字符串
* @param layoutKv kv布局字符串
* @param layoutOut 输出布局字符串
* @param returnSoftmaxLse 是否输出softmax_lse(int64_t)
* @param deterministic 是否确定性计算(int64_t)
* @param executor op执行器
* @return std::array<const aclTensor*, 2> [attnOut, softmaxLse]
* 任意元素为nullptr表示对应输出的InferShape或Launch失败。
*/
const std::array<const aclTensor *, 2> FlashAttn(const aclTensor *q, const aclTensor *k, const aclTensor *v,
const aclTensor *blockTableOptional, const aclTensor *cuSeqlensQOptional, const aclTensor *cuSeqlensKvOptional,
const aclTensor *sequsedQOptional, const aclTensor *sequsedKvOptional, const aclTensor *sinksOptional,
const aclTensor *attnMaskOptional, const aclTensor *metadataOptional, double softmaxScale, int32_t maskMode,
int32_t winLeft, int32_t winRight, const char *layoutQ, const char *layoutKv, const char *layoutOut,
int32_t returnSoftmaxLse, int32_t deterministic, aclOpExecutor *executor);
}
#endif