add-batched-linalg for cuvs - #3110
Conversation
Signed-off-by: Intron7 <severin.dicks@icloud.com>
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change adds strided-batched GEMM support through cuBLASLt, exposes ChangesBatched GEMM
Batched solver wrappers
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp (1)
311-322: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftAccount for batched A spans in the cuBLASLt 13.6 workaround.
- Include
(batch_count - 1) * stride_ain the span check. The current predicate can skip the workaround when the batch span exceeds 2^31 elements.- When
ldais incremented for the heuristic layout, adjuststride_aso consecutive matrices do not overlap. Add batched regression tests.🤖 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 311 - 322, The cuBLASLt 13.6 workaround must account for the full batched A span. Update the predicate controlling use_cublaslt_13_6_workaround to include (batch_count - 1) * stride_a, and when constructing heuristic_a via get_cublaslt_13_6_heuristic_args, adjust stride_a consistently with the incremented lda so batched matrices remain non-overlapping. Add regression coverage for batched spans exceeding the 2^31-element boundary.
🧹 Nitpick comments (2)
cpp/tests/linalg/gemm_batched.cpp (2)
24-27: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd batch-boundary and per-operand stride cases.
The suite fixes the batch size to three and varies only
x_batch_stride. A defect in zero- or one-batch handling, Y or Z stride translation, or Y broadcast handling will pass these tests.Add explicit zero- and one-batch cases. Specify the expected zero-batch behavior. Add independent padding coverage for X, Y, and Z. Add a Y broadcast case. Keep Z matrices separate because output matrices must not alias.
As per path instructions, public batched GEMM tests should cover “zero/edge batches” and “padded and broadcast strides.”
Also applies to: 222-248
🤖 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/tests/linalg/gemm_batched.cpp` around lines 24 - 27, Extend the batched GEMM tests around the existing kBatch cases to cover zero and one batch, asserting the API’s expected zero-batch behavior. Add independent padded-stride cases for X, Y, and Z, plus a Y-broadcast case, and keep each output matrix non-aliasing. Ensure coverage exercises per-operand stride translation rather than varying only x_batch_stride.Source: Path instructions
166-168: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse RAFT mdarrays for flat owning device buffers.
These allocations own contiguous flat device storage and only expose pointers to views. Replace them with one-dimensional
raft::device_mdarrayallocations created fromres.
cpp/tests/linalg/gemm_batched.cpp#L166-L168: Replace the X, Y, and Zrmm::device_uvectorallocations with RAFT mdarrays.cpp/tests/linalg/gemm_batched.cpp#L337-L337: Replace the invalid-stride test buffer with a RAFT mdarray.As per coding guidelines, “Prefer raft mdarray types for owning data over
rmm::device_uvector... when an mdarray fits the use case.”🤖 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/tests/linalg/gemm_batched.cpp` around lines 166 - 168, In cpp/tests/linalg/gemm_batched.cpp lines 166-168, replace the X, Y, and Z rmm::device_uvector allocations with one-dimensional raft::device_mdarray allocations constructed using res, and update their pointer access as needed. At line 337, replace the invalid-stride test buffer with a one-dimensional raft::device_mdarray created from res; no other owning-buffer types need to change.Source: Coding guidelines
🤖 Prompt for all review comments with 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.
Inline comments:
In `@cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp`:
- Around line 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.
In `@cpp/include/raft/linalg/detail/gemm.cuh`:
- Around line 51-58: Update the RAFT_FAIL message in the stride-validation
branch to state that valid matrices require one stride to equal 1 and the other
to be at least the matching matrix extent, covering overlapping cases such as
row_stride == 1 with col_stride < rows.
In `@cpp/include/raft/linalg/gemm.cuh`:
- Around line 337-358: Validate the output operand in the batched GEMM path so
`Z` cannot use a zero batch stride when `batch_count` exceeds one; retain
support for zero stride with a single batch. Update the nearby Doxygen
documentation to state that zero batch stride is allowed for inputs but
disallowed for multi-batch output `Z`, and anchor the guard after `batch_count`
is computed using `z_desc.batch_stride`.
---
Outside diff comments:
In `@cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp`:
- Around line 311-322: The cuBLASLt 13.6 workaround must account for the full
batched A span. Update the predicate controlling use_cublaslt_13_6_workaround to
include (batch_count - 1) * stride_a, and when constructing heuristic_a via
get_cublaslt_13_6_heuristic_args, adjust stride_a consistently with the
incremented lda so batched matrices remain non-overlapping. Add regression
coverage for batched spans exceeding the 2^31-element boundary.
---
Nitpick comments:
In `@cpp/tests/linalg/gemm_batched.cpp`:
- Around line 24-27: Extend the batched GEMM tests around the existing kBatch
cases to cover zero and one batch, asserting the API’s expected zero-batch
behavior. Add independent padded-stride cases for X, Y, and Z, plus a
Y-broadcast case, and keep each output matrix non-aliasing. Ensure coverage
exercises per-operand stride translation rather than varying only
x_batch_stride.
- Around line 166-168: In cpp/tests/linalg/gemm_batched.cpp lines 166-168,
replace the X, Y, and Z rmm::device_uvector allocations with one-dimensional
raft::device_mdarray allocations constructed using res, and update their pointer
access as needed. At line 337, replace the invalid-stride test buffer with a
one-dimensional raft::device_mdarray created from res; no other owning-buffer
types need to change.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8a05dd80-4ce9-4fd6-8005-5d50dd74c3d9
📒 Files selected for processing (7)
cpp/include/raft/linalg/detail/cublas_wrappers.hppcpp/include/raft/linalg/detail/cublaslt_wrappers.hppcpp/include/raft/linalg/detail/cusolver_wrappers.hppcpp/include/raft/linalg/detail/gemm.cuhcpp/include/raft/linalg/gemm.cuhcpp/tests/CMakeLists.txtcpp/tests/linalg/gemm_batched.cpp
| /** | ||
| * 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; | ||
| } |
There was a problem hiding this comment.
🎯 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:
- 1: preferred blas library; cublaslt gemm implementation pytorch/pytorch#122106
- 2: https://github.com/ShawnZhong/CUDA-Programming-Starter-Kit/blob/master/include/CUDA-10.2/cublasLt.h
- 3: https://github.com/openxla/xla/blob/3f96d2c8/xla/stream_executor/cuda/cuda_blas_lt.cc
🏁 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.hppRepository: 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}")
PYRepository: 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
| RAFT_FAIL( | ||
| "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the " | ||
| "matrix strides are [%zu, %zu], one of which must be 1", | ||
| name, | ||
| static_cast<size_t>(rows), | ||
| static_cast<size_t>(cols), | ||
| static_cast<size_t>(row_stride), | ||
| static_cast<size_t>(col_stride)); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the failure message for the overlapping-stride case.
The else branch also triggers when one stride equals 1 but the other stride is smaller than the matching extent, for example row_stride == 1 && col_stride < rows. The text "one of which must be 1" then misleads the caller, because a stride of 1 is present. State both conditions.
📝 Proposed message
RAFT_FAIL(
- "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the "
- "matrix strides are [%zu, %zu], one of which must be 1",
+ "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the "
+ "matrix strides are [%zu, %zu]; one stride must be 1 and the other must be at least the "
+ "size of the other dimension",📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| RAFT_FAIL( | |
| "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the " | |
| "matrix strides are [%zu, %zu], one of which must be 1", | |
| name, | |
| static_cast<size_t>(rows), | |
| static_cast<size_t>(cols), | |
| static_cast<size_t>(row_stride), | |
| static_cast<size_t>(col_stride)); | |
| RAFT_FAIL( | |
| "%s is not a batch of row- or column-major matrices: with extents [batch, %zu, %zu] the " | |
| "matrix strides are [%zu, %zu]; one stride must be 1 and the other must be at least the " | |
| "size of the other dimension", | |
| name, | |
| static_cast<size_t>(rows), | |
| static_cast<size_t>(cols), | |
| static_cast<size_t>(row_stride), | |
| static_cast<size_t>(col_stride)); |
🤖 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/gemm.cuh` around lines 51 - 58, Update the
RAFT_FAIL message in the stride-validation branch to state that valid matrices
require one stride to equal 1 and the other to be at least the matching matrix
extent, covering overlapping cases such as row_stride == 1 with col_stride <
rows.
| RAFT_EXPECTS(x.extent(0) == z.extent(0) && y.extent(0) == z.extent(0), | ||
| "Batch sizes of X, Y and Z should be equal"); | ||
| RAFT_EXPECTS(x.extent(1) == z.extent(1), "Number of rows of X and Z should be equal"); | ||
| RAFT_EXPECTS(y.extent(2) == z.extent(2), "Number of columns of Y and Z should be equal"); | ||
| RAFT_EXPECTS(x.extent(2) == y.extent(1), "Number of columns of X and rows of Y should be equal"); | ||
|
|
||
| const auto x_desc = detail::describe_batched_gemm_operand(x, "X"); | ||
| const auto y_desc = detail::describe_batched_gemm_operand(y, "Y"); | ||
| const auto z_desc = detail::describe_batched_gemm_operand(z, "Z"); | ||
|
|
||
| // NB: the function type constraints only ever allow two view types, so using std::is_same_v is | ||
| // fine | ||
| constexpr auto kDeviceMode = | ||
| std::is_same_v<ScalarViewType, raft::device_scalar_view<ValueType, ScalarIdxType>>; | ||
|
|
||
| // NB: we rely on the implementation of detail::matmul_strided_batched to set defaults | ||
| ValueType* alpha_ptr = nullptr; | ||
| ValueType* beta_ptr = nullptr; | ||
| if (alpha.has_value()) { alpha_ptr = alpha.value().data_handle(); } | ||
| if (beta.has_value()) { beta_ptr = beta.value().data_handle(); } | ||
|
|
||
| const auto batch_count = static_cast<uint64_t>(z.extent(0)); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject an overlapping batch stride for the output operand Z.
describe_batched_gemm_operand leaves the batch stride unconstrained, which is correct for the inputs. Z is the output. If z_desc.batch_stride is 0 and batch_count is greater than 1, every batch index writes the same M x N region concurrently, so the result is nondeterministic. The beta .* Z term also reads memory that other batch indices write. The Doxygen text at lines 304-306 offers a zero batch stride as a documented feature and does not restrict it to inputs, so callers can reach this state through the public contract.
Validate the output stride, and state the restriction in the Doxygen block.
🛡️ Proposed guard
const auto batch_count = static_cast<uint64_t>(z.extent(0));
+
+ if (batch_count > 1) {
+ const auto z_span = static_cast<int64_t>(z_desc.ld) *
+ static_cast<int64_t>(z_desc.col_major ? z.extent(2) : z.extent(1));
+ RAFT_EXPECTS(z_desc.batch_stride >= z_span,
+ "Z batch stride %zd overlaps consecutive output matrices (matrix span %zd); "
+ "broadcasting is only supported for the input operands",
+ static_cast<ptrdiff_t>(z_desc.batch_stride),
+ static_cast<ptrdiff_t>(z_span));
+ }🤖 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/gemm.cuh` around lines 337 - 358, Validate the output
operand in the batched GEMM path so `Z` cannot use a zero batch stride when
`batch_count` exceeds one; retain support for zero stride with a single batch.
Update the nearby Doxygen documentation to state that zero batch stride is
allowed for inputs but disallowed for multi-batch output `Z`, and anchor the
guard after `batch_count` is computed using `z_desc.batch_stride`.
|
It looks like this PR is meant to add the same functionality as #3106 , which is almost ready to merge. Please close this PR if this is the case or clarify the difference otherwise |
|
@achirkin more than happy to close. I was asked to add this in support of NVIDIA/cuvs#2248 if #3106 handles this |
3 participants
This adds batched linalg function for cuvs consumption in GMM @cjnolet