# 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

try:
    from cuda.core import Device, EventOptions
except ImportError:
    from cuda.core.experimental import Device, EventOptions
import numpy as np

from nvmath import memory
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.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)

        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 (pickled)
from a previously planned and autotuned matrix multiplication.""".replace("\n", " "),
        #
        "epilog": """\
Specify an epilog :math:`F` as an object of type :class:`MatmulEpilog` to apply to the result of the matrix
multiplication: :math:`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:: 1.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(),
        #
        "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, using_fp4_ab, 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.
        using_fp4_ab: Whether FP4 operands are being used.
        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 using_fp4_ab:
        scale_mode = cublaslt.MatmulMatrixScale.VEC16_UE4M3
        scale_group_size = 16
        scale_dtype = "float8_e4m3fn"  # UE4M3 format, consistent with input A/B scales
    else:
        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


@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}

    .. 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=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)
        self.memory_space = "cuda"
        self.device_id = utils.get_operands_device_id(operands)
        if self.device_id == "cpu":
            if self.package == "numpy":
                self.package = "cuda"
            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.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.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:
            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 attributes.
        self.workspace_ptr: None | memory.MemoryPointer = None
        self.workspace_size = 0
        self.workspace_allocated_size = 0
        self.workspace_allocated_here = False

        # Attributes to establish stream ordering.
        self.workspace_stream = None
        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.")

    def _free_workspace_memory(self, exception: Exception | None = None) -> bool:
        """
        Free workspace by releasing the MemoryPointer object.
        """
        if self.workspace_ptr is None:
            return True

        self.workspace_ptr = None
        self.workspace_allocated_size = 0
        self.logger.debug("[_free_workspace_memory] The workspace has been released.")

        return True

    def _reset_workspace_allocation_tracking(self):
        """
        Reset workspace allocation tracking attributes to False at the end of the methods
        where workspace memory is potentially allocated. This is necessary to prevent any
        exceptions raised before method entry from using stale tracking values.
        """
        self.workspace_allocated_here = False

    @utils.precondition(_check_valid_matmul)
    def _release_workspace_memory_perhaps(self, release_workspace):
        """
        Free workspace memory if it's larger than the specified limit.
        """
        if not release_workspace:
            return True

        # Establish ordering wrt the computation and free workspace if requested.
        if self.last_compute_event is not None:
            self.workspace_stream.wait(self.last_compute_event)
            self.logger.debug("Established ordering with respect to the computation before releasing the workspace.")
            self.last_compute_event = None

        self.logger.debug("[_release_workspace_memory_perhaps] The workspace memory will be released.")
        return self._free_workspace_memory()

    def _release_workspace_memory_perhaps_wrapper(self, exception: Exception | None = None) -> bool:
        """
        This is used in @atomic.
        """
        self._release_workspace_memory_perhaps(release_workspace=self.workspace_allocated_here)
        self._reset_workspace_allocation_tracking()
        return True

    @utils.precondition(_check_valid_matmul)
    @utils.precondition(_check_planned, "Workspace memory allocation")
    @utils.atomic(_free_workspace_memory, method=True)
    def _allocate_workspace_memory(self, stream_holder: utils.StreamHolder):
        """
        Allocate workspace memory using the specified allocator.
        """

        assert self.workspace_size is not None, "Internal Error."
        assert self.workspace_allocated_here is False, "Internal Error."

        if self.workspace_size == 0:  # For performance, bypass allocator for workspace size == 0.
            self.workspace_ptr = memory.MemoryPointer(0, 0, finalizer=None)
        else:
            self.logger.debug("Allocating workspace for performing the matrix multiplication...")
            with utils.device_ctx(self.device_id), stream_holder.ctx:
                try:
                    if isinstance(self.allocator, memory.BaseCUDAMemoryManagerAsync):
                        self.workspace_ptr = self.allocator.memalloc_async(self.workspace_size, stream_holder.obj)
                    else:
                        self.workspace_ptr = self.allocator.memalloc(self.workspace_size)
                    self.workspace_allocated_here = True
                except TypeError as e:
                    message = (
                        "The method 'memalloc' in the allocator object must conform to the interface in the "
                        "'BaseCUDAMemoryManager' protocol."
                    )
                    raise TypeError(message) from e

        self.workspace_allocated_size = self.workspace_size
        self.workspace_stream = stream_holder.obj
        self.logger.debug(
            f"Finished allocating device workspace of size {formatters.MemoryStr(self.workspace_size)} in the context "
            f"of stream {self.workspace_stream}."
        )

    def _allocate_workspace_memory_perhaps(self, stream_holder: utils.StreamHolder):
        """
        Allocate workspace memory using the specified allocator, if it hasn't already been
        done.
        """

        if self.workspace_ptr is not None and self.workspace_allocated_size >= self.workspace_size:
            return

        return self._allocate_workspace_memory(stream_holder)

    @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)
        """
        # Package validation: Normalize "numpy" to "cuda" to match the behavior in __init__.
        # When operands are NumPy on CPU, self.package is set to "cuda" (execution package),
        # so we must also normalize NumPy scales to "cuda" to allow the same input format.
        # This handles the NumPy <=> CuPy asymmetry where NumPy on CPU is accepted as input
        # but internally converted to CuPy for CUDA execution.
        scale_package = utils.infer_object_package(scale)
        scale_package = "cuda" if scale_package == "numpy" else scale_package
        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 =")

        # 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.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()))
                epilog_package = "cuda" if epilog_package == "numpy" else epilog_package  # Handle the NumPy <=> CuPy asymmetry.
                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.using_fp4_ab, 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)

        if self.preference_ptr is None:
            self.preference_ptr = cublaslt.matmul_preference_create()
        else:
            # We need to create a new preferences object to avoid preferences being set in a
            # cumulative manner if plan() is called multiple times.
            cublaslt.matmul_preference_destroy(self.preference_ptr)
            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))

        self.workspace_size = int(np.max(self.algorithms_buffer["workspace_size"]))
        if self.workspace_size > 0 and self.epilog:
            self.workspace_size += 16  # Workaround for library issue

        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(self.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(self.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":
            package = "cuda" if package == "numpy" else package  # Handle the NumPy <=> CuPy asymmetry.
            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_from_host(
        self,
        quantization_scales: _configuration.MatmulQuantizationScales,
        stream_holder: utils.StreamHolder[utils.AnyStream],
    ):
        """(Private method) Unchecked method to reset quantization residing on the CPU.
        This method updates quantization scales for operands A, B, C, and D.

        Args:
            quantization_scales: MatmulQuantizationScales object
                containing new scale values for operands. Can contain scalars,
                single-element tensors (for tensor-wide scaling),
                or properly shaped uint8 tensors (for block scaling).

            stream_holder: Stream holder for GPU operations.

        This method updates the stored scale values in `self.quantization_scales`
        and the device-side tensor wrappers in `self.quantization_scales_device`.
        It also updates the cuBLAS descriptor pointers via `self.mm_desc_ifc`.
        """
        # Note: when updating the aux scale, we only need to update the "data",
        # but the scale_mode does not need to be update since it is unchanged
        # from initialization. This is guaranteed by reset_operands_unchecked
        # requirement for the newly provided operands to match the
        # original operands' properties.

        qsd = self.quantization_scales_device
        mm_desc_ifc = self.mm_desc_ifc

        def _wrap_and_copy(scale):
            if isinstance(scale, (int, float)):
                return tensor_wrapper.wrap_operand(np.asarray(scale, dtype="float32")).to(self.device_id, stream_holder)
            return tensor_wrapper.wrap_operand(scale).to(self.device_id, stream_holder)

        if quantization_scales.a is not None:
            host_scale = quantization_scales.a
            self.quantization_scales.a = host_scale
            holder = qsd["a"] = _wrap_and_copy(host_scale)
            mm_desc_ifc.set_a_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales.b is not None:
            host_scale = quantization_scales.b
            self.quantization_scales.b = host_scale
            holder = qsd["b"] = _wrap_and_copy(host_scale)
            mm_desc_ifc.set_b_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales.c is not None:
            host_scale = quantization_scales.c
            self.quantization_scales.c = host_scale
            holder = qsd["c"] = _wrap_and_copy(host_scale)
            mm_desc_ifc.set_c_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales.d is not None:
            host_scale = quantization_scales.d
            self.quantization_scales.d = host_scale
            holder = qsd["d"] = _wrap_and_copy(host_scale)
            mm_desc_ifc.set_d_scale_pointer_unchecked(holder.data_ptr)

    def _reset_aux_quantization_scale_unchecked(
        self, aux_quantization_scale, stream: utils.AnyStream | int | None, stream_holder: utils.StreamHolder | None
    ):
        """(Private method) Unchecked method to reset aux_quantization_scale.

        Args:
            aux_quantization_scale: The aux quantization scale (scalar or tensor).
            stream: The stream for GPU operations (may be None).
            stream_holder: The stream holder (may be None, will be created if needed).

        Returns:
            The stream_holder (either the input or newly created one).

        Notes:
            - to minimize overhead, the stream logic is as follows:
              - if the operation does not need a stream, then `stream_holder` is not used.
              - if the operation needs to use a stream, and `stream_holder` is None,
                then a stream holder is created from `stream` and `self.package`.
              - if the operation needs to use a stream, and `stream_holder`
                is not None, then it is used directly.
        """
        # Note: when updating the aux scale, we only need to update the "data",
        # but the scale_mode does not need to be update since it is unchanged
        # from initialization. This is guaranteed by reset_operands_unchecked
        # requirement for the newly provided operands to match the
        # original operands' properties.

        if isinstance(aux_quantization_scale, (int, float)):
            # Scalar scales always need stream holder because we need to do a CPU->GPU copy.
            if stream_holder is None:
                stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)
            scale_op = tensor_wrapper.wrap_operand(np.asarray(aux_quantization_scale, dtype="float32"))
            self.quantization_scales_device["epilog_aux"] = scale_op.to(self.device_id, stream_holder)
            self.mm_desc_ifc.epilogue_aux_scale_pointer = self.quantization_scales_device["epilog_aux"].data_ptr
        else:
            # aux_quantization_scale is a tensor
            member_tensor_holder = self.quantization_scales_device["epilog_aux"]
            if self.memory_space == "cuda":
                member_tensor_holder.tensor = aux_quantization_scale
                self.mm_desc_ifc.epilogue_aux_scale_pointer = member_tensor_holder.data_ptr
            else:
                # wrap and copy to GPU
                if stream_holder is None:
                    stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)
                aux_quantization_scale_wrapped = tensor_wrapper.wrap_operand(aux_quantization_scale)
                member_tensor_holder = aux_quantization_scale_wrapped.to(self.device_id, stream_holder)
                self.mm_desc_ifc.epilogue_aux_scale_pointer = member_tensor_holder.data_ptr

        return stream_holder

    def _reset_quantization_scales_unchecked_on_device(
        self,
        quantization_scales_obj: _configuration.MatmulQuantizationScales,
        stream: utils.AnyStream | int | None,
        stream_holder: utils.StreamHolder | None,
    ):
        """(Private method) Unchecked method to reset quantization residing on GPU.
        This method directly updates tensor references without wrapping overhead.

        Args:
            quantization_scales_obj: Quantization scales object.
            stream_holder: Stream holder (may be None, will be created if needed).

        Returns:
            The stream_holder (either the input or newly created one).

        Notes:
            - to minimize overhead, the stream logic is as follows:
              - if the operation does not need a stream, then `stream_holder` is not used.
              - if the operation needs a stream, and `stream_holder` is None,
                then a stream holder is created from `stream` and `self.package`.
              - if the operation needs to use a stream, and `stream_holder` is
                not None, then it is used directly.
        """
        # Note: when updating the aux scale, we only need to update the "data",
        # but the scale_mode does not need to be update since it is unchanged
        # from initialization. This is guaranteed by reset_operands_unchecked
        # requirement for the newly provided operands to match the
        # original operands' properties.

        mm_desc_ifc = self.mm_desc_ifc

        def _update_scale(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.package)
                holder.tensor = (
                    tensor_wrapper.wrap_operand(np.asarray(scale, dtype="float32")).to(self.device_id, stream_holder).tensor
                )
            else:
                holder.tensor = scale

        if quantization_scales_obj.a is not None:
            scale = quantization_scales_obj.a
            self.quantization_scales.a = scale
            holder = self.quantization_scales_device["a"]
            _update_scale(holder, scale)
            mm_desc_ifc.set_a_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales_obj.b is not None:
            scale = quantization_scales_obj.b
            self.quantization_scales.b = scale
            holder = self.quantization_scales_device["b"]
            _update_scale(holder, scale)
            mm_desc_ifc.set_b_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales_obj.c is not None:
            scale = quantization_scales_obj.c
            self.quantization_scales.c = scale
            holder = self.quantization_scales_device["c"]
            _update_scale(holder, scale)
            mm_desc_ifc.set_c_scale_pointer_unchecked(holder.data_ptr)

        if quantization_scales_obj.d is not None:
            scale = quantization_scales_obj.d
            self.quantization_scales.d = scale
            holder = self.quantization_scales_device["d"]
            _update_scale(holder, scale)
            mm_desc_ifc.set_d_scale_pointer_unchecked(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}
        """

        # After release_operands(), the TensorHolder wrappers of the operands are
        # still alive since only .tensor was cleared. When the operands' memory space
        # is the same as the execution space (i.e. CUDA), the code below can reuse
        # them directly.
        # This is not the case for CPU memory space, where device mirrors
        # need re-allocation. Rather than duplicating all the non-trivial
        # re-initialization logic here, which would add significant overhead and,
        # therefore, undo the performance benefits of this method, we delegate
        # the first reset after a release to reset_operands(); subsequent resets
        # take the fast path below.
        if self._operands_released and self.memory_space == "cpu":
            self.reset_operands(
                a=a,
                b=b,
                c=c,
                alpha=alpha,
                beta=beta,
                quantization_scales=quantization_scales,
                epilog_inputs=epilog_inputs,
                stream=stream,
            )
            return

        # Deal with alpha, beta which are scalars and we don't need the stream holder.
        if alpha is not None:
            self.alpha[0] = alpha

        if beta is not None:
            self.beta[0] = beta

        # Stream holder optimization: Initialize to None and create lazily
        # only when needed. The `stream` parameter is None or the user-provided CUDA stream
        # The `stream_holder` is the internal wrapper
        # around the stream that we reuse across multiple operations
        # in this method to avoid redundant wrapper creation overhead.
        # Each method that needs a stream holder will:
        #   - Create one if stream_holder is None (lazy creation)
        #   - Return the stream_holder so subsequent operations can reuse it
        # This is better than creating a new one for each operation,
        # which is useful when needs to be reused across multiple operations (e.g.,
        # copying scalar quantization scales, copying CPU operands to GPU, etc.).
        # This approach is a good compromise between stream manipulation overhead
        # and code readability, since it allows us to keep the code below
        # quite separated into different methods. One alternative would be to simplify
        # the logic of stream holder creation but that would imply having a
        # single giant method, making the code more complex and hard to reason about.
        stream_holder = None

        # The "unchecked" contract means the caller guarantees that new operands
        # match the original operands' memory space. We exploit this below by
        # using the self.memory_space attribute (set in __init__) to determine
        # the code path without inspecting individual operands:
        #
        # - If self.memory_space == "cuda":
        #     All tensor operands are guaranteed to already be on GPU.
        #     We directly update references avoiding wrapping overhead
        #     and stream holder creation (unless needed for scalar scales).
        #
        # - If self.memory_space == "cpu":
        #     All tensor operands are on CPU and must be copied to GPU.
        #     We create a stream holder upfront and wrap+copy each operand
        #     to the execution device.

        # Handle aux_quantization_scale first (needs special handling because of the
        # particular set of strings that are used to identify/store this internally).
        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)

        # Convert quantization scales to object if needed
        quantization_scales_obj = quantization_scales
        if quantization_scales is not None and isinstance(quantization_scales, dict):
            quantization_scales_obj = _configuration.MatmulQuantizationScales(**quantization_scales)

        # Branch based on memory space
        if self.memory_space == "cuda":
            # Update a, b, c operands by directly updating tensor references
            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]

            # Handle epilog inputs (except aux_quantization_scale handled earlier).
            # The reset_operands_* contract guarantees that operand properties
            # (shape, strides, dtype) are unchanged, so only the data pointer needs
            # updating in the cuBLASLt descriptor while all other attributes
            # (batch stride, leading dimension, data type) remain valid from before.
            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
                        )

            # handle quantization scales (ignore returned stream_holder, not needed anymore)
            if quantization_scales_obj is not None:
                self._reset_quantization_scales_unchecked_on_device(quantization_scales_obj, stream, stream_holder)
        else:
            # Create stream holder (required for CPU->GPU copies)
            if stream_holder is None:
                stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)

            # Wrap and copy a, b, c operands to GPU
            if a is not None:
                a_wrapped = tensor_wrapper.wrap_operand(a)
                self.operands[0] = a_wrapped.to(self.device_id, stream_holder)  # type: ignore[index]
            if b is not None:
                b_wrapped = tensor_wrapper.wrap_operand(b)
                self.operands[1] = b_wrapped.to(self.device_id, stream_holder)  # type: ignore[index]
            if c is not None:
                c_wrapped = tensor_wrapper.wrap_operand(c)
                self.operands[2] = c_wrapped.to(self.device_id, stream_holder)  # type: ignore[index]
                # For inplace operations with CPU operands,
                # hold reference to the CPU operand.
                if self.inplace:
                    self.cpu_c_ref = c_wrapped

            # Handle epilog inputs (except aux_quantization_scale which was handled earlier)
            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)
                        self.epilog_operands[epilog_name] = epilog_input_wrapped.to(self.device_id, stream_holder)
                        self.epilog_input_name_to_handler[epilog_name].update(
                            self.mm_desc_ifc, self.epilog_operands[epilog_name]
                        )

            # handle quantization scales
            if quantization_scales_obj is not None:
                self._reset_quantization_scales_unchecked_from_host(quantization_scales_obj, 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. Only the operands
        explicitly passed are updated; omitted operands retain their current values.

        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.

        .. 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}

        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()

        # 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.package)

        # Update quantization_scales.
        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)

        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],
            )

        # 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}
        """

        # 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")
    @utils.atomic(_release_workspace_memory_perhaps_wrapper, method=True)
    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.package)

        # Allocate workspace if needed.
        self._allocate_workspace_memory_perhaps(stream_holder)

        # 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_ptr(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]
        raw_workspace_ptr = utils.get_ptr_from_memory_pointer(self.workspace_ptr)
        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 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,
                raw_workspace_ptr,
                self.workspace_size,
                stream_holder.ptr,
            )

        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)

        # 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 = EventOptions(enable_timing=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)

        # Establish ordering wrt the computation and free workspace if requested.
        self._release_workspace_memory_perhaps(release_workspace=release_workspace)
        self._reset_workspace_allocation_tracking()

        # 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")
    @utils.atomic(_release_workspace_memory_perhaps_wrapper, method=True)
    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: (Experimental) An algorithm chosen from the sequence returned by
                :meth:`plan` or :py:attr:`algorithms`. By default, the first algorithm in
                the sequence is used.

            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.package)
        if log_info:
            self.logger.info(f"The specified stream for execute() is {stream_holder.obj}.")

        # Allocate workspace if needed.
        self._allocate_workspace_memory_perhaps(stream_holder)

        # 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_ptr(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.using_fp4_ab, 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

        raw_workspace_ptr = utils.get_ptr_from_memory_pointer(self.workspace_ptr)
        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,
                raw_workspace_ptr,
                self.workspace_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.")

        # Establish ordering wrt the computation and free workspace if requested.
        if release_workspace:
            self._release_workspace_memory_perhaps(True)

        # 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 = {}
        self._reset_workspace_allocation_tracking()

        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.
            if self.last_compute_event is not None:
                if self.workspace_stream is not None:
                    self.workspace_stream.wait(self.last_compute_event)
                self.last_compute_event = None

            self._free_workspace_memory()

            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=None,
    preferences=None,
    algorithm=None,
    stream: utils.AnyStream | int | None = None,
):
    """
    Perform the specified matrix multiplication computation :math:`F(\\alpha a @ b + \\beta
    c)`, where :math:`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 (pickled) 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:`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
