There is already support for Numba CUDA JIT but Triton JIT interacts much better with frameworks like PyTorch.