openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.58k stars 403 forks source link

XLA flags: No speed ups on GPUs and segmentation fault #17103

Open AakashKumarNain opened 1 week ago

AakashKumarNain commented 1 week ago

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 commented 6 days ago

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

hawkinsp commented 6 days ago

@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.