Google publiceerde een casestudy over het versnellen van attention bij videodiffusie, een kostbare stap bij het genereren van lange clips met hoge resolutie. De JAX- en Pallas-kernel gebruikt sparse attention en stemt het attention-masker af op de tiles die de TPU verwerkt.
In tests op één TPU v6e-chip duurde de uiteindelijke versie 32.76 milliseconden, tegenover 78.70 milliseconden voor dense Splash attention: een gemelde versnelling van 2.40×. De metingen waren gebaseerd op synthetische invoer en laten routing, het herschikken van tokens en communicatie tussen apparaten buiten beschouwing. Ze tonen dus geen versnelling van de volledige videogeneratiepipeline aan.
Waarom het ertoe doet
Voor teams die videogeneratie op TPU’s optimaliseren, laat het resultaat zien dat een sparse masker trager kan zijn dan dense attention, tenzij de kernel zo is ontworpen dat die de juiste tiles overslaat en efficiënt verwerkt. De bron toont niet wat het effect is op de totale videogeneratietijd.
