How to Streamline Shared Memory Space With the NVSHMEM Communication Library

Akhil Langer, Seth Howell, Jim Dinan, Nvidia | GTC Spring 2023

1. 通信模型

1.1 CPU发起的通信

CPU-Initiated Communication Timeline (Page 2)

特点：

在GPU上进行计算。
由CPU发起通信。
在边界处进行同步。

普遍使用的模型，但存在问题：
* 在关键路径中存在卸载延迟（Offload latencies in critical path）。
* 通信未重叠（Communication is not overlapped）。
* 隐藏这些延迟会增加代码复杂性。
* 不隐藏延迟会限制强扩展性。

1.2 GPU发起的通信

GPU-Initiated Communication Timeline (Page 3)

特点：
* 在GPU上进行计算。
* 由GPU发起通信。

优点：
* 消除了卸载延迟（Eliminates offload latencies）。
* 计算和通信重叠（Compute and communication overlap）。
* 通过线程隐藏延迟（Latencies hidden by threading）。
* 更易于用内联通信表达算法（Easier to express algorithms with inline communication）。

目标： 在提升性能的同时，简化编程。

2. NVSHMEM 概述

2.1 NVSHMEM 简介

OpenSHMEM，为NVIDIA GPU集群上的最佳性能而适配

功能： 将集群中多个GPU的内存聚合到一个分布式全局地址空间中。
- 通过put、get、原子API进行数据访问。
- 提供集合通信API。
通信与CUDA执行模型集成：
1. GPU核函数发起的（GPU kernel-initiated）操作。
2. CUDA流/图上的操作（Operations on CUDA streams/graphs）。
3. CPU发起的（CPU initiated）操作。
可与CPU OpenSHMEM或MPI库结合使用，用于主机内存通信。

2.2 NVSHMEM 对称内存模型

独立分区聚合到全局地址空间

对称对象在每个PE上以相同大小集体分配。
- 对称内存： nvshmem_malloc(...)
- 私有内存： cudaMalloc(...)
读： nvshmem_get(...)
写：： nvshmem_put(...)
原子操作： nvshmem_atomic_add(...)
刷新写操作： nvshmem_quiet()
顺序写操作： nvshmem_fence()
同步： nvshmem_barrier()
轮询： nvshmem_wait_until(...)

2.3 NVLink 通信优化

使用第四代NVLink无缝扩展至256个GPU

NVLink Scale-Up and InfiniBand/RoCE Scale-Out (Page 7)

NVSHMEM无缝扩展：
- 使用NVLink和PCIe进行扩展（Scales-up）。
- 使用InfiniBand、RoCE等进行扩展（Scales-out）。
内部使用CUDA IPC和cuMem API将对等PE的对称内存映射到虚拟地址空间。
- nvshmem_put/get on device → load/store
- nvshmem_put/get_on_stream → cudaMemcpyAsync
- nvshmem_ptr direct pointer bypass → kernel direct load/store

2.4 纵向和横向扩展带宽

理论对等体双向带宽 (GB/s)

PCI Gen4 x16：64
PCI Gen5 x16：128
第二代NVLink (DGX V100)：2400
第三代NVLink (DGX A100)：2400
第四代NVLink (DGX H100)：3600
第四代NVLink (32 DGX H100)：57600

理论网络注入 + 弹出带宽每GPU (GB/s)
* 100Gb ConnectX-5 (DGX V100)：25
* 200Gb ConnectX-6 (DGX A100)：50
* 400Gb ConnectX-7 (DGX H100)：100

2.5 线程级通信

Thread-Level Communication Diagram and Code (Page 9)

允许细粒度通信和重叠。
在DGX系统上高效映射到NVLink网络。

_global_ void stencil_single_step(float *u, float *v, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    compute(u, v, ix, iy);
    // Thread-level data communication API
    if (iy == 1)
        nvshmem_float_p(u+(ny*1)*nx+ix, u[nx+ix], top_pe);
    if (iy == ny)
        nvshmem_float_p(u+ix, u[ny*nx+ix], bottom_pe);
}

for (int iter = 0; iter < N; iter++) {
    swap(u, v);
    stencil_single_step<<<..., stream>>>(u, v, ...);
    nvshmem_barrier_all_on_stream(stream);
}

2.6 线程组通信

Thread-Group Communication Diagram and Code (Page 10)

NVSHMEM操作可以由块/warp中的所有线程发出。
通过IB等网络进行更高效的数据传输。
仍然允许warp间/块间重叠。

_global_ void stencil_single_step(float *u, float *v, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    compute(u, v, ix, iy);
    // Thread block-level communication API
    int boffset = get_block_offet(blockIdx, blockDim);
    if (blockIdx.y == 0)
        nvshmemx_float_put_nbi_block(u+(ny*1)*nx+boffset, u+nx+boffset, blockDim.x, top_pe);
    if (blockIdx.y == (blockDim.y-1))
        nvshmemx_float_put_nbi_block(u+boffset, u+ny*nx+boffset, blockDim.x, bottom_pe);
}

for (int iter = 0; iter < N; iter++) {
    swap(u, v);
    stencil_single_step<<<..., stream>>>(u, v, ...);
    nvshmem_barrier_all_on_stream(stream);
}

2.7 核内同步

In-Kernel Synchronization Diagram and Code (Page 11)

在核函数内跨PE进行集体或点对点同步。
将应用程序的更大一部分卸载到同一个CUDA核函数中。

_global_ void stencil_multi_step(float *u, float *v, int N, int *sync, ...) {
    int ix = get_ix(blockIdx, blockDim, threadIdx);
    int iy = get_iy(blockIdx, blockDim, threadIdx);

    for (int iter = 0; iter < N; iter++) {
        swap(u, v);
        // Thread block-level data exchange (assume even/odd iter buffering)
        int boffset = get_block_offet(blockIdx, blockDim);
        if (blockIdx.y == 0)
            nvshmemx_float_put_nbi_block(u+(ny*1)*nx+boffset, u+nx+boffset, blockDim.x, top_pe);
        if (!((blockIdx.y == (blockDim.y-1))))
            nvshmemx_float_put_nbi_block(u+boffset, u+ny*nx+boffset, blockDim.x, bottom_pe);

        this_grid.sync();
        if ((!itid) nvshmem_barrier(); // Be aware of synchronization costs. Best synchronization approach is application dependent!
        this_grid.sync();
    }
}

2.8 流序操作

NVSHMEM CPU发起的在CUDA流和图上入队的操作

Stream-ordered Operations Diagram (Page 12)

并非总是将所有通信或同步移入CUDA核函数是最佳选择。
Inter-CTA同步（例如grid.sync()）的延迟可能比核函数启动延迟更长。
允许混合细粒度通信和粗粒度同步。

2.9 集体核函数启动

确保使用设备侧核间同步时的进度

Collective Kernel Launch Table (Page 13)

NVSHMEM 用法	CUDA 核函数启动
设备发起的通信	执行配置语法 `<<<...,>>>` 或启动API
设备发起的同步	`nvshmemx_collective_launch`

CUDA的吞吐量计算模型允许（并鼓励）比GPU能容纳的更大的网格。
核间同步要求生产者和消费者线程并发执行。
集体启动通过CUDA协作启动和1PE/GPU的要求来保证共驻。

2.10 与CUDA图的互操作性

CUDA Work in Streams to Graph of Dependencies (Page 14)

CUDA Graphs (CUDA图)：
- 在CUDA 10.0中引入。
- 减少启动开销。
- 工作流优化。
NVSHMEM可与Graphs组合使用：
- 流上操作作为节点。
- CUDA核函数可以使用NVSHMEM。

2.11 与MPI/OpenSHMEM的互操作性

通过基于属性的初始化例程启用

MPI and OpenSHMEM Initialization Code Snippets (Page 15)

MPI：
使用MPI_COMM_WORLD初始化NVSHMEM。

MPI_Init(&argc, &argv);

MPI_Comm mpi_comm = MPI_COMM_WORLD;
nvshmemx_init_attr_t attr;
attr.mpi_comm = &mpi_comm;
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr);

mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
CUDA_CHECK(cudaSetDevice(mype_node));

OpenSHMEM：
使用SHMEM默认上下文初始化NVSHMEM。

shmem_init();

nvshmemx_init_attr_t attr;
nvshmemx_init_attr(NVSHMEMX_INIT_WITH_SHMEM, &attr);

mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE);
CUDA_CHECK(cudaSetDevice(mype_node));

3. NVSHMEM 最新特性

3.1 发布 NVSHMEM 2.9.0

自 NVSHMEM 2.5.0 / GTC Spring 2022 以来新增的特性：

InfiniBand GPUDirect Async (IBGDA) 传输
HPC SDK 集成
- 库版本控制 (Library versioning)
- CUDA 次要版本兼容性支持 (CUDA minor version compatibility support)
- 模块化传输、引导支持 (Modular transport, bootstrap support)
CMake 构建系统 (CMake build system)
集合优化 (Collective Optimizations)
内联设备代码 (Inlined Device Code)
DMABUF 支持
- 与上游 ibverbs 兼容 (Compatibility with upstream ibverbs)
- 无需 nv_peer_mem 模块 (nv_peer_mem module not needed)
Slingshot-11 支持
- Beta 版支持 (Beta support)
- 稳定性与性能的持续改进工作 (Ongoing work on stability and performance)

3.2 GPUDirect Async 内核发起通信

与 Mellanox NICs 更紧密的集成

CPU 代理线程不再用于内核发起的同步操作。
内核可以直接执行：
- 1. 写入内存数据
- 2-4. 提交工作请求至 NIC
- 5-8. 执行通信
- 1. 检查操作完成
并行度显著提升，从而提高了吞吐量。
与 Mellanox Stack 集成：
- 支持 NVIDIA Mellanox 网络
- 直接在 Mellanox 软件栈上运行
- 通过利用原生原子操作避免了对 GDRCopy 的需求

GPUDirect Async Kernel Initiated Communication

3.3 GPU-发起 (IBGDA) 与代理通信性能

在 2kB 和 16kB 传输大小下的全带宽

IBGDA (左图) 性能随着线程块 (CTAs) 的数量扩展；代理 (右图) 在 8 CTAs 时达到扩展极限。
使用 128 CTAs 时，IBGDA 块 Put 吞吐量提高了 5.75 倍；使用 1 CTA 时，IBRC 块 Put 吞吐量提高了 1.8 倍。
权衡：当 CTAs 并行通信时，IBGDA 提供高吞吐量；代理具有较低的操作切换开销。
测试节点：Juwels-Booster Node: 2x AMD EPYC 7402, 4x NVIDIA A100 40GB, 4x NVIDIA Mellanox ConnectX-6 200 Gbps HCAs。

GPU-Initiated (IBGDA) vs Proxy Communication Performance

3.4 标量 Put 合并

将多个 NVSHMEM "p" 操作合并为一个工作请求
* IBGDA 合并 (IBGDA Coalescing) 在 128 个 CTA 时达到 1935 MOPS 的消息速率，远高于 IBRC 的 180 MOPS 和 IBGDA 的 1.7 MOPS。

3.5 二进制分发改进

改进了可移植性和可用性
* HPC_SDK 在多个平台上分发 NVSHMEM 及其他兼容软件包。
* HPC_SDK 中使用的 NVSHMEM 二进制包的近期更改通常也很有用。
* NVSHMEM 现在支持 CUDA 次要版本兼容性。
* 编译器限制：所有静态链接的部分必须使用最旧的编译器版本构建。

模块化传输从主要的 NVSHMEM 库中移除了特定于传输的要求。
- 用于构建传输的文件现已包含在软件包中。
模块化引导接口已得到改进。

3.6 CMAKE 构建系统

CMake 支持构建所有 NVSHMEM 组件 (源码安装、软件包、测试、性能测试、示例)。
替代 Makefiles，支持所有 NVSHMEM_* 环境变量。
也可以通过 CMake 原生方式提供选项。
使用 CMake 命令行选项 "CUDA_ARCHITECTURES" 更改配置的架构，而非 NVCC_GENCODE。
生成的 NVSHMEMConfig.cmake 文件允许轻松地将目标集成到现有 CMake 项目中。
自 v2.9 起，NVSHMEM 正式弃用 Makefiles，转而支持 CMake 构建系统。

3.7 内联设备代码

改进了延迟
* 设备 API 实现：
* P2P 内存复制 (内联)
* 远程通信 (未内联)
* 集合 API 实现 (未内联)

此功能使所有设备 API 实现都能够内联。
通过 NVSHMEM 构建时宏启用：NVSHMEM_ENABLE_ALL_DEVICE_INLINING = 0/1。
所有设备 API 实现已移至 NVSHMEM 头文件。
通过 __forceinline__ 限定符启用内联。
优点：
- 无函数调用开销。
- CUDA 编译器可以全面了解内核代码，从而改进编译器优化、寄存器使用等。

3.8 集合优化

3.8.1 使用 LL 协议改进延迟

典型的消息传递方法：nvshmem_put + nvshmem_fence + nvshmem_signal。
LL (低延迟) 协议：
- 将 4 字节数据与 4 字节标志交织。
- 消息大小加倍但同步成本低。
- 适用于中小型消息的低延迟。
- 乒乓延迟最多提高 75%。

3.8.2 拓扑感知优化

拓扑感知算法：
- 形成子团队：team_node 和 team_same_mype_node。
广播：
- 向节点领导者广播。
- 节点内广播。
规约 (Reduce)：
- 扁平化的 All-to-all：
  - 每个 PE 将数据发送给其他所有 PE。
  - 每个 PE 执行所有数据的本地规约。
- 拓扑感知 All-to-all：
  - 在节点内执行扁平化的 all-to-all 规约。
  - 跨 team_same_mype_node 执行扁平化的 all-to-all 规约。
性能提高 95%。

4. 性能案例研究

4.1 NVSHMEM + cuFFT = cuFFTMp

多节点多 GPU：使用 NVIDIA cuFFTMp FFTs 进行规模化计算。
- 更多信息请访问：https://developer.nvidia.com/blog/multi-node-multi-gpu-using-nvidia-cufftmp-ffts-at-scale
JAX + cuFFTMp：
- JAX-only、JAX+cuFFTMp 和 cuFFTMp 在 Selene 上进行 3D C2C FP32 FFT 的强扩展性测试 (2022 年 12 月)。
- JAX 最新进展 [S51956] - Frederic 等人。
GROMACS + cuFFTMp：
- STMV 强扩展性 (1M 原子)
- BenchPEP-h 强扩展性 (12M 原子)
- NVIDIA 数学库最新进展 [S51176] - Harun Bayraktar 等人。

4.2 NVSHMEM 在 GROMACS 中的应用

如前一幻灯片所述，cuFFTMp 自动使用 NVSHMEM。
在执行分布式 FFTs 时，PME GPU 之间存在全对全通信模式。
GROMACS 中还存在其他可能受益于 NVSHMEM 的通信模式：
- PP (粒子-粒子) Halo 交换
- PME (粒子网格 Ewald) Halo 交换
- PP-PME 通信
原型版本使用 NVSHMEM 进行 PP halo 交换通信。
为了评估性能，使用了反应场 (RF) 案例。
作为 PME 远距离力计算的替代方案。
在扩展到多个节点时，对 PP Halo 交换特别敏感。
在扩展到多个节点时，PP halo 交换具有显著的性能优势。
参考：分子动力学前沿 CUDA 技术 [S51110] - Alan Gray 等人。

4.3 NVSHMEM 在高性能 Linpack (HPL) 中的应用

传统 HPL：
- GPU 用于有限的操作，例如 GEMM。
- 因子分解 (LU 分解) 和通信由 CPU 发起和处理。
HPL + NVSHMEM：
- 无 CPU 版本。
- 所有计算和通信都在 GPU 上完成。
- 矩阵因子分解和 GEMM 在不同的内核中并行完成。
- 因子分解需要频繁的细粒度通信。
- HPL 实现变为无 CPU，因此无需针对不同类型的 CPU 进行优化。
- 性能提升高达 8%。

NVSHMEM in High Performance Linpack (HPL)

5. 即将推出的特性和结论

5.1 NVSHMEM 中的 Grace Hopper 支持

本节详细介绍了 Grace Hopper 超级芯片架构及其在 NVSHMEM 中的支持。

NVIDIA Grace Hopper 超级芯片
- NVIDIA Grace Hopper Superchip Architecture (Page 31)
- Grace Hopper 超级芯片支持 Grace Arm CPU 和 Coherent Hopper GPU。
- Grace CPU: 配备 2 个 LPDDR5X，总计 512 GB 内存，带宽为 546 GB/s。
- Hopper GPU: 配备 2 个 HBM3，总计 96 GB 内存，带宽为 3000 GB/s。
- NVLink C2C: 连接 CPU 和 GPU，提供 900 GB/s 带宽。
- 高速 I/O: 包含 4x 16x PCIe-5 接口，提供 512 GB/s 带宽。
- NVLink 网络: 包含 18x NVLink 4 接口，提供 900 GB/s 带宽，可连接多达 256 个 GPU。
- 架构亮点包括：高速网络 I/O、大容量主内存、高 CPU/GPU 带宽、硬件一致性以及第四代 NVLink。

5.2 第四代 NVLink

第四代 NVLink 进一步增强了互联性能。

每个 NVSwitch 拥有 64 个 NVLink 端口。
端口配备 NVIDIA Scalable Hierarchical Aggregation Reduction Protocol (SHARP)™ 引擎。
提供网络内规约 (in-network reductions) 和多播加速 (multicast acceleration) 功能。
优化 NVSHMEM Collectives:
- 可将大型 all-reduce 操作的带宽加倍。
- 小消息集合操作的延迟更低。
- 屏障延迟降低。

5.3 NVSHMEM 内存空间与上下文

本节阐述了 NVSHMEM 内存空间和上下文，以实现对 Grace CPU 内存的访问并改进通信管理。

内存空间 (Memory Spaces)
- 用户定义与团队关联的对称内存。
- 包括新的对称堆、注册用户缓冲区和继承的对称空间。
- 优势: 增加了灵活性，提供了对额外内存类型的访问。
上下文 (Contexts)
- 管理一组 RMA (远程内存访问) 和 AMO (原子内存操作)。
- Quiet 和 Fence 操作被限定在特定上下文内。
- 优势: 线程/SM 隔离减少了竞争（例如，IBGDA QP 分配给 SM），实现流水线非阻塞通信。
- 图示展示了 nvshmem_team_split 创建 Team，其中包含内存空间参数 team_config_t 和 Memory Space。Memory Space 包含 Symmetric Heap 和 Registered User Buffer(s)。
- nvshmem_ctx_create 创建 Context，与 Resources 通过 ctx_options_t 关联。Context 处理 NVSHMEM RMA and AMO Operations Quiet and Fence。

5.4 总结

NVSHMEM 是一个用于 NVIDIA GPU 集群的 PGAS (Partitioned Global Address Space) 库。

NVSHMEM 可无缝扩展，支持：
- 使用 NVLink 连接 GPU 的节点级 GPU 编程。
- 通过 InfiniBand 或 RoCE 连接的多节点 GPU 集群。
NVSHMEM 提供 Stream/Graph、GPU 内核启动和 CPU 启动的 API。
与 CUDA 编程模型的集成可以提高 GPU 集群的性能并简化扩展。

5.5 NVSHMEM 状态

NVSHMEM 的最新发布信息和关键特性。

最新版本: NVSHMEM 2.9.0，于 2023 年 3 月发布。
可在 https://developer.nvidia.com/nvshmem 下载。
NVSHMEM 2.8 作为 NVIDIA HPC SDK 发布的一部分提供，容器可通过 NGC 获取。
新特性 (New Features):
- Hopper GPU 支持。
- 优化的 GDA-KI 网络传输。
- 设备代码内联。
- CMake 构建系统。
- 改进的集合操作性能。
实现特性 (Implementation Features):
- NVLink 和 PCIe P2P 支持。
- InfiniBand、RoCE 和 Slingshot (beta) 支持。
- X86 和 Power9 支持。
- 可与 MPI 和 OpenSHMEM 互操作。

PaperCache

How to Streamline Shared Memory Space With the NVSHMEM Communication Library

How to Streamline Shared Memory Space With the NVSHMEM Communication Library

目录

1. 通信模型

1.1 CPU发起的通信

1.2 GPU发起的通信

2. NVSHMEM 概述

2.1 NVSHMEM 简介

2.2 NVSHMEM 对称内存模型

2.3 NVLink 通信优化

2.4 纵向和横向扩展带宽

2.5 线程级通信

2.6 线程组通信

2.7 核内同步

2.8 流序操作

2.9 集体核函数启动

2.10 与CUDA图的互操作性

2.11 与MPI/OpenSHMEM的互操作性

3. NVSHMEM 最新特性

3.1 发布 NVSHMEM 2.9.0

3.2 GPUDirect Async 内核发起通信

3.3 GPU-发起 (IBGDA) 与代理通信性能

3.4 标量 Put 合并

3.5 二进制分发改进

3.6 CMAKE 构建系统

3.7 内联设备代码

3.8 集合优化

3.8.1 使用 LL 协议改进延迟

3.8.2 拓扑感知优化

4. 性能案例研究

4.1 NVSHMEM + cuFFT = cuFFTMp

4.2 NVSHMEM 在 GROMACS 中的应用

4.3 NVSHMEM 在高性能 Linpack (HPL) 中的应用

5. 即将推出的特性和结论

5.1 NVSHMEM 中的 Grace Hopper 支持

5.2 第四代 NVLink

5.3 NVSHMEM 内存空间与上下文

5.4 总结

5.5 NVSHMEM 状态

How to Streamline Shared Memory Space With the NVSHMEM Communication Library

目录

1. 通信模型

1.1 CPU发起的通信

1.2 GPU发起的通信

2. NVSHMEM 概述

2.1 NVSHMEM 简介

2.2 NVSHMEM 对称内存模型

2.3 NVLink 通信优化

2.4 纵向和横向扩展带宽

2.5 线程级通信

2.6 线程组通信

2.7 核内同步

2.8 流序操作

2.9 集体核函数启动

2.10 与CUDA图的互操作性

2.11 与MPI/OpenSHMEM的互操作性

3. NVSHMEM 最新特性

3.1 发布 NVSHMEM 2.9.0

3.2 GPUDirect Async 内核发起通信

3.3 GPU-发起 (IBGDA) 与代理通信性能

3.4 标量 Put 合并

3.5 二进制分发改进

3.6 CMAKE 构建系统

3.7 内联设备代码

3.8 集合优化

3.8.1 使用 LL 协议改进延迟

3.8.2 拓扑感知优化

4. 性能案例研究

4.1 NVSHMEM + cuFFT = cuFFTMp

4.2 NVSHMEM 在 GROMACS 中的应用

4.3 NVSHMEM 在高性能 Linpack (HPL) 中的应用

5. 即将推出的特性和结论

5.1 NVSHMEM 中的 Grace Hopper 支持

5.2 第四代 NVLink

5.3 NVSHMEM 内存空间与上下文

5.4 总结

5.5 NVSHMEM 状态

登录

注册

忘记密码

重发验证邮件