xla
[Pallas] Support Flash Attention backward kernels
#6870

Merged

[Pallas] Support Flash Attention backward kernels #6870

alanwaketan merged 4 commits into master from alanwaketan/fa_backward

support test__flash_attention_impl

1ce54b1a

Support test__flash_attention_bwd_dkv

cf9a1d07

Support test__flash_attention_bwd_dkv

e01dead5

Fix linters

99dc7c0a

alanwaketan requested a review from

lsy323 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

JackCaoG commented on 2024-04-02

JackCaoG approved these changes on 2024-04-02

alanwaketan merged c54367c8 into master 1 year ago

Reviewers

JackCaoG

lsy323

Assignees

alanwaketan

Labels

None yet

Milestone

No milestone