Update RMM allocators (#12861)

Authors:
  - Peter Andreas Entschev (https://github.com/pentschev)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)

URL: #12861

python/cudf/cudf/__init__.py

@@ -8,6 +8,8 @@
 from numba import config as numba_config, cuda
 
 import rmm
+from rmm.allocators.cupy import rmm_cupy_allocator
+from rmm.allocators.numba import RMMNumbaManager
 
 from cudf import api, core, datasets, testing
 from cudf.api.extensions import (
@@ -96,8 +98,8 @@
 
     del patch_numba_linker_if_needed
 
-cuda.set_memory_manager(rmm.RMMNumbaManager)
-cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
+cuda.set_memory_manager(RMMNumbaManager)
+cupy.cuda.set_allocator(rmm_cupy_allocator)
 
 try:
     # Numba 0.54: Disable low occupancy warnings