[bhargavchintam]

Description

MXFP8TensorStorage.export_components() returns the FP8 values and E8M0 scales of an MXFP8 tensor in a layout that does not depend on how Transformer Engine stores them. It is for integrations that hand MXFP8 weights to another library, for example RL weight sync to an inference engine. Today NeMo RL #3908, Megatron-Bridge #5917, Miles and verl #8129 each build this by hand from _rowwise_data and _rowwise_scale_inv.

It returns MXFP8Components(data, scale_inv):

• data is _rowwise_data viewed as torch.float8_e4m3fn or torch.float8_e5m2, with the tensor's shape.
• scale_inv holds E8M0 biased exponents as torch.uint8, with shape (*shape[:-1], ceil(shape[-1] / 32)), without padding or GEMM swizzle. Since Improve the dimension checks for the FP8 recipes #2894 the last block can be shorter than 32 values, and it keeps its own scale.

The values are always a view of the tensor's data, so they change when the tensor is updated in place, for example by an optimizer step. Clone them to keep a snapshot. The scales are a view too, except in two cases:

• they are padded in the last dimension (ceil(K/32) is not a multiple of 4), or
• they are swizzled for GEMM (optimize_for_gemm=True).

In those cases they are copied, and the swizzle is undone with one permute in PyTorch, so no new kernel is needed. A tensor without row-scaled data raises ValueError. It works the same for quantized parameters and for tensors returned by a quantizer, with no dequantize or requantize.

On #3644 I asked whether a standalone function that also quantizes a BF16 weight first would be preferred. This PR uses a method. I expect the export logic itself would stay the same in a function, but I have not tried that version.

Fixes #3644

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py: the MXFP8Components NamedTuple and MXFP8TensorStorage.export_components(). MXFP8Tensor inherits it.
• tests/pytorch/test_quantized_tensor.py, 29 cases in all:
  • 4 tests on hand-built storages in TestQuantizedTensor, which need no MXFP8 kernels: layout (including a partial last block, (100, 48)), views, empty tensor, and no row-scaled data.
  • 4 tests in TestMXFP8Tensor:
    • the export decodes to dequantize() exactly for E4M3, E5M2 and 2D quantization;
    • a partial last block (K = 48 and K = 16) keeps its own scale and decodes to dequantize() exactly;
    • optimize_for_gemm=True gives the same export, including for (100, 48);
    • a te.Linear weight from quantized_model_init exports correctly after a forward and backward pass, and after a FusedAdam step the views exported before the step hold exactly the rowwise quantization of the new FP32 master weights.
  • 2 test helpers:
    • _swizzle_mxfp8_rowwise_scale_inv builds the cuBLAS-swizzled scale layout from the kernel's byte-index formula, independently of the code under test;
    • _decode_mxfp8_components decodes the exported values and scales in plain PyTorch.
• docs/features/low_precision_training/mxfp8/mxfp8.rst: an "Exporting values and scales" section.

Not covered here: column-scaled data, per-expert export from GroupedTensor, and the other recipes.

Testing

The GPU runs use one NVIDIA B200 (compute capability 10.0, driver 580.173.02) in nvcr.io/nvidia/pytorch:26.10-py3 (PyTorch 2.15.0a0+875d815502.nv26.10, CUDA 13.4). TE was built from current main (c1c2d4d) with NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS=100a MAX_JOBS=20 python -m pip install -e . -v --no-build-isolation, which gives transformer_engine-2.22.0.dev0+c1c2d4d2 and is_mxfp8_available: (True, ''). This PR's diff was then applied on top, and the tested tree is identical to this branch's head (tree 97416833).

The new tests:

$ python -m pytest -v -p no:cacheprovider tests/pytorch/test_quantized_tensor.py -k export_components
================ 29 passed, 602 deselected, 1 warning in 3.55s =================

The 29 cases

tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape0-False] PASSED [  3%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape0-True] PASSED [  6%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape1-False] PASSED [ 10%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape1-True] PASSED [ 13%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape2-False] PASSED [ 17%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape2-True] PASSED [ 20%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape3-False] PASSED [ 24%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_layout[shape3-True] PASSED [ 27%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_views PASSED [ 31%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_empty PASSED [ 34%]
tests/pytorch/test_quantized_tensor.py::TestQuantizedTensor::test_mxfp8_export_components_requires_rowwise_data PASSED [ 37%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims0-7-False] PASSED [ 41%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims0-8-False] PASSED [ 44%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims0-7-True] PASSED [ 48%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims1-7-False] PASSED [ 51%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims1-8-False] PASSED [ 55%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims1-7-True] PASSED [ 58%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims2-7-False] PASSED [ 62%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims2-8-False] PASSED [ 65%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_matches_dequantize[dims2-7-True] PASSED [ 68%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_partial_block[dims0-7] PASSED [ 72%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_partial_block[dims0-8] PASSED [ 75%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_partial_block[dims1-7] PASSED [ 79%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_partial_block[dims1-8] PASSED [ 82%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_unswizzles_gemm_scales[dims0] PASSED [ 86%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_unswizzles_gemm_scales[dims1] PASSED [ 89%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_unswizzles_gemm_scales[dims2] PASSED [ 93%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_unswizzles_gemm_scales[dims3] PASSED [ 96%]
tests/pytorch/test_quantized_tensor.py::TestMXFP8Tensor::test_mxfp8_export_components_of_quantized_weight PASSED [100%]

With scale_cols put back to K // 32, the 6 partial-block cases fail (both (100, 48) layout cases and all 4 test_mxfp8_export_components_partial_block cases), and the other 23 pass.

The suites around this code, on main and then with this change:

$ python -m pytest -q -p no:cacheprovider tests/pytorch/test_quantized_tensor.py tests/pytorch/mxfp8 tests/pytorch/test_mxfp8_2d_quantize.py
# main (c1c2d4d2)
1231 passed, 1235 skipped, 14 warnings in 102.74s (0:01:42)
# this branch
1260 passed, 1235 skipped, 14 warnings in 101.72s (0:01:41)

That is 29 more passed, with the same skips and the same warnings.

Lint and the license check:

$ TE_PATH=$PWD PYTHON_ONLY=1 bash qa/L0_pytorch_lint/test.sh
Your code has been rated at 10.00/10

pylint reports one informational message, I1101 in transformer_engine/pytorch/attention/fused_mla_q_uproj.py:207, a file this PR does not touch.

$ pre-commit run --files transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py tests/pytorch/test_quantized_tensor.py docs/features/low_precision_training/mxfp8/mxfp8.rst
check for merge conflicts................................................Passed
check for added large files..............................................Passed
fix end of files.........................................................Passed
trim trailing whitespace.................................................Passed
Format python code.......................................................Passed
clang-format.........................................(no files to check)Skipped
Check minimum Python version.............................................Passed
$ TE_PATH=$PWD bash qa/L0_license/test.sh; echo "exit=$?"
exit=0

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings: 14 pytest warnings on main and on this branch for the three suites, and pylint rates 10.00/10.
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes: the three suites around this change pass on a B200 (1260 passed, 1235 skipped, the same skips as main). I did not run the full qa/L0_pytorch_unittest launcher.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issues were found.

Summary

Adds MXFP8TensorStorage.export_components() to return FP8 values and compact, unswizzled E8M0 scales.

• MXFP8 tensors export FP8 values and compact scales for other libraries.

The earlier unnumbered test-coverage finding is fixed: the test now checks saved views after a FusedAdam step.

_{Reviews (3) · Last reviewed commit: "[PyTorch] Export a scale for a partial l..." · Reviewed by Greptile}

[bhargavchintam]

Merged main after #2894. Now that MXFP8 allows partial blocks, a last block of 16 values keeps its own scale, so export_components() returns ceil(K / 32) scale columns instead of K // 32 (ed0362a). New tests cover K = 48 and K = 16, including rows that are not a multiple of 32 and GEMM-swizzled scales. On a B200 with current main, the 29 export cases pass, and the three suites around this code go from 1231 to 1260 passed with the same skips. The description has the output.