AttnFuse：用于将注意力编译到融合 GPU 内核的可组合 DSL

现代人工智能系统建立在 Transformer 架构之上，其核心操作——注意力，占据了大部分计算和内存成本。研究人员不断提出新的注意力变体来提高质量、效率或上下文长度，但每个变体目前都需要专家编写的 GPU 代码才能以可用的速度运行。 PyTorch 最近的 flex\_attention 允许研究人员在 Python 中描述自定义注意力模式并将其编译为融合内核，但其设计仅限于中心矩阵乘法后应用的修改，不包括旋转位置嵌入 (RoPE)，这是每个主要大语言模型使用的位置编码。我们引入了 AttnFuse，这是一种用于注意力的小型 DSL，它可以像 RoPE 一流操作一样进行预乘转换。研究人员构建了十个高级构建块来描述一个变体，AttnFuse 的编译器为整个计算生成一个融合的 GPU 内核。在 RTX 3090 上，AttnFuse 在 RoPE+因果模式上比 flex\_attention 实现了 2.10$\times$ 加速。在 H100 上，它在 PyTorch 手动调整后端的 5\% 内运行完整的 Llama-3-8B 训练步骤。我们的调查揭示了旋转演算：是否融合 RoPE 还是单独应用它取决于 GPU 的计算与带宽比率，以及与测量相匹配的派生交叉。 AttnFuse 证明了一个小型的、特定于注意力的编译器可以缩小灵活的研究代码和生产内核之间的差距。