PR #6785 [Pallas] Make a FlashAttention Wrapper

[Pallas] Make a FlashAttention Wrapper #6785

alanwaketan merged 6 commits into master from alanwaketan/flash_attention_3

tmp

0a44bb6e

tmp

373472de

introduce flash_attention

66026e56

Add test case

fe63be99

Fix the test

522de564

Fix linters

c36fb510

alanwaketan requested a review from

will-cromar 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

alanwaketan added backport_2.3

JackCaoG approved these changes on 2024-03-20

alanwaketan merged fcf24b6c into master 1 year ago

alanwaketan deleted the alanwaketan/flash_attention_3 branch 1 year ago

Reviewers

JackCaoG

will-cromar

Assignees

alanwaketan

Labels

backport_2.3

Milestone

No milestone