Skip to content

Forward-merge release/26.08 into main #3101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GPUtester merged 1 commit into from
Jul 30, 2026
Merged

Forward-merge release/26.08 into main #3101

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 62 additions & 46 deletions cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -19,7 +18,6 @@

#include <cublasLt.h>

#include <array>
#include <type_traits>

namespace raft {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand All @@ -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;
}
};
Expand Down
42 changes: 22 additions & 20 deletions cpp/tests/linalg/gemm_basic.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Loading