xla
[Pallas] Make FlashAttention as torch.autograd.Function
#6886

Merged

[Pallas] Make FlashAttention as torch.autograd.Function #6886

alanwaketan merged 9 commits into master from alanwaketan/fa_autograd

alanwaketan requested a review from

lsy323 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

Initial commit

d18e1082

Fix the test

39e1dead

Add a testcase somewhat

675c8b34

Fix linters

02a15f0b

alanwaketan force pushed from 5b316baa to 02a15f0b 1 year ago

Fix the test

f58c8e78

JackCaoG commented on 2024-04-04

alanwaketan commented on 2024-04-04

JackCaoG commented on 2024-04-04

Imporve lmi

3cc5271b

Fix linters

36027f0f

Address commetns

2c71054b

Minor fix

dcb859c5

JackCaoG approved these changes on 2024-04-04

alanwaketan merged 0c704cf8 into master 1 year ago

Reviewers

JackCaoG

lsy323

Assignees

alanwaketan

Labels

None yet

Milestone

No milestone