nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/linalg/advanced/_configuration.py

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

nvmath/linalg/advanced/helpers/__init__.py

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