Google lanza ML Drift para acelerar la inferencia de IA en hardware heterogéneo

La ejecución local de modelos de inteligencia artificial ofrece ventajas como la mejora de la privacidad, la independencia de la conexión a internet y la reducción de la latencia. No obstante, lograr un rendimiento eficiente en hardware de consumo con diversas GPU, controladores e interfaces de software ha presentado un desafío. Para abordar esta situación, el equipo de Google AI Edge ha desarrollado ML Drift, un motor de computación GPU de código abierto destinado a acelerar la inferencia de IA en dispositivos que abarcan desde teléfonos inteligentes hasta ordenadores de escritorio.

ML Drift representa el sucesor del delegado GPU previamente empleado por TensorFlow Lite y ahora funciona como el motor de aceleración GPU del marco LiteRT de Google. Los desarrolladores también pueden utilizarlo de manera independiente a través de sus API de C++. El software se distribuye bajo la licencia Apache 2.0, y su código fuente, junto con ejemplos para OpenCL y WebGPU, está disponible en GitHub, tal como informa Hackster.io. La publicación del nuevo motor marca el fin del desarrollo de características para el delegado GPU de TensorFlow Lite, y se recomienda a los desarrolladores migrar, proceso facilitado por la retrocompatibilidad.

Una de las principales dificultades para la aceleración GPU fuera de los centros de datos reside en la diversidad de interfaces de programación que emplean las distintas plataformas de hardware, como OpenCL u OpenGL ES en Android, Metal en dispositivos Apple o WebGPU en navegadores web. ML Drift soluciona esta fragmentación mediante una arquitectura de computación unificada que abstrae estas diferencias específicas del hardware. Esto permite a los desarrolladores escribir operaciones GPU utilizando un marco común, evitando la necesidad de mantener implementaciones separadas para cada plataforma compatible.

Elementos clave de esta arquitectura incluyen una técnica de virtualización de tensores que disocia la representación lógica de los tensores de su almacenamiento físico en la GPU. Esta capacidad, combinada con sombreadores generados dinámicamente y una gestión de memoria optimizada, permite a ML Drift adaptar las cargas de trabajo a distintos tipos de hardware con un mínimo impacto en el tiempo de ejecución. El motor también introduce soporte nativo para tensores de cinco dimensiones, lo que posibilita la gestión de cargas de trabajo más complejas, como redes convolucionales 3D y modelos que procesan información espacial o temporal.

Google ha implementado optimizaciones específicas para visión por computador, procesamiento de audio y profundidad, y modelos de lenguaje grandes, incluyendo estrategias de ejecución diferenciadas para el procesamiento de entradas y la generación de tokens de salida. Estas optimizaciones abarcan diseños de memoria especializados, una gestión mejorada de la caché KV y técnicas de cuantificación que reducen el tráfico de memoria. En las pruebas internas de Google, ML Drift demostró una reducción de hasta el 12% en el consumo de memoria en comparación con otros marcos de trabajo.