Comment by simoiacos

Please open a PR! I was too conservative with the supported devices.

If your GPU supports XMX we could also explore using it to improve the prefill kernel, but I don't have the hardware to test it myself.