diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index ac9cc12356..6352789a36 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -53,8 +53,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const __grid_constant__ CUtensorMap tensor_map_output_colwise, e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, const float *noop, float *const dbias_workspace, float *const amax_ptr, - const size_t rows, const size_t cols, const size_t scale_stride_rowwise, - const size_t scale_stride_colwise) { + const bool skip_colwise_data_write, const size_t rows, const size_t cols, + const size_t scale_stride_rowwise, const size_t scale_stride_colwise) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; @@ -509,9 +509,11 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) global_offset_Y, reinterpret_cast(&out_rowwise_data_sh[buff_offset])); } if constexpr (COLWISE_SCALING) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_colwise), global_offset_X, - global_offset_Y, reinterpret_cast(&out_colwise_data_sh[buff_offset])); + if (!skip_colwise_data_write) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_colwise), global_offset_X, + global_offset_Y, reinterpret_cast(&out_colwise_data_sh[buff_offset])); + } } // Create a "bulk async-group" out of the previous bulk copy operation. @@ -872,7 +874,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, kernel<<>>( tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + workspace_ptr, amax_ptr, false, rows, cols, scale_stride_rowwise, scale_stride_colwise); }); break; @@ -889,7 +891,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, kernel<<>>( tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + workspace_ptr, amax_ptr, false, rows, cols, scale_stride_rowwise, scale_stride_colwise); }); break; @@ -906,8 +908,9 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, kernel<<>>( tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + workspace_ptr, amax_ptr, + use_2d_quantization && output->data == output->columnwise_data, rows, cols, + scale_stride_rowwise, scale_stride_colwise); }); break; } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4704261790..33d9e4e5d6 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1531,8 +1531,13 @@ std::pair MXFP8Quantizer::create_tensor( if (columnwise_usage) { const std::vector scale_inv_shape_int64(columnwise_scale_inv_shape.begin(), columnwise_scale_inv_shape.end()); - columnwise_data_tensor = at::empty(shape_int64, uint8_tensor_opts); columnwise_scale_inv_tensor = at::empty(scale_inv_shape_int64, uint8_tensor_opts); + if (with_2d_quantization && rowwise_usage) { + // 2D quantization: rowwise and columnwise data are identical, share the buffer + columnwise_data_tensor = rowwise_data_tensor; + } else { + columnwise_data_tensor = at::empty(shape_int64, uint8_tensor_opts); + } } // Convert tensors to Python diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 54cb281bd6..7567d6be96 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -87,7 +87,9 @@ def inner_tensor_specs( torch.uint8, ) if self.columnwise_usage: - specs["_columnwise_data"] = (shape, torch.uint8) + # 2D quantization: data is identical, reuse rowwise_data instead of allocating a copy + if not (self.with_2d_quantization and self.rowwise_usage): + specs["_columnwise_data"] = (shape, torch.uint8) specs["_columnwise_scale_inv"] = ( tuple(self.get_scale_shape(shape, columnwise=True)), torch.uint8, @@ -263,6 +265,15 @@ def __new__( with_gemm_swizzled_scales: bool, **kwargs, ): + # 2D quantization: columnwise data is identical to rowwise, alias it + if ( + columnwise_data is None + and rowwise_data is not None + and quantizer is not None + and getattr(quantizer, "with_2d_quantization", False) + and getattr(quantizer, "columnwise_usage", False) + ): + columnwise_data = rowwise_data return super().__new__( cls, rowwise_data,