Fast Video Generation with Sliding Tile Attention
TL;DR 精炼摘要
本研究提出滑动平铺注意力(STA)以解决视频生成中的计算瓶颈。STA通过局部滑动窗口注意力的设计,显著减少了计算冗余,实现了58.79%的模型FLOPs利用率。与现有技术相比,STA在不降低生成质量的前提下,将视频生成延迟降低至501秒,展示了其卓越的效率。
摘要
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. We make our codebase public at https://github.com/hao-ai-lab/FastVideo.
思维导图
论文精读
中文精读
1. 论文基本信息
1.1. 标题
中文标题 (English Title): 使用滑动平铺注意力实现快速视频生成 (Fast Video Generation with Sliding Tile Attention)
1.2. 作者
Peiyuan Zhang, Yongqi Chen, Runlong Su, Hangliang Ding, Ion Stoica, Zhengzhong Liu, Hao Zhang 等。 隶属机构: 论文的第一作者 Peiyuan Zhang 和通讯作者 Hao Zhang 隶属于某实验室(可能是论文末尾致谢中提到的 Hao AI Lab),其他作者来自不同机构,例如加州大学圣地亚哥分校 (UCSD) 的 HDSI (Halıcıoğlu Data Science Institute) 等。
1.3. 发表期刊/会议
该论文以预印本 (Preprint) 形式发布于 arXiv,并预计于 2025-02-06T21:17:09.000Z 正式发布。
1.4. 发表年份
2025年(根据原文链接的发布时间判断)。
1.5. 摘要
研究目的: 最先进的视频生成 Diffusion Transformers (DiTs) 采用 3D 全注意力 (3D full attention),计算成本极高。例如,生成一个 5 秒的 720P 视频,仅注意力计算就占总推理时间的 800 秒(总时长 945 秒)。本文旨在解决这一计算瓶颈。
核心方法: 论文引入了 滑动平铺注意力 (Sliding Tile Attention, STA)。STA 利用了预训练视频扩散模型中注意力分数主要集中在局部 3D 窗口内的观察。通过在局部时空区域滑动并执行注意力计算,STA 消除了全注意力 (full attention) 中的冗余。与传统的基于 token 的滑动窗口注意力 (sliding window attention, SWA) 不同,STA 采用新颖的硬件感知滑动窗口设计,以 tile 为单位进行操作,从而在保持表达能力的同时提高硬件效率。
主要结果: 经过精心的 kernel-level 优化,STA 首次实现了高效的 2D/3D 类滑动窗口注意力实现,达到了 58.79% 的模型 FLOPs 利用率 (Model FLOPs Utilization, MFU)。具体而言,STA 相较于 FlashAttention-2 (FA2) 提速 2.8-17 倍,相较于 FlashAttention-3 (FA3) 提速 1.6-10 倍。在领先的视频 DiT 模型 HunyuanVideo 上,STA 在不降低生成质量的情况下,将端到端延迟从 945 秒 (FA3) 降低到 501 秒,且无需额外训练。通过微调 (finetuning),延迟可进一步降低至 268 秒,而 VBench 分数仅下降 0.09%。
关键结论: STA 有效解决了视频 DiTs 中的注意力计算瓶颈,显著提高了视频生成效率,同时保持了高质量输出。
1.6. 原文链接
原文链接: https://arxiv.org/abs/2502.04507 PDF 链接: https://arxiv.org/pdf/2502.04507v3.pdf 发布状态: 预印本 (Preprint),预计于 2025-02-06T21:17:09.000Z 正式发布。
2. 整体概括
2.1. 研究背景与动机
核心问题: 视频生成领域的 Diffusion Transformers (DiTs) 架构,特别是那些采用 3D 全注意力 (3D full attention) 机制的模型,在生成高分辨率、长持续时间视频时面临巨大的计算成本。注意力机制的二次复杂度 () 导致推理速度极慢,严重限制了其实际应用和部署。例如,生成一个 5 秒 720p 的视频可能需要长达 16 分钟,其中大部分时间(800 秒)消耗在注意力计算上,即使使用了 FlashAttention 3 (FA3) 和 H100 GPU 也是如此。
重要性与挑战: 视频数据本身具有高度冗余性,相邻帧之间差异微小,空间上接近的像素关联性更强。当前的 3D 全注意力机制不加区分地让每个 token (词元) 关注所有其他 token,这在处理视频数据时可能导致大量不必要的计算。现有的滑动窗口注意力 (Sliding Window Attention, SWA) 方法虽然旨在利用这种局部性来减少计算量,但由于其在 2D/3D 数据中生成不规则的注意力掩码 (attention mask),导致硬件利用率低下,无法将浮点运算 (FLOPs) 的减少有效转化为实际的运行时加速。这表明,在 kernel-level 上实现高效的 2D/3D SWA 仍然是一个未解决的挑战。
论文切入点与创新思路: 本文假设预训练视频扩散模型中的 3D 全注意力机制携带了大量冗余,这些冗余可以通过利用注意力分数的 3D 局部性 (3D locality) 模式来消除。论文通过可视化和量化发现,注意力分数确实主要集中在局部时空窗口内。在此基础上,论文提出了 滑动平铺注意力 (Sliding Tile Attention, STA),这是一种硬件感知 (hardware-aware) 的注意力机制,它重新思考了滑动窗口计算方式。STA 以 tile (平铺块) 为单位进行操作,而非传统的 token (词元) 单位,并采用新颖的硬件感知滑动窗口设计,旨在生成更规则的注意力模式,从而提高硬件利用率,并有效将稀疏性带来的 FLOPs 减少转化为实际的运行时加速。
2.2. 核心贡献/主要发现
论文的主要贡献包括:
- 识别并量化 3D 局部性和头部特化 (head specialization): 作者识别并量化了最先进的视频 DiTs 中 3D 注意力分数普遍存在的 3D 局部性模式,以及不同注意力头部在局部性模式上的差异(头部特化),揭示了全 3D 注意力中的大量冗余。
- 引入滑动平铺注意力 (STA): 提出了一种基于
tile的滑动窗口注意力机制。该方法通过系统与算法协同设计,将滑动单元从单个token变为一组连续token组成的tile,确保所有查询tile关注相同的键tile集合,从而生成仅包含密集块 (dense blocks) 和空块 (empty blocks) 的注意力图,极大减少了传统 SWA 中低效的混合块 (mixed blocks)。 - 实现优化的内核 (optimized kernel): STA 的优化内核实现了类滑动窗口注意力机制的硬件高效实现,相比 FlashAttention 3 (FA3) 仅引入最小的开销,达到了 58.79% 的模型 FLOPs 利用率 (MFU),首次实现了 2D/3D 类滑动窗口注意力机制将 FLOPs 减少转化为成比例的实际加速。
- 显著加速视频生成: STA 在不降低或仅降低极小质量的情况下,将注意力计算加速超过 10 倍,端到端视频生成速度最高提升 3.53 倍。这包括无需训练 (training-free) 的即插即用 (plug-and-play) 加速,以及通过微调 (finetuning) 实现的更激进的稀疏化和更高效率。
3. 预备知识与相关工作
3.1. 基础概念
为了更好地理解本文,需要了解以下核心概念:
- Diffusion Transformers (DiTs):
Diffusion Transformers(DiTs) 是一种将Transformer架构应用于扩散模型 (Diffusion Models) 的方法。在视频生成领域,DiTs 通过建模从噪声数据到真实数据的逐步去噪过程来生成视频。它们通常将视频帧展平为一系列token(词元),然后使用Transformer捕获这些token之间的时空依赖关系。 - 3D 全注意力 (3D Full Attention): 在视频 DiTs 中,
3D 全注意力机制允许视频中的每个token关注所有其他token,无论它们在空间还是时间上相距多远。这对于捕获复杂的长距离时空依赖关系至关重要,但其计算复杂度随token数量的平方 () 增长,导致高分辨率和长视频的计算成本极高。 - FlashAttention:
FlashAttention是一种高效的注意力机制实现,通过在GPU片上存储器 (SRAM) 上进行tiling(分块) 计算和在线softmax,显著减少了GPU高带宽内存 (HBM) 的访问次数,从而加速了注意力计算并降低了内存消耗。它避免了完整注意力矩阵 和 的实例化。 - 滑动窗口注意力 (Sliding Window Attention, SWA):
滑动窗口注意力是一种稀疏注意力 (sparse attention) 方法,其中每个查询token仅关注其周围固定大小的局部窗口内的键token。这种方法旨在利用数据中的局部性,减少不必要的全局注意力计算,从而降低计算复杂度。通过堆叠多层注意力,感受野 (receptive field) 可以逐步扩大。 - 模型 FLOPs 利用率 (Model FLOPs Utilization, MFU):
MFU是衡量GPU实际计算效率的指标,定义为实际达到的浮点运算速度 (FLOPs/second) 与GPU理论峰值浮点运算速度的比率。高 MFU 意味着GPU的计算资源得到了更充分的利用。
3.2. 前人工作
本文的工作建立在 Transformer 架构和高效注意力实现的基础上,并与多种稀疏注意力及视频生成加速方法进行了比较。
- FlashAttention 家族 (FA2, FA3):
- FlashAttention-2 (FA2) 和 FlashAttention-3 (FA3) 是 FlashAttention 的后续版本,进一步优化了并行性和工作划分,提高了异步处理能力和低精度计算支持,从而实现了更高的计算效率。它们是全注意力机制目前最先进且高效的实现。
- 传统滑动窗口注意力 (SWA) 实现:
- NATTEN (Neighborhood Attention Transformer):一种针对邻域注意力设计的
Transformer,其kernel(内核) 旨在提高 SWA 在 2D/3D 数据上的效率,但仍面临混合块 (mixed blocks) 导致的硬件利用率问题。 - Tiled NATTEN: NATTEN 的优化版本,尝试通过输入
tiling(分块) 和kernel fusion(内核融合) 来提升效率,但仍未能完全解决不规则注意力掩码的问题。 - CLEAR (Conv-like Linearization Revs Pre-trained Diffusion Transformers Up): 通过将全注意力替换为基于圆形窗口的注意力机制,将计算限制在局部窗口内,旨在实现线性注意力。
- Swin Transformer: 引入了基于移位窗口 (shifted window) 的分层视觉
Transformer。它将图像划分为不重叠的窗口进行局部注意力计算,并通过交替移位窗口来引入跨窗口连接。Swin 注意力并非滑动窗口注意力,它在每个注意力层内打破了局部连接性,但通过多层堆叠也能扩大感受野。
- NATTEN (Neighborhood Attention Transformer):一种针对邻域注意力设计的
- 缓存 (Caching) 方法:
- -DiT (Training-free acceleration method tailored for Diffusion Transformers): 通过缓存特征偏移 (feature offsets) 而非完整的特征图来优化推理速度。它采用分阶段缓存策略,存储后期
DiT块的残差 (residuals) 以供早期采样,并缓存早期块的残差以供后期采样。
- -DiT (Training-free acceleration method tailored for Diffusion Transformers): 通过缓存特征偏移 (feature offsets) 而非完整的特征图来优化推理速度。它采用分阶段缓存策略,存储后期
- 分解 3D 注意力:
- 一些方法尝试将 3D 注意力分解为交替进行的空间 (spatial) 和时间 (temporal) 组分(例如 Make-A-Video),以降低计算复杂度。然而,这些方法已被全 3D 注意力在性能上超越,因为它们难以捕获同时在空间和时间维度上偏移的
token之间的相互作用,破坏了 3D 局部性模式。
- 一些方法尝试将 3D 注意力分解为交替进行的空间 (spatial) 和时间 (temporal) 组分(例如 Make-A-Video),以降低计算复杂度。然而,这些方法已被全 3D 注意力在性能上超越,因为它们难以捕获同时在空间和时间维度上偏移的
3.3. 技术演进
视频生成 Transformer 中的注意力机制从最初的全局全连接 (full-connection) 模式,演进到通过 FlashAttention 系列优化内存和计算效率,再到引入各种稀疏化策略(如 SWA、Swin、CLEAR)试图利用数据局部性降低理论计算量。然而,在 2D/3D 场景下,将理论上的 FLOPs 减少转化为实际的 wall-clock (挂钟时间) 加速,仍然是一个挑战。主要原因是传统稀疏注意力 kernel (内核) 会生成不规则的注意力掩码 (attention mask),导致 GPU 硬件利用率低下(出现大量 mixed blocks)。本文的 SLIDING TILE ATTENTION (STA) 正是针对这一 GPU 不友好特性,通过系统与算法协同设计,重新定义滑动窗口的计算模式,使其在硬件层面更加高效。
3.4. 差异化分析
SLIDING TILE ATTENTION (STA) 与上述相关工作的主要区别和创新点在于:
- 与传统 SWA 的区别: 传统的 SWA 关注每个
token的局部窗口,导致不同query token关注不同的key token组,生成大量不规则的mixed blocks。而 STA 将滑动单位从token升级为tile(平铺块)。在一个tile内的所有query都关注相同的key tile集合,从而确保注意力计算始终在密集的FlashAttention块上进行,消除了mixed blocks,极大地提高了GPU硬件效率。 - 与 FlashAttention 的关系: STA 建立在
FlashAttention的基础上,利用其tiling和在线softmax机制。但 STA 更进一步,通过其tile-by-tile的设计,在稀疏注意力场景下,能够将FlashAttention的高效计算模式扩展到 3D 局部窗口,并规避掉mask(掩码) 计算的额外开销。 - 与分解 3D 注意力的区别: 传统的空间/时间分解注意力打破了 3D 局部性,无法捕获时空同时偏移的依赖关系,导致性能下降。STA 通过在 3D 时空立方体内进行滑动,完全保留了 3D 局部性,同时实现了稀疏化,因此能够维持高质量。
- 与 -DiT 的区别: -DiT 是一种基于缓存的训练无关加速方法,主要通过重用特征图来减少计算。STA 则是从注意力机制本身的设计和
kernel(内核) 实现层面进行优化。两者是互补的,未来有望结合使用以获得更大加速。 - 硬件感知设计: STA 的核心优势在于其
hardware-aware(硬件感知) 设计。它将inter-block mask(块间掩码) 逻辑从计算kernel(内核) 中分离出来,交由异步的数据warpgroup(线程组) 处理,使得计算warpgroup始终执行密集的、硬件高效的注意力计算,从而实现了FLOPs减少与wall-clock加速的直接转化。这是现有 2D/3D SWA 实现未能有效解决的问题。
4. 方法论
4.1. 方法原理
SLIDING TILE ATTENTION (STA) 的核心思想是利用视频扩散模型中 3D 注意力分数的局部集中性(即查询 token 倾向于关注其时空邻域内的键 token),并通过一种硬件感知的设计来高效地利用这种局部性。
传统的滑动窗口注意力 (SWA) 在每个 token (词元) 级别上滑动窗口,导致每个 query token 关注的 key token 集合不同,从而在注意力图中产生高度不规则的模式,包含大量低效的“混合块” (mixed blocks)。混合块 指的是一部分注意力分数被保留,一部分被掩盖的块,它们仍然需要完整的计算,并带来额外的掩码计算开销,使得 FLOPs 减少无法有效转化为实际的 wall-clock 加速。
STA 提出了一个全新的滑动窗口计算范式,即以 tile (平铺块) 为单位进行操作。一个 tile 被定义为一组连续的 token,它们在原始 3D 视频空间中形成一个时空立方体。STA 的关键直觉是:如果一个 query tile 中的所有 token 都关注相同的 key tile 集合,那么注意力计算就可以始终在密集的 FlashAttention 块上进行,从而完全消除 混合块,只生成 dense blocks (密集块) 和 empty blocks (空块)。空块 可以被完全跳过,而 密集块 的计算效率最高。这种设计显著改善了注意力计算模式的规则性,使其对 GPU (图形处理单元) 更加友好,从而将理论上的 FLOPs 减少有效地转化为实际的加速。
4.2. 核心方法详解
4.2.1. 滑动平铺注意力 (STA) 机制
STA 机制的核心在于将传统的 token-wise (逐词元) 滑动窗口操作转换为 tile-wise (逐平铺块) 操作,以生成更规则、更高效的注意力模式。
给定一个维度为 (L, L, L) 的视频隐变量 (video latent),以及 FlashAttention 的块大小为 (B, B),STA 设置 tile 的大小 使得 。这表示每个 FlashAttention 块将处理一个完整的 3D tile。为了简化,STA 进一步假设视频尺寸 和窗口尺寸 都是 的整数倍。
视频首先被划分为不重叠的、大小为 (T, T, T) 的 tile。这些 tile 中的 token 被展平为 1D 序列,其顺序保证了同一 tile 内的 token 具有连续的序列索引。这种 token 排序方式(如原文 Figure 8 所示)对于维持 FlashAttention 的高效性至关重要,因为它确保了当一个 tile 关注另一个 tile 时,相应的注意力图会形成一个 dense block (密集块),因为所有参与的序列索引都是连续的。
概念上,STA 以 (T, T, T) 的步长滑动其窗口。对于每一步,它计算中心 query tile 与该窗口内所有 key tile 之间的注意力。由于窗口内包含 个 key tile,这会产生相同数量的 dense attention blocks,而不会产生任何 mixed blocks。
为了定量比较 STA 与 Tiled NATTEN 在生成不同类型块上的优越性,作者提供了以下定理:
定理 3.1 (Tiled NATTEN 的块计数):
考虑一个 Tiled NATTEN 配置,其 tile 大小为 (T, T, T),窗口大小为 (W, W, W),视频大小为 (L, L, L)。设 FA 块大小为 (B, B),其中 。忽略边界效应,dense blocks 的数量由下式给出:
Tiled NATTEN 中 mixed blocks 的数量为:
符号解释:
-
:
Tiled NATTEN中dense blocks(密集块) 的数量。 -
:
Tiled NATTEN中mixed blocks(混合块) 的数量。 -
:视频的每个维度大小(假设是立方体视频,即 )。
-
:
tile的每个维度大小(假设是立方体tile,即 )。 -
:滑动窗口的每个维度大小(假设是立方体窗口,即 )。
-
:向下取整函数。
-
:向上取整函数。
-
max(a, b):取 和 中的最大值。直观解释:
Tiled NATTEN中,只有当窗口大小至少是tile大小的两倍时,一个块才可能完全是dense的。而一个query tile中的最左侧query仍然可以关注到距离它W-1tile远的key,这会产生mixed blocks。
定理 3.2 (SLIDING TILE ATTENTION 的块计数):
在相同符号表示下,如果 是 的整数倍,SLIDING TILE ATTENTION 中 dense blocks 的数量为:
所有剩余的块都是 empty blocks (空块),并且没有 mixed blocks。
符号解释:
-
:
SLIDING TILE ATTENTION中dense blocks(密集块) 的数量。 -
其他符号同定理 3.1。
直观解释: 在
STA中,每个query tile只关注其局部窗口内的key tile。由于窗口内有 个key tile,因此会创建相同数量的dense blocks。所有超出此窗口范围的块都是empty blocks,可以被完全跳过,从而保证了计算效率。
以下是原文 Table 1 所示的两种方法在不同块类型比例上的对比,清晰地展示了 STA 消除 mixed blocks 的能力:
以下是原文 Table 1 的结果:
| Attention | Window Size | Dense Block | Mixed Block |
| Tiled NATTEN | (11,11,11) | 0.06% | 7.17% |
| STA | (12, 12, 12) | 1.56% | 0.0% |
| STA | (20, 20, 20) | 7.23% | 0.0% |
图像 8. 2D SLIDING TILE ATTENTION 可视化
下图(原文 Figure 5)展示了 2D SLIDING TILE ATTENTION 的工作原理。假设 tile 大小为 (2, 2),窗口大小为 (6, 6)。每个查询 tile 在关注所有键 tile 后,将生成九个 4x4 的 dense blocks 在注意力图中。这说明了 tile-by-tile 操作如何生成密集的注意力块。
该图像是一个示意图,左侧展示了传统的锯齿状展平策略,右侧展示了STA的序列展平策略。该图假定图像大小为(9, 9),瓷砖大小为(3, 3)。
Figure 5. 2D SLIDING TILE ATTENTION with tile size (2, 2) and window size (6, 6). After attending to all the key tiles, each query tile will generate nine 4x4 dense blocks in the attention map. We showcase 2D STA for better illustration. 3D STA can be inferred similarly.
图像 11. 序列展平策略对比
下图(原文 Figure 8)对比了传统的锯齿状展平策略与 STA 的序列展平策略。STA 策略确保了 tile 内的 token 具有连续的序列 ID,从而在 FlashAttention 中形成 dense blocks。
该图像是展示不同视频生成方法所需时间的对比图,包含四个视频生成模型的结果,分别为 HunyuanVideo、STA-tf-1.89x、STA-t-2.43x 和 Δ-DiT-1.36x,以及对应生成每个视频所需的时间。
Figure 8. Left: Conventional zigzag flattening strategy. Right: STA' sequence flattening strategy. The plot is given assuming a (9, 9) image with (3, 3) tile size.
图像 12. 2D 滑动窗口注意力可视化
下图(原文 Figure 9)展示了 2D 滑动窗口注意力 (SWA) 的可视化。SWA 是逐 token 滑动的,每个 query token 关注其局部窗口内的 key token。
该图像是一个对比图,展示了不同视频生成方法的生成时间。上方为 Hunyuan 视频(15分钟45秒),下方为使用滑动瓷砖注意力(STA)技术的结果,其中 STA-tf-1.89x 用时8分钟21秒,STA-t-2.43x 用时6分钟29秒,Δ-DiT-1.36x 用时11分钟34秒。
Figure 9. 2D Sliding Window Attention visualization.
图像 7. 注意力图对比
下图(原文 Figure 4)直观展示了 NATTEN、Tiled NATTEN 和 STA 的注意力图差异。NATTEN 生成大量 mixed blocks,效率低下;Tiled NATTEN 稍好,但仍有 mixed blocks;而 STA 主要生成 dense blocks 和 empty blocks,计算模式更规整,硬件效率更高。
该图像是一个条形图,展示了STA与HunyuanVideo及Δ-DiT在200个提示下的人工评估结果。结果显示,STA-tf-1.89x相比HunyuanVideo的胜率为5.0%,平局83.0%,失利12.0%;而与Δ-DiT-1.36x比较,胜率66.5%,平局23.5%,失利10.0%。
Figure 4. The attention map of NATTEN, Tiled NATTEN, and STA. We plot with an image size and a local window. The tile size is set to (a) NATTEN creates many mixed blocks that are very inefficient for Flash Attention computation.(b) Tiled NA but for better illustration, we present the 2D scenario in this plot.
4.2.2. 内核级优化 (Kernel-level optimization)
STA 的高效实现依赖于 kernel-level (内核级) 的优化,特别是借鉴了 FlexAttention、ThunderKittens 和 FlashAttention-3 的思想。
- FlexAttention 集成:
STA可以高效地与FlexAttention结合,FlexAttention提供了跳过empty blocks(空块) 的功能,并避免在dense blocks(密集块) 上添加不必要的块内掩码 (intra-block masks)。 - 生产者-消费者范式 (Consumer-Producer Paradigm): 作者基于
ThunderKittens和FlashAttention-3实现了注意力kernel(内核)。这种实现将threadblock(线程块) 分为compute warpgroups(计算线程组) 和data warpgroups(数据线程组),形成了生产者-消费者范式:-
Data warpgroups(生产者): 负责异步地将KV blocks(键值块) 从HBM(高带宽内存) 加载到SRAM(片上存储器)。对于每个查询块,data warpgroup根据STA的稀疏注意力模式计算inter-block mask(块间掩码),并决定需要加载哪些键值块。由于这是异步操作,计算块间掩码和数据加载的开销可以被隐藏起来,与计算warpgroup的注意力计算重叠。 -
Compute warpgroups(消费者): 负责计算注意力。每个计算warpgroup负责计算一个查询块,该查询块始终驻留在SRAM中(遵循Split-Q设计)。计算warpgroup完全不感知稀疏注意力模式,它只对data warpgroup加载到共享内存中的key value blocks执行密集的注意力计算。一旦循环缓存中的所有数据都被消耗,计算就完成。这种设计使得计算
worker(工作者) 能够完全专注于密集的注意力计算,而将稀疏模式的复杂性转移给异步的数据worker处理,从而实现了极高的GPU利用率和效率。
-
4.2.3. 将 STA 应用于视频扩散模型 (Applying STA to Video Diffusion Model)
STA 可以以两种方式应用于视频扩散模型:无需训练的即插即用,或通过少量微调以实现更激进的稀疏化。
4.2.3.1. 无需训练 (Training-free)
如原文 Figure 3 所示,视频扩散模型展现出明显的 3D 局部性 (3D locality) 和 头部特化 (head specialization) 模式。不同 Transformer 头部具有不同程度的局部性,但这种模式在不同 prompt (提示) 之间大致保持一致。STA 利用这一特性,通过在少量 prompt 上搜索每个头部的最优窗口大小,期望该配置在其他 prompt 上也能表现良好。
图像 6. 注意力局部性与头部特化
下图(原文 Figure 3)左侧展示了在不同扩散步长和 10 个不同 prompt 下,注意力分数在局部窗口内的集中程度。大多数头部显示出高回收率 (high recall),表明存在局部注意力模式。右侧显示了尽管不同头部的回收率不同,但它们在不同 prompt 上的标准差较低,这支持了 头部特化 跨 prompt 一致的观察。
该图像是对比图,展示了不同视频生成模型的生成时间表现。上部分为 HunyuanVideo,其生成时间为 15 分 45 秒;下部分展示了 STA-tf-1.89x、STA-t-2.43x 和 Δ-DiT-1.36x 的生成时间,分别为 8 分 21 秒、6 分 29 秒和 11 分 34 秒。此对比突显了滑动瓷砖注意力(STA)在加速视频生成方面的优势。
Figure 3. Left: Fraction of attention scores within a (12, 24, 24) local window across diffusion steps and 10 different prompts. Most heads show high recall, indicating a local attention pattern. Right: Despite the different recall across heads, the standard deviation across prompts remains low.
算法 1:STA 掩码搜索 (STA Mask Search)
为了找到最优配置,作者开发了一个简单的启发式算法,并通过对 16 个 prompt 的 mask-search loss (掩码搜索损失) 求平均来决定最终配置。
在实际应用中,STA 会在初始 个 timestep (时间步) 保持全注意力 (full attention),然后在剩余的 timestep 应用 STA。
Algorithm 1 STA Mask Search
Input: Transformer model M, Total steps T, Mask pattern list P, Keep first T0 timestep full attn
Output: Dictionary dict that stores selected mask pattern for each head
Initialize dict
for t = T0 + 1 to T do
for each layer head combination (l, h) in M do
O (attn output of original (l, h))
Initialize minimum_loss ← ∞
Initialize best_pattern ← null
for each p in P do
Mask head h for layer l using mask pattern p
O' (attn output of M after masking)
loss ← MSE(O, O')
if loss < minimum_loss then
minimum_loss ← loss
best_pattern ← p
Record best_pattern for (t, l, h) in dict
return dict
算法解释:
- 输入: 一个
Transformer模型 ,总的扩散步数 ,一个预定义的掩码模式列表 ,以及一个指示在前 个时间步保持全注意力的参数。 - 输出: 一个字典
dict,存储了在每个时间步 、每个层 和每个头部 组合下选择的最佳掩码模式。 - 流程:
-
初始化一个空字典
dict。 -
从 开始到 遍历每个时间步 。
-
对于模型 中的每个层-头部组合
(l, h):- 获取原始模型的注意力输出 。
- 初始化
minimum_loss为无穷大,best_pattern为null。 - 遍历掩码模式列表 中的每个模式 :
- 使用当前模式 对层 的头部 应用掩码。
- 获取应用掩码后的模型注意力输出 。
- 计算原始输出 和掩码后输出 之间的均方误差 (Mean Squared Error, MSE) 作为
loss。 - 如果当前
loss小于minimum_loss,则更新minimum_loss和best_pattern。
- 将找到的最佳模式
best_pattern记录到字典dict中,对应于当前(t, l, h)组合。
-
返回填充好的
dict。这个算法通过直接比较原始注意力输出与应用稀疏掩码后的注意力输出的 MSE 损失,来找到最能近似原始行为的稀疏掩码模式,从而在无需训练的情况下保持模型性能。
-
4.2.3.2. 微调 (Finetuning)
除了无需训练的掩码搜索,STA 还可以通过少量微调来适应更激进的稀疏化,从而解锁更高的效率。尽管每个注意力层被限制在局部窗口内,但通过堆叠多层 Transformer,感受野 (receptive field) 仍能扩展,最终使扩散 Transformer 能够生成全局一致的视频。
微调过程中使用了三种不同的损失项:
-
注意力蒸馏损失 (Attention Distillation Loss): 这种损失直接监督
STA的中间注意力模式,使其与原始密集注意力 (dense attention) 的行为匹配。 符号解释:- :注意力蒸馏损失。
- :
Transformer层数。 - :
STA模型(学生模型 )的第 个Transformer层输出,给定噪声隐变量 、时间步 和文本嵌入 。 - :原始全注意力模型(教师模型 )的第 个
Transformer层输出,给定噪声隐变量 、时间步 和文本嵌入 。 - :L2 范数(平方),衡量两个输出之间的差异。 这个损失确保每个稀疏注意力层近似其对应的密集注意力教师。
-
最终层损失 (Final Layer Loss): 该损失旨在对齐学生模型和教师模型的最终输出。 符号解释:
- :最终层损失。
- :
STA模型(学生模型 )的最终输出。 - :原始全注意力模型(教师模型 )的最终输出。 这个损失确保学生模型的整体去噪能力与教师模型保持一致。
-
数据损失 (Data Loss): 此外,还采用了遵循
flow matching(流匹配) 公式的数据损失。 符号解释:- :数据损失。
- :输入帧的
VAE(变分自编码器) 隐变量。 - :在扩散步 时的噪声隐变量。
- :文本嵌入 (text embedding)。
- :通常表示扩散模型的预测目标(如噪声或原始数据)。 这个损失是扩散模型训练中的标准损失,用于指导模型学习正确的去噪方向,使预测结果与真实数据分布对齐。
完整目标函数 (Complete Objective Function): 将这些损失项结合起来,完整的优化目标为: 符号解释:
- :表示通过优化模型参数 来最小化损失。
- :表示对数据 (来自真实数据分布
p(x))、文本嵌入 (来自标准正态分布 ) 和时间步 的期望。 - :是每个损失项的权重系数,用于平衡不同损失项的重要性。
附录 A:3D NATTEN 掩码定义
为了更好地理解 STA 的独特之处,有必要了解 3D NATTEN 的掩码定义(原文 Algorithm 2)。NATTEN 的掩码计算涉及为每个查询 token 确定其窗口中心(考虑到视频边界),然后基于时空距离限制来构建掩码。
Algorithm 2 Mask Definition of 3D NATTEN
Input: Query coordinates (qt, qh, qw), Key coordinates (kt, kh, kw), Video size (Lt, Lh, Lw), Window size (Wt, Wh, Ww)
Compute window center:
qct ← max (min (qt, Lt − 1 − Wt/2), Wt/2)
qch ← max (min (qh, Lh − 1 − Wh/2), Wh/2)
qcw ← max (min (qw, Lw − 1 - Ww/2), Ww/2)
Compute masks:
time_constraint ← |qct − kt| ≤ Wt/2
hori_constraint ← |qch − kh| ≤ Wh/2
vert_constraint ← |qcw − kw| ≤ Ww/2
return time_constraint AND hori_constraint AND vert_constraint
算法解释:
-
输入: 查询
token的时空坐标 ,键token的时空坐标 ,视频的总尺寸 ,以及窗口尺寸 。 -
计算窗口中心: 对于每个查询
token,算法首先计算其注意力窗口的中心 。这个计算考虑了视频的边界,确保窗口中心不会超出视频范围。例如,qct的计算确保时间维度上的中心位于 和 之间。 -
计算掩码: 然后,通过比较键
token的坐标与计算出的窗口中心,并检查它们是否在窗口尺寸的一半以内,来确定时空约束 (time_constraint,hori_constraint,vert_constraint)。 -
返回: 所有三个约束都满足时,掩码为真(即注意力不被掩盖),否则为假。
这种逐
token的掩码计算,使得每个query的有效key集合可能不同,从而导致mixed blocks。
附录 A:3D STA 掩码定义
相比之下,3D STA 采用 tile-based (基于平铺块) 的坐标框架来定义掩码(原文 Algorithm 3)。
Algorithm 3 Mask Definition of 3D STA
Input: Query coordinates (qt, qh, qw), key coordinates (kt, kh, kw), video size (Lt, Lh, Lw), kernel size (Wt, Wh, Ww), tile size (Tt, Th, Tw)
Compute QK coordinates in tiles:
qt_tile ← qt // Tt
qh_tile ← qh // Th
qw_tile ← qw // Tw
kt_tile ← kt // Tt
kh_tile ← kh // Th
kw_tile ← kw // Tw
Compute window size in tiles:
Wt_tile ← Wt // Tt
Wh_tile ← Wh // Th
Ww_tile ← Ww // Tw
Compute window center:
qct ← max (min (qt_tile, (Lt // Tt − 1) − Wt_tile / 2), Wt_tile / 2)
qch ← max (min (qh_tile, (Lh // Th − 1) − Wh_tile / 2), Wh_tile / 2)
qcw ← max (min (qw_tile, (Lw // Tw − 1) − Ww_tile / 2), Ww_tile / 2)
Compute masks:
time_constraint ← |qct − kt_tile| ≤ Wt_tile / 2
hori_constraint ← |qch − kh_tile| ≤ Wh_tile / 2
vert_constraint ← |qcw − kw_tile| ≤ Ww_tile / 2
return time_constraint AND hori_constraint AND vert_constraint
算法解释:
-
输入: 除了
NATTEN的输入外,还包括tile的尺寸 。 -
计算 QK 在
tile中的坐标: 首先将查询和键的原始坐标映射到它们所属的tile坐标,例如qt_tile = qt // Tt。 -
计算窗口在
tile中的尺寸: 将原始窗口尺寸也转换为tile尺寸,例如Wt_tile = Wt // Tt。 -
计算
tile窗口中心: 类似NATTEN,但这里的计算是在tile坐标系下进行的。 -
计算掩码: 同样通过比较键
tile的坐标与计算出的窗口中心,并检查它们是否在tile窗口尺寸的一半以内。 -
返回: 所有三个约束都满足时,掩码为真。
STA的tile-based(基于平铺块) 掩码定义确保了在一个query tile内的所有token都遵循相同的掩码模式,并且只关注其局部窗口内的key tile。这种设计避免了NATTEN中复杂的逐token边界处理,从而产生了更规整、更硬件友好的注意力模式。
4.2.4. 微调细节 (Finetuning Details)
原文附录 B 提供了微调的具体设置:
- 训练数据: 使用 2,000 个从 HunyuanVideo 生成的合成视频,分辨率为 像素,117 帧。
prompt(提示) 来自Mixkit数据集。 - 预处理: 为了减少内存使用和加速训练,
VAE(变分自编码器) 编码的隐变量和文本编码器 (text encoder) 状态被预先计算。 - 训练过程: 微调进行了 1,600 步,批处理大小 (batch size) 为 2,学习率 (learning rate) 为 2e-5。
- 损失函数: 优化使用了公式 (5) 中的损失函数,系数设置为 、、。
- 指导尺度 (Guidance Scale): 为了防止在单一指导尺度上过拟合,模型在奇数步使用指导尺度 1,偶数步使用指导尺度 6。
- 硬件: 整个过程在 8 块 H100 GPU 上运行,结合
FSDP(Fully Sharded Data Parallel) 和context parallelism进行训练(8 小时),以及sequence parallelism进行推理。
5. 实验设置
5.1. 数据集
- HunyuanVideo: 作为主要的评估模型,
STA在 HunyuanVideo (最先进的开源视频DiT) 上进行评估。生成的视频输出包含 117 帧,分辨率为 。经过VAE(变分自编码器) 压缩和tokenization(词元化) 后,对应于形状为 (30, 48, 80) 的隐变量视频。 - FLUX: 为了验证
STA在 2D 图像领域的有效性,它也被应用于领先的图像扩散模型FLUX。 - 微调数据: 微调过程使用了 2,000 个从 HunyuanVideo 合成生成的视频,分辨率为 ,117 帧。这些视频的
prompt(提示) 来源于Mixkit数据集。 - MovieGen Bench: 人工评估时,从
MovieGen Bench中随机抽取 200 个prompt进行pairwise comparisons(成对比较)。 - VBench
prompt:VBench评估使用了其特定的prompt集合。 - COCO-2014: 在图像超分辨率任务中,从
COCO-2014验证数据集中随机抽取 1000 个caption(标题) 进行实验。 - Wan 2.1: 在另一个视频模型
Wan 2.1上也进行了训练无关的性能评估,该模型生成相同分辨率但帧数较短(69 帧)的视频。
5.2. 评估指标
论文从效率和生成质量两个方面评估 STA 的性能。
5.2.1. 效率指标
- MFU (Model FLOPs Utilization,模型浮点运算利用率):
- 概念定义: MFU 衡量的是
GPU实际计算性能与理论峰值性能的比率。它反映了模型在给定硬件上利用其浮点运算能力的效率。高 MFU 值意味着GPU的计算资源得到了更充分的利用,而不是因为内存墙或计算模式不规则而空闲。 - 数学公式:
- 符号解释:
- :模型在实际运行中每秒执行的十亿次浮点运算次数。
- :
GPU理论上每秒能执行的最高十亿次浮点运算次数。
- 概念定义: MFU 衡量的是
- Latency (延迟):
- 概念定义: 指的是完成某个计算任务所需的时间。在本文中,
Latency可以指注意力kernel(内核) 的执行时间(单位毫秒ms)或端到端视频生成所需的wall-clock(挂钟时间)(单位秒 )。 - 数学公式: 无通用数学公式,直接测量时间。
- 符号解释:
Latency(ms):注意力kernel执行时间,单位毫秒。Latency(s):端到端推理时间,单位秒。
- 概念定义: 指的是完成某个计算任务所需的时间。在本文中,
- Speedup (加速比):
- 概念定义:
Speedup衡量的是一个优化方法相对于基线方法在执行时间上的提升倍数。 - 数学公式:
- 符号解释:
- :基线方法完成任务所需的时间。
- :优化方法完成任务所需的时间。
- 概念定义:
- Kernel Efficiency (内核效率):
- 概念定义:
内核效率是稀疏kernel的 MFU 与全注意力kernel的 MFU 之比。它量化了稀疏kernel将理论FLOPs减少转化为实际延迟改进的能力。 - 数学公式:
- 符号解释:
- :稀疏注意力
kernel的 MFU。 - :全注意力
kernel的 MFU。
- :稀疏注意力
- 概念定义:
5.2.2. 质量指标
- 人类评估 (Human Evaluation):
- 概念定义: 通过人类评估者对不同模型生成的视频进行
pairwise comparisons(成对比较),判断哪个视频具有更高的整体质量,或两者质量相当。这是评估视频生成模型视觉感知质量最可靠的方法之一。 - 数学公式: 无通用数学公式,结果通常以胜率 (win rate)、平局率 (tie rate) 和失利率 (loss rate) 呈现。
- 符号解释:
Win Rate:评估者认为该模型生成视频质量优于对比模型的百分比。Tie Rate:评估者认为该模型生成视频质量与对比模型相当的百分比。Loss Rate:评估者认为该模型生成视频质量劣于对比模型的百分比。
- 概念定义: 通过人类评估者对不同模型生成的视频进行
- VBench (Huang et al., 2024):
- 概念定义:
VBench是一个综合性的视频生成模型基准测试套件,旨在评估视频生成模型的多个方面,包括视频质量 (Quality)、语义一致性 (Semantic) 和特定属性(如图像质量、时间连贯性、运动平滑度、物体分类、多物体、人物动作、颜色、空间关系、场景等)。它提供了一个多维度的评估框架。 - 数学公式:
VBench内部由多个子指标组成,每个子指标可能有其特定的计算方法。最终通常给出Quality Score(质量分数)、Semantic Score(语义分数) 和Final Score(最终分数)。 - 符号解释:
VBench Quality:反映视频的视觉保真度、清晰度、审美等。VBench Semantic:反映视频内容与文本提示的语义一致性。VBench Total:综合质量和语义的最终分数。- 其他子指标如
Imaging Quality(图像质量)、Temporal Flickering(时间闪烁)、Motion Smoothness(运动平滑度) 等,各自量化视频的特定视觉属性。
- 概念定义:
- SSIM (Structural Similarity Index Measure,结构相似性指数):
- 概念定义: SSIM 是一种感知指标,用于衡量两幅图像(或视频帧)之间的相似度。它将图像失真视为结构信息的变化,并考虑到亮度、对比度和结构三个方面。SSIM 值范围通常在 0 到 1 之间,1 表示两幅图像完全相同。
- 数学公式: 其中:
- 符号解释:
x, y:两幅待比较的图像(或图像块)。- :图像 和 的平均像素值。
- :图像 和 的像素标准差。
- :图像 和 的协方差。
- :为避免分母为零或接近零的稳定常数。
- :分别用于调整亮度、对比度和结构分量的权重(通常设为 1)。
- :表示该指标值越高越好。
- PSNR (Peak Signal-to-Noise Ratio,峰值信噪比):
- 概念定义: PSNR 是一种衡量图像质量的客观指标,通常用于量化有损压缩或处理后图像的失真程度。它通过计算原始图像和处理后图像之间的均方误差 (MSE) 来评估图像质量。PSNR 值越高表示图像质量越好。
- 数学公式:
其中
MSE(Mean Squared Error,均方误差) 定义为: - 符号解释:
- :图像中像素的最大可能值。对于 8 位灰度图像,通常为 255。
- :原始图像 和重建图像 之间像素值的均方误差。
I(i,j):原始图像在(i,j)处的像素值。K(i,j):重建图像在(i,j)处的像素值。m, n:图像的尺寸(高和宽)。- :表示该指标值越高越好。
- CD-FVD (Content-Dependent Frechet Video Distance,内容依赖的弗雷歇视频距离):
- 概念定义:
CD-FVD是一种用于评估视频生成模型质量的指标,它衡量生成视频分布与真实视频分布之间的相似性。CD-FVD考虑了视频内容,旨在更准确地反映视频的感知质量和多样性。与传统的FID(Frechet Inception Distance) 类似,FVD越低表示生成视频的质量越好,与真实视频的分布越接近。 - 数学公式:
CD-FVD的具体数学公式较为复杂,通常涉及对视频特征嵌入空间中的高斯分布进行匹配。简化来说,它计算两个高斯分布之间的Fréchet distance。 - 符号解释:
- :真实视频的特征嵌入分布。
- :生成视频的特征嵌入分布。
- :真实视频和生成视频特征分布的均值向量。
- :真实视频和生成视频特征分布的协方差矩阵。
- :向量的 L2 范数平方。
- :矩阵的迹。
- :表示该指标值越低越好。
- 概念定义:
5.3. 对比基线
论文将 STA 与以下基线模型和方法进行了比较:
- FlashAttention-2 (FA2) / FlashAttention-3 (FA3): 最先进的全注意力实现,用于提供基线性能。
FA3是FlashAttention系列的最新版本,提供了更高的效率。 - CLEAR (Liu et al., 2024): 一种基于圆形窗口的注意力机制,旨在通过限制查询范围实现线性注意力。
- NATTEN (Hassani et al., 2023): 邻域注意力
Transformer,一种针对 2D/3D SWA 设计的kernel(内核)。 - Tiled NATTEN:
NATTEN的优化版本,通过tiling(分块) 和kernel fusion(内核融合) 改进效率,但仍面临mixed blocks(混合块) 问题。 - Swin Transformer (Liu et al., 2021b): 采用移位窗口 (shifted window) 注意力机制的分层视觉
Transformer。其窗口不重叠,通过交替移位窗口实现跨窗口连接。 - -DiT (Chen et al., 2024): 一种针对扩散
Transformer的训练无关加速方法,通过缓存特征偏移来优化推理速度。
附录 C:基线详情
- Swin Transformer:
Swin Transformer引入了基于移位窗口的分层视觉Transformer。它将图像划分为不重叠的窗口并应用局部注意力,并通过交替移位窗口来连接跨窗口信息。本文将Swin attention应用于HunyuanVideo,并相应地每隔一层进行窗口移位。其局限性在于在一个注意力层内会破坏局部连接性。 - CLEAR:
CLEAR通过将全注意力替换为圆形窗口注意力机制来实现线性注意力,每个查询token只关注半径 内的key-value token。作者使用FlexAttention实现了CLEAR。 - -DiT:
\Delta`-DiT` 通过缓存特征偏移而非完整特征图来加速推理。它采用分阶段缓存策略,缓存后期 `DiT` 块的残差以供早期采样,缓存早期块的残差以供后期采样。论文基于其论文 reimplement 了该方法,并根据速度预算调整了关键超参数 (缓存块数量)、 (残差缓存间隔) 和 (时间步边界)。
# 6. 实验结果与分析
## 6.1. 核心结果分析
### 6.1.1. SLIDING TILE ATTENTION 的效率 (Efficiency of SLIDING TILE ATTENTION)
**评估配置:** 实验以 HunyuanVideo 生成 720P 5 秒视频的场景为基准,所有稀疏 `kernel` (内核) 均保持约 90% 的稀疏度。基线使用 `ThunderKittens` 的 `FA3`。
以下是原文 Table 2 的结果:
Methods Implementation Config Sparsity TFLOPS Latency(ms) MFU Kernel Efficiency Speedup FA 3 ThunderKittens - 0.00% 164.03 265.28 62.49% 100.00% 1.00× FA 3 CUDA - 0.00% 164.03 256.59 64.61% 103.39% 1.03× CLEAR FlexAttention r=16 90.46% 15.65 307.44 5.15% 8.24% 0.86× NATTEN FlexAttention w=(19,25,25) 89.69% 16.91 313.92 5.44% 8.71% 0.85× Tiled NATTEN CUDA w=(19,25,25) 89.69% 16.91 458.36 3.73% 5.97% 0.58× Tiled NATTEN FlexAttention w=(19,25,25) 89.69% 16.91 208.36 8.20% 13.12% 1.27× Swin FlexAttention w=(24,32,32) 87.42% 20.64 47.90 43.55% 69.69% 5.54× STA FlexAttention w=(18,24,24) 91.00% 14.76 36.36 41.03% 65.66% 7.30× STA ThunderKittens w=(30,40,40) 58.33% 68.35 111.73 61.82% 98.93% 2.37× STA ThunderKittens w=(18,24,24) 91.00% 14.76 25.38 58.79% 94.09% 10.45×
**分析:**
* **现有方法的低效性:** `CLEAR` 和 `NATTEN` 及其 `Tiled` 变体,尽管将 `TFLOPs` 显著降低到 15-17 左右(约 90% 稀疏度),但其 `MFU` (模型浮点运算利用率) 和 `Kernel Efficiency` (内核效率) 极低(低至 3.73% MFU,5.97% 内核效率),甚至导致 `wall-clock` (挂钟时间) 上的减速(0.86x 甚至 0.58x 减速)。这证实了这些方法无法有效将 `FLOPs` 减少转化为实际加速的缺陷,主要是由于 `mixed blocks` (混合块) 和 `mask` (掩码) 计算开销。
* **Swin 的表现:** `Swin` 作为一个非滑动窗口注意力,表现相对较好,`MFU` 超过 40%,内核效率超过 60%,实现了 5.54x 的加速。但正如作者在后面讨论的,其效率是以表达能力为代价的。
* **STA 的显著优势:**
* **高 MFU 与内核效率:** `STA` 在 `FlexAttention` 实现下,`MFU` 从 `Tiled NATTEN` 的 8.20% 提升到 41.03%。
* **卓越加速:** 结合 `ThunderKittens` 优化内核(异步数据加载和块间掩码管理),`STA` 在 91% 稀疏度下,`TFLOPs` 为 14.76,延迟仅 25.38ms,`MFU` 达到 58.79%,`Kernel Efficiency` 达到 94.09%,实现了 **10.45x** 的显著加速。
* **高效率下的灵活性:** 即使在较低稀疏度(58.33%)下(即使用更大的窗口),`STA` 仍能实现 2.37x 的加速,其 `MFU` 更是高达 61.82%,几乎与全注意力持平。这表明 `STA` 是第一个能同时实现 3D 局部性和硬件效率的滑动窗口稀疏注意力。
**图像 1. FLOPS分布及注意力内核延迟对比**
下图(原文 Figure 1)直观地展示了注意力计算在整个视频生成过程中的主导地位,以及 `STA` 如何在 `wall-clock` 延迟上实现显著加速。图 (b) 明确指出,现有方法未能将 `FLOPs` 减少转化为实际 `wall-clock` 加速,而 `STA` 实现了与稀疏性成比例的加速。

*该图像是图表,展示了不同方法的FLOPS分布及注意力内核延迟。图(a)显示了在生成115K令牌时,FFN、QKV和ATTN的FLOPS比例。图(b)对比了不同方法在注意力内核延迟上的表现,其中STA显著降低延迟,达到25.38ms。*
Figure 1. (a) Generating a 5s 720P clip in Hunyuan involves processing 115K tokens, making attention the dominant cost. (b) Attention latency comparison: existing methods fail to translate FLOP reduction into wall-clock speedup; STA is hardwareefficient and achieves proportional speedup with sparsity.
**附录 F.1:较低稀疏度下的内核性能**
以下是原文 Table 7 的结果:
Methods Implementation Config Sparsity TFLOPS Latency(ms) MFU Kernel Efficiency Speedup FA 3 ThunderKittens - 0.00% 164.03 265.28 62.49% 100.00% 1.00x FA 3 CUDA - 0.00% 164.03 256.59 64.61% 103.39% 1.03× CLEAR FlexAttention r=32 56.23% 71.80 675.05 10.75% 17.20% 0.39x NATTEN FlexAttention W=(30,41,41) 56.22% 71.81 804.62 9.02% 14.43% 0.33x Tiled NATTEN CUDA w=(29,41,41) 57.68% 69.41 173.57 4.04% 6.47% 0.15x Tiled NATTEN FlexAttention W=(30,41,41) 56.22% 71.81 409.89 17.70% 28.33% 0.65x Swin FlexAttention W=(48,64,64) 55.81% 72.49 127.51 57.46% 91.95% 2.08x STA FlexAttention w=(30,40,40) 58.33% 68.35 174.17 39.66% 63.46% 1.52x STA ThunderKittens w=(30,40,40) 58.33% 68.35 111.73 61.82% 98.93% 2.37x
**分析:** 在约 56% 的稀疏度下,稀疏 `kernel` 的 `MFU` 通常更高,但 `Table 2` 中的主要发现依然成立:`STA` 始终表现出更高的 `MFU` 和 `Speedup`,远超其他滑动窗口注意力实现。例如,在 58.33% 稀疏度下,`STA` (ThunderKittens) 达到 2.37x 加速和 61.82% MFU,而 `Tiled NATTEN` (FlexAttention) 仅有 0.65x 加速和 17.70% MFU。
### 6.1.2. 人工评估 (Human Evaluations)
通过对 `MovieGen Bench` 中的 200 个 `prompt` 进行 `pairwise comparisons` (成对比较),评估了 `HunyuanVideo`、`STA-tf-1.89x` (训练无关 `STA`,1.89x 加速)、`STA-t-2.43x` (微调 `STA`,2.43x 加速) 和两种
\Delta-DiT 变体 (1.36x 和 1.8x 加速) 的视频质量。
图像 10. 人工评估结果对比 下图(原文 Figure 7)展示了人工评估结果:
该图像是示意图,展示了使用滑动平铺注意力(STA)生成视频的效果对比,包括HunyuanVideo和不同模型的生成时间。提示描述了宇航员在石头建筑间行走的场景。各模型的生成时间分别为15分钟45秒、8分钟21秒、6分钟29秒和11分钟34秒。
Figure 7. Human evaluation on 200 prompts from the MovieGen Bench (Polyak et al., 2024). STA achieves a end-to-end speedup while maintaining performance comparable to the original HunyuanVideo. Additionally, STA consistently outperforms . DiT across different inference budgets.
分析:
- STA vs. -DiT:
STA-t-2.43x(2.43x 加速) 显著优于-DiT-1.8x(1.8x 加速),胜率高达 70.0% 对 11.0%。类似地,STA-tf-1.89x(1.89x 加速) 以 66.5% 的胜率超越-DiT-1.36x(1.36x 加速)。这表明STA在更高加速比下仍能保持卓越的视频质量。 - STA vs. HunyuanVideo (原始模型):
STA-tf-1.89x与原始HunyuanVideo相比,取得了 83.0% 的平局率,表明其在大多数情况下能保持与原始模型几乎相当的质量。尽管其胜率比失利率低 7.0 个百分点,但考虑到 1.89x 的显著加速,这一权衡是完全值得的。 - 结论: 人工评估结果有力地证实了
STA在质量与效率权衡方面优于
\Delta`-DiT`。
### 6.1.3. 训练无关结果 (Training-free Results)
在 `VBench` `prompt` 上,评估了训练无关的 `STA` 和
\Delta-DiT 在不同采样步数下的鲁棒性。
以下是原文 Table 3 的结果:
| Model | SSIM ↑ | PSNR ↑ | CD-FVD ↓ | Latency | Speedup |
| steps = 50 | |||||
| Δ-DiT | 72.86 | 18.09 | 122.74 | 693s | 1.36× |
| STA | 87.67 | 28.76 | 66.12 | 501s | 1.89× |
| steps = 25 | |||||
| Δ-DiT | 77.91 | 19.86 | 196.25 | 352s | 1.34× |
| STA | 88.96 | 28.99 | 76.34 | 250s | 1.89× |
| steps = 10 | |||||
| Δ-DiT | 83.19 | 21.20 | 201.24 | 144s | 1.32× |
| STA | 87.84 | 27.14 | 84.80 | 105s | 1.76× |
分析:
- STA 持续优于 -DiT: 无论采样步数如何,训练无关的
STA始终以更高的加速比 (1.76x-1.89x) 取得更优的质量指标。 - 显著的质量差距: 在 50 步采样下,
STA的SSIM比
\Delta`-DiT` 高 14.81 (87.67 vs. 72.86),`PSNR` 高 10.67 (28.76 vs. 18.09)。`CD-FVD` (越低越好) 方面,
\Delta-DiT 比 STA 高 56.62 (122.74 vs. 66.12)。这种差距在 25 步和 10 步采样时甚至更大。
- 定性观察:
\Delta`-DiT` 在所有采样步数下都持续产生视觉质量下降的输出,表现出结构相似性受损和细节丢失,而 `STA` 则能保持与原始模型高度一致的保真度。
**附录 D:Wan 2.1 上的结果**
以下是原文 Table 5 的结果:
Model SSIM↑ PSNR ↑ Latency Speedup STA 85.81 24.42 730s 1.60×
**分析:** 在另一个视频生成模型 `Wan 2.1` 上,`STA` 依然能够实现与 `HunyuanVideo` 上类似的效果,在相同的稀疏度下保持了可比的质量指标 (SSIM 和 PSNR),并提供了 1.60x 的加速。视频帧数较短导致了端到端加速略有降低。
### 6.1.4. 微调结果 (Finetuning Results)
本节评估了在有无微调情况下,直接用稀疏注意力替换全注意力对视频生成质量的影响,特别关注 `VBench` 分数。
以下是原文 Table 4 的结果:
Methods Config VBench Quality VBench Semantic VBench Total Attn Sparsity PFLOPS Latency Speedup FA2 85.34% 72.17% 82.71% 0.00% 574.16 1496s 0.63 × FA3 85.34% 72.17% 82.71% 0.00% 574.16 945s 1.00× w.o training CLEAR r=32 84.41% 74.20% 82.37% 56.23% 280.90 2567s 0.37 × Tiled NATTEN w=(30,41,41) 84.61% 75.00% 82.69% 58.33% 269.92 1858s 0.51× Swin w=(48,64,64) 80.91% 71.35% 79.00% 55.81% 283.11 762s 1.24× Swin w=(30,40,40) 78.84% 72.28% 77.53% 76.49% 175.20 497s 1.90 × STA w=(30,40,40) 84.63% 73.83% 82.46% 58.33% 269.92 527s 1.79× STA w=(18,24,24) 81.47% 77.03% 80.58% 91.00% 99.54 268s 3.53 × w. training Swin w=(30,40,40) 77.50% 67.39% 75.48% 55.81% 283.08 497s 1.90 × STA w=(30,24,40) 85.37% 73.52% 83.00% 75.00% 182.99 388s 2.44× STA w=(18,24,24) 84.76% 74.05% 82.62% 91.00% 99.54 268s 3.53×
**分析:**
* **CLEAR 和 Tiled NATTEN 的问题:** 在无训练情况下,`CLEAR` 和 `Tiled NATTEN` 虽然能保持合理的 `VBench` 分数(82.37% 和 82.68%),与 `FA3` (82.71%) 相近,但它们的端到端推理延迟反而增加(2567s 和 1858s),导致加速比远低于 1x (0.37x 和 0.51x)。这再次强调了它们在将 `FLOPs` 减少转化为实际加速方面的低效性。
* **Swin 的局限性:** `Swin` 虽然能实现适度加速(1.24x-1.90x),但其非重叠窗口设计违反了 3D 局部性,导致 `VBench` 分数显著下降(79.00% 和 77.53%)。更糟糕的是,对其进行微调不仅未能恢复性能,反而进一步降低了 `VBench` 分数(75.48%),说明其架构固有地与视频扩散模型的 3D 局部性不匹配。
* **STA 的卓越表现:**
* **无训练加速:** 在无训练情况下,`STA` 以 配置实现了 91.00% 的注意力稀疏度,`PFLOPS` 降低 5.76 倍,端到端延迟显著降低至 268 秒,实现了 **3.53x** 的实际加速。同时,`VBench` 总分仍保持 80.58%,仅略低于原始模型。
* **微调后进一步提升:** 通过微调,`STA` 的 `VBench` 分数从 80.58% 提升至 82.62%(在 91.00% 稀疏度下),几乎与原始 `FA3` 的性能持平 (82.71%),但加速比仍高达 **3.53x**。在 75.00% 稀疏度下,微调 `STA` 也实现了 2.44x 加速,`VBench` 总分甚至略高于 `FA3` (83.00% vs. 82.71%)。
* **结论:** `STA` 成功解决了现有方法在质量和效率上的两难困境,在实现大幅加速的同时,通过训练无关配置或少量微调,保持了优异的视频生成质量。
**附录 F.2:详细 VBench 结果**
以下是原文 Table 8 的结果:
Model Appearance Style Subject Consistency Background Consistency Temporal Flickering Motion Smoothness Dynamic Degree Aesthetic Quality Imaging Quality Overall Consistency FA3 18.43% 94.22% 96.74% 99.21% 99.15% 75.00% 64.63% 67.97% 25.96% w.o training CLEAR 18.73% 93.63% 96.51% 98.99% 99.01% 68.06% 63.75% 68.35% 26.23% Tiled NATTEN 18.79% 94.59% 96.61% 98.75% 98.85% 70.83% 63.79% 68.16% 26.53% Swin w=(48,64,64) 20.85% 91.74% 95.48% 98.67% 97.77% 77.78% 51.01% 62.22% 25.27% Swin w=(30,40,40) 20.62% 90.33% 93.09% 98.78% 96.53% 75.00% 48.10% 61.89% 25.62% STA w=(30,40,40) 18.79% 94.75% 96.50% 98.82% 98.83% 69.44% 64.18% 68.39% 26.47% STA w=(18,24,24) 21.25% 89.66% 91.64% 98.46% 97.27% 83.33% 59.75% 64.23% 26.61% w. training Swin w=(30,40,40) 20.07% 89.78% 94.93% 98.86% 96.64% 70.83% 44.91% 55.99% 26.00% STA w=(30,24,40) 18.90% 94.90% 97.60% 99.68% 99.23% 73.61% 63.77% 66.21% 26.58% STA w=(18,24,24) 18.90% 94.64% 96.76% 99.22% 99.11% 69.44% 64.52% 66.67% 26.09%
以下是原文 Table 9 的结果:
Model Object Classification Multiple Objects Human Action Color Spatial Relationship Scene Quality Score Semantic Score Final Score FA3 85.76% 70.12% 90.00% 88.66% 71.28% 35.25% 85.34% 72.17% 82.71% w.o training CLEAR 88.13% 77.97% 88.00% 91.10% 77.49% 32.85% 84.41% 74.20% 82.37% Tiled NATTEN 83.54% 72.18% 94.00% 92.28% 81.21% 37.94% 84.61% 75.00% 82.69% Swin w=(48,64,64) 78.16% 58.54% 87.00% 93.68% 77.45% 37.79% 80.91% 71.35% 79.00% Swin w=(30,40,40) 79.19% 60.44% 88.00% 93.68% 77.24% 35.54% 78.84% 72.28% 77.53% STA w=(30,40,40) 80.54% 71.19% 93.00% 89.81% 79.25% 36.77% 84.63% 73.83% 82.47% STA w=(18,24,24) 88.13% 75.46% 91.00% 91.61% 82.52% 42.15% 81.47% 77.03% 80.58% w. training Swin w=(30,40,40) 77.14% 48.86% 73.00% 87.00% 63.38% 39.03% 77.50% 67.39% 75.48% STA w=(30,24,40) 91.77% 68.45% 86.00% 89.59% 72.76% 39.53% 85.37% 73.52% 83.00% STA w=(18,24,24) 92.96% 74.16% 93.00% 84.50% 73.41% 38.23% 84.76% 74.05% 82.62%
**分析:**
* **STA vs. 其他稀疏方法:** `STA` 在 `Imaging Quality` (图像质量) 和 `Multiple Objects` (多物体) 等视频质量指标上优于 `Swin` 注意力,并与 `CLEAR` 和 `Tiled NATTEN` 保持可比或更优的水平。
* **稀疏度对训练无关 `STA` 的影响:** 随着稀疏度增加(例如从 到 ),训练无关的 `STA` 在 `Temporal Flickering` (时间闪烁) 和 `Motion Smoothness` (运动平滑度) 等质量相关指标上出现系统性下降。这反映了过度稀疏化对低级别视觉连贯性的影响。
* **文本嵌入的放大作用:** 另一方面,在更高稀疏度下,`Appearance Style` (外观风格)、`Color` (颜色) 和 `Spatial Relationships` (空间关系) 等语义对齐维度反而有所改善。作者推测这是因为当空间-时间注意力被稀疏化时,文本嵌入在注意力计算中的作用被放大,从而增强了视频与文本提示的语义一致性。
* **微调的功效:** 经过训练的 `STA` 模型在视频质量指标上比其未经训练的对应模型有了显著提升,同时保持了可比的语义连贯性。这强调了微调在不损害文本-视频对齐的情况下,优化低级别视觉保真度的有效性。
### 6.1.5. 图像超分辨率结果 (Results on Image Super-Resolution)
**附录 E:图像超分辨率结果**
以下是原文 Table 6 的结果:
Methods SSIM PSNR Sparsity Latency Speedup 1K →2K CLEAR r=16 0.9291 28.1142 96.12% 13s 1.54× CLEAR r=32 0.9443 29.6722 85.94% 15s 1.33× STA w=(48,72) 0.9357 29.1086 81.25% 14s 1.43× 2K→4K CLEAR r=16 0.9394 29.0463 98.98% 67s 2.90× CLEAR r=32 0.9455 30.0742 96.08% 92s 2.11× STA w=(48,72) 0.9470 30.1939 95.31% 57s 3.40×
**分析:** `STA` 也被应用于使用 `SDEdit` 的图像超分辨率任务,并在 `FLUX` 模型上进行了评估。结果表明,`STA` 在提供显著更高效率的同时,实现了与 `CLEAR` 相当的生成质量。例如,在 2K 到 4K 超分辨率任务中,`STA` 以 95.31% 的稀疏度实现了 3.40x 的加速,同时 `SSIM` 和 `PSNR` (0.9470 和 30.1939) 略优于 (0.9455 和 30.0742) 的 2.11x 加速。这进一步证明了 `STA` 方法在 2D 任务中的通用性和有效性。
### 6.1.6. 定性示例 (Qualitative Examples)
原文附录 G 提供了 `HunyuanVideo`、`STA` 和
\Delta-DiT 生成视频的定性比较。
图像 2. 定性比较 1 下图(原文 Figure 10)展示了一个宇航员在石头建筑间行走的场景。
该图像是多个帧的示意图,展示了注意力在不同视频帧中的局部集中情况。每个帧的绿色框代表查询点,而颜色渐变区域表示对查询的高注意力值,展现了注意力是如何在视频帧中随着时间变化而聚焦在局部区域的。
Figure 10. Qualitative comparisons. While fine-tuning introduces minor shifts in the output distribution of , the model still preserves high video generation quality. Videos generated by -DiT are generally less sharp than those generated by the original HunyuanVideo and STA.
图像 3. 定性比较 2 下图(原文 Figure 11)展示了一个宁静岛屿度假村的场景。
该图像是图表,左侧展示了在多种提示下,注意力分数的平均回收率,右侧显示了不同头部之间的回收率标准差。大多数头部的高回收率表明获得了局部注意力模式,而不同头部的标准差较低则表明稳定性。
Figure 11. Qualitative comparisons. While fine-tuning introduces minor shifts in the output distribution of , the model still preserves high video generation quality. Videos generated by -DiT are generally less sharp than those generated by the original HunyuanVideo and STA.
图像 4. 定性比较 3 下图(原文 Figure 12)展示了一个人们在日落海滩篝火旁聚集的场景。
该图像是图表,展示了 NATTEN、Tiled NATTEN 和 STA 的注意力图。图中包括不同的块类型,如空块、混合块、密集块和 FA 块。从上到下分别为 (a) NATTEN,(b) Tiled NATTEN 和 (c) STA,展示了它们在注意力计算中的表现差异。
Figure 12. Qualitative comparisons. While fne-tuning introduces minor shifts in the output distribution of , the model still preserves high video generation quality. Videos generated by -DiT are generally less sharp than those generated by the original HunyuanVideo and STA.
分析:
- STA-t-2.43x: 微调后的
STA模型 (STA-t-2.43x) 即使引入了输出分布的微小偏移,仍能保持高视频生成质量,表明其在视觉上具有很强的保真度。 - -DiT:
\Delta`-DiT` 生成的视频通常比原始 `HunyuanVideo` 和 `STA` 生成的视频清晰度更低,细节更模糊,这与定量指标 `PSNR` 和 `SSIM` 的结果一致。
* **结论:** `STA` 在视觉质量上显著优于
\Delta-DiT,并能与原始 HunyuanVideo 保持高度可比性,进一步支撑了其高效且高质量的优势。
6.2. 数据呈现 (表格)
本节已将所有论文中的表格转录到对应的小节中进行分析。
6.3. 图像分析
本节已将所有论文中的图像嵌入到对应的小节中进行分析。
7. 总结与思考
7.1. 结论总结
本文成功引入了 滑动平铺注意力 (Sliding Tile Attention, STA),以加速视频扩散模型 (Diffusion Transformers, DiTs) 的生成过程。通过深入分析,作者发现最先进的视频 DiTs 中的 3D 注意力分数普遍集中在局部 3D 窗口内(即 3D 局部性),且这种模式具有 头部特化 现象并跨 prompt 稳定。STA 利用这一关键观察,通过将滑动单位从 token 转换为 tile,并结合硬件感知的设计,消除了传统滑动窗口注意力 (SWA) 中低效的 mixed blocks (混合块)。
在内核级优化方面,STA 借鉴了 FlashAttention-3 和 ThunderKittens 的生产者-消费者范式,将 inter-block mask (块间掩码) 逻辑从计算 kernel (内核) 中解耦,使其计算 warpgroup 始终执行密集的、硬件高效的注意力计算。这使得 STA 成为第一个能够将 FLOPs 减少有效转化为实际 wall-clock (挂钟时间) 加速的 2D/3D 类滑动窗口注意力实现,达到了高达 58.79% 的模型 FLOPs 利用率 (MFU)。
实验结果表明,STA 在领先的视频 DiT 模型 HunyuanVideo 上实现了显著的性能提升:无需训练即可将端到端延迟从 945 秒降低到 501 秒,加速 1.89 倍,且不降低质量。通过微调,延迟可进一步降低至 268 秒,加速高达 3.53 倍,同时 VBench 分数仅下降 0.09%。在人类评估中,STA 也显著优于其他加速方法如
\Delta`-DiT`,证明了其在高质量和高效率之间的优越权衡。此外,`STA` 在图像超分辨率任务中也展现了其通用性。
总而言之,`STA` 为视频生成 `DiTs` 提供了强大的加速解决方案,使其更具实用性和可部署性。
## 7.2. 局限性与未来工作
论文作者指出了其方法在概念上与其他加速技术(如缓存和一致性蒸馏 `consistency distillation`)是正交的。这意味着 `STA` 的优化可以与这些方法独立应用,或结合使用。
**未来工作方向:** 作者计划探索 `STA` 与缓存技术(如
\Delta-DiT)和一致性蒸馏等方法结合使用的有效性,以期实现进一步的效率提升。
7.3. 个人启发与批判
7.3.1. 个人启发
- 3D 局部性与头部特化是关键洞察: 论文的核心在于对视频扩散模型中
3D 局部性和头部特化模式的发现和量化。这表明即使是经过全局注意力训练的模型,其内在机制也倾向于利用局部信息,这为稀疏注意力提供了坚实的理论基础和实践依据。这种数据特性驱动的稀疏化思路非常优雅。 - 硬件感知设计至关重要: 论文深刻揭示了在稀疏注意力中,仅仅减少
FLOPs是不够的,关键在于如何将理论上的FLOPs减少转化为实际的wall-clock加速。STA通过tile-by-tile(逐平铺块) 操作和producer-consumer(生产者-消费者) 异步kernel设计,有效地解决了mixed blocks(混合块) 导致的硬件利用率低下问题。这强调了系统与算法协同设计 (system-algorithm co-design) 在高性能计算中的重要性。对于未来的模型设计,不仅要关注理论复杂度,更要关注其在实际硬件上的执行效率。 - 简易而强大的
tile-wise范式: 将滑动单元从token提升到tile这一看似简单的改变,却极大地规整了注意力模式,消除了mixed blocks,从而带来了显著的效率提升。这表明有时最有效的解决方案可能来源于对基本操作的重新思考和巧妙设计。 - 微调的适应能力: 论文展示了即使在较高稀疏度下,通过少量微调,模型也能有效适应并恢复甚至超越原始性能。这为部署高效模型提供了灵活性:可以在无训练要求下实现即插即用加速,也可以通过相对较小的训练成本进一步优化以达到极致效率。
7.3.2. 潜在问题、未经验证的假设或可以改进的地方
tile和窗口尺寸的约束:STA假设视频尺寸 和窗口尺寸 都是tile尺寸 的整数倍。在实际应用中,如果视频分辨率或窗口设置不满足这些严格的倍数关系,可能需要额外的填充 (padding) 或裁剪 (cropping) 操作,这可能会引入额外的开销或影响局部性捕获的精度。未来的工作可以探索对非整数倍尺寸更鲁棒的设计。- 步的启发式选择: 在训练无关模式下,初始 个
timestep保持全注意力是基于经验启发式方法。这些初始步骤可能在扩散过程中至关重要,但 的具体数值选择可能依赖于模型和任务。进一步的分析可以探索自适应或更优的 选择策略,甚至在所有步骤都应用STA。 head specialization的更深层理解: 论文提到了head specialization(头部特化),即不同注意力头部具有不同局部性模式。虽然这被用于指导窗口大小搜索,但对这种特化模式的深层原因和其在视频生成不同阶段中的作用,仍有待更深入的理论分析。- 泛化到更极端的稀疏性: 尽管
STA已经实现了高达 91% 的稀疏度,但在某些场景下,模型可能需要更极端的稀疏性以进一步降低资源消耗。探索在更高稀疏度下如何维持质量,例如通过更高级的知识蒸馏技术或与量化 (quantization) 方法结合,将是一个有趣的挑战。 - 与其他
Transformer架构的兼容性:STA似乎是为DiTs量身定制的,特别是针对 3D 视频数据。它如何泛化到其他Transformer架构(如ViTs用于图像分类,或LLMs用于文本)以及不同类型的多维数据,还有待进一步验证和适应性研究。
7.3.3. 迁移和应用价值
- 实时视频生成:
STA显著的加速能力使其在需要实时或近实时视频生成的应用中具有巨大潜力,如视频编辑、虚拟现实、增强现实、游戏内容创作等。 - 资源受限环境下的部署: 降低的计算成本和内存需求使得视频生成模型可以在更低端
GPU或边缘设备上进行部署,从而扩展了其可访问性和应用范围。 - 多模态
Transformer: 视频本质上是多模态数据(时间+空间),STA提出的tile-wise稀疏化思想和硬件优化经验,可以为其他高维、多模态Transformer模型(例如 3D 点云、医学图像处理、物理模拟)提供设计灵感。 - 高效训练: 虽然论文主要关注推理加速,但如果
STA能在训练阶段保持性能,它也能显著降低训练大型视频模型所需的计算资源和时间,从而促进研究和开发。
相似论文推荐
基于向量语义检索推荐的相关论文。