Skip to content

add-batched-linalg for cuvs #3110

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

Open
Intron7 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

add-batched-linalg for cuvs #3110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f114b30
add-batched-linalg
Intron7 Aug 11, 2026
ebc2ce4
Merge branch 'main' into feat/add-batched-linalg
Intron7 Aug 12, 2026
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
65 changes: 64 additions & 1 deletion cpp/include/raft/linalg/detail/cublas_wrappers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -930,6 +930,69 @@ inline cublasStatus_t cublastrsm(cublasHandle_t handle,
return cublasDtrsm(handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb);
}

/**
* @defgroup trsmBatched cublas trsmBatched calls
* @{
*/
template <typename T>
cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT
cublasSideMode_t side,
cublasFillMode_t uplo,
cublasOperation_t trans,
cublasDiagType_t diag,
int m,
int n,
const T* alpha,
const T* const Aarray[], // NOLINT
int lda,
T* const Barray[], // NOLINT
int ldb,
int batchCount,
cudaStream_t stream);

template <>
inline cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT
cublasSideMode_t side,
cublasFillMode_t uplo,
cublasOperation_t trans,
cublasDiagType_t diag,
int m,
int n,
const float* alpha,
const float* const Aarray[], // NOLINT
int lda,
float* const Barray[], // NOLINT
int ldb,
int batchCount,
cudaStream_t stream)
{
RAFT_CUBLAS_TRY(cublasSetStream(handle, stream));
return cublasStrsmBatched(
handle, side, uplo, trans, diag, m, n, alpha, Aarray, lda, Barray, ldb, batchCount);
}

template <>
inline cublasStatus_t cublastrsmBatched(cublasHandle_t handle, // NOLINT
cublasSideMode_t side,
cublasFillMode_t uplo,
cublasOperation_t trans,
cublasDiagType_t diag,
int m,
int n,
const double* alpha,
const double* const Aarray[], // NOLINT
int lda,
double* const Barray[], // NOLINT
int ldb,
int batchCount,
cudaStream_t stream)
{
RAFT_CUBLAS_TRY(cublasSetStream(handle, stream));
return cublasDtrsmBatched(
handle, side, uplo, trans, diag, m, n, alpha, Aarray, lda, Barray, ldb, batchCount);
}
/** @} */

/**
* @defgroup dot cublas dot calls
* @{
Expand Down
191 changes: 156 additions & 35 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 @@ -85,18 +85,28 @@ struct matmul_key_t {
uint64_t ldc;
bool trans_a;
bool trans_b;
/** Number of matrices in the batch; 1 means a plain, non-batched matmul. */
uint64_t batch_count = 1;
/** Offsets in elements between consecutive matrices of the batch (ignored if batch_count == 1) */
int64_t stride_a = 0;
int64_t stride_b = 0;
int64_t stride_c = 0;
};

inline auto operator==(const matmul_key_t& a, const matmul_key_t& b) -> bool
{
return a.m == b.m && a.n == b.n && a.k == b.k && a.lda == b.lda && a.ldb == b.ldb &&
a.ldc == b.ldc && a.trans_a == b.trans_a && a.trans_b == b.trans_b;
a.ldc == b.ldc && a.trans_a == b.trans_a && a.trans_b == b.trans_b &&
a.batch_count == b.batch_count && a.stride_a == b.stride_a && a.stride_b == b.stride_b &&
a.stride_c == b.stride_c;
}

struct matmul_key_hash {
inline auto operator()(const matmul_key_t& x) const noexcept -> std::size_t
{
return x.m * x.n * x.k + x.lda * x.ldb * x.ldc + size_t{x.trans_a} + size_t{x.trans_b} * 2;
return x.m * x.n * x.k + x.lda * x.ldb * x.ldc + size_t{x.trans_a} + size_t{x.trans_b} * 2 +
x.batch_count * (static_cast<size_t>(x.stride_a) + static_cast<size_t>(x.stride_b) +
static_cast<size_t>(x.stride_c));
}
};

Expand Down Expand Up @@ -160,12 +170,32 @@ struct cublastlt_matrix_layout {
// NOLINTNEXTLINE
inline operator cublasLtMatrixLayout_t() const noexcept { return res; }

/**
* Describe the matrix as a batch of `batch_count` matrices, `batch_stride` elements apart.
* A `batch_count` of one leaves the layout as a plain, non-batched matrix.
*/
inline void set_batch(uint64_t batch_count, int64_t batch_stride)
{
if (batch_count <= 1) { return; }
const auto count = static_cast<int32_t>(batch_count);
RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute(
res, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &count, sizeof(count)));
RAFT_CUBLAS_TRY(cublasLtMatrixLayoutSetAttribute(
res, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_stride, sizeof(batch_stride)));
}

template <typename T>
static inline auto for_matmul(bool col_major, uint64_t rows, uint64_t cols, uint64_t ld)
-> cublastlt_matrix_layout
static inline auto for_matmul(bool col_major,
uint64_t rows,
uint64_t cols,
uint64_t ld,
uint64_t batch_count = 1,
int64_t batch_stride = 0) -> cublastlt_matrix_layout
{
return cublastlt_matrix_layout{
auto r = cublastlt_matrix_layout{
get_cuda_data_type<T>(), col_major ? rows : cols, col_major ? cols : rows, ld};
r.set_batch(batch_count, batch_stride);
return r;
}
Comment on lines +173 to 199

@coderabbitai coderabbitai Bot Aug 12, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🌐 Web query:

cuBLASLt CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT attribute data type int32_t

💡 Result:

The cuBLASLt attribute CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT uses the int32_t data type [1][2]. This attribute is part of the cublasLtMatrixLayoutAttribute_t enumeration and specifies the number of matrix multiplication operations to perform in a batch [2]. Its default value is 1 [2]. Implementation examples and developer documentation confirm that when setting this attribute using functions such as cublasLtMatrixLayoutSetAttribute, the value should be provided as an int32_t [3][1][2].

Citations:

🏁 Script executed:

#!/bin/bash
set -eu

file="cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp"

printf '%s\n' '--- target implementation ---'
sed -n '130,220p' "$file"

printf '%s\n' '--- related declarations and call sites ---'
rg -n -C 5 'set_batch|matmul_strided_batched|gemm_batched|batch_count' cpp/include cpp/src cpp/tests 2>/dev/null | head -n 1000

printf '%s\n' '--- includes and namespace context ---'
sed -n '1,80p' "$file"

Repository: NVIDIA/raft

Length of output: 35488

🏁 Script executed:

#!/bin/bash
set -eu

printf '%s\n' '--- all non-test call sites ---'
rg -n -C 8 'matmul_strided_batched|for_matmul<|set_batch\(' cpp/include cpp/src --glob '*.{hpp,h,cuh,cpp,cu}' 2>/dev/null

printf '%s\n' '--- validation and expected-error conventions ---'
rg -n -C 4 'RAFT_EXPECTS\(.*numeric_limits|numeric_limits<int32_t>|INT32_MAX|batch_count.*EXPECT|extent\(0\).*RAFT_EXPECTS' cpp/include cpp/src cpp/tests --glob '*.{hpp,h,cuh,cpp,cu}' 2>/dev/null | head -n 500

printf '%s\n' '--- matmul public wrapper context ---'
sed -n '320,405p' cpp/include/raft/linalg/gemm.cuh
sed -n '530,590p' cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp

Repository: NVIDIA/raft

Length of output: 25074

🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path

source = Path("cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp").read_text()
needle = "const auto count = static_cast<int32_t>(batch_count);"
assert needle in source

limit = (1 << 31) - 1
for value in (1, 2, limit, limit + 1, 1 << 32, (1 << 64) - 1):
    unsigned_bits = value & 0xFFFFFFFF
    signed_value = unsigned_bits if unsigned_bits < (1 << 31) else unsigned_bits - (1 << 32)
    print(f"batch_count={value}: int32_t two's-complement result={signed_value}")
PY

Repository: NVIDIA/raft

Length of output: 509

Validate batch_count before the narrowing cast.

CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT requires int32_t, but set_batch casts uint64_t without range validation. Counts above INT32_MAX can configure an incorrect batch. Add a guard before the cast and include <limits> if needed.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp` around lines 173 - 199,
Update cublastlt_matrix_layout::set_batch to validate batch_count does not
exceed int32_t’s maximum before the static_cast<int32_t>. Reject out-of-range
values using the existing error-handling convention, and include <limits> if
required for the validation.

Source: Coding guidelines

};

Expand Down Expand Up @@ -249,9 +279,12 @@ struct matmul_desc {
{
matmul_desc r{
cublastlt_matmul_desc::for_matmul<S, A, B, C, DevicePointerMode>(args.trans_a, args.trans_b),
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)};
cublastlt_matrix_layout::for_matmul<A>(
!(args.trans_a), args.m, args.k, args.lda, args.batch_count, args.stride_a),
cublastlt_matrix_layout::for_matmul<B>(
!(args.trans_b), args.k, args.n, args.ldb, args.batch_count, args.stride_b),
cublastlt_matrix_layout::for_matmul<C>(
true, args.m, args.n, args.ldc, args.batch_count, args.stride_c)};

bool use_cublaslt_13_6_workaround = false;
if constexpr (std::is_same_v<S, float> && std::is_same_v<A, float> &&
Expand All @@ -277,8 +310,12 @@ struct matmul_desc {

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);
const auto heuristic_a = cublastlt_matrix_layout::for_matmul<A>(!(heuristic_args.trans_a),
heuristic_args.m,
heuristic_args.k,
heuristic_args.lda,
heuristic_args.batch_count,
heuristic_args.stride_a);
query_heuristic(heuristic_a, r.c);
} else {
query_heuristic(r.a, r.c);
Expand Down Expand Up @@ -348,6 +385,47 @@ struct coef_wrapper<true, S> {
}
};

/**
* Shared implementation behind all cublasLt matmul wrappers: look up (or create and cache) the
* matmul descriptor for `mm_key` and run it. Batching, if any, is described by `mm_key`.
*/
template <bool DevicePointerMode, typename S, typename A, typename B, typename C>
void matmul_impl(raft::resources const& res,
const matmul_key_t& mm_key,
const S* alpha,
const A* a_ptr,
const B* b_ptr,
const S* beta,
C* c_ptr,
cudaStream_t stream)
{
std::shared_ptr<matmul_desc> mm_desc{nullptr};
auto& cache =
resource::get_custom_resource<matmul_cache<S, A, B, C, DevicePointerMode>>(res)->value;
if (!cache.get(mm_key, &mm_desc)) {
mm_desc.reset(new matmul_desc{matmul_desc::create<S, A, B, C, DevicePointerMode>(res, mm_key)});
cache.set(mm_key, mm_desc);
}
// Allocate alpha and beta pointers if not provided.
coef_wrapper<DevicePointerMode, S> w(alpha, beta, stream);
RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res),
mm_desc->desc,
w.alpha,
a_ptr,
mm_desc->a,
b_ptr,
mm_desc->b,
w.beta,
c_ptr,
mm_desc->c,
c_ptr,
mm_desc->c,
&(mm_desc->heuristics.algo),
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
Expand Down Expand Up @@ -376,32 +454,8 @@ template <bool DevicePointerMode = false, typename S, typename A, typename B, ty
{
common::nvtx::range<common::nvtx::domain::raft> batch_scope(
"linalg::matmul(m = %d, n = %d, k = %d)", m, n, k);
std::shared_ptr<matmul_desc> mm_desc{nullptr};
matmul_key_t mm_key{m, n, k, lda, ldb, ldc, trans_a, trans_b};
auto& cache =
resource::get_custom_resource<matmul_cache<S, A, B, C, DevicePointerMode>>(res)->value;
if (!cache.get(mm_key, &mm_desc)) {
mm_desc.reset(new matmul_desc{matmul_desc::create<S, A, B, C, DevicePointerMode>(res, mm_key)});
cache.set(mm_key, mm_desc);
}
// Allocate alpha and beta pointers if not provided.
coef_wrapper<DevicePointerMode, S> w(alpha, beta, stream);
RAFT_CUBLAS_TRY(cublasLtMatmul(resource::get_cublaslt_handle(res),
mm_desc->desc,
w.alpha,
a_ptr,
mm_desc->a,
b_ptr,
mm_desc->b,
w.beta,
c_ptr,
mm_desc->c,
c_ptr,
mm_desc->c,
&(mm_desc->heuristics.algo),
nullptr,
0,
stream));
matmul_impl<DevicePointerMode, S, A, B, C>(res, mm_key, alpha, a_ptr, b_ptr, beta, c_ptr, stream);
}

/**
Expand Down Expand Up @@ -462,5 +516,72 @@ void matmul(raft::resources const& res,
resource::get_cuda_stream(res));
}

/**
* @brief the wrapper of the strided-batched cublasLt matmul function
* For every batch index i it computes:
* C_i = alpha .* opA(A_i) * opB(B_i) + beta .* C_i
* where X_i is the matrix starting at `x_ptr + i * stride_x`.
*
* All matrices of a batch share the same shape, leading dimension and transpose op; only the
* base pointers differ. A stride of zero broadcasts the same matrix over the whole batch.
*
* @tparam DevicePointerMode whether pointers alpha, beta point to device memory
* @tparam S the type of scale parameters alpha, beta
* @tparam A the element type of matrix A
* @tparam B the element type of matrix B
* @tparam C the element type of matrix C
*
* @param [in] res raft resources
* @param [in] trans_a cublas transpose op for A
* @param [in] trans_b cublas transpose op for B
* @param [in] m number of rows of C
* @param [in] n number of columns of C
* @param [in] k number of rows of opB(B) / number of columns of opA(A)
* @param [in] alpha host or device scalar, if nullptr, the default value 1 will be used
* @param [in] a_ptr such a matrix that the shape of column-major opA(A) is [m, k]
* @param [in] lda leading dimension of A
* @param [in] stride_a offset in elements between consecutive matrices of A
* @param [in] b_ptr such a matrix that the shape of column-major opA(B) is [k, n]
* @param [in] ldb leading dimension of B
* @param [in] stride_b offset in elements between consecutive matrices of B
* @param [in] beta host or device scalar, if nullptr, the default value 0 will be used
* @param [inout] c_ptr column-major matrix of size [m, n]
* @param [in] ldc leading dimension of C
* @param [in] stride_c offset in elements between consecutive matrices of C
* @param [in] batch_count number of matrices in the batch
*/
template <bool DevicePointerMode = false, typename S, typename A, typename B, typename C>
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,
uint64_t batch_count)
{
common::nvtx::range<common::nvtx::domain::raft> batch_scope(
"linalg::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)",
m,
n,
k,
batch_count);
if (batch_count == 0) { return; }
matmul_key_t mm_key{
m, n, k, lda, ldb, ldc, trans_a, trans_b, batch_count, stride_a, stride_b, stride_c};
matmul_impl<DevicePointerMode, S, A, B, C>(
res, mm_key, alpha, a_ptr, b_ptr, beta, c_ptr, resource::get_cuda_stream(res));
}

} // namespace linalg::detail
} // namespace raft
45 changes: 44 additions & 1 deletion cpp/include/raft/linalg/detail/cusolver_wrappers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
*/

Expand Down Expand Up @@ -877,6 +877,49 @@ inline cusolverStatus_t cusolverDnpotrf(cusolverDnHandle_t handle, // NOLINT
}
/** @} */

/**
* @defgroup potrfBatched cusolver potrfBatched operations
* @{
*/
template <typename T>
cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT
cublasFillMode_t uplo,
int n,
T* Aarray[], // NOLINT
int lda,
int* infoArray,
int batchSize,
cudaStream_t stream);

template <>
inline cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT
cublasFillMode_t uplo,
int n,
float* Aarray[], // NOLINT
int lda,
int* infoArray,
int batchSize,
cudaStream_t stream)
{
RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream));
return cusolverDnSpotrfBatched(handle, uplo, n, Aarray, lda, infoArray, batchSize);
}

template <>
inline cusolverStatus_t cusolverDnpotrfBatched(cusolverDnHandle_t handle, // NOLINT
cublasFillMode_t uplo,
int n,
double* Aarray[], // NOLINT
int lda,
int* infoArray,
int batchSize,
cudaStream_t stream)
{
RAFT_CUSOLVER_TRY(cusolverDnSetStream(handle, stream));
return cusolverDnDpotrfBatched(handle, uplo, n, Aarray, lda, infoArray, batchSize);
}
/** @} */

/**
* @defgroup potrs cusolver potrs operations
* @{
Expand Down
Loading
Loading