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

[XLA:GPU] add cuDNN flash attention support in XLA (3rd PR with only rewriter changes) #6872

Closed
wants to merge 16 commits into from
Closed

[XLA:GPU] add cuDNN flash attention support in XLA (3rd PR with only rewriter changes) #6872

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
4b1505c
init flash attention rewriter
Cjkkkk Nov 8, 2023
745c2ad
use while body back pointer to find causal mask
Cjkkkk Nov 22, 2023
5f95a05
add gpu backend to fmha e2e tests && address some format issues
Cjkkkk Nov 27, 2023
4bbd7df
fix rebase error
Cjkkkk Nov 28, 2023
d1d9f2c
Use GPT3_5B model pre rewriter HLo
Cjkkkk Nov 30, 2023
84e9ff3
add flash attention cuDNN version check && restore fwd graph is dbias…
Cjkkkk Nov 30, 2023
3969975
fix case with no bias but also no causal_mask
Cjkkkk Dec 7, 2023
f529af2
remove unused branch
Cjkkkk Dec 7, 2023
9af8953
rebased and address some comments
Cjkkkk Dec 14, 2023
73928b4
make causal mask/bias both optional
Cjkkkk Dec 17, 2023
7c573dc
address some comments and fix wrong layout for softmax stat if O is n…
Cjkkkk Dec 21, 2023
7390db9
add flash attention cross attention with cuDNN > 8.9.4
Cjkkkk Dec 27, 2023
f8db19b
fix fwd graph restore
Cjkkkk Jan 5, 2024
031a11a
fix rebase error
Cjkkkk Jan 5, 2024
f5c6b03
add guard for optional dropout_rate
Cjkkkk Jan 10, 2024
42d55c8
fix gpubackend config
Cjkkkk Jan 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading