PyTorch 2.14 expands GPU kernels, distributed recovery, and dynamic shapes

PyTorch 2.14 incorporates 2,995 commits from 487 contributors since version 2.13. The release announcement centers on four areas: an NVIDIA matmul backend for Inductor, a rebuilt distributed communication stack, native linear algebra on Apple silicon, and one dynamic-shape specification shared across compilation and export.

Many additions carry the API Unstable label, so interfaces and configuration names may change. Production deployments should pin PyTorch and companion packages, test generated kernels on their target hardware, and review the release notes before upgrading.

NVGEMM joins Inductor’s backend race

NVGEMM, a CuTeDSL-generated CUTLASS path, has evolved from the standalone kernel introduced in 2.13 into an Inductor backend. It supports epilogue fusion, which combines a matrix multiplication with subsequent operations such as bias addition, activation, or rescaling in one kernel. That fusion reduces kernel launches and avoids another trip through GPU memory.

During autotuning, Inductor can compare NVGEMM with Triton and ATen for mm, addmm, and scaled_mm. The backend also supports fused low-precision NVFP4 paths on Blackwell GPUs and caches compiled kernels on disk.

Projects can enable it by adding NVGEMM to max_autotune_gemm_backends while using max_autotune. NVGEMM requires nvidia-cutlass-dsl 4.6.0, and its NVFP4 kernels require Blackwell hardware. Triton handles epilogues that NVGEMM cannot express.

Process groups gain recovery hooks

torchcomms, introduced in 2.13, now ships in-tree as the nccl2 backend for PyTorch distributed’s c10d layer. The backend is controlled by the USE_C10D_NCCL build option and implements the full collective contract with nonblocking communicators, eager communicator splitting, and one-sided remote-memory-access windows. An nccl-lazy wrapper retains lazy peer-to-peer initialization for workloads that depend on it.

New reconfiguration interfaces on Backend and ProcessGroup let backends and orchestration layers rebuild a group in place after a rank fails. Earlier recovery paths commonly destroyed the process group and restarted the job, losing warm state across the cluster. The new path also carries abort hooks and hooks that run before and after collectives. Gloo receives the same fault-tolerance plumbing as nccl2.

Flight Recorder now records collective traces through ProcessGroup hooks, extending hang and mismatch diagnostics to Gloo and custom backends. Symmetric memory also gains a one-sided get primitive that reads a peer’s allocation without peer participation, supporting patterns such as embedding lookup, weight transfer, and expert routing.

Metal fills MPS’s linear-algebra gaps

The MPS backend replaces several CPU fallbacks and MPSGraph paths with native Metal kernels:

• SVD, eigh, and lstsq use Jacobi-style kernels for float32 and complex64.
• Cholesky uses a faster panel-factorization algorithm, with reported gains of 1.2× to 2.8× depending on matrix size. The release also fixes a complex-dtype bug that could silently return incorrect results.
• lu_factor and lu_solve use handwritten Metal kernels, with reported speedups above 100× for small batched matrices.
• QR, matrix_exp, and linalg.polar are now available on MPS.

A second prefill-attention kernel uses Apple’s new Metal Performance Primitives in macOS 26.2. It adapts MLX’s approach for M5 chips to earlier Apple silicon generations and supports fp16 and bf16, head dimensions of 64, 96, 128, or 256, and query lengths above 8. The author reports gains of roughly 2× to 4×, with the largest improvements at smaller head dimensions and longer sequences. MPS selects the kernel automatically for eligible shapes and dtypes.

A routing bug previously sent single-token F.linear decode workloads through a slower MPS path, producing a reported 8.5× penalty for bf16 and fp16. PyTorch 2.14 routes those workloads to new GEMV kernels optimized for the vector-matrix products common in autoregressive inference. CTC loss also gains MPS forward and backward implementations, removing its CPU fallback.

Multi-way branches shed their nesting

torch.cond expresses a two-way branch, so developers previously represented n-way dispatch as nested conditionals. That structure enlarges traced graphs, especially in mixture-of-experts models. The new switch prototype selects among multiple branches using an index and deduplicates lifted arguments shared by those branches.

The public name requires care in version 2.14.0. Although the feature is described as torch.switch, the callable ships at torch._higher_order_ops.switch. Its private module path and prototype status make it unsuitable for code that requires a stable public API, as also noted in independent analysis.

torch.while_loop can now be captured in a CUDA Graph through CUDA’s conditional while nodes. The practical benefit is full graph capture for workloads with data-dependent loop counts, including reductions over variable-length index tensors and losses over varying numbers of packed sequences.

One shape contract spans three entry points

Dynamic-shape declarations previously differed across PyTorch’s tracing and compilation interfaces: torch.export accepted a dynamic_shapes dictionary, torch.compile exposed a coarse dynamic= flag, and make_fx relied on a global tracing mode. PyTorch 2.14 introduces a ShapesSpec API that can be attached to a function or module with @dynamic_spec:

from torch.fx.experimental.dynamic_spec import ShapeVar, dynamic_spec

batch = ShapeVar("batch", min=2, max=128)

@dynamic_spec({"x": (batch, 768)})
def forward(x):
    return model(x)

torch.export, torch.compile, and make_fx can consume the same specification. Dimensions declared through this API become unbacked symbolic values. Shape-dependent branches therefore raise data-dependent errors where a guard-based path might have recompiled for a new shape.

Hardware support expands by platform

• AMD: The release adds ROCm 7.14 support.
• Intel: XPU gains native graph capture.
• NVIDIA: Compiler targeting extends to Rubin’s sm_107 architecture.
• Complex tensors: Experimental torch.compile support decomposes eligible operations into real and imaginary components that compiler backends can optimize. This enables compiled paths for more signal-processing and scientific workloads.

Python 3.15 and free-threaded 3.15t receive eager-mode wheels hosted on download.pytorch.org. Invoking torch.compile under Python 3.15 raises a RuntimeError while Dynamo support remains in development. These wheels are currently distributed outside PyPI.

torchvision 0.29 is ABI-stable with PyTorch 2.14. The compatibility guarantee covers PyTorch 2.15, 2.16, and subsequent releases within the stated ABI window, reducing the need to reinstall torchvision after each PyTorch upgrade.

Kernel and compiler changes

• Inductor enables simple_overlap reordering by default. It interleaves collectives with independent computation, improving GPU utilization where the schedule exposes enough work to overlap.
• Scaled dot-product attention dispatches rank-3 inputs to fused CUDA backends. Shapes with missing or flattened batch dimensions previously used the slower math path.
• cuBLASLt joins CUTLASS as a grouped-GEMM backend. It becomes the default for fp16 on Blackwell with CUDA 13.2 or later and Hopper with CUDA 13.3 or later. Its strongest results appear on ragged, mixture-of-experts-style groups; CUTLASS generally performs better on uniform groups.
• CUDA TunableOp registers cuBLASLt heuristic candidates alongside the cuBLAS default. Reported H100 results show an unchanged mean, while individual shapes range from 0.66× to 1.57×. Its value depends on finding poorly served shapes.
• The new AOTInductorModelContainerCreateWithExternalConstants C API lets multiple AOTI model containers share one GPU-resident copy of their weights.
• FlexAttention on AMD RDNA3 GPUs adds sequence-length-aware tile configurations, with reported latency reductions of 2× to 8× for sequence lengths in the low hundreds.

PyTorch is hosting a live Q&A with Andrey Talman, Natalia Gimelshein, and Joe Spisak. PR-level implementation details and the complete change list are available in the GitHub release notes.