[BUG] Regression: `mx.from_fp8` decodes SafeTensors `F8_E4M3` NaNs as finite values

[turingH]

Summary

mx.from_fp8() does not preserve the two canonical NaN encodings of the
SafeTensors F8_E4M3 / PyTorch float8_e4m3fn format:

Raw byte	Expected SafeTensors/PyTorch value	Actual MLX from_fp8 value
0x7f	NaN	+480.0
0xff	NaN	-480.0

This is a regression from the behavior tested in PR #1859, introduced when PR #2985 moved FP8 conversion to the explicit mx.from_fp8() API.

This report is self-contained: the complete minimal reproducer is included
below. It creates one local, four-byte SafeTensors file; it does not download
a model, execute remote code, or run model inference.

Environment

• MLX 0.32.0
• Python 3.14
• SafeTensors 0.8.0
• PyTorch 2.13.0
• Apple M2 Pro, 16-core GPU, Metal backend
• Current MLX main source was inspected at
  e78d894c8f718a805341e1aa1cc835d9c9b8462 (2026-08-10); it has the same
  CPU and Metal decode logic.

Complete minimal reproduction

Save the following as repro.py and run python repro.py in an environment
with mlx, torch, and safetensors installed:

import json
import math
import struct
import tempfile
from pathlib import Path

import mlx.core as mx
import torch
from safetensors import safe_open


def write_fixture(directory: Path) -> Path:
    header = {
        "w": {
            "dtype": "F8_E4M3",
            "shape": [4],
            "data_offsets": [0, 4],
        }
    }
    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
    path = directory / "f8_nan.safetensors"
    path.write_bytes(
        struct.pack("<Q", len(header_bytes))
        + header_bytes
        + bytes([0x00, 0x3C, 0x7F, 0xFF])
    )
    return path


with tempfile.TemporaryDirectory(prefix="mlx-f8-nan-decode-") as tmp:
    path = write_fixture(Path(tmp))

    # MLX's documented native-FP8 path.
    raw = mx.load(path)["w"]
    raw_values = raw.tolist()
    mlx_values = mx.from_fp8(raw, dtype=mx.float32).tolist()

    # Canonical SafeTensors/PyTorch interpretation of the exact same bytes.
    with safe_open(path, framework="pt") as tensors:
        reference_values = tensors.get_tensor("w").to(torch.float32).tolist()

assert raw.dtype == mx.uint8
assert raw_values == [0, 60, 127, 255]
assert mlx_values[:2] == [0.0, 1.5]
assert reference_values[:2] == [0.0, 1.5]
assert mlx_values[2:] == [480.0, -480.0]
assert all(math.isnan(value) for value in reference_values[2:])

print("raw:        ", raw_values)
print("MLX:        ", mlx_values)
print("SafeTensors:", reference_values)

Actual output on the Metal system above:

raw:         [0, 60, 127, 255]
MLX:         [0.0, 1.5, 480.0, -480.0]
SafeTensors: [0.0, 1.5, nan, nan]

Why NaN is the expected result

The SafeTensors PyTorch binding maps F8_E4M3 to
torch.float8_e4m3fn. For that format, 0x7f and 0xff are the two NaN
encodings; the maximum finite values are 0x7e and 0xfe (+448 and
-448). The reproducer obtains the reference through the SafeTensors binding
itself rather than reimplementing the format.

MLX previously implemented and tested the same expectation. PR #1859, titled “Support loading F8_E4M3 from safetensors,” added a fixture containing 0x7f and 0xff and expected both entries to be mx.nan.

Regression provenance

PR #2985 intentionally changed
mx.load() to retain FP8 data as a uint8 carrier and exposed explicit
mx.from_fp8() conversion. That API change is reasonable and is not the
bug reported here.

However, the same PR replaced the prior FromFP8 bit-level conversion, which
handled the NaN encodings, with the current float16-based conversion. It also
changed the prior loader test from the NaN bytes 0x7f/0xff to the ordinary
maximum-finite bytes 0x7e/0xfe; the NaN cases were no longer tested.

Therefore, this can be fixed while keeping PR #2985's intended uint8
carrier API unchanged.

Root cause

The current CPU decoder is effectively:

auto v = Simd<uint16_t, N>(x & 127) << 7;
auto converted = *(Simd<float16_t, N>*)(&v);
out = converted * 256.0;
return select(x & 128, -out, out);

The Metal decoder uses the same transform. Neither special-cases values for
which (x & 0x7f) == 0x7f, so the reserved NaN encodings are interpreted as
finite half-derived values instead.

Relevant current source locations:

• mlx/ops.cpp
  constructs fast::ConvertFP8 for from_fp8.
• mlx/backend/cpu/unary_ops.h
  contains the CPU decoder.
• mlx/backend/metal/kernels/fp8.h
  contains the equivalent Metal conversion.

Verified upper-layer impact (MLX-LM)

This is not only a raw-array API difference. In an MLX-LM checkout, the
following normal model-loading path was exercised with an inert local fixture:

mlx_lm.utils.load_model
  -> mx.load(model.safetensors)
  -> MiniMax Model.sanitize(weights)
  -> mx.from_fp8(weight, dtype=mx.bfloat16)
  -> model.load_weights(...)

The fixture contained one 128 x 128 F8_E4M3 weight block, one F32 scale
of 1.0, and only its first weight byte set to the indicated value. It did
not download or run an external model. The actual result was:

0x7f: loaded=480.0; q_proj MLX-LM=480.0; canonical=nan
0xff: loaded=-480.0; q_proj MLX-LM=-480.0; canonical=nan

Thus a standard MLX-LM sanitizer materializes a finite parameter where the
same immutable SafeTensors artifact has NaN semantics in the canonical reader.
This downstream check is supplementary; the four-byte reproduction above is
complete for the MLX core bug.

Suggested fix

Before or after the ordinary conversion, preserve the two NaN bit patterns:

const auto is_nan = (x & 0x7f) == 0x7f;
// Decode ordinary values as today, then select quiet_NaN for is_nan.

Apply the equivalent behavior to CPU and Metal, and check other backends for
their intended F8_E4M3 semantics.

Acceptance test

Add 0x7f and 0xff back to a mx.from_fp8() regression test, asserting
NaN with equal_nan=True, alongside maximum-finite 0x7e and 0xfe:

bytes:    0x7e    0xfe    0x7f   0xff
expected: +448.0  -448.0  NaN    NaN

The supplied four-byte reproduction can be used verbatim as the test input.

Scope

This is a correctness/interoperability regression, not a report that
mx.load() returning uint8 is unsafe. The raw uint8 carrier is the
intentional API introduced by PR #2985, and raw-byte use (including 0x80) is
outside this report.

The issue can become security-relevant only in a separate workflow that
validates/rejects F8 NaNs using a canonical reader and then deploys the same
artifact through MLX. This report does not claim code execution, a generic
validation bypass, or a security severity.