摘要

 Google AI Edge 团队宣布以 Apache 2.0 许可开源 ML Drift——一个高性能、跨平台的端侧 GPU 计算引擎，专为端侧 AI/ML 推理打造。它抽象掉 OpenGL ES、OpenCL、Metal、WebGPU 等底层 API 的硬件差异，既是 LiteRT 内部的 GPU 加速核心，也能作为独立库嵌入自研的图形或推理运行时。与旧的 TFLite GPU delegate 相比，ML Drift 的关键改动有四项：一是张量虚拟化，把张量的逻辑表示与 GPU 上的物理分配解耦，由动态着色器模板在编译器初始化阶段解析坐标，从此一套着色器模板覆盖所有后端，不再为 OpenGL、OpenCL、Metal 各维护一份；二是可扩展自定义算子框架，附带一份agentic SKILL.md 指南，让编码智能体在几分钟内编写、注册并验证自定义着色器；三是原生 5D 张量支持——旧 delegate 结构上硬编码为 4D，需要 5D 的模型只能靠布局 hack，现在 YOLO 11n、MobileViT v2、Swin Transformer v2 这类三维卷积与时空模型可以直接跑在端侧 GPU 上；四是针对端侧大模型的分阶段优化，自回归 LLM 推理在算力受限的 KV cache 预填充阶段与受内存带宽限制的逐 token 解码阶段之间动态切换内核与布局配置，解码阶段使用卷积对齐的 KV cache 布局并做激进的内核内激活量化。覆盖面上，桌面端为预览状态：WebGPU 后端经 Chromium 的 Dawn 实现把同一套代码编译到浏览器之外，从而原生跑在 Windows 与 Linux 上，macOS 走已成熟的 Metal 后端。落地成绩方面，Chrome 已用 ML Drift 支撑 Gemini Nano 的内置 AI API（Prompt、Summarize、Writer）；YouTube Shorts 的特效在 Android 与 iOS 上平均帧延迟降低最多 40%；Adobe 在移动端的 Select Subject、Select Sky、Adaptive Portrait 上取得最多 30% 的端侧提速；在 Google 自己的 Gemma 测试中，ML Drift 的内存开销比其他框架低最多 12%。需要注意的是，旧的 TFLite GPU delegate 自本次发布起不再获得新功能更新，Google 要求开发者迁移到 LiteRT 的 ML Drift GPU 加速器；Android 上独立版 LiteRT 包已可用，通过 Google Play Services 分发的版本稍后支持，未给出具体时间。

正文（译文）

ML Drift：面向边缘的下一代 GPU AI/ML 推理

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

我们为什么做 ML Drift

数据中心的推理跑在同质、可预测的加速器集群上，而把 GPU 加速的 AI 部署到边缘设备，面对的是巨大的硬件多样性：GPU 架构、驱动版本、底层 API 千差万别，开发者事先并不知道应用最终会跑在哪块硬件上。TFLite GPU delegate 打下了 GPU 加速的地基，但生态已经变了——今天的端侧负载覆盖了从实时计算机视觉、音频、深度处理模型，到高参数量生成式 AI 的广阔区间。这些架构把消费级芯片逼到极限，制造出旧运行时从未被设计去解决的算力与内存瓶颈。这种变化要求一个通用的工程地基，能在经典模型和下一代模型上都交付可移植性与性能。借鉴"Speed is all you need"里讨论过的优化原则，我们做出了一个统一的框架：对传统 ML 稳定可靠，对最先进的生成式 AI 提供峰值性能。

为峰值性能与覆盖面做的结构性改造

• 用张量虚拟化统一着色器：过去为 TFLite GPU delegate 维护优化着色器，需要在 OpenGL、OpenCL 和 Metal 三个后端上分别硬编码逻辑张量到物理 GPU 对象（纹理、缓冲区）的映射。ML Drift 引入张量虚拟化这一核心架构范式，把张量的逻辑表示与它在 GPU 上的物理分配解耦。动态着色器模板在编译器初始化阶段解析并转换坐标，这套统一着色器模型消除了维护多套后端着色器代码库的必要，只带来极小的运行时开销，同时保住了跨平台的模型可移植性。

• 可扩展的自定义算子框架：ML Drift 采用现代化的自定义算子框架，提供直接的注册 API 与底层着色语言访问能力。为了加速这一步，ML Drift 附带一份 agentic SKILL.md 指南，让编码智能体能在几分钟内编写、注册并验证高性能的自定义着色器。这让开发者能把专有的模型模块直接插入执行图，大幅降低部署自有架构的门槛。

• 5D 张量支持：TFLite GPU delegate 在结构上被硬编码为 4D 张量，需要 5D 的模型只能靠布局 hack。ML Drift 在新的 LiteRT ML Drift GPU 加速器里补齐了这个迟到已久的能力，开箱支持 5D 张量。有了它，LiteRT 可以直接在端侧 GPU 上执行三维卷积网络（体积/空间 AI）与时空模型，比如 YOLO 11n、MobileViT v2 和 Swin Transformer v2。

• 经典模型的性能升级：作为旧 GPU 后端的直接继任者，ML Drift 对已有的经典负载是一次严格的性能升级。通过把运行时从硬件特定约束中解耦出来、并提升内核执行效率，原本已经跑在 TFLite GPU 上的模型会立刻获得性能改善。迁移被设计得很直接，旧负载能维持或超过此前的基准。

面向端侧 LLM 的分阶段优化

自回归 LLM 在推理时有两段截然不同的计算负载：算力受限的 KV cache 预填充（prefill）阶段，和受内存带宽限制、逐 token 生成的解码（decode）阶段。ML Drift 会根据当前所处的执行阶段动态切换内核与布局配置。在解码阶段，它使用一套自定义的、卷积对齐的 KV cache 布局，并在内核内做激进的激活量化，绕开冗余的内存往返。

把覆盖面推到更多平台（桌面端预览）

在当下这个agentic coding 兴起的时代，高响应的桌面端推理正变得关键——它要支撑本地编码智能体和工作站工作流。我们最初是为浏览器内加速设计 ML Drift 的 WebGPU 后端，而 Dawn（Chromium 的 WebGPU 实现）让我们能把同一套代码库原生编译到浏览器之外。这让 ML Drift 可以跑在 Windows 和 Linux 上，借助 WebGPU 现代的原生硬件抽象绕开 DirectX 和 Vulkan 的碎片化问题，同时与 macOS 上已有的原生 Metal 后端形成互补。

生产环境里的实际表现

ML Drift 已经在 Google 的多个产品里跑在生产环境：Chrome 用它来运行 Gemini Nano，支撑内置的 Prompt、Summarize 和 Writer API；YouTube Shorts 把基于分割的视觉特效迁移到 ML Drift 后，在 Android 和 iOS 上的平均帧延迟降低最多 40%；Adobe 把它集成进 Lightroom 与 Photoshop 移动端的 AI 功能（Select Subject、Select Sky、Adaptive Portrait），端侧性能提升最多 30%。在 Google 自己的 Gemma 测试中，ML Drift 的内存开销比其他框架低最多 12%。

对开发者的实际含义

如果你的应用通过 TFLite GPU delegate 跑模型，请规划迁移到 LiteRT 的 ML Drift GPU 加速器——GPU 的开发从此只在那里继续。感谢 Apache 2.0 许可，维护自有推理运行时的团队可以只采用这一层 GPU 能力，而不必带走 LiteRT 的其余部分。上述提速数字来自 Google 及其合作伙伴，请在自己的目标设备上实测后再据此决策。