用户缓冲区注册
用户缓冲区注册是一个功能,它允许 MCCL 直接通过用户缓冲区发送/接收/操作数据,无需额外的内部复制(零复制)。它可以加速集体操作,并大大减少资源使用(例如#通道使用)。MCCL 提供了两种注册用户缓冲区的方式;一种是MUSA 图注册,另一种是本地注册。MCCL 要求对于所有 MCCL 通信函数调用(例如,allreduce、sendrecv 等),如果任何 rank 在通信器中传递注册缓冲区给 MCCL 通信函数,同一通信器中的所有其他 rank 也必须传递其注册缓冲区;否则,混合注册和非注册缓冲区可能导致未定义行为。
MTLink Sharp 缓冲区注册
MCCL 支持 MTLink Sharp(MTLS)的用户缓冲区注册;任何支持 MTLS 算法的 MCCL 集体操作(例如,allreduce)都可以利用这一特性。
要启用基于MUSA 图的 MTLS 缓冲区注册,用户必须遵守几个要求:
- 缓冲区是通过
mcclMemAlloc()或合格的分配器分配的(见 内存分配器)。 - MCCL 操作是在每个 rank 的 MUSA 图捕获的流上启动的。
- 在每个 rank 的集体操作中,缓冲区的头部地址偏移是相同的。
注册的缓冲区将在 MUSA 图被销毁时取消注册。这里是一个基于 MUSA 图的缓冲区注册示例:
void* sendbuff;
void* recvbuff;
size_t count = 1 << 25;
CHECK(mcclMemAlloc(&sendbuff, count * sizeof(float)));
CHECK(mcclMemAlloc(&recvbuff, count * sizeof(float)));
musaGraph_t graph;
CHECK(musaStreamBeginCapture(stream, musaStreamCaptureModeThreadLocal));
CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
// 每个 rank 的 sendbuff 和 recvbuff 头部地址偏移相同
CHECK(mcclAllReduce((void*)((float*)sendbuff + 1024), (void*)((float*)recvbuff + 2048), 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamEndCapture(stream, &graph));
musaGraphExec_t instance;
CHECK(musaGraphInstantiate(&instance, graph, NULL, NULL, 0));
CHECK(musaGraphLaunch(instance, stream));
CHECK(musaStreamSynchronize(stream));
CHECK(musaGraphExecDestroy(instance));
CHECK(musaGraphDestroy(graph));
CHECK(mcclMemFree(sendbuff));
CHECK(mcclMemFree(recvbuff));
另一方面,要启用基于本地的 MTLS 缓冲区注册,用户必须遵守以下要求:
-
缓冲区是通过
mcclMemAlloc()或合格的分配器分配的(见 内存分配器)。 -
在每个 rank 调用集体操作之前,使用
mcclCommRegister()注册缓冲区。 -
像往常一样调用 MCCL 集体操作,但同样保持缓冲区头部地址的偏移对每个 rank 都是相同的。
注册的缓冲区将在用户显式调用 mcclCommDeregister() 时取消注册。这里是一个基于本地的缓冲区注册示例:
void* sendbuff;
void* recvbuff;
size_t count = 1 << 25;
void* sendRegHandle;
void* recvRegHandle;
CHECK(mcclMemAlloc(&sendbuff, count * sizeof(float)));
CHECK(mcclMemAlloc(&recvbuff, count * sizeof(float)));
CHECK(mcclCommRegister(comm, sendbuff, count * sizeof(float), &sendRegHandle));
CHECK(mcclCommRegister(comm, recvbuff, count * sizeof(float), &recvRegHandle));
CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
CHECK(mcclAllReduce((void*)((float*)sendbuff + 1024), (void*)((float*)recvbuff + 2048), 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamSynchronize(stream));
CHECK(mcclCommDeregister(comm, sendRegHandle));
CHECK(mcclCommDeregister(comm, recvRegHandle));
CHECK(mcclMemFree(sendbuff));
CHECK(mcclMemFree(recvbuff));
对于基于本地的注册,用户可以在程序开始时注册缓冲区一次,并多次重用缓冲区以利用注册的好处。
为了节省内存,也可以分配一大块缓冲区并只注册一次。sendbuff 和 recvbuff 可以通过大块进一步分配,以实现零复制 MCCL 操作,只要 sendbuff 和 recvbuff 满足偏移要求。以下示例显示了一个用例:
void* buffer;
void* handle;
void* sendbuff;
void* recvbuff;
size_t size = 1 << 29;
CHECK(mcclMemAlloc(&buffer, size));
CHECK(mcclCommRegister(comm, buffer, size, &handle));
// 分配缓冲区块给 sendbuff 和 recvbuff
sendbuff = buffer;
recvbuff = (void*)((uint8_t*)buffer + (1 << 20));
CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamSynchronize(stream));
CHECK(mcclCommDeregister(comm, handle));
CHECK(mcclMemFree(sendbuff));
IB Sharp 缓冲区注册
MCCL 支持 IB Sharp 缓冲区注册,任何支持 IB Sharp 算法的 MCCL 集体操作都可以从该特性中受益,例如 allreduce、reducescatter 和 allgather。目前,MCCL 仅支持每个节点包含 1 个 rank 的通信器的 IB Sharp 缓冲区注册,注册可以减少 MCCL SM 使用量至 1。
通过 MUSA 图启用 IB Sharp 缓冲区注册:
-
使用任何 MUSA 分配器(例如,musaMalloc/mcclMemAlloc)分配发送和接收缓冲区
-
使用 MUSA 图启动 MCCL 集体操作
通过本地注册启用 IB Sharp 缓冲区注册:
-
使用任何 MUSA 分配器(例如,musaMalloc/mcclMemAlloc)分配发送和接收缓冲区
-
在通信器中的每个 rank 使用 mcclCommRegister 注册发送和接收缓冲区
-
启动 MCCL 集体操作
通用缓冲区注册
MCCL 支持 节点内缓冲区注册,它针对所有点对点节点内通信(例如,Allgather Ring)并带来较小的内存压力,更好的通信和计算重叠性能。通过在开始时使用 mcclCommRegister 注册缓冲区或应用 MUSA 图,可以为 MCCL 集体操作和 sendrecv 启用节点内缓冲区注册。
用户缓冲区可以通过 VMM API(即 muMem*)、任何基于 VMM 的分配器(内存分配器)或 mcclMemAlloc 分配。通过传统 MUSA API(例如,musaMalloc)分配的缓冲区也可以用于注册。然而,由于在执行过程中可能发生挂起,在失败或者 abort 过程中可能出现段错误。因此,不推荐使用传统缓冲区进行注册;目前,默认情况下禁用了传统缓冲区注册,用户可以设置 MCCL_LEGACY_MUSA_REGISTER=1 以启用它。
2.4 对称窗口辅助接口
MCCL 在对称窗口能力上提供了两个主机侧辅助接口:
它们都依赖对称窗口注册与通信器能力探测。
mcclWinGetUserPtr
该接口返回窗口最初绑定的用户指针,常用于:
- 验证窗口与业务缓冲区的关联关系
- 在调试或 fallback 路径中恢复应用侧视角的地址
注意:
- 在支持对称窗口辅助能力的平台上,
win必须有效 - 在不支持该能力的平台上,
win == NULL的特定探测路径可能返回成功并把输出置为NULL
参考示例:examples/05_symmetric_memory/03_win_get_user_ptr
mcclGetLsaDevicePointer
该接口用于从主机侧获取对称窗口在指定 lsaRank 视角下的设备地址,适合:
- 主机侧调试 LSA 可见地址
- 构造需要显式设备地址的辅助逻辑
参考示例:examples/05_symmetric_memory/02_get_lsa_device_pointer
内存分配器
为了方便,MCCL 提供了 mcclMemAlloc 函数,帮助用户通过 VMM API 分配缓冲区,之后这些缓冲区可以用于 MCCL 注册。该函数是专门为 MCCL 设计的,因此不建议在应用程序中到处使用由 mcclMemAlloc 分配的缓冲区。
对于高级用户,如果您想为 MTLS UB 创建自己的内存分配器,分配器分配的缓冲区需要满足以下要求:
- 在支持的 GPU 上,分配缓冲区时需要使用共享标志
MU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,同时也要使用MU_MEM_HANDLE_TYPE_FABRIC。 - 缓冲区物理内存大小是 MUMEM 推荐粒度的倍数(即
muMemGetAllocationGranularity(..., MU_MEM_ALLOC_GRANULARITY_RECOMMENDED)) - 缓冲区虚拟头部地址至少与 MUMEM 推荐粒度对齐,大小是 MUMEM 推荐粒度的倍数。
对于使用 VMM API 的通用缓冲区注册,分配器需要满足与 MTLS UB 分配器相同的要求。

