Repository navigation
[REQ] Rapids Memory Manager (RMM) Support #781
Description
Description
https://github.com/rapidsai/rmm
Note that Warp does not support any external memory allocators, which makes this task less trivial.
Context
Improve interop with PyTorch
Activity
coreyjadams commentedon Jun 10, 2025
Thank you for opening this issue!
For posterity: If/when this feature is implemented, the physicsnemo package has several model pipelines that combine pytorch + rapids + warp that could take advantage of this support. In particular, the DoMINO datapipeline + model has warp kernels in both preprocessing and model implementation in pytorch. Preprocessing also uses rapids' cuml package - and combining all of these bumps up against memory limits since pytorch will allocate most of the memory even when it's not in use.
I'm happy to be a tester if/when you need one.
bdice commentedon Jun 10, 2025
Happy to help on the RMM side, reach out if I can assist with a review or design discussion!
coreyjadams commentedon Oct 14, 2025
Hi @shi-eric - just wanted to check, is this something Warp can support?
shi-eric commentedon Nov 13, 2025
We would need @nvlukasz to weigh in on the design. For now I'm pushing this work to the v1.12 release.
10 remaining items
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
coreyjadams commentedon Mar 19, 2026
This is awesome! I will take a look! @loliverhennigh @ktangsali this would probably be really impactful for your functional work and derivatives work in physicsnemo too.
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
shi-eric commentedon Mar 19, 2026
Great, thanks. Keep the feature requests coming!
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
Add pluggable allocator interface and RmmAllocator (NVIDIAGH-781)
shi-eric commentedon Apr 18, 2026
Interface added in b7ced88
Future work should support memory tracking like #1269 did for native allocators
Add pluggable allocator interface and RmmAllocator [NVIDIAGH-781]
Move RmmAllocator to wp.utils.AllocatorRmm [NVIDIAGH-781]
Move RmmAllocator to wp.utils.AllocatorRmm [NVIDIAGH-781]