news 2026/9/18 13:54:16

ascend-transformer-boost AllGatherV 算子源码导读：从路由文件到 HCCL 可变长集合通信实现

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
ascend-transformer-boost AllGatherV 算子源码导读：从路由文件到 HCCL 可变长集合通信实现

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是 ATB（Ascend 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）、是否支持 ACLNN（no）、预估阅读时间（3-5 分钟）；
  • 文件清单：列出该算子全部源文件及角色；
  • 推荐阅读顺序：给出文件级阅读建议与重点关注点；
  • 源码路径：Op 目录、Kernel 目录与参数头文件位置。

完整知识条目则位于.agent/knowledge/ops/communication/all_gatherv/index.md，其状态标记为 complete，并记录了关键结论：Runner 为 HcclRunner（仅 HCCL，无 LCCL 变体），Pipeline 为"集合通信 — 可变长度 AllGather，HCCL 通信库"，相关算子为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.*负责对外暴露算子接口、参数校验与形状推导，并决定创建哪种 Runner；all_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 = 1）；
  • InferShapeImpl()：根据输入推导输出形状；
  • InferShapeCheckImpl()/SetupCheckImpl()：形状与运行时参数校验；
  • CreateRunner()：根据参数创建对应的 Runner；
  • GetParamJson()：将参数序列化为 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 = 1）
2recvCountsint64长度为rankSize的数组，各 rank 接收的元素个数
3rdisplsint64长度为rankSize的数组，各 rank 数据在输出中的偏移
4yfloat16辅助张量，仅用于推导输出 shape（长度为各 rank dim0 之和），不参与实际数据计算

输出output的 dtype 与x保持一致（float16, float, int8, int16, int32, int64, bf16），format 均为 nd。注意y在配置中的 dtype 固定为 float16，测试用例中也以 float16 张量构造。

5.2 形状推导：InferShapeImpl

InferShapeImpl的逻辑非常简洁：

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 / SetupCheckImpl

InferShapeCheckImpl在形状层面做静态校验：

  • 张量 1~4 的dimNum必须等于 1（DIM_NUM_1）；
  • sendCount的长度必须为 1；
  • recvCounts与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_shared<AllGatherVHcclRunner>(param_, !param_.rankTableFile.empty()); } else { return std::make_shared<AllGatherVHcclRunner>(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 表示未传 rank
rankSizeint0通信卡的数量
rankRootint0主通信编号
backendstring"hccl"通信后端，仅支持 "hccl"（Atlas 推理系列产品即 Ascend 310P 仅支持 hccl）
hcclCommHcclCommnullptrHCCL 通信域指针，为空时由加速库创建，非空时复用用户通信域
commModeCommModeCOMM_MULTI_PROCESS通信模式，hccl 多线程只支持外部传入通信域方式
rankTableFilestring空集群信息配置文件路径，适用单机/多机，当前仅支持 hccl 后端
commDomainstring空通信 device 组通信域名，多通信域时使用，当前仅支持 hccl
rsvuint8_t[64]{0}预留参数

约束条件为0 ≤ rank < rankSize且0 ≤ rankRoot < rankSize。头文件同时给出了两条重要使用注意事项：

  1. 多用户并发：多用户使用时需要通过环境变量ATB_SHARE_MEMORY_NAME_SUFFIX区分共享内存，以保证初始化信息同步互不干扰；
  2. 异常退出清理：通信算子异常退出后，需清理残留数据，参考命令为rm -rf /dev/shm/sem.lccl*、rm -rf /dev/shm/sem.hccl*、ipcrm -a。

在all_gatherv_operation.cpp的ParamCheck中，还会对参数做运行时约束：backend必须为 "hccl"（否则报错 "backend must be hccl"），rankSize必须大于 1，并调用OperationUtil::DistributedInitCheck完成分布式初始化检查。

7. Runner 层：HcclAllGatherV 调用链

all_gatherv_hccl_runner.cpp的ExecuteImpl是整个算子的"最后一公里"：它把 ATB 的RunnerVariantPack张量约定翻译为 HCCL 原生接口调用：

HcclResult ret = HcclAllGatherV( runnerVariantPack.inTensors[0].deviceData, // x：本卡输入数据 *(static_cast<int64_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.ini

ops_configs/atb_ops_info.ini中[AllGatherVOperation]段完整声明了算子的输入输出契约（dtype 与 format 的多值对应关系）：

  • input0：x，dtype 支持 float16/float/int8/int16/int32/int64/bf16，format 为 nd；
  • input1：sendCount，dtype 固定 int64；
  • input2：recvCounts，dtype 固定 int64；
  • input3：rdispls，dtype 固定 int64；
  • input4：y，dtype 固定 float16；
  • output0：output，dtype 与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.Process+spawn）方式模拟两卡通信，每个进程绑定一个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.allclose，rtol/atol 均为 1e-4）。

测试头部还给出了运行前置说明：需要export HCCL_WHITELIST_DISABLE=1，通过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在通信类算子中属于"可变长 AllGather"：

  • all_gather（等长版）：所有 rank 发送等长数据，聚合后广播，路由见.agent/knowledge/routing/all_gather.md；
  • all_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），仅供参考

版权声明: 本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若内容造成侵权/违法违规/事实不符，请联系邮箱：809451989@qq.com进行投诉反馈，一经查实，立即删除！
网站建设 2026/9/18 13:53:14

Flutter跨平台开发家庭药箱App的健康报告功能实践

1. 项目概述&#xff1a;家庭药箱管理App的健康报告功能作为一名长期从事移动应用开发的工程师&#xff0c;我最近在基于Flutter for OpenHarmony平台开发一款家庭药箱管理应用时&#xff0c;遇到了一个有趣的需求&#xff1a;如何将零散的健康数据转化为有价值的健康洞察&…

作者头像 李华
网站建设 2026/9/18 13:52:25

容器技术演进：从内核隔离到Docker镜像与Kubernetes编排

简介&#xff1a;一份系统梳理容器发展历史的Word文档&#xff0c;适合正在学习容器与Kubernetes的开发者、运维人员及架构师阅读&#xff0c;帮助理解容器技术真正要解决的问题及其在软件工程演进中的历史定位。资源围绕开发过程&#xff08;瀑布式、敏捷式、DevOps&#xff0…

作者头像 李华
网站建设 2026/9/18 13:51:56

如何十分钟把公众号搬进 RSS 阅读器？WeWe RSS 开源部署指南

如何十分钟把公众号搬进 RSS 阅读器&#xff1f;WeWe RSS 开源部署指南 【免费下载链接】wewe-rss &#x1f917;更优雅的微信公众号订阅方式&#xff0c;支持私有化部署、微信公众号RSS生成&#xff08;基于微信读书&#xff09; 项目地址: https://gitcode.com/GitHub_Tren…

作者头像 李华