XLA flags: No speed ups on GPUs and segmentation fault

[AakashKumarNain]

I am developing some code in Equinox and JAX, and running it on A100 GPUs. As per the JAX GPU tips, some operations should provide better performance.

On my end, I don't see any difference in the performance after enabling these flags. Also, some flags like --xla_gpu_enable_triton_softmax_fusion=true results in segmentation fault. Please check the detailed description in the related issue

[cheshire]

Performance flags are not considered to be stable/supported/safe API to use.

[hawkinsp]

@cheshire I don't think that's a realistic description of the current state. The current state is that because we haven't succeeded at landing many useful flags as defaults, lots of users can and do have to override XLA flags.

While we may not expect performance improvements from doing that, I think the segfault at least is a real bug and shouldn't be summarily closed.

[bchetioui]

@AakashKumarNain thanks for pointing this out to me. Do you have a reproducer for your segmentation fault?

I have not been able to reproduce the crash with the code provided in jax-ml/jax#22705 (comment), at least.

If there is no way to re-trigger the crash, then I'm tempted to close this bug as "WAI" for the performance issue for now. (If this worked out of the box every time, it wouldn't be a flag :))

[AakashKumarNain]

If there is no way to re-trigger the crash, then I'm tempted to close this bug as "WAI" for the performance issue for now. (If this worked out of the box every time, it wouldn't be a flag :))

I will try to reproduce it on my side and will update

[AakashKumarNain]

@bchetioui I tested this with the latest version, and this didn't cause any more segfault. But, it didn't improve any performance either