xla
write _shard_map; refactor flash attention to support 5d inputs.
#8730

Merged

write _shard_map; refactor flash attention to support 5d inputs. #8730

qihqi merged 9 commits into master from hanq_shard_map

qihqi force pushed from d5c7ae7e to af2bbf64 1 year ago

pgmoka commented on 2025-02-24

qihqi force pushed from e113414c to 0a27ca43 1 year ago

qihqi force pushed from 0a27ca43 to 8a5a29f5 1 year ago

qihqi requested a review from

pgmoka 1 year ago

qihqi changed the title ~~write _shard_map; refactor flash attention to use it.~~ write _shard_map; refactor flash attention to support 5d inputs. 1 year ago

qihqi requested a review from

tengyifei 1 year ago

qihqi force pushed from 8a5a29f5 to d885e7b9 1 year ago

qihqi force pushed from d885e7b9 to 8983df7c 1 year ago

tengyifei requested changes on 2025-02-28

pgmoka commented on 2025-02-28

qihqi requested a review from

pgmoka 1 year ago

qihqi requested a review from

tengyifei 1 year ago

write _shard_map; refactor flash attention to use it.

d9536310

fix error & comment

00a78438

add 5d output

576cd3dc

stash

6ccaf033

Fix tests

651fc5c3

lint

33423308

fix unit tests

4677c455

yapf

1119146c

qihqi force pushed from cecefdec to 876b7710 1 year ago

tengyifei approved these changes on 2025-03-01

comments

b14a500f

qihqi force pushed from 876b7710 to b14a500f 1 year ago

qihqi merged cee0820e into master 1 year ago

Reviewers

tengyifei

pgmoka

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

xla write _shard_map; refactor flash attention to support 5d inputs. #8730 Merged

write _shard_map; refactor flash attention to support 5d inputs. #8730

xla
write _shard_map; refactor flash attention to support 5d inputs.
#8730

Merged