在AMD GPU上加速大型语言模型的Flash Attention

Accelerating Large Language Models with Flash Attention on AMD GPUs — ROCm Blogs

引言

在这篇博客文章中，我们将指导您如何在AMD GPU上安装Flash Attention，并提供与在PyTorch中标准SDPA比较其性能的基准测试。我们还将测量Hugging Face中多个大型语言模型（LLM）的端到端预填充延迟。
为了理解Flash Attention及其基准测试结果的重要性，让我们首先深入了解一下推动了变压器架构成功的注意力机制。这种机制是编码器和解码器块的关键组成部分，使得变压器在包括自然语言处理、计算机视觉和音频任务在内的广泛AI领埄中出类拔萃。

尺度点积注意力(Scaled Dot-Product Attention)

Transformer模型中使用的注意力机制被称为尺度点积注意力（SDPA）。SDPA的公式如下，