CANN ATB all_to_allvv2 算子深度解析基于 HCCL 的可变长度 AllToAll 集合通信增强版【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost本文面向需要在华为昇腾Ascend多卡集群上做全交换All-to-All通信的开发者和 Agent系统讲解 ascend-transformer-boost 中all_to_allvv2算子的设计定位、参数结构、张量契约、校验逻辑与底层 HCCL 调用链。读完本文你将掌握该算子的完整数据流从 Operation 参数校验、Shape 推导到HcclAlltoAllV落盘执行能够正确构造 6 个输入张量并配置通信域也能理解它与all_to_all、all_to_allv之间的演进关系。1. 算子定位通信域内的可变长度数据全交换all_to_allvv2是 ATBAscend Transformer Boost推理侧ops_infer提供的一个集合通信算子知识条目将其归类为communication类型、tier: S、type: single单一算子非融合算子底层仅依赖 HCCL 通信库详见 .agent/knowledge/ops/communication/all_to_allvv2/index.md。它解决的是典型的全交换All-to-All问题通信域内的每一张卡都向其他所有卡发送各自定制数量的数据同时从其他所有卡接收定制数量的数据。相比等长版all_to_all其每次通信的数据量count与偏移displs均可按目标 rank 定制属于可变长度variable-lengthAllToAll 语义相比all_to_allv它在参数组织、校验与增强能力上做了演进是当前仓库中的v2 增强版。源码目录src/ops/ops_infer/all_to_allvv2/共 4 个文件operation 定义 hccl runner参数头文件include/atb/infer_op_params.h相关算子all_to_all等长版双后端 HCCL/LCCL见 src/ops/ops_infer/all_to_all/、all_to_allv可变长度版仅 HCCL见 src/ops/ops_infer/all_to_allv/从源码结构看与all_to_all支持 HCCL/LCCL 双后端不同all_to_allvv2只实现了 HCCL 单后端知识条目中的 仅 HCCL 即源于此。2. 参数结构 AllToAllVV2Param 详解算子的超参数定义在 include/atb/infer_op_params.h 的infer::AllToAllVV2Param结构体中全部字段如下字段类型默认值含义与取值说明rankint-1当前卡在通信域内的编号。-1表示未指定单卡场景rankSize 1时会被自动修正为 0多卡场景必须显式指定否则参数校验直接报错rankSizeint0参与通信的卡的数量必须与 sendCounts/recvCounts/sdispls/rdispls 等数组的长度一致rankRootint0主通信编号root rank用于构建通信域时的根节点定位backendstd::stringhccl通信计算类型当前仅支持hccl传入其他值会被 ParamCheck 拒绝hcclCommHcclCommnullptrHCCL 通信域指针。默认为空此时由加速库根据 rank/rankSize/rankTableFile 自行创建若用户希望自行管理通信域则传入该指针加速库将直接复用commModeCommModeCOMM_MULTI_PROCESS通信模式枚举值。注意HCCL 多线程场景只支持外部传入通信域的方式rankTableFilestd::string空集群信息配置文件路径适用单机及多机通信场景当前仅支持 hccl 后端。若单机配置了 rankTable则以 rankTable 初始化通信域commDomainstd::string空通信 device 组用通信域名标识多通信域场景使用当前仅支持 hcclrsv[64]uint8_t[64]全 0预留字段该结构体还提供了operator用于参数相等性比较比较 rank、rankSize、rankRoot、hcclComm、commMode、backend、rankTableFile、commDomain供 ATB 的算子参数缓存/复用机制使用。3. 输入输出张量契约6 入 1 出算子的张量接口由 ops_configs/atb_ops_info.ini 中的[AllToAllVV2Operation]配置定义源码中通过IN_TENSOR_NUM 6、OUT_TENSOR_NUM 1all_to_allvv2_operation.cpp与之对应序号名称dtypeformat说明input0xfloat16 / float / int8 / int32 / int16 / int64 / bf16nd本卡待发送的数据张量device 侧input1sendCountint64nd长度为rankSize的 host 数组指定发给每个目标 rank 的元素个数input2sdisplsint64nd长度为rankSize的 host 数组指定每个发送分片在输入张量中的起始偏移input3recvCountint64nd长度为rankSize的 host 数组指定从每个源 rank 接收的元素个数input4rdisplsint64nd长度为rankSize的 host 数组指定每个接收分片在输出张量中的起始偏移input5tensorForInferShapeint8nd仅用于 shape 推导的辅助张量其dims[0]必须等于sum(recvCounts)dtype 恒为 int8output0output与 input0 相同nd接收并汇聚后的数据张量device 侧几点关键设计count/displs 全部走 host 内存sendCount、sdispls、recvCount、rdispls四个张量在 runner 中通过hostData访问见 all_to_allvv2_hccl_runner.cpp它们是纯 host 侧的控制元数据不占用 device 内存。tensorForInferShape专为 InferShape 服务由于 count/displs 是 host 数据在纯 shape 推导阶段无法可靠读取因此额外引入一个 device 张量使其dims[0]携带sum(recvCounts)的信息从而推导出输出的第二维详见下节。输入数据 dtype 覆盖面广支持 float16、float、int8、int32、int16、int64、bf16 共 7 种输入输出 dtype 一一对应。4. 算子实现剖析校验、Shape 推导与 Runner 决策算子的核心逻辑集中在 all_to_allvv2_operation.cpp包含四条主线。4.1 参数校验 ParamCheck构造函数之外文件内的匿名命名空间定义了ParamCheckL24-L40backend必须等于hccl否则报错 backend must be hccl当rank -1且rankSize 1时报错 multi-card, rank must be specified——即多卡场景不允许省略 rank 号随后调用OperationUtil::DistributedInitCheckAllToAllVV2Param做分布式初始化相关检查。4.2 Shape 推导 InferShapeImplInferShapeImpl 逻辑非常简洁输出张量直接继承输入 0 的描述但强制改写为二维[1, dims0(tensorForInferShape)]即输出 shape 为[1, sum(recvCounts)]。这从侧面说明该算子要求最终输出在逻辑上按一行、sum(recvCounts) 列理解各 rank 的接收数据按rdispls指定的偏移落位到这一行中。4.3 张量校验 SetupCheckImplSetupCheckImpl 承担运行时校验可从源码逐条提取数组长度校验sendCounts、sdispls、recvCounts、rdispls的dims[0]都必须等于param_.rankSize非负校验四个数组中每个元素都必须 0求和溢出校验累加sum(recvCounts)过程中检查 int64 溢出且要求sum(recvCounts) 0接收边界校验对每个 rank要求recvCounts[i] rdispls[i] sum(recvCounts)防止接收分片越出输出范围发送边界校验对每个 rank要求sendCounts[i] sdispls[i] inputCount输入张量总元素数由辅助函数AllToAllVV2CalculateTensorSize计算输出维度校验outTensors[0].shape.dims[1]必须等于sum(recvCounts)。这些检查在 Setup 阶段SetupCheckImpl执行把绝大多数越界、负数和长度不匹配的错误拦截在执行之前避免把非法参数直接传给 HCCL。4.4 Runner 决策 CreateRunnerCreateRunner 依据参数选择执行后端仅当backend hccl时返回AllToAllVV2HcclRunner当hcclComm nullptr时走由加速库创建通信域路径构造参数为(param_, !param_.rankTableFile.empty())即根据是否配置了rankTableFile决定使用 rankTable 初始化还是 rank/rankSize/rankRoot/commDomain 初始化当外部传入hcclComm时走复用用户通信域路径构造参数为(param_, hcclComm)其他 backend 一律返回空指针配合 ParamCheck 的双重保障。此外构造函数会根据硬件平台选择不同的算子 IR 配置L52-L56310P 平台使用AllToAllVV2Operation310p其余平台使用AllToAllVV2Operation两个配置在 ops_configs/atb_ops_info.ini 中均有定义310p 变体仅支持 float16/int8。5. Runner 与底层 HCCL 调用链AllToAllVV2HcclRunner继承自HcclRunnerall_to_allvv2_hccl_runner.h真正的执行发生在 ExecuteImpl其调用链为AllToAllVV2Operation::Execute └─ AllToAllVV2HcclRunner::ExecuteImpl ├─ 校验 hcclComm_ 非空 ├─ 校验 input0.deviceData / output0.deviceData 非空 └─ HcclAlltoAllV( input0.deviceData, // sendbuf input1.hostData, // sendCounts input2.hostData, // sdispls GetHcclDtype(inDtype), // 输入 dtype output0.deviceData, // recvbuf input3.hostData, // recvCounts input4.hostData, // rdispls GetHcclDtype(outDtype), // 输出 dtype hcclComm_.get(), // 通信域 GetExecuteStream(context) // 执行流 )几个实现要点语义映射ATB 层的sendCount/sdispls/recvCount/rdispls与 HCCLHcclAlltoAllV的参数一一对应ATB 在此处的职责是参数组织、校验与通信域管理不涉及 kernel 计算因此该算子没有独立 kernel 目录知识条目中 operation hccl_runner 的结构即源于此。错误处理HcclAlltoAllV返回非HCCL_SUCCESS时通过ConvertHcclResultToStatus转换为 ATB 的Status并记录日志hcclComm_为空或 device 张量为空时分别返回ERROR_INTERNAL_ERROR/ERROR_INVALID_PARAM。注册机制文件末尾通过REG_RUNNER_TYPE(AllToAllVV2HcclRunner)将 Runner 注册进 ATB 的 Runner 工厂all_to_allvv2_hccl_runner.cpp供CreateRunner统一创建。6. 通信域初始化三种方式与适用场景由AllToAllVV2HcclRunner的三个构造函数可见通信域有三种初始化途径方式触发条件初始化依据适用场景rankTable 文件hcclComm nullptr且rankTableFile非空HcclRunner(name, rank, rankTableFile, commDomain)单机/多机集群按集群信息文件建域rank 参数组hcclComm nullptr且rankTableFile为空HcclRunner(name, rank, rankSize, rankRoot, commDomain)单机多卡按 rank/rankSize/rankRoot 建域外部通信域hcclComm ! nullptr直接复用用户传入的HcclComm用户自行管理通信域、多线程等高级场景其中第三种方式对应commMode字段的说明hccl 多线程只支持外部传入通信域方式。日常单卡调试时rank -1且rankSize 1会被SetParam自动修正为rank 0all_to_allvv2_operation.cpp无需显式指定。7. 使用注意事项与参数速查综合源码校验与配置实际使用all_to_allvv2时需满足以下约束平台约束参数头文件注明该算子当前仅支持 Atlas 800I A2 推理产品310P 走独立的AllToAllVV2Operation310p配置且仅支持 float16/int8后端约束backend仅支持hccl不支持 LCCL多卡必须指定 rankrankSize 1时rank不能为-1数组契约sendCount/sdispls/recvCount/rdispls的长度必须等于rankSize元素非负且满足sendCounts[i] sdispls[i] 输入总元素数、recvCounts[i] rdispls[i] sum(recvCounts)输出形状输出固定为二维[1, sum(recvCounts)]tensorForInferShape的dims[0]必须与之相等dtype 对齐输入输出 dtype 一致支持 float16/float/int8/int32/int16/int64/bf16310P 仅 float16/int8。8. 与 all_to_all / all_to_allv 的横向对比仓库中同属 communication 分类的另外两个算子可作为对照详见 .agent/knowledge/ops/communication/all_to_all/index.md 与 .agent/knowledge/ops/communication/all_to_allv/index.md维度all_to_allall_to_allvall_to_allvv2长度语义等长交换可变长度可变长度增强版后端HCCL / LCCL 双后端仅 HCCL仅 HCCL源码文件数644RunnerAllToAllHcclRunner / AllToAllLcclRunnerAllToAllvHcclRunnerAllToAllVV2HcclRunner输入数量较少等长无需 count/displs 数组含 count/displs 控制数组6 输入含 tensorForInferShape 辅助张量可见all_to_allvv2面向可变长度 精确位移控制 强校验 仅 HCCL的增强场景是三者中最精细的版本。若需要阅读更宏观的 Agent 知识索引可参考 .agent/knowledge/README.md 与路由文件 .agent/knowledge/routing/all_to_allvv2.md。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
