{"slug": "pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster", "title": "PyTorch 2.13: FlexAttention on Apple Silicon Is 12x Faster", "summary": "PyTorch 2.13, released July 8, brings FlexAttention to Apple Silicon with up to 12x speedup on sparse attention patterns, such as a 32,768-token sequence with a 256-token sliding window (35ms vs 431ms). The update also introduces fused LinearCrossEntropyLoss for up to 4x peak GPU memory reduction, benefiting fine-tuning on large-vocabulary models. The Metal backend now supports FlexAttention for both prefill and decode paths, including grouped-query attention.", "body_md": "PyTorch 2.13 shipped July 8 and the headline for Mac developers is blunt: FlexAttention now runs on Apple Silicon. The speedup on sparse attention patterns hits 12x over PyTorch’s own SDPA baseline. Not 12% — twelve times. On a 32,768-token sequence with a 256-token sliding window, the measured delta is 35ms versus 431ms. If you have been prototyping transformer variants on a MacBook and watching attention computation eat your afternoon, 2.13 removes the excuse.\n\n## What FlexAttention Actually Is\n\nFlexAttention launched in PyTorch 2.5 with a simple premise: you should not have to write CUDA kernels every time you want a non-standard attention pattern. Sliding window attention, ALiBi positional bias, document masking, causal plus prefix — all of these previously required hand-rolled kernels or slow masked-tensor workarounds. [FlexAttention](https://pytorch.org/blog/flexattention/) lets you express any attention variant as two small Python functions and hands them to `torch.compile`\n\nto fuse into a FlashAttention-grade kernel.\n\nThe catch was that it only ran on CUDA. PyTorch 2.13 fixes that. The team wrote native Metal kernels for both the sparse prefill and decode paths, including grouped-query attention. You can now run the same FlexAttention code on your M-series Mac that you run on your CUDA server — and for sparse patterns, the MPS backend holds up well.\n\n``` python\nimport torch\nfrom torch.nn.attention.flex_attention import flex_attention, create_block_mask\n\ndevice = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n\n# 256-token sliding window\ndef sliding_window(b, h, q_idx, kv_idx):\n    return abs(q_idx - kv_idx) <= 256\n\nblock_mask = create_block_mask(\n    sliding_window, B=1, H=8,\n    Q_LEN=32768, KV_LEN=32768,\n    device=device\n)\n\nq = torch.randn(1, 8, 32768, 64, device=device)\nk = torch.randn(1, 8, 32768, 64, device=device)\nv = torch.randn(1, 8, 32768, 64, device=device)\n\nout = flex_attention(q, k, v, block_mask=block_mask)\n```\n\nThat code runs on Apple Silicon now. The speedup depends entirely on how sparse your pattern is — which is the right tradeoff to make. FlexAttention on MPS is not faster than SDPA on dense causal attention. It is faster when you have a pattern that lets it skip work, which is exactly when you want the shortcut.\n\n## The Performance Numbers in Honest Context\n\nThe 12x figure comes from a specific shape: 1 batch, 8 heads, 32,768 sequence length, 64 head dimension, 256-token sliding window at 0.8% density. Shorter sequences compress the gain: an 8K-length, 64-token window delivers about 4x. Across sliding-window patterns generally, FlexAttention on MPS runs around 20x faster than SDPA and about 3x faster than FlashAttention 2 with a causal mask.\n\nFor standard full causal attention, do not switch. FlexAttention generates its kernels through `torch.compile`\n\n, which does not produce kernels as tight as a hand-tuned FlashAttention 3 implementation. Expect 5–15% slower throughput on pure causal attention. The right mental model: use FlexAttention when your pattern has sparsity to exploit, use SDPA or Flash for everything else. The API does not pretend otherwise.\n\n## Fused LinearCrossEntropyLoss: The Fine-Tuning Win\n\nFlexAttention gets the headline, but `nn.LinearCrossEntropyLoss`\n\nmight matter more if you are fine-tuning rather than training from scratch. It fuses the final projection layer and the cross-entropy loss into a single operation, which cuts peak GPU memory by up to 4x on large-vocabulary models. A model with a 128K-token vocabulary that was OOM-ing on your 24GB card might now fit. On Apple Silicon with its unified memory pool, that headroom is even more valuable — you are working with a fixed budget shared by the entire system.\n\n``` python\nimport torch.nn as nn\n# Drop-in replacement for separate Linear + CrossEntropyLoss\nloss_fn = nn.LinearCrossEntropyLoss()\nloss = loss_fn(hidden_states, weight, labels)\n```\n\n## CuTeDSL, torchcomms, and What to Check Before Upgrading\n\nPyTorch 2.13 adds [CuTeDSL](https://pytorch.org/blog/pytorch-2-13-release-blog/) as a second code-generation backend in Inductor alongside Triton, targeting GEMM and RMSNorm. If you hit compilation latency in Inductor, this is worth testing — it generates CUTLASS-grade kernels and moves compilation to a subprocess pool to clear the GIL bottleneck.\n\ntorchcomms is a new distributed communications backend with better fault tolerance — graceful timeout, partial-group recovery — and structured logging for large-cluster debugging. If you run multi-node training, this improves observability significantly.\n\nTwo breaking changes worth a grep before you upgrade:\n\n- Named tensors are gone. The feature was deprecated; now it is removed. If your codebase uses them, it will break at runtime.\n- Distributed collective renames:\n`all_gather_into_tensor`\n\nbecomes`all_gather_single`\n\n,`reduce_scatter_tensor`\n\nbecomes`reduce_scatter_single`\n\n. Search your distributed training code before upgrading.\n\n## Should You Upgrade Now?\n\nIf you work on Apple Silicon and use any non-standard attention pattern — sliding window, ALiBi, document masking — upgrade now. Run `pip install --upgrade torch torchvision torchaudio`\n\n, swap your device to `\"mps\"`\n\n, and point [flex_attention](https://docs.pytorch.org/docs/stable/nn.attention.flex_attention.html) at a block mask. The gains are real and available today.\n\nIf you are on CUDA-only infrastructure, 2.13 is a solid release but not urgent unless you are hitting Inductor compilation overhead or distributed fault-tolerance gaps. The fused `LinearCrossEntropyLoss`\n\nis worth adding regardless of hardware if you are fine-tuning large-vocabulary models.\n\nPyTorch is hosting a live Q&A for 2.13 on July 22 — that is tomorrow — if you have questions for the team. Check the [official release blog](https://pytorch.org/blog/pytorch-2-13-release-blog/) for registration details. The release comprises 3,328 commits from 526 contributors, which gives some indication of how much ground 2.13 covers beyond what fits in a single post.", "url": "https://wpnews.pro/news/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster", "canonical_source": "https://byteiota.com/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster/", "published_at": "2026-07-21 16:08:50+00:00", "updated_at": "2026-07-21 16:26:50.548634+00:00", "lang": "en", "topics": ["machine-learning", "developer-tools", "ai-infrastructure"], "entities": ["PyTorch", "Apple Silicon", "FlexAttention", "Metal", "SDPA", "FlashAttention", "LinearCrossEntropyLoss", "CuTeDSL"], "alternates": {"html": "https://wpnews.pro/news/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster", "markdown": "https://wpnews.pro/news/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster.md", "text": "https://wpnews.pro/news/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster.txt", "jsonld": "https://wpnews.pro/news/pytorch-2-13-flexattention-on-apple-silicon-is-12x-faster.jsonld"}}