企业官网建设流程全解析

​作者​：昇腾实战派
​知识地图​：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​lse=i∑​exp(lsei​−lsemax​)lsem​=lsemax​+log(lse)O=i∑​Oi​⋅exp(lsei​−lsem​)​​

原始算子的lse和go输入类型为TensorList，sp（切片数）作为独立属性传入。修改目标是将接口改为普通Tensor，sp直接从 tensor shape 的 dim 0 读取，不再作为 attr。

这里建议修改算子工程后对算子进行重命名，以免后续因命名冲突算子无法正确被调用。

op_host

attention_update_def.cpp — 算子定义

• Input("lse")/Input("go")：由DYNAMIC改为REQUIRED，DataType 列表保持不变
• Attr("sp")：删除

attention_update_infershape.cpp — InferShape

原实现基于 TensorList 接口，通过inputNumInferShape / 2推算 sp 并用动态索引取 go shape；新接口下改为固定索引：

// 修改后：直接通过固定 index 取 lse 和 goconstexprint32_tGO_INDEX=1;autolseShape=context->GetInputShape(LSE_INDEX);// [sp, bsh]autogoShape=context->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));// bsh

decode_update_tiling.cpp — 910b/910_93 Tiling

• 删除ATTR_SP_INDEX，不再从 attr 读 sp
• 维度常量重命名，反映新 shape 语义：

constexpruint32_tLSE_SP_DIM=0;// lse [sp, bsh] -> sp at dim 0constexpruint32_tLSE_BSH_DIM=1;// lse [sp, bsh] -> bsh at dim 1constexpruint32_tIN_HD_DIM=2;// go [sp, bsh, hd] -> hd at dim 2constexpruint32_tLOCAL_OUT_INDEX=1;// go is fixed at input index 1

• sp / totalLength / hd 改从 tensor shape 读取（不再依赖 attr）：

constuint32_tsp=static_cast<uint32_t>(lseShape.GetDim(LSE_SP_DIM));constuint32_ttotalLength=static_cast<uint32_t>(lseShape.GetDim(LSE_BSH_DIM));constuint32_thd=static_cast<uint32_t>(inShape.GetDim(IN_HD_DIM));

op_kernel

decode_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_ti=0;i<sp;i++){// lse 第 i 行的起点：跳过前 i 行（每行 totalLength），再加上当前 tile 和块内偏移uint32_tlseOff=static_cast<uint32_t>(i)*totalLength+progress*tileLength+gmStartOffset;// in 第 i 行的起点：每行宽 totalLength * hDimuint32_tinOff=static_cast<uint32_t>(i)*totalLength*hDim+progress*tileLength*hDim+gmStartOffset*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_api

aclnn_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::Contiguous
• l0op::AttentionUpdate调用删除 sp 参数
• 输出处理：lse_m始终由 l0op 层分配，updateType=1时才通过ViewCopy写回

attention_update.h / attention_update.cpp — l0op 层

修改内容：

// 原始：入参为 TensorList，有 sp 参数；返回类型已是 tupleconststd::tuple<constaclTensor*,constaclTensor*>AttentionUpdate(constaclTensorList*lse,constaclTensorList*go,int64_tupdateType,int64_tsp,aclOpExecutor*executor);// 修改后：入参改为 Tensor，删除 sp 参数conststd::tuple<constaclTensor*,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--ops=attention_update--soc=ascend910_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.cpp

bashbuild.sh--run_exampleattention_update eager cust

torch_npu接入（全量编译）

代码准备（这里采用源码编译安装torch_npu）

gitclone https://gitcode.com/Ascend/pytorch.git-bv2.9.0-7.3.0--recursive

op_plugin接入

cdpytorch/third_party/op-plugin

修改点1：op_plugin/config/op_plugin_functions.yaml中attention_update输入类型改为Tensor

修改点2：op_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.#include"op_plugin/AclOpsInterface.h"#include"op_plugin/OpApiInterface.h"#include"op_plugin/utils/op_api_common.h"namespaceop_api{usingnpu_preparation=at_npu::native::OpPreparation;std::tuple<at::Tensor,at::Tensor>npu_attention_update(constat::Tensor&lse,// [sp, bsh]constat::Tensor&local_out,// [sp, bsh, hd]int64_tupdate_type){// out shape: [bsh, hd]，去掉第0维 spautooutput_size=local_out.sizes().slice(1);at::Tensor out=npu_preparation::apply_tensor_without_format(output_size,local_out.options());// lseOut shape: [bsh]，去掉第0维 sp；update_type=0 时不分配at::Tensor lse_out;if(update_type==1){autolseout_size=lse.sizes().slice(1);// [bsh]lse_out=npu_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_api

torchair图模式接入

图模式接入需要把ATen IR接口映射到GE IR

修改点1：cd 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]out=torch.empty(local_out.size(1),local_out.size(2),dtype=local_out.dtype,device=local_out.device)# lse_m: [bsh]lse_m=torch.empty(lse.size(1),dtype=lse.dtype,device=lse.device)returnout,lse_m

修改点2：cd pytorch/third_party/torchair/torchair/python/torchair/_ge_concrete_graph/ge_converter/custom

增加算子converter：npu_attention_update.py

fromtypingimport(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.geimportattr@register_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--python=3.10

安装torch_npu

pipinstalldist/torch_npu-2.9.0.post1+git0a9c637-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中