rapidsai/rmm

NVIDIA RMM (RAPIDS Memory Manager)

개요 – RMM은 GPU 메모리 할당을 위한 공통적이고 유연한 인터페이스를 제공하는 C++/Python 라이브러리입니다. 개발자는 다양한 할당 전략(plain cudaMalloc, 풀 할당자, 관리형 메모리 등)을 플러그인하여 C++ 코드와 Python의 rmm 패키지를 통해 사용할 수 있습니다.

중요성 – 고성능 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)를 사용하여 풀 할당자로 전환할 수 있습니다.

설치 방법 – 사전 빌드된 wheel을 설치하거나(CUDA 13의 경우 pip install librmm-cu13 rmm-cu13 등), CMake와 제공된 build.sh를 사용하여 소스에서 빌드합니다. 저장소에는 전체 문서, 튜토리얼 및 CMake 통합 도우미가 포함되어 있습니다.

사용처 – NVIDIA RAPIDS 생태계의 일부로서, cuDF, cuML 및 기타 GPU 가속 데이터 과학 라이브러리와 같은 하위 프로젝트들이 효율적인 메모리 처리를 위해 RMM에 의존하고 있습니다.