PyTorch/XLA 2.4 为 TPU 和 GPU 上的深度学习提供了显著的改进。该版本引入了对 Pallas 的改进，这是一种自定义的核语言，现在支持 TPU 和 GPU，使用 Python 代码提高性能。新的 API 调用，如 torch_xla.sync()，简化了与现有 PyTorch 工作流的集成。实验性的急切模式允许在目标硬件上立即执行操作，尽管 TPU 需要“mark_step”调用以进行模拟。Pallas 的改进包括 Flash Attention 和 Paged Attention 支持，以及内置的 Megablocks 块稀疏核用于组矩阵乘法。此外，新的 TPU 命令行界面 tpu-info 通过显示利用率和设备信息来帮助调试，类似于 Nvidia 的 nvidia-smi 工具。尽管存在这些变化，但现有的代码仍然兼容，使开发者可以无缝升级。这些更新总体旨在提高 PyTorch/XLA 的可用性和性能。