Add support for E5m3 fused GEMM on Rubin

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

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,

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

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