AttentionWorkScheduler【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer产品支持情况产品是否支持Ascend 950PR/Ascend 950DT√Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×Atlas 推理系列产品√Atlas 训练系列产品√功能说明算子功能Attention和FFN分离部署场景下Attention侧数据扫描算子。该算子接收来自FFNToAttention算子的输出数据并对数据进行逐步扫描确保数据准备就绪。该算子不建议单独使用建议与FFNToAttention和AttentionWorkerCombine算子配合使用形成完整的工作流。接收FFNToAttention算子发送的数据。该数据以ScheduleContext结构体存储。该结构体包含CommonAreaControlAreaAttentionAreaFfnArea域。本接口涉及CommonArea(用于存储配置信息如session_nummicro_batch_nummicro_batch_sizeselected_expert_num)ControlArea(用于上层控制进程是否退出)AttentionArea域(负责管理算子计算过程中所需的核心数据缓冲区与状态信息其中token_info_buf存储了与输入相关的数据信息)。读取ScheduleContext.AttentionArea域中token_info_buf存储的flag信息查看通信数据是否准备就绪。数据全部准备就绪后后续可供AttentionWorkerCombine算子使用。计算公式$$ \text{Initialize:} \quad \text{ready_count} 0, \quad \text{flag_num} \text{micro_batch_size} \times \text{selected_expert_num} $$$$ \text{Check if run_flag is 0:} \quad \text{if run_flag} 0, \quad \text{exit and log} $$$$ \text{Loop:} \quad \text{while run_flag} \neq 0: \quad \text{ready_count} \sum_{i1}^{\text{flag_num}} \mathbf{1}_{{ \text{flag}[i] 1 }}; \quad \text{if ready_count} \text{flag_num}, \quad \text{break} $$$$ \text{Reset flags:} \quad \text{flag}[i] 0 \quad \text{for} \quad i 1, 2, \dots, \text{flag_num} $$$$ \text{Set micro_batch_id:} \quad \text{micro_batch_id} (\text{micro_batch_id} 1) % \text{micro_batch_num} $$备注micro_batch_size、selected_expert_num、run_flag、micro_batch_id是入参ScheduleContext结构体的参数该结构体信息在调用示例中进行展示说明。参数说明参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorscheduleContextRef输入/输出Attention侧接收的待处理数据表示输入scheduleContext信息详细结构见调用示例。不支持空Tensor。INT8ND1维shape固定为(1024)×workspaceSize输出返回需要在Device侧申请的workspace大小。-----executor输出返回op执行器包含了算子计算流程。-----约束说明无。调用说明调用方式样例代码说明aclnn接口test_aclnn_inplace_attention_worker_scheduler通过aclnnInplaceAttentionWorkerScheduler接口方式调用AttentionWorkScheduler算子。【免费下载链接】ops-transformer本项目是CANN提供的transformer类大模型算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-transformer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考