Flex Attention for Inference
summary
FlexAttention ships in PyTorch 2.5.0 and implements most attention variants in a few lines of PyTorch; it adds a compiler-driven FlexDecoding backend that uses a split-key reduce to speed short-query long-K inference, reports performance on par with FlashDecoding and markedly faster than SDPA, and exposes a Paged Attention path via block-mask and score-mode conversion utilities with minimal latency overhead while enabling much larger batch sizes. Features also include Grouped Query Attention, nested jagged tensor support, trainable biases, and a create_block_mask helper (wrap with torch.compile).
AI-written summary · about
chapters (7)
- 0:00Flex Attention intro
- 1:44Inference vs training
- 3:08Flex Decoding
- 4:12Page Attention memory waste
- 5:16Flex + Page Attention
- 7:48Performance and adoption
- 9:14Q&A
related talks
09mFlexAttention + FlashAttention-4: Fast and Flexible
36mML Efficiency for Large Models: From Data Efficiency to Faster Transformers
20mMelanie Warrick, Deep Learning Engineer, Skymind.io
06mGlobalFusion: A Global Attentional Deep Learning Framework for Multisensor Information Fusion
34mTalk by Mahdi Soltanolkotabi (University of Southern California)
18mAdvancing Mamba in PyTorch
42mAttention and Memory in Deep Learning Networks
14mFlexPipe: Maximizing Training Efficiency for Transformer-based Models with...