PyTorch 2.13 shipped July 8 and the headline for Mac developers is blunt: FlexAttention now runs on Apple Silicon. The speedup on sparse attention patterns hits 12x over PyTorch’s own SDPA baseline. Not 12% — twelve times. On a 32,768-token sequence with a 256-token sliding window, the measured delta is 35ms versus 431ms. If you have been prototyping transformer variants on a MacBook and watching attention computation eat your afternoon, 2.13 removes the excuse.
What FlexAttention Actually Is #
FlexAttention launched in PyTorch 2.5 with a simple premise: you should not have to write CUDA kernels every time you want a non-standard attention pattern. Sliding window attention, ALiBi positional bias, document masking, causal plus prefix — all of these previously required hand-rolled kernels or slow masked-tensor workarounds. FlexAttention lets you express any attention variant as two small Python functions and hands them to torch.compile
to fuse into a FlashAttention-grade kernel.
The catch was that it only ran on CUDA. PyTorch 2.13 fixes that. The team wrote native Metal kernels for both the sparse prefill and decode paths, including grouped-query attention. You can now run the same FlexAttention code on your M-series Mac that you run on your CUDA server — and for sparse patterns, the MPS backend holds up well.
import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
device = "mps" if torch.backends.mps.is_available() else "cpu"
def sliding_window(b, h, q_idx, kv_idx):
return abs(q_idx - kv_idx) <= 256
block_mask = create_block_mask(
sliding_window, B=1, H=8,
Q_LEN=32768, KV_LEN=32768,
device=device
)
q = torch.randn(1, 8, 32768, 64, device=device)
k = torch.randn(1, 8, 32768, 64, device=device)
v = torch.randn(1, 8, 32768, 64, device=device)
out = flex_attention(q, k, v, block_mask=block_mask)
That code runs on Apple Silicon now. The speedup depends entirely on how sparse your pattern is — which is the right tradeoff to make. FlexAttention on MPS is not faster than SDPA on dense causal attention. It is faster when you have a pattern that lets it skip work, which is exactly when you want the shortcut.
The Performance Numbers in Honest Context #
The 12x figure comes from a specific shape: 1 batch, 8 heads, 32,768 sequence length, 64 head dimension, 256-token sliding window at 0.8% density. Shorter sequences compress the gain: an 8K-length, 64-token window delivers about 4x. Across sliding-window patterns generally, FlexAttention on MPS runs around 20x faster than SDPA and about 3x faster than FlashAttention 2 with a causal mask.
For standard full causal attention, do not switch. FlexAttention generates its kernels through torch.compile
, which does not produce kernels as tight as a hand-tuned FlashAttention 3 implementation. Expect 5–15% slower throughput on pure causal attention. The right mental model: use FlexAttention when your pattern has sparsity to exploit, use SDPA or Flash for everything else. The API does not pretend otherwise.
Fused LinearCrossEntropyLoss: The Fine-Tuning Win #
FlexAttention gets the headline, but nn.LinearCrossEntropyLoss
might matter more if you are fine-tuning rather than training from scratch. It fuses the final projection layer and the cross-entropy loss into a single operation, which cuts peak GPU memory by up to 4x on large-vocabulary models. A model with a 128K-token vocabulary that was OOM-ing on your 24GB card might now fit. On Apple Silicon with its unified memory pool, that headroom is even more valuable — you are working with a fixed budget shared by the entire system.
import torch.nn as nn
loss_fn = nn.LinearCrossEntropyLoss()
loss = loss_fn(hidden_states, weight, labels)
CuTeDSL, torchcomms, and What to Check Before Upgrading #
PyTorch 2.13 adds CuTeDSL as a second code-generation backend in Inductor alongside Triton, targeting GEMM and RMSNorm. If you hit compilation latency in Inductor, this is worth testing — it generates CUTLASS-grade kernels and moves compilation to a subprocess pool to clear the GIL bottleneck.
torchcomms is a new distributed communications backend with better fault tolerance — graceful timeout, partial-group recovery — and structured logging for large-cluster debugging. If you run multi-node training, this improves observability significantly.
Two breaking changes worth a grep before you upgrade:
- Named tensors are gone. The feature was deprecated; now it is removed. If your codebase uses them, it will break at runtime.
- Distributed collective renames:
all_gather_into_tensor
becomesall_gather_single
,reduce_scatter_tensor
becomesreduce_scatter_single
. Search your distributed training code before upgrading.
Should You Upgrade Now? #
If you work on Apple Silicon and use any non-standard attention pattern — sliding window, ALiBi, document masking — upgrade now. Run pip install --upgrade torch torchvision torchaudio
, swap your device to "mps"
, and point flex_attention at a block mask. The gains are real and available today.
If you are on CUDA-only infrastructure, 2.13 is a solid release but not urgent unless you are hitting Inductor compilation overhead or distributed fault-tolerance gaps. The fused LinearCrossEntropyLoss
is worth adding regardless of hardware if you are fine-tuning large-vocabulary models.
PyTorch is hosting a live Q&A for 2.13 on July 22 — that is tomorrow — if you have questions for the team. Check the official release blog for registration details. The release comprises 3,328 commits from 526 contributors, which gives some indication of how much ground 2.13 covers beyond what fits in a single post.