Google Open-Sources ML Drift, Its On-Device GPU Engine for AI Inference

One engine across four GPU APIs

According to Google, ML Drift hides the hardware and low-level API complexity of on-device GPUs across OpenGL ES, OpenCL, Metal and WebGPU. The aim is to let developers build real-time, interactive ML experiences, from advanced video effects to generative AI, on several platforms.

Google argues that edge deployment differs from datacenter inference, where models run on homogeneous, predictable accelerator clusters. On devices, developers face many GPU architectures, driver versions and low-level APIs, and do not know in advance which hardware an app will run on. The company says the TensorFlow Lite GPU delegate laid the foundation for GPU acceleration, but current workloads, from computer vision and audio to high-parameter generative AI, push consumer silicon to its limits and create bottlenecks that legacy runtimes were not designed to solve.

What changes compared with the TFLite GPU delegate

Google lists several architectural changes in ML Drift:

• Tensor virtualization: a tensor's logical representation is decoupled from its physical allocation on the GPU. Dynamic shader templates resolve coordinates when the compiler initializes, so teams no longer need separate, backend-specific shader codebases. Google says this adds minimal runtime overhead.
• Custom op framework: direct registration APIs with low-level shading language access. ML Drift includes an agentic SKILL.md guide that, according to Google, lets coding agents write, register and verify custom shaders in minutes.
• 5D tensor support: the TFLite GPU delegate is hardcoded to 4D tensors. The new LiteRT ML Drift GPU accelerator supports 5D tensors out of the box, so LiteRT can run 3D convolutional networks and spatiotemporal models such as YOLO 11n, MobileViT v2 and Swin Transformer v2 on edge GPUs.
• Faster classical models: Google calls ML Drift a strict performance upgrade for existing workloads and says migration is designed to be straightforward, with legacy workloads maintaining or exceeding previous benchmarks.
• Stage-aware LLM optimizations: autoregressive LLMs have a compute-bound prefill stage and a memory-bandwidth-bound decode stage. ML Drift switches kernels and layouts depending on the phase. For decoding it uses a convolution-aligned KV cache layout and in-kernel activation quantization to avoid redundant memory roundtrips.

Desktop previews and memory use

ML Drift's WebGPU backend was first built for in-browser acceleration. Because Dawn, Chromium's WebGPU implementation, can compile the same codebase natively, ML Drift can also run on Windows and Linux. It complements the existing native Metal backend on macOS. Google says its main focus remains mobile and edge devices, and describes the desktop results with Gemma models as an early snapshot.

In its Gemma benchmarks, Google says ML Drift achieves up to 12% lower memory overhead than other frameworks, which frees memory for other local tasks.

Already in production

Google says ML Drift already runs daily on millions of devices and powers features in Chrome, YouTube Shorts, Photos, Meet and the recently launched AI Edge Gallery.

• YouTube Shorts: segmentation-based effects moved to ML Drift, delivering up to a 40% reduction in average frame latency on Android and iOS, according to Google.
• Google Photos: integration across computational photography and segmentation pipelines brought up to a 2 second speedup compared with the legacy GPU delegate.
• Chrome: ML Drift is part of the AI runtime that powers hardware-accelerated Gemini Nano models for the Prompt, Summarize and Writer APIs. Google cites Cafe24 using the Prompt API for tag generation.
• Adobe: Lightroom and Photoshop upgraded Select Subject, Select Sky and Adaptive Portrait with ML Drift, achieving up to 30% faster on-device performance on mobile, per Google.

What it means for developers

With the code under Apache 2.0, developers can use ML Drift through LiteRT or as a standalone library for their own graphics and inference runtimes. For teams still on the TFLite GPU delegate, Google positions it as a successor that supports newer model types and generative workloads on one codebase across mobile, web and desktop. The performance figures come from Google and its partners, and the source does not include independent testing.