CANN ascend-transformer-boost LaserAttentionGrad 反向注意力算子源码解析参数、执行链路与 Kernel 结构【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boostLaserAttentionGrad 是 CANN ascend-transformer-boost 中训练场景下的自注意力反向算子用于把前向laser_attention的输出梯度沿 QKV 链路回传到输入。本文以仓库内.agent/knowledge/routing/laser_attention_grad.md路由文档为核心结合src/ops/ops_train/laser_attention_grad/、src/kernels/mixkernels/laser_attention_grad/与include/atb/train_op_params.h的源码实现梳理该算子的文件组织、参数约束、TensorDesc 校验、Runner 执行链路与 Kernel 依赖帮助读者快速上手阅读与使用该算子。一、算子定位训练侧的反向传播实现LaserAttentionGrad 在项目中的定位是训练train类算子对应复杂度评级为 S源码文件数为 4。从include/atb/train_op_params.h中LaserAttentionGradParam的注释可以确认其语义训练场景下使用 LaserAttention 算法实现 self-attention自注意力的计算反向。它与前向算子laser_attention位于 src/ops/ops_train/laser_attention/成对出现执行路径描述为梯度从 QKV 回传到输入。在 .agent/knowledge/ops/attention/laser_attention_grad/index.md 中其计算流水线被归类为single_stage单阶段pipeline_type: single_stage note: LaserAttention 反向传播。梯度从 QKV 回传到输入。平台限制该算子仅支持 Atlas 800I A2 推理产品。这一限制并非口头说明而是硬编码在参数校验逻辑中见 src/ops/ops_train/laser_attention_grad/laser_attention_grad_operation.cppbool is910B atb::GetSingletonatb::Config().Is910B(); if (!is910B) { error.errorDesc Platform is not Atlas 800I A2 inference product, operation is not supported,; ... return false; }即在非 A2 平台上参数校验会直接返回ERROR_INVALID_PARAM并记录 ERROR 日志算子不可用。二、文件清单与推荐阅读顺序路由文档给出的文件清单是理解该算子架构的索引4 个文件分属两层#文件角色1laser_attention_grad_operation.cppOperation 定义2laser_attention_grad_operation.hOperation 定义3laser_attention_grad_ops_runner.cppOps Runner4laser_attention_grad_ops_runner.hOps Runner实际目录位置Op 目录src/ops/ops_train/laser_attention_grad/Kernel 目录src/kernels/mixkernels/laser_attention_grad参数头文件include/atb/train_op_params.h推荐阅读顺序由知识条目补充顺序文件重点关注1laser_attention_grad_operation.h了解输入输出数量、InferShape 签名2laser_attention_grad_operation.cppCreateRunner() 决策逻辑与各 Tensor 校验3laser_attention_grad_ops_runner.h原生 Ops 执行接口与参数比较4laser_attention_grad_ops_runner.cpp原生 Ops 调用链 Kernel 图构建从源码结构看operation.h/.cpp负责算子语义参数、张量形状、InferShape、Runner 装配ops_runner.h/.cpp负责把 ATB 层的参数与张量翻译成 Kernel 图Mki 中间表示并交给底层 Kernel 执行二者职责清晰分层。三、参数结构 LaserAttentionGradParam 详解参数结构体定义在 include/atb/train_op_params.h 的train命名空间下。完整字段、默认值与约束如下参数类型默认值约束headNumint0必须 0且与输入张量 dim1 相等最大 512inputLayoutstd::stringBNSD仅支持BNSD或SBHscaleValuefloat0.08838834764831843取值范围 (0, 1]keepProbfloat1.0暂未使用配置值不生效preTokensint2147483647取值范围 [1, seq_size]且必须是 256 的整数倍nextTokensint1暂未使用配置值不生效sparseModeint0暂未使用配置值不生效innerPreciseint1仅支持配置为 1高精度模式rsv[8]uint8_t{0}预留参数其中scaleValue的默认值 0.08838834764831843 对应 headDim128 时的1/sqrt(headDim)缩放系数即经典 Attention 公式softmax(QK^T / sqrt(d_k))中的缩放因子。3.1 参数校验逻辑laser_attention_grad_operation.cpp 中的ParamCheck函数对上述参数执行严格校验逐条规则如下平台检查必须是 Atlas 800I A2 推理产品Is910B()否则直接拒绝headNum必须大于 0inputLayout必须是BNSD或SBH之一scaleValue必须在 (0, 1] 区间内preTokens必须 1且必须是 256 的整数倍源码中BLOCK_SIZE_ALIGN 256innerPrecise必须等于 1。任何一个条件不满足都会构造ExternalError错误类型ERROR_INVALID_PARAM并写入 ERROR 日志日志中会带上具体参数名与当前取值便于定位问题。四、输入输出张量15 输入、4 输出GetInputNum()返回 15GetOutputNum()返回 4常量IN_TENSOR_NUM 15、OUT_TENSOR_NUM 4见 laser_attention_grad_operation.cpp。4.1 输入张量索引下标从 0 开始下标张量名是否允许空张量说明0query否Q 输入1key否K 输入2value否V 输入3attentionOutGrad否前向输出梯度4pseShift是位置编码偏移须为空张量5dropMask是须为空张量6paddingMask是须为空张量7attenMask是非空时需满足维度约束8softmaxMax否前向保存的 max9softmaxSum否前向保存的 sum形状须与 softmaxMax 相同10softmaxIn否前向 softmax 输入11attentionIn否前向 attention 输入12prefix是须为空张量13actualSeqQLen是须为空张量14actualSeqKVLen是须为空张量空张量权限由GetEmptyInTensorPermissions()显式声明laser_attention_grad_operation.cpppseShift、dropMask、paddingMask、attenMask、prefix、actualSeqQLen、actualSeqKVLen 这 7 个输入允许为空对应前向中未启用的可选能力而 query/key/value/attentionOutGrad/softmaxMax/softmaxSum/softmaxIn/attentionIn 这 8 个核心输入不允许为空。4.2 输出张量下标张量名说明0queryGradQ 的梯度形状与 query 相同1keyGradK 的梯度形状与 key 相同2valueGradV 的梯度形状与 value 相同3dpse预留梯度输出InferShapeImpl的逻辑非常直观queryGrad/keyGrad/valueGrad三个输出的形状分别直接拷贝自query/key/value三个输入的 TensorDescdpse拷贝自query的形状laser_attention_grad_operation.cpp。4.3 关键 TensorDesc 校验规则InferShapeCheckImpl与SetupCheckImpl会根据inputLayout分流到两套校验BNSD 布局走InTensorDescsCheckSBH 布局走InTensorDescsCheckSBH。两套校验都包含以下核心规则BNSD 为例queryinTensor04 维[batch, qHeadNum, seqSize, qHeadDim]dim1 必须等于param.headNum且 512seqSize必须是 256 的整数倍qHeadDim仅支持 128 或 192——headDim128 时要求seqSize preTokensheadDim192DeepSeek 风格时要求seqSize preTokenskeyinTensor14 维[batch, kvHeadNum, kvSize, kHeadDim]batch 与 query 一致qHeadNum必须是kvHeadNum的整数倍支持 MQA/GQA 的 KV head 复用kvSize必须是 256 的整数倍headDim128 时kHeadDim必须为 128headDim192 时kHeadDim必须为 256常量DEEP_SEEK_K_HEAD_DIM_VALUE 256valueinTensor2batch、kvHeadNum、kvSize 分别与 key 对齐vHeadDim必须为 128attentionOutGradinTensor3形状与 query 输出侧对齐dim3 必须为 128attenMaskinTensor7非空时必须是 2 维[seqSize, kvSize]且要求seqSize kvSizesoftmaxSuminTensor9形状必须与 softmaxMax 完全一致TensorDescsEqualprefix / actualSeqQLen / actualSeqKVLeninTensor12~14必须是空张量。SBH 布局下3 维张量源码通过最后一维除以固定 headDim 反推 head 数的方式解析query 为[seqSize, batch, headNum*192]key 为[kvSize, batch, kvHeadNum*256]value 为[kvSize, batch, kvHeadNum*128]其中 192/256/128 分别对应 Q/K/V 的固定 headDim 常量。五、执行链路从 Operation 到 Kernel5.1 CreateRunner 决策逻辑LaserAttentionGradOperation::CreateRunnerlaser_attention_grad_operation.cpp遵循 ATB 的 Runner 池化模式把传入的Context向下转型为ContextBase通过RunnerTypeRegister::GetRunnerTypeIdx(LaserAttentionGradOpsRunner)查询注册的 Runner 类型索引从RunnerPool中MallocRunnerLaserAttentionGradOpsRunner, train::LaserAttentionGradParam(param_)获取池化实例避免反复构造若池化分配失败回退为std::make_sharedLaserAttentionGradOpsRunner(param_)直接创建。Runner 类型通过REG_RUNNER_TYPE(LaserAttentionGradOpsRunner)宏注册见 laser_attention_grad_ops_runner.cpp。5.2 OpsRunner 的 Kernel 图构建LaserAttentionGradOpsRunner::SetupKernelGraphlaser_attention_grad_ops_runner.cpp把 ATB 层的 15 个输入、4 个输出翻译为 MkiKernel 中间表示张量构建单节点 Kernel 图kernelGraph_.nodes.resize(1); KernelGraphNode laserAttentionNode kernelGraph_.nodes.at(0); AtbOps::OpParam::LaserAttentionGrad laserAttentionGradParam; laserAttentionGradParam.headNum param_.headNum; laserAttentionGradParam.inputLayout param_.inputLayout; laserAttentionGradParam.scaleValue param_.scaleValue; laserAttentionGradParam.keepProb param_.keepProb; laserAttentionGradParam.preTokens param_.preTokens; laserAttentionGradParam.nextTokens param_.nextTokens; laserAttentionGradParam.innerPrecise param_.innerPrecise; laserAttentionNode.opDesc {0, LaserAttentionGradOperation, laserAttentionGradParam}; laserAttentionNode.inTensors {queryTensor, keyTensor, valueTensor, attentionOutGradTensor, pseShiftTensor, dropMaskTensor, attenMaskTensor, softmaxMaxTensor, softmaxSumTensor, attentionInTensor}; laserAttentionNode.outTensors {queryGradTensor, keyGradTensor, valueGradTensor};注意此处节点实际只挂接了 10 个输入与 3 个输出paddingMask、softmaxIn、prefix、actualSeqQLen、actualSeqKVLen 以及 dpse 在代码中被显式(void)标记对应校验阶段要求为空张量的可选能力当前版本不参与 Kernel 计算。ATB 参数train::LaserAttentionGradParam在此被逐字段拷贝为底层算子参数AtbOps::OpParam::LaserAttentionGrad完成跨层参数透传。SetParam通过operator定义于 laser_attention_grad_ops_runner.h比较新参数与旧参数仅在参数发生变化时更新param_并置位isParamUpdated_触发 Kernel 图重建避免相同参数下的重复编译。5.3 执行路径总览LaserAttentionGradOperation::CreateRunner() └── → LaserAttentionGradOpsRunner (REG_RUNNER_TYPE 注册) └── SetupKernelGraph() → KernelGraph 单节点 └── Kernel: laser_attention_grad5.4 Kernel 依赖Kernel 实现位于src/kernels/mixkernels/laser_attention_grad/目录结构分为两层op_kernel/算子主体实现包括laser_attention_grad.cpp、cube_backward.hCube 侧反向、vector_backward.hVector 侧反向、address_mapping_cube.h/address_mapping_vector.h地址映射、TransposeCustom.h/TransposeGrad.h/TransposeWithDtype.h转置与数据类型转换、lag_post.h后处理等头文件另有cube_backward_band_op.h与cube_backward_band_op_192.h对应 headDim128 / 192 的带状band算子变体tiling/laser_attention_grad_tiling.cpp/h与tiling_data.h负责运行时 shape 相关的 tile 切分与参数下发。从源码结构看反向计算的典型组成P 矩阵 softmax 输出dV P^T · dO、dK dS^T · Q、dQ dS · V其中 dS P ⊙ (dO·V^T - rowsum)由 Cube 与 Vector 两侧协同完成Cube 侧承担矩阵乘形态的梯度Vector 侧承担 softmax 相关逐元素修正lag_post完成梯度后处理。六、已知限制与关联算子6.1 已知问题#问题状态1仅 Atlas 800I A2 平台支持限制lim除平台限制外从参数与校验逻辑还可归纳出如下使用前提均以源码为准preTokens与seqSize必须为 256 的整数倍BLOCK_SIZE_ALIGN 256headDim192 场景下seqSize必须严格等于preTokenshead 数上限 512Q head 数须为 KV head 数的整数倍keepProb、nextTokens、sparseMode三个参数虽保留在结构体中但当前版本暂未使用配置值未生效。6.2 相关算子算子关系laser_attention前向算子src/ops/ops_train/laser_attention/self_attention推理侧对应但 SA 无单独 backward Opfast_soft_max_grad同属训练侧Softmax 反向src/ops/ops_train/fast_soft_max_grad/七、进一步阅读路由文档本文主体.agent/knowledge/routing/laser_attention_grad.md详细知识条目.agent/knowledge/ops/attention/laser_attention_grad/index.mdOperation 头文件src/ops/ops_train/laser_attention_grad/laser_attention_grad_operation.hOperation 实现src/ops/ops_train/laser_attention_grad/laser_attention_grad_operation.cppOpsRunner 实现src/ops/ops_train/laser_attention_grad/laser_attention_grad_ops_runner.cpp参数头文件include/atb/train_op_params.hKernel 目录src/kernels/mixkernels/laser_attention_grad/【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
