Flex Attention: A Programming Model for Generating Optimized Attention Kernels

Paper · 2024 · Juechu Dong, Boyuan Feng, Driss Guessous et al.

Summary

FlashAttention made attention fast and memory-efficient by fusing it into a single kernel. Its hand-built design, however, covers only a few popular variants. Researchers who want something different, or a combination of variants, are stuck with slow, memory-hungry code; the authors call this a "software lottery." FlexAttention, from Meta and the University of Michigan, addresses this with a compiler-driven programming model in PyTorch. A user describes a variant with two small Python callables: one alters attention scores before softmax (score_mod), and one returns a boolean mask from positional indices (mask_mod). During torch.compile, these functions are traced and fused into the main loop of hand-tuned attention kernel templates. The backward pass is generated automatically. A precomputed BlockMask lets the kernel skip fully masked tiles. On variants that FlashAttention already supports, FlexAttention comes close to or beats it, and it also speeds up real training and inference pipelines. This lowers the cost of trying new attention designs and makes it easy to combine them.

Key points

API: users write score_mod and mask_mod in ordinary PyTorch. Variants combine by nesting score functions or joining masks with boolean logic, which avoids writing a separate kernel for every combination.
Supported variants: examples include ALiBi, Softcapping, Sliding Window, Document Masking, PagedAttention and Neighborhood Attention. Tiled and Morton-curve versions of Neighborhood Attention take fewer than 10 lines.
BlockMask: create_block_mask builds a block-level sparsity map from a mask function using torch.vmap. The same index structure also serves as the page mapping for PagedAttention, with negligible overhead.
Kernel performance: across 7 variants, it reaches 0.68x–1.43x the speed of FAv2 and 0.93x–1.45x of FAKV in decoding. It was also compared against PyTorch SDPA and FAv3.
End-to-end gains: inference in gpt-fast is 2.04x faster at 16k context, and training in torchtune is 2.4x faster.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields, topics and concepts

📚 Shelves 6Computer Science · 18,614 works Artificial Intelligence · 17,965 works Arts and Humanities · 17,450 works Language and Linguistics · 17,296 works Computer Vision and Pattern Recognition · 3,486 works Hardware and Architecture · 167 works

Cite

Cite as Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang and Horace He. 2024. Flex Attention: A Programming Model for Generating Optimized Attention Kernels. arXiv:2412.05496.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2412.05496 · arXiv 2412.05496

@misc{dong2024flex,
  title = {{Flex Attention: A Programming Model for Generating Optimized Attention Kernels}},
  author = {Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang and Horace He},
  year = {2024},
  doi = {10.48550/arxiv.2412.05496},
  eprint = {2412.05496},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2412.05496}
}

References

23 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head checkpoints, 2023. URL https://arxiv.org/abs/ 2305.13245. ↗
Ansel, J., Yang, E., He, H., Gimelshein, N., Jain, A., Voznesensky, M., Bao, B., Bell, P., and et al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programm…
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer, 2020a. URL https://arxiv.org/abs/2004.05150. ↗
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer, 2020b. URL https://arxiv.org/abs/2004.05150. ↗
Chen, T., Moreau, T., Jiang, Z., Zheng, L., Yan, E., Cowan, M., Shen, H., Wang, L., Hu, Y., Ceze, L., Guestrin, C., and Krishnamurthy, A. Tvm: an automated end-to-end optimizing compiler for deep learning. In Proceedings of the 13th USENIX Conference on Operating Systems Design and Implementation, O…
Hassani, A. and Shi, H. Dilated neighborhood attention transformer, 2022. URL https://arxiv.org/ abs/2209.15001. — [arXiv:2209.15001](https://arxiv.org/abs/2209.15001) ↗
Hassani, A., Walton, S., Li, J., Li, S., and Shi, H. Neighborhood attention transformer. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023. find
Hassani, A., Hwu, W.-M., and Shi, H. Faster neighborhood attention: Reducing the o(n 2 ) cost of self attention at the threadblock level, 2024. URL https://arxiv.org/abs/2403.04690. ↗
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023a.
Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023b.
Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation, 2022. URL https://arxiv.org/abs/ 2108.12409. ↗
Rabe, M. N. and Staats, C. Self-attention does not need o(n 2 ) memory, 2022. URL https://arxiv.org/ abs/2112.05682. ↗
Dao, T. Flashattention-2: Faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL https://arxiv.org/ abs/1910.10683. ↗
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. find
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., and Dao, T. Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024. URL https://arxiv.org/abs/2407.08608. ↗
Dao, T., Haziza, D., Massa, F., and Sizov, G. Flashdecoding for long-context inference, 2023. URL https://crfm.stanford.edu/2023/10/12/ flashdecoding.html. Accessed: 2024-09-15. ↗
Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/ stanford_alpaca, 2023. ↗
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., and Akhil Mathur, e. a. The llama 3 herd of models, 2024. URL https://arxiv.org/abs/ 2407.21783. gpt-fast maintainers and contributors. Accelerating generative ai with pytorch ii: Gpt, fast, November 2023. URL https://github.co… ↗
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., and et al., T. M. Gemma 2: Improving open language models at a practical size, 2024. URL https://arxiv.org/abs/2408.00118. torchtune maintainers and contributors. torchtune: Pytorch’s finetuning library, Apri… ↗
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. Attention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, v…
Wang, G., Zeng, J., Xiao, X., Wu, S., Yang, J., Zheng, L., Chen, Z., Bian, J., Yu, D., and Wang, H. Flashmask: Efficient and rich mask extension of flashattention. arXiv preprint arXiv:2410.01359, 2024. find
Wu, M., Cheng, X., Padon, O., and Jia, Z. A multi-level superoptimizer for tensor programs, 2024. URL https://arxiv.org/abs/2405.05751. ↗

Updated 2 November 2026

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/dong-2024-flex-attention-programming-model-generating/