Skip to content

gfx1250 mxfp8: fix dequant #733

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
matthiasdiener wants to merge 1 commit into
base: dev
Choose a base branch
from
Draft

gfx1250 mxfp8: fix dequant #733

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions tests/cpp/operator/test_dequantize_mxfp8.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -391,14 +391,26 @@ 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;

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,
scale_tensor_alignment_X_rowwise);
const size_t blocks_Y_colwise = round_up_to_nearest_multiple(unpadded_blocks_Y_colwise,
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);
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);

#ifdef __HIP_PLATFORM_AMD__
// 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);
}
#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;
Expand Down
12 changes: 12 additions & 0 deletions tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,18 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
round_up_to_nearest_multiple(unpadded_scales_X, scale_tensor_alignment_X_colwise);
}

#ifdef __HIP_PLATFORM_AMD__
// gfx1250 pads MXFP8 scales to a multiple of 4 in both dims; match what the common-code
// scale-shape check (CheckScaleTensorShape) expects for the per-tensor member tensors.
if (getDeviceComputeCapability() == 125) {
const size_t align = mxfp8_gfx1250_scale_tensor_alignment;
per_tensor_scales_first_dim[t] =
round_up_to_nearest_multiple(per_tensor_scales_first_dim[t], align);
per_tensor_scales_last_dim[t] =
round_up_to_nearest_multiple(per_tensor_scales_last_dim[t], align);
}
#endif

const size_t tensor_scales = per_tensor_scales_first_dim[t] * per_tensor_scales_last_dim[t];
total_scales += tensor_scales;
per_tensor_scales_offset[t + 1] = total_scales;
Expand Down
52 changes: 50 additions & 2 deletions transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,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<IType, OType, SCALE_DIM_Y, SCALE_DIM_X, IS_ALIGNED,
WITH_GEMM_SWIZZLED_SCALES>
<<<grid, block, 0, stream>>>(
reinterpret_cast<const IType *>(input_dptr),
reinterpret_cast<OType *>(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();
Expand Down Expand Up @@ -310,6 +347,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(
Expand All @@ -319,11 +364,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<IType, OType, SCALE_DIM_Y, SCALE_DIM_X, IS_ALIGNED>
dequantize_mxfp8_kernel<IType, OType, SCALE_DIM_Y, SCALE_DIM_X, IS_ALIGNED,
WITH_GEMM_SWIZZLED_SCALES>
<<<grid, block, 0, stream>>>(reinterpret_cast<const IType *>(input_data.dptr), reinterpret_cast<OType *>(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,
Expand Down
58 changes: 45 additions & 13 deletions transformer_engine/common/cast/mxfp8/rocm_dequantize_mxfp8.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename IType, typename OType, size_t SCALE_DIM_Y, size_t SCALE_DIM_X, bool IS_ALIGNED>
__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 <typename IType, typename OType, size_t SCALE_DIM_Y, size_t SCALE_DIM_X, bool IS_ALIGNED,
bool WITH_GEMM_SWIZZLED_SCALES>
__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;

Expand All @@ -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;
Expand Down Expand Up @@ -88,7 +98,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
e8m0_t biased_exponent = static_cast<e8m0_t>(127);
if (static_cast<size_t>(scale_offset_Y) < scales_rows &&
static_cast<size_t>(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<size_t>(scale_offset_Y) * scales_stride + scale_offset_X;
}
biased_exponent = scales_ptr[scale_idx];
}
const float block_scale = ptx::exp2f(biased_exponent);
Expand Down Expand Up @@ -142,3 +161,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
}
}

template <typename IType, typename OType, size_t SCALE_DIM_Y, size_t SCALE_DIM_X, bool IS_ALIGNED,
bool WITH_GEMM_SWIZZLED_SCALES>
__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<IType, OType, SCALE_DIM_Y, SCALE_DIM_X, IS_ALIGNED,
WITH_GEMM_SWIZZLED_SCALES>(input_ptr, output_ptr, scales_ptr, rows, cols,
scales_stride, mx_swizzle_padded_dim,
blockIdx.y, blockIdx.x);
}

Loading