Parameterized Stripe Attention for Efficient Video Generation¶
会议: NeurIPS 2026
arXiv: 2609.37001
代码: https://github.com/jxyjason/PSA
领域: 视频理解 / Video Understanding
关键词: 视频生成, Diffusion Transformer, 稀疏注意力, 参数化条带注意力, CUDA Kernel优化
一句话总结¶
针对视频生成 Diffusion Transformer (DiT) 全时空注意力推理开销巨大的瓶颈，本文揭示了注意力图中普遍存在的周期性对角条带结构，提出统一的六参数条带注意力 (PSA) 及硬件感知的统一 CUDA 核函数与无训练离线搜索算法，在 HunyuanVideo 和 Wan 2.1 上取得了 1.37×–1.57× 的端到端推理加速且几乎无画质损失。
研究背景与动机¶
以 Wan 2.1 和 HunyuanVideo 为代表的视频扩散 Transformer (DiT) 架构在文本到视频、图像到视频等生成任务上展现出强大的时空建模能力与生成质量。然而，生成几秒钟的高清视频通常需要几十步去噪迭代，其推理延迟极其昂贵。在长序列时空三维注意力机制下，注意力模块的计算时间占到了模型整体端到端推理时延的 60% 以上，成为阻碍其实际部署与落地应用的首要性能瓶颈。
为了加速 DiT 推理，现有的稀疏注意力方法主要陷入了「灵活性与硬件效率难以兼得」的两难困境。一类方法依赖人工预定义的固定掩码，例如 SVG 仅粗糙地将注意力头分为时间和空间两类，STA 仅考虑局部滑动窗口；这些预设掩码无法覆盖多对角线、跨步长及混合型长程依赖模式，表达能力严重受限。另一类方法则尝试在运行时动态生成稀疏掩码，例如 SVG2 采用在线 K-Means 聚类与 token 重排；但放弃结构规律导致其必须依赖不规则的通用算子，不仅无法利用现代 GPU（如 NVIDIA Hopper 架构）的 Tensor Memory Accelerator (TMA) 硬件特性，还会带来严重的内存 Padding（70% 稀疏度下约 20%）和在线聚类计算开销（约 10%），实际加速比大打折扣。
导致上述困境的核心根源，在于学界缺乏对视频 DiT 注意力分布结构的统一形式化机理认知。本文的切入角度是：视频 DiT 注意力既非无序任意分布，也非简单局部窗口，而是受到视频天然时空冗余与 RoPE-3D 旋转位置编码余弦周期性的共同约束。核心 idea：将视频 DiT 注意力统一建模为沿时间和空间维度的参数化周期性对角条带结构 \(\mathcal{M}(T_w, T_o, T_s, S_w, S_o, S_s)\)，借助统一表达构建单算子硬件感知 CUDA Kernel，并结合免训练离线搜索实现无运行时开销的高效稀疏推理。
方法详解¶
整体框架¶
PSA 方法的整体流程包含三个有机衔接的组成部分：理论驱动的参数化结构建模、硬件感知的统一执行算子（PSA-Kernel）、以及跨提词保持不变性的免训练离线搜索机制（PSA-Search）。首先，基于视频时空局部冗余与 RoPE-3D 的投影分解，本文确立了注意力图由帧间（Inter-frame）与帧内（Intra-frame）两组条带参数联合决定的形式化定义。接着，在模型部署阶段，通过统一的掩码逻辑驱动单一高性能 CUDA 核函数，采用基于 PSABlock 的数据排布与生产者-消费者流水线，实现类似 FlashAttention-3 级别的高硬件利用率。最后，由于注意力掩码在不同 Prompt 间表现出极高的重合度，系统可以在代表性数据集上进行一次性离线搜索，为每个注意力头挑选受误差阈值约束的最大稀疏度配置。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入视频 Token<br/>(T×H×W 时空序列)"] --> B["周期性对角条带参数化<br/>M(Tw, To, Ts, Sw, So, Ss)"]
B --> C["PSABlock 布局重排与 3D-Padding<br/>对齐硬件 SM 与张量块边界"]
C --> D["双 WarpGroup 异步流水线算子<br/>生产者 TMA 预取 + 消费者局部密集计算"]
D --> E["免训练离线搜索策略<br/>逐头贪心最大化稀疏度 (L2 ≤ α)"]
E --> F["输出加速视频 DiT 推理<br/>(1.37×–1.57× 端到端加速)"]关键设计¶
1. 周期性对角条带参数化：统一描述多样化稀疏注意力结构
传统方法在固定预设与完全动态掩码之间两极分化，未能捕捉到 DiT 注意力真正的结构先验。通过对 RoPE-3D 注意力 Logit 的严格数学分解，查询项与键项在通道分组 \(\lambda \in \{t, h, w\}\) 上的内积受余弦项主导。当主导通道 \(m_\lambda^*\) 占优时（在 Wan 2.1 中高达 87.55% 的头满足此条件），注意力 Logit 呈现沿各轴的余弦周期震荡，其轴向周期为 \(\mathcal{T}_\lambda = 2\pi / \theta_{m_\lambda^*} = 2\pi c^{2m_\lambda^*/d}\)。与此同时，视频自身的时空连续冗余在局部半径 \(\varepsilon_t, \varepsilon_h, \varepsilon_w\) 内集中了高权重。这两者共同作用，使二维展开的注意力热图在数学上必然呈现周期性对角条带。
基于此机理，本文提出用六元参数统一描述条带掩码： $\(\mathcal{M}(T_w, T_o, T_s, S_w, S_o, S_s)\)$ 其中 \(T_w, T_o, T_s\) 分别控制帧间条带的带宽（Width）、偏移量（Offset）和跨步周期（Stride）；\(S_w, S_o, S_s\) 则控制帧内条带的对应参数。对于查询 token \(i\)（对应时空坐标 \((t_q, s_q)\)）和键 token \(j\)（对应 \((t_k, s_k)\)），二者能够产生注意力交互的充要条件是同时满足帧间与帧内条件： $\(\mathrm{inter\_cond} \iff ((|t_k - t_q - T_o|) \bmod T_s) \le T_w - 1\)$ $\(\mathrm{intra\_cond} \iff ((|s_k - s_q - S_o|) \bmod S_s) \le S_w - 1\)$ 该参数化统一了帧内型、帧间型、混合型和均匀型四种主要注意力模式（SVG 和 STA 分别退化为其中的极端特例），在所有注意力头上取得了与完整全局注意力高达 0.9831 的平均余弦相似度。
2. PSABlock 布局重排与 3D-Padding：保障硬件对齐与任意分辨率支持
针对细粒度 token 掩码在现代 GPU 上引发的内存访问不连续和分支发散问题，PSA 确立了硬件匹配的最小计算粒度——PSABlock。每个 PSABlock 由时空三维小块构成：\(\text{PSABlock} = \text{PSABlock}_t \times \text{PSABlock}_h \times \text{PSABlock}_w\)，正好能够充分饱和单个流式多处理器（SM）的计算与张量核心资源。
为了实现连续高效的合并内存访问（Coalesced Memory Access），输入 \(Q, K, V\) 张量从原本扁平的 \((T, H, W)\) 重排为按 PSABlock 网格组织的层次化布局 \((T', H', W', \text{PSABlock}_t, \text{PSABlock}_h, \text{PSABlock}_w)\)。同时，为支持任意长宽比与分辨率的视频生成，PSA 引入了轻量级 3D-Padding，将每个维度仅微幅填充至 PSABlock 大小的最小整数倍：\(T' = \lceil T / \text{PSABlock}_t \rceil \cdot \text{PSABlock}_t\)。由于去噪轨迹前 20% 步骤保留密集注意力，3D-Padding 仅需作用于后续步数中，在 720p 分辨率下的整体计算额外开销仅约为 5%，彻底摆脱了 STA 等方法强行绑定特定瓦片分辨率的限制。
3. 双 WarpGroup 异步流水线算子：释放 Hopper 架构极限算力
针对不规则稀疏算子难以发挥硬件极致吞吐的缺陷，本文基于 ThunderKittens 库开发了专为 Hopper 架构定制的高性能 PSA-Kernel。核函数在每个线程块内将工作负载解耦为两组并发协作的 WarpGroup： - 生产者（Producer Warpgroup）：解析掩码参数 \(\mathcal{M}\)，筛选出当前查询块所需的候选 \(K, V\) PSABlock 列表，并通过 TMA 硬件单元以乒乓双缓冲（Ping-Pong Buffering）机制异步将全局显存数据预取至片上 SRAM； - 消费者（Consumer Warpgroups）：在片上共享内存中针对常驻的 \(Q\) 块与预取的 \(K, V\) 块执行高吞吐量的稠密 GEMM 运算、安全 Softmax 规约与局部输出累加。
这一设计在保持上层逻辑灵活掩码的同时，让底层 SM 执行的完全是结构化、硬件高度亲和的局部密集 GEMM，彻底消除了显存存取延迟与空闲计算单元，在 HunyuanVideo 上实现了高达 61.2% 的 MFU（模型算力利用率），逼近 FlashAttention-3 的峰值性能。
4. 免训练离线搜索策略：零推理开销的头自适应配置寻优
为了在推理阶段无需任何动态判断即可赋予每个注意力头最契合的稀疏掩码，本文观察到视频 DiT 的注意力热图虽然在不同去噪时间步、不同网络层和不同头上差异巨大，但在不同文本 Prompt 之间表现出惊人的结构稳定性（余弦相似度通常超过 0.99）。
据此，PSA-Search 算法采用纯离线、无训练的并行评估流程。对于模型中每一个注意力头，算法在离线构造的候选掩码空间中遍历评估，记录稀疏输出相较于全注意力输出的 \(L_2\) 误差损失。在给定的误差容忍上限 \(\alpha\) 下，贪心地挑选出稀疏度最高的掩码配置： $\(\max_{\mathcal{M} \in \mathcal{C_M}} \text{Sparsity}(\mathcal{M}) \quad \text{s.t.} \quad L_2(\text{PSA}(h, \mathcal{M}), \text{FullAttn}(h)) \le \alpha\)$ 借助 DeepSpeed-Ulysses 序列并行支持，8 卡 H100 仅需 2.5 小时即可完成模型全量配置的离线建表。由于各候选掩码的逐头损失已完整持久化，当用户后续微调目标稀疏度或误差阈值 \(\alpha\) 时，仅需两分钟的轻量扫描重排，即可获得全新的全局掩码分配策略。
一个完整示例¶
以 Wan 2.1 14B 模型处理 720×1280 视频（对应 DiT 形状为 \(21 \times 45 \times 80\)）为例： 1. 空间高度维度 \(H=45\) 经 3D-Padding 补齐到 48，匹配 PSABlock 尺寸 \(3 \times 8 \times 16\)； 2. 离线建表发现某深层注意力头具有典型的帧间长程关联与帧内频域带状结构，在误差阈值 \(\alpha=0.007\) 约束下被赋予掩码参数 \((T_w=1, T_o=0, T_s=3, S_w=5, S_o=0, S_s=6)\)； 3. 推理执行时，生成前 20% 的时间步（第 1–10 步）运行原生 FlashAttention-3 保障宏观构图；后 80% 时间步（第 11–50 步）无缝切换至 PSA-Kernel； 4. 生产者 WarpGroup 根据上述参数，跳过无须交互的时空块，仅通过 TMA 加载满足条件的对应 \(K, V\) 块进 SRAM，使该头稀疏度达到 80% 以上； 5. 最终全模型平均稀疏度稳定在 57.57%，单视频生成时延从 FlashAttention-3 的 1214.69 秒缩短至 886.64 秒，提速 1.37×。
实验关键数据¶
主实验¶
在单卡 NVIDIA H100 GPU 上，评估覆盖了 HunyuanVideo（720×1280 与 768×1280 两种配置）和 Wan 2.1 14B 模型。基准对比涵盖了 FlashAttention-2、FlashAttention-3，以及代表性稀疏方法 SVG、SVG2 与 STA。
| 模型与分辨率 | 方法 | 平均稀疏度 (%) | 端到端耗时 (s) / 加速比 | VBench-OC (↑) | VBench-TS (↑) | PSNR (↑) | SSIM (↑) | LPIPS (↓) |
|---|---|---|---|---|---|---|---|---|
| HunyuanVideo (720×1280) | FA3 | 0.00 | 1172.62 (1.00×) | 26.95 | 24.52 | - | - | - |
| HunyuanVideo (720×1280) | SVG | 56.92 | 906.12 (1.29×) | 22.32 | 21.11 | 24.27 | 0.8686 | 0.3400 |
| HunyuanVideo (720×1280) | SVG2 | 61.92 | 866.10 (1.35×) | 25.96 | 23.28 | 29.63 | 0.8910 | 0.1625 |
| HunyuanVideo (720×1280) | PSA (Ours) | 59.74 | 777.28 (1.51×) | 26.00 | 23.63 | 30.35 | 0.9062 | 0.1523 |
| HunyuanVideo (768×1280) | FA3 | 0.00 | 1106.36 (1.00×) | 26.97 | 24.61 | - | - | - |
| HunyuanVideo (768×1280) | STA | 58.37 | 746.73 (1.48×) | 26.19 | 23.39 | 12.52 | 0.6036 | 0.3117 |
| HunyuanVideo (768×1280) | SVG | 56.98 | 797.57 (1.39×) | 23.21 | 22.43 | 23.00 | 0.8654 | 0.1907 |
| HunyuanVideo (768×1280) | SVG2 | 63.23 | 773.47 (1.43×) | 25.31 | 22.74 | 25.80 | 0.7855 | 0.2184 |
| HunyuanVideo (768×1280) | PSA (Ours) | 60.33 | 703.91 (1.57×) | 26.33 | 23.64 | 27.20 | 0.8287 | 0.1568 |
| Wan 2.1 14B (720×1280) | FA3 | 0.00 | 1214.69 (1.00×) | 25.91 | 23.19 | - | - | - |
| Wan 2.1 14B (720×1280) | SVG | 58.51 | 944.81 (1.29×) | 25.29 | 20.32 | 18.59 | 0.7056 | 0.2881 |
| Wan 2.1 14B (720×1280) | SVG2 | 63.40 | 926.78 (1.31×) | 24.69 | 23.33 | 23.47 | 0.8391 | 0.1703 |
| Wan 2.1 14B (720×1280) | PSA (Ours) | 57.57 | 886.64 (1.37×) | 25.29 | 23.59 | 23.48 | 0.8372 | 0.1601 |
消融实验¶
在固定 60% 稀疏度预算下，对参数化条带掩码所包含的各自由度进行系统消融（HunyuanVideo 720×1280）。
| 掩码配置 | 启用参数 | MSE (↓) | PSNR (↑) | SSIM (↑) | LPIPS (↓) | 说明 |
|---|---|---|---|---|---|---|
| SVG-like | \(\mathcal{M}(T_w, S_w)\) | 0.0027 | 26.0230 | 0.8364 | 0.2073 | 仅保留时空带宽，无法表达偏移与周期跨步 |
| STA-like | \(\mathcal{M}(T_w, S_w, S_s)\) | 0.0026 | 26.1466 | 0.8133 | 0.2343 | 仅增加空间跨步，缺少时空偏移和时间周期 |
| Inter-frame 仅限 | \(\mathcal{M}(T_w, T_o, T_s)\) | 0.0019 | 27.8065 | 0.8505 | 0.1757 | 仅保留完整时间条带参数，忽略空间细粒度稀疏 |
| Intra-frame 仅限 | \(\mathcal{M}(S_w, S_o, S_s)\) | 0.0042 | 24.0145 | 0.7838 | 0.2730 | 仅保留完整空间条带参数，破坏时序连续性 |
| Full Model (Ours) | \(\mathcal{M}(T_w, T_o, T_s, S_w, S_o, S_s)\) | 0.0012 | 30.3463 | 0.9062 | 0.1523 | 完整六参数模型，画质显著领先 |
关键发现¶
- 六参数完整建模至关重要：消融实验显示，相较于次优的仅帧间建模（PSNR 27.81），完整六参数模型一举达到 30.35 PSNR（提升超 2.5 dB），证明偏移（Offset）与跨步（Stride）参数对于精确匹配多对角线及混合注意力模式不可或缺。
- 端到端加速与算子效率双领先：在 HunyuanVideo (768×1280) 下，PSA 的单算子执行时间仅为 112.25 ms，MFU 高达 61.21%，相较 SVG/SVG2 算子取得超 1.6× 的微基准加速。在端到端耗时上达到 703.91 秒（1.57× 加速比），超越所有基线。
- 极佳的稀疏度-质量鲁棒性：在严格包夹（Bracketing）实验中，PSA 即使在更高稀疏度（63.84%）下的 PSNR（26.76）仍大幅优于处于更低稀疏度（58.37%）的 STA（PSNR 仅 12.52，发生严重局部模式失真）。在 9 个模型设置的广泛统计中，主导通道的存在比例中位数达到 ~90%，证明该条带结构的普适性。
亮点与洞察¶
- 从数学机理揭示条带注意力物理成因：论文跳出经验式启发规则，从视频天然时空冗余与 RoPE-3D 旋转位置编码通道内积的余弦周期性出发，严谨推导出注意力热图条带周期的闭式解 \(\mathcal{T}_\lambda = 2\pi / \theta_{m_\lambda^*}\)，为后续算法设计提供了坚实的理论立足点。
- 以 PSABlock 规整化抹平灵活性与硬件效率的鸿沟：通过引入与 SM 硬件规格对齐的 PSABlock 抽象，将灵活复杂的稀疏掩码收敛至规则的块级粒度，结合 TMA 异步预取与双缓冲设计，成功在单一 Kernel 内兼顾任意条带形状与接近 FA3 级别的极高硬件算力利用率。
- 免训练离线搜索极大降低落地门槛：敏锐利用了 DiT 注意力分布「跨 Prompt 高度稳定不变」的实证特征，将复杂的配置寻优完全剥离出推理阶段，既免去重新微调或蒸馏的数千 GPU 小时开销，又彻底消除了动态推理中的额外计算损耗。
局限与展望¶
- 模型范式拓展限制：当前研究与验证主要聚焦于视频生成 DiT 领域，虽然在图像生成模型（如 FLUX.1、Hunyuan-DiT）中也发现了类似的主导通道现象，但能否平滑迁移至因果解码语言模型（LLM）或跨模态理解模型仍有待检验。
- 离线搜索算力成本仍可优化：在 8 张 H100 GPU 上离线评估全部注意力头的所有候选掩码仍需 2.5 小时。未来可设计粗粒度筛选到细粒度精修的两阶段搜索算法，进一步压缩离线建表用时。
相关工作与启发¶
- vs SVG / STA (预设掩码范式)：SVG 仅使用粗粒度时空二分头，STA 强行采用局部平铺瓦片窗口。PSA 证明了固定局部假设会丢失 60% 以上头的长程与多对角结构；通过六参数化统一覆盖了任意跨步与偏移，且通过 3D-Padding 摆脱了特定输入分辨率约束。
- vs SVG2 (动态重排范式)：SVG2 假设注意力是任意无规律的，通过动态 K-Means 聚类和 token 置换构建掩码，牺牲了硬件连续性并带来了 ~30% 的非必要开销。PSA 揭示了隐藏的周期性几何结构，证明无需在线动态聚类即可通过离线自适应配置实现更高画质与更优速度。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次从 RoPE-3D 与时空冗余推导出视频 DiT 注意力对角条带周期的数学成因，并提出六参数统一表征。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 HunyuanVideo、Wan 2.1 等百亿参数顶尖模型，微基准 Kernel、端到端延迟、VBench、主导通道普适性及包夹对比完备详实。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰，理论推导严谨自洽，软硬件协同设计的动机交代透彻。
- 价值: ⭐⭐⭐⭐⭐ 为大模型视频生成的高延迟瓶颈提供了免微调、即插即用且兼容任意分辨率的高性能开源解决方案。