CANN ops-transformer 算子 aclnnAttentionWorkerCombine 详解Attention/FFN 分离部署下的 Token 融合加权算子【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer本技术指南围绕 CANN ops-transformer 仓库中的 aclnnAttentionWorkerCombine 算子文档展开完整讲解该算子的功能定位、两段式接口用法、参数语义、约束条件与可运行的调用示例并结合仓库源码算子定义、shape/dtype 推导、tiling 与 kernel 实现深入剖析其底层原理。读者读完可以掌握在 Attention 与 FFN 分离部署场景下如何通过 aclnn 接口完成注意力 token 数据的加权融合与层 ID 更新并将其接入 AttentionWorkerScheduler 等算子组成的完整工作流。功能概述Attention 侧的数据融合算子AttentionWorkerCombine是 CANN ops-transformer 中面向Attention 与 FFN 分离部署场景的 Attention 侧数据融合算子。当大模型以流水线方式将 Attention 计算单元与 FFN 计算单元拆分到不同设备/核上部署时Attention 侧往往需要接收 FFN 侧回传的、由多个计算单元分别处理产生的注意力 token 数据。该算子负责接收 FFN 侧回传的数据数据以ScheduleContext结构体内存排布方式存储。该结构体包含CommonArea、ControlArea、AttentionArea、FfnArea四个域算子从AttentionArea的token_data_buf中读取 token 数据。加权融合结合专家权重expertScales对多路 token 数据进行加权求和输出最终的注意力融合结果y。更新层 ID根据输入layerId计算并输出nextLayerId指示下一个要处理的层 ID。计算公式如下$$ y[i] \sum_{k0}^{K-1} \text{expertScales}[i][k] \times \text{token_data}[i][k] $$$$ \text{nextLayerId} \text{layerId} 1 $$其中K为每个 token 对应的专家数即expertScales的第二维大小i遍历整个 batch 的 token。重要提示该算子不建议单独使用建议与AttentionWorkerScheduler等算子配合使用形成完整的工作流。AttentionWorkerScheduler负责扫描并确认 FFN 侧数据是否准备就绪通过轮询AttentionArea.token_info_buf中的 flag而AttentionWorkerCombine在数据就绪后完成读取、融合与输出两者共同构成 Attention 侧完整的数据接收-消费链路。相关算子说明见 attention_worker_scheduler/README.md。产品支持情况AttentionWorkerCombine在不同 NPU 产品上的支持情况如下与仓库中算子定义op_host/config下的 aicore 配置对应产品是否支持Ascend 950PR / Ascend 950DT不支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品不支持Atlas 训练系列产品不支持从源码角度看attention_worker_combine_def.cpp 中通过this-AICore().AddConfig(ascend910_93)与this-AICore().AddConfig(ascend910b)注册了 aicore 配置与文档中Atlas A3 / A2 系列支持的产品矩阵一致。算子对应不同产品的配置文件位于 op_host/config/ascend910_93 与 op_host/config/ascend910b 目录下各含attention_worker_combine.json与attention_worker_combine_simplified_key.ini。两段式接口与函数原型每个 aclnn 算子都采用两段式接口必须先调用aclnnAttentionWorkerCombineGetWorkspaceSize接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnAttentionWorkerCombine接口执行计算。两段式接口的通用机制可参考 两段式接口说明。第一段接口原型aclnnStatus aclnnAttentionWorkerCombineGetWorkspaceSize( const aclTensor *scheduleContext, const aclTensor *expertScales, const aclTensor *layerId, int64_t hiddenSize, int64_t tokenDtype, int64_t needSchedule, const aclTensor *y, const aclTensor *nextLayerId, uint64_t *workspaceSize, aclOpExecutor **executor)第二段接口原型aclnnStatus aclnnAttentionWorkerCombine( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, const aclrtStream stream)第一段接口负责入参校验、shape/dtype 推导与 tiling 计算并将执行流程封装进executor第二段接口在指定stream上真正下发算子任务。aclnnAttentionWorkerCombineGetWorkspaceSize 参数说明参数表参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorscheduleContext输入Attention 侧接收的调度上下文内含 CommonArea、ControlArea、AttentionArea、FfnArea。算子从 AttentionArea 的 token_data_buf 中读取 token 数据。不支持空 tensor。INT8ND1维shape 固定为 (1024)×expertScales输入专家权重表示每个 token 对应的各专家权重。-FLOATND2维(BS, K)×layerId输入当前的模型层 ID。-INT32ND1维(1)×hiddenSize输入token_data 的隐藏维度大小用于确定输出 y 的第二维大小。-INT64---tokenDtype输入指定 scheduleContext 中 token 数据的原始精度类型。0 表示 FLOAT161 表示 BFLOAT16。取值为 0 或 1。INT64---needSchedule输入指定是否等待 token 数据填充完成后再执行。0 表示不等待1 表示等待。取值为 0 或 1。INT64---y输出最终的注意力合并结果。-FP16、BF16ND2维(BS, hiddenSize)×nextLayerId输出下一个要处理的层 ID。-INT32ND1维(1)×workspaceSize输出返回需要在 Device 侧申请的 workspace 大小。-----executor输出返回 op 执行器包含了算子计算流程。-----几个参数需要注意scheduleContext是1024 字节的固定长度 INT8 张量本质上是ScheduleContext结构体在 Device 侧的字节镜像见下文调用示例中的结构体定义与static_assert(sizeof(ScheduleContext) 1024)。expertScales的 shape 为(BS, K)其中BS为 batch sizeK为每个 token 的专家数它决定了输出y的第一维大小而y的第二维由hiddenSize决定。这一点与 attention_worker_combine_infershape.cpp 中的InferShape4AttentionWorkerCombine实现完全一致yShape-SetDim(0, expertScalesInputShape-GetDim(0))、yShape-SetDim(1, *hiddenSize)。tokenDtype同时决定输出y的数据类型由 InferDtype4AttentionWorkerCombine 可知tokenDtype 1BF16时y为DT_BF16否则为DT_FLOAT16nextLayerId的类型直接继承输入layerId的类型。返回值两段接口均返回aclnnStatus状态码具体参见 aclnn 返回码。第一段接口会完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001输入是空指针。ACLNN_ERR_PARAM_INVALID161002输入数据类型不在支持的范围内。aclnnAttentionWorkerCombine 参数说明第二段接口的参数相对简单均为执行阶段所需的运行时资源参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnAttentionWorkerCombineGetWorkspaceSize获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值同样为aclnnStatus状态码参见 aclnn 返回码。约束说明expertScales的第二维 K ≤ 64。该限制在 tiling 阶段被强制执行见 attention_worker_combine_tiling.cpp 中的OP_CHECK_IF(k K_UPPER_BOUND, ...)校验K_UPPER_BOUND 64。确定性计算aclnnAttentionWorkerCombine为默认确定性实现。scheduleContext为 1D 张量shape 固定为(1024)。expertScales为 2D 张量[BatchSize, K]y为 2D 张量[BatchSize, HiddenSize]layerId与nextLayerId均为 1D 张量。调用示例以下示例代码演示了完整的调用流程该示例同样保存在仓库 examples/test_aclnn_attention_worker_combine.cpp 中可直接对照阅读。具体编译和执行过程请参考 编译与运行样例。1. 头文件与辅助宏#include iostream #include vector #include cstring #include acl/acl.h #include aclnnop/aclnn_attention_worker_combine.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0)2. 辅助函数初始化、环境释放与张量创建int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } 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; } void Finalize(int32_t deviceId, aclrtStream stream) { aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_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, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); std::vectorint64_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; } int CreateAclTensorNoData(const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { uint64_t elemSize sizeof(int8_t); if (dataType ACL_INT32) { elemSize sizeof(int32_t); } if (dataType ACL_FLOAT16) { elemSize sizeof(int16_t); } if (dataType ACL_BF16) { elemSize sizeof(int16_t); } auto size GetShapeSize(shape) * elemSize; auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); std::vectorint64_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; }3. ScheduleContext 结构体定义内存布局ScheduleContext是算子与调度工作流之间的契约它按#pragma pack(push, 1)紧凑排布总大小固定为1024 字节通过static_assert强制校验。四个域分别承担不同职责CommonArea全局配置如session_numAttention 节点数、micro_batch_numbatch 拆分份数、micro_batch_size、selected_expert_numtopK 1、expert_num每层专家数、attn_to_ffn_token_size/ffn_to_attn_token_sizetoken 在对方 window 数据区的存储占用、schedule_mode0 仅调度 FFN1 仅调度 Attention2 同时调度。ControlArearun_flag控制循环退出0 退出1 运行中。AttentionArea核心数据缓冲区与状态信息。token_info_buf存放每个 micro batch 的DataDesc内含每个 token×专家对应的就绪 flagtoken_data_buf存放[M, BS, K1, HS]的 token 数据micro_batch_id记录最新就绪的 micro batch id。FfnAreaFFN 侧回传相关的缓冲区指针与out_num等状态。#pragma pack(push, 1) struct AttentionDataDesc { int32_t flag[0]; }; struct ScheduleContext { struct CommonArea { uint32_t session_num; // Number of attention nodes uint32_t micro_batch_num; uint32_t micro_batch_size; uint32_t selected_expert_num; // topK 1 uint32_t expert_num; // experts per layer uint32_t attn_to_ffn_token_size; uint32_t ffn_to_attn_token_size; int32_t schedule_mode; // 0: Ffn only 1: Attention only int8_t reserve0[96]; }; struct ControlArea { int32_t run_flag; // 0 : exited 1 : running int8_t reserve2[124]; }; struct AttentionArea { uint64_t token_info_buf; // Points to device memory. uint64_t token_info_buf_size; uint64_t token_data_buf; // Points to device memory. uint64_t token_data_buf_size; uint32_t micro_batch_id; // Records the latest ready micro batch id. int8_t reserve5[92]; }; struct FfnArea { uint64_t token_info_buf; uint64_t token_info_buf_size; uint64_t token_data_buf; uint64_t token_data_buf_size; uint64_t polling_index; int8_t reserve3[88]; uint64_t layer_ids_buf; uint64_t layer_ids_buf_size; uint64_t session_ids_buf; uint64_t session_ids_buf_size; uint64_t micro_batch_ids_buf; uint64_t micro_batch_ids_buf_size; uint64_t expert_ids_buf; uint64_t expert_ids_buf_size; uint32_t out_num; int8_t reserve4[60]; }; CommonArea common; ControlArea control; AttentionArea attention; FfnArea ffn; int8_t reserve6[384]; // Padding to 1024 bytes. }; static_assert(sizeof(ScheduleContext) 1024, ScheduleContext size must be 1024 bytes); #pragma pack(pop)该结构与 kernel 侧的定义严格对齐见 attention_worker_combine_common_utils.hkernel 通过GET_OFFSET_B32/GET_OFFSET_B64宏按成员偏移直接读取schedule_context中的字段因此 Host 侧结构体的字段顺序与 padding 必须与 kernel 侧保持一致。4. main 函数完整调用流程int main() { int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 int64_t BS 48; // batch size int64_t K 8; // expert num per token int64_t hiddenSize 20480; int64_t tokenDtype 1; // BF16 int64_t needSchedule 0; // 初始化ScheduleContext ScheduleContext scheduleContext {}; scheduleContext.common.session_num 1; scheduleContext.common.micro_batch_num 1; scheduleContext.common.micro_batch_size BS; scheduleContext.common.selected_expert_num K; scheduleContext.common.expert_num 16; scheduleContext.common.attn_to_ffn_token_size 512; scheduleContext.common.ffn_to_attn_token_size 512; scheduleContext.common.schedule_mode 1; // Attention only scheduleContext.control.run_flag 1; // running scheduleContext.attention.micro_batch_id 0; // 初始化Attention token_info_bufflag置1表示数据就绪 size_t perDataDescSize sizeof(AttentionDataDesc) sizeof(int32_t) * BS * K; size_t tokenInfoBufSize static_castsize_t(scheduleContext.common.micro_batch_num) * perDataDescSize; void* tokenInfoBuf nullptr; ret aclrtMalloc(tokenInfoBuf, tokenInfoBufSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(malloc token info buf failed. ERROR: %d\n, ret); return ret); scheduleContext.attention.token_info_buf reinterpret_castuint64_t(tokenInfoBuf); scheduleContext.attention.token_info_buf_size tokenInfoBufSize; std::vectorint32_t hostFlags(static_castsize_t(BS) * K, 1); ret aclrtMemcpy(tokenInfoBuf, tokenInfoBufSize, hostFlags.data(), static_castsize_t(BS) * K * sizeof(int32_t), ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(cpy token info buf failed. ERROR: %d\n, ret); return ret); // 初始化Attention token_data_buf uint64_t tokenDataSize static_castuint64_t(BS) * K * hiddenSize * sizeof(int16_t); void* tokenDataBuf nullptr; ret aclrtMalloc(tokenDataBuf, tokenDataSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(malloc token data buf failed. ERROR: %d\n, ret); return ret); scheduleContext.attention.token_data_buf reinterpret_castuint64_t(tokenDataBuf); scheduleContext.attention.token_data_buf_size tokenDataSize; std::vectorint16_t hostTokenData(static_castsize_t(BS) * K * hiddenSize, 1); ret aclrtMemcpy(tokenDataBuf, tokenDataSize, hostTokenData.data(), tokenDataSize, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(cpy token data buf failed. ERROR: %d\n, ret); return ret); // 创建scheduleContext aclTensor std::vectorint64_t scheduleContextShape {1024}; void* scheduleContextDeviceAddr nullptr; aclTensor* scheduleContextRef nullptr; std::vectorint8_t hostCtxData(1024, 0); std::memcpy(hostCtxData.data(), scheduleContext, sizeof(ScheduleContext)); ret CreateAclTensor(hostCtxData, scheduleContextShape, scheduleContextDeviceAddr, aclDataType::ACL_INT8, scheduleContextRef); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建expert_scales aclTensor, shape (BS, K) std::vectorint64_t expertScalesShape {BS, K}; std::vectorfloat hostExpertScales(static_castsize_t(BS) * K, 0.125f); void* expertScalesDeviceAddr nullptr; aclTensor* expertScalesRef nullptr; ret CreateAclTensor(hostExpertScales, expertScalesShape, expertScalesDeviceAddr, aclDataType::ACL_FLOAT, expertScalesRef); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建layer_id aclTensor, shape (1) std::vectorint64_t layerIdShape {1}; std::vectorint32_t hostLayerId {0}; void* layerIdDeviceAddr nullptr; aclTensor* layerIdRef nullptr; ret CreateAclTensor(hostLayerId, layerIdShape, layerIdDeviceAddr, aclDataType::ACL_INT32, layerIdRef); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建输出y aclTensor, shape (BS, hiddenSize) std::vectorint64_t yShape {BS, hiddenSize}; void* yDeviceAddr nullptr; aclTensor* yRef nullptr; ret CreateAclTensorNoData(yShape, yDeviceAddr, aclDataType::ACL_BF16, yRef); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建输出next_layer_id aclTensor, shape (1) std::vectorint64_t nextLayerIdShape {1}; void* nextLayerIdDeviceAddr nullptr; aclTensor* nextLayerIdRef nullptr; ret CreateAclTensorNoData(nextLayerIdShape, nextLayerIdDeviceAddr, aclDataType::ACL_INT32, nextLayerIdRef); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnAttentionWorkerCombineGetWorkspaceSize(scheduleContextRef, expertScalesRef, layerIdRef, hiddenSize, tokenDtype, needSchedule, yRef, nextLayerIdRef, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAttentionWorkerCombineGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } ret aclnnAttentionWorkerCombine(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnAttentionWorkerCombine failed. ERROR: %d\n, ret); return ret); // 4.固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出结果将device侧内存上的结果拷贝至host侧 int32_t nextLayerId 0; ret aclrtMemcpy(nextLayerId, sizeof(int32_t), nextLayerIdDeviceAddr, sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy next_layer_id failed. ERROR: %d\n, ret); return ret); LOG_PRINT(next_layer_id %d.\n, nextLayerId); // 6. 释放aclTensor aclDestroyTensor(scheduleContextRef); aclDestroyTensor(expertScalesRef); aclDestroyTensor(layerIdRef); aclDestroyTensor(yRef); aclDestroyTensor(nextLayerIdRef); // 7. 释放device资源 aclrtFree(scheduleContextDeviceAddr); aclrtFree(expertScalesDeviceAddr); aclrtFree(layerIdDeviceAddr); aclrtFree(yDeviceAddr); aclrtFree(nextLayerIdDeviceAddr); aclrtFree(tokenInfoBuf); aclrtFree(tokenDataBuf); if (workspaceSize 0) { aclrtFree(workspaceAddr); } Finalize(deviceId, stream); return 0; }示例中BS48、K8、hiddenSize20480、tokenDtype1BF16、needSchedule0不等待因为示例中数据已由 Host 预先填充。若实际工作流中数据由 FFN 侧异步回传应将needSchedule置为 1算子会轮询token_info_buf中的 flag 直到数据全部就绪。源码级实现解析算子定义op_defattention_worker_combine_def.cpp 中通过OpDef注册了算子的输入输出与属性输入schedule_contextINT8、expert_scalesFLOAT、layer_idINT32均要求 ND 格式并AutoContiguous输出yFLOAT16 / BF16、next_layer_idINT32属性hidden_size必选 Int、token_dtype可选 Int默认 0、need_schedule可选 Int默认 0。这与 aclnn 接口的形参一一对应hiddenSize/tokenDtype/needSchedule三个int64_t参数即映射为上述三个属性。Shape 与 Dtype 推导infershapeattention_worker_combine_infershape.cpp 完成了图模式下的推导逻辑y的 shape 为(expertScales[0], hiddenSize)即第一维取自expertScales的 BS第二维取自属性hiddenSizenextLayerId的 shape 与layerId相同[1]y的数据类型由tokenDtype决定1 → BF16否则 → FP16nextLayerId的类型继承layerId。Tiling 策略按 BS / H / K 三维切分attention_worker_combine_tiling.cpp 中的DoOpTiling根据 UB 容量与 batch 大小在三种切分策略中选择并生成对应的 tiling keyBS 全载 按 BS 分核当hiddenSize对齐后不超过 UB 可容纳的hInFullK上限时K 与 H 均整块载入只对 BS 维按核数切分tiling key10000/10001FP16/BF16按 H 分核BS 分核后核数未达上限且 H 过大时对 H 维切块、跨多核并行并支持needSchedule模式下的循环切块tiling key10010/10011按 K 切分H 可整块容纳但 K 过大时对 K 维分组循环处理tiling key10020/10021。每个 tiling key 在 attention_worker_combine.cpp 的 kernel 入口中分派到对应的模板实现KernelAttentionWorkerCombineSplitBS、KernelAttentionWorkerCombineSplitH、KernelAttentionWorkerCombineSplitK分别以halfFP16与bfloat16_tBF16实例化。此外PostTiling中还会设置DEFAULT_WORKSPACE_SIZE 32字节的 workspace。Kernel 计算逻辑以 attention_worker_combine_split_bs.h 的Process为例其核心流程为按GetBlockIdx()计算本核负责的 BS 区间tailCoreBsLoopNum处理尾块block 0 额外计算nextLayerIddstNextLayerIdGm(0) srcLayerIdGm(0) 1对应文档中的nextLayerId layerId 1若needSchedule 1通过ScanTokenInfo轮询token_info_buf中全部bsLoopNum * (K 1)个 flag累加就绪数直到全部为 1 才继续对应 Scheduler 工作流中的等待语义计算完成后调用ClearTokenInfo将 flag 清零并由 block 0 回写micro_batch_id对每个 BS 循环CopyIn载入(K1)路 token 数据 →Compute完成加权求和 →CopyOut写出y。Compute第 201-233 行的实现正是公式 $y[i] \sum_{k0}^{K-1} \text{expertScales}[i][k] \times \text{token_data}[i][k]$ 的向量化展开对前 K 路数据依次执行CastFP16/BF16 → FP32→ 从expertScales取权重Muls缩放 →Add累加最后加上第K路K1数据布局中索引为 K 的 slot通常为 0 或占位数据再Cast回输出精度写回y。与 AttentionWorkerScheduler 组成的工作流AttentionWorkerCombine是 Attention/FFN 分离部署工作流中的消费端算子建议按以下链路使用FFN 侧计算完成后通过 FfnToAttention 类算子将 token 数据与就绪 flag 写入ScheduleContext的 Attention 域缓冲区AttentionWorkerScheduler见 attention_worker_scheduler/README.md轮询token_info_buf中的 flag确认数据全部就绪并维护micro_batch_id的推进AttentionWorkerCombine读取就绪的token_data_buf结合expertScales完成多路 token 的加权融合输出y并更新nextLayerId驱动模型进入下一层处理。当needSchedule 1时AttentionWorkerCombine自身也会在 kernel 内部执行就绪 flag 的轮询与清零从而与 Scheduler 形成双保险的同步机制在显式调度数据由上层确保就绪的场景下可置needSchedule 0跳过该开销。总结aclnnAttentionWorkerCombine是 CANN ops-transformer 中实现 Attention/FFN 分离部署下 token 融合的关键算子它以固定 1024 字节的ScheduleContext为数据契约从AttentionArea.token_data_buf读取多路注意力 token 数据按expertScales加权求和输出y并顺带完成层 ID 的推进。使用时应牢记两段式接口调用次序、scheduleContext固定 shape(1024)、K ≤ 64以及tokenDtype与输出精度的联动关系并结合AttentionWorkerScheduler组成完整工作流。仓库中同时提供了完整的 aclnn 调用样例 test_aclnn_attention_worker_combine.cpp、图模式构图样例 test_geir_attention_worker_combine.cpp对应算子 IR 定义 attention_worker_combine_proto.h以及 Host 侧单测 test_attention_worker_combine_infershape.cpp 和 test_attention_worker_combine_tiling.cpp可作为继续深入研究的起点。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考