Google 开源 ML Drift 端侧 GPU 推理引擎，接替 TFLite GPU delegate

2026年10月8日

Google AI Edge 团队很高兴地宣布，在 Apache 2.0 许可下开源发布 ML Drift，这是我们专为端侧 AI/ML 推理打造的高性能、跨平台、端侧 GPU 计算引擎。通过抽象 OpenGL ES、OpenCL、Metal 和 WebGPU 等端侧 GPU 的硬件与底层 API 复杂性，ML Drift 赋能开发者跨多个平台构建从高级视频特效到生成式 AI 的实时交互式 ML 体验。作为 LiteRT 中的核心 GPU 加速引擎，ML Drift 也可作为独立库用于自定义图形和推理运行时，提供统一基础，在任何地方都能实现峰值性能。

ML Drift 跨平台的 GPU API 支持矩阵

我们为何构建 ML Drift

与数据中心推理中模型运行在同质、可预测的加速器集群上不同，将 GPU 加速的 AI 部署到边缘设备时，其显著特点是硬件高度多样化。开发者面临着各种 GPU 架构、驱动版本和底层 API，无法预先知道应用将在哪种具体硬件上运行。

TensorFlow Lite GPU 委托奠定了 GPU 加速的基础，但此后生态系统已经演进：当前的端侧工作负载如今涵盖广泛范围，从实时计算机视觉、音频和深度处理模型到高参数生成式 AI。这些架构将消费级芯片推向极限，造成传统运行时并非为应对而设计的计算和内存瓶颈。

这一转变要求建立一个通用工程基础，确保在经典模型和下一代模型上兼具可移植性与性能。借鉴 "Speed is all you need" 中讨论的优化原则，我们构建了一个单一、统一的框架，为传统 ML 提供稳定性，为最先进的生成式 AI 提供峰值性能，确保你的模型运行在我们最高效、最具前瞻性的 GPU 计算引擎上。

为实现峰值性能与覆盖范围的结构性现代化

为了在经典架构和生成式架构上同时实现峰值性能和扩展的模型覆盖范围，ML Drift 在传统 TFLite GPU 委托的基础上引入了核心架构变更和结构性升级。

通过张量虚拟化实现统一着色器： 历史上，为 TFLite GPU 委托维护优化的着色器，需要在 OpenGL、OpenCL 和 Metal 后端中分别将逻辑张量映射硬编码到物理 GPU 对象（如纹理和缓冲区）。ML Drift 引入了张量虚拟化，这是一种核心架构范式，将张量的逻辑表示与其在 GPU 上的物理分配解耦。通过让动态着色器模板在编译器初始化阶段解析和转换坐标，这种统一着色器模型消除了维护不同后端特定着色器代码库的需要，在增加极少运行时开销的同时保持了跨平台模型可移植性。

可扩展自定义算子框架： ML Drift 采用现代化的自定义算子框架，提供直接注册 API 并支持底层着色语言访问。为加速这一过程，ML Drift 内置了智能体 SKILL.md 指南，使编码智能体能够在数分钟内编写、注册并验证高性能自定义着色器。这为开发者提供了细粒度控制能力，可将专用模型模块直接集成到执行图中，大幅降低了部署专有架构的门槛。

5D 张量支持： TFLite GPU 委托在结构上硬编码为 4D 张量，迫使开发者在需要 5D 张量的复杂模型上使用布局技巧。ML Drift 在新的 LiteRT ML Drift GPU 加速器中实现了这一期盼已久的功能，开箱即用地支持 5D 张量。借此，LiteRT 可直接在边缘 GPU 上执行复杂工作负载，如用于体数据/空间 AI 的 3D 卷积网络以及时空模型（例如 YOLO 11n、MobileViT v2 和 Swin Transformer v2）。

经典模型的性能升级： 作为我们旧版 GPU 后端的直接继任者，ML Drift 为现有经典工作负载提供了严格的性能升级。通过将运行时与硬件特定约束解耦并提升内核执行效率，此前已在 TFLite GPU 上运行的模型可获得即时的性能提升。迁移设计力求简单直接，旧有工作负载可保持或超越此前的基准表现。

面向边缘 LLM 的阶段感知优化： 自回归 LLM 在推理期间有两种不同的计算负载：受计算限制的 KV 缓存预填充阶段，以及受内存带宽限制、逐个生成 token 的解码阶段。ML Drift 会根据当前执行阶段动态切换内核和布局配置。在解码阶段，它采用自定义的、卷积对齐的 KV 缓存布局，并应用激进的内核内激活量化，以绕过冗余的内存往返。

预填充 512，解码 128，上下文长度 640，4 位权重（块大小 32），fp16 KV 缓存，argmax 采样，每个 token 同步。 图中展示了预填充（实心柱；左轴）和解码（空心柱；右轴）吞吐量。 Gemma4 E2B 在支持跨 KV 共享的框架上表现更好。

跨平台扩展（桌面预览）： 在当今新兴的智能体编码时代，高响应性的桌面推理对于驱动本地编码智能体和工作站工作流正变得至关重要。虽然我们最初为浏览器内加速而设计 ML Drift 的 WebGPU 后端，但 Dawn（Chromium 的 WebGPU 实现）使我们能够将完全相同的代码库原生编译到浏览器之外。这种统一的 API 模型使 ML Drift 能够在 Windows 和 Linux 上运行，利用 WebGPU 现代的原生硬件抽象绕过 DirectX 和 Vulkan 的碎片化问题，同时补充我们在 macOS 上现有的高性能原生 Metal 后端。

我们的首要重点仍然是在资源约束最紧张的地方——移动和边缘设备——提供轻量级、生产级推理。然而，本地开发需要跨开发者机器的灵活性。以下是 ML Drift 在工作站硬件上运行 Gemma 模型的早期快照，展示了我们的统一运行时如何跨环境扩展。

桌面端 LLM 性能。Prefill 8192，decode 1024，上下文长度 9216， 4 位权重（块大小 32），fp16 KV 缓存，argmax 采样，每个 token 同步。

在边缘设备和工作站环境中，峰值吞吐量只能说明一半问题：内存占用决定了哪些工作负载能够真正并发运行。在我们的 Gemma 基准测试中，ML Drift 的内存开销比其他框架低多达 12%，从而为其他本地任务释放内存。

经过大规模实战检验：从 Google 旗舰产品到行业领导者

ML Drift 已在数百万台设备上每日投入生产运行，为 Google 生态系统中的关键功能提供支持，包括 Chrome、YouTube Shorts、Photos、Meet 以及最近推出的 AI Edge Gallery。

YouTube Shorts 将其基于分割的特效迁移到 ML Drift。这在 Android 和 iOS 上使平均帧延迟降低多达 40%，让创作者能够实时应用高保真视觉效果，而不会掉帧。

抱歉，你的浏览器不支持此视频的播放

为了实现快速、无缝的设备端照片编辑，Google Photos 团队将 ML Drift 集成到其计算摄影和分割流水线中。与旧版 GPU delegate 相比，这带来了多达 2 秒的加速，让高级照片增强功能感觉瞬间完成。

抱歉，你的浏览器不支持此视频的播放

Google Photos Unblur 通过 ML Drift 加速

Google Chrome 将 ML Drift 集成到其 AI 运行时基础设施中，为 Built-In AI API（包括 Prompt、Summarize 和 Writer API）提供原生硬件加速的 Gemini Nano 模型支持。在桌面平台上，ML Drift 在日益壮大的 Web 应用生态系统中实现高效设备端 AI，从电商评论摘要到企业客户工作流：

抱歉，你的浏览器不支持此视频的播放

Cafe24 使用 Prompt API 进行标签生成，并通过 ML Drift 加速。

除了 Google 产品之外，领先的开发者合作伙伴也已采用 ML Drift，为移动用户带来桌面级能力：

Adobe Lightroom 和 Adobe Photoshop 使用 ML Drift 升级了关键 AI 功能，包括 Select Subject、Select Sky 和 Adaptive Portrait，在移动端实现专业级照片编辑的设备端性能提升多达 30%。

抱歉，你的浏览器不支持此视频的播放

Adobe Lightroom 通过移动端 LiteRT ML Drift GPU 加速器在设备端运行

Snap 集成 LiteRT，在 Android 设备上通过 GPU 实现设备端推理。ML Drift GPU 加速器为 Snapchat 镜头中由 ML 驱动的人脸和风格生成器特效带来了 30% 的模型延迟改进。ML Drift 不仅基于现有模型改善用户体验，还扩展了运行大型扩散模型的能力。

抱歉，你的浏览器不支持此视频的播放

与领先硬件合作伙伴共同优化

要最大化边缘 AI 效率，需要软件架构与物理芯片之间的深度协同。在 ML Drift 的开发过程中，Google 与芯片和 IP 合作伙伴密切合作，实现优化，从而最大化现代 GPU 的硬件利用率：

• Arm： 为 Mali 和 Immortalis GPU 协同优化计算内核与纹理缓存局部性，确保在移动设备上实现高吞吐、高能效的推理。
• Intel： ML Drift 利用原生 WebGPU 计算和 Xe Matrix Extensions (XMX)，在搭载 Xe3 显卡的 Intel® Core™ Ultra 处理器上加速客户端 LLM 和视觉工作负载。
• Qualcomm Technologies, Inc.： ML Drift 优化了 OpenCL 内核，以实现 ALU 和内存带宽利用率的峰值，从而在高通® Adreno™ GPU 上实现高效的边缘 AI 加速。

通过与芯片领导者的紧密合作，ML Drift 确保最先进的模型优化能够转化为用户设备上兼顾续航的高帧率性能。

下一步：弃用与迁移

随着 ML Drift 的发布，旧版 TFLite GPU 委托将不再获得新功能更新。我们鼓励所有开发者迁移到 LiteRT ML Drift GPU 加速器，它为现有模型提供完全向后兼容性，同时立即在 Android、iOS、Web 和桌面端解锁现代性能提升。对于使用非捆绑运行时以最小化应用二进制大小的 Android 开发者，ML Drift 加速现已在独立 LiteRT 包中提供，并即将登陆 Google Play Services 中的 LiteRT。

快速上手

无论您是将设备端模型集成到应用中，还是编写自定义 GPU 着色器，ML Drift 都提供两条直接路径：

• 面向应用开发者： 使用 LiteRT ML Drift 加速器 即可获得开箱即用的模型分区、张量虚拟化和优化的 LLM 加速。
• 面向 GPU 与系统工程师： 通过 ML Drift GitHub 仓库 直接使用 ML Drift 的独立 C++ API 构建自定义 `GpuModel` 计算图并编写自定义着色器（可从我们的 Android 上的 OpenCL 和 通过 Dawn 使用 WebGPU 快速入门开始）。

社区反馈与贡献

ML Drift 在 Apache 2.0 许可证下完全开源。我们欢迎贡献、RFC 和社区反馈：

• 源代码：github.com/google-ai-edge/ml-drift
• 问题追踪与功能请求：在 GitHub 上提交 Issue
• 讨论与问答：Google AI Edge 社区论坛
• 贡献：请查看我们的 贡献指南，以添加新的 UCL 算子或硬件特定优化。

致谢：

^{Andrei Kulik, Ankit Goyal, Bob Liu, Byungchul Kim, Champ Yen, Changming Sun, Chris McClanahan, Clark Duvall, Cormac Brick, Daniel Ho, David Neto, Ekaterina Ignasheva, Fengwu Yao, Frank Ban, Gerardo Carranza, Grant Jensen, Jacob Dormuth, Jagadeesh Pakaravoor, James Price, Jeremy Kemp, Jingjiang Li, Jiuqiang Tang, Kelvin Ma, Ken Russell, Kenji Baheux, Lin Chen, Loko Kung, Lu Wang, Marissa Ikonomidis, Matt Kreileder, Matthias Grundmann, Mattias Simonsson, Paweł Andruszkiewicz, Raman Sarokin, Romaric Jodin, Sachin Kotwani, Suleman Shahid, Terry Heo, Tyler Mullen, Vaibhav Jain, Weiyi Wang, Yu-Hui Chen}

^{感谢合作伙伴 Albin Bernhardsson, Aleksei Lebedev, Anuradha Oberoi, Arselan Alvi, Bala Gattu, Enxing Xiong, Erin Di Leva, Hannah Stabingas, Jianhui Dai, Jiawei Shao, Jie Chen, Joseph Hsieh, Komal Desai, M N Suhas, Magzhan Gabidolla, My Linh Van, Ningxin Hu, Padmassri Chandrashekar, Rong Wang, Varun Chari, Viktor Zhou, Zhaoliang Ma}

上一篇

下一页