Linear Algebra#

Overview#

The Linear Algebra module nvmath.linalg in nvmath-python leverages various NVIDIA math libraries to support dense [1] linear algebra computations. As of version 1.0.0, we offer matrix multiplication APIs as well as a dense direct linear solver API.

For matrix multiplication, we offer both a generic API based on the cuBLAS and NVPL libraries and a specialized API (nvmath.linalg.advanced) based on the cuBLASLt library. See Generic and Specialized APIs for motivation. At a high-level, if your use case is predominantly GEMM and requires particular flexibility in matrix data layouts, input and/or compute types, and also in choosing the algorithmic implementation, look at the specialized APIs. Otherwise, look at the generic APIs.

For solving dense square linear systems of the form a @ x = b via LU factorization, we offer a generic dense direct solver (direct_solver() and the stateful DirectSolver).

API Reference#

Generic Linear Algebra APIs (nvmath.linalg)#

The generic linear algebra module bundles host APIs that share a minimal, broadly applicable surface. Beyond that, each operation has its own operand model and configuration; the subsections below list each operation’s surface independently.

Shared Utilities#

ExecutionCPU(*[, num_threads])

A data class for providing CPU execution options.

ExecutionCUDA(*[, device_id])

A data class for providing GPU execution options.

Matrix multiplication#

Accepts structured matrices via matrix qualifiers (general, triangular, symmetric, Hermitian, diagonal) in addition to dense full matrices. Runs on either CPU or CUDA execution.

matmul(a, b, /[, c, alpha, beta, ...])

Perform the specified matrix multiplication computation \(\alpha a @ b + \beta c\).

Matmul(a, b, /[, c, alpha, beta, ...])

Create a stateful object encapsulating the specified matrix multiplication computation \(\alpha a @ b + \beta c\) and the required resources to perform the operation.

matrix_qualifiers_dtype

A NumPy custom dtype which describes a structured matrix.

ComputeType(*values)

See cublasComputeType_t.

DiagonalMatrixQualifier()

A class which constructs and validates matrix_qualifiers_dtype for a diagonal matrix.

GeneralMatrixQualifier()

A class which constructs and validates matrix_qualifiers_dtype for a general rectangular matrix.

HermitianMatrixQualifier()

A class which constructs and validates matrix_qualifiers_dtype for a hermitian matrix.

InvalidMatmulState

SymmetricMatrixQualifier()

A class which constructs and validates matrix_qualifiers_dtype for a symmetric matrix.

TriangularMatrixQualifier()

A class which constructs and validates matrix_qualifiers_dtype for a triangular matrix.

SideMode(*values)

See cublasSideMode_t.

FillMode(*values)

See cublasFillMode_t.

DiagType(*values)

See cublasDiagType_t.

MatmulOptions(*, allocator, blocking, ] =, ...)

A dataclass for providing options to a Matmul object.

Direct linear solver#

Solves the dense system a @ x = b for a general square a and dense b (including implicitly and explicitly batched inputs), via LU factorization on the GPU. Execution is CUDA-only: host-side operands are copied to the selected device for the factorization and solve using the triangular factors. Native CPU execution is not available today but may be added in a future release.

direct_solver(a, b, *[, options, execution, ...])

Solve \(a @ x = b\) for \(x\).

DirectSolver(a, b[, options, execution, stream])

Create a stateful object that encapsulates the specified dense direct linear solver computations and required resources.

InvalidDirectSolverState

Raised when a DirectSolver instance (or a method on it) is used after DirectSolver.free() has run and resources are invalid.

DirectSolverOptions(*, allocator, blocking, ...)

A data class for providing options to direct_solver() and DirectSolver.

Specialized Linear Algebra APIs (nvmath.linalg.advanced)#

The specialized linear algebra module includes a matrix multiplication API which only accepts general matrices, but provides extra functionality such as epilog functions, more options and controls over computational precision, and control over algorithm selection and planning.

matmul(a, b, /[, c, alpha, beta, epilog, ...])

Perform the specified matrix multiplication computation \(\mathcal{F}(\alpha a @ b + \beta c)\), where \(\mathcal{F}\) is the epilog.

matrix_qualifiers_dtype

NumPy dtype object that encapsulates the matrix qualifiers in linalg.advanced.

Algorithm(algorithm)

An interface class to query algorithm capabilities and configure the algorithm.

Matmul(a, b, /[, c, alpha, beta, ...])

Create a stateful object encapsulating the specified matrix multiplication computation \(\alpha a @ b + \beta c\) and the required resources to perform the operation.

MatmulComputeType

alias of ComputeType

MatmulEpilog

alias of Epilogue

MatmulInnerShape(*values)

See cublasLtMatmulInnerShape_t.

MatmulNumericalImplFlags(*values)

These flags can be combined with the | operator: OP_TYPE_FMA | OP_TYPE_TENSOR_HMMA ...

MatmulReductionScheme

alias of ReductionScheme

MatmulEpilogPreferences([aux_type, aux_amax])

A data class for providing epilog options as part of preferences to the Matmul.plan() method and the wrapper function matmul().

MatmulOptions([inplace, compute_type, ...])

A data class for providing options to the Matmul object and the wrapper function matmul().

MatmulPlanPreferences([...])

A data class for providing options to the Matmul.plan() method and the wrapper function matmul().

MatmulQuantizationScales([a, b, c, d])

A data class for providing quantization_scales to Matmul constructor and the wrapper function matmul().

Helpers#

The Specialized Linear Algebra helpers module nvmath.linalg.advanced.helpers provides helper functions to facilitate working with some of the complex features of nvmath.linalg.advanced module.

Matmul helpers (nvmath.linalg.advanced.helpers.matmul)#

BlockScalingFormat(*values)

Block scaling format for microscaling data types.

create_mxfp8_scale(x, exponent[, stream])

invert_mxfp8_scale(mx_scales)

apply_mxfp8_scale(x, scales_1d[, output_dtype])

quantize_to_fp4(x, axis)

unpack_fp4(fp4_tensor, axis)

get_block_scale_offset(index, ...[, axis])

to_block_scale(scale_tensor, ...[, axis, out])

expand_block_scale(scales_1d, ...[, axis, ...])

Footnotes