Google 开源跨平台 GPU 推理引擎 ML Drift

Google 发布了 ML Drift，这是一个开源、跨平台的 GPU 计算引擎，专为端侧 AI 和机器学习推理而设计。它采用 Apache 2.0 许可，抽象了不同 GPU 硬件和底层 API 的复杂性，支持 OpenGL ES、OpenCL、Metal 和 WebGPU。

ML Drift 是驱动 LiteRT 的核心 GPU 加速引擎。通过将其开源，Google 使开发者能够为从传统计算机视觉、音频和深度处理到生成式 AI 的各种应用构建自定义图形和推理运行时。

ML Drift 的主要目标之一是简化对开发者在为边缘设备创建应用时所遇到的多样化 GPU 硬件的访问。

开发者面临着各种各样的 GPU 架构、驱动版本和底层 API，并且无法预先知道应用将在哪种具体硬件上运行。

Google 表示，ML Drift 旨在为可移植性和性能建立新的基础，提供“传统机器学习的稳定性以及最先进生成式 AI 的峰值性能”。

为实现这一目标，ML Drift 在传统 TFLite GPU delegate 的基础上引入了多项技术改进，包括通过张量虚拟化实现统一着色器、可扩展的自定义操作以及五维张量支持。它还针对 Arm、Intel 和 Qualcomm Technologies 的 GPU 加入了硬件特定优化，以“最大化现代 GPU 的硬件利用率”。

ML Drift 通过 Dawn（Chromium 的 WebGPU 实现）扩展了对 Windows 和 Linux 上推理的桌面支持，同时在 macOS 上保留了原生 Metal 后端支持。Dawn 提供了利用 WebGPU 硬件抽象层的统一 API 模型，避免直接与 DirectX 和 Vulkan 等平台特定图形 API 交互。

Google 在开发者级工作站硬件上使用 Gemma 模型对 ML Drift 进行了基准测试，报告称与其他框架（包括 llama.cpp 和 MLX LM）相比，其内存开销低 12%。该公司还强调，ML Drift 已通过 Chrome、YouTube Shorts、Photos 和 Meet 等应用在数百万台设备上运行。

在 YouTube Shorts 中，ML Drift 使 Android 和 iOS 上基于分割的特效平均帧延迟最多降低 40%。在 Google Photos 中，它将计算摄影和分割流水线加速最多两秒。Adobe Lightroom 和 Photoshop 从部分 AI 驱动的照片编辑功能（如选择主体、选择天空和自适应肖像）最高 30% 的性能提升中受益。最后，Snapchat 在 Android 上由 ML 驱动的面部和风格生成特效的模型延迟降低了 30%，从而能够使用更大的扩散模型。

AI/ML 工程师 Pavithra Agraharam 在 LinkedIn 上评论该公告时表示：

无需布局技巧的原生 5D 张量支持是在边缘 GPU 上运行空间和 Transformer 模型的巨大胜利。将逻辑张量与物理内存分配解耦，应能显著减少实时推理期间的内存碎片。

Tensorust 的创始人兼唯一开发者 Matthew Charles Vladislav Busel 质疑“它在近期 iPhone 上与 Core ML 在延迟和内存方面的比较，因为这是我们大多数人目前发布时所依据的基准”，不过 Google 尚未提供任何对比数据。

随着 ML Drift 的发布，Google 将停止更新传统的 TensorFlow Lite GPU delegate。为了与现有模型保持向后兼容，该公司建议迁移到 LiteRT ML Drift GPU 加速器。应用开发者可以使用 LiteRT ML Drift 加速器，而 GPU 和系统工程师可以访问其更低层的 C++ API，以构建自定义 GpuModel 计算图并编写自定义着色器。