Error when lowering pallas kernel: 'jaxlib.triton.dialect' has no attribute 'permute' about jax

[superbobry]

Thanks for confirming @davisyoshida! I will close this bug as fixed.

[superbobry]

Thanks for reporting this, Davis. I sent a PR fixing this. In the meantime you can do

from jaxlib.triton import dialect
dialect.permute = dialect.trans

to workaround the issue.

One caveat, though: the Triton version we have internally crashes if any of the dot operands is transposed. It is possible that the jaxlib version did not get the problematic upstream changes, but if not, we'd have to wait until the issue is fixed in Triton. I started a discussion in their Slack.

[davisyoshida]

the Triton version we have internally crashes if any of the dot operands is transposed

Good to know, I'm sure that just saved me an hour of fruitless debugging. Thanks!

[superbobry]

According to the folks on the Triton Slack channel, dot with transposed inputs is not supported on GPUs below Turing. I was testing on V100, so this seems to explain the crashes.

[davisyoshida]

@superbobry The monkeypatch works for me, and thanks for making it so much easier to get pallas up and running on GPU