Batch matmul fast path in MHAWithCache #449
Open
Conversation
This file contains 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
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
added a commit
to rohan-varma/multimodal
that referenced
this pull request
Aug 17, 2023
Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization in MHAWithCache. Differential Revision: D48418780 fbshipit-source-id: e8001eb870e827b05146221bb66f82939deae0c6
rohan-varma
force-pushed
the
export-D48418780
branch
from
e1233cc
to
dfd2ec6
Compare
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
added a commit
to rohan-varma/multimodal
that referenced
this pull request
Aug 17, 2023
Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization in MHAWithCache. Differential Revision: D48418780 fbshipit-source-id: 0501341832910bf90a7ea1cc902b98f0760548ab
rohan-varma
force-pushed
the
export-D48418780
branch
from
dfd2ec6
to
a2e0a70
Compare
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## main #449 +/- ##
==========================================
- Coverage 69.11% 69.11% -0.01%
==========================================
Files 170 170
Lines 11524 11530 +6
==========================================
+ Hits 7965 7969 +4
- Misses 3559 3561 +2
☔ View full report in Codecov by Sentry. |
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
added a commit
to rohan-varma/multimodal
that referenced
this pull request
Aug 18, 2023
…th (facebookresearch#449) Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization for MHA with cache in a new module `MultiHeadSelfAttentionWithCache`. Note: we are primarily using a new module to avoid breaking checkpoint BC with respect to `MultiHeadAttentionWithCache`. In the future, we should consolidate these MHA implementations. Differential Revision: D48418780 fbshipit-source-id: 5ad930ff27a4b131f8ff1f097a4c9e1548efb587
rohan-varma
force-pushed
the
export-D48418780
branch
from
a2e0a70
to
919dc03
Compare
rohan-varma
added a commit
to rohan-varma/multimodal
that referenced
this pull request
Aug 18, 2023
…th (facebookresearch#449) Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization for MHA with cache in a new module `MultiHeadSelfAttentionWithCache`. Note: we are primarily using a new module to avoid breaking checkpoint BC with respect to `MultiHeadAttentionWithCache`. In the future, we should consolidate these MHA implementations. Differential Revision: D48418780 fbshipit-source-id: eb0691e9d3a4bf729cfd7ca3293585c7d0108403
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
force-pushed
the
export-D48418780
branch
from
919dc03
to
6d67dae
Compare
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
added a commit
to rohan-varma/multimodal
that referenced
this pull request
Aug 18, 2023
…th (facebookresearch#449) Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization for MHA with cache in a new module `MultiHeadSelfAttentionWithCache`. Note: we are primarily using a new module to avoid breaking checkpoint BC with respect to `MultiHeadAttentionWithCache`. In the future, we should consolidate these MHA implementations. Differential Revision: D48418780 fbshipit-source-id: 0b20fb807527109a9a3ad419805e47e0f9ba2c74
rohan-varma
force-pushed
the
export-D48418780
branch
from
6d67dae
to
173699e
Compare
…th (facebookresearch#449) Summary: Pull Request resolved: facebookresearch#449 When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization for MHA with cache in a new module `MultiHeadSelfAttentionWithCache`. Note: we are primarily using a new module to avoid breaking checkpoint BC with respect to `MultiHeadAttentionWithCache`. In the future, we should consolidate these MHA implementations. Differential Revision: D48418780 fbshipit-source-id: 58f00205af26d39f778853c7aa50d560e024b9f8
|
This pull request was exported from Phabricator. Differential Revision: D48418780 |
rohan-varma
force-pushed
the
export-D48418780
branch
from
173699e
to
d459f16
Compare
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: When doing self attention, an optimization is to combine the Q, K, V input projection matrices and do a single matmul, instead of 3. Adding this optimization in MHAWithCache.
Differential Revision: D48418780