Mitigate cuBLASLt 13.6 GEMM bug for inputs greater than 2^31 by divyegala · Pull Request #3098 · NVIDIA/raft

divyegala · 2026-07-29T17:43:25Z

coderabbitai · 2026-07-29T17:46:50Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

📝 Walkthrough

Walkthrough

Changes

GEMM workaround and regression coverage

Layer / File(s)	Summary
cuBLASLt 13.6 heuristic filtering `cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp`	Detects the affected FP32 configuration, requests multiple heuristic candidates, and excludes algorithm 68 when selecting a successful zero-workspace candidate.
GEMM regression tests and test wiring `cpp/tests/linalg/gemm_basic.cpp`, `cpp/tests/linalg/gemm_large.cpp`, `cpp/tests/CMakeLists.txt`, `build.sh`	Adds predicate and large-span GEMM tests and registers `GEMM_LARGE_TEST` for default builds.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: bdice

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the cuBLASLt 13.6 GEMM workaround for large input spans.
Description check	✅ Passed	The description is related to the changeset and links the fix to issue `#3097`.
Linked Issues check	✅ Passed	The code adds the requested cuBLASLt 13.6 workaround, preserves below-boundary behavior, and includes tests for the affected boundary and architectures.
Out of Scope Changes check	✅ Passed	The added build and test updates support the GEMM bug fix and do not appear unrelated to the issue.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

🧹 Nitpick comments (2)

cpp/tests/linalg/gemm_large.cpp (2)
74-84: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Validate an interior result too.

Checking only endpoints can miss a partial/corrupt result in the large output. Copy and assert at least one representative interior element, such as samples / 2.

As per path instructions, “Validate numerical correctness in the test beyond just edge elements as needed to avoid false positives.”
🤖 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_large.cpp` around lines 74 - 84, Extend the output
validation in the large GEMM test to copy and assert a representative interior
element, such as index samples / 2, in addition to the existing endpoint checks.
Update the output_edges storage and asynchronous cudaMemcpyAsync calls while
preserving synchronization and the existing zero-value assertions.
Source: Path instructions

46-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use RAFT owning mdarrays for the test buffers.

These are straightforward owned device buffers; replace rmm::device_uvector with suitable RAFT device vector/mdarray types and pass their data handles to GEMM.

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_large.cpp` around lines 46 - 48, Replace the three
rmm::device_uvector test buffers in the GEMM setup with suitable owning RAFT
device vector/mdarray types, preserving their element counts and stream
allocation. Update the GEMM calls to pass the RAFT containers’ data handles
instead of relying on rmm::device_uvector interfaces.
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.

Nitpick comments:
In `@cpp/tests/linalg/gemm_large.cpp`:
- Around line 74-84: Extend the output validation in the large GEMM test to copy
and assert a representative interior element, such as index samples / 2, in
addition to the existing endpoint checks. Update the output_edges storage and
asynchronous cudaMemcpyAsync calls while preserving synchronization and the
existing zero-value assertions.
- Around line 46-48: Replace the three rmm::device_uvector test buffers in the
GEMM setup with suitable owning RAFT device vector/mdarray types, preserving
their element counts and stream allocation. Update the GEMM calls to pass the
RAFT containers’ data handles instead of relying on rmm::device_uvector
interfaces.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 09766a0a-0585-481a-b532-56b3fb3e8258

📥 Commits

Reviewing files that changed from the base of the PR and between b9cf5e0 and 7852ceb.

📒 Files selected for processing (5)

build.sh
cpp/include/raft/linalg/detail/cublaslt_wrappers.hpp
cpp/tests/CMakeLists.txt
cpp/tests/linalg/gemm_basic.cpp
cpp/tests/linalg/gemm_large.cpp

csadorf

Have you considered constraining our dependency?

csadorf · 2026-07-29T18:06:45Z

+{
+  constexpr uint64_t max_safe_span = (uint64_t{1} << 31) - 1;
+  const auto a_columns             = args.trans_a ? args.m : args.k;
+  return version == 130600 && args.lda != 0 && a_columns > max_safe_span / args.lda;


Did you test other versions? Are we confident that only this very specific version is affected? What about future versions?

No, but the tests should help with that if it fails in a future version. If you prefer being defensive, then I can add a >= check instead?

I think we might have to be defensive here and actually set 130600 as a lower bound unless we get some kind of assurance that this will be fixed in the next patch release.

csadorf · 2026-07-29T18:14:37Z

+ * cuBLASLt 13.6 may select algorithm 68 once A's physical span reaches 2^31 elements, but that
+ * algorithm fails during execution for FP32. Prefer the validated alternatives.
+ */
+inline auto needs_cublaslt_13_6_workaround(const matmul_key_t& args, std::size_t version) noexcept


I'd recommend checking whether the work-around is needed on all archs or only sm_120/sm_121.

This will be a bit tough to verify this close to the release.

Yes, but I think we should at least check sm_90 since restricting the algorithm might lead to a performance regression.

Good point, let me skip the heuristic for non SM120/121 and CI tests should give us an indication for sm80 and sm90.

coderabbitai

Actionable comments posted: 1

🤖 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/tests/linalg/gemm_large.cpp`:
- Around line 30-32: Update the test setup around the device vector allocations
for a, b, and c to query available GPU memory and skip the test when the
required allocation, including a suitable overhead margin, cannot be satisfied.
Ensure the skip occurs before constructing any of these vectors and reports that
insufficient free device memory makes the test unsupported.

🪄 Autofix (Beta)

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: 22149fed-0883-4328-b9d7-95d09c8ab553

📥 Commits

Reviewing files that changed from the base of the PR and between 4b60632 and f083600.

📒 Files selected for processing (2)

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

coderabbitai

Actionable comments posted: 1

🤖 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/tests/linalg/gemm_basic.cpp`:
- Around line 205-216: Add an EXPECT_FALSE assertion for device architecture
(12, 2) in the GemmCublasLt136WorkaroundArchitectures test, alongside the
existing SM12 cases, using the same at_boundary and affected_version inputs.

🪄 Autofix (Beta)

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: b9fa7dc9-3460-4b7e-bfec-9c7b88be42f6

📥 Commits

Reviewing files that changed from the base of the PR and between f083600 and 32a0db4.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (1)

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

dantegd

One quick question, besides it lgtm

dantegd · 2026-07-29T19:23:48Z

+
+    constexpr int workaround_heuristic_results = 2;
+    std::array<cublasLtMatmulHeuristicResult_t, workaround_heuristic_results> heuristic_results{};
+    const int requested_results = use_cublaslt_13_6_workaround ? workaround_heuristic_results : 1;


Is it guaranteed or documented that these affected GEMMs always have a usable non-68 algorithm in the top two heuristic results?

No, but this only skips 68 if it's at the top of the heuristics. If nothing but 68 is returned we will have to fail.

cjnolet · 2026-07-29T22:58:58Z

/merge

The fix in #3098 was incomplete, as according to the cublas team there is no guarantee that `algo13` will be returned along with `algo68` (the failing version for A matrix with > 2^31 elements). Also, the cublas team will patch this bug in version `13.6.1` so we will automatically receive the fix by guarding only against `13.6.0`.

large gemm bug

7852ceb

divyegala requested review from a team as code owners July 29, 2026 17:43

divyegala requested a review from bdice July 29, 2026 17:43

divyegala added bug Something isn't working non-breaking Non-breaking change labels Jul 29, 2026

github-project-automation Bot added this to Unstructured Data Processing Jul 29, 2026

coderabbitai Bot reviewed Jul 29, 2026

View reviewed changes

ease up conditionals

f0d1caf

divyegala changed the base branch from main to release/26.08 July 29, 2026 17:49

remove insufficent mem check and skip

4b60632

csadorf reviewed Jul 29, 2026

View reviewed changes

simplify alternate algorithm selection

f083600

coderabbitai Bot reviewed Jul 29, 2026

View reviewed changes

Comment thread cpp/tests/linalg/gemm_large.cpp

workaround only for sm120 and 121

32a0db4

coderabbitai Bot reviewed Jul 29, 2026

View reviewed changes

Comment thread cpp/tests/linalg/gemm_basic.cpp

dantegd approved these changes Jul 29, 2026

View reviewed changes

add sm100 to workaround

cdb77da

gforsyth approved these changes Jul 29, 2026

View reviewed changes

defensively mitigate

e783775

cjnolet approved these changes Jul 29, 2026

View reviewed changes

cjnolet merged commit fcef0d3 into NVIDIA:release/26.08 Jul 29, 2026
80 checks passed

github-project-automation Bot moved this to Done in Unstructured Data Processing Jul 29, 2026

divyegala mentioned this pull request Jul 30, 2026

cublas team verified workaround for large GEMM algo68 bug #3100

Merged

Uh oh!

Conversation

divyegala commented Jul 29, 2026

Uh oh!

coderabbitai Bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

csadorf left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dantegd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

divyegala Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cjnolet commented Jul 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

coderabbitai Bot commented Jul 29, 2026 •

edited

Loading

divyegala Jul 29, 2026 •

edited

Loading