nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/linalg/_internal/utils.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Utilities.
+"""
+
+__all__ = [
+    "axis_order_in_memory",
+    "calculate_strides",
+    "check_batch_tileable",
+    "get_handle",
+    "pointer_aligned_to",
+]
+
+import typing
+
+from nvmath._internal import threadsafe
+from nvmath.bindings import cublas
+from nvmath.bindings import cublasLt as cublaslt
+from nvmath.internal import utils
+from nvmath.internal._layout import StridedLayout
+
+
+def create_cublas_handle(device_id: int) -> int:
+    """
+    Currently for internal use only.
+    """
+    with utils.device_ctx(device_id):
+        return cublas.create()
+
+
+def create_cublaslt_handle(device_id: int) -> int:
+    """
+    Currently for internal use only.
+    """
+    with utils.device_ctx(device_id):
+        return cublaslt.create()
+
+
+# One per-thread cache per binding. Keying by `device_id` alone (rather than a
+# shared `(binding, device_id)` dict) keeps the two bindings' handles in
+# separate caches, which avoids any chance of a binding-string typo aliasing
+# the wrong handle. cuBLAS handles must not be shared across threads
+# (https://docs.nvidia.com/cuda/cublas/#thread-safety); cuBLASLt handles could
+# be shared, but per-thread caching is a safe superset and keeps both paths
+# uniform.
+_CUBLAS_CACHE = threadsafe.HandleCache[int](
+    create=create_cublas_handle,
+    destroy=cublas.destroy,
+)
+_CUBLASLT_CACHE = threadsafe.HandleCache[int](
+    create=create_cublaslt_handle,
+    destroy=cublaslt.destroy,
+)
+
+
+def get_handle(device_id: int, binding="cublaslt") -> int:
+    """
+    Retrieve the cuBLAS[lt] library handle for the specified device. If one doesn't exist,
+    create, cache, and return the handle.
+
+    Handles are cached per thread in a separate cache per binding, so a single
+    thread can hold distinct handles for ``"cublas"`` and ``"cublaslt"`` on the
+    same device. Cached handles are released automatically when the owning
+    thread exits.
+    """
+    if binding == "cublas":
+        cache = _CUBLAS_CACHE
+    elif binding == "cublaslt":
+        cache = _CUBLASLT_CACHE
+    else:
+        raise ValueError(f"{binding} is not a valid library name.")
+    return cache.get(device_id)
+
+
+def pointer_aligned_to(address):
+    """
+    Return the number of bytes the address is aligned to.
+    """
+    return address & ~(address - 1)
+
+
+def axis_order_in_memory(strides):
+    """
+    Compute the order in which the axes appear in memory.
+    """
+    if len(strides) == 0:
+        return ()
+
+    _, axis_order = zip(*sorted(zip(strides, range(len(strides)), strict=True)), strict=True)
+
+    return axis_order
+
+
+def calculate_strides(shape: typing.Sequence[int], axis_order: typing.Sequence[int], min_stride: int = 1):
+    """
+    Calculate the strides for the provided shape and axis order.
+    """
+    assert len(axis_order) == len(shape), f"axis_order length ({len(axis_order)}) must equal shape length ({len(shape)})"
+    assert len(set(axis_order)) == len(axis_order), f"axis_order must not contain duplicates: {axis_order}"
+    assert set(axis_order) == set(range(len(shape))), f"axis_order must be permutation of range({len(shape)}): {axis_order}"
+
+    strides: list[None | int] = [None] * len(shape)
+
+    stride = min_stride
+    for axis in axis_order:
+        strides[axis] = stride
+        stride *= shape[axis]
+
+    return strides
+
+
+def check_batch_tileable(batch_shape, batch_strides):
+    """
+    Whether the batch layout can be expressed the way cuBLAS strided-batched
+    expects: a single ``batch_count`` and a single stride ``S`` per operand,
+    so that tile ``k`` lives at ``base + k * S`` for ``k = 0, ..., N - 1``.
+
+    Two scenarios -- and only these two -- satisfy that form:
+
+    (a) Regular grid of distinct tiles (``S > 0``): tiles sit at evenly
+        spaced offsets along *some* axis order, with the gap between
+        consecutive tiles allowed to be larger than the tile itself (i.e.
+        padded/gapped batches are fine).
+
+    (b) Replicated tile (``S == 0``): every batch index lands on the same
+        tile, i.e. *every* non-singleton batch dim has stride 0. Typical
+        source: ``torch.expand`` (or equivalent) applied so that the
+        operand has no real batch axis -- e.g. ``A.unsqueeze(0).expand(N,
+        M, K)``.
+
+    A broadcast on *some* batch axis combined with a real (stride > 0)
+    batch axis of size > 1 does **not** qualify for (b) -- it isn't
+    tileable in the form above and is rejected. E.g.
+    ``X.unsqueeze(0).expand(N1, P, M, K)`` with both ``N1 > 1`` *and*
+    ``P > 1`` has batch strides ``(0, M*K)`` and is not tileable.
+    """
+    # Empty / all-size-1 batch: degenerate intersection of (a) and (b),
+    # trivially tileable.
+    if all(n == 1 for n in batch_shape):
+        return True
+
+    layout = StridedLayout(batch_shape, batch_strides, itemsize=1)
+
+    # Scenario (a): tiles on a regular grid in some axis order.
+    # `"K"` lets the grid be in any order (C, F, or a permutation), and
+    # `allow_leading_dim_stride=True` lets the smallest stride exceed 1.
+    if layout.is_dense("K", allow_leading_dim_stride=True):
+        return True
+
+    # Scenario (b): pure broadcast. Flatten collapses to a single 1-D dim
+    # with stride 0 iff every non-size-1 batch stride is 0; mixed
+    # broadcast + non-broadcast dims don't collapse and get rejected.
+    flat = layout.flattened()
+    return flat.ndim == 1 and flat.strides[0] == 0

nvmath/linalg/advanced/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import helpers as helpers  # noqa: F403
+from ._algorithmmod import *  # noqa: F403
+from ._configuration import *  # noqa: F403
+from .matmulmod import *  # noqa: F403

nvmath/linalg/advanced/_algorithmmod.py

@@ -0,0 +1,170 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+An interface class to query algorithm capabilities and configure it.
+"""
+
+__all__ = ["Algorithm"]
+
+import dataclasses
+
+import numpy as np
+import numpy.typing as npt
+
+from nvmath.bindings.cublasLt import MatmulHeuristicResult, matmul_heuristic_result_dtype
+from nvmath.linalg._internal.algo_cap_ifc import AlgoCapInterface
+from nvmath.linalg._internal.algo_config_ifc import AlgoConfigInterface
+from nvmath.linalg._internal.enum_to_tuples import (
+    CLUSTER_SHAPE_TO_ENUM,
+    ENUM_TO_CLUSTER_SHAPE,
+    ENUM_TO_MATMUL_STAGE,
+    ENUM_TO_MATMUL_TILE,
+    MATMUL_STAGE_TO_ENUM,
+    MATMUL_TILE_TO_ENUM,
+)
+from nvmath.linalg.advanced._configuration import AlgorithmCapabilities
+
+
+class Algorithm:
+    """
+    An interface class to query algorithm capabilities and configure the algorithm.
+
+    Note that objects of this type should not be constructed directly by the user.
+    """
+
+    def __init__(self, algorithm: MatmulHeuristicResult) -> None:
+        assert isinstance(algorithm, MatmulHeuristicResult), "Internal error."
+        self.algorithm = algorithm
+        self.cap_ifc = AlgoCapInterface(algorithm)
+        self.config_ifc = AlgoConfigInterface(algorithm)
+
+    @property
+    def capabilities(self):
+        """
+        Return the capabilities of this algorithm as a
+        :class:`nvmath.linalg.advanced.AlgorithmCapabilities` dataclass.
+        """
+        names = [field.name for field in dataclasses.fields(AlgorithmCapabilities)]
+        _capabilities = {}
+        for name in names:
+            # Not all capabilities in the enum are supported in all CTK versions.
+            try:
+                _capabilities[name] = getattr(self.cap_ifc, name)
+            except:
+                pass
+        return AlgorithmCapabilities(**_capabilities)
+
+    @property
+    def algorithm_id(self) -> int:
+        "The ID of the algorithm (integer)."
+        return self.config_ifc.id
+
+    @property
+    def tile(self) -> tuple[int | str, ...] | str:
+        """A tuple representing the tile (see MatmulAlgoConfigAttribute.TILE_ID).
+        The value provided must be one of the `tile_ids` in the algorithm capabilities."""
+        return ENUM_TO_MATMUL_TILE[self.config_ifc.tile_id]
+
+    @tile.setter
+    def tile(self, tile: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.tile_id = MATMUL_TILE_TO_ENUM[tile]
+
+    @property
+    def stages(self) -> tuple[int | str, ...] | str:
+        """A tuple representing the stages (see MatmulAlgoConfigAttribute.STAGES_ID).
+        The value provided must be one of the `stages_ids` in the algorithm capabilities."""
+        return ENUM_TO_MATMUL_STAGE[self.config_ifc.stages_id]
+
+    @stages.setter
+    def stages(self, stages: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.stages_id = MATMUL_STAGE_TO_ENUM[stages]
+
+    @property
+    def split_k(self) -> int:
+        """The number of split-k steps (see MatmulAlgoConfigAttribute.SPLITK_NUM).
+
+        This can be set only if `splitk_support` is 1 in the algorithm capabilities."""
+        return self.config_ifc.splitk_num
+
+    @split_k.setter
+    def split_k(self, number: int) -> None:
+        self.config_ifc.splitk_num = number
+
+    @property
+    def reduction_scheme(self) -> int:
+        """The reduction scheme used (see MatmulAlgoConfigAttribute.REDUCTION_SCHEME).
+
+        The value provided must be consistent with the `reduction_scheme_mask` in the
+        algorithm capabilities."""
+        return self.config_ifc.reduction_scheme
+
+    @reduction_scheme.setter
+    def reduction_scheme(self, scheme_id: int) -> None:
+        self.config_ifc.reduction_scheme = scheme_id
+
+    @property
+    def cta_swizzling(self) -> bool:
+        """A flag indicating CTA swizzling (see MatmulAlgoConfigAttribute.CTA_SWIZZLING).
+
+        This can be set only if `cta_swizzling` is 1 in the algorithm capabilities."""
+        return self.config_ifc.cta_swizzling
+
+    @cta_swizzling.setter
+    def cta_swizzling(self, flag: bool) -> None:
+        self.config_ifc.cta_swizzling = flag
+
+    @property
+    def custom_option(self) -> int:
+        """A value indicating the custom option (see
+        MatmulAlgoConfigAttribute.CUSTOM_OPTION).
+
+        The value provided must be less than `custom_option_max` in the algorithm
+        capabilities."""
+        return self.config_ifc.custom_option
+
+    @custom_option.setter
+    def custom_option(self, value: int) -> None:
+        self.config_ifc.custom_option = value
+
+    @property
+    def inner_shape(self) -> int:
+        """A value indicating the inner shape (see
+        MatmulAlgoConfigAttribute.INNER_SHAPE_ID)."""
+        return self.config_ifc.inner_shape_id
+
+    @inner_shape.setter
+    def inner_shape(self, shape: int) -> None:
+        self.config_ifc.inner_shape_id = shape
+
+    @property
+    def cluster_shape(self) -> tuple[int | str, ...] | str:
+        """A tuple representing the cluster shape (see
+        MatmulAlgoConfigAttribute.CLUSTER_SHAPE_ID).
+
+        The value provided must be one of the `cluster_shape_ids` in the algorithm
+        capabilities."""
+        return ENUM_TO_CLUSTER_SHAPE[self.config_ifc.cluster_shape_id]
+
+    @cluster_shape.setter
+    def cluster_shape(self, shape: tuple[int | str, ...] | str) -> None:
+        self.config_ifc.cluster_shape_id = CLUSTER_SHAPE_TO_ENUM[shape]
+
+    def as_numpy(self) -> npt.NDArray:
+        """Return a copy of this `Algorithm` as a NumPy array.
+
+        Converts the `Algorithm` to a NumPy array of type `matmul_heuristic_result_dtype`
+        which may then be saved to disk using ``np.save(..., allow_pickle=False)``.
+        """
+        return np.frombuffer(buffer=bytes(memoryview(self.algorithm)), dtype=matmul_heuristic_result_dtype)  # type: ignore
+
+    @classmethod
+    def from_numpy(cls, array: npt.NDArray) -> "Algorithm":
+        """Create an `Algorithm` from a NumPy array.
+
+        Creates an `Algorithm` by copying a NumPy array of type
+        `matmul_heuristic_result_dtype`. The array must be of size 1, and may have been
+        loaded from disk using ``np.load(..., allow_pickle=False)``.
+        """
+        return Algorithm(MatmulHeuristicResult.from_data(np.copy(array)))

nvmath/linalg/advanced/_configuration.py

@@ -0,0 +1,351 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = [
+    "MatmulEpilog",
+    "MatmulInnerShape",
+    "MatmulNumericalImplFlags",
+    "MatmulOptions",
+    "MatmulEpilogPreferences",
+    "MatmulPlanPreferences",
+    "MatmulQuantizationScales",
+    "MatmulReductionScheme",
+    "matrix_qualifiers_dtype",
+]
+
+import dataclasses
+from enum import IntEnum
+from logging import Logger
+from typing import Literal
+
+import numpy as _np
+
+from nvmath._utils import CudaDataType
+from nvmath.bindings import cublas  # type: ignore
+from nvmath.bindings import cublasLt as cublaslt  # type: ignore
+from nvmath.internal import enum_utils
+from nvmath.internal.mem_limit import check_memory_str
+from nvmath.internal.tensor_ifc import AnyTensor
+from nvmath.internal.utils import check_or_create_options
+from nvmath.memory import BaseCUDAMemoryManager, BaseCUDAMemoryManagerAsync
+
+MatmulEpilog = cublaslt.Epilogue
+MatmulInnerShape = cublaslt.MatmulInnerShape
+MatmulReductionScheme = cublaslt.ReductionScheme
+
+
+@dataclasses.dataclass
+class MatmulOptions:
+    """A data class for providing options to the :class:`Matmul` object and the wrapper
+    function :func:`matmul`.
+
+    Attributes:
+        inplace: Whether the matrix multiplication is performed in-place (operand C is
+            overwritten); ``inplace=True`` implies that operand C cannot be broadcast (in
+            the batch as well as the N-dimension) since it must be large enough to
+            hold the result of the computation. The default is ``inplace=False``.
+
+            .. experimental:: attribute
+
+        compute_type (nvmath.linalg.ComputeType): CUDA compute type. A suitable compute type
+            will be selected if not specified.
+
+        scale_type (nvmath.CudaDataType): CUDA data type. A suitable data type consistent
+            with the compute type will be selected if not specified.
+
+        result_type (nvmath.CudaDataType): CUDA data type. A requested datatype of the
+            result. If not specified, this type will be determined based on the input types.
+            Non-default result types are only supported for narrow-precision (FP8 and lower)
+            operations.
+
+        result_amax (bool): If set, the absolute maximum (amax) of the result will be
+            returned in the auxiliary output tensor. Only supported for narrow-precision
+            (FP8 and lower) operations.
+
+        block_scaling (bool): If set, block scaling (MXFP8) will be used instead of
+            tensor-wide scaling for FP8 operations. If the result is a narrow-precision
+            (FP8 and lower) data type, scales used for result quantization will be returned
+            in the auxiliary output tensor as ``"d_out_scale"`` in UE8M0 format. For more
+            information on UE8M0 format, see the documentation of
+            :class:`~nvmath.linalg.advanced.MatmulQuantizationScales`.
+            This option is only supported for narrow-precision (FP8 and lower) operations.
+
+        sm_count_target (int) : The number of SMs to use for execution. The default is 0,
+            corresponding to all available SMs.
+
+        fast_accumulation (bool) : Enable or disable FP8 fast accumulation mode. The default
+            is False (disabled).
+
+        device_id: CUDA device ordinal (only used if the operand resides on the CPU). The
+            default value is 0.
+
+        handle: Linear algebra library handle. A handle will be created if one is not
+            provided.
+
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+        memory_limit: Maximum memory available to the MM operation. It can be specified as a
+            value (with optional suffix like K[iB], M[iB], G[iB]) or as a percentage. The
+            default is 80% of the device memory.
+
+        blocking: A flag specifying the behavior of the execution functions and methods,
+            such as :func:`matmul` and :meth:`Matmul.execute`. When ``blocking`` is `True`,
+            the execution methods do not return until the operation is complete. When
+            ``blocking`` is ``"auto"``, the methods return immediately when the inputs are
+            on the GPU. The execution methods always block when the operands are on the CPU
+            to ensure that the user doesn't inadvertently use the result before it becomes
+            available. The default is ``"auto"``.
+
+        allocator: An object that supports the :class:`BaseCUDAMemoryManager` protocol, used
+            to draw device memory. If an allocator is not provided, a memory allocator from
+            the library package will be used (:func:`torch.cuda.caching_allocator_alloc` for
+            PyTorch operands, :func:`cupy.cuda.alloc` otherwise).
+
+    .. seealso::
+       :class:`Matmul`, :func:`matmul`
+    """
+
+    inplace: bool = False
+    compute_type: int | None = None
+    scale_type: int | None = None
+    result_type: int | None = None
+    result_amax: bool = False
+    block_scaling: bool = False
+    sm_count_target: int | None = 0
+    fast_accumulation: bool | None = False
+    device_id: int = 0
+    handle: int | None = None
+    logger: Logger | None = None
+    memory_limit: int | str | None = r"80%"
+    blocking: Literal[True, "auto"] = "auto"
+    allocator: BaseCUDAMemoryManager | None = None
+
+    def __post_init__(self):
+        #  Defer computing the memory limit till we know the device the network is on.
+
+        if self.compute_type is not None:
+            self.compute_type = cublas.ComputeType(self.compute_type)
+
+        if self.scale_type is not None:
+            self.scale_type = CudaDataType(self.scale_type)
+
+        if self.sm_count_target is None:
+            self.sm_count_target = MatmulOptions.sm_count_target
+
+        if self.fast_accumulation is None:
+            self.fast_accumulation = MatmulOptions.fast_accumulation
+
+        if self.device_id is None:
+            self.device_id = 0
+
+        if self.memory_limit is None:
+            self.memory_limit = MatmulOptions.memory_limit
+
+        check_memory_str(self.memory_limit, "memory limit")
+
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for blocking must be either True or 'auto'.")
+
+        if self.allocator is not None and not isinstance(self.allocator, BaseCUDAMemoryManager | BaseCUDAMemoryManagerAsync):
+            raise TypeError("The allocator must be an object of type that fulfills the BaseCUDAMemoryManager protocol.")
+
+
+matrix_qualifiers_dtype = _np.dtype([("structure", object), ("is_conjugate", "<i4")])
+
+
+class MatmulNumericalImplFlags(IntEnum):
+    """
+    These flags can be combined with the | operator: OP_TYPE_FMA | OP_TYPE_TENSOR_HMMA ...
+    """
+
+    OP_TYPE_FMA = 0x01 << 0
+    OP_TYPE_TENSOR_HMMA = 0x02 << 0
+    OP_TYPE_TENSOR_IMMA = 0x04 << 0
+    OP_TYPE_TENSOR_DMMA = 0x08 << 0
+    OP_TYPE_TENSOR_MASK = 0xFE << 0
+    OP_TYPE_MASK = 0xFF << 0
+
+    ACCUMULATOR_16F = 0x01 << 8
+    ACCUMULATOR_32F = 0x02 << 8
+    ACCUMULATOR_64F = 0x04 << 8
+    ACCUMULATOR_32I = 0x08 << 8
+    ACCUMULATOR_TYPE_MASK = 0xFF << 8
+
+    INPUT_TYPE_16F = 0x01 << 16
+    INPUT_TYPE_16BF = 0x02 << 16
+    INPUT_TYPE_TF32 = 0x04 << 16
+    INPUT_TYPE_32F = 0x08 << 16
+    INPUT_TYPE_64F = 0x10 << 16
+    INPUT_TYPE_8I = 0x20 << 16
+    INPUT_TYPE_8F_E4M3 = 0x40 << 16
+    INPUT_TYPE_8F_E5M2 = 0x80 << 16
+    INPUT_TYPE_MASK = 0xFF << 16
+
+    GAUSSIAN = 0x01 << 32
+
+    ALL = (1 << 64) - 1
+
+
+@dataclasses.dataclass
+class MatmulEpilogPreferences:
+    """A data class for providing epilog options as part of ``preferences`` to the
+    :meth:`Matmul.plan` method and the wrapper function :func:`matmul`.
+
+    Attributes:
+        aux_type (nvmath.CudaDataType): The requested datatype of the
+            epilog auxiliary output. If not specified, this type will be determined based on
+            the input types. Non-default auxiliary output types are only supported for
+            narrow-precision operations and certain epilogs. For more details on the
+            supported combinations, see ``CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE`` in
+            cuBLAS documentation. If this option is set to a narrow-precision data type,
+            an additional epilog input ``"aux_quantization_scale"`` needs to be specified.
+
+        aux_amax (bool): If set, the absolute maximum (amax) of the epilog
+            auxiliary output will be returned in the auxiliary output tensor.
+            Only supported when ``aux_type`` option is set to a narrow-precision
+            data type.
+
+    .. seealso::
+       :meth:`Matmul.plan`, :func:`matmul`, :class:`MatmulPlanPreferences`
+    """
+
+    aux_type: int | None = None
+    aux_amax: bool = False
+
+
+@dataclasses.dataclass
+class MatmulPlanPreferences:
+    """A data class for providing options to the :meth:`Matmul.plan` method and the
+    wrapper function :func:`matmul`.
+
+    Attributes:
+        reduction_scheme_mask (:class:`nvmath.linalg.advanced.MatmulReductionScheme`):
+            Enumerators from
+            :class:`nvmath.linalg.advanced.MatmulReductionScheme` combined with
+            bitwise operator ``|``. The default is all reduction schemes.
+
+        max_waves_count (float) : The maximum wave count. Selecting a value greater than 0.
+            will exclude algorithms with device utilization greater than specified. The
+            default is 0.
+
+        numerical_impl_mask (:class:`nvmath.linalg.advanced.MatmulNumericalImplFlags`):
+            Enumerators from
+            :class:`nvmath.linalg.advanced.MatmulNumericalImplFlags` combined with
+            bitwise operator ``|``. The default is all numerical implementation flag
+            choices.
+
+        limit (int) : The number of algorithms to consider. If not specified, a suitable
+            default will be chosen.
+
+        epilog (:class:`nvmath.linalg.advanced.MatmulEpilogPreferences`):
+            Epilog preferences (as an object of class
+            :class:`~nvmath.linalg.advanced.MatmulEpilogPreferences` or a `dict`).
+
+    .. seealso::
+       :meth:`Matmul.plan`, :func:`matmul`
+    """
+
+    reduction_scheme_mask: cublaslt.ReductionScheme | None = cublaslt.ReductionScheme.MASK
+    max_waves_count: float | None = 0.0
+    numerical_impl_mask: MatmulNumericalImplFlags | None = MatmulNumericalImplFlags.ALL
+    limit: int = 8
+    epilog: MatmulEpilogPreferences | None = None
+
+    def __post_init__(self):
+        if self.reduction_scheme_mask is None:
+            self.reduction_scheme_mask = MatmulPlanPreferences.reduction_scheme_mask
+        else:
+            self.reduction_scheme_mask = cublaslt.ReductionScheme(self.reduction_scheme_mask)
+
+        if self.max_waves_count is None:
+            self.max_waves_count = MatmulPlanPreferences.max_waves_count
+
+        if self.numerical_impl_mask is None:
+            self.numerical_impl_mask = MatmulPlanPreferences.numerical_impl_mask
+
+        if self.limit is None:
+            self.limit = MatmulPlanPreferences.limit
+
+        self.epilog = check_or_create_options(MatmulEpilogPreferences, self.epilog, "epilog preferences")
+
+
+@dataclasses.dataclass
+class MatmulQuantizationScales:
+    """A data class for providing quantization_scales to :class:`Matmul` constructor and the
+    wrapper function :func:`matmul`.
+
+    Scales can only be set for narrow-precision (FP8 and lower) matrices.
+
+    **FP8 operations** (``block_scaling=False``, per-tensor scaling):
+
+    * Scale format: **scalar** (integer or float) or single-element tensor
+      of shape ``()`` or ``(1,)``
+    * A single scale value is applied to the entire tensor.
+
+    **MXFP8 operations** (``block_scaling=True``, microscaling with FP8 operands):
+
+    * Scale format: **1D tensor** with layout matching cuBLAS MXFP8 requirements
+    * Scale dtype: ``uint8`` (interpreted as UE8M0 values by cuBLAS)
+    * A scale value :math:`x` causes cuBLAS to multiply the respective
+      block by :math:`2^{x-127}`.
+
+    **NVFP4 operations** (``block_scaling=True``, block scaling with FP4 operands):
+
+    * Scale format: **1D tensor** with layout matching cuBLAS NVFP4 requirements
+    * Scale dtype: ``float8_e4m3fn`` (interpreted by cuBLAS as unsigned UE4M3 values,
+      i.e. the bit sign is ignored).
+    * FP4 only supports block scaling; ``block_scaling=False`` is not supported for FP4.
+
+    .. note::
+       MXFP8 and NVFP4 use "block scaling" where each block of elements has its own
+       scale factor, as opposed to FP8 which uses a single per-tensor scale.
+       When ``block_scaling=True``, tensor scales are required; scalar
+       scales are not allowed.
+
+    .. note::
+       When scales are provided as tensors, they must be from the same package
+       and on the same memory space (CPU or GPU device) as the operands of the matmul.
+
+    Attributes:
+        a (int, float, or Tensor) : Scale for matrix A.
+
+        b (int, float, or Tensor) : Scale for matrix B.
+
+        c (int, float, or Tensor) : Scale for matrix C.
+
+        d (int, float, or Tensor) : Scale for matrix D.
+
+    .. seealso::
+       :class:`Matmul`, :func:`matmul`
+    """
+
+    a: int | float | AnyTensor | None = None
+    b: int | float | AnyTensor | None = None
+    c: int | float | AnyTensor | None = None
+    d: int | float | AnyTensor | None = None
+
+
+_create_options = enum_utils.create_options_class_from_enum
+_algo_cap_enum = cublaslt.MatmulAlgoCapAttribute
+_get_dtype = cublaslt.get_matmul_algo_cap_attribute_dtype
+
+AlgorithmCapabilities = _create_options(
+    "AlgorithmCapabilities", _algo_cap_enum, _get_dtype, "algorithm capabilities", "(?P<option_name>.*)"
+)
+
+
+def algorithm_capabilities_str(self):
+    names = [field.name for field in dataclasses.fields(self)]
+    width = max(len(n) for n in names)
+    s = """Algorithm Capabilities (refer to `MatmulAlgoCapAttribute` for documentation):
+"""
+    for name in names:
+        s += f"    {name:{width}} = {getattr(self, name)}\n"
+    return s
+
+
+AlgorithmCapabilities.__str__ = algorithm_capabilities_str
+
+del _create_options, _algo_cap_enum, _get_dtype

nvmath/linalg/advanced/helpers/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import matmul as matmul  # noqa: F403