Calling cublas primitives (Batched QR)

[lachinov]

Hi,

when calling
Q, R = jax.lax.linalg.qr(A, full_matrices=False, pivoting=False, use_magma=False)
with sufficiently large but not gigantic matrices, e.g. 4x1040x256, the code compiles into cusolve call
%1472:2 = stablehlo.custom_call @cusolver_geqrf_ffi(%1471) {backend_config = "", mhlo.backend_config = {}, mhlo.frontend_attributes = {num_batch_dims = "1"}, operand_layouts = [dense<[1, 2, 0]> : tensor<3xindex>], output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [0], operand_index = 0, operand_tuple_indices = []>], result_layouts = [dense<[1, 2, 0]> : tensor<3xindex>, dense<[1, 0]> : tensor<2xindex>]} : (tensor<4x1040x256xf32>) -> (tensor<4x1040x256xf32>, tensor<4x256xf32>) - validated on the local workstation.

cusolver geqr ignores batching, as I understand.

Manually implementing QR using Householder reflections (as in cublas) improves performance by a factor of 100 on a100 with larger matrices, e.g 64x32768x256. However, there's already an optimized version implemented in cublas, namely geqrfBatched, which would be a better fit.

Is it possible to ask jax or xla compiler to use cublas backend for QR decomposition?

[dfm]

I think that JAX does actually call geqrfBatched for batched systems:

jax/jaxlib/gpu/solver_kernels_ffi.cc

Lines 299 to 305 in ea22c3b

	if (batch > 1 && rows / batch <= 128 && cols / batch <= 128) {
	SOLVER_BLAS_DISPATCH_IMPL(GeqrfBatchedImpl, batch, rows, cols, stream,
	scratch, a, out, tau);
	} else {
	SOLVER_DISPATCH_IMPL(GeqrfImpl, batch, rows, cols, stream, scratch, a, out,
	tau);
	}

although there is a heuristic to avoid batching for smaller batches of large matrices. This decision is made at runtime, not during lowering, so you won't see the cublas name embedded in the HLO. Hope this helps!

[yantaow]

Really? I don't see that behavior hmm..
I still find manually written householder to be faster than jnp.linalg.qr. This is what I have:

45 A = misc.normalized_random([1000, 264, 264], dtype=np.complex128)
46 _fs = [partial(jnp.linalg.qr, mode='reduced'),
47 partial(jax.lax.linalg.qr, full_matrices=False),
48 LA.qr]
49 sts = ['jnp qr', 'lax qr', 'householder']
50
51 for _f, st in zip(_fs, sts):
52 ts = [time.time()]
53 for k in range(4):
54 q = _f(A)[0]
55 if type(q) == jnp.array:
56 q.block_until_ready()
57 ts.append(time.time())
58 print(f'{st} time :', np.diff(ts))

and test result on running qr on random complex (1000, 264,264) matrices. :

jnp qr time : [3.4215 0.4893 0.4848 0.4812]
lax qr time : [0.5012 0.4869 0.487 0.4873]
householder time : [0.0642 0.0013 0.0012 0.0012]

What should I do? (I have cuda13 and jax0.7.2)

[fliingelephant]

QR under Jax’s vmap exhibits poor scaling with batch size, because the step that forms the Q factor orgqr is not batched.

Benchmark

Benchmarked SVD, QR, and Cholesky on 64×64 matrices across different batch sizes. The results show that:

• QR is only partially batched
• SVD (when either matrix dimension > 32) is not batched at all
• Cholesky is well batched

Total Time (milliseconds)

Batch	SVD	QR	Cholesky
1	1.47	0.37	0.10
10	13.9	2.45	0.14
50	68.8	4.45	0.16
100	138	6.97	0.26
200	281	11.9	0.15
500	704	26.9	0.17
1000	1396	52.4	0.25
2000	2771	104	0.38

Scaling Ratio (relative to batch=1, ideal = 1.0)

Batch	SVD	QR	Cholesky
1	1.0	1.0	1.0
10	9.5	6.7	1.4
50	47	12	1.6
100	94	19	2.6
200	192	32	1.5
500	481	73	1.8
1000	953	143	2.6
2000	1892	282	4.0

Analysis

An nsys profile of vmapped QR shows that the factorization itself does use a batched kernel (geqr2_batch_kernel_shmem), evidenced by only 4 kernel launches. However, the computation of the Q matrix (orgqr and its sub-kernels) dominates runtime and is executed sequentially, resulting in ~8000 kernel launches.

Kernel	Instances	Time	Batched?
geqr2_batch_kernel_shmem (GeqrfBatched)	4	0.65 ms	Yes
larft_T_32	8000	53 ms	No
ampere_sgemm_32x128_tn	8000	51 ms	No
ampere_sgemm_32x128_nn	8000	34 ms	No
gemmSN_NN_kernel	8000	27 ms	No
larft_vtv_32	8000	21 ms	No
lacpy_kernel	8000	17 ms	No
orgqr_step1_kernel	8000	14 ms	No

Inspecting the repo confirms that there is currently no batched dispatch for orgqr; it is called sequentially for each matrix in the batch:

jax/jaxlib/gpu/solver_kernels_ffi.cc

Lines 375 to 381 in accb719

	FFI_RETURN_IF_ERROR(
	CheckShape(out->dimensions(), {batch, rows, cols}, "out", "orgqr"));
	SOLVER_DISPATCH_IMPL(OrgqrImpl, batch, rows, cols, size, stream, scratch, a,
	tau, out);
	return ffi::Error::InvalidArgument(absl::StrFormat(
	"Unsupported dtype %s in orgqr", absl::FormatStreamed(dataType)));
	}

jax/jaxlib/gpu/solver_kernels_ffi.cc

Lines 345 to 350 in accb719

	for (auto i = 0; i < batch; ++i) {
	FFI_RETURN_IF_ERROR_STATUS(solver::Orgqr<T>(
	handle.get(), m, n, k, out_data, tau_data, workspace, lwork, info));
	out_data += out_step;
	tau_data += k;
	}

This might explain the scaling behavior observed for batched QR compared to Cholesky.

[fliingelephant]

One possible solution is to use cuSOLVER to perform the raw QR factorization first, and then manually construct Q. See #8542 (comment).

With additional modifications, using the compact WY rep to form Q may even improve further:

def _qr_compactwy(a: jax.Array) -> tuple[jax.Array, jax.Array]:
    """Householder QR via compact WY representation.

    Computes reduced ``(Q, R)`` on trailing matrix axes, batch-polymorphically.
    """
    r, tau = jnp.linalg.qr(a, mode="raw")  # batchable geqrf by CuSOLVER
    q = _householder_wy(r.mT, tau)
    return q, jnp.triu(r.mT[..., : tau.shape[-1], :])


def _householder_wy(r: jax.Array, tau: jax.Array) -> jax.Array:
    """Build reduced ``Q`` from geqrf reflectors in compact WY form.

    Implements ``Q = I - Y T Y^H`` on trailing matrix axes, batch-polymorphically.
    """
    m = r.shape[-2]
    k = tau.shape[-1]
    dtype = r.dtype

    Y = jnp.tril(r[..., :, :k], k=-1) + jnp.eye(m, k, dtype=dtype)
    YHY = jnp.einsum("...ki,...kj->...ij", Y.conj(), Y, optimize=True)
    strict_lower = jnp.tril(jnp.ones((k, k), dtype=dtype), k=-1)
    basis = jnp.eye(k, dtype=dtype)

    def update_column(j: int, T: jax.Array) -> jax.Array:
        mask = strict_lower[j, :]
        yhy_col = YHY[..., :, j] * mask
        t_yhy = jnp.einsum("...ab,...b->...a", T, yhy_col, optimize=True)
        tau_j = tau[..., j][..., None]
        new_col = -tau_j * t_yhy * mask + tau_j * basis[j]
        return jax.lax.dynamic_update_slice_in_dim(T, new_col[..., None], j, axis=-1)

    T = jax.lax.fori_loop(
        0,
        k,
        update_column,
        jnp.zeros(tau.shape[:-1] + (k, k), dtype=dtype),
    )

    return jnp.eye(m, k, dtype=dtype) - jnp.einsum("...ik,...kl,...jl->...ij", Y, T, Y[..., :k, :].conj(), optimize=True)

[fliingelephant]

If you are aggressive in optimization, another option is to do QR via Cholesky. This only works for tall-thin matrices, and is numerically much less stable than Householder QR, but can still be acceptable in certain regimes. In my usage, it seems OK for batched 64×16 Complex128 matrices.

This is significantly faster than Householder QR, whether Q is formed via WY representation or not.

def _qr_cholesky(a: jax.Array) -> tuple[jax.Array, jax.Array]:
    gram = a.mH @ a
    L = jnp.linalg.cholesky(gram)
    q = jax.scipy.linalg.solve_triangular(L, a.mH, lower=True).mH
    return q, L.mH