Most of these optimizations should be easy to add: https://pytorch.org/blog/accelerating-generative-ai-2/