用好视觉Attention局部性，清华、字节提出Token Reorder，无损实现5倍稀疏、4比特...

主要观点总结

本文关注视觉生成任务的“局部性”特性，提出PAROAttention方案，通过Token重排实现算法侧视觉特征提取的局部性与硬件计算的局部性对应，优化注意力图的分布，以减小稀疏和量化的挑战。该方案包括离线确定的稀疏掩码和量化方案，旨在提高算法性能并保持硬件效率。

关键观点总结

关键观点1: 视觉生成任务的“局部性”特性

视觉特征提取具有局部性，不同注意力头在不同情况下呈现一致的局部聚合，可通过Token重排转化为硬件友好的块状模式。

关键观点2: PAROAttention方案介绍

PAROAttention采用Token重排方案，将多样且分散的注意力模式转化为统一的块状模式，便于注意力的稀疏与量化处理。

关键观点3: 稀疏方案的设计

分析动态和静态稀疏方案的优缺点，最终选择静态稀疏方案作为PAROAttention的主要稀疏方案。利用离线确定的稀疏掩码，避免在线产生稀疏掩膜的开销。

关键观点4: 量化方案的设计

通过合适的Token重排缓解Attention Map块内的数据差异，支持更低位宽的量化。

关键观点5: CUDA系统设计

为最小化额外开销，进行算子融合和显存优化。提出的方案与FlashAttention兼容，广泛适用于各种场景。

关键观点6: 实验结果

在主流视频和图片生成模型上测试算法性能保持效果，取得显著成果。硬件加速效果方面，实现了高效的稀疏和量化处理。