rapidsai/rmm

NVIDIA RMM (RAPIDS Memory Manager)

它是什么 – RMM 是一个 C++/Python 库，为分配 GPU 内存提供了一个通用且灵活的接口。它让开发人员能够插入不同的分配策略（plain cudaMalloc、池化分配器、托管内存等），并透过 rmm 包在 C++ 代码和 Python 中使用它们。

为什么它很重要 – 在高性能 GPU 工作负载中，重复调用 cudaMalloc/cudaFree 的成本可能会占据执行时间的大部分。RMM 的池化和分桶分配器能显著降低该开销，并支持 流排序 (stream-ordered) 分配，这意味着内存可以在同一个 CUDA 流上重复使用，而无需额外的同步。

主要功能

• 多种内存资源实现 – cuda_memory_resource、managed_memory_resource、pool_memory_resource、fixed_size_memory_resource、binning_memory_resource。
• 流排序分配 – 在特定的 CUDA 流上进行分配和释放，实现快速重复使用。
• 每个设备的默认资源 – 您可以为每个 GPU 设置全局分配器，或在每次分配时进行覆盖。
• C++ API – CUDA 分配函数的轻量级封装，以及像 rmm::device_buffer、rmm::device_uvector 和 rmm::device_vector 这样的 RAII 容器。
• Python 绑定 – 可透过 pip/conda 安装（librmm-cuXX 和 rmm-cuXX）。
• ABI 版本控制 – 多个版本可以在同一个进程中并存。

典型使用案例

rmm::mr::cuda_memory_resource cuda_mr;
auto pool = rmm::mr::pool_memory_resource{cuda_mr,
                                          rmm::percent_of_free_device_memory(50)};
auto prev = rmm::mr::set_current_device_resource(pool);
// Subsequent allocations use the pool automatically
rmm::device_buffer buf(1024, cuda::stream_ref{cudaStreamDefault}, pool);

在 Python 中，您可以使用 import rmm; rmm.reinitialize(pool_allocator=True) 来切换到池化分配器。

如何获取 – 安装预先构建的 wheels（例如 CUDA 13 使用 pip install librmm-cu13 rmm-cu13）或使用 CMake 和提供的 build.sh 从源代码构建。存储库包含完整的说明文档、教程和 CMake 集成辅助工具。

谁在使用它 – 作为 NVIDIA RAPIDS 生态系统的一部分；下游项目如 cuDF、cuML 以及其他 GPU 加速的数据科学库都依赖 RMM 进行高效的内存处理。