rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00

History

Sergei Lebedev 7929cd8410 [pallas:triton] The lowering now uses PTX instead of Triton IR

This change improves the stability and backward compatibility of Pallas Triton
calls, because unlike PTX, the Triton dialect has no stability guarantees
and does change in practice.

See #25196.

A few notes

* Pallas Triton no longer delegates compilation to PTX to XLA:GPU. Instead,
  compilation is done via a new PjRt extension, which uses its own compilation
  pipeline mirrored after the one in the Triton Python bindings.
* The implementation of the old custom call used by Pallas Triton is
  deprecated and will be removed after 6 months as per
  [compatibility guarantees] [*]

[*]: https://jax.readthedocs.io/en/latest/export/export.html#compatibility-guarantees

PiperOrigin-RevId: 722773884

2025-02-03 13:21:40 -08:00

__init__.py

[pallas:triton] The lowering now uses PTX instead of Triton IR

2025-02-03 13:21:40 -08:00

BUILD.bazel

Replace pjrt_c_api_gpu_plugin.so symlink with XLA dependency.

2024-08-12 13:01:18 -07:00

plugin_pyproject.toml

[PJRT C API] Rename the folder "plugins" to "jax_plugins".

2023-11-14 13:56:13 -08:00

plugin_setup.py

[cuda] Bump nvidia-cuda-nvcc-cu12 dependency to 12.6.85