[Common] Preserve NaN through half-precision MXFP8 amax reductions

tests/cpp/operator/test_cast_mxfp8.cu

@@ -936,6 +936,150 @@ std::string mxfp8_2d_quantization_test_name_generator(
     return name;
 }
 
+TEST(CastMXFP8NaNScaling, RowwiseAndBidimNaNBlocksTakeExceptionalScale) {
+    // Regression test for https://github.com/NVIDIA/TransformerEngine/issues/3550:
+    // the amax reduction of the half-precision MXFP8 rowwise/bidimensional kernels
+    // used max.xorsign.abs (NaN-ignoring), so a NaN input element was dropped
+    // before the exceptional-value handling and its block was quantized with a
+    // finite scale.  Every 32-element block containing NaN must take the
+    // exceptional E8M0 scale (255) and the cast payload must preserve NaN.
+    // The shape reaches the specialized kernels: 256 % 128 == 0 (rowwise) and
+    // 256 % 256 == 0 (bidimensional).
+    if (getDeviceComputeCapability() < blackwellComputeCapability) {
+        GTEST_SKIP();
+    }
+
+    using namespace transformer_engine;
+    using namespace test;
+    using InputType = bf16;
+    using OutputType = fp8e4m3;
+
+    const size_t rows = 128;
+    const size_t cols = 256;
+    constexpr size_t kRowwiseBlockCols = 128;  // specialized rowwise kernel eligibility
+    constexpr size_t kRowBlockCols = 32;       // MXFP8 block width along a row
+
+    Tensor input("input", std::vector<size_t>{rows, cols}, DType::kBFloat16);
+    {
+        // NaN lanes sit at j % 8 == 7 in every row except rows with i % 4 == 3,
+        // which stay NaN-free.  Rowwise blocks (one 128-column stretch of one
+        // row, four 32-element MXFP8 blocks) and colwise blocks (one column)
+        // therefore mix exceptional and finite scales within the same tensor.
+        InputType *data = input.rowwise_cpu_dptr<InputType>();
+        for (size_t i = 0; i < rows; ++i) {
+            for (size_t j = 0; j < cols; ++j) {
+                const size_t idx = i * cols + j;
+                if ((j % 8 == 7) && (i % 4 != 3)) {
+                    // bf16 quiet NaN: exponent all ones, nonzero mantissa
+                    reinterpret_cast<uint16_t &>(data[idx]) = 0x7FC0;
+                } else {
+                    const float magnitude = 1.0f + static_cast<float>(idx % 7);
+                    data[idx] = static_cast<InputType>((idx % 2 == 0) ? magnitude : -magnitude);
+                }
+            }
+        }
+        input.from_cpu();
+    }
+
+    const std::array<size_t, 4> scale_dims_rowwise = get_scale_tensor_dims(rows, cols, 1, 32);
+    const std::array<size_t, 4> scale_dims_colwise = get_scale_tensor_dims(rows, cols, 32, 1);
+    const size_t scales_stride_rowwise = scale_dims_rowwise[3];
+    const size_t scales_stride_colwise = scale_dims_colwise[3];
+
+    // Expected E8M0 scale byte for one 32-element MXFP8 block (mirrors the
+    // reference: std::max over |elements|, then float_to_e8m0).  Blocks
+    // containing NaN take the exceptional code 255.
+    auto expected_scale = [&](const size_t r0, const size_t r1, const size_t c0, const size_t c1) {
+        bool has_nan = false;
+        float block_amax = 0.0f;
+        for (size_t i = r0; i < std::min(r1, rows); ++i) {
+            for (size_t j = c0; j < std::min(c1, cols); ++j) {
+                const float elt = static_cast<float>(input.rowwise_cpu_dptr<InputType>()[i * cols + j]);
+                if (std::isnan(elt)) {
+                    has_nan = true;
+                } else {
+                    block_amax = std::max(block_amax, std::abs(elt));
+                }
+            }
+        }
+        return has_nan ? static_cast<uint8_t>(0xFF)
+                       : static_cast<uint8_t>(float_to_e8m0(
+                             block_amax * Quantized_Limits<OutputType>::max_reciprocal()));
+    };
+    auto is_nan_payload = [](const OutputType &elt) {
+        return (reinterpret_cast<const uint8_t &>(elt) & 0x7F) == 0x7F;
+    };
+
+    // ---- rowwise kernel (MXFP8 1D scaling, rowwise-only layout) ------------
+    Tensor output_rowwise("output_rowwise", std::vector<size_t>{rows, cols}, DType::kFloat8E4M3,
+                          /*rowwise=*/true, /*colwise=*/false, NVTE_MXFP8_1D_SCALING);
+    nvte_quantize(input.data(), output_rowwise.data(), 0);
+    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+
+    // rowwise_cpu_dptr()/columnwise_cpu_dptr() hand back the host mirror without
+    // copying, so the D2H copy has to be requested first.
+    output_rowwise.to_cpu();
+    const fp8e8m0 *scales_rowwise = output_rowwise.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+    const OutputType *out_rowwise = output_rowwise.rowwise_cpu_dptr<OutputType>();
+    for (size_t i = 0; i < rows; ++i) {
+        for (size_t b = 0; b < cols / kRowBlockCols; ++b) {
+            const uint8_t expected = expected_scale(i, i + 1, b * kRowBlockCols, (b + 1) * kRowBlockCols);
+            ASSERT_EQ(static_cast<uint8_t>(scales_rowwise[i * scales_stride_rowwise + b]), expected)
+                << "rowwise block (row " << i << ", block " << b << ")";
+        }
+        const bool row_has_nan = (i % 4 != 3);
+        for (size_t j = 0; j < cols; ++j) {
+            const uint8_t payload = reinterpret_cast<const uint8_t &>(out_rowwise[i * cols + j]);
+            if (row_has_nan) {
+                ASSERT_EQ(payload & 0x7F, 0x7F)
+                    << "rowwise payload (" << i << "," << j << ") must be NaN under the exceptional scale";
+            } else {
+                ASSERT_NE(payload & 0x7F, 0x7F)
+                    << "rowwise payload (" << i << "," << j << ") must stay finite in a NaN-free row";
+            }
+        }
+    }
+
+    // ---- bidimensional kernel (rowwise + colwise layouts) ------------------
+    Tensor output_bidim("output_bidim", std::vector<size_t>{rows, cols}, DType::kFloat8E4M3,
+                        /*rowwise=*/true, /*colwise=*/true, NVTE_MXFP8_1D_SCALING);
+    nvte_quantize(input.data(), output_bidim.data(), 0);
+    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaGetLastError(), cudaSuccess);
+
+    output_bidim.to_cpu();
+    const fp8e8m0 *scales_bidim_rowwise = output_bidim.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+    const fp8e8m0 *scales_bidim_colwise = output_bidim.columnwise_cpu_scale_inv_ptr<fp8e8m0>();
+    const OutputType *out_bidim_rowwise = output_bidim.rowwise_cpu_dptr<OutputType>();
+    const OutputType *out_bidim_colwise = output_bidim.columnwise_cpu_dptr<OutputType>();
+    for (size_t i = 0; i < rows; ++i) {
+        for (size_t b = 0; b < cols / kRowBlockCols; ++b) {
+            const uint8_t expected = expected_scale(i, i + 1, b * kRowBlockCols, (b + 1) * kRowBlockCols);
+            ASSERT_EQ(static_cast<uint8_t>(scales_bidim_rowwise[i * scales_stride_rowwise + b]), expected)
+                << "bidim rowwise block (row " << i << ", block " << b << ")";
+        }
+    }
+    for (size_t j = 0; j < cols; ++j) {
+        for (size_t by = 0; by < rows / 32; ++by) {
+            const uint8_t expected = expected_scale(by * 32, (by + 1) * 32, j, j + 1);
+            ASSERT_EQ(static_cast<uint8_t>(scales_bidim_colwise[by * scales_stride_colwise + j]), expected)
+                << "bidim colwise block (col " << j << ", rows " << by * 32 << "+" << 32 << ")";
+        }
+    }
+    // Bidimensional payloads: NaN under exceptional scales, finite otherwise.
+    for (size_t i = 0; i < rows; ++i) {
+        const bool row_has_nan = (i % 4 != 3);
+        for (size_t j = 0; j < cols; ++j) {
+            ASSERT_EQ(is_nan_payload(out_bidim_rowwise[i * cols + j]), row_has_nan)
+                << "bidim rowwise payload (" << i << "," << j << ")";
+            ASSERT_EQ(is_nan_payload(out_bidim_colwise[i * cols + j]), (j % 8 == 7))
+                << "bidim colwise payload (" << i << "," << j << ")";
+        }
+    }
+}
+
+
 }  // namespace
 
 // Test cases with only cast kernels

transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cuh

@@ -193,7 +193,7 @@ __device__ __forceinline__ uint32_t fold_pair_magnitude(uint32_t pair) {
   ptx::bf16x2 lo, hi, folded;
   reinterpret_cast<uint32_t &>(lo) = pair;
   reinterpret_cast<uint32_t &>(hi) = __byte_perm(pair, pair, 0x1032);
-  ptx::abs_max_2x(folded, lo, hi);
+  ptx::abs_max_nan_2x(folded, lo, hi);
   return reinterpret_cast<const uint32_t &>(folded) & kBf16MagnitudeMask;
 }
 
@@ -364,15 +364,15 @@ __global__ __launch_bounds__(kThreadsPerCta, MIN_BLOCKS_PER_SM) void quantize_bi
     reinterpret_cast<uint32_t &>(amax) = tile[i][0];
 #pragma unroll
     for (int32_t k = 1; k < kWordsPerLane; ++k) {
-      ptx::abs_max_2x(amax, amax, reinterpret_cast<const ptx::bf16x2 &>(tile[i][k]));
+      ptx::abs_max_nan_2x(amax, amax, reinterpret_cast<const ptx::bf16x2 &>(tile[i][k]));
     }
     // Butterfly across the lanes that share this 32-column block.
 #pragma unroll
     for (int32_t d = 1; d < kLanesPerRowBlock; d <<= 1) {
       ptx::bf16x2 partner;
       reinterpret_cast<uint32_t &>(partner) =
           __shfl_xor_sync(0xFFFFFFFFu, reinterpret_cast<const uint32_t &>(amax), d);
-      ptx::abs_max_2x(amax, amax, partner);
+      ptx::abs_max_nan_2x(amax, amax, partner);
     }
 
     const uint32_t reciprocal =
@@ -403,7 +403,7 @@ __global__ __launch_bounds__(kThreadsPerCta, MIN_BLOCKS_PER_SM) void quantize_bi
     reinterpret_cast<uint32_t &>(acc) = tile[0][k];
 #pragma unroll
     for (int32_t i = 1; i < kRowsPerWarp; ++i) {
-      ptx::abs_max_2x(acc, acc, reinterpret_cast<const ptx::bf16x2 &>(tile[i][k]));
+      ptx::abs_max_nan_2x(acc, acc, reinterpret_cast<const ptx::bf16x2 &>(tile[i][k]));
     }
     column_partial[k] = reinterpret_cast<const uint32_t &>(acc);
   }
@@ -423,16 +423,17 @@ __global__ __launch_bounds__(kThreadsPerCta, MIN_BLOCKS_PER_SM) void quantize_bi
     reinterpret_cast<uint32_t &>(acc) = column[0];
 #pragma unroll
     for (int32_t w = 1; w < kWarpsPerCta; ++w) {
-      ptx::abs_max_2x(acc, acc, reinterpret_cast<const ptx::bf16x2 &>(column[w * kColumnSlots]));
+      ptx::abs_max_nan_2x(acc, acc,
+                          reinterpret_cast<const ptx::bf16x2 &>(column[w * kColumnSlots]));
     }
     const uint32_t amax_pair = reinterpret_cast<const uint32_t &>(acc);
 
     uint32_t reciprocal_pair;
-    // Test both halves directly rather than folding them first: the fold uses
-    // max.xorsign.abs, which returns the *other* operand when one is NaN, so a
-    // NaN column paired with a finite one would slip into the packed path.
-    // magnitude + 0x0080 sets bit 15 exactly when magnitude >= 0x7F80, and the
-    // sum cannot carry between halves.
+    // Test both halves directly rather than folding them first: the packed
+    // helper below cannot represent Inf/NaN reciprocals, and columns holding
+    // exceptional values (preserved by the NaN-propagating reductions) must
+    // take the per-half path.  magnitude + 0x0080 sets bit 15 exactly when
+    // magnitude >= 0x7F80, and the sum cannot carry between halves.
     const uint32_t exceptional = ((amax_pair & kMagnitudeMaskPair) + kInfProbePair) & 0x80008000u;
     if (__builtin_expect(exceptional == 0u, 1)) {
       reciprocal_pair = mx_scale_reciprocal_x2<OType>(amax_pair);

transformer_engine/common/cast/mxfp8/specialized/cast_rowwise.cuh

@@ -134,14 +134,14 @@ __device__ __forceinline__ ptx::bf16x2 block_half_amax(const uint32_t (&words)[k
   ptx::bf16x2 level[kInWordsPerLane / 2];
 #pragma unroll
   for (int32_t i = 0; i < kInWordsPerLane / 2; ++i) {
-    ptx::abs_max_2x(level[i], pairs[2 * i], pairs[2 * i + 1]);
+    ptx::abs_max_nan_2x(level[i], pairs[2 * i], pairs[2 * i + 1]);
   }
 #pragma unroll
   for (int32_t i = 0; i < kInWordsPerLane / 4; ++i) {
-    ptx::abs_max_2x(level[i], level[i], level[i + kInWordsPerLane / 4]);
+    ptx::abs_max_nan_2x(level[i], level[i], level[i + kInWordsPerLane / 4]);
   }
   ptx::bf16x2 result;
-  ptx::abs_max_2x(result, level[0], level[1]);
+  ptx::abs_max_nan_2x(result, level[0], level[1]);
   return result;
 }
 
@@ -161,7 +161,7 @@ __device__ __forceinline__ float pair_amax_to_float(ptx::bf16x2 pair) {
   reinterpret_cast<uint32_t &>(a) = bits;
   reinterpret_cast<uint32_t &>(b) = folded;
   ptx::bf16x2 wide;
-  ptx::abs_max_2x(wide, a, b);
+  ptx::abs_max_nan_2x(wide, a, b);
   const uint32_t magnitude = reinterpret_cast<const uint32_t &>(wide) & 0x7FFFu;
   return __int_as_float(magnitude << 16);
 }
@@ -366,7 +366,7 @@ __global__ void __launch_bounds__(THREADS_PER_CTA)
     reinterpret_cast<uint32_t &>(partner) =
         __shfl_xor_sync(0xFFFFFFFFu, reinterpret_cast<const uint32_t &>(half_amax), /*laneMask=*/1);
     ptx::bf16x2 block_amax;
-    ptx::abs_max_2x(block_amax, half_amax, partner);
+    ptx::abs_max_nan_2x(block_amax, half_amax, partner);
 
     const e8m0_t biased_exponent =
         ptx::float_to_e8m0(pair_amax_to_float(block_amax) * Quantized_Limits<OType>::max_norm_rcp);

transformer_engine/common/util/ptx.cuh

@@ -732,6 +732,34 @@ __device__ __forceinline__ void abs_max_2x(fp16x2 &dst, const fp16x2 &p1, const
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
 }
 
+// NaN-preserving variants of abs_max_2x for the MXFP8 scale paths.  Selection
+// is identical (the element with the larger magnitude, xorsign.abs), except a
+// NaN operand propagates to the result (PTX max.NaN) instead of being
+// ignored.  Scale generation relies on this so a block containing NaN reaches
+// the exceptional-value handling (E8M0 255) instead of being quantized with a
+// finite scale.
+__device__ __forceinline__ void abs_max_nan_2x(bf16x2 &dst, const bf16x2 &p1, const bf16x2 &p2) {
+#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+  asm volatile("max.NaN.xorsign.abs.bf16x2 %0, %1, %2;"
+               : "=r"(reinterpret_cast<uint32_t &>(dst))
+               : "r"(reinterpret_cast<const uint32_t &>(p1)),
+                 "r"(reinterpret_cast<const uint32_t &>(p2)));
+#else
+  NVTE_DEVICE_ERROR("abs_max_nan_2x is only supported on SM 8.9+.");
+#endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+}
+
+__device__ __forceinline__ void abs_max_nan_2x(fp16x2 &dst, const fp16x2 &p1, const fp16x2 &p2) {
+#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+  asm volatile("max.NaN.xorsign.abs.f16x2 %0, %1, %2;"
+               : "=r"(reinterpret_cast<uint32_t &>(dst))
+               : "r"(reinterpret_cast<const uint32_t &>(p1)),
+                 "r"(reinterpret_cast<const uint32_t &>(p2)));
+#else
+  NVTE_DEVICE_ERROR("abs_max_nan_2x is only supported on SM 8.9+.");
+#endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+}
+
 __device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
   int32_t pred = 0;