ML Drift统一多平台端侧GPU计算

Google Developers Blog#端侧AI#GPU#跨平台

Editor brief · 编辑速览

Google AI Edge 推出 ML Drift，为端侧 AI 推理提供统一 GPU 计算框架，抽象 OpenGL ES、OpenCL、Metal 和 WebGPU 等底层接口。框架面向移动端与桌面端的实时视频效果、生成式 AI 等应用。

文章思维导图

Knowledge map

拖拽缩放

Full translation

完整中文译文

Google AI Edge 团队很高兴地宣布，ML Drift 正式以 Apache 2.0 许可证开源。它是一款高性能、跨平台的设备端 GPU 计算引擎，专为设备端 AI/ML 推理打造。ML Drift 封装了设备端 GPU 在 OpenGL ES、OpenCL、Metal 和 WebGPU 上的硬件差异与底层 API 复杂性，让开发者能够在多个平台上构建实时、交互式的 ML 体验，从高级视频特效到生成式 AI 都能支持。ML Drift 是 LiteRT 的核心 GPU 加速引擎，同时也提供独立库，供自定义图形与推理运行时使用，为各个平台发挥最佳性能提供统一基础。

为什么要构建 ML Drift

在数据中心进行推理时，模型运行在硬件同构、行为可预测的加速器集群上。而将 GPU 加速的 AI 部署到边缘设备，面对的则是显著的硬件多样性。开发者需要应对各种 GPU 架构、驱动版本和底层 API，而且无法事先知道应用究竟会运行在哪种硬件上。

TensorFlow Lite GPU delegate 为 GPU 加速奠定了基础，但此后的生态已经发生变化：如今的设备端工作负载覆盖了广泛的场景，从实时计算机视觉、音频和深度处理模型，到参数量庞大的生成式 AI。这些架构将消费级芯片推向极限，造成了旧版运行时在设计之初并未考虑解决的计算与内存瓶颈。

这种变化要求我们建立一套通用的工程基础，确保传统模型和新一代模型都能兼顾可移植性与性能。借鉴《Speed is all you need》中讨论的优化原则，我们构建了一个统一框架，既为传统 ML 提供稳定性，又能让最先进的生成式 AI 发挥最佳性能，确保你的模型运行在我们最高效、面向未来的 GPU 计算引擎上。

通过架构现代化提升性能与模型覆盖范围

为了让传统架构和生成式架构都能获得最佳性能，并扩大模型覆盖范围，ML Drift 相比旧版 TFLite GPU delegate 引入了核心架构变更与结构性升级。

通过 Tensor Virtualization 统一 Shader： 过去，为 TFLite GPU delegate 维护优化后的 shader，需要分别在 OpenGL、OpenCL 和 Metal 后端中，将逻辑 tensor 到物理 GPU 对象（如纹理和 buffer）的映射直接硬编码进去。ML Drift 引入了 tensor virtualization 这一核心架构范式，将 tensor 的逻辑表示与它在 GPU 上的物理内存分配解耦。动态 shader 模板会在编译器初始化阶段解析并转换坐标，因此，这种统一的 shader 模型消除了维护多套后端专用 shader 代码的需要，只引入极少的运行时开销，同时保留了模型的跨平台可移植性。

可扩展的 Custom Op 框架： ML Drift 采用了现代化的 custom op 框架，提供直接注册 API，并允许使用底层 shading language。为了加快开发，ML Drift 附带了一份面向 Agent 的 SKILL.md 指南，让 coding agents 能在几分钟内编写、注册并验证高性能 custom shader。开发者由此能够精细控制，将专用模型模块直接集成到执行图中，大幅降低部署自有架构的门槛。

支持 5D Tensor： TFLite GPU delegate 在结构上被硬编码为只支持 4D tensor，迫使开发者在处理需要 5D tensor 的复杂模型时，采用一些绕过限制的布局技巧。ML Drift 在新的 LiteRT ML Drift GPU accelerator 中补上了这项早该提供的功能，开箱即用地支持 5D tensor。这样一来，LiteRT 就能直接在边缘 GPU 上执行复杂工作负载，包括用于体积数据或空间 AI 的 3D 卷积网络，以及时空模型，例如 YOLO 11n、MobileViT v2 和 Swin Transformer v2。

传统模型的性能升级： 作为旧版 GPU 后端的直接继任者，ML Drift 为现有传统工作负载带来了实实在在的性能提升。通过将运行时与硬件特定约束解耦，并提高 kernel 执行效率，原本已经运行在 TFLite GPU 上的模型可以立即获得性能改善。迁移过程力求简单，旧有工作负载的表现能够达到或超过此前的 benchmark。

针对边缘 LLM 不同阶段的优化： 自回归 LLM 在推理过程中有两种不同的计算工作负载：受计算能力限制的 KV cache prefill 阶段，以及受内存带宽限制、逐个生成 token 的 decode 阶段。ML Drift 会根据当前执行阶段，动态切换 kernel 和布局配置。在 decode 阶段，它使用与卷积对齐的自定义 KV cache 布局，并在 kernel 内进行积极的 activation quantization，以避免多余的内存往返访问。

扩展平台支持（桌面端预览）： 在如今兴起的 Agent 编程时代，响应迅速的桌面端推理，对于本地 coding agents 和工作站工作流越来越重要。我们最初为浏览器内加速开发了 ML Drift 的 WebGPU 后端，而 Dawn（Chromium 的 WebGPU 实现）让我们能够将完全相同的代码库编译为浏览器之外的原生程序。这种统一的 API 模型让 ML Drift 可以运行在 Windows 和 Linux 上，利用 WebGPU 现代化的原生硬件抽象，绕过 DirectX 和 Vulkan 的碎片化问题，同时与我们在 macOS 上已有的高性能原生 Metal 后端形成互补。

我们的首要目标，仍然是在资源约束最严格的移动设备和边缘设备上，提供轻量、可用于生产环境的推理能力。不过，本地开发也需要灵活支持不同的开发者电脑。下面展示了 ML Drift 在工作站硬件上运行 Gemma 模型的早期结果，说明我们的统一运行时如何适配不同环境。

无论是边缘设备还是工作站，峰值吞吐量都只反映了一半情况：内存占用决定了哪些工作负载能够真正并发运行。在我们的 Gemma benchmark 中，ML Drift 的内存开销相比其他框架最多降低 12%，为其他本地任务腾出了内存。

经大规模生产验证：从 Google 旗舰产品到行业领军企业

ML Drift 已经投入生产，每天运行在数百万台设备上，为 Google 生态中的关键功能提供支持，包括 Chrome、YouTube Shorts、Photos 和 Meet。在前沿应用中，它也充当了开创性设备端能力的运行时，为 Gallery、Foresight 和 Eloquent 等 Google AI Edge 应用提供支持。

YouTube Shorts 已将基于分割的特效迁移到 ML Drift。迁移后，Android 和 iOS 上的平均帧延迟最多降低了 40%，让创作者能够实时应用高保真视觉特效，同时避免摄像头画面掉帧。