Add opt-in RMM allocation for WholeMemory

[alexbarghi-nv]

Summary

• add C API controls and pylibwholegraph.torch.set_memory_resource() / is_rmm_enabled() to opt supported WholeMemory allocations into RMM
• allocate distributed and hierarchy device storage with CCCL's cuda::device_buffer<cuda::std::byte>, explicitly backed by RMM's current per-device resource
• preserve byte-based sizing, uninitialized storage, CUDA allocation alignment, and legacy default-stream behavior
• retain the allocation-time resource reference for matching deallocation; the Python API also retains supplied memory-resource objects for allocation lifetime safety
• keep chunked, continuous, and NVSHMEM device storage on their specialized allocation paths and emit a rank-0 fallback warning when RMM is enabled
• leave host allocations and the default disabled path unchanged, require matching enablement across ranks, and reset RMM enablement during successful WholeMemory finalization
• add RMM runtime dependencies to conda and Python package metadata, plus Python usage documentation

Test coverage

• verify the default/disabled path does not allocate through RMM
• verify distributed and hierarchy device allocations use the configured RMM resource and release all tracked bytes
• verify an allocation is deallocated through its original resource after RMM's current resource is changed
• verify chunked and continuous device allocations, plus distributed host allocations, remain outside RMM
• verify finalization resets RMM enablement
• verify the Python API installs and retains a valid resource, supports disabling with None, rejects invalid resource types, and reports finalization state correctly

Local testing

Tested locally in the rapids conda environment with CUDA 12.9, GCC 14, RMM 26.10, CCCL 3.4, and NVIDIA V100 GPUs:

• fresh CMake configure and CUDA build of WHOLEMEMORY_HANDLE_TEST (passed)
• complete WHOLEMEMORY_HANDLE_TEST on two peer-connected V100s: 20 tests passed
• focused WholeMemoryHandleRMMTests.* two-GPU run (passed)
• new Python memory-resource API tests: 3 tests passed
• existing Python WholeMemory binding regression: 1 test passed
• repository pre-commit hooks for all changed files, including clang-format, Ruff, RAPIDS dependency generation, CMake formatting/linting, copyright checks, and alpha-spec validation (passed)
• git diff --check (passed)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[alexbarghi-nv]

/ok to test 64ea3f3

[bdice]

Generally looks fine, but I have a comment about how this code could be a little more future-proof.

[alexbarghi-nv]

/ok to test 23bb681

[greptile-apps]

Greptile Summary

Adds opt-in RMM allocation support for WholeMemory.

• Routes supported distributed and hierarchy device allocations through the current per-device RMM resource.
• Adds C and Python controls, fallback behavior, lifecycle tests, dependencies, and usage documentation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
cpp/src/wholememory/memory_handle.cpp	Adds process-wide RMM enablement, RMM-backed device buffers, allocation-path consistency checks, and specialized-path fallbacks.
cpp/src/wholememory/wholememory.cpp	Exposes RMM enablement through the C API and resets it after successful finalization.
python/pylibwholegraph/pylibwholegraph/torch/initialize.py	Exposes boolean RMM controls without retaining obsolete Python memory-resource objects.
cpp/tests/wholememory/wholememory_handle_tests.cpp	Tests opt-in allocation routing, allocation-time resource selection, deallocation accounting, fallback paths, and finalization reset.

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into agent/enable-rm..." | Re-trigger Greptile}

[alexbarghi-nv]

/ok to test da656ce

[alexbarghi-nv]

/ok to test 0004b48

[bdice]

One CMake issue needs to be fixed before we can merge.

[bdice]

Please copy https://github.com/rapidsai/cudf/blob/main/cpp/cmake/thirdparty/get_rmm.cmake, don't use plain find_package here. This must go before get_raft.cmake so that RMM is found by wholegraph (potentially with custom branches/etc.) before RAFT attempts to find RMM.

[alexbarghi-nv]

done; can you re-review @bdice ?

[alexbarghi-nv]

/ok to test 217ff8a

[linhu-nv]

LGTM thanks!

[alexbarghi-nv]

/ok to test 5c9ee13

[bdice]

Nice! Thank you.

[alexbarghi-nv]

/ok to test 08f3ba7