From 25ba9bacbadbd3ad9943dcbb6d8b1d51b0fa9e1e Mon Sep 17 00:00:00 2001 From: tangcc1127 Date: Mon, 17 Aug 2026 12:43:30 +0800 Subject: [PATCH 1/2] [PyTorch] Share MXFP8 data tensor between rowwise and columnwise for 2D quantization For 2D MXFP8 quantization, the rowwise and columnwise FP8 data tensors are byte-identical since they originate from the same 32x32 block scales. This commit shares a single data buffer between the two representations, halving the FP8 weight memory footprint for 2D-quantized weights. The columnwise data tensor is reused as an alias of the rowwise data tensor in the C++ quantizer, the inner_tensor_specs paths, and the MXFP8Tensor constructor. The cuBLAS GEMM path already selects the appropriate pointer via the transA flag and handles the transpose internally, so no GEMM changes are needed. Signed-off-by: tangcc1127 --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 20 +++++++++++-------- transformer_engine/pytorch/csrc/quantizer.cpp | 7 ++++++- .../pytorch/tensor/mxfp8_tensor.py | 13 +++++++++++- 3 files changed, 30 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index ac9cc12356..cd71acf0bb 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -53,7 +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 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; @@ -509,9 +510,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 +875,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 +892,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 +909,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, From fd47bd919eb1d4463943ebb9f05e6d225c951c7e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:28:17 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index cd71acf0bb..6352789a36 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -53,9 +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 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) { + 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;