端侧跑AI模型，最头疼的往往不是模型本身，而是GPU那堆五花八门的底层API。OpenGL ES、OpenCL、Metal、WebGPU，每个后端都得单独伺候。谷歌AI Edge团队这次把ML Drift开源了，Apache 2.0许可，定位很明确：一个高性能、跨平台的端侧GPU计算引擎，专门为端侧AI/ML推理而生。

它不只是LiteRT内部的核心GPU加速引擎，也能作为独立库给自定义图形和推理运行时用。换句话说，谷歌想给端侧GPU推理铺一层统一的地基。

打开网易新闻 查看精彩图片

为什么端侧GPU这么难搞

数据中心里的推理，模型跑在同质、可预测的加速器集群上，硬件环境相对可控。端侧完全是另一回事：GPU架构、驱动版本、底层API千差万别，开发者事先根本不知道应用会跑在什么硬件上。

老的TFLite GPU delegate确实打下了GPU加速的基础，但生态变了。现在的端侧负载从实时计算机视觉、音频、深度处理，一路延伸到高参数量的生成式AI。这些架构把消费级芯片逼到极限，算力和内存瓶颈是旧运行时没设计去解决的。

谷歌的思路是建一个统一框架：传统ML要稳，最前沿的生成式AI要快。

四项架构改动

相比老的TFLite GPU delegate，ML Drift做了几处结构性升级：

  • 张量虚拟化统一着色器：过去要为OpenGL、OpenCL、Metal分别硬编码逻辑张量到物理GPU对象（纹理、缓冲区）的映射。ML Drift把张量的逻辑表示和物理分配解耦，由动态着色器模板在编译器初始化阶段解析坐标，省掉了维护多套后端着色器代码的麻烦，运行时开销极小，跨平台可移植性还在。
  • 可扩展自定义算子框架：提供直接注册API和底层着色语言访问。还带了一份agentic SKILL.md指南，让编码代理能在几分钟内编写、注册、验证自定义着色器。
  • 5D张量支持：TFLite GPU delegate在结构上硬编码为4D张量，复杂模型要5D就得用布局hack。ML Drift在LiteRT的GPU加速器里原生支持5D，3D卷积网络、时空模型（如YOLO 11n、MobileViT v2、Swin Transformer v2）可以直接在端侧GPU上跑。
  • 经典模型性能升级：作为旧GPU后端的直接继任者，已有经典负载能拿到即时性能提升，迁移设计得比较直接，旧负载能保持或超过原有基准。

端侧LLM的分阶段优化

自回归LLM推理时有两段完全不同的计算负载：计算密集的KV cache预填充阶段，和受内存带宽限制的逐token解码阶段。ML Drift会根据当前执行阶段动态切换内核和布局配置。

解码阶段，它用了一套自定义的、卷积对齐的KV cache布局，并在内核内做激进的激活量化，绕开冗余的内存往返。这套针对性优化，是冲着端侧LLM最吃力的环节去的。

桌面端预览

ML Drift的WebGPU后端最初是为浏览器内加速设计的，但Chromium的WebGPU实现Dawn让它能把同一套代码库原生编译到浏览器之外。这让ML Drift能跑在Windows和Linux上，借WebGPU的现代原生硬件抽象绕开DirectX和Vulkan的碎片化，同时和macOS上已有的原生Metal后端互补。

谷歌强调，主要精力仍在移动和边缘设备这类资源最紧张的场景。但本地开发需要跨开发机的灵活性，所以放出了ML Drift在工作站硬件上跑Gemma模型的早期快照，展示统一运行时如何跨设备扩展。