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

[Pallas] Make a FlashAttention Wrapper #6785

Merged
merged 6 commits into from
Mar 20, 2024
Merged

[Pallas] Make a FlashAttention Wrapper #6785

merged 6 commits into from
Mar 20, 2024

Conversation

alanwaketan
Copy link
Collaborator

Summary:
This pull request introduces a FlashAttention wrapper that aims to:

  1. Override some default settings for the best performance out of box.
  2. Ease the UX such that users don't need to do all the custom_kernel paper works.

Test Plan:
PJRT_DEVICE=TPU python test/test_pallas.py -v -k test_flash_attention_wrapper

@alanwaketan alanwaketan self-assigned this Mar 20, 2024
@alanwaketan
Copy link
Collaborator Author

Thanks Jack for approving.

@alanwaketan alanwaketan merged commit fcf24b6 into master Mar 20, 2024
18 checks passed
@alanwaketan alanwaketan deleted the alanwaketan/flash_attention_3 branch March 20, 2024 23:10
@alanwaketan alanwaketan mentioned this pull request Mar 27, 2024
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@JackCaoG JackCaoG JackCaoG approved these changes

@will-cromar will-cromar Awaiting requested review from will-cromar

Assignees

@alanwaketan alanwaketan

Labels
backport_2.3
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@alanwaketan @JackCaoG