What does ML Drift change for on-device AI inference on your phone?

Why the old delegate hit a wall

The TensorFlow Lite GPU delegate gave developers a first taste of on-device acceleration. It was built for a narrower set of workloads: mostly classical computer vision models with predictable tensor shapes. Today the landscape looks different. Applications ship real-time video effects, audio processing, depth estimation, and generative models with billions of parameters. Those architectures push consumer silicon into compute and memory bottlenecks the legacy runtime was never designed to solve. Developers also face a fragmented hardware matrix: OpenGL ES on Android, Metal on iOS, WebGPU in browsers, and a mix of driver versions across every device. The delegate required separate, backend-specific shader codebases for each API. Maintaining them became a drag on performance and model coverage.

Tensor virtualization and the shader problem

ML Drift introduces tensor virtualization. In the old delegate a tensor’s logical shape was hardcoded to physical GPU objects — textures on OpenGL, buffers on Metal — and the mapping had to be rewritten for every backend. The new engine decouples the logical representation from the physical allocation. Dynamic shader templates resolve coordinates during compiler initialization. One unified shader model now covers OpenGL ES, OpenCL, Metal, and WebGPU. The runtime overhead of this indirection is described as minimal. For developers it means a model that runs on an Adreno GPU today will run on an Apple Silicon GPU tomorrow without shader rewrites.

The custom operator framework has also been modernized. Direct registration APIs expose low-level shading language access. An agentic SKILL.md guide lets coding agents author, register, and verify performant custom shaders in minutes. Proprietary model blocks can be integrated straight into the execution graph. The legacy delegate was structurally hardcoded to four-dimensional tensors. Workarounds for five-dimensional data — common in 3D convolutional networks, spatiotemporal models like YOLO 11n, MobileViT v2, and Swin Transformer v2 — are no longer necessary. Five-dimensional support ships out of the box.

Classical models get a free speed boost

Migration is designed to be straightforward. Existing workloads that ran on the TFLite GPU delegate maintain or exceed previous benchmarks. The engine decouples the runtime from hardware-specific constraints and improves kernel execution efficiency. No model changes are required to see the uplift. For classical computer vision and audio pipelines this is a drop-in replacement that immediately unlocks modern performance gains across Android, iOS, Web, and desktop.

Stage-aware optimization for large language models

Autoregressive LLMs have two distinct phases. The prefill stage is compute-bound; the decode stage is memory-bandwidth-bound, generating tokens one by one. ML Drift switches kernels and layout configurations depending on the active phase. During decoding it uses a custom convolution-aligned KV cache layout and applies aggressive in-kernel activation quantization to bypass redundant memory roundtrips. Benchmarks cite prefill 512, decode 128, context length 640, 4-bit weights with block size 32, fp16 KV cache, argmax sampling, sync every token. Gemma 4 E2B performs better on frameworks supporting cross KV sharing. The document does not say how many tokens per second this translates to on a specific phone.

Desktop previews via WebGPU and Dawn

The WebGPU backend was originally engineered for in-browser acceleration. Dawn, Chromium’s WebGPU implementation, allowed the same codebase to compile natively outside the browser. This unified API model lets ML Drift run on Windows and Linux, leveraging WebGPU’s modern hardware abstraction to bypass DirectX and Vulkan fragmentation, while complementing the native Metal backend on macOS. Primary focus remains mobile and edge devices where resource constraints are tightest. Desktop numbers show prefill 8192, decode 1024, context length 9216, 4-bit weights with block size 32, fp16 KV cache, argmax sampling, sync every token. In Gemma benchmarks ML Drift achieves up to 12 percent lower memory overhead than other frameworks, freeing memory for other local tasks.

Already running in production at scale

Millions of devices daily. YouTube Shorts migrated segmentation-based effects to ML Drift and saw up to 40 percent reduction in average frame latency across both Android and iOS. Google Photos integrated the engine across computational photography and segmentation pipelines, delivering up to 2 second speedup compared to the legacy GPU delegate. Chrome uses ML Drift to power Gemini Nano models for Built-In AI APIs including Prompt, Summarize, and Writer APIs. Cafe24 uses the Prompt API for tag generation. Adobe Lightroom and Photoshop upgraded Select Subject, Select Sky, and Adaptive Portrait, achieving up to 30 percent faster on-device performance for pro-quality photo editing on mobile. Snap integrated LiteRT to power on-device inference on GPU across Android devices; the ML Drift GPU accelerator brought 30 percent model latency improvements for face and style generator effects in Snapchat lenses. The engine also expands capabilities to run large diffusion models. The document does not disclose battery life impact or thermal behavior under sustained load.

Silicon partnerships that shape the kernels

Arm co-optimized compute kernels and texture cache locality for Mali and Immortalis GPUs. Intel enabled native WebGPU compute and Xe Matrix Extensions on Core Ultra processors with Xe3 graphics. Qualcomm optimized OpenCL kernels for peak ALU and memory bandwidth utilization on Adreno GPUs. These collaborations ensure that model optimizations translate into battery-conscious, high-frame-rate performance on user devices. The document does not list specific device models or driver versions validated.

Deprecation and the migration path

The legacy TFLite GPU delegate will no longer receive new feature updates. Developers are encouraged to transition to the LiteRT ML Drift GPU accelerator, which offers full backwards compatibility for existing models. For Android developers using unbundled runtimes to minimize app binary size, ML Drift acceleration is available today in standalone LiteRT packages and is coming soon to LiteRT in Google Play Services. Application developers use the LiteRT ML Drift accelerator for out-of-the-box model partitioning, tensor virtualization, and optimized LLM acceleration. GPU and system engineers build directly with ML Drift’s standalone C++ APIs via the GitHub repository to construct custom GpuModel compute graphs and write custom shaders. Quickstarts exist for OpenCL on Android and WebGPU via Dawn. The project is fully open source under Apache 2.0. Contributions, RFCs, and community feedback are welcome through the GitHub issue tracker, the Google AI Edge Community Forum, and the contributing guide.