ATB SelfAttention 融合算子深度解析：知识条目、参数体系与 Runner 分发机制
【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库，基于华为Ascend AI处理器，提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost
本文以 ATB（Ascend 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.md |
| Runner 类型 | ACLNN Runner + OpsRunner（详见路由文件） |
| Pipeline | BMM1 → Softmax → BMM2（详见深度条目） |
从知识条目可以得到该算子最核心的三个技术事实：
- 执行通道是双轨的：SelfAttention 同时支持 ACLNN Runner（走
aclnnFusedInferAttentionScoreV5等 CANN 高层算子接口）与原生 OpsRunner（直接驱动融合算子内核），具体选择由运行平台与参数组合决定。 - 计算管线是经典注意力三阶段：Q·K^T（BMM1）→ Softmax → ·V（BMM2），整个流程被融合进单个算子在昇腾处理器上执行。
- 实现规模与复杂度：路由文件标注该算子分类为 infer、复杂度 XL、共 27 个源文件，属于推理侧最复杂的融合算子之一，并存在 910A、BNSD、Bypass、Encoder、Fusion、PrefixEncoder 六类平台/功能变体。
2. 文件清单与推荐阅读顺序
路由文件 .agent/knowledge/routing/self_attention.md 给出了完整的 27 个文件清单，按角色分为三类：Operation 定义（算子对外行为）、ACLNN Runner（ACLNN 通道）、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 分发决策逻辑； - 随后按平台深入各 Runner：910B 侧读
self_attention_encoder_fusion_ops_runner、self_attention_fusion_bypass_ops_runner（含 BNSD 变体）、self_attention_prefix_encoder_ops_runner；910A 侧读对应的_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 关键枚举
| 枚举 | 取值 | 含义 |
|---|---|---|
CalcType | UNDEFINED/ENCODER/DECODER/PA_ENCODER/PREFIX_ENCODER | 计算场景：FA 的编码器/解码器、PagedAttention 编码器、Prefix Encoder 前缀融合场景 |
KernelType | KERNELTYPE_DEFAULT/KERNELTYPE_HIGH_PRECISION/KERNELTYPE_EXP_M8V2 | 内核精度：默认 fp16 全链路；高精度（BMM 用 fp32 累加）；310P 上的 EXP M8V2 专用内核 |
MaskType | MASK_TYPE_UNDEFINED/NORM/ALIBI/NORM_COMPRESS/ALIBI_COMPRESS/ALIBI_COMPRESS_SQRT/ALIBI_COMPRESS_LEFT_ALIGN/SLIDING_WINDOW_NORM/SLIDING_WINDOW_COMPRESS/CAUSAL_MASK | mask 类型：全 0 mask、倒三角、alibi、各类压缩 mask、SWA 滑动窗口 mask、内部生成因果 mask |
KvCacheCfg | K_CACHE_V_CACHE/K_BYPASS_V_BYPASS | 是否走 KVCache 管理：默认缓存处理；bypass 表示直接传入 KV（不维护 cache） |
ScaleType | SCALE_TYPE_TOR/SCALE_TYPE_LOGN | 是否启用 LogN 缩放 |
QuantType | TYPE_QUANT_UNQUANT/TYPE_DEQUANT_FUSION/TYPE_QUANT_QKV_OFFLINE/TYPE_QUANT_QKV_ONLINE | 量化模式：不量化；反量化融合（预留，当前不可取）；离线/在线 INT8 QKV 量化 |
ClampType | CLAMP_TYPE_UNDEFINED/CLAMP_TYPE_MIN_MAX | 是否对注意力分数做 min/max clamp |
CacheType | CACHE_TYPE_NORM/CACHE_TYPE_SWA | cache 排布：正常 cache；SWA 固定长度 cache（只存后windowSize个 token） |
3.2 核心字段与默认值
| 字段 | 默认值 | 说明 |
|---|---|---|
headNum | 0（需 > 0） | query 头数 |
kvHeadNum | 0 | KV 头数；为 0 时与headNum一致，非 0 时要求headNum % kvHeadNum == 0 |
qScale | 1.0 | query 缩放系数 |
qkScale | 1.0 | 在 Q·K^T 之后乘的缩放值（tor 值） |
batchRunStatusEnable | false | 是否开启动态 batch |
isTriuMask | 0 | 倒三角优化开关，仅 mask 为倒三角时可开启 |
calcType/kernelType/clampType/maskType/kvcacheCfg/scaleType/cacheType | 各枚举首值 | 计算/内核/mask/cache/缩放配置 |
clampMin/clampMax | 0 | clamp 上下界 |
inputLayout | TYPE_BSND | 数据排布，支持 BSND 与 BNSD |
mlaVHeadSize | 0 | 大于 0 时开启 MLA 合并 KVCache 功能，表示合并传入时 V 的 head_size，取值范围 [0, 576] |
windowSize | 0 | 大于 0 时开启 SWA 特性，表示滑动窗口大小，此时 maskType 须为SLIDING_WINDOW_NORM或SLIDING_WINDOW_COMPRESS |
quantType/outDataType | TYPE_QUANT_UNQUANT/ACL_DT_UNDEFINED | 量化类型与输出类型；QKV 量化时 outDataType 只能取ACL_FLOAT16或ACL_BF16 |
rsv[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时不能为 bypass（L72-L76、L110-L114）；quantType不能取TYPE_DEQUANT_FUSION（预留类型）；- QKV 量化（离线/在线）只支持
PA_ENCODER、SCALE_TYPE_TOR、BSND 排布，且输出必须是 fp16/bf16（L91-L104）； - 开启 SWA 后不支持动态 batch、高精度内核、clamp、QKV 量化、LogN 与 BNSD（SWAParamCheck）；
- BNSD 排布与
scaleType非 TOR、量化、Prefix Encoder 互斥（BNSDParamCheck）； - MLA 模式（
mlaVHeadSize > 0）只支持 910B 平台与PA_ENCODER，且mlaVHeadSize ≤ 576（MlaParamCheck）； MASK_TYPE_ALIBI_COMPRESS*系列仅PA_ENCODER/PREFIX_ENCODER可用；MASK_TYPE_UNDEFINED时isTriuMask必须为 0（L115-L126）。
4. 平台与功能变体
路由文件将实现拆分为六类变体，全部体现在 src/ops/ops_infer/self_attention/ 的文件命名中：
| 变体 | 含义 | 代表文件 |
|---|---|---|
| 910A | 昇腾 910A（Atlas 800 训练产品）平台适配 | self_attention_fusion_ops_runner_910a.cpp、self_attention_fusion_bypass_ops_runner_910a.cpp |
| BNSD | BNSD 数据排布适配 | self_attention_fusion_bypass_ops_runner_BNSD.cpp（含_910a变体） |
| Bypass | KVCache Bypass 路径（直接传入 KV） | self_attention_fusion_bypass_ops_runner.cpp |
| Encoder | Encoder 计算路径 | self_attention_encoder_fusion_ops_runner.cpp、atb_acl_self_attention_prefix_encoder.cpp |
| Fusion | 融合算子路径 | self_attention_encoder_fusion_ops_runner.cpp（含_910a变体） |
| PrefixEncoder | Prefix 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 └─> SelfAttentionAclnnRunner（ACLNN 通道） 平台 == 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::GetRunnerTypeIdx+RunnerPool::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.cpp）：
atb::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 即知识条目声明的 BMM1（Q·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 * vHeadSize（InferShapeImpl910B）。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 ≤ 256；MLA 内核场景放宽到 1024、MLA bypass 场景为 576；压缩 Alibi mask 场景收紧到 128（MaxHeadSizeCheck910B）。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]（NormMaskDimCheck）；SWA 压缩 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，对应 README（README.md、README_en.md）中给出的参数模板：
| Demo 文件 | calcType | maskType | 场景 |
|---|---|---|---|
| self_attention_encoder_demo.cpp | ENCODER | MASK_TYPE_NORM | FA Encoder 全量注意力 |
| self_attention_encoder_inference_demo.cpp | ENCODER | MASK_TYPE_UNDEFINED | 免 mask 推理 |
| self_attention_pa_encoder_demo.cpp | PA_ENCODER | MASK_TYPE_NORM | PagedAttention 编码器 |
| self_attention_pa_encoder_qwen_demo.cpp | PA_ENCODER | — | Qwen 系模型 PA 场景 |
| self_attention_prefix_encoder_demo.cpp | PREFIX_ENCODER | — | 前缀缓存融合编码 |
PA_ENCODER 与 PREFIX_ENCODER 场景下，输入除 Q 外还包含 blockTables、seqLen、kvSeqLen（prefix 场景为 8 个输入：query、key、value、blockTables、mask、qSeqLen、kvSeqLen、slopes），tokenOffset 与 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 走 ACLNN（aclnnFusedInferAttentionScoreV5），910B 与 910A/310P 走特化的 OpsRunner，统一经
CreateRunner()分发； - Pipeline：BMM1 → 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），仅供参考
