Geração Rápida de Vídeo com Atenção de Peças Deslizantes

Fast Video Generation with Sliding Tile Attention

February 6, 2025

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

cs.AI

Resumo

Os Transformadores de Difusão (DiTs) com atenção total em 3D são o estado-da-arte na geração de vídeos, mas sofrem com custos computacionais proibitivos - ao gerar um vídeo de 5 segundos em 720P, a atenção sozinha leva 800 dos 945 segundos do tempo total de inferência. Este artigo introduz a atenção de mosaico deslizante (STA) para lidar com esse desafio. O STA aproveita a observação de que os escores de atenção em modelos de difusão de vídeo pré-treinados se concentram predominantemente em janelas 3D localizadas. Ao deslizar e atender sobre a região espacial-temporal local, o STA elimina a redundância da atenção total. Ao contrário da atenção tradicional de janela deslizante por token (SWA), o STA opera azulejo por azulejo com um design de janela deslizante inovador, preservando a expressividade ao mesmo tempo que é eficiente em hardware. Com otimizações cuidadosas em nível de kernel, o STA oferece a primeira implementação eficiente de atenção em janela deslizante 2D/3D, alcançando 58,79% de MFU. Precisamente, o STA acelera a atenção em 2,8-17x em relação ao FlashAttention-2 (FA2) e 1,6-10x em relação ao FlashAttention-3 (FA3). No principal DiT de vídeo, HunyuanVideo, o STA reduz a latência de ponta a ponta de 945s (FA3) para 685s sem degradação de qualidade, sem necessidade de treinamento. Habilitar o ajuste fino reduz ainda mais a latência para 268s com apenas uma queda de 0,09% no 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.