ROCm/jax

Pull Requests Commits

Adatption input sizes based ROCm hardware for gpu_paged_attention_test testcase

phambinhfin committed 90 days ago

36ee5641

Fix v_pages quantization and adjust test params for ROCm compatibility

phambinhfin committed 90 days ago

52652684

Fix hipSparse CSR algorithm mappings for ROCm 7 (#557)

phambinhfin committed 92 days ago

Verified ccf941ee

Make nvidia version data optional for ROCm builds (#552)

phambinhfin committed 102 days ago

Verified fd3e2c6d

Phambinh/jax fix copy (#554)

phambinhfin committed 102 days ago

Verified 3206662a

Skip test_dot_precision() with TF32 for gfx11xx and gfx12xx architecture (#520) (#551)

AleksaArsic committed 112 days ago

Verified 3b02da3a

The ROCm plugin should fail gracefully and allow for other devices to… (#550)

JehandadKhan committed 118 days ago

Verified 0011fb97

Porting v0.6.0 changes to v0.7.1 (#549)

zahiqbal committed 132 days ago

Verified a5d5fa0d

applied rocm-jax jax build fixes patches (#541)

zahiqbal committed 147 days ago

Verified 10d3836e

Prepare for JAX release 0.7.1

jakeharmon8 committed 182 days ago

5712de44

[Mosaic GPU] Separate hardware-agnostic from hardware-specific passes in BUILD files.

bchetioui committed 182 days ago

62482050

[Mosaic GPU] Annotate alignment of GMEM pointers to be 256

apaszke committed 182 days ago

a44869e1

[Mosaic GPU] Second attempt at fixing `test_tcgen05_collective_mma`.

allanrenucci committed 182 days ago

16371553

[Mosaic GPU] Attempt no. 2 at fixing the Mac OS build.

bchetioui committed 182 days ago

ffc0afb0

[Mosaic GPU] Do not predicate collective ops in the lowering.

allanrenucci committed 182 days ago

7ad15086

[Mosaic GPU] Fold non-gather indices into the gather TMA column index

apaszke committed 182 days ago

cbd0844e

Update XLA dependency to use revision http://github.com/openxla/xla/commit/31eb6029e8973337ef6c99810c27e1d807791c11

Google-ML-Automation committed 182 days ago

89628e45

[Mosaic GPU] Include `unistd.h` directly in order to fix build on Mac.

bchetioui committed 182 days ago

b09cc05f

Merge pull request #31151 from mattjj:fix-error-message

Google-ML-Automation committed 182 days ago

e44e8e65

fix a scan tree mismatch error

mattjj committed 182 days ago

cab4a539

[Mosaic:TPU] Mark tpu.relayout as pure with vector input/output, and clean up verifier and apply-vector-layout rule

tlongeri committed 182 days ago

f9dcd81b

Fix splash attention for non-128 divisible v head_dim

rdyro committed 183 days ago

c9a5799a

[Mosaic:TPU][Relayouts] Fix bug after cl/794807848

tlongeri committed 183 days ago

14c77fb1

Merge pull request #31140 from jakevdp:einsum-error

Google-ML-Automation committed 183 days ago

b438f851

Improve the error you get when you fail to construct an offloading policy. Or if the policy returns anything that's not a bool or `Recompute`, `Saveable` or `Offloable(...)`.

yashk2810 committed 183 days ago

d61c10e6

Merge pull request #31142 from jakevdp:ad-comment

Google-ML-Automation committed 183 days ago

1cac6f98

Fix incorrect comment in ad.py

jakevdp committed 183 days ago

335f3652

jnp.einsum: improve error for unsupported optimizations

jakevdp committed 183 days ago

d405ce2a

Merge pull request #31133 from jakevdp:polygamma-doc

Google-ML-Automation committed 183 days ago

31695257

[doc] typo in polygamma docs

jakevdp committed 183 days ago

0928d44a

Older