[PyTorch][Common] Add MXFP8 and NVFP4 QDQ for mixed-precision GEMM

docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst

@@ -288,6 +288,34 @@ Note that in the second example the columnwise direction is high precision but
 holds the value reconstructed from MXFP8, not the original input — see
 `Choosing the columnwise source`_ above.
 
+Output dtype and workspace updates
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Output-dtype handling is opt-in through ``Quantizer.supports_output_dtype``,
+which defaults to ``False``. Identity and Hybrid opt in; native quantizers
+continue to ignore output-dtype requests, even inside a custom recipe.
+
+``IdentityQuantizer.quantize(x, dtype=...)`` selects the held output dtype
+without changing the quantizer's default or adding an autograd edge to its
+stored data. ``HybridQuantizer`` passes a compute-dtype request to its children;
+children that support output-dtype binding retain it for later reconstruction.
+Native quantizers keep their existing behavior, including ignoring this
+argument. Their encodings and default dequantization metadata are unchanged.
+
+Hybrid's ``rowwise_dequantized`` source uses the source tensor's dtype for
+native children, preserving existing native compositions, including internal
+storage. For children that support output-dtype handling, it uses
+``torch.promote_types`` on the source and stored dtypes to avoid narrowing
+decoded forward values. This is the input dtype of the columnwise quantizer,
+not necessarily the backward GEMM dtype. Widening a value already rounded into
+BF16 does not recover the discarded precision.
+
+In-place Identity updates preserve the destination buffer, dtype, shape and
+device. An explicit conflicting ``dtype`` with ``out=`` raises an error for
+Identity and Hybrid. Module weight caches rebuild these workspaces when an
+explicit compute dtype changes. ``HybridQuantizer.make_empty(..., dtype=...)``
+binds the requested dtype on participating custom children just as quantize does.
+
 Example: one format per GEMM
 ----------------------------
 
@@ -362,3 +390,116 @@ API reference
 See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``,
 ``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types.
 See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``.
+
+Experimental QDQ
+----------------
+
+QDQ quantizers model low-precision rounding by quantizing operands and then
+dequantizing them into high-precision ``IdentityTensor`` storage. The decoded
+operands use high-precision GEMMs; QDQ models quantization error for convergence
+experiments, not mixed-format hardware performance.
+
+Forward and backward examples
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The following factories support ``Linear`` and ``GroupedLinear`` with discrete
+expert parameters. With BF16 module compute precision:
+
+* ``nvfp4_qdq_weight_fwd_bf16_bwd_factory`` models ``X[BF16] x W[NVFP4-QDQ]``
+  forward and uses BF16 backward. Backward consumes the original inputs and
+  the decoded forward QDQ weights.
+* ``mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory`` models
+  ``X[MXFP8-QDQ] x W[NVFP4-QDQ]`` forward and uses native MXFP8 backward.
+  Input QDQ uses block-32 E4M3/E8M0;
+  weight QDQ uses block-16 E2M1/E4M3 with a tensor scale. Backward quantizes
+  the original inputs and the decoded forward QDQ weights to MXFP8.
+
+For example::
+
+    from transformer_engine.common.recipe import CustomRecipe
+    from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import (
+        nvfp4_qdq_weight_fwd_bf16_bwd_factory,
+        mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory,
+    )
+
+    recipe = CustomRecipe(qfactory=mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory)
+
+To use BF16 inputs while retaining MXFP8 backward, replace the second factory's
+input rowwise ``MXFP8QDQQuantizer`` with ``IdentityQuantizer``. The input
+columnwise quantizer and its ``columnwise_source="original"`` stay unchanged.
+Both examples use ``columnwise_source="rowwise_dequantized"`` for weights.
+Other module roles use ``IdentityQuantizer``. High-precision arithmetic follows
+the module compute dtype; select BF16 compute for the BF16 example.
+``CustomRecipe`` does not require FP8 hardware to enter ``te.autocast``.
+Its quantizers and GEMMs must support the target device and TE build;
+recipes using only QDQ and high-precision GEMMs can run on Ampere.
+
+Configure ``backend`` or ``nvfp4_options`` with ``functools.partial``. Existing
+Hybrid limitations still apply, and RL convergence has not been validated.
+
+Backends and custom QDQ formats
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Output defaults to the source dtype; a module compute-dtype override changes
+only the decoded dtype. For ``NVFP4QDQQuantizer``, ``backend="auto"`` fuses
+plain 1D E2M1/E4M3 QDQ with tensor or per-row outer scaling on
+GPUs with compute capability 8.0 or higher and a TE build with CUDA 12.8 or
+higher for contiguous, 16-byte-aligned BF16/FP16 tensors with matching output
+dtype and flattened matrix dimensions
+divisible by 16, with ordinary fast math disabled. Other native-valid modes,
+including 4over6, 2D scaling, distributed amax reduction, and mixed dtypes,
+use native ``quantize().dequantize()`` where its kernels are supported.
+``backend="reference"`` forces that composition; ``backend="fused"`` rejects
+unsupported cases. Current weight amax is computed when quantization runs.
+
+For per-token outer scaling, wrap
+``NVFP4Quantizer(row_scaled_nvfp4=True)`` in ``NVFP4QDQQuantizer``. Each flattened
+row uses its own outer amax, while each 16-value block retains its E4M3 scale.
+The fused row-scaled path computes row amax and QDQ in one kernel, without
+intermediate amax, packed-value, or scale buffers. The tensor-scaled fused path
+retains its separate global amax reduction. Both paths preserve native Q+DQ
+rounding and support the same matching BF16/FP16 input and output dtypes.
+
+``custom_recipes.qdq.MXFP8QDQQuantizer`` provides the corresponding rowwise
+E4M3/block-32 E8M0 QDQ operation. Its ``reference`` backend calls native
+``MXFP8Quantizer.quantize().dequantize()``. Its ``fused`` backend supports
+contiguous, 16-byte-aligned BF16/FP16 tensors on GPUs with compute capability
+8.0 or higher, with a matching output dtype, rows divisible by 32, and ordinary
+fast math disabled. ``auto`` falls back to the reference when these requirements
+are not met.
+The fused kernel preserves signed zero. Native specialized FP16 quantization
+may canonicalize negative zero; the decoded values remain numerically equal.
+
+Build Transformer Engine for the target GPU architecture; a build containing
+only ``sm_100a`` kernels cannot run those kernels on other architectures.
+NVFP4 QDQ requires CUDA Toolkit 12.8 or later for the FP4 conversion types.
+On Ampere and Hopper, CUDA emulates the conversions inside the fused kernel.
+The existing native packed Q+DQ reference kernels require compute capability
+10.0 or higher. On older GPUs, use ``auto`` or ``fused`` with matching BF16/FP16
+and the plain fused configuration described above; a request requiring the
+native reference path raises an error. In particular, mixed dtypes, 4over6,
+2D scaling, and distributed amax reduction are not supported there.
+QDQ's conversion requirements are independent of native low-precision GEMM
+support. Recipes using native MXFP8 backward additionally require
+``is_mxfp8_available()``. Numerical and performance validation on a particular
+GPU does not establish performance on other architectures.
+
+Both QDQ classes return high-precision storage with straight-through
+gradients. Use them as rowwise children of ``HybridQuantizer`` to select the
+backward format independently. ``columnwise_source="rowwise_dequantized"``
+requantizes the forward reconstruction for backward; ``"original"`` uses
+the original input. QDQ models quantization error for convergence studies;
+the high-precision GEMM does not model mixed-format hardware performance.
+
+Both formats inherit the experimental ``custom_recipes.qdq.QDQQuantizer``
+base, which supplies allocation, IdentityTensor wrapping, dtype overrides,
+update validation, and straight-through gradients. Custom QDQ subclasses
+implement ``selected_backend(tensor, dtype)``, ``_reference(tensor, dtype)``,
+and ``_compute(tensor, output, noop_flag=None, *, backend=None)``. The reference
+must quantize the original input before decoding to the requested dtype.
+Compute must write into the supplied output, preserving it when the no-op flag
+is one, and honor an explicitly selected backend. Subclasses must also provide
+format-appropriate ``copy()`` and ``is_requantization_safe()`` implementations.
+Storage hooks should not build autograd graphs; Python implementations can
+use ``torch.no_grad()`` since the wrapper supplies the straight-through gradient.
+The base and its extension hooks are experimental and subject to change.

qa/L1_pytorch_hybrid_distributed_unittest/test.sh

@@ -25,6 +25,8 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2_hybrid.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "hybrid" || test_fail "hybrid test_torch_fsdp2.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py"
 
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_nvfp4_qdq_distributed.xml $TE_PATH/tests/pytorch/distributed/test_nvfp4_qdq.py || test_fail "test_nvfp4_qdq.py"
+
 if [ "$RET" -ne 0 ]; then
     echo "Error in the following test cases:$FAILED_CASES"
     exit 1

tests/cpp/operator/CMakeLists.txt

@@ -22,6 +22,7 @@ add_executable(test_operator
                test_dequantize_mxfp8_grouped.cu
                test_fused_group_requantize_mxfp8.cu
                test_dequantize_float8blockwise_grouped.cu
+               test_nvfp4_qdq.cu
                test_dequantize_nvfp4.cu
                test_transpose.cu
                test_cast_transpose.cu

tests/cpp/operator/test_nvfp4_qdq.cu

@@ -0,0 +1,95 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#include <cmath>
+#include <cstring>
+#include <vector>
+
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+#include <transformer_engine/cast.h>
+
+#include "../test_common.h"
+
+using namespace transformer_engine;
+using namespace test;
+
+namespace {
+template <typename T>
+void check_qdq(size_t rows, size_t cols, bool row_scaled = false) {
+  cudaDeviceProp prop;
+  ASSERT_EQ(cudaGetDeviceProperties(&prop, 0), cudaSuccess);
+  if (prop.major < 10) GTEST_SKIP() << "Native packed Q+DQ reference requires SM10.0+";
+  const auto dtype = TypeInfo<T>::dtype;
+  const std::vector<size_t> shape{rows, cols};
+  Tensor input("input", shape, dtype);
+  Tensor packed("packed", shape, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING,
+                DType::kFloat8E4M3);
+  Tensor reference("reference", shape, dtype);
+  Tensor output("output", shape, dtype);
+  Tensor amax("amax", std::vector<size_t>{1}, DType::kFloat32);
+  Tensor noop("noop", std::vector<size_t>{1}, DType::kFloat32);
+  fillCase<fp32>(&input, InputsFillCase::uniform);
+  input.to_cpu();
+  float maximum = 0;
+  for (size_t i = 0; i < rows * cols; ++i) {
+    maximum = std::max(maximum, std::abs(static_cast<float>(input.rowwise_cpu_dptr<T>()[i])));
+  }
+  packed.set_amax(maximum);
+  if (row_scaled) packed.set_row_scaled_nvfp4(true);
+  ASSERT_EQ(cudaMemcpy(amax.rowwise_dptr(), &maximum, sizeof(float), cudaMemcpyHostToDevice),
+            cudaSuccess);
+  QuantizationConfigWrapper config;
+  nvte_quantize_v2(input.data(), packed.data(), config, 0);
+  nvte_dequantize(packed.data(), reference.data(), 0);
+  if (row_scaled) {
+    nvte_nvfp4_qdq_row_scaled(input.data(), output.data(), nullptr, 0);
+  } else {
+    nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), nullptr, 0);
+  }
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+  reference.to_cpu();
+  output.to_cpu();
+  EXPECT_EQ(std::memcmp(reference.rowwise_cpu_dptr<T>(), output.rowwise_cpu_dptr<T>(),
+                        rows * cols * sizeof(T)), 0);
+  const float one = 1;
+  ASSERT_EQ(cudaMemcpy(noop.rowwise_dptr(), &one, sizeof(float), cudaMemcpyHostToDevice),
+            cudaSuccess);
+  ASSERT_EQ(cudaMemset(input.rowwise_dptr(), 0, rows * cols * sizeof(T)), cudaSuccess);
+  if (row_scaled) {
+    nvte_nvfp4_qdq_row_scaled(input.data(), output.data(), noop.data(), 0);
+  } else {
+    nvte_nvfp4_qdq(input.data(), output.data(), amax.data(), noop.data(), 0);
+  }
+  ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
+  output.to_cpu();
+  EXPECT_EQ(std::memcmp(reference.rowwise_cpu_dptr<T>(), output.rowwise_cpu_dptr<T>(),
+                        rows * cols * sizeof(T)), 0);
+}
+}  // namespace
+
+TEST(NVFP4QDQ, BF16) {
+  check_qdq<bf16>(32, 32);
+  check_qdq<bf16>(96, 160);
+  check_qdq<bf16>(512, 1024);
+}
+
+TEST(NVFP4QDQ, FP16) {
+  check_qdq<fp16>(32, 32);
+  check_qdq<fp16>(96, 160);
+}
+
+TEST(NVFP4QDQ, RowScaledBF16) {
+  check_qdq<bf16>(32, 32, true);
+  check_qdq<bf16>(96, 160, true);
+  check_qdq<bf16>(128, 4096, true);
+}
+
+TEST(NVFP4QDQ, RowScaledFP16) {
+  check_qdq<fp16>(32, 32, true);
+  check_qdq<fp16>(96, 160, true);
+  check_qdq<fp16>(32, 16384, true);
+}

tests/pytorch/distributed/run_nvfp4_qdq.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Two-rank NVFP4 QDQ amax reduction and native reference parity."""
+
+import os
+
+import torch
+import torch.distributed as dist
+from transformer_engine.pytorch.custom_recipes.qdq import NVFP4QDQQuantizer
+from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
+import transformer_engine_torch as tex
+
+
+def assert_bits(actual, expected):
+    assert actual.dtype == expected.dtype
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
+
+
+def run_case(options, source_dtype):
+    rank = dist.get_rank()
+    torch.manual_seed(123 + rank)
+    x = torch.randn(128, 256, device="cuda", dtype=source_dtype) * (1 if rank == 0 else 17)
+    output_dtype = torch.bfloat16
+    local_amax = x.float().abs().amax().reshape(1)
+    global_amax = local_amax.clone()
+    dist.all_reduce(global_amax, op=dist.ReduceOp.MAX)
+    assert dist.get_world_size() == 2
+    if rank == 0:
+        assert global_amax.item() > local_amax.item()
+    native = NVFP4Quantizer(rowwise=True, columnwise=False, **options)
+    native.internal = True
+    q = NVFP4QDQQuantizer(native)
+    q.amax_reduction_group = dist.group.WORLD
+    q.with_amax_reduction = True
+    assert q.selected_backend(x, output_dtype) == "reference"
+    out = q.quantize(x, dtype=output_dtype)
+    expected = tex.nvfp4_quantize_with_amax(x, native, global_amax, global_amax).dequantize(
+        dtype=output_dtype
+    )
+    assert_bits(out.dequantize(), expected)
+    saved = out._get_quantizer()
+    assert not saved.with_amax_reduction and saved.amax_reduction_group is None
+    assert q.with_amax_reduction and q.amax_reduction_group is dist.group.WORLD
+    pointer = out._hp_data.data_ptr()
+    x.mul_(2)
+    global_amax.mul_(2)
+    q.update_quantized(x, out)
+    expected = tex.nvfp4_quantize_with_amax(x, native, global_amax, global_amax).dequantize(
+        dtype=output_dtype
+    )
+    assert_bits(out.dequantize(), expected)
+    assert out._hp_data.data_ptr() == pointer
+    # Saved workspaces intentionally update locally, without persistent collectives.
+    out.quantize_(x)
+    assert_bits(out.dequantize(), native(x).dequantize(dtype=output_dtype))
+
+
+def main():
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+    dist.init_process_group("nccl")
+    try:
+        count = 0
+        for options in (
+            {},
+            {"with_2d_quantization": True},
+            {
+                "nvfp4_use_4over6": True,
+                "nvfp4_e4m3_max": 448,
+                "nvfp4_4over6_err_mode": "MSE",
+            },
+        ):
+            for source_dtype in (torch.bfloat16, torch.float32):
+                run_case(options, source_dtype)
+                count += 1
+        dist.barrier()
+        if dist.get_rank() == 0:
+            print(f"NVFP4 QDQ distributed: {count} cases passed on 2 ranks", flush=True)
+    finally:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

tests/pytorch/distributed/test_nvfp4_qdq.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Pytest launcher for distributed NVFP4 QDQ reference parity."""
+
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+import torch
+import transformer_engine.pytorch as te
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="QDQ amax reduction requires two GPUs")
+def test_nvfp4_qdq_amax_reduction():
+    available, reason = te.is_nvfp4_available(return_reason=True)
+    if not available:
+        pytest.skip(reason)
+    script = Path(__file__).with_name("run_nvfp4_qdq.py")
+    subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "torch.distributed.run",
+            "--standalone",
+            "--nproc_per_node=2",
+            str(script),
+        ],
+        env=os.environ.copy(),
+        check=True,
+    )