xla
[Backport] Add flash attention as custom op so dynamo can trace it (#6875)
#6884

Merged

[Backport] Add flash attention as custom op so dynamo can trace it (#6875) #6884

JackCaoG merged 4 commits into r2.3 from JackCaoG/backport_flash_attention_dynamo

Add flash attention as custom op so dynamo can trace it (#6875)

2298cbcb

JackCaoG requested a review from

lsy323 2 years ago

JackCaoG requested a review from

alanwaketan 2 years ago

alanwaketan approved these changes on 2024-04-04

change atol precision

e6f320ec

Change upstream to r2.3 for TPU CI

31329c9a

fix branch name

89c11122

JackCaoG merged 96cc98ff into r2.3 2 years ago

Reviewers

alanwaketan

lsy323

Assignees

No one assigned

Labels

None yet

Milestone

No milestone