Prototype NVFP4 with FP8 UE5M3 block scales

tests/cpp/test_common.cu

@@ -65,6 +65,7 @@ const std::string &typeName(DType type) {
     {DType::kBFloat16, "bfloat16"},
     {DType::kFloat8E4M3, "float8e4m3"},
     {DType::kFloat8E5M2, "float8e5m2"},
+    {DType::kFloat8UE5M3, "float8ue5m3"},
     {DType::kFloat8E8M0, "float8e8m0"},
     {DType::kFloat4E2M1, "float4e2m1"}};
   return name_map.at(type);
@@ -132,12 +133,16 @@ NVTEShape convertShape(const std::vector<size_t>& s) {
 }
 
 std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
-                                                     const NVTEScalingMode scaling_mode) {
+                                                     const NVTEScalingMode scaling_mode,
+                                                     std::optional<DType> scale_type) {
   if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) {
     scale_inv_meta ret;
     ret.shape = {1};
-    ret.type = DType::kFloat32;
-    ret.type_size_bits = typeToNumBits(DType::kFloat32);
+    if (!scale_type) {
+      scale_type = DType::kFloat32;
+    }
+    ret.type = *scale_type;
+    ret.type_size_bits = typeToNumBits(*scale_type);
     return {ret, ret};
   }
   if (scaling_mode == NVTE_MXFP8_1D_SCALING) {
@@ -160,10 +165,13 @@ std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
     size_t scale_dim_X_colwise = DIVUP_TO_MULTIPLE(last_dim, scale_tensor_alignment_X_colwise);
     ret_colwise.shape = {scale_dim_Y_colwise, scale_dim_X_colwise};
 
-    ret_rowwise.type = DType::kFloat8E8M0;
-    ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0);
-    ret_colwise.type = DType::kFloat8E8M0;
-    ret_colwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0);
+    if (!scale_type) {
+      scale_type = DType::kFloat8E8M0;
+    }
+    ret_rowwise.type = *scale_type;
+    ret_rowwise.type_size_bits = typeToNumBits(*scale_type);
+    ret_colwise.type = *scale_type;
+    ret_colwise.type_size_bits = typeToNumBits(*scale_type);
 
     return {ret_rowwise, ret_colwise};
   }
@@ -188,10 +196,13 @@ std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
     size_t scale_dim_X_t = DIVUP_TO_MULTIPLE(DIVUP(first_dim, 16lu), scale_tensor_alignment_X_rowwise);
     ret_colwise.shape = {scale_dim_Y_t, scale_dim_X_t};
 
-    ret_rowwise.type = DType::kFloat8E4M3;
-    ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat8E4M3);
-    ret_colwise.type = DType::kFloat8E4M3;
-    ret_colwise.type_size_bits = typeToNumBits(DType::kFloat8E4M3);
+    if (!scale_type) {
+      scale_type = DType::kFloat8E4M3;
+    }
+    ret_rowwise.type = *scale_type;
+    ret_rowwise.type_size_bits = typeToNumBits(*scale_type);
+    ret_colwise.type = *scale_type;
+    ret_colwise.type_size_bits = typeToNumBits(*scale_type);
 
     return {ret_rowwise, ret_colwise};
   }
@@ -215,10 +226,14 @@ std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
       size_t scale_dim_1 = DIVUP(DIVUP(first_dim, 128lu), 4) * 4;
       ret_colwise.shape = {scale_dim_0, scale_dim_1};
     }
-    ret_rowwise.type = DType::kFloat32;
-    ret_colwise.type = DType::kFloat32;
-    ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat32);
-    ret_colwise.type_size_bits = typeToNumBits(DType::kFloat32);
+
+    if (!scale_type) {
+      scale_type = DType::kFloat32;
+    }
+    ret_rowwise.type = *scale_type;
+    ret_colwise.type = *scale_type;
+    ret_rowwise.type_size_bits = typeToNumBits(*scale_type);
+    ret_colwise.type_size_bits = typeToNumBits(*scale_type);
 
     return {ret_rowwise, ret_colwise};
   }
@@ -241,10 +256,13 @@ std::pair<scale_inv_meta, scale_inv_meta> get_scales(const NVTEShape& shape,
       size_t scale_dim_1 = DIVUP(last_dim, 4) * 4;
       ret_colwise.shape = {scale_dim_0, scale_dim_1};
     }
-    ret_rowwise.type = DType::kFloat32;
-    ret_colwise.type = DType::kFloat32;
-    ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat32);
-    ret_colwise.type_size_bits = typeToNumBits(DType::kFloat32);
+    if (!scale_type) {
+      scale_type = DType::kFloat32;
+    }
+    ret_rowwise.type = *scale_type;
+    ret_colwise.type = *scale_type;
+    ret_rowwise.type_size_bits = typeToNumBits(*scale_type);
+    ret_colwise.type_size_bits = typeToNumBits(*scale_type);
     return {ret_rowwise, ret_colwise};
   }
 
@@ -278,7 +296,8 @@ void Tensor::Buffer::from_cpu() {
 Tensor::Tensor(const std::string& name,
                const NVTEShape &shape, const DType type,
                const bool rowwise, const bool columnwise,
-               const NVTEScalingMode &scaling_mode)
+               const NVTEScalingMode &scaling_mode,
+               const std::optional<DType> scale_type)
   : tensor_(scaling_mode), rowwise_{rowwise}, columnwise_{columnwise}, name_{name} {
   // Initialize RNG
   const size_t seed = create_seed_from_tensor_name(name);
@@ -373,7 +392,7 @@ Tensor::Tensor(const std::string& name,
   case NVTE_NVFP4_1D_SCALING:
     {
       // Block scaling factors
-      auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode());
+      auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode(), scale_type);
       if (rowwise) {
         const auto scale_shape = rowwise_scale_meta.shape;
         const auto scale_dtype = rowwise_scale_meta.type;

tests/cpp/test_common.h

@@ -67,6 +67,9 @@ using bf16 = nv_bfloat16;
 using fp8e4m3 = __nv_fp8_e4m3;
 using fp8e5m2 = __nv_fp8_e5m2;
 using fp8e8m0 = uint8_t;
+#if CUDA_VERSION >= 13040
+using fp8ue5m3 = __nv_fp8_ue5m3;
+#endif
 #if FP4_TYPE_SUPPORTED
 using fp4e2m1 = __nv_fp4_e2m1;
 using fp4e2m1x2 = __nv_fp4x2_e2m1;
@@ -91,7 +94,12 @@ struct BitsNumber {
 template <typename T>
 struct TypeInfo {
 #if FP4_TYPE_SUPPORTED
-    using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3, fp8e5m2, fp8e8m0, fp4e2m1>;
+    using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3,
+                             fp8e5m2, fp8e8m0, fp4e2m1
+#if CUDA_VERSION >= 13040
+                             , fp8ue5m3
+#endif
+                             >;
 #else
     using types = std::tuple<byte, int16, int32, int64, fp32, fp16, bf16, fp8e4m3, fp8e5m2, fp8e8m0>;
 #endif
@@ -151,15 +159,18 @@ class Tensor {
          const NVTEShape &shape, const DType type,
          const bool rowwise = true,
          const bool columnwise = false,
-         const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING);
+         const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING,
+         const std::optional<DType> scale_dtype = std::nullopt);
 
   Tensor(const std::string& name,
          const std::vector<size_t> &shape,
          const DType type,
          const bool rowwise = true,
          const bool columnwise = false,
-         const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING) :
-    Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode) {}
+         const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING,
+         const std::optional<DType> scale_type = std::nullopt) :
+    Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode,
+           scale_type) {}
 
   Tensor() = default;
 
@@ -635,6 +646,16 @@ GroupedBuffers build_grouped_tensor(const std::vector<Tensor*>& tensors,
 #define SWITCH_FP4_TYPE_HANDLE(type, ...) // do nothing
 #endif
 
+#if CUDA_VERSION >= 13040
+#define SWITCH_UE5M3_TYPE_HANDLE(type, ...)     \
+  case DType::kFloat8UE5M3: {                   \
+    using type = fp8ue5m3;                      \
+    { __VA_ARGS__ }                             \
+  } break;
+#else
+#define SWITCH_UE5M3_TYPE_HANDLE(type, ...) // do nothing
+#endif
+
 #define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \
     switch (dtype) { \
         using namespace transformer_engine; \
@@ -693,7 +714,8 @@ GroupedBuffers build_grouped_tensor(const std::vector<Tensor*>& tensors,
             } \
         break; \
         SWITCH_FP4_TYPE_HANDLE(type, __VA_ARGS__) \
-        default: \
+        SWITCH_UE5M3_TYPE_HANDLE(type, __VA_ARGS__) \
+        default:                                            \
             printf("dtype: %d\n", static_cast<int>(dtype)); \
             NVTE_ERROR("Invalid type."); \
     }

tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py

@@ -17,6 +17,55 @@
 recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True)
 
 
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "disable_x, disable_w",
+    [(True, False), (False, True), (True, True)],
+    ids=["x_unit_global_scale", "w_unit_global_scale", "both_unit_global_scale"],
+)
+def test_gemm_with_missing_nvfp4_amax(disable_x: bool, disable_w: bool) -> None:
+    """A null amax contributes a unit global scale to GEMM alpha."""
+    torch.manual_seed(0)
+    x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda")
+    w = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda")
+    unit_scale_amax = 448.0 * 6.0
+    x[0, 0] = unit_scale_amax
+    w[0, 0] = unit_scale_amax
+
+    def quantize(tensor: torch.Tensor, disable_second_level_scale: bool):
+        return NVFP4Quantizer(
+            rowwise=True,
+            columnwise=True,
+            disable_second_level_scale=disable_second_level_scale,
+        )(tensor)
+
+    x_ref, w_ref = quantize(x, False), quantize(w, False)
+    x_test, w_test = quantize(x, disable_x), quantize(w, disable_w)
+
+    def gemm(w_q, x_q):
+        workspace = torch.empty(4, dtype=torch.uint8, device="cuda")
+        return tex.generic_gemm(
+            w_q,
+            True,
+            x_q,
+            False,
+            None,
+            None,
+            TE_DType[torch.bfloat16],
+            None,
+            TE_DType[torch.bfloat16],
+            False,
+            None,
+            False,
+            workspace,
+            workspace.numel(),
+            False,
+            False,
+        )[0]
+
+    torch.testing.assert_close(gemm(w_test, x_test), gemm(w_ref, x_ref), atol=0, rtol=0)
+
+
 def check_nvfp4_gemm_versus_reference(
     x_dtype: torch.dtype,
     w_dtype: torch.dtype,
@@ -641,3 +690,112 @@ def test_nvfp4_row_scaled_gemm_matches_emulated(
         use_4over6=use_4over6,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
     )
+
+
+def _check_ue5m3_gemm_versus_dequantized(
+    M, K, N, x_columnwise, w_columnwise, disable_second_level_scale
+):
+    """Run an NVFP4/UE5M3 GEMM and compare against a dequantized FP32 reference."""
+    torch.manual_seed(0)
+    device, dtype, out_dtype = "cuda", torch.bfloat16, torch.bfloat16
+    x_shape = (K, M) if x_columnwise else (M, K)
+    w_shape = (K, N) if w_columnwise else (N, K)
+    x = torch.randn(x_shape, dtype=dtype, device=device)
+    w = torch.randn(w_shape, dtype=dtype, device=device)
+
+    common = dict(
+        fp4_dtype=tex.DType.kFloat4E2M1,
+        scale_dtype=tex.DType.kFloat8UE5M3,
+        rowwise=True,
+        columnwise=True,
+        with_amax_reduction=False,
+        amax_reduction_group=None,
+        with_rht=False,
+        with_post_rht_amax=False,
+    )
+    # disable_second_level_scale is given per operand, as (x, w).
+    xq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[0])
+    wq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[1])
+    x_q = xq.update_quantized(x, xq.make_empty(x_shape, dtype=dtype, device=device))
+    w_q = wq.update_quantized(w, wq.make_empty(w_shape, dtype=dtype, device=device))
+
+    if disable_second_level_scale[0]:
+        assert x_q._amax_rowwise is None, "disable_second_level_scale should drop the amax"
+    if disable_second_level_scale[1]:
+        assert w_q._amax_rowwise is None, "disable_second_level_scale should drop the amax"
+
+    # Reference: dequantize the orientation each operand is actually read in.
+    x_ref = _dequantize_nvfp4_usage(x_q, columnwise=x_columnwise)
+    w_ref = _dequantize_nvfp4_usage(w_q, columnwise=w_columnwise)
+    # _dequantize_nvfp4_usage returns each operand canonically as (rows, K), so
+    # the reference is the same expression for every layout.
+    ref = x_ref @ w_ref.t()
+
+    if x_columnwise:
+        x_q.update_usage(rowwise_usage=False)
+    if w_columnwise:
+        w_q.update_usage(rowwise_usage=False)
+    transa, transb = not w_columnwise, x_columnwise
+    layout = ("T" if transa else "N") + ("T" if transb else "N")
+    y = general_gemm(w_q, x_q, out_dtype=out_dtype, layout=layout)[0]
+
+    # Both sides see identically quantized operands, so quantization error cancels and
+    # only accumulation order and the bf16 output rounding differ. One bf16 ulp is
+    # already ~4e-3 relative, which no elementwise tolerance survives, so compare the
+    # whole result instead.
+    rel_err = (y.float() - ref).norm() / ref.norm()
+    assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large"
+
+
+ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3)
+@pytest.mark.parametrize(
+    "M, K, N",
+    [
+        (256, 128, 256),
+        (256, 256, 256),
+        (256, 1024, 256),
+        (1024, 1024, 1024),
+        (4096, 512, 3072),
+        (112, 128, 96),
+        (304, 640, 304),
+        (1008, 3072, 992),
+        (256, 64, 256),
+        (128, 128, 112),
+    ],
+)
+@pytest.mark.parametrize(
+    "x_columnwise, w_columnwise",
+    [
+        (False, False),  # TN -- w rowwise, x rowwise   (fprop)
+        (False, True),  # NN -- w colwise, x rowwise   (dgrad)
+        (True, True),  # NT -- w colwise, x colwise   (wgrad)
+    ],
+    ids=["FF", "FT", "TT"],
+)
+@pytest.mark.parametrize(
+    "disable_second_level_scale",
+    [
+        (True, False),
+    ],
+    ids=["TF"],
+)
+def test_nvfp4_ue5m3_gemm_versus_reference(
+    M: int,
+    K: int,
+    N: int,
+    x_columnwise: bool,
+    w_columnwise: bool,
+    disable_second_level_scale: bool,
+):
+    """NVFP4 GEMM with UE5M3 block scales, with and without second-level scaling.
+
+    UE5M3's wider range is what makes dropping the per-tensor global scale
+    viable, so both configurations must match the dequantized reference.
+    """
+    _check_ue5m3_gemm_versus_dequantized(
+        M, K, N, x_columnwise, w_columnwise, disable_second_level_scale
+    )