xla
[Pallas] Integrate FlashAttention with SPMD
#6935

Merged

[Pallas] Integrate FlashAttention with SPMD #6935

alanwaketan merged 9 commits into master from alanwaketan/fa_spmd

Add an e2e test

01a5f871

initial commit

9bf7c12d

Support forward

5a89157b

Enable backward

899735f8

Improve test case

401ca633

Fix linters

51f51a07

alanwaketan requested a review from

yeounoh 1 year ago

alanwaketan requested a review from

jonb377 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

JackCaoG approved these changes on 2024-04-17

jonb377 approved these changes on 2024-04-17

Fix comments

944b9983

Fix tests

a32140b9

Fix tests

8f97ea3f

alanwaketan merged 9f2b82dc into master 1 year ago

Reviewers

jonb377

JackCaoG

yeounoh

Assignees

alanwaketan

Labels

None yet

Milestone

No milestone