From 65867e5de0c9872adf3c59eb2328f0804959542b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Thu, 30 Jul 2026 03:03:31 +0000 Subject: [PATCH 1/3] initial draft of disable 2d Signed-off-by: Varun Thumbe --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 49 ++++++++++++ .../test_nvfp4_group_quantize_graph_safe.py | 53 +++++++++++++ .../nvfp4/test_nvfp4_quantize_exact.py | 58 ++++++++++++++ .../common/cast/dispatch/quantize.cuh | 14 +++- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 10 ++- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 18 +++-- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 7 +- .../quantize_transpose_nvfp4_tuned_1D.cuh | 17 ++-- transformer_engine/common/common.h | 21 +++++ .../common/gemm/cublaslt_grouped_gemm.cu | 43 +++-------- ...cast_col_hadamard_transform_cast_fusion.cu | 22 ++++-- .../group_hadamard_transform_cast_fusion.cu | 15 ++-- ...cast_col_hadamard_transform_cast_fusion.cu | 22 ++++-- .../hadamard_transform_cast_fusion.cu | 14 ++-- ...cast_col_hadamard_transform_cast_fusion.cu | 26 +++++-- transformer_engine/common/recipe/nvfp4.cu | 25 +++--- ...quantize_transpose_vector_blockwise_fp4.cu | 3 +- transformer_engine/pytorch/csrc/common.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 77 ++++++++++++------- transformer_engine/pytorch/csrc/quantizer.cpp | 71 ++++++++++++----- .../pytorch/csrc/type_converters.cpp | 15 ++-- .../pytorch/tensor/nvfp4_tensor.py | 17 +++- .../tensor/storage/nvfp4_tensor_storage.py | 23 +++--- 23 files changed, 454 insertions(+), 168 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 67dbe1d1d5..b836399480 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/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_2d_scaling: bool): + return NVFP4Quantizer( + rowwise=True, + columnwise=True, + disable_2d_scaling=disable_2d_scaling, + )(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, diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index b78c66befd..9b945f632e 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/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_2d_scaling_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_2d_scaling=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, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 4a3d4d3413..036fce5608 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/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,63 @@ 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,row_scaled_nvfp4", + [ + pytest.param(False, False, id="rowwise"), + pytest.param(True, False, id="rowwise_and_columnwise"), + pytest.param(False, True, id="row_scaled"), + ], +) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["standard", "4over6"]) +def test_disable_2d_scaling_uses_only_block_scale( + return_transpose: bool, + row_scaled_nvfp4: 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") + if row_scaled_nvfp4: + # The reference quantizer computes one global scale per row. + x[:, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE + else: + x[0, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE + + common_kwargs = { + "rowwise": True, + "columnwise": return_transpose, + "with_rht": False, + "row_scaled_nvfp4": row_scaled_nvfp4, + "nvfp4_use_4over6": use_4over6, + } + expected = NVFP4Quantizer(**common_kwargs)(x) + actual = NVFP4Quantizer(**common_kwargs, disable_2d_scaling=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_2d_scaling=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) @pytest.mark.parametrize( "M, N", diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index f60cee839d..44c122540f 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -121,8 +121,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); } } @@ -298,8 +301,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); } } diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 13bb01d500..eac28796d7 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/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(E4M3_MAX)); + float amax = transformer_engine::nvfp4::unit_global_scale_amax(); + 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(); float final_scale = static_cast(scale) * amax * factor_inv; #pragma unroll for (int i = 0; i < 4; i++) { @@ -105,7 +109,7 @@ 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(4)); - NVTE_CHECK(!row_scaled_nvfp4 || input.amax.numel() == N, + NVTE_CHECK(!row_scaled_nvfp4 || input.amax.dptr == nullptr || input.amax.numel() == N, "Row-scaled NVFP4 dequantization requires one rowwise amax per row."); TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 50776a3ed6..638c1bd08b 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -476,9 +476,15 @@ __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(); + 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(block_amax, global_amax); @@ -527,7 +533,10 @@ __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() + : amax[0]; const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); @@ -690,7 +699,6 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *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()) { @@ -698,8 +706,6 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, "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( diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..7e5722a5b0 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -724,7 +724,10 @@ __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 = @@ -1460,8 +1463,6 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, 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."); 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 diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index ad21486368..32396911f4 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -232,9 +232,10 @@ __device__ __forceinline__ void colwise_scaling( float S_enc_colwise_block = S_enc_colwise; 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; + S_enc_colwise_block = 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(S_dec_b_fp8, S_enc_rowwise_block); @@ -712,16 +714,11 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(output->has_data(), "NVFP4 output tensor must be allocated."); 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."); - 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."); } const auto [rows, cols] = input.flat_2d_dims(); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..6ae9094b2a 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -722,6 +722,27 @@ struct TypeExtrema { } // namespace detail +namespace nvfp4 { + +/*! \brief Amax that makes NVFP4's global encode scale equal to one. + * + * NVFP4 computes the global encode scale as + * `max(ScaleType) * max(QuantizedType) / amax`. Deriving the sentinel from + * the actual types keeps the null-amax path correct if either format changes. + */ +template +__host__ __device__ constexpr float unit_global_scale_amax() { + return detail::TypeExtrema::max * detail::TypeExtrema::max; +} + +/*! \brief Unit-global-scale amax when the scale format uses a restricted max. */ +template +__host__ __device__ constexpr float unit_global_scale_amax() { + return static_cast(ScaleMax) * detail::TypeExtrema::max; +} + +} // namespace nvfp4 + template struct BitsNumber; diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 3997e5249d..67fbc82686 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1401,21 +1401,23 @@ __global__ void setup_grouped_gemm_kernel( // For NVFP4 on Blackwell+: compute per-group alpha that includes global scale (amax). // A's amax: grouped path indexes a_amax[idx]; discrete path reads amax_ptrs[idx]. if (use_per_group_alpha_beta) { - float a_amax_val = 0.0f; - bool has_a_amax = false; + constexpr float kUnitGlobalScaleAmax = + transformer_engine::nvfp4::unit_global_scale_amax(); + float a_amax_val = kUnitGlobalScaleAmax; if (has_a_multi_tensor) { auto *a_amax_p = static_cast(a_multi_tensor_args.amax_ptrs[idx]); if (a_amax_p != nullptr) { a_amax_val = *a_amax_p; - has_a_amax = true; } } else if (a_amax != nullptr) { a_amax_val = a_amax[idx]; - has_a_amax = true; } - if (has_a_amax && b_amax && nvfp4_computed_alpha) { - constexpr float factor_inv = 1.0f / (6.0f * 6.0f * 448.0f * 448.0f); - nvfp4_computed_alpha[idx] = alpha_ptr[idx] * a_amax_val * b_amax[idx] * factor_inv; + if (nvfp4_computed_alpha != nullptr) { + constexpr float factor_inv = + 1.0f / (kUnitGlobalScaleAmax * kUnitGlobalScaleAmax); + const float b_amax_val = b_amax == nullptr ? kUnitGlobalScaleAmax : b_amax[idx]; + nvfp4_computed_alpha[idx] = alpha_ptr[idx] * a_amax_val * b_amax_val * factor_inv; alpha_ptrs[idx] = &nvfp4_computed_alpha[idx]; } else { alpha_ptrs[idx] = alpha_ptr + idx; @@ -1544,10 +1546,7 @@ inline void launch_grouped_gemm_setup( const bool b_rowwise = B_sel.rowwise; // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete). - const bool a_has_amax = (A_sel.amax != nullptr) || - (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr); - const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && - a_has_amax && (B_sel.amax != nullptr); + const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode); setup_grouped_gemm_kernel<<>>( ws.A_ptrs, ws.B_ptrs, ws.C_ptrs, ws.D_ptrs, ws.a_rows, ws.a_cols, ws.b_rows, ws.b_cols, @@ -1619,14 +1618,6 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT validate_nvfp4_grouped_gemm_support(A_sel, B_sel, use_per_group_alpha_beta); validate_fp8_block_grouped_gemm_support(A_sel, B_sel, sm); - // NVFP4 global-scale alpha requires per-tensor amax for both operands; without it - // the kernel silently drops the (amax_A * amax_B / factor) factor and produces - // numerically wrong output. - if (is_nvfp_scaling(A_sel.scaling_mode)) { - NVTE_CHECK(A_sel.amax != nullptr, "Grouped GEMM: NVFP4 A is missing amax."); - NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); - } - // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); @@ -1775,14 +1766,6 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num A_sel.dptr = nullptr; A_sel.amax = nullptr; - if (nvfp4) { - for (size_t i = 0; i < num_tensors; ++i) { - NVTE_CHECK(a_multi_tensor_args.amax_ptrs[i] != nullptr, "Grouped GEMM: NVFP4 A_list tensor ", - i, " is missing amax."); - } - NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); - } - // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); @@ -1861,12 +1844,6 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, validate_nvfp4_grouped_gemm_support(A_sel, B_sel, use_per_group_alpha_beta); validate_fp8_block_grouped_gemm_support(A_sel, B_sel, sm); - // NVFP4 global-scale alpha requires per-tensor amax for both operands. - if (is_nvfp_scaling(A_sel.scaling_mode)) { - NVTE_CHECK(A_sel.amax != nullptr, "Grouped GEMM: NVFP4 A is missing amax."); - NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); - } - // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index d5dbd0bd82..0e2202651c 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -694,7 +694,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g // g2s load all global_d_amax CUTLASS_PRAGMA_NO_UNROLL for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) { - shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g)); + shared_storage.global_d_amax[g] = + amax_colwise == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : __ldg(amax_colwise + g); } size_t rng_seed = 0; @@ -747,8 +750,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float c_global_amax_val = shared_storage.global_d_amax[group_idx]; float global_encode_scale = c_global_amax_val > 0.0f @@ -948,7 +953,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g // g2s load all global_a_amax for all groups/tensors CUTLASS_PRAGMA_NO_UNROLL for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) { - shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g)); + shared_storage.global_a_amax[g] = + amax_rowwise == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : __ldg(amax_rowwise + g); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { @@ -1004,8 +1012,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g packed_N, M, offsets); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float global_encode_scale = a_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu index e2325dd0fc..af76581183 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -85,8 +85,8 @@ __device__ __forceinline__ int GetTensorId(MultiAmaxHadamardCastFusionArgs *kern // calculate the global encode scale factor for a given global amax. __device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { - constexpr float kFP8E4M3Max = 448.0f; - constexpr float kFP4E2M1Max = 6.0f; + constexpr float kFP8E4M3Max = transformer_engine::detail::TypeExtrema::max; + constexpr float kFP4E2M1Max = transformer_engine::detail::TypeExtrema::max; // If scale is infinity, return max value of float32 float global_encode_scale = cutlass::minimum_with_nan_propagation{}( kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max()); @@ -471,7 +471,8 @@ __global__ static void group_rht_gemm_device( auto thr_r2g = tiled_r2g.get_slice(thread_idx); // NVFP4 non-E8 recipe constants and global scales - static constexpr float fp4_max = 6.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; // get global amax pointer @@ -508,7 +509,10 @@ __global__ static void group_rht_gemm_device( Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); - float global_amax_val = *global_amax_ptr; + constexpr float kUnitGlobalScaleAmax = + transformer_engine::nvfp4::unit_global_scale_amax(); + float global_amax_val = + global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); // Scaling factor for fast math path @@ -529,7 +533,8 @@ __global__ static void group_rht_gemm_device( // TODO(zhongbo): the math operations are very expensive // since the kernel is persistent, we can have a cache for all the possible scaling factors if (tensor_id != new_tensor_id) { - global_amax_val = *global_amax_ptr; + global_amax_val = + global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; global_decode_scale = 1.0f / global_encode_scale; diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu index cab0b38589..de9d6b919d 100644 --- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -682,8 +682,11 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( // g2s load all global_d_amax CUTLASS_PRAGMA_NO_UNROLL for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) { + const auto *amax_ptr = reinterpret_cast(args.global_d_amax_list[g]); shared_storage.global_d_amax[g] = - __ldg(reinterpret_cast(args.global_d_amax_list[g])); + amax_ptr == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : __ldg(amax_ptr); } size_t rng_seed = 0; @@ -729,8 +732,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float c_global_amax_val = shared_storage.global_d_amax[group_idx]; float global_encode_scale = c_global_amax_val > 0.0f @@ -926,8 +931,11 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( // g2s load all global_a_amax for all groups/tensors CUTLASS_PRAGMA_NO_UNROLL for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) { + const auto *amax_ptr = reinterpret_cast(args.global_a_amax_list[g]); shared_storage.global_a_amax[g] = - __ldg(reinterpret_cast(args.global_a_amax_list[g])); + amax_ptr == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : __ldg(amax_ptr); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { @@ -981,8 +989,10 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float global_encode_scale = a_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu index 50b9f63bdd..cfc5efef0d 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -44,8 +44,8 @@ using cute::Shape; // Avoid conflict with transformer_engine::Shape // calculate the global encode scale factor for a given global amax. __device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { - constexpr float kFP8E4M3Max = 448.0f; - constexpr float kFP4E2M1Max = 6.0f; + constexpr float kFP8E4M3Max = transformer_engine::detail::TypeExtrema::max; + constexpr float kFP4E2M1Max = transformer_engine::detail::TypeExtrema::max; // If scale is infinity, return max value of float32 float global_encode_scale = cutlass::minimum_with_nan_propagation{}( kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max()); @@ -278,7 +278,7 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, bool is_dma_warp = (warp_idx == 1); bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7); - if (is_epilogue_warp && elect_one_sync()) { + if (is_epilogue_warp && elect_one_sync() && global_amax != nullptr) { cute::prefetch(raw_pointer_cast(global_amax)); } @@ -414,7 +414,10 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); } else if (is_epilogue_warp) { - const float global_amax_val = *global_amax; + constexpr float kUnitGlobalScaleAmax = + transformer_engine::nvfp4::unit_global_scale_amax(); + const float global_amax_val = + global_amax == nullptr ? kUnitGlobalScaleAmax : *global_amax; static constexpr int FragmentSize = 256 / sizeof_bits_v; tmem_allocation_result_barrier.arrive_and_wait(); @@ -429,7 +432,8 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, auto thr_r2g = tiled_r2g.get_slice(thread_idx); // NVFP4 non-E8 recipe constants and global scales - static constexpr float fp4_max = 6.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; const float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); const float global_decode_scale = 1.0f / global_encode_scale; diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 479922a9bf..0a26e598b0 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -405,10 +405,10 @@ __global__ static void row_col_rht_gemm_device( bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); - if (is_epilogue_col_quant_warp && elect_one_sync()) { + if (is_epilogue_col_quant_warp && elect_one_sync() && c_global_amax != nullptr) { cute::prefetch(raw_pointer_cast(c_global_amax)); } - if (is_epilogue_row_quant_warp && elect_one_sync()) { + if (is_epilogue_row_quant_warp && elect_one_sync() && a_global_amax != nullptr) { cute::prefetch(raw_pointer_cast(a_global_amax)); } @@ -653,7 +653,10 @@ __global__ static void row_col_rht_gemm_device( if constexpr (kEnableRHTColQuant) { using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; - float const c_global_amax_val = *c_global_amax; + float const c_global_amax_val = + c_global_amax == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : *c_global_amax; auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{}); auto bulk_tmem_epilogue_layout = make_layout( acc_epilogue_pipelined_shape, @@ -710,8 +713,10 @@ __global__ static void row_col_rht_gemm_device( auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; float const fp4_max_inv = 1.0f / fp4_max; float const global_encode_scale = c_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( @@ -860,7 +865,10 @@ __global__ static void row_col_rht_gemm_device( cutlass::arch::warpgroup_reg_alloc<136>(); if constexpr (kEnableRowQuant) { using S2RVectorType = uint128_t; - float const a_global_amax_val = *a_global_amax; + float const a_global_amax_val = + a_global_amax == nullptr + ? transformer_engine::nvfp4::unit_global_scale_amax() + : *a_global_amax; int global_thread_idx = threadIdx.x; int local_thread_idx = global_thread_idx % 256; size_t rng_seed = 0; @@ -906,8 +914,10 @@ __global__ static void row_col_rht_gemm_device( cute::Tensor tQApSFA = thr_s2r.partition_D(pSFA_mn); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max = + transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = + transformer_engine::detail::TypeExtrema::max; float const fp4_max_inv = 1.0f / fp4_max; float const global_encode_scale = a_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 576e6139c7..3b353279b6 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -72,9 +72,11 @@ constexpr int kThreadsPerBlock = 256; __global__ void compute_nvfp4_per_tensor_scale_kernel(float alpha_in, const float *amax_A, const float *amax_B, float fp8_max_A, float fp8_max_B, float *alpha_out) { - constexpr float fp4_max = 6.0f; + constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B); - *alpha_out = alpha_in * (*amax_A) * (*amax_B) * factor_inv; + const float amax_A_value = amax_A == nullptr ? fp8_max_A * fp4_max : *amax_A; + const float amax_B_value = amax_B == nullptr ? fp8_max_B * fp4_max : *amax_B; + *alpha_out = alpha_in * amax_A_value * amax_B_value * factor_inv; } template @@ -616,9 +618,9 @@ void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_ro * NVFP4 COMPUTE PER-BLOCK DECODE SCALE KERNEL * * Computes per-block decode scale from block amax and global amax: - * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * global_scale = (max(E4M3) * max(E2M1)) / global_amax * per_block_decode_scale = block_amax * (global_scale * (1 / fp4_max)) - * = block_amax * 448 / global_amax + * = block_amax * max(E4M3) / global_amax * * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh * @@ -636,8 +638,8 @@ __global__ void nvfp4_compute_per_block_scale_kernel( const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= numel) return; - constexpr float fp4_max = 6.0f; - constexpr float fp8_max = 448.0f; + constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN @@ -665,8 +667,8 @@ __global__ void nvfp4_compute_global_scale_kernel( const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_params) return; - constexpr float fp4_max = 6.0f; - constexpr float fp8_max = 448.0f; + constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN @@ -757,8 +759,8 @@ __global__ void nvfp4_fused_scale_kernel( const size_t tile_row = out_row / block_len; // Compute the scale value - constexpr float fp4_max = 6.0f; - constexpr float fp8_max = 448.0f; + constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; @@ -927,9 +929,6 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r const float fp8_max_A = static_cast(tA->nvfp4_e4m3_max); const float fp8_max_B = static_cast(tB->nvfp4_e4m3_max); - // check for not null pointers - NVTE_CHECK(amax_A_ptr != nullptr, "amax_A_ptr is null"); - NVTE_CHECK(amax_B_ptr != nullptr, "amax_B_ptr is null"); NVTE_CHECK(alpha_ptr != nullptr, "alpha_ptr is null"); nvfp4_recipe::compute_nvfp4_per_tensor_scale_kernel<<<1, 1, 0, stream>>>( diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index d5f2fa9a2c..fb7f9a319e 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -415,7 +415,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int kNumThreadsReduce = kScaleBlockDim / kNVecOut; const float global_encode_scale = - kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]); + (kIsE8Scaling || global_amax == nullptr) ? 1.0f + : ComputeGlobalEncodeScaleFP4(global_amax[0]); constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; const float global_decode_scale = 1.0 / global_encode_scale; diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..499bfd6e99 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -354,6 +354,8 @@ class NVFP4Quantizer : public Quantizer { int nvfp4_e4m3_max; // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. bool row_scaled_nvfp4; + // Whether to use only block scaling by fixing the global encode scale to one. + bool disable_2d_scaling; int rht_matrix_random_sign_mask_t; at::Tensor rht_matrix; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..46ce166560 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -345,7 +345,8 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const "group_quantize: varying last dim is not supported with NVFP4."); NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, - nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), true); + nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), + !nvfp4_quantizer_cpp->disable_2d_scaling); break; } case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { @@ -998,6 +999,7 @@ std::tuple, std::vector, bool> bulk_alloc const bool nvfp4_use_4over6 = quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; + const bool disable_2d_scaling = quantizer_cpp_list[0]->disable_2d_scaling; const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); @@ -1028,6 +1030,9 @@ std::tuple, std::vector, bool> bulk_alloc "NVFP4 bulk allocation requires all quantizers in the group to share " "the same with_rht value (tensor 0=", group_with_rht, ", tensor ", i, "=", quantizer_cpp_list[i]->with_rht, ")."); + NVTE_CHECK(quantizer_cpp_list[i]->disable_2d_scaling == disable_2d_scaling, + "NVFP4 bulk allocation requires all quantizers in the group to share " + "the same disable_2d_scaling value."); } bool all_tensors_rht_cast_fusion_eligible = true; for (size_t i = 0; i < num_tensors; ++i) { @@ -1091,18 +1096,22 @@ std::tuple, std::vector, bool> bulk_alloc shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end()); dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); - for (size_t i = 0; i < num_tensors; ++i) { - shapes.emplace_back(amax_shape(rowwise_data_shapes[i], row_scaled_nvfp4)); + if (!disable_2d_scaling) { + for (size_t i = 0; i < num_tensors; ++i) { + shapes.emplace_back(amax_shape(rowwise_data_shapes[i], row_scaled_nvfp4)); + } + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); } - dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); - alignments.insert(alignments.end(), num_tensors, 16); auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); // Split data, scale, and amax tensors for (size_t i = 0; i < num_tensors; ++i) { rowwise_data_list.push_back(tensors[i]); rowwise_scale_list.push_back(tensors[num_tensors + i]); - amax_rowwise_list.push_back(tensors[2 * num_tensors + i]); + if (!disable_2d_scaling) { + amax_rowwise_list.push_back(tensors[2 * num_tensors + i]); + } } } @@ -1145,18 +1154,22 @@ std::tuple, std::vector, bool> bulk_alloc shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end()); dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); - for (size_t i = 0; i < num_tensors; ++i) { - shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); + if (!disable_2d_scaling) { + for (size_t i = 0; i < num_tensors; ++i) { + shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); + } + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); } - dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); - alignments.insert(alignments.end(), num_tensors, 16); auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); // Split data, scale, and amax tensors for (size_t i = 0; i < num_tensors; ++i) { columnwise_data_list.push_back(tensors[i]); columnwise_scale_list.push_back(tensors[num_tensors + i]); - amax_columnwise_list.push_back(tensors[2 * num_tensors + i]); + if (!disable_2d_scaling) { + amax_columnwise_list.push_back(tensors[2 * num_tensors + i]); + } } } @@ -1170,8 +1183,10 @@ std::tuple, std::vector, bool> bulk_alloc (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none()); py::object columnwise_scale = (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); - py::object amax_rowwise = rowwise_usage ? py::cast(amax_rowwise_list[i]) : py::none(); - py::object amax_columnwise = columnwise_usage ? py::cast(amax_columnwise_list[i]) : py::none(); + py::object amax_rowwise = + (rowwise_usage && !disable_2d_scaling) ? py::cast(amax_rowwise_list[i]) : py::none(); + py::object amax_columnwise = + (columnwise_usage && !disable_2d_scaling) ? py::cast(amax_columnwise_list[i]) : py::none(); // Construct Python tensor. tensor_py_list.emplace_back(NVFP4TensorClass( @@ -1200,11 +1215,11 @@ std::tuple, std::vector, bool> bulk_alloc tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max); // Set the amax rowwise and amax columnwise if available - if (rowwise_usage) { + if (rowwise_usage && !disable_2d_scaling) { tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise_list[i])); } - if (columnwise_usage) { + if (columnwise_usage && !disable_2d_scaling) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, std::vector{1}); } @@ -1387,7 +1402,9 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, need_separate_rng_states ? quant_config_list_colwise : quant_config_list; // Compute amaxes - if (quantizer.with_post_rht_amax) { + if (quantizer.disable_2d_scaling) { + // A null amax tells common NVFP4 kernels to use a unit global scale. + } else if (quantizer.with_post_rht_amax) { // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) @@ -1553,19 +1570,21 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, // Columnwise amax will be filled with a fused D2D copy from rowwise amax // Note that the multi compute amax API expects rowwise amax pointer to be not null // So we need to set the pointer accordingly to make colwise-only quantization work - std::vector orig_amax_ptr_list; - for (size_t i = 0; i < num_tensors; i++) { - auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; - orig_amax_ptr_list.push_back(rowwise_amax_ptr); - auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr; - void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; - NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); - output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector{1}); - } - nvte_group_amax(input.data(), reinterpret_cast(nvte_tensor_output_list.data()), - split_sections.data(), num_tensors, stream); - for (size_t i = 0; i < num_tensors; i++) { - output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector{1}); + if (!quantizer.disable_2d_scaling) { + std::vector orig_amax_ptr_list; + for (size_t i = 0; i < num_tensors; i++) { + auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; + orig_amax_ptr_list.push_back(rowwise_amax_ptr); + auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr; + void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; + NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); + output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector{1}); + } + nvte_group_amax(input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + split_sections.data(), num_tensors, stream); + for (size_t i = 0; i < num_tensors; i++) { + output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector{1}); + } } // Quantize tensors individually diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4fff3f92de..146f4534e7 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1876,6 +1876,7 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize NVTE_ERROR("Unsupported NVFP4 4over6 error mode: ", nvfp4_4over6_err_mode); } this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); + this->disable_2d_scaling = quantizer.attr("disable_2d_scaling").cast(); // Get amax reduction group if needed for NVFP4 AG const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -1960,6 +1961,7 @@ std::pair NVFP4Quantizer::create_tensor( "NVFP4 requires tensor dims that are divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool disable_2d_scaling = this->disable_2d_scaling; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -1983,7 +1985,9 @@ std::pair NVFP4Quantizer::create_tensor( const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); + if (!disable_2d_scaling) { + amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); + } } if (columnwise_usage) { const std::vector scale_inv_shape_int64(columnwise_scale_inv_shape.begin(), @@ -1999,7 +2003,9 @@ std::pair NVFP4Quantizer::create_tensor( // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; - amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); + if (!disable_2d_scaling) { + amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); + } } // Convert tensors to Python @@ -2010,8 +2016,9 @@ std::pair NVFP4Quantizer::create_tensor( auto rowwise_scale_inv_py = py_cast(rowwise_scale_inv_tensor, rowwise_usage); auto columnwise_data_py = py_cast(columnwise_data_tensor, columnwise_usage); auto columnwise_scale_inv_py = py_cast(columnwise_scale_inv_tensor, columnwise_usage); - auto amax_rowwise_py = py_cast(amax_rowwise, rowwise_usage); - auto amax_columnwise_py = py_cast(amax_columnwise, columnwise_usage); + auto amax_rowwise_py = py_cast(amax_rowwise, rowwise_usage && !disable_2d_scaling); + auto amax_columnwise_py = + py_cast(amax_columnwise, columnwise_usage && !disable_2d_scaling); // Construct Python NVFP4 tensor py::object out_py; @@ -2079,7 +2086,9 @@ std::pair NVFP4Quantizer::create_tensor( out_cpp.set_rowwise_data(rowwise_data_tensor.data_ptr(), DType::kFloat4E2M1, shape); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, rowwise_scale_inv_shape); - out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + if (!disable_2d_scaling) { + out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + } } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -2090,8 +2099,10 @@ std::pair NVFP4Quantizer::create_tensor( col_data_shape_fp4); out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, columnwise_scale_inv_shape); - out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - getTensorShape(amax_columnwise)); + if (!disable_2d_scaling) { + out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, + getTensorShape(amax_columnwise)); + } } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); @@ -2127,6 +2138,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso std::optional columnwise_amax; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool disable_2d_scaling = this->disable_2d_scaling; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -2144,7 +2156,9 @@ std::pair NVFP4Quantizer::create_grouped_tenso rowwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); const int64_t amax_elements = row_scaled_nvfp4 ? static_cast(logical_first_dim) : static_cast(num_tensors); - rowwise_amax = at::empty({amax_elements}, float_opts); + if (!disable_2d_scaling) { + rowwise_amax = at::empty({amax_elements}, float_opts); + } } if (columnwise_usage) { @@ -2152,7 +2166,9 @@ std::pair NVFP4Quantizer::create_grouped_tenso const auto scale_shape = get_scale_shape(logical_shape_vec, true); const int64_t total_scale_elements = static_cast(product(scale_shape)); columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); - columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + if (!disable_2d_scaling) { + columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + } } GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); @@ -2160,15 +2176,19 @@ std::pair NVFP4Quantizer::create_grouped_tenso out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*rowwise_scale_inv)); - out_cpp.set_amax(rowwise_amax->data_ptr(), DType::kFloat32, getTensorShape(*rowwise_amax)); + if (rowwise_amax.has_value()) { + out_cpp.set_amax(rowwise_amax->data_ptr(), DType::kFloat32, getTensorShape(*rowwise_amax)); + } } if (columnwise_usage) { out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, getTensorShape(*columnwise_data)); out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*columnwise_scale_inv)); - out_cpp.set_columnwise_amax(columnwise_amax->data_ptr(), DType::kFloat32, - getTensorShape(*columnwise_amax)); + if (columnwise_amax.has_value()) { + out_cpp.set_columnwise_amax(columnwise_amax->data_ptr(), DType::kFloat32, + getTensorShape(*columnwise_amax)); + } } if (first_dims.has_value()) { out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); @@ -2286,6 +2306,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const bool with_gemm_swizzled_scales = nvfp4_emits_gemm_swizzled_scales(*this, shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool disable_2d_scaling = this->disable_2d_scaling; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -2313,7 +2334,10 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( tensor.attr("_rowwise_scale_inv") = *rowwise_scale_inv; } const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; - if (!amax_rowwise || amax_rowwise->numel() != amax_rows) { + if (disable_2d_scaling) { + amax_rowwise.reset(); + tensor.attr("_amax_rowwise") = py::none(); + } else if (!amax_rowwise || amax_rowwise->numel() != amax_rows) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed @@ -2356,7 +2380,10 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( tensor.attr("_columnwise_scale_inv") = *columnwise_scale_inv; } const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; - if (!amax_columnwise || amax_columnwise->numel() != amax_cols) { + if (disable_2d_scaling) { + amax_columnwise.reset(); + tensor.attr("_amax_columnwise") = py::none(); + } else if (!amax_columnwise || amax_columnwise->numel() != amax_cols) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed @@ -2384,7 +2411,9 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_rowwise_data(rowwise_data->data_ptr(), DType::kFloat4E2M1, shape); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*rowwise_scale_inv)); - out_cpp.set_amax(amax_rowwise->data_ptr(), DType::kFloat32, getTensorShape(*amax_rowwise)); + if (amax_rowwise.has_value()) { + out_cpp.set_amax(amax_rowwise->data_ptr(), DType::kFloat32, getTensorShape(*amax_rowwise)); + } } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -2395,8 +2424,10 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( col_data_shape_fp4); out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*columnwise_scale_inv)); - out_cpp.set_columnwise_amax(amax_columnwise->data_ptr(), DType::kFloat32, - getTensorShape(*amax_columnwise)); + if (amax_columnwise.has_value()) { + out_cpp.set_columnwise_amax(amax_columnwise->data_ptr(), DType::kFloat32, + getTensorShape(*amax_columnwise)); + } } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); @@ -2484,7 +2515,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou const std::optional& noop_flag, bool compute_amax) { auto reduce_amaxes = [&]() { - if (!this->with_amax_reduction) { + if (!this->with_amax_reduction || this->disable_2d_scaling) { return; } @@ -2611,7 +2642,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) - if (compute_amax) { + if (compute_amax && !this->disable_2d_scaling) { NVTE_SCOPED_GIL_RELEASE({ nvte_hadamard_transform_amax(input.data(), out.data(), 0, this->rht_matrix_random_sign_mask_t, stream); @@ -2624,7 +2655,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou "Use with_post_rht_amax=true instead."); } } else { // Without RHT - if (compute_amax && !row_scaled_nvfp4) { + if (compute_amax && !row_scaled_nvfp4 && !this->disable_2d_scaling) { // Amax pointers auto rowwise_amax_ptr = out.get_amax().data_ptr; auto columnwise_amax_ptr = out.get_columnwise_amax().data_ptr; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index ddb85808a5..6620619ca4 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -143,24 +143,29 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) if (rowwise_usage) { const auto &data = tensor.attr("_rowwise_data").cast(); const auto &scale_inv = tensor.attr("_rowwise_scale_inv").cast(); - const auto &amax_rowwise = tensor.attr("_amax_rowwise").cast(); ret.set_rowwise_data(data.data_ptr(), dtype, convert_shape_back_from_fp4(getTensorShape(data), false)); ret.set_rowwise_scale_inv(scale_inv.data_ptr(), DType::kFloat8E4M3, getTensorShape(scale_inv)); - ret.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + const auto amax_rowwise = tensor.attr("_amax_rowwise"); + if (!amax_rowwise.is_none()) { + const auto &amax = amax_rowwise.cast(); + ret.set_amax(amax.data_ptr(), DType::kFloat32, getTensorShape(amax)); + } } // Column-scaled data if (columnwise_usage) { const auto &data = tensor.attr("_columnwise_data").cast(); const auto &scale_inv = tensor.attr("_columnwise_scale_inv").cast(); - const auto &amax_columnwise = tensor.attr("_amax_columnwise").cast(); ret.set_columnwise_data(data.data_ptr(), DType::kFloat4E2M1, convert_shape_back_from_fp4(getTensorShape(data), false)); ret.set_columnwise_scale_inv(scale_inv.data_ptr(), DType::kFloat8E4M3, getTensorShape(scale_inv)); - ret.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - getTensorShape(amax_columnwise)); + const auto amax_columnwise = tensor.attr("_amax_columnwise"); + if (!amax_columnwise.is_none()) { + const auto &amax = amax_columnwise.cast(); + ret.set_columnwise_amax(amax.data_ptr(), DType::kFloat32, getTensorShape(amax)); + } } // Scale layout diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 4614ea53a6..4881975c54 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -135,6 +135,8 @@ class NVFP4Quantizer(Quantizer): nvfp4_e4m3_max: int """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str + """Whether to disable the global (second-level) NVFP4 scale.""" + disable_2d_scaling: bool """RHT sign mask (0 when sign randomization is disabled)""" rht_matrix_random_sign_mask_t: int @@ -155,6 +157,7 @@ def __init__( nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", with_random_sign_mask: bool = True, + disable_2d_scaling: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = DType.cast(fp4_dtype) @@ -172,6 +175,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") + self.disable_2d_scaling = disable_2d_scaling self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -244,6 +248,7 @@ def copy(self) -> NVFP4Quantizer: nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, with_random_sign_mask=self.rht_matrix_random_sign_mask_t != 0, + disable_2d_scaling=self.disable_2d_scaling, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm @@ -768,7 +773,11 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): rowwise_scale_inv = scale_inv_init_func( tensor._rowwise_scale_inv, *args[1:], **kwargs ) - amax_rowwise = torch.zeros_like(tensor._amax_rowwise, *args[1:], **kwargs) + amax_rowwise = ( + None + if tensor._amax_rowwise is None + else torch.zeros_like(tensor._amax_rowwise, *args[1:], **kwargs) + ) else: rowwise_data, rowwise_scale_inv, amax_rowwise = None, None, None @@ -777,7 +786,11 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): columnwise_scale_inv = scale_inv_init_func( tensor._columnwise_scale_inv, *args[1:], **kwargs ) - amax_columnwise = torch.zeros_like(tensor._amax_columnwise, *args[1:], **kwargs) + amax_columnwise = ( + None + if tensor._amax_columnwise is None + else torch.zeros_like(tensor._amax_columnwise, *args[1:], **kwargs) + ) else: columnwise_data, columnwise_scale_inv, amax_columnwise = ( None, diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 09f040ba67..81d83f52fa 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -90,10 +90,10 @@ class NVFP4TensorStorage(QuantizedTensorStorage): _columnwise_scale_inv: torch.Tensor # Input absolute maximum value (used to compute tensor scale for # row-scaled FP4 data) - _amax_rowwise: torch.Tensor + _amax_rowwise: Optional[torch.Tensor] # Input absolute maximum value (used to compute tensor scale for # column-scaled FP4 data) - _amax_columnwise: torch.Tensor + _amax_columnwise: Optional[torch.Tensor] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] @@ -370,13 +370,16 @@ def update_usage( rowwise_usage = self._rowwise_data is not None if columnwise_usage is None: columnwise_usage = self._columnwise_data is not None + requires_amax = not ( + self._quantizer is not None and self._quantizer.disable_2d_scaling + ) # If both rowwise and columnwise are requested, create columnwise from rowwise if needed if rowwise_usage and columnwise_usage: if ( self._rowwise_data is None or self._rowwise_scale_inv is None - or self._amax_rowwise is None + or (requires_amax and self._amax_rowwise is None) ): raise RuntimeError( "Cannot update to rowwise and columnwise usage because rowwise data is None." @@ -395,7 +398,7 @@ def update_usage( raise RuntimeError( "Requested row-wise usage, but NVFP4Tensor is missing row-scaled scale-inverses" ) - if self._amax_rowwise is None: + if requires_amax and self._amax_rowwise is None: raise RuntimeError( "Requested row-wise usage, but NVFP4Tensor is missing per tensor" " row-scaled scale-inverse" @@ -416,7 +419,7 @@ def update_usage( "Requested column-wise usage, " "but NVFP4Tensor is missing column-scaled scale-inverses" ) - if self._amax_columnwise is None: + if requires_amax and self._amax_columnwise is None: raise RuntimeError( "Requested column-wise usage, " "but NVFP4Tensor is missing per tensor column-scaled scale-inverse" @@ -479,7 +482,9 @@ def _create_columnwise(self): K_tiles, ) - # Also set columnwise amax (same as rowwise since it's just transposed data) - if self._amax_columnwise is None: - self._amax_columnwise = torch.empty_like(self._amax_rowwise) - self._amax_columnwise.copy_(self._amax_rowwise) + # Also set columnwise amax (same as rowwise since it's just transposed data). + # A missing amax represents unit global scaling. + if not self._quantizer.disable_2d_scaling: + if self._amax_columnwise is None: + self._amax_columnwise = torch.empty_like(self._amax_rowwise) + self._amax_columnwise.copy_(self._amax_rowwise) From 96ce61a4bf1cbe84c460e9e9978ab1c5146070df Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 31 Jul 2026 19:16:03 +0000 Subject: [PATCH 2/3] better nameis second level scale, and error out row scaled nvfp4 Signed-off-by: Varun Thumbe --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 4 +- .../test_nvfp4_group_quantize_graph_safe.py | 6 ++- .../nvfp4/test_nvfp4_quantize_exact.py | 38 ++++++++++--------- .../common/cast/dispatch/quantize.cuh | 4 ++ .../common/cast/nvfp4/dequantize_nvfp4.cuh | 4 +- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 2 + .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 2 + .../quantize_transpose_nvfp4_tuned_1D.cuh | 5 +++ transformer_engine/pytorch/csrc/common.h | 2 +- .../pytorch/csrc/extensions/cast.cpp | 33 +++++++++------- transformer_engine/pytorch/csrc/quantizer.cpp | 36 +++++++++--------- .../pytorch/tensor/nvfp4_tensor.py | 16 ++++++-- .../tensor/storage/nvfp4_tensor_storage.py | 4 +- 13 files changed, 95 insertions(+), 61 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index b836399480..fb9c09b933 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -32,11 +32,11 @@ def test_gemm_with_missing_nvfp4_amax(disable_x: bool, disable_w: bool) -> None: x[0, 0] = unit_scale_amax w[0, 0] = unit_scale_amax - def quantize(tensor: torch.Tensor, disable_2d_scaling: bool): + def quantize(tensor: torch.Tensor, disable_second_level_scale: bool): return NVFP4Quantizer( rowwise=True, columnwise=True, - disable_2d_scaling=disable_2d_scaling, + disable_second_level_scale=disable_second_level_scale, )(tensor) x_ref, w_ref = quantize(x, False), quantize(w, False) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index 9b945f632e..42c0fb2ee9 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -45,7 +45,9 @@ def fused_grouped_quantize( @pytest.mark.parametrize( "return_transpose", [False, True], ids=["rowwise", "rowwise_and_columnwise"] ) -def test_grouped_disable_2d_scaling_matches_split_quantize(return_transpose: bool) -> None: +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") @@ -56,7 +58,7 @@ def test_grouped_disable_2d_scaling_matches_split_quantize(return_transpose: boo columnwise=return_transpose, with_rht=True, with_post_rht_amax=True, - disable_2d_scaling=True, + disable_second_level_scale=True, ) grouped = fused_grouped_quantize(x, split_section_tensor, quantizer) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 036fce5608..41fe318d8f 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -242,38 +242,25 @@ def check_quantization_nvfp4_versus_reference( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "return_transpose,row_scaled_nvfp4", - [ - pytest.param(False, False, id="rowwise"), - pytest.param(True, False, id="rowwise_and_columnwise"), - pytest.param(False, True, id="row_scaled"), - ], -) +@pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["standard", "4over6"]) -def test_disable_2d_scaling_uses_only_block_scale( +def test_disable_second_level_scale_uses_only_block_scale( return_transpose: bool, - row_scaled_nvfp4: 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") - if row_scaled_nvfp4: - # The reference quantizer computes one global scale per row. - x[:, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE - else: - x[0, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE + x[0, 0] = NVFP4_AMAX_FOR_UNIT_GLOBAL_SCALE common_kwargs = { "rowwise": True, "columnwise": return_transpose, "with_rht": False, - "row_scaled_nvfp4": row_scaled_nvfp4, "nvfp4_use_4over6": use_4over6, } expected = NVFP4Quantizer(**common_kwargs)(x) - actual = NVFP4Quantizer(**common_kwargs, disable_2d_scaling=True)(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( @@ -292,12 +279,27 @@ def test_disable_2d_scaling_uses_only_block_scale( # 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_2d_scaling=True).update_quantized(x / 2, expected) + 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", diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 44c122540f..cca0709e7e 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/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 " @@ -291,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 " diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index eac28796d7..ad98a2f066 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -109,7 +109,9 @@ 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(4)); - NVTE_CHECK(!row_scaled_nvfp4 || input.amax.dptr == nullptr || input.amax.numel() == N, + 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_TYPE_SWITCH_NON_FP8ONLY( diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 638c1bd08b..6eacb706cc 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -692,6 +692,8 @@ 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(), diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 7e5722a5b0..68d04e8e9f 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1463,6 +1463,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, 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 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 diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 32396911f4..5cdeb2ca02 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -714,11 +714,16 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(output->has_data(), "NVFP4 output tensor must be allocated."); 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 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 does not support disabling " + "second-level scaling."); } const auto [rows, cols] = input.flat_2d_dims(); diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 499bfd6e99..6d7fa32fac 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -355,7 +355,7 @@ class NVFP4Quantizer : public Quantizer { // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. bool row_scaled_nvfp4; // Whether to use only block scaling by fixing the global encode scale to one. - bool disable_2d_scaling; + bool disable_second_level_scale; int rht_matrix_random_sign_mask_t; at::Tensor rht_matrix; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 46ce166560..5f01432688 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -346,7 +346,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), - !nvfp4_quantizer_cpp->disable_2d_scaling); + !nvfp4_quantizer_cpp->disable_second_level_scale); break; } case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { @@ -999,7 +999,8 @@ std::tuple, std::vector, bool> bulk_alloc const bool nvfp4_use_4over6 = quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; - const bool disable_2d_scaling = quantizer_cpp_list[0]->disable_2d_scaling; + const bool disable_second_level_scale = + quantizer_cpp_list[0]->disable_second_level_scale; const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); @@ -1030,9 +1031,10 @@ std::tuple, std::vector, bool> bulk_alloc "NVFP4 bulk allocation requires all quantizers in the group to share " "the same with_rht value (tensor 0=", group_with_rht, ", tensor ", i, "=", quantizer_cpp_list[i]->with_rht, ")."); - NVTE_CHECK(quantizer_cpp_list[i]->disable_2d_scaling == disable_2d_scaling, + NVTE_CHECK(quantizer_cpp_list[i]->disable_second_level_scale == + disable_second_level_scale, "NVFP4 bulk allocation requires all quantizers in the group to share " - "the same disable_2d_scaling value."); + "the same disable_second_level_scale value."); } bool all_tensors_rht_cast_fusion_eligible = true; for (size_t i = 0; i < num_tensors; ++i) { @@ -1096,7 +1098,7 @@ std::tuple, std::vector, bool> bulk_alloc shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end()); dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { for (size_t i = 0; i < num_tensors; ++i) { shapes.emplace_back(amax_shape(rowwise_data_shapes[i], row_scaled_nvfp4)); } @@ -1109,7 +1111,7 @@ std::tuple, std::vector, bool> bulk_alloc for (size_t i = 0; i < num_tensors; ++i) { rowwise_data_list.push_back(tensors[i]); rowwise_scale_list.push_back(tensors[num_tensors + i]); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { amax_rowwise_list.push_back(tensors[2 * num_tensors + i]); } } @@ -1154,7 +1156,7 @@ std::tuple, std::vector, bool> bulk_alloc shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end()); dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); alignments.insert(alignments.end(), num_tensors, 16); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { for (size_t i = 0; i < num_tensors; ++i) { shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); } @@ -1167,7 +1169,7 @@ std::tuple, std::vector, bool> bulk_alloc for (size_t i = 0; i < num_tensors; ++i) { columnwise_data_list.push_back(tensors[i]); columnwise_scale_list.push_back(tensors[num_tensors + i]); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { amax_columnwise_list.push_back(tensors[2 * num_tensors + i]); } } @@ -1184,9 +1186,12 @@ std::tuple, std::vector, bool> bulk_alloc py::object columnwise_scale = (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); py::object amax_rowwise = - (rowwise_usage && !disable_2d_scaling) ? py::cast(amax_rowwise_list[i]) : py::none(); + (rowwise_usage && !disable_second_level_scale) ? py::cast(amax_rowwise_list[i]) + : py::none(); py::object amax_columnwise = - (columnwise_usage && !disable_2d_scaling) ? py::cast(amax_columnwise_list[i]) : py::none(); + (columnwise_usage && !disable_second_level_scale) + ? py::cast(amax_columnwise_list[i]) + : py::none(); // Construct Python tensor. tensor_py_list.emplace_back(NVFP4TensorClass( @@ -1215,11 +1220,11 @@ std::tuple, std::vector, bool> bulk_alloc tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max); // Set the amax rowwise and amax columnwise if available - if (rowwise_usage && !disable_2d_scaling) { + if (rowwise_usage && !disable_second_level_scale) { tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise_list[i])); } - if (columnwise_usage && !disable_2d_scaling) { + if (columnwise_usage && !disable_second_level_scale) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, std::vector{1}); } @@ -1402,7 +1407,7 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, need_separate_rng_states ? quant_config_list_colwise : quant_config_list; // Compute amaxes - if (quantizer.disable_2d_scaling) { + if (quantizer.disable_second_level_scale) { // A null amax tells common NVFP4 kernels to use a unit global scale. } else if (quantizer.with_post_rht_amax) { // We need: @@ -1570,7 +1575,7 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, // Columnwise amax will be filled with a fused D2D copy from rowwise amax // Note that the multi compute amax API expects rowwise amax pointer to be not null // So we need to set the pointer accordingly to make colwise-only quantization work - if (!quantizer.disable_2d_scaling) { + if (!quantizer.disable_second_level_scale) { std::vector orig_amax_ptr_list; for (size_t i = 0; i < num_tensors; i++) { auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 146f4534e7..d4e7b43e35 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1876,7 +1876,8 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize NVTE_ERROR("Unsupported NVFP4 4over6 error mode: ", nvfp4_4over6_err_mode); } this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); - this->disable_2d_scaling = quantizer.attr("disable_2d_scaling").cast(); + this->disable_second_level_scale = + quantizer.attr("disable_second_level_scale").cast(); // Get amax reduction group if needed for NVFP4 AG const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -1961,7 +1962,7 @@ std::pair NVFP4Quantizer::create_tensor( "NVFP4 requires tensor dims that are divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; - const bool disable_2d_scaling = this->disable_2d_scaling; + const bool disable_second_level_scale = this->disable_second_level_scale; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -1985,7 +1986,7 @@ std::pair NVFP4Quantizer::create_tensor( const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); } } @@ -2003,7 +2004,7 @@ std::pair NVFP4Quantizer::create_tensor( // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); } } @@ -2016,9 +2017,10 @@ std::pair NVFP4Quantizer::create_tensor( auto rowwise_scale_inv_py = py_cast(rowwise_scale_inv_tensor, rowwise_usage); auto columnwise_data_py = py_cast(columnwise_data_tensor, columnwise_usage); auto columnwise_scale_inv_py = py_cast(columnwise_scale_inv_tensor, columnwise_usage); - auto amax_rowwise_py = py_cast(amax_rowwise, rowwise_usage && !disable_2d_scaling); + auto amax_rowwise_py = + py_cast(amax_rowwise, rowwise_usage && !disable_second_level_scale); auto amax_columnwise_py = - py_cast(amax_columnwise, columnwise_usage && !disable_2d_scaling); + py_cast(amax_columnwise, columnwise_usage && !disable_second_level_scale); // Construct Python NVFP4 tensor py::object out_py; @@ -2086,7 +2088,7 @@ std::pair NVFP4Quantizer::create_tensor( out_cpp.set_rowwise_data(rowwise_data_tensor.data_ptr(), DType::kFloat4E2M1, shape); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, rowwise_scale_inv_shape); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); } } @@ -2099,7 +2101,7 @@ std::pair NVFP4Quantizer::create_tensor( col_data_shape_fp4); out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, columnwise_scale_inv_shape); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, getTensorShape(amax_columnwise)); } @@ -2138,7 +2140,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso std::optional columnwise_amax; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; - const bool disable_2d_scaling = this->disable_2d_scaling; + const bool disable_second_level_scale = this->disable_second_level_scale; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -2156,7 +2158,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso rowwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); const int64_t amax_elements = row_scaled_nvfp4 ? static_cast(logical_first_dim) : static_cast(num_tensors); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { rowwise_amax = at::empty({amax_elements}, float_opts); } } @@ -2166,7 +2168,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso const auto scale_shape = get_scale_shape(logical_shape_vec, true); const int64_t total_scale_elements = static_cast(product(scale_shape)); columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); - if (!disable_2d_scaling) { + if (!disable_second_level_scale) { columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); } } @@ -2306,7 +2308,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const bool with_gemm_swizzled_scales = nvfp4_emits_gemm_swizzled_scales(*this, shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; - const bool disable_2d_scaling = this->disable_2d_scaling; + const bool disable_second_level_scale = this->disable_second_level_scale; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { @@ -2334,7 +2336,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( tensor.attr("_rowwise_scale_inv") = *rowwise_scale_inv; } const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; - if (disable_2d_scaling) { + if (disable_second_level_scale) { amax_rowwise.reset(); tensor.attr("_amax_rowwise") = py::none(); } else if (!amax_rowwise || amax_rowwise->numel() != amax_rows) { @@ -2380,7 +2382,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( tensor.attr("_columnwise_scale_inv") = *columnwise_scale_inv; } const int64_t amax_cols = row_scaled_nvfp4 ? static_cast(flat_last_dim) : 1; - if (disable_2d_scaling) { + if (disable_second_level_scale) { amax_columnwise.reset(); tensor.attr("_amax_columnwise") = py::none(); } else if (!amax_columnwise || amax_columnwise->numel() != amax_cols) { @@ -2515,7 +2517,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou const std::optional& noop_flag, bool compute_amax) { auto reduce_amaxes = [&]() { - if (!this->with_amax_reduction || this->disable_2d_scaling) { + if (!this->with_amax_reduction || this->disable_second_level_scale) { return; } @@ -2642,7 +2644,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) - if (compute_amax && !this->disable_2d_scaling) { + if (compute_amax && !this->disable_second_level_scale) { NVTE_SCOPED_GIL_RELEASE({ nvte_hadamard_transform_amax(input.data(), out.data(), 0, this->rht_matrix_random_sign_mask_t, stream); @@ -2655,7 +2657,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou "Use with_post_rht_amax=true instead."); } } else { // Without RHT - if (compute_amax && !row_scaled_nvfp4 && !this->disable_2d_scaling) { + if (compute_amax && !row_scaled_nvfp4 && !this->disable_second_level_scale) { // Amax pointers auto rowwise_amax_ptr = out.get_amax().data_ptr; auto columnwise_amax_ptr = out.get_columnwise_amax().data_ptr; diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 4881975c54..72e4e574c4 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -136,7 +136,7 @@ class NVFP4Quantizer(Quantizer): """NVFP4 4over6 candidate-selection error mode.""" nvfp4_4over6_err_mode: str """Whether to disable the global (second-level) NVFP4 scale.""" - disable_2d_scaling: bool + disable_second_level_scale: bool """RHT sign mask (0 when sign randomization is disabled)""" rht_matrix_random_sign_mask_t: int @@ -157,7 +157,7 @@ def __init__( nvfp4_e4m3_max: int = 448, nvfp4_4over6_err_mode: str = "MAE", with_random_sign_mask: bool = True, - disable_2d_scaling: bool = False, + disable_second_level_scale: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = DType.cast(fp4_dtype) @@ -167,6 +167,14 @@ def __init__( self.amax_reduction_group = amax_reduction_group self.with_2d_quantization = with_2d_quantization self.stochastic_rounding = stochastic_rounding + if row_scaled_nvfp4 and disable_second_level_scale: + warnings.warn( + "Row-scaled NVFP4 requires second-level scaling; disabling " + "row_scaled_nvfp4 because disable_second_level_scale=True.", + UserWarning, + stacklevel=2, + ) + row_scaled_nvfp4 = False self.row_scaled_nvfp4 = row_scaled_nvfp4 self.nvfp4_use_4over6 = nvfp4_use_4over6 self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 @@ -175,7 +183,7 @@ def __init__( self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") - self.disable_2d_scaling = disable_2d_scaling + self.disable_second_level_scale = disable_second_level_scale self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -248,7 +256,7 @@ def copy(self) -> NVFP4Quantizer: nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, with_random_sign_mask=self.rht_matrix_random_sign_mask_t != 0, - disable_2d_scaling=self.disable_2d_scaling, + disable_second_level_scale=self.disable_second_level_scale, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 81d83f52fa..4e33b2dd07 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -371,7 +371,7 @@ def update_usage( if columnwise_usage is None: columnwise_usage = self._columnwise_data is not None requires_amax = not ( - self._quantizer is not None and self._quantizer.disable_2d_scaling + self._quantizer is not None and self._quantizer.disable_second_level_scale ) # If both rowwise and columnwise are requested, create columnwise from rowwise if needed @@ -484,7 +484,7 @@ def _create_columnwise(self): # Also set columnwise amax (same as rowwise since it's just transposed data). # A missing amax represents unit global scaling. - if not self._quantizer.disable_2d_scaling: + if not self._quantizer.disable_second_level_scale: if self._amax_columnwise is None: self._amax_columnwise = torch.empty_like(self._amax_rowwise) self._amax_columnwise.copy_(self._amax_rowwise) From 454af70fe658f19489075460faa8b9336aab0f89 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:17:45 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../test_nvfp4_group_quantize_graph_safe.py | 4 +--- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 8 +++----- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 7 +++---- .../quantize_transpose_nvfp4_tuned_1D.cuh | 8 ++++---- .../common/gemm/cublaslt_grouped_gemm.cu | 3 +-- ...cast_col_hadamard_transform_cast_fusion.cu | 12 ++++-------- .../group_hadamard_transform_cast_fusion.cu | 9 +++------ ...cast_col_hadamard_transform_cast_fusion.cu | 12 ++++-------- ...quantize_transpose_vector_blockwise_fp4.cu | 3 +-- .../pytorch/csrc/extensions/cast.cpp | 19 ++++++++----------- transformer_engine/pytorch/csrc/quantizer.cpp | 6 ++---- 11 files changed, 34 insertions(+), 57 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index 42c0fb2ee9..7128bd07ef 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -63,9 +63,7 @@ def test_grouped_disable_second_level_scale_matches_split_quantize( 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] - ) + 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( diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 6eacb706cc..1689fed49f 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -476,8 +476,7 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf block_amax = reduce_group_max_16(group_amax); } - float global_amax = - transformer_engine::nvfp4::unit_global_scale_amax(); + float global_amax = transformer_engine::nvfp4::unit_global_scale_amax(); if (amax != nullptr) { global_amax = amax[0]; } @@ -534,9 +533,8 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, } const float global_amax = - amax == nullptr - ? transformer_engine::nvfp4::unit_global_scale_amax() - : amax[0]; + amax == nullptr ? transformer_engine::nvfp4::unit_global_scale_amax() + : amax[0]; const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 68d04e8e9f..698a755a7a 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -724,10 +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 - ? (amax_rowwise_ptr == nullptr - ? 1.0f - : 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 = diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 5cdeb2ca02..f575f1752b 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -232,10 +232,10 @@ __device__ __forceinline__ void colwise_scaling( float S_enc_colwise_block = S_enc_colwise; 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 && amax_colwise_ptr != nullptr - ? core::compute_global_encode_scaling_factor_FP4( - amax_colwise_ptr[col_idx]) - : 1.0f; + S_enc_colwise_block = + 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); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 67fbc82686..16d688a79b 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1414,8 +1414,7 @@ __global__ void setup_grouped_gemm_kernel( a_amax_val = a_amax[idx]; } if (nvfp4_computed_alpha != nullptr) { - constexpr float factor_inv = - 1.0f / (kUnitGlobalScaleAmax * kUnitGlobalScaleAmax); + constexpr float factor_inv = 1.0f / (kUnitGlobalScaleAmax * kUnitGlobalScaleAmax); const float b_amax_val = b_amax == nullptr ? kUnitGlobalScaleAmax : b_amax[idx]; nvfp4_computed_alpha[idx] = alpha_ptr[idx] * a_amax_val * b_amax_val * factor_inv; alpha_ptrs[idx] = &nvfp4_computed_alpha[idx]; diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index 0e2202651c..a09459e24a 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -750,10 +750,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = - transformer_engine::detail::TypeExtrema::max; - static constexpr float fp8_max = - transformer_engine::detail::TypeExtrema::max; + static constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float c_global_amax_val = shared_storage.global_d_amax[group_idx]; float global_encode_scale = c_global_amax_val > 0.0f @@ -1012,10 +1010,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g packed_N, M, offsets); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = - transformer_engine::detail::TypeExtrema::max; - static constexpr float fp8_max = - transformer_engine::detail::TypeExtrema::max; + static constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float global_encode_scale = a_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu index af76581183..84a5639b35 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -471,8 +471,7 @@ __global__ static void group_rht_gemm_device( auto thr_r2g = tiled_r2g.get_slice(thread_idx); // NVFP4 non-E8 recipe constants and global scales - static constexpr float fp4_max = - transformer_engine::detail::TypeExtrema::max; + static constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; // get global amax pointer @@ -511,8 +510,7 @@ __global__ static void group_rht_gemm_device( constexpr float kUnitGlobalScaleAmax = transformer_engine::nvfp4::unit_global_scale_amax(); - float global_amax_val = - global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; + float global_amax_val = global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); // Scaling factor for fast math path @@ -533,8 +531,7 @@ __global__ static void group_rht_gemm_device( // TODO(zhongbo): the math operations are very expensive // since the kernel is persistent, we can have a cache for all the possible scaling factors if (tensor_id != new_tensor_id) { - global_amax_val = - global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; + global_amax_val = global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; global_decode_scale = 1.0f / global_encode_scale; diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu index de9d6b919d..62c7d4526a 100644 --- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -732,10 +732,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = - transformer_engine::detail::TypeExtrema::max; - static constexpr float fp8_max = - transformer_engine::detail::TypeExtrema::max; + static constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float c_global_amax_val = shared_storage.global_d_amax[group_idx]; float global_encode_scale = c_global_amax_val > 0.0f @@ -989,10 +987,8 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); float a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = - transformer_engine::detail::TypeExtrema::max; - static constexpr float fp8_max = - transformer_engine::detail::TypeExtrema::max; + static constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + static constexpr float fp8_max = transformer_engine::detail::TypeExtrema::max; static constexpr float fp4_max_inv = 1.0f / fp4_max; float global_encode_scale = a_global_amax_val > 0.0f ? cutlass::minimum_with_nan_propagation{}( diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index fb7f9a319e..090798e921 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -415,8 +415,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int kNumThreadsReduce = kScaleBlockDim / kNVecOut; const float global_encode_scale = - (kIsE8Scaling || global_amax == nullptr) ? 1.0f - : ComputeGlobalEncodeScaleFP4(global_amax[0]); + (kIsE8Scaling || global_amax == nullptr) ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]); constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; const float global_decode_scale = 1.0 / global_encode_scale; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 5f01432688..ab59d0cd1f 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -999,8 +999,7 @@ std::tuple, std::vector, bool> bulk_alloc const bool nvfp4_use_4over6 = quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; - const bool disable_second_level_scale = - quantizer_cpp_list[0]->disable_second_level_scale; + const bool disable_second_level_scale = quantizer_cpp_list[0]->disable_second_level_scale; const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); @@ -1031,8 +1030,7 @@ std::tuple, std::vector, bool> bulk_alloc "NVFP4 bulk allocation requires all quantizers in the group to share " "the same with_rht value (tensor 0=", group_with_rht, ", tensor ", i, "=", quantizer_cpp_list[i]->with_rht, ")."); - NVTE_CHECK(quantizer_cpp_list[i]->disable_second_level_scale == - disable_second_level_scale, + NVTE_CHECK(quantizer_cpp_list[i]->disable_second_level_scale == disable_second_level_scale, "NVFP4 bulk allocation requires all quantizers in the group to share " "the same disable_second_level_scale value."); } @@ -1185,13 +1183,12 @@ std::tuple, std::vector, bool> bulk_alloc (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none()); py::object columnwise_scale = (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); - py::object amax_rowwise = - (rowwise_usage && !disable_second_level_scale) ? py::cast(amax_rowwise_list[i]) - : py::none(); - py::object amax_columnwise = - (columnwise_usage && !disable_second_level_scale) - ? py::cast(amax_columnwise_list[i]) - : py::none(); + py::object amax_rowwise = (rowwise_usage && !disable_second_level_scale) + ? py::cast(amax_rowwise_list[i]) + : py::none(); + py::object amax_columnwise = (columnwise_usage && !disable_second_level_scale) + ? py::cast(amax_columnwise_list[i]) + : py::none(); // Construct Python tensor. tensor_py_list.emplace_back(NVFP4TensorClass( diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index d4e7b43e35..3bf423cb21 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1876,8 +1876,7 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize NVTE_ERROR("Unsupported NVFP4 4over6 error mode: ", nvfp4_4over6_err_mode); } this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); - this->disable_second_level_scale = - quantizer.attr("disable_second_level_scale").cast(); + this->disable_second_level_scale = quantizer.attr("disable_second_level_scale").cast(); // Get amax reduction group if needed for NVFP4 AG const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -2017,8 +2016,7 @@ std::pair NVFP4Quantizer::create_tensor( auto rowwise_scale_inv_py = py_cast(rowwise_scale_inv_tensor, rowwise_usage); auto columnwise_data_py = py_cast(columnwise_data_tensor, columnwise_usage); auto columnwise_scale_inv_py = py_cast(columnwise_scale_inv_tensor, columnwise_usage); - auto amax_rowwise_py = - py_cast(amax_rowwise, rowwise_usage && !disable_second_level_scale); + auto amax_rowwise_py = py_cast(amax_rowwise, rowwise_usage && !disable_second_level_scale); auto amax_columnwise_py = py_cast(amax_columnwise, columnwise_usage && !disable_second_level_scale);