From 088de6130a5957c2b2571ba11e5ec7008673f856 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Mon, 10 Aug 2026 18:43:42 +0000 Subject: [PATCH 1/8] gemm_strided_batched --- .../raft/linalg/detail/cublaslt_wrappers.hpp | 88 ++++++++++++++++++- cpp/include/raft/linalg/gemm.cuh | 76 ++++++++++++++++ cpp/tests/linalg/gemm_basic.cpp | 49 +++++++++++ 3 files changed, 211 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index 8cf228f2ed..e509e31822 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -167,6 +167,23 @@ struct cublastlt_matrix_layout { return cublastlt_matrix_layout{ get_cuda_data_type(), col_major ? rows : cols, col_major ? cols : rows, ld}; } + + template + static inline auto for_strided_batched_matmul(bool col_major, + uint64_t rows, + uint64_t cols, + uint64_t ld, + int32_t batch_count, + int64_t batch_stride) -> cublastlt_matrix_layout + { + RAFT_EXPECTS(batch_count > 0, "cuBLASLt batch count must be positive"); + auto layout = for_matmul(col_major, rows, cols, ld); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute( + layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_stride, sizeof(batch_stride))); + return layout; + } }; /** Descriptor for a cublasLt matmul function. */ @@ -197,9 +214,12 @@ struct cublastlt_matmul_desc { inline operator cublasLtMatmulDesc_t() const noexcept { return res; } template - static inline auto for_matmul(bool transpose_a, bool transpose_b) -> cublastlt_matmul_desc + static inline auto for_matmul(bool transpose_a, + bool transpose_b, + cublasComputeType_t compute_type = get_matmul_type()) + -> cublastlt_matmul_desc { - auto desc = cublastlt_matmul_desc{get_matmul_type(), get_cuda_data_type()}; + auto desc = cublastlt_matmul_desc{compute_type, get_cuda_data_type()}; if constexpr (DevicePointerMode) { const cublasPointerMode_t mode = CUBLAS_POINTER_MODE_DEVICE; RAFT_CUBLAS_TRY(cublasLtMatmulDescSetAttribute( @@ -348,6 +368,70 @@ struct coef_wrapper { } }; +/** + * Run a strided-batched cublasLt matmul without an algorithm or workspace preference. + * + * The matrix dimensions and leading dimensions follow the same column-major convention as + * `matmul`. Batch strides are measured in elements. + */ +template +void matmul_strided_batched(raft::resources const& res, + bool trans_a, + bool trans_b, + uint64_t m, + uint64_t n, + uint64_t k, + const S* alpha, + const A* a_ptr, + uint64_t lda, + int64_t stride_a, + const B* b_ptr, + uint64_t ldb, + int64_t stride_b, + const S* beta, + C* c_ptr, + uint64_t ldc, + int64_t stride_c, + int32_t batch_count, + cublasComputeType_t compute_type) +{ + common::nvtx::range batch_scope( + "linalg::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)", + m, + n, + k, + batch_count); + + auto operation = cublastlt_matmul_desc::for_matmul( + trans_a, trans_b, compute_type); + + auto a_layout = cublastlt_matrix_layout::for_strided_batched_matmul( + !trans_a, m, k, lda, batch_count, stride_a); + auto b_layout = cublastlt_matrix_layout::for_strided_batched_matmul( + !trans_b, k, n, ldb, batch_count, stride_b); + auto c_layout = + cublastlt_matrix_layout::for_strided_batched_matmul(true, m, n, ldc, batch_count, stride_c); + + auto stream = resource::get_cuda_stream(res); + coef_wrapper coefficients(alpha, beta, stream); + RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res), + operation, + coefficients.alpha, + a_ptr, + a_layout, + b_ptr, + b_layout, + coefficients.beta, + c_ptr, + c_layout, + c_ptr, + c_layout, + nullptr, + nullptr, + 0, + stream)); +} + /** * Compatibility version of the cublasLt matmul wrapper: It takes the cudaStream_t argument * explicitly rather than through the raft::resources. This function is used by other legacy diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index ada9dfe75f..09023cb267 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -283,6 +283,82 @@ void gemm(raft::resources const& res, } } + +/** + * @brief Strided-batched matrix multiplication using cublasLt. + * + * Computes `C_i = alpha * op(A_i) * op(B_i) + beta * C_i` for `batch_count` matrices. Matrix + * dimensions and leading dimensions use column-major conventions; batch strides are measured in + * elements. The compute type may be selected independently of the input and output storage types. + * + * @tparam A_t element type of A + * @tparam B_t element type of B + * @tparam C_t element type of C + * @tparam S_t element type of alpha and beta + * @tparam DevicePointerMode whether alpha and beta point to device memory + * @param[in] res RAFT resources + * @param[in] trans_a whether to transpose each A matrix + * @param[in] trans_b whether to transpose each B matrix + * @param[in] m number of rows of each output matrix + * @param[in] n number of columns of each output matrix + * @param[in] k shared inner dimension + * @param[in] alpha host or device scalar; defaults to one when null + * @param[in] a pointer to the first A matrix + * @param[in] lda leading dimension of each A matrix + * @param[in] stride_a offset in elements between A matrices + * @param[in] b pointer to the first B matrix + * @param[in] ldb leading dimension of each B matrix + * @param[in] stride_b offset in elements between B matrices + * @param[in] beta host or device scalar; defaults to zero when null + * @param[inout] c pointer to the first C matrix + * @param[in] ldc leading dimension of each C matrix + * @param[in] stride_c offset in elements between C matrices + * @param[in] batch_count number of matrix multiplications + * @param[in] compute_type cublasLt compute type + */ + template + void gemm_strided_batched( + raft::resources const& res, + bool trans_a, + bool trans_b, + uint64_t m, + uint64_t n, + uint64_t k, + const S_t* alpha, + const A_t* a, + uint64_t lda, + int64_t stride_a, + const B_t* b, + uint64_t ldb, + int64_t stride_b, + const S_t* beta, + C_t* c, + uint64_t ldc, + int64_t stride_c, + int32_t batch_count, + cublasComputeType_t compute_type = detail::get_matmul_type()) + { + detail::matmul_strided_batched(res, + trans_a, + trans_b, + m, + n, + k, + alpha, + a, + lda, + stride_a, + b, + ldb, + stride_b, + beta, + c, + ldc, + stride_c, + batch_count, + compute_type); + } + /** @} */ // end of gemm } // namespace linalg diff --git a/cpp/tests/linalg/gemm_basic.cpp b/cpp/tests/linalg/gemm_basic.cpp index 8622eec113..9c70323a8b 100644 --- a/cpp/tests/linalg/gemm_basic.cpp +++ b/cpp/tests/linalg/gemm_basic.cpp @@ -163,6 +163,55 @@ 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, GemmStridedBatched) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + constexpr int64_t stride_a = 8; + constexpr int64_t stride_b = 8; + constexpr int64_t stride_c = 5; + constexpr int32_t batches = 2; + + // Two column-major A (2 x 3) and B (3 x 2) matrices with padding between batches. + std::vector a_host = {1, 4, 2, 5, 3, 6, -1, -1, 2, 1, 0, 3, 1, 4, -1, -1}; + std::vector b_host = {7, 9, 11, 8, 10, 12, -1, -1, 1, 0, 2, 3, 1, 4, -1, -1}; + std::vector c_host(stride_c * batches, -1); + + auto a_device = raft::make_device_vector(res, a_host.size()); + auto b_device = raft::make_device_vector(res, b_host.size()); + auto c_device = raft::make_device_vector(res, c_host.size()); + raft::copy(a_device.data_handle(), a_host.data(), a_host.size(), stream); + raft::copy(b_device.data_handle(), b_host.data(), b_host.size(), stream); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + + raft::linalg::gemm_strided_batched(res, + false, + false, + M, + N, + K, + nullptr, + a_device.data_handle(), + M, + stride_a, + b_device.data_handle(), + K, + stride_b, + nullptr, + c_device.data_handle(), + M, + stride_c, + batches, + CUBLAS_COMPUTE_32F_FAST_TF32); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + + const std::vector expected = {58, 139, 64, 154, -1, 4, 9, 10, 22, -1}; + EXPECT_EQ(c_host, expected); +} + TEST(Raft, GemmCublasLt136WorkaroundPredicate) { constexpr std::size_t affected_version = 130600; From a8d89a27e979e4f5cb46e4f2f41b876fe1ce7eeb Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Mon, 10 Aug 2026 21:42:42 +0000 Subject: [PATCH 2/8] style fix --- cpp/include/raft/linalg/gemm.cuh | 87 ++++++++++++++++---------------- 1 file changed, 43 insertions(+), 44 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index 09023cb267..adb8ed48e3 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #ifndef __GEMM_H @@ -283,7 +283,6 @@ void gemm(raft::resources const& res, } } - /** * @brief Strided-batched matrix multiplication using cublasLt. * @@ -316,48 +315,48 @@ void gemm(raft::resources const& res, * @param[in] batch_count number of matrix multiplications * @param[in] compute_type cublasLt compute type */ - template - void gemm_strided_batched( - raft::resources const& res, - bool trans_a, - bool trans_b, - uint64_t m, - uint64_t n, - uint64_t k, - const S_t* alpha, - const A_t* a, - uint64_t lda, - int64_t stride_a, - const B_t* b, - uint64_t ldb, - int64_t stride_b, - const S_t* beta, - C_t* c, - uint64_t ldc, - int64_t stride_c, - int32_t batch_count, - cublasComputeType_t compute_type = detail::get_matmul_type()) - { - detail::matmul_strided_batched(res, - trans_a, - trans_b, - m, - n, - k, - alpha, - a, - lda, - stride_a, - b, - ldb, - stride_b, - beta, - c, - ldc, - stride_c, - batch_count, - compute_type); - } +template +void gemm_strided_batched( + raft::resources const& res, + bool trans_a, + bool trans_b, + uint64_t m, + uint64_t n, + uint64_t k, + const S_t* alpha, + const A_t* a, + uint64_t lda, + int64_t stride_a, + const B_t* b, + uint64_t ldb, + int64_t stride_b, + const S_t* beta, + C_t* c, + uint64_t ldc, + int64_t stride_c, + int32_t batch_count, + cublasComputeType_t compute_type = detail::get_matmul_type()) +{ + detail::matmul_strided_batched(res, + trans_a, + trans_b, + m, + n, + k, + alpha, + a, + lda, + stride_a, + b, + ldb, + stride_b, + beta, + c, + ldc, + stride_c, + batch_count, + compute_type); +} /** @} */ // end of gemm From 9ee66b1c3e078e6d421405fa64583a3206c097f5 Mon Sep 17 00:00:00 2001 From: Ben Landrum <36489943+landrumb@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:28:09 -0700 Subject: [PATCH 3/8] rewriting public API Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com> --- cpp/include/raft/linalg/gemm.cuh | 38 ++++++++++++++------------------ 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index adb8ed48e3..902ef94596 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -315,27 +315,23 @@ void gemm(raft::resources const& res, * @param[in] batch_count number of matrix multiplications * @param[in] compute_type cublasLt compute type */ -template -void gemm_strided_batched( - raft::resources const& res, - bool trans_a, - bool trans_b, - uint64_t m, - uint64_t n, - uint64_t k, - const S_t* alpha, - const A_t* a, - uint64_t lda, - int64_t stride_a, - const B_t* b, - uint64_t ldb, - int64_t stride_b, - const S_t* beta, - C_t* c, - uint64_t ldc, - int64_t stride_c, - int32_t batch_count, - cublasComputeType_t compute_type = detail::get_matmul_type()) +template , + typename = std::enable_if_t>, + std::is_same>>>> +void gemm_batched(raft::resources const& res, + raft::device_mdspan, LayoutPolicyX> x, + raft::device_mdspan, LayoutPolicyY> y, + raft::device_mdspan, LayoutPolicyZ> z, + std::optional alpha = std::nullopt, + std::optional beta = std::nullopt, + cublasComputeType_t compute_type_override = detail::get_matmul_type())) { detail::matmul_strided_batched(res, trans_a, From 406a785095a0e5e074265582864168b6e2a0c79f Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Tue, 11 Aug 2026 22:04:22 +0000 Subject: [PATCH 4/8] switching to overload for optional compute type --- cpp/include/raft/linalg/gemm.cuh | 215 +++++++++++++++++++++++-------- cpp/tests/linalg/gemm_basic.cpp | 99 ++++++++++---- 2 files changed, 237 insertions(+), 77 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index 902ef94596..419cbe7dc2 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -284,36 +284,27 @@ void gemm(raft::resources const& res, } /** - * @brief Strided-batched matrix multiplication using cublasLt. + * @brief Batched matrix multiplication with an explicit cublasLt compute type. * - * Computes `C_i = alpha * op(A_i) * op(B_i) + beta * C_i` for `batch_count` matrices. Matrix - * dimensions and leading dimensions use column-major conventions; batch strides are measured in - * elements. The compute type may be selected independently of the input and output storage types. + * Computes `Z_i = alpha * X_i * Y_i + beta * Z_i`. The first mdspan extent is the batch + * dimension. `layout_stride` views encode matrix and batch strides in their mappings; contiguous + * batches use `layout_right`, with the matrix size as the batch stride. Use `layout_stride` for + * column-major batches because rank-3 `layout_left` interleaves the batch dimension. * - * @tparam A_t element type of A - * @tparam B_t element type of B - * @tparam C_t element type of C - * @tparam S_t element type of alpha and beta - * @tparam DevicePointerMode whether alpha and beta point to device memory + * @tparam ValueType element type of the input and output matrices + * @tparam IndexType index type + * @tparam LayoutPolicyX layout policy of X + * @tparam LayoutPolicyY layout policy of Y + * @tparam LayoutPolicyZ layout policy of Z + * @tparam ScalarIdxType index type of alpha and beta + * @tparam ScalarViewType scalar view type of alpha and beta * @param[in] res RAFT resources - * @param[in] trans_a whether to transpose each A matrix - * @param[in] trans_b whether to transpose each B matrix - * @param[in] m number of rows of each output matrix - * @param[in] n number of columns of each output matrix - * @param[in] k shared inner dimension - * @param[in] alpha host or device scalar; defaults to one when null - * @param[in] a pointer to the first A matrix - * @param[in] lda leading dimension of each A matrix - * @param[in] stride_a offset in elements between A matrices - * @param[in] b pointer to the first B matrix - * @param[in] ldb leading dimension of each B matrix - * @param[in] stride_b offset in elements between B matrices - * @param[in] beta host or device scalar; defaults to zero when null - * @param[inout] c pointer to the first C matrix - * @param[in] ldc leading dimension of each C matrix - * @param[in] stride_c offset in elements between C matrices - * @param[in] batch_count number of matrix multiplications - * @param[in] compute_type cublasLt compute type + * @param[in] x input matrices X, with shape `[batch_count, m, k]` + * @param[in] y input matrices Y, with shape `[batch_count, k, n]` + * @param[inout] z input/output matrices Z, with shape `[batch_count, m, n]` + * @param[in] alpha scalar multiplier for X * Y; defaults to one when empty + * @param[in] beta scalar multiplier for Z; defaults to zero when empty + * @param[in] compute_type_override cublasLt compute type */ template >, std::is_same>>>> void gemm_batched(raft::resources const& res, - raft::device_mdspan, LayoutPolicyX> x, - raft::device_mdspan, LayoutPolicyY> y, - raft::device_mdspan, LayoutPolicyZ> z, - std::optional alpha = std::nullopt, - std::optional beta = std::nullopt, - cublasComputeType_t compute_type_override = detail::get_matmul_type())) + raft::device_mdspan, LayoutPolicyX> x, + raft::device_mdspan, LayoutPolicyY> y, + raft::device_mdspan, LayoutPolicyZ> z, + std::optional alpha, + std::optional beta, + cublasComputeType_t compute_type_override) +{ + static_assert(!std::is_same_v && + !std::is_same_v && + !std::is_same_v, + "Use layout_stride for column-major batched matrices"); + + RAFT_EXPECTS(x.extent(0) == y.extent(0) && x.extent(0) == z.extent(0), + "All operands must have the same batch count"); + RAFT_EXPECTS(x.extent(1) == z.extent(1), "X and Z must have the same number of rows"); + RAFT_EXPECTS(y.extent(2) == z.extent(2), "Y and Z must have the same number of columns"); + RAFT_EXPECTS(x.extent(2) == y.extent(1), + "The number of columns in X must equal the number of rows in Y"); + + const auto is_row_major = [](auto const& view) { + return view.stride(2) == 1 && view.stride(1) >= view.extent(2); + }; + const auto is_col_major = [](auto const& view) { + return view.stride(1) == 1 && view.stride(2) >= view.extent(1); + }; + const auto batch_stride = [](auto const& view) -> int64_t { + using layout_type = typename std::decay_t::layout_type; + if constexpr (std::is_same_v) { + return static_cast(view.stride(0)); + } else { + return static_cast(view.extent(1)) * static_cast(view.extent(2)); + } + }; + RAFT_EXPECTS(is_row_major(x) || is_col_major(x), + "Each matrix in X must be row-major or column-major"); + RAFT_EXPECTS(is_row_major(y) || is_col_major(y), + "Each matrix in Y must be row-major or column-major"); + RAFT_EXPECTS(is_row_major(z) || is_col_major(z), + "Each matrix in Z must be row-major or column-major"); + + constexpr auto kDeviceMode = + std::is_same_v>; + ValueType* alpha_ptr = nullptr; + ValueType* beta_ptr = nullptr; + if (alpha.has_value()) { alpha_ptr = alpha->data_handle(); } + if (beta.has_value()) { beta_ptr = beta->data_handle(); } + + const bool x_col_major = is_col_major(x); + const bool y_col_major = is_col_major(y); + const bool z_col_major = is_col_major(z); + const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); + const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); + const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); + const auto x_batch_stride = batch_stride(x); + const auto y_batch_stride = batch_stride(y); + const auto z_batch_stride = batch_stride(z); + + if (z_col_major) { + return detail::matmul_strided_batched(res, + !x_col_major, + !y_col_major, + z.extent(1), + z.extent(2), + x.extent(2), + alpha_ptr, + x.data_handle(), + x_ld, + x_batch_stride, + y.data_handle(), + y_ld, + y_batch_stride, + beta_ptr, + z.data_handle(), + z_ld, + z_batch_stride, + z.extent(0), + compute_type_override); + } + + return detail::matmul_strided_batched(res, + y_col_major, + x_col_major, + z.extent(2), + z.extent(1), + x.extent(2), + alpha_ptr, + y.data_handle(), + y_ld, + y_batch_stride, + x.data_handle(), + x_ld, + x_batch_stride, + beta_ptr, + z.data_handle(), + z_ld, + z_batch_stride, + z.extent(0), + compute_type_override); +} + +/** + * @brief Batched matrix multiplication using the default cublasLt compute type. + * + * Computes `Z_i = alpha * X_i * Y_i + beta * Z_i`. The first mdspan extent is the batch + * dimension. `layout_stride` views encode matrix and batch strides in their mappings; contiguous + * batches use `layout_right`, with the matrix size as the batch stride. Use `layout_stride` for + * column-major batches because rank-3 `layout_left` interleaves the batch dimension. + * + * @tparam ValueType element type of the input and output matrices + * @tparam IndexType index type + * @tparam LayoutPolicyX layout policy of X + * @tparam LayoutPolicyY layout policy of Y + * @tparam LayoutPolicyZ layout policy of Z + * @tparam ScalarIdxType index type of alpha and beta + * @tparam ScalarViewType scalar view type of alpha and beta + * @param[in] res RAFT resources + * @param[in] x input matrices X, with shape `[batch_count, m, k]` + * @param[in] y input matrices Y, with shape `[batch_count, k, n]` + * @param[inout] z input/output matrices Z, with shape `[batch_count, m, n]` + * @param[in] alpha scalar multiplier for X * Y; defaults to one when empty + * @param[in] beta scalar multiplier for Z; defaults to zero when empty + */ +template , + typename = std::enable_if_t>, + std::is_same>>>> +void gemm_batched(raft::resources const& res, + raft::device_mdspan, LayoutPolicyX> x, + raft::device_mdspan, LayoutPolicyY> y, + raft::device_mdspan, LayoutPolicyZ> z, + std::optional alpha = std::nullopt, + std::optional beta = std::nullopt) { - detail::matmul_strided_batched(res, - trans_a, - trans_b, - m, - n, - k, - alpha, - a, - lda, - stride_a, - b, - ldb, - stride_b, - beta, - c, - ldc, - stride_c, - batch_count, - compute_type); + return gemm_batched(res, + x, + y, + z, + alpha, + beta, + detail::get_matmul_type()); } /** @} */ // end of gemm diff --git a/cpp/tests/linalg/gemm_basic.cpp b/cpp/tests/linalg/gemm_basic.cpp index 9c70323a8b..8550cb871b 100644 --- a/cpp/tests/linalg/gemm_basic.cpp +++ b/cpp/tests/linalg/gemm_basic.cpp @@ -163,20 +163,21 @@ 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, GemmStridedBatched) +TEST(Raft, GemmBatched) { raft::resources res; auto stream = raft::resource::get_cuda_stream(res); - constexpr int64_t stride_a = 8; - constexpr int64_t stride_b = 8; - constexpr int64_t stride_c = 5; - constexpr int32_t batches = 2; + using index_type = int64_t; + constexpr index_type stride_a = 8; + constexpr index_type stride_b = 8; + constexpr index_type stride_c = 5; + constexpr index_type batch_count = 2; // Two column-major A (2 x 3) and B (3 x 2) matrices with padding between batches. std::vector a_host = {1, 4, 2, 5, 3, 6, -1, -1, 2, 1, 0, 3, 1, 4, -1, -1}; std::vector b_host = {7, 9, 11, 8, 10, 12, -1, -1, 1, 0, 2, 3, 1, 4, -1, -1}; - std::vector c_host(stride_c * batches, -1); + std::vector c_host(stride_c * batch_count, -1); auto a_device = raft::make_device_vector(res, a_host.size()); auto b_device = raft::make_device_vector(res, b_host.size()); @@ -185,31 +186,79 @@ TEST(Raft, GemmStridedBatched) raft::copy(b_device.data_handle(), b_host.data(), b_host.size(), stream); raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); - raft::linalg::gemm_strided_batched(res, - false, - false, - M, - N, - K, - nullptr, - a_device.data_handle(), - M, - stride_a, - b_device.data_handle(), - K, - stride_b, - nullptr, - c_device.data_handle(), - M, - stride_c, - batches, - CUBLAS_COMPUTE_32F_FAST_TF32); + auto a_extents = raft::extent_3d{batch_count, M, K}; + auto b_extents = raft::extent_3d{batch_count, K, N}; + auto c_extents = raft::extent_3d{batch_count, M, N}; + auto a_layout = + raft::make_strided_layout(a_extents, cuda::std::array{stride_a, 1, M}); + auto b_layout = + raft::make_strided_layout(b_extents, cuda::std::array{stride_b, 1, K}); + auto c_layout = + raft::make_strided_layout(c_extents, cuda::std::array{stride_c, 1, M}); + + auto a_view = raft::device_mdspan{ + a_device.data_handle(), a_layout}; + auto b_view = raft::device_mdspan{ + b_device.data_handle(), b_layout}; + auto c_view = raft::device_mdspan{ + c_device.data_handle(), c_layout}; + + std::optional> alpha; + std::optional> beta; + raft::linalg::gemm_batched( + res, a_view, b_view, c_view, alpha, beta, CUBLAS_COMPUTE_32F_FAST_TF32); raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); raft::resource::sync_stream(res); const std::vector expected = {58, 139, 64, 154, -1, 4, 9, 10, 22, -1}; EXPECT_EQ(c_host, expected); + + c_host.assign(stride_c * batch_count, -1); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + raft::linalg::gemm_batched(res, a_view, b_view, c_view); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + EXPECT_EQ(c_host, expected); +} + +TEST(Raft, GemmBatchedContiguous) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + + using index_type = int64_t; + constexpr index_type batch_count = 2; + + std::vector a_host = {1, 2, 3, 4, 5, 6, 2, 0, 1, 1, 3, 4}; + std::vector b_host = {7, 8, 9, 10, 11, 12, 1, 3, 0, 1, 2, 4}; + std::vector c_host(batch_count * M * N, -1); + + auto a_device = raft::make_device_vector(res, a_host.size()); + auto b_device = raft::make_device_vector(res, b_host.size()); + auto c_device = raft::make_device_vector(res, c_host.size()); + raft::copy(a_device.data_handle(), a_host.data(), a_host.size(), stream); + raft::copy(b_device.data_handle(), b_host.data(), b_host.size(), stream); + raft::copy(c_device.data_handle(), c_host.data(), c_host.size(), stream); + + auto a_extents = raft::extent_3d{batch_count, M, K}; + auto b_extents = raft::extent_3d{batch_count, K, N}; + auto c_extents = raft::extent_3d{batch_count, M, N}; + auto a_view = raft::device_mdspan{ + a_device.data_handle(), a_extents}; + auto b_view = raft::device_mdspan{ + b_device.data_handle(), b_extents}; + auto c_view = raft::device_mdspan{ + c_device.data_handle(), c_extents}; + + raft::linalg::gemm_batched(res, a_view, b_view, c_view); + + raft::copy(c_host.data(), c_device.data_handle(), c_host.size(), stream); + raft::resource::sync_stream(res); + + const std::vector expected = {58, 64, 139, 154, 4, 10, 9, 22}; + EXPECT_EQ(c_host, expected); } TEST(Raft, GemmCublasLt136WorkaroundPredicate) From 8d733183bd082eba8c972b53340d64ca045def41 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 12 Aug 2026 01:07:01 +0000 Subject: [PATCH 5/8] addresed coderabbit comment + simplified call --- cpp/include/raft/linalg/gemm.cuh | 29 ++++------------------------- 1 file changed, 4 insertions(+), 25 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index 419cbe7dc2..2c8552104e 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -364,37 +364,16 @@ void gemm_batched(raft::resources const& res, if (alpha.has_value()) { alpha_ptr = alpha->data_handle(); } if (beta.has_value()) { beta_ptr = beta->data_handle(); } - const bool x_col_major = is_col_major(x); - const bool y_col_major = is_col_major(y); const bool z_col_major = is_col_major(z); + const bool x_col_major = is_col_major(x) ^ z_col_major; + const bool y_col_major = is_col_major(y) ^ z_col_major; const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); const auto x_batch_stride = batch_stride(x); const auto y_batch_stride = batch_stride(y); const auto z_batch_stride = batch_stride(z); - - if (z_col_major) { - return detail::matmul_strided_batched(res, - !x_col_major, - !y_col_major, - z.extent(1), - z.extent(2), - x.extent(2), - alpha_ptr, - x.data_handle(), - x_ld, - x_batch_stride, - y.data_handle(), - y_ld, - y_batch_stride, - beta_ptr, - z.data_handle(), - z_ld, - z_batch_stride, - z.extent(0), - compute_type_override); - } + const auto batch_count = static_cast(z.extent(0)); return detail::matmul_strided_batched(res, y_col_major, @@ -413,7 +392,7 @@ void gemm_batched(raft::resources const& res, z.data_handle(), z_ld, z_batch_stride, - z.extent(0), + batch_count, compute_type_override); } From ea5e525425a13acce67679a666b91561fd3d84e1 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 12 Aug 2026 01:09:10 +0000 Subject: [PATCH 6/8] concision --- cpp/include/raft/linalg/gemm.cuh | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index 2c8552104e..3180f5ed96 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -364,16 +364,13 @@ void gemm_batched(raft::resources const& res, if (alpha.has_value()) { alpha_ptr = alpha->data_handle(); } if (beta.has_value()) { beta_ptr = beta->data_handle(); } - const bool z_col_major = is_col_major(z); - const bool x_col_major = is_col_major(x) ^ z_col_major; - const bool y_col_major = is_col_major(y) ^ z_col_major; - const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); - const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); - const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); - const auto x_batch_stride = batch_stride(x); - const auto y_batch_stride = batch_stride(y); - const auto z_batch_stride = batch_stride(z); - const auto batch_count = static_cast(z.extent(0)); + const bool z_col_major = is_col_major(z); + const bool x_col_major = is_col_major(x) ^ z_col_major; + const bool y_col_major = is_col_major(y) ^ z_col_major; + const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); + const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); + const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); + const auto batch_count = static_cast(z.extent(0)); return detail::matmul_strided_batched(res, y_col_major, @@ -384,14 +381,14 @@ void gemm_batched(raft::resources const& res, alpha_ptr, y.data_handle(), y_ld, - y_batch_stride, + batch_stride(y), x.data_handle(), x_ld, - x_batch_stride, + batch_stride(x), beta_ptr, z.data_handle(), z_ld, - z_batch_stride, + batch_stride(z), batch_count, compute_type_override); } From b17fab11df527f4b984d567f5ec8eeca613f7fb8 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 12 Aug 2026 18:00:20 +0000 Subject: [PATCH 7/8] renamed nvtx range --- cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp index e509e31822..fb7fb9047b 100644 --- a/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp +++ b/cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp @@ -396,7 +396,7 @@ void matmul_strided_batched(raft::resources const& res, cublasComputeType_t compute_type) { common::nvtx::range batch_scope( - "linalg::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)", + "linalg::detail::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)", m, n, k, From 75e99fba7aeacb7bb238c6521d0670b073330390 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 12 Aug 2026 18:28:55 +0000 Subject: [PATCH 8/8] fixed incorrect arg ordering for transposed output --- cpp/include/raft/linalg/gemm.cuh | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/linalg/gemm.cuh b/cpp/include/raft/linalg/gemm.cuh index 3180f5ed96..80c083d682 100644 --- a/cpp/include/raft/linalg/gemm.cuh +++ b/cpp/include/raft/linalg/gemm.cuh @@ -364,14 +364,36 @@ void gemm_batched(raft::resources const& res, if (alpha.has_value()) { alpha_ptr = alpha->data_handle(); } if (beta.has_value()) { beta_ptr = beta->data_handle(); } + const bool x_col_major = is_col_major(x); + const bool y_col_major = is_col_major(y); const bool z_col_major = is_col_major(z); - const bool x_col_major = is_col_major(x) ^ z_col_major; - const bool y_col_major = is_col_major(y) ^ z_col_major; const auto x_ld = x_col_major ? x.stride(2) : x.stride(1); const auto y_ld = y_col_major ? y.stride(2) : y.stride(1); const auto z_ld = z_col_major ? z.stride(2) : z.stride(1); const auto batch_count = static_cast(z.extent(0)); + if (z_col_major) { + return detail::matmul_strided_batched(res, + !x_col_major, + !y_col_major, + z.extent(1), + z.extent(2), + x.extent(2), + alpha_ptr, + x.data_handle(), + x_ld, + batch_stride(x), + y.data_handle(), + y_ld, + batch_stride(y), + beta_ptr, + z.data_handle(), + z_ld, + batch_stride(z), + batch_count, + compute_type_override); + } + return detail::matmul_strided_batched(res, y_col_major, x_col_major,