cublas team verified workaround for large GEMM algo68 bug (#3100)

The fix in #3098 was incomplete, as
according to the cublas team there is no guarantee that `algo13` will be
returned along with `algo68` (the failing version for A matrix with >
2^31 elements). Also, the cublas team will patch this bug in version
`13.6.1` so we will automatically receive the fix by guarding only
against `13.6.0`.

Author: divyegala

cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp

@@ -10,7 +10,6 @@
 #include <raft/core/resource/cublaslt_handle.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/custom_resource.hpp>
-#include <raft/core/resource/device_properties.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cache.hpp>
 #include <raft/util/cuda_data_type.hpp>
@@ -19,7 +18,6 @@
 
 #include <cublasLt.h>
 
-#include <array>
 #include <type_traits>
 
 namespace raft {
@@ -102,21 +100,27 @@ struct matmul_key_hash {
 };
 
 /**
- * cuBLASLt 13.6 and later may select algorithm 68 once A's physical span reaches 2^31 elements.
- * That algorithm fails during execution for FP32, so select the next ranked heuristic instead.
+ * cuBLASLt 13.6.0, shipped with CUDA 13.3, may select algorithm 68 once A's physical span reaches
+ * 2^31 elements. That algorithm fails during execution for FP32.
  */
-inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args,
-                                           std::size_t version,
-                                           int device_major,
-                                           int device_minor) noexcept -> bool
+inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args, std::size_t version) noexcept
+  -> bool
 {
   constexpr uint64_t max_safe_span = (uint64_t{1} << 31) - 1;
   const auto a_columns             = args.trans_a ? args.m : args.k;
-  const bool is_affected_architecture =
-    (device_major == 10 && device_minor == 0) ||
-    (device_major == 12 && (device_minor == 0 || device_minor == 1));
-  return version >= 130600 && is_affected_architecture && args.lda != 0 &&
-         a_columns > max_safe_span / args.lda;
+  return version == 130600 && args.lda != 0 && a_columns > max_safe_span / args.lda;
+}
+
+/**
+ * Querying with a physical A leading dimension that is not 16-byte aligned suppresses algorithm 68.
+ * The returned algorithm is then used with the real descriptors.
+ */
+inline auto get_cublaslt_13_6_heuristic_args(const matmul_key_t& args) noexcept -> matmul_key_t
+{
+  constexpr uint64_t fp32_elements_per_16_bytes = 4;
+  auto heuristic_args                           = args;
+  if (heuristic_args.lda % fp32_elements_per_16_bytes == 0) { ++heuristic_args.lda; }
+  return heuristic_args;
 }
 
 inline auto get_cublaslt_algorithm_id(const cublasLtMatmulHeuristicResult_t& heuristic) -> int
@@ -199,6 +203,24 @@ struct cublastlt_matmul_desc {
   }
 };
 
+/** Preference descriptor for a cublasLt matmul heuristic query. */
+struct cublastlt_matmul_preference {
+  cublasLtMatmulPreference_t res{nullptr};
+
+  inline cublastlt_matmul_preference() { RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceCreate(&res)); }
+  inline cublastlt_matmul_preference(const cublastlt_matmul_preference&) = delete;
+  inline auto operator=(const cublastlt_matmul_preference&)
+    -> cublastlt_matmul_preference& = delete;
+
+  inline ~cublastlt_matmul_preference() noexcept
+  {
+    RAFT_CUBLAS_TRY_NO_THROW(cublasLtMatmulPreferenceDestroy(res));
+  }
+
+  // NOLINTNEXTLINE
+  inline operator cublasLtMatmulPreference_t() const noexcept { return res; }
+};
+
 /** Full description of matmul. */
 struct matmul_desc {
   cublastlt_matmul_desc desc;
@@ -219,46 +241,40 @@ struct matmul_desc {
     bool use_cublaslt_13_6_workaround = false;
     if constexpr (std::is_same_v<S, float> && std::is_same_v<A, float> &&
                   std::is_same_v<B, float> && std::is_same_v<C, float>) {
-      const auto& device_properties = resource::get_device_properties(res);
-      use_cublaslt_13_6_workaround  = needs_cublaslt_13_6_workaround(
-        args, cublasLtGetVersion(), device_properties.major, device_properties.minor);
+      use_cublaslt_13_6_workaround = needs_cublaslt_13_6_workaround(args, cublasLtGetVersion());
     }
 
-    constexpr int workaround_heuristic_results = 2;
-    std::array<cublasLtMatmulHeuristicResult_t, workaround_heuristic_results> heuristic_results{};
-    const int requested_results = use_cublaslt_13_6_workaround ? workaround_heuristic_results : 1;
     int algo_count;
-    cublasLtMatmulPreference_t preference;
-    RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceCreate(&preference));
-    RAFT_CUBLAS_TRY(cublasLtMatmulAlgoGetHeuristic(resource::get_cublaslt_handle(res),
-                                                   r.desc,
-                                                   r.a,
-                                                   r.b,
-                                                   r.c,
-                                                   r.c,
-                                                   preference,
-                                                   requested_results,
-                                                   heuristic_results.data(),
-                                                   &algo_count));
-    RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceDestroy(preference));
-
-    RAFT_EXPECTS(algo_count > 0, "cuBLASLt did not return a matmul algorithm");
-    if (!use_cublaslt_13_6_workaround) {
-      r.heuristics = heuristic_results.front();
-      return r;
+    cublastlt_matmul_preference preference;
+    const auto query_heuristic = [&](cublasLtMatrixLayout_t a_layout,
+                                     cublasLtMatrixLayout_t c_layout) {
+      RAFT_CUBLAS_TRY(cublasLtMatmulAlgoGetHeuristic(resource::get_cublaslt_handle(res),
+                                                     r.desc,
+                                                     a_layout,
+                                                     r.b,
+                                                     c_layout,
+                                                     c_layout,
+                                                     preference,
+                                                     1,
+                                                     &r.heuristics,
+                                                     &algo_count));
+    };
+
+    if (use_cublaslt_13_6_workaround) {
+      const auto heuristic_args = get_cublaslt_13_6_heuristic_args(args);
+      const auto heuristic_a    = cublastlt_matrix_layout::for_matmul<A>(
+        !(heuristic_args.trans_a), heuristic_args.m, heuristic_args.k, heuristic_args.lda);
+      query_heuristic(heuristic_a, r.c);
+    } else {
+      query_heuristic(r.a, r.c);
     }
 
+    RAFT_EXPECTS(algo_count > 0, "cuBLASLt did not return a matmul algorithm");
     constexpr int faulty_algorithm = 68;
-    for (int i = 0; i < algo_count; ++i) {
-      const auto& candidate = heuristic_results[i];
-      if (candidate.state == CUBLAS_STATUS_SUCCESS && candidate.workspaceSize == 0 &&
-          get_cublaslt_algorithm_id(candidate) != faulty_algorithm) {
-        r.heuristics = candidate;
-        return r;
-      }
+    if (use_cublaslt_13_6_workaround) {
+      RAFT_EXPECTS(get_cublaslt_algorithm_id(r.heuristics) != faulty_algorithm,
+                   "cuBLASLt 13.6.0 returned faulty algorithm 68 for the workaround query");
     }
-
-    RAFT_FAIL("cuBLASLt 13.6 did not return a safe algorithm for the affected large FP32 GEMM");
     return r;
   }
 };

cpp/tests/linalg/gemm_basic.cpp

@@ -166,26 +166,20 @@ TEST(Raft, GemmPointerModeDeviceDefaults) { test_gemm_pointer_mode_device(false,
 TEST(Raft, GemmCublasLt136WorkaroundPredicate)
 {
   constexpr std::size_t affected_version = 130600;
-  constexpr int affected_device_major    = 12;
-  constexpr int affected_device_minor    = 1;
   const detail::matmul_key_t below_boundary{134217727, 1, 2, 16, 1, 134217727, true, true};
   const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
   const detail::matmul_key_t above_boundary{134217729, 1, 2, 16, 1, 134217729, true, true};
 
   const auto needs_workaround = [&](const auto& args) {
-    return detail::needs_cublaslt_13_6_workaround(
-      args, affected_version, affected_device_major, affected_device_minor);
+    return detail::needs_cublaslt_13_6_workaround(args, affected_version);
   };
 
   EXPECT_FALSE(needs_workaround(below_boundary));
   EXPECT_TRUE(needs_workaround(at_boundary));
   EXPECT_TRUE(needs_workaround(above_boundary));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(
-    at_boundary, 130599, affected_device_major, affected_device_minor));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(
-    at_boundary, 130601, affected_device_major, affected_device_minor));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(
-    at_boundary, 130700, affected_device_major, affected_device_minor));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130599));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130601));
+  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130700));
 
   auto different_output    = at_boundary;
   different_output.trans_b = false;
@@ -204,18 +198,26 @@ TEST(Raft, GemmCublasLt136WorkaroundPredicate)
   EXPECT_FALSE(needs_workaround(invalid_lda));
 }
 
-TEST(Raft, GemmCublasLt136WorkaroundArchitectures)
+TEST(Raft, GemmCublasLt136WorkaroundHeuristicArgs)
 {
-  constexpr std::size_t affected_version = 130600;
-  const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
+  const auto query_lda = [](uint64_t lda) {
+    const detail::matmul_key_t args{134217728, 1, 2, lda, 1, 134217728, true, true};
+    return detail::get_cublaslt_13_6_heuristic_args(args).lda;
+  };
 
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 10, 0));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 0));
-  EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 1));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 7, 5));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 0));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 9));
-  EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 9, 0));
+  EXPECT_EQ(query_lda(12), 13);
+  EXPECT_EQ(query_lda(15), 15);
+  EXPECT_EQ(query_lda(16), 17);
+
+  const detail::matmul_key_t args{134217728, 1, 2, 16, 1, 134217728, true, true};
+  const auto heuristic_args = detail::get_cublaslt_13_6_heuristic_args(args);
+  EXPECT_EQ(heuristic_args.m, args.m);
+  EXPECT_EQ(heuristic_args.n, args.n);
+  EXPECT_EQ(heuristic_args.k, args.k);
+  EXPECT_EQ(heuristic_args.ldb, args.ldb);
+  EXPECT_EQ(heuristic_args.ldc, args.ldc);
+  EXPECT_EQ(heuristic_args.trans_a, args.trans_a);
+  EXPECT_EQ(heuristic_args.trans_b, args.trans_b);
 }
 
 }  // namespace raft::linalg