diff --git a/build.sh b/build.sh index 0b97542b3e..f111ef6a5c 100755 --- a/build.sh +++ b/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # raft build scripts @@ -83,6 +83,7 @@ EXT_HEADERS_TEST_COMPILED_IMPLICIT;\ EXT_HEADERS_TEST_IMPLICIT;\ LABEL_TEST;\ LINALG_TEST;\ +GEMM_LARGE_TEST;\ MATRIX_SELECT_LARGE_TEST;\ MATRIX_SELECT_TEST;\ MATRIX_TEST;\ diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index 2337413fbd..28eb242323 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -18,6 +19,7 @@ #include +#include #include namespace raft { @@ -99,6 +101,33 @@ 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. + */ +inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args, + std::size_t version, + int device_major, + int device_minor) 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; +} + +inline auto get_cublaslt_algorithm_id(const cublasLtMatmulHeuristicResult_t& heuristic) -> int +{ + int algorithm_id{}; + std::size_t bytes_written{}; + RAFT_CUBLAS_TRY(cublasLtMatmulAlgoConfigGetAttribute( + &heuristic.algo, CUBLASLT_ALGO_CONFIG_ID, &algorithm_id, sizeof(algorithm_id), &bytes_written)); + return algorithm_id; +} + /** Descriptor for a column-major cublasLt matrix. */ struct cublastlt_matrix_layout { cublasLtMatrixLayout_t res{nullptr}; @@ -186,6 +215,18 @@ struct matmul_desc { cublastlt_matrix_layout::for_matmul(!(args.trans_a), args.m, args.k, args.lda), cublastlt_matrix_layout::for_matmul(!(args.trans_b), args.k, args.n, args.ldb), cublastlt_matrix_layout::for_matmul(true, args.m, args.n, args.ldc)}; + + 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); + } + + 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)); @@ -196,10 +237,28 @@ struct matmul_desc { r.c, r.c, preference, - 1, - &r.heuristics, + 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; + } + + 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; + } + } + + RAFT_FAIL("cuBLASLt 13.6 did not return a safe algorithm for the affected large FP32 GEMM"); return r; } }; diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 7127312b7b..0825073fd9 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -185,6 +185,8 @@ if(BUILD_TESTS) 100 ) + ConfigureTest(NAME GEMM_LARGE_TEST PATH linalg/gemm_large.cpp GPUS 1 PERCENT 100) + ConfigureTest( NAME MATRIX_TEST diff --git a/cpp/tests/linalg/gemm_basic.cpp b/cpp/tests/linalg/gemm_basic.cpp index 6ab669c1d0..6118458e93 100644 --- a/cpp/tests/linalg/gemm_basic.cpp +++ b/cpp/tests/linalg/gemm_basic.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -163,4 +163,59 @@ TEST(Raft, GemmPointerModeDeviceAlpha) { test_gemm_pointer_mode_device(true, fal TEST(Raft, GemmPointerModeDeviceBeta) { test_gemm_pointer_mode_device(false, true); } TEST(Raft, GemmPointerModeDeviceDefaults) { test_gemm_pointer_mode_device(false, 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); + }; + + 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)); + + auto different_output = at_boundary; + different_output.trans_b = false; + different_output.n = 2; + different_output.ldb = 7; + different_output.ldc = 11; + EXPECT_TRUE(needs_workaround(different_output)); + + const detail::matmul_key_t non_transposed_below{2, 1, 134217727, 16, 134217727, 2, false, false}; + const detail::matmul_key_t non_transposed_at{2, 1, 134217728, 16, 134217728, 2, false, false}; + EXPECT_FALSE(needs_workaround(non_transposed_below)); + EXPECT_TRUE(needs_workaround(non_transposed_at)); + + auto invalid_lda = at_boundary; + invalid_lda.lda = 0; + EXPECT_FALSE(needs_workaround(invalid_lda)); +} + +TEST(Raft, GemmCublasLt136WorkaroundArchitectures) +{ + constexpr std::size_t affected_version = 130600; + const detail::matmul_key_t at_boundary{134217728, 1, 2, 16, 1, 134217728, true, true}; + + 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)); +} + } // namespace raft::linalg diff --git a/cpp/tests/linalg/gemm_large.cpp b/cpp/tests/linalg/gemm_large.cpp new file mode 100644 index 0000000000..9b75421fef --- /dev/null +++ b/cpp/tests/linalg/gemm_large.cpp @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include + +#include +#include + +namespace raft::linalg { + +TEST(Raft, GemmLargeSpan) +{ + constexpr std::int64_t lda = 16; + constexpr std::int64_t informative = 2; + constexpr std::int64_t targets = 1; + constexpr std::array sample_counts{134217727, 134217728, 134217729}; + constexpr auto max_samples = sample_counts.back(); + constexpr auto a_elements = max_samples * lda; + + raft::resources resources; + const auto stream = raft::resource::get_cuda_stream(resources); + + auto a = raft::make_device_vector(resources, a_elements); + auto b = raft::make_device_vector(resources, informative); + auto c = raft::make_device_vector(resources, max_samples); + + RAFT_CUDA_TRY(cudaMemsetAsync(a.data_handle(), 0, a.size() * sizeof(float), stream.value())); + RAFT_CUDA_TRY(cudaMemsetAsync(b.data_handle(), 0, b.size() * sizeof(float), stream.value())); + + const float alpha = 1.0f; + const float beta = 0.0f; + for (const auto samples : sample_counts) { + SCOPED_TRACE(samples); + RAFT_CUDA_TRY(cudaMemsetAsync(c.data_handle(), 0xff, samples * sizeof(float), stream.value())); + raft::linalg::gemm(resources, + true, + true, + static_cast(samples), + targets, + informative, + &alpha, + a.data_handle(), + lda, + b.data_handle(), + targets, + &beta, + c.data_handle(), + static_cast(samples), + stream.value()); + + std::array output_samples{}; + RAFT_CUDA_TRY(cudaMemcpyAsync( + &output_samples[0], c.data_handle(), sizeof(float), cudaMemcpyDeviceToHost, stream.value())); + RAFT_CUDA_TRY(cudaMemcpyAsync(&output_samples[1], + c.data_handle() + samples / 2, + sizeof(float), + cudaMemcpyDeviceToHost, + stream.value())); + RAFT_CUDA_TRY(cudaMemcpyAsync(&output_samples[2], + c.data_handle() + samples - 1, + sizeof(float), + cudaMemcpyDeviceToHost, + stream.value())); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream.value())); + EXPECT_FLOAT_EQ(output_samples[0], 0.0f); + EXPECT_FLOAT_EQ(output_samples[1], 0.0f); + EXPECT_FLOAT_EQ(output_samples[2], 0.0f); + } +} + +} // namespace raft::linalg