ATB SelfAttention 融合算子深度解析：知识条目、参数体系与 Runner 分发机制

ATB SelfAttention 融合算子深度解析知识条目、参数体系与 Runner 分发机制【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本文以 ATBAscend Transformer Boost知识库中 self_attention 知识条目为主体结合仓库内路由文件、算子源码与示例系统梳理 SelfAttention 融合算子的文件结构、参数体系、平台变体、Runner 分发决策链与形状推导规则。读者读完可掌握该算子在 src/ops/ops_infer/self_attention/ 下的代码组织方式、SelfAttentionParam 每个配置项的含义与约束、ACLNN/OpsRunner 双通道的执行路径以及如何基于 example/op_demo/self_attention/ 示例快速搭建 Encoder / PA Encoder / Prefix Encoder 场景的调用代码。1. 知识条目self_attention 的入口与索引在.agent/knowledge/知识库中self_attention 的知识条目位于 .agent/knowledge/ops/attention/self_attention/index.md状态为complete属于 attention 类别、Tier L 的单算子single条目。该条目自身是一份路由 元数据性质的索引其核心信息如下条目内容算子名称self_attention类别 / 级别attention / tier L / 单算子源码路径src/ops/ops_infer/self_attention/路由文件.agent/knowledge/routing/self_attention.mdRunner 类型ACLNN Runner OpsRunner详见路由文件PipelineBMM1 → Softmax → BMM2详见深度条目从知识条目可以得到该算子最核心的三个技术事实执行通道是双轨的SelfAttention 同时支持 ACLNN Runner走aclnnFusedInferAttentionScoreV5等 CANN 高层算子接口与原生 OpsRunner直接驱动融合算子内核具体选择由运行平台与参数组合决定。计算管线是经典注意力三阶段Q·K^TBMM1→ Softmax → ·VBMM2整个流程被融合进单个算子在昇腾处理器上执行。实现规模与复杂度路由文件标注该算子分类为 infer、复杂度 XL、共 27 个源文件属于推理侧最复杂的融合算子之一并存在 910A、BNSD、Bypass、Encoder、Fusion、PrefixEncoder 六类平台/功能变体。2. 文件清单与推荐阅读顺序路由文件 .agent/knowledge/routing/self_attention.md 给出了完整的 27 个文件清单按角色分为三类Operation 定义算子对外行为、ACLNN RunnerACLNN 通道、Ops Runner原生内核通道含各平台变体另有param.cpp/.h、self_attention_runner_utils.cpp/.h等辅助文件。按路由文件推荐的阅读顺序可以按接口 → 决策 → 实现三层递进先读 self_attention_operation.h了解输入输出数量、InferShapeImpl、CreateRunner等核心接口签名以及MlaParamCheck/SWAParamCheck/BNSDParamCheck/PrefixEncoderParamCheck等参数检查入口再读 self_attention_operation.cpp重点看CreateOperation()的参数校验链与CreateRunner()的 Runner 分发决策逻辑随后按平台深入各 Runner910B 侧读self_attention_encoder_fusion_ops_runner、self_attention_fusion_bypass_ops_runner含 BNSD 变体、self_attention_prefix_encoder_ops_runner910A 侧读对应的_910a后缀实现950 侧读 self_attention_aclnn_runner.h 与 self_attention_aclnn_runner.cpp。3. SelfAttentionParam 参数体系全解所有配置都收敛在 include/atb/infer_op_params.h 的atb::infer::SelfAttentionParam结构体中。理解这一结构体是使用该算子的前提下面按枚举与字段两条线展开。3.1 关键枚举枚举取值含义CalcTypeUNDEFINED/ENCODER/DECODER/PA_ENCODER/PREFIX_ENCODER计算场景FA 的编码器/解码器、PagedAttention 编码器、Prefix Encoder 前缀融合场景KernelTypeKERNELTYPE_DEFAULT/KERNELTYPE_HIGH_PRECISION/KERNELTYPE_EXP_M8V2内核精度默认 fp16 全链路高精度BMM 用 fp32 累加310P 上的 EXP M8V2 专用内核MaskTypeMASK_TYPE_UNDEFINED/NORM/ALIBI/NORM_COMPRESS/ALIBI_COMPRESS/ALIBI_COMPRESS_SQRT/ALIBI_COMPRESS_LEFT_ALIGN/SLIDING_WINDOW_NORM/SLIDING_WINDOW_COMPRESS/CAUSAL_MASKmask 类型全 0 mask、倒三角、alibi、各类压缩 mask、SWA 滑动窗口 mask、内部生成因果 maskKvCacheCfgK_CACHE_V_CACHE/K_BYPASS_V_BYPASS是否走 KVCache 管理默认缓存处理bypass 表示直接传入 KV不维护 cacheScaleTypeSCALE_TYPE_TOR/SCALE_TYPE_LOGN是否启用 LogN 缩放QuantTypeTYPE_QUANT_UNQUANT/TYPE_DEQUANT_FUSION/TYPE_QUANT_QKV_OFFLINE/TYPE_QUANT_QKV_ONLINE量化模式不量化反量化融合预留当前不可取离线/在线 INT8 QKV 量化ClampTypeCLAMP_TYPE_UNDEFINED/CLAMP_TYPE_MIN_MAX是否对注意力分数做 min/max clampCacheTypeCACHE_TYPE_NORM/CACHE_TYPE_SWAcache 排布正常 cacheSWA 固定长度 cache只存后windowSize个 token3.2 核心字段与默认值字段默认值说明headNum0需 0query 头数kvHeadNum0KV 头数为 0 时与headNum一致非 0 时要求headNum % kvHeadNum 0qScale1.0query 缩放系数qkScale1.0在 Q·K^T 之后乘的缩放值tor 值batchRunStatusEnablefalse是否开启动态 batchisTriuMask0倒三角优化开关仅 mask 为倒三角时可开启calcType/kernelType/clampType/maskType/kvcacheCfg/scaleType/cacheType各枚举首值计算/内核/mask/cache/缩放配置clampMin/clampMax0clamp 上下界inputLayoutTYPE_BSND数据排布支持 BSND 与 BNSDmlaVHeadSize0大于 0 时开启 MLA 合并 KVCache 功能表示合并传入时 V 的 head_size取值范围 [0, 576]windowSize0大于 0 时开启 SWA 特性表示滑动窗口大小此时 maskType 须为SLIDING_WINDOW_NORM或SLIDING_WINDOW_COMPRESSquantType/outDataTypeTYPE_QUANT_UNQUANT/ACL_DT_UNDEFINED量化类型与输出类型QKV 量化时 outDataType 只能取ACL_FLOAT16或ACL_BF16rsv[64]全 0预留字段3.3 参数之间的硬性约束源码实证CreateOperation()在 self_attention_operation.cpp 中按顺序执行HeadNumCheck→MlaParamCheck→SWAParamCheck→DeviceParamCheck→PrefixEncoderParamCheck→ExpM8v2ParamCheck等十余项校验以下约束均可在源码中直接验证kvcacheCfg只允许K_CACHE_V_CACHE或K_BYPASS_V_BYPASS且calcType PA_ENCODER时不能为 bypassL72-L76、L110-L114quantType不能取TYPE_DEQUANT_FUSION预留类型QKV 量化离线/在线只支持PA_ENCODER、SCALE_TYPE_TOR、BSND 排布且输出必须是 fp16/bf16L91-L104开启 SWA 后不支持动态 batch、高精度内核、clamp、QKV 量化、LogN 与 BNSDSWAParamCheckBNSD 排布与scaleType非 TOR、量化、Prefix Encoder 互斥BNSDParamCheckMLA 模式mlaVHeadSize 0只支持 910B 平台与PA_ENCODER且mlaVHeadSize ≤ 576MlaParamCheckMASK_TYPE_ALIBI_COMPRESS*系列仅PA_ENCODER/PREFIX_ENCODER可用MASK_TYPE_UNDEFINED时isTriuMask必须为 0L115-L126。4. 平台与功能变体路由文件将实现拆分为六类变体全部体现在 src/ops/ops_infer/self_attention/ 的文件命名中变体含义代表文件910A昇腾 910AAtlas 800 训练产品平台适配self_attention_fusion_ops_runner_910a.cpp、self_attention_fusion_bypass_ops_runner_910a.cppBNSDBNSD 数据排布适配self_attention_fusion_bypass_ops_runner_BNSD.cpp含_910a变体BypassKVCache Bypass 路径直接传入 KVself_attention_fusion_bypass_ops_runner.cppEncoderEncoder 计算路径self_attention_encoder_fusion_ops_runner.cpp、atb_acl_self_attention_prefix_encoder.cppFusion融合算子路径self_attention_encoder_fusion_ops_runner.cpp含_910a变体PrefixEncoderPrefix Encoder 融合路径self_attention_prefix_encoder_ops_runner.cpp/.h这些变体不是并列的独立实现而是同一算子在平台 × 计算场景 × cache 模式三维组合下的特化 Runner。例如 910B 上按calcType与kvcacheCfg可组合出EncoderFusionOpsRunner、PrefixEncoderOpsRunner、FusionBypassOpsRunner(BNSD)、FusionOpsRunner四种 Runner而 910A 上则全部落到_910a后缀的实现中。5. Runner 分发决策链从参数到执行器5.1 CreateRunner 的完整决策逻辑Runner 的创建集中在 CreateRunner()决策顺序清晰可读平台 ASCEND_950 └─ SelfAttentionAclnnRunnerACLNN 通道 平台 910B ├─ calcType PA_ENCODER - SelfAttentionEncoderFusionOpsRunner ├─ calcType PREFIX_ENCODER - SelfAttentionPrefixEncoderOpsRunner ├─ kvcacheCfg K_BYPASS_V_BYPASS │ ├─ inputLayout BNSD - SelfAttentionFusionBypassOpsRunnerBNSD │ └─ 否则 - SelfAttentionFusionBypassOpsRunner └─ 否则 - SelfAttentionFusionOpsRunner 其他910A / 310P 等 ├─ PA_ENCODER - SelfAttentionEncoderFusionOpsRunner910A经 RunnerPool ├─ Bypass BNSD - SelfAttentionFusionBypassOpsRunnerBNSD910A经 RunnerPool ├─ Bypass - SelfAttentionFusionBypassOpsRunner910A经 RunnerPool └─ 默认 - SelfAttentionFusionOpsRunner910A经 RunnerPool值得注意的实现细节在非 910B 平台上Runner 通过RunnerTypeRegister::GetRunnerTypeIdxRunnerPool::MallocRunner从池中复用并注册了析构回调归还 Runner避免频繁构造/析构的开销池分配失败时才回退到std::make_shared直接创建L2113-L2140。5.2 ACLNN 通道aclnnFusedInferAttentionScoreV5 封装950 平台走 SelfAttentionAclnnRunner。该类封装了aclnnFusedInferAttentionScoreV5GetWorkspaceSize与aclnnFusedInferAttentionScoreV5两个函数指针通过LoadMethod()动态加载。从函数签名可以看到它对接了完整的 FusedInferAttentionScoreV5 能力面query/key/value、pseShift、attenMask、actualSeqLengths(KV)、反量化/量化 scale 与 offset、blockTable、共享前缀、queryRope、softmaxLse 等一应俱全并以numHeads、scaleValue、preTokens、nextTokens、inputLayout默认TND、numKeyValueHeads、sparseMode、innerPrecise、blockSize等标量控制计算语义self_attention_aclnn_runner.h#L18-L57。5.3 执行流程Setup → Execute无论哪条通道对外暴露的都是统一的 Operation 生命周期。以 Encoder demo 为例self_attention_encoder_demo.cppatb::infer::SelfAttentionParam opParam; opParam.calcType atb::infer::SelfAttentionParam::CalcType::ENCODER; // FA Encoder 场景 opParam.maskType atb::infer::SelfAttentionParam::MaskType::MASK_TYPE_NORM; // 倒三角全量 mask atb::CreateOperation(opParam, encoderOp); // 参数校验 创建 Operation ... encoderOp-Setup(variantPack, workspaceSize, context); // 推理形状、申请 workspace encoderOp-Execute(variantPack, workspacePtr, workspaceSize, context); // 异步下发执行Pipeline 即知识条目声明的 BMM1Q·K^T→ Softmax含 mask、scale→ BMM2·V整个流程由融合内核在单算子内完成避免中间张量落回 HBM。6. InferShape 与形状推导规则InferShapeImpl同样按平台分派L881-L902其核心规则可从源码归纳输出形状非 PA_ENCODER 场景输出 1 个张量。输入 Q 为 4 维[B, S, N, D]时输出被合并为 3 维[B, S, N*D]Q 为 2 维[nTokens, hiddenSize]时输出形状与 Q 一致但最后一维改写为headNum * vHeadSizeInferShapeImpl910B。PA_ENCODER 输出 1 个张量其形状由输入 Q 复制并按mlaVHeadSize或 V 的 head_size 改写最后一维。KVCache 一致性K/V cache 除最后一维外各维必须一致ND 格式NZ 格式则要求整体一致InferShapeDimCheck。batch 一致性tokenOffset、seqLen以及开启动态 batch 时的 batchStatus的第一维必须与 KVCache 的 batch 维一致。headSize 上限910B 上常规场景 headSize ≤ 256MLA 内核场景放宽到 1024、MLA bypass 场景为 576压缩 Alibi mask 场景收紧到 128MaxHeadSizeCheck910B。310P 上要求 headSize 为 16 的倍数且 ≤ 256。mask 形状NORM_COMPRESS压缩 mask 在 910B/950 上为[128, 128]950 为[2048, 2048]310P 的 NZ 格式为[1, 8, 128, 16]或长窗口[1, 128, 2048, 16]NormMaskDimCheckSWA 压缩 mask 在 910B 为[512, 512]310P 为[1, 32, 512, 16]SWAMaskDimCheck。此外输入输出数量由GetInputNum()/GetOutputNum()按参数组合动态计算有 KV 时为 8 个基础输入、bypass 时为 6 个随后按 mask、slopes、LogN、动态 batch、QKV 量化在线量化 4、离线量化 5依次累加GetInputNum。7. 使用示例三种典型场景仓库在 example/op_demo/self_attention/ 提供了多套可直接参考的 demo对应 READMEREADME.md、README_en.md中给出的参数模板Demo 文件calcTypemaskType场景self_attention_encoder_demo.cppENCODERMASK_TYPE_NORMFA Encoder 全量注意力self_attention_encoder_inference_demo.cppENCODERMASK_TYPE_UNDEFINED免 mask 推理self_attention_pa_encoder_demo.cppPA_ENCODERMASK_TYPE_NORMPagedAttention 编码器self_attention_pa_encoder_qwen_demo.cppPA_ENCODER—Qwen 系模型 PA 场景self_attention_prefix_encoder_demo.cppPREFIX_ENCODER—前缀缓存融合编码PA_ENCODER 与 PREFIX_ENCODER 场景下输入除 Q 外还包含 blockTables、seqLen、kvSeqLenprefix 场景为 8 个输入query、key、value、blockTables、mask、qSeqLen、kvSeqLen、slopestokenOffset 与 seqLen 既支持[batch]一维形式也支持[2, batch]形式第 0 维存放 offset 与长度。8. 总结与检索要点围绕 self_attention 知识条目本文完成了从知识索引 → 路由清单 → 参数体系 → 变体矩阵 → Runner 决策 → 形状推导 → 示例代码的全链路梳理。为便于后续检索与引用将关键结论浓缩如下入口知识条目 .agent/knowledge/ops/attention/self_attention/index.md → 路由文件 .agent/knowledge/routing/self_attention.md → 源码 src/ops/ops_infer/self_attention/执行通道950 走 ACLNNaclnnFusedInferAttentionScoreV5910B 与 910A/310P 走特化的 OpsRunner统一经CreateRunner()分发PipelineBMM1 → Softmax → BMM2 单算子融合输出形状由 InferShape 按平台与参数动态推导配置入口include/atb/infer_op_params.h 的SelfAttentionParam所有枚举与字段约束均可在 self_attention_operation.cpp 的十余个*ParamCheck函数中追溯验证快速上手example/op_demo/self_attention/ 下的 Encoder / PA Encoder / Prefix Encoder demo 覆盖了最主流的三种使用场景。说明以上内容基于当前仓库ascend-transformer-boost源码与知识库文档整理涉及的平台能力910A/910B/310P/950与参数限制均以仓库实现为准具体型号的完整支持矩阵请以随版本发布的 CANN 配套文档为准。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考