Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Optimized templated attention to use exp2 #124356

Closed
wants to merge 9 commits into from
Closed

Optimized templated attention to use exp2 #124356

wants to merge 9 commits into from

Conversation

Chillee
Copy link
Contributor

@Chillee Chillee commented Apr 18, 2024 •
edited

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Apr 18, 2024 •
edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/124356

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a3b4bd1 with merge base cebf651 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
10d09fa 
ghstack-source-id: c175f251f60e3d59a479d1cb6be7331a9559ee10
Pull Request resolved: #124356
@github-actions github-actions bot requested a review from ezyang April 18, 2024 04:06
@Chillee
Update on "Optimized templated attention to use exp2"
2b58279 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
cc9998f 
ghstack-source-id: d922d4320e7999e1789cd47dd30bc41c2891aa0b
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
5171e3d 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
f4d1ba7 
ghstack-source-id: c68d3f61b2abdf10f1149ce2db7c1b63824a81b3
Pull Request resolved: #124356
@Chillee
Copy link
Contributor Author

Chillee commented Apr 18, 2024

If I enable the new options (ROWS_ASSUME_SAFE and SCORE_MOD_IS_LINEAR), perf on benchmark goes up to 90% of FAV2.

@ezyang ezyang removed their request for review April 18, 2024 14:37
drisspg
drisspg approved these changes Apr 18, 2024
View reviewed changes
Copy link
Contributor

@drisspg drisspg left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some small nits, perf go ↑

@Chillee
Update on "Optimized templated attention to use exp2"
dce3455 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
55149ee 
ghstack-source-id: 99e051d9f107b3b57a76249b156e1ad8e278d834
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
b84a213 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
112e632 
ghstack-source-id: 326e0eb7d62750c7be4e2dbbd3dc00c8a5dbf40b
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
0020194 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
617143b 
ghstack-source-id: 58f020d5029655ea320709ce40837cdf4c6c405f
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
4c5f4e3 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
aaaafee 
ghstack-source-id: 8f9605cc9a7c6fcf6bf929c9ff9db056bce1235e
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
c344940 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
f22c2db 
ghstack-source-id: 724db559de21eeba6ece47e6fd7e6bfb29b5b979
Pull Request resolved: #124356
@Chillee
Update on "Optimized templated attention to use exp2"
a3b4bd1 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler amjames desertfire chauhang

[ghstack-poisoned]
Chillee added a commit that referenced this pull request Apr 18, 2024
@Chillee
Optimized templated attention to use exp2
88989b1 
ghstack-source-id: 838e82b3455a1dd02e8ddff6dcfd300073d5a8df
Pull Request resolved: #124356
@fmassa
Copy link
Member

fmassa commented Apr 18, 2024

cc @danthe3rd , as IIRC this optimization had potentially some implications wrt accuracy

@Chillee
Copy link
Contributor Author

Chillee commented Apr 18, 2024

@fmassa I guess we'll probably have a variety of configs for controlling it haha.

@Chillee
Copy link
Contributor Author

Chillee commented Apr 18, 2024

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 18, 2024
@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team Raised by workflow job

@drisspg drisspg added the topic: not user facing topic category label Apr 18, 2024
@Chillee
Copy link
Contributor Author

Chillee commented Apr 18, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your 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

Advanced Debugging
Check the merge workflow status
here

This was referenced Apr 19, 2024
Closed
Closed
This was referenced Apr 19, 2024
Open
Open
OnlyFor pushed a commit to OnlyFor/pytorch that referenced this pull request Apr 19, 2024
@Chillee
Optimized templated attention to use exp2
f8dc45a 
ghstack-source-id: ab2e15321d8c7ee262cd6b609dcb8031809b88ab
Pull Request resolved: pytorch#124356
pytorch-bot bot pushed a commit that referenced this pull request Apr 21, 2024
@Chillee @trieuat
Optimized templated attention to use exp2 (#124356)
c8461d4 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">

Pull Request resolved: #124356
Approved by: https://github.com/drisspg
sanketpurandare pushed a commit to sanketpurandare/pytorch that referenced this pull request Apr 22, 2024
@Chillee
Optimized templated attention to use exp2 (pytorch#124356)
bece7ec 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">

Pull Request resolved: pytorch#124356
Approved by: https://github.com/drisspg
andoorve pushed a commit to andoorve/pytorch that referenced this pull request May 1, 2024
@Chillee @andoorve
Optimized templated attention to use exp2 (pytorch#124356)
64c0f24 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">

Pull Request resolved: pytorch#124356
Approved by: https://github.com/drisspg
petrex pushed a commit to petrex/pytorch that referenced this pull request May 3, 2024
@Chillee @petrex
Optimized templated attention to use exp2 (pytorch#124356)
aa4534b 
0.705 (vs. FA2) to 0.860 after this change.

<img width="1270" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/d58f57ba-e50e-44ea-8a8a-4f13b8650adf">

to

<img width="1277" alt="image" src="https://github.com/pytorch/pytorch/assets/6355099/f1945b67-0cfc-463c-a2f6-5812b90677fe">

Pull Request resolved: pytorch#124356
Approved by: https://github.com/drisspg
facebook-github-bot pushed a commit to pytorch/benchmark that referenced this pull request May 8, 2024
@manman-ren @facebook-github-bot
template attention from PT2
1708afa 
Summary:
Based on inductor generated code, but modified to use Triton's tuning
pytorch github: pytorch/pytorch#124369
The base variant is prior to OSS pytorch/pytorch#124356. This PR improves performance for template attention. The second variant is after the PR.

Reviewed By: bertmaher

Differential Revision: D56372010

fbshipit-source-id: 4439113a92fd41b81269af1227deaf5ec52c65dc
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@drisspg drisspg drisspg approved these changes

Assignees
No one assigned
Labels
ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request Merged module: inductor topic: not user facing topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
4 participants
@Chillee @fmassa @pytorchmergebot @drisspg