pytorch/xla

Pull Requests Commits

yaochengji committed 305 days ago

05d97d78

add customized ring in jax

yaochengji committed 307 days ago

f7e27bea

improve collective matmul perf

yaochengji committed 307 days ago

cdc8f584

add jax spmd cm perf

yaochengji committed 309 days ago

bc032217

add jax spmd script

yaochengji committed 310 days ago

673ba516

add perf script

yaochengji committed 310 days ago

658d8573

yaochengji committed 310 days ago

055f66ab

add the correctness check code from siyuan

yaochengji committed 310 days ago

36bbbd11

support reduce scatter

yaochengji committed 310 days ago

0f841b67

yaochengji committed 310 days ago

a25968bf

Fix the ragged kernel bug (#8806)

vanbasten23 committed 311 days ago

Verified 6804c65f

Document xla override (#8802)

pgmoka committed 312 days ago

Verified 0407d6d1

Leverage placeholder tensors for GA API (#8808)

rpsilva-aws committed 312 days ago

Verified 5b7462c6

Manually register einsum on xla (#8801)

pgmoka committed 314 days ago

Verified 1db6c861

Misc fixes (#8804)

qihqi committed 314 days ago

Verified 7aea0585

Re-enable triton tests. (#8803)

ysiraichi committed 315 days ago

Verified 2c70a1c2

Lower `_conj_copy` and `alias` operation. (#8686)

ysiraichi committed 315 days ago

Verified 6c53a1ef

Fix ragged paged attention v2 to resolve the recompilation issue (#8797)

yaochengji committed 316 days ago

Verified d06a9c97

Pin update to 20250303 (#8788)

tengyifei committed 316 days ago

Verified 9b61c1ad

Use placeholder tensor in scan (#8785)

tengyifei committed 316 days ago

Verified a9279287

Use .backward() with in-place grad mutations for the GA API (#8768)

rpsilva-aws committed 316 days ago

Verified 4540d817

Integrate ragged paged attention v2 (#8791)

bythew3i committed 316 days ago

Verified 5644f44a

Minimal support for calling JAX from PyTorch/XLA (#8781)

tengyifei committed 317 days ago

Verified 17270e2f

Build torch_xla wheel in build script (#8782)

zpcore committed 317 days ago

Verified 00fac787

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

qihqi committed 318 days ago

Verified cee0820e

Introduce a GRU module implemented with scan (#8777)

tengyifei committed 318 days ago

Verified 0017ff16

Fix build on Python 3.9 (#8759)

tengyifei committed 321 days ago

Verified c76e9497

Misc changes to make torchax runnable on GPU. (#8756)

qihqi committed 321 days ago

Verified 60001c74

[scan] Make sure inputs into `fn` are not `device_data` IR nodes (#8769)

tengyifei committed 321 days ago

Verified 2feb0ac7

Add sm_scale in ragged attention kernel (#8771)

yaochengji committed 321 days ago

Verified 2675e689

Older