Triton/XLA on ROCm has very bad performance

[Steeve Morin]

Hi,

When running a custom paged attention kernel, we get about ~10x performance hit on a Triton kernel when running on Triton/XLA than on Triton/PyTorch. We are using the kernel at https://github.com/ROCm/aiter/blob/main/aiter/ops/triton/pa_decode.py#L309

We're not sure what can cause this, the grids are the same. We saw that ThreadPerWarp is hardcoded to 32 in XLA (should be 64), but I'm not sure that might explain it.

When trying this Pallas implementation at https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/gpu/paged_attention.py in JAX, we get ~750us on H200 but 20ms (!) on MI300x.

Any help to narrow it down would be appreciated.

Thank you!

[Florian Reichl]

Hi Steve,

I see you've also opened https://github.com/openxla/xla/issues/23574 and are getting support there, correct?

Best,

Florian

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/2febb308-3d20-4436-837f-9a26e040d9b7n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

Florian Reichl  |  Software Engineering Manager  |  Munich

fre...@google.com

Google Germany GmbH

Erika-Mann-Straße 33

80636 München

Geschäftsführer: Paul Manicle, Liana Sebastian

Registergericht und -nummer: Hamburg, HRB 86891

Sitz der Gesellschaft: Hamburg

Diese E-Mail ist vertraulich. Falls Sie diese fälschlicherweise erhalten haben sollten, leiten Sie diese bitte nicht an jemand anderes weiter, löschen Sie alle Kopien und Anhänge davon und lassen Sie mich bitte wissen, dass die E-Mail an die falsche Person gesendet wurde. 

     

This e-mail is confidential. If you received this communication by mistake, please don't forward it to anyone else, please erase all copies and attachments, and please let me know that it has gone to the wrong person.