From 97d22c81ad983de82417964466c61bf72a8e2bb1 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 30 Jul 2026 13:37:43 -0700 Subject: [PATCH] `cublas` team verified workaround for large GEMM `algo68` bug (#3100) The fix in https://github.com/NVIDIA/raft/pull/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`. --- .../raft/linalg/detail/cublaslt_wrappers.hpp | 108 ++++++++++-------- cpp/tests/linalg/gemm_basic.cpp | 42 +++---- 2 files changed, 84 insertions(+), 66 deletions(-) diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index 28eb242323..e3c596a7b5 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -10,7 +10,6 @@ #include #include #include -#include #include #include #include @@ -19,7 +18,6 @@ #include -#include #include 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 && std::is_same_v && std::is_same_v && std::is_same_v) { - 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 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( + !(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; } }; diff --git a/cpp/tests/linalg/gemm_basic.cpp b/cpp/tests/linalg/gemm_basic.cpp index 6118458e93..8622eec113 100644 --- a/cpp/tests/linalg/gemm_basic.cpp +++ b/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