From 25d5b4b7916a1d7f63c47d9c0209935e88fe3553 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 7 Aug 2026 09:58:06 +0000 Subject: [PATCH 01/54] Prototype NVFP4 with UE5M3 scales Co-authored-by: Teddy Do Co-authored-by: Varun Thumbe Signed-off-by: Tim Moon --- .../cpp/operator/test_cast_nvfp4_transpose.cu | 485 +++++++++++++----- tests/cpp/operator/test_dequantize_nvfp4.cu | 137 ++++- tests/cpp/test_common.cu | 14 +- tests/cpp/test_common.h | 19 +- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 49 ++ .../test_nvfp4_group_quantize_graph_safe.py | 53 ++ .../nvfp4/test_nvfp4_quantize_exact.py | 60 +++ tests/pytorch/test_fusible_ops.py | 43 +- tests/pytorch/utils.py | 32 +- .../common/cast/dispatch/quantize.cuh | 40 +- .../common/cast/nvfp4/core_nvfp4.cuh | 156 ++++-- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 87 ++-- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 84 +-- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 257 ++++++---- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 156 +++--- .../quantize_transpose_nvfp4_tuned_1D.cuh | 135 ++--- transformer_engine/common/common.h | 79 ++- .../common/gemm/cublaslt_gemm.cu | 32 +- .../common/gemm/cublaslt_grouped_gemm.cu | 94 ++-- ...cast_col_hadamard_transform_cast_fusion.cu | 139 ++--- .../group_hadamard_transform_cast_fusion.cu | 56 +- ...cast_col_hadamard_transform_cast_fusion.cu | 122 +++-- .../hadamard_transform_cast_fusion.cu | 58 +-- ...cast_col_hadamard_transform_cast_fusion.cu | 108 ++-- .../include/transformer_engine/recipe.h | 37 +- .../transformer_engine/transformer_engine.h | 10 +- transformer_engine/common/recipe/__init__.py | 24 +- transformer_engine/common/recipe/nvfp4.cu | 214 ++++---- .../common/transformer_engine.cpp | 58 ++- ...quantize_transpose_vector_blockwise_fp4.cu | 77 +-- .../common/util/pybind_helper.h | 2 + transformer_engine/pytorch/__init__.py | 1 + transformer_engine/pytorch/constants.py | 6 + transformer_engine/pytorch/csrc/common.h | 10 +- .../pytorch/csrc/extensions/activation.cpp | 6 + .../pytorch/csrc/extensions/bias.cpp | 3 + .../pytorch/csrc/extensions/cast.cpp | 138 +++-- .../pytorch/csrc/extensions/normalization.cpp | 6 + transformer_engine/pytorch/csrc/quantizer.cpp | 116 +++-- .../pytorch/csrc/type_converters.cpp | 44 +- transformer_engine/pytorch/module/base.py | 4 + .../pytorch/ops/fused/grouped_mlp.py | 26 +- transformer_engine/pytorch/quantization.py | 60 ++- .../pytorch/quantized_tensor.py | 5 +- .../pytorch/tensor/grouped_tensor.py | 4 + .../pytorch/tensor/nvfp4_tensor.py | 57 +- .../tensor/storage/grouped_tensor_storage.py | 64 ++- .../tensor/storage/nvfp4_tensor_storage.py | 37 +- 48 files changed, 2410 insertions(+), 1094 deletions(-) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 2d9b3073a2..1527142d21 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -61,15 +62,46 @@ std::vector create_transpose(const InputType* const input, const size return input_t; } -// Compute the global encode scale factor for a given global amax +template +constexpr float nvfp4_encode_scale_max(const int scale_type_max = 448) { + static_assert(std::is_same_v +#if CUDA_VERSION >= 13040 + || std::is_same_v +#endif + , "Unsupported NVFP4 scale type."); + if constexpr (std::is_same_v) { + NVTE_CHECK(scale_type_max == 448 || scale_type_max == 256, + "Unsupported E4M3 scale maximum."); + return static_cast(scale_type_max); +#if CUDA_VERSION >= 13040 + } else { + NVTE_CHECK(scale_type_max == 114688 || scale_type_max == 65536, + "Unsupported UE5M3 scale maximum."); + return static_cast(scale_type_max); +#endif + } +} + +template +constexpr float nvfp4_scale_storage_max() { + if constexpr (std::is_same_v) { + return 448.0f; + } +#if CUDA_VERSION >= 13040 + return 114688.0f; +#endif +} + +// Compute the global encode scale factor for a given global amax. +template float compute_global_encode_scaling_factor_FP4(const float global_amax, const bool use_fast_math, - const int e4m3_max = 448) { - NVTE_CHECK(e4m3_max == 448 || e4m3_max == 256, "Unsupported NVFP4 E4M3 max."); - const float fp8_max = static_cast(e4m3_max); + const int scale_type_max = 448) { + const float fp8_max = nvfp4_encode_scale_max(scale_type_max); constexpr float fp4_max = 6.0f; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; // If scale is infinity, return the max normalized value - const float max_norm_clamp = (use_fast_math && e4m3_max == 448) + const float max_norm_clamp = (use_fast_math && std::is_same_v + && scale_type_max == 448) ? Numeric_Traits::maxNorm : Numeric_Traits::maxNorm; @@ -81,9 +113,10 @@ float compute_global_encode_scaling_factor_FP4(const float global_amax, const bo return global_encode_scale; } +template struct NVFP4FourOverSixQuantization { - fp8e4m3 scale_map4; - fp8e4m3 scale_map6; + ScaleType scale_map4; + ScaleType scale_map6; float reciprocal_map4; float reciprocal_map6; fp4e2m1x2 quantized_map4; @@ -103,7 +136,7 @@ enum class NVFP4ScalingMode { struct NVFP4FourOverSixTestConfig { NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled; - int e4m3_max = 448; + int scale_type_max = 448; bool err_use_fast_math = false; }; @@ -111,17 +144,18 @@ bool use_2d_quantization(const NVFP4ScalingMode scaling_mode) { return scaling_mode == NVFP4ScalingMode::Block2D; } -NVFP4FourOverSixQuantization compute_4over6_quantization_scales( - const float block_amax, const float global_encode_scale) { +template +NVFP4FourOverSixQuantization compute_4over6_quantization_scales( + const float block_amax, const float global_encode_scale, const int scale_type_max) { constexpr float fp4_max = 6.0f; - constexpr float fp8_max = 448.0f; + const float fp8_max = nvfp4_scale_storage_max(); constexpr float scale_expansion_factor = 1.5f; const float base_sf_high_precision = block_amax / fp4_max * global_encode_scale; const float sf_high_precision_map4 = fminf(base_sf_high_precision * scale_expansion_factor, fp8_max); const float sf_high_precision_map6 = fminf(base_sf_high_precision, fp8_max); - const fp8e4m3 scale_map4 = static_cast(sf_high_precision_map4); - const fp8e4m3 scale_map6 = static_cast(sf_high_precision_map6); + const ScaleType scale_map4 = static_cast(sf_high_precision_map4); + const ScaleType scale_map6 = static_cast(sf_high_precision_map6); const float global_decode_scale = 1.0f / global_encode_scale; const float scale_map4_fp32 = static_cast(scale_map4); @@ -142,15 +176,18 @@ NVFP4FourOverSixQuantization compute_4over6_quantization_scales( }; } -fp8e4m3 select_4over6_scale(const NVFP4FourOverSixQuantization& quantization, - const NVFP4FourOverSixCandidate candidate) { +template +ScaleType select_4over6_scale(const NVFP4FourOverSixQuantization& quantization, + const NVFP4FourOverSixCandidate candidate) { if (candidate == NVFP4FourOverSixCandidate::Map4) { return quantization.scale_map4; } return quantization.scale_map6; } -fp4e2m1x2 select_4over6_quantized_pair(const NVFP4FourOverSixQuantization& quantization, +template +fp4e2m1x2 select_4over6_quantized_pair( + const NVFP4FourOverSixQuantization& quantization, const NVFP4FourOverSixCandidate candidate) { if (candidate == NVFP4FourOverSixCandidate::Map4) { return quantization.quantized_map4; @@ -158,8 +195,9 @@ fp4e2m1x2 select_4over6_quantized_pair(const NVFP4FourOverSixQuantization& quant return quantization.quantized_map6; } -NVFP4FourOverSixQuantization quantize_4over6_pair( - const float x, const float y, const NVFP4FourOverSixQuantization& quantization) { +template +NVFP4FourOverSixQuantization quantize_4over6_pair( + const float x, const float y, const NVFP4FourOverSixQuantization& quantization) { const float2 scaled_map4 = {x * quantization.reciprocal_map4, y * quantization.reciprocal_map4}; const fp4e2m1x2 quantized_map4(scaled_map4); @@ -179,24 +217,24 @@ NVFP4FourOverSixQuantization quantize_4over6_pair( } // 1D Scaling: Original implementation with 1x16 blocks -template +template void quantize_nvfp4_1d(float (*OP)(const float), const InputType* const input, fp4e2m1x2* const output, - fp8e4m3* const scales, + ScaleType* const scales, const size_t rows, const size_t cols, const size_t scales_stride, const float global_amax, const bool use_fast_math, const bool use_4over6 = false, - const int e4m3_max = 448, + const int scale_type_max = 448, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { // Compute a global encoding/decoding scaling factor for all S_dec_b - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, - e4m3_max); + const float S_enc = compute_global_encode_scaling_factor_FP4( + global_amax, use_fast_math, scale_type_max); constexpr size_t block_size_X = 16; const size_t blocks_X = divide_round_up(cols, block_size_X); @@ -229,8 +267,9 @@ void quantize_nvfp4_1d(float (*OP)(const float), const size_t scale_idx = i * scales_stride + block_X; if (use_4over6) { - const NVFP4FourOverSixQuantization quantization = - compute_4over6_quantization_scales(block_amax, S_enc); + const NVFP4FourOverSixQuantization quantization = + compute_4over6_quantization_scales(block_amax, S_enc, + scale_type_max); scales[scale_idx] = select_4over6_scale(quantization, four_over_six_candidate); for (size_t j = j_min; j < j_max; j += 2) { @@ -239,7 +278,7 @@ void quantize_nvfp4_1d(float (*OP)(const float), const int cache_idx_y = cache_idx_x + 1; const float cached_x = cache_buffer[cache_idx_x]; const float cached_y = cache_buffer[cache_idx_y]; - const NVFP4FourOverSixQuantization pair_quantization = + const NVFP4FourOverSixQuantization pair_quantization = quantize_4over6_pair(cached_x, cached_y, quantization); output[idx_pair] = select_4over6_quantized_pair(pair_quantization, four_over_six_candidate); @@ -249,7 +288,8 @@ void quantize_nvfp4_1d(float (*OP)(const float), // Compute and store the per-block FP8 decode scale const float S_dec_b = block_amax * (S_enc * (1.0f / 6.0f)); - const fp8e4m3 S_dec_b_fp8 = static_cast(fminf(S_dec_b, Numeric_Traits::maxNorm)); + const ScaleType S_dec_b_fp8 = + static_cast(fminf(S_dec_b, Numeric_Traits::maxNorm)); const float S_dec_b_fp32 = static_cast(S_dec_b_fp8); // Compute "correct" per-block encoding scaling factor @@ -284,27 +324,27 @@ void quantize_nvfp4_1d(float (*OP)(const float), } // Compute 2D mathematical scaling factors (8x8 for 128x128 input) -template +template void compute_2d_mathematical_scales(float (*OP)(const float), const InputType* const input, const size_t rows, const size_t cols, const float global_amax, - std::vector>& math_scales, + std::vector>& math_scales, const bool use_fast_math, const bool use_4over6 = false, - const int e4m3_max = 448, + const int scale_type_max = 448, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, - e4m3_max); + const float S_enc = compute_global_encode_scaling_factor_FP4( + global_amax, use_fast_math, scale_type_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); const size_t blocks_X = divide_round_up(cols, block_size_X); - math_scales.resize(blocks_Y, std::vector(blocks_X)); + math_scales.resize(blocks_Y, std::vector(blocks_X)); for (size_t block_Y = 0; block_Y < blocks_Y; ++block_Y) { for (size_t block_X = 0; block_X < blocks_X; ++block_X) { @@ -327,13 +367,14 @@ void compute_2d_mathematical_scales(float (*OP)(const float), // Compute E4M3 scaling factor for this 16x16 block if (use_4over6) { - const NVFP4FourOverSixQuantization quantization = - compute_4over6_quantization_scales(block_amax, S_enc); + const NVFP4FourOverSixQuantization quantization = + compute_4over6_quantization_scales( + block_amax, S_enc, scale_type_max); math_scales[block_Y][block_X] = select_4over6_scale(quantization, four_over_six_candidate); } else { const float S_dec_b = block_amax / 6.0f * S_enc; - const fp8e4m3 S_dec_b_fp8_map6 = static_cast(S_dec_b); + const ScaleType S_dec_b_fp8_map6 = static_cast(S_dec_b); math_scales[block_Y][block_X] = S_dec_b_fp8_map6; } } @@ -341,28 +382,29 @@ void compute_2d_mathematical_scales(float (*OP)(const float), } // 2D Scaling: NEW implementation with proper replication -template +template void quantize_nvfp4_2d(float (*OP)(const float), const InputType* const input, fp4e2m1x2* const output, - fp8e4m3* const scales, + ScaleType* const scales, const size_t rows, const size_t cols, const size_t scales_stride, const float global_amax, const bool use_fast_math, const bool use_4over6 = false, - const int e4m3_max = 448, + const int scale_type_max = 448, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { // Step 1: Compute mathematical 8x8 scaling factors - std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math, - use_4over6, e4m3_max, four_over_six_candidate); + std::vector> math_scales; + compute_2d_mathematical_scales( + OP, input, rows, cols, global_amax, math_scales, use_fast_math, + use_4over6, scale_type_max, four_over_six_candidate); - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, - e4m3_max); + const float S_enc = compute_global_encode_scaling_factor_FP4( + global_amax, use_fast_math, scale_type_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); @@ -434,11 +476,11 @@ void quantize_nvfp4_2d(float (*OP)(const float), } // Wrapper function that calls appropriate implementation based on 2D flag -template +template void quantize_nvfp4(float (*OP)(const float), const InputType* const input, fp4e2m1x2* const output, - fp8e4m3* const scales, + ScaleType* const scales, const size_t rows, const size_t cols, const size_t scales_stride, @@ -446,25 +488,27 @@ void quantize_nvfp4(float (*OP)(const float), const bool use_fast_math, const bool use_2d_quantization = false, const bool use_4over6 = false, - const int e4m3_max = 448, + const int scale_type_max = 448, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { if (use_2d_quantization) { - quantize_nvfp4_2d(OP, input, output, scales, rows, cols, scales_stride, global_amax, - use_fast_math, use_4over6, e4m3_max, four_over_six_candidate); + quantize_nvfp4_2d( + OP, input, output, scales, rows, cols, scales_stride, global_amax, + use_fast_math, use_4over6, scale_type_max, four_over_six_candidate); } else { - quantize_nvfp4_1d(OP, input, output, scales, rows, cols, scales_stride, global_amax, - use_fast_math, use_4over6, e4m3_max, four_over_six_candidate); + quantize_nvfp4_1d( + OP, input, output, scales, rows, cols, scales_stride, global_amax, + use_fast_math, use_4over6, scale_type_max, four_over_six_candidate); } } -template +template void compute_ref(float (*OP)(const float), const InputType* input, fp4e2m1x2* output, fp4e2m1x2* output_t, - fp8e4m3* scales, - fp8e4m3* scales_t, + ScaleType* scales, + ScaleType* scales_t, const float* amax, const size_t rows, const size_t cols, @@ -474,7 +518,7 @@ void compute_ref(float (*OP)(const float), const bool use_2d_quantization = false, const bool row_scaled_nvfp4 = false, const bool use_4over6 = false, - const int e4m3_max = 448, + const int scale_type_max = 448, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { @@ -485,9 +529,10 @@ void compute_ref(float (*OP)(const float), // Ref impl for 2D quantization if (use_2d_quantization) { // Step 1: Compute mathematical 8×8 scaling factors - std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, *amax, math_scales, use_fast_math, - use_4over6, e4m3_max, four_over_six_candidate); + std::vector> math_scales; + compute_2d_mathematical_scales( + OP, input, rows, cols, *amax, math_scales, use_fast_math, + use_4over6, scale_type_max, four_over_six_candidate); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; @@ -514,12 +559,14 @@ void compute_ref(float (*OP)(const float), // Step 4: Process quantized outputs using the same algorithm as quantize_nvfp4_2d // (This part processes the actual FP4 data using the mathematical scaling factors) - quantize_nvfp4_2d(OP, input, output, nullptr, rows, cols, scales_stride, *amax, - use_fast_math, use_4over6, e4m3_max, - four_over_six_candidate); // scales already filled - quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, *amax, - use_fast_math, use_4over6, e4m3_max, - four_over_six_candidate); // scales_t already filled + quantize_nvfp4_2d( + OP, input, output, nullptr, rows, cols, scales_stride, *amax, + use_fast_math, use_4over6, scale_type_max, + four_over_six_candidate); // scales already filled + quantize_nvfp4_2d( + OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, *amax, + use_fast_math, use_4over6, scale_type_max, + four_over_six_candidate); // scales_t already filled return; } @@ -527,7 +574,7 @@ void compute_ref(float (*OP)(const float), // Ref impl for row-scaling if (row_scaled_nvfp4) { for (size_t row = 0; row < rows; ++row) { - quantize_nvfp4(OP, + quantize_nvfp4(OP, input + row * cols, output + row * (cols / 2), scales + row * scales_stride, @@ -538,19 +585,21 @@ void compute_ref(float (*OP)(const float), use_fast_math, use_2d_quantization, use_4over6, - e4m3_max, + scale_type_max, four_over_six_candidate); } return; } // Ref impl for basic NVFP4 - quantize_nvfp4(OP, input, output, scales, rows, cols, scales_stride, *amax, - use_fast_math, use_2d_quantization, use_4over6, e4m3_max, - four_over_six_candidate); - quantize_nvfp4(OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, *amax, - use_fast_math, use_2d_quantization, use_4over6, e4m3_max, - four_over_six_candidate); + quantize_nvfp4( + OP, input, output, scales, rows, cols, scales_stride, *amax, + use_fast_math, use_2d_quantization, use_4over6, scale_type_max, + four_over_six_candidate); + quantize_nvfp4( + OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, *amax, + use_fast_math, use_2d_quantization, use_4over6, scale_type_max, + four_over_six_candidate); } void compare_nvfp4_tensors(const std::string& name, @@ -687,6 +736,27 @@ bool bitwise_equal(const T& x, const T& y) { return true; } +template +void compare_scaling_factors_exact(const std::string& name, const T* test, const T* ref, + const size_t row_blocks, const size_t col_blocks, + const size_t stride) { + size_t mismatches = 0; + for (size_t row = 0; row < row_blocks; ++row) { + for (size_t col = 0; col < col_blocks; ++col) { + const size_t idx = row * stride + col; + if (!bitwise_equal(test[idx], ref[idx])) { + ++mismatches; + if (mismatches <= 3) { + std::cout << "Bitwise scale mismatch in " << name << " at (" << row << ", " + << col << "): " << static_cast(test[idx]) << " vs " + << static_cast(ref[idx]) << std::endl; + } + } + } + } + EXPECT_EQ(mismatches, 0u) << "Bitwise scale mismatches in " << name; +} + bool nvfp4_output_block_matches(const fp4e2m1x2* const test_data, const fp4e2m1x2* const ref_data, const size_t row, @@ -704,13 +774,14 @@ bool nvfp4_output_block_matches(const fp4e2m1x2* const test_data, return true; } +template void compare_nvfp4_4over6_candidates(const std::string& name, const fp4e2m1* const test_data, - const fp8e4m3* const test_scales, + const ScaleType* const test_scales, const fp4e2m1x2* const ref_data_map4, - const fp8e4m3* const ref_scales_map4, + const ScaleType* const ref_scales_map4, const fp4e2m1x2* const ref_data_map6, - const fp8e4m3* const ref_scales_map6, + const ScaleType* const ref_scales_map6, const size_t rows, const size_t cols, const size_t blocks_X, @@ -771,13 +842,13 @@ void compare_rowwise_amax(Tensor &output, const std::vector &ref_amax) { } } -template +template void performTest(float (*OP)(const float), const std::vector& shape, const bool use_fast_math, const NVFP4ScalingMode scaling_mode = NVFP4ScalingMode::Block1D, const NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled, - const int e4m3_max = 448, + const int scale_type_max = 448, const bool use_4over6_err_use_fast_math = false) { using namespace test; const bool use_4over6 = mode != kNVTENVFP44Over6Disabled; @@ -791,6 +862,7 @@ void performTest(float (*OP)(const float), DType itype = TypeInfo::dtype; DType otype = DType::kFloat4E2M1; + DType scale_type = TypeInfo::dtype; const bool is_2d_quantization = use_2d_quantization(scaling_mode); const bool row_scaled_nvfp4 = scaling_mode == NVFP4ScalingMode::RowScaled1D; @@ -818,22 +890,26 @@ void performTest(float (*OP)(const float), const size_t scales_stride_t = blocks_X_t; Tensor input("input", shape, itype); - Tensor output("output", shape, otype, rowwise, columnwise, NVTE_NVFP4_1D_SCALING); - output.set_nvfp4_e4m3_max(e4m3_max); + Tensor output("output", shape, otype, rowwise, columnwise, NVTE_NVFP4_1D_SCALING, + scale_type); + output.set_nvfp4_e4m3_max(scale_type_max); std::unique_ptr ref_output = std::make_unique(rows * (cols / 2)); std::unique_ptr ref_output_t = std::make_unique(cols * (rows / 2)); - std::unique_ptr ref_scales = std::make_unique(blocks_Y * blocks_X); - std::unique_ptr ref_scales_t = std::make_unique(blocks_Y_t * blocks_X_t); + std::unique_ptr ref_scales = + std::make_unique(blocks_Y * blocks_X); + std::unique_ptr ref_scales_t = + std::make_unique(blocks_Y_t * blocks_X_t); std::unique_ptr ref_output_map6; std::unique_ptr ref_output_t_map6; - std::unique_ptr ref_scales_map6; - std::unique_ptr ref_scales_t_map6; + std::unique_ptr ref_scales_map6; + std::unique_ptr ref_scales_t_map6; fillCase(&input, InputsFillCase::uniform); if (use_4over6 && row_scaled_nvfp4) { - const float target_row_amax = static_cast(e4m3_max) * 6.0f * 8.0f; + const float target_row_amax = + nvfp4_encode_scale_max(scale_type_max) * 6.0f * 8.0f; auto *input_vals = input.rowwise_cpu_dptr(); for (size_t row = 0; row < rows; ++row) { float row_amax = 0.0f; @@ -884,11 +960,8 @@ void performTest(float (*OP)(const float), output.set_row_scaled_nvfp4(row_scaled_nvfp4); } else { // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues - if (use_4over6) { - ref_amax.assign(1, static_cast(e4m3_max) * 6.0f * 8.0f); - } else { - ref_amax.assign(1, 448.0f * 6.0f * 8.0f); - } + ref_amax.assign( + 1, nvfp4_encode_scale_max(scale_type_max) * 6.0f * 8.0f); // Update tensor if (rowwise) { @@ -903,10 +976,10 @@ void performTest(float (*OP)(const float), if (use_4over6) { ref_output_map6 = std::make_unique(rows * (cols / 2)); ref_output_t_map6 = std::make_unique(cols * (rows / 2)); - ref_scales_map6 = std::make_unique(blocks_Y * blocks_X); - ref_scales_t_map6 = std::make_unique(blocks_Y_t * blocks_X_t); + ref_scales_map6 = std::make_unique(blocks_Y * blocks_X); + ref_scales_t_map6 = std::make_unique(blocks_Y_t * blocks_X_t); - compute_ref(OP, + compute_ref(OP, input.rowwise_cpu_dptr(), ref_output.get(), ref_output_t.get(), @@ -921,9 +994,9 @@ void performTest(float (*OP)(const float), is_2d_quantization, row_scaled_nvfp4, use_4over6, - e4m3_max, + scale_type_max, NVFP4FourOverSixCandidate::Map4); - compute_ref(OP, + compute_ref(OP, input.rowwise_cpu_dptr(), ref_output_map6.get(), ref_output_t_map6.get(), @@ -938,10 +1011,10 @@ void performTest(float (*OP)(const float), is_2d_quantization, row_scaled_nvfp4, use_4over6, - e4m3_max, + scale_type_max, NVFP4FourOverSixCandidate::Map6); } else { - compute_ref(OP, + compute_ref(OP, input.rowwise_cpu_dptr(), ref_output.get(), ref_output_t.get(), @@ -955,7 +1028,8 @@ void performTest(float (*OP)(const float), use_fast_math, is_2d_quantization, row_scaled_nvfp4, - use_4over6); + use_4over6, + scale_type_max); } // Initialize stochastic rounding @@ -1002,9 +1076,9 @@ void performTest(float (*OP)(const float), if (use_4over6) { output.to_cpu(); - compare_nvfp4_4over6_candidates("output", + compare_nvfp4_4over6_candidates("output", output.rowwise_cpu_dptr(), - output.rowwise_cpu_scale_inv_ptr(), + output.rowwise_cpu_scale_inv_ptr(), ref_output.get(), ref_scales.get(), ref_output_map6.get(), @@ -1014,9 +1088,9 @@ void performTest(float (*OP)(const float), unpadded_blocks_X, scales_stride); if (!row_scaled_nvfp4) { - compare_nvfp4_4over6_candidates("output_t", + compare_nvfp4_4over6_candidates("output_t", output.columnwise_cpu_dptr(), - output.columnwise_cpu_scale_inv_ptr(), + output.columnwise_cpu_scale_inv_ptr(), ref_output_t.get(), ref_scales_t.get(), ref_output_t_map6.get(), @@ -1032,17 +1106,28 @@ void performTest(float (*OP)(const float), true, false, !row_scaled_nvfp4); size_t scale_mismatches_num = 0; - compare_scaling_factors("scales", output.rowwise_cpu_scale_inv_ptr(), - ref_scales.get(), - unpadded_blocks_Y, unpadded_blocks_X, scales_stride, - scale_mismatches_num); + if constexpr (std::is_same_v) { + compare_scaling_factors( + "scales", output.rowwise_cpu_scale_inv_ptr(), ref_scales.get(), + unpadded_blocks_Y, unpadded_blocks_X, scales_stride, scale_mismatches_num); + } else { + compare_scaling_factors_exact( + "scales", output.rowwise_cpu_scale_inv_ptr(), ref_scales.get(), + unpadded_blocks_Y, unpadded_blocks_X, scales_stride); + } if (!row_scaled_nvfp4) { - compare_scaling_factors("scales_t", - output.columnwise_cpu_scale_inv_ptr(), - ref_scales_t.get(), - unpadded_blocks_Y_t, unpadded_blocks_X_t, - scales_stride_t, scale_mismatches_num); + if constexpr (std::is_same_v) { + compare_scaling_factors( + "scales_t", output.columnwise_cpu_scale_inv_ptr(), + ref_scales_t.get(), unpadded_blocks_Y_t, unpadded_blocks_X_t, + scales_stride_t, scale_mismatches_num); + } else { + compare_scaling_factors_exact( + "scales_t", output.columnwise_cpu_scale_inv_ptr(), + ref_scales_t.get(), unpadded_blocks_Y_t, unpadded_blocks_X_t, + scales_stride_t); + } } } @@ -1283,7 +1368,8 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam transformer_engine::DType, bool, NVFP4ScalingMode, - NVFP4FourOverSixTestConfig>> {}; + NVFP4FourOverSixTestConfig, + transformer_engine::DType>> {}; TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { // Skip tests for pre-Blackwell architectures @@ -1300,6 +1386,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { const bool use_fast_math = std::get<3>(GetParam()); const NVFP4ScalingMode scaling_mode = std::get<4>(GetParam()); const NVFP4FourOverSixTestConfig config = std::get<5>(GetParam()); + const DType scale_type = std::get<6>(GetParam()); // Skip tests if the input tensor is 1D if (tensor_dims.size() < 2) { @@ -1316,10 +1403,21 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { case ActivationType::SReLU: OP = &srelu; break; } - TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, - performTest(OP, tensor_dims, use_fast_math, scaling_mode, config.mode, - config.e4m3_max, - config.err_use_fast_math); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, { + if (scale_type == DType::kFloat8E4M3) { + performTest( + OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_type_max, + config.err_use_fast_math); +#if CUDA_VERSION >= 13040 + } else if (scale_type == DType::kFloat8UE5M3) { + performTest( + OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_type_max, + config.err_use_fast_math); +#endif + } else { + FAIL() << "Unsupported NVFP4 scale dtype " << static_cast(scale_type); + } + } ); } @@ -1358,11 +1456,7 @@ std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param) const NVFP4FourOverSixTestConfig& config = std::get<5>(param); if (config.mode != kNVTENVFP44Over6Disabled) { name += "X4OVER6"; - if (config.e4m3_max == 448) { - name += "XE4M3_MAX_448"; - } else { - name += "XE4M3_MAX_256"; - } + name += "XSCALE_MAX_" + std::to_string(config.scale_type_max); if (config.mode == kNVTENVFP44Over6MinMSE) { name += "XMSE"; } else if (config.mode == kNVTENVFP44Over6MinMAE) { @@ -1374,6 +1468,9 @@ std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param) name += "XERR_USE_FAST_MATH"; } } + if (std::get<6>(param) != DType::kFloat8E4M3) { + name += "X" + test::typeName(std::get<6>(param)); + } return name; } @@ -1386,7 +1483,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16), // input_type ::testing::Values(false), // use_fast_math ::testing::Values(NVFP4ScalingMode::Block1D), // scaling_mode - ::testing::Values(NVFP4FourOverSixTestConfig{})), // four_over_six_config + ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config + ::testing::Values(DType::kFloat8E4M3)), // scale_type [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1400,7 +1498,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kBFloat16, DType::kFloat32), // input_type ::testing::Values(false), // use_fast_math ::testing::Values(NVFP4ScalingMode::RowScaled1D), // scaling_mode - ::testing::Values(NVFP4FourOverSixTestConfig{})), // four_over_six_config + ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config + ::testing::Values(DType::kFloat8E4M3)), // scale_type [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1424,7 +1523,8 @@ INSTANTIATE_TEST_SUITE_P( NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, false}, NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, true}, NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, false}, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, true})), // four_over_six_config + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, true}), // four_over_six_config + ::testing::Values(DType::kFloat8E4M3)), // scale_type [](const testing::TestParamInfo& info) { return test_name(info.param); }); @@ -1477,3 +1577,142 @@ INSTANTIATE_TEST_SUITE_P( } return name; }); + +#if CUDA_VERSION >= 13040 && FP4_TYPE_SUPPORTED +INSTANTIATE_TEST_SUITE_P( + OperatorTestUE5M3, + FusedCastTransposeNVFP4TestSuite, + ::testing::Values( + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block1D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block2D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block1D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 114688, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block2D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 114688, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block1D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 65536, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kBFloat16, false, + NVFP4ScalingMode::Block2D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 65536, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kFloat32, false, + NVFP4ScalingMode::Block1D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, + DType::kFloat8UE5M3}, + FusedCastTransposeNVFP4TestSuite::ParamType{ + ActivationType::Identity, {256, 256}, DType::kFloat32, false, + NVFP4ScalingMode::RowScaled1D, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, + DType::kFloat8UE5M3}), + [](const testing::TestParamInfo& info) { + return test_name(info.param); + }); + +TEST(NVFP4UE5M3ReferenceTest, Grouped) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + constexpr size_t num_outputs = 2; + constexpr size_t rows = 128; + constexpr size_t cols = 256; + constexpr float golden_amax = 114688.0f * 6.0f * 8.0f; + const std::vector input_shape{num_outputs * rows, cols}; + const std::vector output_shape{rows, cols}; + + Tensor input("ue5m3_group_input", input_shape, DType::kBFloat16); + fillCase(&input, InputsFillCase::uniform); + input.to_cpu(); + + Tensor output_storage("ue5m3_group_output_storage", input_shape, DType::kFloat4E2M1, + true, false, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); + const NVTEBasicTensor storage_data = + nvte_get_tensor_param(output_storage.data(), kNVTERowwiseData); + const NVTEBasicTensor storage_scales = + nvte_get_tensor_param(output_storage.data(), kNVTERowwiseScaleInv); + + std::vector output_metadata; + std::vector output_views; + std::vector output_handles; + output_metadata.reserve(num_outputs); + output_views.reserve(num_outputs); + for (size_t i = 0; i < num_outputs; ++i) { + output_metadata.emplace_back( + "ue5m3_group_output_metadata_" + std::to_string(i), output_shape, + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); + output_metadata.back().set_amax(golden_amax); + + const NVTEBasicTensor amax = + nvte_get_tensor_param(output_metadata.back().data(), kNVTEAmax); + auto *data_ptr = + reinterpret_cast(storage_data.data_ptr) + i * rows * cols / 2; + auto *scale_ptr = + reinterpret_cast(storage_scales.data_ptr) + i * rows * (cols / 16); + output_views.emplace_back(NVTE_NVFP4_1D_SCALING); + output_views.back().set_rowwise_data(data_ptr, DType::kFloat4E2M1, output_shape); + output_views.back().set_rowwise_scale_inv( + scale_ptr, DType::kFloat8UE5M3, std::vector{rows, cols / 16}); + output_views.back().set_amax(amax.data_ptr, DType::kFloat32, std::vector{1}); + } + for (auto &output : output_views) { + output_handles.push_back(output.data()); + } + + const size_t split_sections[num_outputs] = {rows, rows}; + QuantizationConfigWrapper config; + config.set_stochastic_rounding(false); + nvte_group_nvfp4_quantize_with_amax(input.data(), output_handles.data(), split_sections, + num_outputs, config, 0); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + ASSERT_EQ(cudaGetLastError(), cudaSuccess); + output_storage.to_cpu(); + + const auto scale_dims = get_scale_tensor_dims(rows, cols, 1, 16); + const auto scale_dims_t = get_scale_tensor_dims(cols, rows, 1, 16); + const size_t scales_stride = scale_dims[3]; + const size_t scales_stride_t = scale_dims_t[3]; + const auto *input_data = input.rowwise_cpu_dptr(); + + for (size_t i = 0; i < num_outputs; ++i) { + std::vector ref_output(rows * cols / 2); + std::vector unused_ref_output_t(cols * rows / 2); + std::vector ref_scales(scale_dims[2] * scale_dims[3]); + std::vector unused_ref_scales_t(scale_dims_t[2] * scale_dims_t[3]); + compute_ref( + &identity, input_data + i * rows * cols, ref_output.data(), + unused_ref_output_t.data(), ref_scales.data(), unused_ref_scales_t.data(), + &golden_amax, rows, cols, scales_stride, scales_stride_t, false, false, false, false, + 114688); + + const auto *test_output = + output_storage.rowwise_cpu_dptr() + i * rows * cols / 2; + const auto *test_scales = + output_storage.rowwise_cpu_scale_inv_ptr() + i * rows * scales_stride; + compare_nvfp4_tensors( + "grouped_output_" + std::to_string(i), + test_output, + reinterpret_cast(ref_output.data()), rows, cols, 0.0, 0.0); + compare_scaling_factors_exact( + "grouped_scales_" + std::to_string(i), test_scales, ref_scales.data(), + scale_dims[0], scale_dims[1], scales_stride); + } +} +#endif diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index 40c1fbd235..d8080db380 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -20,6 +20,7 @@ #endif #include +#include #include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -39,23 +40,23 @@ float2 cvt_fp4x2_to_float2(fp4e2m1x2 fp4_pair) { return {static_cast(h2.x), static_cast(h2.y)}; } -template +template void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, - const fp8e4m3 *scales, + const ScaleType *scales, const std::vector &amax, OType *output, size_t rows, size_t cols, size_t scale_stride, - int e4m3_max) { - const float factor_inv = 1.0f / (6.0f * static_cast(e4m3_max)); + float scale_max) { + const float factor_inv = 1.0f / (6.0f * scale_max); constexpr size_t BLOCK_SIZE = 16; const size_t Mread = cols / BLOCK_SIZE; const size_t bytes_per_block = BLOCK_SIZE / 2; for (size_t row = 0; row < rows; ++row) { for (size_t block = 0; block < Mread; ++block) { - const fp8e4m3 scale = scales[row * scale_stride + block]; + const ScaleType scale = scales[row * scale_stride + block]; const float final_scale = static_cast(scale) * (amax.size() == 1 ? amax[0] : amax[row]) * factor_inv; @@ -94,7 +95,7 @@ struct NVFP4DequantizeTestConfig { // Quantize a high-precision input to NVFP4, then dequantize and compare // against a CPU reference computed from the quantized data. -template +template void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, const bool row_scaled_nvfp4, const NVTENVFP44Over6Mode mode, @@ -105,7 +106,8 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, // Tensors Tensor input("input", std::vector{rows, cols}, otype); Tensor quantized("quantized", std::vector{rows, cols}, - DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, + TypeInfo::dtype); Tensor output("output", std::vector{rows, cols}, otype, true, false); // Fill input with random data @@ -149,16 +151,17 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, quantized.to_cpu(); const uint8_t *fp4_data = reinterpret_cast(quantized.rowwise_cpu_dptr()); - const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr(); + const ScaleType *scales = quantized.rowwise_cpu_scale_inv_ptr(); const auto *amax = quantized.cpu_rowwise_amax_ptr(); const std::vector amax_vals(amax, amax + amax_size); const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape(); const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; std::unique_ptr ref_output = std::make_unique(rows * cols); - compute_ref_dequantize_nvfp4( + const float scale_max = static_cast(e4m3_max); + compute_ref_dequantize_nvfp4( fp4_data, scales, amax_vals, ref_output.get(), - rows, cols, scale_stride, e4m3_max); + rows, cols, scale_stride, scale_max); // Compare results from TE and reference impls auto [atol, rtol] = getTolerances(otype); @@ -166,7 +169,7 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, } // Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path. -template +template void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, const bool row_scaled_nvfp4, const NVTENVFP44Over6Mode mode, @@ -178,7 +181,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, fillCase(&input, InputsFillCase::uniform); Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, - DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, + TypeInfo::dtype); quantized_compact.set_nvfp4_e4m3_max(e4m3_max); ASSERT_EQ(quantized_compact.nvfp4_e4m3_max(), e4m3_max); if (row_scaled_nvfp4) { @@ -203,7 +207,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, // Create tensor with same FP4 data but swizzled scales Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, - DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, + TypeInfo::dtype); quantized_swizzled.set_nvfp4_e4m3_max(e4m3_max); ASSERT_EQ(quantized_swizzled.nvfp4_e4m3_max(), e4m3_max); if (row_scaled_nvfp4) { @@ -325,6 +330,112 @@ INSTANTIATE_TEST_SUITE_P( } ); +#if CUDA_VERSION >= 13040 +TEST(DequantizeNVFP4Test, UE5M3Scales) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + performTest_dequantize_nvfp4( + 32, 64, false, kNVTENVFP44Over6Disabled, 114688); + performTest_dequantize_nvfp4( + 32, 64, true, kNVTENVFP44Over6Disabled, 114688); + performTest_dequantize_nvfp4_swizzled( + 32, 64, false, kNVTENVFP44Over6Disabled, 114688); + performTest_dequantize_nvfp4_swizzled( + 32, 64, true, kNVTENVFP44Over6Disabled, 114688); + performTest_dequantize_nvfp4( + 32, 64, false, kNVTENVFP44Over6MinMAE, 65536); + performTest_dequantize_nvfp4_swizzled( + 32, 64, true, kNVTENVFP44Over6MinMAE, 65536); +} + +TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + Tensor global_amax("global_amax", std::vector{1}, DType::kFloat32); + Tensor global_scale("global_scale", std::vector{1}, DType::kFloat32); + global_amax.rowwise_cpu_dptr()[0] = 12.0f; + global_amax.from_cpu(); + nvte_nvfp4_compute_global_scale( + global_amax.data(), global_scale.data(), 0, kNVTEFloat8UE5M3); + global_scale.to_cpu(); + EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 114688.0f / 12.0f); + + Tensor block_amax("block_amax", std::vector{1, 2}, DType::kFloat32); + Tensor block_scale("block_scale", std::vector{1, 2}, DType::kFloat32); + block_amax.rowwise_cpu_dptr()[0] = 3.0f; + block_amax.rowwise_cpu_dptr()[1] = 6.0f; + block_amax.from_cpu(); + nvte_nvfp4_compute_per_block_scale( + block_amax.data(), block_scale.data(), global_amax.data(), 0, kNVTEFloat8UE5M3); + block_scale.to_cpu(); + EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 114688.0f / 12.0f); + EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 114688.0f / 12.0f); + + Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); + nvte_nvfp4_expand_scale_to_fp8( + block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, 0, kNVTEFloat8UE5M3); + expanded_scale.to_cpu(); + const auto *scales = reinterpret_cast( + expanded_scale.rowwise_cpu_dptr()); + for (size_t row = 0; row < 16; ++row) { + EXPECT_FLOAT_EQ(static_cast(scales[row * 2]), + static_cast(fp8ue5m3(3.0f * 114688.0f / 12.0f))); + EXPECT_FLOAT_EQ(static_cast(scales[row * 2 + 1]), + static_cast(fp8ue5m3(6.0f * 114688.0f / 12.0f))); + } +} + +TEST(NVFP4RecipeTest, UE5M3PerTensorScale) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + Tensor input_a("input_a", std::vector{32, 32}, DType::kFloat4E2M1, + true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); + Tensor input_b("input_b", std::vector{32, 32}, DType::kFloat4E2M1, + true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); + Tensor alpha_out("alpha_out", std::vector{1}, DType::kFloat32); + + constexpr float amax_a = 12.0f; + constexpr float amax_b = 18.0f; + constexpr float alpha_in = 2.0f; + constexpr float fp4_max = 6.0f; + constexpr float ue5m3_max = 114688.0f; + input_a.set_nvfp4_e4m3_max(static_cast(ue5m3_max)); + input_b.set_nvfp4_e4m3_max(static_cast(ue5m3_max)); + input_a.set_amax(amax_a); + input_b.set_tensor_amax_columnwise(amax_b); + + nvte_nvfp4_compute_per_tensor_scale( + input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0); + alpha_out.to_cpu(); + + const float factor_inv = + 1.0f / (fp4_max * fp4_max * ue5m3_max * ue5m3_max); + const float expected = alpha_in * amax_a * amax_b * factor_inv; + EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr()[0], expected); + + input_a.set_nvfp4_e4m3_max(65536); + input_b.set_nvfp4_e4m3_max(65536); + nvte_nvfp4_compute_per_tensor_scale( + input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0); + alpha_out.to_cpu(); + + constexpr float ue5m3_headroom_max = 65536.0f; + const float headroom_factor_inv = + 1.0f / (fp4_max * fp4_max * ue5m3_headroom_max * ue5m3_headroom_max); + const float headroom_expected = alpha_in * amax_a * amax_b * headroom_factor_inv; + EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr()[0], headroom_expected); +} +#endif + class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam , transformer_engine::DType, diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index e1468ef981..95ec7e6679 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -49,6 +49,9 @@ bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2) { } size_t typeToNumBits(DType type) { + if (type == DType::kFloat8UE5M3) { + return 8; + } TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T, { return TypeInfo::size; @@ -65,6 +68,7 @@ const std::string &typeName(DType type) { {DType::kBFloat16, "bfloat16"}, {DType::kFloat8E4M3, "float8e4m3"}, {DType::kFloat8E5M2, "float8e5m2"}, + {DType::kFloat8UE5M3, "float8ue5m3"}, {DType::kFloat8E8M0, "float8e8m0"}, {DType::kFloat4E2M1, "float4e2m1"}}; return name_map.at(type); @@ -278,7 +282,7 @@ void Tensor::Buffer::from_cpu() { Tensor::Tensor(const std::string& name, const NVTEShape &shape, const DType type, const bool rowwise, const bool columnwise, - const NVTEScalingMode &scaling_mode) + const NVTEScalingMode &scaling_mode, const DType scale_dtype) : tensor_(scaling_mode), rowwise_{rowwise}, columnwise_{columnwise}, name_{name} { // Initialize RNG const size_t seed = create_seed_from_tensor_name(name); @@ -374,6 +378,14 @@ Tensor::Tensor(const std::string& name, { // Block scaling factors auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode()); + if (scaling_mode == NVTE_NVFP4_1D_SCALING) { + NVTE_CHECK(scale_dtype == DType::kFloat8E4M3 || + scale_dtype == DType::kFloat8UE5M3); + rowwise_scale_meta.type = scale_dtype; + rowwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype); + colwise_scale_meta.type = scale_dtype; + colwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype); + } if (rowwise) { const auto scale_shape = rowwise_scale_meta.shape; const auto scale_dtype = rowwise_scale_meta.type; diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 11d96c2e60..9156c03d7b 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -67,6 +67,9 @@ using bf16 = nv_bfloat16; using fp8e4m3 = __nv_fp8_e4m3; using fp8e5m2 = __nv_fp8_e5m2; using fp8e8m0 = uint8_t; +#if CUDA_VERSION >= 13040 +using fp8ue5m3 = __nv_fp8_ue5m3; +#endif #if FP4_TYPE_SUPPORTED using fp4e2m1 = __nv_fp4_e2m1; using fp4e2m1x2 = __nv_fp4x2_e2m1; @@ -91,7 +94,12 @@ struct BitsNumber { template struct TypeInfo { #if FP4_TYPE_SUPPORTED - using types = std::tuple; + using types = std::tuple= 13040 + , fp8ue5m3 +#endif + >; #else using types = std::tuple; #endif @@ -151,15 +159,18 @@ class Tensor { const NVTEShape &shape, const DType type, const bool rowwise = true, const bool columnwise = false, - const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING); + const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING, + const DType scale_dtype = DType::kFloat8E4M3); Tensor(const std::string& name, const std::vector &shape, const DType type, const bool rowwise = true, const bool columnwise = false, - const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING) : - Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode) {} + const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING, + const DType scale_dtype = DType::kFloat8E4M3) : + Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode, + scale_dtype) {} Tensor() = default; diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index eb480060e2..639b2f752e 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_second_level_scale: bool): + return NVFP4Quantizer( + rowwise=True, + columnwise=True, + disable_second_level_scale=disable_second_level_scale, + )(tensor) + + x_ref, w_ref = quantize(x, False), quantize(w, False) + x_test, w_test = quantize(x, disable_x), quantize(w, disable_w) + + def gemm(w_q, x_q): + workspace = torch.empty(4, dtype=torch.uint8, device="cuda") + return tex.generic_gemm( + w_q, + True, + x_q, + False, + None, + None, + TE_DType[torch.bfloat16], + None, + TE_DType[torch.bfloat16], + False, + None, + False, + workspace, + workspace.numel(), + False, + False, + )[0] + + torch.testing.assert_close(gemm(w_test, x_test), gemm(w_ref, x_ref), atol=0, rtol=0) + + def check_nvfp4_gemm_versus_reference( x_dtype: torch.dtype, w_dtype: torch.dtype, 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 38bd1b31a0..ae8c4208ea 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_second_level_scale_matches_split_quantize( + return_transpose: bool, +) -> None: + """Grouped NVFP4 skips amax reduction and consumes framework-owned fixed amaxes.""" + split_sections = [128, 128] + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + torch.manual_seed(0) + x = torch.randn((sum(split_sections), 128), dtype=torch.bfloat16, device="cuda") + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=return_transpose, + with_rht=True, + with_post_rht_amax=True, + disable_second_level_scale=True, + ) + + grouped = fused_grouped_quantize(x, split_section_tensor, quantizer) + actual = grouped.split_into_quantized_tensors() + expected = tex.split_quantize(x, split_sections, [quantizer.copy() for _ in split_sections]) + + for actual_tensor, expected_tensor in zip(actual, expected): + torch.testing.assert_close( + actual_tensor._rowwise_data, expected_tensor._rowwise_data, atol=0, rtol=0 + ) + torch.testing.assert_close( + actual_tensor._rowwise_scale_inv, + expected_tensor._rowwise_scale_inv, + atol=0, + rtol=0, + ) + assert actual_tensor._amax_rowwise is None + assert expected_tensor._amax_rowwise is None + if return_transpose: + torch.testing.assert_close( + actual_tensor._columnwise_data, + expected_tensor._columnwise_data, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + actual_tensor._columnwise_scale_inv, + expected_tensor._columnwise_scale_inv, + atol=0, + rtol=0, + ) + assert actual_tensor._amax_columnwise is None + assert expected_tensor._amax_columnwise is None + + def check_grouped_tensor_nvfp4_versus_reference( x_dtype: torch.dtype, M: int, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index fe1a04334e..20c0042fe6 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_E4M3_AMAX_FOR_UNIT_GLOBAL_SCALE = 448.0 * 6.0 @dataclass(frozen=True) @@ -240,6 +241,65 @@ def check_quantization_nvfp4_versus_reference( torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["standard", "4over6"]) +def test_disable_second_level_scale_uses_only_block_scale( + return_transpose: bool, + use_4over6: bool, +) -> None: + """A missing amax makes the global NVFP4 encode scale exactly one.""" + torch.manual_seed(0) + x = torch.randn((128, 128), dtype=torch.bfloat16, device="cuda") + x[0, 0] = NVFP4_E4M3_AMAX_FOR_UNIT_GLOBAL_SCALE + + common_kwargs = { + "rowwise": True, + "columnwise": return_transpose, + "with_rht": False, + "nvfp4_use_4over6": use_4over6, + } + expected = NVFP4Quantizer(**common_kwargs)(x) + actual = NVFP4Quantizer(**common_kwargs, disable_second_level_scale=True)(x) + + torch.testing.assert_close(actual._rowwise_data, expected._rowwise_data, atol=0, rtol=0) + torch.testing.assert_close( + actual._rowwise_scale_inv, expected._rowwise_scale_inv, atol=0, rtol=0 + ) + torch.testing.assert_close(actual.dequantize(), expected.dequantize(), atol=0, rtol=0) + assert actual._amax_rowwise is None + if return_transpose: + torch.testing.assert_close( + actual._columnwise_data, expected._columnwise_data, atol=0, rtol=0 + ) + torch.testing.assert_close( + actual._columnwise_scale_inv, expected._columnwise_scale_inv, atol=0, rtol=0 + ) + assert actual._amax_columnwise is None + + # Reusing an output previously populated by two-level scaling must remove + # its amax buffers so common kernels select the unit-global-scale path. + NVFP4Quantizer(**common_kwargs, disable_second_level_scale=True).update_quantized( + x / 2, expected + ) + assert expected._amax_rowwise is None + if return_transpose: + assert expected._amax_columnwise is None + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +def test_disable_second_level_scale_disables_row_scaled_nvfp4() -> None: + """Row-scaled NVFP4 is incompatible with omitting second-level scales.""" + with pytest.warns(UserWarning, match="Row-scaled NVFP4 requires second-level scaling"): + quantizer = NVFP4Quantizer( + row_scaled_nvfp4=True, + disable_second_level_scale=True, + ) + + assert not quantizer.row_scaled_nvfp4 + assert quantizer.disable_second_level_scale + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..8da1944f53 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -51,6 +51,7 @@ assert_close_grads, dtype_tols, make_recipe, + nvfp4_variant_names, quantization_tols, reset_rng_states, ) @@ -62,6 +63,7 @@ fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( return_reason=True ) +fp8_ue5m3_available, reason_for_no_fp8_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -111,11 +113,10 @@ def maybe_skip_quantization( pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and not nvfp4_available - ): + if quantization in nvfp4_variant_names and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and not fp8_ue5m3_available: + pytest.skip(reason_for_no_fp8_ue5m3) if quantization == "fp8_block_scaling" and not fp8_block_scaling_available: pytest.skip(reason_for_no_fp8_block_scaling) @@ -132,16 +133,13 @@ def maybe_skip_quantization( elif quantization == "fp8_block_scaling": if math.prod(dims[:-1]) % 128 != 0 or dims[-1] % 128 != 0: pytest.skip("FP8 block scaling requires dims that are divisible by 128") - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): + elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") # Check dtype if dtype is not None: - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and dtype != torch.bfloat16 - ): + if quantization in nvfp4_variant_names and dtype != torch.bfloat16: pytest.skip("NVFP4 quantization is only supported with BF16 data") @@ -208,17 +206,24 @@ def make_reference_and_test_tensors( columnwise=True, block_scaling_dim=2 if tensor_type == "weight" else 1, )(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht", "nvfp4_ue5m3", "nvfp4_rht_ue5m3"): tensor_type = "input" if quantizer_role is not None: tensor_type = quantizer_role.tensor_type - with_rht = quantization == "nvfp4_rht" and tensor_type != "weight" + with_rht = ( + quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" + ) + scale_dtype = ( + te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" + else te.DType.kFloat8E4M3 + ) test = NVFP4Quantizer( + scale_dtype=scale_dtype, with_rht=with_rht, with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=False, + with_random_sign_mask=with_rht, )(test) elif quantization == "nvfp4_4over6": tensor_type = "input" @@ -871,6 +876,7 @@ def test_quantize( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="input"), requires_grad=True, ) grad_quantization = quantization @@ -882,6 +888,7 @@ def test_quantize( quantization=grad_quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) @@ -964,6 +971,7 @@ def _test_basic_linear( test_dtype=dtype, test_device=device, test_is_quantized=quantized_input, + quantizer_role=QuantizerRole(tensor_type="input"), ) w_ref, w_test = make_reference_and_test_tensors( (out_features, in_features), @@ -978,6 +986,7 @@ def _test_basic_linear( test_dtype=dtype, test_device=device, test_is_quantized=quantized_grad_output, + quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) @@ -1575,7 +1584,7 @@ def test_add_extra_input( if in_place: if quantization in ("fp8_delayed_scaling", "fp8_current_scaling", "mxfp8"): tols = dtype_tols(x1_test._fp8_dtype) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + elif quantization in nvfp4_variant_names: tols = dtype_tols(x1_test._fp4_dtype) y_test = y_test.to(dtype=torch.float64, device="cpu") dx1_test = x1_test.grad.to(dtype=torch.float64, device="cpu") @@ -1896,7 +1905,7 @@ def test_clamped_swiglu( quantized_compute = quantization is not None if not quantized_compute and (quantize_forward or quantize_backward): pytest.skip("Quantization scheme has not been provided") - maybe_skip_quantization(quantization, dims=in_shape, device=device) + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -1949,7 +1958,7 @@ def test_clamped_swiglu( # Expected numerical error tols = dtype_tols(dtype) - if quantized_compute and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + if quantized_compute and quantization in nvfp4_variant_names: tols = dtype_tols(te.DType.kFloat4E2M1) elif quantized_compute: tols = dtype_tols(te.DType.kFloat8E4M3) @@ -2140,6 +2149,7 @@ def test_grouped_linear( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="input"), requires_grad=input_requires_grad, ) dy_ref, dy_test = make_reference_and_test_tensors( @@ -2147,6 +2157,7 @@ def test_grouped_linear( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) ws_ref, ws_test = [], [] @@ -3588,6 +3599,7 @@ def test_grouped_mlp( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="input"), ) dy_ref, dy_test = make_reference_and_test_tensors( out_shape, @@ -3596,6 +3608,7 @@ def test_grouped_mlp( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) probs_ref, probs_test = make_reference_and_test_tensors( diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 21601d8cdd..a845a48911 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -17,6 +17,7 @@ import torch import transformer_engine +from transformer_engine.common.recipe import Format as RecipeFormat from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import InferenceParams, QuantizedTensor from transformer_engine.pytorch import DType @@ -31,6 +32,17 @@ from transformer_engine.pytorch.module.base import get_dummy_wgrad +# NVFP4 recipe names +nvfp4_variant_names: Tuple[str, ...] = ( + "nvfp4", + "nvfp4_row_scaled", + "nvfp4_4over6", + "nvfp4_rht", + "nvfp4_ue5m3", + "nvfp4_rht_ue5m3", +) + + def str_to_dtype(dtype: str | torch.dtype) -> torch.dtype: """Convert type name to PyTorch dtype""" if isinstance(dtype, torch.dtype): @@ -119,7 +131,7 @@ def quantization_tols(name: str) -> dict[str, float]: "mxfp8_block_scaling", ): return dtype_tols(DType.kFloat8E4M3) - if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): + if name in nvfp4_variant_names: return dtype_tols(DType.kFloat4E2M1) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -130,30 +142,36 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: return None if name in ("fp8", "fp8_delayed_scaling"): return transformer_engine.common.recipe.DelayedScaling( - fp8_format=transformer_engine.common.recipe.Format.E4M3, + fp8_format=RecipeFormat.E4M3, amax_history_len=8, **recipe_kwargs, ) if name == "fp8_current_scaling": return transformer_engine.common.recipe.Float8CurrentScaling( - fp8_format=transformer_engine.common.recipe.Format.E4M3, + fp8_format=RecipeFormat.E4M3, **recipe_kwargs, ) if name == "mxfp8": return transformer_engine.common.recipe.MXFP8BlockScaling( - fp8_format=transformer_engine.common.recipe.Format.E4M3, + fp8_format=RecipeFormat.E4M3, **recipe_kwargs, ) if name == "fp8_block_scaling": return transformer_engine.common.recipe.Float8BlockScaling(**recipe_kwargs) - if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): + if name in nvfp4_variant_names: + with_rht = name in ("nvfp4_rht", "nvfp4_rht_ue5m3") use_4over6 = name == "nvfp4_4over6" + scale_format = ( + RecipeFormat.UE5M3 if name in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + else RecipeFormat.E4M3 + ) kwargs = { - "disable_rht": name != "nvfp4_rht", + "disable_rht": not with_rht, "disable_stochastic_rounding": True, "disable_2d_quantization": not use_4over6, "row_scaled_activation": name == "nvfp4_row_scaled", "nvfp4_4over6": "all" if use_4over6 else "none", + "fp8_format": scale_format, } kwargs.update(recipe_kwargs) return transformer_engine.common.recipe.NVFP4BlockScaling(**kwargs) @@ -172,6 +190,8 @@ def recipe_id(recipe: Optional[Recipe]) -> str: nvfp4_features.append("4Over6") if not recipe.disable_rht: nvfp4_features.append("RHT") + if recipe.fp8_format == RecipeFormat.UE5M3: + nvfp4_features.append("UE5M3") if nvfp4_features: return f"NVFP4{''.join(nvfp4_features)}BlockScaling" return type(recipe).__name__ diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index f60cee839d..a43a43ab00 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -104,13 +104,18 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, - "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + NVTE_CHECK( + nvfp4_use_4over6 || + output_tensor->get_nvfp4_scale_max() == + static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(output_tensor->amax.dptr != nullptr, + "Row-scaled NVFP4 does not support disabling second-level scaling."); NVTE_CHECK( !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " @@ -121,8 +126,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); } } @@ -281,13 +289,17 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens auto dtype = grad_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, - "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + NVTE_CHECK(nvfp4_use_4over6 || + output_tensor->get_nvfp4_scale_max() == + static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(output_tensor->amax.dptr != nullptr, + "Row-scaled NVFP4 does not support disabling second-level scaling."); NVTE_CHECK( !(nvfp4_use_4over6 && output_tensor->has_columnwise_data()), "Row-scaled NVFP4 transpose quantization is not supported with 4over6 mode. The 4over6 " @@ -298,8 +310,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); } } @@ -438,9 +453,12 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou auto dtype = input_tensor->dtype(); const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - for (const auto *output_tensor : output_tensors) { - NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, - "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + if (!nvfp4_use_4over6) { + for (const auto *output_tensor : output_tensors) { + NVTE_CHECK(output_tensor->get_nvfp4_scale_max() + == static_cast(nvfp4::core::scale_max(output_tensors[0]->scale_inv.dtype)), + "NVFP4 quantization with non-default scale max is only supported with 4over6."); + } } NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "2D quantization is not supported for group quantize."); diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 3820430d5b..0ee4589d6f 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -31,58 +31,132 @@ namespace transformer_engine { namespace dispatch { namespace nvfp4 { -using nvfp4_scale_t = fp8e4m3; +// Central runtime-to-compile-time dispatch for NVFP4 scale storage types. +// SWITCH_FP8UE5M3_TYPE_HANDLE adds UE5M3 when the CUDA toolkit supports it. +#define TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(SCALE_DTYPE, SCALE_TYPE, ...) \ + switch (SCALE_DTYPE) { \ + case DType::kFloat8E4M3: { \ + using SCALE_TYPE = fp8e4m3; \ + { __VA_ARGS__ } \ + } break; \ + SWITCH_FP8UE5M3_TYPE_HANDLE(SCALE_TYPE, __VA_ARGS__) \ + default: { \ + NVTE_ERROR("Unsupported NVFP4 scale dtype ", to_string(SCALE_DTYPE), \ + ". Expected Float8E4M3, or Float8UE5M3 when compiled with CUDA 13.4+."); \ + } \ + } + +namespace core { -namespace quantization_and_transposition_SF { #if FP4_TYPE_SUPPORTED -// Used in transpose variant -// Compute per-block E4M3 encoding/decoding scaling factor -__device__ __forceinline__ nvfp4_scale_t compute_decoding_scaling_factor(const float block_amax, - const float S_enc) { - // constexpr float rcp_6f = 1.0f / 6.0f; - // const float S_dec_b = block_amax * rcp_6f; - // const nvfp4_scale_t S_dec_b_fp8 = static_cast(S_dec_b * S_enc); - // return S_dec_b_fp8; - // NOTE: Divide by 6.0f is not elegant and not efficient. - // However, this is part of the emulation code to ensure exact match. - using namespace detail; - constexpr float fp4_max = TypeExtrema::max; // 6.0f; - constexpr float fp4_max_inv = 1.0f / fp4_max; - const float S_dec_b = block_amax * (S_enc * fp4_max_inv); - return static_cast(fminf(S_dec_b, TypeExtrema::max)); +using namespace ptx; + +// Scale-format-specific behavior belongs here rather than in individual kernels. +template +struct NVFP4ScaleTraits { + static constexpr bool is_supported = false; + static constexpr bool supports_fp16_error_path = false; + static constexpr float expected_max = 0.0f; + static constexpr float headroom_max = 0.0f; +}; + +template <> +struct NVFP4ScaleTraits { + // E4M3 scales fit in FP16 and can use the packed E4M3-to-FP16 PTX fast + // path. UE5M3 scales can exceed the FP16 range, so they retain the generic + // FP32 error path. + static constexpr bool is_supported = true; + static constexpr bool supports_fp16_error_path = true; + static constexpr float expected_max = 448.0f; + static constexpr float headroom_max = 256.0f; +}; + +#if CUDA_VERSION >= 13040 +template <> +struct NVFP4ScaleTraits { + static constexpr bool is_supported = true; + static constexpr bool supports_fp16_error_path = false; + static constexpr float expected_max = 114688.0f; + static constexpr float headroom_max = 65536.0f; +}; +#endif + +// Return the effective maximum used to derive the global NVFP4 encode scale. +// SCALE_TYPE_MAX is the resolved maximum for ScaleType (e.g., 448 for E4M3 +// or 114688 for UE5M3). The headroom maximum keeps the 1.5x map-to-4 scale +// used by 4over6 within the scale format's representable range. +template (NVFP4ScaleTraits::expected_max)> +__host__ __device__ constexpr float scale_max() { + using ScaleTraits = NVFP4ScaleTraits; + static_assert(ScaleTraits::is_supported, "Unsupported NVFP4 scale type."); + if constexpr (ScaleTraits::is_supported) { + static_assert(detail::TypeExtrema::max == ScaleTraits::expected_max, + "Unexpected NVFP4 scale type maximum."); + static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || + SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), + "Unsupported NVFP4 scale type maximum."); + static_assert(ScaleTraits::headroom_max * 1.5f <= ScaleTraits::expected_max, + "NVFP4 4over6 scale headroom exceeds scale type maximum."); + return static_cast(SCALE_TYPE_MAX); + } else { + return 0.0f; + } } -#endif // FP4_TYPE_SUPPORTED -} // namespace quantization_and_transposition_SF -namespace quantization_SF { -#if FP4_TYPE_SUPPORTED -// Used in non-transpose variant -// Compute per-block E4M3 encoding/decoding scaling factor -__device__ __forceinline__ fp8e4m3 compute_decoding_scaling_factor(const float block_amax, - const float S_enc) { - using namespace detail; - constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; // 1 / 6.0f - // const float S_dec_b = block_amax * rcp_6f; - // const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); - // return S_dec_b_fp8; - return static_cast(block_amax * (S_enc * fp4_max_inv)); +// Return the full-range maximum for a runtime scale dtype. +inline float scale_max(const DType scale_dtype) { + float result = 0.0f; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, result = scale_max();) + return result; } -#endif // FP4_TYPE_SUPPORTED -} // namespace quantization_SF -namespace core { +// Return and validate a user-provided maximum for a runtime scale dtype. +inline float scale_max(const DType scale_dtype, const int scale_type_max) { + float result = 0.0f; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, { + using ScaleTraits = NVFP4ScaleTraits; + NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || + scale_type_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + result = static_cast(scale_type_max); + }) + return result; +} -#if FP4_TYPE_SUPPORTED -using namespace ptx; +template +__device__ __forceinline__ ScaleType +compute_decoding_scaling_factor(const float block_amax, const float global_encode_scale) { + // Compute the per-block decode scale in the selected scale storage type: + // + // block_decode_scale = block_amax / fp4_max + // stored_decode_scale = block_decode_scale * global_encode_scale + // + // An equivalent, more literal implementation is: + // + // constexpr float rcp_6f = 1.0f / 6.0f; + // const float block_decode_scale = block_amax * rcp_6f; + // return static_cast(block_decode_scale * global_encode_scale); + // + // Keep the multiplication order below to match the emulation code exactly, + // while avoiding a direct division by the FP4 maximum. + using namespace detail; + constexpr float fp4_max = TypeExtrema::max; // 6.0f + constexpr float fp4_max_inv = 1.0f / fp4_max; + const float decode_scale = block_amax * (global_encode_scale * fp4_max_inv); + return static_cast(fminf(decode_scale, TypeExtrema::max)); +} // Compute the global encode scale factor for a given global amax. -// NVFP4 uses the full E4M3 range by default. Some 4over6 tensors dispatch -// E4M3_MAX=256 to leave room for map-to-4 scale expansion. -template +// NVFP4 uses the full scale-type range by default. The explicit SCALE_MAX +// template argument lets recipes such as 4over6 reserve encoding headroom. +template (detail::TypeExtrema::max)> __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const float global_amax) { using namespace detail; - static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); - constexpr float fp8_max = static_cast(E4M3_MAX); + static_assert(SCALE_MAX > 0, "NVFP4 scale maximum must be positive."); + constexpr float fp8_max = static_cast(SCALE_MAX); constexpr float fp4_max = TypeExtrema::max; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; // If scale is infinity, return max value of float32 diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 13bb01d500..08e8993f83 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -21,6 +21,7 @@ #include "../../util/ptx.cuh" #include "../../utils.cuh" #include "../mxfp8/swizzle.cuh" +#include "core_nvfp4.cuh" #if FP4_TYPE_SUPPORTED #include @@ -31,9 +32,10 @@ namespace dispatch { namespace nvfp4 { namespace dequantize_kernel { #if FP4_TYPE_SUPPORTED -template +template __global__ void __launch_bounds__(512) - dequantize_fp4_kernel(const void *const input, OType *output, const fp8e4m3 *const scales, + dequantize_fp4_kernel(const void *const input, OType *output, const ScaleType *const scales, const float *const tensor_amax, const size_t N, const size_t M, const size_t scale_stride, const size_t num_scale_tiles_X) { const size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -62,10 +64,15 @@ __global__ void __launch_bounds__(512) const size_t my_output_index = (x + y * M) * 4; 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)); + ScaleType scale = scales[my_scale_index]; + constexpr float fp4_max = detail::TypeExtrema::max; + constexpr float unit_global_scale_amax = + fp4_max * core::scale_max(); + float amax = unit_global_scale_amax; + if (tensor_amax != nullptr) { + amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; + } + constexpr float factor_inv = 1.0f / unit_global_scale_amax; float final_scale = static_cast(scale) * amax * factor_inv; #pragma unroll for (int i = 0; i < 4; i++) { @@ -81,6 +88,30 @@ __global__ void __launch_bounds__(512) #endif // FP4_TYPE_SUPPORTED } // namespace dequantize_kernel +#if FP4_TYPE_SUPPORTED +template +inline void launch_dequantize(const Tensor &input, Tensor *output, + const bool with_gemm_swizzled_scales, + const bool row_scaled_nvfp4, const size_t N, const size_t Mread, + const size_t blocks, const size_t threads, + const size_t num_scale_tiles_X, cudaStream_t stream) { + using namespace dequantize_kernel; + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->data.dtype, OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + row_scaled_nvfp4, ROW_SCALED_NVFP4, + dequantize_fp4_kernel + <<>>( + input.data.dptr, reinterpret_cast(output->data.dptr), + reinterpret_cast(input.scale_inv.dptr), + reinterpret_cast(input.amax.dptr), N, Mread, + input.scale_inv.shape.back(), num_scale_tiles_X);););); +} +#endif // FP4_TYPE_SUPPORTED + inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace dequantize_kernel; @@ -92,7 +123,8 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; const bool row_scaled_nvfp4 = input.row_scaled_nvfp4; - const int e4m3_max = input.nvfp4_e4m3_max; + const DType scale_dtype = input.scale_inv.dtype; + const int e4m3_max = input.get_nvfp4_scale_max(); constexpr int FP4_BLOCK_SIZE = 16; const auto [N, M] = input.flat_2d_dims(); @@ -105,32 +137,25 @@ 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, + "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( - output->data.dtype, OType, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - row_scaled_nvfp4, ROW_SCALED_NVFP4, - if (e4m3_max == 256) { - dequantize_fp4_kernel - <<>>( - input.data.dptr, reinterpret_cast(output->data.dptr), - reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(input.amax.dptr), N, Mread, - input.scale_inv.shape.back(), num_scale_tiles_X); - } else { - NVTE_CHECK(e4m3_max == 448, "Unsupported NVFP4 E4M3 max (got ", e4m3_max, ")"); - dequantize_fp4_kernel - <<>>( - input.data.dptr, reinterpret_cast(output->data.dptr), - reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(input.amax.dptr), N, Mread, - input.scale_inv.shape.back(), num_scale_tiles_X); - });); // NOLINT(*) - ); // NOLINT(*) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, { + using ScaleTraits = core::NVFP4ScaleTraits; + if (e4m3_max == static_cast(ScaleTraits::expected_max)) { + launch_dequantize(ScaleTraits::expected_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, + threads, num_scale_tiles_X, stream); + } else { + NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + launch_dequantize(ScaleTraits::headroom_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, + threads, num_scale_tiles_X, stream); + } + }) NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index 91c6af26b5..fdb3c92dcf 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -28,7 +28,6 @@ namespace nvfp4 { namespace group_quantize_transpose_kernel { -using namespace quantization_and_transposition_SF; using namespace core; using namespace ptx; @@ -84,10 +83,12 @@ __device__ __forceinline__ int GetTensorIdAndBoundary( return tensor_id_start; } +template __device__ __forceinline__ void UpdateEncodeDecodeScaleFP32(float *amax_ptr, float *s_enc_ptr, float *s_dec_ptr) { - float s_env_value = - (amax_ptr == nullptr) ? 1.0f : compute_global_encode_scaling_factor_FP4(*amax_ptr); + float s_env_value = (amax_ptr == nullptr) + ? 1.0f + : core::compute_global_encode_scaling_factor_FP4(*amax_ptr); float s_dec_value = 1.0 / s_env_value; *s_enc_ptr = s_env_value; *s_dec_ptr = s_dec_value; @@ -167,11 +168,11 @@ constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / 16 template + typename IType, typename ScaleType, bool USE_STOCHASTIC_ROUNDING, bool RETURN_TRANSPOSE> __global__ void __launch_bounds__(THREADS_NUM) group_quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, - nvfp4_scale_t *const scales_ptr, const float *noop, + ScaleType *const scales_ptr, const float *noop, const size_t rows, const size_t cols, const size_t scale_stride, const size_t *rng_state, MultiAmaxCastTransposeFusionArgs kernel_args) { @@ -273,9 +274,9 @@ __global__ void __launch_bounds__(THREADS_NUM) fp4e2m1x2 *out_data_sh = reinterpret_cast(dshmem + in_mem); fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + ScaleType *out_rowwise_scales_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *out_colwise_scales_sh = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer @@ -286,7 +287,7 @@ __global__ void __launch_bounds__(THREADS_NUM) // TODO (zhongbo): finish this float *amax_rowwise_ptr = nullptr; float *amax_colwise_ptr = nullptr; - nvfp4_scale_t *split_rowwise_scale_ptr = nullptr; + ScaleType *split_rowwise_scale_ptr = nullptr; // suppose the amax is fixed for the current 128x128 tile (need 128 padding) bool need_update_tensor_id = true; @@ -296,17 +297,17 @@ __global__ void __launch_bounds__(THREADS_NUM) size_t split_end = kernel_args.split_sections_range[tensor_id + 1]; amax_rowwise_ptr = reinterpret_cast(kernel_args.rowwise_amax_list[tensor_id]); split_rowwise_scale_ptr = - reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); + reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); float S_enc_rowwise = 1.0f; float S_dec_rowwise = 1.0f; - UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); + UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); // TODO (zhongbo): colwise scaling disabled for now because of transpose float S_enc_colwise = 1.0f; float S_dec_colwise = 1.0f; if (amax_colwise_ptr != nullptr) { - UpdateEncodeDecodeScaleFP32(amax_colwise_ptr, &S_enc_colwise, &S_dec_colwise); + UpdateEncodeDecodeScaleFP32(amax_colwise_ptr, &S_enc_colwise, &S_dec_colwise); } else { S_enc_colwise = S_enc_rowwise; S_dec_colwise = S_dec_rowwise; @@ -342,9 +343,9 @@ __global__ void __launch_bounds__(THREADS_NUM) split_start = kernel_args.split_sections_range[tensor_id]; split_end = kernel_args.split_sections_range[tensor_id + 1]; amax_rowwise_ptr = reinterpret_cast(kernel_args.rowwise_amax_list[tensor_id]); - UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); + UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); split_rowwise_scale_ptr = - reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); + reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); // TODO (zhongbo): colwise scaling disabled for now because of transpose // Skip fetching colwise amax pointer and scaling factor updates } @@ -430,9 +431,9 @@ __global__ void __launch_bounds__(THREADS_NUM) in_compute_colwise[i] = elt; } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_colwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_colwise); // Store scaling factors through SHMEM const size_t scale_idx_sh = @@ -603,9 +604,9 @@ __global__ void __launch_bounds__(THREADS_NUM) } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise); // Check boundaries const size_t scales_offset_Y = @@ -711,14 +712,14 @@ __global__ void __launch_bounds__(THREADS_NUM) // TODO(zhongbo): add back when transpose is supported // Vectorized store scaling factors through SHMEM // if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - // using ScalesVec = Vec; + // using ScalesVec = Vec; // const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; // ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); // const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; // const size_t count = // number of scales in Y dimension of this chunk // (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - // nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - // constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + // ScaleType *dst = &scales_t_ptr[scale_idx_global]; + // constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(ScaleType); // if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { // // Fast path: vectorized store when destination is properly aligned // scales_vec.store_to(dst); @@ -764,6 +765,14 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, // also check that the output has not null data pointer NVTE_CHECK(output->data.dptr != nullptr, "Output data pointer is null."); + const DType scale_dtype = output->scale_inv.dtype; + for (const Tensor *group_output : output_list) { + if (group_output->has_data()) { + NVTE_CHECK(group_output->scale_inv.dtype == scale_dtype, + "All grouped NVFP4 scale tensors must have the same dtype (expected ", + to_string(scale_dtype), ", got ", to_string(group_output->scale_inv.dtype), ")."); + } + } // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to // return the transposed data. bool return_transpose = output->has_columnwise_data(); @@ -824,8 +833,6 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, // const size_t scale_stride_transpose = // return_transpose ? output->columnwise_scale_inv.shape[1] : 0; - nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); - const float *noop_ptr = reinterpret_cast(noop->data.dptr); const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; @@ -860,35 +867,38 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); + const size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * typeToSize(scale_dtype); constexpr size_t in_mem = buff_size_aligned_in; constexpr size_t out_data_mem = buff_size_aligned_out; constexpr size_t out_data_transpose_mem = buff_size_aligned_out; - constexpr size_t out_scales_transpose_mem = buff_size_scales; + const size_t out_scales_transpose_mem = buff_size_scales; constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; - constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + const size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = - group_quantize_transpose_nvfp4_kernel; - if constexpr (use_2d_quantization) { NVTE_ERROR("2D quantization is not supported for group quantize transpose."); } - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - kernel<<>>(tensor_map_input, tensor_map_output, - scales_ptr, noop_ptr, rows, cols, - scale_stride, rng_state, kernel_args); + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + auto kernel = + group_quantize_transpose_nvfp4_kernel; + auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>( + tensor_map_input, tensor_map_output, scales_ptr, noop_ptr, rows, cols, scale_stride, + rng_state, kernel_args);) NVTE_CHECK_CUDA(cudaGetLastError()); });); #else diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 50776a3ed6..9e287b2bd6 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -54,16 +54,6 @@ namespace nvfp4 { } \ } -#define TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(E4M3_MAX_VALUE, E4M3_MAX_CONST, ...) \ - if ((E4M3_MAX_VALUE) == 256) { \ - constexpr int E4M3_MAX_CONST = 256; \ - { __VA_ARGS__ } \ - } else { \ - NVTE_CHECK((E4M3_MAX_VALUE) == 448, "Unsupported NVFP4 E4M3 max."); \ - constexpr int E4M3_MAX_CONST = 448; \ - { __VA_ARGS__ } \ - } - namespace quantize_4over6_kernel { constexpr int kThreads = 128; @@ -97,9 +87,10 @@ struct CandidatePair { Candidate map6; }; +template struct ScalePair { - nvfp4_scale_t map4; - nvfp4_scale_t map6; + ScaleType map4; + ScaleType map6; float inv_map4; float inv_map6; float global_encode_scale; @@ -122,19 +113,24 @@ __device__ __forceinline__ float compute_error_rn(const float diff) { } } -template -__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, - const float global_amax) { - static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); +template +__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, + const float global_amax) { + using ScaleTraits = core::NVFP4ScaleTraits; + static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || + SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), + "Unsupported NVFP4 scale type maximum."); constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f + constexpr float fp8_max = detail::TypeExtrema::max; + constexpr int encode_scale_max = static_cast(core::scale_max()); constexpr float expand_to_map4 = 1.5f; - const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = + core::compute_global_encode_scaling_factor_FP4(global_amax); const float base = block_amax / fp4_max * S_enc; - ScalePair scales; - scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); - scales.map6 = static_cast(fminf(base, fp8_max)); + ScalePair scales; + scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); + scales.map6 = static_cast(fminf(base, fp8_max)); const float S_dec = 1.0f / S_enc; scales.inv_map4 = @@ -187,12 +183,12 @@ __device__ __forceinline__ void load_col_group(const IType *tile, const int row_ } } -template +template __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, const float sf, const float global_amax, float *err) { constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = static_cast(E4M3_MAX); + constexpr float fp8_max = core::scale_max(); constexpr float err_denom = fp4_max * fp8_max; const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; const float dequant = __half2float(__ushort_as_half(half_bits)); @@ -201,11 +197,19 @@ __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_ *err = __fadd_rn(*err, compute_error_rn(diff)); } -__device__ __forceinline__ uint8_t fp8_bits(const nvfp4_scale_t sf) { +template +__device__ __forceinline__ uint8_t fp8_bits(const ScaleType sf) { return *reinterpret_cast(&sf); } -__device__ __forceinline__ FP16ErrorScalePair compute_fp16_error_scales(const ScalePair &scales) { +template +__device__ __forceinline__ FP16ErrorScalePair +compute_fp16_error_scales(const ScalePair &scales) { + // This fast error path interprets the packed scale bits as E4M3. UE5M3 + // deliberately does not enable supports_fp16_error_path and instead uses + // the scale-format-independent float error path in + // cvt_fp32_to_fp4_8x_with_error. + static_assert(core::NVFP4ScaleTraits::supports_fp16_error_path); FP16ErrorScalePair result; const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | (static_cast(fp8_bits(scales.map6)) << 8); @@ -257,9 +261,9 @@ __device__ __forceinline__ void accumulate_fp16_scaled_error_pair(const uint32_t *err = __fadd_rn(*err, compute_error_rn(diff1)); } -template +template __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( - const float (&x)[8], const float block_scale_inverse, const nvfp4_scale_t sf, + const float (&x)[8], const float block_scale_inverse, const ScaleType sf, const uint32_t fp16_error_scale, const float global_amax, const float global_encode_scale, float *err) { uint32_t out = 0; @@ -268,6 +272,11 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( uint32_t out_dequant_3 = 0; uint32_t out_dequant_4 = 0; + // ScaleType is not consumed by this PTX. block_scale_inverse applies the + // selected E4M3 or UE5M3 block scale while forming the FP32 operands. These + // instructions only convert the scaled candidates to FP4 E2M1 and back to + // FP16 for error evaluation, so their encoding is identical for both scale + // storage types. constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; if constexpr (is_blackwell) { asm volatile( @@ -295,7 +304,8 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( "Try recompiling with sm_XXXa instead of sm_XXX."); } - if constexpr (Cfg::err_use_fast_math) { + if constexpr (Cfg::err_use_fast_math && + core::NVFP4ScaleTraits::supports_fp16_error_path) { accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, global_encode_scale, err); accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, @@ -306,39 +316,48 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( global_encode_scale, err); } else { const float sf_float = static_cast(sf); - accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_1, x[0], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_1, x[1], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_2, x[2], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_2, x[3], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_3, x[4], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_3, x[5], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_4, x[6], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_4, x[7], sf_float, + global_amax, err); } return out; } -template +template __device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], - const ScalePair &scales, + const ScalePair &scales, const float global_amax) { CandidatePair candidates; candidates.map4.err = 0.0f; candidates.map6.err = 0.0f; FP16ErrorScalePair fp16_error_scales{}; - if constexpr (Cfg::err_use_fast_math) { + if constexpr (Cfg::err_use_fast_math && + core::NVFP4ScaleTraits::supports_fp16_error_path) { fp16_error_scales = compute_fp16_error_scales(scales); } - candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); - candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); return candidates; @@ -380,8 +399,9 @@ __device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &ca return candidates.map6.packed; } -__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, - const bool pick_map4) { +template +__device__ __forceinline__ ScaleType select_scale(const ScalePair &scales, + const bool pick_map4) { if (pick_map4) { return scales.map4; } @@ -449,9 +469,9 @@ __device__ void load_stage_to_shared_async(const IType *input, IType *tile, cons } } -template -__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, +template +__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, ScaleType *scales, const float *amax, const size_t rows, const size_t cols, const size_t stage_row, const size_t tile_col, const size_t scale_stride) { @@ -476,13 +496,21 @@ __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 = + core::scale_max() * detail::TypeExtrema::max; + 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); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + const ScalePair scale_pair = + compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = + make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -492,7 +520,7 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf } const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const ScaleType selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_col_group = global_col / kGroupSize; @@ -501,11 +529,12 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf } } -template -__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales_t, const float *amax, - const size_t rows, const size_t cols, const size_t stage_row, - const size_t tile_col, const size_t scale_stride_t) { +template +__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, ScaleType *scales_t, + const float *amax, const size_t rows, const size_t cols, + const size_t stage_row, const size_t tile_col, + const size_t scale_stride_t) { constexpr int groups = kStageRowGroups * kTileCols; for (int group = threadIdx.x; group < groups; group += blockDim.x) { const int local_row_group = group / kTileCols; @@ -527,9 +556,14 @@ __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 ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + const float global_amax = + amax == nullptr + ? core::scale_max() * detail::TypeExtrema::max + : amax[0]; + const ScalePair scale_pair = + compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = + make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -539,7 +573,7 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, } const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const ScaleType selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_row_group = global_row / kGroupSize; @@ -549,13 +583,14 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, } template + bool ROW_SCALED_NVFP4, typename Cfg, typename ScaleType, int SCALE_TYPE_MAX, + typename IType> __global__ void __launch_bounds__(kThreads) quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, - const float *amax_rowwise, const float *amax_colwise, const size_t rows, - const size_t cols, const size_t scale_stride, - const size_t scale_stride_t, const float *noop) { + ScaleType *scales, ScaleType *scales_t, const float *amax_rowwise, + const float *amax_colwise, const size_t rows, const size_t cols, + const size_t scale_stride, const size_t scale_stride_t, + const float *noop) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -590,7 +625,7 @@ __global__ void __launch_bounds__(kThreads) IType *stage_tile = stage_tiles[stage]; if constexpr (RETURN_IDENTITY) { - quantize_stage_rowwise( + quantize_stage_rowwise( stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); } @@ -599,7 +634,7 @@ __global__ void __launch_bounds__(kThreads) if (columnwise_amax == nullptr) { columnwise_amax = amax_rowwise; } - quantize_stage_colwise( + quantize_stage_colwise( stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, scale_stride_t); } @@ -614,7 +649,8 @@ __global__ void __launch_bounds__(kThreads) #endif } -template +template void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { const size_t rows = input.flat_first_dim(); @@ -626,8 +662,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out const auto *input_ptr = reinterpret_cast(input.data.dptr); auto *output_ptr = reinterpret_cast(output->data.dptr); auto *output_t_ptr = reinterpret_cast(output->columnwise_data.dptr); - auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); - auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); + auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); const auto *amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); const auto *amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); const auto *noop_ptr = reinterpret_cast(noop->data.dptr); @@ -643,7 +679,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { auto kernel = quantize_4over6_kernel; + ROW_SCALED_NVFP4, Cfg, ScaleType, SCALE_TYPE_MAX, + IType>; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, @@ -657,9 +694,9 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out #endif // FP4_TYPE_SUPPORTED -template -void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { +template +void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace quantize_4over6_kernel; @@ -683,6 +720,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(), @@ -690,7 +729,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,23 +736,31 @@ 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( - output->nvfp4_e4m3_max, E4M3_MAX, - TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( - quant_config->nvfp4_4over6_mode, MODE, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { - using Cfg = quantize_4over6_kernel::Config; - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - quantize_4over6_kernel::launch_quantize_4over6( - input, noop, output, stream);); - }););); + using ScaleTraits = core::NVFP4ScaleTraits; + const int scale_type_max = output->get_nvfp4_scale_max(); + NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || + scale_type_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + scale_type_max == static_cast(ScaleTraits::headroom_max), + USE_SCALE_HEADROOM, { + constexpr int SCALE_TYPE_MAX = + static_cast(USE_SCALE_HEADROOM ? ScaleTraits::headroom_max + : ScaleTraits::expected_max); + TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( + quant_config->nvfp4_4over6_mode, MODE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { + using Cfg = quantize_4over6_kernel::Config; + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + quantize_4over6_kernel::launch_quantize_4over6(input, noop, output, + stream);); + });); + }) NVTE_CHECK_CUDA(cudaGetLastError()); #else @@ -722,6 +768,31 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, #endif // FP4_TYPE_SUPPORTED } +template +void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + const bool return_rowwise = output->has_data(); + const bool return_transpose = output->has_columnwise_data(); + NVTE_CHECK(return_rowwise || return_transpose, + "NVFP4 4over6 output tensor must have rowwise or columnwise data."); + const DType scale_dtype = + return_rowwise ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + if (return_rowwise && return_transpose) { + NVTE_CHECK(output->scale_inv.dtype == output->columnwise_scale_inv.dtype, + "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", + to_string(output->scale_inv.dtype), " and ", + to_string(output->columnwise_scale_inv.dtype), ")."); + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, + quantize_4over6_impl( + input, noop, output, quant_config, stream);) +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + } // namespace nvfp4 } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..0a1dc648c2 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -237,7 +237,6 @@ inline void compute_columnwise_amax(const Tensor &input, const Tensor *noop, Ten namespace quantize_transpose_kernel { -using namespace quantization_and_transposition_SF; using namespace core; using namespace ptx; @@ -316,15 +315,14 @@ constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / 16 template __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const __grid_constant__ CUtensorMap tensor_map_output_t, - nvfp4_scale_t *const scales_ptr, - nvfp4_scale_t *const scales_t_ptr, const float *noop, - const float *const amax_rowwise_ptr, + ScaleType *const scales_ptr, ScaleType *const scales_t_ptr, + const float *noop, const float *const amax_rowwise_ptr, const float *const amax_colwise_ptr, const size_t rows, const size_t cols, const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { @@ -421,9 +419,9 @@ __global__ void __launch_bounds__(THREADS_NUM) fp4e2m1x2 *out_data_sh = reinterpret_cast(dshmem + in_mem); fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + ScaleType *out_rowwise_scales_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *out_colwise_scales_sh = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer @@ -432,15 +430,17 @@ __global__ void __launch_bounds__(THREADS_NUM) const bool is_master_thread = (threadIdx.x == 0); // Compute a global encoding/decoding scaling factors for all S_dec_b - const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) - ? 1.0f - : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + const float S_enc_rowwise = + (amax_rowwise_ptr == nullptr) + ? 1.0f + : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); // NOTE: This is to match with how emulation code was written. const float S_dec_rowwise = 1.0 / S_enc_rowwise; - const float S_enc_colwise = (amax_colwise_ptr == nullptr) - ? S_enc_rowwise - : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + const float S_enc_colwise = + (amax_colwise_ptr == nullptr) + ? S_enc_rowwise + : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); const float S_dec_colwise = 1.0 / S_enc_colwise; float thread_amax = 0.0f; @@ -544,9 +544,9 @@ __global__ void __launch_bounds__(THREADS_NUM) in_compute_colwise[i] = elt; } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_colwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_colwise); // Store scaling factors through SHMEM const size_t scale_idx_sh = @@ -719,16 +719,17 @@ __global__ void __launch_bounds__(THREADS_NUM) float block_scale_inverse; if constexpr (ROW_SCALED_NVFP4) { - // 2. Compute E4M3 scaling factor + // 2. Compute block scaling factor const size_t scales_offset_Y = 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]) + scales_offset_Y < rows && amax_rowwise_ptr != nullptr + ? core::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 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); // Check boundaries const size_t scales_offset_X = scales_offset_X_rowwise; @@ -746,9 +747,9 @@ __global__ void __launch_bounds__(THREADS_NUM) fminf(1.0f / (static_cast(S_dec_b_fp8) * S_dec_rowwise_block), float_max); // S_enc_b_fp8 } else { - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise); // Check boundaries const size_t scales_offset_Y = @@ -835,14 +836,14 @@ __global__ void __launch_bounds__(THREADS_NUM) // Vectorized store scaling factors through SHMEM if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - using ScalesVec = Vec; + using ScalesVec = Vec; const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; const size_t count = // number of scales in Y dimension of this chunk (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + ScaleType *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(ScaleType); if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { // Fast path: vectorized store when destination is properly aligned scales_vec.store_to(dst); @@ -859,15 +860,14 @@ __global__ void __launch_bounds__(THREADS_NUM) } template + typename IType, typename ScaleType, bool USE_STOCHASTIC_ROUNDING, bool RETURN_ROWWISE, + bool RETURN_TRANSPOSE, bool WITH_GEMM_SWIZZLED_SCALES = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const __grid_constant__ CUtensorMap tensor_map_output_t, - nvfp4_scale_t *const scales_ptr, - nvfp4_scale_t *const scales_t_ptr, const float *noop, - const float *const amax_rowwise_ptr, + ScaleType *const scales_ptr, ScaleType *const scales_t_ptr, + const float *noop, const float *const amax_rowwise_ptr, const float *const amax_colwise_ptr, const size_t rows, const size_t cols, const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { @@ -963,9 +963,9 @@ __global__ void __launch_bounds__(THREADS_NUM) fp4e2m1x2 *out_data_sh = reinterpret_cast(dshmem + in_mem); fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + ScaleType *out_rowwise_scales_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *out_colwise_scales_sh = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer @@ -974,15 +974,17 @@ __global__ void __launch_bounds__(THREADS_NUM) const bool is_master_thread = (threadIdx.x == 0); // Compute a global encoding/decoding scaling factors for all S_dec_b - const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) - ? 1.0f - : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + const float S_enc_rowwise = + (amax_rowwise_ptr == nullptr) + ? 1.0f + : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); // NOTE: This is to match with how emulation code was written. const float S_dec_rowwise = 1.0 / S_enc_rowwise; - const float S_enc_colwise = (amax_colwise_ptr == nullptr) - ? S_enc_rowwise - : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + const float S_enc_colwise = + (amax_colwise_ptr == nullptr) + ? S_enc_rowwise + : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); const float S_dec_colwise = 1.0 / S_enc_colwise; const size_t warp_id = threadIdx.x / 32; @@ -1155,9 +1157,9 @@ __global__ void __launch_bounds__(THREADS_NUM) } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_colwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_colwise); // // Store scaling factors through SHMEM const size_t scale_idx_sh = @@ -1280,9 +1282,9 @@ __global__ void __launch_bounds__(THREADS_NUM) } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + // 2. Compute block scaling factor + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise); // Check boundaries const size_t scales_offset_Y = @@ -1397,11 +1399,11 @@ __global__ void __launch_bounds__(THREADS_NUM) scales_t_ptr[off] = out_colwise_scales_sh[scale_idx_sh + k]; } } else { - using ScalesVec = Vec; + using ScalesVec = Vec; ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; - nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + ScaleType *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(ScaleType); if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { // Fast path: vectorized store when destination is properly aligned scales_vec.store_to(dst); @@ -1418,9 +1420,9 @@ __global__ void __launch_bounds__(THREADS_NUM) #endif // FP4_TYPE_SUPPORTED } // namespace quantize_transpose_kernel -template -void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { +template +void quantize_transpose_impl(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace quantize_transpose_kernel; using namespace ptx; @@ -1439,7 +1441,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, const bool return_rowwise = output->has_data(); if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { - quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); + quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); return; } @@ -1461,7 +1463,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); } NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, - "Row-scaled NVFP4 quantization requires rowwise amax."); + "Row-scaled NVFP4 does not support disabling second-level scaling."); NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); // In-kernel GEMM-swizzled scale output is only implemented on the 2D quantization @@ -1498,9 +1500,9 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, const size_t scale_stride_transpose = return_transpose ? output->columnwise_scale_inv.shape[1] : 0; - nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); - nvfp4_scale_t *const scales_transpose_ptr = - reinterpret_cast(output->columnwise_scale_inv.dptr); + ScaleType *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); + ScaleType *const scales_transpose_ptr = + reinterpret_cast(output->columnwise_scale_inv.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); @@ -1541,7 +1543,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); + constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(ScaleType); constexpr size_t in_mem = buff_size_aligned_in; @@ -1562,17 +1564,17 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // The 1D kernel always produces rowwise output (no RETURN_ROWWISE); the dispatch only // routes columnwise-only requests here when use_2d_quantization is true. auto kernel = quantize_transpose_nvfp4_kernel; + ScaleType, USE_STOCHASTIC_ROUNDING, + RETURN_TRANSPOSE, ROW_SCALED_NVFP4>; if constexpr (use_2d_quantization) { if (with_gemm_swizzled_scales) { kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, ScaleType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; } else { kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, ScaleType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; } } @@ -1590,6 +1592,32 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, #endif // FP4_TYPE_SUPPORTED } +template +void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + const bool return_rowwise = output->has_data(); + const bool return_transpose = output->has_columnwise_data(); + NVTE_CHECK(return_rowwise || return_transpose, + "NVFP4 output tensor must have rowwise or columnwise data."); + const DType scale_dtype = + return_rowwise ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + if (return_rowwise && return_transpose) { + NVTE_CHECK(output->scale_inv.dtype == output->columnwise_scale_inv.dtype, + "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", + to_string(output->scale_inv.dtype), " and ", + to_string(output->columnwise_scale_inv.dtype), ")."); + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + quantize_transpose_impl(input, noop, output, quant_config, + stream);) +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + } // namespace nvfp4 } // namespace dispatch } // namespace transformer_engine 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..5b7477af89 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 @@ -16,6 +16,8 @@ #include #include +#include + #include "../../../common.h" #include "../../../util/math.h" #include "../../../util/ptx.cuh" @@ -28,7 +30,6 @@ namespace nvfp4 { namespace quantize_transpose_tuned_kernel { -using namespace quantization_and_transposition_SF; using namespace core; using namespace ptx; @@ -140,8 +141,10 @@ using IType3D = IType[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X]; using IType2x3D = IType2[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X / 2]; using OType2x3D = fp4e2m1x2[BUFFS_NUM_OUT][BUFF_OUT_DIM_Y][BUFF_OUT_DIM_X]; using OType2xt3D = fp4e2m1x2[BUFFS_NUM_OUT_TR][BUFF_OUT_TR_DIM_Y][BUFF_OUT_TR_DIM_X]; -using ScalesType2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_Y][SCALES_PER_CHUNK_X]; -using ScalesTypeTr2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_X][SCALES_PER_CHUNK_Y]; +template +using ScalesType2D = ScaleType[TunableConfig::CHUNK_DIM_Y][SCALES_PER_CHUNK_X]; +template +using ScalesTypeTr2D = ScaleType[TunableConfig::CHUNK_DIM_X][SCALES_PER_CHUNK_Y]; using RNG_t = typename transformer_engine::curanddx::detail::philox4x32_native_state< NVTE_BUILD_NUM_PHILOX_ROUNDS>; @@ -160,41 +163,35 @@ __device__ __forceinline__ float get_amax_of_pair(const IType2 pair) { return static_cast(__hmax(__habs(pair.x), __habs(pair.y))); } -// Compute "correct" per-block encoding scaling factor -template -__device__ __forceinline__ SF_TYPE -compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const float S_enc) { - NVTE_DEVICE_ERROR("Unsupported scaling-factor type. Only FP32 and BF16 are supported."); -} - -template <> -__device__ __forceinline__ float compute_nvfp4_scaling_coefficient( - const nvfp4_scale_t S_dec_block, const float S_enc) { - const float S_dec = 1.0f / S_enc; - const float scale_rcp = - fminf(1.0f / (static_cast(S_dec_block) * S_dec), detail::TypeExtrema::max); - return scale_rcp; -} - -template <> -__device__ __forceinline__ bf16 -compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const float S_enc) { - const float scale_rcp = - fminf(S_enc / (static_cast(S_dec_block)), detail::TypeExtrema::max); - return static_cast(scale_rcp); +// Compute "correct" per-block encoding scaling factor. +template +__device__ __forceinline__ SFType +compute_nvfp4_scaling_coefficient(const ScaleType decode_scale, const float global_encode_scale) { + if constexpr (std::is_same_v) { + const float global_decode_scale = 1.0f / global_encode_scale; + return fminf(1.0f / (static_cast(decode_scale) * global_decode_scale), + detail::TypeExtrema::max); + } else if constexpr (std::is_same_v) { + const float scale_rcp = fminf(global_encode_scale / static_cast(decode_scale), + detail::TypeExtrema::max); + return static_cast(scale_rcp); + } else { + NVTE_DEVICE_ERROR("Unsupported scaling-factor type. Only FP32 and BF16 are supported."); + } } -template +template __device__ __forceinline__ void colwise_scaling( const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_tr_ptr, - nvfp4_scale_t *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, + ScaleType *__restrict__ sSFcolwise_ptr, const float S_enc_colwise, const int stage_Y, const int stage_X, const int buff_in, const int buff_out_tr, const float *amax_colwise_ptr, const size_t col_offset, const size_t cols, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn2x = *reinterpret_cast(sIn_ptr); auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); - auto &sSFcolwise = *reinterpret_cast(sSFcolwise_ptr); + auto &sSFcolwise = *reinterpret_cast *>(sSFcolwise_ptr); const int warp = threadIdx.x / THREADS_PER_WARP; const int thread_lane = threadIdx.x % THREADS_PER_WARP; @@ -233,11 +230,12 @@ __device__ __forceinline__ void colwise_scaling( if constexpr (ROW_SCALED_NVFP4) { const size_t col_idx = col_offset + stage_X * TILE_DIM_X + thread_offset_X_colwise + w; S_enc_colwise_block = - col_idx < cols ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx]) - : 1.0f; + col_idx < cols && amax_colwise_ptr != nullptr + ? core::compute_global_encode_scaling_factor_FP4(amax_colwise_ptr[col_idx]) + : 1.0f; } - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax[w], S_enc_colwise_block); + const ScaleType S_dec_b_fp8 = + core::compute_decoding_scaling_factor(block_amax[w], S_enc_colwise_block); // Store scaling factors to SMEM buffer (R2S) sSFcolwise[scale_tr_offset_Y + w][scale_tr_offset_X] = S_dec_b_fp8; @@ -267,17 +265,18 @@ __device__ __forceinline__ void colwise_scaling( } } -template +template __device__ __forceinline__ void rowwise_scaling( const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_ptr, - nvfp4_scale_t *__restrict__ sSFrowwise_ptr, const float S_enc_rowwise, const int stage_Y, + ScaleType *__restrict__ sSFrowwise_ptr, const float S_enc_rowwise, const int stage_Y, const int stage_X, const int buff_in, const int buff_out, const float *amax_rowwise_ptr, const size_t row_offset, const size_t rows, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn = *reinterpret_cast(sIn_ptr); auto &sOut = *reinterpret_cast(sOut_ptr); - auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); + auto &sSFrowwise = *reinterpret_cast *>(sSFrowwise_ptr); const int thread_lane = threadIdx.x % THREADS_PER_WARP; const int bank_group = thread_lane / THREADS_PER_BANK; @@ -319,18 +318,20 @@ __device__ __forceinline__ void rowwise_scaling( } const float block_amax = get_amax_of_pair(thread_amax_2x); - nvfp4_scale_t S_dec_b_fp8; + ScaleType S_dec_b_fp8; scaling_coeff_type SFcoefficient; 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; - S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); + 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 = + core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); SFcoefficient = compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_rowwise_block); } else { - S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + S_dec_b_fp8 = core::compute_decoding_scaling_factor(block_amax, S_enc_rowwise); SFcoefficient = compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_rowwise); } @@ -366,13 +367,13 @@ __device__ __forceinline__ void rowwise_scaling( } } -template +template __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D_kernel( const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, - const __grid_constant__ CUtensorMap tensor_map_output_t, nvfp4_scale_t *const scales_ptr, - nvfp4_scale_t *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, + const __grid_constant__ CUtensorMap tensor_map_output_t, ScaleType *const scales_ptr, + ScaleType *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, const float *const amax_colwise_ptr, const size_t rows, const size_t cols, const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -407,7 +408,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D constexpr int out_mem_rowwise_data = buff_size_aligned_out; constexpr int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; constexpr int out_mem_rowwise_scales = DIVUP_TO_MULTIPLE( - TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned extern __shared__ unsigned char dynamic_shmem[]; @@ -421,13 +422,13 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D auto &sOut = *reinterpret_cast(sOut_ptr); auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); - nvfp4_scale_t *sSFrowwise_ptr = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *sSFcolwise_ptr = reinterpret_cast( + ScaleType *sSFrowwise_ptr = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + ScaleType *sSFcolwise_ptr = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); - auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); - auto &sSFcolwise = *reinterpret_cast(sSFcolwise_ptr); + auto &sSFrowwise = *reinterpret_cast *>(sSFrowwise_ptr); + auto &sSFcolwise = *reinterpret_cast *>(sSFcolwise_ptr); constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -435,12 +436,12 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) ? 1.0f - : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); const float S_enc_colwise = (amax_colwise_ptr == nullptr || ROW_SCALED_NVFP4) ? S_enc_rowwise - : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); __shared__ uint64_t workID_mbar; __shared__ __uint128_t workID_response; @@ -588,12 +589,12 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D ptx::cp_async_bulk_wait_group_read(); // NVFP4 Quantization - rowwise_scaling( + rowwise_scaling( sIn_ptr, sOut_ptr, sSFrowwise_ptr, S_enc_rowwise, stage_Y, stage_X, buff_in, buff_out, amax_rowwise_ptr, block_offset_Y, rows, rng, random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { - colwise_scaling( + colwise_scaling( sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, S_enc_colwise, stage_Y, stage_X, buff_in, buff_out_tr, amax_colwise_ptr, block_offset_X, cols, rng, random_uint4, rnd_idx); } @@ -633,7 +634,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D { // Rowwise { - using ScalesVec = Vec; + using ScalesVec = Vec; // number of scales in X dimension of this chunk const int count = min(SCALES_PER_CHUNK_X, chunk_cols / SCALE_DIM); @@ -650,7 +651,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D // Colwise if constexpr (RETURN_TRANSPOSE) { - using ScalesVec = Vec; + using ScalesVec = Vec; // number of scales in Y dimension of this chunk const int count = min(SCALES_PER_CHUNK_Y, chunk_rows / SCALE_DIM); @@ -688,6 +689,7 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D #endif // FP4_TYPE_SUPPORTED } // namespace quantize_transpose_tuned_kernel +template inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, Tensor *output, const QuantizationConfig *quant_config, cudaStream_t stream) { @@ -713,15 +715,15 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, - "Row-scaled NVFP4 quantization requires rowwise amax."); - + "Row-scaled NVFP4 does not support disabling second-level scaling."); if (return_transpose) { NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Transposed scaling tensor must be allocated"); NVTE_CHECK(!row_scaled_nvfp4 || output->columnwise_amax.dptr != nullptr, - "Row-scaled NVFP4 transpose quantization requires columnwise amax."); + "Row-scaled NVFP4 transpose quantization does not support disabling " + "second-level scaling."); } const auto [rows, cols] = input.flat_2d_dims(); @@ -740,9 +742,9 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, const size_t scale_stride_transpose = return_transpose ? output->columnwise_scale_inv.shape[1] : 0; - nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); - nvfp4_scale_t *const scales_transpose_ptr = - reinterpret_cast(output->columnwise_scale_inv.dptr); + ScaleType *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); + ScaleType *const scales_transpose_ptr = + reinterpret_cast(output->columnwise_scale_inv.dptr); const float *noop_ptr = reinterpret_cast(noop->data.dptr); const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); @@ -784,9 +786,9 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( - TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( - TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(ScaleType), TMA_SHMEM_ALIGNMENT); const int in_mem = buff_size_aligned_in; @@ -808,8 +810,9 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, row_scaled_nvfp4, ROW_SCALED_NVFP4, TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { auto kernel = - quantize_transpose_nvfp4_tuned_1D_kernel; + quantize_transpose_nvfp4_tuned_1D_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index eb4dcc055c..9b309ac267 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -61,6 +61,8 @@ inline std::string to_string(const DType type) { return "Float8E5M2"; case DType::kFloat8E8M0: return "Float8E8M0"; + case DType::kFloat8UE5M3: + return "Float8UE5M3"; case DType::kFloat4E2M1: return "Float4E2M1"; case DType::kInt16: @@ -302,12 +304,15 @@ struct Tensor { * Only meaningful for NVFP4 tensors. */ bool row_scaled_nvfp4 = false; - /*! \brief Global E4M3 scale bound used by NVFP4. + /*! \brief Global scale bound used by NVFP4. * - * Standard NVFP4 uses 448. Some 4over6 tensors use 256 to leave room for - * map-to-4 local scale expansion. + * When negative, use the maximum value of the scale-inverse dtype. + * Some 4over6 tensors use 256 (instead of the E4M3 max of 448) in + * order to leave room for map-to-4 local scale expansion. + * + * TODO: Change to a dtype-agnostic name. */ - int nvfp4_e4m3_max = 448; + int nvfp4_e4m3_max = -1; /*! Map from NVTETensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { @@ -337,7 +342,7 @@ struct Tensor { scaling_mode = NVTE_DELAYED_TENSOR_SCALING; with_gemm_swizzled_scales = false; row_scaled_nvfp4 = false; - nvfp4_e4m3_max = 448; + nvfp4_e4m3_max = -1; } explicit operator NVTETensor() const noexcept { return nvte_tensor; } @@ -447,6 +452,36 @@ struct Tensor { * as a (D1*D2*...*D(n-1), Dn) matrix. */ size_t flat_last_dim() const { return flat_2d_dims()[1]; } + + /*! \brief Global scale bound used by NVFP4. */ + int get_nvfp4_scale_max() const { + NVTE_CHECK(scaling_mode == NVTE_NVFP4_1D_SCALING, + "Attempted to access NVFP4 scale bound for tensor with scaling mode \"", + to_string(scaling_mode), "\"."); + + // Return non-default scale max + if (nvfp4_e4m3_max >= 0) { + return nvfp4_e4m3_max; + } + + // Deduce scale max based on scale-inverse dtype + DType dtype; + if (scale_inv.has_data()) { + dtype = scale_inv.dtype; + } else if (columnwise_scale_inv.has_data()) { + dtype = columnwise_scale_inv.dtype; + } else { + dtype = scale_inv.dtype; + } + switch (dtype) { + case DType::kFloat8E4M3: + return 448; + case DType::kFloat8UE5M3: + return 114688; + default: + NVTE_ERROR("Unsupported scale dtype for NVFP4 tensor (", to_string(dtype), ")"); + } + } }; struct GroupedTensor { @@ -647,6 +682,9 @@ using fp8e5m2 = __nv_fp8_e5m2; #if CUDA_VERSION >= 12080 using fp8e8m0 = __nv_fp8_e8m0; #endif +#if CUDA_VERSION >= 13040 +using fp8ue5m3 = __nv_fp8_ue5m3; +#endif #if FP4_TYPE_SUPPORTED using fp4e2m1 = __nv_fp4_e2m1; using fp4e2m1x2 = __nv_fp4x2_e2m1; @@ -675,6 +713,9 @@ TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp8_e5m2) #if CUDA_VERSION >= 12080 TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp8_e8m0) #endif +#if CUDA_VERSION >= 13040 +TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp8_ue5m3) +#endif #if FP4_TYPE_SUPPORTED TRANSFORMER_ENGINE_TYPE_NAME(__nv_fp4_e2m1) #endif @@ -703,6 +744,14 @@ struct TypeExtrema { static constexpr float max_inverse = 1.0 / max; }; +#if CUDA_VERSION >= 13040 +template <> +struct TypeExtrema { + static constexpr float max = 114688.f; + static constexpr float max_inverse = 1.0 / max; +}; +#endif + template <> struct TypeExtrema { // Hex float format of 1.(7 bits of 1) * 2 ^ 127 @@ -744,6 +793,10 @@ struct TypeInfo { #if CUDA_VERSION >= 12080 , fp8e8m0 +#endif +#if CUDA_VERSION >= 13040 + , + fp8ue5m3 #endif >; #else @@ -751,6 +804,10 @@ struct TypeInfo { #if CUDA_VERSION >= 12080 , fp8e8m0 +#endif +#if CUDA_VERSION >= 13040 + , + fp8ue5m3 #endif >; #endif @@ -792,6 +849,15 @@ struct TypeInfo { #else #define SWITCH_FP4_TYPE_HANDLE(type, ...) // do nothing #endif +#if CUDA_VERSION >= 13040 +#define SWITCH_FP8UE5M3_TYPE_HANDLE(type, ...) \ + case DType::kFloat8UE5M3: { \ + using type = fp8ue5m3; \ + { __VA_ARGS__ } \ + } break; +#else +#define SWITCH_FP8UE5M3_TYPE_HANDLE(type, ...) // do nothing +#endif #define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \ switch (dtype) { \ @@ -837,11 +903,12 @@ struct TypeInfo { { __VA_ARGS__ } \ } break; \ SWITCH_FP4_TYPE_HANDLE(type, __VA_ARGS__) \ + SWITCH_FP8UE5M3_TYPE_HANDLE(type, __VA_ARGS__) \ default: \ NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ ". Expected one of: Byte, Int16, Int32, Int64, Float32, " \ "Float16, BFloat16, Float8E4M3, Float8E5M2, " \ - "Float8E8M0, Float4E2M1."); \ + "Float8E8M0."); \ } #define TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(dtype, type, ...) \ diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index a0529c80c0..c566d16f5b 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -87,6 +87,8 @@ struct GemmParam { transformer_engine::DType Atype = transformer_engine::DType::kNumTypes; transformer_engine::DType Btype = transformer_engine::DType::kNumTypes; void *A_scale_inv = nullptr; + transformer_engine::DType A_scale_inv_type = transformer_engine::DType::kNumTypes; + transformer_engine::DType B_scale_inv_type = transformer_engine::DType::kNumTypes; void *B_scale_inv = nullptr; int lda = 0; // A column strides int ldb = 0; // B column strides @@ -132,6 +134,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = transA; ret.Atype = A.data.dtype; ret.A_scale_inv = A.scale_inv.dptr; + ret.A_scale_inv_type = A.scale_inv.dtype; ret.lda = is_A_transposed ? k : m; if (!is_nvte_non_tn_fp8_gemm_supported && !is_A_transposed) { // Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data. @@ -140,6 +143,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = CUBLAS_OP_T; ret.Atype = A.columnwise_data.dtype; ret.A_scale_inv = A.columnwise_scale_inv.dptr; + ret.A_scale_inv_type = A.columnwise_scale_inv.dtype; ret.lda = k; } else { NVTE_CHECK(!is_fp8_dtype(ret.Atype), "Input A is missing column-wise usage"); @@ -153,6 +157,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = is_A_transposed ? CUBLAS_OP_N : CUBLAS_OP_T; ret.Atype = A.columnwise_data.dtype; ret.A_scale_inv = A.columnwise_scale_inv.dptr; + ret.A_scale_inv_type = A.columnwise_scale_inv.dtype; ret.lda = is_A_transposed ? m : k; } @@ -175,6 +180,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = CUBLAS_OP_T; // NVFP4 gemm is only supported in TN layout. ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype; ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr; + ret.A_scale_inv_type = is_A_transposed ? A.scale_inv.dtype : A.columnwise_scale_inv.dtype; ret.lda = k; } else if (mxfp8) { // MXFP8 GEMM. Either for pure MXFP8 recipe or backward of Hybrid NVFP4 recipe. @@ -190,6 +196,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = transA; ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype; ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr; + ret.A_scale_inv_type = is_A_transposed ? A.scale_inv.dtype : A.columnwise_scale_inv.dtype; ret.lda = is_A_transposed ? k : m; } else if (A.scaling_mode == NVTE_BLOCK_SCALING_1D || A.scaling_mode == NVTE_BLOCK_SCALING_2D) { // FP8 block scaling @@ -203,6 +210,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transA = CUBLAS_OP_T; ret.Atype = is_A_transposed ? A.data.dtype : A.columnwise_data.dtype; ret.A_scale_inv = is_A_transposed ? A.scale_inv.dptr : A.columnwise_scale_inv.dptr; + ret.A_scale_inv_type = is_A_transposed ? A.scale_inv.dtype : A.columnwise_scale_inv.dtype; ret.lda = k; // Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage @@ -223,6 +231,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = transB; ret.Btype = B.data.dtype; ret.B_scale_inv = B.scale_inv.dptr; + ret.B_scale_inv_type = B.scale_inv.dtype; ret.ldb = is_B_transposed ? n : k; if (!is_nvte_non_tn_fp8_gemm_supported && is_B_transposed) { // Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data. @@ -231,6 +240,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = CUBLAS_OP_N; ret.Btype = B.columnwise_data.dtype; ret.B_scale_inv = B.columnwise_scale_inv.dptr; + ret.B_scale_inv_type = B.columnwise_scale_inv.dtype; ret.ldb = k; } else { NVTE_CHECK(!is_fp8_dtype(ret.Btype), "Input B is missing column-wise usage"); @@ -244,6 +254,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = is_B_transposed ? CUBLAS_OP_N : CUBLAS_OP_T; ret.Btype = B.columnwise_data.dtype; ret.B_scale_inv = B.columnwise_scale_inv.dptr; + ret.B_scale_inv_type = B.columnwise_scale_inv.dtype; ret.ldb = is_B_transposed ? k : n; } @@ -264,6 +275,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = CUBLAS_OP_N; // NVFP4 gemm is only supported in TN layout. ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype; ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr; + ret.B_scale_inv_type = is_B_transposed ? B.columnwise_scale_inv.dtype : B.scale_inv.dtype; ret.ldb = k; } else if (mxfp8) { if (is_B_transposed) { @@ -275,6 +287,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = transB; ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype; ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr; + ret.B_scale_inv_type = is_B_transposed ? B.columnwise_scale_inv.dtype : B.scale_inv.dtype; ret.ldb = is_B_transposed ? n : k; } else if (B.scaling_mode == NVTE_BLOCK_SCALING_1D || B.scaling_mode == NVTE_BLOCK_SCALING_2D) { // FP8 block scaling @@ -288,6 +301,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.transB = CUBLAS_OP_N; ret.Btype = is_B_transposed ? B.columnwise_data.dtype : B.data.dtype; ret.B_scale_inv = is_B_transposed ? B.columnwise_scale_inv.dptr : B.scale_inv.dptr; + ret.B_scale_inv_type = is_B_transposed ? B.columnwise_scale_inv.dtype : B.scale_inv.dtype; ret.ldb = k; // Requirements from @@ -553,17 +567,25 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); - // Configure cuBLAS scales - fp8e4m3 *A_scale_inverse = reinterpret_cast(param.A_scale_inv); - fp8e4m3 *B_scale_inverse = reinterpret_cast(param.B_scale_inv); + // Configure cuBLAS scale pointers + void *A_scale_inverse = param.A_scale_inv; + void *B_scale_inverse = param.B_scale_inv; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &A_scale_inverse, sizeof(A_scale_inverse))); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &B_scale_inverse, sizeof(B_scale_inverse))); - scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; - scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + + // Deduce cuBLAS scale mode based on scale dtype + auto get_scale_mode = [] (DType dtype) -> cublasLtMatmulMatrixScale_t { + if (dtype == DType::kFloat8E4M3) { + return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + } + NVTE_ERROR("Unsupported dtype for NVFP4 scales (", to_string(dtype), ")."); + }; + scaling_mode_a = get_scale_mode(param.A_scale_inv_type); + scaling_mode_b = get_scale_mode(param.B_scale_inv_type); #else NVTE_ERROR("FP4 requires cuBLAS 12.8+, but compile-time cuBLAS version is ", CUBLAS_VERSION); #endif // CUBLAS_VERSION >= 120800 diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 3997e5249d..3814aae8b9 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -16,6 +16,7 @@ #include #include "../cast/mxfp8/swizzle.cuh" +#include "../cast/nvfp4/core_nvfp4.cuh" #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/handle_manager.h" @@ -371,6 +372,7 @@ struct GroupedOperandSelection { void *scale_inv = nullptr; // Contiguous array of scales (input) void *amax = nullptr; // Per-tensor amax values (NVFP4 only) transformer_engine::DType dtype = transformer_engine::DType::kNumTypes; + transformer_engine::DType scale_inv_dtype = transformer_engine::DType::kNumTypes; NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; bool with_gemm_swizzled_scales = false; bool trans = false; @@ -776,6 +778,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: auto use_columnwise = [&](bool storage_transposed = true) { sel.dptr = static_cast(t->columnwise_data.dptr); sel.scale_inv = t->columnwise_scale_inv.dptr; + sel.scale_inv_dtype = t->columnwise_scale_inv.dtype; sel.amax = t->columnwise_amax.dptr; sel.dtype = col_dtype; sel.rowwise = false; @@ -787,6 +790,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: auto use_rowwise = [&]() { sel.dptr = static_cast(t->data.dptr); sel.scale_inv = t->scale_inv.dptr; + sel.scale_inv_dtype = t->scale_inv.dtype; sel.amax = t->amax.dptr; sel.dtype = row_dtype; sel.rowwise = true; @@ -889,25 +893,38 @@ inline void set_mxfp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, #endif // CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION } -// Configures cuBLAS for NVFP4 grouped GEMM: sets VEC16_UE4M3 scale mode and scale pointers -// for both A and B. Requires cuBLAS 13.4+. +// Configures cuBLAS for NVFP4 grouped GEMM: sets VEC16_UE4M3 or VEC16_UE5M3 scale mode +// and scale pointers for both A and B. Requires cuBLAS 13.4+. inline void set_nvfp4_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, - void **a_scale_inv_ptrs, void **b_scale_inv_ptrs) { + void **a_scale_inv_ptrs, void **b_scale_inv_ptrs, + transformer_engine::DType a_scale_inv_dtype, + transformer_engine::DType b_scale_inv_dtype) { #if CUBLAS_VERSION >= CUBLAS_NVFP4_GROUPED_GEMM_VERSION NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_NVFP4_GROUPED_GEMM_VERSION, "NVFP4 grouped GEMM requires cuBLAS 13.4+, but run-time cuBLAS version is ", transformer_engine::cuda::cublas_version()); - const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, - &scale_mode, sizeof(scale_mode))); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, - &scale_mode, sizeof(scale_mode))); + + // Configure scale pointers NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs))); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs))); + + // Configure scale mode based on dtype + auto get_scale_mode = [](transformer_engine::DType dtype) -> cublasLtMatmulMatrixScale_t { + if (dtype == transformer_engine::DType::kFloat8E4M3) { + return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + } + NVTE_ERROR("Unsupported dtype for NVFP4 scales (", transformer_engine::to_string(dtype), ")."); + }; + const cublasLtMatmulMatrixScale_t scale_mode_a = get_scale_mode(a_scale_inv_dtype); + const cublasLtMatmulMatrixScale_t scale_mode_b = get_scale_mode(b_scale_inv_dtype); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &scale_mode_a, sizeof(scale_mode_a))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &scale_mode_b, sizeof(scale_mode_b))); #else NVTE_CHECK(false, "NVFP4 grouped GEMM requires cuBLAS 13.4+, but compile-time " @@ -1052,7 +1069,8 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac setup_workspace.b_scale_inv_ptrs); } else if (transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { set_nvfp4_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, - setup_workspace.b_scale_inv_ptrs); + setup_workspace.b_scale_inv_ptrs, + A_sel.scale_inv_dtype, B_sel.scale_inv_dtype); } else if (transformer_engine::is_fp8_block_scaling(A_sel.scaling_mode)) { set_fp8_block_scaling_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, setup_workspace.b_scale_inv_ptrs, A_sel.scaling_mode, @@ -1334,7 +1352,8 @@ __global__ void setup_grouped_gemm_kernel( MultiTensorGroupGemmOutputArgs c_multi_tensor_args, MultiTensorGroupGemmOutputArgs d_multi_tensor_args, // NVFP4: per-group amax values and output buffer for computed alpha - float *a_amax, float *b_amax, float *nvfp4_computed_alpha) { + float *a_amax, float *b_amax, float *nvfp4_computed_alpha, + float a_unit_global_scale_amax, float b_unit_global_scale_amax) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_tensors) return; @@ -1401,21 +1420,21 @@ __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; + float a_amax_val = a_unit_global_scale_amax; 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) { + const float b_amax_val = b_amax == nullptr ? b_unit_global_scale_amax : b_amax[idx]; + const float nvfp4_alpha_factor_inv = + 1.0f / (a_unit_global_scale_amax * b_unit_global_scale_amax); + nvfp4_computed_alpha[idx] = + alpha_ptr[idx] * a_amax_val * b_amax_val * nvfp4_alpha_factor_inv; alpha_ptrs[idx] = &nvfp4_computed_alpha[idx]; } else { alpha_ptrs[idx] = alpha_ptr + idx; @@ -1544,10 +1563,19 @@ 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); + float a_unit_global_scale_amax = 1.0f; + float b_unit_global_scale_amax = 1.0f; + if (needs_nvfp4_alpha) { + constexpr float kFP4Max = + transformer_engine::detail::TypeExtrema::max; + a_unit_global_scale_amax = + transformer_engine::dispatch::nvfp4::core::scale_max(A_sel.scale_inv_dtype) * + kFP4Max; + b_unit_global_scale_amax = + transformer_engine::dispatch::nvfp4::core::scale_max(B_sel.scale_inv_dtype) * + kFP4Max; + } 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, @@ -1560,7 +1588,8 @@ inline void launch_grouped_gemm_setup( B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args, A_sel.amax ? static_cast(A_sel.amax) : nullptr, B_sel.amax ? static_cast(B_sel.amax) : nullptr, - needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr); + needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr, a_unit_global_scale_amax, + b_unit_global_scale_amax); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -1619,14 +1648,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); @@ -1776,11 +1797,8 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num 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."); + const auto& A_tensor0 = *transformer_engine::convertNVTETensorCheck(A_list[0]); + A_sel.scale_inv_dtype = transa ? A_tensor0.scale_inv.dtype : A_tensor0.columnwise_scale_inv.dtype; } // Workspaces: setup (pointer arrays) and cuBLAS @@ -1861,12 +1879,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 0f2456c975..c4ff0f445b 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 @@ -15,6 +15,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -692,7 +693,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : __ldg(amax_colwise + g); } size_t rng_seed = 0; @@ -745,15 +749,12 @@ __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 - 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 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 - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + c_global_amax_val); float global_decode_scale = 1.0f / global_encode_scale; // Scaling factor for fast math path @@ -776,11 +777,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g group_idx = cur_group_idx; c_global_amax_val = shared_storage.global_d_amax[group_idx]; // update amax - global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + c_global_amax_val); global_decode_scale = 1.0f / global_encode_scale; global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; // TODO(zhongbo): double check the logic here @@ -946,7 +945,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : __ldg(amax_rowwise + g); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { @@ -1002,14 +1004,11 @@ __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 - 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 fp4_max_inv = 1.0f / fp4_max; - float global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + a_global_amax_val); float global_decode_scale = 1.0f / global_encode_scale; float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; @@ -1026,11 +1025,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g group_idx = cur_group_idx; a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Update group quantization parameters/scaling - global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + a_global_amax_val); global_decode_scale = 1.0f / global_encode_scale; global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; } @@ -1331,6 +1328,14 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, bool all_has_row_quant = output->has_data(); bool all_has_col_quant = output->has_columnwise_data(); + NVTE_CHECK(all_has_row_quant || all_has_col_quant, + "Output grouped tensor must have rowwise or columnwise quantization."); + const DType scale_dtype = + all_has_row_quant ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + if (all_has_row_quant && all_has_col_quant) { + NVTE_CHECK(output->columnwise_scale_inv.dtype == scale_dtype, + "Rowwise and columnwise NVFP4 scales must use the same dtype."); + } // Stochastic rounding config const bool use_stochastic_rounding = quant_config.stochastic_rounding; @@ -1356,9 +1361,7 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, using TA = cute::bfloat16_t; using TB = cute::bfloat16_t; using TD = cutlass::float_e2m1_t; - using TSFD = cutlass::float_ue4m3_t; using TQA = TD; - using TSFA = TSFD; checkCuDriverContext(stream); @@ -1397,10 +1400,9 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, } TQA *const rowwise_data_base_ptr = reinterpret_cast(output->data.dptr); - TSFA *const rowwise_scale_inv_base_ptr = reinterpret_cast(output->scale_inv.dptr); + void *const rowwise_scale_inv_base_ptr = output->scale_inv.dptr; TQA *const colwise_data_base_ptr = reinterpret_cast(output->columnwise_data.dptr); - TSFA *const colwise_scale_inv_base_ptr = - reinterpret_cast(output->columnwise_scale_inv.dptr); + void *const colwise_scale_inv_base_ptr = output->columnwise_scale_inv.dptr; float *const amax_rowwise_base_ptr = reinterpret_cast(output->amax.dptr); float *const amax_colwise_base_ptr = reinterpret_cast(output->columnwise_amax.dptr); @@ -1418,45 +1420,50 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, const bool use_swizzle_sf_output = output->with_gemm_swizzled_scales; - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_stochastic_rounding, kEnableStochasticRounding, + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, TRANSFORMER_ENGINE_SWITCH_CONDITION( - all_has_col_quant, kEnableRhtColQuant, + use_stochastic_rounding, kEnableStochasticRounding, TRANSFORMER_ENGINE_SWITCH_CONDITION( - all_has_row_quant, kEnableRowQuant, + all_has_col_quant, kEnableRhtColQuant, TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_swizzle_sf_output, kEnableSwizzleSFOutput, + all_has_row_quant, kEnableRowQuant, TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config.use_fast_math, kUseFastMath, - - if constexpr (kEnableRhtColQuant || kEnableRowQuant) { - detail::group_row_col_rht_gemm_ntt_w_sfc_graph_safe< - kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, - kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>( - /*packed_sequence_length=*/first_logical_dim, - /*hidden_size=*/last_logical_dim, - /*num_tensors=*/num_tensors, - /*shape_rep=*/shape_rep, - /*A=*/reinterpret_cast(input_base_ptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*QA=*/reinterpret_cast(rowwise_data_base_ptr), - /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), - /*QA_COLWISE=*/reinterpret_cast(colwise_data_base_ptr), - /*SFA_COLWISE=*/reinterpret_cast(colwise_scale_inv_base_ptr), - /*amax_rowwise=*/reinterpret_cast(amax_rowwise_base_ptr), - /*amax_colwise=*/reinterpret_cast(amax_colwise_base_ptr), - /*offsets=*/offsets_ptr, - /*first_dims=*/first_dims_ptr, - /*rng_state=*/rng_state, - /*tile_scheduler_workspace=*/tile_scheduler_workspace, - /*sm_count=*/sm_count, - /*stream=*/stream, /*k_tile_size=*/k_tile_size); - } else { - NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", - kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); - } - - ););););); + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::group_row_col_rht_gemm_ntt_w_sfc_graph_safe< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TQA, ScaleType, TD, ScaleType, + kUseFastMath>( + /*packed_sequence_length=*/first_logical_dim, + /*hidden_size=*/last_logical_dim, + /*num_tensors=*/num_tensors, + /*shape_rep=*/shape_rep, + /*A=*/reinterpret_cast(input_base_ptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*QA=*/reinterpret_cast(rowwise_data_base_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), + /*QA_COLWISE=*/reinterpret_cast(colwise_data_base_ptr), + /*SFA_COLWISE=*/ + reinterpret_cast(colwise_scale_inv_base_ptr), + /*amax_rowwise=*/reinterpret_cast(amax_rowwise_base_ptr), + /*amax_colwise=*/reinterpret_cast(amax_colwise_base_ptr), + /*offsets=*/offsets_ptr, + /*first_dims=*/first_dims_ptr, + /*rng_state=*/rng_state, + /*tile_scheduler_workspace=*/tile_scheduler_workspace, + /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, + ")."); + } + + );););););) } } // namespace transformer_engine 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 4b1435f9eb..f977c3651e 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 @@ -15,6 +15,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -81,17 +82,6 @@ __device__ __forceinline__ int GetTensorId(MultiAmaxHadamardCastFusionArgs *kern return tensor_id; } -// 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; - // 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()); - // If global amax is 0 or infinity, return 1 - return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale; -} - template struct SharedStorage { static constexpr int AccumulatorPipelineStageCount = 16; @@ -469,7 +459,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 = 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 @@ -506,8 +496,11 @@ __global__ static void group_rht_gemm_device( Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); - float global_amax_val = *global_amax_ptr; - float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); + constexpr float kUnitGlobalScaleAmax = + dispatch::nvfp4::core::scale_max() * TypeExtrema::max; + float global_amax_val = global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; + float global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4(global_amax_val); // Scaling factor for fast math path float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; @@ -527,8 +520,10 @@ __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_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); + global_amax_val = global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; + global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + global_amax_val); global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; global_decode_scale = 1.0f / global_encode_scale; tensor_id = new_tensor_id; @@ -864,12 +859,20 @@ void group_hadamard_transform_cast_fusion_columnwise( MultiAmaxHadamardCastFusionArgs kernel_args; kernel_args.num_tensors = 0; kernel_args.split_sections_range[0] = 0; + DType scale_dtype = DType::kNumTypes; for (size_t i = 0; i < num_tensors; ++i) { NVTE_CHECK(split_sections[i] % 64 == 0, "component ", i, " of split_sections should be 64 multiple"); if (split_sections[i] == 0) { continue; } + const DType output_scale_dtype = output_list[i]->scale_inv.dtype; + if (scale_dtype == DType::kNumTypes) { + scale_dtype = output_scale_dtype; + } else { + NVTE_CHECK(output_scale_dtype == scale_dtype, + "All grouped NVFP4 outputs must use the same scale dtype."); + } kernel_args.global_amax_list[kernel_args.num_tensors] = reinterpret_cast(output_list[i]->amax.dptr); // TODO(zhongbo): should we change API assumption to use columnwise_data instead of data? @@ -899,7 +902,6 @@ void group_hadamard_transform_cast_fusion_columnwise( using TA = cute::bfloat16_t; using TB = cute::bfloat16_t; using TC = cutlass::float_e2m1_t; - using TSFC = cutlass::float_ue4m3_t; checkCuDriverContext(stream); @@ -958,16 +960,18 @@ void group_hadamard_transform_cast_fusion_columnwise( k_tile_size = 512; } - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_stochastic_rounding, kUseStochasticRounding, + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, TSFC, TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config.use_fast_math, kUseFastMath, - detail::group_rht_gemm_ttt_wrapper( - /*m=*/m, /*n=*/n, /*A=*/reinterpret_cast(input.dptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*kernel_args_ptr=*/&kernel_args, /*rng_state=*/rng_state, /*sm_count=*/sm_count, - /*stream=*/stream, /*k_tile_size=*/k_tile_size););); + use_stochastic_rounding, kUseStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + detail::group_rht_gemm_ttt_wrapper( + /*m=*/m, /*n=*/n, /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*kernel_args_ptr=*/&kernel_args, /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size);););) } } // namespace transformer_engine 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 2e6d383ce1..6f560cb232 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 @@ -15,6 +15,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -680,8 +681,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : __ldg(amax_ptr); } size_t rng_seed = 0; @@ -727,15 +731,12 @@ __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 - 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 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 - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + c_global_amax_val); float global_decode_scale = 1.0f / global_encode_scale; // Scaling factor for fast math path @@ -756,11 +757,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( group_idx = cur_group_idx; c_global_amax_val = shared_storage.global_d_amax[group_idx]; // update amax - global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + c_global_amax_val); global_decode_scale = 1.0f / global_encode_scale; global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; cur_N = args.split_sections[group_idx]; @@ -924,8 +923,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : __ldg(amax_ptr); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { @@ -979,14 +981,11 @@ __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 - 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 fp4_max_inv = 1.0f / fp4_max; - float global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + a_global_amax_val); float global_decode_scale = 1.0f / global_encode_scale; float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; @@ -1002,11 +1001,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( group_idx = cur_group_idx; a_global_amax_val = shared_storage.global_a_amax[group_idx]; // Update group quantization parameters/scaling - global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + a_global_amax_val); global_decode_scale = 1.0f / global_encode_scale; global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; } @@ -1320,6 +1317,7 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vectorscale_inv.dtype : output_list[i]->columnwise_scale_inv.dtype; + if (has_row_quant && has_col_quant) { + NVTE_CHECK(output_list[i]->columnwise_scale_inv.dtype == output_scale_dtype, + "Rowwise and columnwise NVFP4 scales must use the same dtype."); + } + if (scale_dtype == DType::kNumTypes) { + scale_dtype = output_scale_dtype; + } else { + NVTE_CHECK(output_scale_dtype == scale_dtype, + "All grouped NVFP4 outputs must use the same scale dtype."); + } void *amax_rowwise_ptr = has_row_quant ? reinterpret_cast(output_list[i]->amax.dptr) : nullptr; void *amax_colwise_ptr = @@ -1386,9 +1396,7 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector( - /*packed_sequence_length=*/m, /*hidden_size=*/n, - /*A=*/reinterpret_cast(input.dptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*QA=*/reinterpret_cast(rowwise_data_base_ptr), - /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), - /*args=*/kernel_args, - /*rng_state=*/rng_state, - /*tile_scheduler_workspace=*/tile_scheduler_workspace, - /*sm_count=*/sm_count, - /*stream=*/stream, /*k_tile_size=*/k_tile_size); - } else { - NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", - kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); - } - - ););););); + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::group_row_col_rht_gemm_ntt_w_sfc< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TQA, ScaleType, TD, ScaleType, + kUseFastMath>( + /*packed_sequence_length=*/m, /*hidden_size=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*QA=*/reinterpret_cast(rowwise_data_base_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), + /*args=*/kernel_args, + /*rng_state=*/rng_state, + /*tile_scheduler_workspace=*/tile_scheduler_workspace, + /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, + ")."); + } + + );););););) } } // namespace transformer_engine 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 433da1f0f0..a0d781b104 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -15,6 +15,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -40,17 +41,6 @@ using namespace cute; using cute::Tensor; // Avoid conflict with transformer_engine::Tensor 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; - // 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()); - // If global amax is 0 or infinity, return 1 - return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale; -} - template = 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)); } @@ -412,7 +402,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 = + dispatch::nvfp4::core::scale_max() * TypeExtrema::max; + 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(); @@ -427,9 +420,11 @@ 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_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4(global_amax_val); const float global_decode_scale = 1.0f / global_encode_scale; // Scaling factor for fast math path @@ -758,7 +753,6 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out using TA = cute::bfloat16_t; using TB = cute::bfloat16_t; using TC = cutlass::float_e2m1_t; - using TSFC = cutlass::float_ue4m3_t; checkCuDriverContext(stream); @@ -819,22 +813,24 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out k_tile_size = 512; } - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_stochastic_rounding, kUseStochasticRounding, + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_inv_t.dtype, TSFC, TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config.use_fast_math, kUseFastMath, - detail::rht_gemm_ttt_wrapper( - /*m=*/m, - /*n=*/n, - /*A=*/reinterpret_cast(input.dptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*C=*/reinterpret_cast(output_t.dptr), - /*SFC=*/reinterpret_cast(scale_inv_t.dptr), - /*global_amax=*/reinterpret_cast(global_amax.dptr), - /*rng_state=*/rng_state, - /*sm_count=*/sm_count, - /*stream=*/stream, - /*k_tile_size=*/k_tile_size););); + use_stochastic_rounding, kUseStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + detail::rht_gemm_ttt_wrapper( + /*m=*/m, + /*n=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*C=*/reinterpret_cast(output_t.dptr), + /*SFC=*/reinterpret_cast(scale_inv_t.dptr), + /*global_amax=*/reinterpret_cast(global_amax.dptr), + /*rng_state=*/rng_state, + /*sm_count=*/sm_count, + /*stream=*/stream, + /*k_tile_size=*/k_tile_size);););) } } // namespace transformer_engine 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 8d8ab20165..9c06f62eb6 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 @@ -15,6 +15,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" @@ -403,10 +404,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)); } @@ -651,7 +652,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : *c_global_amax; auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{}); auto bulk_tmem_epilogue_layout = make_layout( acc_epilogue_pipelined_shape, @@ -708,14 +712,12 @@ __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 - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_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{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float const global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + c_global_amax_val); float const global_decode_scale = 1.0f / global_encode_scale; // Scaling factor for fast math path @@ -858,7 +860,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 + ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + : *a_global_amax; int global_thread_idx = threadIdx.x; int local_thread_idx = global_thread_idx % 256; size_t rng_seed = 0; @@ -904,14 +909,12 @@ __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 - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; + static constexpr float fp4_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{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; + float const global_encode_scale = + dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + a_global_amax_val); float const global_decode_scale = 1.0f / global_encode_scale; // Scaling factor for fast math path @@ -1262,6 +1265,12 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, NVTE_CHECK(has_rowwise_quant || has_columnwise_quant, "Output tensor must have rowwise or columnwise quant."); + const DType scale_dtype = + has_rowwise_quant ? output_.scale_inv.dtype : output_.columnwise_scale_inv.dtype; + if (has_rowwise_quant && has_columnwise_quant) { + NVTE_CHECK(output_.columnwise_scale_inv.dtype == scale_dtype, + "Rowwise and columnwise NVFP4 scales must use the same dtype."); + } // Stochastic rounding config const bool use_stochastic_rounding = quant_config.stochastic_rounding; @@ -1279,9 +1288,7 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, using TA = cute::bfloat16_t; using TB = cute::bfloat16_t; using TD = cutlass::float_e2m1_t; - using TSFD = cutlass::float_ue4m3_t; using TQA = TD; - using TSFA = TSFD; checkCuDriverContext(stream); @@ -1320,38 +1327,43 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, // nvte_swizzle_scaling_factors pass between quantize and GEMM. const bool use_swizzle_sf_output = output_.with_gemm_swizzled_scales; - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_stochastic_rounding, kEnableStochasticRounding, + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, TRANSFORMER_ENGINE_SWITCH_CONDITION( - has_columnwise_quant, kEnableRhtColQuant, + use_stochastic_rounding, kEnableStochasticRounding, TRANSFORMER_ENGINE_SWITCH_CONDITION( - has_rowwise_quant, kEnableRowQuant, + has_columnwise_quant, kEnableRhtColQuant, TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_swizzle_sf_output, kEnableSwizzleSFOutput, + has_rowwise_quant, kEnableRowQuant, TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config.use_fast_math, kUseFastMath, - - if constexpr (kEnableRhtColQuant || kEnableRowQuant) { - detail::row_col_rht_gemm_ntt_w_sfc< - kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, - kEnableSwizzleSFOutput, TA, TB, TD, TSFD, TQA, TSFA, kUseFastMath>( - /*sequence_length=*/m, /*hidden_size=*/n, - /*A=*/reinterpret_cast(input.dptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*D=*/reinterpret_cast(columnwise_data_ptr), - /*SFD=*/reinterpret_cast(columnwise_scale_inv_ptr), - /*QA=*/reinterpret_cast(rowwise_data_ptr), - /*SFA=*/reinterpret_cast(rowwise_scale_inv_ptr), - /*a_global_amax=*/reinterpret_cast(rowwise_amax_ptr), - /*d_global_amax=*/reinterpret_cast(columnwise_amax_ptr), - /*rng_state=*/rng_state, /*sm_count=*/sm_count, - /*stream=*/stream, /*k_tile_size=*/k_tile_size); - } else { - NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", - kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); - } - - ););););); + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::row_col_rht_gemm_ntt_w_sfc< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TD, ScaleType, TQA, ScaleType, + kUseFastMath>( + /*sequence_length=*/m, /*hidden_size=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*D=*/reinterpret_cast(columnwise_data_ptr), + /*SFD=*/reinterpret_cast(columnwise_scale_inv_ptr), + /*QA=*/reinterpret_cast(rowwise_data_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_ptr), + /*a_global_amax=*/reinterpret_cast(rowwise_amax_ptr), + /*d_global_amax=*/ + reinterpret_cast(columnwise_amax_ptr), + /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, + ")."); + } + + );););););) } } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 47539a89a1..b98b87c5ba 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -14,7 +14,10 @@ #include "transformer_engine.h" #ifdef __cplusplus +#define NVTE_NVFP4_SCALE_DTYPE_DEFAULT = kNVTEFloat8E4M3 extern "C" { +#else +#define NVTE_NVFP4_SCALE_DTYPE_DEFAULT #endif /*! \brief Update FP8 scaling factors with delayed scaling recipe. @@ -374,32 +377,36 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s * \param[in] start_offset Starting element offset in the flattened tensor. * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). * \param[in] stream CUDA stream used for the operation. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, cudaStream_t stream); + size_t block_len, cudaStream_t stream, + const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); -/*! \brief Expand tile-level scales to row-level scales and convert to FP8 E4M3, used in partial cast. +/*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. * * Each tile row's scale is repeated block_len times in the output. * * \param[in] input Input tensor with tile scales [tile_rows, tile_cols], float32. - * \param[out] output Output tensor with expanded scales [rows_padded, tile_cols], uint8 (E4M3). + * \param[out] output Output tensor with expanded scales [rows_padded, tile_cols], uint8. * \param[in] tile_rows Number of tile rows. * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Padded row count in output. * \param[in] block_len Block length (typically 16 for NVFP4). * \param[in] stream CUDA stream. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream); + cudaStream_t stream, + const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); /*! \brief Compute per-block decode scale from block amax and global amax. * * Computes: - * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * global_scale = (scale_max * fp4_max) / global_amax * per_block_decode_scale = block_amax / fp4_max * global_scale * * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh. @@ -408,49 +415,57 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s * \param[out] scale Output scale tensor [tile_rows, tile_cols], float32. * \param[in] global_amax Global amax tensor (single element), float32. Avoids D2H transfer. * \param[in] stream CUDA stream. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, cudaStream_t stream); + const NVTETensor global_amax, cudaStream_t stream, + const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); /*! \brief Fused kernel for NVFP4 scale computation. * * Fuses three operations into one kernel: * 1. Compute per-block decode scales from block amax and global amax * 2. Copy global amax to target tensor - * 3. Expand tile-level scales to row-level and convert to FP8 E4M3 + * 3. Expand tile-level scales to row-level and convert to the selected FP8 scale type * * Saves 2 kernel launches per parameter. * * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. * \param[in] global_amax Global amax tensor [1], float32. * \param[out] per_block_scale Output per-block scale [tile_rows, tile_cols], float32 (for partial_cast). - * \param[out] target_scale Output scale tensor [rows_padded, tile_cols], uint8 (E4M3). + * \param[out] target_scale Output scale tensor [rows_padded, tile_cols], uint8. * \param[out] target_amax Output amax tensor [1], float32 (copy of global_amax). * \param[in] tile_rows Number of tile rows. * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Total padded rows in output. * \param[in] block_len Block length (16 for NVFP4). * \param[in] stream CUDA stream. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, cudaStream_t stream); + size_t rows_padded, size_t block_len, cudaStream_t stream, + const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); /*! \brief Compute global encode scale from global amax. * - * Computes: global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * Computes: global_scale = (scale_max * fp4_max) / global_amax * If global_amax <= 0, returns 1.0. * * \param[in] global_amax Input global amax tensor [num_params], float32. * \param[out] global_scale Output global scale tensor [num_params], float32. * \param[in] stream CUDA stream. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - cudaStream_t stream); + cudaStream_t stream, + const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); #ifdef __cplusplus } // extern "C" #endif +#undef NVTE_NVFP4_SCALE_DTYPE_DEFAULT + #endif // TRANSFORMER_ENGINE_RECIPE_H_ diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index aa0405e177..ff09194255 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -34,6 +34,7 @@ enum NVTEDType { kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ + kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ kNVTENumTypes /*!< Number of supported types */ }; @@ -83,11 +84,12 @@ enum NVTETensorParam { * its values are populated during quantization. */ kNVTERowScaledNVFP4 = 8, - /*! Global E4M3 scale bound used by an NVFP4 tensor. + /*! Global scale-bound selector used by an NVFP4 tensor. * * This is part of the tensor data contract. Downstream dequantization and * GEMM scale consumers must use the same bound used during quantization. * Standard NVFP4 uses 448; 4over6 may use 256 for map-to-4 headroom. + * For UE5M3 scales, these settings map to 114688 and 65536, respectively. */ kNVTENVFP4E4M3Max = 9, kNVTENumTensorParams @@ -687,12 +689,16 @@ enum class DType { kFloat8E5M2 = 8, kFloat8E8M0 = 9, kFloat4E2M1 = 10, + kFloat8UE5M3 = 11, kNumTypes }; /*! \brief Check if TE datatype is FP8 * - * Return true if TE datatype is FP8 + * Return whether datatype is FP8 E4M3 or FP8 E5M2. Other FP8 formats + * (E8M0, UE5M3) are not used as primary data encoding, but are + * auxiliary types for block scaling formats. + * * \param[in] t TE Datatype of interest */ inline bool is_fp8_dtype(const DType t) { diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index a89ddba917..525238b3e0 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -16,7 +16,6 @@ _NVFP4_4OVER6_SCOPES = ("none", "weights", "activations", "all") _NVFP4_4OVER6_ERR_MODES = ("MAE", "MSE") - class _FormatHelper(NamedTuple): """ Stores max FP8 values for fprop and bprop a `Format`. @@ -28,26 +27,29 @@ class _FormatHelper(NamedTuple): class Format(Enum): """ - Supported FP8 formats. - Supported FP4 formats. + Low precision data formats. Values ------ E2M1 : - All FP4 tensors are in e2m1 format + FP4 type with e2m1 format E4M3 : - All FP8 tensors are in e4m3 format + FP8 type with e4m3 format E5M2 : - All FP8 tensors are in e5m2 format + FP8 type with e5m2 format HYBRID : FP8 tensors in the forward pass are in e4m3 format, FP8 tensors in the backward pass are in e5m2 format + UE5M3 : + FP8 type with ue5m3 format + """ E2M1 = _FormatHelper(max_fwd=6, max_bwd=6) E4M3 = _FormatHelper(max_fwd=448, max_bwd=448) E5M2 = _FormatHelper(max_fwd=57344, max_bwd=57344) HYBRID = _FormatHelper(max_fwd=E4M3.max_fwd, max_bwd=E5M2.max_bwd) + UE5M3 = _FormatHelper(max_fwd=114688, max_bwd=114688) @dataclass(frozen=True) @@ -261,7 +263,7 @@ def scaling_factor_compute(amax: Tensor, backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: - assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." @@ -312,7 +314,7 @@ class Float8CurrentScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: - assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." @@ -370,7 +372,7 @@ class MXFP8BlockScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: - assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." @@ -457,7 +459,7 @@ def __post_init__(self) -> None: assert ( not self.fp8_dpa and not self.fp8_mha ), "FP8 attention is not supported for Float8BlockScaling." - assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." @@ -571,7 +573,7 @@ class NVFP4BlockScaling(Recipe): def __post_init__(self) -> None: assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" - assert self.fp8_format == Format.E4M3, "Only E4M3 is supported for NVFP4 scaling" + assert self.fp8_format in (Format.E4M3, Format.UE5M3), "Unsupported format for NVFP4 scaling." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 576e6139c7..7047c24e68 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -10,6 +10,7 @@ #include #include +#include "../cast/nvfp4/core_nvfp4.cuh" #include "../common.h" #include "../util/ptx.cuh" #include "../utils.cuh" @@ -70,11 +71,13 @@ constexpr int kThreadsPerBlock = 256; // Kernel to compute alpha *= amax_A * amax_B / factor __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; - 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_B, float scale_max_A, + float scale_max_B, float *alpha_out) { + constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; + const float factor_inv = 1.0f / (fp4_max * fp4_max * scale_max_A * scale_max_B); + const float amax_A_value = amax_A == nullptr ? scale_max_A * fp4_max : *amax_A; + const float amax_B_value = amax_B == nullptr ? scale_max_B * fp4_max : *amax_B; + *alpha_out = alpha_in * amax_A_value * amax_B_value * factor_inv; } template @@ -126,7 +129,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } } -template +template __global__ void __launch_bounds__(kThreadsPerBlock) nvfp4_2d_partial_cast_kernel(const IType *input, uint8_t *output, const float *decode_scale_ptr, const size_t scale_stride_h, const size_t scale_stride_w, @@ -152,7 +155,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) const float global_decode_scale = 1.0f / global_encode_scale; float tile_decode_scale = decode_scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w]; - tile_decode_scale = static_cast(static_cast(tile_decode_scale)); + tile_decode_scale = static_cast(static_cast(tile_decode_scale)); constexpr float kFp32Max = 3.402823466e+38F; float tile_encode_val = (tile_decode_scale > 0.f) ? 1.0f / (tile_decode_scale * global_decode_scale) : kFp32Max; @@ -289,7 +292,7 @@ void nvfp4_2d_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size void nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale, const Tensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, size_t block_len, - cudaStream_t stream) { + DType scale_dtype, cudaStream_t stream) { NVTE_CHECK(block_len == 16, "NVFP4 2D supports 16x16 tiles only (block_len = 16)."); NVTE_CHECK(out.dtype() == DType::kByte, "NVFP4 rowwise data must be uint8."); @@ -305,16 +308,19 @@ void nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale, assert(blocks_y <= std::numeric_limits::max()); dim3 grid(blocks_x, blocks_y); - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - inp.dtype(), inp_dtype, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - w % kTileDim == 0, kWidthAligned, - nvfp4_2d_partial_cast_kernel - <<>>( - reinterpret_cast(inp.data.dptr), - reinterpret_cast(out.data.dptr), - reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w, - reinterpret_cast(global_scale.data.dptr), h, w, start_offset, len);)) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + inp.dtype(), inp_dtype, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + w % kTileDim == 0, kWidthAligned, + nvfp4_2d_partial_cast_kernel + <<>>( + reinterpret_cast(inp.data.dptr), + reinterpret_cast(out.data.dptr), + reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w, + reinterpret_cast(global_scale.data.dptr), h, w, start_offset, + len);))) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -487,7 +493,7 @@ void nvfp4_transpose(const Tensor input, Tensor output, cudaStream_t stream) { * NVFP4 SCALE TRANSPOSE KERNEL * * Transposes tile-level scales from rowwise to columnwise format. - * Scale values are stored as E4M3 (fp8) in uint8 tensors. + * Scale values are stored as raw FP8 scale bytes in uint8 tensors. * * Input (rowwise_scale_inv): [M_padded, K_tiles] where scales are stored * at every 16th row (i.e., row 0, 16, 32, ... contain the actual scales, @@ -502,8 +508,8 @@ void nvfp4_transpose(const Tensor input, Tensor output, cudaStream_t stream) { * --------------------------------------------------------------------------- */ __global__ void nvfp4_scale_transpose_kernel( - const uint8_t *__restrict__ input, // [M_padded, K_tiles], E4M3 stored as uint8 - uint8_t *__restrict__ output, // [K_padded, M_tiles], E4M3 stored as uint8 + const uint8_t *__restrict__ input, // [M_padded, K_tiles], FP8 scale bytes + uint8_t *__restrict__ output, // [K_padded, M_tiles], FP8 scale bytes const size_t M_tiles, // Number of M tiles const size_t K_tiles, // Number of K tiles const size_t input_stride, // K_tiles (input row stride) @@ -532,8 +538,8 @@ __global__ void nvfp4_scale_transpose_kernel( void nvfp4_scale_transpose(const Tensor input, Tensor output, size_t M_tiles, size_t K_tiles, cudaStream_t stream) { - NVTE_CHECK(input.dtype() == DType::kByte, "NVFP4 scale transpose input must be uint8 (E4M3)."); - NVTE_CHECK(output.dtype() == DType::kByte, "NVFP4 scale transpose output must be uint8 (E4M3)."); + NVTE_CHECK(input.dtype() == DType::kByte, "NVFP4 scale transpose input must be uint8."); + NVTE_CHECK(output.dtype() == DType::kByte, "NVFP4 scale transpose output must be uint8."); const auto in_shape = input.shape(); const auto out_shape = output.shape(); @@ -561,17 +567,19 @@ void nvfp4_scale_transpose(const Tensor input, Tensor output, size_t M_tiles, si * --------------------------------------------------------------------------- * NVFP4 SCALE EXPANSION KERNEL * - * Expands tile-level scales to row-level scales and converts to FP8 E4M3, used in partial cast. + * Expands tile-level scales to row-level scales and converts to the selected FP8 scale format, + * used in partial cast. * * Input (per_block_decode_scale): [tile_rows, tile_cols] in float32 - * Output (target_scale): [rows_padded, tile_cols] in uint8 (E4M3) + * Output (target_scale): [rows_padded, tile_cols] in uint8 (E4M3 or UE5M3) * * Each tile row's scale is repeated block_len times in the output. * --------------------------------------------------------------------------- */ +template __global__ void nvfp4_expand_scale_to_fp8_kernel( const float *__restrict__ input, // [tile_rows, tile_cols] - uint8_t *__restrict__ output, // [rows_padded, tile_cols] + ScaleType *__restrict__ output, // [rows_padded, tile_cols] const size_t tile_rows, const size_t tile_cols, const size_t rows_padded, const size_t block_len) { const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y; @@ -587,17 +595,15 @@ __global__ void nvfp4_expand_scale_to_fp8_kernel( scale_val = input[tile_row * tile_cols + out_col]; } - // Convert float32 to FP8 E4M3 - // Clamp to FP8 E4M3 range and convert - fp8e4m3 fp8_val = static_cast(scale_val); - output[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val); + output[out_row * tile_cols + out_col] = static_cast(scale_val); } void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream) { + DType scale_dtype, cudaStream_t stream) { NVTE_CHECK(input.dtype() == DType::kFloat32, "Scale input must be float32."); - NVTE_CHECK(output.dtype() == DType::kByte, "Scale output must be uint8 (E4M3)."); + NVTE_CHECK(output.dtype() == DType::kByte || output.dtype() == scale_dtype, + "Scale output must be byte storage or have the selected NVFP4 scale dtype."); if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return; @@ -605,9 +611,12 @@ void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_ro dim3 block(kBlockDim, kBlockDim); dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim); - nvfp4_expand_scale_to_fp8_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output.data.dptr), tile_rows, tile_cols, rows_padded, block_len); + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4_expand_scale_to_fp8_kernel + <<>>(reinterpret_cast(input.data.dptr), + reinterpret_cast(output.data.dptr), tile_rows, + tile_cols, rows_padded, block_len);) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -616,9 +625,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 = (scale_max * fp4_max) / global_amax * per_block_decode_scale = block_amax * (global_scale * (1 / fp4_max)) - * = block_amax * 448 / global_amax + * = block_amax * scale_max / global_amax * * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh * @@ -628,6 +637,7 @@ void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_ro * Output (global_scale_out): scalar float32 (the computed global encode scale) * --------------------------------------------------------------------------- */ +template __global__ void nvfp4_compute_per_block_scale_kernel( const float *__restrict__ block_amax, // [tile_rows, tile_cols] float *__restrict__ scale, // [tile_rows, tile_cols] @@ -636,18 +646,18 @@ __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 scale_max = dispatch::nvfp4::core::scale_max(); constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN // Read global_amax from device memory (avoids D2H transfer) float global_amax = *global_amax_ptr; - // Compute global encode scale: S_enc = (fp8_max * fp4_max) / global_amax + // Compute global encode scale: S_enc = (scale_max * fp4_max) / global_amax float safe_global_amax = fmaxf(global_amax, tiny); float global_scale = - (global_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; + (global_amax > 0.0f) ? fminf((scale_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; // Compute per-block decode scale: S_dec_b = block_amax * (S_enc * (1 / fp4_max)) float amax_val = block_amax[idx]; @@ -658,6 +668,7 @@ __global__ void nvfp4_compute_per_block_scale_kernel( } // Simple kernel to compute global encode scale from global amax +template __global__ void nvfp4_compute_global_scale_kernel( const float *__restrict__ global_amax, // [num_params] float *__restrict__ global_scale, // [num_params] @@ -665,19 +676,19 @@ __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 scale_max = dispatch::nvfp4::core::scale_max(); constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN float amax = global_amax[idx]; float safe_amax = fmaxf(amax, tiny); - float scale = (amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_amax, flt_max) : 1.0f; + float scale = (amax > 0.0f) ? fminf((scale_max * fp4_max) / safe_amax, flt_max) : 1.0f; global_scale[idx] = scale; } void nvfp4_compute_per_block_scale(const Tensor block_amax, Tensor scale, const Tensor global_amax, - cudaStream_t stream) { + DType scale_dtype, cudaStream_t stream) { NVTE_CHECK(block_amax.dtype() == DType::kFloat32, "Block amax must be float32."); NVTE_CHECK(scale.dtype() == DType::kFloat32, "Scale must be float32."); NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); @@ -689,14 +700,16 @@ void nvfp4_compute_per_block_scale(const Tensor block_amax, Tensor scale, const constexpr int kBlockSize = 256; int grid_size = (numel + kBlockSize - 1) / kBlockSize; - nvfp4_compute_per_block_scale_kernel<<>>( - reinterpret_cast(block_amax.data.dptr), - reinterpret_cast(scale.data.dptr), - reinterpret_cast(global_amax.data.dptr), numel); + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4_compute_per_block_scale_kernel<<>>( + reinterpret_cast(block_amax.data.dptr), + reinterpret_cast(scale.data.dptr), + reinterpret_cast(global_amax.data.dptr), numel);) NVTE_CHECK_CUDA(cudaGetLastError()); } -void nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale, +void nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale, DType scale_dtype, cudaStream_t stream) { NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); NVTE_CHECK(global_scale.dtype() == DType::kFloat32, "Global scale must be float32."); @@ -707,9 +720,11 @@ void nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale, constexpr int kBlockSize = 256; int grid_size = (num_params + kBlockSize - 1) / kBlockSize; - nvfp4_compute_global_scale_kernel<<>>( - reinterpret_cast(global_amax.data.dptr), - reinterpret_cast(global_scale.data.dptr), num_params); + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4_compute_global_scale_kernel<<>>( + reinterpret_cast(global_amax.data.dptr), + reinterpret_cast(global_scale.data.dptr), num_params);) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -720,23 +735,24 @@ void nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale, * Fuses three operations into one kernel: * 1. nvfp4_compute_per_block_scale: compute tile-level decode scales from block amax * 2. target_amax.copy_: copy global amax to target tensor - * 3. nvfp4_expand_scale_to_fp8: expand to row-level and convert to FP8 E4M3 + * 3. nvfp4_expand_scale_to_fp8: expand to row-level and convert to the selected scale format * * Input (block_amax): [tile_rows, tile_cols] float32 * Input (global_amax): [1] float32 * Output (per_block_scale): [tile_rows, tile_cols] float32 (intermediate, for partial_cast) - * Output (target_scale): [rows_padded, tile_cols] uint8 (E4M3) + * Output (target_scale): [rows_padded, tile_cols] uint8 (E4M3 or UE5M3) * Output (target_amax): [1] float32 (copy of global_amax) * * Saves 2 kernel launches per parameter (eliminates nvfp4_compute_per_block_scale and * nvfp4_expand_scale_to_fp8 as separate calls, plus the amax copy). * --------------------------------------------------------------------------- */ +template __global__ void nvfp4_fused_scale_kernel( const float *__restrict__ block_amax, // [tile_rows, tile_cols] const float *__restrict__ global_amax, // [1] float *__restrict__ per_block_scale, // [tile_rows, tile_cols] - for partial_cast - uint8_t *__restrict__ target_scale, // [rows_padded, tile_cols] + ScaleType *__restrict__ target_scale, // [rows_padded, tile_cols] float *__restrict__ target_amax, // [1] const size_t tile_rows, const size_t tile_cols, const size_t rows_padded, const size_t block_len) { @@ -757,8 +773,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 scale_max = dispatch::nvfp4::core::scale_max(); constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; @@ -766,7 +782,7 @@ __global__ void nvfp4_fused_scale_kernel( if (tile_row < tile_rows) { float safe_global_amax = fmaxf(g_amax, tiny); float global_scale = - (g_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; + (g_amax > 0.0f) ? fminf((scale_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; constexpr float fp4_max_inv = 1.0f / fp4_max; const float global_scale_multiplier = global_scale * fp4_max_inv; @@ -780,18 +796,18 @@ __global__ void nvfp4_fused_scale_kernel( } } - // Convert float32 to FP8 E4M3 and write expanded scale - fp8e4m3 fp8_val = static_cast(scale_val); - target_scale[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val); + target_scale[out_row * tile_cols + out_col] = static_cast(scale_val); } void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor per_block_scale, Tensor target_scale, Tensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, cudaStream_t stream) { + size_t rows_padded, size_t block_len, DType scale_dtype, + cudaStream_t stream) { NVTE_CHECK(block_amax.dtype() == DType::kFloat32, "Block amax must be float32."); NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); NVTE_CHECK(per_block_scale.dtype() == DType::kFloat32, "Per-block scale must be float32."); - NVTE_CHECK(target_scale.dtype() == DType::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_scale.dtype() == DType::kByte || target_scale.dtype() == scale_dtype, + "Target scale must be byte storage or have the selected NVFP4 scale dtype."); NVTE_CHECK(target_amax.dtype() == DType::kFloat32, "Target amax must be float32."); NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); @@ -802,13 +818,15 @@ void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor dim3 block(kBlockDim, kBlockDim); dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim); - nvfp4_fused_scale_kernel<<>>( - reinterpret_cast(block_amax.data.dptr), - reinterpret_cast(global_amax.data.dptr), - reinterpret_cast(per_block_scale.data.dptr), - reinterpret_cast(target_scale.data.dptr), - reinterpret_cast(target_amax.data.dptr), tile_rows, tile_cols, rows_padded, - block_len); + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + nvfp4_fused_scale_kernel + <<>>(reinterpret_cast(block_amax.data.dptr), + reinterpret_cast(global_amax.data.dptr), + reinterpret_cast(per_block_scale.data.dptr), + reinterpret_cast(target_scale.data.dptr), + reinterpret_cast(target_amax.data.dptr), tile_rows, + tile_cols, rows_padded, block_len);) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -818,38 +836,40 @@ void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream) { + cudaStream_t stream, const NVTEDType scale_dtype) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8); using namespace transformer_engine; - nvfp4_recipe::nvfp4_expand_scale_to_fp8(*convertNVTETensorCheck(input), - *convertNVTETensorCheck(output), tile_rows, tile_cols, - rows_padded, block_len, stream); + nvfp4_recipe::nvfp4_expand_scale_to_fp8( + *convertNVTETensorCheck(input), *convertNVTETensorCheck(output), tile_rows, tile_cols, + rows_padded, block_len, static_cast(scale_dtype), stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, cudaStream_t stream) { + const NVTETensor global_amax, cudaStream_t stream, + const NVTEDType scale_dtype) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale); using namespace transformer_engine; - nvfp4_recipe::nvfp4_compute_per_block_scale(*convertNVTETensorCheck(block_amax), - *convertNVTETensorCheck(scale), - *convertNVTETensorCheck(global_amax), stream); + nvfp4_recipe::nvfp4_compute_per_block_scale( + *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(scale), + *convertNVTETensorCheck(global_amax), static_cast(scale_dtype), stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - cudaStream_t stream) { + cudaStream_t stream, const NVTEDType scale_dtype) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_global_scale); using namespace transformer_engine; nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax), - *convertNVTETensorCheck(global_scale), stream); + *convertNVTETensorCheck(global_scale), + static_cast(scale_dtype), stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED @@ -896,14 +916,15 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, cudaStream_t stream) { + size_t block_len, cudaStream_t stream, + const NVTEDType scale_dtype) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); using namespace transformer_engine; - nvfp4_recipe::nvfp4_2d_partial_cast(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), - *convertNVTETensorCheck(scale), - *convertNVTETensorCheck(global_scale), h, w, scale_stride_h, - scale_stride_w, start_offset, block_len, stream); + nvfp4_recipe::nvfp4_2d_partial_cast( + *convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), *convertNVTETensorCheck(scale), + *convertNVTETensorCheck(global_scale), h, w, scale_stride_h, scale_stride_w, start_offset, + block_len, static_cast(scale_dtype), stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED @@ -924,17 +945,20 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void *amax_A_ptr = use_rowwise_amax_A ? tA->amax.dptr : tA->columnwise_amax.dptr; void *amax_B_ptr = use_rowwise_amax_B ? tB->amax.dptr : tB->columnwise_amax.dptr; void *alpha_ptr = tOut->data.dptr; - const float fp8_max_A = static_cast(tA->nvfp4_e4m3_max); - const float fp8_max_B = static_cast(tB->nvfp4_e4m3_max); + const DType scale_dtype_A = + use_rowwise_amax_A ? tA->scale_inv.dtype : tA->columnwise_scale_inv.dtype; + const DType scale_dtype_B = + use_rowwise_amax_B ? tB->scale_inv.dtype : tB->columnwise_scale_inv.dtype; + const float scale_max_A = + dispatch::nvfp4::core::scale_max(scale_dtype_A, tA->get_nvfp4_scale_max()); + const float scale_max_B = + dispatch::nvfp4::core::scale_max(scale_dtype_B, tB->get_nvfp4_scale_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>>>( alpha_in, reinterpret_cast(amax_A_ptr), - reinterpret_cast(amax_B_ptr), fp8_max_A, fp8_max_B, + reinterpret_cast(amax_B_ptr), scale_max_A, scale_max_B, reinterpret_cast(alpha_ptr)); NVTE_CHECK_CUDA(cudaGetLastError()); #else @@ -945,14 +969,16 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, cudaStream_t stream) { + size_t rows_padded, size_t block_len, cudaStream_t stream, + const NVTEDType scale_dtype) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_fused_scale); using namespace transformer_engine; nvfp4_recipe::nvfp4_fused_scale( *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax), *convertNVTETensorCheck(per_block_scale), *convertNVTETensorCheck(target_scale), - *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, stream); + *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, + static_cast(scale_dtype), stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 988c32d2b4..172959fa4d 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -173,18 +173,21 @@ void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_i if (t.has_data()) { NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor input ", name, "_scale_inverse must be allocated"); - NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor input ", name, + NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 + || t.scale_inv.dtype == DType::kFloat8UE5M3, + "FP4 scaling factor input ", name, "_scale_inverse has invalid dtype " - "(expected DType::kFloat8E4M3, got ", + "(expected Float8E4M3 or Float8UE5M3, got ", to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor input ", name, "_columnwise_scale_inverse must be allocated"); - NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP8 scaling factor input ", - name, + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 + || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + "FP8 scaling factor input ", name, "_columnwise_scale_inverse has invalid dtype " - "(expected DType::kFloat8E4M3, got ", + "(expected Float8E4M3 or Float8UE5M3, got ", to_string(t.columnwise_scale_inv.dtype), ")"); } } else { @@ -234,18 +237,21 @@ void CheckOutputTensor(const Tensor &t, std::string_view name, bool allow_empty) if (t.has_data()) { NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor output ", name, "_scale_inverse must be allocated"); - NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ", name, + NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 + || t.scale_inv.dtype == DType::kFloat8UE5M3, + "FP4 scaling factor output ", name, "_scale_inverse has invalid dtype " - "(expected Float8E4M3, got ", + "(expected Float8E4M3 or Float8UE5M3, got ", to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor output ", name, "_columnwise_scale_inverse must be allocated"); - NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ", - name, + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 + || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + "FP4 scaling factor output ", name, "_columnwise_scale_inverse has invalid dtype " - "(expected Float8E4M3, got ", + "(expected Float8E4M3 or Float8UE5M3, got ", to_string(t.columnwise_scale_inv.dtype), ")"); } } else { @@ -361,7 +367,26 @@ static void CheckGroupedScaleInv(const GroupedTensor &t, std::string_view name, } else if (is_mxfp8_scaling(t.scaling_mode)) { check_scales(DType::kFloat8E8M0); } else if (is_nvfp4_scaling(t.scaling_mode)) { - check_scales(DType::kFloat8E4M3); + if (t.has_data()) { + NVTE_CHECK(t.scale_inv.has_data(), tensor_type, " ", name, + " rowwise scale_inv must be allocated"); + NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 + || t.scale_inv.dtype == DType::kFloat8UE5M3, + tensor_type, " ", name, + " rowwise scale_inv has invalid dtype " + "(expected Float8E4M3 or Float8UE5M3, got ", + to_string(t.scale_inv.dtype), ")"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_scale_inv.has_data(), tensor_type, " ", name, + " columnwise scale_inv must be allocated"); + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 + || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + tensor_type, " ", name, + " columnwise scale_inv has invalid dtype " + "(expected Float8E4M3 or Float8UE5M3, got ", + to_string(t.columnwise_scale_inv.dtype), ")"); + } } else { // Non-quantized types should not have scale/scale_inv NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv not supported for non-quantized ", tensor_type, @@ -898,8 +923,10 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo break; case kNVTENVFP4E4M3Max: std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); - NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256, - "Unsupported NVFP4 E4M3 max (got ", t.nvfp4_e4m3_max, ")"); + // Need to rename this to nvfp4_scale_type_max + NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || + t.nvfp4_e4m3_max == 114688 || t.nvfp4_e4m3_max == 65536, + "Unsupported NVFP4 scale type max (got ", t.nvfp4_e4m3_max, ")"); break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); @@ -985,7 +1012,10 @@ void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, vo *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); break; case kNVTENVFP4E4M3Max: - std::memcpy(buf, &t->nvfp4_e4m3_max, attr_size); + { + int val = t->get_nvfp4_scale_max(); + std::memcpy(buf, &val, attr_size); + } break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); 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..616748a0eb 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -14,6 +14,7 @@ #include #include +#include "common/cast/nvfp4/core_nvfp4.cuh" #include "common/common.h" #include "common/recipe/recipe_common.cuh" #include "common/transpose/cast_transpose.h" @@ -167,14 +168,6 @@ __device__ __forceinline__ float groupMax(float val, unsigned int groupMask) { return val; } -template -__device__ __forceinline__ ScaleType -ComputeDecodeScaleFP4(const float amax, const float global_encode_scale_multiplier) { - float decode_scale = amax * global_encode_scale_multiplier; - decode_scale = fminf(decode_scale, TypeExtrema::max); - return static_cast(decode_scale); -} - template __device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale, const float global_decode_scale) { @@ -187,19 +180,6 @@ __device__ __forceinline__ float ComputeOutputFP4(IType input, float encode_scal return static_cast(input) * encode_scale; } -__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { - constexpr float fp8_max = TypeExtrema::max; - constexpr float fp4_max = TypeExtrema::max; - float global_encode_scale = fp8_max * fp4_max / global_amax; - // If scale is infinity, return max value of float32 - global_encode_scale = fminf(global_encode_scale, TypeExtrema::max); - // If global amax is 0 or infinity, return 1 - if (global_amax == 0.f || global_encode_scale == 0.f) { - return 1.f; - } - return global_encode_scale; -} - __device__ __forceinline__ uint32_t get_rbits( transformer_engine::curanddx::detail::philox4x32_native_state& rng, // NVTE_BUILD_NUM_PHILOX_ROUNDS rounds of philox4x32 @@ -415,9 +395,10 @@ __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]); - constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; - const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + (kIsE8Scaling || global_amax == nullptr) + ? 1.0f + : dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + global_amax[0]); const float global_decode_scale = 1.0 / global_encode_scale; // Step 2: Cast and store to output_c @@ -510,14 +491,15 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo float row_global_encode_scale = global_encode_scale; if constexpr (kRowScaledNVFP4) { row_global_encode_scale = - row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx]) : 1.0f; + row_idx < num_rows + ? dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4( + global_amax[row_idx]) + : 1.0f; } - const float row_global_encode_scale_multiplier = - kRowScaledNVFP4 ? row_global_encode_scale * fp4_max_inv : global_encode_scale_multiplier; const float row_global_decode_scale = kRowScaledNVFP4 ? 1.0f / row_global_encode_scale : global_decode_scale; - ScaleType scale_inv = - ComputeDecodeScaleFP4(amax, row_global_encode_scale_multiplier); + ScaleType scale_inv = dispatch::nvfp4::core::compute_decoding_scaling_factor( + amax, row_global_encode_scale); float encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale); // Step 2.5: Write scale_inv bool write_scale_inv = is_src_lane; @@ -701,8 +683,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo amax = __shfl_sync(mask, amax, src_lane); } // Step 3.4: Compute scale - ScaleType scale_inv = - ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier); + ScaleType scale_inv = dispatch::nvfp4::core::compute_decoding_scaling_factor( + amax, global_encode_scale); float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); // Step 3.5: Write scale_inv_t bool write_scale_inv = is_src_lane; @@ -772,14 +754,14 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo namespace detail { -void quantize_transpose_vector_blockwise_fp4( +template +void quantize_transpose_vector_blockwise_fp4_impl( const SimpleTensor& input, const SimpleTensor& global_amax, SimpleTensor& scale_inv, SimpleTensor& scale_inv_t, SimpleTensor& output, SimpleTensor& output_t, const float epsilon, const bool return_identity, const bool return_transpose, const bool pow2_scale, const bool swizzled_scale, const bool use_stochastic_rounding, const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, const SimpleTensor& noop_tensor, cudaStream_t stream) { - NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4); #if CUDA_VERSION >= 12080 // pow 2 scale is for MXFP4 since it's using E8M0 scaling @@ -849,8 +831,7 @@ void quantize_transpose_vector_blockwise_fp4( dim3 grid(num_blocks_x, num_blocks_y, 1); - using ScaleType = fp8e4m3; constexpr int kScaleBlockDim = 16; - constexpr bool kPow2Scale = false; + constexpr int kScaleBlockDim = 16; constexpr bool kPow2Scale = false; const bool full_tile = row_length % kTileDim == 0 && num_rows % kTileDim == 0; @@ -914,5 +895,31 @@ void quantize_transpose_vector_blockwise_fp4( #endif // CUDA_VERSION >= 12080 } +void quantize_transpose_vector_blockwise_fp4( + const SimpleTensor& input, const SimpleTensor& global_amax, SimpleTensor& scale_inv, + SimpleTensor& scale_inv_t, SimpleTensor& output, SimpleTensor& output_t, const float epsilon, + const bool return_identity, const bool return_transpose, const bool pow2_scale, + const bool swizzled_scale, const bool use_stochastic_rounding, + const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, + const SimpleTensor& noop_tensor, cudaStream_t stream) { + NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4); + + NVTE_CHECK(return_identity || return_transpose, + "At least one of return_identity or return_transpose must be true."); + const DType scale_dtype = return_identity ? scale_inv.dtype : scale_inv_t.dtype; + if (return_identity && return_transpose) { + NVTE_CHECK(scale_inv.dtype == scale_inv_t.dtype, + "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", + to_string(scale_inv.dtype), " and ", to_string(scale_inv_t.dtype), ")."); + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + quantize_transpose_vector_blockwise_fp4_impl( + input, global_amax, scale_inv, scale_inv_t, output, output_t, epsilon, return_identity, + return_transpose, pow2_scale, swizzled_scale, use_stochastic_rounding, rng_state_tensor, + use_2d_quantization, row_scaled_nvfp4, noop_tensor, stream);) +} + } // namespace detail } // namespace transformer_engine diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index f7ffb5ad8d..4c420cf4cb 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -23,7 +23,9 @@ .value("kBFloat16", transformer_engine::DType::kBFloat16) \ .value("kFloat8E4M3", transformer_engine::DType::kFloat8E4M3) \ .value("kFloat8E5M2", transformer_engine::DType::kFloat8E5M2) \ + .value("kFloat8E8M0", transformer_engine::DType::kFloat8E8M0) \ .value("kFloat4E2M1", transformer_engine::DType::kFloat4E2M1) \ + .value("kFloat8UE5M3", transformer_engine::DType::kFloat8UE5M3) \ .def("__reduce_ex__", \ [](transformer_engine::DType self, pybind11::object /*protocol*/) { \ return pybind11::make_tuple(pybind11::type::of(pybind11::cast(self)), \ diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 2b1803bfb2..685cf69c6a 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -50,6 +50,7 @@ from transformer_engine.pytorch.quantization import is_mxfp8_available from transformer_engine.pytorch.quantization import is_fp8_block_scaling_available from transformer_engine.pytorch.quantization import is_nvfp4_available +from transformer_engine.pytorch.quantization import is_fp8_ue5m3_available from transformer_engine.pytorch.quantization import get_default_recipe from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.quantization import QuantizerRequest diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index 3a145bbb5b..ec54189613 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -28,8 +28,12 @@ class DType(enum.IntEnum): bits (``torch.float8_e4m3fn``). * ``kFloat8E5M2`` -- 8-bit floating point with 5 exponent and 2 mantissa bits (``torch.float8_e5m2``). + * ``kFloat8E8M0`` -- 8-bit unsigned floating point with 8 exponent and 0 + mantissa bits. * ``kFloat4E2M1`` -- 4-bit floating point with 2 exponent and 1 mantissa bits. + * ``kFloat8UE4M3`` -- 8-bit unsigned floating point with 5 exponent and 3 + mantissa bits. The enum mirrors the backend ``transformer_engine_torch.DType`` (pybind11) enum value-for-value, and instances of the two enums compare equal when @@ -43,7 +47,9 @@ class DType(enum.IntEnum): kBFloat16 = int(tex.DType.kBFloat16) kFloat8E4M3 = int(tex.DType.kFloat8E4M3) kFloat8E5M2 = int(tex.DType.kFloat8E5M2) + kFloat8E8M0 = int(tex.DType.kFloat8E8M0) kFloat4E2M1 = int(tex.DType.kFloat4E2M1) + kFloat8UE5M3 = int(tex.DType.kFloat8UE5M3) @classmethod def cast(cls, dtype: "Union[DType, tex.DType]") -> "DType": diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..4aa6d58114 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -50,6 +50,7 @@ #include #include #include +#include #include #include @@ -351,9 +352,13 @@ class NVFP4Quantizer : public Quantizer { // 4over6 candidate-selection mode used when quantizing emitted NVFP4 tensors. NVTENVFP44Over6Mode nvfp4_4over6_mode; // Global E4M3 scale bound used by emitted NVFP4 tensors. - int nvfp4_e4m3_max; + std::optional nvfp4_e4m3_max; + // Dtype of scale_inv tensors (kFloat8E4M3 or kFloat8UE5M3). + DType scale_dtype; // 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_second_level_scale; int rht_matrix_random_sign_mask_t; at::Tensor rht_matrix; @@ -455,6 +460,7 @@ inline size_t typeToNumBits(transformer_engine::DType t) { case transformer_engine::DType::kFloat8E4M3: case transformer_engine::DType::kFloat8E5M2: case transformer_engine::DType::kFloat8E8M0: + case transformer_engine::DType::kFloat8UE5M3: return 8; case transformer_engine::DType::kFloat4E2M1: return 4; @@ -485,6 +491,8 @@ inline at::ScalarType GetATenDType(transformer_engine::DType t) { return at::kFloat8_e5m2; case transformer_engine::DType::kFloat8E8M0: return at::kByte; // e8m0 dtype requires PyTorch 2.7.0+ + case transformer_engine::DType::kFloat8UE5M3: + return at::kByte; default: NVTE_ERROR("Invalid type (", static_cast(t), ")."); } diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 544ff92c1b..643d08875e 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -46,6 +46,9 @@ py::object activation_helper(const at::Tensor& input, py::handle quantizer, int (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; + } else if (nvfp4_quantizer_cpp->disable_second_level_scale) { + // No need for amax + impl = Impl::UNFUSED; } else { impl = Impl::FUSED_ACTIVATION_AMAX_NVFP4; } @@ -159,6 +162,9 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; + } else if (nvfp4_quantizer_cpp->disable_second_level_scale) { + // No need for amax + impl = Impl::UNFUSED; } else { impl = Impl::FUSED_ACTIVATION_AMAX_NVFP4; } diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index 4a78dde388..892dfc8182 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -156,6 +156,9 @@ std::vector dact_dbias( (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; + } else if (nvfp4_quantizer_cpp->disable_second_level_scale) { + // No need for amax + impl = Impl::UNFUSED; } else { impl = Impl::FUSED_DACT_AMAX_NVFP4; } diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 5ce0261c82..ffd9545516 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -346,7 +346,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_second_level_scale); break; } case GroupedQuantizationMode::FP8_CURRENT_SCALING_GROUPED_QUANTIZE: { @@ -616,18 +617,15 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o // Data tensors are stored as flat 1D buffers; use the quantizer's dtype // (e.g. kFloat8E4M3) rather than the raw tensor scalar_type (uint8). const NVTEScalingMode scaling_mode = quantizer->get_scaling_mode(); - const bool is_block_scaling = - (scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D); - const bool is_nvfp4 = (scaling_mode == NVTE_NVFP4_1D_SCALING); - const DType scale_dtype = is_block_scaling ? DType::kFloat32 - : is_nvfp4 ? DType::kFloat8E4M3 - : DType::kFloat8E8M0; + py::object py_scale_dtype = input.attr("scale_inv_dtype"); + const std::optional scale_dtype = py_scale_dtype.cast>(); auto input_cpp = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); if (rowwise_data.has_value()) { input_cpp.set_rowwise_data(rowwise_data->data_ptr(), quantizer->dtype, std::vector{static_cast(rowwise_data->numel())}); if (rowwise_scale_inv.has_value()) { - input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), scale_dtype, + NVTE_CHECK(scale_dtype, "Could not deduce scale dtype"); + input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), *scale_dtype, getTensorShape(*rowwise_scale_inv)); } } @@ -636,7 +634,8 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o columnwise_data->data_ptr(), quantizer->dtype, std::vector{static_cast(columnwise_data->numel())}); if (columnwise_scale_inv.has_value()) { - input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), scale_dtype, + NVTE_CHECK(scale_dtype, "Could not deduce scale dtype"); + input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), *scale_dtype, getTensorShape(*columnwise_scale_inv)); } } @@ -1094,7 +1093,8 @@ std::tuple, std::vector, bool> bulk_alloc const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4; 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 auto 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 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."); @@ -1103,6 +1103,7 @@ std::tuple, std::vector, bool> bulk_alloc } const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; + const auto scale_dtype = quantizer_cpp_list[0]->scale_dtype; // with_gemm_swizzled_scales is a single group-wide boolean baked // into every output tensor. We can safely request it only when @@ -1125,6 +1126,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_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."); } bool all_tensors_rht_cast_fusion_eligible = true; for (size_t i = 0; i < num_tensors; ++i) { @@ -1188,18 +1192,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_second_level_scale) { + 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_second_level_scale) { + amax_rowwise_list.push_back(tensors[2 * num_tensors + i]); + } } } @@ -1242,18 +1250,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_second_level_scale) { + 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_second_level_scale) { + amax_columnwise_list.push_back(tensors[2 * num_tensors + i]); + } } } @@ -1267,14 +1279,19 @@ 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_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( rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, amax_rowwise, - amax_columnwise, MakePythonDType(fp4_dtype), quantizer_py_list[i], - with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, + amax_columnwise, MakePythonDType(fp4_dtype), MakePythonDType(scale_dtype), + quantizer_py_list[i], with_gemm_swizzled_scales, + py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); @@ -1282,28 +1299,31 @@ std::tuple, std::vector, bool> bulk_alloc // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, // then set the amax and amax_columnwise values. { - auto tensor_wrapper = makeTransformerEngineTensor( - rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr, - columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_data_shapes[i] : std::vector{0}, - columnwise_usage ? columnwise_data_shapes[i] : std::vector{0}, fp4_dtype, - /*amax_ptr=*/nullptr, - /*scale_ptr=*/nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr, - columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, - columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode); - tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); - tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4); - tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max); - - // Set the amax rowwise and amax columnwise if available + TensorWrapper tensor_wrapper(NVTE_NVFP4_1D_SCALING); if (rowwise_usage) { - tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, - getTensorShape(amax_rowwise_list[i])); + tensor_wrapper.set_rowwise_data(rowwise_data_list[i].data_ptr(), + fp4_dtype, rowwise_data_shapes[i]); + tensor_wrapper.set_rowwise_scale_inv(rowwise_scale_list[i].data_ptr(), + scale_dtype, rowwise_scale_shapes[i]); + if (!disable_second_level_scale) { + tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, + getTensorShape(amax_rowwise_list[i])); + } } if (columnwise_usage) { - tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, - std::vector{1}); + tensor_wrapper.set_columnwise_data(columnwise_data_list[i].data_ptr(), + fp4_dtype, columnwise_data_shapes[i]); + tensor_wrapper.set_columnwise_scale_inv(columnwise_scale_list[i].data_ptr(), + scale_dtype, columnwise_scale_shapes[i]); + if (!disable_second_level_scale) { + tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, + std::vector{1}); + } + } + tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4); + if (nvfp4_e4m3_max) { + tensor_wrapper.set_nvfp4_e4m3_max(*nvfp4_e4m3_max); } tensor_cpp_list.emplace_back(std::move(tensor_wrapper)); @@ -1484,7 +1504,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_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: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) @@ -1650,19 +1672,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_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; + 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/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index c3dec944e4..43f1d32b8a 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -123,6 +123,9 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; + } else if (nvfp4_quantizer_cpp->disable_second_level_scale) { + // No need for amax + impl = Impl::UNFUSED; } else if (!transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { // TE kernel supports amax output impl = Impl::FUSED_NORM_AMAX_NVFP4; @@ -360,6 +363,9 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; + } else if (nvfp4_quantizer_cpp->disable_second_level_scale) { + // No need for amax + impl = Impl::UNFUSED; } else if (!transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { // TE kernel supports amax output impl = Impl::FUSED_NORM_AMAX_NVFP4; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 3c2d2d9e14..1ee0513f6c 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1883,9 +1883,10 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); this->stochastic_rounding = quantizer.attr("stochastic_rounding").cast(); const bool nvfp4_use_4over6 = quantizer.attr("nvfp4_use_4over6").cast(); - this->nvfp4_e4m3_max = quantizer.attr("nvfp4_e4m3_max").cast(); - NVTE_CHECK(this->nvfp4_e4m3_max == 448 || this->nvfp4_e4m3_max == 256, - "Unsupported NVFP4 E4M3 max: ", this->nvfp4_e4m3_max); + const int e4m3_max = quantizer.attr("nvfp4_e4m3_max").cast(); + if (e4m3_max >= 0) { + this->nvfp4_e4m3_max = e4m3_max; + } const auto nvfp4_4over6_err_mode = quantizer.attr("nvfp4_4over6_err_mode").cast(); if (!nvfp4_use_4over6) { this->nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; @@ -1897,6 +1898,10 @@ 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->scale_dtype = quantizer.attr("scale_dtype").cast(); + NVTE_CHECK(this->scale_dtype == DType::kFloat8E4M3 || this->scale_dtype == DType::kFloat8UE5M3, + "Unsupported NVFP4 scale dtype: ", static_cast(this->scale_dtype)); + 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(); @@ -1981,8 +1986,8 @@ 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_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) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); } @@ -2004,7 +2009,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_second_level_scale) { + amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); + } } if (columnwise_usage) { const std::vector scale_inv_shape_int64(columnwise_scale_inv_shape.begin(), @@ -2020,7 +2027,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_second_level_scale) { + amax_columnwise = at::empty({amax_cols}, bit32_tensor_opts); + } } // Convert tensors to Python @@ -2031,8 +2040,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_second_level_scale); + auto amax_columnwise_py = + py_cast(amax_columnwise, columnwise_usage && !disable_second_level_scale); // Construct Python NVFP4 tensor py::object out_py; @@ -2046,11 +2056,12 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["amax_rowwise"] = amax_rowwise_py; kwargs["amax_columnwise"] = amax_columnwise_py; kwargs["fp4_dtype"] = MakePythonDType(this->dtype); + kwargs["scale_dtype"] = MakePythonDType(this->scale_dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); - kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); + kwargs["nvfp4_e4m3_max"] = py::cast(this->nvfp4_e4m3_max); kwargs["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); @@ -2077,12 +2088,13 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["amax_rowwise"] = amax_rowwise_py; kwargs["amax_columnwise"] = amax_columnwise_py; kwargs["fp4_dtype"] = MakePythonDType(this->dtype); + kwargs["scale_dtype"] = MakePythonDType(this->scale_dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["device"] = py::cast(device); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); - kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); + kwargs["nvfp4_e4m3_max"] = py::cast(this->nvfp4_e4m3_max); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), args.ptr(), kwargs.ptr()); @@ -2098,9 +2110,11 @@ std::pair NVFP4Quantizer::create_tensor( TensorWrapper out_cpp(NVTE_NVFP4_1D_SCALING); if (rowwise_usage) { 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, + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv_tensor.data_ptr(), this->scale_dtype, rowwise_scale_inv_shape); - out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); + if (!disable_second_level_scale) { + 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 @@ -2109,14 +2123,18 @@ std::pair NVFP4Quantizer::create_tensor( auto col_data_shape_fp4 = make_transpose_shape(shape_2d); out_cpp.set_columnwise_data(columnwise_data_tensor.data_ptr(), DType::kFloat4E2M1, col_data_shape_fp4); - out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), this->scale_dtype, columnwise_scale_inv_shape); - out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - getTensorShape(amax_columnwise)); + if (!disable_second_level_scale) { + 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); - out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); + if (this->nvfp4_e4m3_max) { + out_cpp.set_nvfp4_e4m3_max(*this->nvfp4_e4m3_max); + } this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(out_py)}; @@ -2148,8 +2166,8 @@ 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_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) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 grouped quantization requires rowwise usage."); NVTE_CHECK(!columnwise_usage, @@ -2165,7 +2183,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_second_level_scale) { + rowwise_amax = at::empty({amax_elements}, float_opts); + } } if (columnwise_usage) { @@ -2173,23 +2193,29 @@ 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_second_level_scale) { + columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + } } GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); if (rowwise_usage) { 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, + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), this->scale_dtype, 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, + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), this->scale_dtype, 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)); @@ -2228,7 +2254,8 @@ std::pair NVFP4Quantizer::create_grouped_tenso kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); - kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); + kwargs["nvfp4_e4m3_max"] = py::cast(this->nvfp4_e4m3_max); + kwargs["scale_inv_dtype"] = MakePythonDType(this->scale_dtype); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -2307,15 +2334,16 @@ 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_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) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); } tensor.attr("_row_scaled_nvfp4") = row_scaled_nvfp4; tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; tensor.attr("_nvfp4_use_4over6") = py::cast(nvfp4_use_4over6); - tensor.attr("_nvfp4_e4m3_max") = py::cast(nvfp4_e4m3_max); + tensor.attr("_nvfp4_e4m3_max") = py::cast(this->nvfp4_e4m3_max); + tensor.attr("_scale_dtype") = MakePythonDType(this->scale_dtype); // Coerce row-wise data if (rowwise_usage) { @@ -2334,7 +2362,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_second_level_scale) { + 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 @@ -2377,7 +2408,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_second_level_scale) { + 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 @@ -2403,9 +2437,11 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( TensorWrapper out_cpp(NVTE_NVFP4_1D_SCALING); if (rowwise_usage) { 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, + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), this->scale_dtype, 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 @@ -2414,14 +2450,18 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( auto col_data_shape_fp4 = make_transpose_shape(shape_2d); out_cpp.set_columnwise_data(columnwise_data->data_ptr(), DType::kFloat4E2M1, col_data_shape_fp4); - out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), this->scale_dtype, 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); - out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); + if (this->nvfp4_e4m3_max) { + out_cpp.set_nvfp4_e4m3_max(*this->nvfp4_e4m3_max); + } this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(tensor)}; @@ -2505,7 +2545,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_second_level_scale) { return; } @@ -2632,7 +2672,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_second_level_scale) { NVTE_SCOPED_GIL_RELEASE({ nvte_hadamard_transform_amax(input.data(), out.data(), 0, this->rht_matrix_random_sign_mask_t, stream); @@ -2645,7 +2685,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_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/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index ddb85808a5..a53cc3ffe7 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -4,6 +4,9 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include + #include #include #include @@ -135,7 +138,8 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); const bool with_gemm_swizzled_scales = tensor.attr("_with_gemm_swizzled_scales").cast(); const bool row_scaled_nvfp4 = tensor.attr("_row_scaled_nvfp4").cast(); - const int nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast(); + const auto nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast>(); + const DType scale_inv_dtype = tensor.attr("_scale_dtype").cast(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -143,30 +147,37 @@ 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)); + ret.set_rowwise_scale_inv(scale_inv.data_ptr(), scale_inv_dtype, getTensorShape(scale_inv)); + 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, + ret.set_columnwise_scale_inv(scale_inv.data_ptr(), scale_inv_dtype, 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 ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); ret.set_row_scaled_nvfp4(row_scaled_nvfp4); - ret.set_nvfp4_e4m3_max(nvfp4_e4m3_max); + if (nvfp4_e4m3_max) { + ret.set_nvfp4_e4m3_max(*nvfp4_e4m3_max); + } // Quantizer state quantizer->set_quantization_params(&ret); @@ -198,7 +209,7 @@ DType GetTransformerEngineDTypeForScaleInv(py::handle quantizer, at::Tensor scal return DType::kFloat32; } if (IsNVFP4Quantizers(quantizer_ptr)) { - return DType::kFloat8E4M3; + return quantizer.attr("scale_dtype").cast(); } return GetTransformerEngineDType(scale_inv.scalar_type()); } @@ -258,17 +269,20 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { getTensorShape(amax)); } + // Scale inverse dtype + py::object py_scale_inv_dtype = tensor.attr("scale_inv_dtype"); + const std::optional scale_inv_dtype = py_scale_inv_dtype.cast>(); + // Scale inverse if (!tensor.attr("scale_inv").is_none()) { const auto &scale_inv = tensor.attr("scale_inv").cast(); - ret.set_rowwise_scale_inv(scale_inv.data_ptr(), - GetTransformerEngineDTypeForScaleInv(quantizer, scale_inv), - getTensorShape(scale_inv)); + NVTE_CHECK(scale_inv_dtype, "Could not determine dtype of scale_inv buffer."); + ret.set_rowwise_scale_inv(scale_inv.data_ptr(), *scale_inv_dtype, getTensorShape(scale_inv)); } if (!tensor.attr("columnwise_scale_inv").is_none()) { const auto &scale_inv = tensor.attr("columnwise_scale_inv").cast(); - ret.set_columnwise_scale_inv(scale_inv.data_ptr(), - GetTransformerEngineDTypeForScaleInv(quantizer, scale_inv), + NVTE_CHECK(scale_inv_dtype, "Could not determine dtype of scale_inv buffer."); + ret.set_columnwise_scale_inv(scale_inv.data_ptr(), *scale_inv_dtype, getTensorShape(scale_inv)); } diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e9a65c3648..6dc0035204 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -2061,6 +2061,10 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] + if recipe.custom(): + # Custom quantization recipes are compatible with all quantizers + return + weight_tensors = [getattr(self, name) for name in self.weight_names] for i, tensor in enumerate(weight_tensors): if isinstance(tensor, QuantizedTensorStorage): diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 76d51673f0..a44bef0b2d 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -16,6 +16,7 @@ from packaging.version import Version as PkgVersion import transformer_engine_torch as tex +from ....common.recipe import Format as RecipeFormat from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor @@ -808,14 +809,33 @@ def fuse_grouped_mlp_ops( """ if not fused_op_cls.is_supported(): return ops - if recipe is None or not (recipe.mxfp8() or recipe.nvfp4()): + + # Fused kernels are only supported for MXFP8 and NVFP4 + if recipe is None: return ops - # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. - if recipe.nvfp4() and recipe.disable_rht: + elif recipe.custom(): + # Check if custom recipe explicitly enables fusion + if not getattr(recipe, "enable_cutedsl_fused_grouped_mlp", False): + return ops + elif not (recipe.mxfp8() or recipe.nvfp4()): return ops + + # Check for unsupported NVFP4 recipe configs + if recipe.nvfp4(): + if recipe.disable_rht: + # Graph-safe grouped quantize is only supported with RHT + return ops + if ( + recipe.row_scaled_activation + or recipe.nvfp4_4over6 + or recipe.fp8_format == RecipeFormat.UE5M3 + ): + return ops + if activation_op_types is None: activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) + # Scan ops through with sliding window out = [] window, ops = ops[:3], ops[3:] while len(window) == 3: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 07a3b80483..06175faebd 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -39,6 +39,7 @@ "is_mxfp8_available", "is_fp8_block_scaling_available", "is_nvfp4_available", + "is_fp8_ue5m3_available", "get_default_recipe", "get_align_size_for_quantization", "QuantizerRole", @@ -51,6 +52,7 @@ _MXFP8_SUPPORT: Optional[Tuple[bool, str]] = None _NVFP4_SUPPORT: Optional[Tuple[bool, str]] = None _FP8_BLOCK_SCALING_SUPPORT: Optional[Tuple[bool, str]] = None +_FP8_UE5M3_SUPPORT: Optional[Tuple[bool, str]] = None @dataclasses.dataclass(frozen=True) @@ -221,6 +223,21 @@ def check_fp8_block_scaling_support() -> Tuple[bool, str]: return _FP8_BLOCK_SCALING_SUPPORT +@torch.compiler.assume_constant_result +def check_fp8_ue5m3_support() -> Tuple[bool, str]: + """Return if the FP8 UE5M3 format is available.""" + global _FP8_UE5M3_SUPPORT + if _FP8_UE5M3_SUPPORT is None: + def _check_support() -> Tuple[bool, str]: + if get_device_compute_capability() != (10, 7): # Rubin + return False, "Device compute capability 10.7 is required for FP8 UE5M3 support." + if float(torch.version.cuda) < 13.4: + return False, "CUDA 13.4 is required for FP8 UE5M3 support." + return True, "" + _FP8_UE5M3_SUPPORT = _check_support() + return _FP8_UE5M3_SUPPORT + + def check_recipe_support(recipe: Recipe) -> None: """Check if the given recipe is supported.""" if torch.compiler.is_compiling() and isinstance(recipe, DelayedScaling): @@ -385,6 +402,26 @@ def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_nvfp4_support()[0] +def is_fp8_ue5m3_available(return_reason: bool = False) -> Union[bool, Tuple[bool, str]]: + """ + Determine if support is available for the FP8 UE5M3 data type. + + This may be used for NVFP4 scaling factors. + + Parameters + ---------- + return_reason : bool, optional + If ``False`` (default), return only a boolean indicating availability. + If ``True``, return a tuple ``(is_available, reason)`` where ``reason`` provides + a human-readable explanation when required support is not available. The reason + will be an empty string if support is available. + + """ + if return_reason: + return check_fp8_ue5m3_support() + return check_fp8_ue5m3_support()[0] + + @dataclass(slots=True) class FP8GlobalState: """Mutable process-global FP8 state stored on an instance. @@ -1681,8 +1718,16 @@ def _qparams(tensor_type: str): return self.recipe.fp4_quant_fwd_weight return self.recipe.fp4_quant_fwd_inp + scale_dtype = ( + DType.kFloat8UE5M3 + if self.recipe.fp8_format == Format.UE5M3 + else DType.kFloat8E4M3 + ) + def _make(tensor_type: str) -> NVFP4Quantizer: qparams = _qparams(tensor_type) + + # Whether to enable 4over6 nvfp4_use_4over6 = False if tensor_type not in ("grad_output", "grad_input"): if self.recipe.nvfp4_4over6 == "all": @@ -1691,16 +1736,23 @@ def _make(tensor_type: str) -> NVFP4Quantizer: nvfp4_use_4over6 = tensor_type == "weight" elif self.recipe.nvfp4_4over6 == "activations": nvfp4_use_4over6 = tensor_type != "weight" - nvfp4_e4m3_max = 448 + + # Unsupported configs if nvfp4_use_4over6: - # Current 4over6 kernels target RL and post-training quantization paths. - # Pre-training usage still needs a fused RHT + 4over6 quantization kernel. if qparams.random_hadamard_transform: raise ValueError("NVFP4 4over6 quantization does not support RHT.") if qparams.stochastic_rounding: raise ValueError( "NVFP4 4over6 quantization does not support stochastic rounding." ) + if scale_dtype == DType.kFloat8UE5M3: + raise ValueError( + "NVFP4 4over6 quantization is incompatible with UE5M3 scales." + ) + + # Scale max for 4over6 + nvfp4_e4m3_max = None + if nvfp4_use_4over6: if self.recipe.nvfp4_4over6_e4m3_use_256 == "all": nvfp4_e4m3_max = 256 elif self.recipe.nvfp4_4over6_e4m3_use_256 == "weights": @@ -1711,8 +1763,10 @@ def _make(tensor_type: str) -> NVFP4Quantizer: nvfp4_e4m3_max = 256 elif self.recipe.nvfp4_4over6_e4m3_use_256 == "none": nvfp4_e4m3_max = 448 + return NVFP4Quantizer( fp4_dtype=self.dtype, + scale_dtype=scale_dtype, rowwise=True, columnwise=True, with_rht=qparams.random_hadamard_transform, diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index a2e57277d1..c4cf16f069 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -7,6 +7,7 @@ from __future__ import annotations from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints import abc +import enum import warnings import math @@ -681,8 +682,8 @@ def _value_key(self) -> Tuple[Any, ...]: items = [] for name in fields: value = getattr(self, name) - if name == "dtype": - # ``DType`` is an ``IntEnum``; store the int so the key stays + if isinstance(value, enum.IntEnum): + # Store IntEnum values (like DType) as int so that the key stays # plain: hashable and ``repr``-reproducible for FX codegen. value = int(value) items.append((name, value)) diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 0cc03602a1..786316db30 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -12,6 +12,7 @@ from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage +from ..constants import DType def _stride_from_shape(shape: Tuple[int, ...]) -> Tuple[int, ...]: @@ -95,6 +96,7 @@ def __new__( row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, nvfp4_e4m3_max: int = 448, + scale_inv_dtype: Optional[DType] = None, ): if ( shapes is not None @@ -170,6 +172,7 @@ def __new__( row_scaled_nvfp4=row_scaled_nvfp4, nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, + scale_inv_dtype=scale_inv_dtype, ) return instance @@ -204,6 +207,7 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non dst.row_scaled_nvfp4 = src.row_scaled_nvfp4 dst.nvfp4_use_4over6 = src.nvfp4_use_4over6 dst.nvfp4_e4m3_max = src.nvfp4_e4m3_max + dst.scale_inv_dtype = src._scale_inv_dtype def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: """Create a wrapper of the same type and tensor metadata as src.""" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 5589e200ea..f329805fc6 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -115,6 +115,8 @@ class NVFP4Quantizer(Quantizer): """Builder class for NVFP4 tensors with NV block scaling""" dtype: DType + """Scale dtype (e4m3 block scaling factors or ue5m3 for wider dynamic range)""" + scale_dtype: DType """Random Hadamard Transform""" with_rht: bool with_post_rht_amax: bool @@ -135,6 +137,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_second_level_scale: bool """RHT sign mask (0 when sign randomization is disabled)""" rht_matrix_random_sign_mask_t: int @@ -142,6 +146,7 @@ class NVFP4Quantizer(Quantizer): def __init__( self, fp4_dtype: Union[DType, tex.DType] = DType.kFloat4E2M1, + scale_dtype: Union[DType, tex.DType] = DType.kFloat8E4M3, rowwise: bool = True, columnwise: bool = True, with_amax_reduction: bool = False, @@ -152,26 +157,39 @@ def __init__( stochastic_rounding: bool = False, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: Optional[int] = None, nvfp4_4over6_err_mode: str = "MAE", with_random_sign_mask: bool = True, + disable_second_level_scale: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = DType.cast(fp4_dtype) + self.scale_dtype = DType.cast(scale_dtype) + if self.scale_dtype not in (DType.kFloat8E4M3, DType.kFloat8UE5M3): + raise ValueError("scale_dtype must be DType.kFloat8E4M3 or DType.kFloat8UE5M3.") self.with_rht = with_rht self.with_post_rht_amax = with_post_rht_amax self.with_amax_reduction = with_amax_reduction 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 - if self.nvfp4_e4m3_max not in (448, 256): - raise ValueError("nvfp4_e4m3_max must be 448 or 256.") + if nvfp4_use_4over6 and self.scale_dtype == DType.kFloat8UE5M3: + raise ValueError("nvfp4_use_4over6 is incompatible with scale_dtype=DType.kFloat8UE5M3.") + self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_e4m3_max is not None else -1 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_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() ) @@ -230,6 +248,7 @@ def copy(self) -> NVFP4Quantizer: quantizer = NVFP4Quantizer( fp4_dtype=self.dtype, + scale_dtype=self.scale_dtype, rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, @@ -244,6 +263,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_second_level_scale=self.disable_second_level_scale, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm @@ -450,11 +470,12 @@ def __new__( amax_rowwise: Optional[torch.Tensor], amax_columnwise: Optional[torch.Tensor], fp4_dtype: DType, + scale_dtype: DType, quantizer: Quantizer, with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: Optional[int] = None, **kwargs, ): instance = super().__new__( @@ -466,6 +487,7 @@ def __new__( amax_rowwise, amax_columnwise, fp4_dtype, + scale_dtype, quantizer, with_gemm_swizzled_scales, *args, @@ -633,6 +655,7 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m # Pass amax via metadata (scalar, same on all ranks — not all-gathered) metadata = ( self._fp4_dtype, + self._scale_dtype, columnwise_usage, self._amax_rowwise, self._amax_columnwise, @@ -659,6 +682,7 @@ def fsdp_post_all_gather( """ ( fp4_dtype, + scale_dtype, columnwise_usage, amax_rowwise, amax_columnwise, @@ -698,6 +722,7 @@ def fsdp_post_all_gather( shape=logical_shape, dtype=param_dtype, fp4_dtype=fp4_dtype, + scale_dtype=scale_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, columnwise_data=None, @@ -821,7 +846,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 @@ -830,7 +859,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, @@ -842,6 +875,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): shape=tensor.shape, dtype=tensor.dtype, fp4_dtype=tensor._fp4_dtype, + scale_dtype=tensor._scale_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, columnwise_data=columnwise_data, @@ -879,6 +913,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: self._row_scaled_nvfp4, self._nvfp4_use_4over6, self._nvfp4_e4m3_max, + self._scale_dtype, ), ) @@ -1029,7 +1064,8 @@ def _make_nvfp4_tensor_in_reduce_ex( with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: Optional[int] = None, + scale_dtype: DType = DType.kFloat8E4M3, ) -> NVFP4Tensor: """Reconstruct an ``NVFP4Tensor`` from its ``__reduce_ex__`` payload.""" # Infer device from whichever inner buffer is populated so the wrapper @@ -1044,6 +1080,7 @@ def _make_nvfp4_tensor_in_reduce_ex( shape=shape, dtype=dtype, fp4_dtype=fp4_dtype, + scale_dtype=scale_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, columnwise_data=columnwise_data, @@ -1137,6 +1174,7 @@ def forward( amax_columnwise=tensor._amax_columnwise, quantizer=tensor._quantizer, fp4_dtype=tensor._fp4_dtype, + scale_dtype=tensor._scale_dtype, requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, @@ -1183,6 +1221,7 @@ def backward( amax_columnwise=grad._amax_columnwise, quantizer=grad._quantizer, fp4_dtype=grad._fp4_dtype, + scale_dtype=grad._scale_dtype, requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, @@ -1271,6 +1310,7 @@ def forward( amax_columnwise=tensor._amax_columnwise, quantizer=tensor._quantizer, fp4_dtype=tensor._fp4_dtype, + scale_dtype=tensor._scale_dtype, requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, @@ -1317,6 +1357,7 @@ def backward( amax_columnwise=grad._amax_columnwise, quantizer=grad._quantizer, fp4_dtype=grad._fp4_dtype, + scale_dtype=grad._scale_dtype, requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 3473024c03..ba0d74e87c 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -9,9 +9,10 @@ import torch from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...constants import DType, TE_DType -from ..mxfp8_tensor import MXFP8Tensor -from ..nvfp4_tensor import NVFP4Tensor +from ..mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from ..nvfp4_tensor import NVFP4Quantizer, NVFP4Tensor from ..float8_tensor import Float8Tensor from ..float8_blockwise_tensor import Float8BlockwiseQTensor from .float8_tensor_storage import Float8TensorStorage @@ -61,6 +62,7 @@ def _initialize_storage_fields( columnwise_data: Optional[torch.Tensor] = None, scale_inv: Optional[torch.Tensor] = None, columnwise_scale_inv: Optional[torch.Tensor] = None, + scale_inv_dtype: Optional[DType] = None, amax: Optional[torch.Tensor] = None, columnwise_amax: Optional[torch.Tensor] = None, scale: Optional[torch.Tensor] = None, @@ -90,6 +92,7 @@ def _initialize_storage_fields( columnwise_data: Column-wise data buffer (1D flattened) scale_inv: Row-wise scale inverse buffer columnwise_scale_inv: Column-wise scale inverse buffer + scale_inv_dtype: Data type for scale inverse buffers. amax: Row-wise amax buffer columnwise_amax: Column-wise amax buffer scale: Scale buffer (for FP8-DS only) @@ -109,6 +112,7 @@ def _initialize_storage_fields( instance.quantizer = quantizer instance.tensor_shapes = shapes instance.fake_dtype = dtype + instance.scale_inv_dtype = scale_inv_dtype # Data buffers instance.rowwise_data = data @@ -150,6 +154,7 @@ def _initialize_storage_fields( # Hold a reference to the quantized tensors that occupy same storage as the GroupedTensor. # Used as a convenience. instance.quantized_tensors = None + instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance.row_scaled_nvfp4 = row_scaled_nvfp4 instance.nvfp4_use_4over6 = nvfp4_use_4over6 @@ -167,6 +172,7 @@ def __new__( columnwise_data: Optional[torch.Tensor] = None, scale_inv: Optional[torch.Tensor] = None, columnwise_scale_inv: Optional[torch.Tensor] = None, + scale_inv_dtype: Optional[DType] = None, amax: Optional[torch.Tensor] = None, columnwise_amax: Optional[torch.Tensor] = None, scale: Optional[torch.Tensor] = None, @@ -195,6 +201,7 @@ def __new__( columnwise_data=columnwise_data, scale_inv=scale_inv, columnwise_scale_inv=columnwise_scale_inv, + scale_inv_dtype=scale_inv_dtype, amax=amax, columnwise_amax=columnwise_amax, scale=scale, @@ -343,6 +350,40 @@ def nvfp4_e4m3_max(self) -> int: def nvfp4_e4m3_max(self, nvfp4_e4m3_max: int) -> None: self._nvfp4_e4m3_max = nvfp4_e4m3_max + @property + def scale_inv_dtype(self) -> Optional[DType]: + """Data type of scale inverse buffers. + + When explicitly set, that value takes precedence. Otherwise, + the dtype is inferred based on the quantizer and scale-inverse + buffers. + """ + + # Cached value + if self._scale_inv_dtype is not None: + return self._scale_inv_dtype + + # Quantization formats with FP8 scales may store scale-inverse + # in byte buffers rather than the actual dtype. Check + # quantizer directly. + if isinstance(self.quantizer, MXFP8Quantizer): + return DType.kFloat8E8M0 + if isinstance(self.quantizer, NVFP4Quantizer): + return self.quantizer.scale_dtype + + # Get buffer dtype + if self.scale_inv is not None: + return TE_DType[self.scale_inv.dtype] + if self.columnwise_scale_inv is not None: + return TE_DType[self.columnwise_scale_inv.dtype] + + # Tensor has no scale inverse, so no scale inverse dtype + return None + + @scale_inv_dtype.setter + def scale_inv_dtype(self, dtype: Optional[DType]) -> None: + self._scale_inv_dtype = dtype + def prepare_for_saving( self, ) -> Tuple[list[Optional[torch.Tensor]], "GroupedTensorStorage"]: @@ -411,6 +452,7 @@ def clear(self) -> None: self.columnwise_data = None self.scale_inv = None self.columnwise_scale_inv = None + self.scale_inv_dtype = None self.amax = None self.columnwise_amax = None self.scale = None @@ -613,6 +655,7 @@ def copy(self) -> "GroupedTensorStorage": row_scaled_nvfp4=self.row_scaled_nvfp4, nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, + scale_inv_dtype=self._scale_inv_dtype, ) @staticmethod @@ -737,6 +780,7 @@ def make_grouped_tensor( columnwise_data = None scale_inv = None columnwise_scale_inv = None + scale_inv_dtype = None amax = None columnwise_amax = None scale = None @@ -755,6 +799,11 @@ def make_grouped_tensor( # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=dtype, device=device) elif compatible_recipe.mxfp8(): + # Amax buffer for delayed scaling - one per tensor + amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + + scale_inv_dtype = DType.kFloat32 + if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -784,6 +833,7 @@ def make_grouped_tensor( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) elif compatible_recipe.delayed(): + scale_inv_dtype = DType.kFloat8E8M0 if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -799,10 +849,8 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors columnwise_scale_inv_offsets = list(range(num_tensors + 1)) - - # Amax buffer for delayed scaling - one per tensor - amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif compatible_recipe.nvfp4(): + scale_inv_dtype = quantizer.scale_dtype row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 nvfp4_use_4over6 = quantizer.nvfp4_use_4over6 nvfp4_e4m3_max = quantizer.nvfp4_e4m3_max @@ -850,6 +898,8 @@ def make_grouped_tensor( ) columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif compatible_recipe.float8_block_scaling(): + scale_inv_dtype = DType.kFloat32 + if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) @@ -881,6 +931,7 @@ def make_grouped_tensor( non_tn_fp8_gemm_supported = is_non_tn_fp8_gemm_supported() fp8_rowwise_usage = rowwise_usage or non_tn_fp8_gemm_supported fp8_columnwise_usage = columnwise_usage and not non_tn_fp8_gemm_supported + scale_inv_dtype = DType.kFloat32 shared_scale_inv = None if fp8_rowwise_usage or fp8_columnwise_usage: shared_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) @@ -940,6 +991,7 @@ def make_grouped_tensor( row_scaled_nvfp4=row_scaled_nvfp4, nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, + scale_inv_dtype=scale_inv_dtype, ) grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() return grouped_tensor @@ -1064,6 +1116,7 @@ def split_into_quantized_tensors( row_scaled_nvfp4 = self.row_scaled_nvfp4 nvfp4_use_4over6 = self.nvfp4_use_4over6 nvfp4_e4m3_max = self.nvfp4_e4m3_max + scale_inv_dtype = self.scale_inv_dtype if recipe.nvfp4() and row_scaled_nvfp4: cum = 0 nvfp4_rowwise_amax_offsets = [0] @@ -1295,6 +1348,7 @@ def split_into_quantized_tensors( row_scaled_nvfp4=row_scaled_nvfp4, nvfp4_use_4over6=nvfp4_use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, + scale_dtype=scale_inv_dtype, ) result.append(tensor) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 7e0861c967..69e8aace0b 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -87,13 +87,15 @@ class NVFP4TensorStorage(QuantizedTensorStorage): _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")] _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")] # Input absolute maximum values, used to compute the tensor scale - _amax_rowwise: Annotated[torch.Tensor, InnerTensor("amax_rowwise")] - _amax_columnwise: Annotated[torch.Tensor, InnerTensor("amax_columnwise")] + _amax_rowwise: Annotated[Optional[torch.Tensor], InnerTensor("amax_rowwise")] + _amax_columnwise: Annotated[Optional[torch.Tensor], InnerTensor("amax_columnwise")] # Builder class for casting to MXFP8 _quantizer: Optional[Quantizer] # FP4 data type _fp4_dtype: DType + # Data type for block scaling factors + _scale_dtype: DType # Whether scaling factors are in the swizzled format expected by # GEMM _with_gemm_swizzled_scales: bool @@ -102,7 +104,7 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Whether this NVFP4 tensor uses 4over6 map-to-4/map-to-6 block selection _nvfp4_use_4over6: bool # Global E4M3 scale bound used by this NVFP4 tensor - _nvfp4_e4m3_max: int + _nvfp4_e4m3_max: Optional[int] def __new__( cls, @@ -113,13 +115,14 @@ def __new__( amax_rowwise: torch.Tensor, amax_columnwise: torch.Tensor, fp4_dtype: Union[DType, tex.DType], + scale_dtype: Union[DType, tex.DType], quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, *args, fake_dtype: Optional[torch.dtype] = None, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: Optional[int] = None, **kwargs, ): if cls is NVFP4TensorStorage: @@ -131,6 +134,7 @@ def __new__( instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data instance._fp4_dtype = DType.cast(fp4_dtype) + instance._scale_dtype = DType.cast(scale_dtype) instance._quantizer = quantizer.copy() if quantizer is not None else None instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv @@ -139,7 +143,7 @@ def __new__( instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance._row_scaled_nvfp4 = row_scaled_nvfp4 instance._nvfp4_use_4over6 = nvfp4_use_4over6 - instance._nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 + instance._nvfp4_e4m3_max = nvfp4_e4m3_max return instance @@ -162,6 +166,8 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: raise TypeError("copy_from_storage expects NVFP4TensorStorage") if self._fp4_dtype != src._fp4_dtype: raise RuntimeError("FP4 dtype mismatch in copy_from_storage") + if self._scale_dtype != src._scale_dtype: + raise RuntimeError("Scale dtype mismatch in copy_from_storage") if self._with_gemm_swizzled_scales != src._with_gemm_swizzled_scales: raise RuntimeError("Scale layout mismatch in copy_from_storage") if self._row_scaled_nvfp4 != src._row_scaled_nvfp4: @@ -192,6 +198,7 @@ def get_metadata(self) -> Dict[str, Any]: "amax_rowwise": self._amax_rowwise, "amax_columnwise": self._amax_columnwise, "fp4_dtype": self._fp4_dtype, + "scale_dtype": self._scale_dtype, "quantizer": self._quantizer, "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, "row_scaled_nvfp4": self._row_scaled_nvfp4, @@ -328,6 +335,7 @@ def view(self, shape: torch.Size): amax_columnwise=self._amax_columnwise, quantizer=self._quantizer, fp4_dtype=self._fp4_dtype, + scale_dtype=self._scale_dtype, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, row_scaled_nvfp4=self._row_scaled_nvfp4, nvfp4_use_4over6=self._nvfp4_use_4over6, @@ -365,13 +373,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_second_level_scale + ) # 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." @@ -390,7 +401,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" @@ -411,7 +422,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" @@ -474,7 +485,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_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 dede7c7f54fec733426bb0226a622c85a237bb25 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:00:59 +0000 Subject: [PATCH 02/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_fusible_ops.py | 15 +++-- tests/pytorch/utils.py | 3 +- .../common/cast/dispatch/quantize.cuh | 20 +++--- .../common/cast/nvfp4/core_nvfp4.cuh | 19 +++--- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 45 +++++++------- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 30 ++++----- transformer_engine/common/common.h | 22 +++---- .../common/gemm/cublaslt_gemm.cu | 2 +- .../common/gemm/cublaslt_grouped_gemm.cu | 22 +++---- ...cast_col_hadamard_transform_cast_fusion.cu | 4 +- .../transformer_engine/transformer_engine.h | 24 +++---- transformer_engine/common/recipe/__init__.py | 6 +- .../common/transformer_engine.cpp | 62 +++++++++---------- .../pytorch/csrc/extensions/cast.cpp | 30 ++++----- .../pytorch/csrc/type_converters.cpp | 12 ++-- transformer_engine/pytorch/quantization.py | 10 ++- .../pytorch/tensor/nvfp4_tensor.py | 4 +- 17 files changed, 161 insertions(+), 169 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 8da1944f53..cc8ded1cfb 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -206,16 +206,19 @@ def make_reference_and_test_tensors( columnwise=True, block_scaling_dim=2 if tensor_type == "weight" else 1, )(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht", "nvfp4_ue5m3", "nvfp4_rht_ue5m3"): + elif quantization in ( + "nvfp4", + "nvfp4_row_scaled", + "nvfp4_rht", + "nvfp4_ue5m3", + "nvfp4_rht_ue5m3", + ): tensor_type = "input" if quantizer_role is not None: tensor_type = quantizer_role.tensor_type - with_rht = ( - quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" - ) + with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" scale_dtype = ( - te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" - else te.DType.kFloat8E4M3 + te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3 ) test = NVFP4Quantizer( scale_dtype=scale_dtype, diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index a845a48911..353dbf8f60 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -162,8 +162,7 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: with_rht = name in ("nvfp4_rht", "nvfp4_rht_ue5m3") use_4over6 = name == "nvfp4_4over6" scale_format = ( - RecipeFormat.UE5M3 if name in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") - else RecipeFormat.E4M3 + RecipeFormat.UE5M3 if name in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") else RecipeFormat.E4M3 ) kwargs = { "disable_rht": not with_rht, diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index a43a43ab00..f10b165ad2 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -104,11 +104,10 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - NVTE_CHECK( - nvfp4_use_4over6 || - output_tensor->get_nvfp4_scale_max() == - static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), - "NVFP4 quantization with non-default scale max is only supported with 4over6."); + NVTE_CHECK(nvfp4_use_4over6 || + output_tensor->get_nvfp4_scale_max() == + static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); if (row_scaled_nvfp4) { @@ -290,8 +289,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || - output_tensor->get_nvfp4_scale_max() == - static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + output_tensor->get_nvfp4_scale_max() == + static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -455,9 +454,10 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; if (!nvfp4_use_4over6) { for (const auto *output_tensor : output_tensors) { - NVTE_CHECK(output_tensor->get_nvfp4_scale_max() - == static_cast(nvfp4::core::scale_max(output_tensors[0]->scale_inv.dtype)), - "NVFP4 quantization with non-default scale max is only supported with 4over6."); + NVTE_CHECK( + output_tensor->get_nvfp4_scale_max() == + static_cast(nvfp4::core::scale_max(output_tensors[0]->scale_inv.dtype)), + "NVFP4 quantization with non-default scale max is only supported with 4over6."); } } NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 0ee4589d6f..b89dd755a9 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -107,22 +107,21 @@ __host__ __device__ constexpr float scale_max() { // Return the full-range maximum for a runtime scale dtype. inline float scale_max(const DType scale_dtype) { float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, result = scale_max();) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, + result = scale_max();) return result; } // Return and validate a user-provided maximum for a runtime scale dtype. inline float scale_max(const DType scale_dtype, const int scale_type_max) { float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, { - using ScaleTraits = NVFP4ScaleTraits; - NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - result = static_cast(scale_type_max); - }) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { + using ScaleTraits = NVFP4ScaleTraits; + NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || + scale_type_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + result = static_cast(scale_type_max); + }) return result; } diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 08e8993f83..a014244b9b 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -66,8 +66,7 @@ __global__ void __launch_bounds__(512) value.vec = input_vectorized[my_index]; ScaleType scale = scales[my_scale_index]; constexpr float fp4_max = detail::TypeExtrema::max; - constexpr float unit_global_scale_amax = - fp4_max * core::scale_max(); + constexpr float unit_global_scale_amax = fp4_max * core::scale_max(); float amax = unit_global_scale_amax; if (tensor_amax != nullptr) { amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; @@ -91,10 +90,10 @@ __global__ void __launch_bounds__(512) #if FP4_TYPE_SUPPORTED template inline void launch_dequantize(const Tensor &input, Tensor *output, - const bool with_gemm_swizzled_scales, - const bool row_scaled_nvfp4, const size_t N, const size_t Mread, - const size_t blocks, const size_t threads, - const size_t num_scale_tiles_X, cudaStream_t stream) { + const bool with_gemm_swizzled_scales, const bool row_scaled_nvfp4, + const size_t N, const size_t Mread, const size_t blocks, + const size_t threads, const size_t num_scale_tiles_X, + cudaStream_t stream) { using namespace dequantize_kernel; TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( output->data.dtype, OType, @@ -102,9 +101,8 @@ inline void launch_dequantize(const Tensor &input, Tensor *output, with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, TRANSFORMER_ENGINE_SWITCH_CONDITION( row_scaled_nvfp4, ROW_SCALED_NVFP4, - dequantize_fp4_kernel - <<>>( + dequantize_fp4_kernel<<>>( input.data.dptr, reinterpret_cast(output->data.dptr), reinterpret_cast(input.scale_inv.dptr), reinterpret_cast(input.amax.dptr), N, Mread, @@ -141,21 +139,20 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) "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_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, { - using ScaleTraits = core::NVFP4ScaleTraits; - if (e4m3_max == static_cast(ScaleTraits::expected_max)) { - launch_dequantize(ScaleTraits::expected_max)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, - threads, num_scale_tiles_X, stream); - } else { - NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - launch_dequantize(ScaleTraits::headroom_max)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, - threads, num_scale_tiles_X, stream); - } - }) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { + using ScaleTraits = core::NVFP4ScaleTraits; + if (e4m3_max == static_cast(ScaleTraits::expected_max)) { + launch_dequantize(ScaleTraits::expected_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + } else { + NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + launch_dequantize(ScaleTraits::headroom_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + } + }) NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 9e287b2bd6..d5a220d8ef 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -556,10 +556,9 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, S block_amax = reduce_group_max_16(group_amax); } - const float global_amax = - amax == nullptr - ? core::scale_max() * detail::TypeExtrema::max - : amax[0]; + const float global_amax = amax == nullptr ? core::scale_max() * + detail::TypeExtrema::max + : amax[0]; const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); CandidatePair candidates = @@ -678,9 +677,9 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { - auto kernel = quantize_4over6_kernel; + auto kernel = + quantize_4over6_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, @@ -740,14 +739,12 @@ void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *outpu using ScaleTraits = core::NVFP4ScaleTraits; const int scale_type_max = output->get_nvfp4_scale_max(); NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), + scale_type_max == static_cast(ScaleTraits::headroom_max), "Unsupported maximum for NVFP4 scale dtype."); TRANSFORMER_ENGINE_SWITCH_CONDITION( - scale_type_max == static_cast(ScaleTraits::headroom_max), - USE_SCALE_HEADROOM, { - constexpr int SCALE_TYPE_MAX = - static_cast(USE_SCALE_HEADROOM ? ScaleTraits::headroom_max - : ScaleTraits::expected_max); + scale_type_max == static_cast(ScaleTraits::headroom_max), USE_SCALE_HEADROOM, { + constexpr int SCALE_TYPE_MAX = static_cast( + USE_SCALE_HEADROOM ? ScaleTraits::headroom_max : ScaleTraits::expected_max); TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( quant_config->nvfp4_4over6_mode, MODE, TRANSFORMER_ENGINE_SWITCH_CONDITION( @@ -755,10 +752,9 @@ void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *outpu using Cfg = quantize_4over6_kernel::Config; TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( input.dtype(), IType, - quantize_4over6_kernel::launch_quantize_4over6(input, noop, output, - stream);); + quantize_4over6_kernel::launch_quantize_4over6< + use_2d_quantization, Cfg, ScaleType, SCALE_TYPE_MAX, IType>( + input, noop, output, stream);); });); }) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 9b309ac267..44fd2d2d12 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -474,12 +474,12 @@ struct Tensor { dtype = scale_inv.dtype; } switch (dtype) { - case DType::kFloat8E4M3: - return 448; - case DType::kFloat8UE5M3: - return 114688; - default: - NVTE_ERROR("Unsupported scale dtype for NVFP4 tensor (", to_string(dtype), ")"); + case DType::kFloat8E4M3: + return 448; + case DType::kFloat8UE5M3: + return 114688; + default: + NVTE_ERROR("Unsupported scale dtype for NVFP4 tensor (", to_string(dtype), ")"); } } }; @@ -850,10 +850,10 @@ struct TypeInfo { #define SWITCH_FP4_TYPE_HANDLE(type, ...) // do nothing #endif #if CUDA_VERSION >= 13040 -#define SWITCH_FP8UE5M3_TYPE_HANDLE(type, ...) \ - case DType::kFloat8UE5M3: { \ - using type = fp8ue5m3; \ - { __VA_ARGS__ } \ +#define SWITCH_FP8UE5M3_TYPE_HANDLE(type, ...) \ + case DType::kFloat8UE5M3: { \ + using type = fp8ue5m3; \ + { __VA_ARGS__ } \ } break; #else #define SWITCH_FP8UE5M3_TYPE_HANDLE(type, ...) // do nothing @@ -908,7 +908,7 @@ struct TypeInfo { NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ ". Expected one of: Byte, Int16, Int32, Int64, Float32, " \ "Float16, BFloat16, Float8E4M3, Float8E5M2, " \ - "Float8E8M0."); \ + "Float8E8M0."); \ } #define TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(dtype, type, ...) \ diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index c566d16f5b..451155e1f4 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -578,7 +578,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, &B_scale_inverse, sizeof(B_scale_inverse))); // Deduce cuBLAS scale mode based on scale dtype - auto get_scale_mode = [] (DType dtype) -> cublasLtMatmulMatrixScale_t { + auto get_scale_mode = [](DType dtype) -> cublasLtMatmulMatrixScale_t { if (dtype == DType::kFloat8E4M3) { return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; } diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 3814aae8b9..5ed8af0e1e 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1069,8 +1069,8 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac setup_workspace.b_scale_inv_ptrs); } else if (transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { set_nvfp4_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, - setup_workspace.b_scale_inv_ptrs, - A_sel.scale_inv_dtype, B_sel.scale_inv_dtype); + setup_workspace.b_scale_inv_ptrs, A_sel.scale_inv_dtype, + B_sel.scale_inv_dtype); } else if (transformer_engine::is_fp8_block_scaling(A_sel.scaling_mode)) { set_fp8_block_scaling_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, setup_workspace.b_scale_inv_ptrs, A_sel.scaling_mode, @@ -1352,8 +1352,8 @@ __global__ void setup_grouped_gemm_kernel( MultiTensorGroupGemmOutputArgs c_multi_tensor_args, MultiTensorGroupGemmOutputArgs d_multi_tensor_args, // NVFP4: per-group amax values and output buffer for computed alpha - float *a_amax, float *b_amax, float *nvfp4_computed_alpha, - float a_unit_global_scale_amax, float b_unit_global_scale_amax) { + float *a_amax, float *b_amax, float *nvfp4_computed_alpha, float a_unit_global_scale_amax, + float b_unit_global_scale_amax) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_tensors) return; @@ -1433,8 +1433,7 @@ __global__ void setup_grouped_gemm_kernel( const float b_amax_val = b_amax == nullptr ? b_unit_global_scale_amax : b_amax[idx]; const float nvfp4_alpha_factor_inv = 1.0f / (a_unit_global_scale_amax * b_unit_global_scale_amax); - nvfp4_computed_alpha[idx] = - alpha_ptr[idx] * a_amax_val * b_amax_val * nvfp4_alpha_factor_inv; + nvfp4_computed_alpha[idx] = alpha_ptr[idx] * a_amax_val * b_amax_val * nvfp4_alpha_factor_inv; alpha_ptrs[idx] = &nvfp4_computed_alpha[idx]; } else { alpha_ptrs[idx] = alpha_ptr + idx; @@ -1570,11 +1569,9 @@ inline void launch_grouped_gemm_setup( constexpr float kFP4Max = transformer_engine::detail::TypeExtrema::max; a_unit_global_scale_amax = - transformer_engine::dispatch::nvfp4::core::scale_max(A_sel.scale_inv_dtype) * - kFP4Max; + transformer_engine::dispatch::nvfp4::core::scale_max(A_sel.scale_inv_dtype) * kFP4Max; b_unit_global_scale_amax = - transformer_engine::dispatch::nvfp4::core::scale_max(B_sel.scale_inv_dtype) * - kFP4Max; + transformer_engine::dispatch::nvfp4::core::scale_max(B_sel.scale_inv_dtype) * kFP4Max; } setup_grouped_gemm_kernel<<>>( @@ -1797,8 +1794,9 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num A_sel.amax = nullptr; if (nvfp4) { - const auto& A_tensor0 = *transformer_engine::convertNVTETensorCheck(A_list[0]); - A_sel.scale_inv_dtype = transa ? A_tensor0.scale_inv.dtype : A_tensor0.columnwise_scale_inv.dtype; + const auto &A_tensor0 = *transformer_engine::convertNVTETensorCheck(A_list[0]); + A_sel.scale_inv_dtype = + transa ? A_tensor0.scale_inv.dtype : A_tensor0.columnwise_scale_inv.dtype; } // Workspaces: setup (pointer arrays) and cuBLAS 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 6f560cb232..39a49da36b 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 @@ -1331,8 +1331,8 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vectorscale_inv.dtype : output_list[i]->columnwise_scale_inv.dtype; + const DType output_scale_dtype = has_row_quant ? output_list[i]->scale_inv.dtype + : output_list[i]->columnwise_scale_inv.dtype; if (has_row_quant && has_col_quant) { NVTE_CHECK(output_list[i]->columnwise_scale_inv.dtype == output_scale_dtype, "Rowwise and columnwise NVFP4 scales must use the same dtype."); diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index ff09194255..9ce6d7b044 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -23,19 +23,19 @@ extern "C" { * \brief TE datatype. */ enum NVTEDType { - kNVTEByte = 0, /*!< Byte */ - kNVTEInt16 = 1, /*!< 16-bit integer */ - kNVTEInt32 = 2, /*!< 32-bit integer */ - kNVTEInt64 = 3, /*!< 64-bit integer */ - kNVTEFloat32 = 4, /*!< 32-bit float */ - kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ - kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ - kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ - kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ - kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ - kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ + kNVTEByte = 0, /*!< Byte */ + kNVTEInt16 = 1, /*!< 16-bit integer */ + kNVTEInt32 = 2, /*!< 32-bit integer */ + kNVTEInt64 = 3, /*!< 64-bit integer */ + kNVTEFloat32 = 4, /*!< 32-bit float */ + kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ + kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ + kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ + kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ + kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ + kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ - kNVTENumTypes /*!< Number of supported types */ + kNVTENumTypes /*!< Number of supported types */ }; /*! \struct NVTEShape diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 525238b3e0..b8a703fbaa 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -16,6 +16,7 @@ _NVFP4_4OVER6_SCOPES = ("none", "weights", "activations", "all") _NVFP4_4OVER6_ERR_MODES = ("MAE", "MSE") + class _FormatHelper(NamedTuple): """ Stores max FP8 values for fprop and bprop a `Format`. @@ -573,7 +574,10 @@ class NVFP4BlockScaling(Recipe): def __post_init__(self) -> None: assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" - assert self.fp8_format in (Format.E4M3, Format.UE5M3), "Unsupported format for NVFP4 scaling." + assert self.fp8_format in ( + Format.E4M3, + Format.UE5M3, + ), "Unsupported format for NVFP4 scaling." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 172959fa4d..375d21bdfe 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -173,18 +173,18 @@ void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_i if (t.has_data()) { NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor input ", name, "_scale_inverse must be allocated"); - NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 - || t.scale_inv.dtype == DType::kFloat8UE5M3, - "FP4 scaling factor input ", name, - "_scale_inverse has invalid dtype " - "(expected Float8E4M3 or Float8UE5M3, got ", - to_string(t.scale_inv.dtype), ")"); + NVTE_CHECK( + t.scale_inv.dtype == DType::kFloat8E4M3 || t.scale_inv.dtype == DType::kFloat8UE5M3, + "FP4 scaling factor input ", name, + "_scale_inverse has invalid dtype " + "(expected Float8E4M3 or Float8UE5M3, got ", + to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor input ", name, "_columnwise_scale_inverse must be allocated"); - NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 - || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 || + t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, "FP8 scaling factor input ", name, "_columnwise_scale_inverse has invalid dtype " "(expected Float8E4M3 or Float8UE5M3, got ", @@ -237,18 +237,18 @@ void CheckOutputTensor(const Tensor &t, std::string_view name, bool allow_empty) if (t.has_data()) { NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor output ", name, "_scale_inverse must be allocated"); - NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 - || t.scale_inv.dtype == DType::kFloat8UE5M3, - "FP4 scaling factor output ", name, - "_scale_inverse has invalid dtype " - "(expected Float8E4M3 or Float8UE5M3, got ", - to_string(t.scale_inv.dtype), ")"); + NVTE_CHECK( + t.scale_inv.dtype == DType::kFloat8E4M3 || t.scale_inv.dtype == DType::kFloat8UE5M3, + "FP4 scaling factor output ", name, + "_scale_inverse has invalid dtype " + "(expected Float8E4M3 or Float8UE5M3, got ", + to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor output ", name, "_columnwise_scale_inverse must be allocated"); - NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 - || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 || + t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, "FP4 scaling factor output ", name, "_columnwise_scale_inverse has invalid dtype " "(expected Float8E4M3 or Float8UE5M3, got ", @@ -370,18 +370,18 @@ static void CheckGroupedScaleInv(const GroupedTensor &t, std::string_view name, if (t.has_data()) { NVTE_CHECK(t.scale_inv.has_data(), tensor_type, " ", name, " rowwise scale_inv must be allocated"); - NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3 - || t.scale_inv.dtype == DType::kFloat8UE5M3, - tensor_type, " ", name, - " rowwise scale_inv has invalid dtype " - "(expected Float8E4M3 or Float8UE5M3, got ", - to_string(t.scale_inv.dtype), ")"); + NVTE_CHECK( + t.scale_inv.dtype == DType::kFloat8E4M3 || t.scale_inv.dtype == DType::kFloat8UE5M3, + tensor_type, " ", name, + " rowwise scale_inv has invalid dtype " + "(expected Float8E4M3 or Float8UE5M3, got ", + to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { NVTE_CHECK(t.columnwise_scale_inv.has_data(), tensor_type, " ", name, " columnwise scale_inv must be allocated"); - NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 - || t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, + NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3 || + t.columnwise_scale_inv.dtype == DType::kFloat8UE5M3, tensor_type, " ", name, " columnwise scale_inv has invalid dtype " "(expected Float8E4M3 or Float8UE5M3, got ", @@ -924,8 +924,8 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo case kNVTENVFP4E4M3Max: std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); // Need to rename this to nvfp4_scale_type_max - NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || - t.nvfp4_e4m3_max == 114688 || t.nvfp4_e4m3_max == 65536, + NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || t.nvfp4_e4m3_max == 114688 || + t.nvfp4_e4m3_max == 65536, "Unsupported NVFP4 scale type max (got ", t.nvfp4_e4m3_max, ")"); break; default: @@ -1011,12 +1011,10 @@ void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, vo case kNVTERowScaledNVFP4: *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); break; - case kNVTENVFP4E4M3Max: - { - int val = t->get_nvfp4_scale_max(); - std::memcpy(buf, &val, attr_size); - } - break; + case kNVTENVFP4E4M3Max: { + int val = t->get_nvfp4_scale_max(); + std::memcpy(buf, &val, attr_size); + } break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index ffd9545516..466114b7f2 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1287,13 +1287,13 @@ std::tuple, std::vector, bool> bulk_alloc : py::none(); // Construct Python tensor. - tensor_py_list.emplace_back(NVFP4TensorClass( - rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, amax_rowwise, - amax_columnwise, MakePythonDType(fp4_dtype), MakePythonDType(scale_dtype), - quantizer_py_list[i], with_gemm_swizzled_scales, - py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, - py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, - py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); + tensor_py_list.emplace_back( + NVFP4TensorClass(rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, + amax_rowwise, amax_columnwise, MakePythonDType(fp4_dtype), + MakePythonDType(scale_dtype), quantizer_py_list[i], + with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, + py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, + py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); // Construct C++ tensor // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, @@ -1301,20 +1301,20 @@ std::tuple, std::vector, bool> bulk_alloc { TensorWrapper tensor_wrapper(NVTE_NVFP4_1D_SCALING); if (rowwise_usage) { - tensor_wrapper.set_rowwise_data(rowwise_data_list[i].data_ptr(), - fp4_dtype, rowwise_data_shapes[i]); - tensor_wrapper.set_rowwise_scale_inv(rowwise_scale_list[i].data_ptr(), - scale_dtype, rowwise_scale_shapes[i]); + tensor_wrapper.set_rowwise_data(rowwise_data_list[i].data_ptr(), fp4_dtype, + rowwise_data_shapes[i]); + tensor_wrapper.set_rowwise_scale_inv(rowwise_scale_list[i].data_ptr(), scale_dtype, + rowwise_scale_shapes[i]); if (!disable_second_level_scale) { tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise_list[i])); } } if (columnwise_usage) { - tensor_wrapper.set_columnwise_data(columnwise_data_list[i].data_ptr(), - fp4_dtype, columnwise_data_shapes[i]); - tensor_wrapper.set_columnwise_scale_inv(columnwise_scale_list[i].data_ptr(), - scale_dtype, columnwise_scale_shapes[i]); + tensor_wrapper.set_columnwise_data(columnwise_data_list[i].data_ptr(), fp4_dtype, + columnwise_data_shapes[i]); + tensor_wrapper.set_columnwise_scale_inv(columnwise_scale_list[i].data_ptr(), scale_dtype, + columnwise_scale_shapes[i]); if (!disable_second_level_scale) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, std::vector{1}); diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index a53cc3ffe7..710c228648 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -4,13 +4,13 @@ * See LICENSE for license information. ************************************************************************/ -#include -#include - #include #include #include +#include +#include + #include "common.h" #include "pybind.h" @@ -163,8 +163,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const auto &scale_inv = tensor.attr("_columnwise_scale_inv").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(), scale_inv_dtype, - getTensorShape(scale_inv)); + ret.set_columnwise_scale_inv(scale_inv.data_ptr(), scale_inv_dtype, getTensorShape(scale_inv)); const auto amax_columnwise = tensor.attr("_amax_columnwise"); if (!amax_columnwise.is_none()) { const auto &amax = amax_columnwise.cast(); @@ -282,8 +281,7 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { if (!tensor.attr("columnwise_scale_inv").is_none()) { const auto &scale_inv = tensor.attr("columnwise_scale_inv").cast(); NVTE_CHECK(scale_inv_dtype, "Could not determine dtype of scale_inv buffer."); - ret.set_columnwise_scale_inv(scale_inv.data_ptr(), *scale_inv_dtype, - getTensorShape(scale_inv)); + ret.set_columnwise_scale_inv(scale_inv.data_ptr(), *scale_inv_dtype, getTensorShape(scale_inv)); } // Shape metadata diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 06175faebd..9a4493e3df 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -228,12 +228,14 @@ def check_fp8_ue5m3_support() -> Tuple[bool, str]: """Return if the FP8 UE5M3 format is available.""" global _FP8_UE5M3_SUPPORT if _FP8_UE5M3_SUPPORT is None: + def _check_support() -> Tuple[bool, str]: if get_device_compute_capability() != (10, 7): # Rubin return False, "Device compute capability 10.7 is required for FP8 UE5M3 support." if float(torch.version.cuda) < 13.4: return False, "CUDA 13.4 is required for FP8 UE5M3 support." return True, "" + _FP8_UE5M3_SUPPORT = _check_support() return _FP8_UE5M3_SUPPORT @@ -1719,9 +1721,7 @@ def _qparams(tensor_type: str): return self.recipe.fp4_quant_fwd_inp scale_dtype = ( - DType.kFloat8UE5M3 - if self.recipe.fp8_format == Format.UE5M3 - else DType.kFloat8E4M3 + DType.kFloat8UE5M3 if self.recipe.fp8_format == Format.UE5M3 else DType.kFloat8E4M3 ) def _make(tensor_type: str) -> NVFP4Quantizer: @@ -1746,9 +1746,7 @@ def _make(tensor_type: str) -> NVFP4Quantizer: "NVFP4 4over6 quantization does not support stochastic rounding." ) if scale_dtype == DType.kFloat8UE5M3: - raise ValueError( - "NVFP4 4over6 quantization is incompatible with UE5M3 scales." - ) + raise ValueError("NVFP4 4over6 quantization is incompatible with UE5M3 scales.") # Scale max for 4over6 nvfp4_e4m3_max = None diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index f329805fc6..39cc6588a6 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -184,7 +184,9 @@ def __init__( self.row_scaled_nvfp4 = row_scaled_nvfp4 self.nvfp4_use_4over6 = nvfp4_use_4over6 if nvfp4_use_4over6 and self.scale_dtype == DType.kFloat8UE5M3: - raise ValueError("nvfp4_use_4over6 is incompatible with scale_dtype=DType.kFloat8UE5M3.") + raise ValueError( + "nvfp4_use_4over6 is incompatible with scale_dtype=DType.kFloat8UE5M3." + ) self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_e4m3_max is not None else -1 self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): From 343c4bd46f32b2aae540a19efb41253732fe609d Mon Sep 17 00:00:00 2001 From: Kaining Zhong <44538064+kainzhong@users.noreply.github.com> Date: Thu, 13 Aug 2026 04:45:09 -0700 Subject: [PATCH 03/54] [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN (#2) * [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN Signed-off-by: Kaining Zhong * have to pad to 256 to use cuDNN Signed-off-by: Kaining Zhong * fix: need to pass scale_dtype Signed-off-by: Kaining Zhong * route wgrad to cuDNN's wgrad API Signed-off-by: Kaining Zhong * Support grouped linear with NVFP4-UE5M3 NVFP4-UE5M3 grouped GEMM falls back to dense GEMMs. Generalize usage of wgrad kernel and use when tensors sizes are not 256-aligned. Fix inconsistent m,n,k GEMM notation. Remove ue5m3 hacks in op fuser tests. Add ue5m3 to grouped MLP tests. Signed-off-by: Tim Moon * Fix typos Co-authored-by: Codex Signed-off-by: Tim Moon --------- Signed-off-by: Kaining Zhong Signed-off-by: Tim Moon Co-authored-by: Tim Moon Co-authored-by: Codex --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 109 ++++ tests/pytorch/test_fusible_ops.py | 14 +- tests/pytorch/test_grouped_mlp.py | 39 +- .../pytorch/cpp_extensions/gemm.py | 598 ++++++++++++++++-- transformer_engine/pytorch/csrc/extensions.h | 17 +- .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 11 +- .../pytorch/csrc/extensions/pybind.cpp | 15 +- .../pytorch/csrc/extensions/transpose.cpp | 23 +- .../pytorch/module/grouped_linear.py | 99 +-- .../pytorch/ops/basic/grouped_linear.py | 75 ++- .../pytorch/ops/fused/grouped_mlp.py | 380 ++++++++++- transformer_engine/pytorch/tensor/utils.py | 14 +- 12 files changed, 1209 insertions(+), 185 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 639b2f752e..40c7c3bf82 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -690,3 +690,112 @@ def test_nvfp4_row_scaled_gemm_matches_emulated( use_4over6=use_4over6, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) + + +def _check_ue5m3_gemm_versus_dequantized( + M, K, N, x_columnwise, w_columnwise, disable_second_level_scale +): + """Run an NVFP4/UE5M3 GEMM and compare against a dequantized FP32 reference.""" + if M % 256 != 0: + pytest.skip( + "cuDNN's grouped GEMM pads every group to 256 rows, so the UE5M3 path (which " + "routes there while cuBLAS lacks UE5M3 kernels) requires M % 256 == 0." + ) + torch.manual_seed(0) + device, dtype, out_dtype = "cuda", torch.bfloat16, torch.bfloat16 + x_shape = (K, M) if x_columnwise else (M, K) + w_shape = (K, N) if w_columnwise else (N, K) + x = torch.randn(x_shape, dtype=dtype, device=device) + w = torch.randn(w_shape, dtype=dtype, device=device) + + common = dict( + fp4_dtype=tex.DType.kFloat4E2M1, + scale_dtype=tex.DType.kFloat8UE5M3, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + # disable_second_level_scale is given per operand, as (x, w). + xq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[0]) + wq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[1]) + x_q = xq.update_quantized(x, xq.make_empty(x_shape, dtype=dtype, device=device)) + w_q = wq.update_quantized(w, wq.make_empty(w_shape, dtype=dtype, device=device)) + + if disable_second_level_scale[0]: + assert x_q._amax_rowwise is None, "disable_second_level_scale should drop the amax" + if disable_second_level_scale[1]: + assert w_q._amax_rowwise is None, "disable_second_level_scale should drop the amax" + + # Reference: dequantize the orientation each operand is actually read in. + x_ref = _dequantize_nvfp4_usage(x_q, columnwise=x_columnwise) + w_ref = _dequantize_nvfp4_usage(w_q, columnwise=w_columnwise) + # _dequantize_nvfp4_usage returns each operand canonically as (rows, K), so + # the reference is the same expression for every layout. + ref = x_ref @ w_ref.t() + + if x_columnwise: + x_q.update_usage(rowwise_usage=False) + if w_columnwise: + w_q.update_usage(rowwise_usage=False) + transa, transb = not w_columnwise, x_columnwise + layout = ("T" if transa else "N") + ("T" if transb else "N") + y = general_gemm(w_q, x_q, out_dtype=out_dtype, layout=layout)[0] + + # Both sides see identically quantized operands, so quantization error cancels and + # only accumulation order and the bf16 output rounding differ. One bf16 ulp is + # already ~4e-3 relative, which no elementwise tolerance survives, so compare the + # whole result instead. + rel_err = (y.float() - ref).norm() / ref.norm() + assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" + +ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3) +@pytest.mark.parametrize( + "M, K, N", + [ + (256, 128, 256), + (256, 256, 256), + (256, 1024, 256), + (1024, 1024, 1024), + (4096, 512, 3072), + (112, 128, 96), + (304, 640, 304), + (1008, 3072, 992), + (256, 64, 256), + (128, 128, 112), + ], +) +@pytest.mark.parametrize( + "x_columnwise, w_columnwise", + [ + (False, False), # TN -- w rowwise, x rowwise (fprop) + (False, True), # NN -- w colwise, x rowwise (dgrad) + (True, True), # NT -- w colwise, x colwise (wgrad) + ], ids=["FF", "FT", "TT"] +) +@pytest.mark.parametrize( + "disable_second_level_scale", [ + (True, False), + ], ids=["TF"] +) +def test_nvfp4_ue5m3_gemm_versus_reference( + M: int, + K: int, + N: int, + x_columnwise: bool, + w_columnwise: bool, + disable_second_level_scale: bool, +): + """NVFP4 GEMM with UE5M3 block scales, with and without second-level scaling. + + UE5M3's wider range is what makes dropping the per-tensor global scale + viable, so both configurations must match the dequantized reference. + """ + _check_ue5m3_gemm_versus_dequantized( + M, K, N, x_columnwise, w_columnwise, disable_second_level_scale + ) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index cc8ded1cfb..1ae33c0403 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -82,6 +82,8 @@ if nvfp4_available: _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") + if fp8_ue5m3_available: + _quantization_list.append("nvfp4_rht_ue5m3") if fp8_block_scaling_available: _quantization_list.append("fp8_block_scaling") @@ -136,6 +138,11 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") + if ( + quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) + ): + pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") # Check dtype if dtype is not None: @@ -3588,7 +3595,12 @@ def test_grouped_mlp( # Skip invalid configurations with_quantization = quantization is not None - maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + maybe_skip_quantization( + quantization, + dims=in_shape, + device=device, + dtype=dtype, + ) if with_quantization and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") if activation == "scaled_srelu" and quantization == "nvfp4_rht" and bias: diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d09c92ad49..e562dab8d0 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -43,6 +43,7 @@ dtype_tols, make_recipe, MegatronTrainingHelper, + nvfp4_variant_names, quantization_tols, reset_rng_states, ) @@ -51,6 +52,7 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_ue5m3_available, reason_for_no_fp8_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] @@ -73,6 +75,8 @@ _grouped_mlp_quantization_list.append("mxfp8") if nvfp4_available: _grouped_mlp_quantization_list.append("nvfp4_rht") + if fp8_ue5m3_available: + _grouped_mlp_quantization_list.append("nvfp4_rht_ue5m3") @pytest.fixture(autouse=True, scope="function") @@ -102,11 +106,10 @@ def maybe_skip_quantization( pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and not nvfp4_available - ): + if quantization in nvfp4_variant_names and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and not fp8_ue5m3_available: + pytest.skip(reason_for_no_fp8_ue5m3) # Check dims if dims is not None: @@ -118,16 +121,18 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): + elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") + if ( + quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) + ): + pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") # Check dtype if dtype is not None: - if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") - and dtype != torch.bfloat16 - ): + if quantization in nvfp4_variant_names and dtype != torch.bfloat16: pytest.skip("NVFP4 quantization is only supported with BF16 data") @@ -183,17 +188,27 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): + elif quantization in ( + "nvfp4", + "nvfp4_row_scaled", + "nvfp4_rht", + "nvfp4_ue5m3", + "nvfp4_rht_ue5m3", + ): tensor_type = "input" if quantizer_role is not None: tensor_type = quantizer_role.tensor_type - with_rht = quantization == "nvfp4_rht" and tensor_type != "weight" + with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" + scale_dtype = ( + te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3 + ) test = NVFP4Quantizer( + scale_dtype=scale_dtype, with_rht=with_rht, with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=False, + with_random_sign_mask=with_rht, )(test) elif quantization == "nvfp4_4over6": tensor_type = "input" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f3d97b7269..484f1c4503 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -4,13 +4,15 @@ """Python interface for GEMM extensions""" -from typing import Iterable, Literal, Optional, Tuple, Union, List +from typing import Callable, Iterable, Literal, Optional, Tuple, Union, List +import itertools +import math import os import functools import torch import transformer_engine_torch as tex -from ..constants import TE_DType, DType -from ..utils import get_sm_count, _empty_tensor +from ..constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType, DType +from ..utils import ceil_div, get_cached_ones_tensor, get_sm_count, _empty_tensor from ..quantized_tensor import QuantizedTensorStorage, Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer @@ -188,6 +190,496 @@ def _validate_native_gemm_output_quantizer(quantization_params): ) +def validate_or_alloc_output( + buffer: Optional[torch.Tensor], + shape: tuple[int, ...] | list[int], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return the caller's output buffer, or allocate one if it is None. + + The buffer must be a contiguous tensor matching the required + shape, dtype, and device. + + """ + shape = tuple(shape) + if buffer is None: + return torch.empty(shape, dtype=dtype, device=device) + if tuple(buffer.shape) != shape: + raise ValueError(f"Output buffer shape {tuple(buffer.shape)} does not match {shape}.") + if buffer.dtype != dtype: + raise ValueError(f"Output buffer dtype {buffer.dtype} does not match {dtype}.") + if buffer.device != device: + raise ValueError(f"Output buffer device {buffer.device} does not match {device}.") + if not buffer.is_contiguous(): + raise ValueError("Output buffer must be contiguous.") + return buffer + + +@functools.lru_cache(maxsize=None) +def grouped_gemm_wgrad_kernel() -> Callable: + """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_wgrad_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_wgrad_wrapper_sm100 + + +def _cuDNN_wgrad_gemm( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + amax_a: Optional[torch.Tensor], + amax_b: Optional[torch.Tensor], + out_dtype: torch.dtype, + out: torch.Tensor, + accumulate: bool, + alpha: Optional[float] = None, + beta: Optional[float] = None, + bias: Optional[torch.Tensor] = None, +) -> Iterable[Optional[torch.Tensor]]: + """Compute dw = dy^T @ x with cuDNN's purpose-built grouped wgrad kernel.""" + + # Column-wise NVFP4 buffers are physically (features, tokens), FP4-packed + # two values per byte along the token dim. + tokens_packed = a_tensor.shape[-1] + tokens = tokens_packed * 2 + out_features, in_features = out.size() + + # grouped_gemm_wgrad_wrapper_sm100 wants: + # a_tensor (feature_out, tokens) K-major, FP4-packed + # b_tensor (tokens, feature_in) + # sfa (round_up(feature_out, 128), scale_cols) + # sfb (round_up(feature_in, 128), scale_cols) + fp4 = torch.float4_e2m1fn_x2 + a_tensor = a_tensor.view(dtype=fp4).view(out_features, tokens_packed) + b_tensor = b_tensor.view(dtype=fp4).view(in_features, tokens_packed).T + + # Create the scale factor tensors with the logical layout cuDNN expects + # In general_cuDNN_MX_gemm we've already ensured they are swizzled physically + def _sf(scale_inv, features): + leading = ceil_div(features, 128) * 128 + return scale_inv.view(leading, -1).view(dtype=torch.float8_e4m3fn) + + # grouped_gemm_wgrad_wrapper_sm100 expects two separate global_scale + ones = get_cached_ones_tensor(1, dtype=torch.float32, device=a_tensor.device) + denom = 6.0 * 114688.0 # fp4_max * fp8_max(UE5M3) + global_scale_a = ones if amax_a is None else amax_a.to(torch.float32).reshape(1) / denom + global_scale_b = ones if amax_b is None else amax_b.to(torch.float32).reshape(1) / denom + # Fold alpha into one of them if it's given + if alpha is not None and alpha != 1.0: + global_scale_a = global_scale_a * alpha + + out = validate_or_alloc_output(out, (out_features, in_features), out_dtype, a_tensor.device) + grouped_gemm_wgrad_kernel()( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=_sf(sfa, out_features), + sfb_tensor=_sf(sfb, in_features), + offsets_tensor=torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device), + global_scale_a=global_scale_a, + global_scale_b=global_scale_b, + acc_dtype=torch.float32, + wgrad_dtype=out.dtype, + output_mode="dense", + wgrad_tensor=out.view(1, out_features, in_features), + sf_vec_size=NVFP4_BLOCK_SCALING_SIZE, + sf_fp8_dtype_override="e5m3", + input_order="tensor_ragged", + accumulate_on_output=accumulate, + current_stream=torch.cuda.current_stream().cuda_stream, + ) + + # Apply bias + if bias is not None: + out += bias.view(1, in_features) + + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). + return out, None, None, None + + +@functools.lru_cache(maxsize=None) +def grouped_gemm_quant_kernel() -> Callable: + """cuDNN CuTe DSL grouped GEMM kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_quant_wrapper_sm100 + + +def convert_TE_MX_tensor_to_cuDNN_operand( + data: torch.Tensor, + scale_inv: torch.Tensor, + *, + data_dtype: torch.dtype, + scale_dtype: torch.dtype, + valid_M_or_N: int, + k_logical: int, + L: int = 1, + sf_swizzled: bool = False, + use_N_major_for_B: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reshape an plain buffer into the layout cuDNN's grouped GEMM expects. + + cuDNN requirements: + A: (valid_m, K, 1), K-major + B: (N, K, L), K-major (FP8 also supports N-major) + + SFA: (32, 4, ceil(valid_m/128), 4, ceil(ceil(K/sf_vec_size)/4), 1) + SFB: (32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L) + + whereas TE stores flat buffers which can be intepreted as contiguous tensors + with the following layouts: + + Note: K_packed is K/2 for FP4 (two values per byte) and K for FP8 + + A (K-major): (1, valid_m, K_packed) + B (K-major): (L, N, K_packed) -- used for FP4 only now + B (N-major): (L, K, N) -- used for FP8 only now + + SFA (unswizzled): (1, ceil(valid_m/128), 4, 32, ceil(ceil(K/sf_vec_size)/4), 4) + SFB (unswizzled): (L, ceil(N/128), 4, 32, ceil(ceil(K/sf_vec_size)/4), 4) + SFA (swizzled): (1, ceil(valid_m/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) + SFB (swizzled): (L, ceil(N/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) + """ + + if use_N_major_for_B: + assert data_dtype in (torch.float8_e4m3fn, torch.float8_e5m2), \ + f"Using N-major layout for B is only supported for FP8, but got {data_dtype}." + + available_scalings = { + # NVFP4 recipe (UE5M3 rides as E4M3 since torch has no ue5m3 dtype) + (torch.float4_e2m1fn_x2, torch.float8_e4m3fn): NVFP4_BLOCK_SCALING_SIZE, + # MXFP8 recipe + (torch.float8_e4m3fn, torch.float8_e8m0fnu): MXFP8_BLOCK_SCALING_SIZE, + } + assert (data_dtype, scale_dtype) in available_scalings, ( + "Unsupported (data_dtype, scale_dtype) pair for a cuDNN block-scaled operand: " + f"({data_dtype}, {scale_dtype}). Expected NVFP4 (float4_e2m1fn_x2, " + "float8_e4m3fn) or MXFP8 (float8_e4m3fn, float8_e8m0fnu)." + ) + sf_vec_size = available_scalings[(data_dtype, scale_dtype)] + + k_sf_tiles = ceil_div(k_logical, 4 * sf_vec_size) + + if data_dtype == torch.float4_e2m1fn_x2: + k_packed = k_logical // 2 # fp4 packs two values per byte + else: + k_packed = k_logical # fp8 packs one value per byte + + data = data.view(dtype=data_dtype) + if use_N_major_for_B: + # B is stored untransposed, i.e. (L, K, N); permuting to (N, K, L) leaves + # stride 1 on N. Only FP8 accepts this, asserted above. + data = data.view(L, k_packed, valid_M_or_N) + data = data.permute(2, 1, 0) + else: + # (L, N, K) -> (N, K, L), stride 1 on K. + data = data.view(L, valid_M_or_N, k_packed) + data = data.permute(1, 2, 0) + + if sf_swizzled: + scale_inv = scale_inv.view(dtype=scale_dtype) + scale_inv = scale_inv.view( + L, + ceil_div(valid_M_or_N, 128), + k_sf_tiles, + 32, + 4, + 4, + ) + scale_inv = scale_inv.permute(3, 4, 1, 5, 2, 0) + return data, scale_inv + + scale_inv = scale_inv.view(dtype=scale_dtype) + scale_inv = scale_inv.view( + L, + ceil_div(valid_M_or_N, 128), + 4, + 32, + k_sf_tiles, + 4, + ) + scale_inv = scale_inv.permute(3, 2, 1, 5, 4, 0) + return data, scale_inv + + +def general_cuDNN_MX_gemm( + A: torch.Tensor, + B: torch.Tensor, + out_dtype: Optional[torch.dtype] = None, + quantization_params: Optional[Quantizer] = None, + gelu: bool = False, + gelu_in: torch.Tensor = None, + alpha: float = 1.0, + beta: Optional[float] = None, + accumulate: bool = False, + layout: str = "TN", + out: Optional[torch.Tensor] = None, + bias: Optional[torch.Tensor] = None, + use_split_accumulator: bool = False, + grad: bool = False, + ub: Union[tex.CommOverlap, tex.CommOverlapP2P] = None, + ub_type: tex.CommOverlapType = None, + extra_output: Optional[torch.Tensor] = None, + bulk_overlap: bool = False, +) -> Iterable[Optional[torch.Tensor]]: + """Perform GEMM via cuDNN kernels + + The parameters passed are in cuBLAS notation, where + D = alpha * op(B) @ op(A) + beta * C, where the shape is always + (N, M) = (N, K) @ (K, M) + (N, M) + + B: + - "N" is (N, K), which is always TE's rowwise data, and op(B) is B + - "T" is (K, N), which is always TE's colwise data, and op(B) is B.T + A + - "N" is (K, M), which is always TE's colwise data, and op(A) is A + - "T" is (M, K), which is always TE's rowwise data, and op(A) is A.T + + Note: layout string means layout of "A" and "B" respectively. + + TE stores x (token, feature_in), w (feature_out, feature_in) and dy (token, feature_out) in physical rowwise direction. + For cuBLAS: + fprop = x @ wT: token is N, feature_in is K, feature_out is M, so it's TN (x as B, w transposed to wT as A) + dgrad = dy @ w: token is N, feature_out is K, feature_in is M, so it's NN (dy as B, w as A) + wgrad = dyT @ x: feature_out is N, token is K, feature_in is M, so it's NT (dy transposed to dyT as B, x as A) + + We use cuDNN-frontend's APIs here which are supposed to be used for grouped GEMM but we set groups = 1 + so it is effectively a single GEMM. + + Naming convention: uppercase letters (A, B) are used for cuBLAS notation, lowercase letters (a, b) are used for cuDNN notation. + where cuBLAS's B is cuDNN's a, and cuBLAS's A is cuDNN's b (their notation is inverted). + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + + """ + assert isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) and \ + A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3, \ + f"cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." + + assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." + assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." + + # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for + # raising accumulator precision, and the cuDNN kernel always accumulates in + # FP32, so the request is already satisfied either way. + assert ub is None and ub_type is None, "cuDNN GEMM currently does not support CommOverlap." + assert extra_output is None, "cuDNN GEMM currently does not support extra output." + assert bulk_overlap is False, "cuDNN GEMM currently does not support bulk overlap." + + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." + transa = layout[0] == "T" + transb = layout[1] == "T" + + assert out_dtype in (torch.float32, torch.float16, torch.bfloat16), \ + f"cuDNN MX GEMM currently only supports float32, float16, and bfloat16 outputs, but got {out_dtype}." + + device = A.device + + # cuDNN only accepts GEMM-swizzled scale factors -- an unswizzled buffer is + # rejected on its strides -- so swizzle first if the quantizer did not + # (optimize_for_gemm defaults to False). This mirrors what the cuBLAS path + # does in C++ via swizzle_scales_for_gemm. The call is in-place, swizzles + # both orientations, and no-ops when the tensor is already swizzled. + if not A._with_gemm_swizzled_scales: + tex.swizzle_scales_for_gemm_(A) + if not B._with_gemm_swizzled_scales: + tex.swizzle_scales_for_gemm_(B) + + # `grad` only changes behaviour when a bias is supplied: it turns the bias slot + # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs + # that pass grad=True without a bias need nothing special. + assert not (grad and bias is not None), ( + "cuDNN GEMM currently does not support fused bias gradient." + ) + + # Pick the buffer whose block scales run along K. In every case the selected + # buffer is physically (rows, K_packed), so the reshape below is uniform. + # LHS is always (M, K) + if transb: + dataB, sfB, amaxB = B._columnwise_data, B._columnwise_scale_inv, B._amax_columnwise + else: + dataB, sfB, amaxB = B._rowwise_data, B._rowwise_scale_inv, B._amax_rowwise + # RHS is always (K, N) + if transa: + dataA, sfA, amaxA = A._rowwise_data, A._rowwise_scale_inv, A._amax_rowwise + else: + dataA, sfA, amaxA = A._columnwise_data, A._columnwise_scale_inv, A._amax_columnwise + + # Input tensor dims + A_shape = list(dataA.size()) + A_shape[-1] *= 2 + B_shape = list(dataB.size()) + B_shape[-1] *= 2 + + # GEMM dimensions + M_full = A_shape[:-1] if transa else [A_shape[0]] + N_full = [B_shape[0]] if transb else B_shape[:-1] + K_full = [A_shape[-1]] if transa else A_shape[1:] + K_full_b = B_shape[1:] if transb else [B_shape[-1]] + assert K_full == K_full_b, f"Contraction dims disagree: A implies {K_full}, B implies {K_full_b}." + M = math.prod(M_full) + N = math.prod(N_full) + K = math.prod(K_full) + + # Allocate output tensor if needed + out_shape = N_full + M_full + out = validate_or_alloc_output(out, out_shape, out_dtype, device) + + # Trivial cases + if K == 0: + if bias is not None: + out_2d = out.view(N, M) + bias_2d = bias.view(1, M) + if accumulate: + out_2d += bias_2d + else: + out_2d.copy_(bias_2d) + elif not accumulate: + out.zero_() + return out, None, None, None + if M == 0 or N == 0: + return out, None, None, None + + # Route to cuDNN-FE's wgrad API for cases not supported by the + # grouped GEMM (accumulation to output tensor, insufficient + # alignment). The wgrad kernel has no bias epilogue, so any bias + # has to be applied after the GEMM. + if accumulate or N % 256 != 0: + alpha = alpha if alpha is not None else 1.0 + # This path uses cuDNN's wgrad (grouped_gemm_wgrad_wrapper_sm100) which supports grad accumulation + if accumulate: # Accumulate GEMM's result to the out tensor + assert beta in (1.0, None), "beta must be one or None if accumulate is True" + else: # Overwrite GEMM's result to the out tensor + assert beta in (0.0, None), "beta must be zero or None if not accumulate" + _cuDNN_wgrad_gemm( + a_tensor=dataB.view(N, K // 2), + b_tensor=dataA.view(M, K // 2), + sfa=sfB, + sfb=sfA, + amax_a=amaxB, + amax_b=amaxA, + out_dtype=out_dtype, + out=out.view(N, M), + accumulate=accumulate, + alpha=alpha, + bias=bias, + ) + return out, None, None, None + + alpha = alpha if alpha is not None else 1.0 + # cuDNN's general GEMM path (grouped_gemm_quant_wrapper_sm100) doesn't support accumulation + assert accumulate is False, "cuDNN GEMM currently does not support accumulation for this operation." + assert beta in (0.0, None), "beta must be zero or None if not accumulate" + + # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it + N_padded = ceil_div(N, 256) * 256 + if N_padded != N: + src = dataB.reshape(N, K // 2) + buf = src.new_zeros((N_padded, K // 2)) + buf[:N].copy_(src) + dataB = buf + + # Swizzled scales are blocked by 128 rows: + # (1, ceil(M/128), k_sf_tiles, 32, 4, 4) + per_block = ceil_div(K, 4 * NVFP4_BLOCK_SCALING_SIZE) * 32 * 4 * 4 + n_blk, n_blk_padded = ceil_div(N, 128), ceil_div(N_padded, 128) + src_sf = sfB.reshape(-1)[: n_blk * per_block].reshape(n_blk, per_block) + buf_sf = src_sf.new_zeros((n_blk_padded, per_block)) + buf_sf[:n_blk].copy_(src_sf) + sfB = buf_sf + + # cuDNN's own operand names are the other way round: its "a" is the (M, K) + # activation-like operand (TE's B) and its "b" is the (N, K) weight-like one + # (TE's A). + cudnn_a, cudnn_sfa = convert_TE_MX_tensor_to_cuDNN_operand( + dataB, + sfB, + data_dtype=torch.float4_e2m1fn_x2, + scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 + valid_M_or_N=N_padded, + k_logical=K, + L=1, + sf_swizzled=True, # ensured above + ) + cudnn_b, cudnn_sfb = convert_TE_MX_tensor_to_cuDNN_operand( + dataA, + sfA, + data_dtype=torch.float4_e2m1fn_x2, + scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 + valid_M_or_N=M, + k_logical=K, + L=1, + sf_swizzled=True, # ensured above + ) + + # Row-scaled NVFP4 stores one amax per row instead of one per tensor, which + # this path cannot express; general_gemm handles that mode separately. + for name, amax in (("A", amaxA), ("B", amaxB)): + assert amax is None or amax.numel() == 1, ( + f"cuDNN MX GEMM expects a per-tensor amax for {name}, but got {amax.numel()} " + "values. Row-scaled NVFP4 is not supported on this path." + ) + + # Prepare alpha. cuDNN applies the block scales but not TE's per-tensor global + # scale, so alpha carries the product of both operands'. A tensor quantized + # without second-level scaling has no amax and contributes a factor of one. + nvfp4_global_scale = 6.0 * 114688.0 + ones = get_cached_ones_tensor(1, dtype=torch.float32, device=device) + scaleA = ones if amaxA is None else amaxA.to(torch.float32).reshape(1) / nvfp4_global_scale + scaleB = ones if amaxB is None else amaxB.to(torch.float32).reshape(1) / nvfp4_global_scale + alpha_tensor = (alpha * scaleA * scaleB).to(torch.float32) + + if bias is not None: + assert bias.dim() == 1 and bias.shape[0] == M, ( + f"cuDNN MX GEMM expects a ({M},) bias, but got {tuple(bias.shape)}." + ) + # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). + bias = bias.contiguous().as_strided((M, 1), (1, M)) + + # Prepare for output + out = validate_or_alloc_output(out, out_shape, out_dtype, device) + if N_padded != N: + # The kernel writes N_padded rows, so it cannot target `out` directly. + d_buf = torch.empty((N_padded, M), dtype=out_dtype, device=device) + d_tensor = d_buf.as_strided((N_padded, M, 1), (M, 1, N_padded * M)) + else: + d_tensor = out.view(N, M).as_strided((N, M, 1), (M, 1, M * N)) + + gemm_kwargs = { + "a_tensor": cudnn_a, + "sfa_tensor": cudnn_sfa, + "b_tensor": cudnn_b, + "sfb_tensor": cudnn_sfb, + # One group, so the only padded end offset is the full row count. + "padded_offsets": torch.tensor([N_padded], dtype=torch.int32, device=device), + "alpha_tensor": alpha_tensor, + "bias_tensor": bias, + "norm_const_tensor": None, # must be None for FP4 inputs + "acc_dtype": torch.float32, + "d_dtype": out_dtype, # high precision -> no output quantization + "d_tensor": d_tensor, + "cd_major": "n", # only "n" is supported by cuDNN + "sf_vec_size": NVFP4_BLOCK_SCALING_SIZE, # Hardcode to NVFP4 for now + "sf_fp8_dtype_override": "e5m3", # Hardcode for now + "current_stream": torch.cuda.current_stream().cuda_stream, + "discrete_col_sfd": False, + "use_dynamic_sched": True, + } + grouped_gemm_quant_kernel()(**gemm_kwargs) + + if N_padded != N: + # Drop the zero-padded rows. Safe to overwrite rather than accumulate: + # this path asserts accumulate is False above. + out.view(N, M).copy_(d_buf[:N]) + + # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). + return out, None, None, None + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -226,6 +718,36 @@ def general_gemm( beta = validate_gemm_scale(beta, accumulate) workspace = get_cublas_workspace(A.device.index, ub is not None, False) + # Temporary hack to route NVFP4 GEMM with UE5M3 scale factors to + # cuDNN Frontend kernels. UE5M3-specific logic should be removed + # in its entirety once TE supports NVFP4-UE5M3 GEMMs natively. + if ( + isinstance(A, NVFP4TensorStorage) + and isinstance(B, NVFP4TensorStorage) + and A._scale_dtype == DType.kFloat8UE5M3 + and B._scale_dtype == DType.kFloat8UE5M3 + ): + return general_cuDNN_MX_gemm( + A, + B, + out_dtype, + quantization_params, + gelu, + gelu_in, + alpha, + beta, + accumulate, + layout, + out, + bias, + use_split_accumulator, + grad, + ub, + ub_type, + extra_output, + bulk_overlap, + ) + if ub_type is not None: assert ub is not None, ( f"{'AG+GEMM' if ub_type == tex.CommOverlapType.AG else 'GEMM+RS'} overlap requires" @@ -428,71 +950,49 @@ def general_grouped_gemm( if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.") - if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): - assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." - if single_output: - assert ( - m_splits is not None - ), "Row-scaled NVFP4 grouped GEMM requires m_splits with single output." - out_init = out[0] if single_output else None - if single_output: - start_idx = 0 - out_views = [] - for i in range(num_gemms): - size = m_splits[i] - out_views.append(out_init[start_idx : start_idx + size]) - start_idx += size - else: - out_views = out - for i in range(num_gemms): - if out_views[i].numel() == 0: - continue - general_gemm( - A[i], - B[i], - quantization_params=quantization_params[i], - out_dtype=out_views[i].dtype, - out=out_views[i], - gelu=gelu, - accumulate=accumulate, - layout=layout, - bias=bias[i] if use_bias else None, - use_split_accumulator=use_split_accumulator, - grad=grad, - ) - if single_output: - out = out_init - return out, grad_bias, gelu_input + # Determine whether to repeatedly call general_gemm + use_general_gemm_impl = False if isinstance(quantization_params[0], DebugQuantizer): - assert not gelu, "GELU not supported in debug mode" + use_general_gemm_impl = True + elif any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + use_general_gemm_impl = True + elif any( + isinstance(t, NVFP4TensorStorage) and t._scale_dtype == DType.kFloat8UE5M3 + for t in itertools.chain(A, B) + ): + use_general_gemm_impl = True + + # Repeatedly call general_gemm if needed + if use_general_gemm_impl: + out_views = out if single_output: - out_init = out[0] start_idx = 0 - out = [None] * num_gemms + out = out[0] + out_views = [None] * num_gemms for i in range(num_gemms): size = m_splits[i] - out[i] = out_init[start_idx : start_idx + size] + out_views[i] = out[start_idx : start_idx + size] start_idx += size for i in range(num_gemms): - _, bias_or_grad, _, _ = general_gemm( + _, bias_or_grad, gelu_input_i, _ = general_gemm( A[i], B[i], quantization_params=quantization_params[i], - out_dtype=out[0].dtype, + out_dtype=out_views[i].dtype, layout=layout, accumulate=accumulate, - out=out[i], + out=out_views[i], + gelu=gelu, bias=bias[i] if use_bias else None, use_split_accumulator=use_split_accumulator, grad=grad, ) if grad and use_bias: grad_bias[i] = bias_or_grad - if single_output: - out = out_init - - return out, grad_bias if grad else bias, None + if gelu: + gelu_input[i] = gelu_input_i + return out, grad_bias if grad else bias, gelu_input if gelu: gelu_input = [ diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1a9fabb3bb..4f74a81b1a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -217,19 +217,23 @@ void nvfp4_multi_tensor_compute_partial_amax( void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, int64_t tile_cols, int64_t rows_padded, int64_t block_len); -void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax); +void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, + const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len); + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_multi_tensor_fused_scale( std::vector block_amax_list, std::vector global_amax_list, std::vector per_block_scale_list, std::vector target_scale_list, std::vector target_amax_list, std::vector tile_rows_list, - std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len); + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); -void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale); +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, + const DType scale_dtype = DType::kFloat8E4M3); at::Tensor swap_first_dims(at::Tensor tensor, std::optional out = std::nullopt); @@ -489,14 +493,15 @@ void nvfp4_2d_compute_partial_amax(const at::Tensor &tensor, at::Tensor amax, si void nvfp4_2d_partial_cast(const at::Tensor &inp, py::handle out, const at::Tensor &scale, const at::Tensor &global_scale, size_t h, size_t w, size_t start_offset, - size_t block_len); + size_t block_len, const DType scale_dtype = DType::kFloat8E4M3); void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, std::vector out_list, std::vector scale_list, std::vector global_scale_list, std::vector h_list, std::vector w_list, - std::vector start_offset_list, int64_t block_len); + std::vector start_offset_list, int64_t block_len, + const DType scale_dtype = DType::kFloat8E4M3); void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise, at::Tensor amax_colwise, int rows, int cols, size_t start_offset); diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp index 685250d137..8b58299351 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -27,7 +27,7 @@ void nvfp4_2d_compute_partial_amax(const at::Tensor& tensor, at::Tensor amax, si void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tensor& scale, const at::Tensor& global_scale, size_t h, size_t w, size_t start_offset, - size_t block_len) { + size_t block_len, const DType scale_dtype) { TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); TORCH_CHECK(scale.dim() == 2, "scale must be a 2D tensor"); TORCH_CHECK(scale.scalar_type() == at::ScalarType::Float, "scale must be a float tensor"); @@ -45,7 +45,8 @@ void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tens nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), start_offset, block_len, - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, @@ -53,7 +54,8 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, std::vector scale_list, std::vector global_scale_list, std::vector h_list, std::vector w_list, - std::vector start_offset_list, int64_t block_len) { + std::vector start_offset_list, int64_t block_len, + const DType scale_dtype) { TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); const size_t num_tensors = inp_list.size(); @@ -95,7 +97,8 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), - start_offset, static_cast(block_len), stream); + start_offset, static_cast(block_len), stream, + static_cast(scale_dtype)); } } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2173dff8b2..1cd9c11e8b 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -412,15 +412,21 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_compute_per_block_scale", &transformer_engine::pytorch::nvfp4_compute_per_block_scale, "Compute per-block decode scale from block amax and global amax", py::arg("block_amax"), - py::arg("scale"), py::arg("global_amax"), py::call_guard()); + py::arg("scale"), py::arg("global_amax"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_compute_global_scale", &transformer_engine::pytorch::nvfp4_compute_global_scale, "Compute global encode scale from global amax", py::arg("global_amax"), - py::arg("global_scale"), py::call_guard()); + py::arg("global_scale"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_fused_scale", &transformer_engine::pytorch::nvfp4_fused_scale, "Fused kernel: compute per-block decode scale, copy global amax, expand to row-level FP8", py::arg("block_amax"), py::arg("global_amax"), py::arg("per_block_scale"), py::arg("target_scale"), py::arg("target_amax"), py::arg("tile_rows"), py::arg("tile_cols"), - py::arg("rows_padded"), py::arg("block_len"), py::call_guard()); + py::arg("rows_padded"), py::arg("block_len"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::call_guard()); m.def("nvfp4_multi_tensor_fused_scale", &transformer_engine::pytorch::nvfp4_multi_tensor_fused_scale, "Batched fused scale: compute per-block decode scale, copy global amax, expand to FP8 for " @@ -428,6 +434,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("block_amax_list"), py::arg("global_amax_list"), py::arg("per_block_scale_list"), py::arg("target_scale_list"), py::arg("target_amax_list"), py::arg("tile_rows_list"), py::arg("tile_cols_list"), py::arg("rows_padded_list"), py::arg("block_len"), + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("nvfp4_2d_multi_tensor_transpose", &transformer_engine::pytorch::nvfp4_2d_multi_tensor_transpose, @@ -474,12 +481,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Partial cast from master weights for NVFP4 2D", py::arg("inp"), py::arg("out"), py::arg("scale"), py::arg("global_scale"), py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len") = 16, + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("nvfp4_multi_tensor_2d_partial_cast", &transformer_engine::pytorch::nvfp4_multi_tensor_2d_partial_cast, "Batched partial cast from master weights for NVFP4 2D", py::arg("inp_list"), py::arg("out_list"), py::arg("scale_list"), py::arg("global_scale_list"), py::arg("h_list"), py::arg("w_list"), py::arg("start_offset_list"), py::arg("block_len") = 16, + py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("mxfp8_scaling_compute_partial_amax", &transformer_engine::pytorch::mxfp8_scaling_compute_partial_amax, diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 0318978195..4b887c3749 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -145,7 +145,7 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile } void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, - at::Tensor global_amax) { + at::Tensor global_amax, const DType scale_dtype) { init_extension(); // block_amax and scale: [tile_rows, tile_cols], float32 @@ -160,12 +160,14 @@ void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, auto global_amax_cu = makeTransformerEngineTensor(global_amax); nvte_nvfp4_compute_per_block_scale(block_amax_cu.data(), scale_cu.data(), global_amax_cu.data(), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len) { + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + const DType scale_dtype) { init_extension(); // block_amax: [tile_rows, tile_cols], float32 @@ -191,14 +193,16 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor target_scale_cu.data(), target_amax_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), static_cast(block_len), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } void nvfp4_multi_tensor_fused_scale( std::vector block_amax_list, std::vector global_amax_list, std::vector per_block_scale_list, std::vector target_scale_list, std::vector target_amax_list, std::vector tile_rows_list, - std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len) { + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len, + const DType scale_dtype) { init_extension(); const size_t num_tensors = block_amax_list.size(); @@ -242,11 +246,13 @@ void nvfp4_multi_tensor_fused_scale( nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, - rows_padded, static_cast(block_len), stream); + rows_padded, static_cast(block_len), stream, + static_cast(scale_dtype)); } } -void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale) { +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, + const DType scale_dtype) { init_extension(); // global_amax and global_scale: [num_params], float32 @@ -257,7 +263,8 @@ void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale) auto global_scale_cu = makeTransformerEngineTensor(global_scale); nvte_nvfp4_compute_global_scale(global_amax_cu.data(), global_scale_cu.data(), - at::cuda::getCurrentCUDAStream()); + at::cuda::getCurrentCUDAStream(), + static_cast(scale_dtype)); } at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9860d48237..fe871e0182 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -59,7 +59,7 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) -from ..constants import GemmParallelModes, dist_group_type +from ..constants import DType, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias @@ -476,6 +476,7 @@ def _is_grouped_tensor_path_supported( activation_dtype: torch.dtype, input_quantizers: List[Optional[Quantizer]], output_quantizers: List[Optional[Quantizer]], + single_grouped_weight: bool, ) -> bool: """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata. @@ -501,10 +502,11 @@ def _is_grouped_tensor_path_supported( Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - # 1. Filter by environment variable + # Filter by environment variable if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): return False - # 2. Filter out advanced features + + # Filter out advanced features if ( debug or cpu_offloading @@ -513,47 +515,59 @@ def _is_grouped_tensor_path_supported( or save_original_input ): return False - # 3. Filter by compute capability and cuBLAS version - device_capability = get_device_compute_capability() - if not (9, 0) <= device_capability <= (11, 0): - return False - cublaslt_version = tex.get_cublasLt_version() - if cublaslt_version < 130300: - return False - if device_capability < (10, 0) and cublaslt_version < 130400: - return False - # 4. Output quantization is not supported. + + # Output quantization is not supported. if any(q is not None for q in output_quantizers): return False - # 5. Filter by quantization recipes. - if fp8: - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+. - if device_capability < (10, 0) and cublaslt_version < 130500: - return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM - # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast - # emulation. On Blackwell (SM100/SM110, the only other arch that reaches - # this branch) fail loudly rather than silently falling back to the - # unfused path the user explicitly opted out of. - if get_device_compute_capability() >= (10, 0): - raise RuntimeError( - "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the" - " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped" - " FP8 block-scaling path is Hopper-only. Unset" - " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused" - " path (emulated via MXFP8 GEMM on Blackwell)." - ) - return True - # MXFP8 and NVFP4 require Blackwell+. - if not (10, 0) <= device_capability <= (11, 0): + + device_arch = get_device_compute_capability() + + # Unquantized compute + if not fp8: + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ return False - return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all( - isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers - ) - return activation_dtype in (torch.bfloat16, torch.float16) + return activation_dtype in (torch.bfloat16, torch.float16) + + # FP8 current scaling + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + if device_arch[0] == 9 and tex.get_cublasLt_version() < 130500: + # Hopper support for grouped GEMM requires cuBLAS 13.5+ + return False + return True + + # FP8 block scaling + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped GEMM requires Hopper and cuBLAS 13.4+ + return device_arch[0] == 9 and tex.get_cublasLt_version() >= 130400 + + # MXFP8 + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + # MXFP8 grouped quantization requires Blackwell + return (10, 0) <= device_arch <= (11, 0) + + # NVFP4 + if all(isinstance(q, NVFP4Quantizer) for q in input_quantizers): + if not (10, 0) <= device_arch <= (11, 0): + # NVFP4 grouped quantization requires Blackwell + return False + if single_grouped_weight: + # NVFP4 graph-safe grouped quantization only supports discrete weights + return False + for q in input_quantizers: + if not q.with_rht: + # NVFP4 graph-safe grouped quantization requires RHT + return False + if q.scale_dtype != DType.kFloat8E4M3: + # NVFP4 grouped GEMM is only supported with E4M3 scales + return False + return True + + # Fall back to non-graph-safe implementation + return False @staticmethod def _make_grouped_tensor( @@ -903,6 +917,7 @@ def forward( delayed_scaling_input_quantizer, unsafe_requantization_input_quantizer, debug, + single_grouped_weight, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -1003,6 +1018,7 @@ def forward( activation_dtype=activation_dtype, input_quantizers=input_quantizers, output_quantizers=output_quantizers, + single_grouped_weight=single_grouped_weight, ): return _GroupedLinear._forward_grouped_tensor( ctx, @@ -2410,6 +2426,7 @@ def forward( self._delayed_scaling_input_quantizer, self._unsafe_requantization_input_quantizer, debug, + self.single_grouped_weight, ) out, new_workspaces = linear_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e1980d2943..cbff2abe81 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -822,36 +822,55 @@ def _is_graph_safe_path_supported( * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would trigger a fatal error in the cuBLASLt grouped GEMM check. """ - if not (9, 0) <= get_device_compute_capability() <= (11, 0): - return False - if with_quantized_compute: - # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM - # path; the compute-capability range was already checked above. On Hopper it - # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS. - if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): - if ( - get_device_compute_capability() < (10, 0) - and tex.get_cublasLt_version() < 130500 - ): + + device_arch = get_device_compute_capability() + + # Unquantized compute + if not with_quantized_compute: + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + return dtype in (torch.bfloat16, torch.float16) + + # FP8 current scaling + if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers): + if not (9, 0) <= device_arch <= (11, 0): + # cuBLAS supports grouped GEMM on Hopper+ + return False + if device_arch[0] == 9 and tex.get_cublasLt_version() < 130500: + # Hopper support for grouped GEMM requires cuBLAS 13.5+ + return False + return True + + # FP8 block scaling + if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): + # Grouped GEMM requires Hopper and cuBLAS 13.4+ + return device_arch[0] == 9 and tex.get_cublasLt_version() >= 130400 + + # MXFP8 + if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): + # MXFP8 grouped quantization requires Blackwell + return (10, 0) <= device_arch <= (11, 0) + + # NVFP4 + if all(isinstance(q, NVFP4Quantizer) for q in input_quantizers): + if not (10, 0) <= device_arch <= (11, 0): + # NVFP4 grouped quantization requires Blackwell + return False + if single_grouped_weight: + # NVFP4 graph-safe grouped quantization only supports discrete weights + return False + for q in input_quantizers: + if not q.with_rht: + # NVFP4 graph-safe grouped quantization requires RHT return False - return True - if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers): - # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere - # fall back to the split-quantize (MXFP8-emulated) flow. - if get_device_compute_capability() >= (10, 0): + if q.scale_dtype != DType.kFloat8E4M3: + # NVFP4 grouped GEMM is only supported with E4M3 scales return False - return tex.get_cublasLt_version() >= 130400 - # MXFP8 and NVFP4 grouped quantization kernels require Blackwell. - if not (10, 0) <= get_device_compute_capability() <= (11, 0): - return False - if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers): - return True - # NVFP4 graph-safe grouped quantization requires RHT and only supports - # discrete weights; otherwise fall back to the split-quantize flow. - if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers): - return not single_grouped_weight - return False - return dtype in (torch.bfloat16, torch.float16) + return True + + # Fall back to non-graph-safe implementation + return False def _get_grouped_weight_for_gemm( self, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a44bef0b2d..8b03a0587c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -10,7 +10,7 @@ import functools import os from importlib.metadata import PackageNotFoundError, version as get_pkg_version -from typing import Any, Optional +from typing import Any, Literal, Optional import torch from packaging.version import Version as PkgVersion @@ -20,6 +20,7 @@ from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor +from ...cpp_extensions.gemm import convert_TE_MX_tensor_to_cuDNN_operand from ...distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, @@ -251,6 +252,56 @@ def _nvfp4_amax( return torch.cat([amax.view(-1) for amax in amaxes], dim=0) +# TODO(kainingz): remove this temporary workaround after pytorch & tvm-ffi supports e5m3 GEMM +def _nvfp4_sf_dtype_override(quantizer: Optional[Quantizer]) -> Literal["e5m3"] | None: + """Returns a string to indicate the real scale factor dtype for cuDNN. + + Since pytorch doesn't have a native e5m3 dtype, we need let e5m3 pretend to be e4m3 and + use this string to indicate cuDNN to interpret the scale factors as e5m3 correctly when + it enters CuTeDSL region which has e5m3 support. + """ + if quantizer is None or not isinstance(quantizer, NVFP4Quantizer): + return None + if getattr(quantizer, "nvfp4_use_4over6", False): + # We don't use e5m3 for 4over6 + return None + scale_dtype = getattr(quantizer, "scale_dtype", None) + if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + return "e5m3" + # If we don't use e5m3 we don't need to pass this string to override + return None + + +def _nvfp4_scale_max(quantizer: Quantizer) -> float: + """Return the maximum representable magnitude of an NVFP4 scale factor.""" + # 4over6 might override e4m3's max to 256 over default 448 + override_max = getattr(quantizer, "nvfp4_e4m3_max", None) + # NVFP4Quantizer's initialization sets nvfp4_e4m3_max to -1 if no override + if override_max is not None and override_max != -1: + return float(override_max) + scale_dtype = getattr(quantizer, "scale_dtype", None) + if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + return 114688.0 + return 448.0 + + +def _nvfp4_global_scale( + tensors: GroupedTensor | Iterable[NVFP4TensorStorage], + quantizer: Quantizer, + *, + columnwise: bool, + num_groups: int, + device: torch.device, +) -> torch.Tensor: + """Return the per-group global scale factor for an NVFP4 operand.""" + if getattr(quantizer, "disable_second_level_scale", False): + # The second-level scale is disabled, so the global scale is always 1.0. + return get_cached_ones_tensor(num_groups, torch.float32, device) + # 6.0 is NVFP4_FP4_MAX + denom = 6.0 * _nvfp4_scale_max(quantizer) + return _nvfp4_amax(tensors, columnwise=columnwise).to(torch.float32) / denom + + def _single_quantized_tensor_from_grouped( grouped: GroupedTensor, quantizer: Optional[MXFP8Quantizer | NVFP4Quantizer] = None, @@ -303,6 +354,7 @@ def _single_quantized_tensor_from_grouped( with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, ) + # TODO(kainingz): claude told me this doesn't pass the required param scale_dtype. Should check this later return NVFP4Tensor( shape=shape, dtype=grouped.get_dtype(), @@ -475,6 +527,10 @@ def _cudnn_compute_wgrad( out_features, in_features = weight_shape total_tokens = grouped_dy.logical_shape[0] + device = grouped_dy.columnwise_data.device + + dy_quantizer=getattr(grouped_dy, "quantizer", None) + x_quantizer=getattr(grouped_x, "quantizer", None) sfa_leading_dim = round_up_to_nearest_multiple(out_features, 128) sfb_leading_dim = round_up_to_nearest_multiple(in_features, 128) @@ -483,7 +539,6 @@ def _cudnn_compute_wgrad( # A workaround for the case with zero-token experts. # Even for this case, cuteDSL still requires the same # stride requirements for the input and scale tensors. - device = grouped_dy.columnwise_data.device a_tensor = torch.empty_strided( (out_features, 0), (16, 1), @@ -554,9 +609,9 @@ def _cudnn_compute_wgrad( "current_stream": current_stream, } if use_nvfp4: - global_scale_denom = 448.0 * 6.0 + num_groups = offsets.shape[0] if total_tokens == 0: - global_scale_shape = (offsets.shape[0],) + global_scale_shape = (num_groups,) common_wgrad_kwargs["global_scale_a"] = torch.zeros( global_scale_shape, dtype=torch.float32, @@ -568,13 +623,24 @@ def _cudnn_compute_wgrad( device=device, ) else: - common_wgrad_kwargs["global_scale_a"] = ( - _nvfp4_amax(grouped_dy, columnwise=True).to(torch.float32) / global_scale_denom + common_wgrad_kwargs["global_scale_a"] = _nvfp4_global_scale( + grouped_dy, + dy_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, ) - common_wgrad_kwargs["global_scale_b"] = ( - _nvfp4_amax(grouped_x, columnwise=True).to(torch.float32) / global_scale_denom + common_wgrad_kwargs["global_scale_b"] = _nvfp4_global_scale( + grouped_x, + x_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, ) common_wgrad_kwargs["input_order"] = "tensor_ragged" + wgrad_sf_dtype_override = _nvfp4_sf_dtype_override(dy_quantizer) + if wgrad_sf_dtype_override is not None: + common_wgrad_kwargs["sf_fp8_dtype_override"] = wgrad_sf_dtype_override # Prepare wgrad output if single_grouped_weight: @@ -820,20 +886,20 @@ def fuse_grouped_mlp_ops( elif not (recipe.mxfp8() or recipe.nvfp4()): return ops + if activation_op_types is None: + activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) + # Check for unsupported NVFP4 recipe configs if recipe.nvfp4(): if recipe.disable_rht: # Graph-safe grouped quantize is only supported with RHT return ops - if ( - recipe.row_scaled_activation - or recipe.nvfp4_4over6 - or recipe.fp8_format == RecipeFormat.UE5M3 - ): + if recipe.row_scaled_activation or recipe.nvfp4_4over6 != "none": + # 4over6 doesn't used fused kernels + return ops + if recipe.fp8_format == RecipeFormat.UE5M3 and ScaledSReLU in activation_op_types: + # cuDNN has no SReLU support for UE5M3 for now return ops - - if activation_op_types is None: - activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) # Scan ops through with sliding window out = [] @@ -1341,16 +1407,24 @@ def fuser_forward( ) fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor if use_nvfp4: - nvfp4_fp4_max = 6.0 - nvfp4_fp8_max = 448.0 - nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max # cuDNN receives NVFP4 block-scaled inputs without TE's per-group # global scale factors, so alpha supplies the product of the two # operand global scales. fc1_alpha_tensor = ( - _nvfp4_amax(grouped_fc1_x, columnwise=False) - * _nvfp4_amax(grouped_fc1_weight, columnwise=False) - / (nvfp4_global_scale_denom**2) + _nvfp4_global_scale( + grouped_fc1_x, + fc1_input_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc1_weight, + fc1_weight_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) ).to(torch.float32) else: fc1_alpha_tensor = alpha_tensor @@ -1363,6 +1437,8 @@ def fuser_forward( and isinstance(fc2_input_quantizer, NVFP4Quantizer) and fc2_input_quantizer.with_rht and fc2_input_quantizer.with_post_rht_amax + # If we don't have the second-level scaling we don't need the post-RHT amax in the kernel. + and not fc2_input_quantizer.disable_second_level_scale ) activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_supports_hadamard = self._cudnn_act_func == "swiglu" or ( @@ -1389,6 +1465,13 @@ def fuser_forward( "current_stream": current_stream, "use_dynamic_sched": True, } + fc1_sf_dtype_override = _nvfp4_sf_dtype_override(fc1_input_quantizer) + # Only override the dtype if we are using e5m3 and not using the Hadamard kernel, + # since the Hadamard fused GEEM kernel does not support e5m3. + # At the time of writing, the e5m3 recipe doesn't have the second level scaling enabled, + # which naturally leads to use_fc1_act_hadamard=False + if fc1_sf_dtype_override is not None and not use_fc1_act_hadamard: + fc1_activation_kwargs["sf_fp8_dtype_override"] = fc1_sf_dtype_override if use_fc1_act_hadamard_srelu: fc1_activation_kwargs["act_func"] = "srelu" elif self._cudnn_act_func is not None: @@ -1527,7 +1610,8 @@ def fuser_forward( fc2_out_shape = in_shape[:-1] + [fc2_weight_shape[0]] fc2_scales = basic_op_extra_inputs[2][1] if fc2_op._scale_bias else None - if use_nvfp4: + fc2_input_sf_override = _nvfp4_sf_dtype_override(fc2_input_quantizer) + if use_nvfp4 and fc2_input_sf_override is None: fc2_bias_for_gemm = None fc2_bias_scale = None if fc2_bias_packed is not None: @@ -1595,6 +1679,120 @@ def fuser_forward( bias_scale=fc2_bias_scale, ) fc2_out = fc2_out_buf + elif use_nvfp4 and fc2_input_sf_override is not None: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. + fc2_in = fc1_kernel_out["d_tensor"] + fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc2_input_quantizer.optimize_for_gemm = True + + if use_fc1_act_hadamard: # Currently unreachable since e5m3 doesn't use second-level scaling + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), + tensor_offsets=fc2_x_tensor_offsets, + ) + else: + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + ) + + fc2_x_data, fc2_x_scales = convert_TE_MX_tensor_to_cuDNN_operand( + grouped_fc2_x.rowwise_data, + grouped_fc2_x.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=in_shape[0], + k_logical=fc2_weight_shape[1], + sf_swizzled=grouped_fc2_x._with_gemm_swizzled_scales, + ) + + fc2_fwd_alpha_tensor = ( + _nvfp4_global_scale( + grouped_fc2_x, + fc2_input_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc2_weight, + fc2_weight_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + ).to(torch.float32) + + fc2_scales_tensor = ( + fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + if fc2_scales is not None + else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) + ) + fc2_quant_kwargs = { + "a_tensor": fc2_x_data, + "sfa_tensor": fc2_x_scales, + "padded_offsets": split_points, + "alpha_tensor": fc2_fwd_alpha_tensor, + "bias_tensor": fc2_bias_packed, + "norm_const_tensor": None, + "prob_tensor": fc2_scales_tensor, + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": sf_vec_size, + "sf_fp8_dtype_override": fc2_input_sf_override, + "current_stream": current_stream, + "use_dynamic_sched": True, + } + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM (original stays unmodified + # for save_for_backward). + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) + + fc2_w_data, fc2_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc2_weight_for_gemm.rowwise_data, + fc2_weight_for_gemm.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=fc2_weight_shape[0], + k_logical=fc2_weight_shape[1], + L=num_groups, + sf_swizzled=fc2_weight_for_gemm._with_gemm_swizzled_scales, + ) + fc2_quant_kwargs["b_tensor"] = fc2_w_data + fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + "nvfp4", + device, + ) + ) + fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_quant_kwargs["n"] = fc2_weight_shape[0] + fc2_quant_kwargs["b_dtype"] = data_dtype + fc2_quant_kwargs["b_major"] = "k" + + output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) + fc2_quant_kwargs["d_tensor"] = output_buffer.as_strided( + (in_shape[0], fc2_weight_shape[0], 1), + (fc2_weight_shape[0], 1, in_shape[0] * fc2_weight_shape[0]), + ) + self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = output_buffer else: fc2_in_row_data = fc1_kernel_out["d_tensor"] fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) @@ -1780,6 +1978,7 @@ def fuser_forward( ) fc1_ctx.input_quantizers = [fc1_input_quantizer] + fc1_ctx.weight_quantizers = [fc1_weight_quantizer] fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad @@ -1788,6 +1987,7 @@ def fuser_forward( fc2_ctx.input_quantizers = [fc2_input_quantizer] fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] + fc2_ctx.weight_quantizers = [fc2_weight_quantizer] fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad @@ -1888,6 +2088,7 @@ def fuser_backward( # Split grad output tensor and convert dtypes if needed fc2_grad_output_quantizer = fc2_ctx.grad_output_quantizers[0] + fc2_weight_quantizer = fc2_ctx.weight_quantizers[0] fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=fc2_ctx.weight_requires_grad) fc2_grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias @@ -2008,24 +2209,33 @@ def fuser_backward( fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: - nvfp4_fp4_max = 6.0 - nvfp4_fp8_max = 448.0 - nvfp4_global_scale_denom = nvfp4_fp4_max * nvfp4_fp8_max - fc2_dy_amax = _nvfp4_amax(grouped_fc2_dy, columnwise=False) - fc2_weight_col_amax = _nvfp4_amax(grouped_fc2_weight, columnwise=True) + fc2_dy_global_scale = _nvfp4_global_scale( + grouped_fc2_dy, + fc2_grad_output_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + fc2_weight_col_global_scale = _nvfp4_global_scale( + grouped_fc2_weight, + fc2_weight_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, + ) if activation_is_srelu: # DSReLU applies alpha once, so pass the full product of the # two operand global scales. fc2_alpha_tensor = ( - (fc2_dy_amax * fc2_weight_col_amax / (nvfp4_global_scale_denom**2)) + (fc2_dy_global_scale * fc2_weight_col_global_scale) .to(torch.float32) .expand(num_groups) ) else: # DGLU applies alpha to both gate branches, so the wrapper # expects sqrt(product) to recover the same global-scale factor. - fc2_alpha_tensor = ( - torch.sqrt(fc2_dy_amax * fc2_weight_col_amax) / nvfp4_global_scale_denom + fc2_alpha_tensor = torch.sqrt( + fc2_dy_global_scale * fc2_weight_col_global_scale ).expand(num_groups) fc2_beta_tensor = get_cached_ones_tensor(num_groups, torch.float32, device) fc2_norm_const_tensor = None @@ -2054,6 +2264,9 @@ def fuser_backward( dactivation_kernel = self.grouped_gemm_dactivation_kernel() if _cudnn_frontend_supports_single_group_runtime_offsets(): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + fc2_sf_dtype_override = _nvfp4_sf_dtype_override(fc2_grad_output_quantizer) + if fc2_sf_dtype_override is not None: + fc2_dactivation_kwargs["sf_fp8_dtype_override"] = fc2_sf_dtype_override if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func @@ -2271,6 +2484,7 @@ def fuser_backward( # FC1 grad output for dgrad and wgrad GEMMs fc1_dy_tensor_offsets = fc1_out_tensor_offsets fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] + fc1_weight_quantizer = fc1_ctx.weight_quantizers[0] if use_nvfp4: fc1_grad_output_quantizer.set_usage( rowwise=True, @@ -2344,6 +2558,8 @@ def fuser_backward( if is_distributed_weight(fc1_leader): grouped_fc1_weight = materialize_weight_for_backward(fc1_leader) + fc1_dgrad_sf_override = _nvfp4_sf_dtype_override(fc1_grad_output_quantizer) + use_single_group_dense_dgrad = num_groups == 1 if use_single_group_dense_dgrad: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) @@ -2354,7 +2570,7 @@ def fuser_backward( single_grouped_weight=fc1_op.single_grouped_weight, dtype=dtype, ) - elif use_nvfp4: + elif use_nvfp4 and fc1_dgrad_sf_override is None: grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) grouped_grad_input = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[1]), @@ -2371,6 +2587,106 @@ def fuser_backward( grouped_grad_input, layout="NN", ) + elif use_nvfp4: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready + # This assertion should never fail because we set fc1_grad_output_quantizer.optimize_for_gemm = True + assert grouped_fc1_dy._with_gemm_swizzled_scales, ( + "cuDNN NVFP4 dgrad requires GEMM-swizzled grad-output scale factors." + ) + + grad_input_buffer = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) + + dgrad_k = fc1_weight_shape[0] # contraction dim + dgrad_valid_m = out_shape[0] # batch dim + + # Create A and its sf tensor for cuDNN that satisfies its layout requirements + fc1_dgrad_a_data, fc1_dgrad_a_scales = convert_TE_MX_tensor_to_cuDNN_operand( + grouped_fc1_dy.rowwise_data, + grouped_fc1_dy.scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=dgrad_valid_m, + k_logical=dgrad_k, + sf_swizzled=grouped_fc1_dy._with_gemm_swizzled_scales, + ) + + fc1_dgrad_alpha = ( + _nvfp4_global_scale( + grouped_fc1_dy, + fc1_grad_output_quantizer, + columnwise=False, + num_groups=num_groups, + device=device, + ) + * _nvfp4_global_scale( + grouped_fc1_weight, + fc1_weight_quantizer, + columnwise=True, + num_groups=num_groups, + device=device, + ) + ).to(torch.float32) + + fc1_dgrad_kwargs = { + "a_tensor": fc1_dgrad_a_data, + "sfa_tensor": fc1_dgrad_a_scales, + "padded_offsets": split_points, + "alpha_tensor": fc1_dgrad_alpha, + "norm_const_tensor": None, # must be None for FP4 inputs + "acc_dtype": torch.float32, + "d_dtype": dtype, # high precision -> no output quantization + "cd_major": "n", + "sf_vec_size": sf_vec_size, + "sf_fp8_dtype_override": fc1_dgrad_sf_override, + "current_stream": current_stream, + "discrete_col_sfd": False, + "use_dynamic_sched": True, + } + + if fc1_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + fc1_weight_for_gemm = grouped_fc1_weight.copy() + tex.grouped_swizzle_for_gemm( + fc1_weight_for_gemm, rowwise=False, columnwise=True + ) + + # Create B and its sf tensor for cuDNN that satisfies its layout + # requirements. NVFP4 column-wise data is physically transposed, so + # it is already (in_features, out_features) and stays K-major. + fc1_w_data, fc1_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc1_weight_for_gemm.columnwise_data, + fc1_weight_for_gemm.columnwise_scale_inv, + data_dtype=data_dtype, + scale_dtype=scale_view_dtype, + valid_M_or_N=fc1_weight_shape[1], + k_logical=dgrad_k, + L=num_groups, + sf_swizzled=fc1_weight_for_gemm._with_gemm_swizzled_scales, + ) + fc1_dgrad_kwargs["b_tensor"] = fc1_w_data + fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales + else: + fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sfb_buffer = ( + tex.grouped_mlp_experimental.swizzle_scales_and_pack_ptrs_for_discrete_weights( + [w._columnwise_data for w in grouped_fc1_weight], + [w._columnwise_scale_inv for w in grouped_fc1_weight], + "nvfp4", + device, + ) + ) + fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_dgrad_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_dgrad_kwargs["n"] = fc1_weight_shape[1] + fc1_dgrad_kwargs["b_dtype"] = torch.float4_e2m1fn_x2 + # FP4 has no N-major operand support, and the column-wise buffer is + # already transposed, so it is K-major. + fc1_dgrad_kwargs["b_major"] = "k" + + fc1_dgrad_kwargs["d_tensor"] = grad_input_buffer.as_strided( + (out_shape[0], fc1_weight_shape[1], 1), + (fc1_weight_shape[1], 1, out_shape[0] * fc1_weight_shape[1]), + ) + self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + grad_input = grad_input_buffer else: fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index e35d57b363..fa77fdd1e3 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -851,7 +851,17 @@ def _cast_master_weights_to_nvfp4_2d( # This replaces multiple Python tensor operations with a single kernel global_scale_tensor = torch.empty_like(global_amaxes) - tex.nvfp4_compute_global_scale(global_amaxes, global_scale_tensor) + # There should only be one scale dtype for all quantizers in the params list if using the same NVFP4 recipe. + scale_dtypes = {p[0]._get_quantizer().scale_dtype for p in params} + if len(scale_dtypes) != 1: + raise ValueError( + "quantize_master_weights requires a single NVFP4 scale dtype per call, " + f"but got {scale_dtypes}." + ) + # NVFP4Quantizer.scale_dtype is the pure-python constants.DType; the pybind + # entry points want tex.DType. The enum values are shared, so map by value. + scale_dtype = tex.DType(int(scale_dtypes.pop())) + tex.nvfp4_compute_global_scale(global_amaxes, global_scale_tensor, scale_dtype=scale_dtype) global_scale_views = [global_scale_tensor[i : i + 1] for i in range(len(params))] # Collect tensors for batched fused scale kernel @@ -949,6 +959,7 @@ def _cast_master_weights_to_nvfp4_2d( fused_scale_tile_cols_list, fused_scale_rows_padded_list, block_len, + scale_dtype=scale_dtype, ) # Batched multi-tensor call for partial cast @@ -962,6 +973,7 @@ def _cast_master_weights_to_nvfp4_2d( partial_cast_w_list, partial_cast_start_offset_list, block_len, + scale_dtype=scale_dtype, ) From ab3a9b384a2f345b677ad5254424eae16f230055 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 14 Aug 2026 04:39:28 +0000 Subject: [PATCH 04/54] Use custom recipe for NVFP4-UE5M3 tests Signed-off-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 9 ++++--- tests/pytorch/test_grouped_mlp.py | 7 +++++- tests/pytorch/utils.py | 40 ++++++++++++++++++++++++------- 3 files changed, 44 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 1ae33c0403..9f0101c39e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -225,7 +225,11 @@ def make_reference_and_test_tensors( tensor_type = quantizer_role.tensor_type with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" scale_dtype = ( - te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3 + te.DType.kFloat8UE5M3 if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + else te.DType.kFloat8E4M3 + ) + disable_second_level_scale = ( + scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" ) test = NVFP4Quantizer( scale_dtype=scale_dtype, @@ -234,6 +238,7 @@ def make_reference_and_test_tensors( with_2d_quantization=False, stochastic_rounding=False, with_random_sign_mask=with_rht, + disable_second_level_scale=disable_second_level_scale, )(test) elif quantization == "nvfp4_4over6": tensor_type = "input" @@ -886,7 +891,6 @@ def test_quantize( quantization=quantization, test_dtype=dtype, test_device=device, - quantizer_role=QuantizerRole(tensor_type="input"), requires_grad=True, ) grad_quantization = quantization @@ -898,7 +902,6 @@ def test_quantize( quantization=grad_quantization, test_dtype=dtype, test_device=device, - quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 1f797251b1..84509cf34f 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -201,7 +201,11 @@ def make_reference_and_test_tensors( tensor_type = quantizer_role.tensor_type with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" scale_dtype = ( - te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3 + te.DType.kFloat8UE5M3 if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + else te.DType.kFloat8E4M3 + ) + disable_second_level_scale = ( + scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" ) test = NVFP4Quantizer( scale_dtype=scale_dtype, @@ -210,6 +214,7 @@ def make_reference_and_test_tensors( with_2d_quantization=False, stochastic_rounding=False, with_random_sign_mask=with_rht, + disable_second_level_scale=disable_second_level_scale, )(test) elif quantization == "nvfp4_4over6": tensor_type = "input" diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 353dbf8f60..6514d3b1c8 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -17,10 +17,14 @@ import torch import transformer_engine -from transformer_engine.common.recipe import Format as RecipeFormat -from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch import InferenceParams, QuantizedTensor -from transformer_engine.pytorch import DType +from transformer_engine.common.recipe import Format as RecipeFormat, Recipe +from transformer_engine.pytorch import ( + DType, + InferenceParams, + NVFP4Quantizer, + QuantizedTensor, + QuantizerRole, +) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, @@ -158,19 +162,39 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: ) if name == "fp8_block_scaling": return transformer_engine.common.recipe.Float8BlockScaling(**recipe_kwargs) + if name in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3"): + + def make_nvfp4_ue5m3_quantizer(role: QuantizerRole) -> NVFP4Quantizer: + """Quantizer factory for NVFP4-UE5M3 recipe.""" + tensor_type = role.tensor_type if role is not None else "input" + if not tensor_type: + tensor_type = "input" + with_rht = name == "nvfp4_rht_ue5m3" and tensor_type != "weight" + return NVFP4Quantizer( + scale_dtype=DType.kFloat8UE5M3, + with_rht=with_rht, + with_post_rht_amax=with_rht, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=with_rht, + disable_second_level_scale=tensor_type == "input", + ) + + recipe = transformer_engine.common.recipe.CustomRecipe( + qfactory=make_nvfp4_ue5m3_quantizer, + **recipe_kwargs, + ) + recipe.enable_cutedsl_fused_grouped_mlp = True + return recipe if name in nvfp4_variant_names: with_rht = name in ("nvfp4_rht", "nvfp4_rht_ue5m3") use_4over6 = name == "nvfp4_4over6" - scale_format = ( - RecipeFormat.UE5M3 if name in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") else RecipeFormat.E4M3 - ) kwargs = { "disable_rht": not with_rht, "disable_stochastic_rounding": True, "disable_2d_quantization": not use_4over6, "row_scaled_activation": name == "nvfp4_row_scaled", "nvfp4_4over6": "all" if use_4over6 else "none", - "fp8_format": scale_format, } kwargs.update(recipe_kwargs) return transformer_engine.common.recipe.NVFP4BlockScaling(**kwargs) From 92b1063da31904b36b99946b30a129adc3fa68c5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:10:34 +0000 Subject: [PATCH 05/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 11 ++-- tests/pytorch/test_fusible_ops.py | 16 +++--- tests/pytorch/test_grouped_mlp.py | 16 +++--- .../pytorch/cpp_extensions/gemm.py | 53 +++++++++++-------- .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 3 +- .../pytorch/csrc/extensions/pybind.cpp | 3 +- .../pytorch/csrc/extensions/transpose.cpp | 7 ++- .../pytorch/ops/fused/grouped_mlp.py | 24 +++++---- 8 files changed, 76 insertions(+), 57 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 40c7c3bf82..7aeab72f23 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -751,8 +751,10 @@ def _check_ue5m3_gemm_versus_dequantized( rel_err = (y.float() - ref).norm() / ref.norm() assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large" + ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3) @pytest.mark.parametrize( @@ -776,12 +778,15 @@ def _check_ue5m3_gemm_versus_dequantized( (False, False), # TN -- w rowwise, x rowwise (fprop) (False, True), # NN -- w colwise, x rowwise (dgrad) (True, True), # NT -- w colwise, x colwise (wgrad) - ], ids=["FF", "FT", "TT"] + ], + ids=["FF", "FT", "TT"], ) @pytest.mark.parametrize( - "disable_second_level_scale", [ + "disable_second_level_scale", + [ (True, False), - ], ids=["TF"] + ], + ids=["TF"], ) def test_nvfp4_ue5m3_gemm_versus_reference( M: int, diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 9f0101c39e..36e6a8e48e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -138,11 +138,12 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") - if ( - quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") - and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and ( + math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0 ): - pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") + pytest.skip( + "cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors" + ) # Check dtype if dtype is not None: @@ -225,12 +226,11 @@ def make_reference_and_test_tensors( tensor_type = quantizer_role.tensor_type with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" scale_dtype = ( - te.DType.kFloat8UE5M3 if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + te.DType.kFloat8UE5M3 + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") else te.DType.kFloat8E4M3 ) - disable_second_level_scale = ( - scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" - ) + disable_second_level_scale = scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" test = NVFP4Quantizer( scale_dtype=scale_dtype, with_rht=with_rht, diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 84509cf34f..f239518b84 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -125,11 +125,12 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") - if ( - quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") - and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0) + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and ( + math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0 ): - pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors") + pytest.skip( + "cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors" + ) # Check dtype if dtype is not None: @@ -201,12 +202,11 @@ def make_reference_and_test_tensors( tensor_type = quantizer_role.tensor_type with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight" scale_dtype = ( - te.DType.kFloat8UE5M3 if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") + te.DType.kFloat8UE5M3 + if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") else te.DType.kFloat8E4M3 ) - disable_second_level_scale = ( - scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" - ) + disable_second_level_scale = scale_dtype == te.DType.kFloat8UE5M3 and tensor_type == "input" test = NVFP4Quantizer( scale_dtype=scale_dtype, with_rht=with_rht, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 484f1c4503..aae9325cec 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -343,8 +343,10 @@ def convert_TE_MX_tensor_to_cuDNN_operand( """ if use_N_major_for_B: - assert data_dtype in (torch.float8_e4m3fn, torch.float8_e5m2), \ - f"Using N-major layout for B is only supported for FP8, but got {data_dtype}." + assert data_dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ), f"Using N-major layout for B is only supported for FP8, but got {data_dtype}." available_scalings = { # NVFP4 recipe (UE5M3 rides as E4M3 since torch has no ue5m3 dtype) @@ -364,7 +366,7 @@ def convert_TE_MX_tensor_to_cuDNN_operand( if data_dtype == torch.float4_e2m1fn_x2: k_packed = k_logical // 2 # fp4 packs two values per byte else: - k_packed = k_logical # fp8 packs one value per byte + k_packed = k_logical # fp8 packs one value per byte data = data.view(dtype=data_dtype) if use_N_major_for_B: @@ -456,9 +458,12 @@ def general_cuDNN_MX_gemm( functions) should be removed entirely. """ - assert isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) and \ - A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3, \ - f"cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." + assert ( + isinstance(A, NVFP4TensorStorage) + and isinstance(B, NVFP4TensorStorage) + and A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 + and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 + ), f"cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." @@ -474,8 +479,10 @@ def general_cuDNN_MX_gemm( transa = layout[0] == "T" transb = layout[1] == "T" - assert out_dtype in (torch.float32, torch.float16, torch.bfloat16), \ - f"cuDNN MX GEMM currently only supports float32, float16, and bfloat16 outputs, but got {out_dtype}." + assert out_dtype in (torch.float32, torch.float16, torch.bfloat16), ( + "cuDNN MX GEMM currently only supports float32, float16, and bfloat16 outputs, but got" + f" {out_dtype}." + ) device = A.device @@ -492,9 +499,9 @@ def general_cuDNN_MX_gemm( # `grad` only changes behaviour when a bias is supplied: it turns the bias slot # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs # that pass grad=True without a bias need nothing special. - assert not (grad and bias is not None), ( - "cuDNN GEMM currently does not support fused bias gradient." - ) + assert not ( + grad and bias is not None + ), "cuDNN GEMM currently does not support fused bias gradient." # Pick the buffer whose block scales run along K. In every case the selected # buffer is physically (rows, K_packed), so the reshape below is uniform. @@ -520,7 +527,9 @@ def general_cuDNN_MX_gemm( N_full = [B_shape[0]] if transb else B_shape[:-1] K_full = [A_shape[-1]] if transa else A_shape[1:] K_full_b = B_shape[1:] if transb else [B_shape[-1]] - assert K_full == K_full_b, f"Contraction dims disagree: A implies {K_full}, B implies {K_full_b}." + assert ( + K_full == K_full_b + ), f"Contraction dims disagree: A implies {K_full}, B implies {K_full_b}." M = math.prod(M_full) N = math.prod(N_full) K = math.prod(K_full) @@ -551,9 +560,9 @@ def general_cuDNN_MX_gemm( if accumulate or N % 256 != 0: alpha = alpha if alpha is not None else 1.0 # This path uses cuDNN's wgrad (grouped_gemm_wgrad_wrapper_sm100) which supports grad accumulation - if accumulate: # Accumulate GEMM's result to the out tensor + if accumulate: # Accumulate GEMM's result to the out tensor assert beta in (1.0, None), "beta must be one or None if accumulate is True" - else: # Overwrite GEMM's result to the out tensor + else: # Overwrite GEMM's result to the out tensor assert beta in (0.0, None), "beta must be zero or None if not accumulate" _cuDNN_wgrad_gemm( a_tensor=dataB.view(N, K // 2), @@ -572,7 +581,9 @@ def general_cuDNN_MX_gemm( alpha = alpha if alpha is not None else 1.0 # cuDNN's general GEMM path (grouped_gemm_quant_wrapper_sm100) doesn't support accumulation - assert accumulate is False, "cuDNN GEMM currently does not support accumulation for this operation." + assert ( + accumulate is False + ), "cuDNN GEMM currently does not support accumulation for this operation." assert beta in (0.0, None), "beta must be zero or None if not accumulate" # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it @@ -634,9 +645,9 @@ def general_cuDNN_MX_gemm( alpha_tensor = (alpha * scaleA * scaleB).to(torch.float32) if bias is not None: - assert bias.dim() == 1 and bias.shape[0] == M, ( - f"cuDNN MX GEMM expects a ({M},) bias, but got {tuple(bias.shape)}." - ) + assert ( + bias.dim() == 1 and bias.shape[0] == M + ), f"cuDNN MX GEMM expects a ({M},) bias, but got {tuple(bias.shape)}." # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). bias = bias.contiguous().as_strided((M, 1), (1, M)) @@ -662,9 +673,9 @@ def general_cuDNN_MX_gemm( "acc_dtype": torch.float32, "d_dtype": out_dtype, # high precision -> no output quantization "d_tensor": d_tensor, - "cd_major": "n", # only "n" is supported by cuDNN - "sf_vec_size": NVFP4_BLOCK_SCALING_SIZE, # Hardcode to NVFP4 for now - "sf_fp8_dtype_override": "e5m3", # Hardcode for now + "cd_major": "n", # only "n" is supported by cuDNN + "sf_vec_size": NVFP4_BLOCK_SCALING_SIZE, # Hardcode to NVFP4 for now + "sf_fp8_dtype_override": "e5m3", # Hardcode for now "current_stream": torch.cuda.current_stream().cuda_stream, "discrete_col_sfd": False, "use_dynamic_sched": True, diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp index 8b58299351..e9321b9b86 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -45,8 +45,7 @@ void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tens nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), start_offset, block_len, - at::cuda::getCurrentCUDAStream(), - static_cast(scale_dtype)); + at::cuda::getCurrentCUDAStream(), static_cast(scale_dtype)); } void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 0971177ff9..b7cee9d0e1 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -417,8 +417,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); m.def("nvfp4_compute_global_scale", &transformer_engine::pytorch::nvfp4_compute_global_scale, "Compute global encode scale from global amax", py::arg("global_amax"), - py::arg("global_scale"), - py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, + py::arg("global_scale"), py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3, py::call_guard()); m.def("nvfp4_fused_scale", &transformer_engine::pytorch::nvfp4_fused_scale, "Fused kernel: compute per-block decode scale, copy global amax, expand to row-level FP8", diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 4b887c3749..9f34c4d196 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -144,8 +144,8 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile static_cast(block_len), at::cuda::getCurrentCUDAStream()); } -void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, - at::Tensor global_amax, const DType scale_dtype) { +void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, + const DType scale_dtype) { init_extension(); // block_amax and scale: [tile_rows, tile_cols], float32 @@ -193,8 +193,7 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor target_scale_cu.data(), target_amax_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), static_cast(block_len), - at::cuda::getCurrentCUDAStream(), - static_cast(scale_dtype)); + at::cuda::getCurrentCUDAStream(), static_cast(scale_dtype)); } void nvfp4_multi_tensor_fused_scale( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index f6f7c89d68..d21aa92352 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -533,8 +533,8 @@ def _cudnn_compute_wgrad( total_tokens = grouped_dy.logical_shape[0] device = grouped_dy.columnwise_data.device - dy_quantizer=getattr(grouped_dy, "quantizer", None) - x_quantizer=getattr(grouped_x, "quantizer", None) + dy_quantizer = getattr(grouped_dy, "quantizer", None) + x_quantizer = getattr(grouped_x, "quantizer", None) sfa_leading_dim = round_up_to_nearest_multiple(out_features, 128) sfb_leading_dim = round_up_to_nearest_multiple(in_features, 128) @@ -1683,13 +1683,17 @@ def fuser_forward( bias_scale=fc2_bias_scale, ) fc2_out = fc2_out_buf - elif use_nvfp4 and fc2_input_sf_override is not None: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. + elif ( + use_nvfp4 and fc2_input_sf_override is not None + ): # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. fc2_in = fc1_kernel_out["d_tensor"] fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) fc2_input_quantizer.optimize_for_gemm = True - if use_fc1_act_hadamard: # Currently unreachable since e5m3 doesn't use second-level scaling + if ( + use_fc1_act_hadamard + ): # Currently unreachable since e5m3 doesn't use second-level scaling grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( fc2_in, fc2_input_quantizer, @@ -2590,13 +2594,15 @@ def fuser_backward( grouped_grad_input, layout="NN", ) - elif use_nvfp4: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready + elif use_nvfp4: # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready # This assertion should never fail because we set fc1_grad_output_quantizer.optimize_for_gemm = True - assert grouped_fc1_dy._with_gemm_swizzled_scales, ( - "cuDNN NVFP4 dgrad requires GEMM-swizzled grad-output scale factors." - ) + assert ( + grouped_fc1_dy._with_gemm_swizzled_scales + ), "cuDNN NVFP4 dgrad requires GEMM-swizzled grad-output scale factors." - grad_input_buffer = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device) + grad_input_buffer = validate_or_alloc_output( + grad_input_buffer, in_shape, dtype, device + ) dgrad_k = fc1_weight_shape[0] # contraction dim dgrad_valid_m = out_shape[0] # batch dim From 3d25d201851d7cee27566f87654e6919eaa42e7f Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 14 Aug 2026 12:44:14 +0000 Subject: [PATCH 06/54] Add grouped MLP kernel for GGEMM+SwiGLU+RHT+quant Signed-off-by: Tim Moon --- .../pytorch/ops/fused/grouped_mlp.py | 279 ++++++++++-------- 1 file changed, 153 insertions(+), 126 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d21aa92352..847ac63dd9 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -17,7 +17,7 @@ import transformer_engine_torch as tex from ....common.recipe import Format as RecipeFormat -from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType +from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...cpp_extensions.gemm import convert_TE_MX_tensor_to_cuDNN_operand @@ -270,7 +270,7 @@ def _nvfp4_sf_dtype_override(quantizer: Optional[Quantizer]) -> Literal["e5m3"] # We don't use e5m3 for 4over6 return None scale_dtype = getattr(quantizer, "scale_dtype", None) - if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + if scale_dtype is not None and scale_dtype == DType.kFloat8UE5M3: return "e5m3" # If we don't use e5m3 we don't need to pass this string to override return None @@ -284,7 +284,7 @@ def _nvfp4_scale_max(quantizer: Quantizer) -> float: if override_max is not None and override_max != -1: return float(override_max) scale_dtype = getattr(quantizer, "scale_dtype", None) - if scale_dtype is not None and scale_dtype == tex.DType.kFloat8UE5M3: + if scale_dtype is not None and scale_dtype == DType.kFloat8UE5M3: return 114688.0 return 448.0 @@ -878,6 +878,7 @@ def fuse_grouped_mlp_ops( Updated operations with matched triples replaced by fused ops. """ if not fused_op_cls.is_supported(): + assert False ### TODO Remove return ops # Fused kernels are only supported for MXFP8 and NVFP4 @@ -1161,7 +1162,7 @@ def fuser_forward( if unit_activation_scale and num_groups != 1: unit_activation_scale = False - activation_kernel = self.grouped_gemm_activation_kernel() + activation_is_srelu = isinstance(activation_op, ScaledSReLU) supports_single_group_runtime_offsets = ( _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) @@ -1403,13 +1404,11 @@ def fuser_forward( fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) - fc1_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn fc1_prob_tensor = None if not unit_activation_scale: fc1_prob_tensor = ( scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) ) - fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor if use_nvfp4: # cuDNN receives NVFP4 block-scaled inputs without TE's per-group # global scale factors, so alpha supplies the product of the two @@ -1433,27 +1432,36 @@ def fuser_forward( else: fc1_alpha_tensor = alpha_tensor - use_tmem_post_rht_amax = _use_tmem_post_rht_amax() - use_fc1_act_hadamard = False - use_fc1_act_hadamard_srelu = False - use_nvfp4_rht_amax = ( + # Choose kernel implementation for FC1 + act + kernel_impl = "gemm_act" + if ( use_nvfp4 and isinstance(fc2_input_quantizer, NVFP4Quantizer) and fc2_input_quantizer.with_rht - and fc2_input_quantizer.with_post_rht_amax - # If we don't have the second-level scaling we don't need the post-RHT amax in the kernel. - and not fc2_input_quantizer.disable_second_level_scale - ) - activation_is_srelu = isinstance(activation_op, ScaledSReLU) - activation_supports_hadamard = self._cudnn_act_func == "swiglu" or ( - activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() - ) - if use_nvfp4_rht_amax and activation_supports_hadamard: - kernel_getter = getattr(self, "grouped_gemm_act_hadamard_kernel", None) - if kernel_getter is not None: - use_fc1_act_hadamard = kernel_getter() is not None - use_fc1_act_hadamard_srelu = use_fc1_act_hadamard and activation_is_srelu - + ): + if fc2_input_quantizer.disable_second_level_scale: + # Use GEMM + act + RHT + quant kernel if available + kernel_getter = getattr(self, "grouped_gemm_act_hadamard_quant_kernel", None) + if kernel_getter is None or kernel_getter() is None: + # Kernel is not available + pass + elif not activation_is_srelu: + kernel_impl = "gemm_act_rht_quant" + elif fc2_input_quantizer.with_post_rht_amax: + # Use GEMM + act + RHT + amax kernel if available + kernel_getter = getattr(self, "grouped_gemm_act_hadamard_kernel", None) + if kernel_getter is None or kernel_getter() is None: + # Kernel is not available + pass + elif self._cudnn_act_func == "swiglu": + kernel_impl = "gemm_act_rht_amax" + elif ( + activation_is_srelu + and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() + ): + kernel_impl = "gemm_act_rht_amax" + + # Common kernel arguments fc1_activation_kwargs = { "a_tensor": fc1_x_data, "sfa_tensor": fc1_x_scales, @@ -1463,30 +1471,39 @@ def fuser_forward( "prob_tensor": fc1_prob_tensor, "acc_dtype": torch.float32, "c_dtype": torch.bfloat16, - "d_dtype": fc1_d_dtype, "cd_major": "n", "sf_vec_size": sf_vec_size, "current_stream": current_stream, "use_dynamic_sched": True, } - fc1_sf_dtype_override = _nvfp4_sf_dtype_override(fc1_input_quantizer) - # Only override the dtype if we are using e5m3 and not using the Hadamard kernel, - # since the Hadamard fused GEEM kernel does not support e5m3. - # At the time of writing, the e5m3 recipe doesn't have the second level scaling enabled, - # which naturally leads to use_fc1_act_hadamard=False - if fc1_sf_dtype_override is not None and not use_fc1_act_hadamard: - fc1_activation_kwargs["sf_fp8_dtype_override"] = fc1_sf_dtype_override - if use_fc1_act_hadamard_srelu: - fc1_activation_kwargs["act_func"] = "srelu" - elif self._cudnn_act_func is not None: - fc1_activation_kwargs["act_func"] = self._cudnn_act_func - if use_fc1_act_hadamard: - fc1_activation_kwargs["use_tmem_post_rht_amax"] = use_tmem_post_rht_amax - else: - fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor + + # Kernel arguments based on kernel implementation + if kernel_impl == "gemm_act": + fc1_activation_kwargs["norm_const_tensor"] = None if use_nvfp4 else norm_const_tensor + fc1_activation_kwargs["d_dtype"] = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4 if supports_single_group_runtime_offsets: fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 + elif kernel_impl == "gemm_act_rht_amax": + fc1_activation_kwargs["d_dtype"] = torch.bfloat16 + fc1_activation_kwargs["use_tmem_post_rht_amax"] = _use_tmem_post_rht_amax() + elif kernel_impl == "gemm_act_rht_quant": + fc1_activation_kwargs["d_dtype"] = torch.float4_e2m1fn_x2 + fc1_activation_kwargs["rht_dtype"] = torch.float4_e2m1fn_x2 + + # Miscellaneous kernel arguments + if activation_is_srelu: + if kernel_impl in ("gemm_act_rht_amax", "gemm_act_rht_quant"): + fc1_activation_kwargs["act_func"] = "srelu" + elif self._cudnn_act_func is not None: + fc1_activation_kwargs["act_func"] = self._cudnn_act_func + if ( + isinstance(fc1_input_quantizer, NVFP4Quantizer) + and fc1_input_quantizer.scale_dtype == DType.kFloat8UE5M3 + ): + # PyTorch does not have a UE5M3 dtype, so override the + # tensor dtype when using UE5M3 scales + fc1_activation_kwargs["sf_fp8_dtype_override"] = "e5m3" if self._pass_geglu_runtime_params: fc1_activation_kwargs.update( linear_offset=self._cudnn_linear_offset, @@ -1586,10 +1603,18 @@ def fuser_forward( fc1_activation_kwargs["b_dtype"] = data_dtype fc1_activation_kwargs["b_major"] = "k" - if use_fc1_act_hadamard: + # Launch FC1 + act kernel + if kernel_impl == "gemm_act": + fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) + elif kernel_impl == "gemm_act_rht_amax": fc1_kernel_out = self.grouped_gemm_act_hadamard_kernel()(**fc1_activation_kwargs) + elif kernel_impl == "gemm_act_rht_quant": + fc1_kernel_out = self.grouped_gemm_act_hadamard_quant_kernel()(**fc1_activation_kwargs) else: - fc1_kernel_out = activation_kernel(**fc1_activation_kwargs) + raise RuntimeError("Unrecognized kernel variant ({kernel_impl})") + + activation_in = fc1_kernel_out["c_tensor"] + activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) if fc2_is_dist: grouped_fc2_weight = materialize_weight_for_forward(grouped_fc2_weight) @@ -1599,54 +1624,93 @@ def fuser_forward( grouped_fc2_weight._with_gemm_swizzled_scales = False # Unpack kernel outputs - # Note: Fused kernel outputs tensors with non-contiguous - # logical dims. - # Row-wise data logical shape: (sum(m_splits), k, 1) - # Row-wise scale logical shape: (32 (block row), 4 (block row), - # sum(m_splits)/128, 4 (block col), k/128, 1) - # Column-wise data logical shape: (sum(m_splits), k, 1) - # Column-wise scale logical shape: (32 (block col), 4 (block col), - # k/128, 4 (block row), sum(m_splits)/128, 1) - activation_in = fc1_kernel_out["c_tensor"] - activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) - - # FC2 GEMM - fc2_out_shape = in_shape[:-1] + [fc2_weight_shape[0]] - fc2_scales = basic_op_extra_inputs[2][1] if fc2_op._scale_bias else None - - fc2_input_sf_override = _nvfp4_sf_dtype_override(fc2_input_quantizer) - if use_nvfp4 and fc2_input_sf_override is None: - fc2_bias_for_gemm = None - fc2_bias_scale = None - if fc2_bias_packed is not None: - fc2_bias_for_gemm = fc2_op._get_grouped_bias_for_gemm(dtype) - if fc2_scales is not None: - fc2_bias_scale = fc2_scales.reshape(-1) - if fc2_bias_scale.dtype != torch.float32: - fc2_bias_scale = fc2_bias_scale.to(dtype=torch.float32) - - fc2_in = fc1_kernel_out["d_tensor"] - fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + if use_nvfp4: fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) fc2_input_quantizer.optimize_for_gemm = True - if use_fc1_act_hadamard: - grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( + if kernel_impl == "gemm_act": + # Quantize to NVFP4 + fc2_in = fc1_kernel_out["d_tensor"] + fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + grouped_fc2_x = _group_quantize_for_grouped_mlp( fc2_in, fc2_input_quantizer, num_groups, split_sizes, - fc1_kernel_out["amax_tensor"].view(-1), - fc1_kernel_out["post_rht_amax_tensor"].view(-1), tensor_offsets=fc2_x_tensor_offsets, ) - else: - grouped_fc2_x = _group_quantize_for_grouped_mlp( + elif kernel_impl == "gemm_act_rht_amax": + # Quantize to NVFP4 using precomputed amax + fc2_in = fc1_kernel_out["d_tensor"] + fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( fc2_in, fc2_input_quantizer, num_groups, split_sizes, + fc1_kernel_out["amax_tensor"].view(-1), + fc1_kernel_out["post_rht_amax_tensor"].view(-1), tensor_offsets=fc2_x_tensor_offsets, ) + elif kernel_impl == "gemm_act_rht_quant": + # Unpack NVFP4 output + fc2_in_row_data = fc1_kernel_out["d_tensor"] + fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] + fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) + fc2_in_col_data = fc1_kernel_out["rht_tensor"] + fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_col_scale = fc1_kernel_out["sfrht_tensor"] + fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) + grouped_fc2_x = GroupedTensorStorage( + shape=(in_shape[0], fc2_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc2_input_quantizer, + data=fc2_in_row_data.reshape(-1), + columnwise_data=fc2_in_col_data.reshape(-1), + scale_inv=fc2_in_row_scale.reshape(-1), + columnwise_scale_inv=fc2_in_col_scale.reshape(-1), + first_dims=split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + with_gemm_swizzled_scales=True, + ) + else: + # Unpack MXFP8 output + fc2_in_row_data = fc1_kernel_out["d_tensor"] + fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] + fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) + fc2_in_col_data = fc1_kernel_out["d_col_tensor"] + fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] + fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) + grouped_fc2_x = GroupedTensorStorage( + shape=(in_shape[0], fc2_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc2_input_quantizer, + data=fc2_in_row_data.reshape(-1), + columnwise_data=fc2_in_col_data.reshape(-1), + scale_inv=fc2_in_row_scale.reshape(-1), + columnwise_scale_inv=fc2_in_col_scale.reshape(-1), + first_dims=split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + with_gemm_swizzled_scales=True, + ) + + # FC2 GEMM + fc2_out_shape = in_shape[:-1] + [fc2_weight_shape[0]] + fc2_scales = basic_op_extra_inputs[2][1] if fc2_op._scale_bias else None + fc2_input_sf_override = _nvfp4_sf_dtype_override(fc2_input_quantizer) + if use_nvfp4 and fc2_input_sf_override is None: + fc2_bias_for_gemm = None + fc2_bias_scale = None + if fc2_bias_packed is not None: + fc2_bias_for_gemm = fc2_op._get_grouped_bias_for_gemm(dtype) + if fc2_scales is not None: + fc2_bias_scale = fc2_scales.reshape(-1) + if fc2_bias_scale.dtype != torch.float32: + fc2_bias_scale = fc2_bias_scale.to(dtype=torch.float32) fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if ( @@ -1686,32 +1750,6 @@ def fuser_forward( elif ( use_nvfp4 and fc2_input_sf_override is not None ): # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. - fc2_in = fc1_kernel_out["d_tensor"] - fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() - fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) - fc2_input_quantizer.optimize_for_gemm = True - - if ( - use_fc1_act_hadamard - ): # Currently unreachable since e5m3 doesn't use second-level scaling - grouped_fc2_x = _group_quantize_with_amax_for_grouped_mlp( - fc2_in, - fc2_input_quantizer, - num_groups, - split_sizes, - fc1_kernel_out["amax_tensor"].view(-1), - fc1_kernel_out["post_rht_amax_tensor"].view(-1), - tensor_offsets=fc2_x_tensor_offsets, - ) - else: - grouped_fc2_x = _group_quantize_for_grouped_mlp( - fc2_in, - fc2_input_quantizer, - num_groups, - split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - ) - fc2_x_data, fc2_x_scales = convert_TE_MX_tensor_to_cuDNN_operand( grouped_fc2_x.rowwise_data, grouped_fc2_x.scale_inv, @@ -1802,30 +1840,6 @@ def fuser_forward( self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) fc2_out = output_buffer else: - fc2_in_row_data = fc1_kernel_out["d_tensor"] - fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) - fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] - fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) - - fc2_in_col_data = fc1_kernel_out["d_col_tensor"] - fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) - fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] - fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) - - grouped_fc2_x = GroupedTensorStorage( - shape=(in_shape[0], fc2_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc2_input_quantizer, - data=fc2_in_row_data.reshape(-1), - columnwise_data=fc2_in_col_data.reshape(-1), - scale_inv=fc2_in_row_scale.reshape(-1), - columnwise_scale_inv=fc2_in_col_scale.reshape(-1), - first_dims=split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) - use_single_group_dense_fc2 = num_groups == 1 fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) if use_single_group_dense_fc2: @@ -2851,6 +2865,19 @@ def grouped_gemm_act_hadamard_kernel(cls) -> Optional[Callable]: return grouped_gemm_glu_hadamard_wrapper_sm100 + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_act_hadamard_quant_kernel(cls) -> Optional[Callable]: + """Fused grouped GEMM activation kernel that also NVFP4 with RHT.""" + try: + from cudnn import ( + grouped_gemm_glu_hadamard_quant_wrapper_sm100, + ) # pylint: disable=no-name-in-module,import-outside-toplevel + except ImportError: + return None + + return grouped_gemm_glu_hadamard_quant_wrapper_sm100 + @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_dactivation_kernel(cls) -> Callable: From f030e40646193bc1302bc82e4d184f7593a651a2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:45:34 +0000 Subject: [PATCH 07/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 847ac63dd9..2063dab41d 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -878,7 +878,7 @@ def fuse_grouped_mlp_ops( Updated operations with matched triples replaced by fused ops. """ if not fused_op_cls.is_supported(): - assert False ### TODO Remove + assert False ### TODO Remove return ops # Fused kernels are only supported for MXFP8 and NVFP4 @@ -1455,10 +1455,7 @@ def fuser_forward( pass elif self._cudnn_act_func == "swiglu": kernel_impl = "gemm_act_rht_amax" - elif ( - activation_is_srelu - and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() - ): + elif activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): kernel_impl = "gemm_act_rht_amax" # Common kernel arguments From e366e8887eb984047212f83da4f0159e3e78d8c2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:57:34 +0000 Subject: [PATCH 08/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/cpp_extensions/gemm.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 4ac223e7cc..7a2fb538bb 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -963,9 +963,7 @@ def general_grouped_gemm( use_general_gemm_impl = False if isinstance(quantization_params[0], DebugQuantizer): use_general_gemm_impl = True - elif any( - _is_nvfp4_row_scaled_tensor(tensor) for tensor in itertools.chain(A, B) - ): + elif any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in itertools.chain(A, B)): use_general_gemm_impl = True elif any( isinstance(t, NVFP4TensorStorage) and t._scale_dtype == DType.kFloat8UE5M3 From 252ae94052fa8c850ac06f1173df0a3f707efbf7 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 15 Aug 2026 01:40:23 +0000 Subject: [PATCH 09/54] Debug integration with GGEMM+GLU+RHT+quant Signed-off-by: Tim Moon --- .../pytorch/ops/fused/grouped_mlp.py | 42 ++++++++++--------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index f0b00df76b..1abc423ff4 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1445,7 +1445,7 @@ def fuser_forward( if kernel_getter is None or kernel_getter() is None: # Kernel is not available pass - elif not activation_is_srelu: + elif self._cudnn_act_func == "swiglu": kernel_impl = "gemm_act_rht_quant" elif fc2_input_quantizer.with_post_rht_amax: # Use GEMM + act + RHT + amax kernel if available @@ -1473,6 +1473,13 @@ def fuser_forward( "current_stream": current_stream, "use_dynamic_sched": True, } + if ( + isinstance(fc1_input_quantizer, NVFP4Quantizer) + and fc1_input_quantizer.scale_dtype == DType.kFloat8UE5M3 + ): + # PyTorch does not have a UE5M3 dtype, so override the + # tensor dtype when using UE5M3 scales + fc1_activation_kwargs["sf_fp8_dtype_override"] = "e5m3" # Kernel arguments based on kernel implementation if kernel_impl == "gemm_act": @@ -1488,26 +1495,23 @@ def fuser_forward( fc1_activation_kwargs["d_dtype"] = torch.float4_e2m1fn_x2 fc1_activation_kwargs["rht_dtype"] = torch.float4_e2m1fn_x2 - # Miscellaneous kernel arguments + # Kernel arguments based on activation if activation_is_srelu: if kernel_impl in ("gemm_act_rht_amax", "gemm_act_rht_quant"): fc1_activation_kwargs["act_func"] = "srelu" elif self._cudnn_act_func is not None: fc1_activation_kwargs["act_func"] = self._cudnn_act_func - if ( - isinstance(fc1_input_quantizer, NVFP4Quantizer) - and fc1_input_quantizer.scale_dtype == DType.kFloat8UE5M3 - ): - # PyTorch does not have a UE5M3 dtype, so override the - # tensor dtype when using UE5M3 scales - fc1_activation_kwargs["sf_fp8_dtype_override"] = "e5m3" - if self._pass_geglu_runtime_params: - fc1_activation_kwargs.update( - linear_offset=self._cudnn_linear_offset, - geglu_alpha=self._cudnn_geglu_alpha, - glu_clamp_max=self._cudnn_glu_clamp_max, - glu_clamp_min=self._cudnn_glu_clamp_min, - ) + if self._cudnn_act_func == "geglu" and self._pass_geglu_runtime_params: + if kernel_impl == "gemm_act_rht_quant": + fc1_activation_kwargs["glu_alpha"] = self._cudnn_geglu_alpha + fc1_activation_kwargs["glu_limit"] = self._cudnn_glu_clamp_max + else: + fc1_activation_kwargs.update( + linear_offset=self._cudnn_linear_offset, + geglu_alpha=self._cudnn_geglu_alpha, + glu_clamp_max=self._cudnn_glu_clamp_max, + glu_clamp_min=self._cudnn_glu_clamp_min, + ) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. @@ -1651,11 +1655,11 @@ def fuser_forward( elif kernel_impl == "gemm_act_rht_quant": # Unpack NVFP4 output fc2_in_row_data = fc1_kernel_out["d_tensor"] - fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) - fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] + fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1] // 2) + fc2_in_row_scale = fc1_kernel_out["sfd_tensor"] fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) fc2_in_col_data = fc1_kernel_out["rht_tensor"] - fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_col_data = fc2_in_col_data.view(fc2_weight_shape[1], in_shape[0] // 2) fc2_in_col_scale = fc1_kernel_out["sfrht_tensor"] fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) grouped_fc2_x = GroupedTensorStorage( From f2b9c986e24ce8a3bf9091ed9822ab063b4dc946 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 18 Aug 2026 02:26:32 +0000 Subject: [PATCH 10/54] Remove scale max helper functions from NVFP4 cast utils Signed-off-by: Tim Moon --- .../common/cast/dispatch/quantize.cuh | 12 +- .../common/cast/nvfp4/core_nvfp4.cuh | 74 ----- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 35 ++- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 263 +++++++----------- transformer_engine/common/common.h | 1 + .../common/gemm/cublaslt_grouped_gemm.cu | 10 +- ...cast_col_hadamard_transform_cast_fusion.cu | 4 +- .../group_hadamard_transform_cast_fusion.cu | 3 +- ...cast_col_hadamard_transform_cast_fusion.cu | 4 +- .../hadamard_transform_cast_fusion.cu | 3 +- ...cast_col_hadamard_transform_cast_fusion.cu | 4 +- transformer_engine/common/recipe/nvfp4.cu | 23 +- .../common/transformer_engine.cpp | 5 + 13 files changed, 154 insertions(+), 287 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 4f8c019c88..07b402d51f 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -105,8 +105,8 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || - output_tensor->get_nvfp4_scale_max() == - static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + static_cast(output_tensor->get_nvfp4_scale_max()) == + typeToMax(output_tensor->scale_inv.dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -289,8 +289,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || - output_tensor->get_nvfp4_scale_max() == - static_cast(nvfp4::core::scale_max(output_tensor->scale_inv.dtype)), + static_cast(output_tensor->get_nvfp4_scale_max()) == + typeToMax(output_tensor->scale_inv.dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -455,8 +455,8 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou if (!nvfp4_use_4over6) { for (const auto *output_tensor : output_tensors) { NVTE_CHECK( - output_tensor->get_nvfp4_scale_max() == - static_cast(nvfp4::core::scale_max(output_tensors[0]->scale_inv.dtype)), + static_cast(output_tensor->get_nvfp4_scale_max()) == + typeToMax(output_tensors[0]->scale_inv.dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); } } diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index b89dd755a9..4f15a4840c 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -51,80 +51,6 @@ namespace core { #if FP4_TYPE_SUPPORTED using namespace ptx; -// Scale-format-specific behavior belongs here rather than in individual kernels. -template -struct NVFP4ScaleTraits { - static constexpr bool is_supported = false; - static constexpr bool supports_fp16_error_path = false; - static constexpr float expected_max = 0.0f; - static constexpr float headroom_max = 0.0f; -}; - -template <> -struct NVFP4ScaleTraits { - // E4M3 scales fit in FP16 and can use the packed E4M3-to-FP16 PTX fast - // path. UE5M3 scales can exceed the FP16 range, so they retain the generic - // FP32 error path. - static constexpr bool is_supported = true; - static constexpr bool supports_fp16_error_path = true; - static constexpr float expected_max = 448.0f; - static constexpr float headroom_max = 256.0f; -}; - -#if CUDA_VERSION >= 13040 -template <> -struct NVFP4ScaleTraits { - static constexpr bool is_supported = true; - static constexpr bool supports_fp16_error_path = false; - static constexpr float expected_max = 114688.0f; - static constexpr float headroom_max = 65536.0f; -}; -#endif - -// Return the effective maximum used to derive the global NVFP4 encode scale. -// SCALE_TYPE_MAX is the resolved maximum for ScaleType (e.g., 448 for E4M3 -// or 114688 for UE5M3). The headroom maximum keeps the 1.5x map-to-4 scale -// used by 4over6 within the scale format's representable range. -template (NVFP4ScaleTraits::expected_max)> -__host__ __device__ constexpr float scale_max() { - using ScaleTraits = NVFP4ScaleTraits; - static_assert(ScaleTraits::is_supported, "Unsupported NVFP4 scale type."); - if constexpr (ScaleTraits::is_supported) { - static_assert(detail::TypeExtrema::max == ScaleTraits::expected_max, - "Unexpected NVFP4 scale type maximum."); - static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || - SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), - "Unsupported NVFP4 scale type maximum."); - static_assert(ScaleTraits::headroom_max * 1.5f <= ScaleTraits::expected_max, - "NVFP4 4over6 scale headroom exceeds scale type maximum."); - return static_cast(SCALE_TYPE_MAX); - } else { - return 0.0f; - } -} - -// Return the full-range maximum for a runtime scale dtype. -inline float scale_max(const DType scale_dtype) { - float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, - result = scale_max();) - return result; -} - -// Return and validate a user-provided maximum for a runtime scale dtype. -inline float scale_max(const DType scale_dtype, const int scale_type_max) { - float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { - using ScaleTraits = NVFP4ScaleTraits; - NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - result = static_cast(scale_type_max); - }) - return result; -} - template __device__ __forceinline__ ScaleType compute_decoding_scaling_factor(const float block_amax, const float global_encode_scale) { diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index a014244b9b..1ba59f8a5a 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -66,7 +66,7 @@ __global__ void __launch_bounds__(512) value.vec = input_vectorized[my_index]; ScaleType scale = scales[my_scale_index]; constexpr float fp4_max = detail::TypeExtrema::max; - constexpr float unit_global_scale_amax = fp4_max * core::scale_max(); + constexpr float unit_global_scale_amax = fp4_max * SCALE_TYPE_MAX; float amax = unit_global_scale_amax; if (tensor_amax != nullptr) { amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; @@ -139,21 +139,28 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) "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_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { - using ScaleTraits = core::NVFP4ScaleTraits; - if (e4m3_max == static_cast(ScaleTraits::expected_max)) { - launch_dequantize(ScaleTraits::expected_max)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); - } else { - NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - launch_dequantize(ScaleTraits::headroom_max)>( + + if (static_cast(e4m3_max) != typeToMax(scale_dtype)) { + NVTE_CHECK(scale_dtype == DType::kFloat8E4M3, + "NVFP4 dequantization with non-default scale max " + "is only supported with FP8E4M3 scales (found ", + to_string(scale_dtype), ")."); + NVTE_CHECK(e4m3_max == 256, + "NVFP4 dequantization with non-default scale max " + "is only supported with e4m3_max=256 (found ", + e4m3_max, ")."); + launch_dequantize( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + NVTE_CHECK_CUDA(cudaGetLastError()); + } else { + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, + launch_dequantize(TypeInfo::max_finite_value)>( input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, num_scale_tiles_X, stream); - } - }) - NVTE_CHECK_CUDA(cudaGetLastError()); + ); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); + } #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index d5a220d8ef..1dfd35af76 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -54,6 +54,16 @@ namespace nvfp4 { } \ } +#define TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(E4M3_MAX_VALUE, E4M3_MAX_CONST, ...) \ + if ((E4M3_MAX_VALUE) == 256) { \ + constexpr int E4M3_MAX_CONST = 256; \ + { __VA_ARGS__ } \ + } else { \ + NVTE_CHECK((E4M3_MAX_VALUE) == 448, "Unsupported NVFP4 E4M3 max."); \ + constexpr int E4M3_MAX_CONST = 448; \ + { __VA_ARGS__ } \ + } + namespace quantize_4over6_kernel { constexpr int kThreads = 128; @@ -71,6 +81,8 @@ constexpr int kPackedWordsPerGroup = 2; static_assert(kTileRows == kPipelineStages * kStageRows); static_assert(kStageRows % kGroupSize == 0); +using nvfp4_scale_t = fp8e4m3; + template struct Config { static constexpr NVTENVFP44Over6Mode mode = kMode; @@ -87,10 +99,9 @@ struct CandidatePair { Candidate map6; }; -template struct ScalePair { - ScaleType map4; - ScaleType map6; + nvfp4_scale_t map4; + nvfp4_scale_t map6; float inv_map4; float inv_map6; float global_encode_scale; @@ -113,24 +124,19 @@ __device__ __forceinline__ float compute_error_rn(const float diff) { } } -template -__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, - const float global_amax) { - using ScaleTraits = core::NVFP4ScaleTraits; - static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || - SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), - "Unsupported NVFP4 scale type maximum."); +template +__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, + const float global_amax) { + static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = detail::TypeExtrema::max; - constexpr int encode_scale_max = static_cast(core::scale_max()); + constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f constexpr float expand_to_map4 = 1.5f; - const float S_enc = - core::compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); const float base = block_amax / fp4_max * S_enc; - ScalePair scales; - scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); - scales.map6 = static_cast(fminf(base, fp8_max)); + ScalePair scales; + scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); + scales.map6 = static_cast(fminf(base, fp8_max)); const float S_dec = 1.0f / S_enc; scales.inv_map4 = @@ -183,12 +189,12 @@ __device__ __forceinline__ void load_col_group(const IType *tile, const int row_ } } -template +template __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, const float sf, const float global_amax, float *err) { constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = core::scale_max(); + constexpr float fp8_max = static_cast(E4M3_MAX); constexpr float err_denom = fp4_max * fp8_max; const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; const float dequant = __half2float(__ushort_as_half(half_bits)); @@ -197,19 +203,11 @@ __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_ *err = __fadd_rn(*err, compute_error_rn(diff)); } -template -__device__ __forceinline__ uint8_t fp8_bits(const ScaleType sf) { +__device__ __forceinline__ uint8_t fp8_bits(const nvfp4_scale_t sf) { return *reinterpret_cast(&sf); } -template -__device__ __forceinline__ FP16ErrorScalePair -compute_fp16_error_scales(const ScalePair &scales) { - // This fast error path interprets the packed scale bits as E4M3. UE5M3 - // deliberately does not enable supports_fp16_error_path and instead uses - // the scale-format-independent float error path in - // cvt_fp32_to_fp4_8x_with_error. - static_assert(core::NVFP4ScaleTraits::supports_fp16_error_path); +__device__ __forceinline__ FP16ErrorScalePair compute_fp16_error_scales(const ScalePair &scales) { FP16ErrorScalePair result; const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | (static_cast(fp8_bits(scales.map6)) << 8); @@ -261,9 +259,9 @@ __device__ __forceinline__ void accumulate_fp16_scaled_error_pair(const uint32_t *err = __fadd_rn(*err, compute_error_rn(diff1)); } -template +template __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( - const float (&x)[8], const float block_scale_inverse, const ScaleType sf, + const float (&x)[8], const float block_scale_inverse, const nvfp4_scale_t sf, const uint32_t fp16_error_scale, const float global_amax, const float global_encode_scale, float *err) { uint32_t out = 0; @@ -272,11 +270,6 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( uint32_t out_dequant_3 = 0; uint32_t out_dequant_4 = 0; - // ScaleType is not consumed by this PTX. block_scale_inverse applies the - // selected E4M3 or UE5M3 block scale while forming the FP32 operands. These - // instructions only convert the scaled candidates to FP4 E2M1 and back to - // FP16 for error evaluation, so their encoding is identical for both scale - // storage types. constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; if constexpr (is_blackwell) { asm volatile( @@ -304,8 +297,7 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( "Try recompiling with sm_XXXa instead of sm_XXX."); } - if constexpr (Cfg::err_use_fast_math && - core::NVFP4ScaleTraits::supports_fp16_error_path) { + if constexpr (Cfg::err_use_fast_math) { accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, global_encode_scale, err); accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, @@ -316,48 +308,39 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( global_encode_scale, err); } else { const float sf_float = static_cast(sf); - accumulate_dequant_error(out_dequant_1, x[0], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_1, x[1], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_2, x[2], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_2, x[3], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_3, x[4], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_3, x[5], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_4, x[6], sf_float, - global_amax, err); - accumulate_dequant_error(out_dequant_4, x[7], sf_float, - global_amax, err); + accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); } return out; } -template +template __device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], - const ScalePair &scales, + const ScalePair &scales, const float global_amax) { CandidatePair candidates; candidates.map4.err = 0.0f; candidates.map6.err = 0.0f; FP16ErrorScalePair fp16_error_scales{}; - if constexpr (Cfg::err_use_fast_math && - core::NVFP4ScaleTraits::supports_fp16_error_path) { + if constexpr (Cfg::err_use_fast_math) { fp16_error_scales = compute_fp16_error_scales(scales); } - candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); - candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); return candidates; @@ -399,9 +382,8 @@ __device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &ca return candidates.map6.packed; } -template -__device__ __forceinline__ ScaleType select_scale(const ScalePair &scales, - const bool pick_map4) { +__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, + const bool pick_map4) { if (pick_map4) { return scales.map4; } @@ -469,9 +451,9 @@ __device__ void load_stage_to_shared_async(const IType *input, IType *tile, cons } } -template -__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, ScaleType *scales, +template +__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, const float *amax, const size_t rows, const size_t cols, const size_t stage_row, const size_t tile_col, const size_t scale_stride) { @@ -496,21 +478,13 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, Sca block_amax = reduce_group_max_16(group_amax); } - float global_amax = - core::scale_max() * detail::TypeExtrema::max; - if (amax != nullptr) { - global_amax = amax[0]; - } + float global_amax = amax[0]; if constexpr (ROW_SCALED_NVFP4) { - if (amax != nullptr) { - global_amax = amax[global_row]; - } + global_amax = amax[global_row]; } - const ScalePair scale_pair = - compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = - make_candidates(x0, x1, scale_pair, global_amax); + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -520,7 +494,7 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, Sca } const bool pick_map4 = err_map4 < err_map6; - const ScaleType selected_scale = select_scale(scale_pair, pick_map4); + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_col_group = global_col / kGroupSize; @@ -529,12 +503,11 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, Sca } } -template -__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, ScaleType *scales_t, - const float *amax, const size_t rows, const size_t cols, - const size_t stage_row, const size_t tile_col, - const size_t scale_stride_t) { +template +__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales_t, const float *amax, + const size_t rows, const size_t cols, const size_t stage_row, + const size_t tile_col, const size_t scale_stride_t) { constexpr int groups = kStageRowGroups * kTileCols; for (int group = threadIdx.x; group < groups; group += blockDim.x) { const int local_row_group = group / kTileCols; @@ -556,13 +529,9 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, S block_amax = reduce_group_max_16(group_amax); } - const float global_amax = amax == nullptr ? core::scale_max() * - detail::TypeExtrema::max - : amax[0]; - const ScalePair scale_pair = - compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = - make_candidates(x0, x1, scale_pair, global_amax); + const float global_amax = amax[0]; + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -572,7 +541,7 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, S } const bool pick_map4 = err_map4 < err_map6; - const ScaleType selected_scale = select_scale(scale_pair, pick_map4); + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_row_group = global_row / kGroupSize; @@ -582,14 +551,13 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, S } template + bool ROW_SCALED_NVFP4, typename Cfg, int E4M3_MAX, typename IType> __global__ void __launch_bounds__(kThreads) quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, - ScaleType *scales, ScaleType *scales_t, const float *amax_rowwise, - const float *amax_colwise, const size_t rows, const size_t cols, - const size_t scale_stride, const size_t scale_stride_t, - const float *noop) { + nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, + const float *amax_rowwise, const float *amax_colwise, const size_t rows, + const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const float *noop) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -624,7 +592,7 @@ __global__ void __launch_bounds__(kThreads) IType *stage_tile = stage_tiles[stage]; if constexpr (RETURN_IDENTITY) { - quantize_stage_rowwise( + quantize_stage_rowwise( stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); } @@ -633,7 +601,7 @@ __global__ void __launch_bounds__(kThreads) if (columnwise_amax == nullptr) { columnwise_amax = amax_rowwise; } - quantize_stage_colwise( + quantize_stage_colwise( stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, scale_stride_t); } @@ -648,8 +616,7 @@ __global__ void __launch_bounds__(kThreads) #endif } -template +template void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { const size_t rows = input.flat_first_dim(); @@ -661,8 +628,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out const auto *input_ptr = reinterpret_cast(input.data.dptr); auto *output_ptr = reinterpret_cast(output->data.dptr); auto *output_t_ptr = reinterpret_cast(output->columnwise_data.dptr); - auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); - auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); + auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); const auto *amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); const auto *amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); const auto *noop_ptr = reinterpret_cast(noop->data.dptr); @@ -677,9 +644,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { - auto kernel = - quantize_4over6_kernel; + auto kernel = quantize_4over6_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, @@ -693,9 +659,9 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out #endif // FP4_TYPE_SUPPORTED -template -void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { +template +void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace quantize_4over6_kernel; @@ -721,69 +687,44 @@ void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *outpu "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(), + NVTE_CHECK(!output->row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); NVTE_CHECK(!use_2d_quantization || output->has_data(), "NVFP4 4over6 2D quantization requires rowwise output."); if (output->has_data()) { NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); - NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); + NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output data must have FP4 type."); + NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E4M3, + "Output scales must have FP8E4M3 type."); + NVTE_CHECK(output->amax.dptr != nullptr, "Rowwise amax tensor must be allocated."); } if (output->has_columnwise_data()) { NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Transposed scaling tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); + NVTE_CHECK(output->columnwise_scale_inv.dtype == DType::kFloat8E4M3, + "Output scales must have FP8E4M3 type."); + NVTE_CHECK(output->columnwise_amax.dptr != nullptr || output->amax.dptr != nullptr, + "NVFP4 4over6 columnwise quantization requires columnwise amax or rowwise amax."); } - using ScaleTraits = core::NVFP4ScaleTraits; - const int scale_type_max = output->get_nvfp4_scale_max(); - NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - TRANSFORMER_ENGINE_SWITCH_CONDITION( - scale_type_max == static_cast(ScaleTraits::headroom_max), USE_SCALE_HEADROOM, { - constexpr int SCALE_TYPE_MAX = static_cast( - USE_SCALE_HEADROOM ? ScaleTraits::headroom_max : ScaleTraits::expected_max); - TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( - quant_config->nvfp4_4over6_mode, MODE, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { - using Cfg = quantize_4over6_kernel::Config; - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - quantize_4over6_kernel::launch_quantize_4over6< - use_2d_quantization, Cfg, ScaleType, SCALE_TYPE_MAX, IType>( - input, noop, output, stream);); - });); - }) - NVTE_CHECK_CUDA(cudaGetLastError()); -#else - NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); -#endif // FP4_TYPE_SUPPORTED -} - -template -void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED - const bool return_rowwise = output->has_data(); - const bool return_transpose = output->has_columnwise_data(); - NVTE_CHECK(return_rowwise || return_transpose, - "NVFP4 4over6 output tensor must have rowwise or columnwise data."); - const DType scale_dtype = - return_rowwise ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; - if (return_rowwise && return_transpose) { - NVTE_CHECK(output->scale_inv.dtype == output->columnwise_scale_inv.dtype, - "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", - to_string(output->scale_inv.dtype), " and ", - to_string(output->columnwise_scale_inv.dtype), ")."); - } + TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH( + output->nvfp4_e4m3_max, E4M3_MAX, + TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( + quant_config->nvfp4_4over6_mode, MODE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { + using Cfg = quantize_4over6_kernel::Config; + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + quantize_4over6_kernel::launch_quantize_4over6( + input, noop, output, stream);); + }););); - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, - quantize_4over6_impl( - input, noop, output, quant_config, stream);) + NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 3300a82b46..3ed99a19ff 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1248,6 +1248,7 @@ inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { size_t typeToSize(const DType type); size_t typeToNumBits(const DType type); +float typeToMax(const DType type); void CheckNoopTensor(const Tensor &t, std::string_view name); void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_inv_shapes = true); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 4f28569ad6..75d456f809 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -16,7 +16,6 @@ #include #include "../cast/mxfp8/swizzle.cuh" -#include "../cast/nvfp4/core_nvfp4.cuh" #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/handle_manager.h" @@ -1567,12 +1566,9 @@ inline void launch_grouped_gemm_setup( float a_unit_global_scale_amax = 1.0f; float b_unit_global_scale_amax = 1.0f; if (needs_nvfp4_alpha) { - constexpr float kFP4Max = - transformer_engine::detail::TypeExtrema::max; - a_unit_global_scale_amax = - transformer_engine::dispatch::nvfp4::core::scale_max(A_sel.scale_inv_dtype) * kFP4Max; - b_unit_global_scale_amax = - transformer_engine::dispatch::nvfp4::core::scale_max(B_sel.scale_inv_dtype) * kFP4Max; + const float kFP4Max = typeToMax(transformer_engine::DType::kFloat4E2M1); + a_unit_global_scale_amax = typeToMax(A_sel.scale_inv_dtype) * kFP4Max; + b_unit_global_scale_amax = typeToMax(B_sel.scale_inv_dtype) * kFP4Max; } setup_grouped_gemm_kernel<<>>( 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 c4ff0f445b..a52ebcb3eb 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 @@ -695,7 +695,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) { shared_storage.global_d_amax[g] = amax_colwise == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : __ldg(amax_colwise + g); } @@ -947,7 +947,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) { shared_storage.global_a_amax[g] = amax_rowwise == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : __ldg(amax_rowwise + g); } // RNG for stochastic rounding 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 f977c3651e..49a9a243d7 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 @@ -496,8 +496,7 @@ __global__ static void group_rht_gemm_device( Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); - constexpr float kUnitGlobalScaleAmax = - dispatch::nvfp4::core::scale_max() * TypeExtrema::max; + constexpr float kUnitGlobalScaleAmax = TypeExtrema::max * TypeExtrema::max; float global_amax_val = global_amax_ptr == nullptr ? kUnitGlobalScaleAmax : *global_amax_ptr; float global_encode_scale = dispatch::nvfp4::core::compute_global_encode_scaling_factor_FP4(global_amax_val); 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 39a49da36b..ba22bd42f4 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 @@ -684,7 +684,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( const auto *amax_ptr = reinterpret_cast(args.global_d_amax_list[g]); shared_storage.global_d_amax[g] = amax_ptr == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : __ldg(amax_ptr); } @@ -926,7 +926,7 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( const auto *amax_ptr = reinterpret_cast(args.global_a_amax_list[g]); shared_storage.global_a_amax[g] = amax_ptr == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : __ldg(amax_ptr); } // RNG for stochastic rounding 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 a0d781b104..5b4598c0ee 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -402,8 +402,7 @@ 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) { - constexpr float kUnitGlobalScaleAmax = - dispatch::nvfp4::core::scale_max() * TypeExtrema::max; + constexpr float kUnitGlobalScaleAmax = TypeExtrema::max * TypeExtrema::max; const float global_amax_val = global_amax == nullptr ? kUnitGlobalScaleAmax : *global_amax; static constexpr int FragmentSize = 256 / sizeof_bits_v; 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 9c06f62eb6..5ccfbf0c5f 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 @@ -654,7 +654,7 @@ __global__ static void row_col_rht_gemm_device( float const c_global_amax_val = c_global_amax == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : *c_global_amax; auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{}); auto bulk_tmem_epilogue_layout = make_layout( @@ -862,7 +862,7 @@ __global__ static void row_col_rht_gemm_device( using S2RVectorType = uint128_t; float const a_global_amax_val = a_global_amax == nullptr - ? dispatch::nvfp4::core::scale_max() * TypeExtrema::max + ? TypeExtrema::max * TypeExtrema::max : *a_global_amax; int global_thread_idx = threadIdx.x; int local_thread_idx = global_thread_idx % 256; diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 7047c24e68..d180ea21dc 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -10,7 +10,6 @@ #include #include -#include "../cast/nvfp4/core_nvfp4.cuh" #include "../common.h" #include "../util/ptx.cuh" #include "../utils.cuh" @@ -646,8 +645,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 = transformer_engine::detail::TypeExtrema::max; - constexpr float scale_max = dispatch::nvfp4::core::scale_max(); + constexpr float fp4_max = transformer_engine::TypeInfo::max_finite_value; + constexpr float scale_max = transformer_engine::TypeInfo::max_finite_value; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN @@ -676,8 +675,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 = transformer_engine::detail::TypeExtrema::max; - constexpr float scale_max = dispatch::nvfp4::core::scale_max(); + constexpr float fp4_max = transformer_engine::TypeInfo::max_finite_value; + constexpr float scale_max = transformer_engine::TypeInfo::max_finite_value; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; // FLT_MIN @@ -773,8 +772,8 @@ __global__ void nvfp4_fused_scale_kernel( const size_t tile_row = out_row / block_len; // Compute the scale value - constexpr float fp4_max = transformer_engine::detail::TypeExtrema::max; - constexpr float scale_max = dispatch::nvfp4::core::scale_max(); + constexpr float fp4_max = transformer_engine::TypeInfo::max_finite_value; + constexpr float scale_max = transformer_engine::TypeInfo::max_finite_value; constexpr float flt_max = 3.402823466e+38f; constexpr float tiny = 1.17549435e-38f; @@ -945,14 +944,8 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void *amax_A_ptr = use_rowwise_amax_A ? tA->amax.dptr : tA->columnwise_amax.dptr; void *amax_B_ptr = use_rowwise_amax_B ? tB->amax.dptr : tB->columnwise_amax.dptr; void *alpha_ptr = tOut->data.dptr; - const DType scale_dtype_A = - use_rowwise_amax_A ? tA->scale_inv.dtype : tA->columnwise_scale_inv.dtype; - const DType scale_dtype_B = - use_rowwise_amax_B ? tB->scale_inv.dtype : tB->columnwise_scale_inv.dtype; - const float scale_max_A = - dispatch::nvfp4::core::scale_max(scale_dtype_A, tA->get_nvfp4_scale_max()); - const float scale_max_B = - dispatch::nvfp4::core::scale_max(scale_dtype_B, tB->get_nvfp4_scale_max()); + const float scale_max_A = tA->get_nvfp4_scale_max(); + const float scale_max_B = tB->get_nvfp4_scale_max(); NVTE_CHECK(alpha_ptr != nullptr, "alpha_ptr is null"); diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 7081c63bbf..74a482e5e4 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -38,6 +38,11 @@ size_t typeToSize(const DType type) { return typeToNumBits(type) / 8; } +float typeToMax(const DType type) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T, + return TypeInfo::max_finite_value;); // NOLINT(*) +} + std::string to_string(const NVTEScalingMode &mode) { switch (mode) { case NVTE_DELAYED_TENSOR_SCALING: From 868555694dadae711e1a00f466055d18fe0a5f78 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 18 Aug 2026 02:28:43 +0000 Subject: [PATCH 11/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/dispatch/quantize.cuh | 10 ++++------ .../common/cast/nvfp4/dequantize_nvfp4.cuh | 15 +++++++-------- .../common/cast/nvfp4/quantize_4over6_nvfp4.cuh | 5 +++-- ...row_cast_col_hadamard_transform_cast_fusion.cu | 14 ++++++-------- ...row_cast_col_hadamard_transform_cast_fusion.cu | 14 ++++++-------- 5 files changed, 26 insertions(+), 32 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 07b402d51f..d08bd07ef5 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -104,9 +104,8 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - NVTE_CHECK(nvfp4_use_4over6 || - static_cast(output_tensor->get_nvfp4_scale_max()) == - typeToMax(output_tensor->scale_inv.dtype), + NVTE_CHECK(nvfp4_use_4over6 || static_cast(output_tensor->get_nvfp4_scale_max()) == + typeToMax(output_tensor->scale_inv.dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -288,9 +287,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens auto dtype = grad_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - NVTE_CHECK(nvfp4_use_4over6 || - static_cast(output_tensor->get_nvfp4_scale_max()) == - typeToMax(output_tensor->scale_inv.dtype), + NVTE_CHECK(nvfp4_use_4over6 || static_cast(output_tensor->get_nvfp4_scale_max()) == + typeToMax(output_tensor->scale_inv.dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 1ba59f8a5a..bc553d45d0 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -149,16 +149,15 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) "NVFP4 dequantization with non-default scale max " "is only supported with e4m3_max=256 (found ", e4m3_max, ")."); - launch_dequantize( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); + launch_dequantize(input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, + Mread, blocks, threads, num_scale_tiles_X, stream); NVTE_CHECK_CUDA(cudaGetLastError()); } else { - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, - launch_dequantize(TypeInfo::max_finite_value)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); - ); // NOLINT(*) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + launch_dequantize(TypeInfo::max_finite_value)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream);); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } #else diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 1dfd35af76..9a2af69ca1 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -131,7 +131,8 @@ __device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f constexpr float expand_to_map4 = 1.5f; - const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = + core::compute_global_encode_scaling_factor_FP4(global_amax); const float base = block_amax / fp4_max * S_enc; ScalePair scales; @@ -687,7 +688,7 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, "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(), + NVTE_CHECK(!output->row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); NVTE_CHECK(!use_2d_quantization || output->has_data(), "NVFP4 4over6 2D quantization requires rowwise output."); 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 a52ebcb3eb..906beeb8f1 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 @@ -693,10 +693,9 @@ __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] = - amax_colwise == nullptr - ? TypeExtrema::max * TypeExtrema::max - : __ldg(amax_colwise + g); + shared_storage.global_d_amax[g] = amax_colwise == nullptr + ? TypeExtrema::max * TypeExtrema::max + : __ldg(amax_colwise + g); } size_t rng_seed = 0; @@ -945,10 +944,9 @@ __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] = - amax_rowwise == nullptr - ? TypeExtrema::max * TypeExtrema::max - : __ldg(amax_rowwise + g); + shared_storage.global_a_amax[g] = amax_rowwise == nullptr + ? TypeExtrema::max * TypeExtrema::max + : __ldg(amax_rowwise + g); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { 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 ba22bd42f4..c336810aa8 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,10 +682,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( 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] = - amax_ptr == nullptr - ? TypeExtrema::max * TypeExtrema::max - : __ldg(amax_ptr); + shared_storage.global_d_amax[g] = amax_ptr == nullptr + ? TypeExtrema::max * TypeExtrema::max + : __ldg(amax_ptr); } size_t rng_seed = 0; @@ -924,10 +923,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( 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] = - amax_ptr == nullptr - ? TypeExtrema::max * TypeExtrema::max - : __ldg(amax_ptr); + shared_storage.global_a_amax[g] = amax_ptr == nullptr + ? TypeExtrema::max * TypeExtrema::max + : __ldg(amax_ptr); } // RNG for stochastic rounding if constexpr (kEnableStochasticRounding) { From 803dd24ece66ecd8d61ef296db270aed66c4d8be Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 18 Aug 2026 02:30:38 +0000 Subject: [PATCH 12/54] Fix compile error Signed-off-by: Tim Moon --- transformer_engine/common/recipe/nvfp4.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index d180ea21dc..f9d661e2e1 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -10,6 +10,7 @@ #include #include +#include "../cast/nvfp4/core_nvfp4.cuh" #include "../common.h" #include "../util/ptx.cuh" #include "../utils.cuh" From 27f29d07555adca13c6b75ee9010a54293a720c4 Mon Sep 17 00:00:00 2001 From: tdophung Date: Tue, 18 Aug 2026 10:02:33 -0700 Subject: [PATCH 13/54] Fix NVFP4 scale dtype ABI Signed-off-by: Tim Moon --- .../common/include/transformer_engine/recipe.h | 18 +++++------------- transformer_engine/pytorch/constants.py | 3 ++- transformer_engine/pytorch/csrc/extensions.h | 3 ++- .../pytorch/csrc/extensions/pybind.cpp | 4 ++-- 4 files changed, 11 insertions(+), 17 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index b98b87c5ba..0f7c3ae307 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -14,10 +14,7 @@ #include "transformer_engine.h" #ifdef __cplusplus -#define NVTE_NVFP4_SCALE_DTYPE_DEFAULT = kNVTEFloat8E4M3 extern "C" { -#else -#define NVTE_NVFP4_SCALE_DTYPE_DEFAULT #endif /*! \brief Update FP8 scaling factors with delayed scaling recipe. @@ -382,8 +379,7 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, cudaStream_t stream, - const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); + size_t block_len, cudaStream_t stream, const NVTEDType scale_dtype); /*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. * @@ -400,8 +396,7 @@ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTE */ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream, - const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); + cudaStream_t stream, const NVTEDType scale_dtype); /*! \brief Compute per-block decode scale from block amax and global amax. * @@ -419,7 +414,7 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s */ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, const NVTETensor global_amax, cudaStream_t stream, - const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); + const NVTEDType scale_dtype); /*! \brief Fused kernel for NVFP4 scale computation. * @@ -446,7 +441,7 @@ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, cudaStream_t stream, - const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); + const NVTEDType scale_dtype); /*! \brief Compute global encode scale from global amax. * @@ -459,13 +454,10 @@ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). */ void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - cudaStream_t stream, - const NVTEDType scale_dtype NVTE_NVFP4_SCALE_DTYPE_DEFAULT); + cudaStream_t stream, const NVTEDType scale_dtype); #ifdef __cplusplus } // extern "C" #endif -#undef NVTE_NVFP4_SCALE_DTYPE_DEFAULT - #endif // TRANSFORMER_ENGINE_RECIPE_H_ diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index ec54189613..c0e549dde1 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Enums for e2e transformer""" + import enum from types import SimpleNamespace from typing import Union @@ -32,7 +33,7 @@ class DType(enum.IntEnum): mantissa bits. * ``kFloat4E2M1`` -- 4-bit floating point with 2 exponent and 1 mantissa bits. - * ``kFloat8UE4M3`` -- 8-bit unsigned floating point with 5 exponent and 3 + * ``kFloat8UE5M3`` -- 8-bit unsigned floating point with 5 exponent and 3 mantissa bits. The enum mirrors the backend ``transformer_engine_torch.DType`` (pybind11) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index cdedc0f06c..4412c1c5fe 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -215,7 +215,8 @@ void nvfp4_multi_tensor_compute_partial_amax( std::vector w_list, std::vector start_offset_list, int64_t block_len); void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len); + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + DType scale_dtype = DType::kFloat8E4M3); void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, const DType scale_dtype = DType::kFloat8E4M3); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index b7cee9d0e1..7b64e83212 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -406,9 +406,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("input"), py::arg("output"), py::arg("M_tiles"), py::arg("K_tiles"), py::call_guard()); m.def("nvfp4_expand_scale_to_fp8", &transformer_engine::pytorch::nvfp4_expand_scale_to_fp8, - "Expand tile-level scales to row-level scales and convert to FP8 E4M3", py::arg("input"), + "Expand tile-level scales to row-level scales and convert to FP8", py::arg("input"), py::arg("output"), py::arg("tile_rows"), py::arg("tile_cols"), py::arg("rows_padded"), - py::arg("block_len"), py::call_guard()); + py::arg("block_len"), py::arg("scale_dtype") = transformer_engine::DType::kFloat8E4M3); m.def("nvfp4_compute_per_block_scale", &transformer_engine::pytorch::nvfp4_compute_per_block_scale, "Compute per-block decode scale from block amax and global amax", py::arg("block_amax"), From c9dad3e7646bc8c0d0c8448508e6c281276b3d65 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 19 Aug 2026 20:01:45 +0000 Subject: [PATCH 14/54] remove redundant output alloc Signed-off-by: Kaining Zhong --- transformer_engine/pytorch/cpp_extensions/gemm.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 7a2fb538bb..ad4e22c663 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -651,8 +651,6 @@ def general_cuDNN_MX_gemm( # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). bias = bias.contiguous().as_strided((M, 1), (1, M)) - # Prepare for output - out = validate_or_alloc_output(out, out_shape, out_dtype, device) if N_padded != N: # The kernel writes N_padded rows, so it cannot target `out` directly. d_buf = torch.empty((N_padded, M), dtype=out_dtype, device=device) From 2bf42eb91bf98621f285e552de835f3589f33e51 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 19 Aug 2026 20:11:34 +0000 Subject: [PATCH 15/54] no need to pad N now Signed-off-by: Kaining Zhong --- .../pytorch/cpp_extensions/gemm.py | 33 ++----------------- 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index ad4e22c663..7f36828b15 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -586,23 +586,6 @@ def general_cuDNN_MX_gemm( ), "cuDNN GEMM currently does not support accumulation for this operation." assert beta in (0.0, None), "beta must be zero or None if not accumulate" - # cuDNN's grouped quant kernel requires M to be divisible by 256 so we need to pad it - N_padded = ceil_div(N, 256) * 256 - if N_padded != N: - src = dataB.reshape(N, K // 2) - buf = src.new_zeros((N_padded, K // 2)) - buf[:N].copy_(src) - dataB = buf - - # Swizzled scales are blocked by 128 rows: - # (1, ceil(M/128), k_sf_tiles, 32, 4, 4) - per_block = ceil_div(K, 4 * NVFP4_BLOCK_SCALING_SIZE) * 32 * 4 * 4 - n_blk, n_blk_padded = ceil_div(N, 128), ceil_div(N_padded, 128) - src_sf = sfB.reshape(-1)[: n_blk * per_block].reshape(n_blk, per_block) - buf_sf = src_sf.new_zeros((n_blk_padded, per_block)) - buf_sf[:n_blk].copy_(src_sf) - sfB = buf_sf - # cuDNN's own operand names are the other way round: its "a" is the (M, K) # activation-like operand (TE's B) and its "b" is the (N, K) weight-like one # (TE's A). @@ -611,7 +594,7 @@ def general_cuDNN_MX_gemm( sfB, data_dtype=torch.float4_e2m1fn_x2, scale_dtype=torch.float8_e4m3fn, # e5m3 rides as e4m3; torch has no ue5m3 - valid_M_or_N=N_padded, + valid_M_or_N=N, k_logical=K, L=1, sf_swizzled=True, # ensured above @@ -651,12 +634,7 @@ def general_cuDNN_MX_gemm( # cuDNN checks the stride literally, so (1, N) rather than reshape's (1, 1). bias = bias.contiguous().as_strided((M, 1), (1, M)) - if N_padded != N: - # The kernel writes N_padded rows, so it cannot target `out` directly. - d_buf = torch.empty((N_padded, M), dtype=out_dtype, device=device) - d_tensor = d_buf.as_strided((N_padded, M, 1), (M, 1, N_padded * M)) - else: - d_tensor = out.view(N, M).as_strided((N, M, 1), (M, 1, M * N)) + d_tensor = out.view(N, M).as_strided((N, M, 1), (M, 1, M * N)) gemm_kwargs = { "a_tensor": cudnn_a, @@ -664,7 +642,7 @@ def general_cuDNN_MX_gemm( "b_tensor": cudnn_b, "sfb_tensor": cudnn_sfb, # One group, so the only padded end offset is the full row count. - "padded_offsets": torch.tensor([N_padded], dtype=torch.int32, device=device), + "padded_offsets": torch.tensor([N], dtype=torch.int32, device=device), "alpha_tensor": alpha_tensor, "bias_tensor": bias, "norm_const_tensor": None, # must be None for FP4 inputs @@ -680,11 +658,6 @@ def general_cuDNN_MX_gemm( } grouped_gemm_quant_kernel()(**gemm_kwargs) - if N_padded != N: - # Drop the zero-padded rows. Safe to overwrite rather than accumulate: - # this path asserts accumulate is False above. - out.view(N, M).copy_(d_buf[:N]) - # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). return out, None, None, None From b6913b289449f0729f8860ee477982ed2ceb56b5 Mon Sep 17 00:00:00 2001 From: Kaining Zhong Date: Wed, 19 Aug 2026 21:16:41 +0000 Subject: [PATCH 16/54] fix linting errors Signed-off-by: Kaining Zhong --- .../pytorch/cpp_extensions/gemm.py | 4 +-- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/ops/fused/grouped_mlp.py | 27 +++++++++++++------ .../pytorch/tensor/nvfp4_tensor.py | 2 ++ 4 files changed, 24 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 7f36828b15..a2f4041a68 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -235,7 +235,6 @@ def _cuDNN_wgrad_gemm( out: torch.Tensor, accumulate: bool, alpha: Optional[float] = None, - beta: Optional[float] = None, bias: Optional[torch.Tensor] = None, ) -> Iterable[Optional[torch.Tensor]]: """Compute dw = dy^T @ x with cuDNN's purpose-built grouped wgrad kernel.""" @@ -463,10 +462,11 @@ def general_cuDNN_MX_gemm( and isinstance(B, NVFP4TensorStorage) and A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 - ), f"cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." + ), "cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." + assert use_split_accumulator is False, "cuDNN GEMM currently does not support split accumulators." # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for # raising accumulator precision, and the cuDNN kernel always accumulates in diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 92cae29d96..8826284e8d 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -59,7 +59,7 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) -from ..constants import DType, GemmParallelModes, dist_group_type +from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..triton.grouped_dbias_dscales import compute_grouped_dbias diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 1abc423ff4..4e5caaf3a6 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -358,7 +358,6 @@ def _single_quantized_tensor_from_grouped( with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, ) - # TODO(kainingz): claude told me this doesn't pass the required param scale_dtype. Should check this later return NVFP4Tensor( shape=shape, dtype=grouped.get_dtype(), @@ -369,6 +368,7 @@ def _single_quantized_tensor_from_grouped( amax_rowwise=grouped.amax, amax_columnwise=grouped.columnwise_amax, fp4_dtype=fp4_dtype or quantizer.dtype, + scale_dtype=quantizer.scale_dtype, quantizer=quantizer, requires_grad=False, with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, @@ -878,13 +878,12 @@ def fuse_grouped_mlp_ops( Updated operations with matched triples replaced by fused ops. """ if not fused_op_cls.is_supported(): - assert False ### TODO Remove return ops # Fused kernels are only supported for MXFP8 and NVFP4 if recipe is None: return ops - elif recipe.custom(): + if recipe.custom(): # Check if custom recipe explicitly enables fusion if not getattr(recipe, "enable_cutedsl_fused_grouped_mlp", False): return ops @@ -996,6 +995,16 @@ def grouped_gemm_wgrad_kernel(cls) -> Optional[Callable]: return grouped_gemm_wgrad_wrapper_sm100 + @classmethod + def grouped_gemm_act_hadamard_kernel(cls) -> Optional[Callable]: + """Fused grouped GEMM activation kernel that also emits NVFP4 RHT amaxes.""" + return None + + @classmethod + def grouped_gemm_act_hadamard_quant_kernel(cls) -> Optional[Callable]: + """Fused grouped GEMM activation kernel that also quantizes NVFP4 with RHT.""" + return None + @classmethod @functools.lru_cache(maxsize=None) def is_supported(cls) -> bool: @@ -1441,16 +1450,14 @@ def fuser_forward( ): if fc2_input_quantizer.disable_second_level_scale: # Use GEMM + act + RHT + quant kernel if available - kernel_getter = getattr(self, "grouped_gemm_act_hadamard_quant_kernel", None) - if kernel_getter is None or kernel_getter() is None: + if self.grouped_gemm_act_hadamard_quant_kernel() is None: # Kernel is not available pass elif self._cudnn_act_func == "swiglu": kernel_impl = "gemm_act_rht_quant" elif fc2_input_quantizer.with_post_rht_amax: # Use GEMM + act + RHT + amax kernel if available - kernel_getter = getattr(self, "grouped_gemm_act_hadamard_kernel", None) - if kernel_getter is None or kernel_getter() is None: + if self.grouped_gemm_act_hadamard_kernel() is None: # Kernel is not available pass elif self._cudnn_act_func == "swiglu": @@ -1608,11 +1615,13 @@ def fuser_forward( if kernel_impl == "gemm_act": fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) elif kernel_impl == "gemm_act_rht_amax": + # pylint: disable-next=not-callable fc1_kernel_out = self.grouped_gemm_act_hadamard_kernel()(**fc1_activation_kwargs) elif kernel_impl == "gemm_act_rht_quant": + # pylint: disable-next=not-callable fc1_kernel_out = self.grouped_gemm_act_hadamard_quant_kernel()(**fc1_activation_kwargs) else: - raise RuntimeError("Unrecognized kernel variant ({kernel_impl})") + raise RuntimeError(f"Unrecognized kernel variant ({kernel_impl})") activation_in = fc1_kernel_out["c_tensor"] activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) @@ -1675,6 +1684,8 @@ def fuser_forward( tensor_offsets=fc2_x_tensor_offsets, with_gemm_swizzled_scales=True, ) + else: + raise RuntimeError(f"Unrecognized kernel variant ({kernel_impl})") else: # Unpack MXFP8 output fc2_in_row_data = fc1_kernel_out["d_tensor"] diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 39cc6588a6..328dae27bd 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -451,6 +451,8 @@ class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): Columnwise amax tracking tensor. fp4_dtype : DType The FP4 data type used for quantization. + scale_dtype: DType + The FP8 scale factor data type used for quantization. quantizer : Quantizer The quantizer instance used for this tensor. dtype : torch.dtype, default = torch.float32 From f2c26bd0b2ea15fb0b5fe0ef3eec6be83a7c2891 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:17:53 +0000 Subject: [PATCH 17/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/cpp_extensions/gemm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a2f4041a68..08bea8183e 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -466,7 +466,9 @@ def general_cuDNN_MX_gemm( assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." - assert use_split_accumulator is False, "cuDNN GEMM currently does not support split accumulators." + assert ( + use_split_accumulator is False + ), "cuDNN GEMM currently does not support split accumulators." # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for # raising accumulator precision, and the cuDNN kernel always accumulates in From 030e9aff77a6348cb487332d544c986c64a2c185 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 20 Aug 2026 11:22:09 +0000 Subject: [PATCH 18/54] Disable cuDNN GGEMM+GLU+RHT+quant kernel Signed-off-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 2 +- tests/pytorch/test_grouped_mlp.py | 2 +- tests/pytorch/utils.py | 2 +- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 12 ++++-------- 4 files changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3289748d68..fb5cc098a7 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -238,7 +238,7 @@ def make_reference_and_test_tensors( with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=with_rht, + with_random_sign_mask=False, disable_second_level_scale=disable_second_level_scale, )(test) elif quantization == "nvfp4_4over6": diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index f239518b84..e817c0b839 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -213,7 +213,7 @@ def make_reference_and_test_tensors( with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=with_rht, + with_random_sign_mask=False, disable_second_level_scale=disable_second_level_scale, )(test) elif quantization == "nvfp4_4over6": diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 6514d3b1c8..48078378ff 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -176,7 +176,7 @@ def make_nvfp4_ue5m3_quantizer(role: QuantizerRole) -> NVFP4Quantizer: with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, - with_random_sign_mask=with_rht, + with_random_sign_mask=False, disable_second_level_scale=tensor_type == "input", ) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 4e5caaf3a6..683a3c1dcb 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1447,14 +1447,11 @@ def fuser_forward( use_nvfp4 and isinstance(fc2_input_quantizer, NVFP4Quantizer) and fc2_input_quantizer.with_rht + and fc2_input_quantizer.rht_matrix_random_sign_mask_t == 0 ): if fc2_input_quantizer.disable_second_level_scale: - # Use GEMM + act + RHT + quant kernel if available - if self.grouped_gemm_act_hadamard_quant_kernel() is None: - # Kernel is not available - pass - elif self._cudnn_act_func == "swiglu": - kernel_impl = "gemm_act_rht_quant" + # Use GEMM + act + RHT + quant kernel once available + pass elif fc2_input_quantizer.with_post_rht_amax: # Use GEMM + act + RHT + amax kernel if available if self.grouped_gemm_act_hadamard_kernel() is None: @@ -1668,9 +1665,8 @@ def fuser_forward( fc2_in_row_scale = fc1_kernel_out["sfd_tensor"] fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) fc2_in_col_data = fc1_kernel_out["rht_tensor"] - fc2_in_col_data = fc2_in_col_data.view(fc2_weight_shape[1], in_shape[0] // 2) + fc2_in_col_data = fc2_in_col_data.permute(1, 0) fc2_in_col_scale = fc1_kernel_out["sfrht_tensor"] - fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) grouped_fc2_x = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[1]), dtype=dtype, From 9162e984e7e8debdd03dbbb56f178db2cd8d07da Mon Sep 17 00:00:00 2001 From: tdophung Date: Thu, 20 Aug 2026 11:09:46 -0700 Subject: [PATCH 19/54] Guard NVFP4 alpha scaling by scaling mode Signed-off-by: Tim Moon --- transformer_engine/common/gemm/cublaslt_gemm.cu | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 451155e1f4..7ae8d345cd 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -370,12 +370,14 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, const bool gelu = pre_gelu_out != nullptr; const bool use_fp8 = is_fp8_dtype(param.Atype) || is_fp8_dtype(param.Btype); const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype); + const bool nvfp4_tensor_scaling = + is_nvfp_scaling(inputA->scaling_mode) && is_nvfp_scaling(inputB->scaling_mode); // Update scaling factors with NVFP4 tensor scales // TODO: Check whether scales are on CPU/GPU or add API to control. // Currently scales are assumed to be on CPU when amax is provided // and on GPU when not provided, but this is brittle. - if (use_fp4 && + if (use_fp4 && nvfp4_tensor_scaling && ((transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr || (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr)) { // Reserve some workspace for alpha scale From db0957bb05a92cf21ea11bba2c1cab2b5c9d5f43 Mon Sep 17 00:00:00 2001 From: tdophung Date: Thu, 20 Aug 2026 14:11:39 -0700 Subject: [PATCH 20/54] Restore UE5M3 NVFP4 cast support Signed-off-by: Tim Moon --- .../common/cast/dispatch/quantize.cuh | 16 +- .../common/cast/nvfp4/core_nvfp4.cuh | 74 +++++ .../common/cast/nvfp4/dequantize_nvfp4.cuh | 38 ++- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 260 +++++++++++------- 4 files changed, 261 insertions(+), 127 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index d08bd07ef5..04706a40c0 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -104,8 +104,11 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const DType scale_dtype = output_tensor->scale_inv.has_data() + ? output_tensor->scale_inv.dtype + : output_tensor->columnwise_scale_inv.dtype; NVTE_CHECK(nvfp4_use_4over6 || static_cast(output_tensor->get_nvfp4_scale_max()) == - typeToMax(output_tensor->scale_inv.dtype), + typeToMax(scale_dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -287,8 +290,11 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens auto dtype = grad_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const DType scale_dtype = output_tensor->scale_inv.has_data() + ? output_tensor->scale_inv.dtype + : output_tensor->columnwise_scale_inv.dtype; NVTE_CHECK(nvfp4_use_4over6 || static_cast(output_tensor->get_nvfp4_scale_max()) == - typeToMax(output_tensor->scale_inv.dtype), + typeToMax(scale_dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); @@ -452,9 +458,11 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; if (!nvfp4_use_4over6) { for (const auto *output_tensor : output_tensors) { + const DType scale_dtype = output_tensor->scale_inv.has_data() + ? output_tensor->scale_inv.dtype + : output_tensor->columnwise_scale_inv.dtype; NVTE_CHECK( - static_cast(output_tensor->get_nvfp4_scale_max()) == - typeToMax(output_tensors[0]->scale_inv.dtype), + static_cast(output_tensor->get_nvfp4_scale_max()) == typeToMax(scale_dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6."); } } diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 4f15a4840c..b89dd755a9 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -51,6 +51,80 @@ namespace core { #if FP4_TYPE_SUPPORTED using namespace ptx; +// Scale-format-specific behavior belongs here rather than in individual kernels. +template +struct NVFP4ScaleTraits { + static constexpr bool is_supported = false; + static constexpr bool supports_fp16_error_path = false; + static constexpr float expected_max = 0.0f; + static constexpr float headroom_max = 0.0f; +}; + +template <> +struct NVFP4ScaleTraits { + // E4M3 scales fit in FP16 and can use the packed E4M3-to-FP16 PTX fast + // path. UE5M3 scales can exceed the FP16 range, so they retain the generic + // FP32 error path. + static constexpr bool is_supported = true; + static constexpr bool supports_fp16_error_path = true; + static constexpr float expected_max = 448.0f; + static constexpr float headroom_max = 256.0f; +}; + +#if CUDA_VERSION >= 13040 +template <> +struct NVFP4ScaleTraits { + static constexpr bool is_supported = true; + static constexpr bool supports_fp16_error_path = false; + static constexpr float expected_max = 114688.0f; + static constexpr float headroom_max = 65536.0f; +}; +#endif + +// Return the effective maximum used to derive the global NVFP4 encode scale. +// SCALE_TYPE_MAX is the resolved maximum for ScaleType (e.g., 448 for E4M3 +// or 114688 for UE5M3). The headroom maximum keeps the 1.5x map-to-4 scale +// used by 4over6 within the scale format's representable range. +template (NVFP4ScaleTraits::expected_max)> +__host__ __device__ constexpr float scale_max() { + using ScaleTraits = NVFP4ScaleTraits; + static_assert(ScaleTraits::is_supported, "Unsupported NVFP4 scale type."); + if constexpr (ScaleTraits::is_supported) { + static_assert(detail::TypeExtrema::max == ScaleTraits::expected_max, + "Unexpected NVFP4 scale type maximum."); + static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || + SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), + "Unsupported NVFP4 scale type maximum."); + static_assert(ScaleTraits::headroom_max * 1.5f <= ScaleTraits::expected_max, + "NVFP4 4over6 scale headroom exceeds scale type maximum."); + return static_cast(SCALE_TYPE_MAX); + } else { + return 0.0f; + } +} + +// Return the full-range maximum for a runtime scale dtype. +inline float scale_max(const DType scale_dtype) { + float result = 0.0f; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, + result = scale_max();) + return result; +} + +// Return and validate a user-provided maximum for a runtime scale dtype. +inline float scale_max(const DType scale_dtype, const int scale_type_max) { + float result = 0.0f; + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { + using ScaleTraits = NVFP4ScaleTraits; + NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || + scale_type_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + result = static_cast(scale_type_max); + }) + return result; +} + template __device__ __forceinline__ ScaleType compute_decoding_scaling_factor(const float block_amax, const float global_encode_scale) { diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index bc553d45d0..a014244b9b 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -66,7 +66,7 @@ __global__ void __launch_bounds__(512) value.vec = input_vectorized[my_index]; ScaleType scale = scales[my_scale_index]; constexpr float fp4_max = detail::TypeExtrema::max; - constexpr float unit_global_scale_amax = fp4_max * SCALE_TYPE_MAX; + constexpr float unit_global_scale_amax = fp4_max * core::scale_max(); float amax = unit_global_scale_amax; if (tensor_amax != nullptr) { amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; @@ -139,27 +139,21 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) "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."); - - if (static_cast(e4m3_max) != typeToMax(scale_dtype)) { - NVTE_CHECK(scale_dtype == DType::kFloat8E4M3, - "NVFP4 dequantization with non-default scale max " - "is only supported with FP8E4M3 scales (found ", - to_string(scale_dtype), ")."); - NVTE_CHECK(e4m3_max == 256, - "NVFP4 dequantization with non-default scale max " - "is only supported with e4m3_max=256 (found ", - e4m3_max, ")."); - launch_dequantize(input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, - Mread, blocks, threads, num_scale_tiles_X, stream); - NVTE_CHECK_CUDA(cudaGetLastError()); - } else { - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, - launch_dequantize(TypeInfo::max_finite_value)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream);); // NOLINT(*) - NVTE_CHECK_CUDA(cudaGetLastError()); - } + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { + using ScaleTraits = core::NVFP4ScaleTraits; + if (e4m3_max == static_cast(ScaleTraits::expected_max)) { + launch_dequantize(ScaleTraits::expected_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + } else { + NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + launch_dequantize(ScaleTraits::headroom_max)>( + input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + } + }) + NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 9a2af69ca1..d5a220d8ef 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -54,16 +54,6 @@ namespace nvfp4 { } \ } -#define TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(E4M3_MAX_VALUE, E4M3_MAX_CONST, ...) \ - if ((E4M3_MAX_VALUE) == 256) { \ - constexpr int E4M3_MAX_CONST = 256; \ - { __VA_ARGS__ } \ - } else { \ - NVTE_CHECK((E4M3_MAX_VALUE) == 448, "Unsupported NVFP4 E4M3 max."); \ - constexpr int E4M3_MAX_CONST = 448; \ - { __VA_ARGS__ } \ - } - namespace quantize_4over6_kernel { constexpr int kThreads = 128; @@ -81,8 +71,6 @@ constexpr int kPackedWordsPerGroup = 2; static_assert(kTileRows == kPipelineStages * kStageRows); static_assert(kStageRows % kGroupSize == 0); -using nvfp4_scale_t = fp8e4m3; - template struct Config { static constexpr NVTENVFP44Over6Mode mode = kMode; @@ -99,9 +87,10 @@ struct CandidatePair { Candidate map6; }; +template struct ScalePair { - nvfp4_scale_t map4; - nvfp4_scale_t map6; + ScaleType map4; + ScaleType map6; float inv_map4; float inv_map6; float global_encode_scale; @@ -124,20 +113,24 @@ __device__ __forceinline__ float compute_error_rn(const float diff) { } } -template -__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, - const float global_amax) { - static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); +template +__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, + const float global_amax) { + using ScaleTraits = core::NVFP4ScaleTraits; + static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || + SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), + "Unsupported NVFP4 scale type maximum."); constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f + constexpr float fp8_max = detail::TypeExtrema::max; + constexpr int encode_scale_max = static_cast(core::scale_max()); constexpr float expand_to_map4 = 1.5f; const float S_enc = - core::compute_global_encode_scaling_factor_FP4(global_amax); + core::compute_global_encode_scaling_factor_FP4(global_amax); const float base = block_amax / fp4_max * S_enc; - ScalePair scales; - scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); - scales.map6 = static_cast(fminf(base, fp8_max)); + ScalePair scales; + scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); + scales.map6 = static_cast(fminf(base, fp8_max)); const float S_dec = 1.0f / S_enc; scales.inv_map4 = @@ -190,12 +183,12 @@ __device__ __forceinline__ void load_col_group(const IType *tile, const int row_ } } -template +template __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, const float sf, const float global_amax, float *err) { constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = static_cast(E4M3_MAX); + constexpr float fp8_max = core::scale_max(); constexpr float err_denom = fp4_max * fp8_max; const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; const float dequant = __half2float(__ushort_as_half(half_bits)); @@ -204,11 +197,19 @@ __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_ *err = __fadd_rn(*err, compute_error_rn(diff)); } -__device__ __forceinline__ uint8_t fp8_bits(const nvfp4_scale_t sf) { +template +__device__ __forceinline__ uint8_t fp8_bits(const ScaleType sf) { return *reinterpret_cast(&sf); } -__device__ __forceinline__ FP16ErrorScalePair compute_fp16_error_scales(const ScalePair &scales) { +template +__device__ __forceinline__ FP16ErrorScalePair +compute_fp16_error_scales(const ScalePair &scales) { + // This fast error path interprets the packed scale bits as E4M3. UE5M3 + // deliberately does not enable supports_fp16_error_path and instead uses + // the scale-format-independent float error path in + // cvt_fp32_to_fp4_8x_with_error. + static_assert(core::NVFP4ScaleTraits::supports_fp16_error_path); FP16ErrorScalePair result; const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | (static_cast(fp8_bits(scales.map6)) << 8); @@ -260,9 +261,9 @@ __device__ __forceinline__ void accumulate_fp16_scaled_error_pair(const uint32_t *err = __fadd_rn(*err, compute_error_rn(diff1)); } -template +template __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( - const float (&x)[8], const float block_scale_inverse, const nvfp4_scale_t sf, + const float (&x)[8], const float block_scale_inverse, const ScaleType sf, const uint32_t fp16_error_scale, const float global_amax, const float global_encode_scale, float *err) { uint32_t out = 0; @@ -271,6 +272,11 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( uint32_t out_dequant_3 = 0; uint32_t out_dequant_4 = 0; + // ScaleType is not consumed by this PTX. block_scale_inverse applies the + // selected E4M3 or UE5M3 block scale while forming the FP32 operands. These + // instructions only convert the scaled candidates to FP4 E2M1 and back to + // FP16 for error evaluation, so their encoding is identical for both scale + // storage types. constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; if constexpr (is_blackwell) { asm volatile( @@ -298,7 +304,8 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( "Try recompiling with sm_XXXa instead of sm_XXX."); } - if constexpr (Cfg::err_use_fast_math) { + if constexpr (Cfg::err_use_fast_math && + core::NVFP4ScaleTraits::supports_fp16_error_path) { accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, global_encode_scale, err); accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, @@ -309,39 +316,48 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( global_encode_scale, err); } else { const float sf_float = static_cast(sf); - accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); - accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_1, x[0], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_1, x[1], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_2, x[2], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_2, x[3], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_3, x[4], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_3, x[5], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_4, x[6], sf_float, + global_amax, err); + accumulate_dequant_error(out_dequant_4, x[7], sf_float, + global_amax, err); } return out; } -template +template __device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], - const ScalePair &scales, + const ScalePair &scales, const float global_amax) { CandidatePair candidates; candidates.map4.err = 0.0f; candidates.map6.err = 0.0f; FP16ErrorScalePair fp16_error_scales{}; - if constexpr (Cfg::err_use_fast_math) { + if constexpr (Cfg::err_use_fast_math && + core::NVFP4ScaleTraits::supports_fp16_error_path) { fp16_error_scales = compute_fp16_error_scales(scales); } - candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( x0, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); - candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map4, scales.map4, fp16_error_scales.map4, global_amax, scales.global_encode_scale, &candidates.map4.err); - candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( + candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( x1, scales.inv_map6, scales.map6, fp16_error_scales.map6, global_amax, scales.global_encode_scale, &candidates.map6.err); return candidates; @@ -383,8 +399,9 @@ __device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &ca return candidates.map6.packed; } -__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, - const bool pick_map4) { +template +__device__ __forceinline__ ScaleType select_scale(const ScalePair &scales, + const bool pick_map4) { if (pick_map4) { return scales.map4; } @@ -452,9 +469,9 @@ __device__ void load_stage_to_shared_async(const IType *input, IType *tile, cons } } -template -__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, +template +__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, ScaleType *scales, const float *amax, const size_t rows, const size_t cols, const size_t stage_row, const size_t tile_col, const size_t scale_stride) { @@ -479,13 +496,21 @@ __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 = + core::scale_max() * detail::TypeExtrema::max; + 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); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + const ScalePair scale_pair = + compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = + make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -495,7 +520,7 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf } const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const ScaleType selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_col_group = global_col / kGroupSize; @@ -504,11 +529,12 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvf } } -template -__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales_t, const float *amax, - const size_t rows, const size_t cols, const size_t stage_row, - const size_t tile_col, const size_t scale_stride_t) { +template +__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, ScaleType *scales_t, + const float *amax, const size_t rows, const size_t cols, + const size_t stage_row, const size_t tile_col, + const size_t scale_stride_t) { constexpr int groups = kStageRowGroups * kTileCols; for (int group = threadIdx.x; group < groups; group += blockDim.x) { const int local_row_group = group / kTileCols; @@ -530,9 +556,13 @@ __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 ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); - CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + const float global_amax = amax == nullptr ? core::scale_max() * + detail::TypeExtrema::max + : amax[0]; + const ScalePair scale_pair = + compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = + make_candidates(x0, x1, scale_pair, global_amax); float err_map4 = candidates.map4.err; float err_map6 = candidates.map6.err; @@ -542,7 +572,7 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, } const bool pick_map4 = err_map4 < err_map6; - const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const ScaleType selected_scale = select_scale(scale_pair, pick_map4); const uint32_t *selected = select_packed(candidates, pick_map4); const size_t global_row_group = global_row / kGroupSize; @@ -552,13 +582,14 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, } template + bool ROW_SCALED_NVFP4, typename Cfg, typename ScaleType, int SCALE_TYPE_MAX, + typename IType> __global__ void __launch_bounds__(kThreads) quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, - nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, - const float *amax_rowwise, const float *amax_colwise, const size_t rows, - const size_t cols, const size_t scale_stride, - const size_t scale_stride_t, const float *noop) { + ScaleType *scales, ScaleType *scales_t, const float *amax_rowwise, + const float *amax_colwise, const size_t rows, const size_t cols, + const size_t scale_stride, const size_t scale_stride_t, + const float *noop) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) if (noop != nullptr && noop[0] == 1.0f) { return; @@ -593,7 +624,7 @@ __global__ void __launch_bounds__(kThreads) IType *stage_tile = stage_tiles[stage]; if constexpr (RETURN_IDENTITY) { - quantize_stage_rowwise( + quantize_stage_rowwise( stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); } @@ -602,7 +633,7 @@ __global__ void __launch_bounds__(kThreads) if (columnwise_amax == nullptr) { columnwise_amax = amax_rowwise; } - quantize_stage_colwise( + quantize_stage_colwise( stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, scale_stride_t); } @@ -617,7 +648,8 @@ __global__ void __launch_bounds__(kThreads) #endif } -template +template void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { const size_t rows = input.flat_first_dim(); @@ -629,8 +661,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out const auto *input_ptr = reinterpret_cast(input.data.dptr); auto *output_ptr = reinterpret_cast(output->data.dptr); auto *output_t_ptr = reinterpret_cast(output->columnwise_data.dptr); - auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); - auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); + auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); const auto *amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); const auto *amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); const auto *noop_ptr = reinterpret_cast(noop->data.dptr); @@ -645,8 +677,9 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { - auto kernel = quantize_4over6_kernel; + auto kernel = + quantize_4over6_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, @@ -660,9 +693,9 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out #endif // FP4_TYPE_SUPPORTED -template -void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { +template +void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace quantize_4over6_kernel; @@ -695,35 +728,35 @@ 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(is_fp4_dtype(output->data.dtype), "Output data must have FP4 type."); - NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E4M3, - "Output scales must have FP8E4M3 type."); - NVTE_CHECK(output->amax.dptr != nullptr, "Rowwise amax tensor must be allocated."); + NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); } if (output->has_columnwise_data()) { NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Transposed scaling tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); - NVTE_CHECK(output->columnwise_scale_inv.dtype == DType::kFloat8E4M3, - "Output scales must have FP8E4M3 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( - output->nvfp4_e4m3_max, E4M3_MAX, - TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( - quant_config->nvfp4_4over6_mode, MODE, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { - using Cfg = quantize_4over6_kernel::Config; - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - quantize_4over6_kernel::launch_quantize_4over6( - input, noop, output, stream);); - }););); + using ScaleTraits = core::NVFP4ScaleTraits; + const int scale_type_max = output->get_nvfp4_scale_max(); + NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || + scale_type_max == static_cast(ScaleTraits::headroom_max), + "Unsupported maximum for NVFP4 scale dtype."); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + scale_type_max == static_cast(ScaleTraits::headroom_max), USE_SCALE_HEADROOM, { + constexpr int SCALE_TYPE_MAX = static_cast( + USE_SCALE_HEADROOM ? ScaleTraits::headroom_max : ScaleTraits::expected_max); + TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( + quant_config->nvfp4_4over6_mode, MODE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { + using Cfg = quantize_4over6_kernel::Config; + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + quantize_4over6_kernel::launch_quantize_4over6< + use_2d_quantization, Cfg, ScaleType, SCALE_TYPE_MAX, IType>( + input, noop, output, stream);); + });); + }) NVTE_CHECK_CUDA(cudaGetLastError()); #else @@ -731,6 +764,31 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, #endif // FP4_TYPE_SUPPORTED } +template +void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + const bool return_rowwise = output->has_data(); + const bool return_transpose = output->has_columnwise_data(); + NVTE_CHECK(return_rowwise || return_transpose, + "NVFP4 4over6 output tensor must have rowwise or columnwise data."); + const DType scale_dtype = + return_rowwise ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + if (return_rowwise && return_transpose) { + NVTE_CHECK(output->scale_inv.dtype == output->columnwise_scale_inv.dtype, + "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", + to_string(output->scale_inv.dtype), " and ", + to_string(output->columnwise_scale_inv.dtype), ")."); + } + + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, + quantize_4over6_impl( + input, noop, output, quant_config, stream);) +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + } // namespace nvfp4 } // namespace dispatch } // namespace transformer_engine From 982be7b6e64a9430752c51a8986600bb28236abd Mon Sep 17 00:00:00 2001 From: tdophung Date: Thu, 20 Aug 2026 17:22:11 -0700 Subject: [PATCH 21/54] Localize NVFP4 4over6 scale policy Signed-off-by: Tim Moon --- .../common/cast/nvfp4/core_nvfp4.cuh | 74 ------------------- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 31 ++++---- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 73 ++++++++++++------ .../include/transformer_engine/recipe.h | 4 +- transformer_engine/common/recipe/nvfp4.cu | 4 +- .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 6 +- 6 files changed, 73 insertions(+), 119 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index b89dd755a9..4f15a4840c 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -51,80 +51,6 @@ namespace core { #if FP4_TYPE_SUPPORTED using namespace ptx; -// Scale-format-specific behavior belongs here rather than in individual kernels. -template -struct NVFP4ScaleTraits { - static constexpr bool is_supported = false; - static constexpr bool supports_fp16_error_path = false; - static constexpr float expected_max = 0.0f; - static constexpr float headroom_max = 0.0f; -}; - -template <> -struct NVFP4ScaleTraits { - // E4M3 scales fit in FP16 and can use the packed E4M3-to-FP16 PTX fast - // path. UE5M3 scales can exceed the FP16 range, so they retain the generic - // FP32 error path. - static constexpr bool is_supported = true; - static constexpr bool supports_fp16_error_path = true; - static constexpr float expected_max = 448.0f; - static constexpr float headroom_max = 256.0f; -}; - -#if CUDA_VERSION >= 13040 -template <> -struct NVFP4ScaleTraits { - static constexpr bool is_supported = true; - static constexpr bool supports_fp16_error_path = false; - static constexpr float expected_max = 114688.0f; - static constexpr float headroom_max = 65536.0f; -}; -#endif - -// Return the effective maximum used to derive the global NVFP4 encode scale. -// SCALE_TYPE_MAX is the resolved maximum for ScaleType (e.g., 448 for E4M3 -// or 114688 for UE5M3). The headroom maximum keeps the 1.5x map-to-4 scale -// used by 4over6 within the scale format's representable range. -template (NVFP4ScaleTraits::expected_max)> -__host__ __device__ constexpr float scale_max() { - using ScaleTraits = NVFP4ScaleTraits; - static_assert(ScaleTraits::is_supported, "Unsupported NVFP4 scale type."); - if constexpr (ScaleTraits::is_supported) { - static_assert(detail::TypeExtrema::max == ScaleTraits::expected_max, - "Unexpected NVFP4 scale type maximum."); - static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || - SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), - "Unsupported NVFP4 scale type maximum."); - static_assert(ScaleTraits::headroom_max * 1.5f <= ScaleTraits::expected_max, - "NVFP4 4over6 scale headroom exceeds scale type maximum."); - return static_cast(SCALE_TYPE_MAX); - } else { - return 0.0f; - } -} - -// Return the full-range maximum for a runtime scale dtype. -inline float scale_max(const DType scale_dtype) { - float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, - result = scale_max();) - return result; -} - -// Return and validate a user-provided maximum for a runtime scale dtype. -inline float scale_max(const DType scale_dtype, const int scale_type_max) { - float result = 0.0f; - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { - using ScaleTraits = NVFP4ScaleTraits; - NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - result = static_cast(scale_type_max); - }) - return result; -} - template __device__ __forceinline__ ScaleType compute_decoding_scaling_factor(const float block_amax, const float global_encode_scale) { diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index a014244b9b..788b27248c 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -16,6 +16,8 @@ #include #include +#include + #include "../../common.h" #include "../../util/math.h" #include "../../util/ptx.cuh" @@ -66,7 +68,7 @@ __global__ void __launch_bounds__(512) value.vec = input_vectorized[my_index]; ScaleType scale = scales[my_scale_index]; constexpr float fp4_max = detail::TypeExtrema::max; - constexpr float unit_global_scale_amax = fp4_max * core::scale_max(); + constexpr float unit_global_scale_amax = fp4_max * static_cast(SCALE_TYPE_MAX); float amax = unit_global_scale_amax; if (tensor_amax != nullptr) { amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; @@ -122,7 +124,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; const bool row_scaled_nvfp4 = input.row_scaled_nvfp4; const DType scale_dtype = input.scale_inv.dtype; - const int e4m3_max = input.get_nvfp4_scale_max(); + const int scale_type_max = input.get_nvfp4_scale_max(); constexpr int FP4_BLOCK_SIZE = 16; const auto [N, M] = input.flat_2d_dims(); @@ -139,19 +141,20 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) "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."); + const int full_scale_max = static_cast(typeToMax(scale_dtype)); + const bool uses_4over6_headroom = (scale_dtype == DType::kFloat8E4M3 && scale_type_max == 256) || + (scale_dtype == DType::kFloat8UE5M3 && scale_type_max == 65536); + NVTE_CHECK(scale_type_max == full_scale_max || uses_4over6_headroom, "Unsupported maximum ", + scale_type_max, " for NVFP4 scale dtype ", to_string(scale_dtype), "."); TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { - using ScaleTraits = core::NVFP4ScaleTraits; - if (e4m3_max == static_cast(ScaleTraits::expected_max)) { - launch_dequantize(ScaleTraits::expected_max)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); - } else { - NVTE_CHECK(e4m3_max == static_cast(ScaleTraits::headroom_max), - "Unsupported maximum for NVFP4 scale dtype."); - launch_dequantize(ScaleTraits::headroom_max)>( - input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); - } + constexpr int full_max = static_cast(TypeInfo::max_finite_value); + constexpr int headroom_max = std::is_same_v ? 256 : 65536; + TRANSFORMER_ENGINE_SWITCH_CONDITION(scale_type_max == headroom_max, USE_HEADROOM, { + constexpr int SCALE_TYPE_MAX = USE_HEADROOM ? headroom_max : full_max; + launch_dequantize(input, output, with_gemm_swizzled_scales, + row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + }) }) NVTE_CHECK_CUDA(cudaGetLastError()); #else diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index d5a220d8ef..84b25ee871 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -77,6 +77,26 @@ struct Config { static constexpr bool err_use_fast_math = kErrUseFastMath; }; +// Policy that is specific to the 4over6 encoding. The ordinary scale-type +// maximum comes from TypeInfo/typeToMax; only the reduced maximum needed for +// map-to-4 headroom and the E4M3-only FP16 error fast path live here. +template +struct FourOverSixScaleConfig; + +template <> +struct FourOverSixScaleConfig { + static constexpr int headroom_max = 256; + static constexpr bool supports_fp16_error_path = true; +}; + +#if CUDA_VERSION >= 13040 +template <> +struct FourOverSixScaleConfig { + static constexpr int headroom_max = 65536; + static constexpr bool supports_fp16_error_path = false; +}; +#endif + struct Candidate { uint32_t packed[kPackedWordsPerGroup]; float err; @@ -116,16 +136,17 @@ __device__ __forceinline__ float compute_error_rn(const float diff) { template __device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, const float global_amax) { - using ScaleTraits = core::NVFP4ScaleTraits; - static_assert(SCALE_TYPE_MAX == static_cast(ScaleTraits::expected_max) || - SCALE_TYPE_MAX == static_cast(ScaleTraits::headroom_max), + using ScaleConfig = FourOverSixScaleConfig; + constexpr int full_scale_max = static_cast(TypeInfo::max_finite_value); + static_assert(SCALE_TYPE_MAX == full_scale_max || SCALE_TYPE_MAX == ScaleConfig::headroom_max, "Unsupported NVFP4 scale type maximum."); + static_assert(ScaleConfig::headroom_max * 1.5f <= full_scale_max, + "NVFP4 4over6 scale headroom exceeds scale type maximum."); constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f constexpr float fp8_max = detail::TypeExtrema::max; - constexpr int encode_scale_max = static_cast(core::scale_max()); constexpr float expand_to_map4 = 1.5f; const float S_enc = - core::compute_global_encode_scaling_factor_FP4(global_amax); + core::compute_global_encode_scaling_factor_FP4(global_amax); const float base = block_amax / fp4_max * S_enc; ScalePair scales; @@ -188,7 +209,7 @@ __device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_ const float sf, const float global_amax, float *err) { constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f - constexpr float fp8_max = core::scale_max(); + constexpr float fp8_max = static_cast(SCALE_TYPE_MAX); constexpr float err_denom = fp4_max * fp8_max; const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; const float dequant = __half2float(__ushort_as_half(half_bits)); @@ -209,7 +230,7 @@ compute_fp16_error_scales(const ScalePair &scales) { // deliberately does not enable supports_fp16_error_path and instead uses // the scale-format-independent float error path in // cvt_fp32_to_fp4_8x_with_error. - static_assert(core::NVFP4ScaleTraits::supports_fp16_error_path); + static_assert(FourOverSixScaleConfig::supports_fp16_error_path); FP16ErrorScalePair result; const uint32_t packed_scales = static_cast(fp8_bits(scales.map4)) | (static_cast(fp8_bits(scales.map6)) << 8); @@ -305,7 +326,7 @@ __device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error( } if constexpr (Cfg::err_use_fast_math && - core::NVFP4ScaleTraits::supports_fp16_error_path) { + FourOverSixScaleConfig::supports_fp16_error_path) { accumulate_fp16_scaled_error_pair(out_dequant_1, x[0], x[1], fp16_error_scale, global_encode_scale, err); accumulate_fp16_scaled_error_pair(out_dequant_2, x[2], x[3], fp16_error_scale, @@ -345,7 +366,7 @@ __device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], c candidates.map6.err = 0.0f; FP16ErrorScalePair fp16_error_scales{}; if constexpr (Cfg::err_use_fast_math && - core::NVFP4ScaleTraits::supports_fp16_error_path) { + FourOverSixScaleConfig::supports_fp16_error_path) { fp16_error_scales = compute_fp16_error_scales(scales); } candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( @@ -496,8 +517,7 @@ __device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, Sca block_amax = reduce_group_max_16(group_amax); } - float global_amax = - core::scale_max() * detail::TypeExtrema::max; + float global_amax = static_cast(SCALE_TYPE_MAX) * detail::TypeExtrema::max; if (amax != nullptr) { global_amax = amax[0]; } @@ -556,9 +576,9 @@ __device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, S block_amax = reduce_group_max_16(group_amax); } - const float global_amax = amax == nullptr ? core::scale_max() * - detail::TypeExtrema::max - : amax[0]; + const float global_amax = + amax == nullptr ? static_cast(SCALE_TYPE_MAX) * detail::TypeExtrema::max + : amax[0]; const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); CandidatePair candidates = @@ -695,7 +715,8 @@ void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *out template void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { + const QuantizationConfig *quant_config, const DType scale_dtype, + cudaStream_t stream) { #if FP4_TYPE_SUPPORTED using namespace quantize_4over6_kernel; @@ -736,15 +757,18 @@ void quantize_4over6_impl(const Tensor &input, const Tensor *noop, Tensor *outpu NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), "Transposed output must have FP4 type."); } - using ScaleTraits = core::NVFP4ScaleTraits; + using ScaleConfig = FourOverSixScaleConfig; const int scale_type_max = output->get_nvfp4_scale_max(); - NVTE_CHECK(scale_type_max == static_cast(ScaleTraits::expected_max) || - scale_type_max == static_cast(ScaleTraits::headroom_max), + const int full_scale_max = static_cast(typeToMax(scale_dtype)); + NVTE_CHECK(full_scale_max == static_cast(TypeInfo::max_finite_value), + "NVFP4 scale dtype dispatch does not match its datatype maximum."); + NVTE_CHECK(scale_type_max == full_scale_max || scale_type_max == ScaleConfig::headroom_max, "Unsupported maximum for NVFP4 scale dtype."); TRANSFORMER_ENGINE_SWITCH_CONDITION( - scale_type_max == static_cast(ScaleTraits::headroom_max), USE_SCALE_HEADROOM, { - constexpr int SCALE_TYPE_MAX = static_cast( - USE_SCALE_HEADROOM ? ScaleTraits::headroom_max : ScaleTraits::expected_max); + scale_type_max == ScaleConfig::headroom_max, USE_SCALE_HEADROOM, { + constexpr int SCALE_TYPE_MAX = + USE_SCALE_HEADROOM ? ScaleConfig::headroom_max + : static_cast(TypeInfo::max_finite_value); TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( quant_config->nvfp4_4over6_mode, MODE, TRANSFORMER_ENGINE_SWITCH_CONDITION( @@ -781,9 +805,10 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, to_string(output->columnwise_scale_inv.dtype), ")."); } - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, - quantize_4over6_impl( - input, noop, output, quant_config, stream);) + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( + scale_dtype, ScaleType, + quantize_4over6_impl(input, noop, output, quant_config, + scale_dtype, stream);) #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 0f7c3ae307..64e539e938 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -373,13 +373,13 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s * \param[in] scale_stride_w Stride for scale in tile-col dimension. * \param[in] start_offset Starting element offset in the flattened tensor. * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). - * \param[in] stream CUDA stream used for the operation. * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream used for the operation. */ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, cudaStream_t stream, const NVTEDType scale_dtype); + size_t block_len, const NVTEDType scale_dtype, cudaStream_t stream); /*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. * diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index f9d661e2e1..9b1bf67b4c 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -916,8 +916,8 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, cudaStream_t stream, - const NVTEDType scale_dtype) { + size_t block_len, const NVTEDType scale_dtype, + cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); using namespace transformer_engine; diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp index e9321b9b86..f5d3ad9bc7 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -45,7 +45,7 @@ void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tens nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), start_offset, block_len, - at::cuda::getCurrentCUDAStream(), static_cast(scale_dtype)); + static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); } void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, @@ -96,8 +96,8 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), - start_offset, static_cast(block_len), stream, - static_cast(scale_dtype)); + start_offset, static_cast(block_len), + static_cast(scale_dtype), stream); } } From 01f4edc922c6780da2aa719990bd8b7ba261c7de Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 01:25:08 +0000 Subject: [PATCH 22/54] Tweak arg order in C API functions Signed-off-by: Tim Moon --- .../include/transformer_engine/recipe.h | 22 +++++++++---------- transformer_engine/common/recipe/nvfp4.cu | 14 ++++++------ .../pytorch/csrc/extensions/transpose.cpp | 22 ++++++++++--------- 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 64e539e938..2d141e2f16 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -379,7 +379,7 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, const NVTEDType scale_dtype, cudaStream_t stream); + size_t block_len, NVTEDType scale_dtype, cudaStream_t stream); /*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. * @@ -391,12 +391,12 @@ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTE * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Padded row count in output. * \param[in] block_len Block length (typically 16 for NVFP4). - * \param[in] stream CUDA stream. * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. */ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream, const NVTEDType scale_dtype); + NVTEDType scale_dtype, cudaStream_t stream); /*! \brief Compute per-block decode scale from block amax and global amax. * @@ -409,12 +409,12 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. * \param[out] scale Output scale tensor [tile_rows, tile_cols], float32. * \param[in] global_amax Global amax tensor (single element), float32. Avoids D2H transfer. - * \param[in] stream CUDA stream. * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. */ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, cudaStream_t stream, - const NVTEDType scale_dtype); + const NVTETensor global_amax, NVTEDType scale_dtype, + cudaStream_t stream); /*! \brief Fused kernel for NVFP4 scale computation. * @@ -434,14 +434,14 @@ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Total padded rows in output. * \param[in] block_len Block length (16 for NVFP4). - * \param[in] stream CUDA stream. * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. */ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, cudaStream_t stream, - const NVTEDType scale_dtype); + size_t rows_padded, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream); /*! \brief Compute global encode scale from global amax. * @@ -450,11 +450,11 @@ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global * * \param[in] global_amax Input global amax tensor [num_params], float32. * \param[out] global_scale Output global scale tensor [num_params], float32. - * \param[in] stream CUDA stream. * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. */ void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - cudaStream_t stream, const NVTEDType scale_dtype); + NVTEDType scale_dtype, cudaStream_t stream); #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 9b1bf67b4c..2807eb9703 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -836,7 +836,7 @@ void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - cudaStream_t stream, const NVTEDType scale_dtype) { + NVTEDType scale_dtype, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8); using namespace transformer_engine; @@ -849,8 +849,8 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s } void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, cudaStream_t stream, - const NVTEDType scale_dtype) { + const NVTETensor global_amax, NVTEDType scale_dtype, + cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale); using namespace transformer_engine; @@ -863,7 +863,7 @@ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor } void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - cudaStream_t stream, const NVTEDType scale_dtype) { + NVTEDType scale_dtype, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_global_scale); using namespace transformer_engine; @@ -916,7 +916,7 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, const NVTEDType scale_dtype, + size_t block_len, NVTEDType scale_dtype, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); @@ -963,8 +963,8 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, cudaStream_t stream, - const NVTEDType scale_dtype) { + size_t rows_padded, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_fused_scale); using namespace transformer_engine; diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 9f34c4d196..1806208c09 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -122,7 +122,8 @@ void nvfp4_2d_scale_transpose(at::Tensor input, at::Tensor output, int64_t M_til } void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, - int64_t tile_cols, int64_t rows_padded, int64_t block_len) { + int64_t tile_cols, int64_t rows_padded, int64_t block_len, + DType scale_dtype) { init_extension(); // Input: per_block_decode_scale [tile_rows, tile_cols], float32 @@ -141,7 +142,8 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile nvte_nvfp4_expand_scale_to_fp8(input_cu.data(), output_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), - static_cast(block_len), at::cuda::getCurrentCUDAStream()); + static_cast(block_len), static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, @@ -160,8 +162,8 @@ void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at:: auto global_amax_cu = makeTransformerEngineTensor(global_amax); nvte_nvfp4_compute_per_block_scale(block_amax_cu.data(), scale_cu.data(), global_amax_cu.data(), - at::cuda::getCurrentCUDAStream(), - static_cast(scale_dtype)); + static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, @@ -193,7 +195,7 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor target_scale_cu.data(), target_amax_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), static_cast(block_len), - at::cuda::getCurrentCUDAStream(), static_cast(scale_dtype)); + static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); } void nvfp4_multi_tensor_fused_scale( @@ -245,13 +247,13 @@ void nvfp4_multi_tensor_fused_scale( nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, - rows_padded, static_cast(block_len), stream, - static_cast(scale_dtype)); + rows_padded, static_cast(block_len), + static_cast(scale_dtype), stream); } } void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, - const DType scale_dtype) { + DType scale_dtype) { init_extension(); // global_amax and global_scale: [num_params], float32 @@ -262,8 +264,8 @@ void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, auto global_scale_cu = makeTransformerEngineTensor(global_scale); nvte_nvfp4_compute_global_scale(global_amax_cu.data(), global_scale_cu.data(), - at::cuda::getCurrentCUDAStream(), - static_cast(scale_dtype)); + static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { From 3a636233b22da8679e38f9d978a18fc5a827fd35 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 01:54:48 +0000 Subject: [PATCH 23/54] Rename cuDNN GGEMM helper functions for general_gemm Signed-off-by: Tim Moon --- tests/pytorch/utils.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 84 +++++++++++-------- .../pytorch/ops/fused/grouped_mlp.py | 10 +-- 3 files changed, 57 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 48078378ff..9ee82b098d 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -166,7 +166,7 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: def make_nvfp4_ue5m3_quantizer(role: QuantizerRole) -> NVFP4Quantizer: """Quantizer factory for NVFP4-UE5M3 recipe.""" - tensor_type = role.tensor_type if role is not None else "input" + tensor_type = role.tensor_type if role is not None else None if not tensor_type: tensor_type = "input" with_rht = name == "nvfp4_rht_ue5m3" and tensor_type != "weight" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 08bea8183e..ff0c72f2ce 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -217,14 +217,35 @@ def validate_or_alloc_output( @functools.lru_cache(maxsize=None) -def grouped_gemm_wgrad_kernel() -> Callable: - """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs.""" +def _cudnn_grouped_gemm_quant_kernel() -> Callable: + """cuDNN CuTe DSL grouped GEMM kernel for block-scaled inputs. + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + + """ + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_quant_wrapper_sm100 + + +@functools.lru_cache(maxsize=None) +def _cudnn_grouped_gemm_wgrad_kernel() -> Callable: + """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs. + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + + """ from cudnn import grouped_gemm_wgrad_wrapper_sm100 # pylint: disable=no-name-in-module return grouped_gemm_wgrad_wrapper_sm100 - -def _cuDNN_wgrad_gemm( +def _cudnn_wgrad_grouped_gemm_nvfp4_ue5m3( a_tensor: torch.Tensor, b_tensor: torch.Tensor, sfa: torch.Tensor, @@ -237,7 +258,14 @@ def _cuDNN_wgrad_gemm( alpha: Optional[float] = None, bias: Optional[torch.Tensor] = None, ) -> Iterable[Optional[torch.Tensor]]: - """Compute dw = dy^T @ x with cuDNN's purpose-built grouped wgrad kernel.""" + """Compute dw = dy^T @ x for NVFP4-UE5M3 data with cuDNN's grouped wgrad kernel. + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + + """ # Column-wise NVFP4 buffers are physically (features, tokens), FP4-packed # two values per byte along the token dim. @@ -255,7 +283,7 @@ def _cuDNN_wgrad_gemm( b_tensor = b_tensor.view(dtype=fp4).view(in_features, tokens_packed).T # Create the scale factor tensors with the logical layout cuDNN expects - # In general_cuDNN_MX_gemm we've already ensured they are swizzled physically + # Assume scales have already been swizzled in _cudnn_grouped_gemm_nvfp4_ue5m3 def _sf(scale_inv, features): leading = ceil_div(features, 128) * 128 return scale_inv.view(leading, -1).view(dtype=torch.float8_e4m3fn) @@ -270,7 +298,7 @@ def _sf(scale_inv, features): global_scale_a = global_scale_a * alpha out = validate_or_alloc_output(out, (out_features, in_features), out_dtype, a_tensor.device) - grouped_gemm_wgrad_kernel()( + _cudnn_grouped_gemm_wgrad_kernel()( a_tensor=a_tensor, b_tensor=b_tensor, sfa_tensor=_sf(sfa, out_features), @@ -297,15 +325,7 @@ def _sf(scale_inv, features): return out, None, None, None -@functools.lru_cache(maxsize=None) -def grouped_gemm_quant_kernel() -> Callable: - """cuDNN CuTe DSL grouped GEMM kernel for block-scaled inputs.""" - from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module - - return grouped_gemm_quant_wrapper_sm100 - - -def convert_TE_MX_tensor_to_cuDNN_operand( +def _convert_to_cudnn_grouped_gemm_tensor_format( data: torch.Tensor, scale_inv: torch.Tensor, *, @@ -317,7 +337,7 @@ def convert_TE_MX_tensor_to_cuDNN_operand( sf_swizzled: bool = False, use_N_major_for_B: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Reshape an plain buffer into the layout cuDNN's grouped GEMM expects. + """Reshape a plain buffer into the layout cuDNN's grouped GEMM expects. cuDNN requirements: A: (valid_m, K, 1), K-major @@ -339,6 +359,12 @@ def convert_TE_MX_tensor_to_cuDNN_operand( SFB (unswizzled): (L, ceil(N/128), 4, 32, ceil(ceil(K/sf_vec_size)/4), 4) SFA (swizzled): (1, ceil(valid_m/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) SFB (swizzled): (L, ceil(N/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4) + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. This should not be used externally and once native + GEMM support is added then this function (and related helper + functions) should be removed entirely. + """ if use_N_major_for_B: @@ -404,7 +430,7 @@ def convert_TE_MX_tensor_to_cuDNN_operand( return data, scale_inv -def general_cuDNN_MX_gemm( +def _cudnn_grouped_gemm_nvfp4_ue5m3( A: torch.Tensor, B: torch.Tensor, out_dtype: Optional[torch.dtype] = None, @@ -417,7 +443,6 @@ def general_cuDNN_MX_gemm( layout: str = "TN", out: Optional[torch.Tensor] = None, bias: Optional[torch.Tensor] = None, - use_split_accumulator: bool = False, grad: bool = False, ub: Union[tex.CommOverlap, tex.CommOverlapP2P] = None, ub_type: tex.CommOverlapType = None, @@ -460,19 +485,13 @@ def general_cuDNN_MX_gemm( assert ( isinstance(A, NVFP4TensorStorage) and isinstance(B, NVFP4TensorStorage) - and A.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 - and B.get_metadata()["scale_dtype"] == DType.kFloat8UE5M3 + and A._scale_dtype == DType.kFloat8UE5M3 + and B._scale_dtype == DType.kFloat8UE5M3 ), "cuDNN MX GEMM is only used for NVFP4 GEMM with e5m3 scale factors for now." assert quantization_params is None, "cuDNN GEMM currently does not support output quantization." assert gelu is False and gelu_in is None, "cuDNN GEMM currently does not support fused GELU." - assert ( - use_split_accumulator is False - ), "cuDNN GEMM currently does not support split accumulators." - # use_split_accumulator is deliberately not checked: it is a cuBLAS knob for - # raising accumulator precision, and the cuDNN kernel always accumulates in - # FP32, so the request is already satisfied either way. assert ub is None and ub_type is None, "cuDNN GEMM currently does not support CommOverlap." assert extra_output is None, "cuDNN GEMM currently does not support extra output." assert bulk_overlap is False, "cuDNN GEMM currently does not support bulk overlap." @@ -566,7 +585,7 @@ def general_cuDNN_MX_gemm( assert beta in (1.0, None), "beta must be one or None if accumulate is True" else: # Overwrite GEMM's result to the out tensor assert beta in (0.0, None), "beta must be zero or None if not accumulate" - _cuDNN_wgrad_gemm( + _cudnn_wgrad_grouped_gemm_nvfp4_ue5m3( a_tensor=dataB.view(N, K // 2), b_tensor=dataA.view(M, K // 2), sfa=sfB, @@ -591,7 +610,7 @@ def general_cuDNN_MX_gemm( # cuDNN's own operand names are the other way round: its "a" is the (M, K) # activation-like operand (TE's B) and its "b" is the (N, K) weight-like one # (TE's A). - cudnn_a, cudnn_sfa = convert_TE_MX_tensor_to_cuDNN_operand( + cudnn_a, cudnn_sfa = _convert_to_cudnn_grouped_gemm_tensor_format( dataB, sfB, data_dtype=torch.float4_e2m1fn_x2, @@ -601,7 +620,7 @@ def general_cuDNN_MX_gemm( L=1, sf_swizzled=True, # ensured above ) - cudnn_b, cudnn_sfb = convert_TE_MX_tensor_to_cuDNN_operand( + cudnn_b, cudnn_sfb = _convert_to_cudnn_grouped_gemm_tensor_format( dataA, sfA, data_dtype=torch.float4_e2m1fn_x2, @@ -658,7 +677,7 @@ def general_cuDNN_MX_gemm( "discrete_col_sfd": False, "use_dynamic_sched": True, } - grouped_gemm_quant_kernel()(**gemm_kwargs) + _cudnn_grouped_gemm_quant_kernel()(**gemm_kwargs) # Matches general_gemm's contract: (out, bias_grad, gelu_input, extra_output). return out, None, None, None @@ -711,7 +730,7 @@ def general_gemm( and A._scale_dtype == DType.kFloat8UE5M3 and B._scale_dtype == DType.kFloat8UE5M3 ): - return general_cuDNN_MX_gemm( + return _cudnn_grouped_gemm_nvfp4_ue5m3( A, B, out_dtype, @@ -724,7 +743,6 @@ def general_gemm( layout, out, bias, - use_split_accumulator, grad, ub, ub_type, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 683a3c1dcb..1123e42095 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -20,7 +20,7 @@ from ...constants import DType, MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor -from ...cpp_extensions.gemm import convert_TE_MX_tensor_to_cuDNN_operand +from ...cpp_extensions.gemm import _convert_to_cudnn_grouped_gemm_tensor_format from ...distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, @@ -1758,7 +1758,7 @@ def fuser_forward( elif ( use_nvfp4 and fc2_input_sf_override is not None ): # TODO(kainingz): remove this e5m3 workaround once cuBLAS is ready. - fc2_x_data, fc2_x_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc2_x_data, fc2_x_scales = _convert_to_cudnn_grouped_gemm_tensor_format( grouped_fc2_x.rowwise_data, grouped_fc2_x.scale_inv, data_dtype=data_dtype, @@ -1813,7 +1813,7 @@ def fuser_forward( fc2_weight_for_gemm = grouped_fc2_weight.copy() tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) - fc2_w_data, fc2_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc2_w_data, fc2_w_scales = _convert_to_cudnn_grouped_gemm_tensor_format( fc2_weight_for_gemm.rowwise_data, fc2_weight_for_gemm.scale_inv, data_dtype=data_dtype, @@ -2630,7 +2630,7 @@ def fuser_backward( dgrad_valid_m = out_shape[0] # batch dim # Create A and its sf tensor for cuDNN that satisfies its layout requirements - fc1_dgrad_a_data, fc1_dgrad_a_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc1_dgrad_a_data, fc1_dgrad_a_scales = _convert_to_cudnn_grouped_gemm_tensor_format( grouped_fc1_dy.rowwise_data, grouped_fc1_dy.scale_inv, data_dtype=data_dtype, @@ -2683,7 +2683,7 @@ def fuser_backward( # Create B and its sf tensor for cuDNN that satisfies its layout # requirements. NVFP4 column-wise data is physically transposed, so # it is already (in_features, out_features) and stays K-major. - fc1_w_data, fc1_w_scales = convert_TE_MX_tensor_to_cuDNN_operand( + fc1_w_data, fc1_w_scales = _convert_to_cudnn_grouped_gemm_tensor_format( fc1_weight_for_gemm.columnwise_data, fc1_weight_for_gemm.columnwise_scale_inv, data_dtype=data_dtype, From 8e229f15b7d0a230bba3d2a39de490ab7c8e4992 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 03:45:39 +0000 Subject: [PATCH 24/54] Fix compilation error in C++ test Signed-off-by: Tim Moon --- tests/cpp/operator/test_dequantize_nvfp4.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index d8080db380..34a7d87122 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -362,7 +362,7 @@ TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) global_amax.rowwise_cpu_dptr()[0] = 12.0f; global_amax.from_cpu(); nvte_nvfp4_compute_global_scale( - global_amax.data(), global_scale.data(), 0, kNVTEFloat8UE5M3); + global_amax.data(), global_scale.data(), kNVTEFloat8UE5M3, 0); global_scale.to_cpu(); EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 114688.0f / 12.0f); @@ -372,14 +372,14 @@ TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) block_amax.rowwise_cpu_dptr()[1] = 6.0f; block_amax.from_cpu(); nvte_nvfp4_compute_per_block_scale( - block_amax.data(), block_scale.data(), global_amax.data(), 0, kNVTEFloat8UE5M3); + block_amax.data(), block_scale.data(), global_amax.data(), kNVTEFloat8UE5M3, 0); block_scale.to_cpu(); EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 114688.0f / 12.0f); EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 114688.0f / 12.0f); Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); nvte_nvfp4_expand_scale_to_fp8( - block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, 0, kNVTEFloat8UE5M3); + block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, kNVTEFloat8UE5M3, 0); expanded_scale.to_cpu(); const auto *scales = reinterpret_cast( expanded_scale.rowwise_cpu_dptr()); From c8e6ced4ccba7349709d17519763ed50dbf29c3d Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 05:48:37 +0000 Subject: [PATCH 25/54] Treat nvfp4_e4m3_max=0 as unset value Co-authored-by: Codex Signed-off-by: Tim Moon --- 3rdparty/cutlass | 2 +- tests/pytorch/test_recipe.py | 8 +++++--- transformer_engine/common/common.h | 14 +++++++------- .../transformer_engine/transformer_engine.h | 2 +- transformer_engine/common/transformer_engine.cpp | 3 ++- transformer_engine/pytorch/csrc/common.h | 4 ++-- .../pytorch/csrc/extensions/cast.cpp | 6 +++--- transformer_engine/pytorch/csrc/quantizer.cpp | 13 +++++-------- .../pytorch/csrc/type_converters.cpp | 6 +++--- .../pytorch/ops/fused/grouped_mlp.py | 6 +++--- transformer_engine/pytorch/quantization.py | 2 +- .../pytorch/tensor/grouped_tensor.py | 2 +- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 8 ++++---- .../tensor/storage/grouped_tensor_storage.py | 10 +++++----- .../pytorch/tensor/storage/nvfp4_tensor_storage.py | 4 ++-- 15 files changed, 45 insertions(+), 45 deletions(-) diff --git a/3rdparty/cutlass b/3rdparty/cutlass index 57e3cfb47a..75839a6fdf 160000 --- a/3rdparty/cutlass +++ b/3rdparty/cutlass @@ -1 +1 @@ -Subproject commit 57e3cfb47a2d9e0d46eb6335c3dc411498efa198 +Subproject commit 75839a6fdf112e1c1e6e50805c9b323ebb983d07 diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index ccef104a33..40983c8937 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -564,7 +564,7 @@ def expected_use_4over6(tensor_type): def expected_e4m3_max(tensor_type): if not expected_use_4over6(tensor_type): - return 448 + return 0 if nvfp4_4over6_e4m3_use_256 == "all": return 256 if nvfp4_4over6_e4m3_use_256 == "weights": @@ -573,7 +573,9 @@ def expected_e4m3_max(tensor_type): if nvfp4_4over6_e4m3_use_256 == "activations": if tensor_type != "weight": return 256 - return 448 + if nvfp4_4over6_e4m3_use_256 == "none": + return 448 + return 0 forward_quantizers = NVFP4BlockScalingRecipeState( recipe, @@ -624,7 +626,7 @@ def expected_e4m3_max(tensor_type): ).make_quantizers() assert [q.row_scaled_nvfp4 for q in backward_quantizers] == [False, False] assert [q.nvfp4_use_4over6 for q in backward_quantizers] == [False, False] - assert [q.nvfp4_e4m3_max for q in backward_quantizers] == [448, 448] + assert [q.nvfp4_e4m3_max for q in backward_quantizers] == [0, 0] assert [q.nvfp4_4over6_err_mode for q in backward_quantizers] == [nvfp4_4over6_err_mode] * 2 assert [q.stochastic_rounding for q in backward_quantizers] == [True, True] assert [q.with_rht for q in backward_quantizers] == [False, False] diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 3ed99a19ff..c455015cdc 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -306,13 +306,13 @@ struct Tensor { bool row_scaled_nvfp4 = false; /*! \brief Global scale bound used by NVFP4. * - * When negative, use the maximum value of the scale-inverse dtype. + * When zero, use the maximum value of the scale-inverse dtype. * Some 4over6 tensors use 256 (instead of the E4M3 max of 448) in * order to leave room for map-to-4 local scale expansion. * * TODO: Change to a dtype-agnostic name. */ - int nvfp4_e4m3_max = -1; + int nvfp4_e4m3_max = 0; /*! Map from NVTETensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { @@ -342,7 +342,7 @@ struct Tensor { scaling_mode = NVTE_DELAYED_TENSOR_SCALING; with_gemm_swizzled_scales = false; row_scaled_nvfp4 = false; - nvfp4_e4m3_max = -1; + nvfp4_e4m3_max = 0; } explicit operator NVTETensor() const noexcept { return nvte_tensor; } @@ -455,12 +455,12 @@ struct Tensor { /*! \brief Global scale bound used by NVFP4. */ int get_nvfp4_scale_max() const { - NVTE_CHECK(scaling_mode == NVTE_NVFP4_1D_SCALING, - "Attempted to access NVFP4 scale bound for tensor with scaling mode \"", - to_string(scaling_mode), "\"."); + if (scaling_mode != NVTE_NVFP4_1D_SCALING) { + return 0; + } // Return non-default scale max - if (nvfp4_e4m3_max >= 0) { + if (nvfp4_e4m3_max != 0) { return nvfp4_e4m3_max; } diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index d11387f4db..ba9b22f124 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -930,7 +930,7 @@ class TensorWrapper { } int get_nvfp4_e4m3_max() const { - int val = 448; + int val = 0; nvte_get_tensor_param_v2(tensor_, kNVTENVFP4E4M3Max, &val, sizeof(val), nullptr); return val; } diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 74a482e5e4..e941bf996a 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -929,7 +929,8 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo case kNVTENVFP4E4M3Max: std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); // Need to rename this to nvfp4_scale_type_max - NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || t.nvfp4_e4m3_max == 114688 || + NVTE_CHECK(t.nvfp4_e4m3_max == 0 || t.nvfp4_e4m3_max == 448 || + t.nvfp4_e4m3_max == 256 || t.nvfp4_e4m3_max == 114688 || t.nvfp4_e4m3_max == 65536, "Unsupported NVFP4 scale type max (got ", t.nvfp4_e4m3_max, ")"); break; diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 3eeeae5435..4623cd1a93 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -353,8 +353,8 @@ class NVFP4Quantizer : public Quantizer { bool stochastic_rounding; // 4over6 candidate-selection mode used when quantizing emitted NVFP4 tensors. NVTENVFP44Over6Mode nvfp4_4over6_mode; - // Global E4M3 scale bound used by emitted NVFP4 tensors. - std::optional nvfp4_e4m3_max; + // Global E4M3 scale bound used by emitted NVFP4 tensors (0 when inactive). + int nvfp4_e4m3_max = 0; // Dtype of scale_inv tensors (kFloat8E4M3 or kFloat8UE5M3). DType scale_dtype; // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index fd779c66d9..b78ffa5633 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1121,7 +1121,7 @@ std::tuple, std::vector, bool> bulk_alloc const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; - const auto nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; + 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 auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { @@ -1348,8 +1348,8 @@ std::tuple, std::vector, bool> bulk_alloc } tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4); - if (nvfp4_e4m3_max) { - tensor_wrapper.set_nvfp4_e4m3_max(*nvfp4_e4m3_max); + if (nvfp4_e4m3_max != 0) { + tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max); } tensor_cpp_list.emplace_back(std::move(tensor_wrapper)); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index db32be3aac..ac01d49d1d 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1887,10 +1887,7 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); this->stochastic_rounding = quantizer.attr("stochastic_rounding").cast(); const bool nvfp4_use_4over6 = quantizer.attr("nvfp4_use_4over6").cast(); - const int e4m3_max = quantizer.attr("nvfp4_e4m3_max").cast(); - if (e4m3_max >= 0) { - this->nvfp4_e4m3_max = e4m3_max; - } + this->nvfp4_e4m3_max = quantizer.attr("nvfp4_e4m3_max").cast(); const auto nvfp4_4over6_err_mode = quantizer.attr("nvfp4_4over6_err_mode").cast(); if (!nvfp4_use_4over6) { this->nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; @@ -2136,8 +2133,8 @@ std::pair NVFP4Quantizer::create_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); - if (this->nvfp4_e4m3_max) { - out_cpp.set_nvfp4_e4m3_max(*this->nvfp4_e4m3_max); + if (this->nvfp4_e4m3_max != 0) { + out_cpp.set_nvfp4_e4m3_max(this->nvfp4_e4m3_max); } this->set_quantization_params(&out_cpp); @@ -2463,8 +2460,8 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); - if (this->nvfp4_e4m3_max) { - out_cpp.set_nvfp4_e4m3_max(*this->nvfp4_e4m3_max); + if (this->nvfp4_e4m3_max != 0) { + out_cpp.set_nvfp4_e4m3_max(this->nvfp4_e4m3_max); } this->set_quantization_params(&out_cpp); diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 4825083f4f..7f377583d1 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -138,7 +138,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); const bool with_gemm_swizzled_scales = tensor.attr("_with_gemm_swizzled_scales").cast(); const bool row_scaled_nvfp4 = tensor.attr("_row_scaled_nvfp4").cast(); - const auto nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast>(); + const int nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast(); const DType scale_inv_dtype = tensor.attr("_scale_dtype").cast(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -174,8 +174,8 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) // Scale layout ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); ret.set_row_scaled_nvfp4(row_scaled_nvfp4); - if (nvfp4_e4m3_max) { - ret.set_nvfp4_e4m3_max(*nvfp4_e4m3_max); + if (nvfp4_e4m3_max != 0) { + ret.set_nvfp4_e4m3_max(nvfp4_e4m3_max); } // Quantizer state diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d553b6668a..2189b44426 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -306,9 +306,9 @@ def _nvfp4_sf_dtype_override(quantizer: Optional[Quantizer]) -> Literal["e5m3"] def _nvfp4_scale_max(quantizer: Quantizer) -> float: """Return the maximum representable magnitude of an NVFP4 scale factor.""" # 4over6 might override e4m3's max to 256 over default 448 - override_max = getattr(quantizer, "nvfp4_e4m3_max", None) - # NVFP4Quantizer's initialization sets nvfp4_e4m3_max to -1 if no override - if override_max is not None and override_max != -1: + override_max = getattr(quantizer, "nvfp4_e4m3_max", 0) + # NVFP4Quantizer's initialization sets nvfp4_e4m3_max to 0 if no override + if override_max != 0: return float(override_max) scale_dtype = getattr(quantizer, "scale_dtype", None) if scale_dtype is not None and scale_dtype == DType.kFloat8UE5M3: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 34e8f58d87..36f58b02d0 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1765,7 +1765,7 @@ def _make(tensor_type: str) -> NVFP4Quantizer: raise ValueError("NVFP4 4over6 quantization is incompatible with UE5M3 scales.") # Scale max for 4over6 - nvfp4_e4m3_max = None + nvfp4_e4m3_max = 0 if nvfp4_use_4over6: if self.recipe.nvfp4_4over6_e4m3_use_256 == "all": nvfp4_e4m3_max = 256 diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 786316db30..65b4a84822 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -95,7 +95,7 @@ def __new__( with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, scale_inv_dtype: Optional[DType] = None, ): if ( diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 8808c642f3..ac13598b75 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -157,7 +157,7 @@ def __init__( stochastic_rounding: bool = False, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: Optional[int] = None, + nvfp4_e4m3_max: int = 0, nvfp4_4over6_err_mode: str = "MAE", with_random_sign_mask: bool = True, disable_second_level_scale: bool = False, @@ -187,7 +187,7 @@ def __init__( raise ValueError( "nvfp4_use_4over6 is incompatible with scale_dtype=DType.kFloat8UE5M3." ) - self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_e4m3_max is not None else -1 + self.nvfp4_e4m3_max = nvfp4_e4m3_max 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'.") @@ -479,7 +479,7 @@ def __new__( with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: Optional[int] = None, + nvfp4_e4m3_max: int = 0, **kwargs, ): instance = super().__new__( @@ -1063,7 +1063,7 @@ def _make_nvfp4_tensor_in_reduce_ex( with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: Optional[int] = None, + nvfp4_e4m3_max: int = 0, scale_dtype: DType = DType.kFloat8E4M3, ) -> NVFP4Tensor: """Reconstruct an ``NVFP4Tensor`` from its ``__reduce_ex__`` payload.""" diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index ba0d74e87c..3be8d8a337 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -77,7 +77,7 @@ def _initialize_storage_fields( with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, ) -> None: """ Initialize a GroupedTensor. @@ -158,7 +158,7 @@ def _initialize_storage_fields( instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance.row_scaled_nvfp4 = row_scaled_nvfp4 instance.nvfp4_use_4over6 = nvfp4_use_4over6 - instance.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 + instance.nvfp4_e4m3_max = nvfp4_e4m3_max def __new__( cls, @@ -187,7 +187,7 @@ def __new__( with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, ): instance = object.__new__(cls) cls._initialize_storage_fields( @@ -470,7 +470,7 @@ def clear(self) -> None: self.fake_dtype = torch.float32 self.row_scaled_nvfp4 = False self.nvfp4_use_4over6 = False - self.nvfp4_e4m3_max = 448 + self.nvfp4_e4m3_max = 0 def __repr__(self) -> str: """String representation of the GroupedTensorStorage.""" @@ -788,7 +788,7 @@ def make_grouped_tensor( columnwise_scale_inv_offsets = None row_scaled_nvfp4 = False nvfp4_use_4over6 = False - nvfp4_e4m3_max = 448 + nvfp4_e4m3_max = 0 if no_quantization: assert dtype is not None, "dtype must be provided for unquantized GroupedTensor" if rowwise_usage: diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 69e8aace0b..cd1c15b943 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -104,7 +104,7 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Whether this NVFP4 tensor uses 4over6 map-to-4/map-to-6 block selection _nvfp4_use_4over6: bool # Global E4M3 scale bound used by this NVFP4 tensor - _nvfp4_e4m3_max: Optional[int] + _nvfp4_e4m3_max: int def __new__( cls, @@ -122,7 +122,7 @@ def __new__( fake_dtype: Optional[torch.dtype] = None, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: Optional[int] = None, + nvfp4_e4m3_max: int = 0, **kwargs, ): if cls is NVFP4TensorStorage: From 6e8481642c5931a0593873d25dd49293d80d058a Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 06:37:35 +0000 Subject: [PATCH 26/54] Debug torch.compile test failure Signed-off-by: Tim Moon --- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ac13598b75..3b4c7fd2d8 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -378,6 +378,7 @@ def storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]: "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor, "nontensor_kwargs": { "fp4_dtype": self.dtype, + "scale_dtype": self.scale_dtype, "quantizer": self, "with_gemm_swizzled_scales": self.optimize_for_gemm, "row_scaled_nvfp4": self.row_scaled_nvfp4, @@ -474,7 +475,7 @@ def __new__( amax_rowwise: Optional[torch.Tensor], amax_columnwise: Optional[torch.Tensor], fp4_dtype: DType, - scale_dtype: DType, + scale_dtype: DType = DType.kFloat8E4M3, quantizer: Quantizer, with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, From 0c2de5e0a5fb1d9debf631ed1bf160dc15554e02 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 07:39:36 +0000 Subject: [PATCH 27/54] Debug test failures Signed-off-by: Tim Moon --- transformer_engine/common/common.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index c455015cdc..62c21fe42a 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -471,7 +471,7 @@ struct Tensor { } else if (columnwise_scale_inv.has_data()) { dtype = columnwise_scale_inv.dtype; } else { - dtype = scale_inv.dtype; + dtype = DType::kFloat8E4M3; } switch (dtype) { case DType::kFloat8E4M3: diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index ac01d49d1d..bc4d01ae96 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2096,6 +2096,7 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); kwargs["nvfp4_e4m3_max"] = py::cast(this->nvfp4_e4m3_max); + kwargs["scale_inv_dtype"] = MakePythonDType(DType::kFloat8E8M0); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), args.ptr(), kwargs.ptr()); From dd90509ac628d13e4db268a36bb02784fd19abd0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 07:41:39 +0000 Subject: [PATCH 28/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/recipe/nvfp4.cu | 3 +-- transformer_engine/common/transformer_engine.cpp | 5 ++--- transformer_engine/pytorch/cpp_extensions/gemm.py | 1 + transformer_engine/pytorch/csrc/extensions/transpose.cpp | 3 ++- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 2807eb9703..f1666653d1 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -916,8 +916,7 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, NVTEDType scale_dtype, - cudaStream_t stream) { + size_t block_len, NVTEDType scale_dtype, cudaStream_t stream) { #if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); using namespace transformer_engine; diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index e941bf996a..df30a86192 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -929,9 +929,8 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo case kNVTENVFP4E4M3Max: std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); // Need to rename this to nvfp4_scale_type_max - NVTE_CHECK(t.nvfp4_e4m3_max == 0 || t.nvfp4_e4m3_max == 448 || - t.nvfp4_e4m3_max == 256 || t.nvfp4_e4m3_max == 114688 || - t.nvfp4_e4m3_max == 65536, + NVTE_CHECK(t.nvfp4_e4m3_max == 0 || t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || + t.nvfp4_e4m3_max == 114688 || t.nvfp4_e4m3_max == 65536, "Unsupported NVFP4 scale type max (got ", t.nvfp4_e4m3_max, ")"); break; default: diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index ff0c72f2ce..1cdab8e09c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -245,6 +245,7 @@ def _cudnn_grouped_gemm_wgrad_kernel() -> Callable: return grouped_gemm_wgrad_wrapper_sm100 + def _cudnn_wgrad_grouped_gemm_nvfp4_ue5m3( a_tensor: torch.Tensor, b_tensor: torch.Tensor, diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 1806208c09..c6817af5aa 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -142,7 +142,8 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile nvte_nvfp4_expand_scale_to_fp8(input_cu.data(), output_cu.data(), static_cast(tile_rows), static_cast(tile_cols), static_cast(rows_padded), - static_cast(block_len), static_cast(scale_dtype), + static_cast(block_len), + static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); } From d42ccbe50c88e0b919f7a9d5f7a30bbd51ea778d Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 11:22:37 +0000 Subject: [PATCH 29/54] Revert accidental CUTLASS commit change Signed-off-by: Tim Moon --- 3rdparty/cutlass | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cutlass b/3rdparty/cutlass index 75839a6fdf..57e3cfb47a 160000 --- a/3rdparty/cutlass +++ b/3rdparty/cutlass @@ -1 +1 @@ -Subproject commit 75839a6fdf112e1c1e6e50805c9b323ebb983d07 +Subproject commit 57e3cfb47a2d9e0d46eb6335c3dc411498efa198 From b47e1f88b71134da6853d42c57ad5810f9802e68 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 11:44:18 +0000 Subject: [PATCH 30/54] Enable cuDNN GGEMM+GLU+RHT+quant kernel Signed-off-by: Tim Moon --- .../pytorch/ops/fused/grouped_mlp.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 2189b44426..e223a4683c 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1488,8 +1488,15 @@ def fuser_forward( and fc2_input_quantizer.rht_matrix_random_sign_mask_t == 0 ): if fc2_input_quantizer.disable_second_level_scale: - # Use GEMM + act + RHT + quant kernel once available - pass + # Use GEMM + act + RHT + quant kernel if available + if self.grouped_gemm_act_hadamard_quant_kernel() is None: + # Kernel is not available + pass + if fc1_bias_packed is not None: + # Kernel has large numerical error with bias + pass + elif self._cudnn_act_func == "swiglu": + kernel_impl = "gemm_act_rht_quant" elif fc2_input_quantizer.with_post_rht_amax: # Use GEMM + act + RHT + amax kernel if available if self.grouped_gemm_act_hadamard_kernel() is None: @@ -1535,7 +1542,7 @@ def fuser_forward( fc1_activation_kwargs["use_tmem_post_rht_amax"] = _use_tmem_post_rht_amax() elif kernel_impl == "gemm_act_rht_quant": fc1_activation_kwargs["d_dtype"] = torch.float4_e2m1fn_x2 - fc1_activation_kwargs["rht_dtype"] = torch.float4_e2m1fn_x2 + fc1_activation_kwargs["rht_colwise_dtype"] = torch.float4_e2m1fn_x2 # Kernel arguments based on activation if activation_is_srelu: @@ -1707,9 +1714,8 @@ def fuser_forward( fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1] // 2) fc2_in_row_scale = fc1_kernel_out["sfd_tensor"] fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) - fc2_in_col_data = fc1_kernel_out["rht_tensor"] - fc2_in_col_data = fc2_in_col_data.permute(1, 0) - fc2_in_col_scale = fc1_kernel_out["sfrht_tensor"] + fc2_in_col_data = fc1_kernel_out["rht_colwise_tensor"] + fc2_in_col_scale = fc1_kernel_out["sfrht_colwise_tensor"] grouped_fc2_x = GroupedTensorStorage( shape=(in_shape[0], fc2_weight_shape[1]), dtype=dtype, From ae3020fed9ab197d9c29479003b69c62c4857374 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 22 Aug 2026 22:40:12 +0000 Subject: [PATCH 31/54] Remove incorrect scale_inv_dtype arg to NVFP4Tensor constructor Signed-off-by: Tim Moon --- transformer_engine/pytorch/csrc/quantizer.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index bc4d01ae96..ac01d49d1d 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2096,7 +2096,6 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); kwargs["nvfp4_e4m3_max"] = py::cast(this->nvfp4_e4m3_max); - kwargs["scale_inv_dtype"] = MakePythonDType(DType::kFloat8E8M0); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), args.ptr(), kwargs.ptr()); From 1d0d56fd1e5ea54adc77983483b3e55f68a798e6 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:51:51 -0700 Subject: [PATCH 32/54] Fix bug when selecting cuDNN GGEMM+GLU+RHT+quant kernel Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index e223a4683c..8ec3921206 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1492,7 +1492,7 @@ def fuser_forward( if self.grouped_gemm_act_hadamard_quant_kernel() is None: # Kernel is not available pass - if fc1_bias_packed is not None: + elif fc1_bias_packed is not None: # Kernel has large numerical error with bias pass elif self._cudnn_act_func == "swiglu": From b8fa24def0dc4e2bcb0e9b26a4c0305eb1bd214d Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 24 Aug 2026 23:45:40 +0000 Subject: [PATCH 33/54] Restore GGEMM+GLU+RHT+amax kernel with RHT sign mask Signed-off-by: Tim Moon --- transformer_engine/pytorch/ops/fused/grouped_mlp.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 8ec3921206..53dde0dacc 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1485,13 +1485,15 @@ def fuser_forward( use_nvfp4 and isinstance(fc2_input_quantizer, NVFP4Quantizer) and fc2_input_quantizer.with_rht - and fc2_input_quantizer.rht_matrix_random_sign_mask_t == 0 ): if fc2_input_quantizer.disable_second_level_scale: # Use GEMM + act + RHT + quant kernel if available if self.grouped_gemm_act_hadamard_quant_kernel() is None: # Kernel is not available pass + elif fc2_input_quantizer.rht_matrix_random_sign_mask_t != 0: + # Kernel does not apply sign mask in RHT + pass elif fc1_bias_packed is not None: # Kernel has large numerical error with bias pass From 30b8af8812200e97604ffaee1a3822fe1b8d65d8 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 24 Aug 2026 23:47:10 +0000 Subject: [PATCH 34/54] Debug minor test failures Signed-off-by: Tim Moon --- .../common/cast/dispatch/quantize.cuh | 13 +++++++++---- transformer_engine/pytorch/csrc/quantizer.cpp | 1 + .../pytorch/custom_recipes/reference_nvfp4.py | 4 ++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 04706a40c0..c4c9da0111 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -458,12 +458,17 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; if (!nvfp4_use_4over6) { for (const auto *output_tensor : output_tensors) { - const DType scale_dtype = output_tensor->scale_inv.has_data() - ? output_tensor->scale_inv.dtype - : output_tensor->columnwise_scale_inv.dtype; + DType scale_dtype = DType::kFloat8E4M3; + if (output_tensor->scale_inv.has_data()) { + scale_dtype = output_tensor->scale_inv.dtype; + } else if (output_tensor->columnwise_scale_inv.has_data()) { + scale_dtype = output_tensor->columnwise_scale_inv.dtype; + } NVTE_CHECK( static_cast(output_tensor->get_nvfp4_scale_max()) == typeToMax(scale_dtype), - "NVFP4 quantization with non-default scale max is only supported with 4over6."); + "NVFP4 quantization with non-default scale max is only supported with 4over6 " + "(expected ", typeToMax(scale_dtype), ", found ", + output_tensor->get_nvfp4_scale_max(), ")."); } } NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index ac01d49d1d..dcc2ba2000 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1712,6 +1712,7 @@ std::pair MXFP8Quantizer::create_grouped_tenso kwargs["last_dims"] = last_dims.has_value() ? py::cast(*last_dims) : py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; + kwargs["scale_inv_dtype"] = MakePythonDType(DType::kFloat8E8M0); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index 3aae1e9d82..8c5f305554 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -353,7 +353,7 @@ def __init__( quant_tile_shape: Tuple[int, int] = (1, 16), row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, nvfp4_4over6_err_mode: str = "MAE", nvfp4_4over6_err_use_fast_math: bool = False, with_rht: bool = False, @@ -379,7 +379,7 @@ def __init__( self.quant_tile_shape = quant_tile_shape 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 + self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_e4m3_max != 0 else 448 if self.nvfp4_e4m3_max not in (448, 256): raise ValueError("nvfp4_e4m3_max must be 448 or 256.") self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode From d9b633db02e6a34c7034bde28d1c069722aebf0c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:51:49 +0000 Subject: [PATCH 35/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/dispatch/quantize.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index c4c9da0111..bb3d9f1028 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -467,8 +467,8 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou NVTE_CHECK( static_cast(output_tensor->get_nvfp4_scale_max()) == typeToMax(scale_dtype), "NVFP4 quantization with non-default scale max is only supported with 4over6 " - "(expected ", typeToMax(scale_dtype), ", found ", - output_tensor->get_nvfp4_scale_max(), ")."); + "(expected ", + typeToMax(scale_dtype), ", found ", output_tensor->get_nvfp4_scale_max(), ")."); } } NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, From 75bdfe4af83bfa4aefc03e46e0270abd58bbc349 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 25 Aug 2026 07:33:12 +0000 Subject: [PATCH 36/54] Fix incorrect scale dtypes in grouped tensor builder method Signed-off-by: Tim Moon --- .../pytorch/tensor/storage/grouped_tensor_storage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 3be8d8a337..bbd62175ef 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -802,7 +802,7 @@ def make_grouped_tensor( # Amax buffer for delayed scaling - one per tensor amax = torch.empty(num_tensors, dtype=torch.float32, device=device) - scale_inv_dtype = DType.kFloat32 + scale_inv_dtype = DType.kFloat8E8M0 if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) @@ -833,7 +833,7 @@ def make_grouped_tensor( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) elif compatible_recipe.delayed(): - scale_inv_dtype = DType.kFloat8E8M0 + scale_inv_dtype = DType.kFloat32 if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8) data = torch.empty(total_elements, dtype=torch.uint8, device=device) From 58dfc41046b7068dffc29eba80ecf65005988009 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 25 Aug 2026 01:41:43 -0700 Subject: [PATCH 37/54] Avoid redundant amax ptr check in row-scaled NVFP4 quantize Review suggestion from @ptrendx Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/common/cast/dispatch/quantize.cuh | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index bb3d9f1028..5f6973a1c0 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -127,9 +127,7 @@ 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."); - if (output_tensor->amax.dptr != nullptr) { - nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream); - } + 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); @@ -313,9 +311,7 @@ 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."); - if (output_tensor->amax.dptr != nullptr) { - nvfp4::compute_rowwise_amax(*grad_tensor, noop_tensor, output_tensor, stream); - } + 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); From b1a4483f88ce886ace74d950ae0fa4577f232718 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 28 Aug 2026 07:38:27 +0000 Subject: [PATCH 38/54] Clean up C++ unit tests Signed-off-by: Tim Moon --- .../cpp/operator/test_cast_nvfp4_transpose.cu | 152 +++++++----------- tests/cpp/operator/test_dequantize_nvfp4.cu | 150 ++++++++++------- tests/cpp/test_common.cu | 71 ++++---- tests/cpp/test_common.h | 19 ++- 4 files changed, 199 insertions(+), 193 deletions(-) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 1527142d21..7e0873ae8f 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -63,45 +63,39 @@ std::vector create_transpose(const InputType* const input, const size } template -constexpr float nvfp4_encode_scale_max(const int scale_type_max = 448) { +constexpr float get_scale_max(int scale_max = 0) { static_assert(std::is_same_v #if CUDA_VERSION >= 13040 || std::is_same_v #endif , "Unsupported NVFP4 scale type."); if constexpr (std::is_same_v) { - NVTE_CHECK(scale_type_max == 448 || scale_type_max == 256, - "Unsupported E4M3 scale maximum."); - return static_cast(scale_type_max); -#if CUDA_VERSION >= 13040 - } else { - NVTE_CHECK(scale_type_max == 114688 || scale_type_max == 65536, - "Unsupported UE5M3 scale maximum."); - return static_cast(scale_type_max); -#endif - } -} - -template -constexpr float nvfp4_scale_storage_max() { - if constexpr (std::is_same_v) { - return 448.0f; + if (scale_max == 0) { + scale_max = 448; + } + return static_cast(scale_max); } #if CUDA_VERSION >= 13040 - return 114688.0f; + if constexpr (std::is_same_v) { + if (scale_max == 0) { + scale_max = 114688; + } + return static_cast(scale_max); #endif + } + return 0.f; } // Compute the global encode scale factor for a given global amax. template float compute_global_encode_scaling_factor_FP4(const float global_amax, const bool use_fast_math, - const int scale_type_max = 448) { - const float fp8_max = nvfp4_encode_scale_max(scale_type_max); + const int scale_max = 0) { + const float fp8_max = get_scale_max(scale_max); constexpr float fp4_max = 6.0f; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; // If scale is infinity, return the max normalized value - const float max_norm_clamp = (use_fast_math && std::is_same_v - && scale_type_max == 448) + const float max_norm_clamp = (use_fast_math + && fp8_max == get_scale_max()) ? Numeric_Traits::maxNorm : Numeric_Traits::maxNorm; @@ -136,7 +130,7 @@ enum class NVFP4ScalingMode { struct NVFP4FourOverSixTestConfig { NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled; - int scale_type_max = 448; + int scale_max = 0; bool err_use_fast_math = false; }; @@ -146,9 +140,9 @@ bool use_2d_quantization(const NVFP4ScalingMode scaling_mode) { template NVFP4FourOverSixQuantization compute_4over6_quantization_scales( - const float block_amax, const float global_encode_scale, const int scale_type_max) { + const float block_amax, const float global_encode_scale, const int scale_max) { constexpr float fp4_max = 6.0f; - const float fp8_max = nvfp4_scale_storage_max(); + const float fp8_max = get_scale_max(); constexpr float scale_expansion_factor = 1.5f; const float base_sf_high_precision = block_amax / fp4_max * global_encode_scale; const float sf_high_precision_map4 = @@ -228,13 +222,13 @@ void quantize_nvfp4_1d(float (*OP)(const float), const float global_amax, const bool use_fast_math, const bool use_4over6 = false, - const int scale_type_max = 448, + const int scale_max = 0, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { // Compute a global encoding/decoding scaling factor for all S_dec_b const float S_enc = compute_global_encode_scaling_factor_FP4( - global_amax, use_fast_math, scale_type_max); + global_amax, use_fast_math, scale_max); constexpr size_t block_size_X = 16; const size_t blocks_X = divide_round_up(cols, block_size_X); @@ -269,7 +263,7 @@ void quantize_nvfp4_1d(float (*OP)(const float), if (use_4over6) { const NVFP4FourOverSixQuantization quantization = compute_4over6_quantization_scales(block_amax, S_enc, - scale_type_max); + scale_max); scales[scale_idx] = select_4over6_scale(quantization, four_over_six_candidate); for (size_t j = j_min; j < j_max; j += 2) { @@ -333,12 +327,12 @@ void compute_2d_mathematical_scales(float (*OP)(const float), std::vector>& math_scales, const bool use_fast_math, const bool use_4over6 = false, - const int scale_type_max = 448, + const int scale_max = 0, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { const float S_enc = compute_global_encode_scaling_factor_FP4( - global_amax, use_fast_math, scale_type_max); + global_amax, use_fast_math, scale_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); @@ -369,7 +363,7 @@ void compute_2d_mathematical_scales(float (*OP)(const float), if (use_4over6) { const NVFP4FourOverSixQuantization quantization = compute_4over6_quantization_scales( - block_amax, S_enc, scale_type_max); + block_amax, S_enc, scale_max); math_scales[block_Y][block_X] = select_4over6_scale(quantization, four_over_six_candidate); } else { @@ -393,7 +387,7 @@ void quantize_nvfp4_2d(float (*OP)(const float), const float global_amax, const bool use_fast_math, const bool use_4over6 = false, - const int scale_type_max = 448, + const int scale_max = 0, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { @@ -401,10 +395,10 @@ void quantize_nvfp4_2d(float (*OP)(const float), std::vector> math_scales; compute_2d_mathematical_scales( OP, input, rows, cols, global_amax, math_scales, use_fast_math, - use_4over6, scale_type_max, four_over_six_candidate); + use_4over6, scale_max, four_over_six_candidate); const float S_enc = compute_global_encode_scaling_factor_FP4( - global_amax, use_fast_math, scale_type_max); + global_amax, use_fast_math, scale_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); @@ -488,17 +482,17 @@ void quantize_nvfp4(float (*OP)(const float), const bool use_fast_math, const bool use_2d_quantization = false, const bool use_4over6 = false, - const int scale_type_max = 448, + const int scale_max = 0, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { if (use_2d_quantization) { quantize_nvfp4_2d( OP, input, output, scales, rows, cols, scales_stride, global_amax, - use_fast_math, use_4over6, scale_type_max, four_over_six_candidate); + use_fast_math, use_4over6, scale_max, four_over_six_candidate); } else { quantize_nvfp4_1d( OP, input, output, scales, rows, cols, scales_stride, global_amax, - use_fast_math, use_4over6, scale_type_max, four_over_six_candidate); + use_fast_math, use_4over6, scale_max, four_over_six_candidate); } } @@ -518,7 +512,7 @@ void compute_ref(float (*OP)(const float), const bool use_2d_quantization = false, const bool row_scaled_nvfp4 = false, const bool use_4over6 = false, - const int scale_type_max = 448, + const int scale_max = 0, const NVFP4FourOverSixCandidate four_over_six_candidate = NVFP4FourOverSixCandidate::Map6) { @@ -532,7 +526,7 @@ void compute_ref(float (*OP)(const float), std::vector> math_scales; compute_2d_mathematical_scales( OP, input, rows, cols, *amax, math_scales, use_fast_math, - use_4over6, scale_type_max, four_over_six_candidate); + use_4over6, scale_max, four_over_six_candidate); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; @@ -561,11 +555,11 @@ void compute_ref(float (*OP)(const float), // (This part processes the actual FP4 data using the mathematical scaling factors) quantize_nvfp4_2d( OP, input, output, nullptr, rows, cols, scales_stride, *amax, - use_fast_math, use_4over6, scale_type_max, + use_fast_math, use_4over6, scale_max, four_over_six_candidate); // scales already filled quantize_nvfp4_2d( OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, *amax, - use_fast_math, use_4over6, scale_type_max, + use_fast_math, use_4over6, scale_max, four_over_six_candidate); // scales_t already filled return; @@ -585,7 +579,7 @@ void compute_ref(float (*OP)(const float), use_fast_math, use_2d_quantization, use_4over6, - scale_type_max, + scale_max, four_over_six_candidate); } return; @@ -594,11 +588,11 @@ void compute_ref(float (*OP)(const float), // Ref impl for basic NVFP4 quantize_nvfp4( OP, input, output, scales, rows, cols, scales_stride, *amax, - use_fast_math, use_2d_quantization, use_4over6, scale_type_max, + use_fast_math, use_2d_quantization, use_4over6, scale_max, four_over_six_candidate); quantize_nvfp4( OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, *amax, - use_fast_math, use_2d_quantization, use_4over6, scale_type_max, + use_fast_math, use_2d_quantization, use_4over6, scale_max, four_over_six_candidate); } @@ -848,7 +842,7 @@ void performTest(float (*OP)(const float), const bool use_fast_math, const NVFP4ScalingMode scaling_mode = NVFP4ScalingMode::Block1D, const NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled, - const int scale_type_max = 448, + const int scale_max = 0, const bool use_4over6_err_use_fast_math = false) { using namespace test; const bool use_4over6 = mode != kNVTENVFP44Over6Disabled; @@ -892,7 +886,7 @@ void performTest(float (*OP)(const float), Tensor input("input", shape, itype); Tensor output("output", shape, otype, rowwise, columnwise, NVTE_NVFP4_1D_SCALING, scale_type); - output.set_nvfp4_e4m3_max(scale_type_max); + output.set_nvfp4_e4m3_max(scale_max); std::unique_ptr ref_output = std::make_unique(rows * (cols / 2)); std::unique_ptr ref_output_t = std::make_unique(cols * (rows / 2)); @@ -909,7 +903,7 @@ void performTest(float (*OP)(const float), if (use_4over6 && row_scaled_nvfp4) { const float target_row_amax = - nvfp4_encode_scale_max(scale_type_max) * 6.0f * 8.0f; + get_scale_max(scale_max) * 6.0f * 8.0f; auto *input_vals = input.rowwise_cpu_dptr(); for (size_t row = 0; row < rows; ++row) { float row_amax = 0.0f; @@ -961,7 +955,7 @@ void performTest(float (*OP)(const float), } else { // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues ref_amax.assign( - 1, nvfp4_encode_scale_max(scale_type_max) * 6.0f * 8.0f); + 1, get_scale_max(scale_max) * 6.0f * 8.0f); // Update tensor if (rowwise) { @@ -994,7 +988,7 @@ void performTest(float (*OP)(const float), is_2d_quantization, row_scaled_nvfp4, use_4over6, - scale_type_max, + scale_max, NVFP4FourOverSixCandidate::Map4); compute_ref(OP, input.rowwise_cpu_dptr(), @@ -1011,7 +1005,7 @@ void performTest(float (*OP)(const float), is_2d_quantization, row_scaled_nvfp4, use_4over6, - scale_type_max, + scale_max, NVFP4FourOverSixCandidate::Map6); } else { compute_ref(OP, @@ -1029,7 +1023,7 @@ void performTest(float (*OP)(const float), is_2d_quantization, row_scaled_nvfp4, use_4over6, - scale_type_max); + scale_max); } // Initialize stochastic rounding @@ -1406,12 +1400,12 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, { if (scale_type == DType::kFloat8E4M3) { performTest( - OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_type_max, + OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max, config.err_use_fast_math); #if CUDA_VERSION >= 13040 } else if (scale_type == DType::kFloat8UE5M3) { performTest( - OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_type_max, + OP, tensor_dims, use_fast_math, scaling_mode, config.mode, config.scale_max, config.err_use_fast_math); #endif } else { @@ -1456,7 +1450,7 @@ std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param) const NVFP4FourOverSixTestConfig& config = std::get<5>(param); if (config.mode != kNVTENVFP44Over6Disabled) { name += "X4OVER6"; - name += "XSCALE_MAX_" + std::to_string(config.scale_type_max); + name += "XSCALE_MAX_" + std::to_string(config.scale_max); if (config.mode == kNVTENVFP44Over6MinMSE) { name += "XMSE"; } else if (config.mode == kNVTENVFP44Over6MinMAE) { @@ -1582,47 +1576,15 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( OperatorTestUE5M3, FusedCastTransposeNVFP4TestSuite, - ::testing::Values( - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block1D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block2D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block1D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 114688, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block2D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 114688, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block1D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 65536, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kBFloat16, false, - NVFP4ScalingMode::Block2D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 65536, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kFloat32, false, - NVFP4ScalingMode::Block1D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, - DType::kFloat8UE5M3}, - FusedCastTransposeNVFP4TestSuite::ParamType{ - ActivationType::Identity, {256, 256}, DType::kFloat32, false, - NVFP4ScalingMode::RowScaled1D, - NVFP4FourOverSixTestConfig{kNVTENVFP44Over6Disabled, 114688, false}, - DType::kFloat8UE5M3}), + ::testing::Combine( + ::testing::Values(ActivationType::Identity), // activation_dtype + ::testing::ValuesIn(tensor_dims), // tensor_dims + ::testing::Values(DType::kBFloat16, DType::kFloat32), // input_type + ::testing::Values(false), // use_fast_math + ::testing::Values(NVFP4ScalingMode::Block1D, + NVFP4ScalingMode::Block2D), // scaling_mode + ::testing::Values(NVFP4FourOverSixTestConfig{}), // four_over_six_config + ::testing::Values(DType::kFloat8UE5M3)), [](const testing::TestParamInfo& info) { return test_name(info.param); }); diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index 34a7d87122..b56890e968 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -88,9 +89,33 @@ float compute_amax(test::Tensor &t, size_t rows, size_t cols) { return amax; } +template +constexpr float get_scale_max(int scale_max = 0) { + static_assert(std::is_same_v +#if CUDA_VERSION >= 13040 + || std::is_same_v +#endif + , "Unsupported NVFP4 scale type."); + if constexpr (std::is_same_v) { + if (scale_max == 0) { + scale_max = 448; + } + return static_cast(scale_max); + } +#if CUDA_VERSION >= 13040 + if constexpr (std::is_same_v) { + if (scale_max == 0) { + scale_max = 114688; + } + return static_cast(scale_max); + } +#endif + return 0.f; +} + struct NVFP4DequantizeTestConfig { NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled; - int e4m3_max = 448; + int scale_max = 0; }; // Quantize a high-precision input to NVFP4, then dequantize and compare @@ -99,7 +124,7 @@ template void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, const bool row_scaled_nvfp4, const NVTENVFP44Over6Mode mode, - const int e4m3_max) { + const int scale_max) { using namespace test; DType otype = TypeInfo::dtype; @@ -115,8 +140,10 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, // Configure quantized tensor amax size_t amax_size = 1; - quantized.set_nvfp4_e4m3_max(e4m3_max); - ASSERT_EQ(quantized.nvfp4_e4m3_max(), e4m3_max); + if (scale_max != 0) { + quantized.set_nvfp4_e4m3_max(scale_max); + ASSERT_EQ(quantized.nvfp4_e4m3_max(), scale_max); + } if (row_scaled_nvfp4) { quantized.set_row_scaled_nvfp4(true); amax_size = rows; @@ -158,10 +185,9 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; std::unique_ptr ref_output = std::make_unique(rows * cols); - const float scale_max = static_cast(e4m3_max); compute_ref_dequantize_nvfp4( fp4_data, scales, amax_vals, ref_output.get(), - rows, cols, scale_stride, scale_max); + rows, cols, scale_stride, get_scale_max(scale_max)); // Compare results from TE and reference impls auto [atol, rtol] = getTolerances(otype); @@ -173,7 +199,7 @@ template void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, const bool row_scaled_nvfp4, const NVTENVFP44Over6Mode mode, - const int e4m3_max) { + const int scale_max) { using namespace test; DType otype = TypeInfo::dtype; @@ -183,8 +209,10 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, TypeInfo::dtype); - quantized_compact.set_nvfp4_e4m3_max(e4m3_max); - ASSERT_EQ(quantized_compact.nvfp4_e4m3_max(), e4m3_max); + if (scale_max != 0) { + quantized_compact.set_nvfp4_e4m3_max(scale_max); + ASSERT_EQ(quantized_compact.nvfp4_e4m3_max(), scale_max); + } if (row_scaled_nvfp4) { quantized_compact.set_row_scaled_nvfp4(true); } else if (rows > 0 && cols > 0) { @@ -209,8 +237,10 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING, TypeInfo::dtype); - quantized_swizzled.set_nvfp4_e4m3_max(e4m3_max); - ASSERT_EQ(quantized_swizzled.nvfp4_e4m3_max(), e4m3_max); + if (scale_max != 0) { + quantized_swizzled.set_nvfp4_e4m3_max(scale_max); + ASSERT_EQ(quantized_swizzled.nvfp4_e4m3_max(), scale_max); + } if (row_scaled_nvfp4) { quantized_swizzled.set_row_scaled_nvfp4(true); } else { @@ -286,7 +316,22 @@ class DequantizeNVFP4TestSuite : public ::testing::TestWithParam , transformer_engine::DType, bool, - NVFP4DequantizeTestConfig>> {}; + NVFP4DequantizeTestConfig, + transformer_engine::DType>> { +public: + static std::string test_name(const testing::TestParamInfo &info) { + const NVFP4DequantizeTestConfig config = std::get<3>(info.param); + const bool use_4over6 = config.mode != kNVTENVFP44Over6Disabled; + std::string name = std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "X" + + test::typeName(std::get<1>(info.param)) + "X" + + (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + + (use_4over6 ? "FourOverSix" : "Default") + "X" + + "ScaleMax" + std::to_string(config.scale_max) + "X" + + "Scale" + test::typeName(std::get<4>(info.param)); + return name; + } +}; TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) { @@ -298,11 +343,25 @@ TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) const DType output_type = std::get<1>(GetParam()); const bool row_scaled_nvfp4 = std::get<2>(GetParam()); const NVFP4DequantizeTestConfig config = std::get<3>(GetParam()); + const DType scale_type = std::get<4>(GetParam()); TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, - performTest_dequantize_nvfp4( + switch (scale_type) { + case transformer_engine::DType::kFloat8E4M3: + performTest_dequantize_nvfp4( tensor_size.first, tensor_size.second, row_scaled_nvfp4, config.mode, - config.e4m3_max); + config.scale_max); + break; +#if CUDA_VERSION >= 13040 + case transformer_engine::DType::kFloat8UE5M3: + performTest_dequantize_nvfp4( + tensor_size.first, tensor_size.second, row_scaled_nvfp4, config.mode, + config.scale_max); + break; +#endif // CUDA_VERSION >= 13040 + default: + NVTE_ERROR("Invalid scale type (", static_cast(scale_type), "."); + } ); } @@ -315,41 +374,22 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Bool(), ::testing::Values(NVFP4DequantizeTestConfig{}, NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 448}, - NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 256})), - [](const testing::TestParamInfo& info) - { - const NVFP4DequantizeTestConfig config = std::get<3>(info.param); - const bool use_4over6 = config.mode != kNVTENVFP44Over6Disabled; - std::string name = std::to_string(std::get<0>(info.param).first) + "X" + - std::to_string(std::get<0>(info.param).second) + "X" + - test::typeName(std::get<1>(info.param)) + "X" + - (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + - (use_4over6 ? "FourOverSix" : "Default") + "X" + - (config.e4m3_max == 256 ? "E4M3Max256" : "E4M3Max448"); - return name; - } -); + NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 256}), + ::testing::Values(DType::kFloat8E4M3)), + DequantizeNVFP4TestSuite::test_name); #if CUDA_VERSION >= 13040 -TEST(DequantizeNVFP4Test, UE5M3Scales) -{ - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - performTest_dequantize_nvfp4( - 32, 64, false, kNVTENVFP44Over6Disabled, 114688); - performTest_dequantize_nvfp4( - 32, 64, true, kNVTENVFP44Over6Disabled, 114688); - performTest_dequantize_nvfp4_swizzled( - 32, 64, false, kNVTENVFP44Over6Disabled, 114688); - performTest_dequantize_nvfp4_swizzled( - 32, 64, true, kNVTENVFP44Over6Disabled, 114688); - performTest_dequantize_nvfp4( - 32, 64, false, kNVTENVFP44Over6MinMAE, 65536); - performTest_dequantize_nvfp4_swizzled( - 32, 64, true, kNVTENVFP44Over6MinMAE, 65536); -} +INSTANTIATE_TEST_SUITE_P( + OperatorTestUE5M3Scales, + DequantizeNVFP4TestSuite, + ::testing::Combine( + ::testing::ValuesIn(nvfp4_tensor_dims), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(false), + ::testing::Values(NVFP4DequantizeTestConfig{}), + ::testing::Values(DType::kFloat8UE5M3)), + DequantizeNVFP4TestSuite::test_name); TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) { @@ -408,8 +448,6 @@ TEST(NVFP4RecipeTest, UE5M3PerTensorScale) constexpr float alpha_in = 2.0f; constexpr float fp4_max = 6.0f; constexpr float ue5m3_max = 114688.0f; - input_a.set_nvfp4_e4m3_max(static_cast(ue5m3_max)); - input_b.set_nvfp4_e4m3_max(static_cast(ue5m3_max)); input_a.set_amax(amax_a); input_b.set_tensor_amax_columnwise(amax_b); @@ -421,18 +459,6 @@ TEST(NVFP4RecipeTest, UE5M3PerTensorScale) 1.0f / (fp4_max * fp4_max * ue5m3_max * ue5m3_max); const float expected = alpha_in * amax_a * amax_b * factor_inv; EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr()[0], expected); - - input_a.set_nvfp4_e4m3_max(65536); - input_b.set_nvfp4_e4m3_max(65536); - nvte_nvfp4_compute_per_tensor_scale( - input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0); - alpha_out.to_cpu(); - - constexpr float ue5m3_headroom_max = 65536.0f; - const float headroom_factor_inv = - 1.0f / (fp4_max * fp4_max * ue5m3_headroom_max * ue5m3_headroom_max); - const float headroom_expected = alpha_in * amax_a * amax_b * headroom_factor_inv; - EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr()[0], headroom_expected); } #endif @@ -456,7 +482,7 @@ TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, performTest_dequantize_nvfp4_swizzled( tensor_size.first, tensor_size.second, row_scaled_nvfp4, config.mode, - config.e4m3_max); + config.scale_max); ); } @@ -479,7 +505,7 @@ INSTANTIATE_TEST_SUITE_P( test::typeName(std::get<1>(info.param)) + "X" + (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + (use_4over6 ? "FourOverSix" : "Default") + "X" + - (config.e4m3_max == 256 ? "E4M3Max256" : "E4M3Max448") + "X" + + "ScaleMax" + std::to_string(config.scale_max) + "X" + "Swizzled"; return name; } diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index 95ec7e6679..551ebc52e7 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -49,9 +49,6 @@ bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2) { } size_t typeToNumBits(DType type) { - if (type == DType::kFloat8UE5M3) { - return 8; - } TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T, { return TypeInfo::size; @@ -136,12 +133,16 @@ NVTEShape convertShape(const std::vector& s) { } std::pair get_scales(const NVTEShape& shape, - const NVTEScalingMode scaling_mode) { + const NVTEScalingMode scaling_mode, + std::optional scale_type) { if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { scale_inv_meta ret; ret.shape = {1}; - ret.type = DType::kFloat32; - ret.type_size_bits = typeToNumBits(DType::kFloat32); + if (!scale_type) { + scale_type = DType::kFloat32; + } + ret.type = *scale_type; + ret.type_size_bits = typeToNumBits(*scale_type); return {ret, ret}; } if (scaling_mode == NVTE_MXFP8_1D_SCALING) { @@ -164,10 +165,13 @@ std::pair get_scales(const NVTEShape& shape, size_t scale_dim_X_colwise = DIVUP_TO_MULTIPLE(last_dim, scale_tensor_alignment_X_colwise); ret_colwise.shape = {scale_dim_Y_colwise, scale_dim_X_colwise}; - ret_rowwise.type = DType::kFloat8E8M0; - ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0); - ret_colwise.type = DType::kFloat8E8M0; - ret_colwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0); + if (!scale_type) { + scale_type = DType::kFloat8E8M0; + } + ret_rowwise.type = *scale_type; + ret_rowwise.type_size_bits = typeToNumBits(*scale_type); + ret_colwise.type = *scale_type; + ret_colwise.type_size_bits = typeToNumBits(*scale_type); return {ret_rowwise, ret_colwise}; } @@ -192,10 +196,13 @@ std::pair get_scales(const NVTEShape& shape, size_t scale_dim_X_t = DIVUP_TO_MULTIPLE(DIVUP(first_dim, 16lu), scale_tensor_alignment_X_rowwise); ret_colwise.shape = {scale_dim_Y_t, scale_dim_X_t}; - ret_rowwise.type = DType::kFloat8E4M3; - ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat8E4M3); - ret_colwise.type = DType::kFloat8E4M3; - ret_colwise.type_size_bits = typeToNumBits(DType::kFloat8E4M3); + if (!scale_type) { + scale_type = DType::kFloat8E4M3; + } + ret_rowwise.type = *scale_type; + ret_rowwise.type_size_bits = typeToNumBits(*scale_type); + ret_colwise.type = *scale_type; + ret_colwise.type_size_bits = typeToNumBits(*scale_type); return {ret_rowwise, ret_colwise}; } @@ -219,10 +226,14 @@ std::pair get_scales(const NVTEShape& shape, size_t scale_dim_1 = DIVUP(DIVUP(first_dim, 128lu), 4) * 4; ret_colwise.shape = {scale_dim_0, scale_dim_1}; } - ret_rowwise.type = DType::kFloat32; - ret_colwise.type = DType::kFloat32; - ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat32); - ret_colwise.type_size_bits = typeToNumBits(DType::kFloat32); + + if (!scale_type) { + scale_type = DType::kFloat32; + } + ret_rowwise.type = *scale_type; + ret_colwise.type = *scale_type; + ret_rowwise.type_size_bits = typeToNumBits(*scale_type); + ret_colwise.type_size_bits = typeToNumBits(*scale_type); return {ret_rowwise, ret_colwise}; } @@ -245,10 +256,13 @@ std::pair get_scales(const NVTEShape& shape, size_t scale_dim_1 = DIVUP(last_dim, 4) * 4; ret_colwise.shape = {scale_dim_0, scale_dim_1}; } - ret_rowwise.type = DType::kFloat32; - ret_colwise.type = DType::kFloat32; - ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat32); - ret_colwise.type_size_bits = typeToNumBits(DType::kFloat32); + if (!scale_type) { + scale_type = DType::kFloat32; + } + ret_rowwise.type = *scale_type; + ret_colwise.type = *scale_type; + ret_rowwise.type_size_bits = typeToNumBits(*scale_type); + ret_colwise.type_size_bits = typeToNumBits(*scale_type); return {ret_rowwise, ret_colwise}; } @@ -282,7 +296,8 @@ void Tensor::Buffer::from_cpu() { Tensor::Tensor(const std::string& name, const NVTEShape &shape, const DType type, const bool rowwise, const bool columnwise, - const NVTEScalingMode &scaling_mode, const DType scale_dtype) + const NVTEScalingMode &scaling_mode, + const std::optional scale_type) : tensor_(scaling_mode), rowwise_{rowwise}, columnwise_{columnwise}, name_{name} { // Initialize RNG const size_t seed = create_seed_from_tensor_name(name); @@ -377,15 +392,7 @@ Tensor::Tensor(const std::string& name, case NVTE_NVFP4_1D_SCALING: { // Block scaling factors - auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode()); - if (scaling_mode == NVTE_NVFP4_1D_SCALING) { - NVTE_CHECK(scale_dtype == DType::kFloat8E4M3 || - scale_dtype == DType::kFloat8UE5M3); - rowwise_scale_meta.type = scale_dtype; - rowwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype); - colwise_scale_meta.type = scale_dtype; - colwise_scale_meta.type_size_bits = typeToNumBits(scale_dtype); - } + auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode(), scale_type); if (rowwise) { const auto scale_shape = rowwise_scale_meta.shape; const auto scale_dtype = rowwise_scale_meta.type; diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 9156c03d7b..1b5b95d731 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -160,7 +160,7 @@ class Tensor { const bool rowwise = true, const bool columnwise = false, const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING, - const DType scale_dtype = DType::kFloat8E4M3); + const std::optional scale_dtype = std::nullopt); Tensor(const std::string& name, const std::vector &shape, @@ -168,9 +168,9 @@ class Tensor { const bool rowwise = true, const bool columnwise = false, const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING, - const DType scale_dtype = DType::kFloat8E4M3) : + const std::optional scale_type = std::nullopt) : Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode, - scale_dtype) {} + scale_type) {} Tensor() = default; @@ -646,6 +646,16 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, #define SWITCH_FP4_TYPE_HANDLE(type, ...) // do nothing #endif +#if CUDA_VERSION >= 13040 +#define SWITCH_UE5M3_TYPE_HANDLE(type, ...) \ + case DType::kFloat8UE5M3: { \ + using type = fp8ue5m3; \ + { __VA_ARGS__ } \ + } break; +#else +#define SWITCH_UE5M3_TYPE_HANDLE(type, ...) // do nothing +#endif + #define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \ switch (dtype) { \ using namespace transformer_engine; \ @@ -704,7 +714,8 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, } \ break; \ SWITCH_FP4_TYPE_HANDLE(type, __VA_ARGS__) \ - default: \ + SWITCH_UE5M3_TYPE_HANDLE(type, __VA_ARGS__) \ + default: \ printf("dtype: %d\n", static_cast(dtype)); \ NVTE_ERROR("Invalid type."); \ } From 93cbd0a285ec3258c6d9e4da6cd5a2b6754356f8 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 28 Aug 2026 08:54:15 +0000 Subject: [PATCH 39/54] Address some review comments from @ptrendx Signed-off-by: Tim Moon --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 5 --- tests/pytorch/test_fusible_ops.py | 6 --- tests/pytorch/test_grouped_mlp.py | 6 --- .../transformer_engine/transformer_engine.h | 39 ++++++++++--------- .../pytorch/ops/fused/grouped_mlp.py | 2 +- .../tensor/storage/grouped_tensor_storage.py | 6 +-- 6 files changed, 25 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 7aeab72f23..0bd5c259d1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -696,11 +696,6 @@ def _check_ue5m3_gemm_versus_dequantized( M, K, N, x_columnwise, w_columnwise, disable_second_level_scale ): """Run an NVFP4/UE5M3 GEMM and compare against a dequantized FP32 reference.""" - if M % 256 != 0: - pytest.skip( - "cuDNN's grouped GEMM pads every group to 256 rows, so the UE5M3 path (which " - "routes there while cuBLAS lacks UE5M3 kernels) requires M % 256 == 0." - ) torch.manual_seed(0) device, dtype, out_dtype = "cuda", torch.bfloat16, torch.bfloat16 x_shape = (K, M) if x_columnwise else (M, K) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 534432a512..853665ebae 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -139,12 +139,6 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") - if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and ( - math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0 - ): - pytest.skip( - "cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors" - ) # Check dtype if dtype is not None: diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 22eff00f25..6024d1af89 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -267,12 +267,6 @@ def maybe_skip_quantization( elif quantization in nvfp4_variant_names: if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") - if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and ( - math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0 - ): - pytest.skip( - "cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors" - ) # Check dtype if dtype is not None: diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index ba9b22f124..5cfd573187 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -23,19 +23,19 @@ extern "C" { * \brief TE datatype. */ enum NVTEDType { - kNVTEByte = 0, /*!< Byte */ - kNVTEInt16 = 1, /*!< 16-bit integer */ - kNVTEInt32 = 2, /*!< 32-bit integer */ - kNVTEInt64 = 3, /*!< 64-bit integer */ - kNVTEFloat32 = 4, /*!< 32-bit float */ - kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ - kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ - kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ - kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ - kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ - kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ - kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ - kNVTENumTypes /*!< Number of supported types */ + kNVTEByte = 0, /*!< Byte */ + kNVTEInt16 = 1, /*!< 16-bit integer */ + kNVTEInt32 = 2, /*!< 32-bit integer */ + kNVTEInt64 = 3, /*!< 64-bit integer */ + kNVTEFloat32 = 4, /*!< 32-bit float */ + kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ + kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ + kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ + kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ + kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ + kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ + kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ + kNVTENumTypes /*!< Number of supported types */ }; /*! \struct NVTEShape @@ -84,12 +84,15 @@ enum NVTETensorParam { * its values are populated during quantization. */ kNVTERowScaledNVFP4 = 8, - /*! Global scale-bound selector used by an NVFP4 tensor. + /*! Global scale bound used by an NVFP4 tensor. * - * This is part of the tensor data contract. Downstream dequantization and - * GEMM scale consumers must use the same bound used during quantization. - * Standard NVFP4 uses 448; 4over6 may use 256 for map-to-4 headroom. - * For UE5M3 scales, these settings map to 114688 and 65536, respectively. + * When non-zero, this overrides the maximum value of the scaling + * factors. This affects downstream consumers (like GEMM or + * dequantization) because the maximum values of the FP4 type and + * scale type are used to convert amax values to tensor scales. + * + * Standard NVFP4 uses the maximum value of the scale type (448 for + * E4M3). 4over6 may use 256 for map-to-4 headroom. */ kNVTENVFP4E4M3Max = 9, kNVTENumTensorParams diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 53dde0dacc..75c0f60673 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -1504,7 +1504,7 @@ def fuser_forward( if self.grouped_gemm_act_hadamard_kernel() is None: # Kernel is not available pass - elif self._cudnn_act_func in ("swiglu", "situlu"): + elif self._cudnn_act_func in ("swiglu", "situglu"): kernel_impl = "gemm_act_rht_amax" elif activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard(): kernel_impl = "gemm_act_rht_amax" diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index bbd62175ef..df3b14ff3c 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -799,9 +799,6 @@ def make_grouped_tensor( # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=dtype, device=device) elif compatible_recipe.mxfp8(): - # Amax buffer for delayed scaling - one per tensor - amax = torch.empty(num_tensors, dtype=torch.float32, device=device) - scale_inv_dtype = DType.kFloat8E8M0 if rowwise_usage: @@ -849,6 +846,9 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors columnwise_scale_inv_offsets = list(range(num_tensors + 1)) + + # Amax buffer for delayed scaling - one per tensor + amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif compatible_recipe.nvfp4(): scale_inv_dtype = quantizer.scale_dtype row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 From d00cbeb4ae13f9fc7a6e43ee8f477764e81ea68f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:55:41 +0000 Subject: [PATCH 40/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../transformer_engine/transformer_engine.h | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 5cfd573187..76fe24a7ab 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -23,19 +23,19 @@ extern "C" { * \brief TE datatype. */ enum NVTEDType { - kNVTEByte = 0, /*!< Byte */ - kNVTEInt16 = 1, /*!< 16-bit integer */ - kNVTEInt32 = 2, /*!< 32-bit integer */ - kNVTEInt64 = 3, /*!< 64-bit integer */ - kNVTEFloat32 = 4, /*!< 32-bit float */ - kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ - kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ - kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ - kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ - kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ - kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ - kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ - kNVTENumTypes /*!< Number of supported types */ + kNVTEByte = 0, /*!< Byte */ + kNVTEInt16 = 1, /*!< 16-bit integer */ + kNVTEInt32 = 2, /*!< 32-bit integer */ + kNVTEInt64 = 3, /*!< 64-bit integer */ + kNVTEFloat32 = 4, /*!< 32-bit float */ + kNVTEFloat16 = 5, /*!< 16-bit float (E5M10) */ + kNVTEBFloat16 = 6, /*!< 16-bit bfloat (E8M7) */ + kNVTEFloat8E4M3 = 7, /*!< 8-bit float (E4M3) */ + kNVTEFloat8E5M2 = 8, /*!< 8-bit float (E5M2) */ + kNVTEFloat8E8M0 = 9, /*!< 8-bit float (E8M0) */ + kNVTEFloat4E2M1 = 10, /*!< 4-bit float (E2M1) */ + kNVTEFloat8UE5M3 = 11, /*!< 8-bit float (UE5M3) */ + kNVTENumTypes /*!< Number of supported types */ }; /*! \struct NVTEShape From 3a62970a860bedb6285c8ba049296f0d853881f0 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 28 Aug 2026 09:10:36 +0000 Subject: [PATCH 41/54] Respect no-tensor-scaling in grouped tensor helper function Signed-off-by: Tim Moon --- .../pytorch/tensor/storage/grouped_tensor_storage.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index df3b14ff3c..6a48f85c6c 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -854,6 +854,7 @@ def make_grouped_tensor( row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 nvfp4_use_4over6 = quantizer.nvfp4_use_4over6 nvfp4_e4m3_max = quantizer.nvfp4_e4m3_max + disable_second_level_scale = quantizer.disable_second_level_scale if row_scaled_nvfp4: if not rowwise_usage: raise ValueError( @@ -879,7 +880,8 @@ def make_grouped_tensor( total_scale_elements += math.prod(scale_inv_shape) scale_inv_offsets.append(total_scale_elements) scale_inv = torch.empty(total_scale_elements, dtype=torch.uint8, device=device) - amax = torch.empty(total_amax_elements, dtype=torch.float32, device=device) + if not disable_second_level_scale: + amax = torch.empty(total_amax_elements, dtype=torch.float32, device=device) if columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8, FP4 packed) @@ -896,7 +898,8 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) - columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + if not disable_second_level_scale: + columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif compatible_recipe.float8_block_scaling(): scale_inv_dtype = DType.kFloat32 From e404db8c8b29bdfdea27576c023c809d66d60069 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 29 Aug 2026 23:25:31 +0000 Subject: [PATCH 42/54] Debug C++ test compilation error Signed-off-by: Tim Moon --- tests/cpp/operator/test_normalization.cu | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/cpp/operator/test_normalization.cu b/tests/cpp/operator/test_normalization.cu index ea6692dba4..971a6306dc 100644 --- a/tests/cpp/operator/test_normalization.cu +++ b/tests/cpp/operator/test_normalization.cu @@ -261,9 +261,9 @@ TEST_P(NormTestSuite, TestNorm) { const bool cudnn_zero_centered_gamma_in_weight_dtype = std::get<6>(GetParam()); const bool fused_bwd_add = std::get<7>(GetParam()); - TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(input_type, InputType, - TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(output_type, OutputType, - performTest( + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, + if (output_type == DType::kFloat8E4M3) { + performTest( size.first, size.second, zero_centered_gamma, @@ -272,7 +272,20 @@ TEST_P(NormTestSuite, TestNorm) { cudnn_zero_centered_gamma_in_weight_dtype, fused_bwd_add ); - ); + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + output_type, OutputType, + performTest( + size.first, + size.second, + zero_centered_gamma, + norm_type, + use_cudnn, + cudnn_zero_centered_gamma_in_weight_dtype, + fused_bwd_add + ); + ); + } ); } From adef2dc83164df4767c12324b0e184876411e199 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sun, 30 Aug 2026 00:59:07 +0000 Subject: [PATCH 43/54] Make sure GEMM alpha/beta scales are on GPU, even without amaxes Co-authored-by: Codex Signed-off-by: Tim Moon --- transformer_engine/common/gemm/cublaslt_gemm.cu | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 7ae8d345cd..aca0173be9 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -374,12 +374,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, is_nvfp_scaling(inputA->scaling_mode) && is_nvfp_scaling(inputB->scaling_mode); // Update scaling factors with NVFP4 tensor scales - // TODO: Check whether scales are on CPU/GPU or add API to control. - // Currently scales are assumed to be on CPU when amax is provided - // and on GPU when not provided, but this is brittle. - if (use_fp4 && nvfp4_tensor_scaling && - ((transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr || - (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr)) { + if (use_fp4 && nvfp4_tensor_scaling) { // Reserve some workspace for alpha scale NVTE_CHECK(workspaceSize >= 4, "NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has ", From a9715521989c2b8c9ac5ee9f5bf8e93f38a262b5 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 31 Aug 2026 23:10:18 +0000 Subject: [PATCH 44/54] Create new Mcore DDP integration functions rather than breaking backward compatibility Co-authored-by: Codex Signed-off-by: Tim Moon --- tests/cpp/operator/test_dequantize_nvfp4.cu | 44 ++++++- .../include/transformer_engine/recipe.h | 115 +++++++++++++++--- transformer_engine/common/recipe/nvfp4.cu | 63 ++++++++-- .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 15 +-- .../pytorch/csrc/extensions/transpose.cpp | 56 +++++---- 5 files changed, 229 insertions(+), 64 deletions(-) diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index b56890e968..3185aac80f 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -378,6 +378,44 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3)), DequantizeNVFP4TestSuite::test_name); +TEST(NVFP4RecipeTest, LegacyScaleUtilitiesUseE4M3) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + Tensor global_amax("global_amax", std::vector{1}, DType::kFloat32); + Tensor global_scale("global_scale", std::vector{1}, DType::kFloat32); + global_amax.rowwise_cpu_dptr()[0] = 12.0f; + global_amax.from_cpu(); + nvte_nvfp4_compute_global_scale(global_amax.data(), global_scale.data(), 0); + global_scale.to_cpu(); + EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 448.0f / 12.0f); + + Tensor block_amax("block_amax", std::vector{1, 2}, DType::kFloat32); + Tensor block_scale("block_scale", std::vector{1, 2}, DType::kFloat32); + block_amax.rowwise_cpu_dptr()[0] = 3.0f; + block_amax.rowwise_cpu_dptr()[1] = 6.0f; + block_amax.from_cpu(); + nvte_nvfp4_compute_per_block_scale( + block_amax.data(), block_scale.data(), global_amax.data(), 0); + block_scale.to_cpu(); + EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 448.0f / 12.0f); + EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 448.0f / 12.0f); + + Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); + nvte_nvfp4_expand_scale_to_fp8(block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, 0); + expanded_scale.to_cpu(); + const auto *scales = reinterpret_cast( + expanded_scale.rowwise_cpu_dptr()); + for (size_t row = 0; row < 16; ++row) { + EXPECT_FLOAT_EQ(static_cast(scales[row * 2]), + static_cast(fp8e4m3(3.0f * 448.0f / 12.0f))); + EXPECT_FLOAT_EQ(static_cast(scales[row * 2 + 1]), + static_cast(fp8e4m3(6.0f * 448.0f / 12.0f))); + } +} + #if CUDA_VERSION >= 13040 INSTANTIATE_TEST_SUITE_P( @@ -401,7 +439,7 @@ TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) Tensor global_scale("global_scale", std::vector{1}, DType::kFloat32); global_amax.rowwise_cpu_dptr()[0] = 12.0f; global_amax.from_cpu(); - nvte_nvfp4_compute_global_scale( + nvte_nvfp4_compute_global_scale_v2( global_amax.data(), global_scale.data(), kNVTEFloat8UE5M3, 0); global_scale.to_cpu(); EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 114688.0f / 12.0f); @@ -411,14 +449,14 @@ TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) block_amax.rowwise_cpu_dptr()[0] = 3.0f; block_amax.rowwise_cpu_dptr()[1] = 6.0f; block_amax.from_cpu(); - nvte_nvfp4_compute_per_block_scale( + nvte_nvfp4_compute_per_block_scale_v2( block_amax.data(), block_scale.data(), global_amax.data(), kNVTEFloat8UE5M3, 0); block_scale.to_cpu(); EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 114688.0f / 12.0f); EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 114688.0f / 12.0f); Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); - nvte_nvfp4_expand_scale_to_fp8( + nvte_nvfp4_expand_scale_to_fp8_v2( block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, kNVTEFloat8UE5M3, 0); expanded_scale.to_cpu(); const auto *scales = reinterpret_cast( diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 2d141e2f16..9ee26364f6 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -362,6 +362,7 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s * Quantizes elements in [start_offset, start_offset + len) of the flattened tensor * using precomputed per-tile scales. Each 16x16 tile uses its own scale factor. * Used in distributed settings where each rank casts its owned shard. + * Uses E4M3 scale dtype. Use nvte_nvfp4_2d_partial_cast_v2 to specify another scale type. * * \param[in] inp Input tensor (partial shard, high-precision). * \param[out] out Output NVFP4 packed tensor (2 values per byte). @@ -373,15 +374,37 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s * \param[in] scale_stride_w Stride for scale in tile-col dimension. * \param[in] start_offset Starting element offset in the flattened tensor. * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). - * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream used for the operation. */ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, NVTEDType scale_dtype, cudaStream_t stream); + size_t block_len, cudaStream_t stream); -/*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. +/*! \brief Cast a partial shard of a tensor to NVFP4 using 2D tile-based quantization. + * + * This variant allows the NVFP4 scale storage type to be specified explicitly. + * + * \param[in] inp Input tensor (partial shard, high-precision). + * \param[out] out Output NVFP4 packed tensor (2 values per byte). + * \param[in] scale Per-tile scale factors [tile_rows, tile_cols], float32. + * \param[in] global_scale Global scale factor [1], float32. + * \param[in] h Number of rows in the full 2D tensor. + * \param[in] w Number of columns in the full 2D tensor. + * \param[in] scale_stride_h Stride for scale in tile-row dimension. + * \param[in] scale_stride_w Stride for scale in tile-col dimension. + * \param[in] start_offset Starting element offset in the flattened tensor. + * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_nvfp4_2d_partial_cast_v2(const NVTETensor inp, NVTETensor out, const NVTETensor scale, + const NVTETensor global_scale, size_t h, size_t w, + size_t scale_stride_h, size_t scale_stride_w, + size_t start_offset, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream); + +/*! \brief Expand tile-level scales to row-level scales and convert to E4M3 scale type. * * Each tile row's scale is repeated block_len times in the output. * @@ -391,14 +414,29 @@ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTE * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Padded row count in output. * \param[in] block_len Block length (typically 16 for NVFP4). - * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream. */ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - NVTEDType scale_dtype, cudaStream_t stream); + cudaStream_t stream); -/*! \brief Compute per-block decode scale from block amax and global amax. +/*! \brief Expand tile-level scales to row-level scales and convert to the selected FP8 scale type. + * + * \param[in] input Input tensor with tile scales [tile_rows, tile_cols], float32. + * \param[out] output Output tensor with expanded scales [rows_padded, tile_cols], uint8. + * \param[in] tile_rows Number of tile rows. + * \param[in] tile_cols Number of tile columns. + * \param[in] rows_padded Padded row count in output. + * \param[in] block_len Block length (typically 16 for NVFP4). + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_expand_scale_to_fp8_v2(const NVTETensor input, NVTETensor output, + size_t tile_rows, size_t tile_cols, size_t rows_padded, + size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream); + +/*! \brief Compute per-block E4M3 decode scale from block amax and global amax. * * Computes: * global_scale = (scale_max * fp4_max) / global_amax @@ -409,19 +447,31 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. * \param[out] scale Output scale tensor [tile_rows, tile_cols], float32. * \param[in] global_amax Global amax tensor (single element), float32. Avoids D2H transfer. - * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream. */ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, NVTEDType scale_dtype, - cudaStream_t stream); + const NVTETensor global_amax, cudaStream_t stream); -/*! \brief Fused kernel for NVFP4 scale computation. +/*! \brief Compute per-block decode scale from block amax and global amax. + * + * This variant allows the NVFP4 scale storage type to be specified explicitly. + * + * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. + * \param[out] scale Output scale tensor [tile_rows, tile_cols], float32. + * \param[in] global_amax Global amax tensor (single element), float32. Avoids D2H transfer. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_compute_per_block_scale_v2(const NVTETensor block_amax, NVTETensor scale, + const NVTETensor global_amax, NVTEDType scale_dtype, + cudaStream_t stream); + +/*! \brief Fused kernel for NVFP4 E4M3 scale computation. * * Fuses three operations into one kernel: * 1. Compute per-block decode scales from block amax and global amax * 2. Copy global amax to target tensor - * 3. Expand tile-level scales to row-level and convert to the selected FP8 scale type + * 3. Expand tile-level scales to row-level and convert to E4M3 scale type * * Saves 2 kernel launches per parameter. * @@ -434,27 +484,58 @@ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor * \param[in] tile_cols Number of tile columns. * \param[in] rows_padded Total padded rows in output. * \param[in] block_len Block length (16 for NVFP4). - * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream. */ void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, NVTEDType scale_dtype, - cudaStream_t stream); + size_t rows_padded, size_t block_len, cudaStream_t stream); + +/*! \brief Fused kernel for NVFP4 scale computation. + * + * This variant allows the NVFP4 scale storage type to be specified explicitly. + * + * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. + * \param[in] global_amax Global amax tensor [1], float32. + * \param[out] per_block_scale Output per-block scale [tile_rows, tile_cols], float32 (for partial_cast). + * \param[out] target_scale Output scale tensor [rows_padded, tile_cols], uint8. + * \param[out] target_amax Output amax tensor [1], float32 (copy of global_amax). + * \param[in] tile_rows Number of tile rows. + * \param[in] tile_cols Number of tile columns. + * \param[in] rows_padded Total padded rows in output. + * \param[in] block_len Block length (16 for NVFP4). + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_fused_scale_v2(const NVTETensor block_amax, const NVTETensor global_amax, + NVTETensor per_block_scale, NVTETensor target_scale, + NVTETensor target_amax, size_t tile_rows, size_t tile_cols, + size_t rows_padded, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream); /*! \brief Compute global encode scale from global amax. * - * Computes: global_scale = (scale_max * fp4_max) / global_amax + * Computes: global_scale = (scale_max * fp4_max) / global_amax, using E4M3 scale_max. * If global_amax <= 0, returns 1.0. * * \param[in] global_amax Input global amax tensor [num_params], float32. * \param[out] global_scale Output global scale tensor [num_params], float32. - * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream. */ void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - NVTEDType scale_dtype, cudaStream_t stream); + cudaStream_t stream); + +/*! \brief Compute global encode scale from global amax. + * + * This variant allows the NVFP4 scale storage type to be specified explicitly. + * + * \param[in] global_amax Input global amax tensor [num_params], float32. + * \param[out] global_scale Output global scale tensor [num_params], float32. + * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_compute_global_scale_v2(const NVTETensor global_amax, NVTETensor global_scale, + NVTEDType scale_dtype, cudaStream_t stream); #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index f1666653d1..9d07e36beb 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -836,9 +836,17 @@ void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, - NVTEDType scale_dtype, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED + cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8); + nvte_nvfp4_expand_scale_to_fp8_v2(input, output, tile_rows, tile_cols, rows_padded, block_len, + kNVTEFloat8E4M3, stream); +} + +void nvte_nvfp4_expand_scale_to_fp8_v2(const NVTETensor input, NVTETensor output, size_t tile_rows, + size_t tile_cols, size_t rows_padded, size_t block_len, + NVTEDType scale_dtype, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8_v2); using namespace transformer_engine; nvfp4_recipe::nvfp4_expand_scale_to_fp8( *convertNVTETensorCheck(input), *convertNVTETensorCheck(output), tile_rows, tile_cols, @@ -849,10 +857,16 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s } void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, - const NVTETensor global_amax, NVTEDType scale_dtype, - cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED + const NVTETensor global_amax, cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale); + nvte_nvfp4_compute_per_block_scale_v2(block_amax, scale, global_amax, kNVTEFloat8E4M3, stream); +} + +void nvte_nvfp4_compute_per_block_scale_v2(const NVTETensor block_amax, NVTETensor scale, + const NVTETensor global_amax, NVTEDType scale_dtype, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale_v2); using namespace transformer_engine; nvfp4_recipe::nvfp4_compute_per_block_scale( *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(scale), @@ -863,9 +877,15 @@ void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor } void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, - NVTEDType scale_dtype, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED + cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_compute_global_scale); + nvte_nvfp4_compute_global_scale_v2(global_amax, global_scale, kNVTEFloat8E4M3, stream); +} + +void nvte_nvfp4_compute_global_scale_v2(const NVTETensor global_amax, NVTETensor global_scale, + NVTEDType scale_dtype, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + NVTE_API_CALL(nvte_nvfp4_compute_global_scale_v2); using namespace transformer_engine; nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax), *convertNVTETensorCheck(global_scale), @@ -916,9 +936,19 @@ void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, s void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, - size_t block_len, NVTEDType scale_dtype, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED + size_t block_len, cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); + nvte_nvfp4_2d_partial_cast_v2(inp, out, scale, global_scale, h, w, scale_stride_h, + scale_stride_w, start_offset, block_len, kNVTEFloat8E4M3, stream); +} + +void nvte_nvfp4_2d_partial_cast_v2(const NVTETensor inp, NVTETensor out, const NVTETensor scale, + const NVTETensor global_scale, size_t h, size_t w, + size_t scale_stride_h, size_t scale_stride_w, + size_t start_offset, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + NVTE_API_CALL(nvte_nvfp4_2d_partial_cast_v2); using namespace transformer_engine; nvfp4_recipe::nvfp4_2d_partial_cast( *convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), *convertNVTETensorCheck(scale), @@ -962,10 +992,19 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, - size_t rows_padded, size_t block_len, NVTEDType scale_dtype, - cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED + size_t rows_padded, size_t block_len, cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_fused_scale); + nvte_nvfp4_fused_scale_v2(block_amax, global_amax, per_block_scale, target_scale, target_amax, + tile_rows, tile_cols, rows_padded, block_len, kNVTEFloat8E4M3, stream); +} + +void nvte_nvfp4_fused_scale_v2(const NVTETensor block_amax, const NVTETensor global_amax, + NVTETensor per_block_scale, NVTETensor target_scale, + NVTETensor target_amax, size_t tile_rows, size_t tile_cols, + size_t rows_padded, size_t block_len, NVTEDType scale_dtype, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + NVTE_API_CALL(nvte_nvfp4_fused_scale_v2); using namespace transformer_engine; nvfp4_recipe::nvfp4_fused_scale( *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax), diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp index f5d3ad9bc7..4ecd50deb8 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -43,9 +43,10 @@ void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tens const TensorWrapper scale_cu = makeTransformerEngineTensor(scale); const TensorWrapper global_scale_cu = makeTransformerEngineTensor(global_scale); - nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), - h, w, scale.stride(0), scale.stride(1), start_offset, block_len, - static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_2d_partial_cast_v2(inp_cu.data(), out_cu.data(), scale_cu.data(), + global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), + start_offset, block_len, static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, @@ -94,10 +95,10 @@ void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, const TensorWrapper scale_cu = makeTransformerEngineTensor(scale); const TensorWrapper global_scale_cu = makeTransformerEngineTensor(global_scale); - nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), - global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), - start_offset, static_cast(block_len), - static_cast(scale_dtype), stream); + nvte_nvfp4_2d_partial_cast_v2(inp_cu.data(), out_cu.data(), scale_cu.data(), + global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), + start_offset, static_cast(block_len), + static_cast(scale_dtype), stream); } } diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index c6817af5aa..d3dcfe6604 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -127,24 +127,25 @@ void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile init_extension(); // Input: per_block_decode_scale [tile_rows, tile_cols], float32 - // Output: target_scale [rows_padded, tile_cols], uint8 (E4M3) + // Output: target_scale [rows_padded, tile_cols], uint8 scale storage const auto in_shape = getTensorShape(input); const auto out_shape = getTensorShape(output); NVTE_CHECK(in_shape.size() == 2, "NVFP4 expand scale expects 2D input."); NVTE_CHECK(out_shape.size() == 2, "NVFP4 expand scale expects 2D output."); NVTE_CHECK(input.scalar_type() == at::kFloat, "NVFP4 expand scale input must be float32."); - NVTE_CHECK(output.scalar_type() == at::kByte, "NVFP4 expand scale output must be uint8 (E4M3)."); + NVTE_CHECK(output.scalar_type() == at::kByte, + "NVFP4 expand scale output must be uint8 scale storage."); auto input_cu = makeTransformerEngineTensor( input.data_ptr(), std::vector{in_shape[0], in_shape[1]}, DType::kFloat32); auto output_cu = makeTransformerEngineTensor( output.data_ptr(), std::vector{out_shape[0], out_shape[1]}, DType::kByte); - nvte_nvfp4_expand_scale_to_fp8(input_cu.data(), output_cu.data(), static_cast(tile_rows), - static_cast(tile_cols), static_cast(rows_padded), - static_cast(block_len), - static_cast(scale_dtype), - at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_expand_scale_to_fp8_v2( + input_cu.data(), output_cu.data(), static_cast(tile_rows), + static_cast(tile_cols), static_cast(rows_padded), + static_cast(block_len), static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax, @@ -162,9 +163,10 @@ void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at:: auto scale_cu = makeTransformerEngineTensor(scale); auto global_amax_cu = makeTransformerEngineTensor(global_amax); - nvte_nvfp4_compute_per_block_scale(block_amax_cu.data(), scale_cu.data(), global_amax_cu.data(), - static_cast(scale_dtype), - at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_compute_per_block_scale_v2(block_amax_cu.data(), scale_cu.data(), + global_amax_cu.data(), + static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, @@ -176,12 +178,13 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor // block_amax: [tile_rows, tile_cols], float32 // global_amax: [1], float32 // per_block_scale: [tile_rows, tile_cols], float32 (for partial_cast) - // target_scale: [rows_padded, tile_cols], uint8 (E4M3) + // target_scale: [rows_padded, tile_cols], uint8 scale storage // target_amax: [1], float32 NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); NVTE_CHECK(per_block_scale.scalar_type() == at::kFloat, "Per-block scale must be float32."); - NVTE_CHECK(target_scale.scalar_type() == at::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_scale.scalar_type() == at::kByte, + "Target scale must be uint8 scale storage."); NVTE_CHECK(target_amax.scalar_type() == at::kFloat, "Target amax must be float32."); NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); @@ -192,11 +195,12 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor auto target_scale_cu = makeTransformerEngineTensor(target_scale); auto target_amax_cu = makeTransformerEngineTensor(target_amax); - nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), - target_scale_cu.data(), target_amax_cu.data(), - static_cast(tile_rows), static_cast(tile_cols), - static_cast(rows_padded), static_cast(block_len), - static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_fused_scale_v2(block_amax_cu.data(), global_amax_cu.data(), + per_block_scale_cu.data(), target_scale_cu.data(), + target_amax_cu.data(), static_cast(tile_rows), + static_cast(tile_cols), static_cast(rows_padded), + static_cast(block_len), static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } void nvfp4_multi_tensor_fused_scale( @@ -235,7 +239,8 @@ void nvfp4_multi_tensor_fused_scale( NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); NVTE_CHECK(per_block_scale.scalar_type() == at::kFloat, "Per-block scale must be float32."); - NVTE_CHECK(target_scale.scalar_type() == at::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_scale.scalar_type() == at::kByte, + "Target scale must be uint8 scale storage."); NVTE_CHECK(target_amax.scalar_type() == at::kFloat, "Target amax must be float32."); NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); @@ -246,10 +251,11 @@ void nvfp4_multi_tensor_fused_scale( auto target_scale_cu = makeTransformerEngineTensor(target_scale); auto target_amax_cu = makeTransformerEngineTensor(target_amax); - nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), - target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, - rows_padded, static_cast(block_len), - static_cast(scale_dtype), stream); + nvte_nvfp4_fused_scale_v2(block_amax_cu.data(), global_amax_cu.data(), + per_block_scale_cu.data(), target_scale_cu.data(), + target_amax_cu.data(), tile_rows, tile_cols, rows_padded, + static_cast(block_len), + static_cast(scale_dtype), stream); } } @@ -264,9 +270,9 @@ void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale, auto global_amax_cu = makeTransformerEngineTensor(global_amax); auto global_scale_cu = makeTransformerEngineTensor(global_scale); - nvte_nvfp4_compute_global_scale(global_amax_cu.data(), global_scale_cu.data(), - static_cast(scale_dtype), - at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_compute_global_scale_v2(global_amax_cu.data(), global_scale_cu.data(), + static_cast(scale_dtype), + at::cuda::getCurrentCUDAStream()); } at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { From bb59631a12c62b0945c9a339ec30855642c0cb87 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Mon, 31 Aug 2026 23:30:17 +0000 Subject: [PATCH 45/54] Remove unrelated tests from NVFP4 dequantize C++ unit tests Co-authored-by: Codex Signed-off-by: Tim Moon --- tests/cpp/operator/test_dequantize_nvfp4.cu | 109 -------------------- 1 file changed, 109 deletions(-) diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index 3185aac80f..3cf9c12812 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -21,7 +21,6 @@ #endif #include -#include #include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -378,44 +377,6 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(DType::kFloat8E4M3)), DequantizeNVFP4TestSuite::test_name); -TEST(NVFP4RecipeTest, LegacyScaleUtilitiesUseE4M3) -{ - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - - Tensor global_amax("global_amax", std::vector{1}, DType::kFloat32); - Tensor global_scale("global_scale", std::vector{1}, DType::kFloat32); - global_amax.rowwise_cpu_dptr()[0] = 12.0f; - global_amax.from_cpu(); - nvte_nvfp4_compute_global_scale(global_amax.data(), global_scale.data(), 0); - global_scale.to_cpu(); - EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 448.0f / 12.0f); - - Tensor block_amax("block_amax", std::vector{1, 2}, DType::kFloat32); - Tensor block_scale("block_scale", std::vector{1, 2}, DType::kFloat32); - block_amax.rowwise_cpu_dptr()[0] = 3.0f; - block_amax.rowwise_cpu_dptr()[1] = 6.0f; - block_amax.from_cpu(); - nvte_nvfp4_compute_per_block_scale( - block_amax.data(), block_scale.data(), global_amax.data(), 0); - block_scale.to_cpu(); - EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 448.0f / 12.0f); - EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 448.0f / 12.0f); - - Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); - nvte_nvfp4_expand_scale_to_fp8(block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, 0); - expanded_scale.to_cpu(); - const auto *scales = reinterpret_cast( - expanded_scale.rowwise_cpu_dptr()); - for (size_t row = 0; row < 16; ++row) { - EXPECT_FLOAT_EQ(static_cast(scales[row * 2]), - static_cast(fp8e4m3(3.0f * 448.0f / 12.0f))); - EXPECT_FLOAT_EQ(static_cast(scales[row * 2 + 1]), - static_cast(fp8e4m3(6.0f * 448.0f / 12.0f))); - } -} - #if CUDA_VERSION >= 13040 INSTANTIATE_TEST_SUITE_P( @@ -428,76 +389,6 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Values(NVFP4DequantizeTestConfig{}), ::testing::Values(DType::kFloat8UE5M3)), DequantizeNVFP4TestSuite::test_name); - -TEST(NVFP4RecipeTest, UE5M3ScaleUtilities) -{ - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - - Tensor global_amax("global_amax", std::vector{1}, DType::kFloat32); - Tensor global_scale("global_scale", std::vector{1}, DType::kFloat32); - global_amax.rowwise_cpu_dptr()[0] = 12.0f; - global_amax.from_cpu(); - nvte_nvfp4_compute_global_scale_v2( - global_amax.data(), global_scale.data(), kNVTEFloat8UE5M3, 0); - global_scale.to_cpu(); - EXPECT_FLOAT_EQ(global_scale.rowwise_cpu_dptr()[0], 6.0f * 114688.0f / 12.0f); - - Tensor block_amax("block_amax", std::vector{1, 2}, DType::kFloat32); - Tensor block_scale("block_scale", std::vector{1, 2}, DType::kFloat32); - block_amax.rowwise_cpu_dptr()[0] = 3.0f; - block_amax.rowwise_cpu_dptr()[1] = 6.0f; - block_amax.from_cpu(); - nvte_nvfp4_compute_per_block_scale_v2( - block_amax.data(), block_scale.data(), global_amax.data(), kNVTEFloat8UE5M3, 0); - block_scale.to_cpu(); - EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[0], 3.0f * 114688.0f / 12.0f); - EXPECT_FLOAT_EQ(block_scale.rowwise_cpu_dptr()[1], 6.0f * 114688.0f / 12.0f); - - Tensor expanded_scale("expanded_scale", std::vector{16, 2}, DType::kByte); - nvte_nvfp4_expand_scale_to_fp8_v2( - block_scale.data(), expanded_scale.data(), 1, 2, 16, 16, kNVTEFloat8UE5M3, 0); - expanded_scale.to_cpu(); - const auto *scales = reinterpret_cast( - expanded_scale.rowwise_cpu_dptr()); - for (size_t row = 0; row < 16; ++row) { - EXPECT_FLOAT_EQ(static_cast(scales[row * 2]), - static_cast(fp8ue5m3(3.0f * 114688.0f / 12.0f))); - EXPECT_FLOAT_EQ(static_cast(scales[row * 2 + 1]), - static_cast(fp8ue5m3(6.0f * 114688.0f / 12.0f))); - } -} - -TEST(NVFP4RecipeTest, UE5M3PerTensorScale) -{ - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP(); - } - - Tensor input_a("input_a", std::vector{32, 32}, DType::kFloat4E2M1, - true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); - Tensor input_b("input_b", std::vector{32, 32}, DType::kFloat4E2M1, - true, true, NVTE_NVFP4_1D_SCALING, DType::kFloat8UE5M3); - Tensor alpha_out("alpha_out", std::vector{1}, DType::kFloat32); - - constexpr float amax_a = 12.0f; - constexpr float amax_b = 18.0f; - constexpr float alpha_in = 2.0f; - constexpr float fp4_max = 6.0f; - constexpr float ue5m3_max = 114688.0f; - input_a.set_amax(amax_a); - input_b.set_tensor_amax_columnwise(amax_b); - - nvte_nvfp4_compute_per_tensor_scale( - input_a.data(), true, input_b.data(), false, alpha_in, alpha_out.data(), 0); - alpha_out.to_cpu(); - - const float factor_inv = - 1.0f / (fp4_max * fp4_max * ue5m3_max * ue5m3_max); - const float expected = alpha_in * amax_a * amax_b * factor_inv; - EXPECT_FLOAT_EQ(alpha_out.rowwise_cpu_dptr()[0], expected); -} #endif class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam From 85cf34faddcf2045658d8a0eacc39cbfc3766ac4 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 1 Sep 2026 01:52:05 +0000 Subject: [PATCH 46/54] Add reference impl and test for NVFP4-UE5M3 Co-authored-by: Codex Signed-off-by: Tim Moon --- .../nvfp4/test_nvfp4_quantize_exact.py | 46 +++- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 17 +- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 16 +- .../common/transformer_engine.cpp | 2 +- .../pytorch/custom_recipes/reference_nvfp4.py | 204 ++++++++++++++---- 5 files changed, 221 insertions(+), 64 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 20c0042fe6..0865326ae1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -17,23 +17,26 @@ recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) +ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True) NVFP4_E4M3_AMAX_FOR_UNIT_GLOBAL_SCALE = 448.0 * 6.0 +NVFP4_SUPPORTED_SCALE_MAX_VALUES = (0, 256, 448, 114688) @dataclass(frozen=True) class NVFP44Over6TestConfig: id: str use_4over6: bool = True - e4m3_max: int = 448 + e4m3_max: int = 0 err_mode: str = "MAE" err_use_fast_math: bool = False NVFP4_4OVER6_CONFIGS = [ NVFP44Over6TestConfig(id="nvfp4", use_4over6=False), - NVFP44Over6TestConfig(id="4over6-mae-e4m3-448-exact", err_mode="MAE"), + NVFP44Over6TestConfig(id="4over6-mae-e4m3-448-exact", e4m3_max=448, err_mode="MAE"), NVFP44Over6TestConfig( id="4over6-mae-e4m3-448-err-fast", + e4m3_max=448, err_mode="MAE", err_use_fast_math=True, ), @@ -44,9 +47,10 @@ class NVFP44Over6TestConfig: err_mode="MAE", err_use_fast_math=True, ), - NVFP44Over6TestConfig(id="4over6-mse-e4m3-448-exact", err_mode="MSE"), + NVFP44Over6TestConfig(id="4over6-mse-e4m3-448-exact", e4m3_max=448, err_mode="MSE"), NVFP44Over6TestConfig( id="4over6-mse-e4m3-448-err-fast", + e4m3_max=448, err_mode="MSE", err_use_fast_math=True, ), @@ -115,12 +119,18 @@ def check_quantization_nvfp4_versus_reference( with_2d_quantization: bool, row_scaled_nvfp4: bool = False, use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, nvfp4_4over6_err_mode: str = "MAE", nvfp4_4over6_err_use_fast_math: bool = False, + scale_dtype: te.DType = te.DType.kFloat8E4M3, ) -> None: - if nvfp4_e4m3_max != 448 and not use_4over6: - pytest.skip("E4M3 max 256 is only meaningful for 4over6") + scale_dtype = te.DType.cast(scale_dtype) + if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES: + raise ValueError( + "nvfp4_e4m3_max must be 0, 256, 448, or 114688." + ) + if use_4over6 and scale_dtype == te.DType.kFloat8UE5M3: + pytest.skip("NVFP4 4over6 is incompatible with UE5M3 scales") maybe_skip_row_scaled_unsupported_quantization( row_scaled_nvfp4, return_transpose, with_2d_quantization, use_4over6, x_dtype, M, N ) @@ -149,6 +159,7 @@ def check_quantization_nvfp4_versus_reference( nvfp4_use_4over6=use_4over6, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, + scale_dtype=scale_dtype, ) if use_4over6: @@ -169,6 +180,8 @@ def check_quantization_nvfp4_versus_reference( ) x_nvfp4_sut = nvfp4_quantizer.update_quantized(x, x_nvfp4_sut) + assert x_nvfp4_sut._scale_dtype == scale_dtype + # Extract data from NVFP4Tensor assert x_nvfp4_sut._rowwise_data is not None qx: torch.Tensor = x_nvfp4_sut._rowwise_data.view(dtype=torch.uint8) @@ -197,6 +210,7 @@ def check_quantization_nvfp4_versus_reference( nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, nvfp4_4over6_err_use_fast_math=nvfp4_4over6_err_use_fast_math, + scale_dtype=scale_dtype, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -495,6 +509,26 @@ def test_nvfp4_quantization_extrema_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.skipif(not ue5m3_available, reason=reason_for_no_ue5m3) +@pytest.mark.parametrize( + "with_2d_quantization", + [False, True], + ids=["1d_quantization", "2d_quantization"], +) +def test_nvfp4_ue5m3_quantization_versus_reference(with_2d_quantization: bool) -> None: + check_quantization_nvfp4_versus_reference( + x_dtype=torch.bfloat16, + M=128, + N=128, + return_transpose=True, + swizzled_scale=False, + use_cpp_allocator=False, + with_2d_quantization=with_2d_quantization, + scale_dtype=te.DType.kFloat8UE5M3, + ) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 788b27248c..2f208dde41 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -142,20 +142,21 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) NVTE_CHECK(!row_scaled_nvfp4 || input.amax.numel() == N, "Row-scaled NVFP4 dequantization requires one rowwise amax per row."); const int full_scale_max = static_cast(typeToMax(scale_dtype)); - const bool uses_4over6_headroom = (scale_dtype == DType::kFloat8E4M3 && scale_type_max == 256) || - (scale_dtype == DType::kFloat8UE5M3 && scale_type_max == 65536); + const bool uses_4over6_headroom = scale_dtype == DType::kFloat8E4M3 && scale_type_max == 256; NVTE_CHECK(scale_type_max == full_scale_max || uses_4over6_headroom, "Unsupported maximum ", scale_type_max, " for NVFP4 scale dtype ", to_string(scale_dtype), "."); - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { - constexpr int full_max = static_cast(TypeInfo::max_finite_value); - constexpr int headroom_max = std::is_same_v ? 256 : 65536; - TRANSFORMER_ENGINE_SWITCH_CONDITION(scale_type_max == headroom_max, USE_HEADROOM, { - constexpr int SCALE_TYPE_MAX = USE_HEADROOM ? headroom_max : full_max; + if (uses_4over6_headroom) { + launch_dequantize(input, output, with_gemm_swizzled_scales, + row_scaled_nvfp4, N, Mread, blocks, threads, + num_scale_tiles_X, stream); + } else { + TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { + constexpr int SCALE_TYPE_MAX = static_cast(TypeInfo::max_finite_value); launch_dequantize(input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, Mread, blocks, threads, num_scale_tiles_X, stream); }) - }) + } NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 84b25ee871..5aeb97961d 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -89,14 +89,6 @@ struct FourOverSixScaleConfig { static constexpr bool supports_fp16_error_path = true; }; -#if CUDA_VERSION >= 13040 -template <> -struct FourOverSixScaleConfig { - static constexpr int headroom_max = 65536; - static constexpr bool supports_fp16_error_path = false; -}; -#endif - struct Candidate { uint32_t packed[kPackedWordsPerGroup]; float err; @@ -798,6 +790,8 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, "NVFP4 4over6 output tensor must have rowwise or columnwise data."); const DType scale_dtype = return_rowwise ? output->scale_inv.dtype : output->columnwise_scale_inv.dtype; + NVTE_CHECK(scale_dtype == DType::kFloat8E4M3, + "NVFP4 4over6 is only supported with FP8E4M3 scales."); if (return_rowwise && return_transpose) { NVTE_CHECK(output->scale_inv.dtype == output->columnwise_scale_inv.dtype, "Rowwise and columnwise NVFP4 scale tensors must have the same dtype (got ", @@ -805,10 +799,8 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, to_string(output->columnwise_scale_inv.dtype), ")."); } - TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH( - scale_dtype, ScaleType, - quantize_4over6_impl(input, noop, output, quant_config, - scale_dtype, stream);) + quantize_4over6_impl(input, noop, output, quant_config, + scale_dtype, stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index df30a86192..e1c854000a 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -930,7 +930,7 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); // Need to rename this to nvfp4_scale_type_max NVTE_CHECK(t.nvfp4_e4m3_max == 0 || t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256 || - t.nvfp4_e4m3_max == 114688 || t.nvfp4_e4m3_max == 65536, + t.nvfp4_e4m3_max == 114688, "Unsupported NVFP4 scale type max (got ", t.nvfp4_e4m3_max, ")"); break; default: diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index 8c5f305554..b397fab8fd 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -11,9 +11,16 @@ from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils +from transformer_engine.pytorch.constants import DType from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +NVFP4_FP4_MAX = 6.0 +NVFP4_E4M3_SCALE_MAX = 448.0 +NVFP4_UE5M3_SCALE_MAX = 114688.0 +NVFP4_SUPPORTED_SCALE_MAX_VALUES = (0, 256, 448, 114688) + + def nvfp4_ref_rht_2d_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). @@ -141,6 +148,104 @@ def cast_to_e4m3(decode_scale, global_amax): return decode_scale.to(torch.float8_e4m3fn) +def _cast_to_ue5m3(decode_scale: torch.Tensor) -> torch.Tensor: + """Cast positive scale values to UE5M3 byte storage.""" + bias = 15 + mantissa_bits = 3 + max_code = 0xFE + min_normal = torch.tensor(2.0**-14, device=decode_scale.device, dtype=torch.float32) + subnormal_step = torch.tensor(2.0**-17, device=decode_scale.device, dtype=torch.float32) + x = torch.nan_to_num( + decode_scale.to(torch.float32), nan=0.0, posinf=NVFP4_UE5M3_SCALE_MAX + ) + x = torch.clamp(x, min=0.0, max=NVFP4_UE5M3_SCALE_MAX) + + subnormal_code = torch.round(x / subnormal_step).to(torch.int32) + + normal_x = torch.clamp(x, min=min_normal) + exponent = torch.floor(torch.log2(normal_x)) + exponent_value = torch.pow(torch.tensor(2.0, device=x.device, dtype=torch.float32), exponent) + mantissa = torch.round((normal_x / exponent_value - 1.0) * (1 << mantissa_bits)).to( + torch.int32 + ) + exponent_field = exponent.to(torch.int32) + bias + exponent_field = exponent_field + (mantissa == (1 << mantissa_bits)).to(torch.int32) + mantissa = torch.where(mantissa == (1 << mantissa_bits), torch.zeros_like(mantissa), mantissa) + normal_code = torch.bitwise_or( + torch.bitwise_left_shift(exponent_field, mantissa_bits), mantissa + ) + normal_code = torch.where( + exponent_field > 31, torch.full_like(normal_code, max_code), normal_code + ) + normal_code = torch.where( + normal_code > max_code, torch.full_like(normal_code, max_code), normal_code + ) + + code = torch.where(x < min_normal, subnormal_code, normal_code) + return torch.clamp(code, min=0, max=max_code).to(torch.uint8) + + +def _cast_to_nvfp4_scale(decode_scale: torch.Tensor, scale_dtype: DType) -> torch.Tensor: + scale_dtype = DType.cast(scale_dtype) + if scale_dtype == DType.kFloat8E4M3: + scale_max = torch.tensor( + NVFP4_E4M3_SCALE_MAX, device=decode_scale.device, dtype=torch.float32 + ) + return torch.clamp(decode_scale, min=-scale_max, max=scale_max).to(torch.float8_e4m3fn) + if scale_dtype == DType.kFloat8UE5M3: + return _cast_to_ue5m3(decode_scale) + raise ValueError(f"Unsupported NVFP4 scale dtype: {scale_dtype}.") + + +def _ue5m3_to_float32(scale: torch.Tensor) -> torch.Tensor: + """Decode UE5M3 byte storage to FP32.""" + bias = 15 + mantissa_bits = 3 + code = scale.contiguous().view(torch.uint8).to(torch.int32) + exponent_field = torch.bitwise_right_shift(code, mantissa_bits) + mantissa = torch.bitwise_and(code, (1 << mantissa_bits) - 1) + + subnormal = mantissa.to(torch.float32) * torch.tensor( + 2.0**-17, device=scale.device, dtype=torch.float32 + ) + significand = 1.0 + mantissa.to(torch.float32) / float(1 << mantissa_bits) + normal = torch.ldexp(significand, exponent_field - bias) + return torch.where(exponent_field == 0, subnormal, normal) + + +def _nvfp4_scale_to_float32(scale: torch.Tensor, scale_dtype: DType) -> torch.Tensor: + scale_dtype = DType.cast(scale_dtype) + if scale_dtype == DType.kFloat8E4M3: + if scale.dtype == torch.uint8: + return scale.contiguous().view(torch.float8_e4m3fn).to(torch.float32) + return scale.to(torch.float32) + if scale_dtype == DType.kFloat8UE5M3: + return _ue5m3_to_float32(scale) + raise ValueError(f"Unsupported NVFP4 scale dtype: {scale_dtype}.") + + +def _nvfp4_scale_dtype_max(scale_dtype: DType) -> float: + scale_dtype = DType.cast(scale_dtype) + if scale_dtype == DType.kFloat8E4M3: + return NVFP4_E4M3_SCALE_MAX + if scale_dtype == DType.kFloat8UE5M3: + return NVFP4_UE5M3_SCALE_MAX + raise ValueError(f"Unsupported NVFP4 scale dtype: {scale_dtype}.") + + +def _validate_nvfp4_scale_max(nvfp4_e4m3_max: int) -> None: + if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES: + raise ValueError( + "nvfp4_e4m3_max must be 0, 256, 448, or 114688." + ) + + +def _nvfp4_effective_scale_max(scale_dtype: DType, nvfp4_e4m3_max: int) -> float: + if nvfp4_e4m3_max != 0: + return float(nvfp4_e4m3_max) + return _nvfp4_scale_dtype_max(scale_dtype) + + def high_precision_gemm_ref( a: torch.Tensor, b: torch.Tensor, @@ -222,7 +327,8 @@ class NVFP4TensorRef(QuantizedTensorStorage): global_amax_row: Optional[torch.Tensor] = None global_amax_col: Optional[torch.Tensor] = None nvfp4_use_4over6: bool = False - nvfp4_e4m3_max: int = 448 + nvfp4_e4m3_max: int = 0 + scale_dtype: DType = DType.kFloat8E4M3 dtype: Optional[torch.dtype] = None device: Optional[torch.device] = None @@ -273,12 +379,21 @@ def _scale_inv(self): def _scale_inv(self, value): self.scale = value + @property + def _scale_dtype(self): + return self.scale_dtype + + @_scale_dtype.setter + def _scale_dtype(self, value): + self.scale_dtype = DType.cast(value) + def __repr__(self): return ( f"{self.__class__.__name__}(" f"dtype={self.dtype}, " f"device={self.device}, " f"quant_dtype={self.quant_dtype}, " + f"scale_dtype={self.scale_dtype}, " f"original_shape={self.original_shape}" ")" ) @@ -356,10 +471,14 @@ def __init__( nvfp4_e4m3_max: int = 0, nvfp4_4over6_err_mode: str = "MAE", nvfp4_4over6_err_use_fast_math: bool = False, + scale_dtype: Union[DType, int] = DType.kFloat8E4M3, with_rht: bool = False, with_random_sign_mask: bool = True, ): nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() + scale_dtype = DType.cast(scale_dtype) + if scale_dtype not in (DType.kFloat8E4M3, DType.kFloat8UE5M3): + raise ValueError("scale_dtype must be DType.kFloat8E4M3 or DType.kFloat8UE5M3.") if row_scaled_nvfp4: if not rowwise: raise ValueError("Row-scaled NVFP4 reference quantization requires rowwise usage.") @@ -368,6 +487,8 @@ def __init__( raise ValueError(f"Unsupported NVFP4 4over6 error mode: {nvfp4_4over6_err_mode}.") if pow_2_scales: raise ValueError("4over6 is only supported for NVFP4 (non-pow2) mode.") + if scale_dtype == DType.kFloat8UE5M3: + raise ValueError("4over6 is incompatible with UE5M3 scales.") if quant_tile_shape not in ((1, 16), (16, 16)): raise ValueError("4over6 reference quantization only supports 1x16 or 16x16 tiles.") super().__init__(rowwise=rowwise, columnwise=columnwise) @@ -379,11 +500,11 @@ def __init__( self.quant_tile_shape = quant_tile_shape self.row_scaled_nvfp4 = row_scaled_nvfp4 self.nvfp4_use_4over6 = nvfp4_use_4over6 - self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_e4m3_max != 0 else 448 - if self.nvfp4_e4m3_max not in (448, 256): - raise ValueError("nvfp4_e4m3_max must be 448 or 256.") + _validate_nvfp4_scale_max(nvfp4_e4m3_max) + self.nvfp4_e4m3_max = nvfp4_e4m3_max self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode self.nvfp4_4over6_err_use_fast_math = nvfp4_4over6_err_use_fast_math + self.scale_dtype = scale_dtype self.with_rht = with_rht self.with_random_sign_mask = with_random_sign_mask @@ -658,9 +779,10 @@ def _quantize_blockwise_reference( pow_2_scales: bool, row_scaled_nvfp4: bool = False, nvfp4_use_4over6: bool = False, - nvfp4_e4m3_max: int = 448, + nvfp4_e4m3_max: int = 0, nvfp4_4over6_err_mode: str = "MAE", nvfp4_4over6_err_use_fast_math: bool = False, + scale_dtype: DType = DType.kFloat8E4M3, eps: float, # pylint: disable=unused-argument ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -669,6 +791,10 @@ def _quantize_blockwise_reference( f"_quantize_blockwise_reference expects a 2D tensor, got {x.ndim}D with shape" f" {x.shape}" ) + scale_dtype = DType.cast(scale_dtype) + if nvfp4_use_4over6 and scale_dtype == DType.kFloat8UE5M3: + raise ValueError("4over6 is incompatible with UE5M3 scales.") + _validate_nvfp4_scale_max(nvfp4_e4m3_max) using_2d_quantization = tile_len_x == 16 and tile_len_y == 16 m, n = x.shape # Compute vec_max based on the original x (before reshape) @@ -691,11 +817,10 @@ def _quantize_blockwise_reference( torch.float32 ) # (128, 8, 1) x = x.view(m, n // tile_len_x, tile_len_x) - FLOAT4_E2M1_MAX = torch.tensor(6.0, device=x.device, dtype=torch.float32) - FLOAT8_E4M3_MAX = torch.tensor(448.0, device=x.device, dtype=torch.float32) - global_scale_e4m3_max = float(nvfp4_e4m3_max if nvfp4_use_4over6 else 448) - GLOBAL_SCALE_E4M3_MAX = torch.tensor( - global_scale_e4m3_max, device=x.device, dtype=torch.float32 + FLOAT4_E2M1_MAX = torch.tensor(NVFP4_FP4_MAX, device=x.device, dtype=torch.float32) + global_scale_max = _nvfp4_effective_scale_max(scale_dtype, nvfp4_e4m3_max) + GLOBAL_SCALE_MAX = torch.tensor( + global_scale_max, device=x.device, dtype=torch.float32 ) decode_scale = torch.div(vec_max, FLOAT4_E2M1_MAX) @@ -709,7 +834,7 @@ def _quantize_blockwise_reference( if row_scaled_nvfp4: global_amax = global_amax.to(torch.float32).view(m, 1, 1) - global_encode_scale = torch.div(GLOBAL_SCALE_E4M3_MAX * FLOAT4_E2M1_MAX, global_amax) + global_encode_scale = torch.div(GLOBAL_SCALE_MAX * FLOAT4_E2M1_MAX, global_amax) global_encode_scale = torch.min( global_encode_scale, torch.tensor( @@ -742,7 +867,7 @@ def _quantize_blockwise_reference( tile_len_y, nvfp4_4over6_err_mode, nvfp4_4over6_err_use_fast_math, - nvfp4_e4m3_max, + int(global_scale_max), ) global_encode_scale_multiplier = global_encode_scale * torch.reciprocal(FLOAT4_E2M1_MAX) @@ -758,11 +883,13 @@ def _quantize_blockwise_reference( dtype=torch.float32, ), ) - decode_scale = torch.clamp(decode_scale, min=-FLOAT8_E4M3_MAX, max=FLOAT8_E4M3_MAX) - decode_scale = decode_scale.to(torch.float8_e4m3fn) + decode_scale = _cast_to_nvfp4_scale(decode_scale, scale_dtype) encode_scale = torch.min( - torch.div(1.0, decode_scale.to(torch.float32) * global_decode_scale), + torch.div( + 1.0, + _nvfp4_scale_to_float32(decode_scale, scale_dtype) * global_decode_scale, + ), torch.tensor( torch.finfo(torch.float32).max, device=decode_scale.device, @@ -916,6 +1043,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, nvfp4_4over6_err_use_fast_math=self.nvfp4_4over6_err_use_fast_math, + scale_dtype=self.scale_dtype, eps=self.eps, ) if transpose_scales: @@ -944,6 +1072,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ nvfp4_e4m3_max=self.nvfp4_e4m3_max, nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, nvfp4_4over6_err_use_fast_math=self.nvfp4_4over6_err_use_fast_math, + scale_dtype=self.scale_dtype, eps=self.eps, ) @@ -985,6 +1114,7 @@ def quantize( global_amax_col=global_amax_col, nvfp4_use_4over6=self.nvfp4_use_4over6, nvfp4_e4m3_max=self.nvfp4_e4m3_max, + scale_dtype=self.scale_dtype, dtype=tensor.dtype, device=tensor.device, quant_dtype=self.dtype, @@ -1034,6 +1164,7 @@ def update_quantized( dst.global_amax_col = global_amax_col dst.nvfp4_use_4over6 = self.nvfp4_use_4over6 dst.nvfp4_e4m3_max = self.nvfp4_e4m3_max + dst.scale_dtype = self.scale_dtype dst.dtype = src.dtype dst.quant_dtype = self.dtype dst.original_shape = original_shape @@ -1136,19 +1267,6 @@ def qgemm( "qresult_w.global_amax_col must be set for non-pow_2_scales NVFP4 GEMM" ) - sx = sx.to(torch.float32) - sw = sw.to(torch.float32) - - qresult_x_nvfp4_use_4over6 = getattr( - qresult_x, - "nvfp4_use_4over6", - getattr(qresult_x, "_nvfp4_use_4over6", self.nvfp4_use_4over6), - ) - qresult_w_nvfp4_use_4over6 = getattr( - qresult_w, - "nvfp4_use_4over6", - getattr(qresult_w, "_nvfp4_use_4over6", self.nvfp4_use_4over6), - ) qresult_x_e4m3_max = getattr( qresult_x, "nvfp4_e4m3_max", @@ -1159,15 +1277,27 @@ def qgemm( "nvfp4_e4m3_max", getattr(qresult_w, "_nvfp4_e4m3_max", self.nvfp4_e4m3_max), ) - if qresult_x_nvfp4_use_4over6: - fp8_max_x = float(qresult_x_e4m3_max) - else: - fp8_max_x = 448.0 - if qresult_w_nvfp4_use_4over6: - fp8_max_w = float(qresult_w_e4m3_max) - else: - fp8_max_w = 448.0 - factor = 6.0 * 6.0 * fp8_max_x * fp8_max_w + qresult_x_scale_dtype = DType.cast( + getattr( + qresult_x, + "scale_dtype", + getattr(qresult_x, "_scale_dtype", self.scale_dtype), + ) + ) + qresult_w_scale_dtype = DType.cast( + getattr( + qresult_w, + "scale_dtype", + getattr(qresult_w, "_scale_dtype", self.scale_dtype), + ) + ) + sx = _nvfp4_scale_to_float32(sx, qresult_x_scale_dtype) + sw = _nvfp4_scale_to_float32(sw, qresult_w_scale_dtype) + _validate_nvfp4_scale_max(qresult_x_e4m3_max) + _validate_nvfp4_scale_max(qresult_w_e4m3_max) + fp8_max_x = _nvfp4_effective_scale_max(qresult_x_scale_dtype, qresult_x_e4m3_max) + fp8_max_w = _nvfp4_effective_scale_max(qresult_w_scale_dtype, qresult_w_e4m3_max) + factor = NVFP4_FP4_MAX * NVFP4_FP4_MAX * fp8_max_x * fp8_max_w if gemm_type == gemm.GEMMType.WGRAD: partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col From bdc1d97f165165c6781248feedf12554a79699d6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:54:25 +0000 Subject: [PATCH 47/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../nvfp4/test_nvfp4_quantize_exact.py | 4 +-- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 5 ++-- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 4 +-- .../include/transformer_engine/recipe.h | 7 +++-- transformer_engine/common/recipe/nvfp4.cu | 4 +-- .../pytorch/csrc/extensions/transpose.cpp | 26 ++++++++----------- .../pytorch/custom_recipes/reference_nvfp4.py | 16 +++--------- 7 files changed, 25 insertions(+), 41 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 0865326ae1..51b5b46c15 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -126,9 +126,7 @@ def check_quantization_nvfp4_versus_reference( ) -> None: scale_dtype = te.DType.cast(scale_dtype) if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES: - raise ValueError( - "nvfp4_e4m3_max must be 0, 256, 448, or 114688." - ) + raise ValueError("nvfp4_e4m3_max must be 0, 256, 448, or 114688.") if use_4over6 and scale_dtype == te.DType.kFloat8UE5M3: pytest.skip("NVFP4 4over6 is incompatible with UE5M3 scales") maybe_skip_row_scaled_unsupported_quantization( diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 2f208dde41..ac30fe3f91 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -146,9 +146,8 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) NVTE_CHECK(scale_type_max == full_scale_max || uses_4over6_headroom, "Unsupported maximum ", scale_type_max, " for NVFP4 scale dtype ", to_string(scale_dtype), "."); if (uses_4over6_headroom) { - launch_dequantize(input, output, with_gemm_swizzled_scales, - row_scaled_nvfp4, N, Mread, blocks, threads, - num_scale_tiles_X, stream); + launch_dequantize(input, output, with_gemm_swizzled_scales, row_scaled_nvfp4, N, + Mread, blocks, threads, num_scale_tiles_X, stream); } else { TRANSFORMER_ENGINE_NVFP4_SCALE_TYPE_SWITCH(scale_dtype, ScaleType, { constexpr int SCALE_TYPE_MAX = static_cast(TypeInfo::max_finite_value); diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh index 5aeb97961d..7311e50be8 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -799,8 +799,8 @@ void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, to_string(output->columnwise_scale_inv.dtype), ")."); } - quantize_4over6_impl(input, noop, output, quant_config, - scale_dtype, stream); + quantize_4over6_impl(input, noop, output, quant_config, scale_dtype, + stream); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 9ee26364f6..b3feeba6b7 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -431,10 +431,9 @@ void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, s * \param[in] scale_dtype NVFP4 scale storage type (E4M3 or UE5M3). * \param[in] stream CUDA stream. */ -void nvte_nvfp4_expand_scale_to_fp8_v2(const NVTETensor input, NVTETensor output, - size_t tile_rows, size_t tile_cols, size_t rows_padded, - size_t block_len, NVTEDType scale_dtype, - cudaStream_t stream); +void nvte_nvfp4_expand_scale_to_fp8_v2(const NVTETensor input, NVTETensor output, size_t tile_rows, + size_t tile_cols, size_t rows_padded, size_t block_len, + NVTEDType scale_dtype, cudaStream_t stream); /*! \brief Compute per-block E4M3 decode scale from block amax and global amax. * diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 9d46add5ac..14a0264074 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -938,8 +938,8 @@ void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTE size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, size_t block_len, cudaStream_t stream) { NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); - nvte_nvfp4_2d_partial_cast_v2(inp, out, scale, global_scale, h, w, scale_stride_h, - scale_stride_w, start_offset, block_len, kNVTEFloat8E4M3, stream); + nvte_nvfp4_2d_partial_cast_v2(inp, out, scale, global_scale, h, w, scale_stride_h, scale_stride_w, + start_offset, block_len, kNVTEFloat8E4M3, stream); } void nvte_nvfp4_2d_partial_cast_v2(const NVTETensor inp, NVTETensor out, const NVTETensor scale, diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index d3dcfe6604..938b0f83b7 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -164,8 +164,7 @@ void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at:: auto global_amax_cu = makeTransformerEngineTensor(global_amax); nvte_nvfp4_compute_per_block_scale_v2(block_amax_cu.data(), scale_cu.data(), - global_amax_cu.data(), - static_cast(scale_dtype), + global_amax_cu.data(), static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); } @@ -183,8 +182,7 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); NVTE_CHECK(per_block_scale.scalar_type() == at::kFloat, "Per-block scale must be float32."); - NVTE_CHECK(target_scale.scalar_type() == at::kByte, - "Target scale must be uint8 scale storage."); + NVTE_CHECK(target_scale.scalar_type() == at::kByte, "Target scale must be uint8 scale storage."); NVTE_CHECK(target_amax.scalar_type() == at::kFloat, "Target amax must be float32."); NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); @@ -195,12 +193,11 @@ void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor auto target_scale_cu = makeTransformerEngineTensor(target_scale); auto target_amax_cu = makeTransformerEngineTensor(target_amax); - nvte_nvfp4_fused_scale_v2(block_amax_cu.data(), global_amax_cu.data(), - per_block_scale_cu.data(), target_scale_cu.data(), - target_amax_cu.data(), static_cast(tile_rows), - static_cast(tile_cols), static_cast(rows_padded), - static_cast(block_len), static_cast(scale_dtype), - at::cuda::getCurrentCUDAStream()); + nvte_nvfp4_fused_scale_v2(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), + target_scale_cu.data(), target_amax_cu.data(), + static_cast(tile_rows), static_cast(tile_cols), + static_cast(rows_padded), static_cast(block_len), + static_cast(scale_dtype), at::cuda::getCurrentCUDAStream()); } void nvfp4_multi_tensor_fused_scale( @@ -251,11 +248,10 @@ void nvfp4_multi_tensor_fused_scale( auto target_scale_cu = makeTransformerEngineTensor(target_scale); auto target_amax_cu = makeTransformerEngineTensor(target_amax); - nvte_nvfp4_fused_scale_v2(block_amax_cu.data(), global_amax_cu.data(), - per_block_scale_cu.data(), target_scale_cu.data(), - target_amax_cu.data(), tile_rows, tile_cols, rows_padded, - static_cast(block_len), - static_cast(scale_dtype), stream); + nvte_nvfp4_fused_scale_v2( + block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), + target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, rows_padded, + static_cast(block_len), static_cast(scale_dtype), stream); } } diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index b397fab8fd..381baf0bb1 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -155,9 +155,7 @@ def _cast_to_ue5m3(decode_scale: torch.Tensor) -> torch.Tensor: max_code = 0xFE min_normal = torch.tensor(2.0**-14, device=decode_scale.device, dtype=torch.float32) subnormal_step = torch.tensor(2.0**-17, device=decode_scale.device, dtype=torch.float32) - x = torch.nan_to_num( - decode_scale.to(torch.float32), nan=0.0, posinf=NVFP4_UE5M3_SCALE_MAX - ) + x = torch.nan_to_num(decode_scale.to(torch.float32), nan=0.0, posinf=NVFP4_UE5M3_SCALE_MAX) x = torch.clamp(x, min=0.0, max=NVFP4_UE5M3_SCALE_MAX) subnormal_code = torch.round(x / subnormal_step).to(torch.int32) @@ -165,9 +163,7 @@ def _cast_to_ue5m3(decode_scale: torch.Tensor) -> torch.Tensor: normal_x = torch.clamp(x, min=min_normal) exponent = torch.floor(torch.log2(normal_x)) exponent_value = torch.pow(torch.tensor(2.0, device=x.device, dtype=torch.float32), exponent) - mantissa = torch.round((normal_x / exponent_value - 1.0) * (1 << mantissa_bits)).to( - torch.int32 - ) + mantissa = torch.round((normal_x / exponent_value - 1.0) * (1 << mantissa_bits)).to(torch.int32) exponent_field = exponent.to(torch.int32) + bias exponent_field = exponent_field + (mantissa == (1 << mantissa_bits)).to(torch.int32) mantissa = torch.where(mantissa == (1 << mantissa_bits), torch.zeros_like(mantissa), mantissa) @@ -235,9 +231,7 @@ def _nvfp4_scale_dtype_max(scale_dtype: DType) -> float: def _validate_nvfp4_scale_max(nvfp4_e4m3_max: int) -> None: if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES: - raise ValueError( - "nvfp4_e4m3_max must be 0, 256, 448, or 114688." - ) + raise ValueError("nvfp4_e4m3_max must be 0, 256, 448, or 114688.") def _nvfp4_effective_scale_max(scale_dtype: DType, nvfp4_e4m3_max: int) -> float: @@ -819,9 +813,7 @@ def _quantize_blockwise_reference( x = x.view(m, n // tile_len_x, tile_len_x) FLOAT4_E2M1_MAX = torch.tensor(NVFP4_FP4_MAX, device=x.device, dtype=torch.float32) global_scale_max = _nvfp4_effective_scale_max(scale_dtype, nvfp4_e4m3_max) - GLOBAL_SCALE_MAX = torch.tensor( - global_scale_max, device=x.device, dtype=torch.float32 - ) + GLOBAL_SCALE_MAX = torch.tensor(global_scale_max, device=x.device, dtype=torch.float32) decode_scale = torch.div(vec_max, FLOAT4_E2M1_MAX) if pow_2_scales: From dc40106b706de49a14ca9fb729128fc530d3ed5b Mon Sep 17 00:00:00 2001 From: Michal Futrega Date: Tue, 1 Sep 2026 21:47:18 +0200 Subject: [PATCH 48/54] Make NVFP4-UE5M3 grouped GEMM CUDA-graph safe (#4) * Make NVFP4-UE5M3 grouped GEMM CUDA-graph safe The UE5M3 grouped GEMM builds its one-element int32 offsets tensor inline at two call sites: torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device) # wgrad torch.tensor([N], dtype=torch.int32, device=device) # fprop Each of these materializes an unpinned CPU tensor and copies it host-to-device. CUDA rejects that during graph capture: RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned. Please use tensor.pin_memory() or allocate the tensor with pin_memory=True. Any model capturing these GEMMs in a CUDA graph therefore fails at capture time. This was hit end-to-end on DeepSeek-V3 671B, where every convergence config captures full forward/backward graphs; it aborts during warmup capture and no training step completes. Route both sites through a cached helper, mirroring get_cached_ones_tensor in this same file -- which both functions already use for their `ones` tensor a few lines earlier, and whose docstring notes it keeps "stable data pointers across CUDA graph replays". Caching by (value, device) moves the host-to-device copy to warmup and keeps the pointer stable across replays. The offsets are constant for a given shape, so this is numerically neutral. Signed-off-by: Michal Futrega * Review suggestions Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Michal Futrega Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/cpp_extensions/gemm.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 1cdab8e09c..be2a0981e9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -231,6 +231,15 @@ def _cudnn_grouped_gemm_quant_kernel() -> Callable: return grouped_gemm_quant_wrapper_sm100 +@functools.lru_cache +def _get_cached_offsets_tensor( + value: int, + device: torch.device, +) -> torch.Tensor: + """Return a cached int32 one-element offsets tensor.""" + return torch.full((1,), value, dtype=torch.int32, device=device) + + @functools.lru_cache(maxsize=None) def _cudnn_grouped_gemm_wgrad_kernel() -> Callable: """cuDNN CuTe DSL grouped wgrad kernel for block-scaled inputs. @@ -304,7 +313,7 @@ def _sf(scale_inv, features): b_tensor=b_tensor, sfa_tensor=_sf(sfa, out_features), sfb_tensor=_sf(sfb, in_features), - offsets_tensor=torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device), + offsets_tensor=_get_cached_offsets_tensor(int(tokens), a_tensor.device), global_scale_a=global_scale_a, global_scale_b=global_scale_b, acc_dtype=torch.float32, @@ -664,7 +673,7 @@ def _cudnn_grouped_gemm_nvfp4_ue5m3( "b_tensor": cudnn_b, "sfb_tensor": cudnn_sfb, # One group, so the only padded end offset is the full row count. - "padded_offsets": torch.tensor([N], dtype=torch.int32, device=device), + "padded_offsets": _get_cached_offsets_tensor(int(N), device), "alpha_tensor": alpha_tensor, "bias_tensor": bias, "norm_const_tensor": None, # must be None for FP4 inputs From 4be085f0479fa076723dd8e630d0da33a253b2ae Mon Sep 17 00:00:00 2001 From: tdophung Date: Tue, 1 Sep 2026 16:34:34 -0700 Subject: [PATCH 49/54] Fix JAX NVFP4 GEMM device scalar handling Signed-off-by: tdophung --- transformer_engine/common/gemm/cublaslt_gemm.cu | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index aca0173be9..d491cab9d6 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -372,9 +372,12 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype); const bool nvfp4_tensor_scaling = is_nvfp_scaling(inputA->scaling_mode) && is_nvfp_scaling(inputB->scaling_mode); + const bool have_nvfp4_amax = + (transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr || + (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr; // Update scaling factors with NVFP4 tensor scales - if (use_fp4 && nvfp4_tensor_scaling) { + if (use_fp4 && nvfp4_tensor_scaling && have_nvfp4_amax) { // Reserve some workspace for alpha scale NVTE_CHECK(workspaceSize >= 4, "NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has ", From 1bf887c1ce85128de672ee14084da0b3faa4c589 Mon Sep 17 00:00:00 2001 From: tdophung Date: Tue, 1 Sep 2026 22:54:52 -0700 Subject: [PATCH 50/54] Replace NVFP4 amax heuristic with explicit alpha/beta residency flag b3415c19 gated the NVFP4 per-tensor scale branch in cublas_gemm() on have_nvfp4_amax, an OR over the two operands' amax pointers. That branch is also what moves alpha/beta from host into device workspace, and CUBLASLT_POINTER_MODE_DEVICE is set unconditionally for every NVFP4 GEMM. So when both operands had a null amax the branch was skipped, leaving alpha/beta as host stack pointers that cuBLAS then dereferenced as device addresses -- garbage output, not a mis-scaled result. That is exactly the shape of the CI failure: only the both_unit_global_scale parametrization of test_gemm_with_missing_nvfp4_amax went red, while the x_/w_ cases passed because the OR still held. Note that nvte_nvfp4_compute_per_tensor_scale with a null amax is an exact identity on alpha (it substitutes scale_max * fp4_max), so skipping it never saved meaningful work -- its real job on that path is materializing the scalars on device. Making the call is always safe; skipping it is not. Replace the heuristic with kNVTEMatmulConfigAlphaBetaOnDevice, which states pointer residency directly instead of inferring it. Default false, so PyTorch (nvte_cublas_gemm/_scaled pass host pointers, and generic_gemm reaches nvte_cublas_gemm_v2 without setting the flag) is restored to its pre-b3415c19 behavior. JAX sets it when is_nvfp4_scaling(), matching the device buffers it binds for alpha/beta, and folds the per-tensor scale into alpha itself. Also document in the JAX helper that its block-scale bound is hardcoded to E4M3's 448 while the C++ get_nvfp4_scale_max() returns 114688 for UE5M3, which will need reconciling once UE5M3 reaches the JAX path. Verified on GB200: tests/pytorch/nvfp4 9310 passed / 0 failed / 6852 skipped (CI at b3415c19 was 9309 passed / 1 failed / 6852 skipped), tests/pytorch/mxfp8 371 passed, and JAX nvfp4+gemm 382 passed. Co-Authored-By: Claude Opus 5 Signed-off-by: tdophung --- transformer_engine/common/gemm/config.cpp | 6 +++++ transformer_engine/common/gemm/config.h | 4 +++- .../common/gemm/cublaslt_gemm.cu | 22 +++++++++---------- .../common/include/transformer_engine/gemm.h | 9 ++++++++ transformer_engine/jax/cpp_extensions/gemm.py | 8 +++++++ .../jax/csrc/extensions/gemm.cpp | 1 + 6 files changed, 38 insertions(+), 12 deletions(-) diff --git a/transformer_engine/common/gemm/config.cpp b/transformer_engine/common/gemm/config.cpp index de533909f6..6e470f1b4f 100644 --- a/transformer_engine/common/gemm/config.cpp +++ b/transformer_engine/common/gemm/config.cpp @@ -67,6 +67,9 @@ void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA case kNVTEMatmulConfigSMCount: *reinterpret_cast(buf) = static_cast(config_.sm_count); break; + case kNVTEMatmulConfigAlphaBetaOnDevice: + bool_to_uint8(config_.alpha_beta_on_device, buf); + break; default: NVTE_ERROR("Unsupported NVTEMatmulConfigAttribute (got ", static_cast(attr), ")"); } @@ -116,6 +119,9 @@ void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA case kNVTEMatmulConfigSMCount: config_.sm_count = static_cast(*reinterpret_cast(buf)); break; + case kNVTEMatmulConfigAlphaBetaOnDevice: + uint8_to_bool(buf, config_.alpha_beta_on_device); + break; default: NVTE_ERROR("Unsupported NVTEMatmulConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/common/gemm/config.h b/transformer_engine/common/gemm/config.h index eed47e23d9..eaefcf6c54 100644 --- a/transformer_engine/common/gemm/config.h +++ b/transformer_engine/common/gemm/config.h @@ -22,6 +22,7 @@ struct MatmulConfig { NVTETensor epilogue_aux_tensor = nullptr; bool use_split_accumulator = false; int sm_count = 0; + bool alpha_beta_on_device = false; static constexpr size_t attr_sizes[] = { sizeof(NVTETensor), // bias_tensor @@ -30,7 +31,8 @@ struct MatmulConfig { sizeof(uint8_t), // with_dgelu_epilogue sizeof(NVTETensor), // epilogue_aux_tensor sizeof(uint8_t), // use_split_accumulator - sizeof(int32_t) // sm_count + sizeof(int32_t), // sm_count + sizeof(uint8_t) // alpha_beta_on_device }; }; diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index d491cab9d6..159709eead 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -329,9 +329,9 @@ using cublasHandleManager = detail::HandleManagerrow_scaled_nvfp4 && !inputB->row_scaled_nvfp4, "cuBLAS GEMM does not support row-scaled NVFP4 inputs."); @@ -372,12 +372,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, const bool use_fp4 = is_fp4_dtype(param.Atype) || is_fp4_dtype(param.Btype); const bool nvfp4_tensor_scaling = is_nvfp_scaling(inputA->scaling_mode) && is_nvfp_scaling(inputB->scaling_mode); - const bool have_nvfp4_amax = - (transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr || - (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr; // Update scaling factors with NVFP4 tensor scales - if (use_fp4 && nvfp4_tensor_scaling && have_nvfp4_amax) { + if (use_fp4 && nvfp4_tensor_scaling && !alpha_beta_on_device) { // Reserve some workspace for alpha scale NVTE_CHECK(workspaceSize >= 4, "NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has ", @@ -850,7 +847,8 @@ void nvte_cublas_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, cons // Launch GEMM cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N, (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0], - &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream); + &alpha, &beta, false, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, + stream); } void nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTETensor A, @@ -911,7 +909,8 @@ void nvte_cublas_gemm_v2(int transa, int transb, const float *alpha, const NVTET cublas_gemm(A_tensor, B_tensor, D_tensor, epilogue_bias_tensor, epilogue_aux_tensor, transa ? CUBLAS_OP_T : CUBLAS_OP_N, transb ? CUBLAS_OP_T : CUBLAS_OP_N, with_grad_epilogue, workspace_ptr, workspace_size, alpha, beta, - config_.use_split_accumulator, config_.sm_count, 0, 0, false, nullptr, stream); + config_.alpha_beta_on_device, config_.use_split_accumulator, config_.sm_count, 0, 0, + false, nullptr, stream); } void nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor D, @@ -938,7 +937,8 @@ void nvte_cublas_gemm_scaled(const NVTETensor A, const NVTETensor B, NVTETensor // Launch GEMM cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N, (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0], - &alpha, &beta, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, stream); + &alpha, &beta, false, use_split_accumulator, math_sm_count, 0, 0, false, nullptr, + stream); } void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor D, @@ -984,7 +984,7 @@ void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor "Atomic GEMM only supports delayed scaling."); cublas_gemm(inputA, inputB, outputD, biasTensor, outputGelu, (transa) ? CUBLAS_OP_T : CUBLAS_OP_N, (transb) ? CUBLAS_OP_T : CUBLAS_OP_N, grad, wspace->data.dptr, wspace->data.shape[0], - alpha_ptr, beta_ptr, use_split_accumulator, math_sm_count, m_split, n_split, + alpha_ptr, beta_ptr, true, use_split_accumulator, math_sm_count, m_split, n_split, gemm_producer, inputCounter, stream); #endif } diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index a99e0946ef..fcd5dee1f1 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -54,6 +54,8 @@ enum NVTEMatmulConfigAttribute { kNVTEMatmulConfigUseSplitAccumulator = 5, /*! Number of streaming multiprocessors to use in GEMM kernel. */ kNVTEMatmulConfigSMCount = 6, + /*! Whether alpha and beta are device pointers. Default: false. */ + kNVTEMatmulConfigAlphaBetaOnDevice = 7, kNVTEMatmulConfigNumAttributes }; @@ -565,6 +567,13 @@ class MatmulConfigWrapper { nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigSMCount, &val, sizeof(val)); } + /*! \brief Set whether alpha and beta are device pointers. */ + void set_alpha_beta_on_device(bool alpha_beta_on_device) { + const auto val = static_cast(alpha_beta_on_device); + nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigAlphaBetaOnDevice, &val, + sizeof(val)); + } + private: /*! \brief Wrapped NVTEMatmulConfig. */ NVTEMatmulConfig config_ = nullptr; diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 68c72d5059..0574fe051e 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -208,6 +208,14 @@ def has_rht_applied(q: AbstractBaseTensor) -> bool: def _get_nvfp4_tensor_scale_inv(amax): + # NOTE: SCALE_DTYPE_MAX is hardcoded to the E4M3 block-scale bound (448). JAX sets + # kNVTEMatmulConfigAlphaBetaOnDevice for NVFP4, so the C++ GEMM skips + # nvte_nvfp4_compute_per_tensor_scale entirely and consumes this alpha as-is -- meaning + # this is the only place the block-scale bound is applied. The C++ equivalent, + # Tensor::get_nvfp4_scale_max() in common/common.h, instead deduces the bound from the + # scale-inverse dtype and returns 114688 for kFloat8UE5M3. So once UE5M3 block scales are + # wired into the JAX path, this must deduce SCALE_DTYPE_MAX from the scale dtype to match; + # leaving it at 448 would scale alpha wrong by 256x per operand. DATA_DTYPE_MAX = jnp.finfo(jnp.float4_e2m1fn.dtype).max.astype(jnp.float32) SCALE_DTYPE_MAX = jnp.finfo(jnp.float8_e4m3fn.dtype).max.astype(jnp.float32) return amax / (DATA_DTYPE_MAX * SCALE_DTYPE_MAX) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index a36bf6cd22..34e77539e9 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -325,6 +325,7 @@ Error_Type GemmV2FFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale transformer_engine::MatmulConfigWrapper matmul_config; matmul_config.set_use_split_accumulator(config.use_split_accumulator); matmul_config.set_sm_count(num_math_sm); + matmul_config.set_alpha_beta_on_device(is_nvfp4_scaling(config.scaling_mode)); if (fuse_bias) matmul_config.set_bias_tensor(bias_.data()); if (config.collective_op == JAXX_Collective_Op::NONE) { From 9ca3313a26d1966af1ff0aa566ae4a4d5a6e5f98 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:01:08 -0700 Subject: [PATCH 51/54] Document bugs with `alpha_beta_on_device` GEMM config Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/common/gemm/cublaslt_gemm.cu | 1 + .../common/include/transformer_engine/gemm.h | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 159709eead..d7170e4bc9 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -374,6 +374,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, is_nvfp_scaling(inputA->scaling_mode) && is_nvfp_scaling(inputB->scaling_mode); // Update scaling factors with NVFP4 tensor scales + // TODO Fix bug where amax is ignored when alpha and beta are on device if (use_fp4 && nvfp4_tensor_scaling && !alpha_beta_on_device) { // Reserve some workspace for alpha scale NVTE_CHECK(workspaceSize >= 4, diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index fcd5dee1f1..fee00d08dd 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -54,7 +54,13 @@ enum NVTEMatmulConfigAttribute { kNVTEMatmulConfigUseSplitAccumulator = 5, /*! Number of streaming multiprocessors to use in GEMM kernel. */ kNVTEMatmulConfigSMCount = 6, - /*! Whether alpha and beta are device pointers. Default: false. */ + /*! Whether alpha and beta are device pointers. Default: false. + * + * Known bugs: only supported with NVFP4, NVFP4 amaxes are ignored when alpha and beta are device pointers. + * + * \todo Generalize to all tensor formats + * \todo Correctly handle NVFP4 amaxes when alpha and beta are device pointers + */ kNVTEMatmulConfigAlphaBetaOnDevice = 7, kNVTEMatmulConfigNumAttributes }; From 608ce6be49d5a596fed62c683519f5b8f5314c2b Mon Sep 17 00:00:00 2001 From: Michal Futrega Date: Wed, 2 Sep 2026 23:01:27 +0200 Subject: [PATCH 52/54] [PyTorch] Do not swizzle operand scales in place in the NVFP4-UE5M3 GEMM (#5) * Do not swizzle operand scales in place in the NVFP4-UE5M3 GEMM The cuDNN UE5M3 GEMM wrapper swizzled A/B scales in place and marked the tensors as swizzled. For weights that mark persisted on the parameter. The master-weight cast (_cast_master_weights_to_nvfp4_2d, used by Megatron's fp4 param gather after every optimizer step) then rewrote the scales in the unswizzled layout without clearing the mark, so every later GEMM and dequantize read unswizzled scales as swizzled. Gaussian weights hide it (near-uniform scales, ~0.2 relative error); real DeepSeek-V3 weights are destroyed (2.6 relative error, loss 12.98 vs 7.86 for E4M3 under identical routing). Swizzle clones of the scale tensors instead and leave the operands untouched, matching what the cuBLAS path does in C++. * Test that the NVFP4-UE5M3 GEMM leaves operand scales untouched * Avoid suppressing errors in swizzle function. Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 37 ++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 48 +++++++++++++++---- 2 files changed, 75 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 0bd5c259d1..3ecd452db5 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -799,3 +799,40 @@ def test_nvfp4_ue5m3_gemm_versus_reference( _check_ue5m3_gemm_versus_dequantized( M, K, N, x_columnwise, w_columnwise, disable_second_level_scale ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3) +@pytest.mark.parametrize("M, K, N", [(256, 256, 256), (1024, 3072, 992)]) +def test_nvfp4_ue5m3_gemm_leaves_operands_unswizzled(M: int, K: int, N: int): + """The UE5M3 GEMM must not swizzle operand scales in place. + + Weights are persistent: the master-weight cast rewrites their scales in + the unswizzled layout, so a swizzled mark left behind by a GEMM makes + every later GEMM and dequantize read the scales in the wrong layout. + """ + torch.manual_seed(0) + device, dtype = "cuda", torch.bfloat16 + x = torch.randn((M, K), dtype=dtype, device=device) + w = torch.randn((N, K), dtype=dtype, device=device) + quantizer = NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + scale_dtype=tex.DType.kFloat8UE5M3, + rowwise=True, + columnwise=True, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=True, + ) + x_q = quantizer.update_quantized(x, quantizer.make_empty(x.shape, dtype=dtype, device=device)) + w_q = quantizer.update_quantized(w, quantizer.make_empty(w.shape, dtype=dtype, device=device)) + before = [ + (t._rowwise_scale_inv.clone(), t._columnwise_scale_inv.clone()) for t in (w_q, x_q) + ] + y1 = general_gemm(w_q, x_q, out_dtype=dtype, layout="TN")[0] + y2 = general_gemm(w_q, x_q, out_dtype=dtype, layout="TN")[0] + for t, (rowwise, columnwise) in zip((w_q, x_q), before): + assert not t._with_gemm_swizzled_scales, "GEMM marked its operand as swizzled" + assert torch.equal(t._rowwise_scale_inv, rowwise), "GEMM changed rowwise scales" + assert torch.equal(t._columnwise_scale_inv, columnwise), "GEMM changed columnwise scales" + assert torch.equal(y1, y2), "repeated GEMM on the same operands differs" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index be2a0981e9..1da9722b29 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -440,6 +440,36 @@ def _convert_to_cudnn_grouped_gemm_tensor_format( return data, scale_inv +def _gemm_swizzled_scales( + tensor: NVFP4TensorStorage, +) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Return GEMM-swizzled (rowwise, columnwise) scales without mutating the tensor. + + ``tex.swizzle_scales_for_gemm_`` swizzles in place and marks the tensor as + swizzled. That mark must not persist on weights: the master-weight cast + (``_cast_master_weights_to_nvfp4_2d``) rewrites their scales in the + unswizzled layout and leaves the mark alone, so the next GEMM would read + unswizzled scales as swizzled. Swizzle clones instead and restore the + original attributes. + + This function is a temporary hack until TE supports NVFP4-UE5M3 + GEMMs natively. + """ + if tensor._with_gemm_swizzled_scales: + return tensor._rowwise_scale_inv, tensor._columnwise_scale_inv + compact_rowwise = tensor._rowwise_scale_inv + compact_columnwise = tensor._columnwise_scale_inv + swizzled_rowwise = None if compact_rowwise is None else compact_rowwise.clone() + swizzled_columnwise = None if compact_columnwise is None else compact_columnwise.clone() + tensor._rowwise_scale_inv = swizzled_rowwise + tensor._columnwise_scale_inv = swizzled_columnwise + tex.swizzle_scales_for_gemm_(tensor) + tensor._rowwise_scale_inv = compact_rowwise + tensor._columnwise_scale_inv = compact_columnwise + tensor._with_gemm_swizzled_scales = False + return swizzled_rowwise, swizzled_columnwise + + def _cudnn_grouped_gemm_nvfp4_ue5m3( A: torch.Tensor, B: torch.Tensor, @@ -520,12 +550,10 @@ def _cudnn_grouped_gemm_nvfp4_ue5m3( # cuDNN only accepts GEMM-swizzled scale factors -- an unswizzled buffer is # rejected on its strides -- so swizzle first if the quantizer did not # (optimize_for_gemm defaults to False). This mirrors what the cuBLAS path - # does in C++ via swizzle_scales_for_gemm. The call is in-place, swizzles - # both orientations, and no-ops when the tensor is already swizzled. - if not A._with_gemm_swizzled_scales: - tex.swizzle_scales_for_gemm_(A) - if not B._with_gemm_swizzled_scales: - tex.swizzle_scales_for_gemm_(B) + # does in C++ via swizzle_scales_for_gemm, and like that path it leaves the + # operands untouched. + sfA_rowwise, sfA_columnwise = _gemm_swizzled_scales(A) + sfB_rowwise, sfB_columnwise = _gemm_swizzled_scales(B) # `grad` only changes behaviour when a bias is supplied: it turns the bias slot # into a bias-gradient output, which cuDNN has no epilogue for. Backward GEMMs @@ -538,14 +566,14 @@ def _cudnn_grouped_gemm_nvfp4_ue5m3( # buffer is physically (rows, K_packed), so the reshape below is uniform. # LHS is always (M, K) if transb: - dataB, sfB, amaxB = B._columnwise_data, B._columnwise_scale_inv, B._amax_columnwise + dataB, sfB, amaxB = B._columnwise_data, sfB_columnwise, B._amax_columnwise else: - dataB, sfB, amaxB = B._rowwise_data, B._rowwise_scale_inv, B._amax_rowwise + dataB, sfB, amaxB = B._rowwise_data, sfB_rowwise, B._amax_rowwise # RHS is always (K, N) if transa: - dataA, sfA, amaxA = A._rowwise_data, A._rowwise_scale_inv, A._amax_rowwise + dataA, sfA, amaxA = A._rowwise_data, sfA_rowwise, A._amax_rowwise else: - dataA, sfA, amaxA = A._columnwise_data, A._columnwise_scale_inv, A._amax_columnwise + dataA, sfA, amaxA = A._columnwise_data, sfA_columnwise, A._amax_columnwise # Input tensor dims A_shape = list(dataA.size()) From 7a6c88be4cff532a3656e5b0e1325d86c729106b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:02:15 +0000 Subject: [PATCH 53/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/include/transformer_engine/gemm.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index fee00d08dd..6e0d2d44a4 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -57,7 +57,7 @@ enum NVTEMatmulConfigAttribute { /*! Whether alpha and beta are device pointers. Default: false. * * Known bugs: only supported with NVFP4, NVFP4 amaxes are ignored when alpha and beta are device pointers. - * + * * \todo Generalize to all tensor formats * \todo Correctly handle NVFP4 amaxes when alpha and beta are device pointers */ From e4fa0f502f9630af7541c12d939209b88f8a21a4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:51:12 +0000 Subject: [PATCH 54/54] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 3ecd452db5..d3494b947d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -826,9 +826,7 @@ def test_nvfp4_ue5m3_gemm_leaves_operands_unswizzled(M: int, K: int, N: int): ) x_q = quantizer.update_quantized(x, quantizer.make_empty(x.shape, dtype=dtype, device=device)) w_q = quantizer.update_quantized(w, quantizer.make_empty(w.shape, dtype=dtype, device=device)) - before = [ - (t._rowwise_scale_inv.clone(), t._columnwise_scale_inv.clone()) for t in (w_q, x_q) - ] + before = [(t._rowwise_scale_inv.clone(), t._columnwise_scale_inv.clone()) for t in (w_q, x_q)] y1 = general_gemm(w_q, x_q, out_dtype=dtype, layout="TN")[0] y2 = general_gemm(w_q, x_q, out_dtype=dtype, layout="TN")[0] for t, (rowwise, columnwise) in zip((w_q, x_q), before):