Skip to content

Commit 52a9fbf

Browse files
GPUtesterGPUtester
authored
Merge pull request #3101 from NVIDIA/release/26.08
Forward-merge release/26.08 into main
2 parents edc3997 + 97d22c8 commit 52a9fbf

2 files changed

Lines changed: 84 additions & 66 deletions

File tree

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

Lines changed: 62 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
#include <raft/core/resource/cublaslt_handle.hpp>
1111
#include <raft/core/resource/cuda_stream.hpp>
1212
#include <raft/core/resource/custom_resource.hpp>
13-
#include <raft/core/resource/device_properties.hpp>
1413
#include <raft/core/resources.hpp>
1514
#include <raft/util/cache.hpp>
1615
#include <raft/util/cuda_data_type.hpp>
@@ -19,7 +18,6 @@
1918

2019
#include <cublasLt.h>
2120

22-
#include <array>
2321
#include <type_traits>
2422
#include <utility>
2523

@@ -103,21 +101,27 @@ struct matmul_key_hash {
103101
};
104102

105103
/**
106-
* cuBLASLt 13.6 and later may select algorithm 68 once A's physical span reaches 2^31 elements.
107-
* That algorithm fails during execution for FP32, so select the next ranked heuristic instead.
104+
* cuBLASLt 13.6.0, shipped with CUDA 13.3, may select algorithm 68 once A's physical span reaches
105+
* 2^31 elements. That algorithm fails during execution for FP32.
108106
*/
109-
inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args,
110-
std::size_t version,
111-
int device_major,
112-
int device_minor) noexcept -> bool
107+
inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args, std::size_t version) noexcept
108+
-> bool
113109
{
114110
constexpr uint64_t max_safe_span = (uint64_t{1} << 31) - 1;
115111
const auto a_columns = args.trans_a ? args.m : args.k;
116-
const bool is_affected_architecture =
117-
(device_major == 10 && device_minor == 0) ||
118-
(device_major == 12 && (device_minor == 0 || device_minor == 1));
119-
return version >= 130600 && is_affected_architecture && args.lda != 0 &&
120-
a_columns > max_safe_span / args.lda;
112+
return version == 130600 && args.lda != 0 && a_columns > max_safe_span / args.lda;
113+
}
114+
115+
/**
116+
* Querying with a physical A leading dimension that is not 16-byte aligned suppresses algorithm 68.
117+
* The returned algorithm is then used with the real descriptors.
118+
*/
119+
inline auto get_cublaslt_13_6_heuristic_args(const matmul_key_t& args) noexcept -> matmul_key_t
120+
{
121+
constexpr uint64_t fp32_elements_per_16_bytes = 4;
122+
auto heuristic_args = args;
123+
if (heuristic_args.lda % fp32_elements_per_16_bytes == 0) { ++heuristic_args.lda; }
124+
return heuristic_args;
121125
}
122126

123127
inline auto get_cublaslt_algorithm_id(const cublasLtMatmulHeuristicResult_t& heuristic) -> int
@@ -214,6 +218,24 @@ struct cublastlt_matmul_desc {
214218
}
215219
};
216220

221+
/** Preference descriptor for a cublasLt matmul heuristic query. */
222+
struct cublastlt_matmul_preference {
223+
cublasLtMatmulPreference_t res{nullptr};
224+
225+
inline cublastlt_matmul_preference() { RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceCreate(&res)); }
226+
inline cublastlt_matmul_preference(const cublastlt_matmul_preference&) = delete;
227+
inline auto operator=(const cublastlt_matmul_preference&)
228+
-> cublastlt_matmul_preference& = delete;
229+
230+
inline ~cublastlt_matmul_preference() noexcept
231+
{
232+
RAFT_CUBLAS_TRY_NO_THROW(cublasLtMatmulPreferenceDestroy(res));
233+
}
234+
235+
// NOLINTNEXTLINE
236+
inline operator cublasLtMatmulPreference_t() const noexcept { return res; }
237+
};
238+
217239
/** Full description of matmul. */
218240
struct matmul_desc {
219241
cublastlt_matmul_desc desc;
@@ -234,46 +256,40 @@ struct matmul_desc {
234256
bool use_cublaslt_13_6_workaround = false;
235257
if constexpr (std::is_same_v<S, float> && std::is_same_v<A, float> &&
236258
std::is_same_v<B, float> && std::is_same_v<C, float>) {
237-
const auto& device_properties = resource::get_device_properties(res);
238-
use_cublaslt_13_6_workaround = needs_cublaslt_13_6_workaround(
239-
args, cublasLtGetVersion(), device_properties.major, device_properties.minor);
259+
use_cublaslt_13_6_workaround = needs_cublaslt_13_6_workaround(args, cublasLtGetVersion());
240260
}
241261

242-
constexpr int workaround_heuristic_results = 2;
243-
std::array<cublasLtMatmulHeuristicResult_t, workaround_heuristic_results> heuristic_results{};
244-
const int requested_results = use_cublaslt_13_6_workaround ? workaround_heuristic_results : 1;
245262
int algo_count;
246-
cublasLtMatmulPreference_t preference;
247-
RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceCreate(&preference));
248-
RAFT_CUBLAS_TRY(cublasLtMatmulAlgoGetHeuristic(resource::get_cublaslt_handle(res),
249-
r.desc,
250-
r.a,
251-
r.b,
252-
r.c,
253-
r.c,
254-
preference,
255-
requested_results,
256-
heuristic_results.data(),
257-
&algo_count));
258-
RAFT_CUBLAS_TRY(cublasLtMatmulPreferenceDestroy(preference));
259-
260-
RAFT_EXPECTS(algo_count > 0, "cuBLASLt did not return a matmul algorithm");
261-
if (!use_cublaslt_13_6_workaround) {
262-
r.heuristics = heuristic_results.front();
263-
return r;
263+
cublastlt_matmul_preference preference;
264+
const auto query_heuristic = [&](cublasLtMatrixLayout_t a_layout,
265+
cublasLtMatrixLayout_t c_layout) {
266+
RAFT_CUBLAS_TRY(cublasLtMatmulAlgoGetHeuristic(resource::get_cublaslt_handle(res),
267+
r.desc,
268+
a_layout,
269+
r.b,
270+
c_layout,
271+
c_layout,
272+
preference,
273+
1,
274+
&r.heuristics,
275+
&algo_count));
276+
};
277+
278+
if (use_cublaslt_13_6_workaround) {
279+
const auto heuristic_args = get_cublaslt_13_6_heuristic_args(args);
280+
const auto heuristic_a = cublastlt_matrix_layout::for_matmul<A>(
281+
!(heuristic_args.trans_a), heuristic_args.m, heuristic_args.k, heuristic_args.lda);
282+
query_heuristic(heuristic_a, r.c);
283+
} else {
284+
query_heuristic(r.a, r.c);
264285
}
265286

287+
RAFT_EXPECTS(algo_count > 0, "cuBLASLt did not return a matmul algorithm");
266288
constexpr int faulty_algorithm = 68;
267-
for (int i = 0; i < algo_count; ++i) {
268-
const auto& candidate = heuristic_results[i];
269-
if (candidate.state == CUBLAS_STATUS_SUCCESS && candidate.workspaceSize == 0 &&
270-
get_cublaslt_algorithm_id(candidate) != faulty_algorithm) {
271-
r.heuristics = candidate;
272-
return r;
273-
}
289+
if (use_cublaslt_13_6_workaround) {
290+
RAFT_EXPECTS(get_cublaslt_algorithm_id(r.heuristics) != faulty_algorithm,
291+
"cuBLASLt 13.6.0 returned faulty algorithm 68 for the workaround query");
274292
}
275-
276-
RAFT_FAIL("cuBLASLt 13.6 did not return a safe algorithm for the affected large FP32 GEMM");
277293
return r;
278294
}
279295
};

‎cpp/tests/linalg/gemm_basic.cpp‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -166,26 +166,20 @@ TEST(Raft, GemmPointerModeDeviceDefaults) { test_gemm_pointer_mode_device(false,
166166
TEST(Raft, GemmCublasLt136WorkaroundPredicate)
167167
{
168168
constexpr std::size_t affected_version = 130600;
169-
constexpr int affected_device_major = 12;
170-
constexpr int affected_device_minor = 1;
171169
const detail::matmul_key_t below_boundary{134217727, 1, 2, 16, 1, 134217727, true, true};
172170
const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
173171
const detail::matmul_key_t above_boundary{134217729, 1, 2, 16, 1, 134217729, true, true};
174172

175173
const auto needs_workaround = [&](const auto& args) {
176-
return detail::needs_cublaslt_13_6_workaround(
177-
args, affected_version, affected_device_major, affected_device_minor);
174+
return detail::needs_cublaslt_13_6_workaround(args, affected_version);
178175
};
179176

180177
EXPECT_FALSE(needs_workaround(below_boundary));
181178
EXPECT_TRUE(needs_workaround(at_boundary));
182179
EXPECT_TRUE(needs_workaround(above_boundary));
183-
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(
184-
at_boundary, 130599, affected_device_major, affected_device_minor));
185-
EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(
186-
at_boundary, 130601, affected_device_major, affected_device_minor));
187-
EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(
188-
at_boundary, 130700, affected_device_major, affected_device_minor));
180+
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130599));
181+
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130601));
182+
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, 130700));
189183

190184
auto different_output = at_boundary;
191185
different_output.trans_b = false;
@@ -204,18 +198,26 @@ TEST(Raft, GemmCublasLt136WorkaroundPredicate)
204198
EXPECT_FALSE(needs_workaround(invalid_lda));
205199
}
206200

207-
TEST(Raft, GemmCublasLt136WorkaroundArchitectures)
201+
TEST(Raft, GemmCublasLt136WorkaroundHeuristicArgs)
208202
{
209-
constexpr std::size_t affected_version = 130600;
210-
const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true};
203+
const auto query_lda = [](uint64_t lda) {
204+
const detail::matmul_key_t args{134217728, 1, 2, lda, 1, 134217728, true, true};
205+
return detail::get_cublaslt_13_6_heuristic_args(args).lda;
206+
};
211207

212-
EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 10, 0));
213-
EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 0));
214-
EXPECT_TRUE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 12, 1));
215-
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 7, 5));
216-
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 0));
217-
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 8, 9));
218-
EXPECT_FALSE(detail::needs_cublaslt_13_6_workaround(at_boundary, affected_version, 9, 0));
208+
EXPECT_EQ(query_lda(12), 13);
209+
EXPECT_EQ(query_lda(15), 15);
210+
EXPECT_EQ(query_lda(16), 17);
211+
212+
const detail::matmul_key_t args{134217728, 1, 2, 16, 1, 134217728, true, true};
213+
const auto heuristic_args = detail::get_cublaslt_13_6_heuristic_args(args);
214+
EXPECT_EQ(heuristic_args.m, args.m);
215+
EXPECT_EQ(heuristic_args.n, args.n);
216+
EXPECT_EQ(heuristic_args.k, args.k);
217+
EXPECT_EQ(heuristic_args.ldb, args.ldb);
218+
EXPECT_EQ(heuristic_args.ldc, args.ldc);
219+
EXPECT_EQ(heuristic_args.trans_a, args.trans_a);
220+
EXPECT_EQ(heuristic_args.trans_b, args.trans_b);
219221
}
220222

221223
} // namespace raft::linalg

0 commit comments

Comments
 (0)