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

[TangChangcheng]

Summary

For 2D MXFP8 quantization, the rowwise and columnwise FP8 data tensors are byte-identical since they originate from the same 32x32 block scales. Currently, two separate [M, K] uint8 buffers are allocated and written with identical data.

This PR shares a single data buffer between the two representations, halving the FP8 weight memory footprint for 2D-quantized weights.

Changes

• C++ quantizer (quantizer.cpp): When with_2d_quantization && rowwise_usage && columnwise_usage, reuse rowwise_data_tensor as columnwise_data_tensor instead of allocating a separate buffer.
• Python inner_tensor_specs (mxfp8_tensor.py): Skip _columnwise_data allocation when 2D and rowwise is already enabled.
• MXFP8Tensor.__new__ (mxfp8_tensor.py): When columnwise_data is None and 2D quantization is active, alias it to rowwise_data.

No GEMM or kernel changes are needed: cuBLAS already selects the appropriate pointer via the transA flag and handles the transpose internally.

Related

• Follow-up to Add 2d quant for mxfp8 #2634 (Add 2d quant for mxfp8)

[greptile-apps]

Greptile Summary

The PR reduces dual-direction 2D MXFP8 weight storage by sharing the byte-identical rowwise data buffer with the columnwise representation.

• Aliases rowwise and columnwise data in the native C++ tensor allocator.
• Omits the redundant columnwise allocation from Python allocation specifications.
• Restores the alias when constructing dual-use Python MXFP8 tensor wrappers.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/quantizer.cpp	Shares the rowwise data allocation with the columnwise representation for dual-use 2D MXFP8 tensors while retaining direction-specific scales.
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Aligns Python allocation and wrapper construction with the shared-data representation, including the columnwise-usage guard that resolves the prior rowwise-only availability issue.

_{Reviews (2): Last reviewed commit: "[PyTorch] Share MXFP8 data tensor betwee..." | Re-trigger Greptile}

[ptrendx]

In general I agree with this optimization (although I would like this change to also include a change to the kernel to skip writing the transposed data if the pointers for columnwise and rowwise tensors are the same).

One problem is that currently Megatron does not have a good handling of the 2D MXFP8 weights and so enabling this optimization unconditionally could possibly make it unusable there. @kunlunl @zhongbozhu could you comment on this - if we still technically have both rowwise and columnwise tensors on the pyTorch side, but they are backed by the same actual buffer, is the distributed optimizer still going to work or do we need more work for that?