attention_update算子优化实践

​作者​昇腾实战派​知识地图​https://blog.csdn.net/Lumos_Lovegood/article/details/161601003背景概述在大模型推理场景中序列并行SP技术将长序列切分为多个分块每个分块独立计算局部注意力后需要通过 AttentionUpdate 算子将局部结果归约为全局结果。原始算子采用 TensorList 输入在开启较大上下文并行CP时Tensor 到 TensorList 的转换会引入大量 slice 操作影响推理性能。本文从性能瓶颈定位出发详细介绍了将算子输入从 TensorList 改为普通 Tensor 的改造方案涵盖算子开发、接入流程及问题排查最终实现了显著的性能提升。算子开发方案设计开发实现算子功能说明AttentionUpdate的功能是将序列并行SP各分块上 PagedAttention算子的局部结果归约合并为全局结果。具体来说每个 SP 分块独立计算了局部注意力输出两个中间量- lse_i该分块的 log-sum-exp 值代表注意力权重的归一化因子- O_i该分块的局部输出由于各分块的 softmax 是在局部 sequence上算的分母不一致不能直接相加。该算子通过以下步骤修正并合并l s e m a x max ⁡ i l s e i l s e ∑ i exp ⁡ ( l s e i − l s e m a x ) l s e m l s e m a x log ⁡ ( l s e ) O ∑ i O i ⋅ exp ⁡ ( l s e i − l s e m ) \begin{align*} lse_{max} \max_i lse_i \\[6pt] lse \sum_{i} \exp(lse_i - lse_{max}) \\[6pt] lse_m lse_{max} \log(lse) \\[6pt] O \sum_{i} O_i \cdot \exp(lse_i - lse_m) \end{align*}​lsemax​imax​lsei​lsei∑​exp(lsei​−lsemax​)lsem​lsemax​log(lse)Oi∑​Oi​⋅exp(lsei​−lsem​)​​原始算子的lse和go输入类型为TensorListsp切片数作为独立属性传入。修改目标是将接口改为普通Tensorsp直接从 tensor shape 的 dim 0 读取不再作为 attr。这里建议修改算子工程后对算子进行重命名以免后续因命名冲突算子无法正确被调用。op_hostattention_update_def.cpp — 算子定义Input(lse)/Input(go)由DYNAMIC改为REQUIREDDataType 列表保持不变Attr(sp)删除attention_update_infershape.cpp — InferShape原实现基于 TensorList 接口通过inputNumInferShape / 2推算 sp 并用动态索引取 go shape新接口下改为固定索引// 修改后直接通过固定 index 取 lse 和 goconstexprint32_tGO_INDEX1;autolseShapecontext-GetInputShape(LSE_INDEX);// [sp, bsh]autogoShapecontext-GetInputShape(GO_INDEX);// [sp, bsh, hd]// out: [bsh, hd]outShape-SetDimNum(2);outShape-SetDim(0,goShape-GetDim(1));// bshoutShape-SetDim(1,goShape-GetDim(2));// hd// lseOut: [bsh]lseOutShape-SetDimNum(1);lseOutShape-SetDim(0,lseShape-GetDim(1));// bshdecode_update_tiling.cpp — 910b/910_93 Tiling删除ATTR_SP_INDEX不再从 attr 读 sp维度常量重命名反映新 shape 语义constexpruint32_tLSE_SP_DIM0;// lse [sp, bsh] - sp at dim 0constexpruint32_tLSE_BSH_DIM1;// lse [sp, bsh] - bsh at dim 1constexpruint32_tIN_HD_DIM2;// go [sp, bsh, hd] - hd at dim 2constexpruint32_tLOCAL_OUT_INDEX1;// go is fixed at input index 1sp / totalLength / hd 改从 tensor shape 读取不再依赖 attrconstuint32_tspstatic_castuint32_t(lseShape.GetDim(LSE_SP_DIM));constuint32_ttotalLengthstatic_castuint32_t(lseShape.GetDim(LSE_BSH_DIM));constuint32_thdstatic_castuint32_t(inShape.GetDim(IN_HD_DIM));op_kerneldecode_update.h — 910b/910_93 Kernel 核心实现删除项GetTensorPtr()辅助函数原用于解析 TensorList 中指针间接层读取首地址偏移量还原实际数据指针成员变量lsePtr、inPtr__gm__ uint64_t*类型原保存各切片数据地址Init() — GM Buffer 绑定方式原 TensorList 接口下通过GetTensorPtr取出各切片地址在CopyIn里逐次绑定。新接口传入完整 Tensor在Init里一次性绑定// lse: [sp, bsh]整体大小 sp * totalLengthlseGm.SetGlobalBuffer((__gm__ lseType*)lse,sp*totalLength);// in: [sp, bsh, hd]整体大小 sp * totalLength * hDiminGm.SetGlobalBuffer((__gm__ outType*)in,sp*totalLength*hDim);CopyIn() — 按步长寻址各 sp 切片原来可以直接用 TensorList[i] 取到第 i 个切片的指针。新接口需要手动计算偏移for(int32_ti0;isp;i){// lse 第 i 行的起点跳过前 i 行每行 totalLength再加上当前 tile 和块内偏移uint32_tlseOffstatic_castuint32_t(i)*totalLengthprogress*tileLengthgmStartOffset;// in 第 i 行的起点每行宽 totalLength * hDimuint32_tinOffstatic_castuint32_t(i)*totalLength*hDimprogress*tileLength*hDimgmStartOffset*hDim;DataCopy(lseLocal[curLength*i],lseGm[lseOff],curLength);DataCopy(inLocal[curLength*hDim*i],inGm[inOff],curLength*hDim);}公式解释i * totalLength跳到第 i 个 sp 切片的行首lse 视图中每行有bsh totalLength个元素progress * tileLength当前处理的是第几个 tile每个 tile 处理tileLength个 bsh 位置gmStartOffset该 AI Core 负责的块在整行内的起始偏移op_apiaclnn_attention_update.h — ACLNN 头文件修改内容// 原始aclnnStatusaclnnAttentionUpdateGetWorkspaceSize(constaclTensorList*lse,constaclTensorList*localOut,int64_tupdateType,aclTensor*out,aclTensor*lseOut,...);// 修改后aclnnStatusaclnnAttentionUpdateGetWorkspaceSize(constaclTensor*lse,constaclTensor*localOut,int64_tupdateType,aclTensor*out,aclTensor*lseOut,...);aclnn_attention_update.cpp — ACLNN 实现修改内容参数类型aclTensorList*→aclTensor*CheckSp_95改为从lse-GetViewShape().GetDim(0)读取 sp不再依赖 TensorList 的Size()CheckShape_95简化为直接调用CheckShape合并了原来重复的校验逻辑去掉逐切片Contiguous循环改为对整体 tensor 调用一次l0op::Contiguousl0op::AttentionUpdate调用删除 sp 参数输出处理lse_m始终由 l0op 层分配updateType1时才通过ViewCopy写回attention_update.h / attention_update.cpp — l0op 层修改内容// 原始入参为 TensorList有 sp 参数返回类型已是 tupleconststd::tupleconstaclTensor*,constaclTensor*AttentionUpdate(constaclTensorList*lse,constaclTensorList*go,int64_tupdateType,int64_tsp,aclOpExecutor*executor);// 修改后入参改为 Tensor删除 sp 参数conststd::tupleconstaclTensor*,constaclTensor*AttentionUpdate(constaclTensor*lse,constaclTensor*go,int64_tupdateType,aclOpExecutor*executor);outshape[go.dim(1), go.dim(2)]即[bsh, hd]lseOutshape[lse.dim(1)]即[bsh]算子接入算子编译安装编译算子bashbuild.sh--pkg--opsattention_update--socascend910_93出包后安装./build_out/cann-ops-transformer-custom_linux-aarch64.run这里有一个很重要的环境变量建议放在你的执行脚本中export LD_LIBRARY_PATH/usr/local/cann-8.5.0-beta.1/opp/vendors/custom_transformer/op_api/lib:${LD_LIBRARY_PATH}运行测试示例示例位置attention_update/examples/test_aclnn_attention_update.cppbashbuild.sh--run_exampleattention_update eager custtorch_npu接入全量编译代码准备这里采用源码编译安装torch_npugitclone https://gitcode.com/Ascend/pytorch.git-bv2.9.0-7.3.0--recursiveop_plugin接入cdpytorch/third_party/op-plugin修改点1op_plugin/config/op_plugin_functions.yaml中attention_update输入类型改为Tensor修改点2op_plugin/ops/opapi/AttentionUpdateKernelNpuOpApi.cpp// Copyright (c) 2025 Huawei Technologies Co., Ltd// All rights reserved.//// Licensed under the BSD 3-Clause License (the License);// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an AS IS BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License.#includeop_plugin/AclOpsInterface.h#includeop_plugin/OpApiInterface.h#includeop_plugin/utils/op_api_common.hnamespaceop_api{usingnpu_preparationat_npu::native::OpPreparation;std::tupleat::Tensor,at::Tensornpu_attention_update(constat::Tensorlse,// [sp, bsh]constat::Tensorlocal_out,// [sp, bsh, hd]int64_tupdate_type){// out shape: [bsh, hd]去掉第0维 spautooutput_sizelocal_out.sizes().slice(1);at::Tensor outnpu_preparation::apply_tensor_without_format(output_size,local_out.options());// lseOut shape: [bsh]去掉第0维 spupdate_type0 时不分配at::Tensor lse_out;if(update_type1){autolseout_sizelse.sizes().slice(1);// [bsh]lse_outnpu_preparation::apply_tensor_without_format(lseout_size,lse.options());}EXEC_NPU_CMD(aclnnAttentionUpdate,lse,local_out,update_type,out,lse_out);returnstd::make_tuple(out,lse_out);}}// namespace op_apitorchair图模式接入图模式接入需要把ATen IR接口映射到GE IR修改点1cd pytorch/third_party/op-plugin/op_plugin/python/meta/_meta_registrations.py把原有的impl(m, npu_attention_update)替换成以下代码impl(m,npu_attention_update)defnpu_attention_update_meta(lse,local_out,update_type):# out: [bsh, hd]outtorch.empty(local_out.size(1),local_out.size(2),dtypelocal_out.dtype,devicelocal_out.device)# lse_m: [bsh]lse_mtorch.empty(lse.size(1),dtypelse.dtype,devicelse.device)returnout,lse_m修改点2cd pytorch/third_party/torchair/torchair/python/torchair/_ge_concrete_graph/ge_converter/custom增加算子converternpu_attention_update.pyfromtypingimport(Any,Callable,ContextManager,Iterator,List,Literal,NamedTuple,Optional,Sequence,Tuple,TypeVar,Union,overload,)importtorchfromtorchimportGenerator,contiguous_format,inf,strided,SymIntfromtorch.typesimportDevice,Number,_bool,_complex,_device,_dtype,_float,_int,_layout,_qscheme,_sizeimporttorchairfromtorchair._ge_concrete_graphimportge_apisasgefromtorchair._ge_concrete_graph.fx2ge_converterimportdeclare_supported,register_fx_node_ge_converterfromtorchair.ge._ge_graphimportTensor,TensorSpec,DataTypefromtorchair._ge_concrete_graph.supported_declarationimport_TypedTensor,F32,F16,F64,I32,I16,I64,I8,U8,\ BOOL,Supportfromtorchair._ge_concrete_graph.utilsimportdtype_promotefromtorchair.geimportattrregister_fx_node_ge_converter(torch.ops.npu.npu_attention_update.default)defconvert_npu_attention_update(lse:Tensor,local_out:Tensor,update_type:int,meta_outputs:List[TensorSpec]None,):returntorchair.ge.custom_op(AttentionUpdate,inputs{lse:lse,go:local_out,},attrs{update_type:attr.Int(update_type),},outputs[output,lse_m])编译出包cdpytorchbashci/build.sh--python3.10安装torch_npupipinstalldist/torch_npu-2.9.0.post1git0a9c637-cp310-cp310-linux_aarch64.whl安装后验证单算子测试pytest-vtest_npu_attention_update.py!优化效果通过对AttentionUpdate算子的改进在decode阶段取得了性能收益相对普通dp收益8%。修改前修改后问题记录编译算子测试文件报错问题定位编译报错显示算子调用的接口仍然是修改输入类型前的接口问题根因c编译顺序问题原本CUST_INCLUDE_PATH顺序在INCLUDE_PATH后由于算子未进行重命名编译优先从INCLUDE_PATH找到了修改入参类型前的接口导致编译报错vLLM整网拉起实验算子执行卡住问题定位通过vllm整网和单算子执行的plog进行比对发现vllm整网执行与单算子执行时链接到的op api是不同的整网plog这里aclnnAttentionUpdateGetWorkspaceSize只进不出一般为入参类型有问题单算子plog有进有出功能正常仔细看plog发现整网执行该算子时明显调用的是修改类型前的接口与单算子调用不一致问题根因vLLM 启动时bootstrap_custom_op_env()将vllm_ascend/_cann_ops_custom前插到ASCEND_CUSTOM_OPP_PATH该路径下存在与自定义算子同名的旧实现CANN runtime 按路径顺序优先命中旧路径后不再继续查找导致自定义算子被跳过、fallback 到旧实现。解决1对算子接口进行重命名在vllm自定义算子路径查找不到后就会自动去cann的自定义算子路径查找整网功能打通解决2直接将该算子增量编译至vllm-ascend中