Forward-merge release/26.08 into main

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;\

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 <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>
@@ -18,6 +19,7 @@
 
 #include <cublasLt.h>
 
+#include <array>
 #include <type_traits>
 
 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<A>(!(args.trans_a), args.m, args.k, args.lda),
       cublastlt_matrix_layout::for_matmul<B>(!(args.trans_b), args.k, args.n, args.ldb),
       cublastlt_matrix_layout::for_matmul<C>(true, args.m, args.n, args.ldc)};
+
+    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);
+    }
+
+    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));
@@ -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;
   }
 };

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

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

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 <raft/core/device_mdarray.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/linalg/gemm.cuh>
+
+#include <gtest/gtest.h>
+
+#include <array>
+#include <cstdint>
+
+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<std::int64_t, 3> 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<float, std::int64_t>(resources, a_elements);
+  auto b = raft::make_device_vector<float, std::int64_t>(resources, informative);
+  auto c = raft::make_device_vector<float, std::int64_t>(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<int>(samples),
+                       targets,
+                       informative,
+                       &alpha,
+                       a.data_handle(),
+                       lda,
+                       b.data_handle(),
+                       targets,
+                       &beta,
+                       c.data_handle(),
+                       static_cast<int>(samples),
+                       stream.value());
+
+    std::array<float, 3> 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