diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py index 8b02e181c..69e01351c 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/_blockscaled_api.py @@ -24,7 +24,7 @@ from ..backend_utils import rubin_single_group_offsets_kwarg from cuda.bindings import driver as cuda import os -from typing import Tuple, Optional +from typing import Literal, Tuple, Optional import cutlass import cutlass.cute as cute @@ -128,6 +128,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -167,6 +168,12 @@ def __init__( :param mma_tiler_mn: MMA tiler shape (M, N) :param cluster_shape_mn: Cluster shape (M, N) :param sf_vec_size: Scale factor vector size + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. :param vector_f32: Use vectorized f32 operations :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor @@ -259,6 +266,7 @@ def __init__( else: self.cluster_shape_mn = cluster_shape_mn self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override self.vector_f32 = vector_f32 self.m_aligned = m_aligned self.discrete_col_sfd = discrete_col_sfd @@ -495,6 +503,30 @@ def check_support(self) -> bool: f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) + # torch has no e5m3 dtype and TVM-FFI cannot marshal FloatNV8E5M3FNU, so e5m3 + # scale factors arrive as e4m3 storage of the same width and the Rubin kernel + # reinterprets them. That reinterpretation is the only real override; every + # other format the kernel reads straight off sfa.element_type. + + # e5m3 is the only override currently supported + self._value_error_if( + self.sf_fp8_dtype_override not in (None, "e5m3"), + f"sf_fp8_dtype_override must be None or 'e5m3', got {self.sf_fp8_dtype_override!r}", + ) + if self.sf_fp8_dtype_override == "e5m3": + # Only allow e5m3 to pretend to be e4m3fn + self._value_error_if( + self.sf_dtype != torch.float8_e4m3fn, + f"sf_fp8_dtype_override='e5m3' requires the NVFP4 recipe -- FP4 A/B with " + f"torch.float8_e4m3fn scale factors at sf_vec_size 16 -- but got " + f"ab_dtype={self.ab_dtype}, sf_dtype={self.sf_dtype}, sf_vec_size={self.sf_vec_size}", + ) + # Only allow e5m3 for rubin kernels + self._value_error_if( + not self._is_rubin_kernel, + f"sf_fp8_dtype_override='e5m3' requires Rubin (SM107), got device type {self._device_type!r}", + ) + self._check_dtype( self.acc_dtype, dtype=torch.float32, @@ -711,6 +743,11 @@ def compile(self) -> None: act_func=self.act_func, use_dynamic_sched=self.use_dynamic_sched, **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets), + # Only the Rubin kernel accepts sf_fp8_dtype_override, and check_support + # rejects "e5m3" unless _is_rubin_kernel -- the same flag that selected + # self._kernel. The kernel maps the string to FloatNV8E5M3FNU itself, so + # that internal-only type is never named outside the Rubin module. + **({"sf_fp8_dtype_override": self.sf_fp8_dtype_override} if self.sf_fp8_dtype_override == "e5m3" else {}), ) hardware_info = cutlass.utils.HardwareInfo() diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py index e5eebd985..4a4d80736 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/api.py @@ -31,7 +31,7 @@ from cuda.bindings import driver as cuda import logging import os -from typing import Any, Tuple, Optional, overload +from typing import Any, Literal, Tuple, Optional, overload import cutlass @@ -107,6 +107,7 @@ class DgluCall: mma_tiler_mn: Tuple[int, int] = (256, 256) cluster_shape_mn: Optional[Tuple[int, int]] = None sf_vec_size: int = 16 + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False m_aligned: int = 256 discrete_col_sfd: bool = False @@ -189,6 +190,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -230,6 +232,7 @@ def check_support(self) -> bool: ("sample_amax", kwargs["sample_amax"]), ("sample_norm_const", kwargs["sample_norm_const"]), ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ("sf_fp8_dtype_override", kwargs["sf_fp8_dtype_override"]), ("discrete_col_sfd", kwargs["discrete_col_sfd"] if kwargs["discrete_col_sfd"] else None), ("geglu_alpha", kwargs["geglu_alpha"] if kwargs["geglu_alpha"] != 1.702 else None), ("glu_clamp_max", kwargs["glu_clamp_max"] if kwargs["glu_clamp_max"] != 7.0 else None), @@ -468,6 +471,15 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: mma_tiler_mn: MMA tiler shape cluster_shape_mn: Cluster shape sf_vec_size: Scale factor vector size + sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors as + E5M3 instead of the encoding implied by ``sfa_tensor.dtype``. ``None`` + (default) infers as usual -- E4M3 for NVFP4, E8M0 for MXFP4/MXFP8 -- + and is the only accepted value on the BF16 backend, which has no + scale factors. ``"e5m3"`` selects an unsigned 5-exponent-bit, + 3-mantissa-bit format that trades two mantissa bits for one exponent + bit to widen the scale range; it is Rubin-only, requires the NVFP4 + recipe, and the scale tensors are still passed as + ``torch.float8_e4m3fn`` because torch has no e5m3 dtype. vector_f32: Use vectorized f32 m_aligned: M alignment (must be 256) discrete_col_sfd: Generate discrete col-major scale factor tensor @@ -525,6 +537,7 @@ def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: mma_tiler_mn = call.mma_tiler_mn cluster_shape_mn = call.cluster_shape_mn sf_vec_size = call.sf_vec_size + sf_fp8_dtype_override = call.sf_fp8_dtype_override vector_f32 = call.vector_f32 m_aligned = call.m_aligned discrete_col_sfd = call.discrete_col_sfd @@ -698,6 +711,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -737,6 +751,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -775,6 +790,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, @@ -811,6 +827,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, @@ -939,6 +956,7 @@ def _normalize_dglu_call( ("sfb_ptrs", call.sfb_ptrs), ("norm_const_tensor", call.norm_const_tensor), ("sf_vec_size", call.sf_vec_size if call.sf_vec_size != 16 else None), + ("sf_fp8_dtype_override", call.sf_fp8_dtype_override), ( "discrete_col_sfd", call.discrete_col_sfd if call.discrete_col_sfd else None, @@ -1209,6 +1227,7 @@ def grouped_gemm_dglu_wrapper_sm100( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -1255,6 +1274,7 @@ def grouped_gemm_dglu_wrapper_sm100( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, diff --git a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py index f820e5c3a..20b60e65a 100644 --- a/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py +++ b/python/cudnn/gemm/cutedsl/grouped/dglu/moe_blockscaled_grouped_gemm_dglu_rubin.py @@ -14,7 +14,7 @@ MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. """ -from typing import Type, Tuple, Union, Optional +from typing import Literal, Type, Tuple, Union, Optional from functools import partial import cuda.bindings.driver as cuda @@ -105,7 +105,17 @@ class BlockScaledMoEGroupedGemmDgluKernel: :note: Supported combinations of A/B data types, SF data typs and SF vector size: - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 - - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN/FloatNV8E5M3FNU + sf_vec_size: 16 + + :note: FloatNV8E5M3FNU scale factors are Rubin-only and reachable solely through + the FP4xFP4 atom (SM107MmaMXF4NVF4Op); the FP8 atom accepts Float8E8M0FNU only. + torch has no e5m3 dtype, so the frontend passes such scale factors as + torch.float8_e4m3fn storage and overrides the CuTe element type at compile + time -- see ``sf_fp8_dtype_override`` in ``_blockscaled_api.py``. + ``can_implement`` below does not model this: it reports E5M3 as unsupported + because its shared validator is arch-agnostic and other callers are SM100, + where E5M3 scales really are invalid. The block-scaled API validates the + combination itself and never calls ``can_implement``. :note: Supported accumulator data types: - Float32 @@ -228,6 +238,7 @@ def __init__( weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, use_dynamic_sched: bool = False, act_func: str = "dswiglu", + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, use_single_group_runtime_offsets: bool = False, ): """Initializes the configuration for a Blackwell blockscaled grouped GEMM dGLU kernel. @@ -260,6 +271,13 @@ def __init__( :type cluster_shape_mn: Tuple[int, int] :param expert_cnt: Number of experts (compile-time constant). :type expert_cnt: int + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. + :type sf_fp8_dtype_override: Optional[Literal["e5m3"]] :raises ValueError: If FIX_PAD_SIZE is not divisible by mma_tiler_mn[0]. """ @@ -289,6 +307,7 @@ def __init__( raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") self.sf_vec_size = sf_vec_size + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None self.expert_cnt = expert_cnt self.acc_dtype: Type[cutlass.Numeric] = acc_dtype self.use_2cta_instrs = use_2cta_instrs @@ -753,7 +772,14 @@ def __call__( self.b_dtype: Type[cutlass.Numeric] = a.element_type # B must match A dtype self.c_dtype: Type[cutlass.Numeric] = c.element_type self.d_dtype: Type[cutlass.Numeric] = d.element_type - self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + # Scale factors may arrive under a stand-in element type: FloatNV8E5M3FNU has + # no torch dtype and TVM-FFI cannot marshal it, so e5m3 scales are passed as + # Float8E4M3FN storage of the same width and reinterpreted here. This must + # happen before _setup_attributes(), which picks the MMA atom off sf_dtype. + if cutlass.const_expr(self.sf_dtype_override is not None): + self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override + else: + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py index fd9e48c4f..24576990d 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/_blockscaled_api.py @@ -24,7 +24,7 @@ from ..moe_utils import MoEWeightMode from cuda.bindings import driver as cuda import os -from typing import Tuple, Optional +from typing import Literal, Tuple, Optional import cutlass import cutlass.cute as cute @@ -122,6 +122,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -155,6 +156,12 @@ def __init__( :param mma_tiler_mn: MMA tiler shape (M, N) :param cluster_shape_mn: Cluster shape (M, N) :param sf_vec_size: Scale factor vector size + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. :param vector_f32: Use vectorized f32 operations :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Generate discrete col-major scale factor tensor @@ -235,6 +242,7 @@ def __init__( else: self.cluster_shape_mn = cluster_shape_mn self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override self.vector_f32 = vector_f32 self.m_aligned = m_aligned self.discrete_col_sfd = discrete_col_sfd @@ -442,6 +450,30 @@ def check_support(self) -> bool: f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) + # torch has no e5m3 dtype and TVM-FFI cannot marshal FloatNV8E5M3FNU, so e5m3 + # scale factors arrive as e4m3 storage of the same width and the Rubin kernel + # reinterprets them. That reinterpretation is the only real override; every + # other format the kernel reads straight off sfa.element_type. + + # e5m3 is the only override currently supported + self._value_error_if( + self.sf_fp8_dtype_override not in (None, "e5m3"), + f"sf_fp8_dtype_override must be None or 'e5m3', got {self.sf_fp8_dtype_override!r}", + ) + if self.sf_fp8_dtype_override == "e5m3": + # Only allow e5m3 to pretend to be e4m3fn + self._value_error_if( + self.sf_dtype != torch.float8_e4m3fn, + f"sf_fp8_dtype_override='e5m3' requires the NVFP4 recipe -- FP4 A/B with " + f"torch.float8_e4m3fn scale factors at sf_vec_size 16 -- but got " + f"ab_dtype={self.ab_dtype}, sf_dtype={self.sf_dtype}, sf_vec_size={self.sf_vec_size}", + ) + # Only allow e5m3 for rubin kernels + self._value_error_if( + not self._is_rubin_kernel, + f"sf_fp8_dtype_override='e5m3' requires Rubin (SM107), got device type {self._device_type!r}", + ) + self._check_dtype( self.acc_dtype, dtype=torch.float32, @@ -642,6 +674,11 @@ def compile(self) -> None: enable_bias=self._has_bias, use_dynamic_sched=self.use_dynamic_sched, **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets), + # Only the Rubin kernel accepts sf_fp8_dtype_override, and check_support + # rejects "e5m3" unless _is_rubin_kernel -- the same flag that selected + # self._kernel. The kernel maps the string to FloatNV8E5M3FNU itself, so + # that internal-only type is never named outside the Rubin module. + **({"sf_fp8_dtype_override": self.sf_fp8_dtype_override} if self.sf_fp8_dtype_override == "e5m3" else {}), ) hardware_info = cutlass.utils.HardwareInfo() diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/api.py b/python/cudnn/gemm/cutedsl/grouped/glu/api.py index f38ecfa66..4e0a5e653 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/api.py @@ -30,7 +30,7 @@ from cuda.bindings import driver as cuda import logging import os -from typing import Any, Tuple, Optional, overload +from typing import Any, Literal, Tuple, Optional, overload import cutlass @@ -107,6 +107,7 @@ class GluCall: mma_tiler_mn: Tuple[int, int] = (256, 256) cluster_shape_mn: Optional[Tuple[int, int]] = None sf_vec_size: int = 16 + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False m_aligned: int = 256 discrete_col_sfd: bool = False @@ -181,6 +182,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -218,6 +220,7 @@ def check_support(self) -> bool: ("sample_amax", kwargs["sample_amax"]), ("sample_norm_const", kwargs["sample_norm_const"]), ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ("sf_fp8_dtype_override", kwargs["sf_fp8_dtype_override"]), ("discrete_col_sfd", kwargs["discrete_col_sfd"] if kwargs["discrete_col_sfd"] else None), ), block_scaled_dtype_pairs=_block_scaled_dtype_pairs(), @@ -454,6 +457,15 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: mma_tiler_mn: MMA tiler shape cluster_shape_mn: Cluster shape sf_vec_size: Scale factor vector size + sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors as + E5M3 instead of the encoding implied by ``sfa_tensor.dtype``. ``None`` + (default) infers as usual -- E4M3 for NVFP4, E8M0 for MXFP4/MXFP8 -- + and is the only accepted value on the BF16 backend, which has no + scale factors. ``"e5m3"`` selects an unsigned 5-exponent-bit, + 3-mantissa-bit format that trades two mantissa bits for one exponent + bit to widen the scale range; it is Rubin-only, requires the NVFP4 + recipe, and the scale tensors are still passed as + ``torch.float8_e4m3fn`` because torch has no e5m3 dtype. vector_f32: Use vectorized f32 m_aligned: M alignment (must be 256) discrete_col_sfd: Generate discrete col-major scale factor tensor @@ -515,6 +527,7 @@ def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: mma_tiler_mn = call.mma_tiler_mn cluster_shape_mn = call.cluster_shape_mn sf_vec_size = call.sf_vec_size + sf_fp8_dtype_override = call.sf_fp8_dtype_override vector_f32 = call.vector_f32 m_aligned = call.m_aligned discrete_col_sfd = call.discrete_col_sfd @@ -672,6 +685,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -713,6 +727,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -751,6 +766,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, @@ -780,6 +796,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, @@ -905,6 +922,7 @@ def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: ("sfb_ptrs", call.sfb_ptrs), ("norm_const_tensor", call.norm_const_tensor), ("sf_vec_size", call.sf_vec_size if call.sf_vec_size != 16 else None), + ("sf_fp8_dtype_override", call.sf_fp8_dtype_override), ( "discrete_col_sfd", call.discrete_col_sfd if call.discrete_col_sfd else None, @@ -1172,6 +1190,7 @@ def grouped_gemm_glu_wrapper_sm100( use_single_group_runtime_offsets: bool = False, current_stream: Optional[cuda.CUstream] = None, generate_c: bool = False, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, ) -> TupleDict: """Dispatch grouped GEMM GLU once from an immutable normalized call.""" framework = detect_framework(a_tensor) @@ -1214,6 +1233,7 @@ def grouped_gemm_glu_wrapper_sm100( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, diff --git a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py index f3a4ae259..36ef0d61f 100644 --- a/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py +++ b/python/cudnn/gemm/cutedsl/grouped/glu/moe_blockscaled_grouped_gemm_glu_rubin.py @@ -15,7 +15,7 @@ MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. """ -from typing import Type, Tuple, Union, Optional +from typing import Literal, Type, Tuple, Union, Optional import cuda.bindings.driver as cuda @@ -84,7 +84,17 @@ class BlockScaledMoEGroupedGemmGluKernel: :note: Supported combinations of A/B data types, SF data typs and SF vector size: - MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32 - MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32 - - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16 + - NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN/FloatNV8E5M3FNU + sf_vec_size: 16 + + :note: FloatNV8E5M3FNU scale factors are Rubin-only and reachable solely through + the FP4xFP4 atom (SM107MmaMXF4NVF4Op); the FP8 atom accepts Float8E8M0FNU only. + torch has no e5m3 dtype, so the frontend passes such scale factors as + torch.float8_e4m3fn storage and overrides the CuTe element type at compile + time -- see ``sf_fp8_dtype_override`` in ``_blockscaled_api.py``. + ``can_implement`` below does not model this: it reports E5M3 as unsupported + because its shared validator is arch-agnostic and other callers are SM100, + where E5M3 scales really are invalid. The block-scaled API validates the + combination itself and never calls ``can_implement``. :note: Supported accumulator data types: - Float32 @@ -212,6 +222,7 @@ def __init__( act_func: str = "swiglu", enable_bias: bool = False, generate_c: bool = True, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, use_single_group_runtime_offsets: bool = False, ): """Initializes the configuration for a Blackwell blockscaled grouped GEMM GLU kernel. @@ -244,6 +255,13 @@ def __init__( :type cluster_shape_mn: Tuple[int, int] :param expert_cnt: Number of experts (compile-time constant). :type expert_cnt: int + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. + :type sf_fp8_dtype_override: Optional[Literal["e5m3"]] :raises ValueError: If FIX_PAD_SIZE is not divisible by mma_tiler_mn[0]. """ @@ -273,6 +291,7 @@ def __init__( self.sf_vec_size = sf_vec_size self.expert_cnt = expert_cnt + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None self.use_single_group_runtime_offsets = use_single_group_runtime_offsets self.acc_dtype: Type[cutlass.Numeric] = acc_dtype self.use_2cta_instrs = use_2cta_instrs @@ -735,7 +754,10 @@ def __call__( self.b_dtype: Type[cutlass.Numeric] = a.element_type self.c_dtype: Type[cutlass.Numeric] = c.element_type self.d_dtype: Type[cutlass.Numeric] = d.element_type - self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + if cutlass.const_expr(self.sf_dtype_override is not None): + self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override + else: + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(c) diff --git a/python/cudnn/gemm/cutedsl/grouped/quant/api.py b/python/cudnn/gemm/cutedsl/grouped/quant/api.py index f229770c0..02e3cce12 100644 --- a/python/cudnn/gemm/cutedsl/grouped/quant/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/quant/api.py @@ -11,7 +11,7 @@ from __future__ import annotations import os -from typing import Optional, Tuple +from typing import Literal, Optional, Tuple import cutlass import cutlass.cute as cute @@ -96,6 +96,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -130,6 +131,12 @@ def __init__( :param mma_tiler_mn: MMA tiler shape (M, N) :param cluster_shape_mn: Cluster shape (M, N) :param sf_vec_size: Scale factor vector size + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. :param vector_f32: Use vectorized f32 operations :param m_aligned: Alignment for group M dimension :param discrete_col_sfd: Enable discrete col-major scale factor tensor @@ -217,6 +224,7 @@ def __init__( else: self.cluster_shape_mn = cluster_shape_mn self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override self.vector_f32 = vector_f32 self.m_aligned = m_aligned self.discrete_col_sfd = discrete_col_sfd @@ -417,6 +425,30 @@ def check_support(self) -> bool: f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", ) + # torch has no e5m3 dtype and TVM-FFI cannot marshal FloatNV8E5M3FNU, so e5m3 + # scale factors arrive as e4m3 storage of the same width and the Rubin kernel + # reinterprets them. That reinterpretation is the only real override; every + # other format the kernel reads straight off sfa.element_type. + + # e5m3 is the only override currently supported + self._value_error_if( + self.sf_fp8_dtype_override not in (None, "e5m3"), + f"sf_fp8_dtype_override must be None or 'e5m3', got {self.sf_fp8_dtype_override!r}", + ) + if self.sf_fp8_dtype_override == "e5m3": + # Only allow e5m3 to pretend to be e4m3fn + self._value_error_if( + self.sf_dtype != cutlass.Float8E4M3FN, + f"sf_fp8_dtype_override='e5m3' requires the NVFP4 recipe -- FP4 A/B with " + f"torch.float8_e4m3fn scale factors at sf_vec_size 16 -- but got " + f"ab_dtype={self.ab_dtype}, sf_dtype={self.sf_dtype}, sf_vec_size={self.sf_vec_size}", + ) + # Only allow e5m3 for rubin kernels + self._value_error_if( + not self._is_rubin_kernel, + f"sf_fp8_dtype_override='e5m3' requires Rubin (SM107), got device type {self._device_type!r}", + ) + self._check_dtype( self.acc_dtype, dtype=cutlass.Float32, @@ -584,6 +616,11 @@ def compile(self) -> None: weight_mode=self.weight_mode, use_dynamic_sched=self.use_dynamic_sched, **rubin_single_group_offsets_kwarg(self._is_rubin_kernel, self.use_single_group_runtime_offsets), + # Only the Rubin kernel accepts sf_fp8_dtype_override, and check_support + # rejects "e5m3" unless _is_rubin_kernel -- the same flag that selected + # self._kernel. The kernel maps the string to FloatNV8E5M3FNU itself, so + # that internal-only type is never named outside the Rubin module. + **({"sf_fp8_dtype_override": self.sf_fp8_dtype_override} if self.sf_fp8_dtype_override == "e5m3" else {}), ) if self._is_rubin_kernel: # The Rubin quant kernel supports optional C materialization, but @@ -1270,6 +1307,7 @@ def grouped_gemm_quant_wrapper_sm100( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, vector_f32: bool = False, m_aligned: int = 256, discrete_col_sfd: bool = False, @@ -1315,6 +1353,14 @@ def grouped_gemm_quant_wrapper_sm100( mma_tiler_mn: MMA tiler shape cluster_shape_mn: Cluster shape sf_vec_size: Scale factor vector size + sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors as + E5M3 instead of the encoding implied by ``sfa_tensor.dtype``. ``None`` + (default) infers as usual -- E4M3 for NVFP4, E8M0 for MXFP4/MXFP8. + ``"e5m3"`` selects an unsigned 5-exponent-bit, 3-mantissa-bit format + that trades two mantissa bits for one exponent bit to widen the scale + range; it is Rubin-only, requires the NVFP4 recipe, and the scale + tensors are still passed as ``torch.float8_e4m3fn`` because torch has + no e5m3 dtype. vector_f32: Use vectorized f32 m_aligned: M alignment (must be 256) discrete_col_sfd: Enable discrete col-major scale factor tensor @@ -1543,6 +1589,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -1581,6 +1628,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn, cluster_shape_mn, sf_vec_size, + sf_fp8_dtype_override, vector_f32, m_aligned, discrete_col_sfd, @@ -1616,6 +1664,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, @@ -1644,6 +1693,7 @@ def dynamic_m_tensor_signature( mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, m_aligned=m_aligned, discrete_col_sfd=discrete_col_sfd, diff --git a/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py b/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py index 3df331bb8..10862358f 100644 --- a/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py +++ b/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py @@ -15,7 +15,7 @@ MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. """ -from typing import Type, Tuple, Union, Optional +from typing import Literal, Type, Tuple, Union, Optional from enum import Enum import cuda.bindings.driver as cuda @@ -85,6 +85,12 @@ class BlockScaledMoEGroupedGemmQuantKernel: :param generate_c: Generate C output tensor. :param enable_bias: Fuse bias addition. :param expert_cnt: Number of experts. + :param sf_fp8_dtype_override: Reinterpret the FP8-format block scale factors + as E5M3 instead of the E4M3 implied by their storage dtype. ``None`` + (default) leaves the format inferred, as every caller did before this + knob existed. ``"e5m3"`` requires Rubin and the NVFP4 recipe, and the + scale tensors are still supplied as ``torch.float8_e4m3fn`` because + torch has no e5m3 dtype -- only the CuTe element type is overridden. :param weight_mode: ``MoEWeightMode.DENSE`` or ``MoEWeightMode.DISCRETE``. :param use_dynamic_sched: Enable dynamic tile scheduling. """ @@ -185,6 +191,7 @@ def __init__( weight_mode: MoEWeightMode = MoEWeightMode.DENSE, use_dynamic_sched: bool = False, epilogue_type: int = EpilogueType.NONE.value, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, use_single_group_runtime_offsets: bool = False, ): # Hardware MMA instruction M: 2CTA → 256, 1CTA → 128 @@ -216,6 +223,7 @@ def __init__( raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") self.sf_vec_size = sf_vec_size + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None self.expert_cnt = expert_cnt self.use_single_group_runtime_offsets = use_single_group_runtime_offsets self.acc_dtype: Type[cutlass.Numeric] = acc_dtype @@ -727,7 +735,14 @@ def __call__( self.b_dtype: Type[cutlass.Numeric] = a.element_type self.c_dtype: Type[cutlass.Numeric] = c.element_type self.d_dtype: Type[cutlass.Numeric] = d.element_type - self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + # Scale factors may arrive under a stand-in element type: FloatNV8E5M3FNU has + # no torch dtype and TVM-FFI cannot marshal it, so e5m3 scales are passed as + # Float8E4M3FN storage of the same width and reinterpreted here. This must + # happen before _setup_attributes(), which picks the MMA atom off sf_dtype. + if cutlass.const_expr(self.sf_dtype_override is not None): + self.sf_dtype: Type[cutlass.Numeric] = self.sf_dtype_override + else: + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(c) self.d_layout = utils.LayoutEnum.from_tensor(d) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py index e8008274b..121976e5d 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_bf16_api.py @@ -7,7 +7,7 @@ import os import weakref -from typing import Optional, Tuple, Union +from typing import Literal, Optional, Tuple, Union import cutlass import cutlass.cute as cute @@ -64,6 +64,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, accumulate_on_output: bool = False, input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, ) -> None: @@ -106,6 +107,7 @@ def __init__( self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) self.accumulate_on_output = accumulate_on_output self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override self._scale_controls = ( sample_sfa, sample_sfb, @@ -264,6 +266,8 @@ def check_support(self) -> bool: raise ValueError("BF16 wgrad forbids scale and global-scale tensors") if self.sf_vec_size != 16: raise ValueError(f"BF16 wgrad requires sf_vec_size=16, got {self.sf_vec_size}") + if self.sf_fp8_dtype_override is not None: + raise ValueError(f"BF16 wgrad forbids sf_fp8_dtype_override, got {self.sf_fp8_dtype_override!r}") if self.offsets_desc.shape != (self.expert_cnt,): raise ValueError(f"sample_offsets must have shape {(self.expert_cnt,)}, got {self.offsets_desc.shape}") if self.offsets_desc.stride != (1,): diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py index 1c0094169..782fc3820 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/_blockscaled_api.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import Optional, Tuple, Union +from typing import Literal, Optional, Tuple, Union import cutlass import cutlass.cute as cute @@ -63,6 +63,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, accumulate_on_output: bool = False, input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, ): @@ -104,6 +105,7 @@ def __init__( self.global_scale_a_desc = self._make_tensor_desc(sample_global_scale_a, name="sample_global_scale_a") self.global_scale_b_desc = self._make_tensor_desc(sample_global_scale_b, name="sample_global_scale_b") self.sf_vec_size = sf_vec_size + self.sf_fp8_dtype_override = sf_fp8_dtype_override tokens_sum_a = self.a_desc.shape[1] tokens_sum_b = self.b_desc.shape[0] self._value_error_if( @@ -224,6 +226,30 @@ def check_support(self) -> bool: "sample_sfb", extra_error_msg="sample_sfb must have dtype float8_e8m0fnu or float8_e4m3fn", ) + # torch has no e5m3 dtype and TVM-FFI cannot marshal FloatNV8E5M3FNU, so e5m3 + # scale factors arrive as e4m3 storage of the same width and the Rubin kernel + # reinterprets them. That reinterpretation is the only real override; every + # other format the kernel reads straight off sfa.element_type. + + # e5m3 is the only override currently supported + self._value_error_if( + self.sf_fp8_dtype_override not in (None, "e5m3"), + f"sf_fp8_dtype_override must be None or 'e5m3', got {self.sf_fp8_dtype_override!r}", + ) + if self.sf_fp8_dtype_override == "e5m3": + # Only allow e5m3 to pretend to be e4m3fn + self._value_error_if( + self.sfa_desc.dtype != torch.float8_e4m3fn, + f"sf_fp8_dtype_override='e5m3' requires the NVFP4 recipe -- FP4 A/B with " + f"torch.float8_e4m3fn scale factors at sf_vec_size 16 -- but got " + f"ab_dtype={self.a_desc.dtype}, sf_dtype={self.sfa_desc.dtype}, sf_vec_size={self.sf_vec_size}", + ) + # Only allow e5m3 for rubin kernels + self._value_error_if( + not self._is_rubin_kernel, + f"sf_fp8_dtype_override='e5m3' requires Rubin (SM107), got device type {self._device_type!r}", + ) + self._check_rubin_quantization_support() self._check_dtype(self.offsets_desc, torch.int32, "sample_offsets", extra_error_msg="sample_offsets must be int32") self._check_dtype( @@ -304,6 +330,7 @@ def compile(self) -> None: expert_cnt=self.expert_cnt, weight_mode=self.weight_mode, input_order=self.input_order, + sf_fp8_dtype_override=self.sf_fp8_dtype_override, ) hardware_info = cutlass.utils.HardwareInfo() diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py index 6ddf55588..948d7bc39 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/api.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import Any, Optional, Tuple, overload +from typing import Any, Literal, Optional, Tuple, overload import os from cuda.bindings import driver as cuda @@ -109,6 +109,7 @@ def __init__( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, accumulate_on_output: bool = False, input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, ) -> None: @@ -133,6 +134,7 @@ def check_support(self) -> bool: ("sample_global_scale_a", kwargs["sample_global_scale_a"]), ("sample_global_scale_b", kwargs["sample_global_scale_b"]), ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ("sf_fp8_dtype_override", kwargs["sf_fp8_dtype_override"]), ), block_scaled_dtype_pairs=_block_scaled_dtype_pairs(), ) @@ -248,6 +250,7 @@ def grouped_gemm_wgrad_wrapper_sm100( mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, sf_vec_size: int = 16, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, accumulate_on_output: bool = False, input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, current_stream: Optional[cuda.CUstream] = None, @@ -285,6 +288,7 @@ def grouped_gemm_wgrad_wrapper_sm100( ("global_scale_a", global_scale_a), ("global_scale_b", global_scale_b), ("sf_vec_size", sf_vec_size if sf_vec_size != 16 else None), + ("sf_fp8_dtype_override", sf_fp8_dtype_override), ), block_scaled_dtype_pairs=_block_scaled_dtype_pairs(), ) @@ -324,6 +328,7 @@ def grouped_gemm_wgrad_wrapper_sm100( tuple(mma_tiler_mn), tuple(cluster_shape_mn) if cluster_shape_mn is not None else None, sf_vec_size, + sf_fp8_dtype_override, accumulate_on_output, input_order, int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")), @@ -355,6 +360,7 @@ def _sample_wgrad_expert(): mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, + sf_fp8_dtype_override=sf_fp8_dtype_override, accumulate_on_output=accumulate_on_output, input_order=input_order, ) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py index 588f6342e..523b63640 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad.py @@ -23,7 +23,7 @@ """ from importlib.metadata import PackageNotFoundError, version -from typing import Type, Tuple, Optional +from typing import Literal, Type, Tuple, Optional import cuda.bindings.driver as cuda @@ -93,8 +93,10 @@ def __init__( expert_cnt: int = 1, weight_mode: MoEWeightMode = MoEWeightMode.DENSE, input_order: WGradInputOrder = WGradInputOrder.Tensor2D, + sf_fp8_dtype_override: Optional[Literal["e5m3"]] = None, ): self.sf_vec_size = sf_vec_size + self.sf_dtype_override: Optional[Type[cutlass.Numeric]] = cutlass.FloatNV8E5M3FNU if sf_fp8_dtype_override == "e5m3" else None self.expert_cnt = expert_cnt self.acc_dtype = acc_dtype self.use_2cta_instrs = use_2cta_instrs @@ -415,7 +417,14 @@ def __call__( self.a_dtype = a_gemm.element_type self.b_dtype = b_gemm.element_type self.c_dtype = c_gemm.element_type - self.sf_dtype = sfa_gemm.element_type + # Scale factors may arrive under a stand-in element type: FloatNV8E5M3FNU has + # no torch dtype and TVM-FFI cannot marshal it, so e5m3 scales are passed as + # Float8E4M3FN storage of the same width and reinterpreted here. This must + # happen before _setup_attributes(), which picks the MMA atom off sf_dtype. + if cutlass.const_expr(self.sf_dtype_override is not None): + self.sf_dtype = self.sf_dtype_override + else: + self.sf_dtype = sfa_gemm.element_type self.a_major_mode = utils.LayoutEnum.from_tensor(a_gemm).mma_major_mode() self.b_major_mode = utils.LayoutEnum.from_tensor(b_gemm).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(c_gemm) diff --git a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py index 40a638081..fd5d42a41 100644 --- a/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py +++ b/python/cudnn/gemm/cutedsl/grouped/wgrad/moe_blockscaled_grouped_gemm_wgrad_rubin.py @@ -98,7 +98,20 @@ def _setup_attributes(self) -> None: valid_quantization = ( self.a_dtype is cutlass.Float4E2M1FN and self.b_dtype is cutlass.Float4E2M1FN - and ((self.sf_dtype is cutlass.Float8E4M3FN and self.sf_vec_size == 16) or (self.sf_dtype is cutlass.Float8E8M0FNU and self.sf_vec_size == 32)) + and ( + ( + self.sf_dtype is cutlass.Float8E4M3FN + and self.sf_vec_size == 16 + ) + or ( + self.sf_dtype is cutlass.Float8E8M0FNU + and self.sf_vec_size == 32 + ) + or ( + self.sf_dtype is cutlass.FloatNV8E5M3FNU + and self.sf_vec_size == 16 + ) + ) ) or ( self.a_dtype in (cutlass.Float8E4M3FN, cutlass.Float8E5M2) and self.b_dtype is self.a_dtype @@ -106,7 +119,10 @@ def _setup_attributes(self) -> None: and self.sf_vec_size == 32 ) if not valid_quantization: - raise ValueError("Rubin wgrad supports NVFP4, MXFP4, MXFP8-E4M3, " "or MXFP8-E5M2 block scaling.") + raise ValueError( + "Rubin wgrad supports NVFP4 (E4M3 or E5M3 scales), MXFP4, " + "MXFP8-E4M3, or MXFP8-E5M2 block scaling." + ) if self.acc_dtype is not cutlass.Float32: raise ValueError("Rubin wgrad requires Float32 accumulators.") if self.a_dtype.width == 4 and (self.a_major_mode != OperandMajorMode.K or self.b_major_mode != OperandMajorMode.K): @@ -125,10 +141,21 @@ def _setup_attributes(self) -> None: self.mma_tiler = (*self.mma_inst_shape_mn, mma_tiler_k) self.mma_tiler_sfb = (*self.mma_inst_shape_mn_sfb, mma_tiler_k) - use_sf_window = self.sf_vec_size == 16 and self.sf_dtype is cutlass.Float8E4M3FN and self.mma_tiler[1] == 256 and self.mma_tiler[2] == 512 - self.sf_window_k = self.instruction_k * 2 if use_sf_window else self.mma_tiler[2] - self.num_mma_instructions_per_sf_window = self.sf_window_k // self.instruction_k - self.num_sf_windows_per_ab_stage = self.mma_tiler[2] // self.sf_window_k + use_sf_window = ( + self.sf_vec_size == 16 + and self.sf_dtype in (cutlass.Float8E4M3FN, cutlass.FloatNV8E5M3FNU) + and self.mma_tiler[1] == 256 + and self.mma_tiler[2] == 512 + ) + self.sf_window_k = ( + self.instruction_k * 2 if use_sf_window else self.mma_tiler[2] + ) + self.num_mma_instructions_per_sf_window = ( + self.sf_window_k // self.instruction_k + ) + self.num_sf_windows_per_ab_stage = ( + self.mma_tiler[2] // self.sf_window_k + ) tiled_mma = self._create_tiled_mma() tiled_mma_sfb = self._create_tiled_mma_sfb() diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index d2c45ccc2..770e64af7 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -18,7 +18,11 @@ grouped_gemm_swiglu_init, allocate_grouped_gemm_input_tensors as allocate_grouped_gemm_input_tensors_base, ) +from fe_api.test_fe_api_utils import reencode_sf_tensor_as_ue5m3 +from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import _skip_unless_e5m3_supported from fe_api.grouped_gemm.test_grouped_gemm_dswiglu_utils import ( + GROUPED_GEMM_DSWIGLU_COMMON_MARKS, + GROUPED_GEMM_DSWIGLU_FP4_TYPE_MARKS, with_grouped_gemm_dswiglu_params_fp4, with_grouped_gemm_dswiglu_params_fp8, with_grouped_gemm_dswiglu_params_dbias_fp4, @@ -213,7 +217,7 @@ def test_grouped_gemm_dglu_class_bf16_rejects_single_group_runtime_offsets(): @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_scheduler_modes -@with_grouped_gemm_dswiglu_params_fp4 +@with_grouped_gemm_dswiglu_params_fp4(with_e5m3=True) def test_grouped_gemm_dglu_dense_compile_execute_fp4( ab_dtype, c_dtype, @@ -229,6 +233,7 @@ def test_grouped_gemm_dglu_dense_compile_execute_fp4( discrete_col_sfd, use_dynamic_sched, request, + sf_fp8_dtype_override, ): _test_grouped_gemm_dglu_dense_compile_execute( ab_dtype=ab_dtype, @@ -241,6 +246,7 @@ def test_grouped_gemm_dglu_dense_compile_execute_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, @@ -294,7 +300,7 @@ def test_grouped_gemm_dglu_dense_compile_execute_fp8( @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_scheduler_modes -@with_grouped_gemm_dswiglu_params_fp4 +@with_grouped_gemm_dswiglu_params_fp4(with_e5m3=True) def test_grouped_gemm_dglu_dense_wrapper_fp4( ab_dtype, c_dtype, @@ -310,6 +316,7 @@ def test_grouped_gemm_dglu_dense_wrapper_fp4( discrete_col_sfd, use_dynamic_sched, request, + sf_fp8_dtype_override, ): _test_grouped_gemm_dglu_dense_wrapper( ab_dtype=ab_dtype, @@ -322,6 +329,7 @@ def test_grouped_gemm_dglu_dense_wrapper_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, @@ -592,6 +600,7 @@ def _test_grouped_gemm_dglu_dense_compile_execute( use_dynamic_sched=False, omit_prob=False, use_single_group_runtime_offsets=False, + sf_fp8_dtype_override=None, ): try: from cudnn import GroupedGemmDgluSm100 @@ -648,6 +657,12 @@ def _test_grouped_gemm_dglu_dense_compile_execute( input_mutator(inputs, cfg) # Use the new unified dGLU API in dense mode + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + api = GroupedGemmDgluSm100( sample_a=inputs["a_tensor"], sample_c=inputs["c_tensor"], @@ -673,6 +688,7 @@ def _test_grouped_gemm_dglu_dense_compile_execute( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -738,6 +754,7 @@ def _test_grouped_gemm_dglu_dense_wrapper( use_dynamic_sched=False, omit_prob=False, use_single_group_runtime_offsets=False, + sf_fp8_dtype_override=None, ): try: from cudnn import grouped_gemm_dglu_wrapper_sm100 @@ -793,6 +810,12 @@ def _test_grouped_gemm_dglu_dense_wrapper( if input_mutator is not None: input_mutator(inputs, cfg) + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + try: for _ in range(2): # Run twice to test caching path if not omit_prob: @@ -818,6 +841,7 @@ def _test_grouped_gemm_dglu_dense_wrapper( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -2140,3 +2164,150 @@ def invalidate_runtime_offset(inputs, _cfg): input_mutator=invalidate_runtime_offset, use_single_group_runtime_offsets=True, ) + + +def _dglu_nvfp4_inputs(request, sf_vec_size=16, sf_dtype=torch.float8_e4m3fn, ab_dtype=torch.float4_e2m1fn_x2): + cfg = grouped_gemm_swiglu_init( + request=request, + ab_dtype=ab_dtype, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, + vector_f32=False, + discrete_col_sfd=False, + b_major="k", + ) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + b_major=cfg["b_major"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + ) + inputs, outputs = allocate_grouped_gemm_dswiglu_tensors( + tensor_m=inputs["tensor_m"], + n=cfg["n"], + l=cfg["l"], + ab_dtype=cfg["ab_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + generate_dbias=False, + input_tensors=inputs, + ) + return cfg, inputs, outputs + + +def _run_dglu_wrapper(cfg, inputs, outputs, sf_fp8_dtype_override): + """Call the wrapper directly; the harnesses turn ValueError into a skip.""" + outputs["dprob_tensor"].zero_() + return cudnn.grouped_gemm_dglu_wrapper_sm100( + a_tensor=inputs["a_tensor"], + c_tensor=inputs["c_tensor"], + sfa_tensor=inputs["sfa_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + beta_tensor=inputs["beta_tensor"], + prob_tensor=inputs["prob_tensor"], + dprob_tensor=outputs["dprob_tensor"], + generate_dbias=False, + b_tensor=inputs["b_tensor"], + sfb_tensor=inputs["sfb_tensor"], + norm_const_tensor=inputs.get("norm_const_tensor"), + acc_dtype=cfg["acc_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + discrete_col_sfd=cfg["discrete_col_sfd"], + act_func="dswiglu", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + "sf_fp8_dtype_override,overrides,expected", + [ + pytest.param("e5m3", dict(sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), "requires the NVFP4 recipe", id="mxfp4_e8m0_carrier"), + pytest.param( + "e5m3", + dict(ab_dtype=torch.float8_e4m3fn, sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), + "requires the NVFP4 recipe", + id="fp8_ab", + ), + pytest.param("e4m3", {}, "sf_fp8_dtype_override must be", id="e4m3_is_not_an_override"), + pytest.param("e5m2", {}, "sf_fp8_dtype_override must be", id="unknown_format"), + ], +) +def test_grouped_gemm_dglu_rejects_unsupported_sf_fp8_dtype(request, sf_fp8_dtype_override, overrides, expected): + """e5m3 is only reachable through the Rubin FP4xFP4 atom with e4m3-carried scales.""" + if sf_fp8_dtype_override == "e5m3": + _skip_unless_e5m3_supported() + cfg, inputs, outputs = _dglu_nvfp4_inputs(request, **overrides) + with pytest.raises(ValueError, match=expected): + _run_dglu_wrapper(cfg, inputs, outputs, sf_fp8_dtype_override) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_dglu_e5m3_is_not_cached_as_e4m3(request): + """sf_fp8_dtype_override must take part in the compile cache key. + + Identical scale-factor bytes decode differently under E4M3 and UE5M3, so if + the override were missing from the key the second call would reuse the first + kernel and silently return E4M3 results. + """ + _skip_unless_e5m3_supported() + cfg, inputs, outputs = _dglu_nvfp4_inputs(request) + d_e4m3 = _run_dglu_wrapper(cfg, inputs, outputs, None)["d_row_tensor"].float().clone() + d_e5m3 = _run_dglu_wrapper(cfg, inputs, outputs, "e5m3")["d_row_tensor"].float().clone() + torch.cuda.synchronize() + assert not torch.equal( + d_e4m3, d_e5m3 + ), "e5m3 and e4m3 produced identical output from identical scale-factor bytes; sf_fp8_dtype_override is likely missing from the compile cache key" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_dglu_bf16_rejects_sf_fp8_dtype_override(): + """The BF16 backend has no scale factors, so any explicit override is an error. + + The None case matters as much as the rejection: it pins down that merely + adding the parameter did not break BF16 dispatch. + """ + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("Requires SM100+ for grouped GEMM dGLU BF16 kernel.") + problem = make_grouped_gemm_dglu_bf16_problem(discrete=False, b_major="k") + kwargs = dict( + a_tensor=problem["a"], + c_tensor=problem["c"], + sfa_tensor=None, + padded_offsets=problem["offsets"], + alpha_tensor=problem["alpha"], + beta_tensor=problem["beta"], + prob_tensor=problem["prob"], + dprob_tensor=problem["dprob"], + d_dtype=torch.bfloat16, + b_tensor=problem["b"], + sfb_tensor=None, + ) + # None is accepted and dispatches to BF16 as usual. + cudnn.grouped_gemm_dglu_wrapper_sm100(**kwargs, sf_fp8_dtype_override=None) + with pytest.raises(ValueError, match="BF16 forbids scale control sf_fp8_dtype_override"): + cudnn.grouped_gemm_dglu_wrapper_sm100(**kwargs, sf_fp8_dtype_override="e5m3") diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py index 7daae6403..f6630b7af 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dswiglu_utils.py @@ -7,6 +7,8 @@ """ import torch +import functools + import pytest from typing import Optional, Tuple, List, Dict, Any from fe_api.test_fe_api_utils import ( @@ -97,23 +99,40 @@ ] ) -GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4 = ( +GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4_BASE = ( GROUPED_GEMM_DSWIGLU_FP4_TYPE_MARKS + GROUPED_GEMM_DSWIGLU_COMMON_MARKS + [ - pytest.mark.parametrize( - "sf_vec_size,sf_dtype", - [ - (16, torch.float8_e8m0fnu), - (16, torch.float8_e4m3fn), - (32, torch.float8_e8m0fnu), - (32, torch.float8_e4m3fn), - ], - ), pytest.mark.parametrize("discrete_col_sfd", [False]), ] ) +GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4_WITH_E5M3 = GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4_BASE + [ + pytest.mark.parametrize( + "sf_vec_size,sf_dtype,sf_fp8_dtype_override", + [ + (16, torch.float8_e8m0fnu, None), + (16, torch.float8_e4m3fn, None), + (32, torch.float8_e8m0fnu, None), + (32, torch.float8_e4m3fn, None), + (16, torch.float8_e4m3fn, "e5m3"), + ], + ids=["v16_e8m0", "v16_e4m3", "v32_e8m0", "v32_e4m3", "v16_e5m3"], + ), +] + +GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4 = GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4_BASE + [ + pytest.mark.parametrize( + "sf_vec_size,sf_dtype", + [ + (16, torch.float8_e8m0fnu), + (16, torch.float8_e4m3fn), + (32, torch.float8_e8m0fnu), + (32, torch.float8_e4m3fn), + ], + ), +] + GROUPED_GEMM_DSWIGLU_PARAM_MARKS_DBIAS_FP8 = ( GROUPED_GEMM_DSWIGLU_FP8_TYPE_MARKS + GROUPED_GEMM_DSWIGLU_COMMON_MARKS @@ -140,9 +159,12 @@ ) -def with_grouped_gemm_dswiglu_params_fp4(func): +def with_grouped_gemm_dswiglu_params_fp4(func=None, *, with_e5m3: bool = False): """Decorator to apply grouped GEMM dSwiGLU FP4 test parameters.""" - for mark in reversed(GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4): + if func is None: + return functools.partial(with_grouped_gemm_dswiglu_params_fp4, with_e5m3=with_e5m3) + param_marks = GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4_WITH_E5M3 if with_e5m3 else GROUPED_GEMM_DSWIGLU_PARAM_MARKS_FP4 + for mark in reversed(param_marks): func = mark(func) return func diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py index 5285d70b9..9d0a9e73a 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py @@ -12,8 +12,11 @@ import pytest import cudnn from test_utils import torch_fork_set_rng -from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES +from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES, reencode_sf_tensor_as_ue5m3 +from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import _skip_unless_e5m3_supported from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import ( + GROUPED_GEMM_SWIGLU_COMMON_MARKS, + GROUPED_GEMM_SWIGLU_FP4_TYPE_MARKS, grouped_gemm_swiglu_init, with_grouped_gemm_swiglu_params_fp4, with_grouped_gemm_swiglu_params_fp8, @@ -210,8 +213,9 @@ def test_grouped_gemm_glu_dense_compile_execute_fp8( @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_scheduler_modes -@with_grouped_gemm_swiglu_params_fp4 +@with_grouped_gemm_swiglu_params_fp4(with_e5m3=True) def test_grouped_gemm_glu_dense_wrapper_fp4( + sf_fp8_dtype_override, ab_dtype, c_dtype, d_dtype, @@ -239,6 +243,7 @@ def test_grouped_gemm_glu_dense_wrapper_fp4( vector_f32=vector_f32, discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, + sf_fp8_dtype_override=sf_fp8_dtype_override, request=request, ) @@ -676,6 +681,7 @@ def _test_grouped_gemm_glu_dense_wrapper( use_dynamic_sched=False, omit_prob=False, use_single_group_runtime_offsets=False, + sf_fp8_dtype_override=None, ): try: from cudnn import grouped_gemm_glu_wrapper_sm100 @@ -717,6 +723,13 @@ def _test_grouped_gemm_glu_dense_wrapper( if input_mutator is not None: input_mutator(inputs, cfg) + if sf_fp8_dtype_override == "e5m3": + # Re-encode the scale factors in place. The generator emits values that are + # exact in both e4m3 and ue5m3, so sfa_ref/sfb_ref stay valid and the shared + # reference check below is reused unchanged -- only the bytes differ. + inputs["sfa_tensor"] = reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + inputs["sfb_tensor"] = reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + try: for _ in range(2): # Run twice to test caching path outputs = grouped_gemm_glu_wrapper_sm100( @@ -738,6 +751,7 @@ def _test_grouped_gemm_glu_dense_wrapper( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -960,7 +974,16 @@ def counted_compile(self): @torch_fork_set_rng(seed=0) @with_scheduler_modes @pytest.mark.parametrize("act_func", ["swiglu", "geglu"]) -def test_grouped_gemm_glu_discrete_compile_execute_fp4(act_func, use_dynamic_sched, request): +@pytest.mark.parametrize( + "sf_vec_size,sf_dtype,sf_fp8_dtype_override", + [ + (32, torch.float8_e8m0fnu, None), + (16, torch.float8_e4m3fn, None), + (16, torch.float8_e4m3fn, "e5m3"), + ], + ids=["mxfp4", "nvfp4_e4m3", "nvfp4_e5m3"], +) +def test_grouped_gemm_glu_discrete_compile_execute_fp4(sf_vec_size, sf_dtype, sf_fp8_dtype_override, act_func, use_dynamic_sched, request): _test_grouped_gemm_glu_discrete_compile_execute( ab_dtype=torch.float4_e2m1fn_x2, c_dtype=torch.bfloat16, @@ -969,8 +992,9 @@ def test_grouped_gemm_glu_discrete_compile_execute_fp4(act_func, use_dynamic_sch acc_dtype=torch.float32, mma_tiler_mn=(256, 256), cluster_shape_mn=(2, 1), - sf_vec_size=32, - sf_dtype=torch.float8_e8m0fnu, + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=False, discrete_col_sfd=False, act_func=act_func, @@ -1074,6 +1098,7 @@ def _test_grouped_gemm_glu_discrete_compile_execute( b_major="k", enable_bias=False, use_dynamic_sched=False, + sf_fp8_dtype_override=None, ): try: from cudnn import GroupedGemmGluSm100 @@ -1113,6 +1138,14 @@ def _test_grouped_gemm_glu_discrete_compile_execute( enable_bias=enable_bias, ) + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place. The generator emits values exact + # in both E4M3 and UE5M3, so the reference stays valid; in place is required + # because sfb_ptrs already holds the per-expert device addresses. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + for sfb in inputs["sfb_list"]: + reencode_sf_tensor_as_ue5m3(sfb) + outputs = allocate_discrete_output_tensors( tensor_m=inputs["tensor_m"], n=cfg["n"], @@ -1146,6 +1179,7 @@ def _test_grouped_gemm_glu_discrete_compile_execute( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -1950,3 +1984,160 @@ def invalidate_runtime_offset(inputs, _cfg): input_mutator=invalidate_runtime_offset, use_single_group_runtime_offsets=True, ) + + +_NVFP4_E5M3_CFG = dict( + ab_dtype=torch.float4_e2m1fn_x2, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=16, + sf_dtype=torch.float8_e4m3fn, + vector_f32=False, + discrete_col_sfd=False, +) + + +def _make_glu_inputs(request, **overrides): + """Build a dense NVFP4 problem without running the reference check.""" + cfg = grouped_gemm_swiglu_init(request, **{**_NVFP4_E5M3_CFG, **overrides}) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + enable_bias=cfg["enable_bias"], + ) + return cfg, inputs + + +def _run_glu_wrapper(cfg, inputs, sf_fp8_dtype_override): + """Call the wrapper directly. + + The shared harness turns ValueError into pytest.skip and always runs the + reference check, neither of which suits the rejection and cache tests below. + """ + from cudnn import grouped_gemm_glu_wrapper_sm100 + + return grouped_gemm_glu_wrapper_sm100( + a_tensor=inputs["a_tensor"], + sfa_tensor=inputs["sfa_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + b_tensor=inputs["b_tensor"], + sfb_tensor=inputs["sfb_tensor"], + norm_const_tensor=inputs.get("norm_const_tensor"), + prob_tensor=inputs.get("prob_tensor"), + acc_dtype=cfg["acc_dtype"], + c_dtype=cfg["c_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + discrete_col_sfd=cfg["discrete_col_sfd"], + act_func="swiglu", + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_e5m3_is_not_cached_as_e4m3(request): + """sf_fp8_dtype_override must take part in the compile cache key. + + Identical scale-factor bytes decode to different values under E4M3 and + UE5M3, so if sf_fp8_dtype_override were omitted from the key the second call would + reuse the first kernel and silently return E4M3 results. + """ + _skip_unless_e5m3_supported() + + # One problem, one set of scale-factor bytes, two interpretations. Any + # difference in the output can only come from sf_fp8_dtype_override. + cfg, inputs = _make_glu_inputs(request) + d_e4m3 = _run_glu_wrapper(cfg, inputs, None)["d_tensor"].float().clone() + d_e5m3 = _run_glu_wrapper(cfg, inputs, "e5m3")["d_tensor"].float().clone() + + assert not torch.equal( + d_e5m3, d_e4m3 + ), "e5m3 and e4m3 produced identical output from identical scale-factor bytes; sf_fp8_dtype_override is likely missing from the compile cache key" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + "sf_fp8_dtype_override,overrides,expected", + [ + pytest.param( + "e5m3", + dict(ab_dtype=torch.float8_e4m3fn, sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), + "requires the NVFP4 recipe", + id="fp8_ab", + ), + pytest.param( + "e5m3", + dict(sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), + "requires the NVFP4 recipe", + id="mxfp4_e8m0_carrier", + ), + pytest.param("e4m3", {}, "sf_fp8_dtype_override must be", id="e4m3_is_not_an_override"), + pytest.param("e5m2", {}, "sf_fp8_dtype_override must be", id="unknown_format"), + ], +) +def test_grouped_gemm_glu_rejects_unsupported_sf_fp8_dtype(request, sf_fp8_dtype_override, overrides, expected): + """e5m3 is only reachable through the Rubin FP4xFP4 atom with e4m3-carried scales.""" + if sf_fp8_dtype_override == "e5m3": + _skip_unless_e5m3_supported() + cfg, inputs = _make_glu_inputs(request, **overrides) + with pytest.raises(ValueError, match=expected): + _run_glu_wrapper(cfg, inputs, sf_fp8_dtype_override) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_mxfp4_ignores_sf_fp8_dtype_override_default(request): + """Leaving the override at None must keep inferring E8M0 for MXFP4.""" + _test_grouped_gemm_glu_dense_wrapper( + request=request, + sf_fp8_dtype_override=None, + **{**_NVFP4_E5M3_CFG, "sf_vec_size": 32, "sf_dtype": torch.float8_e8m0fnu}, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_glu_bf16_rejects_sf_fp8_dtype_override(request): + """The BF16 backend has no scale factors, so any explicit override is an error. + + This is what the None default buys: with an "e4m3" default the value would + always be non-None and could not be distinguished from an explicit request. + """ + try: + from cudnn import grouped_gemm_glu_wrapper_sm100 + except ImportError: + pytest.skip("cudnn optional dependencies not installed") + + problem = make_grouped_gemm_glu_bf16_problem(discrete=False, b_major="k") + kwargs = dict( + a_tensor=problem["a"], + sfa_tensor=None, + padded_offsets=problem["offsets"], + alpha_tensor=problem["alpha"], + bias_tensor=problem["bias"], + prob_tensor=problem["prob"], + b_tensor=problem["b"], + sfb_tensor=None, + ) + # None is accepted and dispatches to BF16 as usual. + grouped_gemm_glu_wrapper_sm100(**kwargs, sf_fp8_dtype_override=None) + with pytest.raises(ValueError, match="BF16 forbids scale control sf_fp8_dtype_override"): + grouped_gemm_glu_wrapper_sm100(**kwargs, sf_fp8_dtype_override="e5m3") diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py index 9eec31812..4ae078163 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py @@ -19,6 +19,8 @@ from fe_api.grouped_gemm.test_discrete_grouped_gemm_swiglu_utils import ( allocate_discrete_input_tensors, ) +from fe_api.test_fe_api_utils import reencode_sf_tensor_as_ue5m3 +from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import _skip_unless_e5m3_supported from fe_api.grouped_gemm.test_grouped_gemm_quant_utils import ( grouped_gemm_quant_init, with_grouped_gemm_quant_params_fp4, @@ -126,6 +128,7 @@ def test_grouped_gemm_quant_compile_execute_fp4( discrete_col_sfd, use_dynamic_sched, request, + sf_fp8_dtype_override, ): _test_grouped_gemm_quant_compile_execute( ab_dtype=ab_dtype, @@ -137,6 +140,7 @@ def test_grouped_gemm_quant_compile_execute_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, @@ -198,6 +202,7 @@ def test_grouped_gemm_quant_wrapper_fp4( discrete_col_sfd, use_dynamic_sched, request, + sf_fp8_dtype_override, ): _test_grouped_gemm_quant_wrapper( ab_dtype=ab_dtype, @@ -209,6 +214,7 @@ def test_grouped_gemm_quant_wrapper_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=vector_f32, discrete_col_sfd=discrete_col_sfd, use_dynamic_sched=use_dynamic_sched, @@ -1079,6 +1085,7 @@ def _test_grouped_gemm_quant_compile_execute( cfg_overrides=None, input_mutator=None, use_dynamic_sched=False, + sf_fp8_dtype_override=None, ): """Test GroupedGemmQuant API with explicit check_support, compile, and execute paths.""" try: @@ -1131,6 +1138,12 @@ def _test_grouped_gemm_quant_compile_execute( if input_mutator is not None: input_mutator(inputs, cfg) + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + api = GroupedGemmQuantSm100( sample_a=inputs["a_tensor"], sample_b=inputs["b_tensor"], @@ -1150,6 +1163,7 @@ def _test_grouped_gemm_quant_compile_execute( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -1208,6 +1222,7 @@ def _test_grouped_gemm_quant_wrapper( enable_bias=False, provide_d_tensor=False, use_single_group_runtime_offsets=False, + sf_fp8_dtype_override=None, ): """Test GroupedGemmQuant API via the wrapper function (with caching).""" try: @@ -1259,6 +1274,12 @@ def _test_grouped_gemm_quant_wrapper( device=inputs["a_tensor"].device, ) + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so the fp32 reference stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) + try: for _ in range(2): # Run twice to test caching path outputs = grouped_gemm_quant_wrapper_sm100( @@ -1279,6 +1300,7 @@ def _test_grouped_gemm_quant_wrapper( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, vector_f32=cfg["vector_f32"], m_aligned=cfg["m_aligned"], discrete_col_sfd=cfg["discrete_col_sfd"], @@ -1557,3 +1579,100 @@ def _test_grouped_gemm_quant_discrete_wrapper( skip_ref=cfg["skip_ref"], ) return inputs, outputs, cfg + + +def _quant_nvfp4_inputs(request, sf_vec_size=16, sf_dtype=torch.float8_e4m3fn, ab_dtype=torch.float4_e2m1fn_x2): + cfg = grouped_gemm_quant_init( + request, + ab_dtype=ab_dtype, + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + cd_major="n", + acc_dtype=torch.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, + vector_f32=False, + discrete_col_sfd=False, + ) + inputs = allocate_grouped_gemm_input_tensors( + n=cfg["n"], + k=cfg["k"], + l=cfg["l"], + group_m_list=cfg["group_m_list"], + ab_dtype=cfg["ab_dtype"], + sf_dtype=cfg["sf_dtype"], + sf_vec_size=cfg["sf_vec_size"], + m_aligned=cfg["m_aligned"], + ) + return cfg, inputs + + +def _run_quant_wrapper(cfg, inputs, sf_fp8_dtype_override): + """Call the wrapper directly; the harnesses turn ValueError into a skip.""" + from cudnn import grouped_gemm_quant_wrapper_sm100 + + return grouped_gemm_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + norm_const_tensor=inputs.get("norm_const_tensor"), + acc_dtype=cfg["acc_dtype"], + d_dtype=cfg["d_dtype"], + cd_major=cfg["cd_major"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + vector_f32=cfg["vector_f32"], + m_aligned=cfg["m_aligned"], + discrete_col_sfd=cfg["discrete_col_sfd"], + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + "sf_fp8_dtype_override,overrides,expected", + [ + pytest.param("e5m3", dict(sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), "requires the NVFP4 recipe", id="mxfp4_e8m0_carrier"), + pytest.param( + "e5m3", + dict(ab_dtype=torch.float8_e4m3fn, sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), + "requires the NVFP4 recipe", + id="fp8_ab", + ), + pytest.param("e4m3", {}, "sf_fp8_dtype_override must be", id="e4m3_is_not_an_override"), + pytest.param("e5m2", {}, "sf_fp8_dtype_override must be", id="unknown_format"), + ], +) +def test_grouped_gemm_quant_rejects_unsupported_sf_fp8_dtype(request, sf_fp8_dtype_override, overrides, expected): + """e5m3 is only reachable through the Rubin FP4xFP4 atom with e4m3-carried scales.""" + if sf_fp8_dtype_override == "e5m3": + _skip_unless_e5m3_supported() + cfg, inputs = _quant_nvfp4_inputs(request, **overrides) + with pytest.raises(ValueError, match=expected): + _run_quant_wrapper(cfg, inputs, sf_fp8_dtype_override) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_quant_e5m3_is_not_cached_as_e4m3(request): + """sf_fp8_dtype_override must take part in the compile cache key. + + Identical scale-factor bytes decode differently under E4M3 and UE5M3, so if + the override were missing from the key the second call would reuse the first + kernel and silently return E4M3 results. + """ + _skip_unless_e5m3_supported() + cfg, inputs = _quant_nvfp4_inputs(request) + d_e4m3 = _run_quant_wrapper(cfg, inputs, None)["d_tensor"].float().clone() + d_e5m3 = _run_quant_wrapper(cfg, inputs, "e5m3")["d_tensor"].float().clone() + torch.cuda.synchronize() + assert not torch.equal( + d_e4m3, d_e5m3 + ), "e5m3 and e4m3 produced identical output from identical scale-factor bytes; sf_fp8_dtype_override is likely missing from the compile cache key" diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py index e9e25bdb7..f5cc00577 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py @@ -105,13 +105,16 @@ (1, 1), ], ), - pytest.mark.parametrize("sf_vec_size", [16, 32]), pytest.mark.parametrize( - "sf_dtype", + "sf_vec_size,sf_dtype,sf_fp8_dtype_override", [ - torch.float8_e8m0fnu, - torch.float8_e4m3fn, + (16, torch.float8_e8m0fnu, None), + (16, torch.float8_e4m3fn, None), + (32, torch.float8_e8m0fnu, None), + (32, torch.float8_e4m3fn, None), + (16, torch.float8_e4m3fn, "e5m3"), ], + ids=["v16_e8m0", "v16_e4m3", "v32_e8m0", "v32_e4m3", "v16_e5m3"], ), pytest.mark.parametrize("vector_f32", [True, False]), pytest.mark.parametrize("discrete_col_sfd", [False]), diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_utils.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_utils.py index f7c1bec64..d2d331111 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_utils.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_swiglu_utils.py @@ -9,6 +9,8 @@ """ import torch +import functools + import pytest from typing import Optional, Tuple, List, Dict, Any from fe_api.test_fe_api_utils import ( @@ -94,24 +96,41 @@ ] ) -GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4 = ( +GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4_BASE = ( GROUPED_GEMM_SWIGLU_FP4_TYPE_MARKS + GROUPED_GEMM_SWIGLU_COMMON_MARKS + [ pytest.mark.parametrize("mma_tiler_mn", [(256, 256), (128, 256)]), - pytest.mark.parametrize( - "sf_vec_size,sf_dtype", - [ - (16, torch.float8_e8m0fnu), - (16, torch.float8_e4m3fn), - (32, torch.float8_e8m0fnu), - (32, torch.float8_e4m3fn), - ], - ), pytest.mark.parametrize("discrete_col_sfd", [False]), ] ) +GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4 = GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4_BASE + [ + pytest.mark.parametrize( + "sf_vec_size,sf_dtype", + [ + (16, torch.float8_e8m0fnu), + (16, torch.float8_e4m3fn), + (32, torch.float8_e8m0fnu), + (32, torch.float8_e4m3fn), + ], + ), +] + +GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4_WITH_E5M3 = GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4_BASE + [ + pytest.mark.parametrize( + "sf_vec_size,sf_dtype,sf_fp8_dtype_override", + [ + (16, torch.float8_e8m0fnu, None), + (16, torch.float8_e4m3fn, None), + (32, torch.float8_e8m0fnu, None), + (32, torch.float8_e4m3fn, None), + (16, torch.float8_e4m3fn, "e5m3"), + ], + ids=["v16_e8m0", "v16_e4m3", "v32_e8m0", "v32_e4m3", "v16_e5m3"], + ), +] + GROUPED_GEMM_SWIGLU_PARAM_MARKS_BIAS_FP8 = ( GROUPED_GEMM_SWIGLU_FP8_TYPE_MARKS + GROUPED_GEMM_SWIGLU_COMMON_MARKS @@ -140,9 +159,12 @@ ) -def with_grouped_gemm_swiglu_params_fp4(func): +def with_grouped_gemm_swiglu_params_fp4(func=None, *, with_e5m3: bool = False): """Decorator to apply grouped GEMM SwiGLU FP4 test parameters.""" - for mark in reversed(GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4): + if func is None: + return functools.partial(with_grouped_gemm_swiglu_params_fp4, with_e5m3=with_e5m3) + param_marks = GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4_WITH_E5M3 if with_e5m3 else GROUPED_GEMM_SWIGLU_PARAM_MARKS_FP4 + for mark in reversed(param_marks): func = mark(func) return func diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py index ada88d0c8..85fefa5ea 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py @@ -8,7 +8,9 @@ import cudnn from test_utils import torch_fork_set_rng +from fe_api.test_fe_api_utils import reencode_sf_tensor_as_ue5m3 from fe_api.grouped_gemm.test_grouped_gemm_wgrad_utils import ( + _skip_unless_e5m3_supported, grouped_gemm_wgrad_init, with_grouped_gemm_wgrad_params_fp4, with_grouped_gemm_wgrad_params_fp8, @@ -64,6 +66,7 @@ def _test_grouped_gemm_wgrad_dense_compile_execute( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override=None, ): cfg = grouped_gemm_wgrad_init( ab_dtype=ab_dtype, @@ -75,6 +78,12 @@ def _test_grouped_gemm_wgrad_dense_compile_execute( sf_dtype=sf_dtype, ) inputs = allocate_grouped_gemm_wgrad_tensors(cfg) + + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so inputs["ref_result"] stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) wgrad_tensor = allocate_grouped_gemm_wgrad_output(cfg) op = cudnn.GroupedGemmWgradSm100( @@ -90,6 +99,7 @@ def _test_grouped_gemm_wgrad_dense_compile_execute( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, ) try: assert op.check_support() @@ -121,6 +131,7 @@ def test_grouped_gemm_wgrad_dense_compile_execute_fp4( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override, ): _test_grouped_gemm_wgrad_dense_compile_execute( ab_dtype=ab_dtype, @@ -130,6 +141,7 @@ def test_grouped_gemm_wgrad_dense_compile_execute_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -169,6 +181,7 @@ def _test_grouped_gemm_wgrad_dense_wrapper( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override=None, ): cfg = grouped_gemm_wgrad_init( ab_dtype=ab_dtype, @@ -180,6 +193,12 @@ def _test_grouped_gemm_wgrad_dense_wrapper( sf_dtype=sf_dtype, ) inputs = allocate_grouped_gemm_wgrad_tensors(cfg) + + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so inputs["ref_result"] stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) try: for _ in range(2): # Run twice to test caching path result = cudnn.grouped_gemm_wgrad_wrapper_sm100( @@ -196,6 +215,7 @@ def _test_grouped_gemm_wgrad_dense_wrapper( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, ) except (ValueError, NotImplementedError) as e: pytest.skip(f"Unsupported testcase: {e}") @@ -214,6 +234,7 @@ def test_grouped_gemm_wgrad_dense_wrapper_fp4( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override, ): _test_grouped_gemm_wgrad_dense_wrapper( ab_dtype=ab_dtype, @@ -223,6 +244,7 @@ def test_grouped_gemm_wgrad_dense_wrapper_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -263,6 +285,7 @@ def _test_grouped_gemm_wgrad_discrete_compile_execute( sf_vec_size, sf_dtype, accumulate_on_output=False, + sf_fp8_dtype_override=None, ): cfg = grouped_gemm_wgrad_init( ab_dtype=ab_dtype, @@ -274,6 +297,12 @@ def _test_grouped_gemm_wgrad_discrete_compile_execute( sf_dtype=sf_dtype, ) inputs = allocate_grouped_gemm_wgrad_tensors(cfg) + + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so inputs["ref_result"] stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) wgrad_tensor = allocate_grouped_gemm_wgrad_output(cfg, accumulate_on_output=accumulate_on_output) expected = inputs["ref_result"] if accumulate_on_output: @@ -297,6 +326,7 @@ def _test_grouped_gemm_wgrad_discrete_compile_execute( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, accumulate_on_output=accumulate_on_output, ) try: @@ -329,6 +359,7 @@ def test_grouped_gemm_wgrad_discrete_compile_execute_fp4( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override, ): _test_grouped_gemm_wgrad_discrete_compile_execute( ab_dtype=ab_dtype, @@ -338,6 +369,7 @@ def test_grouped_gemm_wgrad_discrete_compile_execute_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -375,6 +407,7 @@ def test_grouped_gemm_wgrad_discrete_accumulate_compile_execute_fp4( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override, ): _test_grouped_gemm_wgrad_discrete_compile_execute( ab_dtype=ab_dtype, @@ -384,6 +417,7 @@ def test_grouped_gemm_wgrad_discrete_accumulate_compile_execute_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, accumulate_on_output=True, ) @@ -425,6 +459,7 @@ def _test_grouped_gemm_wgrad_discrete_wrapper( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override=None, ): cfg = grouped_gemm_wgrad_init( ab_dtype=ab_dtype, @@ -436,6 +471,12 @@ def _test_grouped_gemm_wgrad_discrete_wrapper( sf_dtype=sf_dtype, ) inputs = allocate_grouped_gemm_wgrad_tensors(cfg) + + if sf_fp8_dtype_override == "e5m3": + # Rewrite the scale bytes as UE5M3 in place; values are exact in both + # formats so inputs["ref_result"] stays valid. + reencode_sf_tensor_as_ue5m3(inputs["sfa_tensor"]) + reencode_sf_tensor_as_ue5m3(inputs["sfb_tensor"]) try: for _ in range(2): # Run twice to test caching path result = cudnn.grouped_gemm_wgrad_wrapper_sm100( @@ -452,6 +493,7 @@ def _test_grouped_gemm_wgrad_discrete_wrapper( mma_tiler_mn=cfg["mma_tiler_mn"], cluster_shape_mn=cfg["cluster_shape_mn"], sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, ) except (ValueError, NotImplementedError) as e: pytest.skip(f"Unsupported testcase: {e}") @@ -470,6 +512,7 @@ def test_grouped_gemm_wgrad_discrete_wrapper_fp4( cluster_shape_mn, sf_vec_size, sf_dtype, + sf_fp8_dtype_override, ): _test_grouped_gemm_wgrad_discrete_wrapper( ab_dtype=ab_dtype, @@ -479,6 +522,7 @@ def test_grouped_gemm_wgrad_discrete_wrapper_fp4( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, sf_dtype=sf_dtype, + sf_fp8_dtype_override=sf_fp8_dtype_override, ) @@ -607,7 +651,11 @@ def test_grouped_gemm_wgrad_dynamic_tokens_compile_execute_fp4( sf_vec_size, sf_dtype, output_mode, + sf_fp8_dtype_override, # noqa: ARG001 ): + if sf_fp8_dtype_override is not None: + pytest.skip("Skip e5m3 test. This test is not for numerical correctness and covering e5m3's gain is marginal.") + _test_grouped_gemm_wgrad_dynamic_tokens_compile_execute( ab_dtype=ab_dtype, wgrad_dtype=wgrad_dtype, @@ -790,3 +838,113 @@ def test_grouped_gemm_wgrad_dense_wrapper_tensor_ragged_fp4(): torch.cuda.synchronize() check_ref_grouped_gemm_wgrad(result["wgrad_tensor"], inputs["ref_result"], cfg["tolerance"]) + + +def _wgrad_nvfp4_inputs(sf_vec_size=16, sf_dtype=torch.float8_e4m3fn, ab_dtype=torch.float4_e2m1fn_x2): + cfg = grouped_gemm_wgrad_init( + ab_dtype=ab_dtype, + wgrad_dtype=torch.bfloat16, + acc_dtype=torch.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=sf_vec_size, + sf_dtype=sf_dtype, + ) + return cfg, allocate_grouped_gemm_wgrad_tensors(cfg) + + +def _run_wgrad_wrapper(cfg, inputs, sf_fp8_dtype_override): + """Call the wrapper directly; the harnesses turn ValueError into a skip.""" + return cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + offsets_tensor=inputs["offsets_tensor"], + output_mode="dense", + global_scale_a=inputs["global_scale_a"], + global_scale_b=inputs["global_scale_b"], + acc_dtype=cfg["acc_dtype"], + wgrad_dtype=cfg["wgrad_dtype"], + mma_tiler_mn=cfg["mma_tiler_mn"], + cluster_shape_mn=cfg["cluster_shape_mn"], + sf_vec_size=cfg["sf_vec_size"], + sf_fp8_dtype_override=sf_fp8_dtype_override, + ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +@pytest.mark.parametrize( + "sf_fp8_dtype_override,overrides,expected", + [ + pytest.param("e5m3", dict(sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), "requires the NVFP4 recipe", id="mxfp4_e8m0_carrier"), + pytest.param( + "e5m3", + dict(ab_dtype=torch.float8_e4m3fn, sf_vec_size=32, sf_dtype=torch.float8_e8m0fnu), + "requires the NVFP4 recipe", + id="fp8_ab", + ), + pytest.param("e4m3", {}, "sf_fp8_dtype_override must be", id="e4m3_is_not_an_override"), + pytest.param("e5m2", {}, "sf_fp8_dtype_override must be", id="unknown_format"), + ], +) +def test_grouped_gemm_wgrad_rejects_unsupported_sf_fp8_dtype(sf_fp8_dtype_override, overrides, expected): + """e5m3 is only reachable through the Rubin FP4xFP4 atom with e4m3-carried scales.""" + if sf_fp8_dtype_override == "e5m3": + _skip_unless_e5m3_supported() + cfg, inputs = _wgrad_nvfp4_inputs(**overrides) + with pytest.raises(ValueError, match=expected): + _run_wgrad_wrapper(cfg, inputs, sf_fp8_dtype_override) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_wgrad_e5m3_is_not_cached_as_e4m3(): + """sf_fp8_dtype_override must take part in the compile cache key. + + Identical scale-factor bytes decode differently under E4M3 and UE5M3, so if + the override were missing from the key the second call would reuse the first + kernel and silently return E4M3 results. + """ + _skip_unless_e5m3_supported() + cfg, inputs = _wgrad_nvfp4_inputs() + w_e4m3 = _run_wgrad_wrapper(cfg, inputs, None)["wgrad_tensor"].float().clone() + w_e5m3 = _run_wgrad_wrapper(cfg, inputs, "e5m3")["wgrad_tensor"].float().clone() + torch.cuda.synchronize() + assert not torch.equal( + w_e4m3, w_e5m3 + ), "e5m3 and e4m3 produced identical output from identical scale-factor bytes; sf_fp8_dtype_override is likely missing from the compile cache key" + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_grouped_gemm_wgrad_bf16_rejects_sf_fp8_dtype_override(): + """The BF16 backend has no scale factors, so any explicit override is an error. + + The None case is the important one: wgrad forwards **kwargs to both backends, + so merely adding the parameter once broke every BF16 call with a TypeError, + which a rejection-only test would not have caught. + """ + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("Requires SM100+ for grouped GEMM WGrad BF16 kernel.") + problem = make_grouped_gemm_wgrad_bf16_problem(discrete=False) + kwargs = dict( + a_tensor=problem["a"], + b_tensor=problem["b"], + sfa_tensor=None, + sfb_tensor=None, + offsets_tensor=problem["offsets"], + output_mode="dense", + wgrad_tensor=problem["output"], + wgrad_ptrs=problem["output_ptrs"], + acc_dtype=torch.float32, + wgrad_dtype=problem["output_dtype"], + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + input_order=problem["input_order"], + ) + # None is accepted and dispatches to BF16 as usual. + cudnn.grouped_gemm_wgrad_wrapper_sm100(**kwargs, sf_fp8_dtype_override=None) + with pytest.raises(ValueError, match="BF16 forbids scale control sf_fp8_dtype_override"): + cudnn.grouped_gemm_wgrad_wrapper_sm100(**kwargs, sf_fp8_dtype_override="e5m3") diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py index ef39eee1a..89062706f 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad_utils.py @@ -9,6 +9,25 @@ import torch from fe_api.test_fe_api_utils import ceil_div + +def _skip_unless_e5m3_supported(): + """E5M3 scales need Rubin plus an internal cutlass-dsl build. + + Shared by the glu, dglu, quant and wgrad e5m3 guard tests -- those are the + only four APIs that expose ``sf_fp8_dtype_override``. + """ + try: + import cutlass + + from cudnn.api_base import get_device_type + except ImportError: + pytest.skip("cudnn optional dependencies not installed") + if get_device_type() != "rubin": + pytest.skip("e5m3 scale factors require Rubin (SM107)") + if not hasattr(cutlass, "FloatNV8E5M3FNU"): + pytest.skip("cutlass-dsl build does not provide FloatNV8E5M3FNU") + + GROUPED_GEMM_WGRAD_PARAM_MARKS_FP4 = [ pytest.mark.parametrize("ab_dtype", [torch.float4_e2m1fn_x2]), pytest.mark.parametrize("wgrad_dtype", [torch.bfloat16]), @@ -17,6 +36,7 @@ pytest.mark.parametrize("cluster_shape_mn", [(1, 1), (2, 1)]), pytest.mark.parametrize("sf_vec_size", [16]), pytest.mark.parametrize("sf_dtype", [torch.float8_e4m3fn]), + pytest.mark.parametrize("sf_fp8_dtype_override", [None, "e5m3"], ids=["sf_e4m3", "sf_e5m3"]), ] GROUPED_GEMM_WGRAD_PARAM_MARKS_FP8 = [ diff --git a/test/python/fe_api/test_fe_api_utils.py b/test/python/fe_api/test_fe_api_utils.py index 91de1bb5c..3c6f4ac6a 100644 --- a/test/python/fe_api/test_fe_api_utils.py +++ b/test/python/fe_api/test_fe_api_utils.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import functools + import torch import pytest from test_low_precision_matmul import ( @@ -117,6 +119,72 @@ def create_sf_layout_tensor(l, mn, nk, sf_vec_size): return cute_f32_torch_tensor_cpu, sf_k +_UE5M3_BIAS = 15 +_UE5M3_NAN_BYTE = 0xFF + + +def ue5m3_decode_byte(byte: int) -> float: + """Decode a single UE5M3 byte to float.""" + if byte == _UE5M3_NAN_BYTE: + return float("nan") + exp = (byte >> 3) & 0x1F + mant = byte & 0x07 + if exp == 0: + return (2.0 ** (1 - _UE5M3_BIAS)) * (mant / 8.0) + return (2.0 ** (exp - _UE5M3_BIAS)) * (1.0 + mant / 8.0) + + +@functools.lru_cache(maxsize=None) +def _ue5m3_lut(device) -> torch.Tensor: + """All finite UE5M3 values, indexed by byte (0x00..0xFE). + + Memoized per device: building it is a 255-iteration Python loop plus a + host-to-device copy, which otherwise dominates the encode for tensors of + scale-factor size. Callers share one tensor, so treat it as read-only -- + never mutate it or apply an in-place op to it. + """ + return torch.tensor( + [ue5m3_decode_byte(b) for b in range(_UE5M3_NAN_BYTE)], + dtype=torch.float32, + device=device, + ) + + +def f32_to_ue5m3_bytes(values: torch.Tensor) -> torch.Tensor: + """Round-to-nearest-even encode a float tensor to UE5M3, returned as uint8.""" + lut = _ue5m3_lut(values.device) + flat = values.detach().to(torch.float32).reshape(-1) + + is_nan = flat.isnan() + max_finite = lut[-1] + # UE5M3 can't express negative values or values larger than its maximum representable value (the last one from the LUT). + # NaN compares False against both bounds, so it passes through here and picks up the NaN byte at the end. + out_of_range = (flat < 0) | (flat > max_finite) + if out_of_range.any(): + offenders = flat[out_of_range] + raise ValueError(f"{offenders.numel()} value(s) outside the finite UE5M3 range [0, {max_finite.item()}], e.g. {offenders[0].item()}") + + # Find the next larger and next smaller LUT entries for each value. For NAN it returns len(lut) which is clamped + hi = torch.searchsorted(lut, flat).clamp(max=lut.numel() - 1) + lo = (hi - 1).clamp(min=0) + # Round to nearest + d_lo = flat - lut[lo] + d_hi = lut[hi] - flat + idx = torch.where(d_hi < d_lo, hi, lo) + # lo and hi are adjacent, so exactly one of them is even; ties take that one. + idx = torch.where(d_lo == d_hi, torch.where(lo % 2 == 0, lo, hi), idx) + + return idx.to(torch.uint8).masked_fill(is_nan, _UE5M3_NAN_BYTE).reshape(values.shape) + + +def reencode_sf_tensor_as_ue5m3(sf_tensor: torch.Tensor) -> torch.Tensor: + """Rewrite an e4m3-valued scale-factor tensor's bytes as UE5M3, in place.""" + assert sf_tensor.dtype == torch.float8_e4m3fn, f"expected e4m3 storage, got {sf_tensor.dtype}" + encoded = f32_to_ue5m3_bytes(sf_tensor.to(torch.float32)) + sf_tensor.view(torch.uint8).copy_(encoded) + return sf_tensor + + # Create scale factor tensor SFA/SFB def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): cute_f32_torch_tensor_cpu, sf_k = create_sf_layout_tensor(l, mn, k, sf_vec_size)