企业官网建设流程全解析

CANN ops-transformer 算子 aclnnAttentionUpdate 接口详解：SP 域序列并行 Attention 局部结果到全局结果的合并更新

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer

导读

aclnnAttentionUpdate是 CANN ops-transformer 算子库（attention/attention_update）中专门面向序列并行（Sequence Parallel，SP）场景的融合更新算子：它将各 SP 域 PagedAttention（PA）类算子输出的局部 log-sum-exp（lse）与局部 Attention 输出（localOut）合并为全局结果，是序列并行大模型推理/训练中跨域归一化的关键一环。本文以 attention/attention_update/docs/aclnnAttentionUpdate.md 为骨架，结合仓库内算子定义、tiling、kernel 与测试源码，完整讲解其数学原理、两段式接口原型、参数约束、错误码语义、C++ 调用示例与底层实现机制，帮助开发者快速上手并深入理解该算子在 NPU 上的执行方式。

功能定位：序列并行下的局部结果合并

在长序列（超长 context）大模型场景中，单个 NPU 无法容纳整条序列的 KV Cache，通常采用序列并行（SP）把序列切成多个分片（SP 域），每个分片由独立的计算单元执行 PagedAttention 等算子。此时每个 SP 域只会得到局部的softmax 统计量（lse_i，log-sum-exp）和局部的Attention 输出（O_i，即localOut），这些局部结果相互独立，不能直接拼接，必须先做跨域合并才能得到与整条序列语义等价的全局输出。

aclnnAttentionUpdate正是完成这一合并任务的算子：它接收sp个 SP 域的局部 lse 与局部 Attention 输出，通过"求全局最大值 → 重新指数归一化 → 加权求和"的流程，输出全局 lse（可选）与全局 Attention 输出，即：

• 接口功能：将各 SP 域 PA 算子的输出的中间结果lse、localOut两个局部变量结果更新成全局结果（出自 aclnnAttentionUpdate.md）。

对应地，算子内部输入在 attention_update_def.cpp 中定义为动态个数的lse（1 维）与go（2 维，即 localOut）张量列表，并通过属性sp声明张量个数。

数学原理：跨 SP 域的 softmax 合并公式

设第i个 SP 域的局部统计量为lse_i、局部输出为O_i（i = 1 … sp），算子按如下 4 步将局部量合并为全局量：

$$ lse_{max} = \text{max}_i, lse_i $$

$$ lse = \sum_i \text{exp}(lse_i - lse_{max}) $$

$$ lse_m = lse_{max} + \text{log}(lse) $$

$$ O = \sum_i O_i \cdot \text{exp}(lse_i - lse_m) $$

其中：

• lse_max是所有 SP 域局部 lse 的最大值，用于数值稳定性（防止指数溢出）；
• lse是重归一化后的指数和；
• lse_m是合并后的全局log-sum-exp，即最终的全局归一化常数；
• O是加权求和得到的全局 Attention 输出，每个局部输出O_i的权重为exp(lse_i - lse_m)。

这一过程在数学上等价于对多个子序列 softmax 结果做"log-sum-exp 形式"的精确合并，不损失精度。仓库测试中的 CPU 参考实现 executor_aclnnAttentionUpdate.py 也严格按照该公式构造（先torch.exp再按总和归一化、加权求和），用于对 NPU 结果做一致性比对，可作为公式理解的辅助参考。

产品支持情况

根据文档 aclnnAttentionUpdate.md，各产品形态支持情况如下：

与算子定义源码中AICore().AddConfig注册的硬件配置一致：ascend910b（Atlas A2 系列）、ascend910_93（Atlas A3 系列）、ascend950三个平台（见 attention_update_def.cpp），其中ascend950走独立的 regbase（arch35）实现路径。

两段式接口：先获取 workspace，再执行计算

aclnnAttentionUpdate采用 CANN 单算子 API 通用的两段式接口（机制说明见 docs/zh/context/two_phase_api.md）：

1. 先调用aclnnAttentionUpdateGetWorkspaceSize完成入参校验，并获取计算所需 workspace 大小及封装了计算流程的执行器；
2. 按返回的workspaceSize在 Device 侧申请内存后，再调用aclnnAttentionUpdate执行计算。

两段接口的函数原型如下（出自 aclnnAttentionUpdate.md）：

aclnnStatus aclnnAttentionUpdateGetWorkspaceSize( const aclTensorList *lse, const aclTensorList *localOut, int64_t updateType, aclTensor *out, aclTensor *lseOut, uint64_t *workspaceSize, aclOpExecutor **executor)

aclnnStatus aclnnAttentionUpdate( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

注意：第二段接口aclnnAttentionUpdate不可重复调用，同一executor只能执行一次计算，否则行为异常。

aclnnAttentionUpdateGetWorkspaceSize 参数详解

第一段接口共 7 个入参/出参，完整参数说明如下（整理自 aclnnAttentionUpdate.md）：

补充说明：

• lse / localOut 的 tensorList 长度必须等于 sp，且两个列表长度必须一致，lse[i]与localOut[i]一一对应同一个 SP 域；
• localOut支持 FLOAT32/FLOAT16/BFLOAT16，而lse恒为 FLOAT32（softmax 统计量统一用高精度 FP32 表达），lseOut同样为 FLOAT32；
• 第一维batch * seqLen * headNum即文档与源码中的bsh语义（batch × seq × headNum），它也是 tiling 分核的基本单位；
• "非连续 Tensor"列标记为 x，表示不支持非连续张量：源码中第一段接口会通过l0op::Contiguous对每个输入张量做连续性规整后再送入计算图（见 aclnn_attention_update.cpp）；
• updateType与lseOut存在强绑定关系，详见下文错误码部分。

返回值与错误码语义

两段接口均返回aclnnStatus，通用返回码说明见 docs/zh/context/aclnn_return_code.md。第一段接口aclnnAttentionUpdateGetWorkspaceSize完成入参校验，出现以下场景时报错（整理自 aclnnAttentionUpdate.md）：

这些校验逻辑在源码中有完整的对应实现：

• 空指针检查：CheckNotNull逐元素校验lse、localOut列表内的每个 tensor 以及out（aclnn_attention_update.cpp）；
• 数据类型检查：CheckDtypeValid校验 lse 为 FLOAT、localOut/out 在 FLOAT/FLOAT16/BF16 范围内（aclnn_attention_update.cpp）；Ascend 950 的 regbase 分支CheckDtypeValid_95还额外要求 localOut 列表内各张量 dtype 一致、out 与 localOut[0] 一致（aclnn_attention_update.cpp）；
• sp 范围与列表长度一致性：CheckSp（sp ∈ [1, 128]）与CheckSp_95（Ascend 950 上 sp ∈ [1, 16]）（aclnn_attention_update.cpp）；
• shape 约束：CheckShape/CheckShape_95校验 lse 为 1 维、localOut/out 为 2 维、所有张量 shape 相同、localOut[1]为 8 的倍数且 ≤ 512、各张量第一维bsh一致、updateType == 1时 lseOut 为 1 维且第一维一致（aclnn_attention_update.cpp）；
• updateType 与 lseOut 绑定关系：CheckUpdateTypeAndLseOut（aclnn_attention_update.cpp）。

约束说明

使用aclnnAttentionUpdate需遵守以下约束（整理自 aclnnAttentionUpdate.md）：

• 确定性计算：aclnnAttentionUpdate默认为确定性实现（相关概念见 docs/zh/context/determinism_compute.md），相同输入可复现相同结果；
• sp 取值范围：
  • Atlas A2 训练系列产品 / Atlas A2 推理系列产品、Atlas A3 训练系列产品 / Atlas A3 推理系列产品：[1, 128]；
  • Ascend 950PR / Ascend 950DT：[1, 16]；
• headDim 取值范围：[8, 512]且是 8 的倍数；
• 支持空 Tensor：当lse[0]或localOut[0]为空张量时，第一段接口直接返回workspaceSize = 0并跳过计算（见 aclnn_attention_update.cpp）；
• 不支持非连续 Tensor：如上文所述，接口内部会先做Contiguous处理。

上述约束与 tiling 源码中的硬校验一一对应：attention_update_tiling.cpp中定义了D_MIN = 8、D_MAX = 512、D_DIVIDE_8 = 8、ATTR_SP_MAX = 16，并在CheckInputParams（updateType ∈ {0,1}、sp ∈ [1,16]）、CheckInputDim（H 维在 [8,512] 且为 8 的倍数、所有 lse/go 首维一致）、CheckInputDtype（go ∈ {FLOAT, FLOAT16, BF16}，lse 恒为 FLOAT）中执行（见 attention_update_tiling.cpp）。tiling 层 sp 上限取 16，而 aclnn 接口层在 A2/A3 上放宽到 128，两处范围差异由接口层的连续化与计算图拼接逻辑（多个l0op::AttentionUpdate组合）衔接。

完整调用示例与逐步讲解

文档给出了完整的可直接参考的调用示例（见 aclnnAttentionUpdate.md），仓库中另有可独立编译运行的简化版样例 examples/test_aclnn_attention_update.cpp。整体调用流程（含编译运行环境的准备请参考 docs/zh/context/compile_and_run_sample.md）如下：

1. 环境初始化（固定写法）

int Init(int32_t deviceId, aclrtStream* stream) { auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); ret = aclrtSetDevice(deviceId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return 0; }

依次完成aclInit、aclrtSetDevice、aclrtCreateStream，这是所有 aclnn 单算子调用的固定前置步骤。

2. 构造输入与输出

以sp = 2、bsh = 256、headDim = 128为例：

std::vector<int64_t> lseShape = {256}; // [batch * seqLen * headNum] std::vector<int64_t> localOutShape = {256, 128}; // [batch * seqLen * headNum, headDim] std::vector<int64_t> outShape = {256, 128}; int64_t updateType = 0; // 0：不输出 lseOut；1：输出 lseOut void* lseDeviceAddr[2] = {nullptr, nullptr}; void* localOutDeviceAddr[2] = {nullptr, nullptr}; void* outDeviceAddr = nullptr; std::vector<aclTensor*> lse = {nullptr, nullptr}; std::vector<aclTensor*> localOut = {nullptr, nullptr}; aclTensor* out = nullptr;

对每个张量调用aclrtMalloc申请 Device 侧内存、aclrtMemcpy拷入 host 数据，再通过aclCreateTensor创建aclTensor（连续张量需按 shape 计算 stride）：

template <typename T> int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size = GetShapeSize(shape) * sizeof(T); auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, ...); ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, ...); std::vector<int64_t> stride(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { stride[i] = shape[i + 1] * stride[i + 1]; } *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, stride.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; }

由于lse与localOut是 tensorList，需分别调用aclCreateTensorList将张量指针数组封装为aclTensorList：

aclTensorList *lseList = aclCreateTensorList(lse.data(), lse.size()); aclTensorList *localOutList = aclCreateTensorList(localOut.data(), localOut.size());

3. 两段式调用

先调用第一段接口获取 workspaceSize 与 executor（注意updateType == 0时lseOut传nullptr）：

uint64_t workspaceSize = 0; aclOpExecutor* executor; ret = aclnnAttentionUpdateGetWorkspaceSize(lseList, localOutList, updateType, out, nullptr, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAttentionUpdateGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);

按返回大小申请 workspace 后调用第二段接口：

void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, ...); } ret = aclnnAttentionUpdate(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnAttentionUpdate failed. ERROR: %d\n", ret); return ret);

4. 同步与结果回拷

ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, ...); auto size = GetShapeSize(outShape); std::vector<float> outData(size, 0); ret = aclrtMemcpy(outData.data(), outData.size() * sizeof(outData[0]), outDeviceAddr, size * sizeof(outData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, ...); for (int64_t i = 0; i < size; i++) { LOG_PRINT("out result[%ld] is: %f\n", i, outData[i]); }

最后通过aclrtDestroyStream、aclrtResetDevice、aclFinalize完成资源释放（完整释放逻辑见示例代码Finalize函数）。

源码级实现原理

算子图定义

在 attention_update_def.cpp 中：

• 输入lse（DYNAMIC，FLOAT/ND）与go（DYNAMIC，FLOAT/FLOAT16/BF16/ND）为动态个数的 tensorList，个数由属性sp决定；
• 输出output（与 go 同 dtype）与lse_m（FLOAT）；
• 属性update_type（默认 0）与sp（必填）；
• ascend950配置走attention_update_apt（regbase/动态编译路径，见ExtendCfgInfo("opFile.value", "attention_update_apt")），并开启DynamicRankSupportFlag、DynamicShapeSupportFlag、PrecisionReduceFlag。

接口层组装

第一段接口aclnnAttentionUpdateGetWorkspaceSize在完成校验后，将sp个输入张量逐个Contiguous规整，组装为 tensorList，交给l0op::AttentionUpdate构造计算图，并用l0op::ViewCopy将中间结果拷贝到用户提供的out/lseOut张量，最终通过uniqueExecutor->GetWorkspaceSize()返回所需 workspace 大小（见 aclnn_attention_update.cpp）。

Tiling 策略

attention_update_tiling.cpp实现了基于 bsh（batch × seq × headNum）维度的多核切分（attention_update_tiling.cpp）：

• 按CeilDiv(bshSize, totalCoreNum)计算每核处理量perCoreCount，尾核处理lastCoreCount；
• 依据 UB 容量与 double buffer（DOUBLE_BUFFER_NUM = 2）开销计算单次内循环可承载的 bsh 数bshInLoop，其中输入因子为sp × lse + sp × dAlign × go的双缓冲占用，中间计算用 FP32 提升精度，输出同样开双缓冲；
• 预留sp * ubBlockSize空间防止 lse 对齐搬入 UB 后占用膨胀；
• 通过context->SetBlockDim(usedCoreNum)设置核数，workspace 固定申请16 * 1024 * 1024字节（SYS_WORKSPACE_SIZE，见 attention_update_tiling.cpp）。

tiling data 中完整记录 sp、d、usedCoreNum、perCoreCount、lastCoreCount、perCoreLoops、lastCoreLoops、perCorePerLoopCount、bshInLoop 等切分参数（定义见 attention_update_tiling.h），并按 tiling key20000 + updateType注册（空张量场景为 10000）。

Kernel 执行

NPU kernel 入口在 op_kernel/attention_update.cpp，仅使用 AIV（KERNEL_TYPE_AIV_ONLY），按 tiling key 实例化不同精度的DecodeUpdate<float, goType>模板执行。Ascend 950 的 regbase 实现（arch35/attention_update_with_lse_regbase.h）展示了核心计算过程：

• 通过ListTensorDesc读取 sp 个 lse/go 全局张量地址；
• ComputeMaxVF用向量寄存器对 sp 个 lse 逐元素求Max（将+INF先替换为-INF做保护），随后逐域做Sub → Exp → Add累加、Log，得到全局lse_m并回写；
• ComputeOutVF按UNROLL_NUM = 2展开 sp 域，用MulAddDst完成O_i × exp(lse_i - lse_m)的加权累加，得到全局输出 O；
• 数据搬入搬出全部采用双缓冲（BUFFER_NUM = 2）与DataCopyPad对齐，d 维按 UB block 对齐（CeilAlign(d, goBlockNum)）。

测试验证

仓库为该算子提供了完善的 ST/UT 覆盖，可用于验证正确性与精度：

• ST 用例：atk_aclnnAttentionUpdate.json中配置了sp = 2、bsh = 20480、headDim = 512、updateType = 0的边界用例，要求high_precision精度（见 tests/st/aclnnAttentionUpdate/atk_aclnnAttentionUpdate.json）；
• CPU 参考实现：executor_aclnnAttentionUpdate.py用 PyTorch 严格按合并公式实现 reference，供 NPU 结果比对（executor_aclnnAttentionUpdate.py）；
• bf16/fp16 专项：tests/st/aclnnAttentionUpdate_bf16fp16/目录覆盖半精度 localOut 场景；
• UT 测试：tests/ut/下分别对 tiling（test_attention_update_tiling.cpp）、kernel（test_attention_update.cpp）、aclnn 接口层（test_aclnn_attention_update.cpp）提供单元测试。

小结

aclnnAttentionUpdate是 CANN ops-transformer 中面向序列并行 Attention 的"收尾"算子，通过 log-sum-exp 形式的跨 SP 域合并，将各域的局部 lse 与局部输出精确融合为全局结果。使用时需牢记：lse 恒为 FP32、localOut 支持 FP16/BF16/FP32，headDim 必须为 [8, 512] 内 8 的倍数，sp 上限在 A2/A3 上为 128、在 Ascend 950 上为 16，updateType与lseOut的传参必须严格配对。掌握两段式调用流程与上述约束后，即可在序列并行推理/训练链路中正确接入该算子，实现跨域 Attention 结果的全局归一化合并。

【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer