FVAttn:面向视频生成的运行时负载均衡自适应稀疏注意力 FVAttn: Adaptive Sparse Attention with Runtime Load Balancing for Video Generation
训练免修的稀疏注意力系统,用P2P头部迁移和空闲补偿消除多GPU视频生成的负载不均
前置知识
视频扩散Transformer(Video DiT)
视频扩散Transformer是一类将扩散模型的去噪过程实现为Transformer架构的生成模型,用于生成高分辨率视频。它的核心特征是将视频编码为长序列的时空token,然后通过Transformer的自注意力机制建模token之间的关系。由于视频的分辨率、帧数和时长持续增长,需要处理的时空token序列急剧膨胀,自注意力因其对序列长度的二次复杂度成为推理的主要瓶颈。例如在单卡NVIDIA H20上用Wan2.2-14B I2V生成5秒720p视频仍需约两小时,其中注意力占74.1%的总推理时间。
本文的全部动机和优化目标都建立在视频DiT注意力计算开销过大这一事实之上,理解视频DiT如何工作、为什么注意力会成为瓶颈,是理解本文为什么要做稀疏注意力和负载均衡的前提。
稀疏注意力与Top-p路由
稀疏注意力是减少注意力计算量的技术,核心思想是:对于每个query block,只有少数key block对输出有显著贡献,可以跳过低价值的query-key交互。训练免(training-free)的稀疏方法在推理时构建块级掩码。Top-k为每个query block固定保留k个最高分key block,负载规则但保真度受限;Top-p则保留累积归一化重要性达到阈值p的最小block集合,能根据注意力分布形状自适应分配计算量——集中分布保留少量block,平坦或多峰分布保留更多,保真度更好但产生异构负载。本文用Top-p路由加Top-k安全下限的组合作为稀疏路由前端。
Top-p路由的保真度优势和它引入的负载异构性是本文矛盾的根源:Top-p越自适应,跨头的工作量差异就越大,正是这种异构性在多GPU下被放大成负载不均,是理解整个FVAttn设计的钥匙。
序列并行(Ulysses)
序列并行是处理长序列的多GPU并行策略,其中Ulysses风格通过一次序列到头的All-to-All通信,将序列分片转换为头分片:每个GPU持有完整序列但只负责一部分注意力头,独立计算本地注意力。对于密集注意力,每个头的计算量几乎相同,因此这种布局天然均衡。但动态稀疏注意力打破了这个性质,因为每个头的实际计算量由当前的稀疏掩码决定,不同头保留的block数不同,导致某些GPU(rank)负载过重成为拖尾者,其他rank完成后被迫闲置等待。本文的负载均衡因子ρ衡量这种不均程度,ρ=1为完美均衡。
FVAttn的所有机制都建立在Ulysses序列并行之上,不理解序列到头的All-to-All如何分配工作、为什么密集注意力均衡而稀疏注意力不均衡,就无法理解负载不均从何而来,也无法理解RLB为什么要做头部迁移。
步数蒸馏(Step Distillation)
步数蒸馏是一种将需要多步去噪的扩散模型压缩为少步模型的工业部署技术,可将去噪步数大幅减少(例如本文的LightX2V 4步蒸馏LoRA配置),提供高达25倍的推理加速。然而步数蒸馏只减少了去噪步数,并没有改变每步的序列长度和由此带来的二次注意力开销,因此在少步机制下注意力仍是主要瓶颈。更关键的是,少步生成中相邻步之间的稀疏模式变化极大,本文测量到在4步蒸馏设置下,相同头密度的最大相邻步变化达97%,相同rank负载的变化达44%,这使得基于历史稀疏性的预测调度非常脆弱。
步数蒸馏是本文的设定背景,它解释了为什么传统的基于历史 profiling 的负载布局方案会失效——少步生成下相邻步的负载剧烈波动,必须基于当前实际观测的负载来调度,这是FVAttn选择'运行时修复当前负载'而非'预测未来布局'的根本原因。
负载均衡因子ρ
负载均衡因子ρ是本文用来量化多GPU负载不均程度的核心指标,定义为ρ = max_r L_r / ( (1/R) Σ_{r=1}^{R} L_r ),其中R是GPU数量(rank数),L_r是rank r的工作量,等于该rank上所有本地注意力头工作量之和 L_r = Σ_{h∈H_r} L_h,而单头工作量 L_h = mean_i Σ_j M_{h,i,j}(即该头掩码中被选中的query-key block对的平均数量)。ρ接近1表示负载均衡,ρ远大于1表示很多GPU在等待最慢的rank而浪费时间。本文报告的核心改进之一就是把平均ρ从1.34降到1.08。
ρ是贯穿全文的核心度量,是定义问题严重程度、触发不同机制(分段阈值k1=1.05、k2=1.10)和评估RLB效果的依据,理解ρ的定义才能理解为什么本文认为迁移20%的头部就能解决大部分问题。
研究动机
视频扩散Transformer在生成高分辨率长视频时,自注意力因对时空token序列的二次复杂度成为推理瓶颈。在单卡NVIDIA H20上,用Wan2.2-14B I2V生成5秒720p视频仍需约两小时,注意力占总推理时间的74.1%。工业界广泛采用的步数蒸馏虽可减少去噪步数、提供最高25倍加速,却不改变每步的序列长度和注意力开销,因此少步机制下注意力仍是主要瓶颈。作为互补优化,训练免的Top-p稀疏注意力能根据每个头的注意力分布自适应分配计算量、提升掩码保真度,但这种自适应性引入了新的系统瓶颈:Top-p路由让不同头保留数量各异的block,在Ulysses序列并行下多个高密度头若被映射到同一rank,会造成工作负载倾斜,使该rank在同步注意力阶段成为拖尾者,迫使其他rank完成后闲置等待,稀疏带来的理论计算节省被分布式执行开销部分抵消。
本文的目标是本文要解决的核心问题是:如何在保留自适应稀疏注意力保真度优势的同时,实现均衡且高效的多GPU分布式执行。具体目标是构建一个训练免的稀疏注意力系统FVAttn,在多GPU序列并行下提升自适应稀疏注意力的分布式执行效率,消除Top-p路由带来的头级工作量异构性所导致的rank级拖尾问题,从而同时改善负载均衡、注意力延迟、DiT延迟和视频质量指标,并在多个视频DiT工作负载上推动质量-效率Pareto前沿。作者还希望运行时机制的开销极小且可被计算重叠隐藏,使可见的额外开销控制在毫秒级。
与已有工作不同的是,本文的独特切入角度有三点。第一,关键观察:Top-p引入的负载倾斜高度局部化,绝大多数工作量偏差集中在少数重负载头上,因此消除注意力拖尾不需要昂贵的全局重分区或重新设计整个并行布局,只需在当前稀疏掩码物化后识别并迁移少数过载头即可恢复大部分关键路径效率。第二,时机选择:现有系统多在当前注意力计算前用历史稀疏性、离线profiling或近似相似性信号来选择头或block布局,但这在少步视频生成中非常脆弱(相邻步最大变化头级97%、rank级44%);而真正的当前稀疏掩码工作量只有在序列到头All-to-All完成后才可观测,FVAttn选择等待掩码物化后基于已观测的真实工作量做受限的P2P修复,而非预测。第三,迁移规模:默认8-GPU设置下每个rank只发送和接收最多一个本地头(约占其头集合的20%),即可把平均ρ从1.34降到1.08,进一步增加迁移预算收益递减。
核心方法
FVAttn的整体思路是'路由决定算什么,RLB决定在哪算,SASA把闲置算力转成更多掩码覆盖'。系统从一个系统瓶颈出发:Top-p路由提升掩码保真度,但其头级稀疏差异在Ulysses序列并行下变成rank级工作量不均,拖慢同步注意力阶段。因此FVAttn在当前稀疏掩码物化后执行两阶段运行时调度。完整流水线为:序列到头All-to-All后,每个稀疏注意力层依次执行(1)构建自适应稀疏掩码——稀疏路由前端沿Hilbert曲线重排block使空间相邻token落入邻近block,用池化Q/K算block级相似度,主要用Top-p的CDF阈值自适应决定保留block数,并用Top-k安全下限防止过度稀疏,同一路由还产出RLB所需的每头密度和SASA所需的key-block重要性排序;(2)RLB测量已实现的每头工作量并用受限P2P通信修复不均;(3)SASA把非关键rank的剩余空闲转化为额外掩码预算;(4)执行block稀疏注意力;(5)反向All-to-All并恢复头顺序。增益模型为 $G = K - P - C - A$,其中K为关键路径收益,P为通信、C为调度、A为增广开销。
核心创新点是在掩码物化后基于已观测真实每头工作量做受限P2P头部迁移,而非像现有方案在计算前用历史或近似信号做全局重分区。与预测调度(通信$P\approx 0$但$L_h$跨步剧烈变化使预测$K$缩水甚至为负)、或掩码后全局重分区(逼近完美均衡但需大规模搬移使$P$过大)相比,FVAttn走第三条路:在受限P2P空间内寻高增益调度。RLB把迁移表为rank上排列$\sigma$,每rank最多迁一个本地头(约20%),后均衡量 $L_r^{new}=L_r-L_{h_r}+L_{h_{\sigma^{-1}(r)}}$,在$S_{20\%}$内按词典序 $\sigma^*=\arg\min(\max_r L_r^{new},\,\mathrm{Var}_r(L_r^{new}))$ 搜索。SASA把RLB后非关键rank剩余空闲 $\Delta L_r=L_{max}^{new}-L_r^{new}$ 按$n_1=0.8$折为增广预算 $B_r=n_1\Delta L_r$,仅在$n_2=0.07$之上启用,把闲置算力变额外高价值block而不延关键路径。
方法步骤详情
方法分五步。步骤一(稀疏路由前端):block沿Hilbert曲线重排使空间相邻token落入邻近block;用池化Q/K算block级相似度,主要用Top-p的CDF阈值自适应定保留block数,辅以Top-k安全下限;同一路由顺带产出每头密度(供RLB)和key-block重要性排序(供SASA)。步骤二(RLB):各rank通过一次轻量all-gather收集全局密度;因信息相同且搜索确定,每GPU独立运行相同搜索提取收发伙伴,避免rank-0求解-广播;R=8下构建满足一收一发约束的配对拓扑枚举$S_{20\%}$选最优,更大规模用单向环。步骤三(SASA):按 $B_r=n_1\Delta L_r$($n_1=0.8$)分配预算,按路由重要性顺序追加未选且排除迁出头的高价值block,就地更新掩码。步骤四(分段):$\rho<k_1{=}1.05$跳过;$k_1\le\rho<k_2{=}1.10$仅SASA;$\rho\ge k_2$先RLB后SASA。步骤五(重叠):CPU-GPU重叠用V量化窗口吸收CPU搜索;计算-通信重叠把P2P迁移与未迁头V量化和SASA增广重叠。
技术新颖性
技术新颖性体现在四个层面。第一,问题识别:作者首次系统指出少步视频DiT分布式推理中自适应稀疏注意力的一个运行时系统瓶颈——提升保真度的Top-p路由引发头级工作量异构,在序列并行下变成rank级拖尾。第二,RLB设计:这是掩码后运行时修复机制,用掩码物化后的已实现每头工作量做通信预算受限的P2P头部迁移以缩短当前注意力关键路径,区别于掩码前预测布局或全局重分区。第三,SASA设计:它把非关键rank的剩余空闲转化为额外高价值稀疏block,改进掩码覆盖却不延长关键路径——不等同于全局提高Top-p阈值(那会抬高所有rank负载并延长makespan),而是仅在有空的rank上追加block,把同步闲置变成有用计算。第四,统一增益模型 G = K − P − C − A 把算法目标(K大)和系统目标(P、C、A尽量不可见)统一在一个框架内,并通过两条异步重叠路径把可见额外开销压到约0.7ms(1.87%),为这类运行时调度提供了可分析的收益分解范式。
实验结果
实验在8×NVIDIA H20(NVLink,CUDA 13.1)评估Wan2.2 I2V/Animate与Wan2.1 T2V(4步蒸馏)。整体上FVAttn在I2V/Animate/T2V三个工作负载改善质量-效率Pareto前沿,展现'同速更高质、同质更高速':I2V的Top-p=0.95与SpargeAttention同速但四项保真指标全面领先,Top-p=0.90则更快且VBench更高。注意力级相对FlashAttention的完整栈达37.54ms、4.41×加速,平均负载不均$\rho$从1.34降到1.08。机制级证据对应三个设计目标:RLB降关键路径负载(45.91→41.32ms、$\rho$ 1.34→1.08,且优于db-SP);SASA填空闲为有用block(PSNR +0.458dB、SSIM +1.7%、LPIPS −13.9%,延迟仅增0.03s验证 $B_r\le\Delta L_r$);Overlap/Opt隐藏开销(41.33→37.54ms、4.00→4.41×)。完整栈额外可见开销仅0.7ms(1.87%);Animate、T2V趋势一致。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Wan2.2-14B I2V 4步蒸馏端到端(质量导向Top-p=0.95) | DiT延迟 / PSNR | 19.10s / 23.801(VBench 88.8%) | SpargeAttention Top-p=0.95:19.92s / 22.400(VBench 88.3%) | 同速下PSNR +1.404dB,VBench +0.5% |
| Wan2.2-14B I2V 4步蒸馏端到端(速度导向Top-p=0.90) | DiT延迟 / VBench | 18.30s / 88.8%(2.11×加速) | SpargeAttention Top-p=0.90:19.23s / 88.2%(2.01×加速) | 延迟降0.93s,VBench +0.6%,质量与速度双升 |
| 注意力运行时(I2V) | 注意力延迟 / 加速 / 均衡ρ | 37.54ms / 4.41× / 1.01(完整栈) | FlashAttention:165.60ms / 1.00× / 1.00 | 延迟降128.06ms,加速4.41× |
| Wan2.2 Animate 4步蒸馏 | DiT加速 / PSNR | FVAttn Top-p=0.90:2.50× / 21.512 | Jenga:1.66× / 21.534;SpargeAttention Top-k=0.60:2.35× / 21.188 | 更高加速,PSNR与Jenga相当 |
| 运行时机制开销(I2V,Top-p=0.95+RLB+SASA+Overlap/Opt) | 可见额外开销占比 | 0.7ms / 1.87% | 独立串行开销2.6ms / 6.28% | 重叠后可见开销降至1.87% |
局限与改进
作者承认的局限主要有三方面。其一,任务适用性:FVAttn主要针对长时空序列视频生成,对于短序列或注意力稀疏性弱的任务(如许多图像生成),稀疏注意力的节省以及负载均衡/空闲复用收益都会相应下降。其二,硬件依赖:收益取决于硬件的通信-计算比,在H20这类通信与内存带宽充足的服务器上,轻量P2P迁移和重叠能有效隐藏开销;但在PCIe连接、算力更强但卡间通信较弱的设备上,头部迁移收益可能受限,迁移预算、SASA触发阈值和重叠策略都需要重新校准。其三,并行与设定范围:本文主要在Ulysses风格序列并行和少步视频DiT推理下验证,Ring Attention、USP、更多样硬件拓扑和训练时稀疏注意力的适用性留待未来。从外部观察,我还注意到:实验全部基于腾讯自家的Wan系列模型,缺乏对开源非Wan模型的泛化验证;默认仅迁移20%头部(一收一发约束)虽高效,但对极高不均场景可能不是最优;R=8时的匹配枚举策略对更大规模GPU的扩展性表述较概括;此外Top-p阈值的选定(0.90/0.95)和分段阈值k1=1.05、k2=1.10在不同模型/硬件间需要经验调参,自动化程度有限。
独立分析的弱点
独立分析有以下几个弱点。第一,泛化性待验证:所有实验都基于Wan2.2 I2V、Wan2.2 Animate和Wan2.1 T2V,未在如HunyuanVideo、CogVideoX等异构架构视频DiT上验证,Top-p路由的负载局部性假设是否在所有视频模型成立存疑——改进方向是用更多模型族做泛化基准。第二,超参敏感性:分段阈值k1=1.05/k2=1.10、SASA系数n1=0.8/n2=0.07、Top-p值都靠经验设定,作者也承认需按硬件重校准;可引入自适应或学习式超参策略,例如根据在线测量的ρ分布自动调整触发阈值。第三,迁移约束偏保守:一收一发(每rank最多迁一个头)在极端不均时可能留有大量未利用的修复空间;可设计预算自适应的多头迁移,在通信成本与均衡收益间动态权衡。第四,规模扩展性:R=8用配对枚举,更大规模改用单向环,但环拓扑的迁移路径较长可能拖高P;可结合拓扑感知(如NVLink子组)做分层迁移。第五,单靠消融表4无法完全解释41.33→37.54ms的下降,因为Overlap/Opt还移除了底层注意力路径的同步与执行低效——这部分应更透明地拆解。
未来方向
作者明确提出的未来方向包括:将FVAttn扩展到Ring Attention和USP等更通用的序列并行策略;探索在更多样硬件拓扑下的适用性;以及训练时稀疏注意力的应用。基于成果可延伸的方向有:第一,把运行时修复思想推广到其他动态路由场景,如MoE专家负载均衡——MoE同样存在专家分配不均导致的token丢弃或通信瓶颈,RLB式的掩码后P2P修复可借鉴。第二,将SASA的'把同步空闲转成有用计算'思想推广到任意存在同步气泡的并行训练/推理,例如流水并行的bubble filling。第三,增益模型 G = K − P − C − A 可作为通用运行时调度的分析框架,用于指导其他通信-计算重叠系统的设计。第四,结合在线学习让阈值(k1/k2、n1/n2)随工作负载自适应演化,减少人工校准。第五,与稠密量化注意力(如SageAttention)正交组合,探索稀疏+量化的叠加收益。
复现评估
复现评估如下。代码与开源:论文未提供代码仓库链接或开源承诺,FVAttn作为系统工作,P2P迁移、block稀疏内核和重叠调度涉及底层CUDA/NCCL实现,复现门槛较高。硬件:实验全部在8×NVIDIA H20(NVLink、CUDA 13.1)上,H20属受限出口高端数据中心GPU,普通机构难获取,且作者明确PCIe设备收益可能下降、换硬件需重校准超参。模型与数据:基于Wan2.2-14B I2V/Animate和Wan2.1 T2V,用LightX2V 4步蒸馏LoRA;评测用VBench官方协议(T2V 946例、I2V 1118例),输入与推理设置各方法一致,利于公平比较。关键超参(k1=1.05、k2=1.10、n1=0.8、n2=0.07、迁移20%头部、Top-p=0.90/0.95、前25%步走密集)都已给出,算法行为理论可复现。总体算法层可复现性中等,但完整系统级复现(含NCCL P2P、kernel overlap)难度高、硬件门槛大,强烈建议作者开源参考实现以降低社区成本。
论文图表
该图展示FVAttn的整体定位:将高效稀疏注意力与运行时负载均衡结合,在保持视觉质量的同时加速视频生成。它直观呈现了FVAttn如何在视频生成的去噪过程中,把稀疏路由、负载均衡和质量保持三者整合在一起,为读者建立'系统既快又保质'的第一印象。
作为论文封面图,它概括了全文的核心主张和质量-效率兼得的设计哲学,是理解FVAttn价值主张的入口。
该图量化了4步蒸馏视频生成中相邻步的工作负载剧烈波动:在相同头密度下最大相邻步变化达97%,在相同rank负载下达44%。这组数据直接说明了为什么基于历史profiling的预测式负载布局在少步生成中会失效——历史信号与当前步的真实负载偏差极大。
这张图是FVAttn选择'掩码物化后基于当前观测负载修复'而非'预测未来布局'的核心实证依据,是理解整个方法动机的关键证据。
该表给出Wan2.2-14B I2V和Wan2.1-14B T2V 4步蒸馏的端到端结果,包含VBench、PSNR、SSIM、LPIPS、CLIP-Sim、DiT延迟和相对FlashAttention的加速。I2V上FVAttn(Top-p=0.95)DiT 19.10s、PSNR 23.801、VBench 88.8%、加速2.02×,对比SpargeAttention(Top-p=0.95)的19.92s/22.400/88.3%/1.94×展现同速更高质;Top-p=0.90时18.30s/23.473/88.8%/2.11×实现同质更高速。表中还包含受控的SpargeAttention+db-SP/+RLB+SASA行以隔离运行时调度效果。T2V上FVAttn(Top-p=0.90)2.32×加速、VBench 81.6%。
这是论文最核心的端到端结果表,证明FVAttn在多个工作负载上持续改善质量-效率Pareto前沿,是评估方法整体价值的主表。
该表分解Top-p=0.95+RLB+SASA+Overlap/Opt配置下各运行时组件的开销:密度交换、平衡计划搜索、P2P头部迁移、空闲增广、头顺序恢复及其他(元数据/内核启动)。独立串行总开销约2.6ms(6.28%),经Overlap/Opt后可见总开销仅约0.7ms(1.87%)。其中SASA保留最大可见开销(~0.2ms,0.53%),因为部分掩码更新仍在关键路径上。参考注意力延迟从41.4ms降到37.5ms。
这张表把增益模型G=K−P−C−A中的P、C、A各项具体量化,并证明重叠能把可见开销压到1.87%,是理解FVAttn为何'额外机制几乎免费'的系统级证据,对评估方法实用价值至关重要。