arXiv:2511.02043cs.LGcs.PF2025-11被引 2

Flashlight让PyTorch自动生成高效注意力核,支持任意注意力变体。

Flashlight: PyTorch Compiler Extensions to Accelerate Attention Variants

  • 基于PyTorch编译流程,自动融合与分块实现注意力计算。
  • 在多种注意力变体上性能媲美或超越FlexAttention。
  • 适合快速实验新注意力模型的开发者,无需手写优化内核。

注意力是大语言模型的核心组件,已有诸多高效实现方案。例如,FlashAttention通过分块和内核融合优化注意力计算。近年来,涌现出多种注意力变体以提升模型质量或效率,但其高效支持仍困难,因常需专用内核或手工调优。FlexAttention通过静态编程模板支持部分变体的FlashAttention式内核。本文提出Flashlight,一个原生集成于PyTorch生态的编译框架,可自动为任意注意力程序生成融合、分块的FlashAttention风格内核,无需静态模板或预定义特殊化。Flashlight利用PyTorch编译流程透明地融合与分块注意力计算,实现多样化注意力模式的高效执行。不仅支持所有可在FlexAttention中表达的变体,还处理更通用的、依赖数据的注意力形式,超出FlexAttention能力范围。结果表明,Flashlight生成的内核性能可媲美或优于FlexAttention,同时保留原生PyTorch代码的灵活性,使开发者能快速探索新型注意力模型而不牺牲性能。

原文摘要 · Abstract (English)

Attention is a fundamental building block of large language models (LLMs), so there have been many efforts to implement it efficiently. For example, FlashAttention leverages tiling and kernel fusion to optimize attention. Recently, a number of variants of attention have been introduced to enhance model quality or efficiency. Supporting them efficiently remains difficult since they usually require specialized kernels or hand-tuned implementations. FlexAttention recently addressed part of this gap by using static programming templates to support FlashAttention-like kernels for a subset of attention variants. In this paper, we introduce Flashlight, a compiler-native framework within the PyTorch ecosystem that automatically generates fused, FlashAttention-style kernels for arbitrary attention-based programs, without relying on static templates or predefined kernel specializations. Flashlight leverages PyTorch's compilation workflow to fuse and tile attention computations transparently, enabling efficient execution for diverse attention patterns. Not only does it support all variants expressible in the FlexAttention model but it also handles more general, data-dependent attention formulations that are beyond the capabilities of FlexAttention. Our results show that Flashlight produces kernels with competitive or superior performance to FlexAttention, while offering the flexibility of native PyTorch code, enabling developers to rapidly explore new attention models without sacrificing performance.

注意力机制编译优化PyTorch高性能计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。