Google publicó un caso de estudio sobre cómo acelerar la atención en la difusión de video, una etapa costosa al generar clips largos y de alta resolución. Su kernel en JAX y Pallas utiliza atención dispersa y alinea la máscara de atención con los bloques que procesa la TPU.
En las pruebas realizadas con un solo chip TPU v6e, la versión final tardó 32.76 milisegundos, frente a los 78.70 milisegundos de la atención densa Splash: una aceleración informada de 2.40×. Las mediciones utilizaron entradas sintéticas y excluyen el enrutamiento, la reorganización de tokens y la comunicación entre dispositivos, por lo que no demuestran una mejora de velocidad de principio a fin en la generación de video.
Por qué importa
Para los equipos que optimizan la generación de video en TPU, el resultado muestra que una máscara dispersa puede ser más lenta que la atención densa si el kernel no está diseñado para omitir y procesar eficientemente los bloques adecuados; la fuente no muestra el impacto en el tiempo total de generación de video.
