[FlexAttention] Allow num_warps 8 since when block size >=128 #143299
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
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/143299
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit a5420d8 with merge base 7ab3177 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
yanboliang
approved these changes
Dec 16, 2024
|
@pytorchbot merge |
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 |
AmdSampsa
pushed a commit
to AmdSampsa/pytorch
that referenced
this pull request
Dec 17, 2024
Fixes pytorch#143331 Pull Request resolved: pytorch#143344 Approved by: https://github.com/Chillee ghstack dependencies: pytorch#143299
aditew01
pushed a commit
to aditew01/pytorch
that referenced
this pull request
Dec 18, 2024
…h#143299) # Summary Fixes pytorch#143290 We already strip bad configs here: https://github.com/pytorch/pytorch/blob/e0e763e33135d2ad25c613007aa5f2fee6d2cc24/torch/_inductor/kernel/flex_attention.py#L2299 So this shouldn't be needed. Confirming that the 64 x 128 case is valid otherwise we can just change the default config Pull Request resolved: pytorch#143299 Approved by: https://github.com/yanboliang
aditew01
pushed a commit
to aditew01/pytorch
that referenced
this pull request
Dec 18, 2024
Fixes pytorch#143331 Pull Request resolved: pytorch#143344 Approved by: https://github.com/Chillee ghstack dependencies: pytorch#143299
Git-Hub-Chris
pushed a commit
to Git-Hub-Chris/PyTorch
that referenced
this pull request
Jan 19, 2025
ghstack-source-id: 6992e3a Pull Request resolved: pytorch/pytorch#143299
dimdi-y
pushed a commit
to poolsideai/pytorch
that referenced
this pull request
Jan 20, 2025
…h#143299) # Summary Fixes pytorch#143290 We already strip bad configs here: https://github.com/pytorch/pytorch/blob/e0e763e33135d2ad25c613007aa5f2fee6d2cc24/torch/_inductor/kernel/flex_attention.py#L2299 So this shouldn't be needed. Confirming that the 64 x 128 case is valid otherwise we can just change the default config Pull Request resolved: pytorch#143299 Approved by: https://github.com/yanboliang
dimdi-y
pushed a commit
to poolsideai/pytorch
that referenced
this pull request
Jan 20, 2025
Fixes pytorch#143331 Pull Request resolved: pytorch#143344 Approved by: https://github.com/Chillee ghstack dependencies: pytorch#143299
dimdi-y
pushed a commit
to poolsideai/pytorch
that referenced
this pull request
Jan 20, 2025
…h#143299) # Summary Fixes pytorch#143290 We already strip bad configs here: https://github.com/pytorch/pytorch/blob/e0e763e33135d2ad25c613007aa5f2fee6d2cc24/torch/_inductor/kernel/flex_attention.py#L2299 So this shouldn't be needed. Confirming that the 64 x 128 case is valid otherwise we can just change the default config Pull Request resolved: pytorch#143299 Approved by: https://github.com/yanboliang
dimdi-y
pushed a commit
to poolsideai/pytorch
that referenced
this pull request
Jan 20, 2025
Fixes pytorch#143331 Pull Request resolved: pytorch#143344 Approved by: https://github.com/Chillee ghstack dependencies: pytorch#143299
|
Adding to 2.6.1 as requested by Runway |
|
Is that not worth making a patch release 2.6.1? |
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.
Stack from ghstack (oldest at bottom):
Summary
Fixes #143290
We already strip bad configs here:
pytorch/torch/_inductor/kernel/flex_attention.py
Line 2299 in e0e763e
So this shouldn't be needed. Confirming that the 64 x 128 case is valid otherwise we can just change the default config
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov @Chillee @yanboliang @BoyuanFeng