diff --git a/CMakeLists.txt b/CMakeLists.txt
index bd21e9b67fb2..966f0921c0de 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -112,6 +112,9 @@ endif (ENABLE_ALL_WARNINGS)
 if (BUILD_STATIC_LIB AND (R_LIB OR JVM_BINDINGS))
   message(SEND_ERROR "Cannot build a static library libxgboost.a when R or JVM packages are enabled.")
 endif (BUILD_STATIC_LIB AND (R_LIB OR JVM_BINDINGS))
+if (PLUGIN_RMM AND (NOT BUILD_WITH_CUDA_CUB))
+  message(SEND_ERROR "Cannot build with RMM using cub submodule.")
+endif (PLUGIN_RMM AND (NOT BUILD_WITH_CUDA_CUB))
 
 #-- Sanitizer
 if (USE_SANITIZER)
diff --git a/Jenkinsfile b/Jenkinsfile
index a038d27264dd..7ed17a966577 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -249,7 +249,7 @@ def BuildCUDA(args) {
       docker_args = "--build-arg CUDA_VERSION_ARG=${args.cuda_version}"
       sh """
       rm -rf build/
-      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh --conda-env=gpu_test -DUSE_CUDA=ON -DUSE_NCCL=ON -DPLUGIN_RMM=ON ${arch_flag}
+      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh --conda-env=gpu_test -DUSE_CUDA=ON -DUSE_NCCL=ON -DPLUGIN_RMM=ON -DBUILD_WITH_CUDA_CUB=ON ${arch_flag}
       ${dockerRun} ${container_type} ${docker_binary} ${docker_args} bash -c "cd python-package && rm -rf dist/* && python setup.py bdist_wheel --universal"
       ${dockerRun} ${container_type} ${docker_binary} ${docker_args} python tests/ci_build/rename_whl.py python-package/dist/*.whl ${commit_id} manylinux2014_x86_64
       """
diff --git a/python-package/xgboost/data.py b/python-package/xgboost/data.py
index 1b7b966d22be..e4b5a690359a 100644
--- a/python-package/xgboost/data.py
+++ b/python-package/xgboost/data.py
@@ -435,7 +435,11 @@ def _cudf_array_interfaces(data) -> Tuple[list, bytes]:
     interface is finished.
 
     """
-    from cudf.utils.dtypes import is_categorical_dtype
+    try:
+        from cudf.api.types import is_categorical_dtype
+    except ImportError:
+        from cudf.utils.dtypes import is_categorical_dtype
+
     cat_codes = []
     interfaces = []
     if _is_cudf_ser(data):
@@ -461,7 +465,10 @@ def _transform_cudf_df(
     feature_types: Optional[List[str]],
     enable_categorical: bool,
 ):
-    from cudf.utils.dtypes import is_categorical_dtype
+    try:
+        from cudf.api.types import is_categorical_dtype
+    except ImportError:
+        from cudf.utils.dtypes import is_categorical_dtype
 
     if feature_names is None:
         if _is_cudf_ser(data):
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index c2c7ff07b463..61e0fd553659 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -708,6 +708,10 @@ constexpr std::pair<int, int> CUDAVersion() {
 #endif  // defined(__CUDACC_VER_MAJOR__)
 }
 
+constexpr std::pair<int32_t, int32_t> ThrustVersion() {
+  return std::make_pair(THRUST_MAJOR_VERSION, THRUST_MINOR_VERSION);
+}
+
 namespace detail {
 template <typename T>
 using TypedDiscardCTK114 = thrust::discard_iterator<T>;
@@ -721,9 +725,9 @@ class TypedDiscard : public thrust::discard_iterator<T> {
 
 template <typename T>
 using TypedDiscard =
-    std::conditional_t<((CUDAVersion().first == 11 &&
-                         CUDAVersion().second >= 4) ||
-                        CUDAVersion().first > 11),
+    std::conditional_t<((ThrustVersion().first == 1 &&
+                         ThrustVersion().second >= 12) ||
+                        ThrustVersion().first > 1),
                        detail::TypedDiscardCTK114<T>, detail::TypedDiscard<T>>;
 
 /**
diff --git a/tests/ci_build/Dockerfile.gpu b/tests/ci_build/Dockerfile.gpu
index b86e2133705b..72f7b7853c9c 100644
--- a/tests/ci_build/Dockerfile.gpu
+++ b/tests/ci_build/Dockerfile.gpu
@@ -19,7 +19,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 cudf=21.08* rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda dask-cudf cupy=9.1* \
+        python=3.8 cudf=21.10* rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=21.10* dask-cudf=21.10* cupy=9.1* \
         numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis
 
 ENV GOSU_VERSION 1.10
diff --git a/tests/ci_build/Dockerfile.rmm b/tests/ci_build/Dockerfile.rmm
index 65ec7266fbfa..3d4906d4a25a 100644
--- a/tests/ci_build/Dockerfile.rmm
+++ b/tests/ci_build/Dockerfile.rmm
@@ -29,7 +29,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with RMM
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG
+        python=3.8 rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG
 
 ENV GOSU_VERSION 1.10