Flex Attention: A Programming Model for Generating Optimized Attention Kernels

• URL: http://arxiv.org/abs/2412.05496v1
• Date: Sat, 07 Dec 2024 01:46:38 GMT
• Title: Flex Attention: A Programming Model for Generating Optimized Attention Kernels
• Authors: Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, Horace He,
• Abstract summary: We introduce FlexAttention, a novel compiler-driven programming model that allows implementing the majority of attention variants in a few lines of PyTorch code. We demonstrate how FlexAttention allows for easy composition of attention variants, solving the explosion of attention variants.
• Score: 5.489362130813523
• License:
• Abstract: Over the past 7 years, attention has become one of the most important primitives in deep learning. The primary approach to optimize attention is FlashAttention, which fuses the operation together, drastically improving both the runtime and the memory consumption. However, the importance of FlashAttention combined with its monolithic nature poses a problem for researchers aiming to try new attention variants -- a "software lottery". This problem is exacerbated by the difficulty of writing efficient fused attention kernels, resisting traditional compiler-based approaches. We introduce FlexAttention, a novel compiler-driven programming model that allows implementing the majority of attention variants in a few lines of idiomatic PyTorch code. We demonstrate that many existing attention variants (e.g. Alibi, Document Masking, PagedAttention, etc.) can be implemented via FlexAttention, and that we achieve competitive performance compared to these handwritten kernels. Finally, we demonstrate how FlexAttention allows for easy composition of attention variants, solving the combinatorial explosion of attention variants.

Related papers

• Attention Entropy is a Key Factor: An Analysis of Parallel Context Encoding with Full-attention-based Pre-trained Language Models [49.84163262868945]
  Large language models have shown remarkable performance across a wide range of language tasks, owing to their exceptional capabilities in context modeling. The most commonly used method of context modeling is full self-attention, as seen in standard decoder-only Transformers. We propose parallel context encoding, which splits the context into sub-pieces and encodes them parallelly.
  arXiv  Detail & Related papers  (2024-12-21T09:04:51Z)
• Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers [58.5711048151424]
  We introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome computational and memory obstacles. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods.
  arXiv  Detail & Related papers  (2024-06-24T15:55:59Z)
• An end-to-end attention-based approach for learning on graphs [8.552020965470113]
  transformer-based architectures for learning on graphs are motivated by attention as an effective learning mechanism. We propose a purely attention-based approach consisting of an encoder and an attention pooling mechanism. Despite its simplicity, the approach outperforms fine-tuned message passing baselines and recently proposed transformer-based methods on more than 70 node and graph-level tasks.
  arXiv  Detail & Related papers  (2024-02-16T16:20:11Z)
• FLatten Transformer: Vision Transformer using Focused Linear Attention [80.61335173752146]
  Linear attention offers a much more efficient alternative with its linear complexity. Current linear attention approaches either suffer from significant performance degradation or introduce additional computation overhead. We propose a novel Focused Linear Attention module to achieve both high efficiency and expressiveness.
  arXiv  Detail & Related papers  (2023-08-01T10:37:12Z)
• Faster Causal Attention Over Large Sequences Through Sparse Flash Attention [45.18552512844457]
  We extend FlashAttention to accommodate a large class of attention sparsity patterns. We increase the training speed of a transformer language model by $2.0times$ and $3.3times$ for sequences of respectively $8k$ and $16k$ tokens.
  arXiv  Detail & Related papers  (2023-06-01T21:33:59Z)
• Unlocking Pixels for Reinforcement Learning via Implicit Attention [61.666538764049854]
  We make use of new efficient attention algorithms, recently shown to be highly effective for Transformers. This allows our attention-based controllers to scale to larger visual inputs, and facilitate the use of smaller patches. In addition, we propose a new efficient algorithm approximating softmax attention with what we call hybrid random features.
  arXiv  Detail & Related papers  (2021-02-08T17:00:26Z)
• Bayesian Attention Modules [65.52970388117923]
  We propose a scalable version of attention that is easy to implement and optimize. Our experiments show the proposed method brings consistent improvements over the corresponding baselines.
  arXiv  Detail & Related papers  (2020-10-20T20:30:55Z)
• Learning fine-grained search space pruning and heuristics for combinatorial optimization [5.72274610208488]
  We propose a framework for leveraging machine learning techniques to scale-up exact optimization algorithms. Our framework learns the relatively simpler task of pruning the elements in order to reduce the size of the problem instances. We show that our framework can prune a large fraction of the input graph and still detect almost all of the maximum cliques.
  arXiv  Detail & Related papers  (2020-01-05T13:10:39Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.