
























Abstract: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.
| Subjects: | Machine Learning (cs.LG); Performance (cs.PF) |
| Cite as: | arXiv:2511.02043 [cs.LG] |
| (or arXiv:2511.02043v4 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2511.02043 arXiv-issued DOI via DataCite |
From: Bozhi You [view email]
[v1]
Mon, 3 Nov 2025 20:25:19 UTC (1,652 KB)
[v2]
Thu, 6 Nov 2025 00:10:15 UTC (1,652 KB)
[v3]
Fri, 7 Nov 2025 17:26:02 UTC (1,652 KB)
[v4]
Wed, 20 May 2026 23:03:49 UTC (3,595 KB)
此内容由惯性聚合(RSS阅读器)自动聚合整理,仅供阅读参考。 原文来自 — 版权归原作者所有。