mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-07-17 09:07:47 +00:00
455 lines
17 KiB
Python
455 lines
17 KiB
Python
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
|
|
# 3. Neither the name of the copyright holder nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
from __future__ import annotations
|
|
|
|
from enum import IntEnum
|
|
from functools import partial
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
import cutlass
|
|
|
|
ArrayLike = torch.Tensor | np.ndarray | jnp.ndarray
|
|
DtypeLike = cutlass.Numeric | torch.dtype | np.dtype | jnp.dtype
|
|
import cutlass
|
|
import cutlass.torch
|
|
|
|
import cutlass.operators.utils.dtype as dtype_utils
|
|
from cutlass.operators.utils.tensor import (
|
|
is_jax_array,
|
|
is_numpy_tensor,
|
|
is_torch_tensor,
|
|
)
|
|
|
|
|
|
class ReferenceCheckWarning(UserWarning):
|
|
"""Warning related to reference check numerics."""
|
|
|
|
|
|
class ClampMode(IntEnum):
|
|
NoClamp = 0 # bool(False) -> no clamp
|
|
Clamp = 1 # bool(True) -> clamp
|
|
ClampOnRestrictedDtypes = 2 # Clamp on narrow precision dtype
|
|
|
|
|
|
def get_clamp_torch_dtypes():
|
|
"""Get the dtypes that will be clamped under ClampOnRestrictedDtypes.
|
|
These dtypes are narrow precision float datatypes that when used
|
|
in kernels as output dtypes, typically use the saturation-to-finite
|
|
modifier in mma ops. Therefore, to maintain same behavior between
|
|
external numerical libraries and operators, we restrict the dtypes' behavior
|
|
to clamp the reference tensor output to its max/min finite range.
|
|
"""
|
|
import torch
|
|
|
|
clamp_torch_dtypes = {
|
|
torch.float8_e4m3fn,
|
|
torch.float8_e5m2,
|
|
}
|
|
return clamp_torch_dtypes
|
|
|
|
|
|
def get_clamp_jax_dtypes():
|
|
"""Get the dtypes that will be clamped under ClampOnRestrictedDtypes."""
|
|
import jax.numpy as jnp
|
|
|
|
clamp_jax_dtypes = {
|
|
jnp.dtype(jnp.float8_e4m3fn),
|
|
jnp.dtype(jnp.float8_e5m2),
|
|
}
|
|
return clamp_jax_dtypes
|
|
|
|
|
|
def get_clamp_cutlass_dtypes():
|
|
"""Get the cutlass dtypes that will be clamped under ClampOnRestrictedDtypes."""
|
|
clamp_cutlass_dtypes = {
|
|
cutlass.Float8E4M3FN,
|
|
cutlass.Float8E5M2,
|
|
}
|
|
return clamp_cutlass_dtypes
|
|
|
|
|
|
def get_torch_default_tolerances():
|
|
"""PyTorch per-dtype default ``(rtol, atol)`` tolerances.
|
|
|
|
Snapshotted from
|
|
https://docs.pytorch.org/docs/stable/testing.html#torch.testing.assert_close
|
|
"""
|
|
import torch
|
|
|
|
return {
|
|
torch.float16: (1e-3, 1e-5),
|
|
torch.bfloat16: (1.6e-2, 1e-5),
|
|
torch.float32: (1.3e-6, 1e-5),
|
|
torch.float64: (1e-7, 1e-7),
|
|
}
|
|
|
|
|
|
def check_as_numpy_or_jax(
|
|
out: np.ndarray | jnp.ndarray,
|
|
reference: np.ndarray | jnp.ndarray,
|
|
output_dtype: np.dtype | jnp.dtype | cutlass.Numeric,
|
|
rtol: float,
|
|
atol: float,
|
|
equal_nan: bool,
|
|
msg: str,
|
|
skip_check: bool,
|
|
clamp_mode: ClampMode,
|
|
) -> np.ndarray:
|
|
"""Helper functions for assert_close_with_reference_conversion when the
|
|
reference uses numpy or jax arrays. Returns converted ref as numpy array.
|
|
"""
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
convert_fn = np.ndarray.astype
|
|
base = np.float32
|
|
|
|
safe_emulate_fn = partial(convert_fn, dtype=base)
|
|
|
|
# convert from jax to numpy array
|
|
# (numpy and jax share dtypes if jax is imported with numpy)
|
|
reference = np.asarray(reference)
|
|
out = np.asarray(out)
|
|
|
|
if isinstance(output_dtype, type) and issubclass(output_dtype, cutlass.Numeric):
|
|
embedded_dtype = dtype_utils.jax_type_from_cutlass_type(output_dtype)
|
|
else:
|
|
# raises a TypeError if cannot convert
|
|
embedded_dtype = jnp.dtype(output_dtype)
|
|
|
|
if clamp_mode == ClampMode.ClampOnRestrictedDtypes:
|
|
clamp = embedded_dtype in get_clamp_jax_dtypes()
|
|
elif clamp_mode == ClampMode.Clamp:
|
|
clamp = True
|
|
elif clamp_mode == ClampMode.NoClamp:
|
|
clamp = False
|
|
else:
|
|
raise ValueError(f"Unsupported clamp mode: {clamp_mode}")
|
|
|
|
# Check that out tensor dtype matches expected dtype
|
|
if out.dtype != embedded_dtype:
|
|
raise ValueError(
|
|
f"Result tensor dtype {out.dtype} does not match target dtype {output_dtype} ({embedded_dtype})"
|
|
)
|
|
|
|
if clamp:
|
|
if jnp.issubdtype(embedded_dtype, jnp.floating):
|
|
embedded_info = jnp.finfo(embedded_dtype)
|
|
elif jnp.issubdtype(embedded_dtype, jnp.integer):
|
|
embedded_info = jnp.iinfo(embedded_dtype)
|
|
else:
|
|
raise ValueError(f"Unsupported dtype {embedded_dtype}")
|
|
|
|
# clamp the value of the reference tensor to the finite range of the output dtype
|
|
reference = np.clip(reference, min=embedded_info.min, max=embedded_info.max)
|
|
|
|
# Convert reference tensor if it doesn't match expected dtype
|
|
if reference.dtype != embedded_dtype:
|
|
reference = convert_fn(reference, embedded_dtype)
|
|
|
|
if not skip_check:
|
|
np.testing.assert_allclose(
|
|
safe_emulate_fn(out),
|
|
safe_emulate_fn(reference),
|
|
rtol=rtol,
|
|
atol=atol,
|
|
equal_nan=equal_nan,
|
|
err_msg=msg,
|
|
)
|
|
return reference
|
|
|
|
|
|
def check_as_torch(
|
|
out: torch.Tensor,
|
|
reference: torch.Tensor,
|
|
output_dtype: torch.dtype,
|
|
rtol: float,
|
|
atol: float,
|
|
equal_nan: bool,
|
|
msg: str,
|
|
skip_check: bool,
|
|
clamp_mode: ClampMode,
|
|
) -> torch.Tensor:
|
|
"""Helper functions for assert_close_with_reference_conversion when the
|
|
reference uses PyTorch tensors. Returns converted ref as torch tensor.
|
|
"""
|
|
import torch
|
|
|
|
convert_fn = torch.Tensor.to
|
|
base = torch.float32
|
|
safe_emulate_fn = partial(convert_fn, dtype=base)
|
|
|
|
# Normalize output dtype to torch dtype
|
|
if isinstance(output_dtype, torch.dtype):
|
|
embedded_dtype = output_dtype
|
|
elif isinstance(output_dtype, type) and issubclass(output_dtype, cutlass.Numeric):
|
|
embedded_dtype = cutlass.torch.dtype(output_dtype)
|
|
else:
|
|
raise ValueError(
|
|
f"Output dtype requires a torch.dtype or cutlass.Numeric, got {output_dtype}"
|
|
)
|
|
|
|
if clamp_mode == ClampMode.ClampOnRestrictedDtypes:
|
|
clamp = embedded_dtype in get_clamp_torch_dtypes()
|
|
elif clamp_mode == ClampMode.Clamp:
|
|
clamp = True
|
|
elif clamp_mode == ClampMode.NoClamp:
|
|
clamp = False
|
|
else:
|
|
raise ValueError(f"Unsupported clamp mode: {clamp_mode}")
|
|
|
|
# Check that out tensor dtype matches expected dtype
|
|
if out.dtype != embedded_dtype:
|
|
raise ValueError(
|
|
f"Result tensor dtype {out.dtype} does not match target dtype {output_dtype} ({embedded_dtype})"
|
|
)
|
|
|
|
if clamp:
|
|
if embedded_dtype.is_floating_point:
|
|
embedded_info = torch.finfo(embedded_dtype)
|
|
elif embedded_dtype.is_complex:
|
|
raise ValueError(f"Complex dtype {embedded_dtype} is not supported")
|
|
else:
|
|
embedded_info = torch.iinfo(embedded_dtype)
|
|
|
|
# clamp the value of the reference tensor to the finite range of the output dtype
|
|
reference = torch.clamp(reference, min=embedded_info.min, max=embedded_info.max)
|
|
|
|
# Convert reference tensor if it doesn't match expected dtype
|
|
if reference.dtype != embedded_dtype:
|
|
reference = convert_fn(reference, embedded_dtype)
|
|
|
|
# Do comparisons in float32 because testing.assert_close and testing.assert_allclose
|
|
# is not safe, doesn't support atol and rtol on lower-bitwidth datatypes
|
|
if not skip_check:
|
|
torch.testing.assert_close(
|
|
safe_emulate_fn(out),
|
|
safe_emulate_fn(reference),
|
|
rtol=rtol,
|
|
atol=atol,
|
|
equal_nan=equal_nan,
|
|
msg=msg,
|
|
)
|
|
return reference
|
|
|
|
|
|
def assert_close_with_reference_conversion(
|
|
result_tensors: ArrayLike | list[ArrayLike],
|
|
accumulator_reference_tensor: ArrayLike | list[ArrayLike],
|
|
output_dtypes: DtypeLike | list[DtypeLike],
|
|
rtols: float | list[float] = None,
|
|
atols: float | list[float] = None,
|
|
equal_nan: bool = False,
|
|
msg: str = None,
|
|
skip_check: bool = False,
|
|
clamp: ClampMode | int = ClampMode.ClampOnRestrictedDtypes,
|
|
) -> torch.Tensor | np.ndarray | list[torch.Tensor | np.ndarray]:
|
|
"""
|
|
Asserts that the result tensor is similar to the reference tensor.
|
|
This function automatically handles:
|
|
- Converting the reference tensor to the correct output dtype
|
|
- Handling overflow as saturation instead of NaN or Inf
|
|
When reference checking multiple tensors at a time, they must be passed in the same order.
|
|
|
|
Args:
|
|
result_tensors (ArrayLike | list[ArrayLike]): The output tensors to compare against the reference. Can be a torch.Tensor, np.ndarray, or jnp.ndarray
|
|
or a list of these types.
|
|
accumulator_reference_tensor (ArrayLike | list[ArrayLike]): The reference tensor in accumulator dtype. Must be the same type as the result tensor.
|
|
output_dtypes (DtypeLike | list[DtypeLike]): The dtypes of the output, can be a cutlass.Numeric, torch.dtype, np.dtype, or jnp.dtype
|
|
or a list of these types.
|
|
rtols (float | list[float]): The relative tolerance to use. Optional, if no value is set then will default to 0.0.
|
|
atols (float | list[float]): The absolute tolerance to use. Optional, if no value is set then will default to 0.0.
|
|
equal_nan (bool): If True, will compare NaN values as equal
|
|
msg (str): The message to display if the assertion fails
|
|
skip_check (bool): If True, will skip the ref check and return the reference tensors directly
|
|
clamp (ClampMode | int): Configures whether to clamp the reference tensor to the range of the output dtype
|
|
|
|
Returns:
|
|
The reference tensors as torch.Tensor or np.array, after being converted to output dtype
|
|
|
|
Raises:
|
|
ValueError: If the number of result tensors, references, and output dtypes are not the same
|
|
ValueError: If the number of rtol or atol values is >1 but not the same as the number of result tensors
|
|
ValueError: If the reference tensor and result tensor types don't match
|
|
ValueError: If the reference tensor is not a torch tensor, numpy array, or jax array
|
|
ValueError: If the clamp mode is not supported
|
|
"""
|
|
|
|
# Argument list preprocessing
|
|
accumulator_reference_tensor = (
|
|
[accumulator_reference_tensor]
|
|
if not isinstance(accumulator_reference_tensor, list)
|
|
else accumulator_reference_tensor
|
|
)
|
|
result_tensors = (
|
|
[result_tensors] if not isinstance(result_tensors, list) else result_tensors
|
|
)
|
|
output_dtypes = (
|
|
[output_dtypes] if not isinstance(output_dtypes, list) else output_dtypes
|
|
)
|
|
rtols = [rtols] if not isinstance(rtols, list) else rtols
|
|
atols = [atols] if not isinstance(atols, list) else atols
|
|
|
|
num_to_check = len(result_tensors)
|
|
dtypes_match = num_to_check == len(output_dtypes)
|
|
rtols_match = num_to_check == len(rtols)
|
|
atols_match = num_to_check == len(atols)
|
|
references_match = num_to_check == len(accumulator_reference_tensor)
|
|
|
|
if not all([dtypes_match, references_match]):
|
|
raise ValueError(
|
|
"The number of result tensors, references, and output dtypes must be the same"
|
|
)
|
|
|
|
if not rtols_match:
|
|
if len(rtols) == 1:
|
|
rtols = (
|
|
rtols * num_to_check
|
|
) # expands length of list to match number of tensors
|
|
else:
|
|
raise ValueError(
|
|
"The number of rtol values must be the same as the number of result tensors"
|
|
)
|
|
if not atols_match:
|
|
if len(atols) == 1:
|
|
atols = (
|
|
atols * num_to_check
|
|
) # expands length of list to match number of tensors
|
|
else:
|
|
raise ValueError(
|
|
"The number of atol values must be the same as the number of result tensors"
|
|
)
|
|
|
|
clamp_mode = ClampMode(int(clamp))
|
|
|
|
# When rtol/atol are not explicitly passed, they default to 0.0 (exact
|
|
# match). Tests should compute the reference at the kernel's precision
|
|
# rather than relax tolerances; the exception is hardware fast-math
|
|
# transcendentals (tanh/exp/sigmoid), which need the dtype-default
|
|
# tolerance from get_torch_default_tolerances.
|
|
#
|
|
# Call reference check for each set of reference, result, output_dtype, and rtol/atol
|
|
for i, (reference, out, output_dtype, rtol, atol) in enumerate(
|
|
zip(accumulator_reference_tensor, result_tensors, output_dtypes, rtols, atols)
|
|
):
|
|
rtol = rtol if rtol is not None else 0.0
|
|
atol = atol if atol is not None else 0.0
|
|
|
|
if is_torch_tensor(reference):
|
|
if is_numpy_tensor(out) or is_jax_array(out):
|
|
raise ValueError(
|
|
"Reference tensor is a torch tensor and result tensor is not, this is not supported"
|
|
)
|
|
accumulator_reference_tensor[i] = check_as_torch(
|
|
out,
|
|
reference,
|
|
output_dtype,
|
|
rtol,
|
|
atol,
|
|
equal_nan,
|
|
msg,
|
|
skip_check,
|
|
clamp_mode,
|
|
)
|
|
elif is_numpy_tensor(reference) or is_jax_array(reference):
|
|
if is_torch_tensor(out):
|
|
raise ValueError(
|
|
"Reference tensor is a jax array and result tensor is a torch tensor, this is not supported"
|
|
)
|
|
accumulator_reference_tensor[i] = check_as_numpy_or_jax(
|
|
out,
|
|
reference,
|
|
output_dtype,
|
|
rtol,
|
|
atol,
|
|
equal_nan,
|
|
msg,
|
|
skip_check,
|
|
clamp_mode,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported tensor type: {type(reference)}")
|
|
|
|
return accumulator_reference_tensor
|
|
|
|
|
|
def reference_scaled_mm(
|
|
A: torch.Tensor,
|
|
B: torch.Tensor,
|
|
scale_A: torch.Tensor,
|
|
scale_B: torch.Tensor,
|
|
out_dtype: torch.dtype,
|
|
transform_sf: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
|
|
):
|
|
"""
|
|
Computes a reference scaled mm operation. Currently, torch._scaled_mm does not support batch mode.
|
|
When a batch mode is present, this function iterates through each problem in the batch.
|
|
|
|
:param A: The A tensor
|
|
:type A: torch.Tensor
|
|
:param B: The B tensor
|
|
:type B: torch.Tensor
|
|
:param scale_A: The scale factor tensor for operand A
|
|
:type scale_A: torch.Tensor
|
|
:param scale_B: The scale factor tensor for operand B
|
|
:type scale_B: torch.Tensor
|
|
:param out_dtype: The output dtype
|
|
:type out_dtype: torch.dtype
|
|
:param transform_sf: A function to transform the scale factor tensors to the correct shape for the scaled mm operation
|
|
:type transform_sf: Callable
|
|
:return: The reference scaled mm operation
|
|
"""
|
|
if len(A.shape) == 2:
|
|
return torch._scaled_mm(
|
|
A, B, scale_a=scale_A, scale_b=scale_B, out_dtype=out_dtype
|
|
)
|
|
else:
|
|
# torch._scaled_mm does not support batch mode. Iterate through each problem in the batch
|
|
L, M, N = A.shape[0], A.shape[1], B.shape[2]
|
|
scale_A = scale_A.view(L, -1)
|
|
scale_B = scale_B.view(L, -1)
|
|
reference = torch.empty((L, M, N), device=A.device, dtype=out_dtype)
|
|
for l_idx in range(L):
|
|
reference[l_idx] = torch._scaled_mm(
|
|
A[l_idx],
|
|
B[l_idx],
|
|
scale_a=transform_sf(scale_A[l_idx]),
|
|
scale_b=transform_sf(scale_B[l_idx]),
|
|
out_dtype=out_dtype,
|
|
)
|
|
return reference
|