rapidsai/rmm

What is RAPIDS RMM?

RMM (RAPIDS Memory Manager) is a C++/Python library that gives you a common, flexible way to allocate and free GPU memory. It lets you plug in different allocation strategies (plain cudaMalloc, pooled allocators, fixed‑size blocks, managed memory, etc.) and use them from both C++ code and Python via the rmm package.

Why does it matter?

GPU kernels run fastest when memory is allocated efficiently. The default CUDA allocator (cudaMalloc/cudaFree) can be slow for many small allocations, and using pinned host memory or stream‑ordered allocations can further improve performance. RMM abstracts these details behind a simple interface so you can:

• Choose the best allocator for your workload (e.g., a pool that re‑uses freed blocks).
• Change the allocator globally or per‑device without touching the rest of your code.
• Write code that works with any allocator because the API is the same.

Core concepts

Concept	What it does
Memory resource	An object that implements allocate/deallocate following the CCCL memory‑resource concept. RMM ships several ready‑made resources (cuda_memory_resource, managed_memory_resource, pool_memory_resource, fixed_size_memory_resource, binning_memory_resource).
Stream‑ordered allocation	Allocation and deallocation are tied to a CUDA stream, allowing the allocator to reuse memory without extra synchronization.
Current‑device resource	A global (per‑CUDA‑device) resource that is used when a function does not receive an explicit allocator. You can get or set it with get_current_device_resource_ref() / set_current_device_resource().
Device data structures	Helper classes such as rmm::device_buffer, rmm::device_uvector, and rmm::device_vector that allocate their storage through the chosen memory resource.
CUDA stream wrappers	rmm::cuda_stream (RAII owning stream) and rmm::cuda_stream_view (non‑owning view) make it easy to work with streams, and rmm::cuda_stream_pool lets you reuse a pool of streams.

How you use it (C++)

#include <rmm/mr/device/pool_memory_resource.hpp>
#include <rmm/device_uvector.hpp>

int main(){
    // Default allocator is plain cudaMalloc
    rmm::mr::cuda_memory_resource cuda_mr;
    // Create a pool that starts with 50 % of free device memory
    auto init = rmm::percent_of_free_device_memory(50);
    rmm::mr::pool_memory_resource pool_mr{cuda_mr, init};

    // Make the pool the default for this device
    auto old = rmm::mr::set_current_device_resource(pool_mr);

    // Allocate a vector using the pool (stream‑ordered)
    rmm::device_uvector<float> vec(1024, rmm::cuda_stream_default);
    // …use vec in kernels…
}

The same concepts are exposed to Python via the rmm package, which wraps the C++ library and provides rmm.DeviceBuffer, rmm.DeviceArray, etc.

Getting it

• Python – install pre‑built wheels that match your CUDA version:

  pip install librmm-cu13 rmm-cu13   # for CUDA 13
  pip install librmm-cu12 rmm-cu12   # for CUDA 12
  

• Conda – conda install -c rapidsai -c conda-forge librmm rmm
• From source – clone the repo, create the rmm_dev conda environment, and run cmake && make && make install (or use the provided build.sh).

Who should care?

• Data‑science and ML engineers building GPU‑accelerated pipelines (e.g., RAPIDS cuDF, cuML) that need fast, predictable memory behavior.
• Library authors who want to give users the ability to plug in custom allocators without rewriting code.
• Performance‑critical applications (simulation, graphics, deep‑learning training loops) where allocation overhead matters.

---

Bottom line: RAPIDS RMM is a low‑level, high‑performance memory‑management toolkit for CUDA that lets you swap allocation strategies, use stream‑ordered pools, and integrate easily with C++ or Python GPU code.