Round mxfp8 block scales up to avoid saturation

mlx/backend/cpu/quantized.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2026 Apple Inc.
 
 #include "mlx/backend/common/quantized.h"
 #include "mlx/backend/common/unary.h"
@@ -1061,6 +1061,15 @@ uint8_t to_fp8_e8m0(float x) {
   return static_cast<uint8_t>(n + 127);
 }
 
+// Smallest E8M0 >= x, so a block's largest elements do not saturate.
+uint8_t to_fp8_e8m0_round_up(float x) {
+  uint8_t bits = to_fp8_e8m0(x);
+  if (bits < 0xFE && dequantize_scale<float, 32>(bits) < x) {
+    bits += 1;
+  }
+  return bits;
+}
+
 uint8_t to_fp4_e2m1(float x) {
   if (std::isnan(x)) {
     return 0x7;
@@ -1112,7 +1121,7 @@ void fp_quantize_dequantize(
     if (group_size == 16) {
       scale = dequantize_scale<float, 16>(detail::ToFP8()(scale));
     } else {
-      scale = dequantize_scale<float, 32>(to_fp8_e8m0(scale));
+      scale = dequantize_scale<float, 32>(to_fp8_e8m0_round_up(scale));
     }
 
     for (int j = 0; j < group_size; ++j) {

mlx/backend/metal/kernels/fp8.h

@@ -78,3 +78,14 @@ struct fp8_e8m0 {
 
   uint8_t bits;
 };
+
+// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down
+// leaves the block's largest elements outside the element range, where they
+// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0.
+inline float mx_scale_round_up(float x) {
+  fp8_e8m0 s(x);
+  if (s.bits < 0xFE && float(s) < x) {
+    s.bits += 1;
+  }
+  return float(s);
+}

mlx/backend/metal/kernels/fp_quantized.h

@@ -2122,7 +2122,7 @@ template <typename T, int group_size, int bits, bool has_global_scale>
 
   float scale_dec_b;
   float w_thread = w[index];
-  if (use_mx_scale) {
+  if constexpr (use_mx_scale) {
     scale_dec_b = simd_max(abs(w_thread));
   } else {
     float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0);
@@ -2132,6 +2132,8 @@ template <typename T, int group_size, int bits, bool has_global_scale>
   scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
   if constexpr (has_global_scale) {
     scale_dec_b *= scale_enc;
+  } else if constexpr (use_mx_scale) {
+    scale_dec_b = mx_scale_round_up(scale_dec_b);
   }
 
   using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;
@@ -2214,7 +2216,7 @@ template <typename T, int group_size, int bits, bool has_global_scale>
 
   float scale_dec_b;
   float w_thread = w[index];
-  if (use_mx_scale) {
+  if constexpr (use_mx_scale) {
     scale_dec_b = simd_max(abs(w_thread));
   } else {
     float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0);
@@ -2224,6 +2226,8 @@ template <typename T, int group_size, int bits, bool has_global_scale>
   scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
   if constexpr (has_global_scale) {
     scale_dec_b *= scale_enc;
+  } else if constexpr (use_mx_scale) {
+    scale_dec_b = mx_scale_round_up(scale_dec_b);
   }
 
   using ScaleType = metal::conditional_t<use_mx_scale, fp8_e8m0, fp8_e4m3>;

mlx/ops.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023-2026 Apple Inc.
 
 // Required for using M_PI in MSVC.
 #define _USE_MATH_DEFINES
@@ -5164,11 +5164,17 @@ std::vector<array> fp_quantize(
     } else {
       // convert to e8m0
       auto z = array(0, scales.dtype());
-      scales = where(
-          equal(scales, z, s),
-          z,
-          astype(round(log2(scales, s), s), int32, s),
+      // Round the scale up so the block maximum stays representable,
+      // matching the CUDA backend.
+      auto exponent = astype(round(log2(scales, s), s), int32, s);
+      auto decoded =
+          power(array(2.0f, float32), astype(exponent, float32, s), s);
+      exponent = where(
+          less(decoded, astype(scales, float32, s), s),
+          add(exponent, array(1, int32), s),
+          exponent,
           s);
+      scales = where(equal(scales, z, s), z, exponent, s);
 
       wq = divide(wq, power(array(2.0f, w.dtype()), scales, s), s);
       scales = astype(add(scales, array(127, int32), s), uint8, s);

python/tests/test_quantized.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2026 Apple Inc.
 
+import math
 import os
 import platform
 import subprocess
@@ -121,6 +122,28 @@ def test_mxfp8_quantize_dequantize(self):
         w_hat = mx.dequantize(w_q, scales, mode="mxfp8")
         self.assertTrue(mx.all(w_hat == 0))
 
+    def test_mxfp8_block_scale_does_not_saturate(self):
+        # E4M3 has three mantissa bits, so an in-range element loses at most
+        # half a step, 6.25%. More than that means the block scale rounded
+        # below amax/448 and the block maximum saturated.
+        mx.random.seed(0)
+        group_size = 32
+        n_blocks = 512
+
+        # Sweep the block magnitude across one binade so both scale rounding
+        # directions are covered.
+        w = mx.random.normal(shape=(n_blocks, group_size))
+        w = w * mx.exp(mx.arange(n_blocks) / n_blocks * math.log(2.0)).reshape(-1, 1)
+
+        w_q, scales = mx.quantize(w, group_size=group_size, mode="mxfp8")
+        w_hat = mx.dequantize(w_q, scales, group_size=group_size, mode="mxfp8")
+
+        # Quantization is monotone in |w|, so a block's largest output is the
+        # reconstruction of its largest input.
+        amax = mx.max(mx.abs(w), axis=1)
+        rel = mx.abs(amax - mx.max(mx.abs(w_hat), axis=1)) / amax
+        self.assertLess(mx.max(rel).item(), 0.0626)
+
     def test_nvfp4_quantize_dequantize(self):
         lut = mx.array(
             [