ML Drift: Google's open-source on-device GPU engine, and when it beats the generalists

ML Drift is the on-device GPU inference engine Google made open source. It is the same technology that runs models inside Chrome, YouTube, Photos, and Meet, and you can now run it on your own GPU under the Apache 2.0 license.
Think of the difference between a Formula 1 engine and the engine in a family car. Ollama and llama.cpp are the family-car engine: versatile, reliable, able to run almost any model almost anywhere. ML Drift is the F1 engine: designed for one specific context — the GPU in a smartphone, a browser, or an edge device — where every kilobyte of memory and every millisecond of latency counts.
What ML Drift is
ML Drift is a cross-platform GPU compute engine, open source under Apache 2.0. Google's on-device AI team built it and positions it as the successor to the TensorFlow Lite GPU delegate: that delegate will no longer receive new features, and ML Drift is the official path for on-device GPU inference.
It integrates with LiteRT, Google's on-device AI toolkit (the evolution of TensorFlow Lite), and works on practically every platform with a GPU:
| Platform | Backend |
|---|---|
| Android | OpenCL |
| iOS, iPadOS, macOS | Metal |
| Browser | WebGPU (via Dawn) |
| Windows and Linux | via Dawn — early preview |
| Legacy GPUs | OpenGL ES 3.1+ |
Under the hood
Four technical choices explain the difference from generalist engines:
1. Tensor virtualization. ML Drift abstracts tensors away from the model: instead of compiling every operation for every GPU, the engine maps tensors onto a virtual view that backends execute natively. The result is less platform-specific code and models that behave more consistently across devices.
2. 5D tensors. Vision models (YOLO 11n, MobileViT v2, Swin v2 in the repo's examples) work with five-dimensional tensors; ML Drift handles them natively, where LLM-first engines often chop them into pieces.
3. Stage-aware KV cache. For generative models, prefill and decode have different memory profiles: prefill is heavy on incoming data, decode runs at low intensity. ML Drift manages the two stages with separate caching strategies instead of a single compromise strategy.
4. Multiple precisions. FP16, INT8, and INT4: the same model can run at three quantization levels depending on the device's VRAM and the quality you need.
ML Drift, Ollama, llama.cpp: who wins where
They are not competitors: they operate at different levels.
| ML Drift | Ollama | llama.cpp | |
|---|---|---|---|
| What it is | on-device GPU engine | generalist LLM runtime | generalist LLM engine |
| Where it excels | latency and memory on device | simplicity: one command for everything | hardware and model versatility |
| Formats | FP16 / INT8 / INT4 | GGUF (via llama.cpp) | GGUF, F16 |
| Typical use case | small models on mobile and in the browser | conversational model on your desktop | benchmarks and any existing GPU |
When ML Drift beats the generalists: when the constraint is memory or latency on a device. When you need to run a 24-70B model on a desktop, or want the full GGUF ecosystem, you stay on Ollama and llama.cpp. To see how the two generalist engines work, start with Ollama vs llama.cpp: under the hood.
The numbers at release
The numbers Google reports in the announcement, all on production hardware:
- Gemma: up to 12% lower memory overhead than the legacy delegate
- YouTube Shorts: up to 40% lower frame latency (Android and iOS)
- Google Photos: up to a 2-second speedup over the previous delegate
- Adobe Lightroom and Photoshop: up to 30% faster on-device processing
- Snap: 30% lower model latency
The repository claims an order of magnitude over existing open-source GPU inference engines — a project's own figure, to take as such until independent benchmarks land. The technology is documented in the technical paper arXiv 2505.00232.
And the limits?
- Desktop is an early preview. Windows and Linux run via Dawn: the support exists, but Google labels it early stage.
- On-device models. The benchmarks show the Gemma family: ML Drift is tuned for models that fit on a device, not for replacing a cluster. It is not a vLLM.
- No GGUF ecosystem. With FP16/INT8/INT4 you do not carry over the quantized model library you use with Ollama: model selection is narrower.
- The TFLite GPU delegate is not evolving anymore. If you have apps built on it, migrating to ML Drift/LiteRT is the path Google recommends.
Trying it
The repository is the starting point: it contains 5D tensor examples (YOLO 11n, MobileViT v2, Swin v2), the custom op framework, and a SKILL.md designed to have an AI agent work on the engine's code. App integration goes through LiteRT; in the browser, through WebGPU via Dawn. If it does not work: check that your GPU exposes a supported backend (OpenCL, Metal, WebGPU) and that the model is converted to one of the supported precision formats — those are the two points where ports fail most often. If you want to squeeze more speed out of inference without changing hardware, look at Speculative Decoding and MTP: speed for free.
Sources
- Official announcement: ML Drift, next-gen GPU AI/ML inference at the edge
- Repository google-ai-edge/ml-drift
- Technical paper: arXiv 2505.00232
In summary
ML Drift is Google's on-device GPU engine, open source under Apache 2.0, the successor to the TFLite GPU delegate: virtualized tensors, a KV cache separated for prefill and decode, FP16/INT8/INT4. The production numbers are real — up to 40% lower latency on YouTube, 30% on Adobe and Snap — and the strength is the device, not the datacenter. Desktop is an early preview, the model ecosystem is narrower than GGUF, and it is not a replacement for Ollama, llama.cpp, or vLLM: it is a different machine, built for a different job.
Article generated with AI
In these reading paths
Enjoyed this article?
Subscribe to the newsletter for new tutorials and guides.
Related Articles

Cosmos 3 Edge: the DROID policy for robotics that runs on the edge
Cosmos3-Edge-Policy-DROID: NVIDIA's 4B robotic policy. How much hardware you need to run it in real time on a Jetson, and where it stops.

Frontier vs local AI: the honest comparison nobody makes
Frontier models are still more powerful on complex cases, but cost is the last thing that decides: the table, the break-even point and the honest answer to "front-load and then ride for free?".

