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.

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. python 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 Standard / CUDA pip install --upgrade torch CUDA 12.8 pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu128 ROCm 7.2 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 https://pytorch.org/blog/pytorch-2-13-release-blog/ , GitHub release notes https://github.com/pytorch/pytorch/releases/tag/v2.13.0 , and the official install page https://pytorch.org/get-started/locally/ . 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 https://byteiota.com/vllm-v0-21-kv-offload-spec-decode-upgrade/ , check their compatibility notes before upgrading. Otherwise, install it now.