viable/strict/1788439121: [MPS] migrate log softmax to metal (#195441)
PyTorch 在 viable/strict/1788439121 版本中合并了 PR #195441，将 log_softmax 迁移到 Metal (MPS) 后端，包含 4 个 kernel，针对不同维度大小优化。基准测试显示多种模型（如 wav2vec2、GPT-OSS、Whisper、Qwen3、ViT、DeepSeek、BERT、SegFormer、GLM）的 log_softmax 延迟显著降低，例如 wav2vec2-xlsr-53-pt 从 34.9us 降至 5.1us（6.8 倍加速），bert-base-uncased 从 8927.1us 降至 5079.3us（1.76 倍加速）。
发展脉络
- 首次出现viable/strict/1788439121: [MPS] migrate log softmax to metal (#195441)PyTorch Core
- 当前判断PyTorch 持续优化 MPS 后端，表明 Apple Silicon 在 AI 推理和训练中的重要性日益增加。随着 Mac 成为开发者常用平台，MPS 性能提升有助于降低开发者使用门槛，促进本地模型开发。该 PR 是 PyTorch 社区对硬件多样化支持的体现，可能推动更多模型在 Apple 设备上高效运行。可验证的下一信号：观察 PyTorch 官方文档或发布说明中 MPS 支持算子的增加情况，以及 Apple 在 AI 框架上的投入。Agent Pulse · 分析
PyTorch 在 2026 年 9 月 3 日发布的 viable/strict/1788439121 版本中，合并了 PR #195441，将 log_softmax 操作迁移到 Metal Performance Shaders (MPS) 后端。该 PR 是系列优化的一部分，旨在减少现代模型推理和训练中常用操作的内存使用。PR 实现了 4 个 kernel：log_softmax_row 用于小维度（≤128），log_softmax_partial 和 log_softmax_finalize 用于大维度（>2048）的两遍 kernel，以及 log_softmax 通用 kernel 处理 strided 张量。基准测试显示，在多种模型上 log_softmax 延迟显著降低，例如 wav2vec2-xlsr-53-pt 加速 6.8 倍，gpt-oss-20b 加速 4.85 倍，whisper-large-v3-turbo 加速 4.13 倍，Qwen3 加速 3.84 倍，ViT-base 加速 3.69 倍，DeepSeek-V4-Flash 加速 3.08 倍，bert-base-uncased 加速 1.76 倍。这些改进主要针对 Apple Silicon 上的 MPS 后端，对使用 PyTorch 在 Mac 上进行模型推理和训练的开发者有直接影响。
该 PR 展示了针对特定硬件后端（MPS）进行算子级优化的典型方法：通过分析 log_softmax 的归约维度大小，设计多个专用 kernel 以充分利用 GPU 并行性。对于小归约维度（≤128），使用单 kernel 减少启动开销；对于大归约维度（>2048），采用两遍 kernel 避免中间张量存储；通用 kernel 处理非连续张量。这种优化策略可推广到其他类似算子，如 softmax、layer norm 等。可验证的下一信号：检查 PyTorch 后续版本是否将类似优化应用于其他算子，或是否在 MPS 后端实现更多融合 kernel。
PyTorch 持续优化 MPS 后端，表明 Apple Silicon 在 AI 推理和训练中的重要性日益增加。随着 Mac 成为开发者常用平台，MPS 性能提升有助于降低开发者使用门槛，促进本地模型开发。该 PR 是 PyTorch 社区对硬件多样化支持的体现，可能推动更多模型在 Apple 设备上高效运行。可验证的下一信号：观察 PyTorch 官方文档或发布说明中 MPS 支持算子的增加情况，以及 Apple 在 AI 框架上的投入。
对于依赖 PyTorch 在 Mac 上进行模型开发的团队，该优化可显著降低推理延迟，提升开发效率和用户体验。例如，语音识别模型 wav2vec2 的 log_softmax 延迟降低 6.8 倍，可加快实时处理。对于提供本地 AI 解决方案的公司，MPS 性能提升有助于产品在 Apple 设备上的流畅运行，增强竞争力。可验证的下一信号：观察相关模型在 Mac 上的端到端推理速度是否提升。
随着 MPS 后端的持续优化，预计更多 PyTorch 算子将获得类似加速，使得在 Apple Silicon 上运行大型模型的性能接近甚至超越 CUDA 平台。这将促进边缘 AI 和本地推理的发展，减少对云服务的依赖。可验证的下一信号：关注 PyTorch 后续版本中 MPS 相关 PR 的数量和性能提升幅度。