Flashlight：加速注意力变体的 PyTorch 编译器扩展
机器学习
2026-05-22 v4 性能
摘要
注意力是大型语言模型（LLM）的基本构件，因此已有许多努力来实现其高效性。例如，FlashAttention 利用分块和内核融合来优化注意力。最近，一些注意力变体被引入以增强模型质量或效率。高效支持它们仍然困难，因为它们通常需要专业内核或手工调优的实现。FlexAttention 最近通过使用静态编程模板来支持类似 FlashAttention 的内核，以覆盖部分注意力变体。在本文中，我们引入了 Flashlight，这是一个基于编译器的框架，集成在 PyTorch 生态系统中，能够自动为任意基于注意力的程序生成融合、类似 FlashAttention 风格的内核，而无需依赖静态模板或预定义的内核专化。Flashlight 利用 PyTorch 的编译工作流程来透明地融合和分块注意力计算，从而支持多样化的注意力模式。不仅它支持所有可表达于 FlexAttention 模型的变体，还处理超出 FlexAttention 能力的更通用、数据依赖的注意力公式。我们的结果表明，Flashlight 产生的内核在性能上与 FlexAttention 具竞争力或更优，同时提供了原生 PyTorch 代码的灵活性，使开发人员能够在不牺牲性能的前提下快速探索新的注意力模型。
引用
@article{arxiv.2511.02043,
title = {Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants},
author = {Bozhi You and Irene Wang and Zelal Su Mustafaoglu and Abhinav Jangda and Angélica Moreira and Roshan Dathathri and Divya Mahajan and Keshav Pingali},
journal= {arXiv preprint arXiv:2511.02043},
year = {2026}
}