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.