[PyTorch] Share MXFP8 data tensor between rowwise and columnwise for 2D quantization

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<uint64_t *>(&out_rowwise_data_sh[buff_offset]));
       }
       if constexpr (COLWISE_SCALING) {
-        ptx::cp_async_bulk_tensor_2d_shared_to_global(
-            reinterpret_cast<const uint64_t *>(&tensor_map_output_colwise), global_offset_X,
-            global_offset_Y, reinterpret_cast<uint64_t *>(&out_colwise_data_sh[buff_offset]));
+        if (!skip_colwise_data_write) {
+          ptx::cp_async_bulk_tensor_2d_shared_to_global(
+              reinterpret_cast<const uint64_t *>(&tensor_map_output_colwise), global_offset_X,
+              global_offset_Y, reinterpret_cast<uint64_t *>(&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<<<grid, block_size, dshmem_size, stream>>>(
                         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<<<grid, block_size, dshmem_size, stream>>>(
                         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<<<grid, block_size, dshmem_size, stream>>>(
                         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;
                 }

transformer_engine/pytorch/csrc/quantizer.cpp

@@ -1531,8 +1531,13 @@ std::pair<TensorWrapper, py::object> MXFP8Quantizer::create_tensor(
   if (columnwise_usage) {
     const std::vector<int64_t> 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

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,