pytorch

Key Features and Improvements

• NVGEMM Integration: Introduces CuTeDSL-generated CUTLASS kernels in Inductor, supporting epilogue fusion, scaling, and NVFP4 GEMM, and group-reduction epilogues, with autotuning alongside Triton and ATen.
• Control Flow Improvements: Adds torch.switch, which generalizes torch.cond to multi-way branching, and enables torch.while_loop to be captured in CUDA graphs.
• Dynamic Shapes Declaration: The @dynamic_spec decorator, shared across torch.compile, torch.export, and make_fx, allows dynamic shapes to be specified declaratively.
• Complex Tensor Compilation Support: torch.compile experimentally supports complex tensors by decomposing complex operations into real and imaginary operations, enabling optimization by compiler backends.
• Distributed Training Updates:
  • A new nccl2 backend, ported from torchcomms, has been added, implementing the full collective communication contract including non-blocking communicators and eager communicator splitting.
  • Fault tolerance has been elevated to a first-class concept in c10d, introducing in-place process group reconfiguration, one-sided RMA windows, and a Flight Recorder that works across all backends.
• Apple Silicon Optimization: Supports native linear algebra operations such as Jacobi kernel-based SVD, eigh, QR, and Cholesky, along with a rewrite of 5-stage reduction and further migration from MPSGraph to Metal kernels.