Use of new JAX Performance Flags?

[adam-hartshorne]

Could setting any of these flags cause issues with functions under Torch2Jax (as you know, I use Keops a lot)

https://github.com/google/jax/blob/664e834784a35117c244030c28862177dcfc76f0/docs/gpu_performance_tips.md

[rdyro]

The flags specifically are these if I understand correctly:

• '--xla_gpu_enable_triton_softmax_fusion=true '
• '--xla_gpu_triton_gemm_any=True '
• '--xla_gpu_enable_async_collectives=true '
• '--xla_gpu_enable_latency_hiding_scheduler=true '
• '--xla_gpu_enable_highest_priority_async_stream=true '

Triton should no cause any issues.

Async options, the latter 3 flags are trickier to judge. torch2jax uses a conservative synchronization approach, so I would not expect it to cause any issues.

My best guess is those flags should probably not cause any issues.

I'll rerun the test suite with these flags later this week to give a more definitive answer. I'll leave the issue open until then.

[rdyro]

The tests pass with those flags included, so hopefully, you will have no problems.

Please reopen this issue if you encounter any issues.