pytorch/xla

Pull Requests Commits

fix complication and runtime check

yaochengji committed 240 days ago

e99636c3

yaochengji committed 240 days ago

67068c40

Fix ragged paged attention v2 to resolve the recompilation issue

yaochengji committed 240 days ago

7dcc8788

Integrate ragged paged attention v2 (#8791)

bythew3i committed 240 days ago

Verified 5644f44a

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

tengyifei committed 242 days ago

Verified 17270e2f

Build torch_xla wheel in build script (#8782)

zpcore committed 242 days ago

Verified 00fac787

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

qihqi committed 243 days ago

Verified cee0820e

Introduce a GRU module implemented with scan (#8777)

tengyifei committed 243 days ago

Verified 0017ff16

Fix build on Python 3.9 (#8759)

tengyifei committed 245 days ago

Verified c76e9497

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

qihqi committed 245 days ago

Verified 60001c74

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

tengyifei committed 245 days ago

Verified 2feb0ac7

Add sm_scale in ragged attention kernel (#8771)

yaochengji committed 246 days ago

Verified 2675e689

Revert "Change num_seqs type from int to torch.Tensor" (#8767)

yaochengji committed 246 days ago

Verified 2e4f073a

Enable default buffer donation for gradient accumulation (#8758)

rpsilva-aws committed 246 days ago

Verified 6712eb91

Change num_seqs type from int to torch.Tensor (#8736)

vanbasten23 committed 246 days ago

Verified b18a65fa

Lower `as_strided_copy` use fast path with `slice` (#8734)

zpcore committed 247 days ago

Verified 1ab8216a

Add start_trace and stop_trace API in profiler (#8743)

lsy323 committed 247 days ago

Verified b4ba17b5

Add --network=host to TPU docker build command (#8735)

bhavya01 committed 247 days ago

Verified 71d0ce65

Increase tolerance in test_scan_xla_patched_linear (#8749)

tengyifei committed 247 days ago

Verified ac430abf

Extend buffer donation aliasing APIs (#8721)

rpsilva-aws committed 247 days ago

Verified 07976424

Introduce apply_xla_patch_to_nn_linear and test that in a scan (#8739)

tengyifei committed 248 days ago

Verified 6f020aad

Add padding ragged paged attention test (#8741)

vanbasten23 committed 248 days ago

Verified 5e1478d6

Improve the ragged kernel benchmarking script. (#8733)

vanbasten23 committed 249 days ago

Verified 79d1e86b

Add autograd function for mark_sharding (#8723)

bhavya01 committed 253 days ago

Verified 1acc987a

Add information to xla.launch (#8724)

pgmoka committed 253 days ago

Verified a08b7f86

Split page indices in the ragged paged attention. (#8688)

vanbasten23 committed 254 days ago

Verified 0c4e82d6

Bump CUDNN version in docker image. (#8718)

ysiraichi committed 255 days ago

Verified 60160233

Bump XLA commit pin. (#8714)

ysiraichi committed 256 days ago

Verified ca248c7e

Support splitting physical axis in SPMD mesh (#8698)

tengyifei committed 259 days ago

Verified 8a712495

update pin (#8677)

lsy323 committed 261 days ago

Verified 6f423d0b

Older