Google hat eine Fallstudie zur Beschleunigung der Attention in der Videodiffusion veröffentlicht – einem rechenintensiven Schritt bei der Erzeugung langer Clips mit hoher Auflösung. Der JAX- und Pallas-Kernel nutzt Sparse Attention und richtet die Attention-Maske an den Kacheln aus, die der TPU berechnet.
In Tests auf einem TPU-v6e-Chip benötigte die finale Version 32,76 Millisekunden, gegenüber 78,70 Millisekunden für die dichte Splash Attention – laut Google eine Beschleunigung um den Faktor 2,40. Die Messungen basierten auf synthetischen Eingaben und schlossen Routing, die Umordnung von Tokens sowie die Kommunikation zwischen Geräten aus. Sie belegen daher keine Beschleunigung der Videogenerierung insgesamt.
Warum es wichtig ist
Für Teams, die die Videogenerierung auf TPUs optimieren, zeigt das Ergebnis: Eine Sparse-Maske kann langsamer sein als dichte Attention, wenn der Kernel nicht darauf ausgelegt ist, die richtigen Kacheln zu überspringen und effizient zu verarbeiten. Die Quelle zeigt nicht, wie sich das auf die Gesamtdauer der Videogenerierung auswirkt.
