Metal Performance Shaders produces from bf16/fp16 softmax for large attention tensors on , forcing the entire local generative-AI ecosystem to fall back to fp32. This is a . Environment
Minimal reproduction Decomposed softmax on MPS — diffs = x - maxes produces NaN even though every element is identical (result should be all zeros):
Real-world impact Running ComfyUI (the dominant local generative-AI UI) on M5 Max:
Root cause (per PyTorch MPS maintainers) PyTorch maintainers (@drisspg, @albanD) have traced this to in the MPS/MPSGraph fused kernels. The NaN originates in the x - maxes subtraction inside softmax, then propagates through the attention block and the entire network. PyTorch cannot fix this — it is in the Metal/MPS kernel layer. Not isolated
The M4-era fix ("fixed on macOS 15.1") did not survive onto M5, indicating the MPS fused-kernel precision fix was either reverted or not ported to the M5 GPU architecture. Request
I ran your minimal case on an M5 Max under macOS 26.5.1, through MPSGraph. With finite values in range, reductionMaximum followed by subtraction returned exactly zero in bf16, fp16, and fp32. The fused softmax returned exactly 1/S, where S is the row length, at sizes up to 33,554,432 elements.
I did find two cases that produce the result you describe.
The first is a row that is entirely -inf. The maximum is then -inf as well, and (-inf) - (-inf) is NaN under IEEE 754. That is what a fully masked attention row looks like.
The second is the value of the mask constant. -1e9 is finite in bf16 and fp32. The largest finite fp16 value is 65,504, and on this machine -1e9 stored in fp16 came out as -inf. The row then reduces to the first case. That would also explain why fp32 appears to fix the problem.
I went through the MPS and MPSGraph documentation and the SDK header comments. Neither states what to expect from either case. The closest thing to such a statement is chapter 7 of the Metal Shading Language Specification (https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf). It sets numerical expectations for the shading language rather than for MPSGraph. Even there, the text says only that Metal is "compliant to a subset of the IEEE 754 standard," without saying which parts of the standard are in that subset.
Please consider filing a bug report asking for clarification in the documentation. A resolution may involve changes to Apple's software or documentation. If you post the Feedback number here I'll check the status next time I do a sweep of forums posts where I've suggested bug reports.
Bug Reporting: How and Why? has tips on creating your bug report.