ascend-transformer-boost AllGatherV 算子源码导读从路由文件到 HCCL 可变长集合通信实现【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost导读本文围绕 ascend-transformer-boost 仓库中all_gatherv算子的源码路由文档展开系统梳理该算子的文件清单、推荐阅读顺序、参数结构与底层执行链路。AllGatherVOperation是 ATBAscend Transformer Boost推理侧infer提供的高性能集合通信算子用于将多卡数据按 rank 顺序在第一维上聚合后广播到所有卡且支持每张卡数据不等长。读完本文你将掌握如何定位并阅读该算子的 4 个核心文件理解其 5 输入 1 输出的张量约定、AllGatherVParam参数语义、InferShape 与 Setup 校验逻辑以及最终如何委托 HCCL 的HcclAllGatherV完成底层通信。1. 路由文件在 ATB 知识体系中的定位.agent/knowledge/routing/all_gatherv.md是 ATB Agent 知识体系中路由routing层的一个入口文件其作用是为开发者和 Agent 提供一条最短的源码阅读路径。它并不重复讲解算子细节而是用统一模板给出四类信息元信息分类infer、复杂度S、文件数4、Runner 类型OpsRunner / Operation、是否支持 ACLNNno、预估阅读时间3-5 分钟文件清单列出该算子全部源文件及角色推荐阅读顺序给出文件级阅读建议与重点关注点源码路径Op 目录、Kernel 目录与参数头文件位置。完整知识条目则位于.agent/knowledge/ops/communication/all_gatherv/index.md其状态标记为 complete并记录了关键结论Runner 为 HcclRunner仅 HCCL无 LCCL 变体Pipeline 为集合通信 — 可变长度 AllGatherHCCL 通信库相关算子为all_gather等长版与all_gathervv2增强版。读者可将路由文件视为地图将本文章节视为逐文件的深度讲解。2. 文件清单与角色划分按照路由文件all_gatherv算子在src/ops/ops_infer/all_gatherv/目录下恰好由 4 个文件组成#文件角色1all_gatherv_hccl_runner.cpp源码Runner 实现2all_gatherv_hccl_runner.h头文件Runner 声明3all_gatherv_operation.cppOperation 定义4all_gatherv_operation.hOperation 定义从文件职责看这一结构遵循了 ATB 算子Operation Runner的两层分离设计all_gatherv_operation.*负责对外暴露算子接口、参数校验与形状推导并决定创建哪种 Runnerall_gatherv_hccl_runner.*则封装具体的 HCCL 通信执行逻辑。整个算子没有任何独立的 Kernel如.cce/.cppkernel 实现因为它是纯通信算子计算完全委托给 HCCL 通信库——这一点与路由文件中Runner 类型: OpsRunner,Operation的标注一致。需要说明的是路由文件的 Kernel 目录 字段记录为src/kernels/mixkernels/laser_attention但从当前仓库源码结构看all_gatherv目录下并未发现对应的 kernel 实现可以推断该字段属于模板中的历史遗留信息实际阅读时应以src/ops/ops_infer/all_gatherv/下的 4 个文件为准。3. 推荐阅读顺序路由文件给出的阅读顺序对快速理解该算子非常有效结合源码可以进一步明确每个文件的关注点顺序文件重点关注1all_gatherv_operation.h了解输入输出数量、InferShape 签名2all_gatherv_operation.cppCreateRunner()决策逻辑3all_gatherv_hccl_runner.cpp辅助文件HCCL 调用实现4all_gatherv_hccl_runner.h辅助文件Runner 类声明按此顺序阅读的理由是先通过头文件建立接口视图算子有多少输入输出、继承了哪些虚函数再通过实现文件建立行为视图参数如何校验、形状如何推导、Runner 如何选择最后进入 Runner 实现了解真实的通信调用。4. 源码路径速查Op 目录src/ops/ops_infer/all_gatherv/参数头文件include/atb/infer_op_params.h其中定义了struct AllGatherVParam位于第 1090 行附近需要再次强调该算子没有独立的 Kernel 实现路由文件中指向src/kernels/mixkernels/laser_attention的 Kernel 目录字段对all_gatherv不适用阅读时应忽略。5. Operation 层源码解析5.1 接口视图5 输入 1 输出all_gatherv_operation.h中AllGatherVOperation继承自OperationBase并重写了以下关键接口GetInputNum()/GetOutputNum()返回输入数量 5、输出数量 1对应实现文件中的IN_TENSOR_NUM 5、OUT_TENSOR_NUM 1InferShapeImpl()根据输入推导输出形状InferShapeCheckImpl()/SetupCheckImpl()形状与运行时参数校验CreateRunner()根据参数创建对应的 RunnerGetParamJson()将参数序列化为 JSON便于打印与调试。五个输入张量在all_gatherv_operation.cpp中被定义为常量IN_TENSOR_0~IN_TENSOR_4结合ops_configs/atb_ops_info.ini中[AllGatherVOperation]段的配置可以确认它们的语义序号名称数据类型来自 atb_ops_info.ini语义0xfloat16, float, int8, int16, int32, int64, bf16当前 rank 待聚合的输入数据1sendCountint64标量当前 rank 发送的元素个数SENDCOUNT_LENGTH_1 12recvCountsint64长度为rankSize的数组各 rank 接收的元素个数3rdisplsint64长度为rankSize的数组各 rank 数据在输出中的偏移4yfloat16辅助张量仅用于推导输出 shape长度为各 rank dim0 之和不参与实际数据计算输出output的 dtype 与x保持一致float16, float, int8, int16, int32, int64, bf16format 均为 nd。注意y在配置中的 dtype 固定为 float16测试用例中也以 float16 张量构造。5.2 形状推导InferShapeImplInferShapeImpl的逻辑非常简洁outTensorDescs.at(OUT_TENSOR_0) inTensorDescs.at(IN_TENSOR_0); outTensorDescs.at(OUT_TENSOR_0).shape.dims[DIM_0] inTensorDescs.at(IN_TENSOR_4).shape.dims[DIM_0];即输出张量以输入x为模板仅将第 0 维替换为辅助张量y的第 0 维长度即所有 rank 的 dim0 之和。这正符合 AllGatherV 的语义各 rank 数据按 rank 顺序在第一维聚合。5.3 形状与参数校验InferShapeCheckImpl / SetupCheckImplInferShapeCheckImpl在形状层面做静态校验张量 1~4 的dimNum必须等于 1DIM_NUM_1sendCount的长度必须为 1recvCounts与rdispls的长度必须等于param_.rankSize。SetupCheckImpl则在运行时结合真实数据做更严格的校验包括输出 dim0 必须等于y的 dim0其余维度与x一致sendCount不能超过输入x的元素总数recvCounts[i]、rdispls[i]不能为负数且recvCounts[i] rdispls[i]不能越界含溢出检查各 rank 的recvCounts之和必须大于 0当前 rank 的sendCount必须等于recvCounts[rank]自身发送量等于其他卡从本卡接收的量。这些校验与all_gatherv_hccl_runner.cpp中HcclAllGatherV的参数要求一一对应保证在调用底层通信接口前数据已合法。5.4 Runner 创建CreateRunner 决策逻辑CreateRunner是该算子的核心决策点路由文件也专门提示关注此处。源码逻辑如下if (param_.backend hccl) { if (param_.hcclComm nullptr) { return std::make_sharedAllGatherVHcclRunner(param_, !param_.rankTableFile.empty()); } else { return std::make_sharedAllGatherVHcclRunner(param_, param_.hcclComm); } } ATB_LOG(FATAL) ... backend param_.backend does not exist.;可以看到两种路径加速库自建通信域hcclComm为空时若rankTableFile非空则通过 rank 表文件初始化通信域HcclRunner(rank, rankTableFile, commDomain)否则通过rank/rankSize/rankRoot/commDomain四元组初始化用户托管通信域hcclComm非空时直接复用用户传入的 HCCL 通信域HcclRunner(name, hcclComm)这是 hccl 多线程场景下的推荐方式。AllGatherVHcclRunner的三个构造函数默认、rankTableFile、外部 hcclComm与基类HcclRunner声明于src/atb/runner/hccl_runner.h的初始化方式一一对应。6. AllGatherVParam 参数详解参数结构定义在include/atb/infer_op_params.h的struct AllGatherVParam中其字段与默认值如下字段类型默认值说明rankint-1当前卡所属通信编号默认 -1 表示未传 rankrankSizeint0通信卡的数量rankRootint0主通信编号backendstringhccl通信后端仅支持 hcclAtlas 推理系列产品即 Ascend 310P 仅支持 hcclhcclCommHcclCommnullptrHCCL 通信域指针为空时由加速库创建非空时复用用户通信域commModeCommModeCOMM_MULTI_PROCESS通信模式hccl 多线程只支持外部传入通信域方式rankTableFilestring空集群信息配置文件路径适用单机/多机当前仅支持 hccl 后端commDomainstring空通信 device 组通信域名多通信域时使用当前仅支持 hcclrsvuint8_t[64]{0}预留参数约束条件为0 ≤ rank rankSize且0 ≤ rankRoot rankSize。头文件同时给出了两条重要使用注意事项多用户并发多用户使用时需要通过环境变量ATB_SHARE_MEMORY_NAME_SUFFIX区分共享内存以保证初始化信息同步互不干扰异常退出清理通信算子异常退出后需清理残留数据参考命令为rm -rf /dev/shm/sem.lccl*、rm -rf /dev/shm/sem.hccl*、ipcrm -a。在all_gatherv_operation.cpp的ParamCheck中还会对参数做运行时约束backend必须为 hccl否则报错 backend must be hcclrankSize必须大于 1并调用OperationUtil::DistributedInitCheck完成分布式初始化检查。7. Runner 层HcclAllGatherV 调用链all_gatherv_hccl_runner.cpp的ExecuteImpl是整个算子的最后一公里它把 ATB 的RunnerVariantPack张量约定翻译为 HCCL 原生接口调用HcclResult ret HcclAllGatherV( runnerVariantPack.inTensors[0].deviceData, // x本卡输入数据 *(static_castint64_t *(runnerVariantPack.inTensors[1].hostData)), // sendCount runnerVariantPack.outTensors[0].deviceData, // output runnerVariantPack.inTensors[2].hostData, // recvCounts runnerVariantPack.inTensors[3].hostData, // rdispls GetHcclDtype(runnerVariantPack.inTensors[0].desc.dtype), // 数据类型映射 hcclComm_.get(), GetExecuteStream(runnerVariantPack.context)); // 执行流执行前还会做两项运行时检查hcclComm_非空以及输入输出deviceData非空HcclAllGatherV返回非HCCL_SUCCESS时通过ConvertHcclResultToStatus转换为 ATB 状态码。可见sendCount、recvCounts、rdispls均通过hostData以host 侧内存传入它们是标量与数组形式的控制信息而非 device 张量x与output通过deviceData传入真正参与通信的设备数据Runner 类末尾通过REG_RUNNER_TYPE(AllGatherVHcclRunner)注册到 ATB Runner 工厂。这与知识条目中Runner: HcclRunner仅 HCCL无 LCCL 变体的记录完全吻合。8. 算子配置atb_ops_info.iniops_configs/atb_ops_info.ini中[AllGatherVOperation]段完整声明了算子的输入输出契约dtype 与 format 的多值对应关系input0xdtype 支持 float16/float/int8/int16/int32/int64/bf16format 为 ndinput1sendCountdtype 固定 int64input2recvCountsdtype 固定 int64input3rdisplsdtype 固定 int64input4ydtype 固定 float16output0outputdtype 与x一致7 种format 为 nd。该配置是算子注册与参数绑定如OPERATION_PARAM_FUNCS(AllGatherVOperation, infer::AllGatherVParam)的重要依据也是理解张量角色最快的入口。9. 测试用例验证仓库在tests/apitest/opstest/python/operations/all_gatherv/test_all_gatherv_operation.py中提供了完整的端到端测试。测试要点包括以rankSize 2的多进程multiprocessing.Processspawn方式模拟两卡通信每个进程绑定一个npu设备通过torch.classes.OperationTorch.OperationTorch(AllGatherVOperation)创建算子以 JSON 形式设置{rank: rank, rankSize: world_size, rankRoot: 0, backend: hccl}sendCount按 rank 取不同值如[4, 2]recvCounts与rdispls按 rank 构造验证不等长聚合辅助张量y的长度取所有输入 dim0 之和用于推导输出 shape代码注释明确说明y 用来推导 outputshape长度应为所有 inputtensor 的 dim0 之和通过golden_compare将聚合结果与 golden 张量比对torch.allclosertol/atol 均为 1e-4。测试头部还给出了运行前置说明需要export HCCL_WHITELIST_DISABLE1通过python3 -m unittest test_all_gather_operation.py执行且需先source set_env.sh并设置ATB_HOME_PATH环境变量。此外tests/framework/c/atb_torch/operation/hosttensor_binders/all_gatherv_binder.h/.cpp与tests/apitest/opstest/csv/all_gatherv.csv也提供了框架层绑定与用例数据可作为补充参考。10. 相关算子与阅读延伸all_gatherv在通信类算子中属于可变长 AllGatherall_gather等长版所有 rank 发送等长数据聚合后广播路由见.agent/knowledge/routing/all_gather.mdall_gathervv2增强版在 all_gatherv 基础上提供更丰富的控制能力。若希望系统了解 ATB 的集合通信架构可继续阅读 Runner 基类实现src/atb/runner/hccl_runner.h、src/atb/runner/hccl_runner.cpp以及知识体系总索引.agent/knowledge/README.md。对于开发同类通信算子的场景建议完全复用all_gatherv的Operation 参数校验 → CreateRunner 决策 → HcclRunner 封装三段式结构先保证AllGatherVParam在头文件中语义完备再在InferShapeCheckImpl/SetupCheckImpl中守住形状与越界边界最后在 Runner 的ExecuteImpl中完成 host 控制信息与 device 数据的分离传递即可安全、高效地接入 HCCL 通信能力。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考