Add torch.compile compatibility to FP8 SDPA using FA3#172622
Closed
howardzhang-cv wants to merge 4 commits intogh/howardzhang-cv/7/basefrom
Closed
Add torch.compile compatibility to FP8 SDPA using FA3#172622howardzhang-cv wants to merge 4 commits intogh/howardzhang-cv/7/basefrom
howardzhang-cv wants to merge 4 commits intogh/howardzhang-cv/7/basefrom
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/172622
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 7cb5426 with merge base 32642ba ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
drisspg
reviewed
Jan 17, 2026
| sdpa_constraint, | ||
| warn=False, | ||
| ) | ||
| make_fallback( |
Contributor
There was a problem hiding this comment.
does the constraint work for the FP8 V layout ?
Contributor
There was a problem hiding this comment.
I think there is actually a section in this file on the layout constraints needed for inductor we should also apply for this overload
drisspg
reviewed
Jan 17, 2026
| ) | ||
| # Directly call the internal flash attention operator which has descale support | ||
| result = torch._scaled_dot_product_flash_attention( | ||
| # Use the .low_p OpOverload directly for better torch.compile compatibility |
drisspg
reviewed
Jan 17, 2026
| seed = torch.empty((2), dtype=torch.uint64, device="meta") | ||
| offset = torch.empty((), dtype=torch.uint64, device="meta") | ||
|
|
||
| return ( |
Contributor
There was a problem hiding this comment.
how much of this is resuable from the other meta funcs, can we just call them directly here?
drisspg
reviewed
Jan 17, 2026
| seqused_k: Tensor | None = None, | ||
| alibi_slopes: Tensor | None = None, | ||
| ): | ||
| print( |
SergeyTyshkevich
pushed a commit
to SergeyTyshkevich/chart2
that referenced
this pull request
Jan 19, 2026
Summary: Added meta registration for new scaled_dot_product_flash_attention.low_p overload Added inductor lowering fallback for new overload Directly call op overload in _scaled_dot_product_attention_fp8 instead of python builtin function call ghstack-source-id: 9605c0f Pull-Request: pytorch/pytorch#172622
drisspg
approved these changes
Jan 21, 2026
Contributor
Author
|
@pytorchbot merge |
Collaborator
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
suncapitalllc007-star
pushed a commit
to suncapitalllc007-star/pytorch
that referenced
this pull request
Jan 25, 2026
Summary: Added meta registration for new scaled_dot_product_flash_attention.low_p overload Added inductor lowering fallback for new overload Directly call op overload in _scaled_dot_product_attention_fp8 instead of python builtin function call ghstack-source-id: baa029c Pull-Request: pytorch/pytorch#172622
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
3 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.
Stack from ghstack (oldest at bottom):
Summary:
Added meta registration for new scaled_dot_product_flash_attention.low_p overload
Added inductor lowering fallback for new overload
Directly call op overload in _scaled_dot_product_attention_fp8 instead of python builtin function call
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo