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..cc8ded1cfb 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,27 @@ 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 +879,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 +891,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 +974,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 +989,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 +1587,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 +1908,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 +1961,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 +2152,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 +2160,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 +3602,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 +3611,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..353dbf8f60 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,35 @@ 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 +189,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..f10b165ad2 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -104,13 +104,17 @@ 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 +125,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 +288,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 +309,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 +452,13 @@ 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..b89dd755a9 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -31,58 +31,131 @@ 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..a014244b9b 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,14 @@ __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 +87,29 @@ __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 +121,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 +135,24 @@ 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..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; @@ -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,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; @@ -539,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; @@ -549,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; @@ -590,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); } @@ -599,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); } @@ -614,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(); @@ -626,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); @@ -642,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, @@ -657,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; @@ -683,6 +719,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 +728,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 +735,28 @@ 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< + use_2d_quantization, Cfg, ScaleType, SCALE_TYPE_MAX, IType>( + input, noop, output, stream);); + });); + }) NVTE_CHECK_CUDA(cudaGetLastError()); #else @@ -722,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 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..44fd2d2d12 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..451155e1f4 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..5ed8af0e1e 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,20 @@ __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 +1562,17 @@ 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 +1585,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 +1645,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 +1794,9 @@ 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 +1877,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..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 @@ -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..9ce6d7b044 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -23,18 +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) */ - 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 @@ -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..b8a703fbaa 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -28,26 +28,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 +264,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 +315,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 +373,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 +460,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 +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 == 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..375d21bdfe 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, - "_scale_inverse has invalid dtype " - "(expected DType::kFloat8E4M3, 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, "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, - "_scale_inverse has invalid dtype " - "(expected Float8E4M3, 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, "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), ")"); @@ -984,9 +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: - std::memcpy(buf, &t->nvfp4_e4m3_max, 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/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..466114b7f2 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,43 +1279,51 @@ 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, - 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, // 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..710c228648 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -8,6 +8,9 @@ #include #include +#include +#include + #include "common.h" #include "pybind.h" @@ -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,36 @@ 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, - getTensorShape(scale_inv)); - ret.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, - getTensorShape(amax_columnwise)); + 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(); + 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 +208,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,18 +268,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), - getTensorShape(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)); } // Shape metadata 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..9a4493e3df 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,23 @@ 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 +404,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 +1720,14 @@ 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,21 @@ 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 +1761,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..39cc6588a6 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,41 @@ 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 +250,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 +265,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 +472,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 +489,7 @@ def __new__( amax_rowwise, amax_columnwise, fp4_dtype, + scale_dtype, quantizer, with_gemm_swizzled_scales, *args, @@ -633,6 +657,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 +684,7 @@ def fsdp_post_all_gather( """ ( fp4_dtype, + scale_dtype, columnwise_usage, amax_rowwise, amax_columnwise, @@ -698,6 +724,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 +848,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 +861,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 +877,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 +915,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 +1066,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 +1082,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 +1176,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 +1223,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 +1312,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 +1359,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)