Enable NVFP4 grouped MLP GLU RHT amax path (#3073)

sraman-rgb · pre-commit-ci[bot] · web-flow · commit 3f640735485e · 2026-06-04T18:50:40.000-07:00
* Enable NVFP4 grouped MLP GLU RHT amax path Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address NVFP4 GLU RHT amax review comments Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> * Deduplicate grouped NVFP4 quantize helper Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route precomputed amax through NVFP4 quantize Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> * Handle empty NVFP4 precomputed amax reduction Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> * Fix NVFP4 amax quantize binding signature Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> --------- Signed-off-by: Siddhartha Raman S <sraman@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
@@ -380,6 +380,8 @@ class NVFP4Quantizer : public Quantizer {
 
   void quantize(const TensorWrapper& input, TensorWrapper& out,
                 const std::optional<TensorWrapper>& noop_flag = std::nullopt) override;
+  void quantize_impl(const TensorWrapper& input, TensorWrapper& out,
+                     const std::optional<TensorWrapper>& noop_flag, bool compute_amax);
 
   /*! @brief Quantize to NVFP4, skipping local amax computation
    *
@@ -392,8 +394,6 @@ class NVFP4Quantizer : public Quantizer {
   std::vector<size_t> get_scale_shape(const std::vector<size_t>& shape, bool columnwise) const;
 
  private:
-  void quantize_impl(const TensorWrapper& input, TensorWrapper& out,
-                     const std::optional<TensorWrapper>& noop_flag, bool compute_amax);
   void quantize_with_rht_unfused_helper(const TensorWrapper& input, TensorWrapper& out,
                                         TensorWrapper& rht_output_t_cpp,
                                         QuantizationConfigWrapper& quant_config,
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
@@ -335,12 +335,23 @@ py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector
 py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output,
                     std::optional<at::Tensor> noop_flag);
 
+py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
+                                    const at::Tensor &rowwise_amax,
+                                    const at::Tensor &columnwise_amax);
+
 py::object dequantize(const py::handle &input, DType otype);
 
 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> tensor_offsets);
 
+py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
+                                          const size_t num_tensors,
+                                          std::optional<at::Tensor> first_dims,
+                                          const at::Tensor &rowwise_amax,
+                                          const at::Tensor &columnwise_amax,
+                                          std::optional<at::Tensor> tensor_offsets);
+
 py::object group_dequantize(const py::handle &input, DType otype);
 
 py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -31,6 +31,18 @@ std::vector<size_t> get_tensor_shape(const TensorWrapper &tensor) {
   return std::vector<size_t>(shape.data, shape.data + shape.ndim);
 }
 
+void allreduce_nvfp4_amax_tensors(NVFP4Quantizer *nvfp4_quantizer_cpp,
+                                  std::vector<at::Tensor> &&amax_tensors) {
+  if (!nvfp4_quantizer_cpp->with_amax_reduction || amax_tensors.empty()) {
+    return;
+  }
+  c10d::AllreduceCoalescedOptions opts;
+  opts.reduceOp = c10d::ReduceOp::MAX;
+  NVTE_SCOPED_GIL_RELEASE({
+    nvfp4_quantizer_cpp->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait();
+  });
+}
+
 }  // namespace
 
 py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output,
@@ -71,6 +83,51 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob
   return output_py;
 }
 
+py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
+                                    const at::Tensor &rowwise_amax,
+                                    const at::Tensor &columnwise_amax) {
+  using namespace transformer_engine::pytorch::detail;
+  init_extension();
+
+  NVTE_CHECK(tensor.dim() >= 2, "Tensor must be at least 2D");
+  NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(),
+             "Precomputed amax tensors must be CUDA tensors.");
+  NVTE_CHECK(
+      rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat,
+      "Precomputed amax tensors must be float32.");
+  NVTE_CHECK(rowwise_amax.numel() == 1 && columnwise_amax.numel() == 1,
+             "nvfp4_quantize_with_amax expects scalar rowwise and columnwise amaxes.");
+
+  auto quantizer_cpp = convert_quantizer(quantizer);
+  NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()),
+             "nvfp4_quantize_with_amax only supports NVFP4 quantizers.");
+  NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());
+
+  auto input_contiguous = tensor.contiguous();
+  auto input_cpp = makeTransformerEngineTensor(input_contiguous);
+
+  const auto shape = get_tensor_shape(input_cpp);
+  const auto fake_dtype = input_cpp.dtype();
+  auto [output_cpp, output_py] = quantizer_cpp->create_tensor(shape, fake_dtype);
+
+  if (output_cpp.get_amax().data_ptr != nullptr) {
+    output_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32, getTensorShape(rowwise_amax));
+    output_py.attr("_amax_rowwise") = py::cast(rowwise_amax);
+  }
+  if (output_cpp.get_columnwise_amax().data_ptr != nullptr) {
+    output_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32,
+                                   getTensorShape(columnwise_amax));
+    output_py.attr("_amax_columnwise") = py::cast(columnwise_amax);
+  }
+
+  nvfp4_quantizer_cpp->quantize_impl(input_cpp, output_cpp, std::nullopt, false);
+  if (quantizer_cpp->optimize_for_gemm && !output_cpp.get_with_gemm_swizzled_scales()) {
+    inplace_swizzle_scale_for_gemm(output_py);
+  }
+
+  return output_py;
+}
+
 py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector<size_t> &shape,
                                          at::ScalarType dtype, at::Device device, bool pin_memory) {
   auto quantizer_cpp = convert_quantizer(quantizer);
@@ -84,14 +141,19 @@ namespace {
 // helper functions for NVFP4 grouped quantization (cuda graph safe with shapes stored in device without D2H copy)
 void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor,
                                GroupedTensorWrapper &grouped_output_tensor,
-                               NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream) {
+                               NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream,
+                               bool compute_amax) {
   size_t num_tensors = grouped_input_tensor.num_tensors();
 
   // assert the 2D scaling case, since 2D scaling grouped quant kernel is not ready yet
   NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization,
              "2D scaling grouped quant kernel is not ready yet");
   NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled,
              "NVFP4 4over6 quantization is not supported for grouped quantization.");
+  NVTE_CHECK(nvfp4_quantizer_cpp->with_rht,
+             "graph safe grouped quant kernel for non-RHT path is not ready yet");
+  NVTE_CHECK(nvfp4_quantizer_cpp->with_post_rht_amax,
+             "grouped NVFP4 RHT quantization expects post-RHT amax buffers.");
 
   auto quant_config_cpp = QuantizationConfigWrapper();
 
@@ -122,37 +184,24 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor,
     quant_config_cpp.set_use_fast_math(true);
   }
 
-  // so far, only the RHT path has grouped kernel support
-  // grouped kernels for non-RHT path will be added later
-
-  if (nvfp4_quantizer_cpp->with_rht) {
-    // post-RHT amax or not
-    if (nvfp4_quantizer_cpp->with_post_rht_amax) {
-      NVTE_SCOPED_GIL_RELEASE({
-        nvte_group_hadamard_transform_amax_graph_safe(
-            grouped_input_tensor.data(), grouped_output_tensor.data(), 0,
-            nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream);
-      });
-    } else {
-      NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet");
-    }
-
-    // RHT cast fusion
-    auto tile_scheduler_workspace_torch =
-        at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32));
-    auto nvte_tile_scheduler_workspace =
-        makeTransformerEngineTensor(tile_scheduler_workspace_torch);
-
-    auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix);
+  if (compute_amax) {
     NVTE_SCOPED_GIL_RELEASE({
-      nvte_group_hadamard_transform_cast_fusion_graph_safe(
-          grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(),
-          quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream);
+      nvte_group_hadamard_transform_amax_graph_safe(
+          grouped_input_tensor.data(), grouped_output_tensor.data(), 0,
+          nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream);
     });
-
-  } else {
-    NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet");
   }
+
+  // RHT cast fusion
+  auto tile_scheduler_workspace_torch = at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32));
+  auto nvte_tile_scheduler_workspace = makeTransformerEngineTensor(tile_scheduler_workspace_torch);
+
+  auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix);
+  NVTE_SCOPED_GIL_RELEASE({
+    nvte_group_hadamard_transform_cast_fusion_graph_safe(
+        grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(),
+        quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream);
+  });
 }
 
 }  // namespace
@@ -214,7 +263,7 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const
       // NVFP4 grouped quantization
       NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());
       group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp,
-                                nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream());
+                                nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream(), true);
       break;
     }
     case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: {
@@ -234,6 +283,79 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const
   return py::reinterpret_borrow<py::object>(grouped_output_py);
 }
 
+py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
+                                          const size_t num_tensors,
+                                          std::optional<at::Tensor> first_dims,
+                                          const at::Tensor &rowwise_amax,
+                                          const at::Tensor &columnwise_amax,
+                                          std::optional<at::Tensor> tensor_offsets) {
+  using namespace transformer_engine::pytorch::detail;
+  init_extension();
+
+  NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D");
+  NVTE_CHECK(rowwise_amax.is_cuda() && columnwise_amax.is_cuda(),
+             "Precomputed amax tensors must be CUDA tensors.");
+  NVTE_CHECK(
+      rowwise_amax.scalar_type() == at::kFloat && columnwise_amax.scalar_type() == at::kFloat,
+      "Precomputed amax tensors must be float32.");
+  NVTE_CHECK(rowwise_amax.numel() == static_cast<int64_t>(num_tensors),
+             "Rowwise amax must contain one value per group.");
+  NVTE_CHECK(columnwise_amax.numel() == static_cast<int64_t>(num_tensors),
+             "Columnwise amax must contain one value per group.");
+
+  std::vector<size_t> logical_shape;
+  for (const auto &d : tensor.sizes()) {
+    logical_shape.push_back(d);
+  }
+  const auto logical_first_dim = logical_shape[0];
+  const auto logical_last_dim = logical_shape[1];
+
+  bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0;
+
+  auto quantizer_cpp = convert_quantizer(quantizer);
+  NVTE_CHECK(IsNVFP4Quantizers(quantizer.ptr()),
+             "nvfp4_group_quantize_with_amax only supports NVFP4 quantizers.");
+  NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast<NVFP4Quantizer *>(quantizer_cpp.get());
+
+  auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape);
+  grouped_input_tensor.set_rowwise_data(
+      tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor));
+
+  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
+      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
+      py::reinterpret_borrow<py::object>(quantizer), first_dims, tensor_offsets, logical_first_dim,
+      logical_last_dim);
+
+  if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) {
+    grouped_output_tensor_cpp.set_amax(rowwise_amax.data_ptr(), DType::kFloat32,
+                                       getTensorShape(rowwise_amax));
+    grouped_output_py.attr("amax") = py::cast(rowwise_amax);
+  }
+  if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) {
+    grouped_output_tensor_cpp.set_columnwise_amax(columnwise_amax.data_ptr(), DType::kFloat32,
+                                                  getTensorShape(columnwise_amax));
+    grouped_output_py.attr("columnwise_amax") = py::cast(columnwise_amax);
+  }
+
+  std::vector<at::Tensor> amax_tensors;
+  if (grouped_output_tensor_cpp.get_amax().data_ptr != nullptr) {
+    amax_tensors.push_back(rowwise_amax);
+  }
+  if (grouped_output_tensor_cpp.get_columnwise_amax().data_ptr != nullptr) {
+    amax_tensors.push_back(columnwise_amax);
+  }
+  allreduce_nvfp4_amax_tensors(nvfp4_quantizer_cpp, std::move(amax_tensors));
+
+  if (empty_input_buffer) {
+    return py::reinterpret_borrow<py::object>(grouped_output_py);
+  }
+
+  group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, nvfp4_quantizer_cpp,
+                            at::cuda::getCurrentCUDAStream(), false);
+
+  return py::reinterpret_borrow<py::object>(grouped_output_py);
+}
+
 py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 const size_t num_tensors, std::optional<at::Tensor> first_dims,
                                 std::optional<at::Tensor> tensor_offsets) {
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -164,6 +164,15 @@ void init_router_bindings(pybind11::module &m) {
         py::arg("grad_aux_loss"), "Fused aux loss bwd");
 }
 
+void bind_quantize_with_amax_extensions(py::module_ &m) {
+  m.def("nvfp4_quantize_with_amax", nvfp4_quantize_with_amax, py::arg("tensor"),
+        py::arg("quantizer"), py::arg("rowwise_amax"), py::arg("columnwise_amax"));
+  m.def("nvfp4_group_quantize_with_amax", nvfp4_group_quantize_with_amax, py::arg("tensor"),
+        py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
+        py::arg("rowwise_amax"), py::arg("columnwise_amax"),
+        py::arg("tensor_offsets") = py::none());
+}
+
 }  // namespace transformer_engine::pytorch
 
 #include "common/util/pybind_helper.h"
@@ -195,6 +204,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"),
         py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("tensor_offsets") = py::none());
+  transformer_engine::pytorch::bind_quantize_with_amax_extensions(m);
   m.def("group_dequantize", transformer_engine::pytorch::group_dequantize,
         "Dequantize group tensor", py::arg("input"), py::arg("otype"));
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp
@@ -2331,8 +2331,40 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper(
 void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& out,
                                    const std::optional<TensorWrapper>& noop_flag,
                                    bool compute_amax) {
+  auto reduce_amaxes = [&]() {
+    if (!this->with_amax_reduction) {
+      return;
+    }
+
+    std::vector<at::Tensor> amax_tensors;
+    auto make_amax_tensor = [](void* data_ptr) {
+      NVTE_CHECK(data_ptr != nullptr, "Could not find amax pointer for NVFP4 amax reduction.");
+      return at::from_blob(
+          data_ptr, std::vector<int64_t>{1},
+          [](void*) {},  // deleter doing nothing since it doesn't own the data
+          at::device(at::kCUDA).dtype(torch::kFloat32));
+    };
+    if (rowwise_usage) {
+      amax_tensors.push_back(make_amax_tensor(out.get_amax().data_ptr));
+    }
+    if (columnwise_usage) {
+      amax_tensors.push_back(make_amax_tensor(out.get_columnwise_amax().data_ptr));
+    }
+    if (amax_tensors.empty()) {
+      return;
+    }
+
+    c10d::AllreduceCoalescedOptions opts;
+    opts.reduceOp = c10d::ReduceOp::MAX;
+    NVTE_SCOPED_GIL_RELEASE(
+        { this->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); });
+  };
+
   // Nothing to be done if input is empty
   if (input.numel() == 0) {
+    if (!compute_amax) {
+      reduce_amaxes();
+    }
     return;
   }
 
@@ -2431,10 +2463,12 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou
       // We need:
       // 1. Rowwise amax = amax for input
       // 2. Columnwise amax = amax for RHT(input.t)
-      NVTE_SCOPED_GIL_RELEASE({
-        nvte_hadamard_transform_amax(input.data(), out.data(), 0,
-                                     this->rht_matrix_random_sign_mask_t, stream);
-      });
+      if (compute_amax) {
+        NVTE_SCOPED_GIL_RELEASE({
+          nvte_hadamard_transform_amax(input.data(), out.data(), 0,
+                                       this->rht_matrix_random_sign_mask_t, stream);
+        });
+      }
     } else {
       // raise error since it's not supported yet
       NVTE_ERROR(
@@ -2467,27 +2501,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou
     }
   }
 
-  // amax reduction
-  if (this->with_amax_reduction) {
-    std::vector<at::Tensor> amax_tensors;
-    // push amax tensors inside if they need to be reduced
-    auto make_amax_tensor = [](void* data_ptr) {
-      return at::from_blob(
-          data_ptr, std::vector<int64_t>{1},
-          [](void*) {},  // deleter doing nothing since it doesn't own the data
-          at::device(at::kCUDA).dtype(torch::kFloat32));
-    };
-    if (rowwise_usage) {
-      amax_tensors.push_back(make_amax_tensor(out.get_amax().data_ptr));
-    }
-    if (columnwise_usage) {
-      amax_tensors.push_back(make_amax_tensor(out.get_columnwise_amax().data_ptr));
-    }
-    c10d::AllreduceCoalescedOptions opts;
-    opts.reduceOp = c10d::ReduceOp::MAX;
-    NVTE_SCOPED_GIL_RELEASE(
-        { this->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); });
-  }
+  reduce_amaxes();
 
   // Fast math toggle: RHT transform can be accelerated
   // What math is accelerated? Only the high precision math, so numerical impact is minimal
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py
