Google libera ML Drift, un motor GPU de código abierto para IA en el edge que unifica OpenCL, Metal, WebGPU y OpenGL ES

El equipo de Google AI Edge ha publicado en código abierto, bajo licencia Apache 2.0, ML Drift: un motor de cómputo GPU multiplataforma diseñado específicamente para la inferencia de IA/ML en el dispositivo (on-device). Su función es abstraer la complejidad del hardware y de las APIs de bajo nivel de las GPU en distintos sistemas operativos y navegadores —OpenGL ES, OpenCL, Metal y WebGPU—, de modo que los desarrolladores puedan construir experiencias de IA en tiempo real, desde efectos de vídeo avanzados hasta IA generativa, sin tener que reescribir shaders para cada backend. ML Drift pasa a ser el motor de aceleración GPU que hay dentro de LiteRT (el runtime ligero de Google, sucesor de TensorFlow Lite), pero también se distribuye como librería independiente para quienes quieran construir sus propios motores gráficos o de inferencia.

La motivación detrás del proyecto es la enorme fragmentación del hardware de borde (edge) frente a la homogeneidad de los centros de datos: los desarrolladores no saben de antemano en qué GPU, con qué controlador y bajo qué API de bajo nivel se ejecutará su aplicación. El antiguo 'TFLite GPU delegate' sentó las bases de la aceleración por GPU en el dispositivo, pero el ecosistema ha evolucionado: hoy conviven modelos clásicos de visión por computador, audio y profundidad con arquitecturas generativas de muchos parámetros, que llevan el silicio de consumo al límite de cómputo y memoria. Google explica que necesitaba una base de ingeniería unificada capaz de dar estabilidad a los modelos clásicos y, a la vez, máximo rendimiento a la IA generativa más exigente.

Para lograrlo, ML Drift introduce varios cambios estructurales respecto al antiguo delegate de TFLite GPU. El más relevante es la 'virtualización de tensores': antes había que programar a mano la correspondencia entre la representación lógica de un tensor y su asignación física (texturas, buffers) por separado para OpenGL, OpenCL y Metal, lo que obligaba a mantener bases de código de shaders distintas para cada backend. Ahora, unas plantillas de shaders dinámicas resuelven y traducen esas coordenadas durante la fase de inicialización del compilador, lo que elimina la necesidad de mantener código específico por plataforma y añade una sobrecarga de ejecución mínima, preservando la portabilidad de los modelos entre plataformas.

Otra novedad es un marco de operadores personalizados (custom ops) modernizado, con APIs de registro directo y acceso a bajo nivel del lenguaje de shading. Para facilitar esto, ML Drift incluye una guía 'agentic SKILL.md' pensada para que agentes de codificación (coding agents) puedan escribir, registrar y verificar shaders personalizados de forma rápida, lo que da a los desarrolladores un control granular para integrar bloques de modelos propios directamente en el grafo de ejecución.

También se añade soporte nativo para tensores de 5 dimensiones. El antiguo delegate de TFLite GPU estaba limitado estructuralmente a tensores de 4D, lo que forzaba a usar trucos de disposición de datos (layout hacks) para modelos complejos que requerían 5D. Con soporte 5D de fábrica, LiteRT puede ejecutar directamente en GPUs de borde cargas de trabajo como redes convolucionales 3D para IA volumétrica/espacial y modelos espaciotemporales, citándose como ejemplos YOLO 11n, MobileViT v2 y Swin Transformer v2.

Para los modelos clásicos que ya corrían sobre el antiguo backend, ML Drift se presenta como una actualización de rendimiento estricta: al desacoplar el runtime de las restricciones específicas de cada hardware y mejorar la eficiencia de ejecución de los kernels, esas cargas de trabajo existentes obtienen mejoras inmediatas de rendimiento, con una migración diseñada para ser sencilla y benchmarks que se mantienen o superan a los previos.

Un apartado especialmente técnico es la optimización 'consciente de la fase' (stage-aware) para LLMs en el borde. Los modelos de lenguaje autorregresivos tienen dos cargas de cómputo muy distintas durante la inferencia: la fase de precarga de la caché KV (prefill), que está limitada por el cómputo, y la fase de decodificación, que genera tokens uno a uno y está limitada por el ancho de banda de memoria. ML Drift cambia dinámicamente de kernels y configuraciones de disposición de datos según la fase activa; durante la decodificación utiliza una disposición de caché KV personalizada, alineada con convoluciones, y aplica cuantización agresiva de activaciones dentro del propio kernel para evitar idas y vueltas redundantes a memoria. El artículo incluye una gráfica de referencia con la configuración de prueba (prefill de 512 tokens, decodificación de 128, contexto de 640, pesos de 4 bits con bloque de tamaño 32, caché KV en fp16, muestreo por argmax y sincronización en cada token), señalando que el modelo Gemma4 E2B rinde mejor en frameworks que soportan el uso compartido de caché KV (cross KV sharing).

Más allá del móvil, Google está llevando ML Drift al escritorio. Aunque el backend WebGPU se diseñó originalmente para acelerar la inferencia dentro del navegador, el hecho de que Dawn (la implementación de WebGPU de Chromium) permita compilar exactamente ese mismo código fuera del navegador hace posible ejecutar ML Drift de forma nativa en Windows y Linux, aprovechando la abstracción de hardware moderna de WebGPU para evitar la fragmentación entre DirectX y Vulkan, y complementando el backend nativo de Metal ya existente en macOS. Google enmarca esto en el auge de los agentes de codificación locales y los flujos de trabajo en estaciones de trabajo, aunque insiste en que su prioridad sigue siendo la inferencia ligera en móvil y borde, donde las restricciones de recursos son más estrictas. Para el escritorio publica una instantánea temprana con modelos Gemma (prefill de 8192, decodificación de 1024, contexto de 9216, mismas condiciones de cuantización y muestreo que antes) y subraya que el pico de rendimiento no lo es todo, porque la huella de memoria determina qué cargas pueden ejecutarse a la vez: en sus benchmarks con Gemma, ML Drift logra hasta un 12% menos de sobrecarga de memoria que otros frameworks, lo que deja más memoria libre para otras tareas locales.

En cuanto a despliegue real, Google afirma que ML Drift ya funciona en producción en millones de dispositivos cada día, impulsando funciones clave en Chrome, YouTube Shorts, Google Photos y Meet, además de actuar como runtime de aplicaciones punteras de Google AI Edge como Gallery, Foresight y Eloquent. Se citan varios casos concretos: YouTube Shorts migró sus efectos basados en segmentación a ML Drift y logró hasta un 40% de reducción en la latencia media de fotograma tanto en Android como en iOS, permitiendo aplicar efectos visuales de alta fidelidad en tiempo real sin perder fotogramas de cámara. Google Photos integró ML Drift en sus canalizaciones de fotografía computacional y segmentación, consiguiendo hasta 2 segundos de mejora frente al antiguo delegate de GPU, lo que se ilustra con la función de desenfoque (Unblur). Google Chrome, por su parte, integró ML Drift en su infraestructura de runtime de IA para acelerar por hardware los modelos Gemini Nano que sustentan las API integradas de IA del navegador (Prompt, Summarize y Writer), con un ejemplo de uso de la API Prompt por parte de Cafe24 para generación de etiquetas.

Fuera del propio ecosistema de Google, varios socios de desarrollo también han adoptado la tecnología. Adobe actualizó funciones clave de IA en Lightroom y Photoshop —como Seleccionar Sujeto, Seleccionar Cielo y Retrato Adaptativo— con ML Drift, logrando hasta un 30% más de velocidad en edición fotográfica profesional en el dispositivo móvil. Snap integró LiteRT para potenciar la inferencia en GPU en dispositivos Android, y el acelerador GPU de ML Drift aportó mejoras de latencia de modelo del 30% en los efectos de generación facial y de estilo de los lentes de Snapchat; Google añade que ML Drift no solo mejora modelos existentes, sino que también permite ejecutar modelos de difusión de mayor tamaño.

El desarrollo de ML Drift también ha implicado una colaboración estrecha con fabricantes de silicio. Con Arm se co-optimizaron los kernels de cómputo y la localidad de la caché de texturas para las GPU Mali e Immortalis, buscando inferencia eficiente energéticamente en dispositivos móviles. Con Intel, ML Drift aprovecha el cómputo nativo de WebGPU y las extensiones Xe Matrix Extensions (XMX) para acelerar cargas de LLM y visión en procesadores Intel Core Ultra con gráficos Xe3. Con Qualcomm, se optimizaron los kernels de OpenCL para maximizar el uso de las unidades aritmético-lógicas (ALU) y el ancho de banda de memoria en las GPU Adreno.

En el plano práctico, Google anuncia que el antiguo delegate de GPU de TFLite deja de recibir nuevas funciones y anima a todos los desarrolladores a migrar al acelerador GPU ML Drift de LiteRT, que ofrece compatibilidad total con los modelos existentes y desbloquea de inmediato las mejoras de rendimiento en Android, iOS, web y escritorio. Para quienes usan runtimes 'unbundled' en Android (para minimizar el tamaño del binario de la app), la aceleración de ML Drift ya está disponible en los paquetes independientes de LiteRT, y llegará próximamente a LiteRT dentro de Google Play Services.

Para empezar, Google plantea dos caminos: los desarrolladores de aplicaciones pueden usar el acelerador ML Drift de LiteRT para obtener de serie el particionado de modelos, la virtualización de tensores y la aceleración de LLM optimizada; los ingenieros de GPU y de sistemas pueden construir directamente con las APIs C++ independientes de ML Drift desde su repositorio en GitHub, creando grafos de cómputo 'GpuModel' personalizados y escribiendo shaders propios, con guías rápidas de introducción tanto para OpenCL en Android como para WebGPU vía Dawn. Todo el proyecto es de código abierto bajo licencia Apache 2.0, con repositorio en github.com/google-ai-edge/ml-drift, un rastreador de incidencias en GitHub para reportar problemas o pedir funciones, un foro comunitario de Google AI Edge para debates y preguntas, y una guía de contribución para quienes quieran añadir nuevos operadores UCL u optimizaciones específicas de hardware.

En conjunto, el anuncio es relevante porque consolida en un único motor —con una sola base de shaders virtualizados— algo que hasta ahora exigía mantener implementaciones distintas por API gráfica y por plataforma, al tiempo que añade capacidades que el ecosistema de IA generativa en el borde venía reclamando: soporte de tensores de alta dimensionalidad, gestión específica de las dos fases de inferencia de los LLM y una vía de expansión hacia el escritorio basada en WebGPU nativo. El hecho de que ya esté desplegado en productos de gran escala de Google y en aplicaciones de terceros como Adobe o Snap, además de contar con optimizaciones conjuntas con Arm, Intel y Qualcomm, sugiere que no se trata de un proyecto experimental, sino de una pieza de infraestructura ya probada en producción que ahora se abre a toda la comunidad de desarrolladores de IA en el dispositivo.