9月4号PyTorch官网发布推文，代码编辑宣称实现了100%基于OpenAI的Triton语言推理的大语言模型，这是否意味着去CUDA时代的开端？
![]()
Triton 提供了一条途径，使大模型能够在不同类型的 GPU 上运行—— NVIDIA、AMD，以及未来的 Intel 和其他基于 GPU 的加速器。它还在 Python 中为 GPU 编程提供了更高的抽象层，使我们能够比使用供应商特定的 API 更快地编写高性能内核。
PyTorch官方对当前流行的大模型（例如Meta的Llama3-8B和IBM的Granite-8B）在FP16精度推理，进行了Triton和CUDA实现的对比。
在NVIDIA H100和A100，Llama3-8B和Granite-8B分别基于Triton和CUDA推理速度
基于Triton内核的模型推理性能在NVIDIA H100上约为CUDA内核的0.76~0.78， 在A100上则约为CUDA内核的0.62~0.82。
大语言模型一般均为Transformer架构，下图显示了典型的Transformer包含的操作。![]()
Llama3中Transformer的核心操作包含图中每一个方块，其中主要操作有RMSNorm层标准化，Fused QKV、Attention和Projection等矩阵乘法操作，SiLU激活函数。其中一些函数如RMSNorm、SiLU等，使用Torch.compile可以直接替换为Triton内核，使用英伟达内核可视化工具Nsight Systems，可以看到生成的内核，其中在矩阵乘法和注意力间较小的深绿色内核为Triton内核。
图中越长的内核表示计算需要时间越长，可以看出矩阵乘法和注意力占据了约80%的时长。所以，PyTorch官方手写了基于Triton的矩阵乘法和注意力内核。
官方选择在AMD Flash内核上进行torch算子的封装，下图展示了编译后的Transformer内核流程，现在Llama3可以运行在100%由Triton编写的运行管道中。
编写完成后，两种大语言模型均可以完成端到端推理，官方对比了两种模型在两种GPU下不同内核的单个Token推理速度。![]()
端到端推理性能对比
总结来说，Triton 模型在 H100 上可以获得 CUDA 模型高达78%的性能，在 A100 上可以获得高达82% 的性能。端到端推理性能的差异官方也通过了内核性能的差异来解释，通过细致的对比每个过程内核的延迟，可以看出性能差距的来源。![]()
内核推理性能对比
- Triton matmul 内核比 CUDA 慢1.2-1.4 倍
- AMD 的 Triton Flash Attention 内核比 CUDA SDPA 慢1.6 倍
最近的新工作FlashAttention3和FlexAttention提供了更好的利用底层硬件的方法，为此官方也分别测试了在Triton上不同注意力机制的推理速度的对比。![]()
NVIDIA H100上FlexAttention性能对比
可以明显看出Flex Attention显示出更长的上下文长度和解码问题形状的希望。
未来，PyTorch 团队计划探索进一步优化 matmuls的方法，以便更好地利用硬件，并为基于 Triton 的方法实现更大的加速。
对于Flash Attention，PyTorch 团队计划探索 FlexAttention 和 FlashAttention-3 等内核中使用到的技术，以帮助进一步缩小 Triton 与 CUDA 之间的差距。同时还将探索端到端 FP8 LLM 推理。
结合大模型通用知识及在交通、社会治理、安全生产、自然资源等行业领域的知识，深入业务场景，精确捕获用户意图，为用户提供智能问答、数据分析、报表生成、工作任务理解与执行等一系列服务
📞若您有相关需求，欢迎点击下方链接与我们沟通洽谈
🗨️也可以在公众号后台给我们留言
![]()
- https://pytorch.org/blog/cuda-free-inference-for-llms/
![]()
