[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 the memory footprint of two-direction MXFP8 2D weights by sharing their byte-identical data tensor while retaining separate directional scales.

• Reuses the rowwise allocation as columnwise data when both usages are enabled.
• Avoids the redundant columnwise CUDA data write when both descriptors point to the shared buffer.
• Restricts Python-side aliasing to quantizers that actually request columnwise usage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the earlier rowwise-only issue is fixed because both columnwise buffers remain absent and the alias now requires columnwise usage.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh	Suppresses the duplicate columnwise data store only when 2D quantization uses identical output pointers.
transformer_engine/pytorch/csrc/quantizer.cpp	Shares the rowwise allocation with columnwise storage for two-direction 2D MXFP8 tensors while keeping directional scales separate.
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Aligns Python allocation and construction with shared two-direction data while leaving rowwise-only tensors without columnwise storage.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Input["FP16/BF16 weight"] --> Quant["2D MXFP8 quantization"]
  Quant --> Data["Shared FP8 data buffer"]
  Quant --> RowScale["Rowwise scale-inverse"]
  Quant --> ColScale["Columnwise scale-inverse"]
  Data --> RowRep["Rowwise representation"]
  RowScale --> RowRep
  Data --> ColRep["Columnwise representation"]
  ColScale --> ColRep

Loading

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | 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?

[TangChangcheng]

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).

The latest commit included the change to the kernel to skip writing the duplicate data.