xla
Support collective matmul optimization in mp
#8855

Merged

Support collective matmul optimization in mp #8855

yaochengji merged 10 commits into master from chengji/enable-cm

yaochengji requested a review from

tengyifei 228 days ago

yaochengji requested a review from

lsy323 228 days ago

yaochengji requested a review from

zpcore 228 days ago

tengyifei requested changes on 2025-03-19

init

b3d9739e

add tests

ed1f9f7e

fix cpp format

5ffbd7a2

format python

676f5b8d

yaochengji force pushed from 150779b9 to 676f5b8d 228 days ago

update libtpu date to 20250320

b84f5c12

pgmoka commented on 2025-03-20

bump libtpu to 0.0.12 nightly

e3302034

Update jax to 0.5.4

61f35290

Run the MP tests on TPU

ff5e6603

seperate tests in another file

338db2d3

fix lint

82d089b9

tengyifei approved these changes on 2025-03-21

yaochengji merged 2ad5b5d3 into master 226 days ago

Reviewers

tengyifei

pgmoka

lsy323

zpcore

Assignees

No one assigned

Labels

None yet

Milestone

No milestone