MXFP4 weight QAT recipes on MXFP8 and FP8 block-scaling hosts

transformer_engine/common/cast/cast.cu

@@ -16,6 +16,7 @@
 #include "../utils.cuh"
 #include "dispatch/dequantize.cuh"
 #include "dispatch/quantize.cuh"
+#include "mxfp4/fake_quantize_mxfp4.cuh"
 #include "transformer_engine/transpose.h"
 
 void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
@@ -47,6 +48,37 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output,
   dispatch::quantize_fwd_helper<IS_ACT, Empty, nullptr>(input, output, quant_config, stream);
 }
 
+void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
+  NVTE_API_CALL(nvte_mxfp4_fake_quantize);
+  using namespace transformer_engine;
+  const Tensor &input_tensor = *convertNVTETensorCheck(input);
+  Tensor &output_tensor = *convertNVTETensorCheck(output);
+  NVTE_CHECK(input_tensor.dtype() == DType::kBFloat16 || input_tensor.dtype() == DType::kFloat32,
+             "MXFP4 fake-quantize supports BF16/FP32 inputs, got ", to_string(input_tensor.dtype()));
+  NVTE_CHECK(output_tensor.dtype() == input_tensor.dtype(),
+             "MXFP4 fake-quantize output dtype must match the input dtype.");
+  NVTE_CHECK(input_tensor.data.shape == output_tensor.data.shape,
+             "MXFP4 fake-quantize output shape must match the input shape.");
+  const size_t numel = input_tensor.numel();
+  NVTE_CHECK(numel % 32 == 0, "MXFP4 fake-quantize needs the element count divisible by 32.");
+  NVTE_CHECK(!input_tensor.data.shape.empty() && input_tensor.data.shape.back() % 32 == 0,
+             "MXFP4 fake-quantize needs the innermost dimension divisible by 32 "
+             "(1x32 blocks must not cross rows).");
+  if (numel == 0) {
+    return;
+  }
+  if (input_tensor.dtype() == DType::kBFloat16) {
+    dispatch::mxfp4::fake_quantize_mxfp4_launch<bf16>(
+        reinterpret_cast<const bf16 *>(input_tensor.data.dptr),
+        reinterpret_cast<bf16 *>(output_tensor.data.dptr), numel, stream);
+  } else {
+    dispatch::mxfp4::fake_quantize_mxfp4_launch<fp32>(
+        reinterpret_cast<const fp32 *>(input_tensor.data.dptr),
+        reinterpret_cast<fp32 *>(output_tensor.data.dptr), numel, stream);
+  }
+  NVTE_CHECK_CUDA(cudaGetLastError());
+}
+
 void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) {
   NVTE_API_CALL(nvte_dequantize);
   using namespace transformer_engine;

transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh

@@ -0,0 +1,129 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+
+#ifndef TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_
+#define TRANSFORMER_ENGINE_CAST_MXFP4_FAKE_QUANTIZE_MXFP4_CUH_
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include "../../utils.cuh"
+
+namespace transformer_engine {
+namespace dispatch {
+namespace mxfp4 {
+namespace fake_quantize_kernel {
+
+constexpr size_t MXFP4_BLOCK_SIZE = 32;
+constexpr size_t THREADS_PER_CHUNK = 256;
+constexpr float E2M1_MAX = 6.0f;
+constexpr int SCALE_EXP_MIN = -126;
+constexpr int SCALE_EXP_MAX = 125;
+
+__device__ __forceinline__ float round_e2m1_magnitude(const float y) {
+  if (y <= 2.0f) {
+    return rintf(y * 2.0f) * 0.5f;
+  }
+  if (y <= 4.0f) {
+    return rintf(y);
+  }
+  return rintf(y * 0.5f) * 2.0f;
+}
+
+__device__ __forceinline__ int block_scale_exponent_from_bits(const int bits) {
+  const int exp_field = bits >> 23;
+  const int mantissa = bits & 0x7FFFFF;
+  int e = exp_field - 129 + (mantissa > 0x400000 ? 1 : 0);
+  if (exp_field == 0) {
+    e = SCALE_EXP_MIN;
+  }
+  return max(SCALE_EXP_MIN, min(SCALE_EXP_MAX, e));
+}
+
+__device__ __forceinline__ float exp2_from_int(const int e) {
+  return __int_as_float((e + 127) << 23);
+}
+
+__device__ __forceinline__ float mul_rn_no_ftz(const float a, const float b) {
+  float r;
+  asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b));
+  return r;
+}
+
+template <typename IType>
+__global__ void __launch_bounds__(THREADS_PER_CHUNK)
+    fake_quantize_mxfp4_kernel(const IType *__restrict__ input, IType *__restrict__ output,
+                               const size_t num_vectors) {
+  constexpr size_t VEC_ELTS = 16 / sizeof(IType);
+  constexpr int LANES_PER_BLOCK = MXFP4_BLOCK_SIZE / VEC_ELTS;
+
+  const size_t stride = static_cast<size_t>(gridDim.x) * blockDim.x;
+  for (size_t vec_idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       vec_idx < num_vectors; vec_idx += stride) {
+    Vec<IType, VEC_ELTS> in_vec;
+    in_vec.load_from(input, vec_idx);
+
+    float elts[VEC_ELTS];
+    int amax_bits = 0;
+#pragma unroll
+    for (int i = 0; i < static_cast<int>(VEC_ELTS); ++i) {
+      elts[i] = static_cast<float>(in_vec.data.elt[i]);
+      amax_bits = max(amax_bits, __float_as_int(elts[i]) & 0x7fffffff);
+    }
+    const unsigned lane = threadIdx.x & 31u;
+    const unsigned group_mask =
+        (LANES_PER_BLOCK == 32)
+            ? 0xffffffffu
+            : (((1u << LANES_PER_BLOCK) - 1u) << (lane & ~(LANES_PER_BLOCK - 1u)));
+#pragma unroll
+    for (int offset = LANES_PER_BLOCK / 2; offset > 0; offset >>= 1) {
+      amax_bits = max(amax_bits, __shfl_xor_sync(group_mask, amax_bits, offset));
+    }
+    const int nonfinite = amax_bits >= 0x7f800000;
+
+    Vec<IType, VEC_ELTS> out_vec;
+    if (nonfinite) {
+      const float qnan = __int_as_float(0x7fc00000);
+#pragma unroll
+      for (int i = 0; i < static_cast<int>(VEC_ELTS); ++i) {
+        out_vec.data.elt[i] = static_cast<IType>(qnan);
+      }
+    } else {
+      const int e = (amax_bits > 0) ? block_scale_exponent_from_bits(amax_bits) : 0;
+      const float scale = exp2_from_int(e);
+      const float inv_scale = exp2_from_int(-e);
+#pragma unroll
+      for (int i = 0; i < static_cast<int>(VEC_ELTS); ++i) {
+        const float av = __int_as_float(__float_as_int(elts[i]) & 0x7fffffff);
+        const float y = fminf(mul_rn_no_ftz(av, inv_scale), E2M1_MAX);
+        const float q = copysignf(round_e2m1_magnitude(y), elts[i]);
+        out_vec.data.elt[i] = static_cast<IType>(mul_rn_no_ftz(q, scale));
+      }
+    }
+    out_vec.store_to(output, vec_idx);
+  }
+}
+
+}
+
+template <typename IType>
+void fake_quantize_mxfp4_launch(const IType *input, IType *output, const size_t numel,
+                                cudaStream_t stream) {
+  using namespace fake_quantize_kernel;
+  constexpr size_t VEC_ELTS = 16 / sizeof(IType);
+  const size_t num_vectors = numel / VEC_ELTS;
+  const size_t blocks_needed = (num_vectors + THREADS_PER_CHUNK - 1) / THREADS_PER_CHUNK;
+  const int grid = static_cast<int>(blocks_needed < 65535 ? blocks_needed : 65535);
+  fake_quantize_mxfp4_kernel<IType>
+      <<<grid, THREADS_PER_CHUNK, 0, stream>>>(input, output, num_vectors);
+}
+
+}
+}
+}
+
+#endif

transformer_engine/common/include/transformer_engine/cast.h

@@ -418,6 +418,15 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input,
  */
 void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream);
 
+/*! \brief Project a BF16/FP32 tensor onto the MXFP4 (E2M1) grid with 1x32
+ *         power-of-two block scales, writing back dequantized values.
+ *
+ *  \param[in]     input     Input tensor (BF16/FP32, innermost dim divisible by 32).
+ *  \param[in,out] output    Output tensor with the same dtype and shape.
+ *  \param[in]     stream    CUDA stream used for the operation.
+ */
+void nvte_mxfp4_fake_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream);
+
 /*! \brief Casts input grouped tensor from reduced to higher precision.
  *         In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise
  *         data of the output tensor, regardless of whether the row- or columnwise scaling is used.

transformer_engine/common/recipe/__init__.py

@@ -168,6 +168,11 @@ def custom(cls):
         """Whether the given recipe is custom."""
         return issubclass(cls, CustomRecipe)
 
+    @classmethod
+    def mxfp4_qat(cls):
+        """Whether the given recipe applies MXFP4 weight QAT on top of its base recipe."""
+        return issubclass(cls, (MXFP4QATMXFP8BlockScaling, MXFP4QATFloat8BlockScaling))
+
 
 @dataclass(repr=False)
 class DelayedScaling(Recipe):
@@ -481,6 +486,43 @@ def _make_repr(self) -> str:
         )
 
 
+@dataclass(repr=False)
+class MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling):
+    """
+    MXFP8 recipe with MXFP4 weight quantization-aware training.
+
+    Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid
+    before the regular MXFP8 weight quantization; the rowwise encoding of the
+    projected weight is lossless. Activations, gradients and
+    ``backward_override`` behave as in the base recipe. bf16/fp32 only.
+    """
+
+
+@dataclass(repr=False)
+class MXFP4QATFloat8BlockScaling(Float8BlockScaling):
+    """
+    Float8 block-scaling recipe with MXFP4 weight quantization-aware training.
+
+    Weights are projected onto the MXFP4 (E2M1, 1x32 power-of-two scale) grid
+    before the regular 128x128 blockwise FP8 weight quantization (lossless
+    within each tile's FP8 dynamic-range headroom). Activations, gradients and
+    ``backward_override`` behave as in the base recipe. bf16/fp32 only.
+    """
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.w_block_scaling_dim != 2:
+            raise ValueError(
+                "MXFP4 QAT requires 128x128 (2D) weight scaling blocks, got "
+                f"w_block_scaling_dim={self.w_block_scaling_dim}."
+            )
+        if not self.fp8_quant_fwd_weight.power_2_scale:
+            raise ValueError(
+                "MXFP4 QAT requires power-of-two weight scales; "
+                "NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is incompatible."
+            )
+
+
 @dataclass(repr=False)
 class NVFP4BlockScaling(Recipe):
     """

transformer_engine/common/util/ptx.cuh

@@ -380,6 +380,8 @@ __device__ __forceinline__ bf16 exp2f_rcp<bf16>(e8m0_t biased_exp) {
 }
 
 __device__ __forceinline__ float exp2f(e8m0_t biased_exp) {
+  if (biased_exp == 0) return __int_as_float(0x00400000);
+  if (biased_exp == 255) return __int_as_float(0x7fffffff);
   return __int_as_float(biased_exp << FP32_MANTISSA_BITS);
 }
 

transformer_engine/pytorch/csrc/extensions.h

@@ -347,6 +347,8 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz
 
 py::object dequantize(const py::handle &input, DType otype);
 
+at::Tensor mxfp4_fake_quantize(const at::Tensor &input);
+
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                           std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                           std::optional<at::Tensor> tensor_offsets,

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -556,6 +556,25 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                         py::cast(std::move(dbias_torch)));
 }
 
+at::Tensor mxfp4_fake_quantize(const at::Tensor &input) {
+  TORCH_CHECK(input.is_cuda(), "mxfp4_fake_quantize: input must be a CUDA tensor.");
+  TORCH_CHECK(input.dim() == 2, "mxfp4_fake_quantize: input must be 2D, got dim=", input.dim());
+  TORCH_CHECK(input.size(1) % 32 == 0,
+              "mxfp4_fake_quantize: innermost dimension must be divisible by 32, got ",
+              input.size(1));
+  const c10::cuda::CUDAGuard device_guard(input.device());
+  auto input_contiguous = input.contiguous();
+  if (reinterpret_cast<uintptr_t>(input_contiguous.data_ptr()) % 16 != 0) {
+    input_contiguous = input_contiguous.clone();
+  }
+  auto output = at::empty_like(input_contiguous);
+  auto input_cpp = makeTransformerEngineTensor(input_contiguous);
+  auto output_cpp = makeTransformerEngineTensor(output);
+  nvte_mxfp4_fake_quantize(input_cpp.data(), output_cpp.data(),
+                           at::cuda::getCurrentCUDAStream());
+  return output;
+}
+
 py::object dequantize(const py::handle &input, transformer_engine::DType otype) {
   init_extension();
 

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -199,6 +199,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 
   m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"),
         py::arg("output") = py::none(), py::arg("noop") = py::none());
+  m.def("mxfp4_fake_quantize", &transformer_engine::pytorch::mxfp4_fake_quantize,
+        "Project a BF16/FP32 tensor onto the MXFP4 grid (QAT fake-quantization)",
+        py::arg("input"), py::call_guard<py::gil_scoped_release>());
   m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"),
         py::arg("otype"));
   m.def("create_empty_quantized_tensor",

transformer_engine/pytorch/module/base.py

@@ -46,6 +46,7 @@
     _fsdp_gather_tensors,
 )
 from ..constants import dist_group_type
+from ..mxfp4_qat import mxfp4_fake_quantize
 from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS
 from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer
 from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer
@@ -800,8 +801,26 @@ def quantize_weight(
         ``_fp8_workspaces``.
     """
 
+    _mxfp4_qat_active = (
+        FP8GlobalStateManager.get_fp8_recipe().mxfp4_qat()
+        if FP8GlobalStateManager.is_fp8_enabled()
+        else False
+    )
+    if _mxfp4_qat_active and workspace_dtype == torch.float16:
+        raise NotImplementedError(
+            "MXFP4 QAT does not support fp16 as the activation/dequantize dtype: "
+            "the MXFP4 grid (values up to 6*2^125) exceeds fp16 range. Use bf16 "
+            "or fp32."
+        )
+
     # Already-quantized weight (primary FP8 parameters)
     if isinstance(tensor, QuantizedTensor):
+        if _mxfp4_qat_active:
+            raise NotImplementedError(
+                "MXFP4 QAT recipes do not support primary quantized weights: the "
+                "high-precision master weight is required to project onto the "
+                "MXFP4 grid."
+            )
         update_rowwise = True if quantizer.rowwise_usage else None
         update_columnwise = True if quantizer.columnwise_usage else None
         tensor.update_usage(
@@ -833,6 +852,8 @@ def quantize_weight(
         if update_workspace:
             if tensor is None:
                 raise ValueError("tensor kwarg must be provided to update FP8 workspace")
+            if _mxfp4_qat_active:
+                tensor = mxfp4_fake_quantize(tensor)
             if hasattr(workspace, "quantize_"):
                 workspace.quantize_(tensor, noop_flag=skip_update_flag)
             else:
@@ -842,6 +863,8 @@ def quantize_weight(
     # Cache miss — create new workspace
     if tensor is None or quantizer is None:
         raise ValueError("tensor and quantizer kwargs must be provided to construct FP8 workspace")
+    if _mxfp4_qat_active:
+        tensor = mxfp4_fake_quantize(tensor)
     if cache:
         # Ensure the tensor in the cache is an instance of torch.Tensor,
         # as it persists beyond a single forward pass.
@@ -1506,9 +1529,12 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
             meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe()
 
         _current_recipe = meta["recipe"]
-        if _original_recipe is not None and not (
-            issubclass(_current_recipe.__class__, _original_recipe.__class__)
-            or issubclass(_original_recipe.__class__, _current_recipe.__class__)
+        if _original_recipe is not None and (
+            not (
+                issubclass(_current_recipe.__class__, _original_recipe.__class__)
+                or issubclass(_original_recipe.__class__, _current_recipe.__class__)
+            )
+            or _current_recipe.mxfp4_qat() != _original_recipe.mxfp4_qat()
         ):
             warnings.warn(
                 f"Recipe type changed from {_original_recipe.__class__.__name__} "

transformer_engine/pytorch/module/layernorm_linear.py

@@ -32,6 +32,7 @@
     _2X_ACC_WGRAD,
 )
 from ..quantization import FP8GlobalStateManager, QuantizerRole
+from ..mxfp4_qat import mxfp4_fake_quantize
 from ..utils import (
     assert_dim_for_fp8_exec,
     cast_if_needed,
@@ -810,6 +811,8 @@ def backward(
                     # fsdp2 quantized-tensor hooks when workspace was not saved.
                     weight = saved_weight
                 elif ctx.weight_quantizer is not None:
+                    if ctx.fp8_recipe is not None and ctx.fp8_recipe.mxfp4_qat():
+                        saved_weight = mxfp4_fake_quantize(saved_weight)
                     ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True)
                     weight = ctx.weight_quantizer(saved_weight)