ROCm/jax

Pull Requests Commits

Fixed perf regression introduced by checkin "fix correctness, concurrency, multi-GPU and FFI safety in aiter mha"

zahiqbal committed 156 days ago

06fd7b81

aiter libs alternative path to local dir, --repo_env=LOCAL_AITER_PATH=/path/to/aiter-libs

zahiqbal committed 160 days ago

46d314d4

fix correctness, concurrency, multi-GPU and FFI safety in aiter mha

zahiqbal committed 162 days ago

1d827fcb

code review updates

zahiqbal committed 163 days ago

ca02de70

corrected mha asm_v3 vs ck kernel calls

zahiqbal committed 164 days ago

b616cf94

upgraded aiter libs

zahiqbal committed 165 days ago

121df61e

handling aiter-mha libs in wheel at rocm/jax fork

zahiqbal committed 165 days ago

f461fcf3

code cleanup and formatting after code review

zahiqbal committed 169 days ago

720431d3

aiter mha kernels (ASM+CK) integration

zahiqbal committed 179 days ago

a6041a3b

Add support for parallelizing lapack invocations on the batch dimension.

hawkinsp committed 180 days ago

bfd9e1fb

[Mosaic TPU] Allow inferring divisibility through divisions

apaszke committed 180 days ago

3ba0e1f2

[Pallas:MGPU] Make sure that the warp axis returns indices in [0, 3)

apaszke committed 180 days ago

5bd8a46e

Merge pull request #36559 from superbobry:main

Google-ML-Automation committed 180 days ago

e0696314

Upgraded to Pyrefly 0.60.0

superbobry committed 180 days ago

34bf5455

[Mosaic GPU] Enable layout inference for collective scaled MMA.

allanrenucci committed 180 days ago

03cf9812

[mGPU] Use CPU target config, that is passed from XLA side via FFI.

khasanovaa committed 180 days ago

8708d738

add a test for a simple elementwise unary kernel

ImanHosseini committed 180 days ago

6e73777c

[Mosaic GPU] Add an `AnyOf` constraint to layout inference.

allanrenucci committed 180 days ago

7141987a

Merge pull request #36549 from jakevdp:pyrefly-059

Google-ML-Automation committed 180 days ago

cbc13839

[pallas:mgpu] Allow `commit_smem_to_gmem_group` in a single warp context.

chr1sj0nes committed 180 days ago

9bf6adf3

Remove the tools/build_defs.bzl

gnecula committed 180 days ago

254515d8

Update XLA dependency to use revision http://github.com/openxla/xla/commit/1924eec184be486bce582db95964636696ea3f03

Google-ML-Automation committed 180 days ago

8552ae99

[Mosaic GPU][NFC] Add `orders_tensor_core` to DialectBarrierRef.

allanrenucci committed 180 days ago

beed3313

Merge pull request #36525 from gnecula:consts_2

Google-ML-Automation committed 180 days ago

8c1da538

Make sure vary_unreduced_cast and reduced_vary_cast order the axes wrt the context mesh. Also code-golf some code in shard_map

yashk2810 committed 180 days ago

6f55c374

Merge pull request #36551 from jakevdp:eig-autodiff-docs

Google-ML-Automation committed 180 days ago

a6285551

[pmap] Create new ifrt array when rewrapping global arrays as host local in multihost pmap.

danielsuo committed 180 days ago

6bad8fcf

Merge pull request #36535 from jakevdp:numpy-dtype-dep

Google-ML-Automation committed 180 days ago

1b046188

[doc] update docs for eig() autodiff

jakevdp committed 180 days ago

18623265

[pyrefly] update pyrefly to v0.59.1 & fix/silence errors

jakevdp committed 180 days ago

f1935ff1

Older