Testing PyTorch 2.13 MPS FlexAttention on M1 Max: Up to 7.83x Faster for Sparse Attention

Hello, everyone.

Attention becomes expensive very quickly as more text is given to an AI model. Can a Mac GPU make it faster when every token is restricted to looking only at nearby tokens?

Today, I am comparing FlexAttention, newly available on Apple Silicon in PyTorch 2.13, with standard SDPA on an M1 Max.

The short answer is that, with 32,768 tokens and a 256-token local window, FlexAttention took 75.27 ms while SDPA took 589.05 ms: a 7.83x speedup. For ordinary causal attention, however, SDPA was about 19x faster. FlexAttention was not universally faster; it helped with long, extremely sparse attention.

FlexAttention was announced as a prototype with PyTorch 2.5 in October 2024. It lets developers express an attention rule as a short Python function, which torch.compile

turns into a specialized fused kernel.

Attention determines which tokens in the input should influence one another. The number of possible comparisons grows rapidly with sequence length. FlexAttention can describe rules such as “look only into the past” or “look only at the previous 256 tokens,” allowing unnecessary comparisons to be skipped. A pattern with many skipped comparisons is called sparse.

PyTorch 2.13, announced on July 8, 2026, added Metal/MPS kernels for FlexAttention on Apple Silicon. PyTorch reports up to roughly 12x higher performance than SDPA for sparse patterns. The API and kernel options remain unstable in 2.13.

This experiment does not use a pretrained AI model. It is a kernel benchmark that sends the same randomly generated attention inputs through two implementations.

Both paths receive the same query, key, and value tensors. Loosely speaking, these represent what to search for, what to match against, and what content to retrieve.

same random inputs (query / key / value)
  ├─ FlexAttention
  │    Python mask rule -> BlockMask -> torch.compile -> Metal kernel
  │
  └─ SDPA
       dense mask --------------------> MPS backend

              -> compare output error and forward time

A BlockMask

groups the regions to compute into 128×128-token blocks. FlexAttention can skip unnecessary blocks. The SDPA path received an ordinary boolean mask describing the same rule.

The complete code and JSON report are available in the pytorch-2-13-flexattention-mps lab in kiarina/labs.

You will need an Apple Silicon Mac, mise

, and uv

. The default task also creates a 32,768×32,768 boolean mask, which alone occupies 1 GiB. Stop other GPU workloads and run it on a machine with sufficient memory.

git clone --depth 1 --filter=blob:none --sparse \
  https://github.com/kiarina/labs.git
cd labs
git sparse-checkout set .gitignore .mise/tasks Makefile mise.toml \
  2026/07/21/pytorch-2-13-flexattention-mps
mise -C 2026/07/21/pytorch-2-13-flexattention-mps run

To skip the long case, run this inside the lab:

uv run python benchmark.py --quick

Query, key, and value were generated independently at random. Both implementations received the same tensors. I timed only the forward pass, excluding mask creation and compilation. Because MPS runs asynchronously, each measurement explicitly waited for the GPU before and after the call. The reported values are medians from ten trials after three warm-ups.

machine: MacBook Pro (Apple M1 Max, 32 GPU cores, 64 GB)
OS: macOS 26.5.2
Python: 3.13.7
PyTorch: 2.13.0
shape: batch 1, 8 heads, head dimension 64
dtype: bfloat16
FlexAttention: torch.compile(..., dynamic=False)
BlockMask: 128×128 blocks
CPU fallback: disabled

A sliding window lets each position see the previous W tokens. With a window of 256, for example, each token is limited to the most recent 256 tokens even in a very long input.

An SDPA / Flex

value above one means that FlexAttention was faster. Density is the percentage of all token pairs that the pattern actually allows.

At 8,192 tokens, FlexAttention won through a window of 256. At a window of 1,024, SDPA became 2.22x faster. The crossover on this shape lies between 3.08% and 11.72% token density.

Causal attention allows each token to see the entire past, giving it a density of about 50%. SDPA has a specialized path for this common pattern and was 18.75x faster than FlexAttention. Simply replacing ordinary causal attention is counterproductive.

For the 32,768 / window 256 case, one of ten FlexAttention trials rose to 147.89 ms. Its full range of 74.96–147.89 ms still did not overlap with SDPA's 586.08–590.03 ms range, so the performance ordering was unambiguous.

PyTorch reports 4.15x for 8,192 / window 64 and approximately 12.3x for 32,768 / window 256. The M1 Max produced 2.15x and 7.83x, respectively.

The expected direction was reproduced: sparser patterns were faster, and the gap grew at longer sequence lengths. The speedups did not reach the official figures. Because the release blog does not identify the Apple Silicon model used for those figures, the difference cannot be attributed to hardware alone.

In addition to steady-state forward time, FlexAttention needs BlockMask construction and a first compiled call.

The 32,768 / window 256 case saves about 514 ms per call, so reusing the same mask recovers the setup cost in roughly three forwards. The 8,192 / window 64 case needs about 14. A mask used only once should not be selected based on the steady-state 7.83x or 2.15x figure alone.

The first-call figures include the state of the host's compile cache. BlockMask construction was eager; I did not benchmark compiling mask construction itself.

On identical MPS bfloat16 inputs, maximum absolute error between FlexAttention and SDPA ranged from 0.0078125 to 0.015625, while mean absolute error ranged from 0.000079 to 0.000363.

I also compared a smaller input against float32 SDPA on the CPU.

Both MPS paths had the same maximum error, and I did not observe a large discrepancy specific to FlexAttention. This is a numerical comparison on synthetic inputs, not a quality evaluation of a complete model.

A probe with requires_grad=True

failed with FlexAttention does not support backward on MPS

. The PyTorch 2.13 MPS path is forward-inference only. The deterministic backward feature added in 2.13 applies to CUDA and does not add MPS training support.

FlexAttention was 7.83x faster when 32,768 tokens were limited to the previous 256. It benefits when a large share of the work can be skipped.

Standard SDPA was about 19x faster for causal attention over the full past. The pattern and density must be measured.

Mask construction and compilation add an initial wait. FlexAttention is better suited to work that reuses them across layers or inference calls.

Seeing the gap reach 7.83x for extremely sparse, long attention on an M1 Max was a good result. Writing the mask rule in Python and letting PyTorch produce a specialized Metal kernel substantially lowers the barrier to experimenting with custom attention on a Mac.

The advantage disappeared quickly as the window widened, and the initial setup time was not trivial. Despite the “Flex” name, it should not be treated as a universal optimization; density and reuse count need to be measured. It looks useful for local LLM inference over long documents with nearby context or experiments that operate on sparse relationships.