Feature request: lower jnp.exp2 to tl.math.exp2 in pallas

[hr0nix]

The flash attention implementation in the triton repo uses exp2, but I can't use the same trick with pallas and I suspect it might be important.

[sharadmv]

This might be a pretty simple PR if you are up for contributing. Otherwise I will get around to it in a week or so.

If interested, take a look at jax_triton/pallas/triton_lowering.py.

[sharadmv]

Maybe look at an existing unary primitive lowering such as exp or tanh.

[hr0nix]

The only problem is that I'm unsure whether I can support exp2 in the same way, given that in jax it's not a lax primitive, but rather a function implemented in terms of other lax primitives.

[sharadmv]

Interesting...we might want to change that in JAX then! cc: @froystig @mattjj

[hr0nix]

@froystig @mattjj Hey guys! Can you comment on whether this can be fixed in JAX in the foreseeable future?

[mattjj]

Sure, let's do it!

[mattjj]

jax-ml/jax#16883 fixes the JAX side. Once we merge that I can probably update jax-triton too (unless you want to take it @hr0nix )

[hr0nix]

Oh, I was too slow to react. Thanks a lot for a quick fix, guys!

[jon-chuang]

Hello guys, I tried using exp2 for this trick in Triton:

https://github.com/openai/triton/blob/cc45c3826fecabc2d6882e4bd94464ab39c8f730/python/tutorials/06-fused-attention.py#L186

However, the bwd tests fail -

For just 13 out of 49K elements! How odd! It may simply be the extreme values which result in large floating point errors.

In my case, I implemented the trick for both the backward and fowards passes in https://github.com/google/jax/blob/main/jax/experimental/pallas/ops/attention.py. Just implementing for forward or backward alone resulted in many more floating point errors.

Just as an anecdote - do let me know if any of you guys succeed with this.

[jon-chuang]

I found the issue! Scaling before/after the float32 conversion makes all the difference!

Now implemented in: jax-ml/jax#17328