pytorch - 💡(How to fix) Fix CUDA linalg: replace hard LU M >= 512 cuSOLVER cutoff with batch-aware heuristic [1 participants]

🐛 Describe the bug

The CUDA LU factorization heuristic for the cuBLAS/cuSOLVER path uses a hard matrix-size cutoff for batched square inputs:

https://github.com/pytorch/pytorch/blob/37bc32b8736cca7afb41d4794e32ed65bbc83521/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp#L877-L882

if (m != n || (batch_size == 1 || m >= 512)) {
  lu_factor_looped_cusolver(input, pivots, infos, compute_pivots);
} else {
  lu_factor_batched_cublas(input, pivots, infos, compute_pivots);
}

For batched square M x M inputs this means:

• M < 512: cuBLAS getrfBatched
• M >= 512: looped cuSOLVER getrf

This cutoff is too coarse. It is good for small batches, but for larger batches cuBLAS batched remains substantially faster beyond M=512. The result is a visible performance cliff in torch.linalg.lu_factor_ex and in callers such as torch.linalg.solve_ex.

This is not a cuBLAS API limit. NVIDIA cuBLAS docs have long documented that cublas<t>getrfBatched supports arbitrary dimension, while describing it as intended for small matrices. The right cutoff is a performance heuristic, and it depends strongly on batch size. It may also depend on dtype, CUDA version, and GPU generation.

Repro

Build without MAGMA, or otherwise exercise the cuBLAS/cuSOLVER CUDA linalg path, then run:

import torch

torch.cuda.set_device(0)
for batch in (8, 16, 32, 64, 128):
    for n in (511, 512):
        A = torch.randn((batch, n, n), device="cuda", dtype=torch.float32)
        B = torch.randn((batch, n, 1), device="cuda", dtype=torch.float32)
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        torch.linalg.solve_ex(A, B, check_errors=False)
        end.record()
        torch.cuda.synchronize()
        print(batch, n, start.elapsed_time(end), "ms")

Local results

Environment:

• PyTorch 2.13.0a0+git37bc32b
• Commit 37bc32b8736cca7afb41d4794e32ed65bbc83521
• CUDA runtime 13.2
• GPU: NVIDIA RTX 6000 Ada Generation
• Build option relevant to repro: USE_MAGMA=0

torch.linalg.solve_ex(A, B, check_errors=False), B.shape == (batch, M, 1):

The cliff starts at M=512 because the LU factorization switches from cuBLAS batched to looped cuSOLVER. For larger batch sizes this is the wrong backend choice.

A standalone CUDA float32 GETRF sweep on the same machine produced these approximate crossover regions:

The exact table above should not be blindly hard-coded from one Ada GPU and one dtype, but it shows that the current scalar cutoff is not robust. A simple batch-aware heuristic, or a small benchmark-derived table, would avoid large regressions for common batched sizes.

Expected behavior

The cuBLAS/cuSOLVER LU heuristic should account for batch size when choosing between cuBLAS getrfBatched and looped cuSOLVER getrf for batched square matrices.

Before changing the heuristic, we should benchmark at least:

• dtypes: float32, float64, complex64, complex128
• batch sizes: small and large, e.g. 1, 2, 4, 8, 16, 32, 64, 128
• matrix sizes around the crossover bands, especially 511, 512, 513
• GPU generations: at least one Ampere/Hopper/Blackwell if available, plus the Ada result above
• CUDA versions currently supported in CI/binaries, not only CUDA 13.2
• lu_factor_ex directly and a representative caller such as solve_ex

Related

The looped cuSOLVER path also allocates GETRF workspace once per matrix through the PyTorch wrapper, which is tracked separately in:

https://github.com/pytorch/pytorch/issues/181997

That allocator issue is independent but makes the hard M >= 512 switch more visible.

cc @jerryzh168 @ptrblck @msaroufim @eqy @tinglvv @nWEIdia @csarofeen @jianyuh @nikitaved @mruberry @walterddr @xwang233 @Lezcano