Google pokazuje, jak przyspieszyć uwagę w dyfuzji wideo na TPU 2,4x
Inżynierowie Google opisują, jak zaimplementowali rzadką uwagę przestrzenno-czasową (sparse spatio-temporal attention) dla modeli dyfuzji wideo na chipach TPU v6e, przekształcając teoretyczną rzadkość podejścia Sparse VideoGen (SVG) w realne przyspieszenia sprzętowe. Dzięki serii optymalizacji kernela (specjalizacja kafelków pełnych/brzegowych, strojenie rozmiaru kafelków oraz wyrównanie kafelków maski) zredukowali opóźnienie kernela uwagi z 96,37ms (naiwna wersja rzadka) do 32,76ms, co daje 2,40x przyspieszenie względem gęstej Splash Attention na pojedynczym chipie TPU v6e, przy 75,6K tokenów, 10 głowicach i wymiarze głowicy 128.