Skip to content

Fix FP8 NaN preservation in to_fp8 and E8M0 scale dequantization - #4756

Merged
zcbenz merged 4 commits into
ml-explore:mainfrom
FaizPalwala:fix-fp8-nan-handling
Oct 28, 2026
Merged

zcbenz merged 4 commits into
ml-explore:mainfrom
FaizPalwala:fix-fp8-nan-handling

Conversation

@FaizPalwala

Copy link
Copy Markdown
Contributor
  • ☑️ I understand it is strictly prohibited to use AI to write PR description

Summary

Fixes #4751.

In ToFP8 (mlx/backend/cpu/unary_ops.h) and fp8_e4m3 (mlx/backend/metal/kernels/fp8.h), the saturation check if (f_bits >= fp8_max) unconditionally assigned 0x7E. Because float32 NaN has f_bits > 0x7F800000 > fp8_max, all NaNs saturated to finite 448.0 (0x7E) instead of the standard FP8 E4M3fn NaN bit pattern (0x7F / 0xFF).

Rationale

  1. OCP E4M3fn Specification & PyTorch Parity: In the E4M3fn format, there is no infinity encoding—overflow and infinities saturate to the maximum finite value (0x7E / 448.0), while 0x7F / 0xFF are reserved for NaNs. The conversion in MLX was adapted from PyTorch's c10/util/Float8_e4m3fn.h, which explicitly gates saturation with (f_bits > 0x7F800000) ? 0x7F : 0x7E. This PR restores that check across both CPU and Metal paths.
  2. OCP Microscaling (MX) Spec v1.0 for E8M0: For E8M0 block scales, 0xFF is reserved as unrepresentable / NaN. Decoding 0xFF via static_cast<uint16_t>(bits) << 23 evaluated to 0x7F800000 (+Infinity) in float32. This PR decodes 0xFF to quiet NaN (0x7FC00000 / 0x7FC0) in both Metal and CPU, and adds a NaN guard in fp_dequantize so power(2.0, scales - 127) does not overflow to +inf on CPU.

Changes

  1. Added (f_bits > 0x7F800000) ? 0x7F : 0x7E check in Metal fp8_e4m3 and CPU SIMD ToFP8 prior to applying the sign bit.
  2. In fp8_e8m0 (Metal) and dequantize_scale (CPU), decode 0xFF to quiet NaN (0x7FC00000 / 0x7FC0).
  3. In mlx/ops.cpp:fp_dequantize, substitute NAN for 0xFF scales to prevent power(2.0, 128) overflow on CPU.
  4. Added unit tests in test_ops.py and test_quantized.py validating NaN preservation, round-trip, overflow saturation, and E8M0 NaN scale decoding across CPU and GPU.
  • AI usage disclosure: AI assistance was used to scaffold unit test coverage across device and dtype permutations; all kernel patches, specifications, and test passes were directly reviewed and validated.

@zcbenz
zcbenz force-pushed the fix-fp8-nan-handling branch 2 times, most recently from 95a6b2e to 4886af9 Compare October 28, 2026 00:06
@zcbenz
zcbenz force-pushed the fix-fp8-nan-handling branch from 4886af9 to 261f2ef Compare October 28, 2026 08:13
@zcbenz
zcbenz merged commit 7d2d94a into ml-explore:main Oct 28, 2026
30 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] NaN inputs convert to 448 in FP8 (should be NaN)

2 participants