PR #6872: [XLA:GPU] add cuDNN flash attention support in XLA (3rd PR with only rewriter changes) #7593
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
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
force-pushed
the
test_588714363
branch
from
a46a369
to
633fc7f
Compare
…with only rewriter changes) Imported from GitHub PR #6872 This is the 3nd PR of splitting #5910 with only rewriter changes 1st PR #6293 merged. 2nd PR #6657 merged. * Add pattern match for causal mask * Add paxml dropout pattern match * Add flash attention fusion * Add flash attention support cuDNN version guard * Add tests for flash attention rewriter/e2e Copybara import of the project: -- 490d0a3 by cjkkkk <ske@nvidia.com>: init flash attention rewriter -- 90e765f by cjkkkk <ske@nvidia.com>: use while body back pointer to find causal mask -- a3e5905 by cjkkkk <ske@nvidia.com>: add gpu backend to fmha e2e tests && address some format issues -- c82c064 by cjkkkk <ske@nvidia.com>: fix rebase error -- 2f30df0 by cjkkkk <ske@nvidia.com>: Use GPT3_5B model pre rewriter HLo -- 47aceb1 by cjkkkk <ske@nvidia.com>: add flash attention cuDNN version check && restore fwd graph is dbias/mask is not supported Merging this change closes #6872 FUTURE_COPYBARA_INTEGRATE_REVIEW=#6872 from Cjkkkk:flash_attention_rewriter 47aceb1 PiperOrigin-RevId: 588714363
copybara-service
bot
force-pushed
the
test_588714363
branch
from
633fc7f
to
2e39469
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.
PR #6872: [XLA:GPU] add cuDNN flash attention support in XLA (3rd PR with only rewriter changes)
Imported from GitHub PR #6872
This is the 3nd PR of splitting #5910 with only rewriter changes
1st PR #6293 merged.
2nd PR #6657 merged.
Copybara import of the project:
--
490d0a3 by cjkkkk ske@nvidia.com:
init flash attention rewriter
--
90e765f by cjkkkk ske@nvidia.com:
use while body back pointer to find causal mask
--
a3e5905 by cjkkkk ske@nvidia.com:
add gpu backend to fmha e2e tests && address some format issues
--
c82c064 by cjkkkk ske@nvidia.com:
fix rebase error
--
2f30df0 by cjkkkk ske@nvidia.com:
Use GPT3_5B model pre rewriter HLo
--
47aceb1 by cjkkkk ske@nvidia.com:
add flash attention cuDNN version check && restore fwd graph is dbias/mask is not supported
Merging this change closes #6872
FUTURE_COPYBARA_INTEGRATE_REVIEW=#6872 from Cjkkkk:flash_attention_rewriter 47aceb1