Google ha pubblicato un caso di studio su come accelerare l’attenzione nei modelli di diffusione video, una fase onerosa nella generazione di clip lunghe e ad alta risoluzione. Il kernel, realizzato con JAX e Pallas, usa l’attenzione sparsa e allinea la maschera di attenzione ai tile elaborati dalla TPU.
Nei test su un singolo chip TPU v6e, la versione finale ha impiegato 32.76 millisecondi, contro i 78.70 millisecondi dell’attenzione densa Splash: un’accelerazione dichiarata di 2.40×. Le misurazioni hanno usato input sintetici ed escludono il routing, il riordinamento dei token e la comunicazione tra dispositivi; non dimostrano quindi un’accelerazione dell’intero processo di generazione video.
Perché conta
Per i team che ottimizzano la generazione video sulle TPU, il risultato mostra che una maschera sparsa può essere più lenta dell’attenzione densa se il kernel non è progettato per saltare i tile giusti ed elaborarli in modo efficiente. La fonte non indica l’impatto sui tempi complessivi di generazione video.
