Generación rápida de video con Atención de Azulejos 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

Resumen

Los Transformadores de Difusión (DiTs) con atención completa en 3D han logrado un estado del arte en generación de video, pero sufren de un costo computacional prohibitivo: al generar un video de 720P de solo 5 segundos, la atención sola representa 800 de los 945 segundos totales de tiempo de inferencia. Este artículo presenta la atención de ventana deslizante (STA) para abordar este desafío. STA aprovecha la observación de que los puntajes de atención en modelos de difusión de video preentrenados se concentran predominantemente en ventanas 3D localizadas. Al deslizar y atender sobre la región espacio-temporal local, STA elimina la redundancia de la atención completa. A diferencia de la atención tradicional de ventana deslizante por token (SWA), STA opera de manera baldosa con un diseño de ventana deslizante novedoso consciente del hardware, preservando la expresividad al mismo tiempo que es eficiente en hardware. Con optimizaciones cuidadosas a nivel de kernel, STA ofrece la primera implementación eficiente de atención tipo ventana deslizante en 2D/3D, logrando un 58.79% de MFU. Precisamente, STA acelera la atención de 2.8 a 17 veces en comparación con FlashAttention-2 (FA2) y de 1.6 a 10 veces en comparación con FlashAttention-3 (FA3). En el principal DiT de video, HunyuanVideo, STA reduce la latencia de extremo a extremo de 945s (FA3) a 685s sin degradación de calidad, sin necesidad de entrenamiento. Habilitar el ajuste fino reduce aún más la latencia a 268s con solo una disminución del 0.09% en 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.