diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 38cef08ce..993d1cea9 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -66,6 +66,7 @@ run_test_config(){ run_default_fa 1 test_float8_current_scaling_exact.py run_default_fa 1 test_float8blockwisetensor.py run_default_fa 1 test_float8_blockwise_scaling_exact.py + run_default_fa 1 test_float8_blockwise_gemm_exact.py run_default_fa 1 test_quantized_tensor.py test $_fus_attn = auto -o $_fus_attn = ck && run 1 test_cpu_offloading.py test $_fus_attn = auto -o $_fus_attn = ck -o $_fus_attn = aotriton && NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 run 3 test_cpu_offloading_v1.py @@ -75,6 +76,7 @@ run_test_config(){ run_default_fa 1 test_gemm_autotune.py run 1 test_gqa.py run 1 test_grouped_linear.py + NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 test_grouped_tensor.py run 1 test_jit.py NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 test_multi_tensor.py run 1 test_numerics.py diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index a13d4f37d..ecaaa83e9 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -46,6 +46,7 @@ if(USE_ROCM) list(REMOVE_ITEM test_cuda_sources test_grouped_gemm.cu) list(APPEND test_cuda_sources + test_multi_tensor_transpose_bhsd.cu test_dequantize_nvfp4.cu test_cublaslt_gemm.cu test_cast_mxfp4_transpose.cu diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 27eb7ad51..40c40d86f 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include #include @@ -188,6 +190,146 @@ NVFP4FourOverSixQuantization quantize_4over6_pair( }; } +// CUDA keeps upstream's looser rule; its fast-math configs score in packed fp16. +#ifdef __HIP_PLATFORM_AMD__ +constexpr bool kAcceptEitherCandidate = false; +#else +constexpr bool kAcceptEitherCandidate = true; +#endif + +struct NVFP4FourOverSixDecision { + NVFP4FourOverSixCandidate candidate = NVFP4FourOverSixCandidate::Map6; + // Candidate errors too close for FP32 summation order to resolve. + bool ambiguous = false; +}; + +// Mirrors fp4_roundtrip_err() in quantize_transpose_vector_blockwise_fp4.cu, +// operand order included. +float fp4_roundtrip_err_ref(const float x, const float block_scale_inverse, const float sf, + const float global_amax, const float err_denom, const bool use_mse) { + const float2 scaled = {x * block_scale_inverse, 0.0f}; + const fp4e2m1x2 quantized(scaled); + const float dequant = static_cast(cvt_fp4x2_to_double2(quantized).x); + const float val = dequant * sf * global_amax / err_denom; + const float diff = val - x; + return use_mse ? diff * diff : std::fabs(diff); +} + +// Lower round-trip error wins; ties go to map-to-6. +NVFP4FourOverSixDecision decide_4over6_candidate(const std::vector& block, + const float block_amax, + const float S_enc, + const int e4m3_max, + const bool use_mse) { + const NVFP4FourOverSixQuantization quantization = + compute_4over6_quantization_scales(block_amax, S_enc); + const float err_denom = 6.0f * static_cast(e4m3_max); + const float global_amax = err_denom / S_enc; + + double err_map4 = 0.0; + double err_map6 = 0.0; + for (const float x : block) { + err_map4 += fp4_roundtrip_err_ref(x, quantization.reciprocal_map4, + static_cast(quantization.scale_map4), global_amax, + err_denom, use_mse); + err_map6 += fp4_roundtrip_err_ref(x, quantization.reciprocal_map6, + static_cast(quantization.scale_map6), global_amax, + err_denom, use_mse); + } + + // Worst-case FP32 summation drift of the kernel's warp reduction. + const double tie_band = 4.0 * static_cast(block.size()) * + static_cast(FLT_EPSILON) * std::max(err_map4, err_map6); + + NVFP4FourOverSixDecision decision; + decision.ambiguous = std::fabs(err_map4 - err_map6) <= tie_band; + decision.candidate = + err_map4 < err_map6 ? NVFP4FourOverSixCandidate::Map4 : NVFP4FourOverSixCandidate::Map6; + return decision; +} + +// Laid out like the scale tensor so the checker can index it by scale_idx. +template +std::vector compute_4over6_expected_decisions( + float (*OP)(const float), + const InputType* const input, + const size_t rows, + const size_t cols, + const size_t scales_stride, + const float* const amax, + const bool use_fast_math, + const bool use_2d_quantization, + const bool row_scaled_nvfp4, + const bool use_mse, + const int e4m3_max) { + + constexpr size_t block_size_Y = 16; + constexpr size_t block_size_X = 16; + const size_t blocks_X = divide_round_up(cols, block_size_X); + std::vector decisions(rows * scales_stride); + + // Same numerical truncation the reference quantizers apply. + auto activated = [OP, input, cols](const size_t i, const size_t j) { + const float act_elt = OP(static_cast(input[i * cols + j])); + return static_cast(static_cast(act_elt)); + }; + + if (use_2d_quantization) { + const float S_enc = compute_global_encode_scaling_factor_FP4(*amax, use_fast_math, + e4m3_max); + const size_t blocks_Y = divide_round_up(rows, block_size_Y); + for (size_t block_Y = 0; block_Y < blocks_Y; ++block_Y) { + for (size_t block_X = 0; block_X < blocks_X; ++block_X) { + const size_t i_min = block_Y * block_size_Y; + const size_t i_max = std::min(i_min + block_size_Y, rows); + const size_t j_min = block_X * block_size_X; + const size_t j_max = std::min(j_min + block_size_X, cols); + + std::vector block; + block.reserve(block_size_Y * block_size_X); + float block_amax = 0.0f; + for (size_t i = i_min; i < i_max; ++i) { + for (size_t j = j_min; j < j_max; ++j) { + const float elt = activated(i, j); + block.push_back(elt); + block_amax = std::max(block_amax, std::abs(elt)); + } + } + + const NVFP4FourOverSixDecision decision = + decide_4over6_candidate(block, block_amax, S_enc, e4m3_max, use_mse); + // Block scale and its candidate replicate down the block's rows. + for (size_t i = i_min; i < i_max; ++i) { + decisions[i * scales_stride + block_X] = decision; + } + } + } + return decisions; + } + + for (size_t i = 0; i < rows; ++i) { + const float S_enc = compute_global_encode_scaling_factor_FP4( + row_scaled_nvfp4 ? amax[i] : *amax, use_fast_math, e4m3_max); + for (size_t block_X = 0; block_X < blocks_X; ++block_X) { + const size_t j_min = block_X * block_size_X; + const size_t j_max = std::min(j_min + block_size_X, cols); + + std::vector block; + block.reserve(block_size_X); + float block_amax = 0.0f; + for (size_t j = j_min; j < j_max; ++j) { + const float elt = activated(i, j); + block.push_back(elt); + block_amax = std::max(block_amax, std::abs(elt)); + } + + decisions[i * scales_stride + block_X] = + decide_4over6_candidate(block, block_amax, S_enc, e4m3_max, use_mse); + } + } + return decisions; +} + // 1D Scaling: Original implementation with 1x16 blocks template void quantize_nvfp4_1d(float (*OP)(const float), @@ -735,6 +877,7 @@ void compare_nvfp4_4over6_candidates(const std::string& name, const fp8e4m3* const ref_scales_map4, const fp4e2m1x2* const ref_data_map6, const fp8e4m3* const ref_scales_map6, + const std::vector& decisions, const size_t rows, const size_t cols, const size_t blocks_X, @@ -742,6 +885,13 @@ void compare_nvfp4_4over6_candidates(const std::string& name, constexpr int max_mismatches_to_print = 3; const auto* const test_data_pairs = reinterpret_cast(test_data); size_t total_mismatches = 0; + size_t wrong_candidate = 0; + size_t expected_map4 = 0; + size_t expected_map6 = 0; + // Both candidates encode identically; no selection is observable. + size_t identical_candidates = 0; + // Candidate errors too close for the host to call. + size_t unpinned_blocks = 0; for (size_t row = 0; row < rows; ++row) { for (size_t block_x = 0; block_x < blocks_X; ++block_x) { @@ -754,34 +904,72 @@ void compare_nvfp4_4over6_candidates(const std::string& name, bitwise_equal(test_scales[scale_idx], ref_scales_map6[scale_idx]); const bool data_matches_map6 = nvfp4_output_block_matches(test_data_pairs, ref_data_map6, row, cols, block_x); + const bool matches_map4 = scale_matches_map4 && data_matches_map4; + const bool matches_map6 = scale_matches_map6 && data_matches_map6; + + const bool candidates_differ = + !bitwise_equal(ref_scales_map4[scale_idx], ref_scales_map6[scale_idx]) || + !nvfp4_output_block_matches(ref_data_map4, ref_data_map6, row, cols, block_x); + + const NVFP4FourOverSixDecision& decision = decisions[scale_idx]; + bool matched = false; + const char* expectation = nullptr; + if (decision.ambiguous || kAcceptEitherCandidate) { + // Too close to call from the host; either encoding is a correct answer. + if (candidates_differ) { + ++unpinned_blocks; + } else { + ++identical_candidates; + } + matched = matches_map4 || matches_map6; + expectation = "map-to-4 or map-to-6 (candidate errors within the tie band)"; + } else if (decision.candidate == NVFP4FourOverSixCandidate::Map4) { + ++expected_map4; + matched = matches_map4; + expectation = "map-to-4"; + } else { + ++expected_map6; + matched = matches_map6; + expectation = "map-to-6"; + } - if ((scale_matches_map4 && data_matches_map4) || - (scale_matches_map6 && data_matches_map6)) { + if (matched) { continue; } ++total_mismatches; + const bool matched_other = matches_map4 || matches_map6; + if (matched_other) { + ++wrong_candidate; + } if (total_mismatches <= max_mismatches_to_print) { std::cout << "Error in tensor " << name << ": 4over6 block mismatch at row " - << row << ", block_x " << block_x - << ". The output did not match either map-to-4 or map-to-6 exactly." - << std::endl; + << row << ", block_x " << block_x << ". Expected " << expectation + << "; the output " + << (matched_other ? "matched the other candidate instead" + : "matched neither candidate exactly") + << "." << std::endl; } } } std::cout << "=== SUMMARY for tensor " << name << " ===" << std::endl; std::cout << "Total 4over6 blocks checked: " << (rows * blocks_X) << std::endl; + std::cout << "Blocks pinned to map-to-4: " << expected_map4 + << ", pinned to map-to-6: " << expected_map6 + << ", left unpinned by the tie band: " << unpinned_blocks + << ", encoding identically either way: " << identical_candidates << std::endl; if (total_mismatches > 0) { std::cout << "STATUS: FAILED for output" << std::endl; - std::cout << "Total mismatched 4over6 blocks found: " << total_mismatches << std::endl; + std::cout << "Total mismatched 4over6 blocks found: " << total_mismatches + << " (" << wrong_candidate << " picked the wrong candidate)" << std::endl; std::cout << "============================" << std::endl; GTEST_FAIL() << "Found " << total_mismatches << " 4over6 block mismatches in tensor " - << name; + << name << " (" << wrong_candidate << " picked the wrong candidate)"; } std::cout << "STATUS: PASSED for output" << std::endl; - std::cout << "Each 4over6 block matched either map-to-4 or map-to-6 exactly" << std::endl; + std::cout << "Each 4over6 block matched the candidate the reference selected" << std::endl; std::cout << "============================" << std::endl; } @@ -823,7 +1011,8 @@ void performTest(float (*OP)(const float), #ifdef __HIP_PLATFORM_AMD__ if (te_fp8_fnuz()) GTEST_SKIP() << "NVFP4 not supported on gfx942 (fnuz)"; - if (use_4over6) GTEST_SKIP() << "NVFP4 4over6 not supported on ROCm"; + if (use_4over6 && use_4over6_err_use_fast_math) + GTEST_SKIP() << "NVFP4 4over6 fast-math error mode is not supported on ROCm"; #endif const size_t rows = first_dimension(shape); @@ -863,6 +1052,8 @@ void performTest(float (*OP)(const float), std::unique_ptr ref_output_t_map6; std::unique_ptr ref_scales_map6; std::unique_ptr ref_scales_t_map6; + std::vector expected_decisions; + std::vector expected_decisions_t; fillCase(&input, InputsFillCase::uniform); @@ -974,6 +1165,19 @@ void performTest(float (*OP)(const float), use_4over6, e4m3_max, NVFP4FourOverSixCandidate::Map6); + + // Reference choice: the candidate with the smaller round-trip error. + const bool use_mse = mode == kNVTENVFP44Over6MinMSE; + expected_decisions = compute_4over6_expected_decisions( + OP, input.rowwise_cpu_dptr(), rows, cols, scales_stride, ref_amax.data(), + use_fast_math, is_2d_quantization, row_scaled_nvfp4, use_mse, e4m3_max); + if (!row_scaled_nvfp4) { + const std::vector input_t = + create_transpose(input.rowwise_cpu_dptr(), rows, cols); + expected_decisions_t = compute_4over6_expected_decisions( + OP, input_t.data(), cols, rows, scales_stride_t, ref_amax.data(), use_fast_math, + is_2d_quantization, row_scaled_nvfp4, use_mse, e4m3_max); + } } else { compute_ref(OP, input.rowwise_cpu_dptr(), @@ -997,8 +1201,10 @@ void performTest(float (*OP)(const float), hipDeviceProp_t prop; hipGetDeviceProperties(&prop, 0); const bool is_gfx950 = prop.major == 9 && prop.minor == 5; - for (bool use_stochastic_rounding : (is_gfx950 ? std::vector{false, true} - : std::vector{false})) { + // The quantize dispatch refuses 4over6 combined with stochastic rounding. + for (bool use_stochastic_rounding : (is_gfx950 && !use_4over6 + ? std::vector{false, true} + : std::vector{false})) { #endif // Initialize stochastic rounding @@ -1056,6 +1262,7 @@ void performTest(float (*OP)(const float), ref_scales.get(), ref_output_map6.get(), ref_scales_map6.get(), + expected_decisions, rows, cols, unpadded_blocks_X, @@ -1068,6 +1275,7 @@ void performTest(float (*OP)(const float), ref_scales_t.get(), ref_output_t_map6.get(), ref_scales_t_map6.get(), + expected_decisions_t, cols, rows, unpadded_blocks_X_t, diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu index 4a18bb589..bfec25328 100644 --- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu @@ -1,4 +1,6 @@ /************************************************************************* + * This file was modified for portability to AMDGPU + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -229,6 +231,40 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors, &offsets_tensor, sizeof(offsets_tensor)); } + bool bad_last_dim = false; + for (size_t t = 0; t < num_tensors; ++t) { + if (last_dims_h[t] % 16 != 0) bad_last_dim = true; + } +#ifdef __HIP_PLATFORM_AMD__ + auto free_device = [&]() { + cudaFree(in_data_d); + cudaFree(out_grouped_d); + cudaFree(in_scales_d); + cudaFree(first_dims_d); + cudaFree(last_dims_d); + cudaFree(offsets_d); + }; + + if (bad_last_dim) { + EXPECT_THROW(nvte_group_dequantize(in_group_tensor, out_group_tensor, 0), std::runtime_error); + free_device(); + return; + } + + const bool uniform_last_dim = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + bool ragged_colwise = false; + if (!rowwise && uniform_last_dim && num_tensors > 1) { + for (size_t t = 0; t < num_tensors; ++t) { + if (first_dims_h[t] % 32 != 0) ragged_colwise = true; + } + } + if (ragged_colwise) { + EXPECT_THROW(nvte_group_dequantize(in_group_tensor, out_group_tensor, 0), std::runtime_error); + free_device(); + return; + } +#endif + // Run grouped dequantize nvte_group_dequantize(in_group_tensor, out_group_tensor, 0); cudaDeviceSynchronize(); @@ -310,15 +346,14 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors, int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset, tensor_elts * sizeof(OutputType)); - if (result != 0) { - // Find first mismatch for error reporting - for (size_t i = 0; i < tensor_elts; ++i) { - if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) { - GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i - << " (global offset " << (data_offset + i) << "): grouped=" - << static_cast(out_grouped_h[data_offset + i]) - << " vs reference=" << static_cast(out_ref_h[data_offset + i]); - } + // memcmp is the assertion; != would let -0.0 vs +0.0 and NaN payloads through. + ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t; + for (size_t i = 0; i < tensor_elts; ++i) { + if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) { + GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i + << " (global offset " << (data_offset + i) << "): grouped=" + << static_cast(out_grouped_h[data_offset + i]) + << " vs reference=" << static_cast(out_ref_h[data_offset + i]); } } } @@ -349,6 +384,18 @@ std::vector> input_configs = { {VARYING_FIRST_DIM, 3, 768, 96, 256, 256, 256}, {VARYING_LAST_DIM, 2, 160, 384, 128, 256}, {VARYING_LAST_DIM, 3, 96, 512, 128, 128, 256}, +#ifdef __HIP_PLATFORM_AMD__ + // Per-tensor rows not scale-block aligned: rejected on ROCm. + {VARYING_FIRST_DIM, 2, 96, 128, 48, 48}, + // Last dim not 16-byte aligned: vectorized loads would shift the tile. + {SAME_BOTH_DIMS, 1, 128, 120}, + // Launcher derives fewer blocks than the tile count, so the grid-stride loop must iterate. + {VARYING_LAST_DIM, 2, 256, 256, 96, 160}, + // Rows are scale-block aligned but not chunk aligned: valid, must not be rejected. + {VARYING_FIRST_DIM, 2, 128, 128, 64, 64}, + // Uniform shape whose per-tensor rows are not scale-block aligned. + {SAME_BOTH_DIMS, 2, 96, 128}, +#endif }; std::vector scaling_directions = { @@ -366,10 +413,12 @@ class GroupedDequantizeMXFP8TestSuite >> {}; TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) { +#ifndef __HIP_PLATFORM_AMD__ // Skip tests for pre-Blackwell architectures if (getDeviceComputeCapability() < blackwellComputeCapability) { GTEST_SKIP(); } +#endif using namespace transformer_engine; using namespace test; @@ -419,6 +468,19 @@ TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) { (shape_rep == VARYING_FIRST_DIM || shape_rep == VARYING_BOTH_DIMS); const bool last_dim_varies = (shape_rep == VARYING_LAST_DIM || shape_rep == VARYING_BOTH_DIMS); +#ifdef __HIP_PLATFORM_AMD__ + // ROCm rejects the shapes below instead of skipping them; see the EXPECT_THROW paths. + const bool uniform_last = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + if (first_dim_varies && shape_rep != VARYING_FIRST_DIM && (first_dims[t] % 128 != 0)) { + GTEST_SKIP(); + } + if (last_dim_varies && shape_rep != VARYING_LAST_DIM && (last_dims[t] % 128 != 0)) { + GTEST_SKIP(); + } + if (!rowwise && !uniform_last && (first_dims[t] % 32 != 0)) { + GTEST_SKIP(); + } +#else if (first_dim_varies && (first_dims[t] % 128 != 0)) { GTEST_SKIP(); } @@ -429,10 +491,10 @@ TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) { if (last_dims[t] % 16 != 0) { GTEST_SKIP(); } - // For colwise: first dim must be divisible by 32 if (!rowwise && (first_dims[t] % 32 != 0)) { GTEST_SKIP(); } +#endif // For rowwise: last dim must be divisible by 32 if (rowwise && (last_dims[t] % 32 != 0)) { GTEST_SKIP(); diff --git a/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu b/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu new file mode 100644 index 000000000..b859108f7 --- /dev/null +++ b/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu @@ -0,0 +1,145 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +struct PermuteCase { + size_t B, S, H, D_in, D_out; +}; + +// out[b][h][s][d] = in[b][s][h][d] (BSHD) or in[s][b][h][d] (SBHD), zero-padded +// past D_in when D_out is larger. +template +void compute_ref(const T *in, T *out, const PermuteCase &c, bool is_bshd) { + for (size_t b = 0; b < c.B; ++b) { + for (size_t h = 0; h < c.H; ++h) { + for (size_t s = 0; s < c.S; ++s) { + const size_t in_off = + (is_bshd ? ((b * c.S + s) * c.H + h) : ((s * c.B + b) * c.H + h)) * c.D_in; + T *dst = out + ((b * c.H + h) * c.S + s) * c.D_out; + for (size_t d = 0; d < c.D_out; ++d) { + dst[d] = (d < c.D_in) ? in[in_off + d] : static_cast(0); + } + } + } + } +} + +// Shared-memory transpose needs 32*(32*D_pad+4) bytes; query the device budget. +template +bool exceeds_smem_budget(const std::vector &cases) { + size_t d_in_max = 0; + for (const auto &c : cases) d_in_max = std::max(d_in_max, c.D_in); + const size_t d_bytes = d_in_max * sizeof(T); + if (d_bytes % 4 == 0) return false; // vectorized path, no shared memory + const size_t d_pad = (d_bytes + 3u) & ~size_t(3); + const size_t needed = 32 * (32 * d_pad + 4); + int max_smem = 0; + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0)); + return needed > static_cast(max_smem); +} + +// D_in * sizeof(T) % 4 selects vectorized vs shared-memory path; TMA is CUDA-only. +template +void performTest(const std::vector &cases, bool is_bshd, DType dtype) { + if (exceeds_smem_budget(cases)) { + GTEST_SKIP() << "shared-memory transpose exceeds this device's shared memory"; + } + const size_t n = cases.size(); + std::vector> ins, outs; + std::vector in_h, out_h; + std::vector> refs(n); + + for (size_t i = 0; i < n; ++i) { + const auto &c = cases[i]; + const std::vector in_shape = is_bshd ? std::vector{c.B, c.S, c.H, c.D_in} + : std::vector{c.S, c.B, c.H, c.D_in}; + const std::vector out_shape{c.B, c.H, c.S, c.D_out}; + + auto in = std::make_unique("in_" + std::to_string(i), in_shape, dtype); + auto out = std::make_unique("out_" + std::to_string(i), out_shape, dtype); + fillUniform(in.get()); + // Poison output so unwritten positions cannot match the expected zero pad. + NVTE_CHECK_CUDA(cudaMemset(out->rowwise_dptr(), 0xCD, c.B * c.H * c.S * c.D_out * sizeof(T))); + + in->to_cpu(); + refs[i].resize(c.B * c.H * c.S * c.D_out); + compute_ref(in->rowwise_cpu_dptr(), refs[i].data(), c, is_bshd); + + in_h.push_back(in->data()); + out_h.push_back(out->data()); + ins.emplace_back(std::move(in)); + outs.emplace_back(std::move(out)); + } + + nvte_multi_tensor_transpose_to_bhsd( + in_h.data(), out_h.data(), n, + is_bshd ? NVTE_QKV_Format::NVTE_BSHD : NVTE_QKV_Format::NVTE_SBHD, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + NVTE_CHECK_CUDA(cudaGetLastError()); + + // A pure layout permute, so byte equality is the right assertion. + for (size_t i = 0; i < n; ++i) { + outs[i]->to_cpu(); + compareResults("bhsd_" + std::to_string(i), + reinterpret_cast(outs[i]->rowwise_cpu_dptr()), + reinterpret_cast(refs[i].data()), refs[i].size() * sizeof(T)); + } +} + +class MultiTensorTransposeBhsdTestSuite + : public ::testing::TestWithParam, bool>> {}; + +TEST_P(MultiTensorTransposeBhsdTestSuite, TestFp16) { + const auto cases = std::get<0>(GetParam()); + const bool is_bshd = std::get<1>(GetParam()); + performTest(cases, is_bshd, DType::kFloat16); +} + +TEST_P(MultiTensorTransposeBhsdTestSuite, TestByte) { + const auto cases = std::get<0>(GetParam()); + const bool is_bshd = std::get<1>(GetParam()); + performTest(cases, is_bshd, DType::kByte); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, MultiTensorTransposeBhsdTestSuite, + ::testing::Combine(::testing::Values( + // Single tensor, D_in*elem a multiple of 4 -> vectorized fallback. + std::vector{{2, 64, 4, 64, 64}}, + // Odd D_in with 1-byte elements: shared-memory transpose. + std::vector{{2, 48, 3, 33, 33}}, + // D_out > D_in exercises the zero pad on the vectorized path. + std::vector{{1, 32, 2, 40, 64}}, + // Same pad via the shared-memory transpose, within gfx942 budget. + std::vector{{1, 40, 2, 21, 32}}, + // Multiple tensors per launch; B must match across the group. + std::vector{ + {2, 64, 4, 64, 64}, {2, 33, 2, 16, 24}, {2, 16, 1, 8, 8}}, + // S below one tile and H = 1, the degenerate grid. + std::vector{{1, 1, 1, 16, 16}}), + ::testing::Bool()), + [](const testing::TestParamInfo &info) { + return "case" + std::to_string(info.index) + "_n" + + std::to_string(std::get<0>(info.param).size()) + + (std::get<1>(info.param) ? "_bshd" : "_sbhd"); + }); + +} // namespace diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index b86091df4..6144a484f 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -2108,7 +2108,6 @@ def f(x): assert_allclose(actual, expected, dtype=dtype) -@pytest.mark.skipif(is_hip_extension(), reason="Standalone TopK (nvte_topk) is not supported on ROCm") @pytest.mark.parametrize("dtype", [jnp.bfloat16, jnp.float32]) @pytest.mark.parametrize( "problem_size", diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index d7e02c55c..b3ef196bf 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -39,8 +39,6 @@ def check_nvfp4_gemm_versus_reference( ): if nvfp4_e4m3_max != 448 and not use_4over6: pytest.skip("E4M3 max 256 is only meaningful for 4over6") - if IS_HIP_EXTENSION and use_4over6: - pytest.skip("NVFP4 4over6 is not supported on ROCm") te_dtype = te.DType.kFloat4E2M1 # Setup device and random seed @@ -467,7 +465,7 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( ids=["rowxrow", "colxrow", "colxcol"], ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) -@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) @pytest.mark.parametrize("nvfp4_e4m3_max", [448, 256], ids=["e4m3_448", "e4m3_256"]) @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_gemm_versus_reference( @@ -530,7 +528,7 @@ def test_nvfp4_gemm_versus_reference( @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) @pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) -@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) @pytest.mark.skipif(IS_HIP_EXTENSION, reason="Grouped NVFP4 GEMM is not supported on ROCm") def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( @@ -579,7 +577,7 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) -@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_row_scaled_gemm_matches_emulated( M: int, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index d60e48f75..0a01f0811 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -18,7 +18,6 @@ from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.common.recipe import NVFP4BlockScaling - recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -84,8 +83,6 @@ def maybe_skip_row_scaled_unsupported_quantization( M: int | None = None, N: int | None = None, ) -> None: - if IS_HIP_EXTENSION and use_4over6: - pytest.skip("NVFP4 4over6 is not supported on ROCm") if use_4over6 and with_2d_quantization: if x_dtype != torch.bfloat16 or M is None or N is None or M % 32 != 0 or N % 32 != 0: pytest.skip("NVFP4 2D 4over6 exact tests require the optimized BF16 kernel path") @@ -151,6 +148,8 @@ def check_quantization_nvfp4_versus_reference( ) if use_4over6: + if IS_HIP_EXTENSION and nvfp4_4over6_err_use_fast_math: + pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm") with nvfp4_4over6_err_fast_math(nvfp4_4over6_err_use_fast_math): if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -363,6 +362,8 @@ def test_nvfp4_quantization_extrema_versus_reference( ) if nvfp4_4over6_config.use_4over6: + if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math: + pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm") with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math): if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -508,6 +509,8 @@ def test_nvfp4_quantization_boundary_values( ) if nvfp4_4over6_config.use_4over6: + if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math: + pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm") with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math): if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -639,6 +642,8 @@ def test_nvfp4_quantization_noncontiguous_inputs( ) if nvfp4_4over6_config.use_4over6: + if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math: + pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm") with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math): if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x_nc) @@ -811,3 +816,45 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer: # Sanity: column-only path must not allocate a rowwise output. assert out_col_only._rowwise_data is None + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("M, N", [(128, 128), (256, 256), (512, 1024), (320, 256)]) +@pytest.mark.parametrize( + "with_2d_quantization", [True, False], ids=["2d_quantization", "1d_quantization"] +) +def test_nvfp4_4over6_columnwise_only_matches_both_directions( + M: int, + N: int, + with_2d_quantization: bool, +): + """Columnwise-only 4over6 must match the columnwise half of both-directions. + + Without a rowwise pass the candidate errors have to be reduced by the + columnwise-only path itself; if that reduction is skipped, selection reads an + uninitialised buffer and silently falls back to map-to-6. + """ + if with_2d_quantization and not IS_HIP_EXTENSION: + pytest.skip("CUDA routes 4over6 to a kernel that requires rowwise output for 2D") + device = "cuda" + torch.manual_seed(0) + torch.cuda.manual_seed(0) + x = torch.randn((M, N), dtype=torch.bfloat16, device=device) + + def _quantize(*, rowwise: bool): + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=True, + with_2d_quantization=with_2d_quantization, + nvfp4_use_4over6=True, + )(x) + + both = _quantize(rowwise=True) + col_only = _quantize(rowwise=False) + + torch.testing.assert_close(col_only._columnwise_data, both._columnwise_data, atol=0, rtol=0) + torch.testing.assert_close( + col_only._columnwise_scale_inv, both._columnwise_scale_inv, atol=0, rtol=0 + ) + assert col_only._rowwise_data is None diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 83d54ab09..4ca7691a4 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -76,8 +76,7 @@ _quantization_list.append("mxfp8") if nvfp4_available: _quantization_list.append("nvfp4") - if not IS_HIP_EXTENSION: # NVFP4 4over6 is not supported on ROCm - _quantization_list.append("nvfp4_4over6") + _quantization_list.append("nvfp4_4over6") @pytest.fixture(autouse=True, scope="function") diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e646b72c3..0c9b3b506 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -310,8 +310,7 @@ set(cuda_only_cuda_sources fused_attn/fused_attn_f16_arbitrary_seqlen.cu fused_attn/fused_attn_fp8.cu fused_attn/utils.cu - swizzle/swizzle_block_scaling.cu - util/topk.cu) #CUDA-only + swizzle/swizzle_block_scaling.cu) #CUDA-only list(APPEND transformer_engine_cuda_arch_specific_sources fused_attn/flash_attn.cu @@ -443,6 +442,25 @@ else() #USE_ROCM # process source code files include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/hipify/hipify.cmake") TE_Hipify(${CMAKE_CURRENT_SOURCE_DIR}) + + # TE_Hipify runs at configure time, so its inputs must be registered or edits go unnoticed. + # Generated _hip.* are excluded to avoid re-triggering configure on their own output. + set(_hipify_deps) + foreach(_src ${transformer_engine_SOURCES}) + get_filename_component(_abs "${CMAKE_CURRENT_SOURCE_DIR}/${_src}" ABSOLUTE) + list(APPEND _hipify_deps "${_abs}") + endforeach() + file(GLOB_RECURSE _hipify_hdrs + "${CMAKE_CURRENT_SOURCE_DIR}/*.h" + "${CMAKE_CURRENT_SOURCE_DIR}/*.cuh" + "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp") + foreach(_hdr ${_hipify_hdrs}) + if(NOT _hdr MATCHES "_hip\\.(h|cuh|hpp)$") + list(APPEND _hipify_deps "${_hdr}") + endif() + endforeach() + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_hipify_deps}) + TE_GetHipifiedSources("${transformer_engine_SOURCES}" ${CMAKE_CURRENT_SOURCE_DIR} te_hip_sources) message("${message_line}") diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 40ceb59d7..ca4d0e7ad 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -43,13 +43,13 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t case NVTE_MXFP8_1D_SCALING: { #ifndef __HIP_PLATFORM_AMD__ if (is_supported_by_CC_100()) { -#endif //#ifndef __HIP_PLATFORM_AMD__ +#endif //#ifndef __HIP_PLATFORM_AMD__ mxfp8::dequantize(input, output, stream); #ifndef __HIP_PLATFORM_AMD__ } else { NVTE_ERROR("MXFP8 Dequantization is NOT supported by architectures < 10.0"); } -#endif //#ifndef __HIP_PLATFORM_AMD__ +#endif //#ifndef __HIP_PLATFORM_AMD__ break; } case NVTE_NVFP4_1D_SCALING: { @@ -70,13 +70,13 @@ inline void group_dequantize_helper(const GroupedTensor &input, GroupedTensor *o case NVTE_MXFP8_1D_SCALING: { #ifndef __HIP_PLATFORM_AMD__ if (is_supported_by_CC_100()) { +#endif //#ifndef __HIP_PLATFORM_AMD__ mxfp8::group_dequantize(&input, output, stream); +#ifndef __HIP_PLATFORM_AMD__ } else { NVTE_ERROR("MXFP8 Grouped Dequantization is NOT supported by architectures < 10.0"); } -#else - NVTE_ERROR("MXFP8 Grouped Dequantization is not supported on ROCm."); -#endif +#endif //#ifndef __HIP_PLATFORM_AMD__ break; } default: diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 4b31f1d18..4afaf29a5 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -139,6 +139,11 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, "Non-4over6 NVFP4 quantization requires E4M3 max 448."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); +#ifdef __HIP_PLATFORM_AMD__ + // Refuse the fast-math error path rather than silently scoring with the exact one. + NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math, + "NVFP4 4over6 fast-math error mode is not supported on ROCm."); +#endif if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); @@ -154,9 +159,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, (output_tensor->has_data() || (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization)); - // Launch NVFP4 quantize kernel. 4over6 and the optimized quantize_transpose kernels are - // CUDA-only (Blackwell); ROCm falls through to the portable blockwise path below, which - // supports row-scaled NVFP4. + // Launch NVFP4 quantize kernel. ROCm uses the portable blockwise path below. #ifndef __HIP_PLATFORM_AMD__ if (nvfp4_use_4over6) { if (quant_config_cpp.nvfp4_2d_quantization) { @@ -191,6 +194,8 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, /*row_scaled_nvfp4=*/row_scaled_nvfp4, /*noop_tensor=*/noop_tensor->data, + /*nvfp4_e4m3_max=*/output_tensor->nvfp4_e4m3_max, + /*nvfp4_4over6_mode=*/quant_config_cpp.nvfp4_4over6_mode, /*stream=*/stream); #ifndef __HIP_PLATFORM_AMD__ } @@ -318,6 +323,11 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens "Non-4over6 NVFP4 quantization requires E4M3 max 448."); NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, "NVFP4 4over6 quantization does not support stochastic rounding."); +#ifdef __HIP_PLATFORM_AMD__ + // Refuse the fast-math error path rather than silently scoring with the exact one. + NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math, + "NVFP4 4over6 fast-math error mode is not supported on ROCm."); +#endif NVTE_CHECK(!output_tensor->row_scaled_nvfp4, "Backward NVFP4 quantization does not support row-scaled outputs."); // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and @@ -328,9 +338,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens (output_tensor->has_data() || (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization)); - // Launch NVFP4 quantize kernel. 4over6 and the optimized quantize_transpose kernels are - // CUDA-only (Blackwell); ROCm falls through to the portable blockwise path below, which - // supports row-scaled NVFP4. + // Launch NVFP4 quantize kernel. ROCm uses the portable blockwise path below. #ifndef __HIP_PLATFORM_AMD__ if (nvfp4_use_4over6) { if (quant_config_cpp.nvfp4_2d_quantization) { @@ -365,6 +373,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, /*row_scaled_nvfp4=*/false, /*noop_tensor=*/noop_tensor->data, + /*nvfp4_e4m3_max=*/output_tensor->nvfp4_e4m3_max, + /*nvfp4_4over6_mode=*/quant_config_cpp.nvfp4_4over6_mode, /*stream=*/stream); #ifndef __HIP_PLATFORM_AMD__ } diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index fbe1985cb..62a7f61e5 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -26,10 +26,27 @@ #include "../../utils.cuh" #include "group_quantize_mxfp8.cuh" +#else + +#include +#include + +#include + +#include "../../common.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "./rocm_vectorized_2d.cuh" + +#endif // #ifndef __HIP_PLATFORM_AMD__ + namespace transformer_engine { namespace dispatch { namespace mxfp8 { namespace group_dequantize_kernel { +#ifdef __HIP_PLATFORM_AMD__ +#include "rocm_group_dequantize_mxfp8.cuh" +#else constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; __device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; @@ -373,10 +390,145 @@ __global__ void __launch_bounds__(128) destroy_barriers(mbar, is_master_thread); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +#endif // #ifdef __HIP_PLATFORM_AMD__ } // namespace group_dequantize_kernel inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, cudaStream_t stream) { +#ifdef __HIP_PLATFORM_AMD__ + using namespace group_dequantize_kernel; + + const bool use_rowwise_scaling = input->has_data(); + const bool use_colwise_scaling = input->has_columnwise_data(); + NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling, + "Input tensor must have either rowwise or columnwise data."); + NVTE_CHECK(!(use_rowwise_scaling && use_colwise_scaling), + "Dequantize only supports rowwise or columnwise scaling, not both simultaneously."); + + NVTE_CHECK(!input->with_gemm_swizzled_scales, "Input must have scales in compact format."); + NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); + NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4."); + NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type."); + + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (input->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (input->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (input->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (input->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + const bool is_single_tensor = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); + + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + const size_t elts_total = first_logical_dim * last_logical_dim; + if (elts_total == 0) return; + + const size_t num_tensors = input->num_tensors; + + size_t blocks_X = 0; + size_t blocks_Y = 1; + if (is_single_tensor) { + blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); + blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); + } else { + NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0, + "Last dimension of a grouped tensor should be divisible by 128."); + if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) { + blocks_X = DIVUP(first_logical_dim, CHUNK_DIM_Y) * (last_logical_dim / CHUNK_DIM_X); + } else { + blocks_X = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); + } + } + + const dim3 grid(blocks_X, blocks_Y); + const dim3 block(THREADS_PER_CHUNK); + + const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + + // Only the uniform-last-dim path indexes colwise scales from the group's global row. + if (use_colwise_scaling && last_dims_ptr == nullptr && num_tensors > 1) { + constexpr size_t COLWISE_SCALE_ROWS = 32; + std::vector host_first_dims(num_tensors); + if (first_dims_ptr != nullptr) { + NVTE_CHECK_CUDA(cudaMemcpyAsync(host_first_dims.data(), first_dims_ptr, + num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost, + stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + } else { + NVTE_CHECK(first_logical_dim % num_tensors == 0, "Uniform grouped tensor has ", + first_logical_dim, " rows, which is not divisible by ", num_tensors, " tensors."); + host_first_dims.assign(num_tensors, static_cast(first_logical_dim / num_tensors)); + } + for (size_t t = 0; t < num_tensors; t++) { + NVTE_CHECK(static_cast(host_first_dims[t]) % COLWISE_SCALE_ROWS == 0, + "Columnwise grouped dequantize requires each per-tensor first dimension to be " + "divisible by ", + COLWISE_SCALE_ROWS, "; tensor ", t, " has ", host_first_dims[t], "."); + } + } + + // Row strides must be a whole number of vectors; VECTOR_WIDTH is 16 elements. + constexpr size_t VECTOR_ELEMS = 16; + if (last_dims_ptr == nullptr) { + NVTE_CHECK(last_logical_dim % VECTOR_ELEMS == 0, + "Grouped dequantize requires the last dimension to be divisible by ", VECTOR_ELEMS, + "; got ", last_logical_dim, "."); + } else { + std::vector host_last_dims(num_tensors); + NVTE_CHECK_CUDA(cudaMemcpyAsync(host_last_dims.data(), last_dims_ptr, + num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + for (size_t t = 0; t < num_tensors; t++) { + NVTE_CHECK(static_cast(host_last_dims[t]) % VECTOR_ELEMS == 0, + "Grouped dequantize requires each per-tensor last dimension to be divisible by ", + VECTOR_ELEMS, "; tensor ", t, " has ", host_last_dims[t], "."); + } + } + + const e8m0_t *const scales_ptr = + use_rowwise_scaling ? reinterpret_cast(input->scale_inv.dptr) + : reinterpret_cast(input->columnwise_scale_inv.dptr); + + const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data; + + const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; + const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1; + + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + scale_dim_Y_colwise, SCALE_DIM_Y, + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( + scale_dim_X_rowwise, SCALE_DIM_X, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->dtype(), OType, + // Only the single-tensor layout has one row stride, so only it can be + // proven vector-aligned on the host; groups take the general path. + TRANSFORMER_ENGINE_SWITCH_CONDITION( + is_single_tensor && !(last_logical_dim % (32 * sizeof(OType))), IS_ALIGNED, + grouped_dequantize_mxfp8_kernel<<>>( + reinterpret_cast(input_data.dptr), + reinterpret_cast(output->data.dptr), scales_ptr, + first_logical_dim, last_logical_dim, num_tensors, offsets_ptr, + first_dims_ptr, last_dims_ptr);); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); +#else using namespace group_dequantize_kernel; checkCuDriverContext(stream); @@ -490,12 +642,11 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, }); // NOLINT(*) ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); +#endif // #ifdef __HIP_PLATFORM_AMD__ } } // namespace mxfp8 } // namespace dispatch } // namespace transformer_engine -#endif // __HIP_PLATFORM_AMD__ - #endif // TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh new file mode 100644 index 000000000..5392cb21a --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh @@ -0,0 +1,212 @@ +/************************************************************************* + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. + * + * License for AMD contributions = MIT. See LICENSE for more information + ************************************************************************/ + +#pragma once +// drop-in rocm replacement for mxfp8 grouped dequantize kernel +//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 128; + +constexpr size_t ELEMS_PER_THREAD = 16; +constexpr size_t BUFFER_DIM_Y = 16; +constexpr size_t BUFFER_DIM_X = CHUNK_DIM_X; // 128 +constexpr size_t SHMEM_DIM_Y = BUFFER_DIM_Y; // 16 +constexpr size_t SHMEM_DIM_X = BUFFER_DIM_X; // 128 + +constexpr size_t THREADS_PER_CHUNK_X_ROWWISE = CHUNK_DIM_X / ELEMS_PER_THREAD; // 8 = 128 / 16 +constexpr size_t THREADS_PER_CHUNK_X_COLWISE = CHUNK_DIM_X; // 128 +constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y; // 8 = 128 / 16 +static_assert(ITERATIONS >= 1); + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + grouped_dequantize_mxfp8_kernel(const IType *input_ptr, OType *output_ptr, + const e8m0_t *scales_ptr, const size_t first_logical_dim, + const size_t last_logical_dim, const size_t num_tensors, + const int64_t *const offsets_ptr, + const int64_t *const first_dims_ptr, + const int64_t *const last_dims_ptr) { + constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; + + constexpr size_t SCALES_ROWWISE_PER_CHUNK_Y = CHUNK_DIM_Y; // 128 + constexpr size_t SCALES_ROWWISE_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM_X; // 4 = 128 / 32 + + constexpr size_t SCALES_COLWISE_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM_Y; // 4 = 128 / 32 + constexpr size_t SCALES_COLWISE_PER_CHUNK_X = CHUNK_DIM_X; // 128 + + constexpr size_t THREADS_PER_SCALE_X_ROWWISE = + DIVUP(SCALE_DIM_X, ELEMS_PER_THREAD); // 2 = 32 / 16 + constexpr size_t VECTOR_WIDTH = IS_ALIGNED ? 8 : 16; + + const int tid_rowwise_Y = threadIdx.x / THREADS_PER_CHUNK_X_ROWWISE; + const int tid_rowwise_X = threadIdx.x % THREADS_PER_CHUNK_X_ROWWISE; + const int tid_colwise_X = threadIdx.x % THREADS_PER_CHUNK_X_COLWISE; + + const int thread_offset_Y = tid_rowwise_Y; + const int thread_offset_X_rowwise = tid_rowwise_X * ELEMS_PER_THREAD; + + alignas(128) __shared__ IType in_sh[SHMEM_DIM_Y][SHMEM_DIM_X]; + alignas(128) __shared__ OType out_sh[SHMEM_DIM_Y][SHMEM_DIM_X]; + + // Mirrors the launcher's is_single_tensor: null last_dims means a 2D grid, not a tile stride. + const bool uniform_last_dim = (last_dims_ptr == nullptr); + size_t total_tiles = 1; + size_t tiles_stride = 1; + if (!uniform_last_dim) { + total_tiles = 0; + for (size_t t = 0; t < num_tensors; t++) { + const size_t t_rows = + (first_dims_ptr != nullptr) ? static_cast(first_dims_ptr[t]) : first_logical_dim; + const size_t t_cols = static_cast(last_dims_ptr[t]); + total_tiles += DIVUP(t_rows, CHUNK_DIM_Y) * DIVUP(t_cols, CHUNK_DIM_X); + } + tiles_stride = gridDim.x; + } + + for (size_t block_tile = uniform_last_dim ? 0 : blockIdx.x; block_tile < total_tiles; + block_tile += tiles_stride) { + size_t rows = 0; + size_t cols = 0; + int block_id_Y = 0; + int block_id_X = 0; + size_t tensor_base_elts = 0; + size_t scales_base_offset = 0; + + if (uniform_last_dim) { + // Uniform last dim: the group is one contiguous 2D tensor and, because ROCm scale + // tensors are unpadded, the per-tensor scale blocks stack into one contiguous array. + rows = first_logical_dim; + cols = last_logical_dim; + block_id_Y = blockIdx.y; + block_id_X = blockIdx.x; + } else { + // Find the tensor owning this tile. + size_t tiles_before = 0; + size_t tensor_id = 0; + for (; tensor_id < num_tensors; tensor_id++) { + const size_t t_rows = (first_dims_ptr != nullptr) + ? static_cast(first_dims_ptr[tensor_id]) + : first_logical_dim; + const size_t t_cols = static_cast(last_dims_ptr[tensor_id]); + const size_t t_tiles = DIVUP(t_rows, CHUNK_DIM_Y) * DIVUP(t_cols, CHUNK_DIM_X); + if (block_tile < tiles_before + t_tiles) { + const size_t local_tile = block_tile - tiles_before; + const size_t tiles_x = DIVUP(t_cols, CHUNK_DIM_X); + rows = t_rows; + cols = t_cols; + block_id_Y = static_cast(local_tile / tiles_x); + block_id_X = static_cast(local_tile % tiles_x); + break; + } + tiles_before += t_tiles; + tensor_base_elts += t_rows * t_cols; + scales_base_offset += USE_ROWWISE_SCALING ? t_rows * DIVUP(t_cols, SCALE_DIM_X) + : DIVUP(t_rows, SCALE_DIM_Y) * t_cols; + } + if (tensor_id >= num_tensors) continue; + if (offsets_ptr != nullptr) tensor_base_elts = static_cast(offsets_ptr[tensor_id]); + } + + if (rows == 0 || cols == 0) continue; + + // Per-tile views of the group buffers; the kernel arguments stay untouched so that the + // bases do not accumulate across tiles. + const IType *const tile_input_ptr = input_ptr + tensor_base_elts; + OType *const tile_output_ptr = output_ptr + tensor_base_elts; + const e8m0_t *const tile_scales_ptr = scales_ptr + scales_base_offset; + + const size_t scales_stride = USE_ROWWISE_SCALING ? DIVUP(cols, SCALE_DIM_X) : cols; + + const int chunk_offset_Y = block_id_Y * CHUNK_DIM_Y; + const int chunk_offset_X = block_id_X * CHUNK_DIM_X; + + const int scales_rowwise_chunk_offset_Y = block_id_Y * SCALES_ROWWISE_PER_CHUNK_Y; + const int scales_rowwise_chunk_offset_X = block_id_X * SCALES_ROWWISE_PER_CHUNK_X; + const int scales_colwise_chunk_offset_Y = block_id_Y * SCALES_COLWISE_PER_CHUNK_Y; + const int scales_colwise_chunk_offset_X = block_id_X * SCALES_COLWISE_PER_CHUNK_X; + + for (int iter = 0; iter < ITERATIONS; iter++) { + const int chunk_it_offset_y = chunk_offset_Y + iter * BUFFER_DIM_Y; + const int chunk_it_offset_x = chunk_offset_X; + + copy_2d_to_shared(&in_sh[0][0], tile_input_ptr, + chunk_it_offset_x, chunk_it_offset_y, cols, + SHMEM_DIM_Y, SHMEM_DIM_X, rows, cols); + __syncthreads(); + + const int scale_offset_Y = + USE_ROWWISE_SCALING + ? (scales_rowwise_chunk_offset_Y + iter * BUFFER_DIM_Y + tid_rowwise_Y) + : (scales_colwise_chunk_offset_Y + (iter * BUFFER_DIM_Y) / SCALE_DIM_Y); + + const int scale_offset_X = + USE_ROWWISE_SCALING + ? (scales_rowwise_chunk_offset_X + tid_rowwise_X / THREADS_PER_SCALE_X_ROWWISE) + : (scales_colwise_chunk_offset_X + tid_colwise_X); + + const size_t scales_rows = USE_ROWWISE_SCALING ? rows : DIVUP(rows, SCALE_DIM_Y); + const size_t scales_cols = USE_ROWWISE_SCALING ? DIVUP(cols, SCALE_DIM_X) : cols; + + e8m0_t biased_exponent = static_cast(127); + if (static_cast(scale_offset_Y) < scales_rows && + static_cast(scale_offset_X) < scales_cols) { + const size_t scale_idx = scale_offset_Y * scales_stride + scale_offset_X; + biased_exponent = tile_scales_ptr[scale_idx]; + } + const float block_scale = ptx::exp2f(biased_exponent); + + if constexpr (USE_ROWWISE_SCALING) { + Vec in; + Vec out; + + const int shmem_offset_y = thread_offset_Y; + const int shmem_offset_x = thread_offset_X_rowwise; + in.load_from(&in_sh[shmem_offset_y][shmem_offset_x]); + +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; j++) { +#if defined(__gfx1250__) + // FIXME: Force E4M3 OCP interpretation because HIP headers do not declare + // which type gfx1250 supports. This can be removed once HIP headers are updated. + const float elt = + std::is_same_v + ? static_cast(*reinterpret_cast<__hip_fp8_e4m3 *>(&in.data.elt[j])) + : static_cast(in.data.elt[j]); + out.data.elt[j] = static_cast(block_scale * elt); +#else + out.data.elt[j] = static_cast(block_scale * static_cast(in.data.elt[j])); +#endif + } + out.store_to(&out_sh[shmem_offset_y][shmem_offset_x]); + } else { +#pragma unroll + for (int i = 0; i < BUFFER_DIM_Y; i++) { +#if defined(__gfx1250__) + // FIXME: Force E4M3 OCP interpretation because HIP headers do not declare + // which type gfx1250 supports. This can be removed once HIP headers are updated. + const float elt = std::is_same_v + ? static_cast( + *reinterpret_cast<__hip_fp8_e4m3 *>(&in_sh[i][tid_colwise_X])) + : static_cast(in_sh[i][tid_colwise_X]); +#else + const float elt = static_cast(in_sh[i][tid_colwise_X]); +#endif + out_sh[i][tid_colwise_X] = static_cast(block_scale * elt); + } + } + + __syncthreads(); + + bulk_tensor_2d_shared_to_global( + &out_sh[0][0], tile_output_ptr, chunk_it_offset_x, chunk_it_offset_y, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, rows, cols); + + __syncthreads(); + } + } +} diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu index 14704d8cd..296c4b0ad 100644 --- a/transformer_engine/common/fused_attn/flash_attn.cu +++ b/transformer_engine/common/fused_attn/flash_attn.cu @@ -161,7 +161,6 @@ void prepare_flash_attn_bwd(Tensor q, Tensor k, Tensor v, Tensor qkv, cudaStream // multi_tensor_transpose_to_bhsd: BSHD/SBHD -> BHSD // ============================================================================ -#ifndef __HIP_PLATFORM_AMD__ // Disabled on ROCm namespace multi_tensor_transpose_to_bhsd { using flash_attention::Vec; @@ -178,9 +177,11 @@ struct PermuteParams { PermuteSlot slots[kMaxPermuteTensors]; }; +#ifndef __HIP_PLATFORM_AMD__ // CUtensorMap has no ROCm counterpart struct TmaMapParams { CUtensorMap maps[kMaxPermuteTensors]; }; +#endif // ---------- path 3: fallback_not_vec_aligned ---------- @@ -369,6 +370,8 @@ __launch_bounds__(fallback_permute_threads) __global__ // ---------- path 1: TMA ---------- +#ifndef __HIP_PLATFORM_AMD__ // TMA: CUtensorMap, cp.async.bulk, mbarrier, SM 10.0+ + constexpr int tma_permute_threads = 128; constexpr int tma_permute_s_tile_default = 32; @@ -589,6 +592,8 @@ static void create_strided_tensor_map(CUtensorMap &map, void *ptr, DType dtype, } } +#endif // !__HIP_PLATFORM_AMD__ (path 1: TMA) + void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_tensors, NVTE_QKV_Format original_format, cudaStream_t stream) { using namespace transformer_engine; @@ -605,7 +610,15 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_ size_t s_max = 0, h_max = 0, s_min = SIZE_MAX; size_t d_in_max = 0, d_out_max = 0; bool any_not_vec_aligned = false; +#ifdef __HIP_PLATFORM_AMD__ + // Both readers below are compiled out on ROCm, so this only keeps the variable + // defined and states the intent: had they been reachable, the sm >= 100 test + // would not have excluded us -- sm_arch() returns the raw gfx number, so + // gfx1250 reports 125 and satisfies it. + bool all_tma_ok = false; +#else bool all_tma_ok = true; +#endif for (size_t i = 0; i < num_tensors; ++i) { const size_t H = outputs[i].shape()[1]; @@ -623,6 +636,7 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_ if (inner < 32 || inner % 16 != 0) all_tma_ok = false; } +#ifndef __HIP_PLATFORM_AMD__ if (all_tma_ok) { const int sm = cuda::sm_arch(cuda::current_device()); if (sm < 100) { @@ -641,12 +655,14 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_ } } } +#endif // !__HIP_PLATFORM_AMD__ // Dispatch order: // 1. TMA path: SM 10.0+, D_in*elem >= 32 && 16-aligned, supported dtype, // and s_tile*D_in*elem is uint4-aligned. // 2. Fallback path (vec-aligned): vectorized loads/stores when D_in*elem % 4 == 0. // 3. Fallback path (not-vec-aligned): shared-memory transpose when D_in*elem % 4 != 0. +#ifndef __HIP_PLATFORM_AMD__ if (all_tma_ok) { const size_t s_tile = std::min(static_cast(tma_permute_s_tile_default), s_min); bool tma_aligned = true; @@ -686,6 +702,7 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_ return; } } +#endif // !__HIP_PLATFORM_AMD__ (TMA dispatch) if (!any_not_vec_aligned) { const unsigned int permute_s_splits = std::max( @@ -730,7 +747,6 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_ } } // namespace multi_tensor_transpose_to_bhsd -#endif // =================================================================================== // multi_tensor_pad_last_dim: pad the last dim of multiple tensors to certain alignment @@ -872,7 +888,6 @@ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs size_t num_tensors, NVTE_QKV_Format original_format, cudaStream_t stream) { NVTE_API_CALL(nvte_multi_tensor_transpose_to_bhsd); -#ifndef __HIP_PLATFORM_AMD__ NVTE_CHECK(original_format == NVTE_QKV_Format::NVTE_BSHD || original_format == NVTE_QKV_Format::NVTE_SBHD, "nvte_multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently " @@ -890,14 +905,6 @@ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs multi_tensor_transpose_to_bhsd::multi_tensor_transpose_to_bhsd( in_vec.data() + offset, out_vec.data() + offset, batch, original_format, stream); } -#else - (void)inputs; - (void)outputs; - (void)num_tensors; - (void)original_format; - (void)stream; - NVTE_ERROR("nvte_multi_tensor_transpose_to_bhsd is not supported on ROCm."); -#endif } void nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors, diff --git a/transformer_engine/common/gemm/kittens/CMakeLists.txt b/transformer_engine/common/gemm/kittens/CMakeLists.txt index f09a7c7ca..6fee3125a 100644 --- a/transformer_engine/common/gemm/kittens/CMakeLists.txt +++ b/transformer_engine/common/gemm/kittens/CMakeLists.txt @@ -57,6 +57,9 @@ else() add_library(kittens_gemm_cdna4 OBJECT cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp) set_source_files_properties(cdna4/mxfp8_gemm.cpp PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only") + # -ffast-math folds fpext(fptrunc(x)) back to x, so any deliberate narrow-then-widen + # in this file has to be made opaque to the optimizer -- see round_to_out_dtype() in + # cdna4/blockwise_fp8_gemm_helper.cuh. set_source_files_properties(cdna4/blockwise_fp8_gemm.cpp PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only -ffast-math") set_target_properties(kittens_gemm_cdna4 PROPERTIES diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh index 413b126c8..086419a65 100644 --- a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh +++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh @@ -94,20 +94,22 @@ __device__ inline float read_elem(const void *p, int dtype, int idx) { } template -__device__ inline float rtne_cast_roundtrip(float v) { - if constexpr (std::is_same_v) { - return v; +__device__ inline OType convert_out(float v) { + if constexpr (std::is_same_v) { + return kittens::base_types::convertor::convert(rtne_bias(v)); } else { - return static_cast(kittens::base_types::convertor::convert(rtne_bias(v))); + return kittens::base_types::convertor::convert(v); } } +// Deliberate: the reference rounds to the output type before the beta*C add +// (blockwise_fp8_gemm_reference.py::qgemm). convert_out keeps rtne_bias at store granularity. template -__device__ inline OType convert_out(float v) { - if constexpr (std::is_same_v) { - return kittens::base_types::convertor::convert(rtne_bias(v)); +__device__ inline float round_to_out_dtype(float v) { + if constexpr (std::is_same_v) { + return v; } else { - return kittens::base_types::convertor::convert(v); + return kittens::base_types::convertor::convert(convert_out(v)); } } @@ -184,7 +186,7 @@ __device__ inline void apply_epilogue( float x = v[r]; if constexpr (HAS_BIAS) x += bias_v; if constexpr (HAS_BETA) { - x = rtne_cast_roundtrip(x); + x = round_to_out_dtype(x); x += beta * static_cast(c_in[m_g * N + col]); } if constexpr (HAS_GELU) { diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh index 7207a97fd..bc9be2284 100644 --- a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh +++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh @@ -83,14 +83,21 @@ __device__ inline float read_elem(const void *p, int dtype, int idx) { return reinterpret_cast(p)[idx]; } +// Deliberate: the reference rounds to the output type before the beta*C add +// (blockwise_fp8_gemm_reference.py::qgemm). This TU is built with -ffast-math, which folds +// fpext(fptrunc(x)) back to x, so the empty asm is required to keep the conversion. template __device__ inline float round_to_out_dtype(float v) { if constexpr (std::is_same_v) { return v; } else if constexpr (std::is_same_v) { - return __bfloat162float(__float2bfloat16(v)); + uint16_t bits = __builtin_bit_cast(uint16_t, __float2bfloat16(v)); + asm volatile("" : "+v"(bits)); + return __bfloat162float(__builtin_bit_cast(kittens::bf16, bits)); } else { - return __half2float(__float2half(v)); + uint16_t bits = __builtin_bit_cast(uint16_t, __float2half(v)); + asm volatile("" : "+v"(bits)); + return __half2float(__builtin_bit_cast(__half, bits)); } } diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu index 082c96ee7..e200a41bf 100644 --- a/transformer_engine/common/gemm/rocm_gemm.cu +++ b/transformer_engine/common/gemm/rocm_gemm.cu @@ -535,8 +535,10 @@ static void dequant_fp4_gemm_inputs( const void** alpha_ptr_out, hipStream_t stream) { const float fp4_max = 6.0f; - const float fp8_max = te_fp8_fnuz() ? 240.0f : 448.0f; - const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max * fp8_max); + // Read the bound per operand as upstream does: 4over6 tensors carry 256, not 448. + const float fp8_max_A = te_fp8_fnuz() ? 240.0f : static_cast(inputA.nvfp4_e4m3_max); + const float fp8_max_B = te_fp8_fnuz() ? 240.0f : static_cast(inputB.nvfp4_e4m3_max); + const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B); const float* amax_A = (transa == CUBLAS_OP_T) ? reinterpret_cast(inputA.amax.dptr) diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index c462b3014..bbc5e6cd8 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -1,4 +1,6 @@ /************************************************************************* + * This file was modified for portability to AMDGPU + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -68,7 +70,8 @@ void quantize_transpose_vector_blockwise_fp4( 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); + const SimpleTensor &noop_tensor, const int nvfp4_e4m3_max, + const NVTENVFP44Over6Mode nvfp4_4over6_mode, cudaStream_t stream); } // namespace transformer_engine::detail 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 e6481e09a..a8dadf228 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -191,6 +191,33 @@ constexpr int smem_size_for_tile(int tile_dim) { static __device__ constexpr unsigned int WARP_REDUCE_AMAX_GROUP_MASKS[8] = { 0x01010101, 0x02020202, 0x04040404, 0x08080808, 0x10101010, 0x20202020, 0x40404040, 0x80808080}; +#ifdef __HIP_PLATFORM_AMD__ +// Sum the 16 row errors of a 2D block in upstream's tree order: pair rows 8 +// apart, then 4, then 2, then 1 (quantize_4over6_nvfp4.cuh::reduce_group_sum_16). +// fp32 addition is not associative and 4over6 selection compares two sums that +// are frequently equal to within an ulp, so the association is part of the +// contract rather than an implementation detail. Upstream can run this tree in +// warp shuffles because one warp there owns 16 consecutive rows; here a warp +// owns only 4, so the rows paired first sit in different warps and the whole +// set has to be staged in shared memory and reduced by one thread. +template +__device__ __forceinline__ float reduceBlockError(const CType* staged) { + float v[kBlockSize]; +#pragma unroll + for (int i = 0; i < kBlockSize; ++i) { + v[i] = static_cast(staged[i]); + } +#pragma unroll + for (int stride = kBlockSize / 2; stride > 0; stride /= 2) { +#pragma unroll + for (int i = 0; i < stride; ++i) { + v[i] = __fadd_rn(v[i], v[i + stride]); + } + } + return v[0]; +} +#endif + // max for every group_size elements in warp template __device__ __forceinline__ float groupMax(float val, unsigned int groupMask) { @@ -220,18 +247,100 @@ __device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale, TypeExtrema::max); } +#ifdef __HIP_PLATFORM_AMD__ +// Round-trip one value through FP4 at a given scale and return its contribution +// to the block error. +// E2M1 holds 16 values, so decode by table rather than through the packed-pair +// conversion intrinsic, whose element ordering is easy to get wrong here. +__device__ __forceinline__ float fp4_decode_e2m1(const __hip_fp4_storage_t q) { + constexpr float kMagnitude[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + const float m = kMagnitude[q & 0x7]; + return (q & 0x8) ? -m : m; +} + +// Rounding is pinned with the _rn intrinsics and the operand order upstream uses +// (quantize_4over6_nvfp4.cuh:199). Selection turns on comparing two sums that are +// often close, so an FMA contraction here flips near-ties and the result stops +// being bit-comparable with the reference. +__device__ __forceinline__ float fp4_roundtrip_err(const float x, const float block_scale_inverse, + const float sf, const float global_amax, + const float err_denom, const bool use_mse) { + const __hip_fp4_storage_t q = + __hip_cvt_float_to_fp4(__fmul_rn(x, block_scale_inverse), __HIP_E2M1, hipRoundNearest); + const float dequant = fp4_decode_e2m1(q); + const float val = __fdiv_rn(__fmul_rn(__fmul_rn(dequant, sf), global_amax), err_denom); + const float diff = __fsub_rn(val, x); + return use_mse ? __fmul_rn(diff, diff) : fabsf(diff); +} + +// 4over6: a block scale expanded by 1.5x maps the block max onto FP4 code 4 +// rather than code 6, spacing the codes more finely for blocks whose values sit +// away from the top of the range. Score both encodings of the block and keep the +// one that reconstructs it with less error. +// +// Takes the block as a plain array: the rowwise and columnwise call sites walk +// shared memory differently, so a shared loop cannot serve both. +// +// Split in two so a caller that has already reduced the block error can rebuild the +// two scales without walking the block again. +template +__device__ __forceinline__ void Compute4Over6Scales(const float amax, + const float global_decode_scale, + const float global_encode_scale, + ScaleType* sf6_out, ScaleType* sf4_out, + float* enc6_out, float* enc4_out) { + // Both candidates come off upstream's base expression, amax/6*S expanded by 1.5 for + // map-to-4 (quantize_4over6_nvfp4.cuh:129). Folding the /6 into the multiplier the way + // the non-4over6 path does rounds differently and moves blocks across E4M3 codes. + const float base = __fmul_rn(__fdiv_rn(amax, 6.0f), global_encode_scale); + const ScaleType sf6 = static_cast(fminf(base, TypeExtrema::max)); + const ScaleType sf4 = + static_cast(fminf(__fmul_rn(base, 1.5f), TypeExtrema::max)); + *sf6_out = sf6; + *sf4_out = sf4; + *enc6_out = ComputeEncodeScaleFP4(sf6, global_decode_scale); + *enc4_out = ComputeEncodeScaleFP4(sf4, global_decode_scale); +} + +template +__device__ __forceinline__ void Compute4Over6Candidates( + const float* x, const int n, const float amax, const float global_decode_scale, + const bool use_mse, const float global_encode_scale, const float true_global_amax, + const float err_denom, ScaleType* sf6_out, ScaleType* sf4_out, float* enc6_out, float* enc4_out, + float* err6_out, float* err4_out) { + Compute4Over6Scales(amax, global_decode_scale, global_encode_scale, sf6_out, sf4_out, + enc6_out, enc4_out); + const ScaleType sf6 = *sf6_out; + const ScaleType sf4 = *sf4_out; + const float enc6 = *enc6_out; + const float enc4 = *enc4_out; + // Dequantisation is val = dequant * sf * global_amax / (fp4_max * e4m3_max). Take the + // amax the encode scale was derived from; dividing err_denom back out costs a bit, and + // selection compares sums that are often equal to within one. + const float global_amax = true_global_amax; + + float err6 = 0.0f; + float err4 = 0.0f; + for (int i = 0; i < n; ++i) { + err6 = __fadd_rn(err6, fp4_roundtrip_err(x[i], enc6, static_cast(sf6), global_amax, + err_denom, use_mse)); + err4 = __fadd_rn(err4, fp4_roundtrip_err(x[i], enc4, static_cast(sf4), global_amax, + err_denom, use_mse)); + } + *err6_out = err6; + *err4_out = err4; +} + +#endif // __HIP_PLATFORM_AMD__ + template __device__ __forceinline__ float ComputeOutputFP4(IType input, float encode_scale) { return static_cast(input) * encode_scale; } -__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { -#if defined(__HIP_PLATFORM_AMD__) && !defined(__HIP_DEVICE_COMPILE__) - // On AMD host, TypeExtrema::max is non-constexpr (runtime FNUZ detection) - const float fp8_max = TypeExtrema::max; -#else - constexpr float fp8_max = TypeExtrema::max; -#endif +__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax, + const float e4m3_max) { + const float fp8_max = e4m3_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 @@ -433,14 +542,15 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01, template + bool kApplyStochasticRounding, bool kIs2DBlockScaling, bool kRowScaledNVFP4, + bool kUse4Over6> __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel( const IType* const input, const float* global_amax, OType* const output_c, OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t, const size_t row_length, const size_t num_rows, const size_t scale_stride_x, const size_t scale_stride_y, const size_t scale_t_stride_x, const size_t scale_t_stride_y, const size_t kScaleBlockDim, const float epsilon, const size_t* rng_state, - const float* noop_ptr) { + const float* noop_ptr, const float e4m3_max, const bool use_mse) { constexpr int kNVecContainer = kNVecOut / kNFP4PerContainer; using SMemVec = Vec; using OVec = Vec; @@ -478,12 +588,21 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo // for 128x128 block, 2D block scaling means there will be 8x8 amax values for nvfp4, 4x4 for 2D mxfp4 // use constexpr to define the size, when not using 2D, use minimal size 1x1 constexpr int kFP4BlockScalingSize = 16; + constexpr float kFp4Max = 6.0f; // TypeExtrema::max constexpr int k2DBlockAmaxDim = kIs2DBlockScaling ? (kTileDim / kFP4BlockScalingSize) : 1; constexpr int kNumRowsPerWarp = kThreadsPerWarp / kNumThreadsStore; // 4 constexpr int k2DBlockAmaxReduceDim = kIs2DBlockScaling ? (kFP4BlockScalingSize / kNumRowsPerWarp) : 1; __shared__ CType amax_smem_red[k2DBlockAmaxDim][k2DBlockAmaxDim][k2DBlockAmaxReduceDim]; __shared__ CType amax_smem[k2DBlockAmaxDim][k2DBlockAmaxDim]; + // 4over6 scores its two candidates over the same block the scale serves, so the + // error is reduced over the same 16x16 block the 2D amax is. Unlike the amax it + // stages all 16 rows: max is associative, the fp32 sum is not, so the rows cannot + // be pre-combined per warp without diverging from upstream's tree. + constexpr int k4Over6Dim = kUse4Over6 ? k2DBlockAmaxDim : 1; + constexpr int k4Over6RowDim = kUse4Over6 ? kFP4BlockScalingSize : 1; + __shared__ CType err_smem_red[2][k4Over6Dim][k4Over6Dim][k4Over6RowDim]; + __shared__ CType err_smem[2][k4Over6Dim][k4Over6Dim]; // Step 1: Load input to shared memory { @@ -532,12 +651,17 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int kNumThreadsReduce = kScaleBlockDim / kNVecOut; #ifdef __HIP_PLATFORM_AMD__ - const float global_encode_scale = - (kIsE8Scaling || global_amax == nullptr) ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]); + const float global_encode_scale = (kIsE8Scaling || global_amax == nullptr) + ? 1.0f + : ComputeGlobalEncodeScaleFP4(global_amax[0], e4m3_max); #else const float global_encode_scale = - kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]); + kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0], e4m3_max); #endif + // Mirrors the global_encode_scale guard exactly: where that falls back to 1.0f the + // dequantisation denominator is err_denom, so the equivalent amax is err_denom itself. + const float block_true_global_amax = + (kIsE8Scaling || global_amax == nullptr) ? kFp4Max * e4m3_max : global_amax[0]; constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; const float global_decode_scale = 1.0 / global_encode_scale; @@ -632,15 +756,68 @@ __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 ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx], e4m3_max) : 1.0f; + } + float row_true_global_amax = block_true_global_amax; + if constexpr (kRowScaledNVFP4) { + row_true_global_amax = row_idx < num_rows ? global_amax[row_idx] : kFp4Max * e4m3_max; } 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); - float encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale); + ScaleType scale_inv; + float encode_scale; +#ifdef __HIP_PLATFORM_AMD__ + // 4over6 lives here only on ROCm; CUDA dispatches it to upstream's own kernel, + // so the helpers below are not defined there. + if constexpr (kUse4Over6) { + float blk[kNVecOut]; +#pragma unroll + for (int i = 0; i < kNVecOut / kNVecSMem; ++i) { +#pragma unroll + for (int j = 0; j < kNVecSMem; ++j) { + blk[i * kNVecSMem + j] = static_cast(smem_vec[i].data.elt[j]); + } + } + ScaleType sf6, sf4; + float enc6, enc4, err6, err4; + Compute4Over6Candidates( + blk, kNVecOut, amax, row_global_decode_scale, use_mse, row_global_encode_scale, + row_true_global_amax, kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4); + if constexpr (kIs2DBlockScaling) { + // The scale spans a 16x16 block, so the error must too. Stage this row's + // error and let one thread run upstream's tree over the 16; indices + // recomputed rather than hoisted so the non-4over6 code is untouched. + constexpr int kNumRowsPerIter2 = kThreadsPerBlock / kNumThreadsStore; + const int warp_i = threadIdx.x / kThreadsPerWarp; + const int tx = threadIdx.x % kNumThreadsStore; + const int ty = (threadIdx.x / kNumThreadsStore) % kNumRowsPerWarp; + const int row_i = iter * kNumRowsPerIter2 + warp_i * kNumRowsPerWarp + ty; + const int blk_i = row_i / kFP4BlockScalingSize; + const int row_in_blk = row_i % kFP4BlockScalingSize; + err_smem_red[0][blk_i][tx][row_in_blk] = err6; + err_smem_red[1][blk_i][tx][row_in_blk] = err4; + __syncthreads(); + if (row_in_blk == 0) { + err_smem[0][blk_i][tx] = + reduceBlockError(err_smem_red[0][blk_i][tx]); + err_smem[1][blk_i][tx] = + reduceBlockError(err_smem_red[1][blk_i][tx]); + } + __syncthreads(); + err6 = err_smem[0][blk_i][tx]; + err4 = err_smem[1][blk_i][tx]; + } + const bool pick4 = err4 < err6; // tie keeps map-to-6, matching upstream + scale_inv = pick4 ? sf4 : sf6; + encode_scale = pick4 ? enc4 : enc6; + } else // NOLINT(readability/braces) +#endif + { + scale_inv = ComputeDecodeScaleFP4(amax, row_global_encode_scale_multiplier); + encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale); + } // Step 2.5: Write scale_inv bool write_scale_inv = is_src_lane; if constexpr (!kAligned) { @@ -752,6 +929,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x] = amax_2d; } __syncthreads(); + // No 4over6 work here: Step 3 scores the columnwise candidates on their own + // column groups, so it only needs amax_smem from this pass. r_s += r_stride; } } @@ -779,6 +958,55 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo // This mask represents which threads should do the reduction together. const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane; const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0; +#ifdef __HIP_PLATFORM_AMD__ + // Step 3.0: 4over6 block errors for the whole tile. + // Upstream scores the columnwise candidates on their own column groups + // (quantize_4over6_nvfp4.cuh::quantize_stage_colwise) instead of reusing what the + // rowwise pass computed. Both passes cover the same 16x16 block but sum its errors + // along different axes, and fp32 addition is not associative, so the two can land on + // different candidates for one block; that is the algorithm, not a bug. + // + // Done here rather than inside the loop below because a block's 16 columns are split + // across both smem_idx halves, so no single pass of that loop sees a whole block. + // Keeping it separate also leaves nothing live across the barriers. + if constexpr (kUse4Over6 && kIs2DBlockScaling) { + const int row_blk = threadIdx.x % kNumThreadsStore; + int c = c_s; +#pragma unroll + for (int iter = 0; iter < num_iterations; ++iter, c += c_stride) { +#pragma unroll + for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) { + const int col_in_tile = c * kNVecSMem + smem_idx; + float blk[kNVecOut]; +#pragma unroll + for (int i = 0; i < kNVecOut; ++i) { + blk[i] = static_cast(smem[(r_s + i) * kSMemCol + c].data.elt[smem_idx]); + } + ScaleType sf6, sf4; + float enc6, enc4, err6, err4; + Compute4Over6Candidates( + blk, kNVecOut, amax_smem[row_blk][col_in_tile / kFP4BlockScalingSize], + global_decode_scale, use_mse, global_encode_scale, block_true_global_amax, + kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4); + err_smem_red[0][row_blk][col_in_tile / kFP4BlockScalingSize] + [col_in_tile % kFP4BlockScalingSize] = err6; + err_smem_red[1][row_blk][col_in_tile / kFP4BlockScalingSize] + [col_in_tile % kFP4BlockScalingSize] = err4; + } + } + __syncthreads(); + // One thread per (row block, column block) runs upstream's tree over the 16 staged + // column errors; the columns paired first are 8 apart and so live in different + // warps, exactly as the rows do in the rowwise pass. + for (int idx = threadIdx.x; idx < k2DBlockAmaxDim * k2DBlockAmaxDim; idx += blockDim.x) { + const int rb = idx / k2DBlockAmaxDim; + const int cb = idx % k2DBlockAmaxDim; + err_smem[0][rb][cb] = reduceBlockError(err_smem_red[0][rb][cb]); + err_smem[1][rb][cb] = reduceBlockError(err_smem_red[1][rb][cb]); + } + __syncthreads(); + } +#endif #pragma unroll for (int iter = 0; iter < num_iterations; ++iter) { SMemVec smem_vec[kNVecOut]; @@ -823,9 +1051,39 @@ __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); - float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); + ScaleType scale_inv; + float encode_scale; +#ifdef __HIP_PLATFORM_AMD__ + if constexpr (kUse4Over6) { + ScaleType sf6, sf4; + float enc6, enc4, err6, err4; + if constexpr (kIs2DBlockScaling) { + // Step 3.0 already scored this block; only the two scales are needed here. + const int row_blk = threadIdx.x % kNumThreadsStore; + const int col_blk = (c_s * kNVecSMem + smem_idx) / kFP4BlockScalingSize; + err6 = err_smem[0][row_blk][col_blk]; + err4 = err_smem[1][row_blk][col_blk]; + Compute4Over6Scales(amax, global_decode_scale, global_encode_scale, &sf6, + &sf4, &enc6, &enc4); + } else { + float blk[kNVecOut]; +#pragma unroll + for (int i = 0; i < kNVecOut; ++i) { + blk[i] = static_cast(smem_vec[i].data.elt[smem_idx]); + } + Compute4Over6Candidates( + blk, kNVecOut, amax, global_decode_scale, use_mse, global_encode_scale, + block_true_global_amax, kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4); + } + const bool pick4 = err4 < err6; // tie keeps map-to-6, matching upstream + scale_inv = pick4 ? sf4 : sf6; + encode_scale = pick4 ? enc4 : enc6; + } else // NOLINT(readability/braces) +#endif + { + scale_inv = ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier); + encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); + } // Step 3.5: Write scale_inv_t bool write_scale_inv = is_src_lane; if constexpr (!kAligned) { @@ -886,7 +1144,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo } } } -} +} // NOLINT(readability/fn_size) 4over6 grows upstream's kernel past 500 lines } // namespace } // namespace quantize_transpose_nvfp4 @@ -900,7 +1158,8 @@ void quantize_transpose_vector_blockwise_fp4( 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) { + const SimpleTensor& noop_tensor, const int nvfp4_e4m3_max, + const NVTENVFP44Over6Mode nvfp4_4over6_mode, cudaStream_t stream) { NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4); #if defined(__HIP_PLATFORM_AMD__) || CUDA_VERSION >= 12080 @@ -929,6 +1188,14 @@ void quantize_transpose_vector_blockwise_fp4( return; } + // gfx942 stores E4M3 as FNUZ (max 240), gfx950 as OCP (max 448). The quantizer must use the + // same bound the GEMM's dequant does (rocm_gemm.cu:539), or the two disagree by 448/240. +#ifdef __HIP_PLATFORM_AMD__ + const float e4m3_max_for_arch = te_fp8_fnuz() ? 240.0f : static_cast(nvfp4_e4m3_max); +#else + const float e4m3_max_for_arch = static_cast(nvfp4_e4m3_max); +#endif + size_t scale_stride_x = 0; size_t scale_stride_y = 0; @@ -1009,48 +1276,60 @@ void quantize_transpose_vector_blockwise_fp4( TRANSFORMER_ENGINE_SWITCH_CONDITION( row_scaled_nvfp4, kRowScaledNVFP4, + // 4over6 only has a body on ROCm, so pin the switch + // false on CUDA rather than instantiate every + // specialization twice for a branch it cannot take. +#ifdef __HIP_PLATFORM_AMD__ + TRANSFORMER_ENGINE_SWITCH_CONDITION( + nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, kUse4Over6, +#else + TRANSFORMER_ENGINE_SWITCH_CONDITION( + false, kUse4Over6, +#endif + #ifdef __HIP_PLATFORM_AMD__ - size_t smem_bytes = - smem_size_for_tile(tile_dim) * sizeof(InputType); + size_t smem_bytes = + smem_size_for_tile(tile_dim) * sizeof(InputType); #else - size_t smem_bytes = kSMemSize * sizeof(InputType); + size_t smem_bytes = kSMemSize * sizeof(InputType); #endif - auto kernel = block_scaled_1d_cast_transpose_kernel< - kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned, - float, InputType, OutputType, ScaleType, kSwizzledScale, - kApplyStochasticRounding, kIs2DBlockScaling, - kRowScaledNVFP4>; - if (smem_bytes >= 48 * 1024) { - cudaError_t err = cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_bytes); - NVTE_CHECK(err == cudaSuccess, - "Failed to set dynamic shared memory size."); + auto kernel = block_scaled_1d_cast_transpose_kernel< + kReturnIdentity, kReturnTranspose, kPow2Scale, + kAligned, float, InputType, OutputType, ScaleType, + kSwizzledScale, kApplyStochasticRounding, + kIs2DBlockScaling, kRowScaledNVFP4, kUse4Over6>; + if (smem_bytes >= 48 * 1024) { + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes); + NVTE_CHECK(err == cudaSuccess, + "Failed to set dynamic shared memory size."); #ifdef __HIP_PLATFORM_AMD__ - } kernel<<>>( + } kernel<<>>( #else - } kernel<<>>( + } kernel<<>>( #endif - reinterpret_cast(input.dptr), - reinterpret_cast(global_amax.dptr), - reinterpret_cast(output.dptr), - reinterpret_cast(output_t.dptr), - reinterpret_cast(scale_inv.dptr), - reinterpret_cast(scale_inv_t.dptr), - row_length, num_rows, scale_stride_x, scale_stride_y, - scale_t_stride_x, scale_t_stride_y, kScaleBlockDim, - epsilon, rng_state, - noop_ptr);) // kRowScaledNVFP4 - ) // kIs2DBlockScaling - ) // kApplyStochasticRounding - ) // kSwizzledScale - ) // kAligned - ) // kReturnTranspose - ) // kReturnIdentity - ) // OutputType - ) // InputType + reinterpret_cast(input.dptr), + reinterpret_cast(global_amax.dptr), + reinterpret_cast(output.dptr), + reinterpret_cast(output_t.dptr), + reinterpret_cast(scale_inv.dptr), + reinterpret_cast(scale_inv_t.dptr), + row_length, num_rows, scale_stride_x, scale_stride_y, + scale_t_stride_x, scale_t_stride_y, kScaleBlockDim, + epsilon, rng_state, noop_ptr, e4m3_max_for_arch, + nvfp4_4over6_mode == + kNVTENVFP44Over6MinMSE);)) // kUse4Over6 / kRowScaledNVFP4 + ) // kIs2DBlockScaling + ) // kApplyStochasticRounding + ) // kSwizzledScale + ) // kAligned + ) // kReturnTranspose + ) // kReturnIdentity + ) // OutputType + ) // InputType NVTE_CHECK_CUDA(cudaGetLastError()); #else diff --git a/transformer_engine/common/util/standalone_topk.cuh b/transformer_engine/common/util/standalone_topk.cuh index 3d19cbfcf..190027acc 100644 --- a/transformer_engine/common/util/standalone_topk.cuh +++ b/transformer_engine/common/util/standalone_topk.cuh @@ -1,4 +1,6 @@ /************************************************************************* + * This file was modified for portability to AMDGPU + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -7,7 +9,12 @@ #pragma once #include +#ifndef __HIP_PLATFORM_AMD__ +// Only scan_warp_version needs cg::reduce, and that is excluded below on ROCm; +// HIP keeps the reduce overloads behind an amd_detail header the umbrella one +// does not pull in. #include +#endif #include #include @@ -15,6 +22,12 @@ #include namespace cg = cooperative_groups; +#ifdef __HIP_PLATFORM_AMD__ +// hipify rewrites the CUB class templates but not Traits or the BLOCK_* +// enumerators, so resolve the rest through the namespace. +namespace cub = hipcub; +#endif + // Workspace pointer-alignment helpers. inline size_t calc_aligned_size(const std::vector &sizes) { const size_t ALIGN_BYTES = 256; @@ -485,6 +498,9 @@ __device__ void choose_bucket(Counter *counter, const IdxT *histogram, } } +// Dead: the call site below is commented out upstream. Excluded on ROCm rather +// than ported -- it is a wave32 reduction over a 32-bit mask, which HIP rejects. +#ifndef __HIP_PLATFORM_AMD__ template __device__ void scan_warp_version(cg::thread_block_tile const &warp, volatile IdxT *histogram, Counter *counter, const IdxT k, @@ -569,6 +585,7 @@ __device__ void scan_warp_version(cg::thread_block_tile const &warp, // } } } +#endif // !__HIP_PLATFORM_AMD__ (scan_warp_version) // For one-block version, last_filter() could be called when pass < num_passes // - 1. So `pass` could not be constexpr template diff --git a/transformer_engine/common/util/topk.cu b/transformer_engine/common/util/topk.cu index 21018a494..33f03fc91 100644 --- a/transformer_engine/common/util/topk.cu +++ b/transformer_engine/common/util/topk.cu @@ -1,4 +1,6 @@ /************************************************************************* + * This file was modified for portability to AMDGPU + * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -49,7 +51,13 @@ void nvte_topk(cudaStream_t stream, const NVTETensor keys_in, const NVTETensor l } while (0) if (dtype == DType::kBFloat16) { +#ifdef __HIP_PLATFORM_AMD__ + // hipCUB specializes NumericTraits for hip_bfloat16, not __hip_bfloat16, + // and the radix select needs those traits to twiddle the key bits. + DISPATCH_TOPK(hip_bfloat16); +#else DISPATCH_TOPK(__nv_bfloat16); +#endif } else if (dtype == DType::kFloat32) { DISPATCH_TOPK(float); } else { diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 3fc054303..725ac26e8 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -151,9 +151,7 @@ pybind11::dict Registrations() { #endif // NVTE_WITH_NCCL_EP // TopK -#ifndef USE_ROCM // Disabled on ROCm dict["te_topk_ffi"] = EncapsulateFFI(TopkHandler); -#endif return dict; } @@ -177,9 +175,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_norm_bwd_workspace_sizes", &GetNormBackwardWorkspaceSizes); m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes); m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes); -#ifndef USE_ROCM m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); -#endif m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); #ifndef USE_ROCM diff --git a/transformer_engine/jax/csrc/extensions/topk.cpp b/transformer_engine/jax/csrc/extensions/topk.cpp index 0cc68a45a..450ff08b3 100644 --- a/transformer_engine/jax/csrc/extensions/topk.cpp +++ b/transformer_engine/jax/csrc/extensions/topk.cpp @@ -1,6 +1,4 @@ /************************************************************************* - * This file was modified for portability to AMDGPU - * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. @@ -14,9 +12,6 @@ namespace transformer_engine { namespace jax { -// Disabled on ROCm -#ifndef USE_ROCM - // --------------------------------------------------------------------------- // JAX FFI handler // --------------------------------------------------------------------------- @@ -105,7 +100,5 @@ pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k) { return pybind11::make_tuple(std::make_pair(work_shape, workspace_tensor.dtype())); } -#endif - } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index 0a47e600d..c6b312434 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -13,6 +13,12 @@ #include "transformer_engine/transformer_engine.h" #ifdef USE_ROCM +#include + +#include +#include +#include + #include "common/common.h" #endif @@ -367,10 +373,25 @@ at::Tensor allocate_amax_workspace(const TensorWrapper& input_tensor) { return at::empty(0, at::CUDA(at::kFloat)); } - const auto N = input_tensor.numel(); - size_t workspace_blocks = nvte_amax_workspace_num_blocks(N); - - return at::empty(workspace_blocks, at::CUDA(at::kFloat)); + // A per-call tensor returns to the caching allocator at the caller's closing brace while + // amax_kernel and amax_final_reduce are still only enqueued. Stream-ordered reuse covers + // that in eager mode, but graph capture replaces stream order with a DAG, so the freed + // block is aliased in the graph mempool and replays read partials they never wrote. + // One buffer per (device, stream), sized to the block cap: same stream serializes. + static std::mutex amax_ws_mutex; + static std::map, at::Tensor> amax_ws_cache; + + const int device_id = at::cuda::current_device(); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + std::lock_guard lock(amax_ws_mutex); + auto& ws = amax_ws_cache[std::make_pair(device_id, stream)]; + if (!ws.defined()) { + // Any N past the block cap yields the cap itself; avoid SIZE_MAX so DIVUP cannot overflow. + ws = at::empty(nvte_amax_workspace_num_blocks(static_cast(1) << 40), + at::CUDA(at::kFloat)); + } + return ws; } #endif