Comparison with XFormers or other acceleration libraries

[sh0416]

I just experience that xformers library actually provides notable improvement when compiling the binary with given specific hardware architecture.

It seems better than jax.

Could jax make these kind of improvement when we build from source?

[mattjj]

Thanks for the question!

The numerical code that JAX executes is just-in-time (JIT) compiled for whatever specific hardware you run it on, right before it's executed. That includes things like autotuning for your particular GPU, for example.

So if I understand correctly, I don't think there's a change to be made here.

[sh0416]

Then, is there any other acceleration library for enhancing the speed by leveraging the specific hardware architecture?

For example, PyTorch has different numerous acceleration library such as triton, xformers, fastertransformer, megatron-lm, flexgen, fairscale.

I just want to integrate them with jax.

[mattjj]

Well, that's kind of a broad question, but maybe the jax-triton package would be interesting to look at?

[sh0416]

That would help! Thank you mattjj!