🐛 Describe the bug

In the CUDA linalg.lu_factor_ex / linalg.solve_ex path, the looped cuSOLVER LU backend allocates cuSOLVER GETRF workspace once per matrix in a batch.

Current flow:

• lu_factor_looped_cusolver loops over batch items and calls at::cuda::solver::getrf(...) for each matrix: https://github.com/pytorch/pytorch/blob/37bc32b8736cca7afb41d4794e32ed65bbc83521/aten/src/ATen/native/cuda/linalg/BatchLinearAlgebraLib.cpp#L1647-L1675
• The getrf wrapper queries lwork and allocates workspace via CUDACachingAllocator::allocate(...) inside each call: https://github.com/pytorch/pytorch/blob/37bc32b8736cca7afb41d4794e32ed65bbc83521/aten/src/ATen/native/cuda/linalg/CUDASolver.cpp#L23-L31

That means a batch of size B performs B allocator allocate/free cycles for cuSOLVER workspace. With CUDA caching enabled these may usually hit the cache after warmup, but cache misses or PYTORCH_NO_CUDA_MEMORY_CACHING=1 turn this into real cudaMalloc / cudaFree calls. This also makes the allocator overhead visible at dispatch boundaries where PyTorch switches from cuBLAS batched LU to looped cuSOLVER LU.

This does not appear to be a cuSOLVER-internal allocation requirement: cusolverDn<t>getrf takes a caller-provided workspace pointer, and a standalone cuSOLVER benchmark that reused one workspace across the loop did not show hidden CUDA allocation calls.

Repro

Build without MAGMA, or otherwise force the cuSOLVER/cuBLAS CUDA linalg path, and run:

import torch

torch.cuda.set_device(0)
A = torch.randn((16, 512, 512), device="cuda")
torch.cuda.synchronize()
LU, pivots, info = torch.linalg.lu_factor_ex(A, check_errors=False)
torch.cuda.synchronize()
print(LU.shape, pivots.shape, info[:3].tolist())

To make allocator calls explicit, run with PYTORCH_NO_CUDA_MEMORY_CACHING=1 and trace CUDA API calls, or use an LD_PRELOAD interposer that counts cudaMalloc / cudaFree.

Observed locally on main commit 37bc32b8736cca7afb41d4794e32ed65bbc83521, CUDA 13.2, RTX 6000 Ada, USE_MAGMA=0:

The 512 case includes one cuSOLVER workspace allocation per batch item from the PyTorch wrapper.

Expected behavior

lu_factor_looped_cusolver should allocate GETRF workspace once per lu_factor call and reuse it across all matrices in the batch.

One possible implementation direction:

• expose a cuSOLVER GETRF buffer-size helper, or add a getrf wrapper overload that accepts caller-owned workspace;
• query lwork once for the fixed (m, n, dtype) call;
• allocate the workspace once in lu_factor_looped_cusolver;
• pass the same workspace to each per-matrix cusolverDn<t>getrf call.

The legacy cuSOLVER buffer-size API takes the matrix pointer, so the implementation can conservatively query using the first matrix pointer. For a fixed dtype and (m, n, lda), the required workspace should not vary across batch items.

Additional context

The allocator behavior is especially visible because the no-MAGMA / cuSOLVER-preferred LU heuristic currently switches batched square matrices from cuBLAS getrfBatched to looped cuSOLVER at M >= 512:

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

That threshold is a separate performance heuristic issue, but hoisting/reusing cuSOLVER workspace is independently useful and should reduce allocator traffic for all batched looped-cuSOLVER LU calls.

Versions

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

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