Flash attention - #931

Open

oliverdutton wants to merge 9 commits into

google-deepmind:mainfrom

oliverdutton:flash_attention

Commits on Apr 19, 2024

feat: flash attention
oliverdutton
committed
f86ad29

Commits on Apr 20, 2024

feat: allow bfloat16 flash attention
oliverdutton
committed
d4fdf81

Commits on Apr 21, 2024

fix: support batch dim broadcast mask
oliverdutton
committed
789d876
fix: add null callback, lowers to nan avoiding MHLO
oliverdutton
committed
ef77bf6

Commits on Apr 22, 2024

Commits on Apr 23, 2024

fix: mask=None case bug
oliverdutton
committed
1787c30