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

主要观点总结

本文介绍了PyTorch借助OpenAI开发的Triton语言编写内核来加速LLM推理，实现和CUDA类似甚至更佳的性能。文章详细阐述了CUDA、Triton以及两者在LLM推理中的应用，通过实例说明了使用Triton内核实现CUDA-Free的计算过程，并进行了基准测试。最后，文章还展望了未来的优化方向和进一步研究的内容。

关键观点总结

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

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

关键观点2: PyTorch使用Triton内核实现CUDA-Free推理

PyTorch以Llama3-8B和IBM的Granite-8B Code为例，使用Triton内核实现了FP16推理。通过手写Triton内核替换矩阵乘法和注意力内核等关键操作，实现了端到端的推理。

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

在英伟达H100和A100上进行的基准测试表明，Triton内核的性能可以达到CUDA内核的76%～82%。但相比CUDA，Triton内核在某些操作上还存在性能差距，如矩阵乘法和flash attention的内核延迟。

关键观点4: 未来展望

未来，PyTorch计划进一步优化矩阵乘法等关键原语的内核性能，并继续探索FlexAttention和FlashAttention-3等新技术，以缩小Triton和CUDA间的性能差距。

免责声明