Decode the e4m3 NaN encodings in from_fp8 - #4164
Open
axiom-of-choice wants to merge 1 commit into
Open
Conversation
0x7f and 0xff are e4m3's only NaN encodings, and both decoded to +/-480. The decode reinterprets (bits & 127) << 7 as a float16 and scales by 256, which is exact for every finite e4m3 value. It does not carry NaN: the shift leaves those two patterns on exponent field 15 rather than 31, so float16 reads them as a normal number, and 1.875 * 256 = 480. That is also above the format's largest finite magnitude of 448, so a NaN weight silently became an out-of-range value rather than a signal. Special-case the two patterns in both the CPU and Metal decoders. Checked exhaustively against torch.float8_e4m3fn over all 256 byte values: 0 mismatches on CPU and Metal, where before the change those two bytes were the only ones that disagreed. The existing round-trip test only covers finite values, which is why this went unnoticed. The new test pins both NaN encodings across float16, bfloat16 and float32 on both streams, and pins 448 as finite so a future fix cannot swallow the maximum. to_fp8 does not preserve NaN either, saturating it to 0x7e (448) the way it saturates overflow; torch maps both to NaN. That is a separate call, and its saturating behaviour looks deliberate, so this change leaves it alone.
zcbenz
added
the
await verification
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Proposed changes
Fixes #4135.
0x7fand0xffare e4m3's only NaN encodings, andfrom_fp8decoded both as+/-480.0.The decode reinterprets
(bits & 127) << 7as a float16 and scales by 256, which is exact for every finite e4m3 value. It does not carry NaN: the shift leaves those two patterns on exponent field 15 rather than 31, so float16 reads them as a normal number, and1.875 * 256 = 480.Worth noting beyond the issue: 480 is above the format's largest finite magnitude of 448, so a NaN weight did not merely lose its NaN-ness, it silently became an out-of-range value that reads as ordinary data.
Both decoders had the same defect.
mlx/backend/metal/kernels/fp8.h'sfp8_e4m3is not a native Metal type, it performs the same shift-and-scale as the CPU path, so CPU and GPU agreed on the wrong answer.Verification
Exhaustive over all 256 byte values against
torch.float8_e4m3fn, on both streams:0x7f,0xff)0x7f,0xff)Those two bytes were the only ones that ever disagreed, so nothing finite moves.
The new test fails without the change with
array([480, -480], dtype=float16). It pins both NaN encodings across float16, bfloat16 and float32 on both streams, and pins 448 as finite so a future change cannot swallow the maximum. The pre-existing round-trip test only covers finite values, which is why this went unnoticed.Full suites on an M2 Pro (macOS 26.5.2, Metal build): 759 passed / 46 skipped / 10804 subtests in Python, and 261 C++ cases with 3521 assertions.
Out of scope
to_fp8does not preserve NaN either, saturating it to0x7e(448) the way it saturates overflow, where torch maps both to NaN. That is a separate call and its saturating behaviour looks deliberate, so this PR leaves it alone. Happy to follow up if you want it changed.Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes