rapidsai/rmm

NVIDIA RMM

NOTE: For the latest stable README.md ensure you are on the main branch.

Resources

• RMM Reference Documentation: Python and C++ API references, tutorials, and topic guides.
• Installation: Instructions for installing RMM.
• GitHub Repository: Download the RMM source code.
• Issue Tracker: Report issues or request features.

Overview

Achieving optimal performance in GPU-centric workflows frequently requires customizing how host and device memory are allocated. For example, using "pinned" host memory for asynchronous host <-> device memory transfers, or using a device memory pool sub-allocator to reduce the cost of dynamic device memory allocation.

The goal of RMM is to provide:

• A common interface that allows customizing memory allocation on device and host
• A collection of implementations of the interface
• A collection of data structures that use the interface for memory allocation

For information on the interface RMM provides and how to use RMM in your C++ code, see below.

For a walkthrough of the design of RMM, read Fast, Flexible Allocation for NVIDIA CUDA with RMM on the NVIDIA Developer Blog.

Installation

System Requirements

Please see the Installation Guide for NVIDIA CUDA-X libraries for data science for information about supported operating systems, GPU drivers, and CUDA versions.

pip

Stable releases of librmm and rmm are available on PyPI. Match the package suffix to the major CUDA version supported by your installed driver.

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

Development versions are available as nightly releases:

# CUDA 13
pip install --pre \
  --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
  librmm-cu13
pip install --pre \
  --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
  rmm-cu13
CUDA 12
pip install --pre \
  --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
  librmm-cu12
pip install --pre \
  --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple \
  rmm-cu12

conda

Stable releases of librmm and rmm are available from the rapidsai channel. Development versions are available from the rapidsai-nightly channel.

# Stable
conda install -c rapidsai -c conda-forge librmm
conda install -c rapidsai -c conda-forge rmm
Nightly
conda install -c rapidsai-nightly -c conda-forge librmm
conda install -c rapidsai-nightly -c conda-forge rmm

Building from Source

Get RMM Dependencies

Compiler requirements:

• gcc version 13.3+
• nvcc version 12.9+
• cmake version 4.0+

CUDA/GPU runtime requirements:

• CUDA 12.2+. You can obtain CUDA from

https://developer.nvidia.com/cuda-downloads

GPU Support:

• RMM is tested and supported only on Volta architecture and newer (Compute Capability 7.0+).

Python requirements:

• rapids-build-backend (available from PyPI or the rapidsai conda channel)
• scikit-build-core
• cuda-python
• cython

For more details, see pyproject.toml

Script to build RMM from source

To install RMM from source, ensure the dependencies are met and follow the steps below:

• Clone the repository:

$ git clone https://github.com/rapidsai/rmm.git
$ cd rmm

• Create the conda development environment rmm_dev:

# create the conda environment (assuming in base rmm directory)
$ conda env create --name rmm_dev --file conda/environments/all_cuda-133_arch-$(uname -m).yaml
activate the environment
$ conda activate rmm_dev

• Build and install librmm using cmake & make. CMake depends on the nvcc executable being on

your path or defined in CUDACXX environment variable.

$ mkdir build                                       # make a build directory
$ cd build                                          # enter the build directory
$ cmake .. -DCMAKE_INSTALL_PREFIX=/install/path     # configure cmake ... use $CONDA_PREFIX if you're using a conda environment
$ make -j                                           # compile the library librmm.so ... '-j' will start a parallel job using the number of physical cores available on your system
$ make install                                      # install the library librmm.so to '/install/path'

• Building and installing librmm and rmm using build.sh. build.sh creates a build directory

at the root of the git repository. build.sh depends on the nvcc executable being on your path or defined in the CUDACXX environment variable.

$ ./build.sh -h                                     # Display help and exit
$ ./build.sh -n librmm                              # Build librmm without installing
$ ./build.sh -n rmm                                 # Build rmm without installing
$ ./build.sh -n librmm rmm                          # Build librmm and rmm without installing
$ ./build.sh librmm rmm                             # Build and install librmm and rmm

• To run tests (Optional):

$ cd build (if you are not already in the build directory)
$ make test

• Build, install, and test the rmm python package, in the python folder:

# In the root rmm directory
$ python -m pip wheel ./python/librmm
$ python -m pip install --find-links=. -e ./python/rmm
$ pytest -v

Done! You are ready to develop for the RMM project.

Caching third-party dependencies

RMM uses CPM.cmake to handle third-party dependencies like CCCL, GoogleTest, GoogleBenchmark. In general you won't have to worry about it. If CMake finds an appropriate version on your system, it uses it (you can help it along by setting CMAKE_PREFIX_PATH to point to the installed location). Otherwise those dependencies will be downloaded as part of the build.

If you frequently start new builds from scratch, consider setting the environment variable CPM_SOURCE_CACHE to an external download directory to avoid repeated downloads of the third-party dependencies.

ABI versioning

RMM symbols are placed in an inline namespace derived from the RMM major and minor version. The public API remains available through the rmm:: namespace, while static RMM libraries built for different ABI versions can coexist in one process. Process-global state is shared by RMM copies with the same ABI version and kept separate across different ABI versions.

Using RMM in a downstream CMake project

The installed RMM library provides a set of config files that makes it easy to integrate RMM into your own CMake project. Add the following to CMakeLists.txt:

find_package(rmm [VERSION])
...
target_link_libraries( (PRIVATE|PUBLIC|INTERFACE) rmm::rmm)

This links librmm, makes RMM headers available, and pulls in transitive dependencies. If RMM is not installed in a default location, use CMAKE_PREFIX_PATH or rmm_ROOT to point to its location.

Using CPM to manage RMM

RMM uses CPM.cmake to manage its dependencies, including CCCL, and you can use CPM for your project's dependency on RMM.

There is an issue with using CPM's single-argument compact syntax for RMM/CCCL as it transitively marks targets as SYSTEM dependencies. This causes the CCCL headers pulled in through CPM to be of lower priority to the preprocessor than the (potentially outdated) CCCL headers provided by the CUDA SDK. To avoid this issue, use CPM's multi-argument syntax instead:

CPMAddPackage(NAME rmm [VERSION]
              GITHUB_REPOSITORY rapidsai/rmm
              SYSTEM OFF
              SOURCE_SUBDIR cpp)
...
target_link_libraries( (PRIVATE|PUBLIC|INTERFACE) rmm::rmm)

Using RMM in C++

The first goal of RMM is to provide a common interface for device memory allocation. This allows both _users_ and _implementers_ of custom allocation logic to program to a single interface.

RMM's memory resources use CCCL's memory resource concepts. Resource APIs accept either concrete resource objects, non-owning resource refs such as rmm::device_async_resource_ref, or owning type-erased resources such as cuda::mr::any_resource.

Memory Resources

A device memory resource satisfies the CCCL resource concept and provides stream-ordered allocation and deallocation:

void* allocate(cuda::stream_ref stream,
               std::size_t bytes,
               std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT);
void deallocate(cuda::stream_ref stream,
                void* ptr,
                std::size_t bytes,
                std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept;

RMM also uses rmm::device_async_resource_ref, an alias for cuda::mr::resource_ref, as a lightweight non-owning reference to a device resource. RMM uses cuda::mr::any_resource as an owning type-erased resource. RMM resources with non-trivial state are value types with shared ownership of their internal state, so copying a resource object is inexpensive and keeps the underlying state alive.

Stream-ordered Memory Allocation

RMM memory resources provide stream-ordered memory allocation. This allows optimizations such as re-using memory deallocated on the same stream without the overhead of synchronization.

A call to resource.allocate(stream_a, bytes) returns a pointer that is valid to use on stream_a. Using the memory on a different stream (say stream_b) is Undefined Behavior unless the two streams are first synchronized, for example by using cudaStreamSynchronize(stream_a) or by recording a CUDA event on stream_a and then calling cudaStreamWaitEvent(stream_b, event).

The stream specified to deallocate should be a stream on which it is valid to use the deallocated memory immediately for another allocation. Typically this is the stream on which the allocation was last used before the call to deallocate. The passed stream may be used internally by a memory resource for managing available memory with minimal synchronization, and it may also be synchronized at a later time, for example using a call to cudaStreamSynchronize().

For this reason, it is Undefined Behavior to destroy a CUDA stream that is passed to deallocate. If the stream on which the allocation was last used has been destroyed before calling deallocate or it is known that it will be destroyed, it is likely better to synchronize the stream (before destroying it) and then pass a different stream to deallocate (e.g. the default stream).

Note that device memory data structures such as rmm::device_buffer and rmm::device_uvector follow these stream-ordered memory allocation semantics and rules.

For further information about stream-ordered memory allocation semantics, read Using the NVIDIA CUDA Stream-Ordered Memory Allocator on the NVIDIA Developer Blog.

Available Device Resources

RMM provides several device memory resources to satisfy various user requirements. For more detailed information about these resources, see their respective documentation.

cuda_memory_resource

Allocates and frees device memory using cudaMalloc and cudaFree.

managed_memory_resource

Allocates and frees device memory using cudaMallocManaged and cudaFree.

Note that NVIDIA Virtual GPU Software (vGPU, for use with virtual machines or hypervisors) does not support managed_memory_resource by default. To support this, Unified Memory must be enabled for vGPU.

pool_memory_resource

A coalescing, best-fit pool sub-allocator.

fixed_size_memory_resource

A memory resource that can only allocate a single fixed size. Average allocation and deallocation cost is constant.

binning_memory_resource

Configurable to use multiple upstream memory resources for allocations that fall within different bin sizes. Often configured with multiple bins backed by fixed_size_memory_resources and a single pool_memory_resource for allocations larger than the largest bin size.

Default Resources and Per-device Resources

RMM users commonly need to configure a resource object to use for all allocations where another resource has not explicitly been provided. A common example is configuring a pool_memory_resource to use for all allocations to get fast dynamic allocation.

To enable this use case, RMM provides the concept of a resource for the currently active CUDA device. This resource is used when another is not explicitly provided.

Accessing and modifying this resource is done through two functions:

• device_async_resource_ref get_current_device_resource_ref()
• Returns a non-owning reference to the resource for the active CUDA device.
• The initial resource is an instance of cuda_memory_resource.
• This function is thread safe with respect to concurrent calls to it and

set_current_device_resource().

• For more explicit control, you can use get_per_device_resource_ref(), which takes a device ID.
• cuda::mr::any_resource set_current_device_resource(cuda::mr::any_resource new_mr)
• Updates the resource for the active CUDA device to new_mr
• Returns the previous resource as an owning type-erased resource
• This function is thread safe with respect to concurrent calls to it and

get_current_device_resource_ref()

• For more explicit control, you can use set_per_device_resource(), which takes a device ID.
• To reset to the initial resource, call reset_current_device_resource().

Example

rmm::mr::cuda_memory_resource cuda_mr;
// Construct a resource that uses a coalescing best-fit pool allocator
// With the pool initially half of available device memory
auto initial_size = rmm::percent_of_free_device_memory(50);
rmm::mr::pool_memory_resource pool_mr{cuda_mr, initial_size};
auto previous = rmm::mr::set_current_device_resource(pool_mr);
auto mr = rmm::mr::get_current_device_resource_ref();

Multiple Devices

A memory resource should only be used when the active CUDA device is the same device that was active when the resource was created. Otherwise behavior is undefined.

If a memory resource is used with a stream associated with a different CUDA device than the device for which the memory resource was created, behavior is undefined.

Creating a memory resource for each device requires care to set the current device before creating each resource. Here is an example loop that creates pool_memory_resource objects for each device and sets them as the per-device resource for that device.

for (int i = 0; i < N; ++i) {
  cudaSetDevice(i); // set device i before creating MR
  auto initial_size = rmm::percent_of_free_device_memory(50);
  rmm::mr::pool_memory_resource pool{rmm::mr::cuda_memory_resource{}, initial_size};
  // Set the per-device resource for device i
  set_per_device_resource(cuda_device_id{i}, pool);
}

Note that the CUDA device that is current when creating a memory resource must also be current any time that resource is used to deallocate memory, including in a destructor. The RAII class rmm::device_buffer and classes that use it as a backing store (rmm::device_scalar and rmm::device_uvector) handle this by storing the active device when the constructor is called, and then ensuring that the stored device is active whenever an allocation or deallocation is performed (including in the destructor). The user must therefore only ensure that the device active during _creation_ of an rmm::device_buffer matches the active device of the memory resource being used.

Here is an _incorrect_ example that creates a memory resource on device 0 and then uses it to allocate a device_buffer on device 1:

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  {
    RMM_CUDA_TRY(cudaSetDevice(1));
    // Invalid, current device is 1, but MR is only valid for device 0
    rmm::device_buffer buf(16, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}, mr);
  }
}

A correct example creates the device buffer with device 0 active. After that it is safe to switch devices and let the buffer go out of scope and destruct with a different device active. For example, this code is correct:

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  rmm::device_buffer buf(16, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}, mr);
  RMM_CUDA_TRY(cudaSetDevice(1));
  ...
  // No need to switch back to device 0 before ~buf runs
}

Use of rmm::device_vector with multiple devices

rmm::device_vector uses an rmm::mr::thrust_allocator to enable thrust::device_vector to allocate and deallocate memory using RMM. As such, the usual rules for usage of the backing memory resource apply: the active device must match the active device at resource construction time. To facilitate use in an RAII setting, rmm::mr::thrust_allocator records the active device at construction time and ensures that device is active whenever it allocates or deallocates memory. Usage of rmm::device_vector with multiple devices is therefore the same as rmm::device_buffer. One must _create_ device_vectors with the correct device active, but it is safe to destroy them with a different active device.

For example, recapitulating the previous example using rmm::device_vector:

{
  RMM_CUDA_TRY(cudaSetDevice(0));
  auto mr = rmm::mr::cuda_memory_resource{};
  rmm::device_vector vec(
    16, rmm::mr::thrust_allocator(cuda::stream_ref{cudaStream_t{cudaStreamDefault}}, mr));
  RMM_CUDA_TRY(cudaSetDevice(1));
  ...
  // No need to switch back to device 0 before ~vec runs
}

[!NOTE]

Although allocation and deallocation in the thrust_allocator run with the correct active device,

modification of rmm::device_vector might necessitate a kernel launch, and this must run with the

correct device active. For example, .resize() might both allocate _and_ launch a kernel to

initialize new elements: the user must arrange for this kernel launch to occur with the correct

device for the memory resource active.

cuda::stream_ref, cuda_stream_view, and cuda_stream

cuda::stream_ref, provided by ``, is the preferred non-owning CUDA stream wrapper. rmm::cuda_stream_view is deprecated; migrate existing code to cuda::stream_ref. The deprecated wrapper remains convertible to and from cuda::stream_ref for compatibility.

rmm::cuda_stream is a simple owning wrapper around a CUDA cudaStream_t. This class provides RAII semantics (constructor creates the CUDA stream, destructor destroys it). An rmm::cuda_stream can never represent the CUDA default stream or per-thread default stream; it only ever represents a single non-default stream. rmm::cuda_stream cannot be copied, but can be moved.

cuda_stream_pool

rmm::cuda_stream_pool provides fast access to a pool of CUDA streams. This class can be used to create a set of cuda_stream objects whose lifetime is equal to the cuda_stream_pool. Using the stream pool can be faster than creating the streams on the fly. The size of the pool is configurable. Depending on this size, multiple calls to cuda_stream_pool::get_stream() may return cuda::stream_ref instances that represent identical CUDA streams.

Thread Safety

All current device memory resources are thread safe unless documented otherwise. More specifically, calls to memory resource allocate() and deallocate() methods are safe with respect to calls to either of these functions from other threads. They are _not_ thread safe with respect to construction and destruction of the memory resource object.

Note that a class thread_safe_resource_adapter is provided which can be used to adapt a memory resource that is not thread safe to be thread safe (as described above). This adapter is not needed with any current RMM device memory resources.

Allocators

C++ interfaces commonly allow customizable memory allocation through an Allocator object. RMM provides several Allocator and Allocator-like classes.

polymorphic_allocator

A stream-ordered allocator similar to std::pmr::polymorphic_allocator. Unlike the standard C++ Allocator interface, the allocate and deallocate functions take a cuda::stream_ref indicating the stream on which the (de)allocation occurs.

stream_allocator_adaptor

stream_allocator_adaptor can be used to adapt a stream-ordered allocator to present a standard Allocator interface to consumers that may not be designed to work with a stream-ordered interface.

Example:

rmm::cuda_stream stream;
rmm::mr::polymorphic_allocator stream_alloc;
// Constructs an adaptor that forwards all (de)allocations to stream_alloc on stream.
auto adapted = rmm::mr::stream_allocator_adaptor(stream_alloc, stream);
// Allocates storage for 100 ints using stream_alloc on stream
auto p = adapted.allocate(100);
...
// Deallocates using stream_alloc on stream
adapted.deallocate(p, 100);

thrust_allocator

thrust_allocator is a device memory allocator that uses the strongly typed thrust::device_ptr, making it usable with containers like thrust::device_vector.

See below for more information on using RMM with Thrust.

Device Data Structures

device_buffer

An untyped, uninitialized RAII class for stream ordered device memory allocation.

Example

cuda::stream_ref s{...};
// Allocates at least 100 bytes on stream s using the default resource
rmm::device_buffer b{100, s};
void* p = b.data();                   // Raw, untyped pointer to underlying device memory
kernel<<<..., s.get()>>>(b.data());   // b is only safe to use on s
rmm::mr::cuda_memory_resource mr;
// Allocates at least 100 bytes on stream s using the resource mr
rmm::device_buffer b2{100, s, mr};

device_uvector

A typed, uninitialized RAII class for allocation of a contiguous set of elements in device memory. Similar to a thrust::device_vector, but as an optimization, does not default initialize the contained elements. This optimization restricts the types T to trivially copyable types.

Example

cuda::stream_ref s{...};
// Allocates uninitialized storage for 100 int32_t elements on stream s using the
// default resource
rmm::device_uvector<int32_t> v(100, s);
// Initializes the elements to 0
thrust::uninitialized_fill(thrust::cuda::par.on(s.get()), v.begin(), v.end(), int32_t{0});
rmm::mr::cuda_memory_resource mr;
// Allocates uninitialized storage for 100 int32_t elements on stream s using the resource mr
rmm::device_uvector<int32_t> v2{100, s, mr};

device_scalar

A