[Common, PyTorch] Write both NVFP4 amax outputs in the reduction

tests/pytorch/nvfp4/test_nvfp4_dual_amax.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Check both NVFP4 amax outputs, including reuse with a decreasing maximum."""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch import NVFP4Quantizer
+
+
+@pytest.mark.parametrize("direction", ["row", "column", "both"])
+@pytest.mark.parametrize("shape", [(128, 128), (160, 192)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("with_2d_quantization", [False, True])
+def test_nvfp4_amax_matches_input_after_reuse(direction, shape, dtype, with_2d_quantization):
+    """The second amax must reset and track the same input maximum as the first."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required")
+    available, reason = te.is_nvfp4_available(return_reason=True)
+    if not available:
+        pytest.skip(reason)
+    q = NVFP4Quantizer(
+        rowwise=direction != "column",
+        columnwise=direction != "row",
+        with_2d_quantization=with_2d_quantization,
+    )
+    q.optimize_for_gemm = False
+    out = q.make_empty(shape, dtype=dtype)
+    # The maximum decreases, reaches zero and rises again; stale atomic-max state fails.
+    for magnitude in (16.0, 0.5, 0.0, 2.0):
+        x = torch.full(shape, -magnitude, device="cuda", dtype=dtype)
+        q.update_quantized(x, out)
+        expected = x.float().abs().amax().reshape(1)
+        for name in ("_amax_rowwise", "_amax_columnwise"):
+            actual = getattr(out, name)
+            if actual is not None:
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)

transformer_engine/common/include/transformer_engine/recipe.h

@@ -76,7 +76,9 @@ void nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction(
 /*! \brief Compute an FP8 tensor's amax.
  *
  *  The amax (maximum absolute value) of the input tensor is computed
- *  and written to the amax buffer of the output tensor.
+ *  and written to the amax buffer of the output tensor. For an NVFP4
+ *  output with distinct row-wise and column-wise amax buffers, the
+ *  amax is written to both.
  *
  *  \param[in]     input            Input tensor. Must be unquantized.
  *  \param[in,out] output           Output tensor. Must be an FP8 tensor with per-tensor scaling.
@@ -88,7 +90,8 @@ void nvte_compute_amax(const NVTETensor input, NVTETensor output, cudaStream_t s
  *
  *  The amax (maximum absolute value) of the input tensor is computed
  *  and written to the amax buffer of the output tensor, using the provided
- *  quantization configuration.
+ *  quantization configuration. For an NVFP4 output with distinct row-wise
+ *  and column-wise amax buffers, the amax is written to both.
  *  One useful config is the noop tensor, which is needed by cuda graph.
  *
  *  \param[in]     input            Input tensor. Must be unquantized.

transformer_engine/common/recipe/current_scaling.cu

@@ -21,16 +21,20 @@ namespace {
 
 constexpr int amax_kernel_threads = 512;
 
-__launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const float *noop_ptr) {
+__launch_bounds__(1) __global__
+    void zero_amax_kernel(float *amax_ptr, float *amax2_ptr, const float *noop_ptr) {
   if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {
     return;
   }
   *amax_ptr = 0;
+  if (amax2_ptr != nullptr) {
+    *amax2_ptr = 0;
+  }
 }
 
 template <int nvec, bool aligned, typename InputType>
 __launch_bounds__(amax_kernel_threads) __global__
-    void amax_kernel(const InputType *input, float *amax, const size_t N,
+    void amax_kernel(const InputType *input, float *amax, float *amax2, const size_t N,
                      const size_t num_aligned_elements, const float *noop_ptr) {
   if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {
     return;
@@ -66,14 +70,17 @@ __launch_bounds__(amax_kernel_threads) __global__
   max = reduce_max<amax_kernel_threads / THREADS_PER_WARP>(max, warp_id);
   if (threadIdx.x == 0) {
     atomicMaxFloat(amax, max);
+    if (amax2 != nullptr) {
+      atomicMaxFloat(amax2, max);
+    }
   }
 }
 
 template <int nvec, typename InputType>
 void launch_amax_kernel(const InputType *input, float *amax, const size_t N, const float *noop_ptr,
-                        cudaStream_t stream) {
-  // Zero out amax so we can update with atomic max
-  zero_amax_kernel<<<1, 1, 0, stream>>>(amax, noop_ptr);
+                        cudaStream_t stream, float *amax2 = nullptr) {
+  // Zero out amax (and the optional second output) so we can update with atomic max
+  zero_amax_kernel<<<1, 1, 0, stream>>>(amax, amax2, noop_ptr);
   NVTE_CHECK_CUDA(cudaGetLastError());
 
   // Return immediately if tensor is empty
@@ -94,18 +101,18 @@ void launch_amax_kernel(const InputType *input, float *amax, const size_t N, con
   // Launch kernel
   switch (align) {
     case Alignment::SAME_ALIGNED:
-      amax_kernel<nvec, true, InputType>
-          <<<num_blocks, threads, 0, stream>>>(input, amax, N, num_aligned_elements, noop_ptr);
+      amax_kernel<nvec, true, InputType><<<num_blocks, threads, 0, stream>>>(
+          input, amax, amax2, N, num_aligned_elements, noop_ptr);
       break;
     case Alignment::SAME_UNALIGNED:
-      amax_kernel<nvec, false, InputType>
-          <<<num_blocks, threads, 0, stream>>>(input, amax, N, num_aligned_elements, noop_ptr);
+      amax_kernel<nvec, false, InputType><<<num_blocks, threads, 0, stream>>>(
+          input, amax, amax2, N, num_aligned_elements, noop_ptr);
       break;
     case Alignment::DIFFERENT: {
       // This case is a logic error, since there is only one pointer (input)
       // in the alignment check. Still safe to process without vectorization.
       amax_kernel<1, true, InputType>
-          <<<num_blocks, threads, 0, stream>>>(input, amax, N, N, noop_ptr);
+          <<<num_blocks, threads, 0, stream>>>(input, amax, amax2, N, N, noop_ptr);
       break;
     }
   }
@@ -170,10 +177,17 @@ void compute_amax_impl(const NVTETensor input_, const NVTETensor output_, cudaSt
   // Compute amax
   float *amax_ptr = reinterpret_cast<float *>(
       (output.amax.dptr != nullptr) ? output.amax.dptr : output.columnwise_amax.dptr);
+  // NVFP4 outputs with a distinct column-wise amax get the same value written there as well,
+  // which replaces a separate device-to-device copy by the caller.
+  float *amax2_ptr = nullptr;
+  if (output.scaling_mode == NVTE_NVFP4_1D_SCALING && output.columnwise_amax.dptr != nullptr &&
+      output.columnwise_amax.dptr != amax_ptr) {
+    amax2_ptr = reinterpret_cast<float *>(output.columnwise_amax.dptr);
+  }
   TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(
-      input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); launch_amax_kernel<nvec>(
-          reinterpret_cast<const IType *>(input.data.dptr), amax_ptr, input.data.numel(), noop_ptr,
-          stream););  // NOLINT(*)
+      input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType);
+      launch_amax_kernel<nvec>(reinterpret_cast<const IType *>(input.data.dptr), amax_ptr,
+                               input.data.numel(), noop_ptr, stream, amax2_ptr););  // NOLINT(*)
 }
 
 void group_compute_amax_impl(const NVTEGroupedTensor input_, NVTEGroupedTensor output_,

transformer_engine/pytorch/csrc/quantizer.cpp

@@ -2730,21 +2730,13 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou
       void* amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr;
       NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer");
 
-      // Compute amax of input tensor
+      // Compute amax of input tensor. When both the row-wise and the column-wise amax are
+      // allocated, the amax kernel updates both (identical values), so no device-to-device copy
+      // is needed afterwards.
       out.set_amax(amax_ptr, DType::kFloat32, std::vector<size_t>{1});
       NVTE_SCOPED_GIL_RELEASE(
           { nvte_compute_amax_with_config(input.data(), out.data(), quant_config, stream); });
       out.set_amax(rowwise_amax_ptr, DType::kFloat32, std::vector<size_t>{1});
-
-      // Make sure row-wise and column-wise amaxes match
-      if (rowwise_amax_ptr != amax_ptr && rowwise_amax_ptr != nullptr) {
-        NVTE_CHECK_CUDA(cudaMemcpyAsync(rowwise_amax_ptr, amax_ptr, sizeof(float),
-                                        cudaMemcpyDeviceToDevice, stream));
-      }
-      if (columnwise_amax_ptr != amax_ptr && columnwise_amax_ptr != nullptr) {
-        NVTE_CHECK_CUDA(cudaMemcpyAsync(columnwise_amax_ptr, amax_ptr, sizeof(float),
-                                        cudaMemcpyDeviceToDevice, stream));
-      }
     }
   }