Repository navigation
[MPS] Fix torch.nextafter for bfloat16 - #190481
Closed
sahil-mangla wants to merge 2 commits into
Closed
sahil-mangla wants to merge 2 commits into
sahil-mangla wants to merge 2 commits into
Conversation
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`. The root cause was that `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. We fix this by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, similar to the CPU implementation in `c10/util/BFloat16-math.h`. We also fix a build issue in `scripts/write_metallib_headers.py` where output directories were not created before trying to write metallib wrapper headers, which caused file not found errors during local compilation. This PR was authored with the assistance of an AI assistant. Test Plan: Verify the fix by running: ```bash python -c "import torch; a=torch.tensor(1.0, dtype=torch.bfloat16, device='mps'); b=torch.tensor(2.0, dtype=torch.bfloat16, device='mps'); print(torch.nextafter(a, b).cpu())" ``` And verify it matches the CPU output: ```bash python -c "import torch; a=torch.tensor(1.0, dtype=torch.bfloat16); b=torch.tensor(2.0, dtype=torch.bfloat16); print(torch.nextafter(a, b))" ```
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/190481
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 7da0d06 with merge base b838c64 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
Bot
added
ciflow/mps
|
|
malfet
approved these changes
Jul 20, 2026
malfet
left a comment
malfet
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, but let's make implementation a bit more compact and do not change metallic header generation
Comment on lines
+14
to
+16
| out_path = Path(output_filename) | ||
| out_path.parent.mkdir(parents=True, exist_ok=True) | ||
| with open(out_path, "w") as out: |
malfet
changed the title
Shorten the bfloat16 nextafter specialization by selecting the increment/decrement direction with the `(from < to) == (from > 0)` idiom instead of computing the absolute-value bit patterns. This folds the bit-equal and +0/-0 cases into a single float `from == to` compare and removes the nested zero sub-case, with no change in behavior (NaN, +/-0, subnormal, and finite<->inf boundaries all match CPU). Also revert the unrelated `write_metallib_headers.py` mkdir change so this PR stays focused on the nextafter fix. Test Plan: ``` python test/test_mps.py -k test_nextafter ``` Co-Authored-By: Claude Opus 4.8 <[email protected]>
Contributor
Author
|
@malfet Seems like the changes applied have some cases failing do you want me to go through them and fix them? |
Collaborator
|
@pytorchbot merge -f "MPS tests are green, lint just need a rebase" |
Collaborator
|
@sahil-mangla those are false positives, and could be ignored |
Collaborator
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
zjliu-amd
pushed a commit
to zjliu-amd/pytorch
that referenced
this pull request
Jul 21, 2026
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`, because `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. Fixed by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, closely mirroring the CPU implementation in `c10/util/BFloat16-math.h` and extend nextafter unit test to cover bfloat as well Fixes pytorch#190476 Pull Request resolved: pytorch#190481 Approved by: https://github.com/malfet Co-authored-by: Nikita Shulga <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
youlituao
pushed a commit
to youlituao/pytorch
that referenced
this pull request
Jul 28, 2026
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`, because `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. Fixed by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, closely mirroring the CPU implementation in `c10/util/BFloat16-math.h` and extend nextafter unit test to cover bfloat as well Fixes pytorch#190476 Pull Request resolved: pytorch#190481 Approved by: https://github.com/malfet Co-authored-by: Nikita Shulga <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
youlituao
pushed a commit
to youlituao/pytorch
that referenced
this pull request
Jul 28, 2026
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`, because `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. Fixed by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, closely mirroring the CPU implementation in `c10/util/BFloat16-math.h` and extend nextafter unit test to cover bfloat as well Fixes pytorch#190476 Pull Request resolved: pytorch#190481 Approved by: https://github.com/malfet Co-authored-by: Nikita Shulga <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
aws-kingrj
pushed a commit
to amazon-contributing/upstream-to-pytorch
that referenced
this pull request
Jul 29, 2026
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`, because `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. Fixed by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, closely mirroring the CPU implementation in `c10/util/BFloat16-math.h` and extend nextafter unit test to cover bfloat as well Fixes pytorch#190476 Pull Request resolved: pytorch#190481 Approved by: https://github.com/malfet Co-authored-by: Nikita Shulga <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
MelindaTan
pushed a commit
to MelindaTan/pytorch
that referenced
this pull request
Jul 31, 2026
On MPS, `torch.nextafter` was returning its input unchanged for `bfloat16`, because `nextafter_functor` in `BinaryKernel.metal` called `::metal::nextafter`, which does not have a `bfloat` overload. This caused the arguments to be promoted to `float`, stepped by one fp32 ULP, and then cast back to `bfloat`, which rounded the step away. Fixed by implementing a bitwise manual stepping overload for `bfloat` inside `nextafter_functor` in `BinaryKernel.metal` using its `ushort` representation, closely mirroring the CPU implementation in `c10/util/BFloat16-math.h` and extend nextafter unit test to cover bfloat as well Fixes pytorch#190476 Pull Request resolved: pytorch#190481 Approved by: https://github.com/malfet Co-authored-by: Nikita Shulga <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
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
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.
On MPS,
torch.nextafterwas returning its input unchanged forbfloat16, becausenextafter_functorinBinaryKernel.metalcalled::metal::nextafter, which does not have abfloatoverload.This caused the arguments to be promoted to
float, stepped by one fp32 ULP, and then cast back tobfloat, which rounded the step away.Fixed by implementing a bitwise manual stepping overload for
bfloatinsidenextafter_functorinBinaryKernel.metalusing itsushortrepresentation, closely mirroring the CPU implementation inc10/util/BFloat16-math.hand extend nextafter unit test to cover bfloat as wellFixes #190476