Google open-sources ML Drift for on-device GPU inference

The Google AI Edge Team has released ML Drift as open source under the Apache 2.0 license. ML Drift is a high-performance cross-platform on-device GPU compute engine built specifically for AI/ML inference. By abstracting hardware and low-level API complexities across OpenGL ES, OpenCL, Metal, and WebGPU, it enables developers to build real-time interactive ML experiences from advanced video effects to generative AI across multiple platforms. ML Drift serves as the core GPU acceleration engine within LiteRT and is also available as a standalone library for custom graphics and inference runtimes.

Edge deployment faces hardware diversity

Unlike datacenter inference where models run on homogeneous predictable accelerator clusters, deploying GPU-accelerated AI to edge devices means confronting significant hardware diversity. Developers face a wide range of GPU architectures, driver versions, and low-level APIs without knowing which specific hardware the application will run on. The TensorFlow Lite GPU delegate established foundational GPU acceleration, but the ecosystem has evolved: current on-device workloads span real-time computer vision, audio, and depth processing models to high-parameter generative AI. These architectures push consumer silicon to its limits, creating compute and memory bottlenecks that legacy runtimes were not designed to solve.

Tensor virtualization unifies shader code across backends

ML Drift introduces core architectural changes over the legacy TFLite GPU delegate. Historically, maintaining optimized shaders required hardcoding logical tensor mappings directly to physical GPU objects such as textures and buffers separately across OpenGL, OpenCL, and Metal backends. ML Drift introduces tensor virtualization, a paradigm that decouples a tensor's logical representation from its physical allocation on the GPU. Dynamic shader templates resolve and translate coordinates during the compiler's initialization phase. This unified shader model eliminates the need to maintain disparate backend-specific shader codebases while adding minimal runtime overhead and preserving cross-platform model portability.

Custom op framework provides direct shading language access

ML Drift adopts a modernized custom op framework with direct registration APIs and low-level shading language access. A guide enables coding agents to author, register, and verify performant custom shaders in minutes. This gives developers granular control to integrate specialized model blocks directly into the execution graph, reducing the barrier to deploying proprietary architectures.

5D tensor support enables volumetric models

The TFLite GPU delegate was structurally hardcoded to 4D tensors, forcing developers to use layout hacks for complex models requiring 5D tensors. ML Drift enables 5D tensor support out-of-the-box. LiteRT can now execute complex workloads like 3D convolutional networks for volumetric and spatial AI and spatiotemporal models such as YOLO 11n, MobileViT v2, and Swin Transformer v2 directly on edge GPUs.

Classical models get immediate performance upgrades

As the direct successor to the legacy GPU backend, ML Drift provides a strict performance upgrade for existing classical workloads. By decoupling the runtime from hardware-specific constraints and improving kernel execution efficiency, models already running on TFLite GPU gain immediate performance improvements. Migration is straightforward with legacy workloads maintaining or exceeding previous benchmarks.

Stage-aware optimizations target autoregressive LLMs

Autoregressive LLMs have two different computational workloads during inference: the compute-bound KV cache prefill stage and the memory-bandwidth-bound decode stage generating tokens one by one. ML Drift dynamically switches kernels and layout configurations depending on the active execution phase. During decoding, it uses a custom convolution-aligned KV cache layout and applies aggressive in-kernel activation quantization to bypass redundant memory roundtrips.

Desktop previews extend WebGPU beyond the browser

Highly responsive desktop inference is becoming vital for local coding agents and workstation workflows. ML Drift's WebGPU backend was originally engineered for in-browser acceleration. Dawn, Chromium's WebGPU implementation, allowed the same codebase to be compiled natively outside the browser. This unified API model enables ML Drift to run on Windows and Linux, leveraging WebGPU's modern native hardware abstraction to bypass DirectX and Vulkan fragmentation while complementing the existing high-performance native Metal backend on macOS.

Benchmarks show lower memory overhead

Across edge and workstation environments, peak throughput only tells half the story: memory footprint dictates what workloads can run concurrently. In Gemma benchmarks, ML Drift achieves up to 12% lower memory overhead than other frameworks, freeing memory for other local tasks.

Production deployments across Google products

ML Drift runs in production across millions of devices daily, powering key features across Google's ecosystem including Chrome, YouTube Shorts, Photos, and Meet. It acts as the runtime for Google AI Edge apps like Gallery, Foresight, and Eloquent.

  • YouTube Shorts migrated segmentation-based effects to ML Drift, delivering up to a 40% reduction in average frame latency across Android and iOS, enabling creators to apply high-fidelity visual effects in real time without dropping camera frames.
  • Google Photos integrated ML Drift across computational photography and segmentation pipelines, delivering up to 2 sec speedup compared to the legacy GPU delegate.
  • Google Chrome integrated ML Drift into its AI runtime infrastructure to power native hardware-accelerated Gemini Nano models for Built-In AI APIs including Prompt, Summarize, and Writer APIs.

Partner adoptions bring desktop capabilities to mobile

Leading developer partners have adopted ML Drift to bring desktop-grade capabilities to mobile users.

  • Adobe Lightroom and Adobe Photoshop upgraded key AI features including Select Subject, Select Sky, and Adaptive Portrait with ML Drift, achieving up to 30% faster on-device performance for pro-quality photo editing on mobile.
  • Snap integrated LiteRT to power on-device inference on GPU across Android devices. ML Drift GPU accelerator brought 30% model latency improvements for ML-powered face and style generator effects in Snapchat lenses. ML Drift expands capabilities to run even large diffusion models.

Silicon partnerships align software to hardware

Maximizing edge AI efficiency requires deep alignment between software architectures and physical silicon. Throughout ML Drift's development, Google worked closely with silicon and IP partners to enable optimizations that maximize hardware utilization across modern GPUs. This collaboration ensures state-of-the-art model optimizations translate into battery-conscious high-frame-rate performance on user devices.

Migration path from legacy TFLite GPU delegate

With the launch of ML Drift, 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 while immediately unlocking modern performance gains across Android, iOS, Web, and desktop. 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.

Two paths to get started

Whether integrating on-device models into an app or authoring custom GPU shaders, ML Drift provides two straightforward paths. The project welcomes community feedback and contributions. The release acknowledges contributions from Google team members and partners.