[REVIEW] Addition of gemm_strided_batched

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

@@ -167,6 +167,23 @@ struct cublastlt_matrix_layout {
     return cublastlt_matrix_layout{
       get_cuda_data_type<T>(), col_major ? rows : cols, col_major ? cols : rows, ld};
   }
+
+  template <typename T>
+  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<T>(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 <typename S, typename A, typename B, typename C, bool DevicePointerMode = false>
-  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<S, A, B, C>())
+    -> cublastlt_matmul_desc
   {
-    auto desc = cublastlt_matmul_desc{get_matmul_type<S, A, B, C>(), get_cuda_data_type<S>()};
+    auto desc = cublastlt_matmul_desc{compute_type, get_cuda_data_type<S>()};
     if constexpr (DevicePointerMode) {
       const cublasPointerMode_t mode = CUBLAS_POINTER_MODE_DEVICE;
       RAFT_CUBLAS_TRY(cublasLtMatmulDescSetAttribute(
@@ -348,6 +368,70 @@ struct coef_wrapper<true, S> {
   }
 };
 
+/**
+ * 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 <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,
+                            int32_t batch_count,
+                            cublasComputeType_t compute_type)
+{
+  common::nvtx::range<common::nvtx::domain::raft> batch_scope(
+    "linalg::detail::matmul_strided_batched(m = %d, n = %d, k = %d, batch_count = %d)",
+    m,
+    n,
+    k,
+    batch_count);
+
+  auto operation = cublastlt_matmul_desc::for_matmul<S, A, B, C, DevicePointerMode>(
+    trans_a, trans_b, compute_type);
+
+  auto a_layout = cublastlt_matrix_layout::for_strided_batched_matmul<A>(
+    !trans_a, m, k, lda, batch_count, stride_a);
+  auto b_layout = cublastlt_matrix_layout::for_strided_batched_matmul<B>(
+    !trans_b, k, n, ldb, batch_count, stride_b);
+  auto c_layout =
+    cublastlt_matrix_layout::for_strided_batched_matmul<C>(true, m, n, ldc, batch_count, stride_c);
+
+  auto stream = resource::get_cuda_stream(res);
+  coef_wrapper<DevicePointerMode, S> 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

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,6 +283,186 @@ void gemm(raft::resources const& res,
   }
 }
 
+/**
+ * @brief Batched matrix multiplication with an explicit 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
+ * @param[in] compute_type_override cublasLt compute type
+ */
+template <typename ValueType,
+          typename IndexType,
+          typename LayoutPolicyX,
+          typename LayoutPolicyY,
+          typename LayoutPolicyZ,
+          typename ScalarIdxType  = std::uint32_t,
+          typename ScalarViewType = raft::host_scalar_view<ValueType, ScalarIdxType>,
+          typename                = std::enable_if_t<std::disjunction_v<
+                           std::is_same<ScalarViewType, raft::host_scalar_view<ValueType, ScalarIdxType>>,
+                           std::is_same<ScalarViewType, raft::device_scalar_view<ValueType, ScalarIdxType>>>>>
+void gemm_batched(raft::resources const& res,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyX> x,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyY> y,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyZ> z,
+                  std::optional<ScalarViewType> alpha,
+                  std::optional<ScalarViewType> beta,
+                  cublasComputeType_t compute_type_override)
+{
+  static_assert(!std::is_same_v<LayoutPolicyX, raft::layout_left> &&
+                  !std::is_same_v<LayoutPolicyY, raft::layout_left> &&
+                  !std::is_same_v<LayoutPolicyZ, raft::layout_left>,
+                "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<decltype(view)>::layout_type;
+    if constexpr (std::is_same_v<layout_type, raft::layout_stride>) {
+      return static_cast<int64_t>(view.stride(0));
+    } else {
+      return static_cast<int64_t>(view.extent(1)) * static_cast<int64_t>(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<ScalarViewType, raft::device_scalar_view<ValueType, ScalarIdxType>>;
+  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 batch_count = static_cast<int32_t>(z.extent(0));
+
+  if (z_col_major) {
+    return detail::matmul_strided_batched<kDeviceMode>(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<kDeviceMode>(res,
+                                                     y_col_major,
+                                                     x_col_major,
+                                                     z.extent(2),
+                                                     z.extent(1),
+                                                     x.extent(2),
+                                                     alpha_ptr,
+                                                     y.data_handle(),
+                                                     y_ld,
+                                                     batch_stride(y),
+                                                     x.data_handle(),
+                                                     x_ld,
+                                                     batch_stride(x),
+                                                     beta_ptr,
+                                                     z.data_handle(),
+                                                     z_ld,
+                                                     batch_stride(z),
+                                                     batch_count,
+                                                     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 ValueType,
+          typename IndexType,
+          typename LayoutPolicyX,
+          typename LayoutPolicyY,
+          typename LayoutPolicyZ,
+          typename ScalarIdxType  = std::uint32_t,
+          typename ScalarViewType = raft::host_scalar_view<ValueType, ScalarIdxType>,
+          typename                = std::enable_if_t<std::disjunction_v<
+                           std::is_same<ScalarViewType, raft::host_scalar_view<ValueType, ScalarIdxType>>,
+                           std::is_same<ScalarViewType, raft::device_scalar_view<ValueType, ScalarIdxType>>>>>
+void gemm_batched(raft::resources const& res,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyX> x,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyY> y,
+                  raft::device_mdspan<ValueType, raft::extent_3d<IndexType>, LayoutPolicyZ> z,
+                  std::optional<ScalarViewType> alpha = std::nullopt,
+                  std::optional<ScalarViewType> beta  = std::nullopt)
+{
+  return gemm_batched(res,
+                      x,
+                      y,
+                      z,
+                      alpha,
+                      beta,
+                      detail::get_matmul_type<ValueType, ValueType, ValueType, ValueType>());
+}
+
 /** @} */  // end of gemm
 
 }  // namespace linalg