用于高效视频生成的参数化条纹注意力

扩散Transformer (DiT) 可以生成高质量的视频，但会遭受严重的推理延迟，这主要归因于计算成本高昂的完整时空注意力。虽然稀疏注意力方法提供了潜在的解决方案，但现有方法面临着固有的灵活性-效率困境：预定义的掩码缺乏捕获不同注意力模式的灵活性，而运行时确定的掩码会带来开销并牺牲硬件效率。我们认为缺乏统一的 DiT 注意力结构表征是现有方法的一个关键限制，并确定视频 DiT 注意力在时间和空间维度上都表现出 周期性对角条纹结构。为了在单个高效内核中正式编码这些结构化模式，我们提出了 {\bf PSA}，一种参数化的条带注意力，它形式化了观察到的条带规律性，统一了不同的注意力模式以实现高效的掩模生成。这种统一的表示使单个硬件高效的 CUDA 内核能够处理所有稀疏模式，从而实现 FlashAttention-3 级模型 FLOP 利用率。为了确定最佳的稀疏性配置，我们提出了一种免训练的离线搜索算法，该算法可以在每个注意头的指定误差容限下自动最大化稀疏性。在 HunyuanVideo 和 Wan~2.1 上的实验表明，PSA 比 FlashAttention-3 基线实现了 1.57$\times$ 和 1.37$\times$ 端到端加速，并且视觉质量下降是可以接受的。