Schnelle Videogenerierung mit Schiebe-Kachel-Aufmerksamkeit

Fast Video Generation with Sliding Tile Attention

February 6, 2025

Autoren: Peiyuan Zhang, Yongqi Chen, Runlong Su, Hangliang Ding, Ion Stoica, Zhenghong Liu, Hao Zhang

cs.AI

Zusammenfassung

Diffusions-Transformer (DiTs) mit 3D-Voll-Aufmerksamkeit erreichen einen Spitzenzustand in der Videogenerierung, leiden jedoch unter prohibitiven Rechenkosten - bei der Erzeugung eines nur 5 Sekunden langen 720P-Videos beansprucht die Aufmerksamkeit allein 800 von insgesamt 945 Sekunden Inferenzzeit. Dieser Artikel stellt das Schiebe-Kachel-Aufmerksamkeit (STA) vor, um diese Herausforderung anzugehen. STA nutzt die Beobachtung, dass die Aufmerksamkeitswerte in vorab trainierten Video-Diffusionsmodellen hauptsächlich in lokalisierten 3D-Fenstern konzentriert sind. Durch Verschieben und Aufmerksamkeit auf den lokalen raumzeitlichen Bereich beseitigt STA Redundanzen der Voll-Aufmerksamkeit. Im Gegensatz zur traditionellen tokenweisen Schiebefenster-Aufmerksamkeit (SWA) arbeitet STA kachelweise mit einem neuartigen, hardwarebewussten Schiebefensterdesign, das die Ausdruckskraft bewahrt und gleichzeitig hardwareeffizient ist. Mit sorgfältigen Optimierungen auf Kernebene bietet STA die erste effiziente Implementierung einer 2D/3D-Schiebefenster-ähnlichen Aufmerksamkeit und erreicht 58,79% MFU. Präzise beschleunigt STA die Aufmerksamkeit um das 2,8-17-fache gegenüber FlashAttention-2 (FA2) und das 1,6-10-fache gegenüber FlashAttention-3 (FA3). Bei dem führenden Video-DiT, HunyuanVideo, reduziert STA die Latenz von Ende zu Ende von 945s (FA3) auf 685s, ohne Qualitätsverlust und ohne Training zu benötigen. Durch die Aktivierung des Feintunings wird die Latenz weiter auf 268s gesenkt, bei nur einem 0,09%igen Rückgang auf VBench.

English

Diffusion Transformers (DiTs) with 3D full attention power state-of-the-art video generation, but suffer from prohibitive compute cost -- when generating just a 5-second 720P video, attention alone takes 800 out of 945 seconds of total inference time. This paper introduces sliding tile attention (STA) to address this challenge. STA leverages the observation that attention scores in pretrained video diffusion models predominantly concentrate within localized 3D windows. By sliding and attending over the local spatial-temporal region, STA eliminates redundancy from full attention. Unlike traditional token-wise sliding window attention (SWA), STA operates tile-by-tile with a novel hardware-aware sliding window design, preserving expressiveness while being hardware-efficient. With careful kernel-level optimizations, STA offers the first efficient 2D/3D sliding-window-like attention implementation, achieving 58.79% MFU. Precisely, STA accelerates attention by 2.8-17x over FlashAttention-2 (FA2) and 1.6-10x over FlashAttention-3 (FA3). On the leading video DiT, HunyuanVideo, STA reduces end-to-end latency from 945s (FA3) to 685s without quality degradation, requiring no training. Enabling finetuning further lowers latency to 268s with only a 0.09% drop on VBench.