From f926c56f85cfda71db4e19548f1705e27176c184 Mon Sep 17 00:00:00 2001 From: tangcc1127 Date: Mon, 17 Aug 2026 12:43:30 +0800 Subject: [PATCH] [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 --- transformer_engine/pytorch/csrc/quantizer.cpp | 7 ++++++- transformer_engine/pytorch/tensor/mxfp8_tensor.py | 13 ++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) 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,