PyTorch 2.13 brings FlexAttention to Apple Silicon: a 12x sparse-attention speedup

Summary. PyTorch 2.13 shipped on July 8, 2026, and for anyone training or running models on a Mac the headline is one line in the release notes: FlexAttention now runs on Apple Silicon through the Metal (MPS) backend, with up to a ~12x speedup over scaled dot-product attention (SDPA) on sparse masks. The release notes give the exact case: a 1x8x32768x64 attention shape with a 256-element sliding window (0.8% density) runs in about 35 ms with FlexAttention versus about 431 ms with SDPA, a ~12.3x gain. A smaller 8192-length, 64-window case lands at ~4.15x. The 2.13 release carries 3,328 commits from 526 contributors since 2.12. It matters because an entry M4 Mac mini starts at $599 in the US and ₹1,18,900 for the 512GB model in India (Apple), so a lot of local machine-learning work now happens on Metal rather than on rented NVIDIA GPUs. The catch: the MPS path is marked API Unstable, and for plain LLM inference MLX is often still faster.

This guide covers what actually shipped, what the benchmark does and does not say, how to write a sparse mask in about two lines of Python, and when FlexAttention on PyTorch MPS is the right tool versus MLX.

What shipped for Apple Silicon in PyTorch 2.13

Two changes drive Mac performance in this release. First, FlexAttention gained hand-written Metal kernels for both the sparse prefill and decode paths, including grouped-query attention (GQA) and captured buffers. Second, a broad set of MPS operations moved off Apple's MPSGraph framework onto hand-written Metal compute kernels, including copy and cast, uniform and normal random generation, comparisons, reductions such as sum and mean, cumsum and cumprod, sort, embedding backward, and scatter and gather. The native path removes MPSGraph's per-operation compilation cost and cuts kernel launch latency, according to the PyTorch 2.13 release notes.

FlexAttention is PyTorch's unified way to express a custom attention pattern as a plain Python function that the compiler fuses into a fast kernel. As the release notes put it, "instead of writing a custom CUDA kernel for every attention variant you need, FlexAttention will write a 2-line Python function, and the compiler builds a fast kernel for you automatically." That idea shipped for CUDA in 2024 as FlexAttention; 2.13 brings the same programming model to the Mac.

Here is the fuller list of what landed, and who each piece helps.

The benchmark, and what it does not say

The 12x figure is specific, and reading it precisely keeps expectations honest. It measures a sliding-window mask that is 99.2% empty, so FlexAttention skips almost all of the key-value blocks that SDPA still walks through. That is exactly the workload where block sparsity pays off.

The release notes are direct that "dense patterns continue to favor SDPA, as expected." So the win is not universal. If your attention is full or near-full, SDPA remains the better call, and FlexAttention gives you no free speedup. The gain scales with how sparse the mask is and how long the sequence runs, which is why the 32768-token case beats the 8192-token case by a wide margin.

One more caveat sits in the release notes: the MPS FlexAttention work is labelled API Unstable. That is the right maturity signal for a first landing. Pin your PyTorch version, keep a test that compares FlexAttention output against SDPA on a small case, and expect the surface to shift in later releases.

Why sparse attention is the workload that benefits

Sliding-window attention is a good, real example of the pattern FlexAttention accelerates. Mistral 7B, released in 2023, uses a 4096-token sliding window on top of a 32,768-token maximum context, so each token attends to at most 4096 neighbours rather than the whole sequence (Mistral 7B paper). Stacked across 32 layers, that local window still reaches a theoretical span near 131K tokens, while the key-value cache can be capped at the window size, which halves cache memory at a sequence length of 8192 per Mistral's own writeup.

The point for a Mac developer: any model or research idea that uses local, windowed, block, or otherwise structured-sparse attention is now a candidate for a real speedup on Apple Silicon, without hand-writing a Metal kernel. Document masking, causal-plus-window schemes, and prefix-LM patterns all fit the same mould. For retrieval-heavy systems, the attention pattern often interacts with how you store context; if you are weighing where embeddings live, our note on pgvector versus a dedicated vector database covers that trade-off.

How to use FlexAttention on a Mac

The API mirrors the CUDA version. You describe the sparsity as a mask_mod function, pre-compute a BlockMask with create_block_mask, then call flex_attention. A sliding-window causal mask is about two lines of logic.

            import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

device = "mps"  # Apple Silicon GPU via Metal
WINDOW = 256

def sliding_window(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    within = q_idx - kv_idx <= WINDOW
    return causal & within

# B, H, Q_LEN, KV_LEN
block_mask = create_block_mask(sliding_window, 1, 8, 32768, 32768, device=device)

q = torch.randn(1, 8, 32768, 64, device=device, dtype=torch.float16)
k = torch.randn_like(q)
v = torch.randn_like(q)

out = flex_attention(q, k, v, block_mask=block_mask)
          

Two practical notes. Keep a correctness check: run the same inputs through torch.nn.functional.scaled_dot_product_attention with an equivalent mask and compare, because an API-unstable kernel deserves a guard. And profile your real shape before committing; the speedup depends on your density and sequence length, not on the headline number. The FlexAttention documentation has the full score_mod and BlockMask reference.

FlexAttention on PyTorch MPS versus MLX

FlexAttention closes a real gap, but it does not make PyTorch the fastest path for every on-device job. For transformer LLM inference on Apple Silicon, MLX, Apple's own array framework, generally leads. On an M4 Pro, one comparison found MLX running object-detection inference 1.1x to 2.6x faster than PyTorch MPS, and the ExecuTorch MLX delegate reported 3x to 6x higher throughput on generative workloads than existing ExecuTorch delegates on macOS (PyTorch ExecuTorch MLX delegate). PyTorch MPS, meanwhile, holds up well for training and keeps the whole CUDA-era ecosystem.

The reasonable read: prototype and train in PyTorch, use FlexAttention when your attention is sparse, and reach for MLX when you are shipping latency-sensitive inference into a macOS or iOS product. If self-hosting cost is the real question behind "should this run locally," our breakdown of self-hosting GPU cost versus an API frames the maths.

What else in 2.13 is worth your attention

Beyond the Mac story, three changes affect day-to-day work. nn.LinearCrossEntropyLoss fuses the final linear projection and cross-entropy into one module and processes the vocabulary in chunks, cutting peak memory up to ~4x for large-vocabulary training as a drop-in replacement. torch.load now reads .safetensors files directly, so the format's memory-mapped, code-execution-free loading no longer needs a separate library. And PyTorch wheels now cover Python 3.15, including the free-threaded 3.15t build, though only on Linux for now and only via download.pytorch.org rather than PyPI.

A few deprecations deserve a note before you upgrade: named tensors were hard-removed, all_gather_into_tensor and reduce_scatter_tensor were renamed to the _single scheme, the Bazel build was dropped, and CUDA 13.0 is the default. Run your test suite against 2.13 in a branch before you move a training pipeline.

India-specific considerations

For teams in India, the economics push toward local Apple Silicon for a lot of experimentation. Apple lists the M4 Mac mini from $599 in the US and ₹1,18,900 for the 512GB, 16GB model in India, a one-time cost against per-hour cloud GPU rental. A single Mac with generous unified memory can hold a mid-sized model for prototyping without a cloud bill, and FlexAttention now removes one of the reasons you had to move sparse-attention work off the Mac. Where data residency and privacy matter, keeping model work on a local device also sits comfortably with the Digital Personal Data Protection Act, 2023 (DPDP), since experimentation data need not leave the machine. Production training at scale still belongs on multi-GPU clusters, where the 2.13 FSDP2 and torchcomms changes apply.

FAQ

How eCorpIT can help

eCorpIT is a Gurugram-based technology organisation, founded in 2021 and assessed at CMMI Level 5, with senior-led teams building AI and on-device machine-learning features across iOS, macOS, and cloud. We help teams decide what should run locally on Apple Silicon versus in the cloud, port and profile PyTorch workloads on the MPS backend, and design private, DPDP-aware inference. If you are planning on-device AI or weighing PyTorch against MLX for a product, talk to us.

References

1. PyTorch Foundation, "PyTorch 2.13 Release Blog," July 8, 2026 - https://pytorch.org/pytorch-2-13-release-blog/

1. PyTorch, "PyTorch 2.13.0 Release" (GitHub release notes) - https://github.com/pytorch/pytorch/releases/tag/v2.13.0

1. PyTorch, "FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention" - https://pytorch.org/flexattention/

1. PyTorch, "torch.nn.attention.flex_attention" documentation - https://docs.pytorch.org/docs/main/nn.attention.flex_attention.html

1. Apple, "Accelerated PyTorch training on Mac" (Metal) - https://developer.apple.com/metal/pytorch/

1. PyTorch, "Running PyTorch Models on Apple Silicon GPUs with the ExecuTorch MLX Delegate" - https://pytorch.org/running-pytorch-models-on-apple-silicon-gpus-with-the-executorch-mlx-delegate/

1. Towards Data Science, "PyTorch and MLX for Apple Silicon" - https://towardsdatascience.com/pytorch-and-mlx-for-apple-silicon-4f35b9f60e39/

1. Jiang et al., "Mistral 7B" (arXiv:2310.06825) - https://arxiv.org/pdf/2310.06825

1. Mistral AI, "Announcing Mistral 7B" - https://mistral.ai/news/announcing-mistral-7b/

1. Apple (India), "Buy Mac mini, M4 chip, 16GB memory, 512GB storage" - https://www.apple.com/in/shop/buy-mac/mac-mini/m4-chip-10-core-cpu-10-core-gpu-16gb-memory-512gb-storage

1. Hugging Face, "Mistral" model documentation - https://huggingface.co/docs/transformers/en/model_doc/mistral

1. MetalCloud, "MLX vs PyTorch: When to Use Apple's ML Framework" - https://metalcloud.space/mlx-vs-pytorch-comparison/

_Last updated: July 18, 2026._