- 高性能计算
- 通信
- 深度学习
- 分布式训练
【免费下载链接】nccl
Optimized primitives for collective multi-GPU communication
本文讲解 NCCL 开源仓库docs/examples/07_kernel_fusion中的 Multimem 示例（02_rmsnorm_multimem）：如何利用 SM 9.0+（Hopper H100 及以上）的硬件多播内存（multimem）指令，把"reduce-scatter → RMSNorm → all-gather"三段式分布式归一化逻辑融合进一次GPU kernel 启动。读完本文，你将掌握ncclGetLsaMultimemPointer、multimemLoadSum、multimemStore三类设备端 API 的用法、Device Communicator 的资源声明方式，以及它与纯 LSA（Load Store Accessible）版本在内核写法上的本质差异。
背景：分布式 RMSNorm 的三阶段通信模式
RMSNorm（Root Mean Square Normalization）是 Transformer 中常用的归一化手段。对向量 $\mathbf{x}=[x_1,\dots,x_n]$，其定义为：
RMS(x) = sqrt( (1/n) * Σ x_i² + ε ) y_j = x_j / RMS(x)其中 ε 是数值稳定常数（示例取1e-6，与 PyTorch RMSNorm 一致，ε 加在根号内部）。在张量并行场景下，前一层（如 FFN、Attention）沿隐藏维被切分到多张 GPU，每张卡只持有所有 token 的局部贡献；而 RMS 统计量依赖完整激活，因此在归一化之前必须先跨卡求和，归一化之后又必须把结果广播回去。这构成了本文示例内核的三个逻辑阶段：
- Phase 1（Reduce-Scatter）：聚合所有 GPU 上同一 token 的局部贡献，还原出完整激活；
- Phase 2（RMSNorm）：对完整激活做归一化；
- Phase 3（All-Gather）：把归一化结果广播回所有 GPU，供下一层使用。
该模式在 NCCL 传统集体通信 API 下等价于ncclReduceScatter+ 自定义 kernel +ncclAllGather三次独立启动；而 07_kernel_fusion 分类 README 详细阐述了其数学推导、4×8 矩阵的具体算例、集体通信等价实现与融合收益，是所有子示例（LSA / Multimem / GIN / Hybrid）共享的公共背景，建议先阅读。
什么是 Multimem：硬件多播内存操作
Multimem 是 NCCL Device API 提供的一类基于CUDA 硬件多播内存指令（SM 9.0+）的访问能力：单个多播指针即可同时覆盖一个 LSA 团队内的所有 peer。
ncclGetLsaMultimemPointer(window, offset, devComm)：返回一个多播指针。对该指针执行 load 会从所有 LSA peer读取，执行 store 会写入所有 LSA peer（core.h 中的声明）。multimemLoadSum(addr)：对应 PTX 指令multimem.ld_reduce.global.add.f32，一次加载并完成跨 peer 的浮点加法归约。multimemStore(addr, val)：对应 PTX 指令multimem.st.global.b32，把值广播写入所有 peer。
相比纯 LSA 需要显式遍历每个 peer（for (peer...) ncclGetLsaPointer(...)），Multimem 用一个指针消掉了 peer 循环，内核逻辑显著简化。这两个内联函数的完整实现见共享头文件 rmsnorm_utils.cuh，且它们以__CUDA_ARCH__ >= 900为编译期开关——低于 SM 9.0 时直接assert(false)，从源头杜绝在旧架构上误用。
目标架构与拓扑约束
| 项目 | 约束 |
|---|---|
| 拓扑 | 单节点 / 单个 NVLink 域 |
| GPU 算力 | Compute Capability 9.0+（Hopper H100 及以上） |
| LSA 团队 | 必须满足nLsaTeams == 1（所有 rank 能互相直接访问） |
| 其他 | 需要 NCCL Device API 支持 |
Device Communicator 资源需求
Multimem 内核通过ncclDevCommCreate创建设备通信器，并在ncclDevCommRequirements中声明两类资源：
| 需求字段 | 取值 | 作用 |
|---|---|---|
lsaBarrierCount | tokens_per_gpu | 每个 block 一个 LSA barrier，用于阶段间跨 GPU 同步 |
lsaMultimem | true | 开启 Multimem，使ncclGetLsaMultimemPointer与 PTX multimem 指令可用（要求 SM 9.0+） |
lsaMultimem字段定义于 nccl_device/core.h 的 ncclDevCommRequirements 结构。在主机端创建通信器时（见 main.cu）：
ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; reqs.lsaBarrierCount = tokens_per_gpu; // One barrier per token reqs.lsaMultimem = true; // Enable multimem support NCCLCHECK(ncclDevCommCreate(comm, &reqs, &devComm));三阶段融合内核详解
内核RMSNormMultimem（完整代码见 main.cu）以grid =tokens_per_gpu、block = 256 线程的配置启动，即每个 CUDA block 负责一个 token 的完整三阶段流水。
Phase 1：Reduce-Scatter（multimemLoadSum）
每个 block 用单个多播指针加载并累加自己负责的 token 在所有 peer 上的数据，归约结果仅写回本地：
ncclCoopCta coop = ncclCoopCta(); // LSA barrier session：每个 block 一个实例（索引 blockIdx.x），开启 multimem ncclLsaBarrierSession<ncclCoopCta> bar { coop, devComm, ncclTeamTagLsa(), blockIdx.x, /*multimem=*/true }; // 初始跨 GPU 同步 bar.sync(coop, cuda::memory_order_acquire); const int rank = devComm.rank; const int token_idx = rank * gridDim.x + blockIdx.x; // 全局 token 索引 const int window_offset = token_idx * hidden_dim * sizeof(float); float* local_pointer = (float*)ncclGetLocalPointer(window, window_offset); float* multimem_pointer = reinterpret_cast<float*>( ncclGetLsaMultimemPointer(window, window_offset, devComm)); // 从所有 peer 加载并求和；本 rank 只需要自己 token 的归约结果，因此只写本地 for (int i = threadIdx.x; i < hidden_dim; i += blockDim.x) { local_pointer[i] = multimemLoadSum(multimem_pointer + i); } coop.sync();注意这里没有peer 循环：multimemLoadSum一条指令完成"读全部 peer + 求和"。每个 rank 只保留属于自己的tokens_per_gpu个 token 的归约结果，这正是 reduce-scatter 的语义。
Phase 2：RMSNorm（blockRMSNorm）
与纯 LSA 示例完全一致，调用共享辅助函数 blockRMSNorm 对本地归约数据做归一化。其内部算法为：每线程计算各自维度的平方和 → 写入共享内存 →coop.sync()后做树形块内归约 → 得到rms_scale = rsqrtf(mean_sq + eps)→ 逐元素缩放。共享内存缓冲区大小恰好为threads_per_block个 float：
blockRMSNorm(local_pointer, hidden_dim, eps, reduction_buffer, coop);Phase 3：All-Gather（multimemStore）
归一化完成后，每个线程把自己负责的维度通过多播 store 广播到所有 peer：
for (int i = threadIdx.x; i < hidden_dim; i += blockDim.x) { float val = local_pointer[i]; // 本地只读一次 multimemStore(multimem_pointer + i, val); } bar.sync(coop, cuda::memory_order_release);同步策略
与 LSA 版本相同：Phase 1 前的初始 barrier 使用cuda::memory_order_acquire，保证能看到所有 peer 写入的初始数据；Phase 3 后的最终 barrier 使用cuda::memory_order_release，保证广播写对所有 peer 可见。两处 barrier 都通过ncclLsaBarrierSession（按blockIdx.x索引，每个 block 独立实例）完成跨 GPU 的细粒度同步。
与纯 LSA 版本的内核对比
同一目录下的 01_rmsnorm_lsa 示例 用ncclGetLsaPointer(window, offset, peer)显式遍历每个 peer：Phase 1 需要for (peer...) sum += peer_token_data[i]的嵌套循环（见其 main.cu），Phase 3 同样需要逐 peer 写入（见 main.cu 的 L123-L128）。Multimem 版本把这两处 peer 循环都压缩成了单条指令级操作，代码更短、语义更直接；代价是要求硬件支持（SM 9.0+）且归约操作只支持求和（见下文实现约束）。
主机端程序流程（8 步）
完整实现见 main.cu 的 rms_norm 函数，流程如下：
- 初始化 NCCL 通信器：rank 0 生成
ncclUniqueId并广播，随后ncclCommInitRank建通信器。Multimem 示例在建通信器之前先用cudaGetDeviceProperties检查 GPU 算力（prop.major < 9直接报错退出），再用ncclCommQueryProperties校验三项能力（main.cu L170-L191）：props.deviceApiSupport：Device API 支持；props.multimemSupport：multimem 支持（字段定义于 core.h）；props.nLsaTeams == 1：必须是单一 LSA 团队。
- 分配主机内存并初始化数据：
h_data（输入，用initialize_data以seed=42+rank填充随机值）与h_data_validation（接收 GPU 输出用于比对）。此时才分配主机缓冲，是为了让不支持配置的提前退出无需释放资源。 - 分配设备内存并注册窗口：
ncclMemAlloc+ncclCommWindowRegister(..., NCCL_WIN_COLL_SYMMETRIC)注册对称窗口，拷入数据并创建 CUDA stream。 - 创建设备通信器：如前述设置
lsaBarrierCount = tokens_per_gpu、lsaMultimem = true后ncclDevCommCreate。 - 启动融合内核：
RMSNormMultimem<<<tokens_per_gpu, threads_per_block, shared_mem_size, stream>>>，其中shared_mem_size = threads_per_block * sizeof(float)；随后cudaGetLastError+cudaStreamSynchronize。 - 与 CPU 参考比对：
rms_norm_generate在主机侧按同样公式生成期望结果，GPU 输出拷回后用verify_results按绝对容差1e-5、相对容差1e-4逐元素校验。 - 清理资源：
ncclDevCommDestroy→ncclCommWindowDeregister→ncclMemFree→ncclCommFinalize/ncclCommDestroy→cudaStreamDestroy→ 释放主机缓冲。 - 报告结果：打印
SUCCESS/FAILED。
主机侧的run_example生命周期封装在公共工具 utils.cc：pthread 模式下按NTHREADS环境变量或 GPU 数量创建线程，每线程一个 rank；MPI 模式下每个进程一个 rank。
构建与运行
构建
在示例目录 c/ 下（Makefile 见 c/Makefile）：
make [MPI=1] [MPI_HOME=<path-to-mpi>] [NCCL_HOME=<path-to-nccl>] [CUDA_HOME=<path-to-cuda>]或在分类目录下与其余示例统一构建：
# 从 07_kernel_fusion 目录 make -C 02_rmsnorm_multimem/c # 单独构建 make MPI=1 # 带 MPI 支持统一构建 make clean # 清理构建产物环境要求：NCCL 2.30 及以上（Device API 类型与ncclDevCommRequirements字段在旧版本有差异）、CUDA Toolkit 12.0+（Makefile 的 help 目标注明 CUDA 12.2+）、SM 9.0+ GPU。编译需包含共享头文件路径（-I ../../../common/include -I ../../include，后者提供rmsnorm_utils.cuh）；若出现rms_norm_generate/blockRMSNorm/verify_results未解析符号，通常是rmsnorm_utils.cuh包含路径未配置正确。
运行
pthread 模式（默认），默认跑全部可用 GPU，也可用NTHREADS指定线程数：
[NTHREADS=N] ./rmsnorm_multimemMPI 模式（需make MPI=1）：
mpirun -np <num_processes> ./rmsnorm_multimem默认配置参数
| 参数 | 默认值 | 说明 |
|---|---|---|
sequence_length | 4096 | 全部 GPU 的 token 总数，必须能被 rank 数整除（sequence_length % total_ranks == 0，代码中有显式校验） |
hidden_size | 1024 | 每个 token 的隐藏维 |
threads_per_block | 256 | 每 block 线程数，同时也是共享内存归约缓冲的 float 数 |
eps | 1e-6f | RMSNorm 数值稳定常数 |
tokens_per_gpu | sequence_length / total_ranks | 派生参数：每卡 token 数 = grid 维度 = LSA barrier 数 |
这些参数在 main.cu L118-L126 中定义，可按需修改。
预期输出
运行成功后按 rank 依次输出：
- 算力与 multimem 能力检查结果（SM 9.0+ 与
multimemSupport）； - 融合 RMSNorm 内核以
tokens_per_gpu个 block 启动的配置打印（rank 数、序列长度、隐藏维、每卡 token 数、每 block 线程数）； - 与 CPU 参考结果的逐元素比对，最终打印
SUCCESS: Example completed successfully!或FAILED: Incorrect results detected on rank N。
性能特征
- 最佳场景：Hopper+ GPU 搭配 NVLink/NVSwitch；
- 访存方式：硬件多播归约（reduce）与多播存储（store），数据直接从 peer 内存进入寄存器，无中间拷贝，也无需 peer 循环带来的额外指令开销；
- 扩展性：受限于单节点 / 单 NVLink 域（这是 LSA/multimem 的固有拓扑约束；跨节点场景请参考同目录的 GIN 与 Hybrid 示例）。
需要跨节点融合时，可参考 03_rmsnorm_gin（纯 GIN，多节点）与 04_rmsnorm_hybrid（LSA 用于节点内、GIN 用于节点间的混合方案）。
源码级的实现约束（进阶）
从 src/include/nccl_device/impl/multimem__funcs.h 的LoadImpl模板可以确认 multimem 归约加载的两条硬性约束：
- 归约算子仅支持求和（OpSum）：
static_assert(!UseMultimem || std::is_same<RedOp, OpSum<...>>::value, "Multimem sources only support OpSum - use LSA sources for custom RedOp")——需要自定义归约（如 max/min）时只能退回 LSA source； - 不支持 int8/uint8：
static_assert(!UseMultimem || !std::is_same<PackEltType, int8_t>::value && ...)，8 位整型需改用 LSA source。
此外，multimem 在 NCCL 库内为 float/double/half/bf16/fp8/int32/uint32/int64/uint64 等类型提供了从 32 位到 128 位的向量化multimem.ld_reduce特化（如multimem.ld_reduce.global.add.v2.f32、.v4.f32、acc::f32.f16x2等），其中 fp8 特化仅在特定架构（sm_100a/101a/120a/121a 等）上可用；multimemStore则按字节宽度（1/2/4/8/16 字节）选择st.global.b8/b16或multimem.st.global.b32/b64/v4.f32。这些细节说明：Multimem 虽然在内核里写法简洁，但其能力边界（求和算子、位宽、架构）需要在使用时仔细核对。
常见问题排查
- 报错 "compute capability 9.0 or higher"：Multimem 需要 SM 9.0+（Hopper H100 或更新），可用
nvidia-smi确认 GPU 型号，旧卡请改用 01_rmsnorm_lsa 示例。 sequence_length无法被 GPU 数整除：调整代码中的sequence_length为其整数倍。- LSA barrier 失败：用
nvidia-smi topo -m确认 NVLink 连通性；多节点/跨域场景应改用 GIN/Hybrid 示例。 - 编译报未解析符号：检查
rmsnorm_utils.cuh的 include 路径是否配置正确。 - 诊断信息不足：设置
NCCL_DEBUG=INFO获取更详细的 NCCL 日志。
小结
Multimem 示例展示了 NCCL Device API 中一类极具"硬件加速"色彩的通信方式：在 SM 9.0+ 上，一次多播 load 即完成跨 peer 归约、一次多播 store 即完成跨 peer 广播，把分布式 RMSNorm 的 reduce-scatter 与 all-gather 阶段以最少的指令融入单次内核启动。若想进一步对比，建议同时阅读 01_rmsnorm_lsa（显式 peer 迭代的基线写法）、03_rmsnorm_gin 与 04_rmsnorm_hybrid（跨节点方案），并结合 07_kernel_fusion 分类 README 的数学背景与主机流程理解全貌。
- 高性能计算
- 通信
- 深度学习
- 分布式训练
【免费下载链接】nccl
Optimized primitives for collective multi-GPU communication
相关推荐
NCCL 纯 LSA 融合内核实战：单内核实现分布式 RMSNorm（Reduce-Scatter → RMSNorm → All-Gather）
NCCL 纯 LSA 融合内核实战：单内核实现分布式 RMSNorm（Reduce Scatter → RMSNorm → All Gather） 导读 本文以
高性能计算通信深度学习分布式训练NCCL 设备端 API 实战：用 LSA 在单 GPU 内核中融合分布式 RMSNorm（C/CUDA 示例解析）
NCCL 设备端 API 实战：用 LSA 在单 GPU 内核中融合分布式 RMSNorm（C/CUDA 示例解析） 导读 本文讲解 NCCL 仓库中 docs
高性能计算通信深度学习分布式训练NCCL 定制 AllReduce 内核实战：基于 Device API 与 NVLink 硬件多播（Multimem）的单/双程集合通信
NCCL 定制 AllReduce 内核实战：基于 Device API 与 NVLink 硬件多播（Multimem）的单/双程集合通信 本文围绕 NCCL
高性能计算通信深度学习分布式训练
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
