nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/linalg/advanced/matmulmod.py

@@ -0,0 +1,3734 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["MatmulComputeType", "Matmul", "matmul"]
+
+import copy
+import dataclasses
+import functools
+import logging
+import operator
+import random
+import typing
+from collections import namedtuple
+from collections.abc import MutableSequence, Sequence
+from dataclasses import dataclass
+
+import numpy as np
+from cuda.core import Device
+
+from nvmath import memory
+from nvmath._internal.workspace import Workspace
+from nvmath._utils import CudaDataType
+from nvmath.bindings import cublas
+from nvmath.bindings import cublasLt as cublaslt  # type: ignore
+from nvmath.internal import formatters, tensor_wrapper, typemaps, utils
+from nvmath.internal._layout import StridedLayout
+from nvmath.linalg._internal import matmul_desc_ifc, matmul_pref_ifc, matrix_layout_ifc
+from nvmath.linalg._internal.epilog_protocol import (
+    BATCHED_EPILOG_MINIMUM_VERSIONS_MAP,
+    EPILOG_INPUT_HANDLERS_MAP,
+    EPILOG_MINIMUM_VERSIONS_MAP,
+    EPILOG_OUTPUT_HANDLERS_MAP,
+    EpilogOutputHandler,
+)
+from nvmath.linalg._internal.typemaps import (
+    COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE,
+    NAMES_TO_DEFAULT_COMPUTE_TYPE,
+    NAMES_TO_DEFAULT_SCALE_TYPE,
+    SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE,
+    SUPPORTED_TYPES,
+)
+from nvmath.linalg._internal.utils import (
+    axis_order_in_memory,
+    calculate_strides,
+    check_batch_tileable,
+    get_handle,
+    pointer_aligned_to,
+)
+from nvmath.linalg.advanced import _algorithmmod, _configuration
+
+MatmulComputeType = cublas.ComputeType
+
+EpilogInputTraits = namedtuple("EpilogInputTraits", ["dtype", "extents", "strides"])
+
+
+@dataclass
+class MatrixLayout:
+    """An internal data class for capturing the tensor layout."""
+
+    shape: Sequence[int]
+    strides: Sequence[int]
+    is_conjugate: bool = False  # Used to support is_conjugate via conjugate_transpose.
+
+
+@dataclass
+class LayoutTraits:
+    """An internal data class for capturing the matrix multiplication traits."""
+
+    order: cublaslt.Order
+    ld: int
+    batch_offset: int  # Based on strides
+    is_conjugate: bool  # Used to support is_conjugate via conjugate_transpose.
+    mm_shape: Sequence[int] | None = None
+    mm_strides: Sequence[int] | None = None
+
+    def get_mm_layout(self, transpose=False):
+        if self.is_conjugate:
+            transpose = True
+        if not transpose:
+            return *self.mm_shape, self.ld, self.order
+
+        # Use of transpose is supported only for A and B for two specific use cases till the
+        # C library directly supports these use cases:
+        #
+        #   1. When A or B has the conjugate qualifier, we transpose it internally and then
+        #      use conjugate transpose in the MM (A @ B.conj() == A @ B.T.H).
+        #
+        #   2. When the epilog is BGRADB, we transpose B internally and use transpose in the
+        #      MM since this epilog requires B to be transposed (A @ B == A @ B.T.T).
+        #
+        #  This requires that the layout order be ROW or COL (no special layouts such as
+        #  structured or hierarchical).
+        assert self.mm_shape is not None and self.mm_strides is not None, "Internal Error."
+        assert self.ld != 0, "Internal Error."
+
+        mm_shape = self.mm_shape[1], self.mm_shape[0]
+        if self.order == cublaslt.Order.ROW:
+            order = cublaslt.Order.COL
+            ld = max(self.mm_shape[1], self.mm_strides[0])
+        elif self.order == cublaslt.Order.COL:
+            order = cublaslt.Order.ROW
+            ld = max(self.mm_shape[0], self.mm_strides[1])
+        else:
+            raise AssertionError("Internal Error. Invalid layout order.")
+
+        return *mm_shape, ld, order
+
+
+@dataclass
+class MMTraits:
+    """An internal data class for capturing the matrix multiplication traits. The
+    result traits are captured separately, because we need to wait for the
+    epilog to be provided.
+    """
+
+    M: int
+    N: int
+    K: int
+    d_mm_shape: Sequence[int]
+    inplace: bool
+    c_layout: MatrixLayout
+    a_layout_traits: LayoutTraits
+    b_layout_traits: LayoutTraits
+    c_layout_traits: LayoutTraits
+    batch_count: int
+    batch_shape: Sequence[int]
+    batch_axis_order: Sequence[int]
+
+
+@dataclass
+class ResultTraits:
+    """An internal data class for capturing the result matrix's traits."""
+
+    d_layout_traits: LayoutTraits
+    result_shape: Sequence[int]
+    result_strides: Sequence[int]
+
+
+def get_matrix_layout_traits(
+    mm_shape: Sequence[int],
+    mm_strides: Sequence[int],
+    batch_strides: Sequence[int],
+    col_bcast: bool,
+    ordering: cublaslt.Order | None = None,
+    orientation: cublaslt.Order | None = None,
+) -> tuple[cublaslt.Order, int, int]:
+    """
+    The 'ordering' option specifies the layout order, if it's not None, as in the case of
+    the D matrix whose layout is determined by the other operands' layout. It is required if
+    the matrix is degenerate (a vector or scalar: len(mm_shape) < 2).
+
+    The 'orientation' option (ROW or COL) is required to infer the correct leading dimension
+    for degenerate matrices (a vector or scalar: len(mm_shape) < 2).
+    """
+    if len(mm_shape) < 2:  # The result D can be a scalar or vector.
+        assert ordering is not None, "Internal Error: 'ordering' must be specified for degenerate matrices."
+        assert orientation is not None, "Internal Error: 'orientation' must be specified for degenerate matrices."
+        batch_offset = min(batch_strides) if batch_strides else 0
+        order = ordering
+        if len(mm_shape) < 1:
+            ld = 1
+        elif order != orientation:
+            # For a ROW vector in COL order, the LD should be 1 since we promote the row
+            # vector to a matrix as (1, M). Similarly, for a COL vector in ROW order, the LD
+            # should be 1 as well since we promote the column vector to a matrix as (M, 1).
+            ld = 1
+        else:
+            ld = max(mm_strides[0], mm_shape[0])
+        return order, ld, batch_offset
+
+    M, N = mm_shape
+
+    if ordering is not None:
+        order = ordering
+        message = f"Internal Error: incompatible ordering '{ordering}' and strides {mm_strides}"
+        if order == cublaslt.Order.ROW:
+            assert mm_strides[0] >= mm_strides[1] and mm_strides[1] == 1, message
+        else:
+            assert mm_strides[1] >= mm_strides[0] and mm_strides[0] == 1, message
+    else:
+        # Important: start with the first dimension so that cases like (M, 1) : (1, 1) or
+        # (1, M) : (1, 1) in CuTe notation map to COL.
+        if mm_strides[0] == 1:
+            order = cublaslt.Order.COL
+        elif mm_strides[1] == 1:
+            order = cublaslt.Order.ROW
+        else:
+            if M == 1:
+                order = cublaslt.Order.COL
+            elif N == 1:
+                order = cublaslt.Order.ROW
+            else:
+                raise ValueError("Unsupported layout.")
+
+    # We need to handle broadcast dimensions with zero-stride for the c matrix.
+    if col_bcast and N == 1:
+        ld = 0
+    else:
+        ld = max(M, mm_strides[1]) if order == cublaslt.Order.COL else max(N, mm_strides[0])
+
+    # Batch dimensions should be contiguous in memory, which we have already checked. The
+    # batch_offset should be based on the lowest stride in the batch dimension to account
+    # for embedded matrices.
+    batch_offset = min(batch_strides) if batch_strides else 0
+
+    return order, ld, batch_offset
+
+
+def get_mm_traits(a_layout, b_layout, c_layout, inplace, logger):
+    """
+    First check A and B compatibility:
+
+    1. Check MM compatibility (K):
+        a. First pad A and/or B MM dimensions if 1-D according to NumPy convention.
+        b. The padding is used to determine M, N, and K but should not appear in the output
+           dimensions.
+        c. If both A and B are N-D, the dimensions must match.
+    2. Check batch dimensions:
+        a. One of A or B can have missing batch extents, in which case it is broadcast,
+           otherwise
+        b. A and B must have the same batch ordering.
+        c. In addition, the batch dimensions must be tileable (contiguous in memory).
+
+    Then check C:
+
+    C can be None. If C is passed in, it must be a matrix. The batching rule is the
+    same as above.
+
+    For the inplace option (result == C), C cannot be broadcast and so must have batch
+    dimensions (if batched).
+    """
+    inplace_addendum = ""
+    if inplace:
+        inplace_addendum = "This is required for inplace operations."
+
+    a_shape, a_strides = list(a_layout.shape), list(a_layout.strides)
+    b_shape, b_strides = list(b_layout.shape), list(b_layout.strides)
+
+    a_batch_shape, a_mm_shape = a_shape[:-2], a_shape[-2:]
+    b_batch_shape, b_mm_shape = b_shape[:-2], b_shape[-2:]
+
+    a_batch_strides, a_mm_strides = a_strides[:-2], a_strides[-2:]
+    b_batch_strides, b_mm_strides = b_strides[:-2], b_strides[-2:]
+    d_mm_shape = []
+    if len(a_mm_shape) == 1:
+        s, d = a_mm_shape[0], a_mm_strides[0]
+        a_mm_shape = [1] + a_mm_shape
+        a_mm_strides = [s * d] + a_mm_strides
+    else:
+        d_mm_shape.append(a_mm_shape[0])  # The first mode for d applies only when a is not a vector.
+
+    if len(b_mm_shape) == 1:
+        s, d = b_mm_shape[0], b_mm_strides[0]
+        b_mm_shape = b_mm_shape + [1]
+        b_mm_strides = b_mm_strides + [s * d]
+    else:
+        d_mm_shape.append(b_mm_shape[1])  # The second mode for d applies only when b is not a vector.
+
+    logger.debug(f"The MM shape for operand A is {a_mm_shape} with strides {a_mm_strides}.")
+    logger.debug(f"The MM shape for operand B is {b_mm_shape} with strides {b_mm_strides}.")
+    logger.debug(f"The MM shape for operand D is {d_mm_shape}.")
+
+    M0, K0 = a_mm_shape
+    K1, N0 = b_mm_shape
+    if K0 != K1:
+        raise ValueError(
+            f"The 'K' extent must match for the operands: K={K0} in operand A is not equal to K={K1} in operand B."
+        )
+
+    # Check if batch dimensions of A and B are tileable as well as compatible.
+    batch_shape, batch_axis_order = [], ()
+    if len(a_batch_shape) > 0:
+        if not check_batch_tileable(a_batch_shape, a_batch_strides):
+            message = (
+                f"The batch layout for A corresponding to shape = {a_batch_shape} and strides = {a_batch_strides} "
+                "is currently not supported because it is not tileable."
+            )
+            raise ValueError(message)
+        logger.debug(
+            f"The batch layout for A corresponding to shape = {a_batch_shape} and strides = {a_batch_strides} IS tileable."
+        )
+        batch_shape = a_batch_shape
+        batch_axis_order = a_batch_axis_order = axis_order_in_memory(a_batch_strides)
+
+    if len(b_batch_shape) > 0:
+        if not check_batch_tileable(b_batch_shape, b_batch_strides):
+            message = (
+                f"The batch layout for B corresponding to shape = {b_batch_shape} and strides = {b_batch_strides} "
+                "is currently not supported because it is not tileable."
+            )
+            raise ValueError(message)
+        logger.debug(
+            f"The batch layout for B corresponding to shape = {b_batch_shape} and strides = {b_batch_strides} IS tileable."
+        )
+        batch_shape = b_batch_shape
+        batch_axis_order = b_batch_axis_order = axis_order_in_memory(b_batch_strides)
+
+    if len(a_batch_shape) > 0 and len(b_batch_shape) > 0:
+        if a_batch_shape != b_batch_shape:
+            raise ValueError(f"The batch dimensions of operands A {a_batch_shape} and B {b_batch_shape} must match.")
+        if a_batch_axis_order != b_batch_axis_order:
+            raise ValueError(f"The batch order of operands A {a_batch_axis_order} and B {b_batch_axis_order} must match.")
+
+    logger.debug(f"The batch shape is {batch_shape} with batch axis order {batch_axis_order}.")
+
+    batch_count = functools.reduce(operator.mul, batch_shape, 1)
+
+    # Create matrix layout traits.
+    a_order, a_ld, a_batch_offset = get_matrix_layout_traits(a_mm_shape, a_mm_strides, a_batch_strides, col_bcast=False)
+    a_layout_traits = LayoutTraits(
+        order=a_order,
+        ld=a_ld,
+        batch_offset=a_batch_offset,
+        is_conjugate=a_layout.is_conjugate,
+        mm_shape=a_mm_shape,
+        mm_strides=a_mm_strides,
+    )
+    logger.debug(f"The layout order for operand A is {a_order.name}, with LD {a_ld}, and batch offset {a_batch_offset}.")
+
+    b_order, b_ld, b_batch_offset = get_matrix_layout_traits(b_mm_shape, b_mm_strides, b_batch_strides, col_bcast=False)
+    b_layout_traits = LayoutTraits(
+        order=b_order,
+        ld=b_ld,
+        batch_offset=b_batch_offset,
+        is_conjugate=b_layout.is_conjugate,
+        mm_shape=b_mm_shape,
+        mm_strides=b_mm_strides,
+    )
+    logger.debug(f"The layout order for operand B is {b_order.name}, with LD {b_ld}, and batch offset {b_batch_offset}.")
+
+    # Process matrix c, if provided.
+    c_layout_traits = None
+    if c_layout is not None:
+        # 1. C cannot be a vector.
+        # 2. C can be a matrix of dimension (M, N) or (M, 1), broadcast in the latter case
+        #    and has to have contiguous strides.
+        # 3. C can be batched matrices of dimension (..., M, N) or (..., M, 1), broadcast in
+        #    the latter case and has to have contiguous strides.
+        c_shape, c_strides = list(c_layout.shape), list(c_layout.strides)
+
+        c_batch_shape, c_mm_shape = c_shape[:-2], c_shape[-2:]
+        c_batch_strides, c_mm_strides = c_strides[:-2], c_strides[-2:]
+        if len(c_mm_shape) == 1:
+            raise ValueError(f"C cannot be a vector. C shape: {c_mm_shape}")
+        logger.debug(f"The MM shape for operand C is {c_mm_shape} with strides {c_mm_strides}.")
+
+        Mc, Nc = c_mm_shape
+        if Mc != M0:
+            raise ValueError(f"The M dimension of the C matrix ({Mc}) must match the M dimension of A.")
+
+        if (inplace and Nc != N0) or (Nc != 1 and Nc != N0):
+            raise ValueError(f"The N dimension of the C matrix ({Nc}) must match the N dimension of B. {inplace_addendum}")
+
+        # For the inplace option, C must be batched if an operand is batched.
+        if inplace or len(c_batch_shape) > 0:
+            if c_batch_shape != batch_shape:
+                raise ValueError(
+                    f"The batch dimension of operand C {c_batch_shape} must match with that of the other operands "
+                    f"{batch_shape}. {inplace_addendum}"
+                )
+
+            if batch_shape and (c_batch_axis_order := axis_order_in_memory(c_batch_strides)) != batch_axis_order:
+                raise ValueError(
+                    f"The batch axis order of operand C {c_batch_axis_order} must match with that of the other "
+                    f"operands {batch_axis_order}."
+                )
+
+            if batch_shape and not check_batch_tileable(c_batch_shape, c_batch_strides):
+                message = (
+                    f"The batch layout for C corresponding to shape = {c_batch_shape} and strides = "
+                    f"{c_batch_strides} is currently not supported because it is not tileable."
+                )
+                raise ValueError(message)
+
+        # Inplace writes back into C, so the C layout must be unique.
+        # Otherwise distinct logical indices alias the same memory cell,
+        # producing overlapping writes / undefined behavior.
+        if inplace and not StridedLayout(c_shape, c_strides, itemsize=1).is_unique:
+            raise ValueError(
+                f"The layout for C with shape = {c_shape} and strides = {c_strides} is not "
+                f"supported for inplace operations: the index-to-memory mapping is not "
+                f"injective, which would cause overlapping writes into the same memory."
+            )
+
+        c_order, c_ld, c_batch_offset = get_matrix_layout_traits(
+            c_mm_shape, c_mm_strides, c_batch_strides, col_bcast=not inplace
+        )
+        c_layout_traits = LayoutTraits(order=c_order, ld=c_ld, batch_offset=c_batch_offset, is_conjugate=c_layout.is_conjugate)
+        logger.debug(f"The layout order for operand C is {c_order.name}, with LD {c_ld}, and batch offset {c_batch_offset}.")
+
+    return MMTraits(
+        M=M0,
+        N=N0,
+        K=K0,
+        d_mm_shape=d_mm_shape,
+        inplace=inplace,
+        c_layout=c_layout,
+        batch_count=batch_count,
+        batch_shape=batch_shape,
+        batch_axis_order=batch_axis_order,
+        a_layout_traits=a_layout_traits,
+        b_layout_traits=b_layout_traits,
+        c_layout_traits=c_layout_traits,
+    )
+
+
+def get_result_traits(mm_traits: MMTraits, epilog_ordering: cublaslt.Order | None, logger: logging.Logger) -> ResultTraits:
+    """
+    epilog_ordering = value of type cublaslt.Order or None.
+
+    The result layout is determined from:
+        - the ordering of operand c, if it is provided, or
+        - the epilog requirement, if it exists, or
+        - the ordering of operand a.
+
+    The result batch dimensions must have the same extents and axis order as the inputs. The
+    MM layout can be C or F.
+    """
+
+    # For inplace=True, the result traits are essentially the same as for operand C except
+    # for conjugation.
+    if mm_traits.inplace:
+        result_shape, result_strides = mm_traits.c_layout.shape, mm_traits.c_layout.strides
+        d_layout_traits = dataclasses.replace(mm_traits.c_layout_traits, is_conjugate=False)
+        return ResultTraits(result_shape=result_shape, result_strides=result_strides, d_layout_traits=d_layout_traits)
+
+    # The result shape is the batch shape + d_mm_shape.
+    result_shape = (*mm_traits.batch_shape, *mm_traits.d_mm_shape)
+
+    if mm_traits.c_layout_traits is not None:
+        result_ordering = mm_traits.c_layout_traits.order
+    elif epilog_ordering is not None:
+        result_ordering = epilog_ordering
+    else:
+        result_ordering = mm_traits.a_layout_traits.order
+
+    if result_ordering == cublaslt.Order.ROW:
+        d_order = list(range(len(mm_traits.d_mm_shape) - 1, -1, -1))
+    elif result_ordering == cublaslt.Order.COL:
+        d_order = list(range(len(mm_traits.d_mm_shape)))
+    else:
+        raise AssertionError("Internal Error.")
+
+    result_axis_order = [len(mm_traits.batch_axis_order) + a for a in d_order] + list(mm_traits.batch_axis_order)
+
+    # Calculate the result strides.
+    result_strides = calculate_strides(result_shape, result_axis_order)
+
+    # For degenerate matrices, we need to specify the result orientation.
+    result_orientation = None
+    if len(mm_traits.d_mm_shape) < 2:
+        if mm_traits.M == 1:
+            result_orientation = cublaslt.Order.ROW
+        elif mm_traits.N == 1:
+            result_orientation = cublaslt.Order.COL
+
+    # The result's traits.
+    d_batch_strides, d_mm_strides = (
+        result_strides[: len(mm_traits.batch_shape)],
+        result_strides[len(mm_traits.batch_shape) :],
+    )
+    order, d_ld, d_batch_offset = get_matrix_layout_traits(
+        mm_traits.d_mm_shape,
+        d_mm_strides,
+        d_batch_strides,
+        col_bcast=False,
+        ordering=result_ordering,
+        orientation=result_orientation,
+    )
+    assert order == result_ordering, (
+        f"Internal Error: d_order = {order.name}, result_ordering = {result_ordering.name}, mm_traits = {mm_traits}."
+    )
+    d_layout_traits = LayoutTraits(order=order, ld=d_ld, batch_offset=d_batch_offset, is_conjugate=False)
+    logger.debug(f"The layout order for operand D is {order.name}, with LD {d_ld}, and batch offset {d_batch_offset}.")
+
+    return ResultTraits(result_shape=result_shape, result_strides=result_strides, d_layout_traits=d_layout_traits)
+
+
+SHARED_MM_DOCUMENTATION = utils.COMMON_SHARED_DOC_MAP.copy()
+SHARED_MM_DOCUMENTATION.update(
+    {
+        "a": """\
+A tensor representing the first operand to the matrix multiplication (see `Semantics`_). The currently supported types
+are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
+        #
+        "b": """\
+A tensor representing the second operand to the matrix multiplication (see `Semantics`_). The currently supported types
+are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
+        #
+        "c": """\
+(Optional) A tensor representing the operand to add to the matrix multiplication result (see `Semantics`_). The currently
+supported types are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
+        #
+        "c_admonitions": """
+                .. versionchanged:: 0.3.0
+                    In order to avoid broadcasting behavior ambiguity, nvmath-python no longer
+                    accepts a 1-D (vector) `c`. Use a singleton dimension to convert your input
+                    array to 2-D.
+""",
+        #
+        "alpha": """\
+The scale factor for the matrix multiplication term as a real or complex number. The default is
+:math:`1.0`.""".replace("\n", " "),
+        #
+        "beta": """\
+The scale factor for the matrix addition term as a real or complex number. A value for `beta` must be provided if
+operand `c` is specified.""".replace("\n", " "),
+        #
+        "quantization_scales": """\
+Specify scale factors for the matrix multiplication as a :class:`~nvmath.linalg.advanced.MatmulQuantizationScales`
+object. Alternatively, a `dict` containing the parameters for the
+:class:`~nvmath.linalg.advanced.MatmulQuantizationScales`
+constructor can also be provided. The scale factors can be provided as scalars or tensors.
+If a scale factor is provided as a tensor, it must be from the same package and on the same
+memory space (CPU or GPU device) as the operands of the matmul.
+If a scale factor is provided as a scalar, and the execution space is GPU,
+a CPU->GPU copy is inevitable. To avoid this copy, provide
+the quantization scale as one-element array on the GPU.
+Allowed and required only for narrow-precision (FP8 and lower) operations.""".replace("\n", " "),
+        #
+        "algorithms": """\
+A sequence of :class:`Algorithm` objects that can be directly provided to bypass planning. The algorithm objects must be
+compatible with the matrix multiplication. A typical use for this option is to provide algorithms serialized (numpy)
+from a previously planned and autotuned matrix multiplication.""".replace("\n", " "),
+        #
+        "epilog": """\
+Specify an epilog :math:`\\mathcal{F}` as an object of type :class:`MatmulEpilog` to apply to the result of the matrix
+multiplication: :math:`\\mathcal{F}(\\alpha a @ b + \\beta c)`. The default is no epilog. See `cuBLASLt documentation
+<https://docs.nvidia.com/cuda/cublas/#cublasltepilogue-t>`_ for the list of available epilogs.""".replace("\n", " "),
+        #
+        "epilog_inputs": """\
+Specify the additional inputs needed for the selected epilog as a dictionary, where the key is the epilog input name and
+the value is the epilog input. The epilog input must be a tensor with the same package and in the same memory space as
+the operands (see the constructor for more information on the operands). If the required epilog inputs are not provided,
+an exception is raised that lists the required epilog inputs. Some epilog inputs are generated by other epilogs. For
+example, the epilog input for :class:`MatmulEpilog.DRELU` is generated by matrix multiplication with the same operands
+using :class:`MatmulEpilog.RELU_AUX`. """.replace("\n", " "),
+        #
+        "qualifiers": """\
+If desired, specify the matrix qualifiers as a :class:`numpy.ndarray` of
+:class:`~nvmath.linalg.advanced.matrix_qualifiers_dtype` objects of length 3 corresponding to the operands `a`, `b`, and
+`c`. See :ref:`matrix-tensor-qualifiers` for the motivation behind
+qualifiers.""".replace("\n", " "),
+        #
+        "options": """\
+Specify options for the matrix multiplication as a :class:`~nvmath.linalg.advanced.MatmulOptions` object. Alternatively,
+a `dict` containing the parameters for the ``MatmulOptions`` constructor can also be provided. If not specified, the
+value will be set to the default-constructed ``MatmulOptions`` object.""".replace("\n", " "),
+        #
+        "preferences": """\
+This parameter specifies the preferences for planning as a :class:`MatmulPlanPreferences` object. Alternatively, a
+dictionary containing the parameters for the :class:`MatmulPlanPreferences` constructor can also be provided. If not
+specified, the value will be set to the default-constructed :class:`MatmulPlanPreferences` object.
+""".replace("\n", " "),
+        #
+        "result": """\
+The result of the specified matrix multiplication (epilog applied), which remains on the same device and belongs to the
+same package as the input operands. If an epilog (like :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_AUX`) that
+results in extra output is used, or an extra output is requested (for example by setting
+:attr:`~nvmath.linalg.advanced.MatmulOptions.result_amax` option in ``options`` argument),
+a tuple is returned with the first element being the matrix multiplication result (epilog applied) and the second element
+being the auxiliary output provided as a `dict`. """.replace("\n", " "),
+        #
+        "narrow_precision": """\
+
+        .. _narrow_precision:
+
+        Matrix multiplication with narrow-precision operands is supported, in FP8, MXFP8, and NVFP4 formats.
+
+        **FP8 and MXFP8**
+
+        FP8 and MXFP8 use ``float8_e4m3fn`` or ``float8_e5m2`` data types. The difference is the scaling mode:
+        FP8 (``block_scaling=False``) uses per-tensor scaling where a single scalar scale is applied to each operand;
+        MXFP8 (``block_scaling=True``) uses microscaling with 32-element blocks arranged in 128x128 tiles.
+
+        .. note::
+
+            FP8 and MXFP8 matrix multiplication requires **CUDA Toolkit 12.8 or newer**.
+            **FP8 requires a device with compute capability 8.9 or higher** (Ada, Hopper, Blackwell or newer architecture).
+            **MXFP8 requires a device with compute capability 10.0 or higher** (Blackwell or newer architecture).
+            Please refer to the `compute capability table <https://developer.nvidia.com/cuda-gpus>`_
+            to check the compute capability of your device.
+
+        For FP8 operations:
+
+        * For each operand a scaling factor needs to be specified via ``quantization_scales`` argument.
+        * Maximum absolute value of the result (amax) can be requested via
+          :attr:`~nvmath.linalg.advanced.MatmulOptions.result_amax` option in ``options`` argument.
+        * Custom result type (both FP8 and non-FP8) can be requested via
+          :attr:`~nvmath.linalg.advanced.MatmulOptions.result_type` option in ``options`` argument.
+
+        For MXFP8 operations:
+
+        * 1-D (vector) operands are not supported. Both ``a`` and ``b`` must be at least 2-D matrices.
+        * Broadcasting of batch dimensions is not supported. The batch shapes of ``a`` and ``b`` must match exactly.
+        * All operand dimensions (M, N, K) must be multiples of 128.
+        * :attr:`~nvmath.linalg.advanced.MatmulOptions.block_scaling` option must be set to ``True`` and
+          block scaling factors need to be specified via ``quantization_scales`` argument.
+          Utilities in :mod:`nvmath.linalg.advanced.helpers.matmul` can be used to create and modify
+          block scaling factors, see e.g. :func:`~nvmath.linalg.advanced.helpers.matmul.create_mxfp8_scale`.
+        * When the result type is a narrow-precision data type, the auxiliary output ``"d_out_scale"``
+          will be returned containing the scales used for result quantization.
+
+        *Layout Requirements*
+
+        Due to the requirements of narrow-precision GEMM kernels, the contracting dimension
+        K must be contiguous (stride-1) for both operands. The following layout constraints
+        apply to both FP8 and MXFP8:
+
+        * Operand ``a`` must be ``(..., M, K)`` with ``stride[-1] == 1`` and
+          ``stride[-2] >= K`` (row-major).  The leading dimension (``stride[-2]``) can be
+          larger than ``K`` to support sliced or padded views.
+
+        * Operand ``b`` must be ``(..., K, N)`` with ``stride[-2] == 1`` and
+          ``stride[-1] >= K`` (column-major).  The leading dimension (``stride[-1]``) can
+          be larger than ``K`` to support sliced or padded views.
+
+        .. attention::
+
+            Epilog support for MXFP8 is still evolving in the underlying cuBLASLt library,
+            so not every combination of epilog, data type, and layout is guaranteed to work.
+            If running into unsupported combinations, a cuBLASLt error will be raised
+            either at planning time or at execution time that will reveal the
+            root cause. These gaps are expected to be filled in future cuBLASLt releases.
+
+        For more details on the FP8 and MXFP8 formats in cuBLAS,
+        see the `cublasLtMatmul documentation <https://docs.nvidia.com/cuda/cublas/#cublasltmatmul>`_.
+
+        **NVFP4**
+
+        .. versionadded:: 0.9.0
+            NVFP4 support.
+
+        NVFP4 uses ``float4_e2m1fn_x2`` data type with block scaling (16-element blocks arranged in 128x64 tiles).
+
+        .. note::
+
+            NVFP4 matrix multiplication currently requires **CUDA Toolkit 12.8 or newer**,
+            **a device with compute capability 10.0 or higher** (Blackwell or newer architecture),
+            and **PyTorch 2.9 or newer** for ``float4_e2m1fn_x2`` dtype support.
+            Please refer to the `compute capability table <https://developer.nvidia.com/cuda-gpus>`_
+            to check the compute capability of your device.
+
+        For NVFP4 operations:
+
+        * 1-D (vector) operands are not supported. Both ``a`` and ``b`` must be at least 2-D matrices.
+        * Broadcasting of batch dimensions is not supported. The batch shapes of ``a`` and ``b`` must match exactly.
+        * The outer dimensions of ``a`` and ``b`` (M and N) must be multiples of 128,
+          and the contracting dimension K must be a multiple of 64.
+        * :attr:`~nvmath.linalg.advanced.MatmulOptions.block_scaling` option must be set to ``True`` and
+          block scaling factors need to be specified via ``quantization_scales`` argument.
+        * When the result type is a narrow-precision data type, the auxiliary output ``"d_out_scale"``
+          will be returned containing the scales used for result quantization.
+
+        *Layout and Packing Requirements*
+
+        FP4 data is per-byte packed: ``float4_e2m1fn_x2`` stores 2 FP4 values per byte.
+        The block scaling (VEC16_UE4M3) assigns one scale factor per 16 consecutive elements
+        along the innermost (stride-1) dimension of each operand. The layout requirements
+        below ensure that this innermost dimension corresponds to the contracting dimension K
+        for both operands.
+
+        * Operand ``a`` must be ``(..., M, K//2)`` with ``stride[-1] == 1`` and
+          ``stride[-2] >= K//2``, i.e., row-wise packed along K. Note that the
+          leading dimension (``stride[-2]``) can be larger than ``K//2`` to support
+          sliced views, as long as the stride remains 16-byte aligned.
+
+        * Operand ``b`` must be ``(..., K//2, N)`` with ``stride[-2] == 1`` and
+          ``stride[-1] >= K//2``, i.e., column-wise packed along K. Note that the
+          leading dimension (``stride[-1]``) can be larger than ``K//2`` to support
+          sliced views, as long as the stride remains 16-byte aligned.
+
+        If your data has the stride-1 axis along a dimension other than K,
+        you must repack it before calling :func:`matmul`.
+
+        When the result type is also FP4, the output is packed along a dimension that
+        depends on the result layout order:
+
+        * **Row-major result**: packed along N — shape ``(..., M, N//2)``, strides ``(..., N//2, 1)``.
+        * **Column-major result**: packed along M — shape ``(..., M//2, N)``, strides ``(..., 1, M//2)``.
+
+        The result layout order is determined by the following priority:
+
+        1. If ``c`` is provided, the result inherits ``c``'s layout order.
+        2. Otherwise, if the epilog requests a specific layout, that layout is used.
+        3. Otherwise, the result inherits ``a``'s layout order as a fallback.
+
+        *Epilog Support*
+
+        NVFP4 matmul supports epilogs. The following have been verified:
+
+        * ``RELU``, ``GELU`` -- with both row-major and column-major output.
+        * ``BIAS``, ``RELU_BIAS``, ``GELU_BIAS`` -- with column-major output only
+          (``BIAS`` with ``float16`` C/D requires cuBLASLt >= 13.0).
+
+        .. attention::
+
+            Epilog support for NVFP4 is still evolving in the underlying cuBLASLt library,
+            so not every combination of epilog, data type, and layout is guaranteed to work.
+            If running into unsupported combinations, a cuBLASLt error will be raised
+            either at planning time or at execution time that will reveal the
+            root cause. These gaps are expected to be filled in future cuBLASLt releases.
+
+        *Helper Functions*
+
+        The :mod:`nvmath.linalg.advanced.helpers.matmul` module provides helpers for working with
+        FP4 encoding/decoding and NVFP4 block scales, see e.g.
+        :func:`~nvmath.linalg.advanced.helpers.matmul.quantize_to_fp4`,
+        :func:`~nvmath.linalg.advanced.helpers.matmul.unpack_fp4`,
+        :func:`~nvmath.linalg.advanced.helpers.matmul.get_block_scale_offset`,
+        :func:`~nvmath.linalg.advanced.helpers.matmul.to_block_scale`,
+        :func:`~nvmath.linalg.advanced.helpers.matmul.expand_block_scale`.
+
+        For more details on the NVFP4 format in cuBLAS,
+        see the `cublasLtMatmul documentation <https://docs.nvidia.com/cuda/cublas/#cublasltmatmul>`_.
+        For usage examples, see the relevant files in the
+        `examples/linalg/advanced/matmul
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
+        directory.
+""".strip(),
+        #
+        "thread_safety": """\
+A :class:`Matmul` instance should only be used by its creating thread; it is not safe to use
+concurrently from multiple threads. The free function :func:`matmul` creates a fresh
+instance per call and is safe to call concurrently from multiple threads.
+""".replace("\n", " "),
+        #
+        "semantics": """\
+        .. _semantics:
+
+        The semantics of the matrix multiplication follows :external:py:data:`numpy.matmul` semantics, with some restrictions on
+        broadcasting. In addition, the semantics for the fused matrix addition are described below.
+
+        .. note::
+
+            For narrow-precision formats (FP8, MXFP8, NVFP4), some of the rules below are
+            restricted — see the :ref:`narrow-precision <narrow_precision>` section for details.
+
+        * For in-place matrix multiplication (where the result is written into `c`) the result has the same shape as `c`.
+        * If arguments `a` and `b` are matrices, they are multiplied according to the rules of matrix multiplication.
+        * If argument `a` is 1-D, it is promoted to a matrix by prefixing ``1`` to its dimensions. After matrix
+          multiplication, the prefixed ``1`` is removed from the result's dimensions if the operation is not in-place.
+        * If argument `b` is 1-D, it is promoted to a matrix by appending ``1`` to its dimensions. After matrix
+          multiplication, the appended ``1`` is removed from the result's dimensions if the operation is not in-place.
+        * If `a` or `b` is N-D (N > 2), then the operand is treated as a batch of matrices. If both `a` and `b` are N-D,
+          their batch dimensions must match. If exactly one of `a` or `b` is N-D, the other operand is broadcast.
+        * The operand for the matrix addition `c` may be a matrix of shape (M, 1) or (M, N), or the batched versions
+          (..., M, 1) or (..., M, N). Here M and N are the dimensions of the result of the matrix multiplication. If N = 1, the
+          columns of `c` are broadcast for the addition; the rows of `c` are never broadcast. If batch dimensions are not
+          present, `c` is broadcast across batches as needed. If the operation is in-place, `c` cannot be broadcast since
+          it must be large enough to hold the result.
+        * Similarly, when operating on a batch, auxiliary outputs are 3-D for all epilogs. Therefore, epilogs that return 1-D
+          vectors of length N in non-batched mode return 3-D matrices of size (batch, N, 1) in batched mode.
+
+        For narrow-precision operations (FP8 and lower), further restrictions apply;
+        see the narrow-precision support section below.
+""".strip(),
+    }
+)
+
+
+class InvalidMatmulState(Exception):
+    pass
+
+
+def _check_extents(shape: tuple, name: str):
+    if any(e <= 0 for e in shape):
+        message = f"The specified extents {shape} for operand {name} are not valid. The extents must be strictly positive. "
+        raise ValueError(message)
+
+
+def _detect_and_validate_ab_operands_for_fp4(a, b, logger) -> bool:
+    """
+    Detect whether A and B are FP4 operands, and if so validate their format.
+
+    Returns True if both operands are FP4 (float4_e2m1fn_x2), False if neither
+    is FP4, and raises ValueError if exactly one is FP4 (mixed operands are not
+    supported).
+
+    When both operands are FP4, the function validates layout and dimensions.
+    FP4 data is always packed: float4_e2m1fn_x2 stores 2 FP4 values per byte.
+    A and B must be in the following format:
+
+    - A must be (M, K//2) with stride[-1]==1 and stride[-2] >= K//2, i.e. row-wise
+      packed along K. The leading dimension can be larger than K//2, namely
+      padded/sliced views are supported.
+    - B must be (K//2, N) with stride[-2]==1 and stride[-1] >= K//2, i.e. column-wise
+      packed along K. The leading dimension can be larger than K//2, namely
+      padded/sliced views are supported.
+
+    Why packing must be along K:
+
+    cuBLASLt hardcodes transa=OP_T for MXP4, and expects both A and B to be provided
+    in column-major layout. In column-major,
+    the stride-1 (innermost) dimension corresponds to the rows of the stored matrix.
+    Combined with OP_T, these rows map to the contracting dimension K from the user's
+    logical perspective. This means:
+
+    - User's row-major A (M, K//2) with stride[-1]==1: cuBLASLt interprets this
+      as a column-major (K//2, M) matrix (K//2 rows along the stride-1 axis,
+      M columns). With transa=OP_T, the logical GEMM operand becomes (M, K),
+      with the physical stride-1 axis aligned to the contraction dimension K.
+    - User's column-major B (K//2, N) with stride[-2]==1: cuBLASLt interprets
+      this as a column-major (K//2, N) matrix directly. With transb=OP_N
+      (the default), the logical GEMM operand is (K, N), again with the
+      physical stride-1 axis aligned to K.
+
+    NVFP4 block scaling (VEC16_UE4M3) assigns one scale per 16 elements along the
+    innermost (stride-1) dimension. Because of the layout above, this innermost
+    dimension is K, so the scales are naturally aligned with the contraction axis.
+
+    Dimension requirements:
+    - M (outer dim of A): must be a multiple of 128
+    - N (outer dim of B): must be a multiple of 128
+    - K (contraction dim): must be a multiple of 64
+    """
+    a_is_fp4 = a.dtype == "float4_e2m1fn_x2"
+    b_is_fp4 = b.dtype == "float4_e2m1fn_x2"
+
+    # Neither is FP4 means this is not an FP4 operation
+    if not a_is_fp4 and not b_is_fp4:
+        return False
+
+    # One is FP4 and the other is not: unsupported
+    if a_is_fp4 != b_is_fp4:
+        raise ValueError(
+            f"Mixed FP4/non-FP4 A/B operands are not supported. "
+            f"Got A dtype={a.dtype}, B dtype={b.dtype}. "
+            f"Both A and B operands must be float4_e2m1fn_x2 or neither."
+        )
+
+    logger.info(f"FP4 validation: A shape={a.shape}, strides={a.strides}")
+    logger.info(f"FP4 validation: B shape={b.shape}, strides={b.strides}")
+
+    if len(a.shape) < 2:
+        raise ValueError(f"FP4 operand A must be at least 2-D, got shape {a.shape}.")
+    if len(b.shape) < 2:
+        raise ValueError(f"FP4 operand B must be at least 2-D, got shape {b.shape}.")
+
+    # Validate A: Must be row-major (..., M, K//2) with strides (..., K//2, 1)
+    if a.strides[-1] != 1:
+        raise ValueError(f"FP4 operand A must be row-major with stride[-1]=1. Got shape {a.shape} with strides {a.strides}.")
+
+    # Validate B: Must be column-major (..., K//2, N) with strides (..., 1, K//2)
+    if b.strides[-2] != 1:
+        raise ValueError(
+            f"FP4 operand B must be column-major with stride[-2]=1. "
+            f"Got shape {b.shape} with strides {b.strides}. "
+            f"Use torch.as_strided() to create column-major view."
+        )
+
+    # Extract dimensions from packed tensors:
+    # A is (M, K//2) -> M = a.shape[-2], K = a.shape[-1] * 2
+    # B is (K//2, N) -> N = b.shape[-1], K = b.shape[-2] * 2
+    M = a.shape[-2]
+    K_from_a = a.shape[-1] * 2  # Unpack: K//2 -> K
+    K_from_b = b.shape[-2] * 2  # Unpack: K//2 -> K
+    N = b.shape[-1]
+
+    # K should match between A and B
+    if K_from_a != K_from_b:
+        raise ValueError(
+            f"Dimension mismatch: K from A ({K_from_a}) != K from B ({K_from_b}). "
+            f"A shape is {a.shape} (M, K//2), B shape is {b.shape} (K//2, N)."
+        )
+    K = K_from_a
+
+    logger.info(f"FP4 dimensions: M={M}, N={N}, K={K}")
+
+    # Validate block scaling dimension requirements
+    errors = []
+    if M % 128 != 0:
+        errors.append(f"M={M} must be divisible by 128")
+    if N % 128 != 0:
+        errors.append(f"N={N} must be divisible by 128")
+    if K % 64 != 0:
+        errors.append(f"K={K} must be divisible by 64")
+
+    if errors:
+        raise ValueError(
+            f"FP4 block scaling dimension requirements not met: {', '.join(errors)}. "
+            f"NVFP4 requires M and N divisible by 128, K divisible by 64."
+        )
+
+    return True
+
+
+def _create_fp4_ab_layouts(a_shape, a_strides, a_is_conjugate, b_shape, b_strides, b_is_conjugate):
+    """
+    Create MatrixLayout objects for packed FP4 operands A and B.
+
+    cuBLASLt always computes A^T @ B (transa=OP_T is hardcoded). The user provides
+    A as (M, K//2) row-wise packed and B as (K//2, N) column-wise packed. We convert
+    these to logical shapes/strides that cuBLASLt expects:
+
+    - A is reinterpreted as (K//2, M) internally (swap last two dims), so that
+      cuBLASLt computes (A^T)^T @ B = A @ B.
+    - B is kept as is, but strides are converted from packed to logical.
+
+    Since float4_e2m1fn_x2 packs 2 values per byte, logical strides are 2x the
+    packed strides along the K dimension.
+
+    Args:
+        a_shape: Physical shape of A, e.g. (..., M, K//2).
+        a_strides: Physical strides of A.
+        a_is_conjugate: Whether A is conjugated.
+        b_shape: Physical shape of B, e.g. (..., K//2, N).
+        b_strides: Physical strides of B.
+        b_is_conjugate: Whether B is conjugated.
+
+    Returns:
+        (a_layout, b_layout) — MatrixLayout objects with logical shapes/strides.
+    """
+    M, K_packed = a_shape[-2], a_shape[-1]
+    K_logical = K_packed * 2
+
+    # A: swap last two stride dimensions (transpose for cuBLASLt), then convert to logical
+    a_transposed_strides = (*a_strides[:-2], a_strides[-1], a_strides[-2])
+    a_logical_shape = (*a_shape[:-2], M, K_logical)
+    a_batch_strides = tuple(s * 2 for s in a_transposed_strides[:-2])
+    a_logical_strides = (*a_batch_strides, a_transposed_strides[-1] * 2, a_transposed_strides[-2])
+
+    # B: keep dimension order, convert packed strides to logical
+    N = b_shape[-1]
+    b_logical_shape = (*b_shape[:-2], K_logical, N)
+    b_batch_strides = tuple(s * 2 for s in b_strides[:-2])
+    b_logical_strides = (*b_batch_strides, b_strides[-2], b_strides[-1] * 2)
+
+    a_layout = MatrixLayout(a_logical_shape, a_logical_strides, a_is_conjugate)
+    b_layout = MatrixLayout(b_logical_shape, b_logical_strides, b_is_conjugate)
+    return a_layout, b_layout
+
+
+def _create_d_out_scale_and_scale_mode(result_class, num_output_elements, d_dtype_width, device_id, stream_holder):
+    """
+    Create the d_out_scale tensor and determine the scale mode for block-scaled output.
+    Block-scaled narrow-precision output (FP4 or FP8) requires a scale tensor where each
+    scale covers a group of output elements:
+
+    - FP4 (VEC16_UE4M3): each scale covers 16 elements, stored as ``float8_e4m3fn``
+    - FP8 (VEC32_UE8M0): each scale covers 32 elements, stored as ``uint8``.
+
+    Args:
+        result_class: The tensor class used to create the output tensor.
+        num_output_elements: Total number of output elements.
+        d_dtype_width: Result dtype width.
+        device_id: The CUDA device id.
+        stream_holder: The stream holder for tensor allocation.
+
+    Returns:
+        (d_out_scale_tensor, d_out_scale_mode) — the allocated scale tensor and the
+        corresponding ``cublaslt.MatmulMatrixScale`` mode.
+    """
+    if d_dtype_width == 4:
+        scale_mode = cublaslt.MatmulMatrixScale.VEC16_UE4M3
+        scale_group_size = 16
+        scale_dtype = "float8_e4m3fn"  # UE4M3 format, consistent with input A/B scales
+    else:
+        assert d_dtype_width == 8, "Internal error."
+        scale_mode = cublaslt.MatmulMatrixScale.VEC32_UE8M0
+        scale_group_size = 32
+        scale_dtype = "uint8"  # UE8M0 format
+
+    num_scales = num_output_elements // scale_group_size
+    d_out_scale = utils.create_empty_tensor(
+        result_class,
+        (num_scales,),
+        scale_dtype,
+        device_id,
+        stream_holder,
+        verify_strides=False,
+    )
+    return d_out_scale, scale_mode
+
+
+def _realloc_if_needed_and_copy_to_mirror(src_holder, dest_holder, destination_device_id, stream_holder):
+    """
+    Copy a CPU source operand into its device-side mirror (``dest_holder``).
+
+    Two cases:
+
+    * **Buffer exists** (``dest_holder.tensor is not None``): copy the
+      source data into the existing device buffer in-place via ``copy_()``.
+    * **Buffer released** (``dest_holder.tensor is None``, e.g. after
+      ``release_operands()``): call ``src_holder.to()`` to allocate a new
+      device buffer and copy the data into it, then assign the underlying
+      raw tensor to ``dest_holder.tensor``. The ``dest_holder`` wrapper
+      object itself is never replaced, so its concrete type is preserved.
+
+    Args:
+        src_holder: Wrapped CPU source operand.
+        dest_holder: Device-side tensor wrapper whose ``.tensor`` may be None.
+            Note: ``dest_holder.device_id`` is not usable when ``.tensor`` is
+            None, hence the explicit ``destination_device_id`` parameter.
+        destination_device_id: CUDA device ordinal (int, never ``"cpu"``)
+            since this function is only used for CPU-to-GPU mirroring.
+        stream_holder: Stream wrapper for the copy / allocation.
+
+    Returns:
+        True if a new buffer was allocated (data pointer changed),
+        False if the existing buffer was reused via in-place copy.
+    """
+    if dest_holder.tensor is not None:
+        dest_holder.copy_(src_holder, stream_holder=stream_holder)
+        return False
+    dest_holder.tensor = src_holder.to(destination_device_id, stream_holder).tensor
+    return True
+
+
+@utils.docstring_decorator(SHARED_MM_DOCUMENTATION, skip_missing=False)
+class Matmul:
+    """
+    Create a stateful object encapsulating the specified matrix multiplication computation
+    :math:`\\alpha a @ b + \\beta c` and the required resources to perform the operation.  A
+    stateful object can be used to amortize the cost of preparation (planning in the case of
+    matrix multiplication) across multiple executions (also see the :ref:`Stateful APIs
+    <host api types>` section).
+
+    The function-form API :func:`matmul` is a convenient alternative to using stateful
+    objects for *single* use (the user needs to perform just one matrix multiplication, for
+    example), in which case there is no possibility of amortizing preparatory costs. The
+    function-form APIs are just convenience wrappers around the stateful object APIs.
+
+    Using the stateful object typically involves the following steps:
+
+    1. **Problem Specification**: Initialize the object with a defined operation and
+       options.
+    2. **Preparation**: Use :meth:`plan` to determine the best algorithmic implementation
+       for this specific matrix multiplication operation.
+    3. **Execution**: Perform the matrix multiplication computation with :meth:`execute`.
+    4. **Resource Management**: Ensure all resources are released either by explicitly
+       calling :meth:`free` or by managing the stateful object within a context manager.
+
+    Detailed information on what's happening in the various phases described above can be
+    obtained by passing in a :class:`logging.Logger` object to :class:`MatmulOptions` or by
+    setting the appropriate options in the root logger object, which is used by default:
+
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     format="%(asctime)s %(levelname)-8s %(message)s",
+        ...     datefmt="%m-%d %H:%M:%S",
+        ... )
+
+    A user can select the desired logging level and, in general, take advantage of all of
+    the functionality offered by the Python `logging` module.
+
+    Args:
+        a: {a}
+
+        b: {b}
+
+        c: {c}
+            {c_admonitions}
+
+        alpha: {alpha}
+
+        beta: {beta}
+
+        qualifiers: {qualifiers}
+
+        options: {options}
+
+        stream: {stream}
+
+        quantization_scales: {quantization_scales}
+
+    Semantics:
+        {semantics}
+
+    Narrow-precision support:
+        {narrow_precision}
+
+    Thread Safety:
+        {thread_safety}
+
+    .. seealso::
+        :meth:`autotune`, :meth:`plan`, :meth:`reset_operands`,
+        :meth:`reset_operands_unchecked`, :meth:`execute`
+
+    Examples:
+
+        >>> import numpy as np
+        >>> import nvmath
+
+        Create two 2-D float64 ndarrays on the CPU:
+
+        >>> M, N, K = 1024, 1024, 1024
+        >>> a = np.random.rand(M, K)
+        >>> b = np.random.rand(K, N)
+
+        We will define a matrix multiplication operation followed by a RELU epilog function
+        using the specialized matrix multiplication interface.
+
+        Create a Matmul object encapsulating the problem specification above:
+
+        >>> mm = nvmath.linalg.advanced.Matmul(a, b)
+
+        Options can be provided above to control the behavior of the operation using the
+        `options` argument (see :class:`MatmulOptions`).
+
+        Next, plan the operation. The epilog is specified, and optionally, preferences can
+        be specified for planning:
+
+        >>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
+        >>> algorithms = mm.plan(epilog=epilog)
+
+        Certain epilog choices (like :attr:`nvmath.linalg.advanced.MatmulEpilog.BIAS`)
+        require additional input provided using the `epilog_inputs` argument to
+        :meth:`plan`.
+
+        Now execute the matrix multiplication, and obtain the result `r1` as a NumPy
+        ndarray.
+
+        >>> r1 = mm.execute()
+
+        Finally, free the object's resources. To avoid having to explicitly make this
+        call, it's recommended to use the Matmul object as a context manager as shown below,
+        if possible.
+
+        >>> mm.free()
+
+        Note that all :class:`Matmul` methods execute on the current stream by default.
+        Alternatively, the `stream` argument can be used to run a method on a specified
+        stream.
+
+        Let's now look at the same problem with CuPy ndarrays on the GPU.
+
+        >>> import cupy as cp
+        >>> a = cp.random.rand(M, K)
+        >>> b = cp.random.rand(K, N)
+
+        Create an Matmul object encapsulating the problem specification described earlier
+        and use it as a context manager.
+
+        >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
+        ...     algorithms = mm.plan(epilog=epilog)
+        ...
+        ...     # Execute the operation to get the first result.
+        ...     r1 = mm.execute()
+        ...
+        ...     # Update operands 'a' and 'b' in-place (see reset_operands() for an
+        ...     # alternative).
+        ...     a[:] = cp.random.rand(M, K)
+        ...     b[:] = cp.random.rand(K, N)
+        ...
+        ...     # Execute the operation to get the new result.
+        ...     r2 = mm.execute()
+
+
+        All the resources used by the object are released at the end of the block.
+
+        Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
+        directory.
+    """
+
+    def __init__(
+        self,
+        a,
+        b,
+        /,
+        c=None,
+        *,
+        alpha=None,
+        beta=None,
+        qualifiers=None,
+        quantization_scales=None,
+        options: _configuration.MatmulOptions | dict[str, typing.Any] | None = None,
+        stream: utils.AnyStream | int | None = None,
+    ):
+        options = utils.check_or_create_options(_configuration.MatmulOptions, options, "Matrix multiplication options")
+        assert options is not None
+        self.options = options
+
+        if c is None and options.inplace:
+            raise ValueError("The operation cannot be inplace if operand C is not provided.")
+
+        self.logger = options.logger if options.logger is not None else logging.getLogger()
+
+        if options.inplace and options.result_type is not None:
+            self.logger.warning(
+                f"Matmul: The provided result type {options.result_type} in options is ignored since \
+the operation is in-place."
+            )
+        self.inplace = options.inplace
+
+        def check_dtype(dtype, operand_name):
+            if dtype not in SUPPORTED_TYPES:
+                raise ValueError(f"The dtype of operand {operand_name} ({dtype}) is not supported.")
+
+        # The matrix multiplication has two required operands 'a' and 'b', and one optional
+        # operand 'c'.
+        a = tensor_wrapper.wrap_operand(a)
+        b = tensor_wrapper.wrap_operand(b)
+        check_dtype(a.dtype, "A")
+        check_dtype(b.dtype, "B")
+        self.logger.info("= SPECIFICATION PHASE =")
+        if self.inplace:
+            self.logger.info("The MM operation will be performed in-place (the result will be written into operand C).")
+        self.logger.info(f"The data type of operand A is '{a.dtype}', and that of operand B is '{b.dtype}'.")
+
+        # Detect FP4 A/B operands, and validate compatibility/format.
+        # This check is done as early as possible.
+        self.using_fp4_ab = _detect_and_validate_ab_operands_for_fp4(a, b, self.logger)
+        if self.using_fp4_ab and not options.block_scaling:
+            # FP4 only supports block scaling, at least for now
+            raise ValueError(
+                f"When using FP4 (float4_e2m1fn_x2) A, B operands, block_scaling=True is required. "
+                f"Got block_scaling={options.block_scaling}."
+            )
+
+        self.num_operands = 2
+        if c is not None:
+            self.num_operands = 3
+            c = tensor_wrapper.wrap_operand(c)
+            if len(c.shape) < 2:
+                raise ValueError(
+                    "In order to avoid broadcasting behavior ambiguity, `c` must be at least 2-D. "
+                    "Use a singleton dimension to convert your input array to 2-D."
+                )
+            check_dtype(c.dtype, "C")
+            self.logger.info(f"The data type of operand C is {c.dtype}.")
+
+        if c is not None and beta is None:
+            raise ValueError("A value for beta must be provided if operand C is provided.")
+
+        if (a.dtype, b.dtype) not in NAMES_TO_DEFAULT_SCALE_TYPE:
+            raise ValueError(f"Unsupported combination of dtypes for operands A {a.dtype} and B {b.dtype}.")
+
+        # Currently, a.dtype != b.dtype is only supported for FP8 (different FP8 kinds are
+        # allowed), so we assume that A and B have equal width.
+        self.input_type_width = typemaps.NAME_TO_DATA_WIDTH[a.dtype]
+
+        assert self.num_operands == 2 or self.num_operands == 3, "Internal Error."
+
+        _check_extents(a.shape, "a")
+        _check_extents(b.shape, "b")
+        if c is not None:
+            _check_extents(c.shape, "c")
+
+        # Infer the library package & device ID the operands belong to.
+        operands = [a, b]
+        if self.num_operands == 3:
+            operands.append(c)
+        self.operands: MutableSequence[utils.TensorHolder] | None = operands
+
+        self.package = utils.get_operands_package(operands)
+        # Package used for internal operations (stream + allocator). NumPy uses cuda.core.
+        self.internal_op_package = "cuda" if self.package == "numpy" else self.package
+        self.memory_space = "cuda"
+        self.device_id = utils.get_operands_device_id(operands)
+        if self.device_id == "cpu":
+            self.memory_space = "cpu"
+            self.device_id = options.device_id
+        self.logger.info(
+            f"The input operands' memory space is {self.memory_space}, and the execution space is on device {self.device_id}."
+        )
+
+        # Allocate device memory (in stream context) if needed.
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+        self.logger.info(f"The specified stream for the Matmul ctor is {stream_holder.obj}.")
+
+        # Copy operands to device (and store reference to CPU operand), if needed.
+        self.cpu_c_ref = None
+        if self.memory_space == "cpu":
+            if self.inplace:
+                self.cpu_c_ref = self.operands[2]  # Hold reference, needed for inplace operations.
+            self.operands = list(tensor_wrapper.to(self.operands, self.device_id, stream_holder))
+
+        # Set qualifiers.
+        self.qualifiers = qualifiers if qualifiers is not None else np.zeros((3,), dtype=_configuration.matrix_qualifiers_dtype)
+        if self.qualifiers.dtype != _configuration.matrix_qualifiers_dtype:
+            raise ValueError(
+                "The qualifiers must be specified as a NumPy array of length 3 corresponding to the operands A, B, and "
+                "C of type 'matrix_qualifiers_dtype'."
+            )
+        # Set qualifiers based on torch lazy conjugation flag if not provided.
+        self.qualifiers[0]["is_conjugate"] = self.qualifiers[0]["is_conjugate"] ^ self.operands[0].is_conjugate
+        self.qualifiers[1]["is_conjugate"] = self.qualifiers[1]["is_conjugate"] ^ self.operands[1].is_conjugate
+        self.lazy_conjugation = (self.operands[0].is_conjugate, self.operands[1].is_conjugate, False)
+        if self.num_operands == 3:
+            self.qualifiers[2]["is_conjugate"] = self.qualifiers[2]["is_conjugate"] ^ self.operands[2].is_conjugate
+        if self.qualifiers[2]["is_conjugate"]:
+            raise ValueError("The conjugate flag is currently not supported for operand C.")
+
+        # Set blocking or non-blocking behavior.
+        self.blocking = self.options.blocking is True or self.memory_space == "cpu"
+        if self.blocking:
+            self.call_prologue = "This call is blocking and will return only after the operation is complete."
+        else:
+            self.call_prologue = (
+                "This call is non-blocking and will return immediately after the operation is launched on the device."
+            )
+
+        # The result class is that of the first wrapped device operand.
+        self.result_class = self.operands[0].__class__
+
+        # Set memory allocator.
+        self.allocator = (
+            options.allocator
+            if options.allocator is not None
+            else memory._MEMORY_MANAGER[self.internal_op_package](self.device_id, self.logger)
+        )
+
+        # Set memory limit.
+        self.memory_limit = utils.get_memory_limit_from_device_id(self.options.memory_limit, self.device_id)
+        self.logger.info(f"The memory limit is {formatters.MemoryStr(self.memory_limit)}.")
+
+        # Set handle. We don't destroy handles we create.
+        if options.handle is not None:
+            self.handle = options.handle
+        else:
+            self.handle = get_handle(self.device_id)
+
+        # Determine the data types for a and b.
+        self.a_dtype = typemaps.NAME_TO_DATA_TYPE[a.dtype]
+        self.b_dtype = typemaps.NAME_TO_DATA_TYPE[b.dtype]
+        self.a_dtype_name = a.dtype
+        self.b_dtype_name = b.dtype
+
+        self.is_complex = "complex" in self.a_dtype_name or "complex" in self.b_dtype_name
+
+        # Determine the data types for c and d.
+        self.d_dtype = None if self.inplace else options.result_type
+        if self.num_operands == 3:
+            self.c_dtype = typemaps.NAME_TO_DATA_TYPE[c.dtype]
+            if self.d_dtype is None:
+                self.d_dtype = self.c_dtype
+
+        elif self.num_operands == 2:
+            if self.d_dtype is None:
+                self.d_dtype = self.a_dtype
+
+            # c_dtype matches d_dtype, except if output is FP4/FP8, then C uses float16
+            if self.d_dtype in (CudaDataType.CUDA_R_4F_E2M1, CudaDataType.CUDA_R_8F_E5M2, CudaDataType.CUDA_R_8F_E4M3):
+                self.c_dtype = CudaDataType.CUDA_R_16F
+            else:
+                self.c_dtype = self.d_dtype
+
+        self.c_dtype_name = typemaps.DATA_TYPE_TO_NAME[self.c_dtype]
+        self.d_dtype_name = typemaps.DATA_TYPE_TO_NAME[self.d_dtype]
+        self.c_dtype_width = typemaps.NAME_TO_DATA_WIDTH[self.c_dtype_name]
+        self.d_dtype_width = typemaps.NAME_TO_DATA_WIDTH[self.d_dtype_name]
+
+        self.logger.info(f"The data type for the result D is '{self.d_dtype_name}'.")
+
+        def assert_valid_compute_type(compute_type):
+            if compute_type not in COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["real"]:
+                message = f"Unsupported compute type. The compute type '{repr(compute_type)}' is currently not supported."
+                raise ValueError(message)
+
+        # Determine the scale type.
+        if options.scale_type is None:
+            if options.compute_type is not None:
+                assert_valid_compute_type(options.compute_type)
+                if self.is_complex:
+                    scale_type_map = COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["complex"]
+                else:
+                    scale_type_map = COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["real"]
+                self.scale_type = scale_type_map[options.compute_type]
+            else:
+                self.scale_type = NAMES_TO_DEFAULT_SCALE_TYPE[(self.a_dtype_name, self.b_dtype_name)]
+            self.scale_type_name = typemaps.DATA_TYPE_TO_NAME[self.scale_type]
+        else:
+            self.scale_type = options.scale_type
+            if self.scale_type not in SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE:
+                message = f"Unsupported scale type. The data type '{repr(self.scale_type)}' is currently not supported."
+                raise ValueError(message)
+            self.scale_type_name = typemaps.DATA_TYPE_TO_NAME[self.scale_type]
+        self.logger.info(f"The scale type is '{self.scale_type_name}'.")
+
+        # Determine the compute type.
+        if options.compute_type is None:
+            if options.scale_type is not None:
+                self.compute_type = SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE[options.scale_type]
+            else:
+                self.compute_type = NAMES_TO_DEFAULT_COMPUTE_TYPE[(self.a_dtype_name, self.b_dtype_name)]
+        else:
+            self.compute_type = options.compute_type
+        assert_valid_compute_type(self.compute_type)
+        self.logger.info(f"The compute type is {self.compute_type.name}.")
+
+        def is_supported(atype, btype, compute_type, scale_type):
+            ct = cublas.ComputeType
+            st = CudaDataType
+            abtype = atype if atype == btype else (atype, btype)
+            if compute_type in (ct.COMPUTE_16F, ct.COMPUTE_16F_PEDANTIC):
+                return scale_type == st.CUDA_R_16F and abtype == "float16"
+            elif compute_type == ct.COMPUTE_32F_PEDANTIC:
+                if scale_type == st.CUDA_R_32F:
+                    return abtype in ("float32", "bfloat16", "float16", "float8_e4m3fn", "float8_e5m2")
+                elif scale_type == st.CUDA_C_32F:
+                    return abtype == "complex64"
+            elif compute_type == ct.COMPUTE_32F:
+                if scale_type == st.CUDA_R_32F:
+                    return abtype in (
+                        "float32",
+                        "bfloat16",
+                        "float16",
+                        "float8_e4m3fn",
+                        "float8_e5m2",
+                        ("float8_e4m3fn", "float8_e5m2"),
+                        ("float8_e5m2", "float8_e4m3fn"),
+                    )
+                elif scale_type == st.CUDA_C_32F:
+                    return abtype == "complex64"
+            elif compute_type in (
+                ct.COMPUTE_32F_FAST_16F,
+                ct.COMPUTE_32F_FAST_16BF,
+                ct.COMPUTE_32F_FAST_TF32,
+                ct.COMPUTE_32F_EMULATED_16BFX9,
+            ):
+                if scale_type == st.CUDA_R_32F:
+                    return abtype == "float32"
+                if scale_type == st.CUDA_C_32F:
+                    return abtype == "complex64"
+            elif compute_type in (ct.COMPUTE_64F, ct.COMPUTE_64F_PEDANTIC, ct.COMPUTE_64F_EMULATED_FIXEDPOINT):
+                if scale_type == st.CUDA_R_64F:
+                    return abtype == "float64"
+                if scale_type == st.CUDA_C_64F:
+                    return abtype == "complex128"
+            return False
+
+        def is_supported_nvfp4(ctype, dtype, compute_type, scale_type):
+            """
+            Validate type combinations based on cuBLAS documentation.
+            https://docs.nvidia.com/cuda/cublas/index.html#id105,
+            see Table 4 and related text.
+            """
+            ct = cublas.ComputeType
+            st = CudaDataType
+
+            # Check compute/scale types
+            if not (compute_type == ct.COMPUTE_32F and scale_type == st.CUDA_R_32F):
+                raise ValueError(
+                    f"Selected scale_type={repr(self.scale_type)} compute_type={repr(self.compute_type)} "
+                    f"are not supported. Only the following combination is supported: "
+                    f"A and B must both be float4_e2m1fn_x2, compute_type=COMPUTE_32F, scale_type=CUDA_R_32F."
+                )
+
+            # Check valid C/D type combinations (Table 4)
+            valid_cd_combos = {
+                ("bfloat16", "bfloat16"),
+                ("bfloat16", "float4_e2m1fn_x2"),
+                ("float16", "float16"),
+                ("float16", "float4_e2m1fn_x2"),
+                ("float32", "float32"),
+            }
+            if (ctype, dtype) not in valid_cd_combos:
+                raise ValueError(
+                    f"Invalid C/D type combination for FP4: ctype={self.c_dtype_name}, dtype={self.d_dtype_name}. "
+                    f"Valid combinations are: (bfloat16, bfloat16), (bfloat16, float4_e2m1fn_x2), "
+                    f"(float16, float16), (float16, float4_e2m1fn_x2), (float32, float32). "
+                    f"See cuBLAS documentation Table 4."
+                )
+
+        if self.using_fp4_ab:
+            # Validate remaining type combinations (A and B are already
+            # known to be float4_e2m1fn_x2)
+            is_supported_nvfp4(self.c_dtype_name, self.d_dtype_name, self.compute_type, self.scale_type)
+        elif not is_supported(self.a_dtype_name, self.b_dtype_name, self.compute_type, self.scale_type):
+            raise ValueError(
+                f"Selected scale_type={repr(self.scale_type)} compute_type={repr(self.compute_type)} "
+                + f"are not supported for data types {self.a_dtype_name} (A) and {self.b_dtype_name} (B)."
+            )
+
+        # Set alpha and beta.
+        self.alpha = np.zeros((1,), dtype=self.scale_type_name)
+        try:
+            self.alpha[0] = alpha if alpha is not None else 1
+        except (ValueError, TypeError) as e:
+            raise ValueError(f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'.") from e
+
+        self.beta = np.zeros((1,), dtype=self.scale_type_name)
+        if beta is not None and self.num_operands == 2:
+            self.logger.warning(f"Matmul: The provided beta value {beta} is ignored since operand C is not specified.")
+        try:
+            self.beta[0] = beta if beta is not None and self.num_operands == 3 else 0
+        except (ValueError, TypeError) as e:
+            raise ValueError(f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'.") from e
+
+        # Set narrow-precision (FP8 and lower) quantization_scales.
+        if self.input_type_width <= 8:
+            self.quantization_scales = self._validate_operand_scales(quantization_scales, all_required=True)
+        elif quantization_scales is not None:
+            self.logger.warning(
+                "Matmul: The provided scales are ignored, since they are only applicable to narrow-precision (FP8 and lower) "
+                "operations."
+            )
+
+        if self.options.result_amax and self.d_dtype_width > 8:
+            raise ValueError("result_amax=True is allowed only for narrow-precision (FP8 and lower) results")
+
+        # Check operands alignment if needed
+        if self.input_type_width <= 8:
+            for operand, operand_name in zip(self.operands, "ABC", strict=False):
+                if operand.data_ptr % 16 != 0:
+                    raise ValueError(
+                        f"For narrow-precision (FP8 and lower) multiplication, operand {operand_name} should be aligned to 16 "
+                        "bytes."
+                    )
+
+        # Capture physical operand extents and strides for consistency check when resetting
+        # operands. In the case of FP4, these must be the packed (physical) values,
+        # not the logical ones , because reset() receives packed FP4 tensors
+        # and must compare against the same physical format.
+        # Hence, these two do not need special handling for FP4.
+        self.operand_extents = tuple(o.shape for o in self.operands)
+        self.operand_strides = tuple(o.strides for o in self.operands)
+
+        if self.using_fp4_ab:
+            a_layout, b_layout = _create_fp4_ab_layouts(
+                self.operands[0].shape,
+                self.operands[0].strides,
+                self.qualifiers[0]["is_conjugate"],
+                self.operands[1].shape,
+                self.operands[1].strides,
+                self.qualifiers[1]["is_conjugate"],
+            )
+            self.logger.info(f"A: packed shape {self.operands[0].shape} -> logical shape {a_layout.shape}")
+            self.logger.info(f"B: packed shape {self.operands[1].shape} -> logical shape {b_layout.shape}")
+        else:
+            self.logger.info(f"A: shape {self.operands[0].shape}, strides {self.operands[0].strides}")
+            self.logger.info(f"B: shape {self.operands[1].shape}, strides {self.operands[1].strides}")
+            a_layout = MatrixLayout(self.operands[0].shape, self.operands[0].strides, self.qualifiers[0]["is_conjugate"])
+            b_layout = MatrixLayout(self.operands[1].shape, self.operands[1].strides, self.qualifiers[1]["is_conjugate"])
+
+        # C is never FP4 (only A and B can be float4_e2m1fn_x2),
+        # so no packed-to-logical conversion needed.
+        c_layout = MatrixLayout(self.operands[2].shape, self.operands[2].strides) if self.num_operands == 3 else None  # type: ignore[union-attr]
+
+        # Enforce equal batch shape for A and B if block_scaling=True.
+        if self.options.block_scaling and a_layout.shape[:-2] != b_layout.shape[:-2]:
+            raise ValueError(
+                "When block_scaling=True, the batch dimensions of A and B must match (broadcasting is not supported)."
+            )
+
+        # Get the operation traits.
+        self.mm_traits = get_mm_traits(a_layout, b_layout, c_layout, self.inplace, self.logger)
+        self.result_traits = None  # Wait till planning to determine this based on the epilog.
+        self.logger.info(
+            f"The matrix multiplication attributes are M = {self.mm_traits.M}, N = {self.mm_traits.N}, and "
+            f"K = {self.mm_traits.K}."
+        )
+        self.logger.info(
+            f"The batch count is {self.mm_traits.batch_count}, and the batch shape is {self.mm_traits.batch_shape} "
+            f"with batch axis order {self.mm_traits.batch_axis_order}."
+        )
+
+        # Create and set the operation descriptor.
+        self.mm_desc = cublaslt.matmul_desc_create(self.compute_type, self.scale_type)
+
+        self.mm_desc_ifc = matmul_desc_ifc.MatmulDescInterface(self.mm_desc)
+        self.mm_desc_ifc.compute_type = self.compute_type
+        self.mm_desc_ifc.scale_type = self.scale_type
+
+        # Guard SM count target and fast accumulation flag.
+        version = cublaslt.get_version()
+        if options.sm_count_target > 0:  # type: ignore[operator]
+            if version < 111103:
+                raise ValueError(f"The 'sm_count_target' option is not supported in cuBLASLt version {version}.")
+            self.mm_desc_ifc.sm_count_target = options.sm_count_target
+            self.logger.info(f"The SM count target is {options.sm_count_target}.")
+
+        if options.fast_accumulation:
+            if version < 111103:
+                raise ValueError(f"The 'fast_accumulation' option is not supported in cuBLASLt version {version}.")
+            self.mm_desc_ifc.fast_accum = options.fast_accumulation
+            self.logger.info(f"The flag for fast accumulation mode is {options.fast_accumulation}.")
+
+        if self.input_type_width == 8 and version < 120800:
+            raise ValueError(
+                f"FP8 is not supported for cuBLASLt version {version}. cuBLASLt version 12.8 or higher is required."
+            )
+
+        # Planning preferences
+        self.preferences = None
+
+        # Epilog attributes.
+        self.epilog = None
+
+        # Epilog attributes: name-to-operand.
+        self.epilog_operands: dict[str, typing.Any] = {}
+
+        # Epilog attributes: epilog input name-to-handler.
+        self.epilog_input_name_to_handler: dict[str, typing.Any] = {}
+
+        # Epilog attributes: name-to-output tensor.
+        self.epilog_outputs: dict[str, typing.Any] = {}
+
+        # Keep track of epilog input traits for resetting operands.
+        self.epilog_inputs_traits: dict[str, typing.Any] = {}
+
+        # Keep track of epilog output handlers to allocate output in execute().
+        self.epilog_output_handlers: list[EpilogOutputHandler] = []
+
+        # Non-epilog aux outputs. Currently, only used for quantization outputs (amax etc.)
+        self.aux_outputs = None
+
+        # Plan attributes.
+        self.preference_ptr = None
+        self.a_layout_ptr, self.b_layout_ptr, self.c_layout_ptr, self.d_layout_ptr = None, None, None, None
+        self.flop_count = 0
+        self.mm_planned = False
+
+        # Algorithm attributes.
+        self.algorithms_buffer = None
+        self.algorithm_objects = None
+        self.cached_best_algorithm_struct = None
+
+        # Workspace lifecycle is managed by `Workspace` (allocate-on-demand,
+        # reuse, mid-call release, exception cleanup, stream-ordered free).
+        self.workspace = Workspace(
+            self.allocator,
+            self.logger,
+            label="device workspace",
+            device_id=self.device_id,
+        )
+
+        # Last event recorded by cuda_call_ctx; consumed on workspace release
+        # so the free is ordered after the compute that touched the buffer.
+        self.last_compute_event = None
+
+        # Track whether the user has called release_operands().
+        self._operands_released = False
+
+        # Device-side array with the quantization_scales
+        self.quantization_scales_device: dict[str, utils.TensorHolder] = {}
+
+        self.valid_state = True
+        self.logger.info("The Matmul operation has been created.")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.free()
+
+    def _check_valid_matmul(self, *args, **kwargs):
+        """
+        Check if the Matmul object is alive and well.
+        """
+        if not self.valid_state:
+            raise InvalidMatmulState("The Matmul object cannot be used after resources are free'd")
+
+    def _check_valid_operands(self, *args, **kwargs):
+        """
+        Check if the operands are available for the operation.
+        """
+        what = kwargs["what"]
+        if self._operands_released:
+            raise RuntimeError(
+                f"{what} cannot be performed after the operands have been released. "
+                f"Use reset_operands() to provide new operands before performing the {what.lower()}."
+            )
+
+    def _free_plan_resources(self, exception: Exception | None = None) -> bool:
+        """
+        Free resources allocated in planning.
+        """
+
+        # Destroy matrix layouts.
+        if self.a_layout_ptr is not None:
+            cublaslt.matrix_layout_destroy(self.a_layout_ptr)
+            self.a_layout_ptr = None
+        if self.b_layout_ptr is not None:
+            cublaslt.matrix_layout_destroy(self.b_layout_ptr)
+            self.b_layout_ptr = None
+        if self.c_layout_ptr != self.d_layout_ptr and self.c_layout_ptr is not None:
+            cublaslt.matrix_layout_destroy(self.c_layout_ptr)
+        self.c_layout_ptr = None
+        if self.d_layout_ptr is not None:
+            cublaslt.matrix_layout_destroy(self.d_layout_ptr)
+            self.d_layout_ptr = None
+
+        if self.preference_ptr is not None:
+            cublaslt.matmul_preference_destroy(self.preference_ptr)
+            self.preference_ptr = None
+
+        self.mm_planned = False
+        return True
+
+    def _check_planned(self, *args, **kwargs):
+        what = kwargs["what"]
+        if not self.mm_planned:
+            raise RuntimeError(f"{what} cannot be performed before plan() has been called.")
+
+    @utils.precondition(_check_valid_matmul)
+    def applicable_algorithm_ids(self, limit=8):
+        """Obtain the algorithm IDs that are applicable to this matrix multiplication.
+
+        Args:
+            limit: The maximum number of applicable algorithm IDs that is desired
+
+        Returns:
+            A sequence of algorithm IDs that are applicable to this matrix multiplication
+            problem specification, in random order.
+        """
+        ...
+        algo_ids = cublaslt.matmul_algo_get_ids(
+            self.handle,
+            self.compute_type,
+            self.scale_type,
+            self.a_dtype,
+            self.b_dtype,
+            self.c_dtype,
+            self.d_dtype,
+            limit,
+        )
+        return algo_ids
+
+    def _validate_scalar_scale(self, operand: str):
+        """
+        Validates a scalar scale.
+        """
+        if self.options.block_scaling:
+            raise ValueError(f"A scalar tensor-wide scale factor is not allowed for {operand.upper()} when block_scaling=True.")
+
+    def _validate_tensor_scale(self, scale, operand: str, operand_size=None):
+        """
+        Validates a tensor scale.
+
+        Args:
+            scale: The tensor scale to validate
+            operand: The operand name (a, b, c, d)
+            operand_size: Size of the operand (needed for block scaling shape validation)
+        """
+        scale_package = utils.infer_object_package(scale)
+        if scale_package != self.package:
+            raise TypeError(
+                f"The quantization scaling tensor for {operand.upper()} must belong to the same package as the operands."
+            )
+
+        # Wrap temporarily since this is needed for validation
+        scale_wrapped = tensor_wrapper.wrap_operand(scale)
+
+        # Device/memory space validation
+        expected_device_id = "cpu" if self.memory_space == "cpu" else self.device_id
+        if expected_device_id != scale_wrapped.device_id:
+            raise ValueError(
+                f'The scale for {operand.upper()} is on device "{scale_wrapped.device_id}", '
+                f'but it should be on device "{expected_device_id}" to match the operands memory space.'
+            )
+
+        # Shape and dtype validation for non-block-scaling
+        if not self.options.block_scaling:
+            if scale_wrapped.shape not in ((1,), ()):
+                raise ValueError(
+                    f"The scale for {operand.upper()} must be of shape (1,) or (). Got {scale_wrapped.shape} instead."
+                )
+            if scale_wrapped.dtype != "float32":
+                raise ValueError(f"The scale for {operand.upper()} must be float32 type. Got {scale_wrapped.dtype} instead.")
+
+        # Shape and dtype validation for block-scaling
+        elif self.input_type_width == 8:
+            # FP8: 32-element blocks with 128x128 tiled layout (VEC32_UE8M0)
+            # Dtype validation (always possible)
+            if scale_wrapped.dtype != "uint8":
+                raise ValueError(f"Block scales for {operand.upper()} must be uint8 tensor.")
+
+            # Shape validation (only if operand_size is available)
+            if operand_size is not None:
+                expected_shape = (operand_size // 32,)
+                if scale_wrapped.shape != expected_shape:
+                    raise ValueError(
+                        f"Scales for {operand.upper()} should have shape {expected_shape}. Got {scale_wrapped.shape}."
+                    )
+        elif self.using_fp4_ab:
+            # FP4: 16-element blocks with 128x64 tiled layout (VEC16_UE4M3)
+            # Dtype validation (always possible)
+            if scale_wrapped.dtype != "float8_e4m3fn":
+                raise ValueError(
+                    f"Block scales for {operand.upper()} must be float8_e4m3fn tensor (UE4M3 format). "
+                    f"Got {scale_wrapped.dtype}."
+                )
+
+            # Shape validation (only if operand_size is available)
+            if operand_size is not None:
+                expected_shape = (operand_size // 16,)
+                if scale_wrapped.shape != expected_shape:
+                    raise ValueError(
+                        f"Scales for {operand.upper()} should have shape {expected_shape}. Got {scale_wrapped.shape}."
+                    )
+        else:
+            raise ValueError("block_scaling == True is not supported for non-FP8/FP4 types.")
+
+    def _validate_operand_scales(self, quantization_scales, all_required):
+        """
+        Validates quantization scales, wrapping them into a MatmulQuantizationScales
+        object if needed.
+
+        Args:
+            quantization_scales: The quantization scales to validate.
+            all_required: Whether all scales are required.
+
+        Returns:
+            A MatmulQuantizationScales object with the validated quantization scales.
+        """
+        if quantization_scales is None:
+            raise ValueError(
+                "Scales are required for narrow-precision (FP8 and lower) operations. "
+                "Please set `quantization_scales` argument."
+            )
+
+        # wrap the quantization scales into a MatmulQuantizationScales object if needed
+        # otherwise, return the quantization scales as is
+        quantization_scales = utils.check_or_create_options(
+            _configuration.MatmulQuantizationScales, quantization_scales, "Scale factors"
+        )
+
+        # Validate which scales are required/allowed
+        expected_scales = "AB"
+        if self.d_dtype_width <= 8 and not self.options.block_scaling:
+            expected_scales += "D"
+        elif quantization_scales.d is not None:
+            if self.options.block_scaling:
+                raise ValueError("Quantization scaling is not supported for D when `block_scaling` option is enabled.")
+            if self.d_dtype_width > 8:
+                raise ValueError(
+                    "Quantization scaling is not supported for D when it is not a narrow-precision (FP8 and lower) type."
+                )
+
+        if self.num_operands == 3 and self.c_dtype_width <= 8:
+            expected_scales += "C"
+        elif quantization_scales.c is not None:
+            raise ValueError(
+                "Quantization scaling is not supported for C when it is not a narrow-precision (FP8 and lower) type."
+            )
+
+        if all_required:
+            for operand in expected_scales:
+                if getattr(quantization_scales, operand.lower()) is None:
+                    raise ValueError(f"Scale for {operand.upper()} is not specified")
+
+        # Validate each scale by delegating to scalar/tensor specific validators
+        for operand in "abcd":
+            scale = getattr(quantization_scales, operand)
+            if scale is None:
+                continue
+
+            if isinstance(scale, (int, float)):
+                self._validate_scalar_scale(operand)
+            else:
+                # For block scaling, pass operand size for shape validation
+                if self.options.block_scaling and operand in ("a", "b"):
+                    operand_idx = 0 if operand == "a" else 1
+                    operand_size = self.operands[operand_idx].size  # type: ignore[union-attr,index]
+
+                    # For FP4, data is always packed: double the size to get logical size
+                    if self.using_fp4_ab:
+                        operand_size *= 2  # Packed tensor has half the elements
+                        self.logger.debug(f"FP4 scale validation for {operand}: adjusted operand_size={operand_size}")
+                else:
+                    operand_size = None
+                self._validate_tensor_scale(scale, operand, operand_size)
+
+        return quantization_scales
+
+    def _validate_epilog_aux_scale(self, aux_quantization_scale, *, required):
+        is_narrow_aux = (
+            self.preferences.epilog.aux_type is not None
+            and typemaps.NAME_TO_DATA_WIDTH[typemaps.DATA_TYPE_TO_NAME[self.preferences.epilog.aux_type]] <= 8
+        )
+        if aux_quantization_scale is not None and not is_narrow_aux:
+            raise ValueError(
+                "Scales for epilog auxiliary output are not supported when `preferences.epilog.aux_type` is not set to a "
+                "narrow-precision type."
+            )
+        elif aux_quantization_scale is None and is_narrow_aux and required:
+            raise ValueError(
+                '"aux_quantization_scale" epilog input is required when `preferences.epilog.aux_type` is set to a '
+                "narrow-precision type."
+            )
+
+        # Validate scalar vs tensor scale (same as for operand scales)
+        if aux_quantization_scale is not None:
+            if isinstance(aux_quantization_scale, (int, float)):
+                self._validate_scalar_scale("epilog_aux")
+            else:
+                # No operand_size for epilog_aux scales
+                self._validate_tensor_scale(aux_quantization_scale, "epilog_aux", operand_size=None)
+
+    def _prepare_validated_scalar_scale(self, scale: int | float, operand: str, stream_holder: utils.StreamHolder):
+        """
+        Converts validated scalar to float32 tensor and copies to GPU.
+        Assumes validation already done in _validate_scalar_scale.
+        """
+        # If it's a scalar, copy to GPU. Float32 is the only type allowed by
+        # cublasLtMatmulScale_t for tensor-wide scaling.
+        self.logger.debug(f"Scale for {operand.upper()} will be copied to device {self.device_id}.")
+        scale_op = tensor_wrapper.wrap_operand(np.asarray([scale], dtype="float32"))
+        self.quantization_scales_device[operand] = scale_op.to(self.device_id, stream_holder)
+
+    def _prepare_validated_tensor_scale(self, scale, operand: str, stream_holder: utils.StreamHolder):
+        """
+        Wraps validated tensor and copies to GPU.
+        Assumes all validation already done in _validate_tensor_scale (called in __init__).
+        This is pure preparation - no validation here.
+
+        Note: We wrap the tensor a second time here (first wrap was for validation).
+        This is acceptable because wrapping is cheap and we get early error detection.
+        """
+        # Wrap the scale (second time - first was for validation in __init__)
+        self.quantization_scales_device[operand] = tensor_wrapper.wrap_operand(scale)
+
+        # Copy to GPU if on CPU (no validation, just preparation)
+        if self.quantization_scales_device[operand].device in (None, "cpu"):
+            self.logger.debug(f"Scale for {operand.upper()} will be copied to device {self.device_id}.")
+            self.quantization_scales_device[operand] = self.quantization_scales_device[operand].to(
+                self.device_id, stream_holder
+            )
+
+    def _prepare_single_validated_scale(self, scale, operand: str, cublas_operand: str, stream_holder: utils.StreamHolder):
+        """
+        Prepares a single validated scale and sets its pointer/mode in mm_desc_ifc.
+        Used for both operand scales (a,b,c,d) and epilog scales (epilog_aux).
+        Assumes validation already done.
+        """
+        if scale is None:
+            return
+
+        # Delegate to specific preparer (validation already done)
+        if isinstance(scale, (int, float)):
+            self._prepare_validated_scalar_scale(scale, operand, stream_holder)
+        else:
+            self._prepare_validated_tensor_scale(scale, operand, stream_holder)
+
+        # Set pointer and mode in descriptor
+        setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_pointer", self.quantization_scales_device[operand].data_ptr)
+
+        if self.options.block_scaling:
+            # MXFP8 uses 32-element blocks with UE8M0, FP4 uses 16-element blocks with UE4M3
+            if self.using_fp4_ab:
+                # FP4: 16-element block scaling with UE4M3 format
+                self.logger.debug(f"Using VEC16_UE4M3 scale mode for operand {operand.upper()}.")
+                setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_mode", cublaslt.MatmulMatrixScale.VEC16_UE4M3)
+            else:
+                # FP8: 32-element block scaling with UE8M0 format
+                self.logger.debug(f"Using VEC32_UE8M0 scale mode for operand {operand.upper()}.")
+                setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_mode", cublaslt.MatmulMatrixScale.VEC32_UE8M0)
+        else:
+            self.logger.debug(f"Using SCALAR_32F scale mode for operand {operand.upper()}.")
+            setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_mode", cublaslt.MatmulMatrixScale.SCALAR_32F)
+
+    def _prepare_operand_quantization_scales(self, scales, stream_holder: utils.StreamHolder):
+        """
+        Prepares validated operand scales (a,b,c,d).
+        Assumes scales are validated and wrapped into a MatmulQuantizationScales object.
+        """
+        for operand in "abcd":
+            scale = getattr(scales, operand)
+            self._prepare_single_validated_scale(scale, operand, cublas_operand=operand, stream_holder=stream_holder)
+
+    @utils.precondition(_check_valid_matmul)
+    @utils.atomic(_free_plan_resources, method=True)
+    def plan(
+        self, *, preferences=None, algorithms=None, epilog=None, epilog_inputs=None, stream: utils.AnyStream | int | None = None
+    ):  # Epilog inputs require as many inputs (with specific shapes etc) as required by the epilogue. It's a dict.
+        """
+        Plan the matrix multiplication operation, considering the epilog (if provided).
+
+        Args:
+            preferences: {preferences}
+
+            algorithms: {algorithms}
+
+            epilog: {epilog}
+
+            epilog_inputs: {epilog_inputs}
+
+            stream: {stream}
+
+        Returns:
+            A sequence of :class:`nvmath.linalg.advanced.Algorithm` objects that are
+            applicable to this matrix multiplication problem specification, heuristically
+            ordered from fastest to slowest.
+
+        Notes:
+            Epilogs that have ``BIAS`` in their name need an epilog input with the key
+            ``'bias'``. Epilogs that have ``DRELU`` need an epilog input with the key
+            ``'relu_aux'``, which is produced in a "forward pass" epilog like ``RELU_AUX``
+            or ``RELU_AUX_BIAS``. Similarly, epilogs with ``DGELU`` in their name require an
+            epilog input with the key ``'gelu_aux'``, produced in the corresponding forward
+            pass operation.
+
+        Examples:
+
+            >>> import numpy as np
+            >>> import nvmath
+
+            Create two 3-D float64 ndarrays on the CPU representing batched matrices, along
+            with a bias vector:
+
+            >>> batch = 32
+            >>> M, N, K = 1024, 1024, 1024
+            >>> a = np.random.rand(batch, M, K)
+            >>> b = np.random.rand(batch, K, N)
+            >>> # The bias vector will be broadcast along the columns, as well as along the
+            >>> # batch dimension.
+            >>> bias = np.random.rand(M)
+
+            We will define a matrix multiplication operation followed by a
+            :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS` epilog function.
+
+            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
+            ...     # Plan the operation with RELU_BIAS epilog and corresponding epilog
+            ...     # input.
+            ...     p = nvmath.linalg.advanced.MatmulPlanPreferences(limit=8)
+            ...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS
+            ...     epilog_inputs = {{"bias": bias}}
+            ...     # The preferences can also be provided as a dict: {{'limit': 8}}
+            ...     algorithms = mm.plan(
+            ...         preferences=p,
+            ...         epilog=epilog,
+            ...         epilog_inputs=epilog_inputs,
+            ...     )
+            ...
+            ...     # Execute the matrix multiplication, and obtain the result `r` as a
+            ...     # NumPy ndarray.
+            ...     r = mm.execute()
+
+            Some epilogs like :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_AUX` produce
+            auxiliary output.
+
+            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
+            ...     # Plan the operation with RELU_AUX epilog>
+            ...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_AUX
+            ...     algorithms = mm.plan(epilog=epilog)
+            ...
+            ...     # Execute the matrix multiplication, and obtain the result `r` along
+            ...     # with the auxiliary output.
+            ...     r, auxiliary = mm.execute()
+
+            The auxiliary output is a Python `dict` with the names of each auxiliary output
+            as keys.
+
+        Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
+        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
+        directory.
+        """
+        log_info = self.logger.isEnabledFor(logging.INFO)
+
+        self.logger.info("= PLANNING PHASE =")
+
+        # Release layout objects and other plan resources from a previous plan() call
+        # to prevent leaks on re-plan. This is a no-op on the first call.
+        self._free_plan_resources()
+
+        # Clear epilog operands, since different epilogs can be provided in different calls.
+        # We don't need to worry about ordering, since it's the user's responsibility to
+        # order calls that accept a stream argument. This applies to CPU operands as well,
+        # even though we move them to the GPU, since the execution is blocking.
+        self.epilog_operands = {}  # Clear operands in case of repeated planning.
+        self.epilog_input_name_to_handler = {}  # Clear input name to handler map as well,
+        self.epilog_inputs_traits = {}  # ... and the input traits as well.
+
+        preferences = utils.check_or_create_options(
+            _configuration.MatmulPlanPreferences, preferences, "Matrix multiplication plan preferences"
+        )
+        self.preferences = preferences
+
+        mm_traits = self.mm_traits
+
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+        self.logger.info(f"The specified stream for the matrix multiplication plan is {stream_holder.obj}.")
+
+        # Base FLOP count.
+        self.flop_count = 2 * mm_traits.M * mm_traits.N * mm_traits.K
+        self.logger.info(f"The base matrix multiplication FLOP count is {formatters.FLOPSStr(self.flop_count, 'FLOP')}.")
+
+        if epilog is None and epilog_inputs is not None:
+            self.logger.warning(
+                f"Matmul: The provided epilog inputs {epilog_inputs.keys()} are ignored since an epilog is not specified."
+            )
+
+        self.epilog = epilog
+        epilog_ordering = None
+        if epilog is not None:
+            assert epilog in EPILOG_INPUT_HANDLERS_MAP, "Not supported."
+            self.logger.info(f"The specified epilog is {epilog.name}.")
+
+            epilog_minimum_versions = EPILOG_MINIMUM_VERSIONS_MAP[epilog]
+            batched_epilog_minimum_versions = BATCHED_EPILOG_MINIMUM_VERSIONS_MAP[epilog]
+            version = cublaslt.get_version()
+            if version < epilog_minimum_versions["cublaslt"]:
+                message = (
+                    f"The epilog {epilog.name} requires cublaslt >= {epilog_minimum_versions['cublaslt']}; "
+                    f"you have version {version}. Update to CUDA Toolkit >= {epilog_minimum_versions['ctk']}."
+                )
+                raise ValueError(message)
+
+            if len(mm_traits.batch_shape) > 0 and version < batched_epilog_minimum_versions["cublaslt"]:
+                message = (
+                    f"The epilog {epilog.name} supports batching in "
+                    f"cublaslt >= {batched_epilog_minimum_versions['cublaslt']}; "
+                    f"you have version {version}. Update to CUDA Toolkit >= {epilog_minimum_versions['ctk']}."
+                )
+                raise ValueError(message)
+            if (
+                self.mm_traits.c_layout_traits is not None
+                and self.mm_traits.c_layout_traits.order == cublaslt.Order.ROW
+                and epilog
+                in [
+                    cublaslt.Epilogue.BGRADA,
+                    cublaslt.Epilogue.BGRADB,
+                ]
+            ):
+                msg = f"The epilog {epilog.name} requires input matrix 'c' to be F-contiguous (column-major)."
+                raise ValueError(msg)
+            if (
+                version < 120804
+                # A has one row
+                and self.mm_traits.M == 1
+                # C is broadcast
+                and self.mm_traits.c_layout_traits is not None
+                and self.mm_traits.c_layout_traits.ld == 0
+                # Using both an bias epilog and C
+                and self.epilog & _configuration.MatmulEpilog.BIAS > 0
+            ):
+                message = (
+                    "When matrix 'a' has one row, "
+                    "simultaneously broadcasting matrix 'c' and using a BIAS epilog requires cublaslt >= 120804; "
+                    f"You have version {version}. Update to CUDA Toolkit >= 12.8.1."
+                )
+                raise ValueError(message)
+
+            # Take a copy of the user-provided inputs.
+            if epilog_inputs is not None:
+                epilog_inputs = epilog_inputs.copy()
+            else:
+                epilog_inputs = {}
+
+            # Get the dtype of auxiliary buffer
+            aux_dtype_name = (
+                typemaps.DATA_TYPE_TO_NAME[self.preferences.epilog.aux_type]  # type: ignore[attr-defined]
+                if self.preferences.epilog.aux_type is not None  # type: ignore[attr-defined]
+                else None
+            )
+
+            # Extract aux quantization scale from the inputs.
+            aux_quantization_scale = epilog_inputs.pop("aux_quantization_scale", None)
+            self._validate_epilog_aux_scale(aux_quantization_scale, required=True)
+            self._prepare_single_validated_scale(
+                aux_quantization_scale, "epilog_aux", cublas_operand="epilogue_aux", stream_holder=stream_holder
+            )
+
+            epilog_input_handler_types = EPILOG_INPUT_HANDLERS_MAP[epilog]
+            if epilog_input_handler_types:
+                epilog_input_handlers = [
+                    handler_type(self.logger, mm_traits, epilog, self.c_dtype_name, self.d_dtype_name, aux_dtype_name)
+                    for handler_type in epilog_input_handler_types
+                ]
+
+                # Check if the epilog requires a specific result layout, and if the
+                # requirement is consistent for all the handlers.
+                epilog_input_handlers_ordering = {h.order for h in epilog_input_handlers}
+                assert len(epilog_input_handlers_ordering) == 1, "Internal error."
+                epilog_ordering = epilog_input_handlers_ordering.pop()
+
+                required_epilog_input_names = {h.name for h in epilog_input_handlers}
+
+                self.logger.info(f"The epilog requires the following additional inputs: {required_epilog_input_names}.")
+                if required_epilog_input_names != set(epilog_inputs.keys()):
+                    raise ValueError(
+                        f"The epilog {epilog.name} requires the following input tensors: "
+                        f"{required_epilog_input_names}. The provided tensor names are: {epilog_inputs.keys()}"
+                    )
+
+                # Wrap epilog inputs.
+                for name in epilog_inputs:
+                    epilog_inputs[name] = tensor_wrapper.wrap_operand(epilog_inputs[name])
+
+                # Check if epilog inputs all belong to the same package, which is the same
+                # as the package of the MM operands.
+                epilog_package = utils.get_operands_package(list(epilog_inputs.values()))
+                if self.package != epilog_package:
+                    message = f"Library package mismatch for epilog: '{self.package}' => '{epilog_package}'"
+                    raise TypeError(message)
+
+                # When we get here, epilog_inputs should only contain tensors because
+                # we popped aux_quantization_scale from the dictionary above.
+                assert all(isinstance(v, tensor_wrapper.TensorHolder) for v in epilog_inputs.values()), "Internal error."
+                # Since all epilog inputs are tensors, we can ensure they all are on
+                # the same memory space as the operands provided at initialization,
+                # as per the documentation for the epilog_inputs parameter.
+                device_id = utils.get_operands_device_id(list(epilog_inputs.values()))
+                expected_device = "cpu" if self.memory_space == "cpu" else self.device_id
+                if device_id != expected_device:
+                    raise ValueError(
+                        f"The epilog inputs must be in the same memory space as the operands. "
+                        f"Expected device '{expected_device}' (operands' memory space is '{self.memory_space}'), "
+                        f"but epilog inputs are on device '{device_id}'."
+                    )
+
+                # Move epilog inputs to the GPU, if needed.
+                if device_id == "cpu":
+                    for e in required_epilog_input_names:
+                        self.logger.debug(f"The epilog input {e} will be copied to device{self.device_id}.")
+                        self.epilog_operands[e] = epilog_inputs[e].to(self.device_id, stream_holder)
+                else:
+                    for e in required_epilog_input_names:
+                        self.epilog_operands[e] = epilog_inputs[e]
+
+                # First validate all epilog inputs. Use the GPU tensors in case metadata has
+                # changed.
+                for handler in epilog_input_handlers:
+                    handler.validate(epilog_inputs[handler.name])
+
+                # Finally, update the MM descriptor. Note that we pass in
+                # self.epilog_operands (which are on the GPU).
+                for handler in epilog_input_handlers:
+                    handler.update(self.mm_desc_ifc, self.epilog_operands[handler.name])
+                    self.epilog_input_name_to_handler[handler.name] = handler
+
+                # Capture the epilog operands traits for consistency checks when resetting
+                # operands.
+                self.epilog_inputs_traits = {
+                    name: EpilogInputTraits(
+                        dtype=self.epilog_operands[name].dtype,
+                        extents=self.epilog_operands[name].shape,
+                        strides=self.epilog_operands[name].strides,
+                    )
+                    for name in self.epilog_operands
+                }
+
+            epilog_output_handler_types = EPILOG_OUTPUT_HANDLERS_MAP[epilog]
+            if epilog_output_handler_types:
+                self.epilog_output_handlers = epilog_output_handlers = [
+                    handler_type(self.logger, mm_traits, epilog, self.c_dtype_name, self.d_dtype_name, aux_dtype_name)
+                    for handler_type in epilog_output_handler_types
+                ]
+                # Check if the epilog requires a specific result layout, and if the
+                # requirement is consistent for all the handlers.
+                epilog_output_handlers_ordering = {h.order for h in epilog_output_handlers}
+                assert len(epilog_output_handlers_ordering) == 1, "Internal error."
+                op_epilog_ordering = epilog_output_handlers_ordering.pop()
+                if epilog_ordering is None:
+                    epilog_ordering = op_epilog_ordering
+                else:
+                    assert epilog_ordering == op_epilog_ordering, "Internal error."
+
+                # Update the MM descriptor, except for the device pointer.
+                for ohandler in epilog_output_handlers:
+                    ohandler.update(self.mm_desc_ifc)
+
+            # Set the epilog. At this point, we're sure that the epilog inputs, if any, are
+            # valid and have been set.
+            self.mm_desc_ifc.epilogue = epilog
+
+        # Fill the result traits, now that we know the epilog.
+        self.result_traits = result_traits = get_result_traits(mm_traits, epilog_ordering, self.logger)  # type: ignore[assignment]
+        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()."
+        self.logger.info(
+            f"The layout order for the result D is {self.result_traits.d_layout_traits.order.name}, with LD "
+            f"{self.result_traits.d_layout_traits.ld}, and batch offset "
+            f"{self.result_traits.d_layout_traits.batch_offset}."
+        )
+
+        # Internally transpose operand A if required (conjugate flag) and create layout.
+        transpose = False
+        if mm_traits.a_layout_traits.is_conjugate and self.is_complex:
+            self.mm_desc_ifc.transa = cublas.Operation.C
+            transpose = True
+            self.logger.debug(
+                "To conjugate A, the operand A will be internally transposed and the matrix multiplication will be "
+                "performed with OP_C for operand A."
+            )
+        if self.input_type_width <= 8:
+            # narrow-precision (FP8 and lower) data types are only supported for transa=OP_T
+            self.mm_desc_ifc.transa = cublas.Operation.T
+            transpose = True
+            self.logger.debug(
+                "For narrow-precision (FP8 and lower) multiplication, the operand A will be internally transposed and the "
+                "matrix multiplication will be performed with OP_T for operand A."
+            )
+        m, n, ld, a_order = mm_traits.a_layout_traits.get_mm_layout(transpose=transpose)
+        self.a_layout_ptr = cublaslt.matrix_layout_create(self.a_dtype, rows=m, cols=n, ld=ld)
+        self.logger.debug(f"Layout for A: rows = {m}, cols = {n}, ld = {ld}.")
+
+        # Internally transpose operand B if required (conjugate flag, or epilog is BGRADB)
+        # and create layout.
+        transpose = False
+        if mm_traits.b_layout_traits.is_conjugate and self.is_complex:
+            self.mm_desc_ifc.transb = cublas.Operation.C
+            transpose = True
+            self.logger.debug(
+                "To conjugate B, the operand B will be internally transposed and the matrix multiplication will be "
+                "performed with OP_C for operand B."
+            )
+        elif epilog == _configuration.MatmulEpilog.BGRADB:
+            self.mm_desc_ifc.transb = cublas.Operation.T
+            transpose = True
+            self.logger.debug(
+                "For BGRADB epilog, the operand B will be internally transposed and the matrix multiplication will be "
+                "performed with OP_T for operand B."
+            )
+        m, n, ld, b_order = mm_traits.b_layout_traits.get_mm_layout(transpose=transpose)
+        self.b_layout_ptr = cublaslt.matrix_layout_create(self.b_dtype, rows=m, cols=n, ld=ld)
+        self.logger.debug(f"Layout for B: rows = {m}, cols = {n}, ld = {ld}.")
+
+        self.d_layout_ptr = cublaslt.matrix_layout_create(
+            self.d_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=result_traits.d_layout_traits.ld
+        )
+
+        layout_a_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.a_layout_ptr)
+        layout_a_ifc.order = a_order
+        layout_a_ifc.batch_count = mm_traits.batch_count
+        layout_a_ifc.strided_batch_offset = mm_traits.a_layout_traits.batch_offset
+
+        layout_b_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.b_layout_ptr)
+        layout_b_ifc.order = b_order
+        layout_b_ifc.batch_count = mm_traits.batch_count
+        layout_b_ifc.strided_batch_offset = mm_traits.b_layout_traits.batch_offset
+
+        layout_d_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.d_layout_ptr)
+        layout_d_ifc.order = result_traits.d_layout_traits.order
+        layout_d_ifc.batch_count = mm_traits.batch_count
+        layout_d_ifc.strided_batch_offset = result_traits.d_layout_traits.batch_offset
+
+        if self.num_operands == 2:  # By defn, this cannot be inplace.
+            if self.c_dtype == self.d_dtype:
+                # If C and D have equal types, reuse the layout.
+                self.c_layout_ptr = self.d_layout_ptr
+            else:
+                # Otherwise, create a D-like layout, but with different type.
+                self.c_layout_ptr = cublaslt.matrix_layout_create(
+                    self.c_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=result_traits.d_layout_traits.ld
+                )
+                layout_c_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.c_layout_ptr)
+                layout_c_ifc.order = result_traits.d_layout_traits.order
+                layout_c_ifc.batch_count = mm_traits.batch_count
+                layout_c_ifc.strided_batch_offset = result_traits.d_layout_traits.batch_offset
+        else:
+            # For inplace operation, use the same layout for C and D.
+            if self.inplace:
+                self.c_layout_ptr = self.d_layout_ptr
+            else:
+                self.c_layout_ptr = cublaslt.matrix_layout_create(
+                    self.c_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=mm_traits.c_layout_traits.ld
+                )
+                layout_c_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.c_layout_ptr)
+                layout_c_ifc.order = mm_traits.c_layout_traits.order
+                layout_c_ifc.batch_count = mm_traits.batch_count
+                layout_c_ifc.strided_batch_offset = mm_traits.c_layout_traits.batch_offset
+
+        # FP8 block scaling dimension requirements
+        if (
+            self.input_type_width == 8
+            and self.options.block_scaling
+            and (mm_traits.M % 128 != 0 or mm_traits.N % 128 != 0 or mm_traits.K % 128 != 0)
+        ):
+            raise ValueError(f"M={mm_traits.M} N={mm_traits.N} K={mm_traits.K} must be divisible by 128 for FP8 block_scaling.")
+
+        # Note: FP4 block scaling dimension requirements (M,N % 128, K % 64) are checked
+        # earlier in __init__ via validate_fp4_ab_dimensions().
+
+        # General FP8 alignment
+        if self.input_type_width == 8 and (mm_traits.M % 16 != 0 or mm_traits.N % 16 != 0 or mm_traits.K % 16 != 0):
+            raise ValueError(f"M={mm_traits.M} N={mm_traits.N} K={mm_traits.K} must be divisible by 16 for FP8 operations")
+
+        # General FP4 alignment
+        if self.using_fp4_ab and (mm_traits.M % 16 != 0 or mm_traits.N % 16 != 0 or mm_traits.K % 16 != 0):
+            raise ValueError(f"M={mm_traits.M} N={mm_traits.N} K={mm_traits.K} must be divisible by 16 for FP4 operations")
+
+        if self.options.block_scaling and self.d_dtype_width <= 8:  # FP8 and FP4
+            self.mm_desc_ifc.alpha_vector_batch_stride = 1  # Workaround for library caching issue
+
+            # cublasLtMatmulAlgoGetHeuristic requires the scale pointer to be set.
+            num_output_elements = mm_traits.M * mm_traits.N * mm_traits.batch_count
+            d_out_scale, d_out_scale_mode = _create_d_out_scale_and_scale_mode(
+                self.result_class, num_output_elements, self.d_dtype_width, self.device_id, stream_holder
+            )
+            self.aux_outputs = {"d_out_scale": d_out_scale}
+            self.mm_desc_ifc.d_out_scale_pointer = d_out_scale.data_ptr
+            self.mm_desc_ifc.d_out_scale_mode = d_out_scale_mode
+
+        limit = preferences.limit
+        if algorithms is None:
+            num_algorithms = np.empty((1,), dtype=np.int32)
+            self.algorithms_buffer = cublaslt.MatmulHeuristicResult(limit)
+        else:
+            assert all(isinstance(algo, _algorithmmod.Algorithm) for algo in algorithms), (
+                "The algorithms passed to plan() are of wrong type."
+            )
+            num_algorithms = len(algorithms)
+
+        self.preference_ptr = cublaslt.matmul_preference_create()
+
+        if self.input_type_width <= 8:
+            self._prepare_operand_quantization_scales(self.quantization_scales, stream_holder)
+
+        if algorithms is None:
+            # Set preferences.
+            preference_ifc = matmul_pref_ifc.MatmulPreferenceInterface(self.preference_ptr)
+            preference_ifc.max_workspace_bytes = self.memory_limit
+            preference_ifc.reduction_scheme_mask = preferences.reduction_scheme_mask
+            preference_ifc.max_waves_count = preferences.max_waves_count
+            preference_ifc.impl_mask = preferences.numerical_impl_mask
+
+            # Set minimum alignments.
+            a_ptr, b_ptr = self.operands[0].data_ptr, self.operands[1].data_ptr
+            preference_ifc.min_alignment_a_bytes = min(256, pointer_aligned_to(a_ptr))
+            preference_ifc.min_alignment_b_bytes = min(256, pointer_aligned_to(b_ptr))
+            self.logger.debug(f"The minimum alignment for operand A is {preference_ifc.min_alignment_a_bytes} bytes.")
+            self.logger.debug(f"The minimum alignment for operand B is {preference_ifc.min_alignment_b_bytes} bytes.")
+            if self.num_operands == 3:
+                c_ptr = self.operands[2].data_ptr
+                preference_ifc.min_alignment_c_bytes = min(256, pointer_aligned_to(c_ptr))
+                self.logger.debug(f"The minimum alignment for operand C is {preference_ifc.min_alignment_c_bytes} bytes.")
+            # The result alignment should be 256 bytes.
+            self.logger.debug("The minimum alignment for the result D is the default 256 bytes.")
+
+            self.logger.info("Starting matrix multiplication planning...")
+            assert isinstance(self.device_id, int), self.device_id
+            assert stream_holder is not None
+            with utils.cuda_call_ctx(stream_holder, blocking=True, timing=log_info) as (
+                _,
+                elapsed,
+            ):
+                cublaslt.matmul_algo_get_heuristic(
+                    self.handle,
+                    self.mm_desc,
+                    self.a_layout_ptr,
+                    self.b_layout_ptr,
+                    self.c_layout_ptr,
+                    self.d_layout_ptr,
+                    self.preference_ptr,
+                    limit,
+                    self.algorithms_buffer.ptr,
+                    num_algorithms.ctypes.data,
+                )
+
+            num_algorithms = num_algorithms[0]
+            if num_algorithms == 0:
+                raise RuntimeError("Planning failed to find any suitable algorithm.")
+            assert self.algorithms_buffer is not None, (
+                "Internal Error. self.algorithms_buffer should have been set by self.plan()."
+            )
+            self.algorithms_buffer = self.algorithms_buffer[:num_algorithms]
+
+            # Create algorithm objects.
+            self.algorithm_objects = tuple(_algorithmmod.Algorithm(a) for a in self.algorithms_buffer)
+        else:
+            self.algorithm_objects = tuple(algorithms)
+            self.algorithms_buffer = cublaslt.MatmulHeuristicResult(len(algorithms))
+            for i, algo in enumerate(algorithms):
+                # we wrap it too well that it's hard to copy-construct...
+                self.algorithms_buffer[i] = algo.algorithm._data
+
+        # Cache the first (best) algorithm struct.
+        self.cached_best_algorithm_struct = self.algorithms_buffer[0]["algo"]
+
+        # Create the map from object to buffer.
+        self.algorithm_object_to_buffer = dict(zip(self.algorithm_objects, self.algorithms_buffer, strict=True))
+
+        workspace_size = int(np.max(self.algorithms_buffer["workspace_size"]))
+        if workspace_size > 0 and self.epilog:
+            workspace_size += 16  # Workaround for library issue
+        self.workspace.set_size(workspace_size)
+
+        if algorithms is None:
+            self.logger.info(
+                f"The plan found {num_algorithms} suitable algorithms within the requested limit of {limit} "
+                f"algorithms, with a workspace requirement of {formatters.MemoryStr(workspace_size)}."
+            )
+        else:
+            self.logger.info(
+                f"The plan is using {num_algorithms} algorithm passed through the algorithms argument, with a "
+                f"workspace requirement of {formatters.MemoryStr(workspace_size)}."
+            )
+
+        self.mm_planned = True
+        if algorithms is None and elapsed.data is not None:
+            self.logger.info(f"The matrix multiplication planning phase took {elapsed.data:.3f} ms to complete.")
+
+        return self.algorithm_objects
+
+    @property
+    def algorithms(self):
+        """
+        After planning using :meth:`plan()`, get the sequence of algorithm objects to
+        inquire their capabilities, configure them, or serialize them for later use.
+
+        Returns:
+            A sequence of :class:`nvmath.linalg.advanced.Algorithm` objects that are
+            applicable to this matrix multiplication problem specification.
+        """
+        return self.algorithm_objects
+
+    def _check_and_set_operand(
+        self,
+        operand,
+        operand_name,
+        mm_desc_ifc,
+        stream_holder,
+        *,
+        operand_index=None,
+        epilog_name=None,
+        package=None,
+        dtype=None,
+        extents=None,
+        strides=None,
+    ):
+        """
+        Check to make sure that the provided operand is consistent with the one it's
+        updating, and update it.
+        """
+        assert (operand_index is None) ^ (epilog_name is None), "Internal Error."
+        assert self.operands is not None, "Internal Error."
+
+        # Make sure that the data type and extents match.
+        utils.check_attribute_match(dtype, operand.dtype, "data type")
+        utils.check_attribute_match(extents, operand.shape, "extents")
+
+        package = utils.infer_object_package(operand.tensor)
+
+        # Conjugate flag of the provided operands must match the original qualifiers
+        if operand_index is not None and self.lazy_conjugation[operand_index] != operand.is_conjugate:
+            raise ValueError(f"The provided operand {operand_name} has different conjugate flag than the original operand")
+
+        device_id = operand.device_id
+        # When memory_space is "cpu", operands should be on CPU even though
+        # self.device_id is the execution device.
+        expected_device_id = "cpu" if self.memory_space == "cpu" else self.device_id
+        if expected_device_id != device_id:
+            raise ValueError(
+                f'The operand {operand_name} is on device "{device_id}", but it should be on device '
+                f'"{expected_device_id}" to match the original operand.'
+            )
+
+        if device_id == "cpu":
+            if self.package != package:
+                message = f"Library package mismatch: '{self.package}' => '{package}'"
+                raise TypeError(message)
+
+            # Check if we have a GPU buffer to update into.
+            if operand_index is not None:
+                o = self.operands[operand_index]
+            else:
+                o = self.epilog_operands[epilog_name]
+            if o is None:  # No buffer, create one.
+                # Copy operand across memory spaces (CPU to GPU).
+                o = operand.to(self.device_id, stream_holder)
+                if operand_index is not None:
+                    self.operands[operand_index] = o
+                else:
+                    self.epilog_operands[epilog_name] = o
+                    # Update the epilog pointer, since we're starting afresh.
+                    self.epilog_input_name_to_handler[epilog_name].update(mm_desc_ifc, o)
+            else:
+                # In-place copy to existing device pointer because the new operand is on the
+                # CPU.
+                o.copy_(operand, stream_holder=stream_holder)
+            if self.memory_space == "cpu" and operand_index == 2:
+                # Hold references, needed for inplace operations.
+                self.cpu_c_ref = operand
+        else:
+            if self.package != package:
+                message = f"Library package mismatch: '{self.package}' => '{package}'"
+                raise TypeError(message)
+
+            utils.check_attribute_match(strides, operand.strides, "strides")
+
+            # Finally, replace the original operand by the new one.
+            if operand_index is not None:
+                self.operands[operand_index] = operand
+            else:
+                self.epilog_operands[epilog_name] = operand
+                # Update the epilog pointer, since we're starting afresh.
+                self.epilog_input_name_to_handler[epilog_name].update(mm_desc_ifc, operand)
+
+        self.logger.info(f"Operand '{operand_name}' has been reset to the new value.")
+
+        return
+
+    def _reset_quantization_scales_unchecked(self, quantization_scales_obj, stream, stream_holder):
+        """
+        Unchecked reset of quantization scales for operands A, B, C, and D.
+
+        Works for both CPU and CUDA memory spaces.  For each provided scale the
+        function updates the host-side value in ``self.quantization_scales``, updates
+        the existing device-side holder in ``self.quantization_scales_device``
+        (copy in-place or rebind ``.tensor``), and refreshes the descriptor pointer.
+
+        A stream holder is created lazily only when a CPU-to-GPU copy is needed
+        (scalar scales, or tensor scales in CPU memory space).
+
+        Returns the stream_holder (may be newly created).
+        """
+        mm_desc_ifc = self.mm_desc_ifc
+        qsd = self.quantization_scales_device
+
+        def _update(holder, scale):
+            nonlocal stream_holder
+            if isinstance(scale, (int, float)):
+                if stream_holder is None:
+                    stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+                src = tensor_wrapper.wrap_operand(np.asarray(scale, dtype="float32"))
+                _realloc_if_needed_and_copy_to_mirror(src, holder, self.device_id, stream_holder)
+            elif self.memory_space == "cpu":
+                if stream_holder is None:
+                    stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+                _realloc_if_needed_and_copy_to_mirror(tensor_wrapper.wrap_operand(scale), holder, self.device_id, stream_holder)
+            else:
+                holder.tensor = scale
+
+        if quantization_scales_obj.a is not None:
+            self.quantization_scales.a = quantization_scales_obj.a
+            holder = qsd["a"]
+            _update(holder, quantization_scales_obj.a)
+            mm_desc_ifc.set_a_scale_pointer_unchecked(holder.data_ptr)
+
+        if quantization_scales_obj.b is not None:
+            self.quantization_scales.b = quantization_scales_obj.b
+            holder = qsd["b"]
+            _update(holder, quantization_scales_obj.b)
+            mm_desc_ifc.set_b_scale_pointer_unchecked(holder.data_ptr)
+
+        if quantization_scales_obj.c is not None:
+            self.quantization_scales.c = quantization_scales_obj.c
+            holder = qsd["c"]
+            _update(holder, quantization_scales_obj.c)
+            mm_desc_ifc.set_c_scale_pointer_unchecked(holder.data_ptr)
+
+        if quantization_scales_obj.d is not None:
+            self.quantization_scales.d = quantization_scales_obj.d
+            holder = qsd["d"]
+            _update(holder, quantization_scales_obj.d)
+            mm_desc_ifc.set_d_scale_pointer_unchecked(holder.data_ptr)
+
+        return stream_holder
+
+    def _reset_aux_quantization_scale_unchecked(self, aux_quantization_scale, stream, stream_holder):
+        """
+        Unchecked reset of aux_quantization_scale.
+
+        Works for both CPU and CUDA memory spaces.  Updates the existing
+        device-side holder (copy in-place or rebind ``.tensor``) and refreshes
+        the descriptor pointer.
+        A stream holder is created lazily only when a CPU-to-GPU copy is needed.
+
+        Returns the stream_holder (may be newly created).
+        """
+        holder = self.quantization_scales_device["epilog_aux"]
+
+        if isinstance(aux_quantization_scale, (int, float)):
+            if stream_holder is None:
+                stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+            src = tensor_wrapper.wrap_operand(np.asarray(aux_quantization_scale, dtype="float32"))
+            _realloc_if_needed_and_copy_to_mirror(src, holder, self.device_id, stream_holder)
+        elif self.memory_space == "cpu":
+            if stream_holder is None:
+                stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+            _realloc_if_needed_and_copy_to_mirror(
+                tensor_wrapper.wrap_operand(aux_quantization_scale), holder, self.device_id, stream_holder
+            )
+        else:
+            holder.tensor = aux_quantization_scale
+
+        self.mm_desc_ifc.epilogue_aux_scale_pointer = holder.data_ptr
+        return stream_holder
+
+    def reset_operands_unchecked(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        quantization_scales=None,
+        epilog_inputs=None,
+        stream: utils.AnyStream | int | None = None,
+    ):
+        """
+        {reset_operands_unchecked}
+        """
+        if alpha is not None:
+            self.alpha[0] = alpha
+
+        if beta is not None:
+            self.beta[0] = beta
+
+        stream_holder = None
+
+        if epilog_inputs is not None and "aux_quantization_scale" in epilog_inputs:
+            aux_quantization_scale = epilog_inputs.get("aux_quantization_scale")
+            stream_holder = self._reset_aux_quantization_scale_unchecked(aux_quantization_scale, stream, stream_holder)
+
+        quantization_scales_obj = quantization_scales
+        if quantization_scales is not None and isinstance(quantization_scales, dict):
+            quantization_scales_obj = _configuration.MatmulQuantizationScales(**quantization_scales)
+
+        if self.memory_space == "cuda":
+            if a is not None:
+                self.operands[0].tensor = a  # type: ignore[index, union-attr]
+            if b is not None:
+                self.operands[1].tensor = b  # type: ignore[index, union-attr]
+            if c is not None:
+                self.operands[2].tensor = c  # type: ignore[index, union-attr]
+
+            if epilog_inputs is not None:
+                for epilog_name, epilog_value in epilog_inputs.items():
+                    if epilog_name != "aux_quantization_scale":
+                        self.epilog_operands[epilog_name].tensor = epilog_value
+                        self.epilog_input_name_to_handler[epilog_name].update_pointer(
+                            self.mm_desc_ifc, self.epilog_operands[epilog_name].data_ptr
+                        )
+
+        else:
+            if stream_holder is None:
+                stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+
+            if a is not None:
+                a_wrapped = tensor_wrapper.wrap_operand(a)
+                _realloc_if_needed_and_copy_to_mirror(a_wrapped, self.operands[0], self.device_id, stream_holder)  # type: ignore[index, union-attr]
+            if b is not None:
+                b_wrapped = tensor_wrapper.wrap_operand(b)
+                _realloc_if_needed_and_copy_to_mirror(b_wrapped, self.operands[1], self.device_id, stream_holder)  # type: ignore[index, union-attr]
+            if c is not None:
+                c_wrapped = tensor_wrapper.wrap_operand(c)
+                _realloc_if_needed_and_copy_to_mirror(c_wrapped, self.operands[2], self.device_id, stream_holder)  # type: ignore[index, union-attr]
+                if self.inplace:
+                    self.cpu_c_ref = c_wrapped
+
+            if epilog_inputs is not None:
+                for epilog_name, epilog_value in epilog_inputs.items():
+                    if epilog_name != "aux_quantization_scale":
+                        epilog_input_wrapped = tensor_wrapper.wrap_operand(epilog_value)
+                        reallocated = _realloc_if_needed_and_copy_to_mirror(
+                            epilog_input_wrapped, self.epilog_operands[epilog_name], self.device_id, stream_holder
+                        )
+                        if reallocated:
+                            self.epilog_input_name_to_handler[epilog_name].update_pointer(
+                                self.mm_desc_ifc, self.epilog_operands[epilog_name].data_ptr
+                            )
+
+        if quantization_scales_obj is not None:
+            self._reset_quantization_scales_unchecked(quantization_scales_obj, stream, stream_holder)
+
+        self._operands_released = False
+
+    @utils.precondition(_check_valid_matmul)
+    def reset_operands(
+        self,
+        *,
+        a=None,
+        b=None,
+        c=None,
+        alpha=None,
+        beta=None,
+        quantization_scales=None,
+        epilog_inputs=None,
+        stream: utils.AnyStream | int | None = None,
+    ):
+        """
+        Reset one or more operands held by this :class:`Matmul` instance.
+
+        .. versionchanged:: 0.9
+            All parameters are now keyword-only.
+
+        Args:
+            a: {a}
+
+            b: {b}
+
+            c: {c}
+                {c_admonitions}
+
+            alpha: {alpha}
+
+            beta: {beta}
+
+            epilog_inputs: {epilog_inputs}
+
+            stream: {stream}
+
+            quantization_scales: {quantization_scales}
+
+        Semantics:
+            - Only the operands explicitly passed are updated. At least one operand
+              is required (all of them after :meth:`release_operands`), otherwise
+              a :class:`ValueError` is raised.
+
+            - This method will perform various checks on the new operands to make sure:
+
+              - The shapes, strides, datatypes match those of the old ones.
+              - The packages that the operands belong to match those of the old ones.
+              - If input tensors are on GPU, the device must match.
+
+        Examples:
+
+            >>> import cupy as cp
+            >>> import nvmath
+
+            Create two 3-D float64 ndarrays on the GPU:
+
+            >>> M, N, K = 128, 128, 256
+            >>> a = cp.random.rand(M, K)
+            >>> b = cp.random.rand(K, N)
+
+            Create an matrix multiplication object as a context manager
+
+            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
+            ...     # Plan the operation.
+            ...     algorithms = mm.plan()
+            ...
+            ...     # Execute the MM to get the first result.
+            ...     r1 = mm.execute()
+            ...
+            ...     # Reset the operands to new CuPy ndarrays.
+            ...     a_new = cp.random.rand(M, K)
+            ...     b_new = cp.random.rand(K, N)
+            ...     mm.reset_operands(a=a_new, b=b_new)
+            ...
+            ...     # Execute to get the new result corresponding to the updated operands.
+            ...     r2 = mm.execute()
+
+            With :meth:`reset_operands`, minimal overhead is achieved as problem
+            specification and planning are only performed once.
+
+            For the particular example above, the operands are on the GPU, so calling
+            :meth:`reset_operands` only updates internal references and is efficient. An
+            alternative would be to modify the existing operands in-place (e.g.
+            ``a[:]=a_new`` and ``b[:]=b_new``), but that would copy data and have
+            performance implications. When using in-place updates, the operand memory space
+            must be accessible from the execution space.
+
+            For more details, please refer to `inplace update example
+            <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul/example05_stateful_inplace.py>`_.
+
+        .. seealso::
+            :meth:`release_operands`, :meth:`reset_operands_unchecked`
+        """
+
+        if c is not None and self.num_operands == 2:
+            raise ValueError(
+                "The matrix multiplication problem specification does not include operand C, so it cannot be reset."
+            )
+
+        # If operands have been released, all required operands must be provided.
+        if self._operands_released:
+            # Check that all main operands are provided
+            all_main_provided = a is not None and b is not None and (self.num_operands != 3 or c is not None)
+
+            # Check epilog_inputs if required
+            epilog_names = self.epilog_inputs_traits.keys()
+            epilog_ok = True
+            if epilog_names:
+                if epilog_inputs is None:
+                    epilog_ok = False
+                elif epilog_names != epilog_inputs.keys():
+                    raise ValueError(
+                        f"The epilog inputs {epilog_names} are required. "
+                        f"The provided epilog input names are {epilog_inputs.keys()}."
+                    )
+
+            scales_ok = True
+            needs_scales = self.input_type_width <= 8
+            if needs_scales and quantization_scales is None:
+                scales_ok = False
+
+            if not all_main_provided or not epilog_ok or not scales_ok:
+                raise ValueError(
+                    "After release_operands(), all required operands must be provided to reset_operands(). "
+                    f"Required: a, b{', c' if self.num_operands == 3 else ''}"
+                    f"{', quantization_scales' if needs_scales else ''}"
+                    f"{', epilog_inputs' if epilog_names else ''}"
+                )
+
+            # Initialize operands and epilog_operands to prepare for new values
+            self.operands = [None] * self.num_operands  # type: ignore[list-item]
+            epilog_names = self.epilog_inputs_traits.keys()
+            self.epilog_operands = dict.fromkeys(epilog_names)
+            if needs_scales:
+                self.quantization_scales = _configuration.MatmulQuantizationScales()
+        elif all(arg is None for arg in (a, b, c, alpha, beta, quantization_scales, epilog_inputs)):
+            # All arguments are None: there is nothing to update, so reject the call.
+            raise ValueError("reset_operands() requires at least one operand to be provided.")
+
+        # Update alpha.
+        if alpha is not None:
+            try:
+                self.alpha[0] = alpha
+            except (ValueError, TypeError) as e:
+                raise ValueError(
+                    f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'."
+                ) from e
+
+        # Update beta.
+        if beta is not None:
+            if self.num_operands == 2:
+                self.logger.warning(f"Matmul: The provided beta value {beta} is ignored since operand C is not specified.")
+            else:
+                try:
+                    self.beta[0] = beta
+                except (ValueError, TypeError) as e:
+                    raise ValueError(
+                        f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'."
+                    ) from e
+
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+
+        if epilog_inputs is not None and "aux_quantization_scale" in epilog_inputs:
+            epilog_inputs = epilog_inputs.copy()
+            aux_quantization_scale = epilog_inputs.pop("aux_quantization_scale")
+            self._validate_epilog_aux_scale(aux_quantization_scale, required=False)
+            self._prepare_single_validated_scale(
+                aux_quantization_scale, "epilog_aux", cublas_operand="epilogue_aux", stream_holder=stream_holder
+            )
+
+        # Reset the provided operands.
+        if a is not None:
+            a = tensor_wrapper.wrap_operand(a)
+            index = 0
+            self._check_and_set_operand(
+                a,
+                "A",
+                self.mm_desc_ifc,
+                stream_holder,
+                operand_index=index,
+                dtype=self.a_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+        if b is not None:
+            b = tensor_wrapper.wrap_operand(b)
+            index = 1
+            self._check_and_set_operand(
+                b,
+                "B",
+                self.mm_desc_ifc,
+                stream_holder,
+                operand_index=index,
+                dtype=self.b_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+        if c is not None:  # If we get here, we know that C is one of the operands in the problem specification.
+            c = tensor_wrapper.wrap_operand(c)
+            index = 2
+            self._check_and_set_operand(
+                c,
+                "C",
+                self.mm_desc_ifc,
+                stream_holder,
+                operand_index=index,
+                dtype=self.c_dtype_name,
+                extents=self.operand_extents[index],
+                strides=self.operand_strides[index],
+            )
+
+        # Update quantization_scales after operands are set, so _validate_operand_scales
+        # can read self.operands[i].size for block-scaling shape checks.
+        if quantization_scales is not None:
+            quantization_scales = self._validate_operand_scales(quantization_scales, all_required=False)
+            if quantization_scales.a is not None:
+                self.quantization_scales.a = quantization_scales.a
+            if quantization_scales.b is not None:
+                self.quantization_scales.b = quantization_scales.b
+            if quantization_scales.c is not None:
+                self.quantization_scales.c = quantization_scales.c
+            if quantization_scales.d is not None:
+                self.quantization_scales.d = quantization_scales.d
+            self._prepare_operand_quantization_scales(self.quantization_scales, stream_holder)
+
+        # Reset the provided epilog inputs.
+        if epilog_inputs is not None:
+            for name in epilog_inputs:
+                epilog_input = tensor_wrapper.wrap_operand(epilog_inputs[name])
+                self._check_and_set_operand(
+                    epilog_input,
+                    name,
+                    self.mm_desc_ifc,
+                    stream_holder,
+                    epilog_name=name,
+                    dtype=self.epilog_inputs_traits[name].dtype,
+                    extents=self.epilog_inputs_traits[name].extents,
+                    strides=self.epilog_inputs_traits[name].strides,
+                )
+
+        self._operands_released = False
+
+    @utils.precondition(_check_valid_matmul)
+    def release_operands(self):
+        """
+        {release_operands}
+        """
+        if self._operands_released:
+            self.logger.info("Operands have already been released; nothing to do.")
+            return
+
+        # CUDA memory space:
+        #   self.operands, self.epilog_operands, and
+        #   self.quantization_scales_device hold direct user references;
+        #   self.quantization_scales holds the user-provided scales object.
+        # CPU memory space:
+        #   self.operands, self.epilog_operands, and
+        #   self.quantization_scales_device hold internal device mirrors;
+        #   self.quantization_scales holds the user-provided scales object;
+        #   self.cpu_c_ref holds a direct reference to the user's CPU operand C.
+        # In both cases, release all of them.
+        # Note that if/when possible, we keep the TensorHolder wrappers and
+        # container structures alive and only release the internal tensor reference.
+        # This is useful when reset_operands_unchecked is called subsequently
+        # because it can reuse the existing wrappers, saving overhead.
+
+        self.operands[0].tensor = None  # A
+        self.operands[1].tensor = None  # B
+        if self.num_operands == 3:
+            self.operands[2].tensor = None  # C
+        for op in self.epilog_operands.values():
+            op.tensor = None
+        for op in self.quantization_scales_device.values():
+            op.tensor = None
+
+        # For the quant scales, the attribute itself might not exist,
+        # so we need to check for that first.
+        if hasattr(self, "quantization_scales"):
+            self.quantization_scales.a = None
+            self.quantization_scales.b = None
+            self.quantization_scales.c = None
+            self.quantization_scales.d = None
+
+        if self.memory_space == "cpu" and self.cpu_c_ref is not None:
+            self.cpu_c_ref.tensor = None
+
+        self._operands_released = True
+        self.logger.info("User-provided operands have been released.")
+
+    @utils.precondition(_check_valid_matmul)
+    @utils.precondition(_check_planned, "Autotuning")
+    @utils.precondition(_check_valid_operands, "Autotuning")
+    def autotune(
+        self, iterations=10, prune=None, release_workspace=False, stream: utils.AnyStream | int | None = None
+    ):  # Prune means keep top N of the algorithms only.
+        """
+        Autotune the matrix multiplication to order the algorithms from the fastest measured
+        execution time to the slowest. Once autotuned, the optimally-ordered algorithm
+        sequence can be accessed using :py:attr:`algorithms`.
+
+        .. note::
+            This function will benchmark each of the algorithms and order the algorithms
+            based on the benchmark results. The measurements can be impacted by factors
+            such as GPU temperature, clock settings, or power consumption. Autotuning
+            in an unstable environment can result in a suboptimal algorithm ordering.
+            If you experience performance problems, consider omitting the autotuning.
+
+        Args:
+            iterations: The number of autotuning iterations to perform.
+
+            prune: An integer N, specifying the top N fastest algorithms to retain after
+                autotuning. The default is to retain all algorithms.
+
+            release_workspace: {release_workspace}
+
+            stream: {stream}
+        """
+        self.logger.info("= AUTOTUNING PHASE =")
+        # Measure time taken for autotuning.
+        from timeit import default_timer as timer
+
+        self.logger.info("Starting autotuning...")
+        start = timer()
+
+        assert self.algorithm_objects is not None, "Internal Error. self.algorithm_objects should have been set by self.plan()."
+        num_algorithms = len(self.algorithm_objects)
+        if num_algorithms == 1:
+            self.logger.info("Skipping the autotuning, because only one algorithm has been returned in the planning phase.")
+            return
+        limit = min(prune, num_algorithms) if prune is not None else num_algorithms
+        self.logger.info(
+            f"The number of algorithms in the plan is {num_algorithms}, from which the top {limit} will be retained."
+        )
+        self.logger.info(f"The requested number of iterations is {iterations}.")
+
+        # Autotune setup.
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+
+        # Create empty tensors for auxiliary output.
+        epilog_outputs = {}
+        for handler in self.epilog_output_handlers:
+            name = handler.name
+            shape, strides, dtype_name = handler.attributes()
+            epilog_outputs[name] = aux = utils.create_empty_tensor(
+                self.result_class,
+                shape,
+                dtype_name,
+                self.device_id,
+                stream_holder,
+                verify_strides=False,
+                strides=strides,
+            )
+
+            # Update the data pointer in the MM descriptor.
+            handler.update_pointer(self.mm_desc_ifc, aux.data_ptr)
+
+        # Take a copy of `c` for autotuning to avoid clobbering it if inplace=True.
+        c = None
+        if self.inplace:
+            # `c` is guaranteed to exist. Clone `c` and copy into it, since `tensor.to`
+            # returns the original tensor when the device ID is the same.
+            c = self.operands[2]
+            c = utils.create_empty_tensor(
+                c.__class__, c.shape, c.dtype, self.device_id, stream_holder, verify_strides=False, strides=c.strides
+            )
+            c.copy_(self.operands[2], stream_holder)
+            assert c.data_ptr != self.operands[2].data_ptr, "Internal error."
+        elif self.num_operands == 3:
+            c = self.operands[2]
+
+        # Create empty tensor for the result, if the operation is not in-place.
+        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()."
+        if self.inplace:
+            result = c
+        else:
+            result = utils.create_empty_tensor(
+                self.result_class,
+                self.result_traits.result_shape,
+                self.d_dtype_name,
+                self.device_id,
+                stream_holder,
+                verify_strides=False,
+                strides=self.result_traits.result_strides,
+            )
+        result_ptr = result.data_ptr
+
+        a, b = self.operands[0], self.operands[1]
+        alpha_ptr, a_ptr, b_ptr, beta_ptr = self.alpha.ctypes.data, a.data_ptr, b.data_ptr, self.beta.ctypes.data
+        # If `c` is not provided, set c_ptr to nullptr.
+        if self.num_operands == 3:
+            c_ptr = c.data_ptr
+        else:
+            assert self.beta[0] == 0.0, "Internal error: beta must be zero if `c` is not specified."
+            c_ptr = 0
+
+        def flush_cache():
+            """
+            Write data to a temporary buffer to flush the L2 cache.
+            """
+
+            @functools.cache
+            def get_l2_cache_size(device_id):
+                device = Device(device_id)
+                return device.properties.l2_cache_size
+
+            l2_cache_size = get_l2_cache_size(self.device_id)
+            cpu_buffer = np.zeros(l2_cache_size, dtype=np.uint8)
+            tensor_wrapper.wrap_operand(cpu_buffer).to(device_id=self.device_id, stream_holder=stream_holder)
+
+        with self.workspace.allocate_perhaps(
+            stream_holder,
+            get_last_event=lambda: self.last_compute_event,
+        ) as ws:
+
+            def execute_matmul(algorithm_ptr):
+                cublaslt.matmul(
+                    self.handle,
+                    self.mm_desc,
+                    alpha_ptr,
+                    a_ptr,
+                    self.a_layout_ptr,
+                    b_ptr,
+                    self.b_layout_ptr,
+                    beta_ptr,
+                    c_ptr,
+                    self.c_layout_ptr,
+                    result_ptr,
+                    self.d_layout_ptr,
+                    algorithm_ptr,
+                    ws.raw_ptr,
+                    ws.size,
+                    stream_holder.ptr,
+                )
+
+            # Tune.
+            with utils.cuda_call_ctx(stream_holder, blocking=False, timing=False) as (
+                self.last_compute_event,
+                _,
+            ):
+                gpu_times = np.empty(shape=(len(self.algorithms_buffer), iterations), dtype=float)
+                algorithm_idxs = list(range(len(self.algorithms_buffer)))
+                timing_enabled_options = {utils.timing_enabled_name(): True}
+                start0 = stream_holder.obj.device.create_event(options=timing_enabled_options)
+                end0 = stream_holder.obj.device.create_event(options=timing_enabled_options)
+                for i in range(iterations):
+                    random.shuffle(algorithm_idxs)
+                    for algorithm_idx in algorithm_idxs:
+                        algorithm_ptr = self.algorithms_buffer[algorithm_idx]["algo"].ctypes.data
+                        stream_holder.obj.record(start0)
+                        execute_matmul(algorithm_ptr=algorithm_ptr)
+                        stream_holder.obj.record(end0)
+                        # FIXME: @dching Calling sync() here could slow down tuning by
+                        # forcing the device to wait for the host to record the elapsed
+                        # time. It may be faster to store references to 2 * len(iterations)
+                        # * len(algorithms_buffer) Events and compute the elapsed time at
+                        # the end.
+                        end0.sync()
+                        gpu_times[algorithm_idx, i] = end0 - start0
+                        # Avoid potential overflow for inplace operations.
+                        if self.inplace:
+                            c.copy_(self.operands[2], stream_holder)
+                        flush_cache()
+
+            gpu_times = np.median(gpu_times, axis=1)
+
+            if release_workspace:
+                ws.release(self.last_compute_event)
+                self.last_compute_event = None
+
+        # Get the sort order based on the GPU times.
+        sorted_gpu_times, sort_order = zip(*sorted(zip(gpu_times, range(num_algorithms), strict=True)), strict=True)
+
+        # Reorder the algorithms buffer and algorithm objects according to the sort order,
+        # and prune it.
+        sort_order = sort_order[:limit]
+        self.algorithms_buffer = self.algorithms_buffer[list(sort_order)]
+        self.algorithm_objects = tuple(self.algorithm_objects[i] for i in sort_order)
+
+        # Update cached first (best) algorithm struct after tuning.
+        self.cached_best_algorithm_struct = self.algorithms_buffer[0]["algo"]
+
+        # Create the map from object to buffer.
+        self.algorithm_object_to_buffer = dict(zip(self.algorithm_objects, self.algorithms_buffer, strict=True))
+
+        gpu_times_str = ", ".join(f"{t:0.3f}" for t in gpu_times)
+        self.logger.info(f"The autotuned GPU times (in milliseconds) are: {gpu_times_str}.")
+        self.logger.info(f"The corresponding sort order is: {sort_order}.")
+        orig_flop_rate = self.flop_count / gpu_times[0] * 1000
+        if sort_order[0] != 0:
+            self.logger.info(
+                f"Autotuning found that the algorithm originally ranked {sort_order[0]} is the best out of the "
+                f"{num_algorithms} in the plan, and moved it to rank 0."
+            )
+            new_flop_rate = self.flop_count / sorted_gpu_times[0] * 1000
+            self.logger.info(
+                f"Autotuning has improved performance from {formatters.FLOPSStr(orig_flop_rate, 'FLOP/s')} to "
+                f"{formatters.FLOPSStr(new_flop_rate, 'FLOP/s')}."
+            )
+        else:
+            self.logger.info(
+                f"Autotuning found that the algorithm ranked best by the plan heuristics remains the best out of the "
+                f"{num_algorithms} algorithms in the plan."
+            )
+            self.logger.info(f"The best performance remains at {formatters.FLOPSStr(orig_flop_rate, 'FLOP/s')}.")
+
+        end = timer()
+        self.logger.info(f"The autotuning took {(end - start) * 1000.0:.3f} ms to complete.")
+
+    def _create_result_tensor(self, stream_holder: utils.StreamHolder, log_debug: bool):
+        """
+        Create the output result tensor for the matmul operation.
+
+        We need special treatment for FP4 output (``float4_e2m1fn_x2``),
+        because the logical shape ``(M, N)`` from ``result_traits`` is converted
+        to a packed shape where one dimension is halved
+        (since two FP4 values are packed per byte).
+
+        For all other dtypes, the logical shape is used directly.
+        """
+        assert self.result_traits is not None, (
+            "Internal error: result_traits must be set by plan() before creating the result tensor."
+        )
+
+        if self.d_dtype != CudaDataType.CUDA_R_4F_E2M1:
+            # Non-FP4 output: no packing needed, use logical shape as-is.
+            result_shape = self.result_traits.result_shape
+            result_strides = self.result_traits.result_strides
+
+            if log_debug:
+                self.logger.debug(
+                    f"The output tensor shape = {result_shape} with strides = "
+                    f"{result_strides} and data type '{self.d_dtype_name}'."
+                )
+
+            self.result = utils.create_empty_tensor(
+                self.result_class,
+                result_shape,
+                self.d_dtype_name,
+                self.device_id,
+                stream_holder,
+                verify_strides=False,
+                strides=result_strides,
+            )
+            return
+
+        # If we are here, we are dealing with FP4 output that must be packed.
+        # Convert logical shape/strides to packed (physical) shape/strides.
+        # Batching dimensions (leading) are unchanged.
+        result_ndim = len(self.result_traits.result_shape)
+        assert result_ndim >= 2, "Internal error: result must be at least 2D."
+        if self.result_traits.d_layout_traits.order == cublaslt.Order.ROW:
+            packed_axis = result_ndim - 1
+        else:
+            packed_axis = result_ndim - 2
+        result_packed_shape = tuple(e if i != packed_axis else e // 2 for i, e in enumerate(self.result_traits.result_shape))
+        result_packed_strides = tuple(
+            s // 2 if i != packed_axis else s for i, s in enumerate(self.result_traits.result_strides)
+        )
+
+        if log_debug:
+            self.logger.debug(
+                f"FP4 output: logical shape = {tuple(self.result_traits.result_shape)}, "
+                f"logical strides = {tuple(self.result_traits.result_strides)} -> "
+                f"packed shape = {result_packed_shape}, packed strides = {result_packed_strides}, "
+                f"data type '{self.d_dtype_name}'."
+            )
+
+        self.result = utils.create_empty_tensor(
+            self.result_class,
+            result_packed_shape,
+            self.d_dtype_name,
+            self.device_id,
+            stream_holder,
+            verify_strides=False,
+            strides=result_packed_strides,
+        )
+
+    @utils.precondition(_check_valid_matmul)
+    @utils.precondition(_check_planned, "Execution")
+    @utils.precondition(_check_valid_operands, "Execution")
+    def execute(self, *, algorithm=None, release_workspace=False, stream: utils.AnyStream | int | None = None):
+        """
+        Execute a prepared (planned and possibly autotuned) matrix multiplication.
+
+        Args:
+            algorithm: An algorithm chosen from the sequence returned by
+                :meth:`plan` or :py:attr:`algorithms`. By default, the first algorithm in
+                the sequence is used.
+
+                .. experimental:: parameter
+
+            release_workspace: {release_workspace}
+
+            stream: {stream}
+
+        Returns:
+           {result}
+        """
+        log_info = self.logger.isEnabledFor(logging.INFO)
+        log_debug = self.logger.isEnabledFor(logging.DEBUG)
+
+        if log_info:
+            self.logger.info("= EXECUTION PHASE =")
+        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.internal_op_package)
+        if log_info:
+            self.logger.info(f"The specified stream for execute() is {stream_holder.obj}.")
+
+        # Create empty tensors for auxiliary output.
+        for handler in self.epilog_output_handlers:
+            name = handler.name
+            shape, strides, dtype_name = handler.attributes()
+            if log_debug:
+                self.logger.debug(f"Beginning auxiliary output tensor '{name}' creation...")
+                self.logger.debug(f"The '{name}' tensor shape = {shape} with strides = {strides} and data type '{dtype_name}'.")
+            self.epilog_outputs[name] = aux = utils.create_empty_tensor(
+                self.result_class,
+                shape,
+                dtype_name,
+                self.device_id,
+                stream_holder,
+                verify_strides=False,
+                strides=strides,
+            )
+            if log_debug:
+                self.logger.debug(f"The auxiliary output tensor '{name}' has been created.")
+            if self.preferences.epilog.aux_amax:  # type: ignore[attr-defined]
+                if "float8" not in dtype_name:
+                    raise ValueError("epilog.aux_amax=True is not supported when epilog output type is not FP8.")
+                self.epilog_outputs[f"{name}_amax"] = utils.create_empty_tensor(
+                    self.result_class,
+                    (1,),
+                    "float32",  # This is the only type allowed by cuBLAS for AMAX.
+                    self.device_id,
+                    stream_holder,
+                    verify_strides=False,
+                )
+                self.mm_desc_ifc.epilogue_aux_amax_pointer = self.epilog_outputs[f"{name}_amax"].data_ptr
+
+            # Update the data pointer in the MM descriptor.
+            handler.update_pointer(self.mm_desc_ifc, aux.data_ptr)
+
+        # Create empty tensor for the result, if the operation is not in-place.
+        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()"
+        if self.inplace:
+            if log_debug:
+                self.logger.debug("The operation is in-place (operand C will be overwritten).")
+            self.result = self.operands[2]
+        else:
+            if log_debug:
+                self.logger.debug("Beginning output (empty) tensor creation...")
+            self._create_result_tensor(stream_holder, log_debug)
+            if log_debug:
+                self.logger.debug("The output (empty) tensor has been created.")
+
+        self.aux_outputs = {}
+
+        if self.options.result_amax:
+            self.aux_outputs["result_amax"] = utils.create_empty_tensor(
+                self.result_class,
+                (1,),
+                "float32",  # This is the only type allowed by cuBLAS for AMAX.
+                self.device_id,
+                stream_holder,
+                verify_strides=False,
+            )
+            self.mm_desc_ifc.amax_d_pointer = self.aux_outputs["result_amax"].data_ptr
+
+        if self.options.block_scaling and self.d_dtype_width <= 8:  # FP8 and FP4
+            num_output_elements = (
+                self.mm_traits.batch_count * self.result_traits.result_shape[-1] * self.result_traits.result_shape[-2]
+            )
+            d_out_scale, _ = _create_d_out_scale_and_scale_mode(
+                self.result_class, num_output_elements, self.d_dtype_width, self.device_id, stream_holder
+            )
+            self.aux_outputs["d_out_scale"] = d_out_scale
+            self.mm_desc_ifc.d_out_scale_pointer = d_out_scale.data_ptr
+
+        # Select the first (best) algorithm if one is not provided.
+        if algorithm is None:
+            algorithm_struct = self.cached_best_algorithm_struct
+            if log_info:
+                self.logger.info(
+                    "The highest ranked algorithm in the plan (algorithm id = "
+                    f"{self.algorithm_objects[0].algorithm_id}) will be used."
+                )
+        else:
+            if algorithm not in self.algorithm_objects:
+                raise ValueError("Algorithm passed to execute() has to be included in the plan() algorithms")
+            algorithm_struct = self.algorithm_object_to_buffer[algorithm]["algo"]
+            if log_info:
+                self.logger.info(f"The specified algorithm (algorithm id = {algorithm.algorithm_id}) will be used.")
+
+        a, b = self.operands[0], self.operands[1]
+
+        # If `c` is not provided, set c_ptr to nullptr.
+        if self.num_operands == 3:
+            c_ptr = self.operands[2].data_ptr
+        else:
+            assert self.beta[0] == 0.0, "Internal error: beta must be zero if `c` is not specified."
+            c_ptr = 0
+
+        with self.workspace.allocate_perhaps(
+            stream_holder,
+            get_last_event=lambda: self.last_compute_event,
+        ) as ws:
+            if log_info:
+                self.logger.info("Starting matrix multiplication...")
+                self.logger.info(f"{self.call_prologue}")
+            with utils.cuda_call_ctx(stream_holder, self.blocking, timing=log_info) as (
+                self.last_compute_event,
+                elapsed,
+            ):
+                cublaslt.matmul(
+                    self.handle,
+                    self.mm_desc,
+                    self.alpha.ctypes.data,
+                    a.data_ptr,
+                    self.a_layout_ptr,
+                    b.data_ptr,
+                    self.b_layout_ptr,
+                    self.beta.ctypes.data,
+                    c_ptr,
+                    self.c_layout_ptr,
+                    self.result.data_ptr,
+                    self.d_layout_ptr,
+                    algorithm_struct.ctypes.data,
+                    ws.raw_ptr,
+                    ws.size,
+                    stream_holder.ptr,
+                )
+
+            if log_info and elapsed.data is not None:
+                self.logger.info(f"The matrix multiplication calculation took {elapsed.data:.3f} ms to complete.")
+
+            if release_workspace:
+                ws.release(self.last_compute_event)
+                self.last_compute_event = None
+
+        # Return the result and auxiliary outputs, if present.
+        all_outputs = self.epilog_outputs | self.aux_outputs
+        if self.memory_space == "cpu":
+            if self.inplace:
+                # Overwrite operand C.
+                assert self.cpu_c_ref is not None, "Internal error."
+                c = self.cpu_c_ref
+                c.copy_(self.result, stream_holder=stream_holder)
+                out = c.tensor
+            else:
+                out = self.result.to("cpu", stream_holder=stream_holder).tensor
+            # Copy auxiliary output to CPU.
+            aux = {name: all_outputs[name].to("cpu", stream_holder=stream_holder).tensor for name in all_outputs}
+        else:
+            out = self.result.tensor
+            # Return the unwrapped epilog output tensor(s).
+            aux = {name: all_outputs[name].tensor for name in all_outputs}
+
+        # Release internal reference to the result to permit recycling of memory.
+        self.result = None
+        self.aux_outputs = {}
+        self.epilog_outputs = {}
+
+        if aux:
+            return out, aux
+
+        return out
+
+    def free(self):
+        """Free Matmul resources.
+
+        It is recommended that the :class:`Matmul` object be used within a context, but if
+        it is not possible then this method must be called explicitly to ensure that the
+        matrix multiplication resources (especially internal library objects) are properly
+        cleaned up.
+        """
+
+        if not self.valid_state:
+            return
+
+        try:
+            # Ensure ordering with respect to the last computation to avoid
+            # race conditions when releasing internal resources.
+            self.workspace.release(self.last_compute_event)
+            self.last_compute_event = None
+
+            self._free_plan_resources()
+
+            # Destroy matmul descriptor.
+            if self.mm_desc is not None:
+                cublaslt.matmul_desc_destroy(self.mm_desc)
+                self.mm_desc = None
+
+            # Note: self.handle is obtained via get_handle and doesn't need
+            # explicit destruction, it's cached globally
+            # and cleaned up at program exit.
+            # Set all attributes to None except for logger and valid_state.
+            _keep = {"logger", "valid_state"}
+            for attr in list(vars(self)):
+                if attr not in _keep:
+                    setattr(self, attr, None)
+
+        except Exception as e:
+            self.logger.critical("Internal error: only part of the Matmul object's resources have been released.")
+            self.logger.critical(str(e))
+            raise e
+        finally:
+            self.valid_state = False
+
+        self.logger.info("The Matmul object's resources have been released.")
+
+
+@utils.docstring_decorator(SHARED_MM_DOCUMENTATION, skip_missing=False)
+def matmul(
+    a,
+    b,
+    /,
+    c=None,
+    *,
+    alpha=None,
+    beta=None,
+    epilog=None,
+    epilog_inputs=None,
+    qualifiers=None,
+    quantization_scales=None,
+    options: _configuration.MatmulOptions | dict[str, typing.Any] | None = None,
+    preferences=None,
+    algorithm=None,
+    stream: utils.AnyStream | int | None = None,
+):
+    """
+    Perform the specified matrix multiplication computation
+    :math:`\\mathcal{{F}}(\\alpha a @ b + \\beta c)`, where :math:`\\mathcal{{F}}` is
+    the epilog. This function-form is a wrapper around the stateful
+    :class:`Matmul` object APIs and is meant for *single* use (the user needs to perform
+    just one matrix multiplication, for example), in which case there is no possibility of
+    amortizing preparatory costs.
+
+    Detailed information on what's happening within this function can be obtained by passing
+    in a :class:`logging.Logger` object to :class:`MatmulOptions` or by setting the
+    appropriate options in the root logger object, which is used by default:
+
+        >>> import logging
+        >>> logging.basicConfig(
+        ...     level=logging.INFO,
+        ...     format="%(asctime)s %(levelname)-8s %(message)s",
+        ...     datefmt="%m-%d %H:%M:%S",
+        ... )
+
+    A user can select the desired logging level and, in general, take advantage of all of
+    the functionality offered by the Python `logging` module.
+
+    Args:
+        a: {a}
+
+        b: {b}
+
+        c: {c}
+            {c_admonitions}
+
+        alpha: {alpha}
+
+        beta: {beta}
+
+        epilog: {epilog}
+
+        epilog_inputs: {epilog_inputs}
+
+        qualifiers: {qualifiers}
+
+        options: {options}
+
+        preferences: {preferences}
+
+        algorithm: An object of type :class:`Algorithm` objects can be directly provided to
+            bypass planning, if desired. The algorithm object must be compatible with the
+            matrix multiplication. A typical use for this option is to provide an algorithm
+            that has been serialized (numpy) from a previously planned and autotuned
+            matrix multiplication.
+
+        stream: {stream}
+
+        quantization_scales: {quantization_scales}
+
+    Returns:
+        {result}
+
+    Semantics:
+        {semantics}
+
+    Narrow-precision support:
+        {narrow_precision}
+
+    .. seealso::
+        :class:`Matmul`, :class:`MatmulOptions`, :class:`MatmulEpilog`,
+        :class:`MatmulPlanPreferences`
+
+    Examples:
+
+        >>> import cupy as cp
+        >>> import nvmath
+
+        Create three float32 ndarrays on the GPU:
+
+        >>> M, N, K = 128, 64, 256
+        >>> a = cp.random.rand(M, K, dtype=cp.float32)
+        >>> b = cp.random.rand(K, N, dtype=cp.float32)
+        >>> c = cp.random.rand(M, N, dtype=cp.float32)
+
+        Perform the operation :math:`\\alpha a @ b + \\beta c` using :func:`matmul`. The
+        result `r` is also a CuPy float32 ndarray:
+
+        >>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74)
+
+        An epilog can be used as well. Here we perform
+        :math:`\\operatorname{{RELU}}(\\alpha a @ b + \\beta c)`:
+
+        >>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
+        >>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74, epilog=epilog)
+
+        Options can be provided to customize the operation:
+
+        >>> compute_type = nvmath.linalg.advanced.MatmulComputeType.COMPUTE_32F_FAST_TF32
+        >>> o = nvmath.linalg.advanced.MatmulOptions(compute_type=compute_type)
+        >>> r = nvmath.linalg.advanced.matmul(a, b, options=o)
+
+        See `MatmulOptions` for the complete list of available options.
+
+        The package current stream is used by default, but a stream can be explicitly
+        provided to the Matmul operation. This can be done if the operands are computed on a
+        different stream, for example:
+
+        >>> s = cp.cuda.Stream()
+        >>> with s:
+        ...     a = cp.random.rand(M, K)
+        ...     b = cp.random.rand(K, N)
+        >>> r = nvmath.linalg.advanced.matmul(a, b, stream=s)
+
+        The operation above runs on stream `s` and is ordered with respect to the input
+        computation.
+
+        Create  NumPy ndarrays on the CPU.
+
+        >>> import numpy as np
+        >>> a = np.random.rand(M, K)
+        >>> b = np.random.rand(K, N)
+
+        Provide the NumPy ndarrays to :func:`matmul`, with the result also being a NumPy
+        ndarray:
+
+        >>> r = nvmath.linalg.advanced.matmul(a, b)
+
+    Notes:
+        - This function is a convenience wrapper around :class:`Matmul` and is
+          specifically meant for *single* use.
+
+    Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
+    <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
+    directory.
+    """
+
+    # Set algorithm limit to 1, but take a copy first if needed.
+    if isinstance(preferences, _configuration.MatmulPlanPreferences):
+        preferences = copy.copy(preferences)
+
+    preferences = utils.check_or_create_options(
+        _configuration.MatmulPlanPreferences, preferences, "Matrix multiplication plan preferences"
+    )
+    preferences.limit = 1
+
+    if algorithm is None:
+        algorithms = None
+    else:
+        algorithms = [algorithm]  # The type of algorithm should be algorithm.Algorithm and will be checked in plan()
+
+    with Matmul(
+        a,
+        b,
+        c=c,
+        alpha=alpha,
+        beta=beta,
+        qualifiers=qualifiers,
+        options=options,
+        stream=stream,
+        quantization_scales=quantization_scales,
+    ) as mm:
+        mm.plan(preferences=preferences, epilog=epilog, epilog_inputs=epilog_inputs, stream=stream, algorithms=algorithms)
+
+        r = mm.execute(stream=stream)
+
+    return r