FusedInferAttentionScore 融合推理注意力算子全解析：基于 CANN ops-transformer 的 PFA/IFA 统一 FlashAttention 实现

FusedInferAttentionScore 融合推理注意力算子全解析基于 CANN ops-transformer 的 PFA/IFA 统一 FlashAttention 实现【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformerFusedInferAttentionScore下文简称 FIA是 CANN ops-transformer 算子库中面向大模型推理场景的融合 FlashAttention 算子通过一套算子同时覆盖 prefill全量与 decode增量两种计算模式并衍生出 V4、V5 两代 aclnn 接口以满足 Atlas A2/A3 与 Ascend 950 系列不同产品的推理需求。本文以 FusedInferAttentionScore 官方 README 为主体骨架结合仓库内 aclnnFusedInferAttentionScoreV4 接口文档、aclnnFusedInferAttentionScoreV5 接口文档、op_api/op_host/op_kernel 源码与测试用例系统讲解其产品支持矩阵、计算原理、参数语义、约束边界与两段式调用实战帮助开发者在推理引擎vLLM 类框架或自研调度器中正确选型与调用。产品支持情况与版本选型FIA 算子并非在所有昇腾产品上都可用调用前首先需要确认目标硬件的支持情况。根据 README 中的产品支持矩阵产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品×Atlas 训练系列产品×需要特别说明的是支持的产品不同对应的 aclnn 接口版本也不同Atlas A2/A3 训练与推理系列产品通过 aclnnFusedInferAttentionScoreV4 调用 FIA 算子配套示例为 test_aclnn_fused_infer_attention_score_v4_gqa_noquant.cppAscend 950PR/Ascend 950DT通过 aclnnFusedInferAttentionScoreV5 调用 FIA 算子配套示例位于 examples/arch35/test_aclnn_fused_infer_attention_score_v5.cpp。两个接口的详细产品支持矩阵可分别参见对应接口文档。从仓库源码结构看接口与硬件的对应关系同样体现在代码组织上op_host 目录下按arch22Atlas A2/A3 架构承载 V4 tiling 逻辑、arch35Ascend 950 架构承载 V4/V5 tiling 逻辑和arch38分目录维护各自的 tiling 实现op_kernel 目录同样按arch22、arch35分目录存放 kernel 代码这印证了不同产品代际走不同接口与不同内核的事实。从源码结构看V5 是面向 950 系列的演进版本在 V4 基础上新增了qStartIdxOptional、kvStartIdxOptional、pseType等参数。功能说明一个算子融合 PFA 与 IFA 两种推理场景FIA 的核心设计目标是适配增量decode与全量prefill推理场景的 FlashAttention 算子。全量计算场景等价于 PromptFlashAttention此时 query 的序列长度 Q_S 1模型一次性并行处理 prompt/context 中所有 token增量计算场景等价于 IncreFlashAttention此时 Q_S 1模型逐一生成新 token依赖已缓存的 KV 向量完成自注意力计算。decode 场景下特有的KV Cache机制是理解 FIA 的关键背景采样时Transformer 模型以给定的 prompt/context 作为初始输入进行推理可并行处理随后逐一生成额外的 token 来继续完善生成的序列体现模型的自回归性质。在采样过程中模型对当前序列中的每个项目无论是 prompt/context 还是生成的 token都会提取键值KV向量这些向量被存储在一个矩阵中通常被称为 KV Cache。FIA 通过 key/value 的aclTensorList输入、blockTableblockSize的 PagedAttention 支持等机制直接消费推理框架维护的 KV Cache 布局避免每步 decode 重复计算历史 token 的 K/V这是大模型推理性能优化的常用技术。计算原理与公式FIA 实现的是 self-attention自注意力。其原理是假设有一个长度为 n 的输入样本序列 xx 的每个元素都是一个 d 维向量可以将每个 d 维向量看作一个 token embedding将这条序列经过 3 个权重矩阵变换得到 3 个维度为 n*d 的矩阵。其中 Q、K、V 为输入样本的重要属性元素是输入样本经过空间变换得到且可以统一到一个特征空间中。公式及算子名称中的 Attention 为 self-attention 的简写。self-attention 的计算公式一般定义如下$$ Attention(Q,K,V)Score(Q,K)V $$FIA 算子中 Score 函数采用 Softmax 函数self-attention 计算公式为$$ Attention(Q,K,V)Softmax(\frac{QK^T}{\sqrt{d}})V $$其中 Q 和 K^T 的乘积代表输入 x 的注意力为避免该值变得过大通常除以 d 的平方根进行缩放并对每行进行 softmax 归一化与 V 相乘后得到一个 n*d 的矩阵。V4 接口Atlas A2/A3 产品线的计算公式还支持在 softmax 前叠加 alibi fullmask 能力$$ Attention(Q,K,V)Softmax(\frac{QK^T}{\sqrt{d}} FullMask)V $$其中 FullMask 即通过 pseShiftalibi 位置编码在特定约束下实现的 fullmask 掩码相关约束在约束说明章节详述。版本功能演进从接口文档对比可以看出 FIA 家族接口的演进脉络V3基础版本支持 decode prefill 场景V4相对 V3新增dequantScaleQueryOptional、learnableSinkOptional、queryQuantMode参数另外新增 alibi 的 fullmask 能力V5相对 V4Ascend 950 系列新增qStartIdxOptional、kvStartIdxOptional、pseType参数并扩展了 FLOAT8_E4M3FN/HIFLOAT8/FLOAT4_E2M1 等量化数据类型支持。核心输入输出参数说明FIA 算子的核心输入输出为 query、key、value、attentionOut 四个 Tensor其数据类型与格式约束如下表源自 README参数名输入/输出描述数据类型数据格式query输入公式中的输入 QFLOAT16、BFLOAT16、INT8NDkey输入公式中的输入 KFLOAT16、BFLOAT16、INT8、INT4NDvalue输入公式中的输入 VFLOAT16、BFLOAT16、INT8、INT4NDattentionOut输出公式中的输出FLOAT16、BFLOAT16、INT8ND在 V4/V5 的 aclnn 接口层面key 与 value 均以aclTensorList形式传入支持张量列表便于非连续 KV Cache 场景query 与 attentionOut 以单个aclTensor传入。V5 在 950 系列上进一步扩展了 query 的 FLOAT8_E4M3FN/HIFLOAT8、key/value 的 FLOAT4_E2M1 等数据类型。维度解读约定query、key、value 数据排布格式支持从多种维度解读其中各维度约定为BBatch输入样本批量大小SSeq-Length输入样本序列长度HHidden-Size隐藏层的大小NHead-Num多头数DHead-Dim隐藏层最小的单元尺寸且满足 DH/NT所有 Batch 输入样本序列长度的累加和。同时约定Q_S 表示 query shape 中的 SKV_S 表示 key/value shape 中的 SQ_N 表示 num_query_headsKV_N 表示 num_key_value_headsP 表示 Softmax(QK^T/√d) 的计算结果。aclnn 两段式调用接口FIA 的算子执行接口遵循 CANN 标准的两段式接口设计见 V4/V5 接口文档即必须先调用 GetWorkspaceSize 接口获取入参并根据计算流程计算所需 workspace 大小再调用执行接口完成计算。以 V4 为例两段接口原型如下aclnnStatus aclnnFusedInferAttentionScoreV4GetWorkspaceSize( const aclTensor *query, const aclTensorList *key, const aclTensorList *value, const aclTensor *pseShiftOptional, const aclTensor *attenMaskOptional, const aclIntArray *actualSeqLengthsOptional, const aclIntArray *actualSeqLengthsKvOptional, const aclTensor *deqScale1Optional, const aclTensor *quantScale1Optional, const aclTensor *deqScale2Optional, const aclTensor *quantScale2Optional, const aclTensor *quantOffset2Optional, const aclTensor *antiquantScaleOptional, const aclTensor *antiquantOffsetOptional, const aclTensor *blockTableOptional, const aclTensor *queryPaddingSizeOptional, const aclTensor *kvPaddingSizeOptional, const aclTensor *keyAntiquantScaleOptional, const aclTensor *keyAntiquantOffsetOptional, const aclTensor *valueAntiquantScaleOptional, const aclTensor *valueAntiquantOffsetOptional, const aclTensor *keySharedPrefixOptional, const aclTensor *valueSharedPrefixOptional, const aclIntArray *actualSharedPrefixLenOptional, const aclTensor *queryRopeOptional, const aclTensor *keyRopeOptional, const aclTensor *keyRopeAntiquantScaleOptional, const aclTensor *dequantScaleQueryOptional, const aclTensor *learnableSinkOptional, int64_t numHeads, double scaleValue, int64_t preTokens, int64_t nextTokens, char *inputLayout, int64_t numKeyValueHeads, int64_t sparseMode, int64_t innerPrecise, int64_t blockSize, int64_t antiquantMode, bool softmaxLseFlag, int64_t keyAntiquantMode, int64_t valueAntiquantMode, int64_t queryQuantMode, const aclTensor *attentionOut, const aclTensor *softmaxLse, uint64_t *workspaceSize, aclOpExecutor **executor) aclnnStatus aclnnFusedInferAttentionScoreV4( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, const aclrtStream stream)V5 的函数原型与 V4 结构一致仅在learnableSinkOptional之后追加qStartIdxOptional、kvStartIdxOptional两个aclIntArray输入并在 queryQuantMode 之后追加pseType属性参数。完整原型见 aclnnFusedInferAttentionScoreV5.md。第一段接口GetWorkspaceSize完成入参校验并计算 workspace 大小出现异常时返回对应错误码返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 query、key、value、attentionOut 是空指针ACLNN_ERR_PARAM_INVALID161002query、key、value、pseShift、attenMaskOptional、attentionOut 的数据类型和数据格式不在支持的范围内ACLNN_ERR_RUNTIME_ERROR361001API 内存调用 npu runtime 的接口异常关键属性参数语义在 30 余个接口参数中以下属性参数Attribute对算子行为起决定性作用理解其语义是正确调用的前提以下内容综合 V4/V5 接口文档整理numHeads / numKeyValueHeadsquery 与 key/value 的 head 个数。需要满足 numHeads 整除 numKeyValueHeads即支持 GQA分组查询注意力/MQA多查询注意力配置。V5 文档建议不特意指定时 numKeyValueHeads 传入 0表示与 query head 数相等。V4 中在 BSND、TND、BNSD 等带 N 轴排布场景下还需要与 shape 中 key/value 的 N 轴值相同。scaleValue公式中根号 d 的倒数即缩放系数数据类型 DOUBLE。V5 建议用户不特意指定时传入 1.0调用示例中通常传 1/sqrt(headDims)。preTokens / nextTokens用于稀疏计算表示 attention 需要和前/后几个 Token 计算关联。V5 建议不特意指定时传入 2147483647即不做稀疏裁剪。inputLayout标识输入 query、key、value 的数据排布格式。当字段包含 _ 时表示输入 layout_输出 layout。V4 支持 BSH、BSND、TND、BNSD、NTD、BSH_BNSD、BSND_BNSD、BNSD_BSND、NTD_TND、BSH_NBSD、BSND_NBSD、BNSD_NBSD 等V5 在此基础上扩展了 TND_NTD 等。V5 建议不特意指定时传入 BSH。sparseMode稀疏 mask 的模式取值 0~9其中 5/6/7/8prefix/global/dilated/block_local在推理接口中均不支持。V5 建议不特意指定时传入 0defaultMask 模式。innerPrecise高精度/高性能及是否进行行无效修正的选择0高精度模式不做行无效修正V5 中为保留配置值1高性能模式不做行无效修正V5 中为保留配置值2高精度模式做行无效修正3高性能模式做行无效修正。注意BFLOAT16 和 INT8 不区分高精度与高性能但行无效修正对 FLOAT16、BFLOAT16 和 INT8 均生效。当计算过程中参与计算的 mask 部分存在某整行全为 1 的情况时精度可能会有损失此时可将参数配置为 2 或 3 开启行无效功能提升精度代价是性能下降。如果算子可判断出存在无效行场景如 sparseMode 为 3 且 Sq Skv会自动开启无效行计算。blockSizePagedAttention 中 KV 存储每个 block 中最大的 token 个数。不传时按 0 处理即不启用 PagedAttention。softmaxLseFlag是否输出 softmaxLse支持 S 轴外切增加输出V5 建议默认传入 false。antiquantMode / keyAntiquantMode / valueAntiquantMode / queryQuantModeKV 伪量化/反量化模式编号。以 V5 为例0 为 per-channel含 per-tensor1 为 per-token2 为 per-tensor 叠加 per-head3 为 per-token 叠加 per-head4/5 为结合 PagedAttention 管理 scale/offset 的 per-token叠加 per-head模式6 为 per-token-groupvalue 侧另有 8per-channel-group。可选输入参数组V5 视角V5 接口文档将全部入参划分为多个特性参数组这一组织方式既是文档结构也与 op_host 目录下 checkers 子目录的校验器一一对应。从源码目录 op_host/checkers 可以看出仓库为每个特性参数组维护了独立的校验器fia_checker.cpp总入口、common_checker.cpp、mask_checker.cpp、pse_checker.cpp、rope_checker.cpp、dequant_checker.cpp、post_quant_checker.cpp、paged_attention_checker.cpp、left_padding_checker.cpp、system_prefix_checker.cpp、actual_seq_len_checker.cpp、softmax_lse_checker.cpp、learnable_sink_checker.cpp等印证了约束按参数组分层校验的实现事实。各参数组的核心语义如下公共参数组query/key/value/numHeads/scaleValue/inputLayout/numKeyValueHeads/innerPrecise/attentionOutPSE 参数组pseShift、qStartIdxOptional、kvStartIdxOptional、pseTypepseType 为 0 时外部传入 pse先 mul 再 addpseType1 仅支持训练场景推理场景不支持Mask 参数组attenMaskOptional、preTokens、nextTokens、sparseModeActualSeqLens 参数组actualSeqLengthsOptional、actualSeqLengthsKvOptionalINT64shape 为 (1)/(B)/(B)伪量化全量化参数组deqScale1/quantScale1/deqScale2/antiquantScale/antiquantOffset/keyAntiquantScale/keyAntiquantOffset/valueAntiquantScale/valueAntiquantOffset/dequantScaleQuery/antiquantMode/keyAntiquantMode/valueAntiquantMode/queryQuantModePostQuant 参数组quantScale2Optional、quantOffset2Optional输出后量化支持对称/非对称Paged Attention 参数组blockTableOptional、blockSizeLeftPadding 参数组queryPaddingSizeOptional、kvPaddingSizeOptionalSystemPrefix 参数组keySharedPrefixOptional、valueSharedPrefixOptional、actualSharedPrefixLenOptional系统公共前缀 KV 复用用于多轮对话场景Rope 参数组queryRopeOptional、keyRopeOptional、keyRopeAntiquantScaleOptionalMLA 结构专用queryRope/keyRope 的 shape 中 d 为 64其余维度与 query/key 一致keyRopeAntiquantScale 为预留参数当前版本不生效LearnableSink 参数组learnableSinkOptional通过可学习的 Sink Token 吸收 Attention Scoreshape 为 (Q_N)仅支持非量化场景与 V_D128/64SoftmaxLSE 参数组softmaxLseFlag、softmaxLse。约束说明空入参与 Shape 边界空入参处理该接口与 PyTorch 配合使用时需要保证 CANN 相关包与 PyTorch 相关包的版本匹配。算子内部需要判断参数 query 是否为空如果是空则直接返回。参数 query 不为空 Tensor、参数 key/value 为空 tensor即 S2 为 0则 attentionOut 填充为全零。attentionOut 为空 Tensor 时AscendCLNN 框架会处理。V4/V5 接口文档进一步细化了空 Tensor 语义空 Tensor 指必选输入和输出的 shapeSize 为 0。在空 Tensor 场景下若 attentionOut 为空则返回空否则返回全 0若有 lse 且 lse 为空时返回空lse 不为空则返回全 inf。属于空 Tensor 时跳过校验流程否则走正常校验流程。key/value 的 TensorList 约束参数 key、value 中对应 Tensor 的 shape 需要完全一致非连续场景下 key、value 的 tensorlist 中的 batch 只能为 1个数等于 query 的 BN 和 D 需要相等由于 tensorlist 限制非连续场景下 B 不能大于 256。Q_S 1prefill 场景的功能使用限制当 Q_S 大于 1 时query、key、value 输入的功能使用限制如下支持 B 轴小于等于 65536如果输入类型为 INT8 且 D 轴不是 32 字节对齐则 B 轴的最大支持值为 128若输入类型为 FLOAT16 或 BFLOAT16 且 D 轴不是 16 字节对齐B 轴同样仅支持到 128支持 N 轴小于等于 256支持 D 轴小于等于 512。inputLayout 为 BSH 或者 BSND 时建议 N*D 小于 65535S 支持小于等于 2097152020M。部分长序列场景下如果计算量过大可能会导致算子执行超时aicore error 类型报错errorStr 为timeout or trap error此场景下建议做 S 切分处理。注这里计算量会受 B、S、N、D 等的影响值越大计算量越大。典型的会超时的长序列即 B、S、N、D 的乘积较大场景包括但不限于BQ_NQ_SDKV_NKV_S1202097152256120971521220971520256220971520201209715225612097152110209715251212097152D 轴限制query、key、value 或 attentionOut 类型包含 INT8 时D 轴需要 32 对齐类型包含 INT4 时D 轴需要 64 对齐类型全为 FLOAT16、BFLOAT16 时D 轴需 16 对齐。Q_S 1decode 场景的功能使用限制当 Q_S 等于 1 时query、key、value 输入的功能使用限制如下支持 B 轴小于等于 65536支持 N 轴小于等于 256支持 D 轴小于等于 512query、key、value 输入类型均为 INT8 的场景暂不支持在 INT4 伪量化场景下aclnn 单算子调用支持 KV INT4 输入或者 INT4 拼接成 INT32 输入建议通过 dynamicQuant 生成 INT4 格式的数据因为 dynamicQuant 就是一个 INT32 包括 8 个 INT4在 INT4 伪量化场景下若 KV INT4 拼接成 INT32 输入那么 KV 的 N、D 或者 H 是实际值的八分之一prefix 同理key、value 在特定数据类型下存在对 D 轴的限制key、value 输入类型为 INT4INT32时D 轴需要 64 对齐INT32 仅支持 D 8 对齐。常用可选特性的约束要点综合 V4/V5 接口文档以下约束直接影响工程落地值得重点记忆PagedAttention开启必要条件是 blockTable 存在且有效同时 key、value 按照 blockTable 中的索引在一片连续内存中排布该场景下 key、value 的 inputLayout 参数无效。blockSize 需传入非 0 值A2 上非量化场景一般要求 16 对齐、最大 1024。PagedAttention 不支持 tensorlist 场景不支持左 padding 场景且必须传入 actualSeqLengthsKv。A2 上支持的 KV Cache layout 有 BnBsHBlockNum, BlockSize, H、BnNBsDBlockNum, N, BlockSize, D、NZBlockNum, N, D/16, BlockSize, 16三种BnNBsD 排布性能通常优于 BnBsH建议优先选择。当 KV Cache 排布为 BnBsH 且 KV_N*D 超过 65535 时受硬件指令约束会被拦截可通过开启 GQA减小 KV_N或改为 BnNBsD 排布解决。左 paddingquery 左 padding 场景下 query 的搬运起点计算公式为 Q_S - queryPaddingSize - actualSeqLengths搬运终点为 Q_S - queryPaddingSize且需与 actualSeqLengths 参数一起开启否则默认为右 padding仅支持 Q_S1 的场景且不支持 PagedAttention 与 tensorlist。系统前缀prefixkeySharedPrefix 与 valueSharedPrefix 要么都为空、要么都不为空shape 第一维 batch 必须为 1且公共前缀的 S 加上 key/value 的 S 的结果需满足 key/value 的 S 的限制不支持 PagedAttention、左 padding、tensorlist 场景。MLA 场景queryRope/keyRope 不为空queryRope、keyRope 要求同时配置或同时不配置不支持只配置其中一个query d512 的 Decode MLA 场景要求 K_N1key/value 支持五维 shape [blockNum, N, D/16, blockSize, 16]全量化为 INT8 时为 [blockNum, N, D/32, blockSize, 32]Decode MLA 场景不支持 tensorlist、左 padding、伪量化、prefix。mask 的 sparseMode 取值sparseMode0 为 defaultMask未传 attenMask 时不做 mask 操作忽略 preTokens/nextTokens传入时需传完整矩阵1 为 allMask必须传入完整矩阵2/3/4 分别为 leftUpCausal/rightDownCausal/band 模式需要传入优化后的 2048×2048 下三角矩阵对角线全 09 为 treeMask推测解码场景的树形注意力掩码非量化支持 GQA 和 MLA全量化仅支持 MLA输入 shape 为 (B,Q_S,Q_S) 或 TND 下的 (∑Q_Si²,)。TND 场景actualSeqLengths 与 actualSeqLengthsKv 必须传入query d512 时仅支持 TND、TND_NTD 且 numKeyValueHeads 必须为 1、sparseMode 仅支持 0/3/4/9不支持 SoftMaxLse、左 padding、tensorlist、pse、prefix、伪量化、全量化、后量化。调用实战以 aclnnFusedInferAttentionScoreV4 为例V4 接口的完整调用示例见 test_aclnn_fused_infer_attention_score_v4_gqa_noquant.cpp接口文档 aclnnFusedInferAttentionScoreV4.md 中也给出了可编译运行的完整代码。以下为核心调用骨架省略头文件与辅助函数int main() { // 1. 固定写法device/stream 初始化 int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); // 内部为 aclInit/aclrtSetDevice/aclrtCreateStream // 2. 构造输入输出GQA 配置numHeads2, numKeyValueHeads1decode 场景 Q_S1 int64_t batchSize 1; int64_t numHeads 2; int64_t numKeyValueHeads 1; int64_t sequenceLengthQ 1; int64_t sequenceLengthKV 16; int64_t headDims 128; std::vectorint64_t queryShape {batchSize, numHeads, sequenceLengthQ, headDims}; // BNSD std::vectorint64_t keyShape {batchSize, numKeyValueHeads, sequenceLengthKV, headDims}; std::vectorint64_t valueShape {batchSize, numKeyValueHeads, sequenceLengthKV, headDims}; std::vectorint64_t attentionOutShape {batchSize, numHeads, sequenceLengthQ, headDims}; // 通过 aclrtMalloc aclrtMemcpy 将 host 数据搬运到 device再 aclCreateTensor 创建 aclTensor aclTensor *queryTensor nullptr, *keyTensor nullptr, *valueTensor nullptr, *attentionOutTensor nullptr; // ... CreateAclTensor(...) 创建各 tensor ... // key/value 以 TensorList 传入 int kvTensorNum 1; aclTensor *tensorsOfKey[kvTensorNum] {keyTensor}; auto tensorKeyList aclCreateTensorList(tensorsOfKey, kvTensorNum); aclTensor *tensorsOfValue[kvTensorNum] {valueTensor}; auto tensorValueList aclCreateTensorList(tensorsOfValue, kvTensorNum); // 3. 构造属性参数 double scaleValue 1 / sqrt(headDims); // 缩放系数 int64_t preTokens 2147483647; // 不做稀疏裁剪 int64_t nextTokens 2147483647; string sInputLayout BNSD; // 输入排布 char inputLayout[sInputLayout.length() 1]; strcpy(inputLayout, sInputLayout.c_str()); int64_t sparseMode 0; // defaultMask不传 attenMask 则不做 mask int64_t innerPrecise 0; // 高精度、不做行无效修正 int64_t blockSize 0; // 不启用 PagedAttention int64_t antiquantMode 0; bool softmaxLseFlag false; int keyAntiquantMode 0; int64_t valueAntiquantMode 0; int64_t queryQuantMode 0; // 4. 两段式调用第一段获取 workspaceSize uint64_t workspaceSize 0; aclOpExecutor *executor; ret aclnnFusedInferAttentionScoreV4GetWorkspaceSize( queryTensor, tensorKeyList, tensorValueList, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, numHeads, scaleValue, preTokens, nextTokens, inputLayout, numKeyValueHeads, sparseMode, innerPrecise, blockSize, antiquantMode, softmaxLseFlag, keyAntiquantMode, valueAntiquantMode, queryQuantMode, attentionOutTensor, nullptr, workspaceSize, executor); // 5. 根据 workspaceSize 申请 device 内存 void *workspaceAddr nullptr; if (workspaceSize 0U) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 6. 第二段调用执行计算 ret aclnnFusedInferAttentionScoreV4(workspaceAddr, workspaceSize, executor, stream); // 7. 同步等待执行结束 ret aclrtSynchronizeStream(stream); // 8. 将结果从 device 拷贝回 host 并校验、释放资源 // aclrtMemcpy(DEVICE_TO_HOST) - 打印 result - aclDestroyTensor/aclrtFree/aclrtDestroyStream/aclFinalize }调用要点总结未启用的可选功能一律传nullptr例如不传 attenMask、pseShift、blockTable 时对应位置传空指针key/value 必须构造为aclTensorList即使只有一个 tensor也需要用aclCreateTensorList包装workspace 内存由第一段接口计算出大小后由用户自行aclrtMalloc申请用完需aclrtFree释放结果需要通过aclrtSynchronizeStream同步后再以aclrtMemcpyDEVICE_TO_HOST拷回验证。Ascend 950 系列上的 V5 调用方式与之相同两段式仅需替换头文件aclnn_fused_infer_attention_score_v5.h与接口名并注意 V5 中 queryQuantMode 之后多出的pseType参数位置参考 examples/arch35/test_aclnn_fused_infer_attention_score_v5.cpp。V5 示例还提供了 full_quant、antiquant、antiquant_mxfp4 等扩展场景的独立样例位于 examples/arch35 目录。源码级实现纵深从 aclnn 到 kernelop_api 层接口分发FIA 的对外 aclnn 接口实现在 op_api 目录按版本拆分为aclnn_fused_infer_attention_score.cppV1、aclnn_fused_infer_attention_score_v2.cpp、aclnn_fused_infer_attention_score_v3.cpp、aclnn_fused_infer_attention_score_v4.cpp、aclnn_fused_infer_attention_score_v5.cpp以及公共的fused_infer_attention_score_base_aclnn.cpp与fused_infer_attention_score_inner.h。从源码结构看各版本接口通过共享基础实现base_aclnn来复用参数校验、workspace 计算与 executor 构造流程版本差异集中在参数列表与特性开关上。op_host 层tiling 与校验op_host 承载算子编排tiling与入参校验逻辑tiling 实现fused_infer_attention_score_tiling.cpp 为总入口按架构分目录维护arch22 下含fia_tiling_nonquant.cpp、fia_tiling_nonquant_mla.cpp、fia_tiling_empty_tensor.cpp等Atlas A2/A3 的 V3/V4 tilingarch35 下含fia_tiling_nonquant_gqa.cpp、fia_tiling_nonquant_mla.cpp、fia_tiling_fullquant_gqa.cpp、fia_tiling_fullquant_mla.cpp、fia_tiling_fullquant_mx.cpp等Ascend 950 的 V4/V5 tiling覆盖全量化与 MxFP8 场景arch38 下为新一代架构的 tiling 入口。fia_tiling_info.cpp与fused_infer_attention_score_tiling_info_parser.cpp负责 tiling 数据的组织与解析fused_infer_attention_score_tiling_cache.h提供 tiling 缓存能力。入参校验checkersop_host/checkers 目录按特性参数组实现了分层校验器fia_checker 为总入口与 V5 文档中的特性参数组章节一一对应包括 mask、pse、rope、dequant、post_quant、paged_attention、left_padding、system_prefix、actual_seq_len、softmax_lse、learnable_sink 等实现单参数约束 → 存在性约束 → 一致性约束 → 特性交叉约束四档校验体系。op_kernel 层非量化/全量化/MLA 多内核op_kernel 按架构与量化模式组织内核实现kernel 侧文件为 .cpp/.h/.hpparch22fused_infer_attention_score_v3.cppAtlas A2/A3 的 V3 内核及配套的fia_kernel_nonquant.h、fia_kernel_nonquant_mla.h、fia_kernel_empty_tensor.h、fia_block_cube_nonquant.h、fia_block_vec_nonquant.h等arch35fia_kernel_noquant_gqa.h、fia_kernel_noquant_mla.h、fia_kernel_fullquant_gqa.h、fia_kernel_fullquant_mla.h、fia_kernel_fullquant_mx.h、fia_kernel_antiquant.h以及 flash_attention 系列的 block_cube/block_vec 模板头文件实现 GQA 非量化、MLAPrefill/Decode、INT8 全量化、MxFP8/反量化等多条内核路径顶层 fused_infer_attention_score.cpp 与fused_infer_attention_score_apt.cpp负责内核入口与调度。从源码结构看FIA 的内核设计围绕QK^T 矩阵乘BMM1→ Softmax/掩码 → PV 矩阵乘BMM2的 FlashAttention 流水展开fia_block_cube_*承担 Cube矩阵乘侧的 QK/PV 分块计算fia_block_vec_*承担 Vector向量侧的 online softmax、mask 应用与结果搬运fia_block_vec_flashdecode_*专门针对 decodeQ_S1场景做了优化这与 FIA同时适配 prefill 与 decode的算子定位相吻合。测试体系仓库为 FIA 提供了多层次的测试验证UT单元测试tests/ut 下按 arch22/arch35 维护 tiling 单测test_fused_infer_attention_score_tiling.cpp及配套 csv 参数表以及 dtype/shape 的 infershape 单测ST系统测试tests/st/aclnnFusedInferAttentionScoreV4 提供 ATK 测试描述json与 executor 脚本pytest 精度/性能测试tests/pytest 下按 fia_fullquant_gqa_test、fia_fullquant_mla_test、fia_fullquant_mxfp8_test 等场景组织 golden 计算、结果比对与参数集debug/func_rdv/perf_rdvtiny_rdv 用例集tests/tiny_rdv/case_json 沉淀了大量真实模型场景的 JSON 用例从用例命名可以看到覆盖范围例如 llama2-70b 的 GQAPAPagedAttention用例、pangu_xiaoyi_38B 的 actq/actkv伪量化用例、kaiyuanSD15 的 BNSD 长序列用例、FIA Decode MLA 用例、VIT 的 qkv 非对齐用例等是理解算子支持矩阵的宝贵参考。总结与选型建议FusedInferAttentionScore 是 CANN ops-transformer 中面向大模型推理的核心融合注意力算子其价值在于一套算子两用同时覆盖 prefillPromptFlashAttention与 decodeIncreFlashAttention两类推理场景简化推理引擎的算子栈特性齐备支持 GQA/MQA/MLA、PagedAttention 分页 KV Cache、INT8/INT4/MxFP8 等量化与反量化、alibi/pse 位置编码、稀疏 mask含推测解码 treeMask、系统前缀复用、RingAttention 所需的 softmaxLse 输出等推理侧高频特性接口随硬件演进Atlas A2/A3 使用 aclnnFusedInferAttentionScoreV4Ascend 950 系列使用 aclnnFusedInferAttentionScoreV5工程接入时应按目标硬件选择对应接口与头文件。工程落地时的三条关键建议先核对产品支持矩阵与接口文档中的约束说明尤其是空入参处理、D 轴对齐FP16 16 对齐 / INT8 32 对齐 / INT4 64 对齐、B 轴上限65536非对齐时收缩到 128与长序列超时S 上限 20M超算力场景建议 S 切分这几类高频坑位decode 场景优先考虑 PagedAttention配置 blockTable blockSize actualSeqLengthsKvKV Cache 优先选用 BnNBsD 排布以获得更优性能需要 softmaxLse如 RingAttention 或 MLA 训练反向时注意 softmaxLseFlag 与 softmaxLse 的配套使用规则并确认目标产品与 layout 下支持该特性。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考