Linear Algebra#

Overview#

The Linear Algebra module nvmath.linalg in nvmath-python leverages various NVIDIA math libraries to support multiple linear algebra computations. As of the initial Beta release, we offer the specialized matrix multiplication API based on the cuBLASLt library.

API Reference#

Generic Linear Algebra APIs (nvmath.linalg)#

Generic APIs will be available in a later release.

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

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

Perform the specified matrix multiplication computation \(F(\alpha a @ b + \beta c)\), where \(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(value[, names, module, ...])

See cublasLtMatmulInnerShape_t.

MatmulNumericalImplFlags(value[, names, ...])

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([compute_type, scale_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)#

create_mxfp8_scale(x, exponent[, stream])

Create MXFP8 block scale with the same value for the whole tensor x.

invert_mxfp8_scale(scale)

Compute a reciprocal of MXFP8 block scale.

apply_mxfp8_scale(x, scale)

Apply MXFP8 block scale factors to tensor x.

get_mxfp8_scale_offset(x, index)

Computes the offset of MXFP8 scale used for element x[index].