ROCm/jax

Pull Requests Commits

Use dq_ref type instead of fixed when converting dq

vgokhale committed 2 years ago

61126d58

Resolve conflicts for transposed loading

vgokhale committed 2 years ago

a44a4a6c

Add support for transposing tensor being loaded

Rahul Batra committed 2 years ago

97e2934c

Add FA bwd split kernels

Rahul Batra committed 2 years ago

35922952

Convert P to bf16 instead of f16

vgokhale committed 2 years ago

21e75404

Update FA to v2

vgokhale committed 2 years ago

d908637d

[ROCm] Jax-Triton returns hsaco path only for compiled kernel

Rahul Batra committed 2 years ago

6abbcf46

[ROCm]: Update docker file for ROCm6.0w

Rahul Batra committed 2 years ago

66fb5915

Lower block size of attention op to 32

mqhc2020 committed 2 years ago

50e25c15

Extract hsaco path from compiled kernel tuple

mqhc2020 committed 2 years ago

7965e1b2

[ROCm]: Update docker file for ROCm6.0w

Rahul Batra committed 2 years ago

ee929e56

Set blk size to 64 for attn opn and disable some pallas tests

Rahul Batra committed 2 years ago

933702a5

[ROCm]: Add ROCm Target

Rahul Batra committed 2 years ago

c99438d5

[ROCm]: Updates for ROCm6.0

Rahul Batra committed 2 years ago

172141f5

Add gfx942 to the list of supported devices

yuychang committed 2 years ago

70ae2a32

Implement workaround for GPU-related error handling in prng_kernels.cc

yuychang committed 2 years ago

599e55b5

Enhance GPU platform support for MI300 series

yuychang committed 2 years ago

d9789b6c

Add gfx940 and gfx941 to the list of supported devices

yuychang committed 2 years ago

4a8cfcae

[ROCm]: Disable some export tests

Rahul Batra committed 2 years ago

1a4f5ff8

[ROCm]: Comment out GPU_LOWERING_ENABLED in sparse_test

Rahul Batra committed 2 years ago

ef68c4ff

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

e6a04026

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

19da4c01

[ROCm]: Update build script

Rahul Batra committed 2 years ago

9fa53699

Merge pull request #18325 from skye:version

a-googler committed 2 years ago

e22b8bbc

Update versions for jax + jaxlib 0.4.20 release

skye committed 2 years ago

68138191

Fix typing annotations for `@jax.named_call`

Conchylicultor committed 2 years ago

81ac67f3

Update XLA dependency to use revision

a-googler committed 2 years ago

1c66ac53

Properly pack and unpack int4 arrays on CPU in PJRT.

reedwm committed 2 years ago

d41078fb

[Mosaic] Add support for generalized reduction unrolling to infer_vector_layout

apaszke committed 2 years ago

5d289615

Add support for dynamic indices in VMEM loads and stores

apaszke committed 2 years ago

f17f5492

Older