{"slug": "testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse", "title": "Testing PyTorch 2.13 MPS FlexAttention on M1 Max: Up to 7.83x Faster for Sparse Attention", "summary": "A developer benchmarked PyTorch 2.13's FlexAttention on an M1 Max Mac, finding up to 7.83x speedup over standard SDPA for sparse attention with 32,768 tokens and a 256-token local window. FlexAttention, introduced as a prototype in PyTorch 2.5, allows custom attention rules via short Python functions compiled into fused kernels, and PyTorch 2.13 added Metal/MPS kernels for Apple Silicon.", "body_md": "Hello, everyone.\n\nAttention 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?\n\nToday, I am comparing FlexAttention, newly available on Apple Silicon in PyTorch 2.13, with standard SDPA on an M1 Max.\n\nThe 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.\n\n[FlexAttention](https://pytorch.org/blog/pytorch2-5/) 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`\n\nturns into a specialized fused kernel.\n\nAttention 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**.\n\n[PyTorch 2.13](https://pytorch.org/blog/pytorch-2-13-release-blog/), 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.\n\nThis experiment does not use a pretrained AI model. It is a kernel benchmark that sends the same randomly generated attention inputs through two implementations.\n\n| Technology | Role |\n|---|---|\n| PyTorch 2.13.0 / FlexAttention | Run custom attention rules as compiled kernels |\n| SDPA | Standard PyTorch attention implementation used as the baseline |\n| MPS / Metal | Execute calculations on the Apple Silicon GPU |\n\nBoth 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.\n\n```\nsame random inputs (query / key / value)\n  ├─ FlexAttention\n  │    Python mask rule -> BlockMask -> torch.compile -> Metal kernel\n  │\n  └─ SDPA\n       dense mask --------------------> MPS backend\n\n              -> compare output error and forward time\n```\n\nA `BlockMask`\n\ngroups 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.\n\nThe complete code and JSON report are available in the [pytorch-2-13-flexattention-mps lab in kiarina/labs](https://github.com/kiarina/labs/tree/main/2026/07/21/pytorch-2-13-flexattention-mps).\n\nYou will need an Apple Silicon Mac, `mise`\n\n, and `uv`\n\n. 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.\n\n```\ngit clone --depth 1 --filter=blob:none --sparse \\\n  https://github.com/kiarina/labs.git\ncd labs\ngit sparse-checkout set .gitignore .mise/tasks Makefile mise.toml \\\n  2026/07/21/pytorch-2-13-flexattention-mps\nmise -C 2026/07/21/pytorch-2-13-flexattention-mps run\n```\n\nTo skip the long case, run this inside the lab:\n\n```\nuv run python benchmark.py --quick\n```\n\nQuery, 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.\n\n```\nmachine: MacBook Pro (Apple M1 Max, 32 GPU cores, 64 GB)\nOS: macOS 26.5.2\nPython: 3.13.7\nPyTorch: 2.13.0\nshape: batch 1, 8 heads, head dimension 64\ndtype: bfloat16\nFlexAttention: torch.compile(..., dynamic=False)\nBlockMask: 128×128 blocks\nCPU fallback: disabled\n```\n\nA 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.\n\nAn `SDPA / Flex`\n\nvalue above one means that FlexAttention was faster. Density is the percentage of all token pairs that the pattern actually allows.\n\n| Pattern | Sequence | Window | Token / block density | Flex median | SDPA median | SDPA / Flex |\n|---|---|---|---|---|---|---|\n| causal | 8,192 | — | 50.01% / 50.78% | 231.22 ms | 12.33 ms |\n0.05x |\n| local | 8,192 | 64 | 0.78% / 3.10% | 11.75 ms |\n25.23 ms | 2.15x |\n| local | 8,192 | 256 | 3.08% / 4.61% | 19.14 ms |\n25.25 ms | 1.32x |\n| local | 8,192 | 1,024 | 11.72% / 13.18% | 56.24 ms | 25.30 ms |\n0.45x |\n| local | 8,192 | 4,096 | 37.50% / 38.67% | 169.24 ms | 25.30 ms |\n0.15x |\n| local | 32,768 | 256 | 0.78% / 1.17% | 75.27 ms |\n589.05 ms | 7.83x |\n\nAt 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.\n\nCausal 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.\n\nFor 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.\n\nPyTorch 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.\n\nThe 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.\n\nIn addition to steady-state forward time, FlexAttention needs BlockMask construction and a first compiled call.\n\n| Pattern | Sequence / window | BlockMask build | First compiled call |\n|---|---|---|---|\n| causal | 8,192 / — | 564.54 ms | 546.57 ms |\n| local | 8,192 / 64 | 95.90 ms | 94.31 ms |\n| local | 8,192 / 256 | 100.54 ms | 102.34 ms |\n| local | 8,192 / 1,024 | 95.02 ms | 130.75 ms |\n| local | 8,192 / 4,096 | 96.65 ms | 244.62 ms |\n| local | 32,768 / 256 | 1,120.85 ms | 162.57 ms |\n\nThe 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.\n\nThe first-call figures include the state of the host's compile cache. BlockMask construction was eager; I did not benchmark compiling mask construction itself.\n\nOn 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.\n\nI also compared a smaller input against float32 SDPA on the CPU.\n\n| MPS implementation | Maximum absolute error | Mean absolute error |\n|---|---|---|\n| FlexAttention bfloat16 | 0.012440 | 0.000553 |\n| SDPA bfloat16 | 0.012440 | 0.000677 |\n\nBoth 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.\n\nA probe with `requires_grad=True`\n\nfailed with `FlexAttention does not support backward on MPS`\n\n. 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.\n\nFlexAttention 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.\n\nStandard SDPA was about 19x faster for causal attention over the full past. The pattern and density must be measured.\n\nMask construction and compilation add an initial wait. FlexAttention is better suited to work that reuses them across layers or inference calls.\n\nSeeing 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.\n\nThe 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.", "url": "https://wpnews.pro/news/testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse", "canonical_source": "https://dev.to/kiarina/testing-pytorch-213-mps-flexattention-on-m1-max-up-to-783x-faster-for-sparse-attention-1p7d", "published_at": "2026-07-21 04:03:12+00:00", "updated_at": "2026-07-21 04:32:19.925535+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["PyTorch", "Apple Silicon", "M1 Max", "FlexAttention", "SDPA", "Metal", "MPS"], "alternates": {"html": "https://wpnews.pro/news/testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse", "markdown": "https://wpnews.pro/news/testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse.md", "text": "https://wpnews.pro/news/testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse.txt", "jsonld": "https://wpnews.pro/news/testing-pytorch-2-13-mps-flexattention-on-m1-max-up-to-7-83x-faster-for-sparse.jsonld"}}