diff --git a/tests/cpp/operator/test_dequantize_mxfp8.cu b/tests/cpp/operator/test_dequantize_mxfp8.cu index 82cfdd591..33416f74d 100644 --- a/tests/cpp/operator/test_dequantize_mxfp8.cu +++ b/tests/cpp/operator/test_dequantize_mxfp8.cu @@ -391,6 +391,25 @@ void performTest_x1_swizzled(const size_t rows, const size_t unpadded_blocks_Y_colwise = divide_round_up(rows, block_size_rows); const size_t unpadded_blocks_X_colwise = cols; +#ifdef __HIP_PLATFORM_AMD__ + size_t blocks_Y_rowwise = round_up_to_nearest_multiple(unpadded_blocks_Y_rowwise, + scale_tensor_alignment_Y_rowwise); + size_t blocks_X_rowwise = round_up_to_nearest_multiple(unpadded_blocks_X_rowwise, + scale_tensor_alignment_X_rowwise); + size_t blocks_Y_colwise = round_up_to_nearest_multiple(unpadded_blocks_Y_colwise, + scale_tensor_alignment_Y_colwise); + size_t blocks_X_colwise = round_up_to_nearest_multiple(unpadded_blocks_X_colwise, + scale_tensor_alignment_X_colwise); + // gfx1250 pads MXFP8 scales to a multiple of 4 in both dims. The GEMM-swizzle producer reads + // the compact scales with this padded stride, so the reference/fill must use the same layout. + if (getDeviceComputeCapability() == 125) { + const size_t align = mxfp8_gfx1250_scale_tensor_alignment; + blocks_Y_rowwise = round_up_to_nearest_multiple(blocks_Y_rowwise, align); + blocks_X_rowwise = round_up_to_nearest_multiple(blocks_X_rowwise, align); + blocks_Y_colwise = round_up_to_nearest_multiple(blocks_Y_colwise, align); + blocks_X_colwise = round_up_to_nearest_multiple(blocks_X_colwise, align); + } +#else const size_t blocks_Y_rowwise = round_up_to_nearest_multiple(unpadded_blocks_Y_rowwise, scale_tensor_alignment_Y_rowwise); const size_t blocks_X_rowwise = round_up_to_nearest_multiple(unpadded_blocks_X_rowwise, @@ -399,6 +418,7 @@ void performTest_x1_swizzled(const size_t rows, scale_tensor_alignment_Y_colwise); const size_t blocks_X_colwise = round_up_to_nearest_multiple(unpadded_blocks_X_colwise, scale_tensor_alignment_X_colwise); +#endif const size_t blocks_num_rowwise = blocks_Y_rowwise * blocks_X_rowwise; const size_t blocks_num_colwise = blocks_Y_colwise * blocks_X_colwise; diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu index c5e420cfb..8868bdc28 100644 --- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu @@ -309,6 +309,15 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors, std::vector single_shape = {M, K}; std::vector scale_shape_vec = {per_tensor_scales_first_dim[t], per_tensor_scales_last_dim[t]}; +#ifdef __HIP_PLATFORM_AMD__ + // gfx1250's single-tensor dequantize shape check expects scales padded to a multiple of 4; + // the compact buffer is unchanged (the kernel reads the compact stride from cols). + if (getDeviceComputeCapability() == 125) { + const size_t align = mxfp8_gfx1250_scale_tensor_alignment; + scale_shape_vec = {round_up_to_nearest_multiple(scale_shape_vec[0], align), + round_up_to_nearest_multiple(scale_shape_vec[1], align)}; + } +#endif TensorWrapper input_w(NVTE_MXFP8_1D_SCALING); if (rowwise) { diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index 2df5dd51f..17187c9e5 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -27,6 +27,7 @@ #include "swizzle.cuh" #ifdef __HIP_PLATFORM_AMD__ +#include "../../util/cuda_runtime.h" #include "./rocm_vectorized_2d.cuh" #endif @@ -245,6 +246,43 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) #endif // #ifndef __HIP_PLATFORM_AMD__ } // namespace dequantize_kernel +#ifdef __HIP_PLATFORM_AMD__ +// Launch the single-tensor ROCm MXFP8 dequantize kernel. Also reused by the +// grouped dequantize path, which invokes it once per tensor. +inline void launch_dequantize_mxfp8_rocm(const void *input_dptr, void *output_dptr, + const e8m0_t *scales_ptr, DType input_dtype, + DType output_dtype, size_t rows, size_t cols, + size_t scale_dim_Y_colwise, size_t scale_dim_X_rowwise, + size_t scales_stride, bool with_gemm_swizzled_scales, + size_t mx_swizzle_padded_dim, cudaStream_t stream) { + using namespace dequantize_kernel; + const dim3 block(THREADS_PER_CHUNK); + const dim3 grid(DIVUP(cols, CHUNK_DIM_X), DIVUP(rows, CHUNK_DIM_Y)); + 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, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + !(cols % (32 * sizeof(OType))), IS_ALIGNED, + dequantize_mxfp8_kernel + <<>>( + reinterpret_cast(input_dptr), + reinterpret_cast(output_dptr), scales_ptr, rows, cols, + scales_stride, mx_swizzle_padded_dim););); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) +} +#endif // __HIP_PLATFORM_AMD__ + inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { using namespace dequantize_kernel; bool use_rowwise_scaling = input.has_data(); @@ -270,6 +308,13 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; +#ifdef __HIP_PLATFORM_AMD__ + // The ROCm kernel only decodes the gfx1250 MX pre-swizzle scale layout; other ROCm + // architectures emit the generic 128x4 GEMM swizzle, which this kernel cannot read. + NVTE_CHECK(!with_gemm_swizzled_scales || cuda::sm_arch() == 125, + "Dequantizing GEMM-swizzled MXFP8 scales is only supported on gfx1250."); +#endif + // TODO: Make more general const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1; @@ -310,6 +355,14 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const dim3 block(THREADS_PER_CHUNK); const dim3 grid(chunks_X, chunks_Y); +#ifdef __HIP_PLATFORM_AMD__ + // The MX pre-swizzle producer (swizzle_scaling_factors_mx) lays scales out using the actual + // padded scale-tensor dimension, so read it back with the same value: rowwise uses the scale + // rows, colwise uses the columnwise-scale cols. + const size_t mx_swizzle_padded_dim = + use_rowwise_scaling ? input.scale_inv.shape[0] : input.columnwise_scale_inv.shape[1]; +#endif + TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( scale_dim_Y_colwise, SCALE_DIM_Y, TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH( @@ -319,11 +372,14 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( output->dtype(), OType, #ifdef __HIP_PLATFORM_AMD__ + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, TRANSFORMER_ENGINE_SWITCH_CONDITION( !(cols % (32 * sizeof(OType))), IS_ALIGNED, - dequantize_mxfp8_kernel + dequantize_mxfp8_kernel <<>>(reinterpret_cast(input_data.dptr), reinterpret_cast(output->data.dptr), scales_ptr, - rows, cols, scales_stride);); // NOLINT(*) + rows, cols, scales_stride, mx_swizzle_padded_dim););); // NOLINT(*) #else // #ifdef __HIP_PLATFORM_AMD__ TRANSFORMER_ENGINE_SWITCH_CONDITION( with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, diff --git a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh index 45f79012d..bd077a437 100644 --- a/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh @@ -24,12 +24,22 @@ constexpr size_t THREADS_PER_CHUNK_X_COLWISE = CHUNK_DIM_X; 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) - dequantize_mxfp8_kernel(const IType *input_ptr, - OType *output_ptr, - const e8m0_t *const scales_ptr, const size_t rows, const size_t cols, - const size_t scales_stride) { +// MX pre-swizzle inverse index (matches swizzle_scaling_factors_mx in swizzle.cu): +// dst = (j / 4) * (padded_dim * 4) + i * 4 + (j % 4) +// The grouped-by-4 dimension is j; the contiguous dimension is i. +__device__ __forceinline__ size_t mx_preswizzle_scale_idx(size_t i, size_t j, size_t padded_dim) { + constexpr size_t GROUP = 4; // MX_PRESWIZZLE_GROUP_SIZE + return (j / GROUP) * (padded_dim * GROUP) + i * GROUP + (j % GROUP); +} + +template +__device__ __forceinline__ void + dequantize_mxfp8_chunk(const IType *input_ptr, + OType *output_ptr, + const e8m0_t *const scales_ptr, const size_t rows, const size_t cols, + const size_t scales_stride, const size_t mx_swizzle_padded_dim, + const int block_id_Y, const int block_id_X) { constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; constexpr bool USE_COLWISE_SCALING = SCALE_DIM_Y > 1; @@ -43,13 +53,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) DIVUP(SCALE_DIM_X, ELEMS_PER_THREAD); // 2 = 32 / 16 constexpr size_t VECTOR_WIDTH = IS_ALIGNED ? 8 : 16; - const int chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const int chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + 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 = blockIdx.y * SCALES_ROWWISE_PER_CHUNK_Y; - const int scales_rowwise_chunk_offset_X = blockIdx.x * SCALES_ROWWISE_PER_CHUNK_X; - const int scales_colwise_chunk_offset_Y = blockIdx.y * SCALES_COLWISE_PER_CHUNK_Y; - const int scales_colwise_chunk_offset_X = blockIdx.x * SCALES_COLWISE_PER_CHUNK_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; const int tid_rowwise_Y = threadIdx.x / THREADS_PER_CHUNK_X_ROWWISE; const int tid_rowwise_X = threadIdx.x % THREADS_PER_CHUNK_X_ROWWISE; @@ -88,7 +98,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) e8m0_t biased_exponent = static_cast(127); if (static_cast(scale_offset_Y) < scales_rows && static_cast(scale_offset_X) < scales_cols) { - const int scale_idx = scale_offset_Y * scales_stride + scale_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = USE_ROWWISE_SCALING + ? mx_preswizzle_scale_idx(scale_offset_Y, scale_offset_X, + mx_swizzle_padded_dim) + : mx_preswizzle_scale_idx(scale_offset_X, scale_offset_Y, + mx_swizzle_padded_dim); + } else { + scale_idx = static_cast(scale_offset_Y) * scales_stride + scale_offset_X; + } biased_exponent = scales_ptr[scale_idx]; } const float block_scale = ptx::exp2f(biased_exponent); @@ -142,3 +161,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + dequantize_mxfp8_kernel(const IType *input_ptr, + OType *output_ptr, + const e8m0_t *const scales_ptr, const size_t rows, const size_t cols, + const size_t scales_stride, const size_t mx_swizzle_padded_dim) { + dequantize_mxfp8_chunk(input_ptr, output_ptr, scales_ptr, rows, cols, + scales_stride, mx_swizzle_padded_dim, + blockIdx.y, blockIdx.x); +} +