xla
[Pallas] Improve FlashAttention segment_ids test case
#7034

Merged

[Pallas] Improve FlashAttention segment_ids test case #7034

alanwaketan merged 6 commits into master from alanwaketan/segment_ids_2

tmp

407e83f2

some minor changes

c28f379d

Fix everything

d81d8699

Enable backward as well

3470ac3b

Delete customer code

9d6955ee

Improve test case

faf570e9

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

JackCaoG approved these changes on 2024-05-07

alanwaketan merged 887d3446 into master 1 year ago

Reviewers

JackCaoG

Assignees

alanwaketan

Labels

None yet

Milestone

No milestone