PyTorch官宣：告别CUDA，GPU推理迎来Triton加速新时代

主要观点总结

本文介绍了PyTorch通过OpenAI开发的Triton语言编写内核来加速LLM推理，实现了和CUDA类似甚至更佳的性能。文章详细阐述了CUDA和Triton的定义及其优势，并介绍了PyTorch如何实现CUDA-free推理。同时，文章还讨论了模型推理的过程，包括内核架构、模型推理实现方式、基准测试等。

关键观点总结

关键观点1: CUDA与Triton的介绍与优势

CUDA是英伟达开发的用于并行计算平台和应用程序的编程API，能让开发者通过GPU开展高性能计算。Triton是一种用于编写高效自定义深度学习基元的语言和编译器，旨在实现更高的可移植性和更高的抽象层。

关键观点2: PyTorch实现CUDA-free推理的方式

PyTorch通过100%使用Triton内核实现了FP16推理，使用OpenAI的Triton语言执行计算任务。这有助于解决CUDA版本与PyTorch匹配的问题，提高了模型的可移植性。

关键观点3: 模型推理的过程

模型推理包括内核架构、模型推理实现方式等。其中，内核架构主要由矩阵乘法、注意力机制等组成。模型推理实现方式则需要手写Triton内核或使用torch.compile自动生成。

关键观点4: 基准测试结果

在英伟达H100和A100上进行的基准测试表明，Triton内核的性能可以达到CUDA内核的76%～82%。然而，仍然存在一些性能差距，需要进一步优化矩阵乘法和注意力机制等关键原语的内核性能。

关键观点5: 未来展望

未来，PyTorch计划进一步优化矩阵乘法，探索FlexAttention和FlashAttention-3等新技术，并探讨端到端FP8 LLM推理的可能性。这些优化和新技术有望进一步缩小Triton和CUDA间的性能差距。

免责声明