cd /news/machine-learning/pytorch-2-13-flexattention-on-apple-… · home topics machine-learning article
[ARTICLE · art-60925] src=byteiota.com ↗ pub= topic=machine-learning verified=true sentiment=↑ positive

PyTorch 2.13: FlexAttention on Apple Silicon and 4x LLM Memory Savings

PyTorch 2.13 introduces FlexAttention support on Apple Silicon, delivering up to 12x speedup for sparse attention patterns on M-series Macs, and a new fused LinearCrossEntropyLoss operator that reduces peak GPU memory by up to 4x during large-vocabulary model fine-tuning. The release also adds a deterministic backward pass for FlexAttention on CUDA, a CuTeDSL backend for the Inductor compiler, and a new distributed communications backend called torchcomms.

read4 min views1 publishedJul 15, 2026
PyTorch 2.13: FlexAttention on Apple Silicon and 4x LLM Memory Savings
Image: Byteiota (auto-discovered)

PyTorch 2.13 dropped this week with two features that change real workflows. If you train on a Mac, FlexAttention now runs on Apple Silicon with up to 12x speedup over standard attention on sparse patterns. If you fine-tune large-vocabulary models, a new fused loss operator cuts peak GPU memory by up to 4x — without changing your model architecture. The rest of the release is also worth knowing.

FlexAttention Finally Runs on Apple Silicon #

FlexAttention is PyTorch’s API for writing custom attention patterns as plain Python functions. Instead of dropping down to CUDA or Metal to implement sliding window attention, document-level masking, or any other sparse variant, you write a score-modification function and PyTorch compiles it into a fused kernel. The API has been available on CUDA for a while. As of 2.13, it works on MPS — the Metal backend that runs on M-series Macs.

The Benchmark

The performance numbers are not subtle. A 256-element sliding window pattern on a 32,768-length sequence runs in about 35ms with FlexAttention on MPS, compared to 431ms with PyTorch’s standard scaled dot-product attention — a 12.3x difference at 0.8% sparsity density. Mac developers who previously had to fall back to CPU for custom attention patterns now have hardware acceleration for them.

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 < 256

mask = create_block_mask(sliding_window, B=1, H=8, Q_LEN=32768, KV_LEN=32768)
out = flex_attention(q, k, v, block_mask=mask)

Deterministic Backward on CUDA

FlexAttention on CUDA also gains a deterministic backward pass in this release. Previously, the backward path used atomic operations — fast but nondeterministic, meaning two runs with identical inputs could produce slightly different gradients. The new compute_dq_write_order

path pre-computes a stable write ordering, eliminating gradient variance without a meaningful performance penalty.

LinearCrossEntropyLoss: The Memory Fix Nobody Asked for Out Loud #

The cross-entropy loss step at the end of a language model forward pass has always been a peak memory landmine. Before computing the loss, the model must produce logits — a tensor of shape [batch, sequence_length, vocab_size]

. For a model with a 128,000-token vocabulary (common in modern multilingual LLMs), this single tensor can consume more GPU memory than all your model parameters combined.

PyTorch 2.13 ships nn.LinearCrossEntropyLoss

, a fused operator that combines the final output projection with the cross-entropy computation. It processes hidden states in chunks, masks ignored tokens before projecting, and never materializes the full logit matrix. The result: up to 4x less peak memory during loss computation, with no change to model behavior or numerical output. The operator supports label smoothing, weight tying, z-loss regularization, torch.compile

, and tensor parallelism.

This is one of those improvements that makes you wonder why it was not in PyTorch years ago. The logit materialization problem has been a known pain point for anyone fine-tuning large-vocabulary models on consumer hardware. The fix is now in core — no third-party libraries required.

CuTeDSL: A Second Code Path for Inductor #

PyTorch’s Inductor compiler previously generated GPU code exclusively through Triton. In 2.13, it gains a second backend: CuTeDSL, a Python-native DSL that gives Inductor access to CUTLASS-grade GEMM kernels. The two backends coexist — CuTeDSL targets matrix multiplication and normalization operations where it can compile faster than Triton in certain configurations. For most projects this is transparent. For teams where Triton compilation times are a bottleneck, it is a useful alternative.

Distributed and Platform Updates #

On the distributed side, 2.13 ships torchcomms

, a new communications backend for PyTorch Distributed that improves fault tolerance and debuggability at cluster scale. FSDP2 now optionally overlaps reduce-scatter and all-gather operations through a dedicated process group, increasing throughput for multi-node training.

Platform coverage also expands. ExecuTorch is integrated into PyTorch Core, making on-device inference a first-class part of the framework. ROCm gains AOTriton 0.12b with native HIP CMake. ARM gets Armv9-A torch.compile

targeting. Python 3.15 wheels are available for Linux, including free-threaded builds. One notable removal: Python 3.13t (free-threaded) support has been dropped — the dependency conflict with 3.15t builds made it untenable.

How to Upgrade #

pip install --upgrade torch

pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu128

pip install --upgrade torch --index-url https://download.pytorch.org/whl/rocm7.2

The upgrade is straightforward for most projects. Two things to check: named tensors were previously deprecated and are now hard-removed, so migrate any code that uses them. Distributed collective names have been updated to align with the torchcomms API — check renamed calls if you run multi-node training. Full details in the official release blog, GitHub release notes, and the official install page.

For Apple Silicon users and anyone training large-vocabulary models, 2.13 is a meaningful step up. If you work with vLLM or other inference frameworks that sit on top of PyTorch, check their compatibility notes before upgrading. Otherwise, install it now.

── more in #machine-learning 4 stories · sorted by recency
── more on @pytorch 3 stories trending now
sponsored brought to you by zahid.host 4,200+ EU-deployed projects
reading about agents? ship yours in a single git push.

Run your AI side-project on zahid.host

EU-based hosting, git-push deploys, automatic HTTPS, no cold starts. Free tier with a custom domain — perfect for shipping the agent you just read about.

$git push zahid main
Live at https://your-agent.zahid.host
Get free account → Pricing
from €0/mo · no card required
LIVE [news/pytorch-2-13-flexatt…] indexed:0 read:4min 2026-07-15 ·