动机
最近，OpenAI Triton后端已经合并了AMD硬件PR 3643,这是迄今为止唯一带有vLLM源代码部分的闪存注意力后端。OpenAI Triton的一些优势是平台优越性和性能可移植性。因此，我们(@tdoublep和我自己)想调查一下这段代码是否能在不同的平台上同样有效地工作，即NVIDIA GPUs。
我们的实验表明，在不同NVIDIA硬件(A100、L40、H100)上使用由AMD贡献的代码与默认选项(flash_attn)相比，具有竞争力的预填充性能。对于使用Tensor并行性的较小数量的头(可能的情况),甚至更快。
为了这次实验，我们使用了code contributed by AMD,但用更适合不同GPU的选项替换了自动调优器选项。然而，我们并没有改变实际的Triton代码。
因此，我们是否可以考虑仅使用Triton后端？虽然这在所有情况下都不一定会带来性能优势，但有几个额外的技术动机：
- 主要使用Triton代码，将减少对手写cuda代码的依赖(例如，different variants of the forward kernel in flash_attn的约3500 LoC vs约500 LoC的Triton代码)。
- 它将提高vLLM的平台可移植性。
- 因此，有助于vLLM的未来兼容性。
建议的更改
我们建议添加一个新的后端，该后端在NVIDIA和AMD平台上运行闪存注意力Triton代码。我们将从现有的闪存注意力开始，但我们还希望讨论其他内核的选项。
我们还将为Triton自动调优器提供我们的附加选项，以实现上面蓝色曲线的结果。
回复期限
- 无响应*
抄送列表
@hongxiayang@WoosukKwon@Yard1@jpvillam-amd
其他事项
- 无响应*
3条答案
按热度按时间o2gm4chl1#
我认为这是一个好主意
问题 --> 对于其他(非注意力)自定义操作，你会怎么做？
我们还没有以我们为注意力所采用的方式对这些参数化进行处理，可能需要更新这个
csga3l582#
感谢您编写并分享这个RFC,同意这是一个好主意，为了便于在各种GPU上移植，需要更少的代码行(如您所提到：LoC为3500 vs 500),同时不影响性能。
efzxgjgh3#
感谢您的反馈！
@robertgshaw2-neuralmagic
问题 --> 对于其他(非注意力)自定义操作，您会如何处理？
我们还没有以类似注意力的方式对这些参数化进行建模，可能需要更新这个。
是的，确切地说，这是在这个背景下的下一个/更大的问题。
我们还在Triton中实现了一些其他内核，目标是实现一个“仅Triton”的vLLM部署。一般来说，Triton内核的性能具有竞争力，甚至比vLLM中的cuda内核更快。例如，RMSNorm可以快70%(图中的
wR表示带有残差的RMSNorm):目前，在开发和调试过程中，我的代码根据环境变量在cuda和triton内核之间进行选择...但我认为这应该以更好的方式完成；).
但也许这可以是一个逐步的过程：首先，从“仅Triton”后端开始(对于我们来说，分页注意力是一个开放性问题),因为这种抽象已经存在。然后(或并行)考虑一种架构/抽象来实现其他内核的参数化选择？