谷歌开源ML Drift：跨平台GPU推理引擎，内存开销低12%

谷歌近日开源了ML Drift，这是一款跨平台的GPU计算引擎，面向端侧AI与机器学习推理场景。该引擎采用Apache 2.0许可，把不同GPU硬件和底层API的差异做了封装，可支持OpenGL ES、OpenCL、Metal和WebGPU。

ML Drift同时也是驱动LiteRT的核心GPU加速引擎。开源之后，开发者能借助它为各类应用搭建自定义的图形渲染与推理运行时，覆盖传统计算机视觉、音频与深度处理，以及生成式AI等方向。

谷歌希望ML Drift能帮开发者绕开设备端应用开发中GPU硬件五花八门这一难题。现实中，GPU架构、驱动版本、底层API各不相同，而应用最终跑在哪种硬件上，开发者事先根本没法确定。

按照谷歌的说法，ML Drift要为可移植性和性能打下新基础，兼顾“传统 ML 的稳定性和前沿生成式 AI 的峰值性能”。

和旧的TFLite GPU委托相比，ML Drift带来了几处技术升级：用张量虚拟化实现统一着色器、支持可扩展的自定义算子、原生支持五维张量。此外，它针对Arm、Intel和高通技术公司的GPU做了硬件专属优化，以便把现代GPU的硬件资源用足。

桌面端推理方面，ML Drift借助Dawn（Chromium的WebGPU实现）把支持范围延伸到了Windows和Linux，macOS上则继续保留原生Metal后端。Dawn提供统一的API模型，可复用WebGPU的硬件抽象层，开发者不必直接对接DirectX、Vulkan这类平台专属图形API。

谷歌在开发者级工作站硬件上用Gemma模型对ML Drift做了基准测试，报告称其内存开销比llama.cpp、MLX LM等其他框架低12%。谷歌还提到，ML Drift已在Chrome、YouTube Shorts、Photos和Meet等应用中跑在数百万台设备上。

具体到应用效果：在YouTube Shorts上，Android和iOS平台基于分割算法的特效平均帧延迟最多下降40%；谷歌Photos的计算摄影与图像分割流水线最多提速2秒；Adobe Lightroom和Photoshop的部分AI照片编辑功能（如选择主体、选择天空和自适应人像）性能最高提升30%；Snapchat在Android上基于ML驱动的人脸和风格生成特效模型延迟降低了30%，因而得以采用更大的扩散模型。

AI/ML工程师Pavithra Agraharam在LinkedIn的相关动态下评论道：

对空间模型和 Transformer 模型在边缘 GPU 上跑这件事，原生支持 5D 张量、不必再靠布局 Hack，算是个大好消息。逻辑张量与物理内存分配一旦分开，在线推理时内存碎片化就能少很多。

Tensorust创始人兼开发者Matthew Charles Vladislav Busel则提出疑问：“它与 Core ML 在新 iPhone 上的延迟和内存表现，毕竟这是目前大多数人的交付目标平台”。不过谷歌并未给出任何对比数据。

伴随ML Drift上线，旧版TensorFlow Lite GPU委托将不再更新。谷歌建议迁移至LiteRT ML Drift GPU加速器，以保持对现有模型的向后兼容。应用开发者可以直接使用LiteRT ML Drift加速器，GPU和系统工程师则能通过其底层C++ API构建自定义GpuModel计算图，并编写自定义着色器。