FlexAttention: A Programming Model for Generating Optimized Attention Kernels

Attention variants beyond softmax attention (Alibi, document masking, sliding window, PagedAttention) are hard to implement fast because each needs a bespoke fused kernel, so the field suffers a 'software lottery' where a variant cannot be tried without writing a kernel.

A compiler-driven programming model inside PyTorch. The user writes a score_mod (an elementwise function of score index) and a block_mask (a sparse mask expressed as a nested-stride BlockMask) in a few lines of idiomatic PyTorch; torch.compile traces this into a Triton kernel that tiles Q/KV through SRAM like FlashAttention. Composable, so variants combine without a combinatorial explosion of hand-written kernels.

PARTIAL - the paper's own claim (abstract) is QUALITATIVE, not a multiple: it demonstrates that existing attention variants (Alibi, document masking, PagedAttention) CAN be implemented via FlexAttention and achieves 'competitive performance compared to these handwritten kernels', plus a compositionality result. No single speedup figure with a named baseline appears in the abstract. The decision-relevant companion number is the next-paper delta: FlexAttention was superseded by torch.compile fusing it into the default path, which is why adoption is in-upstream-engine.

arXiv ID 2412.05496 verified via abs-page meta tags on 2026-10-03; the abs page's citation_title is 'FlexAttention: A Programming Model for Generating Fused Attention Kernels' while the record name says '...Optimized Attention Kernels' - same paper, the title was revised between versions, so the record name was left as the dead agent wrote it and the arXiv discrepancy is noted here rather than silently overwritten. CORRECTION TO THE SOURCE RECORD: it says 'arXiv preprint; no journal_ref', but OpenReview note 2QMYV4bA0R returns venue='MLSys 2025 withshepherding', venueid='MLSys.org/2025/Conference'. venue=mlsys, year=2025. IMPORTANT NAMING HAZARD, confirmed against the dead agent's own warning: there is a DIFFERENT paper, arXiv 2407.20228 'FlexAttention for Efficient High-Resolution Vision-Language Models' (ECCV 2024), which appears in OpenReview as a Public_Article and is unrelated. Do not conflate. Adoption=in-upstream-engine verified directly in source: torch/nn/attention/flex_attention.py is 104 KB in the main branch of pytorch/pytorch and defines flex_attention, create_block_mask and and_masks; vLLM also ships vllm/v1/attention/backends/flex_attention.py. This is the strongest adoption evidence of any attention record in this batch. VENUE_TRACK BACKFILL 2026-10-04 (paper-fields pass, agent w3p): venue_track 'main'. This is STRUCTURAL rather than researched: mlsys has no Findings volume and no workshop track in which a paper could appear, so 'main' is the only possible honest value and no per-paper evidence was needed. The field is annotated here so that 'findings' and 'workshop' remain MEANINGFUL - a filter for findings=true must return only the papers that actually are findings, which requires every other paper to be explicitly NOT findings.