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

tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py

@@ -123,6 +123,7 @@ def check_quantization_nvfp4_versus_reference(
     nvfp4_4over6_err_mode: str = "MAE",
     nvfp4_4over6_err_use_fast_math: bool = False,
     scale_dtype: te.DType = te.DType.kFloat8E4M3,
+    reuse_output: bool = False,
 ) -> None:
     scale_dtype = te.DType.cast(scale_dtype)
     if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES:
@@ -252,6 +253,20 @@ def check_quantization_nvfp4_versus_reference(
 
     torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0)
 
+    if reuse_output:
+        assert not use_cpp_allocator
+        assert not row_scaled_nvfp4
+        # Reusing an output whose maximum decreases verifies that every amax buffer is reset.
+        for magnitude in (16.0, 0.5, 0.0, 2.0):
+            x = torch.full((M, N), -magnitude, dtype=x_dtype, device=device)
+            nvfp4_quantizer.update_quantized(x, x_nvfp4_sut)
+            expected_amax = x.float().abs().amax().reshape(1)
+            torch.testing.assert_close(x_nvfp4_sut._amax_rowwise, expected_amax, atol=0.0, rtol=0.0)
+            if return_transpose:
+                torch.testing.assert_close(
+                    x_nvfp4_sut._amax_columnwise, expected_amax, atol=0.0, rtol=0.0
+                )
+
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"])
@@ -359,6 +374,12 @@ def test_quantization_block_tiling_versus_reference(
     row_scaled_nvfp4: bool,
     nvfp4_4over6_config: NVFP44Over6TestConfig,
 ) -> None:
+    reuse_output = (
+        not use_cpp_allocator
+        and not row_scaled_nvfp4
+        and not nvfp4_4over6_config.use_4over6
+        and (M, N) in ((128, 128), (256, 272))
+    )
     check_quantization_nvfp4_versus_reference(
         x_dtype=x_dtype,
         M=M,
@@ -372,6 +393,7 @@ def test_quantization_block_tiling_versus_reference(
         nvfp4_e4m3_max=nvfp4_4over6_config.e4m3_max,
         nvfp4_4over6_err_mode=nvfp4_4over6_config.err_mode,
         nvfp4_4over6_err_use_fast_math=nvfp4_4over6_config.err_use_fast_math,
+        reuse_output=reuse_output,
     )
 
 
@@ -807,29 +829,42 @@ def test_nvfp4_quantization_noncontiguous_inputs(
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize(
-    "M, N",
+    "with_2d_quantization, M, N",
     [
-        # Aligned tiles
-        (128, 128),
-        (256, 256),
-        (512, 512),
-        (2048, 2048),
-        # Padded tiles (non-multiple of kTileDim=128)
-        (256, 272),
-        (304, 304),
-        (320, 256),
+        # Existing 2D coverage: aligned and padded tiles.
+        (True, 128, 128),
+        (True, 256, 256),
+        (True, 512, 512),
+        (True, 2048, 2048),
+        (True, 256, 272),
+        (True, 304, 304),
+        (True, 320, 256),
+        # Focused 1D coverage for aligned and padded output reuse.
+        (False, 128, 128),
+        (False, 256, 272),
+    ],
+    ids=[
+        "2d-128x128",
+        "2d-256x256",
+        "2d-512x512",
+        "2d-2048x2048",
+        "2d-256x272",
+        "2d-304x304",
+        "2d-320x256",
+        "1d-128x128",
+        "1d-256x272",
     ],
 )
 @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str)
-def test_nvfp4_2d_columnwise_only_matches_both_directions(
+def test_nvfp4_columnwise_only_matches_both_directions(
     x_dtype: torch.dtype,
+    with_2d_quantization: bool,
     M: int,
     N: int,
 ):
-    """Bitwise check: 2D NVFP4 with columnwise-only must produce the same
-    columnwise data/scales as the columnwise half of (rowwise + columnwise) 2D.
+    """Columnwise-only NVFP4 must match the columnwise half of both-directions.
 
-    Covers both kernels depending on the (dtype, shape) routing:
+    The 2D cases cover both kernels depending on the (dtype, shape) routing:
     - bf16 with rows % 32 == 0 and cols % 32 == 0 routes to the optimized
       ``quantize_transpose_nvfp4_2D_kernel`` (instantiated with RETURN_ROWWISE=false),
       validating that gating the rowwise pass/store leaves the shared
@@ -853,15 +888,15 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer:
             amax_reduction_group=None,
             with_rht=False,
             with_post_rht_amax=False,
-            with_2d_quantization=True,
+            with_2d_quantization=with_2d_quantization,
             row_scaled_nvfp4=False,
         )
 
     # Reference: produce both directions in a single kernel call.
     q_both = _make_quantizer(rowwise=True, columnwise=True)
     out_both = q_both(x)
 
-    # SUT: produce columnwise only (the path that hits the new amax-only pass).
+    # SUT: produce columnwise only.
     q_col_only = _make_quantizer(rowwise=False, columnwise=True)
     out_col_only = q_col_only(x)
 
@@ -901,5 +936,15 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer:
         out_col_only._amax_columnwise, out_both._amax_columnwise, atol=0, rtol=0
     )
 
+    if (M, N) in ((128, 128), (256, 272)):
+        # Reusing a columnwise-only output exercises the path where its amax is primary.
+        for magnitude in (16.0, 0.5, 0.0, 2.0):
+            x = torch.full((M, N), -magnitude, dtype=x_dtype, device=device)
+            q_col_only.update_quantized(x, out_col_only)
+            expected_amax = x.float().abs().amax().reshape(1)
+            torch.testing.assert_close(
+                out_col_only._amax_columnwise, expected_amax, atol=0.0, rtol=0.0
+            )
+
     # Sanity: column-only path must not allocate a rowwise output.
     assert out_col_only._rowwise_data is None

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 *rowwise_amax_ptr, float *colwise_amax_ptr, const float *noop_ptr) {
   if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) {
     return;
   }
-  *amax_ptr = 0;
+  *rowwise_amax_ptr = 0;
+  if (colwise_amax_ptr != nullptr) {
+    *colwise_amax_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));
-      }
     }
   }