[BUG] RAFT GEMM fails at a 2^31-element input span with cuBLASLt 13.6

[csadorf]

raft::linalg::gemm fails with CUBLAS_STATUS_EXECUTION_FAILED when the nominal input span reaches exactly 2^31 elements under cuBLASLt 13.6. The adjacent below-boundary case succeeds.

The reproducer uses only RAFT and CUDA runtime allocations. It does not depend on a downstream RAPIDS library or Python package.

Boundary results

Case	Nominal input span	Selected algorithm	Result
134,217,727 × 16	2^31 - 16 elements	13, option 62	Pass
134,217,728 × 16	exactly 2^31 elements	68	CUBLAS_STATUS_EXECUTION_FAILED

The GEMM descriptors reported by CUBLASLT_LOG_LEVEL=5 for the failing case are:

A: type=R_32F rows=2 cols=134217728 ld=16
B: type=R_32F rows=1 cols=2 ld=1
C: type=R_32F rows=134217728 cols=1 ld=134217728
D: type=R_32F rows=134217728 cols=1 ld=134217728
computeType=COMPUTE_32F
transa=OP_T
transb=OP_T
algoId=68
workSpaceSizeInBytes=0

At 134,217,727 rows, cuBLASLt 13.6 selects algorithm 13 with customOption=62, and the operation succeeds. At 134,217,728 rows, the heuristic selects algorithm 68, which is returned as supported but fails during cublasLtMatmul.

The exact power-of-two boundary suggests a 32-bit span/index limitation in algorithm 68 or its heuristic eligibility logic.

Environment

• Image: rapidsai/base@sha256:7d3527573dcfec08e119774118f9edfaa96b542bf0f58d168ea883bd7356df03
• libraft: 26.08.00a59, build cuda13_260729_eec9ba69
• RAFT revision: eec9ba69be3e7735306bd7779204d2265bd53f37
• libcublas: 13.6.0.2
• cuda-cudart: 13.3.29
• GPU: NVIDIA RTX PRO 6000 Blackwell Workstation Edition
• Driver: 595.84
• OS: Ubuntu 24.04 x86_64

The reproducer requires approximately 9 GiB of available device memory.

Steps to reproduce

1. Save the source below as reproducer.cu.

RAFT-only reproducer source

#include <raft/core/resource/cuda_stream.hpp>
#include <raft/core/resources.hpp>
#include <raft/linalg/gemm.cuh>

#include <cuda_runtime_api.h>

#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>

namespace {

void check_cuda(cudaError_t status, char const* operation)
{
  if (status != cudaSuccess) {
    throw std::runtime_error(std::string{operation} + ": " +
                             cudaGetErrorString(status));
  }
}

}  // namespace

int main(int argc, char** argv)
{
  constexpr std::int64_t default_n_samples = 134217728;
  constexpr int lda                       = 16;
  constexpr int n_informative             = 2;
  constexpr int n_targets                 = 1;

  const auto n_samples =
    argc > 1 ? static_cast<std::int64_t>(std::stoll(argv[1]))
             : default_n_samples;

  if (n_samples <= 0 || n_samples > std::numeric_limits<int>::max()) {
    std::cerr << "n_samples must be in [1, INT_MAX]\n";
    return EXIT_FAILURE;
  }

  const auto a_elements =
    static_cast<std::uint64_t>(n_samples) * lda;

  std::cout << "n_samples=" << n_samples
            << ", lda=" << lda
            << ", nominal_A_elements=" << a_elements << '\n';

  float* a = nullptr;
  float* b = nullptr;
  float* c = nullptr;

  try {
    check_cuda(
      cudaMalloc(reinterpret_cast<void**>(&a),
                 a_elements * sizeof(float)),
      "cudaMalloc(A)");

    check_cuda(
      cudaMalloc(reinterpret_cast<void**>(&b),
                 n_informative * sizeof(float)),
      "cudaMalloc(B)");

    check_cuda(
      cudaMalloc(reinterpret_cast<void**>(&c),
                 n_samples * sizeof(float)),
      "cudaMalloc(C)");

    raft::resources resources;
    const auto stream = raft::resource::get_cuda_stream(resources);

    check_cuda(
      cudaMemsetAsync(a, 0, a_elements * sizeof(float), stream.value()),
      "cudaMemsetAsync(A)");

    check_cuda(
      cudaMemsetAsync(b, 0, n_informative * sizeof(float), stream.value()),
      "cudaMemsetAsync(B)");

    check_cuda(
      cudaMemsetAsync(c, 0, n_samples * sizeof(float), stream.value()),
      "cudaMemsetAsync(C)");

    const float alpha = 1.0f;
    const float beta  = 0.0f;

    // Exact layout:
    //
    // A: rows=2, cols=n_samples, ld=16, OP_T
    // B: rows=1, cols=2,         ld=1,  OP_T
    // C: rows=n_samples, cols=1, ld=n_samples
    raft::linalg::gemm(resources,
                       true,
                       true,
                       static_cast<int>(n_samples),
                       n_targets,
                       n_informative,
                       &alpha,
                       a,
                       lda,
                       b,
                       n_targets,
                       &beta,
                       c,
                       static_cast<int>(n_samples),
                       stream.value());

    check_cuda(
      cudaStreamSynchronize(stream.value()),
      "cudaStreamSynchronize");

    std::cout << "PASS\n";
  } catch (std::exception const& error) {
    std::cerr << "FAIL: " << error.what() << '\n';
    cudaFree(c);
    cudaFree(b);
    cudaFree(a);
    return EXIT_FAILURE;
  }

  check_cuda(cudaFree(c), "cudaFree(C)");
  check_cuda(cudaFree(b), "cudaFree(B)");
  check_cuda(cudaFree(a), "cudaFree(A)");

  return EXIT_SUCCESS;
}

2. Build and run both boundary cases in the affected image:

docker run --rm --gpus all --user root \
  -v "$PWD/reproducer.cu:/tmp/reproducer.cu:ro" \
  --entrypoint bash \
  rapidsai/base@sha256:7d3527573dcfec08e119774118f9edfaa96b542bf0f58d168ea883bd7356df03 \
  -lc '
    set -eux

    apt-get update
    apt-get install -y --no-install-recommends g++

    nvcc -std=c++20 -O2 \
      /tmp/reproducer.cu \
      -o /tmp/reproducer \
      -ccbin /usr/bin/g++ \
      -I/opt/conda/targets/x86_64-linux/include/cccl \
      -I/opt/conda/include \
      -L/opt/conda/lib \
      -Xlinker -rpath \
      -Xlinker /opt/conda/lib \
      -lraft \
      -lrmm \
      -lcublasLt \
      -lcublas \
      -lcudart

    CUBLASLT_LOG_LEVEL=5 /tmp/reproducer 134217727
    CUBLASLT_LOG_LEVEL=5 /tmp/reproducer 134217728
  '

Actual behavior

n_samples=134217727, lda=16, nominal_A_elements=2147483632
PASS

n_samples=134217728, lda=16, nominal_A_elements=2147483648
FAIL: cuBLAS error encountered at:
.../raft/linalg/detail/cublaslt_wrappers.hpp line=299:
call='cublasLtMatmul(...)',
Reason=13:CUBLAS_STATUS_EXECUTION_FAILED

Expected behavior

Both valid GEMM configurations should succeed. If algorithm 68 cannot support this layout, it should not be returned by cublasLtMatmulAlgoGetHeuristic.

Additional controls

The same exact-boundary operation succeeds with cuBLASLt 13.1.1. That version selects algorithm 13 with customOption=14.

RAFT currently requests one heuristic result and executes it without an alternative-algorithm fallback:

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

Lines 192 to 201 in eec9ba6

	RAFT_CUBLAS_TRY(cublasLtMatmulAlgoGetHeuristic(resource::get_cublaslt_handle(res),
	r.desc,
	r.a,
	r.b,
	r.c,
	r.c,
	preference,
	1,
	&r.heuristics,
	&algo_count));

The primary defect appears to be in cuBLASLt 13.6, but RAFT could potentially mitigate it by requesting alternative algorithms, falling back to classic cuBLAS for large physical spans, or chunking the operation.

[csadorf]

I can also reproduce this on an NVIDIA DGX Spark, so the failure is not limited to the RTX PRO 6000/x86_64 configuration:

• GPU: NVIDIA GB10, compute capability 12.1
• CPU: 20-core Arm (10× Cortex-X925 and 10× Cortex-A725)
• Architecture: aarch64
• Memory: 128 GB unified memory
• OS: Ubuntu 24.04.4 LTS
• Driver: 580.142
• CUDA: 13.0

The reproducer shows the same boundary behavior: 134217727 × 16 passes, while 134217728 × 16 fails with CUBLAS_STATUS_EXECUTION_FAILED.