[REVIEW] Addition of gemm_strided_batched - #3106
Open
landrumb wants to merge 2 commits into
Open
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe change adds cuBLASLt strided-batched matrix multiplication support, exposes it through Strided-batched GEMM
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 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: 2
🧹 Nitpick comments (1)
cpp/tests/linalg/gemm_basic.cpp (1)
188-212: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtend the batched coverage to the remaining argument combinations.
The test covers only the non-transposed, default-coefficient, host-pointer path. The following cases stay untested and each one exercises distinct code in
matmul_strided_batched:
trans_a/trans_bset to true, which changes the layout row/column swap and thelda/ldbrequirements.- Explicit non-null
alphaandbeta, which verifies that existingCpadding and values are combined correctly.DevicePointerMode = true, which exercises thecoef_wrapperdevice allocation path.- A zero batch stride, which is the cuBLASLt broadcast case.
Do you want me to generate these additional test cases?
As per path instructions: "Tests should use synthetic data, validate numerical results and padding, and cover important edge cases."
🤖 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_basic.cpp` around lines 188 - 212, Extend the batched GEMM coverage around gemm_strided_batched to add synthetic cases for transposed A/B operands with corresponding lda/ldb values, non-null alpha and beta while validating C padding and accumulation, DevicePointerMode enabled to exercise device coefficient allocation, and zero batch strides for broadcast behavior. Assert numerical outputs and untouched padding for each case while preserving the existing baseline test.Source: Path instructions
🤖 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 378-416: Update matmul_strided_batched to validate positive m, n,
k, and batch_count; require lda and ldc to meet the row requirements for the
configured layouts and ldb to meet the corresponding B-matrix requirement;
validate non-null matrix and scalar pointers and permitted pointer locations
before constructing cuBLASLt layouts. Validate strides for invalid values while
explicitly allowing zero strides for broadcasting, and reject any precondition
failure before forwarding inputs to cuBLASLt.
- Around line 417-432: Update the strided-batched matmul flow around
cublasLtMatmul to create or reuse a valid initialized cublasLtMatmulAlgo_t
instead of passing nullptr. Cache heuristic algorithm results and their
descriptors using strides, batch_count, and compute_type in the key, then
configure and pass workspace when supported while preserving the existing
operation and coefficient parameters.
---
Nitpick comments:
In `@cpp/tests/linalg/gemm_basic.cpp`:
- Around line 188-212: Extend the batched GEMM coverage around
gemm_strided_batched to add synthetic cases for transposed A/B operands with
corresponding lda/ldb values, non-null alpha and beta while validating C padding
and accumulation, DevicePointerMode enabled to exercise device coefficient
allocation, and zero batch strides for broadcast behavior. Assert numerical
outputs and untouched padding for each case while preserving the existing
baseline test.
🪄 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: 821297de-e419-4c57-93c2-210c59d3c380
📒 Files selected for processing (3)
cpp/include/raft/linalg/detail/cublaslt_wrappers.hppcpp/include/raft/linalg/gemm.cuhcpp/tests/linalg/gemm_basic.cpp
achirkin
requested changes
Aug 11, 2026
achirkin
left a comment
achirkin
left a comment
Contributor
There was a problem hiding this comment.
Thanks for implementing this! Let's refine the API (see below)
Comment on lines
+318
to
+338
| template <typename A_t, typename B_t, typename C_t, typename S_t, bool DevicePointerMode = false> | ||
| 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<S_t, A_t, B_t, C_t>()) |
Contributor
There was a problem hiding this comment.
Please rewrite the public API to use mdspans - we should not add new legacy overloads to public functions.
Also note, we can encode the _strided_ part as in the layout policy - it doesn't need to be a part of the name; I'd suggest we make it here generic and pass the stride = matrix size if the layout is not strided.
Suggested change
| template <typename A_t, typename B_t, typename C_t, typename S_t, bool DevicePointerMode = false> | |
| 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<S_t, A_t, B_t, C_t>()) | |
| 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, | |
| cublasComputeType_t compute_type_override = detail::get_matmul_type<ValueType, ValueType, ValueType, ValueType>())) |
| uint64_t ldc, | ||
| int64_t stride_c, | ||
| int32_t batch_count, | ||
| cublasComputeType_t compute_type = detail::get_matmul_type<S_t, A_t, B_t, C_t>()) |
Contributor
There was a problem hiding this comment.
Also I like the way you suggest to optionally pass the cublas compute type override - looks like a good compromise between the user-friendliness and flexibility.
However, maybe we could sacrifice a bit of brevity add replace the optional default argument with two overloads (with/without argument) - to avoid detail namespace from appearing in the public headers?
achirkin
added
feature request
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds the wrapper for strided and batched matrix multiplication with cublasLt requested by @achirkin in the reviews for NVIDIA/cuvs#2352 , where this operation is used for batched row-wise dot products. TF32 support is included.
The associated test covers most of the functionality besides eg transposed inputs, but can be expanded if reviewer(s) prefer.