ML Drift: Next-Gen GPU AI/ML Inference at the Edge

Introduction to ML Drift

Google AI Edge has introduced ML Drift, an open-source GPU compute engine designed for on-device AI/ML inference. It abstracts hardware complexities across multiple platforms, enabling developers to build real-time, interactive ML experiences. ML Drift is part of LiteRT and supports a wide range of GPU architectures, including OpenGL ES, OpenCL, Metal, and WebGPU. This framework is intended to address the challenges of deploying AI models on edge devices, where hardware diversity and resource constraints are significant.

Why ML Drift Was Built

Unlike datacenter inference, which operates on homogeneous, predictable accelerator clusters, edge devices present a complex landscape of varying GPU architectures, driver versions, and low-level APIs. The TensorFlow Lite GPU delegate provided foundational GPU acceleration, but the evolving ecosystem of on-device workloads—ranging from real-time computer vision to generative AI—has necessitated a more robust and flexible solution. ML Drift aims to provide a unified engineering foundation that ensures portability and performance across both classical and next-generation models.

Structural Modernization for Performance and Coverage

To achieve peak performance and expanded model coverage, ML Drift introduces core architectural changes over the legacy TFLite GPU delegate. One of the key innovations is tensor virtualization, which decouples a tensor's logical representation from its physical allocation on the GPU. This approach eliminates the need for maintaining separate shader codebases for different backends, reducing runtime overhead while preserving cross-platform model portability.

ML Drift also introduces an extensible custom op framework that provides direct registration APIs with low-level shading language access. This allows developers to integrate specialized model blocks directly into the execution graph, reducing the barrier to deploying proprietary architectures. Additionally, ML Drift supports 5D tensors, a feature that was previously limited to 4D tensors in the TFLite GPU delegate. This enables the execution of complex workloads such as 3D convolutional networks for volumetric AI and spatiotemporal models directly on edge GPUs.

Performance Upgrades and Edge LLMs

ML Drift provides a strict performance upgrade for existing classical workloads by decoupling the runtime from hardware-specific constraints and improving kernel execution efficiency. It also includes stage-aware optimizations for edge LLMs, dynamically switching kernels and layout configurations depending on the active execution phase. During the decoding stage, ML Drift utilizes a custom, convolution-aligned KV cache layout and applies aggressive in-kernel activation quantization to bypass redundant memory roundtrips.

Broadening Across Platforms

ML Drift's WebGPU backend was originally designed for in-browser acceleration, but with the help of Dawn (Chromium's WebGPU implementation), the same codebase can now run natively outside the browser on Windows and Linux. This unified API model enables ML Drift to run on multiple platforms, leveraging WebGPU's modern, native hardware abstraction to bypass DirectX and Vulkan fragmentation. On macOS, ML Drift continues to use its high-performance native Metal backend.

Industry Adoption and Real-World Impact

ML Drift is already in production across millions of devices, powering key features in Google's ecosystem, including Chrome, YouTube Shorts, Photos, and Meet. It also acts as the runtime for pioneering on-device capabilities in Google AI Edge apps like Gallery, Foresight, and Eloquent. Adobe Lightroom and Photoshop have integrated ML Drift to enhance on-device performance for AI features such as Select Subject and Adaptive Portrait. Snapchat has also adopted ML Drift to improve the performance of ML-powered face and style generator effects in lenses.

Collaboration with Hardware Partners

Google has worked closely with silicon and IP partners to optimize ML Drift for modern GPUs. Collaborations with Arm, Intel, and Qualcomm Technologies have enabled compute kernel optimizations and texture cache locality improvements for Mali and Immortalis GPUs, Xe Matrix Extensions (XMX) for Intel Core Ultra processors, and OpenCL kernel optimizations for Qualcomm Adreno GPUs. These partnerships ensure that ML Drift maximizes hardware utilization across a wide range of edge devices.

Source: developers.googleblog.com. Analysis based on publicly available information; no hands-on testing is claimed.