导读：PyTorch 2.9 正式发布，其中包括 C++/CUDA 扩展的稳定 ABI 更新、多 GPU 编程的对称内存支持、torch.compile 的图中断控制、更广泛的硬件加速支持（ROCm, XPU, CUDA 13），以及 FlexAttention 在 Intel GPU 和 X86 CPU 上的优化，同时增强了 Arm 平台的性能和覆盖范围。

在人工智能与深度学习飞速发展的今天，高效、灵活且易于使用的深度学习框架是推动技术进步的关键。PyTorch 作为业界领先的开源机器学习库，持续不断地进行迭代更新，以满足日益增长的计算需求和开发者社区的期望。

近日，PyTorch 基金会正式发布了 PyTorch 2.9 版本，带来了多项重要更新和优化，旨在提升性能、增强跨平台兼容性，并为开发者提供更强大的工具集。

本文将深入解读 PyTorch 2.9 的核心亮点，帮助开发者全面了解此次升级的意义和潜在影响。

Release Blog 内容解读

PyTorch 2.9 的发布标志着其在多个关键领域取得了显著进展。本次更新共包含 3216 个提交，由 452 位贡献者共同完成，充分体现了 PyTorch 社区的活力与协作。以下是本次发布的核心内容概览：

1. 稳定的 libtorch ABI 更新： 对于需要构建自定义 C++/CUDA 扩展的开发者而言，这是一个重大利好。PyTorch 2.9 进一步完善了其 C++ 接口的稳定性（stable ABI），允许开发者使用一个版本的 PyTorch 构建的扩展，在另一个版本的 PyTorch 中也能兼容运行。这大大降低了第三方库的维护成本和用户的集成难度。新版本引入了设备工具（如 Device Guard 和 Stream）以及更多的 torch::stable::Tensor API（如默认构造函数、is_cpu、scalar_type 和 get_device_index），并开放了更多稳定的 ATen 操作。目前，Flash-Attention 3 已成功基于此 ABI 实现。

2. 对称内存（Symmetric Memory）编程： 为了简化和优化多 GPU 之间的通信，PyTorch 2.9 引入了对称内存编程模型。该模型支持在 GPU 内核中直接进行通信（如 put/get 操作），并能实现超低延迟的远程内存访问，甚至支持单向操作，无需远程 GPU 配合。这为构建高度定制化的多 GPU 内核提供了前所未有的灵活性，尤其是在处理大规模模型和复杂通信模式时。新版本支持分配可远程直接访问的对称张量，并提供了加速的集体通信操作，如 one_shot_all_reduce、two_shot_all_reduce_ 等，以及针对 MoE 模型优化的 all_to_all_v 操作。

3. torch.compile 的图中断控制： torch.compile 作为 PyTorch 2.x 系列的核心功能，旨在通过图编译来提升模型性能。在 2.9 版本中，torch._dynamo.error_on_graph_break() 上下文管理器/装饰器允许开发者更精细地控制图中断（graph break）的处理方式。开发者可以选择在遇到图中断时直接报错，或尝试恢复执行。这为调试和优化编译过程提供了更大的灵活性，尤其是在处理动态形状或复杂控制流的代码时。

4. 扩展的轮子（Wheel）变体支持： PyTorch 持续致力于提升其在不同硬件平台上的可用性。2.9 版本扩展了对 AMD (ROCm)、Intel (XPU) 和 NVIDIA CUDA 13 的轮子变体支持。这意味着开发者可以在更多异构计算环境中更便捷地安装和使用 PyTorch。虽然目前 ROCm 和 XPU 平台仅支持 Linux，但 NVIDIA CUDA 轮子已同时支持 Windows 和 Linux。这一进展是 WheelNext 倡议的一部分，旨在构建更统一和高效的 Python 包生态系统。

5. FlexAttention 优化：

• Intel GPU 支持：
   FlexAttention 的前向和后向传播现在支持 Intel GPU，使得开发者在 Intel 硬件上也能享受到注意力机制的性能提升，无需修改现有代码。
• X86 CPU 上的 Flash Decoding：
   针对 LLM 推理中常用的 Flash Decoding 技术，PyTorch 2.9 在 X86 CPU 的 Inductor 后端实现了基于 FlexAttention 的优化。通过对 KV 序列进行分区和归约，该优化能够显著提升 CPU 利用率，尤其是在小批量、小头数或短 Q 序列长度但长 KV 序列长度的场景下，有助于改善长上下文 LLM 的推理性能。

6. Arm 平台改进与优化： PyTorch 2.9 在 Arm 平台上也带来了多项关键改进，包括：

• torch.compile 性能提升：
   在 Arm CPU (aarch64) 上，torch.compile 模式下的 TorchBench、HuggingFace 和 TIMM 测试套件性能均优于 Eager 模式。
• 算子优化：
   针对 AArch64 平台，对卷积、激活和量化算子进行了扩展和优化，以加速模型执行。
• CI 覆盖范围扩大：
   新增了基于 Arm Neoverse V2 的 AWS Graviton 4 实例，扩展了 Arm 平台的持续集成测试覆盖。

亮点分析

1. C++/CUDA 扩展的 ABI 稳定性：
    极大地降低了第三方库的集成门槛，提升了 PyTorch 生态系统的整体健壮性。
2. 对称内存编程：
    为高性能多 GPU 计算提供了强大的新工具，有望在分布式训练和大规模模型推理中发挥关键作用。
3. torch.compile 图中断控制：
    增强了 torch.compile 的灵活性和可调试性，使开发者能更好地驾驭复杂的模型编译过程。
4. 跨平台硬件加速扩展：
    显著提升了 PyTorch 在 AMD、Intel 和最新 NVIDIA GPU 上的可用性，推动了 AI 计算的普惠性。
5. FlexAttention 在 CPU 和 Intel GPU 上的优化：
    进一步拓展了高性能注意力机制的应用范围，尤其是在 CPU 推理和 Intel 硬件生态中，为 LLM 推理带来了新的性能增长点。

应用场景

PyTorch 2.9 的更新将广泛应用于以下场景：

• 高性能计算与分布式训练：
   对称内存和优化的多 GPU 通信将加速大规模模型（如大型语言模型、多模态模型）的训练过程。
• 第三方库与框架开发：
   稳定的 C++ ABI 使得第三方库（如 Flash-Attention）能够更轻松地与 PyTorch 集成，并保持跨版本兼容性。
• 异构计算环境：
   扩展的轮子支持使得开发者可以在更多种类的硬件上部署和运行 PyTorch 模型，包括 AMD 和 Intel 的加速器。
• LLM 推理优化：
   FlexAttention 在 CPU 和 Intel GPU 上的优化，以及 Flash Decoding 的支持，将显著提升长上下文 LLM 的推理速度和效率，降低部署成本。
• 嵌入式与边缘计算：
   Arm 平台的持续优化和改进，将有助于 PyTorch 在移动设备、嵌入式系统等资源受限的环境中获得更好的性能表现。

PyTorch 2.9 的发布是 PyTorch 发展历程中的又一个重要里程碑。它不仅在核心性能和稳定性上进行了扎实的提升，更在多 GPU 编程、跨平台支持和特定优化（如 FlexAttention）方面带来了令人兴奋的创新。

这些更新将进一步巩固 PyTorch 在深度学习领域的领先地位，并为开发者提供更强大、更灵活的工具，以应对未来日益复杂的 AI 应用挑战。社区的积极贡献和持续迭代，预示着 PyTorch 在推动人工智能技术民主化和普及化方面将扮演越来越重要的角色。

参考资料

• https://pytorch.org/blog/pytorch-2-9/