PyTorch 2.13 Brings FlexAttention to Apple Silicon — and It's Faster Than You'd Expect

PyTorch 2.13, released on July 8, 2026, quietly shipped one of the more practically useful updates for developers running models on Mac hardware: native FlexAttention support for the Apple Silicon MPS (Metal Performance Shaders) backend. For sparse attention patterns, the speedup over standard Scaled Dot-Product Attention (SDPA) reaches roughly 12x. That's not a rounding error — it's a meaningful change in what's feasible to run locally on a MacBook or Mac Studio.

This post walks through what FlexAttention actually is, why the MPS port matters, what the performance numbers look like in practice, and what else landed in 2.13 that's worth knowing about.

What FlexAttention Is (and Isn't)

FlexAttention was introduced in PyTorch 2.5 as a way to write custom attention variants without dropping into CUDA. The core idea is simple: instead of implementing a new attention kernel from scratch, you write a short Python function that describes how attention scores should be modified or masked, and PyTorch's compiler turns that into a fused, block-sparse kernel automatically.

There are two main hooks:

• score_mod(score, b, h, q_idx, kv_idx) — modifies attention logits before softmax. Useful for things like ALiBi positional biases or soft-capping.
• mask_mod(b, h, q_idx, kv_idx) — returns a boolean indicating whether a query-key pair should attend to each other. This is where sparsity patterns like sliding-window or causal masking live.

When you call create_block_mask with a mask_mod function, PyTorch builds a BlockMask object that encodes the sparsity structure. The compiled kernel then skips entire blocks of the attention matrix that are masked out — rather than computing them and zeroing the result afterward. That's where the speedup comes from: less work, not just faster work.

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

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

block_mask = create_block_mask(sliding_window, B=None, H=None, Q_LEN=32768, KV_LEN=32768, device="mps")
output = flex_attention(q, k, v, block_mask=block_mask)

Before 2.13, this API existed on CUDA but not on MPS. Apple Silicon users were stuck with dense SDPA or maintaining their own Metal kernels.

The MPS Port: What Changed Under the Hood

The PyTorch 2.13 release blog describes two related changes that together make FlexAttention on MPS work well.

First, the team wrote hand-optimized Metal kernels specifically for sparse prefill and decode paths. These kernels support grouped-query attention (GQA) and captured buffers, which means they cover the attention patterns used in most modern transformer architectures. The kernels are designed to skip masked blocks at the hardware level, which is what produces the large speedup on sparse patterns.

Second, and more broadly, PyTorch 2.13 migrated a wide range of MPS operations — including copy/cast, reductions, sort, and scatter/gather — from Apple's MPSGraph framework to hand-written Metal compute kernels. MPSGraph is a higher-level abstraction that adds per-operation compilation overhead and limits fine-grained control over memory layout and thread dispatch. Moving to native Metal eliminates that overhead and gives PyTorch more direct control over how work is scheduled on the GPU.

The result is that even operations unrelated to attention benefit from lower kernel launch latency in 2.13.

The Numbers

The benchmark that appears in the release notes uses a 1×8×32768×64 attention shape with a 256-element sliding window — that's a sequence length of 32,768 tokens with a window density of about 0.8%. At that sparsity level:

• FlexAttention (MPS): ~35 ms
• SDPA (MPS): ~431 ms
• Speedup: ~12.3x

The caveat is important: this speedup is specific to sparse patterns. For dense attention — where every query attends to every key — there's no sparsity to exploit, and SDPA remains the better choice. The more sparse your attention pattern, the more FlexAttention helps.

For context, sliding-window attention is used in models like Mistral and Phi-3, and is increasingly common in long-context architectures where full quadratic attention is too expensive. If you're fine-tuning or running inference on those models locally, this update is directly relevant.

One Practical Limitation

The MPS implementation of FlexAttention is currently marked API Unstable in the release notes. That means the interface may change in future versions. The practical advice is to pin your PyTorch version if you're building something that depends on this, and to run a correctness check against SDPA on small inputs before trusting results from the new kernels.

It's also worth noting that Apple's own MLX framework often outperforms PyTorch MPS for latency-sensitive LLM inference tasks. PyTorch MPS is generally the better fit for training workflows and for code that needs to stay within the PyTorch ecosystem — for example, if you're using torch.compile, Hugging Face Transformers, or other libraries that assume PyTorch as the backend.

What Else Landed in 2.13

FlexAttention on MPS is the headline for Apple Silicon users, but 2.13 includes several other changes worth scanning:

CuTeDSL backend for Inductor. A new code-generation path based on NVIDIA's CuTe library, optimized for GEMM and RMSNorm. It compiles faster than Triton for these operations and gives more control over tensor layouts.

torch.compiler.set_default_backend. A new API that lets you set a process-wide default compiler backend, so you don't have to pass backend= to every torch.compile() call. Useful for large codebases with many compilation sites.

nn.LinearCrossEntropyLoss. A fused module that combines linear projection and cross-entropy loss. When used with torch.compile, it can cut peak GPU memory by up to 4x for large-vocabulary models — relevant for anyone training language models with large output vocabularies.

Armv9-A support. torch.compile on AArch64 now correctly targets Armv9-A CPUs like AWS Graviton4, propagating SVE (Scalable Vector Extension) support through Inductor codegen.

Why This Matters for Local Development

The broader trend here is that the gap between what's practical on consumer hardware and what requires a data center GPU is narrowing. FlexAttention on MPS is one piece of that: it makes a class of attention patterns that were previously slow on Mac hardware significantly faster, without requiring any changes to model code beyond switching the device.

For researchers prototyping long-context models, or developers building applications that run inference locally, this is a meaningful quality-of-life improvement. The full release notes are worth reading if you're on Apple Silicon — there's more there than the headline feature.

---

Sources: PyTorch 2.13 Release Blog · PyTorch 2.13 GitHub Release · FlexAttention Introduction · eCorpIT: FlexAttention on Apple Silicon