MCCL over ACE
1. 概述
异步通信引擎 Async Copy Engine (ACE) 是 MTGPU 的一个独立硬件单元,用于在不同 内存域之间实现高带宽、低延迟的异步数据搬运。

运行原理
在典型的多芯片(Multi-Die)GPU 结构中,每个 GPU 包含多个计算芯片,ACE 的任务是为这些芯片之间及 GPU 之间提供高效的数据通路,主要包括:
- Host 内存 ↔ 显存(H2D/D2H)
- 本地 GPU 显存 ↔ 远程 GPU 显存
- 芯片内部 Die 间数据搬移(D2D - Die to Die)
对称内存 (Symmetric Memory)
MCCL over ACE 使用对称内存机制,允许每个 GPU 直接访问其他 GPU 的内存地址空间。
Rank 0: [Base + 0*stride] ───┐
Rank 1: [Base + 1*stride] ───┼─→ 对称内存窗口
Rank 2: [Base + 2*stride] ───┘
关键函数:
mcclMemAlloc(): 分配对称内存mcclCommWindowRegister(): 注册对称内存窗口
图执行 (Graph Execution)
MCCL over ACE 预先构建 MUSA Graph,包含所有通信节点(同步、拷贝、原子操作),运行时动态更新参数后执行。
优势:
- 减少 API 调用开销
- 硬件级流水线优化
- 更低的通信延迟
2. 支持的集合通信原语
AllReduce
AllReduce 操作在设备间对数据(例如,求和、最小值、最大值)进行归约,并将结果存储在每个 rank 的接收缓冲区中。
以 sum 为例,在 k 个 rank 之间的求和 AllReduce 操作中,每个 rank 将提供一个包含 N 个值的数组,并在包含 N 个值的数组 out 中接收相同的结果,其中 out[i] = in0[i] + in1[i] + ... + in(k-1)[i]。

AllGather
AllGather 操作从 k 个 rank 中收集 N 个值到大小为 k*N 的输出缓冲区,并将该结果分发到所有 rank。输出按 rank 索引排序。因此,AllGather 操作会受到不同 rank 到设备映射的影响。

ReduceScatter
ReduceScatter 操作执行与归约相同的操作,但结果在 rank 之间以等大小的块分散,每个 rank 根据其 rank 索引获得一块数据。由于 rank 决定数据布局,ReduceScatter 操作会受到不同 rank 到设备映射的影响。

AlltoAll
在 k 个进程组(rank)之间的 AllToAll 操作中,每个进程组会提供一个大小为 k×N 个值的输入缓冲区:其中第 j 个包含 N 个值的数据块(chunk)会发送至目标进程组 j;每个进程组会接收一个大小为 k×N 个值的输出缓冲区:其中第 i 个包含 N 个值的数据块源自源进程组 i。
ACE 版本的 mcclAllToAllv 接口支持非对称输入。

3. API 参考
初始化与销毁
mcclAceCommInitRank
初始化 MCCL over ACE 通信器。
mcclResult_t mcclAceCommInitRank(
mcclComm_t* comm, // [out] 通信器句柄
int nranks, // [in] 进程总数
mcclUniqueId commId, // [in] 唯一 ID
int rank // [in] 当前进程 rank
);
与 SHADER 版本的区别:
mcclAceCommInitRank 和 mcclCommInitRank 的主要区别在于是否启用 ACE。
// SHADER 版本
mcclResult_t mcclCommInitRank(mcclComm_t* newcomm, int nranks,
mcclUniqueId commId, int myrank) {
...
// isAceComm 默认为 false
MCCLCHECK(mcclCommInitRankDev(...));
...
}
// ACE 版本
mcclResult_t mcclAceCommInitRank(mcclComm_t* newcomm, int nranks,
mcclUniqueId commId, int myrank) {
...
// 显式传递 isAceComm = true
MCCLCHECK(mcclCommInitRankDev(..., true));
...
}
示例:
mcclUniqueId id;
if (rank == 0) {
mcclGetUniqueId(&id);
}
MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);
mcclComm_t comm;
mcclAceCommInitRank(&comm, worldSize, id, worldRank);
mcclCommDestroy
销毁通信器。
mcclResult_t mcclCommDestroy(mcclComm_t comm);
内存管理
mcclMemAlloc
分配通信显存。
mcclResult_t mcclMemAlloc(
void** ptr,
size_t size
);
mcclMemFree
释放通信显存。
mcclResult_t mcclMemFree(void* ptr);
mcclCommWindowRegister
注册为对称内存。
mcclResult_t mcclCommWindowRegister(
mcclComm_t comm, // [in] 通信器
void* localPtr, // [in] 本地内存指针
size_t size, // [in] 内存大小
void** symmetricHandle, // [out] 对称内存句柄
int flags // [in] 标志
);
flags 参数支持:
/* Window Registration flags */
#define MCCL_WIN_DEFAULT 0x00 // 普通注册,一般不用
#define MCCL_WIN_COLL_SYMMETRIC 0x01 // 对称内存注册(ACE 集合操作必需)
集合操作
mcclAllReduce
mcclResult_t mcclAllReduce(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count,
mcclDataType_t datatype,
mcclRedOp_t op,
mcclComm_t comm,
musaStream_t stream
);
两种模式:
-
Out-of-place (
sendbuff != recvbuff):sendbuff: 对称内存recvbuff: 普通内存- 每个 rank 收集所有进程的数据,pull mode
-
In-place (
sendbuff == recvbuff):recvbuff: 对称内存,总大小 = count * nranks * sizeof(datatype)- 每个 rank 的数据已在对应位置
支持的归约操作:
| 操作 | 支持的数据类型 |
|---|---|
| mcclSum (求和) | mcclInt32, mcclUint64, mcclFloat16, mcclBfloat16, mcclFloat32, mcclFloat64 |
| mcclMin (最小值) | mcclInt32, mcclInt64, mcclUint32, mcclUint64 |
| mcclMax (最大值) | mcclInt32, mcclInt64, mcclUint32, mcclUint64 |
重要说明:
- 浮点类型(
mcclFloat16/32/64,mcclBfloat16)仅支持mcclSum操作 - 整数类型支持所有三种操作,但
mcclSum不支持mcclInt64
mcclAllGather
mcclResult_t mcclAllGather(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count,
mcclDataType_t datatype,
mcclComm_t comm,
musaStream_t stream
);
mcclReduceScatter
mcclResult_t mcclReduceScatter(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count, // [in] 每个进程接收的元素数量
mcclDataType_t datatype,
mcclRedOp_t op, // [in] 归约操作
mcclComm_t comm,
musaStream_t stream
);
内存要求:
sendbuff: 对称内存,大小 = count * nranks * sizeof(datatype)recvbuff: 普通内存,大小 = count * sizeof(datatype)
mcclAllToAllv
mcclAllToAllv 接口支持非对称,recvcounts[i] 直接支持非均匀。
mcclResult_t mcclAllToAllv(
const void* sendbuff, // [in] 发送缓冲区 (对称内存)
const size_t sendcounts[], // [in] 发送计数数组
const size_t sdispls[], // [in] 发送偏移数组
void* recvbuff, // [out] 接收缓冲区
const size_t recvcounts[], // [in] 接收计数数组
const size_t rdispls[], // [in] 接收偏移数组
mcclDataType_t datatype,
mcclComm_t comm,
musaStream_t stream
);
4. 示例代码
AllReduce 示例
AllReduce 是一个用于将所有进程的数据进行归约操作(如求和)的集合操作,结果存回所有 rank。
关键要点:
- 必须使用
mcclMemAlloc分配通信显存 - 必须注册为对称内存:调用
mcclCommWindowRegister
#include <mpi.h>
#include <mccl.h>
#include <musa_runtime.h>
#include <vector>
#include <cassert>
#include <cstdio>
int main(int argc, char* argv[]) {
// ========== 1. 初始化 MPI ==========
CHECK_MPI(MPI_Init(&argc, &argv));
int worldRank, worldSize;
CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &worldRank));
CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &worldSize));
// ========== 2. 设置 GPU 设备 ==========
int ngpus = 0;
CHECK_MUSA(musaGetDeviceCount(&ngpus));
int dev = worldRank % ngpus; // 轮询分配 GPU
CHECK_MUSA(musaSetDevice(dev));
// 创建 MUSA stream
musaStream_t stream;
CHECK_MUSA(musaStreamCreate(&stream));
// ========== 3. 初始化 MCCL 通信器 ==========
mcclUniqueId id;
if (worldRank == 0) {
CHECK_MCCL(mcclGetUniqueId(&id));
}
CHECK_MPI(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD));
mcclComm_t comm;
CHECK_MCCL(mcclAceCommInitRank(&comm, worldSize, id, worldRank));
// ========== 4. 执行集合操作 ==========
// ========== 4.1 分配和初始化内存 ==========
size_t count = 16 * 1024 * 1024; // 16M 个元素
std::vector<uint64_t> data(count, 1); // 初始化为 1
void* buff = nullptr;
void* handle = nullptr;
// 分配对称内存
CHECK_MCCL(mcclMemAlloc(&buff, count * sizeof(uint64_t)));
// 拷贝数据到 GPU
CHECK_MUSA(musaMemcpy(buff, data.data(),
count * sizeof(uint64_t),
musaMemcpyHostToDevice));
// 注册对称内存窗口
CHECK_MCCL(mcclCommWindowRegister(comm, buff,
count * sizeof(uint64_t),
&handle,
MCCL_WIN_COLL_SYMMETRIC));
// ========== 4.2 执行 AllReduce ==========
// 多次调用测试
for (int i = 0; i < 4; i++) {
CHECK_MCCL(mcclAllReduce(buff, buff, // in-place!
count,
mcclUint64,
mcclSum, // 求和操作
comm,
stream));
}
CHECK_MUSA(musaStreamSynchronize(stream));
// ========== 4.3 验证结果 ==========
uint64_t* host_data = (uint64_t*)malloc(count * sizeof(uint64_t));
CHECK_MUSA(musaMemcpy(host_data, buff,
count * sizeof(uint64_t),
musaMemcpyDeviceToHost));
// 每次 AllReduce: 1*worldSize -> worldSize
// 4 次后:worldSize^4
uint64_t expected = 1;
for (int i = 0; i < 4; i++) {
expected *= worldSize;
}
for (size_t i = 0; i < count; i++) {
assert(host_data[i] == expected);
}
free(host_data);
printf("Rank %d: AllReduce PASSED!\n", worldRank);
// ========== 4.4 清理 ==========
CHECK_MCCL(mcclMemFree(buff));
// ========== 5. 清理资源 ==========
CHECK_MUSA(musaStreamDestroy(stream));
CHECK_MCCL(mcclCommDestroy(comm));
CHECK_MPI(MPI_Finalize());
return 0;
}