Google open-sources ML Drift, a GPU engine for edge AI that unifies OpenCL, Metal, WebGPU and OpenGL ES

Google's AI Edge team has open-sourced ML Drift under the Apache 2.0 license: a cross-platform GPU compute engine designed specifically for on-device AI/ML inference. Its job is to abstract away the complexity of hardware and of low-level GPU APIs across different operating systems and browsers — OpenGL ES, OpenCL, Metal and WebGPU — so that developers can build real-time AI experiences, from advanced video effects to generative AI, without having to rewrite shaders for each backend. ML Drift becomes the GPU acceleration engine inside LiteRT (Google's lightweight runtime, the successor to TensorFlow Lite), but it is also distributed as a standalone library for those who want to build their own graphics or inference engines.

The motivation behind the project is the enormous fragmentation of edge hardware compared with the homogeneity of data centers: developers do not know in advance which GPU, with which driver and under which low-level API, their application will run on. The old 'TFLite GPU delegate' laid the groundwork for on-device GPU acceleration, but the ecosystem has moved on: classic computer vision, audio and depth models now coexist with many-parameter generative architectures that push consumer silicon to its compute and memory limits. Google explains that it needed a unified engineering foundation capable of delivering stability for classic models and, at the same time, maximum performance for the most demanding generative AI.

To achieve that, ML Drift introduces several structural changes compared with the old TFLite GPU delegate. The most significant is 'tensor virtualization': previously, the mapping between a tensor's logical representation and its physical allocation (textures, buffers) had to be hand-coded separately for OpenGL, OpenCL and Metal, which meant maintaining different shader codebases for each backend. Now, dynamic shader templates resolve and translate those coordinates during the compiler's initialization phase, which removes the need to maintain platform-specific code and adds minimal runtime overhead, preserving model portability across platforms.

Another new feature is a modernized custom ops framework, with direct registration APIs and low-level access to the shading language. To make this easier, ML Drift includes an 'agentic SKILL.md' guide designed so that coding agents can quickly write, register and verify custom shaders, giving developers granular control to integrate their own model blocks directly into the execution graph.

Native support for 5-dimensional tensors has also been added. The old TFLite GPU delegate was structurally limited to 4D tensors, which forced the use of layout hacks for complex models requiring 5D. With built-in 5D support, LiteRT can run workloads such as 3D convolutional networks for volumetric/spatial AI and spatiotemporal models directly on edge GPUs, with YOLO 11n, MobileViT v2 and Swin Transformer v2 cited as examples.

For classic models already running on the old backend, ML Drift is presented as a strict performance upgrade: by decoupling the runtime from hardware-specific constraints and improving kernel execution efficiency, those existing workloads get immediate performance gains, with a migration designed to be straightforward and benchmarks that match or beat the previous ones.

A particularly technical section covers 'stage-aware' optimization for LLMs at the edge. Autoregressive language models have two very different compute loads during inference: the KV cache prefill phase, which is compute-bound, and the decode phase, which generates tokens one by one and is memory-bandwidth-bound. ML Drift dynamically switches kernels and data layout configurations depending on the active phase; during decoding it uses a custom, convolution-aligned KV cache layout and applies aggressive activation quantization inside the kernel itself to avoid redundant round trips to memory. The article includes a reference chart with the test configuration (512-token prefill, 128-token decode, 640 context, 4-bit weights with block size 32, fp16 KV cache, argmax sampling and per-token synchronization), noting that the Gemma4 E2B model performs better in frameworks that support cross KV sharing.

Beyond mobile, Google is bringing ML Drift to the desktop. Although the WebGPU backend was originally designed to accelerate inference inside the browser, the fact that Dawn (Chromium's WebGPU implementation) allows exactly that same code to be compiled outside the browser makes it possible to run ML Drift natively on Windows and Linux, leveraging WebGPU's modern hardware abstraction to sidestep the fragmentation between DirectX and Vulkan, and complementing the native Metal backend already available on macOS. Google frames this against the rise of local coding agents and workstation workflows, though it insists its priority remains lightweight inference on mobile and edge, where resource constraints are tighter. For desktop it is publishing an early snapshot with Gemma models (8192 prefill, 1024 decode, 9216 context, the same quantization and sampling conditions as before) and stresses that peak performance is not everything, because the memory footprint determines which workloads can run at the same time: in its Gemma benchmarks, ML Drift achieves up to 12% less memory overhead than other frameworks, leaving more memory free for other local tasks.

On actual deployment, Google says ML Drift is already running in production on millions of devices every day, powering key features in Chrome, YouTube Shorts, Google Photos and Meet, as well as acting as the runtime for flagship Google AI Edge apps such as Gallery, Foresight and Eloquent. Several specific cases are cited: YouTube Shorts migrated its segmentation-based effects to ML Drift and achieved up to a 40% reduction in average frame latency on both Android and iOS, making it possible to apply high-fidelity visual effects in real time without dropping camera frames. Google Photos built ML Drift into its computational photography and segmentation pipelines, achieving up to 2 seconds of improvement over the old GPU delegate, illustrated with the Unblur feature. Google Chrome, for its part, integrated ML Drift into its AI runtime infrastructure to hardware-accelerate the Gemini Nano models that underpin the browser's built-in AI APIs (Prompt, Summarize and Writer), with an example of Cafe24 using the Prompt API for tag generation.

Outside Google's own ecosystem, several development partners have also adopted the technology. Adobe updated key AI features in Lightroom and Photoshop — such as Select Subject, Select Sky and Adaptive Portrait — with ML Drift, making professional photo editing on mobile devices up to 30% faster. Snap integrated LiteRT to power GPU inference on Android devices, and ML Drift's GPU accelerator delivered 30% model latency improvements in the face generation and style effects of Snapchat Lenses; Google adds that ML Drift not only improves existing models but also makes it possible to run larger diffusion models.

Developing ML Drift has also involved close collaboration with silicon makers. With Arm, compute kernels and texture cache locality were co-optimized for Mali and Immortalis GPUs, aiming for energy-efficient inference on mobile devices. With Intel, ML Drift takes advantage of native WebGPU compute and the Xe Matrix Extensions (XMX) to accelerate LLM and vision workloads on Intel Core Ultra processors with Xe3 graphics. With Qualcomm, OpenCL kernels were optimized to maximize use of the arithmetic logic units (ALUs) and memory bandwidth on Adreno GPUs.

On the practical side, Google announces that the old TFLite GPU delegate will stop receiving new features and encourages all developers to migrate to LiteRT's ML Drift GPU accelerator, which offers full compatibility with existing models and immediately unlocks performance gains on Android, iOS, web and desktop. For those using 'unbundled' runtimes on Android (to minimize app binary size), ML Drift acceleration is already available in the standalone LiteRT packages, and will come soon to LiteRT in Google Play Services.

To get started, Google outlines two paths: app developers can use LiteRT's ML Drift accelerator to get model partitioning, tensor virtualization and optimized LLM acceleration out of the box; GPU and systems engineers can build directly with ML Drift's standalone C++ APIs from its GitHub repository, creating custom 'GpuModel' compute graphs and writing their own shaders, with quick-start guides for both OpenCL on Android and WebGPU via Dawn. The entire project is open source under the Apache 2.0 license, with a repository at github.com/google-ai-edge/ml-drift, a GitHub issue tracker for reporting problems or requesting features, a Google AI Edge community forum for discussion and questions, and a contribution guide for those who want to add new UCL operators or hardware-specific optimizations.

Overall, the announcement matters because it consolidates into a single engine — with one virtualized shader codebase — something that until now required maintaining separate implementations per graphics API and per platform, while adding capabilities the edge generative AI ecosystem had been calling for: support for high-dimensional tensors, specific handling of the two LLM inference phases, and a path to the desktop based on native WebGPU. The fact that it is already deployed across large-scale Google products and in third-party apps such as Adobe's or Snap's, on top of joint optimizations with Arm, Intel and Qualcomm, suggests this is not an experimental project but a piece of infrastructure already proven in production that is now being opened up to the entire on-device AI developer community.