PyTorch 2.2: FlashAttention-2 in SDPA, AOTInductor
PyTorch 2.2 brought FlashAttention-2 into scaled_dot_product_attention for about 2× speedups, and introduced AOTInductor for compiling models ahead of time for non-Python deployment.
What happened
PyTorch 2.2 was released, with 3,628 commits from 521 contributors since 2.1.
The technical details
From the release notes:
scaled_dot_product_attention(SDPA) now supports FlashAttention-2, with about 2× speedups over previous versions- AOTInductor: an ahead-of-time extension of TorchInductor that compiles and deploys PyTorch programs for non-Python, server-side environments
device_mesh: a new abstraction intorch.distributedfor initialising and representing process groupsTORCH_LOGS: a standard, configurable logging mechanism- better
torch.compilesupport for optimizers, plus Inductor fusion and layout improvements - PyTorch 2.2.x is the last release supporting macOS x86
Why it matters
Attention is the core cost of Transformers. In my self-attention study post I checked my hand-written attention against F.scaled_dot_product_attention. That’s the same function, and releases like this make it faster without any code change.
Sources
- PyTorch 2.2 release notes — GitHub · pytorch/pytorch
- Dao (2023), FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning — arXiv
This is a summary written from the linked public sources. Figures and numbers are as reported there.