Fix FP8 Rowwise Gemm Compilation with Auto-functionalize V2 #3457
Conversation
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
Summary: X-link: facebookresearch/FBGEMM#541 Torch recently introduced auto_functionalized_v2, which makes custom functions pickier about how they are defined. Specifically, torch no longer allows optional preallocated outputs. A custom function must either allocate a tensor and return it, or directly write to a preallocated output and return nothing. This conflicts with our impelemtnation of f8f8bf16_rowwise and could cause confusing behaviors or errors when compiled. The only solution is to split into two functions with correct signatures. This diff adds `f8f8bf16_rowwise_out`, which is a very thin wrapper that allows preallocated outputs. It's a bit annoying, but this should allow both versions of the function to compile correctly. Differential Revision: D66795225
|
This pull request was exported from Phabricator. Differential Revision: D66795225 |
✅ Deploy Preview for pytorch-fbgemm-docs ready!
To edit notification comments on pull requests, go to your Netlify site configuration. |
|
This pull request has been merged in 03129ae. |
q10
pushed a commit
to q10/FBGEMM
that referenced
this pull request
Apr 10, 2025
) Summary: X-link: pytorch#3457 X-link: fairinternal/moviegen#800 Pull Request resolved: facebookresearch/FBGEMM#541 Torch recently introduced auto_functionalized_v2, which makes custom functions pickier about how they are defined. Specifically, torch no longer allows optional preallocated outputs. A custom function must either allocate a tensor and return it, or directly write to a preallocated output and return nothing. This conflicts with our impelemtnation of f8f8bf16_rowwise and could cause confusing behaviors or errors when compiled. The only solution is to split into two functions with correct signatures. This diff adds `f8f8bf16_rowwise_out`, which is a very thin wrapper that allows preallocated outputs. It's a bit annoying, but this should allow both versions of the function to compile correctly. Reviewed By: henryhu6, henrylhtsang, jiawenliu64 Differential Revision: D66795225 fbshipit-source-id: 7874421554f65e1c4d04a814c4d8115116f402d4
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.
Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/541
Torch recently introduced auto_functionalized_v2, which makes custom functions pickier about how they are defined. Specifically, torch no longer allows optional preallocated outputs. A custom function must either allocate a tensor and return it, or directly write to a preallocated output and return nothing.
This conflicts with our impelemtnation of f8f8bf16_rowwise and could cause confusing behaviors or errors when compiled. The only solution is to split into two functions with correct signatures. This diff adds
f8f8bf16_rowwise_out, which is a very thin wrapper that allows preallocated outputs. It's a bit annoying, but this should allow both versions of the function to compile correctly.Differential Revision: D66795225