[Pytorch][Common] Option to Disable 2nd level Scale in NVFP4

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,

tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py

@@ -41,6 +41,59 @@ def fused_grouped_quantize(
     return grouped_output
 
 
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "return_transpose", [False, True], ids=["rowwise", "rowwise_and_columnwise"]
+)
+def test_grouped_disable_second_level_scale_matches_split_quantize(
+    return_transpose: bool,
+) -> None:
+    """Grouped NVFP4 skips amax reduction and consumes framework-owned fixed amaxes."""
+    split_sections = [128, 128]
+    split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda")
+    torch.manual_seed(0)
+    x = torch.randn((sum(split_sections), 128), dtype=torch.bfloat16, device="cuda")
+    quantizer = NVFP4Quantizer(
+        rowwise=True,
+        columnwise=return_transpose,
+        with_rht=True,
+        with_post_rht_amax=True,
+        disable_second_level_scale=True,
+    )
+
+    grouped = fused_grouped_quantize(x, split_section_tensor, quantizer)
+    actual = grouped.split_into_quantized_tensors()
+    expected = tex.split_quantize(x, split_sections, [quantizer.copy() for _ in split_sections])
+
+    for actual_tensor, expected_tensor in zip(actual, expected):
+        torch.testing.assert_close(
+            actual_tensor._rowwise_data, expected_tensor._rowwise_data, atol=0, rtol=0
+        )
+        torch.testing.assert_close(
+            actual_tensor._rowwise_scale_inv,
+            expected_tensor._rowwise_scale_inv,
+            atol=0,
+            rtol=0,
+        )
+        assert actual_tensor._amax_rowwise is None
+        assert expected_tensor._amax_rowwise is None
+        if return_transpose:
+            torch.testing.assert_close(
+                actual_tensor._columnwise_data,
+                expected_tensor._columnwise_data,
+                atol=0,
+                rtol=0,
+            )
+            torch.testing.assert_close(
+                actual_tensor._columnwise_scale_inv,
+                expected_tensor._columnwise_scale_inv,
+                atol=0,
+                rtol=0,
+            )
+            assert actual_tensor._amax_columnwise is None
+            assert expected_tensor._amax_columnwise is None
+
+
 def check_grouped_tensor_nvfp4_versus_reference(
     x_dtype: torch.dtype,
     M: int,

tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py

@@ -17,6 +17,7 @@
 
 
 recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True)
+NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE = 448.0 * 6.0
 
 
 @dataclass(frozen=True)
@@ -240,6 +241,65 @@ def check_quantization_nvfp4_versus_reference(
     torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0)
 
 
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"])
+@pytest.mark.parametrize("use_4over6", [False, True], ids=["standard", "4over6"])
+def test_disable_second_level_scale_uses_only_block_scale(
+    return_transpose: bool,
+    use_4over6: bool,
+) -> None:
+    """A missing amax makes the global NVFP4 encode scale exactly one."""
+    torch.manual_seed(0)
+    x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda")
+    x[0, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE
+
+    common_kwargs = {
+        "rowwise": True,
+        "columnwise": return_transpose,
+        "with_rht": False,
+        "nvfp4_use_4over6": use_4over6,
+    }
+    expected = NVFP4Quantizer(**common_kwargs)(x)
+    actual = NVFP4Quantizer(**common_kwargs, disable_second_level_scale=True)(x)
+
+    torch.testing.assert_close(actual._rowwise_data, expected._rowwise_data, atol=0, rtol=0)
+    torch.testing.assert_close(
+        actual._rowwise_scale_inv, expected._rowwise_scale_inv, atol=0, rtol=0
+    )
+    torch.testing.assert_close(actual.dequantize(), expected.dequantize(), atol=0, rtol=0)
+    assert actual._amax_rowwise is None
+    if return_transpose:
+        torch.testing.assert_close(
+            actual._columnwise_data, expected._columnwise_data, atol=0, rtol=0
+        )
+        torch.testing.assert_close(
+            actual._columnwise_scale_inv, expected._columnwise_scale_inv, atol=0, rtol=0
+        )
+        assert actual._amax_columnwise is None
+
+    # Reusing an output previously populated by two-level scaling must remove
+    # its amax buffers so common kernels select the unit-global-scale path.
+    NVFP4Quantizer(**common_kwargs, disable_second_level_scale=True).update_quantized(
+        x / 2, expected
+    )
+    assert expected._amax_rowwise is None
+    if return_transpose:
+        assert expected._amax_columnwise is None
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_disable_second_level_scale_disables_row_scaled_nvfp4() -> None:
+    """Row-scaled NVFP4 is incompatible with omitting second-level scales."""
+    with pytest.warns(UserWarning, match="Row-scaled NVFP4 requires second-level scaling"):
+        quantizer = NVFP4Quantizer(
+            row_scaled_nvfp4=True,
+            disable_second_level_scale=True,
+        )
+
+    assert not quantizer.row_scaled_nvfp4
+    assert quantizer.disable_second_level_scale
+
+
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize(
     "M, N",

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -111,6 +111,8 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
       if (row_scaled_nvfp4) {
         NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization,
                    "Row-scaled NVFP4 quantization does not support 2D quantization.");
+        NVTE_CHECK(output_tensor->amax.dptr != nullptr,
+                   "Row-scaled NVFP4 does not support disabling second-level scaling.");
         NVTE_CHECK(
             !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()),
             "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 "
@@ -121,8 +123,11 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
                 (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0),
             "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are "
             "multiples of 32.");
-        nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream);
-        if (output_tensor->has_columnwise_data()) {
+        if (output_tensor->amax.dptr != nullptr) {
+          nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream);
+        }
+        if (output_tensor->has_columnwise_data() &&
+            output_tensor->columnwise_amax.dptr != nullptr) {
           nvfp4::compute_columnwise_amax(*input_tensor, noop_tensor, output_tensor, stream);
         }
       }
@@ -288,6 +293,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
       if (row_scaled_nvfp4) {
         NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization,
                    "Row-scaled NVFP4 quantization does not support 2D quantization.");
+        NVTE_CHECK(output_tensor->amax.dptr != nullptr,
+                   "Row-scaled NVFP4 does not support disabling second-level scaling.");
         NVTE_CHECK(
             !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()),
             "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 "
@@ -298,8 +305,11 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
                 (dtype == DType::kBFloat16 && rows % 32 == 0 && cols % 32 == 0),
             "Row-scaled NVFP4 transpose quantization requires BF16 input and dimensions that are "
             "multiples of 32.");
-        nvfp4::compute_rowwise_amax(*grad_tensor, noop_tensor, output_tensor, stream);
-        if (output_tensor->has_columnwise_data()) {
+        if (output_tensor->amax.dptr != nullptr) {
+          nvfp4::compute_rowwise_amax(*grad_tensor, noop_tensor, output_tensor, stream);
+        }
+        if (output_tensor->has_columnwise_data() &&
+            output_tensor->columnwise_amax.dptr != nullptr) {
           nvfp4::compute_columnwise_amax(*grad_tensor, noop_tensor, output_tensor, stream);
         }
       }

transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh

@@ -63,9 +63,13 @@ __global__ void __launch_bounds__(512)
   fp4vec value;
   value.vec = input_vectorized[my_index];
   fp8e4m3 scale = scales[my_scale_index];
-  float amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0];
   static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max.");
-  constexpr float factor_inv = 1.0f / (6.0f * static_cast<float>(E4M3_MAX));
+  float amax = transformer_engine::nvfp4::unit_global_scale_amax<E4M3_MAX, fp4e2m1>();
+  if (tensor_amax != nullptr) {
+    amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0];
+  }
+  constexpr float factor_inv =
+      1.0f / transformer_engine::nvfp4::unit_global_scale_amax<E4M3_MAX, fp4e2m1>();
   float final_scale = static_cast<float>(scale) * amax * factor_inv;
 #pragma unroll
   for (int i = 0; i < 4; i++) {
@@ -105,6 +109,8 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream)
   const size_t threads = 512;
   const size_t blocks = DIVUP(total, threads);
   const size_t num_scale_tiles_X = DIVUP(Mread, static_cast<size_t>(4));
+  NVTE_CHECK(!row_scaled_nvfp4 || input.amax.dptr != nullptr,
+             "Row-scaled NVFP4 does not support disabling second-level scaling.");
   NVTE_CHECK(!row_scaled_nvfp4 || input.amax.numel() == N,
              "Row-scaled NVFP4 dequantization requires one rowwise amax per row.");
 

transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh

@@ -476,9 +476,14 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf
       block_amax = reduce_group_max_16(group_amax);
     }
 
-    float global_amax = amax[0];
+    float global_amax = transformer_engine::nvfp4::unit_global_scale_amax<E4M3_MAX, fp4e2m1>();
+    if (amax != nullptr) {
+      global_amax = amax[0];
+    }
     if constexpr (ROW_SCALED_NVFP4) {
-      global_amax = amax[global_row];
+      if (amax != nullptr) {
+        global_amax = amax[global_row];
+      }
     }
 
     const ScalePair scale_pair = compute_scale_pair<E4M3_MAX>(block_amax, global_amax);
@@ -527,7 +532,9 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t,
       block_amax = reduce_group_max_16(group_amax);
     }
 
-    const float global_amax = amax[0];
+    const float global_amax =
+        amax == nullptr ? transformer_engine::nvfp4::unit_global_scale_amax<E4M3_MAX, fp4e2m1>()
+                        : amax[0];
     const ScalePair scale_pair = compute_scale_pair<E4M3_MAX>(block_amax, global_amax);
     CandidatePair candidates = make_candidates<Cfg, E4M3_MAX>(x0, x1, scale_pair, global_amax);
 
@@ -683,23 +690,22 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output,
              ".");
   NVTE_CHECK(!output->row_scaled_nvfp4 || !use_2d_quantization,
              "Row-scaled NVFP4 quantization does not support 2D quantization.");
+  NVTE_CHECK(!output->row_scaled_nvfp4 || output->amax.dptr != nullptr,
+             "Row-scaled NVFP4 does not support disabling second-level scaling.");
   NVTE_CHECK(!output->row_scaled_nvfp4 || !output->has_columnwise_data(),
              "Row-scaled NVFP4 quantization does not produce columnwise output.");
   NVTE_CHECK(!use_2d_quantization || output->has_data(),
              "NVFP4 4over6 2D quantization requires rowwise output.");
 
   if (output->has_data()) {
     NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated.");
-    NVTE_CHECK(output->amax.dptr != nullptr, "Rowwise amax tensor must be allocated.");
     NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type.");
   }
   if (output->has_columnwise_data()) {
     NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr,
                "Transposed scaling tensor must be allocated.");
     NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype),
                "Transposed output must have FP4 type.");
-    NVTE_CHECK(output->columnwise_amax.dptr != nullptr || output->amax.dptr != nullptr,
-               "NVFP4 4over6 columnwise quantization requires columnwise amax or rowwise amax.");
   }
 
   TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(

transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

@@ -724,7 +724,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
               scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE;
           const float S_enc_rowwise_block =
               scales_offset_Y < rows
-                  ? compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[scales_offset_Y])
+                  ? (amax_rowwise_ptr == nullptr ? 1.0f
+                                                 : compute_global_encode_scaling_factor_FP4(
+                                                       amax_rowwise_ptr[scales_offset_Y]))
                   : 1.0f;
           const float S_dec_rowwise_block = 1.0f / S_enc_rowwise_block;
           const nvfp4_scale_t S_dec_b_fp8 =
@@ -1461,7 +1463,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output,
     NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated");
   }
   NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr,
-             "Row-scaled NVFP4 quantization requires rowwise amax.");
+             "Row-scaled NVFP4 does not support disabling second-level scaling.");
   NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(),
              "Row-scaled NVFP4 quantization does not produce columnwise output.");
   // In-kernel GEMM-swizzled scale output is only implemented on the 2D quantization

transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh

@@ -233,8 +233,9 @@ __device__ __forceinline__ void colwise_scaling(
     if constexpr (ROW_SCALED_NVFP4) {
       const size_t col_idx = col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w;
       S_enc_colwise_block =
-          col_idx < cols ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx])
-                         : 1.0f;
+          col_idx < cols && amax_colwise_ptr != nullptr
+              ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx])
+              : 1.0f;
     }
     const nvfp4_scale_t S_dec_b_fp8 =
         compute_decoding_scaling_factor(block_amax[w], S_enc_colwise_block);
@@ -324,8 +325,9 @@ __device__ __forceinline__ void rowwise_scaling(
     if constexpr (ROW_SCALED_NVFP4) {
       const size_t row_idx = row_offset + stage_Y * TILE_DIM_Y + it_offset_Y_rowwise;
       const float S_enc_rowwise_block =
-          row_idx < rows ? core::compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[row_idx])
-                         : 1.0f;
+          row_idx < rows && amax_rowwise_ptr != nullptr
+              ? core::compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[row_idx])
+              : 1.0f;
       S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block);
       SFcoefficient =
           compute_nvfp4_scaling_coefficient<scaling_coeff_type>(S_dec_b_fp8, S_enc_rowwise_block);
@@ -713,15 +715,15 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop,
   NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type.");
   NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated");
   NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr,
-             "Row-scaled NVFP4 quantization requires rowwise amax.");
-
+             "Row-scaled NVFP4 does not support disabling second-level scaling.");
   if (return_transpose) {
     NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype),
                "Transposed output must have FP4 type.");
     NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr,
                "Transposed scaling tensor must be allocated");
     NVTE_CHECK(!row_scaled_nvfp4 || output->columnwise_amax.dptr != nullptr,
-               "Row-scaled NVFP4 transpose quantization requires columnwise amax.");
+               "Row-scaled NVFP4 transpose quantization does not support disabling "
+               "second-level scaling.");
   }
 
   const auto [rows, cols] = input.flat_2d_dims();