Calling Triton kernels from Pallas

[vvvm23]

Recently there have been a lot of interesting kernels written in triton appearing, there exists jax-triton to call triton kernels using jax arrays. However, jax-triton is not very portable and easily installable versus jax itself, which does not require jax_triton or triton anymore to compile Pallas kernels to GPU.

This might be a slightly naive ask, but given we can compile Pallas kernels with jax alone and Pallas uses Triton IR MLIR (which triton also uses), is there a way (hacky it may be) of simply calling Triton kernels using Pallas directly?

Thanks~

[hawkinsp]

Can you try installing jax-triton from git head? We just relaxed some of its installation constraints in the last few days. It should be fairly easy to install now, because you can use the released JAX and the released Triton: we're eager that that remain true, too. Does that work? We should probably make a release on pypi soon, too.

It would be difficult to compile Triton IR from a Triton release with the copy of Triton inside JAX, because JAX closely tracks Triton head, and there is no guarantee of compatibility between your copy of Triton and ours. So we think jax-triton (now easier to install) should be the way: that way the coupling between JAX, Triton and jax-triton is much looser.

Let us know!