fix(fp8): floor tiny amax so scale_inv stays at least FLT_MIN

qa/L0_pytorch_unittest/test.sh

@@ -53,6 +53,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwiset
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_current_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_current_scaling_exact.py || test_fail "test_float8_current_scaling_exact.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_tiny_amax_scale_host.xml $TE_PATH/tests/cpp/operator/test_tiny_amax_scale_host.py || test_fail "test_tiny_amax_scale_host.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_qk_norm.xml $TE_PATH/tests/pytorch/test_qk_norm.py || test_fail "test_qk_norm.py"

tests/cpp/operator/CMakeLists.txt

@@ -67,3 +67,7 @@ if("$ENV{NVTE_TEST_DEBUG_CPP_UT_IN_PARALLEL}" STREQUAL "1")
 else()
   add_test(NAME test_operator COMMAND test_operator)
 endif()
+
+add_executable(test_tiny_amax_scale_host test_tiny_amax_scale_host.cpp)
+target_compile_options(test_tiny_amax_scale_host PRIVATE -O2 -Wall -Wextra)
+add_test(NAME test_tiny_amax_scale_host COMMAND test_tiny_amax_scale_host)

tests/cpp/operator/test_cast_current_scaling.cu

@@ -4,6 +4,9 @@
  * See LICENSE for license information.
  ************************************************************************/
 
+#include <cfloat>
+#include <cmath>
+#include <cstdint>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
@@ -58,12 +61,17 @@ void compute_amax_scale_ref(const InputType *data,
   float scale = 1.f;
   float scale_inv = 1.f;
 
-  if (isinf(clamp_amax) || clamp_amax == 0.f) {
+  if (isinf(clamp_amax) || isnan(clamp_amax) || clamp_amax == 0.f) {
       *scale_ptr = scale;
       *scale_inv_ptr = scale_inv;
       return;
   }
 
+  const float amax_floor = max_fp8 * FLT_MIN;
+  if (clamp_amax < amax_floor) {
+    clamp_amax = amax_floor;
+  }
+
   // use ieee_div in CPU
   scale = max_fp8 / clamp_amax;
 
@@ -213,3 +221,65 @@ INSTANTIATE_TEST_SUITE_P(
     }
     return name;
   });
+
+namespace {
+
+using namespace test;
+
+uint32_t float_bits(float value) {
+  uint32_t bits = 0;
+  std::memcpy(&bits, &value, sizeof(bits));
+  return bits;
+}
+
+// Call the device scale kernel, then quantize so scale_inv is stored.
+void check_device_scale(DType otype, float amax, uint32_t expect_scale, uint32_t expect_inv,
+                        bool expect_nonzero_code) {
+  const std::vector<size_t> shape = {128, 128};
+  const size_t n = 128 * 128;
+  Tensor input("tiny_input", shape, DType::kFloat32);
+  Tensor output("tiny_output", shape, otype, true, false);
+  float *data = input.rowwise_cpu_dptr<float>();
+  for (size_t i = 0; i < n; ++i) {
+    data[i] = amax;
+  }
+  input.from_cpu();
+  output.set_amax(amax);
+
+  QuantizationConfigWrapper config;
+  nvte_compute_scale_from_amax(output.data(), config, 0);
+  nvte_quantize(input.data(), output.data(), 0);
+  cudaDeviceSynchronize();
+  ASSERT_EQ(cudaGetLastError(), cudaSuccess) << cudaGetErrorString(cudaGetLastError());
+  EXPECT_EQ(float_bits(output.scale()), expect_scale);
+  EXPECT_EQ(float_bits(output.rowwise_scale_inv()), expect_inv);
+
+  output.to_cpu();
+  bool nonzero = false;
+  if (otype == DType::kFloat8E4M3) {
+    const auto *codes = output.rowwise_cpu_dptr<fp8e4m3>();
+    for (size_t i = 0; i < n; ++i) {
+      uint8_t raw = 0;
+      std::memcpy(&raw, codes + i, sizeof(raw));
+      nonzero = nonzero || raw != 0;
+    }
+  } else {
+    const auto *codes = output.rowwise_cpu_dptr<fp8e5m2>();
+    for (size_t i = 0; i < n; ++i) {
+      uint8_t raw = 0;
+      std::memcpy(&raw, codes + i, sizeof(raw));
+      nonzero = nonzero || raw != 0;
+    }
+  }
+  EXPECT_EQ(nonzero, expect_nonzero_code);
+}
+
+}  // namespace
+
+TEST(CastCSTinyAmax, DeviceScaleInvIsFltMin) {
+  const float tiny = 7.7e-37f;
+  for (DType otype : {DType::kFloat8E4M3, DType::kFloat8E5M2}) {
+    check_device_scale(otype, tiny, 0x7e800000u, 0x00800000u, true);
+    check_device_scale(otype, 0.f, 0x3f800000u, 0x3f800000u, false);
+  }
+}

tests/cpp/operator/test_cast_current_scaling_grouped.cu

@@ -5,6 +5,7 @@
  ************************************************************************/
 
 #include <algorithm>
+#include <cfloat>
 #include <cmath>
 #include <cstdint>
 #include <limits>
@@ -57,12 +58,17 @@ void compute_amax_scale_ref(const InputType *data, const size_t size, float *ama
 
   float scale = 1.0f;
   float scale_inv = 1.0f;
-  if (std::isinf(clamp_amax) || clamp_amax == 0.0f || std::isnan(clamp_amax)) {
+  if (std::isinf(clamp_amax) || std::isnan(clamp_amax) || clamp_amax == 0.0f) {
     *scale_ptr = scale;
     *scale_inv_ptr = scale_inv;
     return;
   }
 
+  const float amax_floor = max_fp8 * FLT_MIN;
+  if (clamp_amax < amax_floor) {
+    clamp_amax = amax_floor;
+  }
+
   scale = max_fp8 / clamp_amax;
   if (std::isinf(scale)) {
     scale = std::numeric_limits<float>::max();

tests/cpp/operator/test_cast_float8blockwise.cu

@@ -4,6 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
+#include <cfloat>
 #include <cuda_bf16.h>
 #include <cuda_fp8.h>
 #include <cuda_runtime.h>
@@ -51,15 +52,22 @@ void scales_from_amax(float amax, const QuantizationOptions& opts, float* qscale
   float quant_type_max_val = Quantized_Limits<OutputType>::max();
   float eps = opts.amax_epsilon;
   amax = std::max(amax, eps);
-  float qscale = quant_type_max_val / amax;
-  if (std::isinf(qscale)) {
-    qscale = input_type_max_val;
-  }
-  if (std::isnan(qscale) || amax == 0) {
-    qscale = 1.0;
+  const float guard = quant_type_max_val * FLT_MIN;
+  float qscale = 1.0f;
+  if (!(std::isinf(amax) || std::isnan(amax) || amax == 0.0f)) {
+    if (amax < guard) {
+      amax = guard;
+    }
+    qscale = quant_type_max_val / amax;
+    if (std::isinf(qscale)) {
+      qscale = input_type_max_val;
+    }
+    if (std::isnan(qscale)) {
+      qscale = 1.0f;
+    }
   }
 
-  if (opts.force_pow_2_scales && qscale != 0.0) {
+  if (opts.force_pow_2_scales && qscale != 0.0f) {
     uint32_t scale_bits = *reinterpret_cast<uint32_t*>(&qscale);
     // Scale must be positive, shift it
     uint8_t exp = scale_bits >> 23;

tests/cpp/operator/test_cast_transpose_current_scaling.cu

@@ -4,6 +4,8 @@
  * See LICENSE for license information.
  ************************************************************************/
 
+#include <cfloat>
+#include <cmath>
 #include <cstring>
 #include <iomanip>
 #include <iostream>
@@ -62,12 +64,17 @@ void compute_amax_scale_ref(const InputType *data,
   float scale = 1.f;
   float scale_inv = 1.f;
 
-  if (isinf(clamp_amax) || clamp_amax == 0.f) {
+  if (isinf(clamp_amax) || isnan(clamp_amax) || clamp_amax == 0.f) {
       *scale_ptr = scale;
       *scale_inv_ptr = scale_inv;
       return;
   }
 
+  const float amax_floor = max_fp8 * FLT_MIN;
+  if (clamp_amax < amax_floor) {
+    clamp_amax = amax_floor;
+  }
+
   // use ieee_div in CPU
   scale = max_fp8 / clamp_amax;
 

tests/cpp/operator/test_tiny_amax_scale_host.cpp

@@ -0,0 +1,179 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+// Host IEEE-754 model of compute_scale_from_amax without fast-math.
+
+#include <cfloat>
+#include <cmath>
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <limits>
+#include <string>
+#include <vector>
+
+namespace {
+
+float bits_to_f32(uint32_t bits) {
+  float value;
+  std::memcpy(&value, &bits, sizeof(value));
+  return value;
+}
+
+uint32_t f32_to_bits(float value) {
+  uint32_t bits;
+  std::memcpy(&bits, &value, sizeof(bits));
+  return bits;
+}
+
+bool is_subnormal(float value) {
+  const uint32_t bits = f32_to_bits(value);
+  return (bits & 0x7f800000u) == 0u && (bits & 0x007fffffu) != 0u;
+}
+
+bool is_finite_normal(float value) {
+  return std::isfinite(value) && !is_subnormal(value) && value != 0.f;
+}
+
+struct Scales {
+  float scale;
+  float scale_inv;
+};
+
+Scales compute_scale(float amax, float max_fp8, bool force_pow_2, bool floor_tiny,
+                     float epsilon) {
+  if (amax < epsilon) {
+    amax = epsilon;
+  }
+  float scale = 1.f;
+  const bool early = std::isinf(amax) || std::isnan(amax) || amax == 0.f;
+  if (!early) {
+    if (floor_tiny) {
+      const float amax_floor = max_fp8 * FLT_MIN;
+      if (amax < amax_floor) {
+        amax = amax_floor;
+      }
+    }
+    scale = max_fp8 / amax;
+    if (std::isinf(scale)) {
+      scale = std::numeric_limits<float>::max();
+    }
+    if (force_pow_2) {
+      uint32_t scale_bits = f32_to_bits(scale);
+      scale_bits &= 0xFF800000u;
+      scale = bits_to_f32(scale_bits);
+    }
+  }
+  Scales out;
+  out.scale = scale;
+  out.scale_inv = 1.0f / scale;
+  return out;
+}
+
+void require(bool condition, const std::string &message) {
+  if (!condition) {
+    std::fprintf(stderr, "FAIL: %s\n", message.c_str());
+    std::exit(1);
+  }
+}
+
+void emit(const char *fmt, int pow2, const char *label, float amax, const Scales &fixed,
+          const Scales &legacy) {
+  std::printf("%s %d %s %08x %08x %08x %08x %08x\n", fmt, pow2, label, f32_to_bits(amax),
+              f32_to_bits(fixed.scale), f32_to_bits(fixed.scale_inv), f32_to_bits(legacy.scale),
+              f32_to_bits(legacy.scale_inv));
+}
+
+void check_format(const char *name, float max_fp8, uint32_t expected_threshold_bits) {
+  const float threshold = max_fp8 * FLT_MIN;
+  require(f32_to_bits(threshold) == expected_threshold_bits, std::string(name) + " threshold bits");
+  require(static_cast<double>(threshold) / static_cast<double>(max_fp8) ==
+              static_cast<double>(FLT_MIN),
+          std::string(name) + " exact quotient");
+
+  const float below = bits_to_f32(expected_threshold_bits - 1u);
+  const float above = bits_to_f32(expected_threshold_bits + 1u);
+  require(static_cast<double>(below) / static_cast<double>(max_fp8) < static_cast<double>(FLT_MIN),
+          std::string(name) + " -1 ULP is mathematically below FLT_MIN");
+  require(static_cast<double>(above) / static_cast<double>(max_fp8) > static_cast<double>(FLT_MIN),
+          std::string(name) + " +1 ULP is mathematically above FLT_MIN");
+
+  const float tiny = 7.7e-37f;
+  require(tiny > 0.f && tiny < threshold, std::string(name) + " 7.7e-37 is below the guard");
+
+  const float samples[] = {0.f, tiny, below, threshold, above};
+  const char *labels[] = {"zero", "tiny_7_7e_37", "minus_1ulp", "exact", "plus_1ulp"};
+
+  for (int pow2 = 0; pow2 <= 1; ++pow2) {
+    for (int i = 0; i < 5; ++i) {
+      const Scales fixed = compute_scale(samples[i], max_fp8, pow2, true, 0.f);
+      const Scales legacy = compute_scale(samples[i], max_fp8, pow2, false, 0.f);
+      emit(name, pow2, labels[i], samples[i], fixed, legacy);
+
+      if (i == 0) {
+        require(fixed.scale == 1.f && fixed.scale_inv == 1.f, std::string(name) + " zero scale");
+        require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv,
+                std::string(name) + " zero matches legacy");
+      } else if (i < 3) {
+        // Below the floor, scale is 2^126 and scale_inv is FLT_MIN.
+        require(f32_to_bits(fixed.scale) == 0x7e800000u, std::string(name) + " floored scale 2^126");
+        require(f32_to_bits(fixed.scale_inv) == 0x00800000u, std::string(name) + " floored inv FLT_MIN");
+        require(samples[i] * fixed.scale > 1.f, std::string(name) + " tiny value stays above 1");
+        // Power-of-two rounding already collapses the -1 ULP sample onto 2^126.
+        if (pow2 && i == 2) {
+          require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv,
+                  std::string(name) + " pow2 -1 ULP matches legacy");
+        } else {
+          require(is_subnormal(legacy.scale_inv), std::string(name) + " legacy inv subnormal");
+          require(f32_to_bits(fixed.scale_inv) != f32_to_bits(legacy.scale_inv),
+                  std::string(name) + " floor changes the inverse");
+        }
+      } else if (i == 3) {
+        require(f32_to_bits(fixed.scale) == 0x7e800000u, std::string(name) + " exact scale 2^126");
+        require(f32_to_bits(fixed.scale_inv) == 0x00800000u, std::string(name) + " exact inv FLT_MIN");
+        require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv,
+                std::string(name) + " exact matches legacy");
+      } else {
+        require(is_finite_normal(fixed.scale) && is_finite_normal(fixed.scale_inv),
+                std::string(name) + " +1 ULP normal");
+        require(fixed.scale_inv >= FLT_MIN, std::string(name) + " +1 ULP inv >= FLT_MIN");
+        require(fixed.scale == legacy.scale && fixed.scale_inv == legacy.scale_inv,
+                std::string(name) + " +1 ULP matches legacy");
+      }
+    }
+  }
+
+  // Epsilon above the guard keeps the historical floor behavior.
+  const float epsilon = 1.0e-6f;
+  const Scales lifted = compute_scale(tiny, max_fp8, false, true, epsilon);
+  require(lifted.scale == max_fp8 / epsilon, "epsilon floor still applies before the guard");
+
+  // Normal magnitudes, zeros, inf, and NaN stay bitwise-identical.
+  const float normals[] = {1.0e-20f, 1.0e-5f, 0.125f, 1.f, 448.f, 1.0e10f,
+                           std::numeric_limits<float>::infinity(),
+                           std::numeric_limits<float>::quiet_NaN(), -0.f, 0.f};
+  for (float amax : normals) {
+    for (int pow2 = 0; pow2 <= 1; ++pow2) {
+      const Scales fixed = compute_scale(amax, max_fp8, pow2, true, 0.f);
+      const Scales legacy = compute_scale(amax, max_fp8, pow2, false, 0.f);
+      require(f32_to_bits(fixed.scale) == f32_to_bits(legacy.scale), "normal scale bits");
+      require(f32_to_bits(fixed.scale_inv) == f32_to_bits(legacy.scale_inv), "normal inv bits");
+    }
+  }
+}
+
+}  // namespace
+
+int main() {
+  std::printf("THRESHOLD e4m3 %08x\n", f32_to_bits(448.0f * FLT_MIN));
+  std::printf("THRESHOLD e5m2 %08x\n", f32_to_bits(57344.0f * FLT_MIN));
+  std::printf("FLT_MIN %08x\n", f32_to_bits(FLT_MIN));
+  std::printf("TINY %08x\n", f32_to_bits(7.7e-37f));
+  check_format("e4m3", 448.0f, 0x04e00000u);
+  check_format("e5m2", 57344.0f, 0x08600000u);
+  std::printf("HOST_OK\n");
+  return 0;
+}