Check nvidia smi version before loading decomp API. (#11464)

Author: trivialfis

src/common/cuda_dr_utils.cc

@@ -18,7 +18,7 @@
 #include "xgboost/string_view.h"  // for StringView
 
 namespace xgboost::cudr {
-CuDriverApi::CuDriverApi() {
+CuDriverApi::CuDriverApi(std::int32_t cu_major, std::int32_t cu_minor, std::int32_t kdm_major) {
   // similar to dlopen, but without the need to release a handle.
   auto safe_load = [](xgboost::StringView name, auto **fnptr) {
     cudaDriverEntryPointQueryResult status;
@@ -41,7 +41,12 @@ CuDriverApi::CuDriverApi() {
   safe_load("cuDeviceGetAttribute", &this->cuDeviceGetAttribute);
   safe_load("cuDeviceGet", &this->cuDeviceGet);
 #if defined(CUDA_HW_DECOM_AVAILABLE)
-  safe_load("cuMemBatchDecompressAsync", &this->cuMemBatchDecompressAsync);
+  // CTK 12.8
+  if (((cu_major == 12 && cu_minor >= 8) || cu_major > 12) && (kdm_major >= 570)) {
+    safe_load("cuMemBatchDecompressAsync", &this->cuMemBatchDecompressAsync);
+  } else {
+    this->cuMemBatchDecompressAsync = nullptr;
+  }
 #endif  // defined(CUDA_HW_DECOM_AVAILABLE)
   CHECK(this->cuMemGetAllocationGranularity);
 }
@@ -76,9 +81,17 @@ void CuDriverApi::ThrowIfError(CUresult status, StringView fn, std::int32_t line
 }
 
 [[nodiscard]] CuDriverApi &GetGlobalCuDriverApi() {
+  std::int32_t cu_major = -1, cu_minor = -1;
+  GetDrVersionGlobal(&cu_major, &cu_minor);
+
+  std::int32_t kdm_major = -1, kdm_minor = -1;
+  if (!GetVersionFromSmiGlobal(&kdm_major, &kdm_minor)) {
+    kdm_major = -1;
+  }
+
   static std::once_flag flag;
   static std::unique_ptr<CuDriverApi> cu;
-  std::call_once(flag, [&] { cu = std::make_unique<CuDriverApi>(); });
+  std::call_once(flag, [&] { cu = std::make_unique<CuDriverApi>(cu_major, cu_minor, kdm_major); });
   return *cu;
 }
 
@@ -154,5 +167,24 @@ void MakeCuMemLocation(CUmemLocationType type, CUmemLocation *loc) {
 
   return Invalid();
 }
+
+[[nodiscard]] bool GetVersionFromSmiGlobal(std::int32_t *p_major, std::int32_t *p_minor) {
+  static std::once_flag flag;
+  static std::int32_t major = -1, minor = -1;
+  static bool result = false;
+  std::call_once(flag, [&] { result = GetVersionFromSmi(&major, &minor); });
+
+  *p_major = major;
+  *p_minor = minor;
+  return result;
+}
+
+void GetDrVersionGlobal(std::int32_t *p_major, std::int32_t *p_minor) {
+  static std::once_flag once;
+  static std::int32_t major{0}, minor{0};
+  std::call_once(once, [] { xgboost::curt::DrVersion(&major, &minor); });
+  *p_major = major;
+  *p_minor = minor;
+}
 }  // namespace xgboost::cudr
 #endif

src/common/cuda_dr_utils.h

@@ -82,7 +82,7 @@ struct CuDriverApi {
 
 #endif  // defined(CUDA_HW_DECOM_AVAILABLE)
 
-  CuDriverApi();
+  CuDriverApi(std::int32_t cu_major, std::int32_t cu_minor, std::int32_t kdm_major);
 
   void ThrowIfError(CUresult status, StringView fn, std::int32_t line, char const *file) const;
 };
@@ -124,4 +124,14 @@ void MakeCuMemLocation(CUmemLocationType type, CUmemLocation *loc);
  * @return Whether the system call is successful.
  */
 [[nodiscard]] bool GetVersionFromSmi(std::int32_t *p_major, std::int32_t *p_minor);
+
+/**
+ * @brief Cache the result from @ref GetVersionFromSmi in a global variable
+ */
+[[nodiscard]] bool GetVersionFromSmiGlobal(std::int32_t *p_major, std::int32_t *p_minor);
+
+/**
+ * @brief Cache the result from @ref DrVersion in a global variable
+ */
+void GetDrVersionGlobal(std::int32_t *p_major, std::int32_t *p_minor);
 }  // namespace xgboost::cudr

src/common/device_compression.cu

@@ -242,7 +242,8 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
     CHECK(out.empty());
     return;
   }
-  if (GetGlobalDeStatus().avail) {
+  if (GetGlobalDeStatus().avail &&
+      cudr::GetGlobalCuDriverApi().cuMemBatchDecompressAsync != nullptr) {
     // Invoke the DE.
 #if defined(CUDA_HW_DECOM_AVAILABLE)
     std::size_t error_index;

src/common/device_helpers.cu

@@ -1,10 +1,7 @@
 /**
  * Copyright 2024-2025, XGBoost contributors
  */
-#include <mutex>  // for once_flag, call_once
-
 #include "../common/cuda_dr_utils.h"  // for GetVersionFromSmi
-#include "cuda_rt_utils.h"            // for RtVersion
 #include "device_helpers.cuh"
 #include "device_vector.cuh"  // for GrowOnlyVirtualMemVec
 #include "xgboost/windefs.h"  // for xgboost_IS_WIN
@@ -18,25 +15,22 @@ namespace {
 // Check whether cuda virtual memory can be used.
 // Host NUMA allocation requires driver that supports CTK >= 12.5 to be stable
 [[nodiscard]] bool CheckVmAlloc() {
-  static bool vm_flag = true;
-  static std::once_flag once;
+  std::int32_t major{0}, minor{0};
+  xgboost::cudr::GetDrVersionGlobal(&major, &minor);
 
-  std::call_once(once, [] {
-    std::int32_t major{0}, minor{0};
-    xgboost::curt::DrVersion(&major, &minor);
-    if (IsSupportedDrVer(major, minor)) {
-      // The result from the driver api is not reliable. The system driver might not match
-      // the CUDA driver in some obscure cases.
-      //
-      // https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
-      // Ver                 Linux       Win
-      // CUDA 12.5 Update 1  >=555.42.06 >=555.85
-      // CUDA 12.5 GA        >=555.42.02 >=555.85
-      vm_flag = xgboost::cudr::GetVersionFromSmi(&major, &minor) && major >= 555;
-    } else {
-      vm_flag = false;
-    }
-  });
+  bool vm_flag = true;
+  if (IsSupportedDrVer(major, minor)) {
+    // The result from the driver api is not reliable. The system driver might not match
+    // the CUDA driver in some obscure cases.
+    //
+    // https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
+    // Ver                 Linux       Win
+    // CUDA 12.5 Update 1  >=555.42.06 >=555.85
+    // CUDA 12.5 GA        >=555.42.02 >=555.85
+    vm_flag = xgboost::cudr::GetVersionFromSmiGlobal(&major, &minor) && major >= 555;
+  } else {
+    vm_flag = false;
+  }
   return vm_flag;
 }
 }  // namespace