Accelerating Spatio-Temporal Attention for Video Diffusion on TPUs

在TPU上加速视频扩散的时空注意力机制 - Google Developers Blog

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

在TPU上加速视频扩散的时空注意力机制

2026年9月30日

Ravisri Valluri

博士软件工程师实习生

Sagar Chapara

机器学习工程师

Rishabh Manoj

高级机器学习工程师

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

一个将理论稀疏性转化为TPU端到端加速的实际案例研究。

为什么视频扩散速度慢？

视频扩散模型通常生成视频速度较慢，主要有两个原因：大量的去噪步骤，以及每个去噪步骤的高昂成本。在长视频序列场景中，自注意力机制可能成为单个去噪步骤延迟的主要贡献者之一。仅以81帧720p视频为例，当前主流开源视频生成模型的序列长度范围从50K到400K不等。

表1：不同分辨率下自注意力延迟及其在单层Transformer块延迟中的占比（81帧视频在八块TPU v6e芯片上的表现）。

以720p（高清）到1440p（2K）的分辨率升级为例，序列长度会增加至四倍。由于全注意力机制的计算复杂度与序列长度呈平方关系，其在单层延迟中的占比可能从55.5%增长到88.2%。

加速注意力机制

由于注意力机制是单个步骤中延迟的最主要来源，任何激进的推理优化都应优先针对注意力机制。幸运的是，视频扩散中的注意力具有高度结构化特征：许多查询-键交互携带的注意力质量非常有限，因此大量成对计算可以被跳过。密集注意力会无差别计算所有交互，而稀疏注意力通过掩码仅保留最重要的交互并丢弃其余部分。下图展示了视频扩散模型的注意力矩阵示例，用以说明注意力质量的集中程度。

图1：空间头（第12头）的注意力质量分布。在以帧优先排序（左图）时，注意力集中在局部帧；在时间排序（右图）时，相同质量却分散在不同token中——这说明token排序必须与头类型匹配。

然而，注意力模式在不同头、不同层甚至不同步骤中会有所变化。稀疏视频生成（Sparse VideoGen，SVG）是一项开创性工作，它识别并利用了一种注意力模式：不同头通常可以归类为空间头或时间头。重要的是，这些模式并非任意稀疏结构，而是具有高度规则的几何特性。在空间头中，一个块主要关注同一帧或相邻帧中的其他块；在时间头中，一个块则关注跨大量帧的小范围空间区域。下图展示了空间头（第一行）和时间头（第二行）中单个查询对应的注意力质量分布。在空间头中，查询会广泛关注其所在帧及相邻帧的所有token；在时间头中，其关注范围则局限于更小的空间区域，但覆盖更多帧。

图2：注意力头专业化（第6步，第32层）。青色框标记了查询令牌在帧F06上的空间位置（y,x）。在空间头（顶部），94.8%的注意力质量在查询帧上广泛分布，用于二维合成。在时间头（底部），注意力集中在先前帧（F04–F05）的同一空间（y,x）管线上，以追踪局部运动。

基于这一结构，SVG在推理时动态分析注意力头，并将每个头路由到空间或时间掩码中。它是通过采样少量查询，在密集、空间和时间注意力下计算它们的注意力输出，并选择与密集基线输出偏差最小的稀疏掩码来实现的。这使SVG能够在给定稀疏度水平下保持更高的质量。有关分析和路由算法的细节，请参阅SVG论文。如图3所示，两种掩码也保留了对第一帧（F00）的全部注意力——作为锚定全局场景外观的注意力汇点，同时还有一个用于帧间交互的局部带宽。

图3：按帧优先顺序排列的稀疏注意力掩码。左：空间掩码（连续的局部带宽+第一帧汇点F00）。右：时间掩码，在按帧优先顺序排列时表现为细粒度的对角条纹——这促使我们后续对令牌进行重新排列。

在TPU上实现稀疏注意力

下表总结了我们比较的自定义JAX和Pallas Splash注意力内核实现。所有时间均使用单个TPU v6e设备上的合成BF16输入，包含75.6K令牌、10个头和128的头维度。参考查询和键/值瓦片大小分别为3328和1536。稀疏变体保留了约38.87%的查询-键对。这些是孤立的注意力内核时间；路由、令牌排列和设备间通信不在此次测量范围内。

表2：单个TPU v6e芯片（75,600个令牌，10个头，头维度128）上稀疏Splash注意力内核优化的进展。

图4：与表2对应的瓦片遍历策略（从左到右：B0密集、B1精确稀疏、B2分类瓦片和B3平衡瓦片舍入）。

逻辑稀疏性不等于物理稀疏性

图5：稀疏局部窗口注意力的硬件瓦片分类。瓦片被划分为完整瓦片（通过无掩码的快速路径执行）、边界瓦片（需要逐元素坐标掩码）和跳过瓦片（完全省略以节省计算和带宽）。

尽管SVG有巧妙的方法选择注意力掩码，但这种理论上的稀疏性需要在实际硬件上转化为加速效果。现代注意力实现不会显式生成完整的注意力矩阵乘法Q·Kᵀ；而是将其划分为如Q[i₁:i₂]·K[j₁:j₂]ᵀ这样的瓦片，并在这些瓦片间累积必要的统计信息以计算注意力输出SoftMax(QKᵀ / √d)V。在访问的瓦片内，被排除的查询-键对的分数在softmax之前会被设置为负无穷，使它们对注意力权重的贡献为零。这强制应用了掩码，但本身并未避免最初计算这些分数。

掩码将瓷砖分为三种类型。完整瓷砖仅包含保留的查询-键对，不需要块内掩码。边界瓷砖同时包含保留和排除的对，因此需要逐元素掩码。空瓷砖（在图中标记为 Skipped tiles）不包含任何保留对，可以完全跳过。逻辑稀疏性统计排除的对，但硬件节省取决于我们能跳过哪些瓷砖以及在访问的瓷砖中剩余的工作量。

Splash 注意力机制已经支持稀疏掩码并可以跳过空瓷砖。然而，这本身并不能保证加速：内核处理访问瓷砖的方式同样重要。在我们最初的稀疏原型（B1：朴素块遍历）中，掩码保留了 39% 的查询-键对，即约 61% 的逻辑稀疏性。但内核访问了 44% 的外层瓷砖，因为边界瓷砖也包含将被掩码排除的对。它跳过了其余瓷砖，但仍对每个访问的瓷砖内的精确掩码进行计算。与密集型 Splash 的 78.70 毫秒相比，此实现耗时 96.37 毫秒：尽管跳过了超过一半的瓷砖，但延迟仍高出 22%。

当没有需要掩码的内容时避免块内掩码

计算精确掩码需要确定哪些查询-键坐标是有效的，并将该条件应用于注意力分数。在 TPU 上，即使瓷砖中的每对都是有效的，对每个访问的瓷砖执行逐元素掩码计算会成为向量处理单元（VPU）的瓶颈，并导致矩阵乘法单元（MXU）停滞。我们可以通过在执行注意力之前识别完整瓷砖和边界瓷砖并为它们提供独立路径，来避免这种不必要的计算。

为了解决这个问题，我们的第二次迭代（B2：完整/边界瓷砖专业化）为完整瓷砖提供完全无需掩码的快速路径，仅在边界瓷砖上执行精确坐标掩码。这两条路径共同贡献于相同的注意力输出，其部分结果通过对应的 softmax 归一化统计进行合并。保留的查询-键对保持不变。延迟从 96.37 毫秒降至 54.12 毫秒——相比朴素稀疏遍历（B1）减少了 44%，比密集型 Splash 快 31%。这一对比展示了在保持精确稀疏掩码的同时对瓷砖执行进行专业化带来的优势。

瓷砖尺寸调整可能带来权衡

较小的瓷砖可以更紧密地跟随掩码边界，减少需要块内掩码的瓷砖比例。然而，瓷砖尺寸也会影响内核对计算和数据移动的分组方式，因此减少边界处理并不一定减少延迟。

图 6：大瓷砖与小瓷砖的边界对比。深色线条追踪 SVG 掩码边界（包括左侧的垂直首帧汇流）。较小的瓷砖更紧密地跟随边界，以牺牲计算瓷砖效率为代价，减少了黄色边界瓷砖的数量。

为了说明这种权衡，我们在保持键/值瓷砖尺寸（BKV）固定为 1536 的同时，调整查询瓷砖尺寸（BQ）。每种配置均已使用独立的完整路径和边界路径。当 BQ = 1024 时，仅有 16.95% 的执行瓷砖需要内部掩码，但延迟为 66.88 毫秒。将 BQ 提高到 3328 时，这一比例上升至 27.55%，但延迟降至 54.12 毫秒。进一步将 BQ 提高到 4864 时，延迟再次上升至 58.62 毫秒。图表显示了这种权衡：掩码工作量最少的配置并非最快的。这些百分比描述的是需要内部掩码的工作量/瓷砖比例，而非计算掩码所花费的时间。

Figure 7: 折线图显示TPU v6e上四种查询瓦片尺寸下内核延迟（毫秒）与需要内瓦片掩码的工作百分比关系，表明当BQ等于3328时延迟达到最小值。

将稀疏掩码与最终执行对齐

我们可以通过将SVG掩码与内核执行的瓦片对齐来进一步减少边界掩码。向外四舍五入边界会包含一些之前被排除的配对；向内四舍五入会移除一些之前保留的配对。在保持查询-键配对预算大致不变的同时，这种权衡会略微修改掩码。与之前的优化不同，这会改变对输出有贡献的交互方式。

四舍五入消除了在每个访问瓦片中强制执行精确SVG边界的需求。由于75,600个token不是3328 × 1536瓦片尺寸的精确倍数，超出真实序列长度的边缘瓦片仍然包含填充，且填充位置必须贡献精确的零注意力权重。因此我们重用完整/边界区分，现在只需对包含填充的瓦片进行掩码。

在最终配置（B3：瓦片对齐的稀疏遍历）中，这将需要掩码的瓦片比例从27.55%降低至仅2.22%。将掩码严格限制在边界填充上可将延迟降至32.76毫秒——相比密集型Splash注意力实现了2.40倍的加速。

这些结果共同表明为什么稀疏注意力需要同时具备合适的掩码和高效的执行策略：跳过空瓦片、避免对完整瓦片进行掩码，并将整体掩码与硬件计算的瓦片对齐。

在TPU上支持动态掩码和token布局

我们现在有了更快的稀疏注意力内核，但将其集成到分布式推理中需要处理动态token布局和不同头的掩码。

Figure 8: 动态时空注意力的端到端执行流程。动态路由评估采样查询以生成每个头的布尔标志，这些标志在静态形状张量上选择token布局，然后分发到稀疏Splash注意力内核并恢复输出顺序。

使用固定张量形状支持动态路由

路由步骤为每个头返回布尔标志以选择空间或时间注意力。这些标志在保持张量维度固定的同时，在原始布局和置换布局之间进行选择。与将头动态分割为可变大小的空间和时间子集不同，运行时更改头的分配不会改变张量形状或触发代价高昂的XLA重新编译。相同的标志在注意力后将时间头恢复到原始token顺序，因此下游层接收到预期的顺序。空间头保留其原始顺序。

安排token以实现高效时间访问

Figure 9: token内存布局转换。以帧优先顺序（F,H,W）存储时，帧1的空间token会连续存储。转换为时间优先顺序（H,W,F）后，来自同一空间位置的连续帧（B1–B4）的token会被分组到连续内存中，以实现高效的时间内核遍历。

将token从以帧优先（F, H, W）的排列方式转换为时间顺序（H, W, F）后，同一空间位置在不同帧中的token会被排列在一起——这种转换将时间掩码（如图3右侧所示的分散对角条纹）变为连续的单条带，便于分块稀疏核高效遍历。我们尝试过避免这种排列，改在核内部处理时间访问，但在时间头基准测试中，延迟从44.40ms增加到74.14ms，尽管更快的路径包含额外的排列和恢复开销。这与核内部额外的索引和数据重排操作一致，尽管我们未单独测量这些成本。

对本地头执行布局转换

在分布式推理中，执行这种排列的位置至关重要。当40个头分布在4个设备上时，每个设备最初持有序列四分之一部分的全部40个头。通过全连接交换（all-to-all exchange）重新分配后，每个设备将持有10个头，且每个设备都拥有完整的序列。我们称这种布局为头本地（head-local）布局。

在交换前进行排列可能需要额外通信，因为一个头的序列仍然分布在多个设备上。交换后，设备本地重排所需的所有token都已就绪。因此我们选择在输入头交换后执行排列，运行稀疏注意力，然后在返回交换前恢复原始token顺序。路由计算在交换前完成，每个头都携带自己的路由标志。

在TPU v6e上针对单个捕获步骤和层的匹配实验中，将token重排移至头本地区域，使总注意力延迟从75.64ms降至46.81ms，减少了38%。这包括路由、排列、通信和稀疏核的开销。这说明了布局转换的位置与稀疏核本身同样重要。

端到端结果与启示

在优化稀疏核及其相关开销后，我们测量这些改进如何转化为更快的去噪效果。我们评估了720p视频生成（81帧，40次去噪步骤，8个TPU v6e芯片）。表格比较了三种保留逐渐减少注意力对的调度方案，使用相同提示和种子，以及针对每种配置的本地密集控制。去噪时间取三次预热运行的中位数，加速比计算为密集时间除以SVG时间。FFmpeg PSNR衡量编码SVG与密集视频的相似度，数值越高表示输出越接近。

表3：端到端720p去噪延迟（81帧，40步在8个TPU v6e芯片上）及不同SVG稀疏度调度下的输出质量。

激进调度将去噪时间从153.50秒减少到119.86秒，加速1.28倍，延迟降低22%。这些结果反映了单独调优的密集和稀疏配置的组合实现，包括它们的分布式布局。由于去噪还包括稀疏注意力之外的工作，因此收益小于孤立的核加速。

随着序列长度增加，自注意力在transformer块延迟中的占比持续上升（从720p的55.5%上升到1080p的72.5%和1440p（2K）的88.2%）——稀疏注意力的端到端加速效果会随着分辨率显著提升。在不同分辨率下评估激进调度方案，可以观察到这些节省在更长序列上如何叠加：

表4：在八块TPU v6e芯片上，激进的SVG稀疏计划下，不同分辨率的端到端去噪延迟和质量扩展情况。

在1080p（171K tokens）分辨率下，激进的计划将去噪时间从683秒缩短至457秒（在24.66 dB PSNR下实现1.49倍加速）。在1440p（302K tokens，序列长度为720p的四倍）分辨率下，去噪时间从2,471秒降至1,461秒——在保持24.05 dB PSNR的同时实现1.69倍端到端加速，每生成一段视频节省超过16分钟。

Sorry, your browser doesn't support playback for this video

图10：密集注意力（左）与激进SVG稀疏计划（右，1.28倍加速，24.80 dB PSNR）下720p 81帧视频生成的并排对比。

关键结论：TPU上的高效稀疏注意力取决于硬件能够避免的工作量，而不仅仅是掩码移除了多少注意力对。跳过空白图块、限制图块内掩码以及优化token排列以提高访问效率，帮助我们将算法稀疏性转化为实际的端到端生成加速。

posted in:

• AI

• Cloud

• Tutorials

• Case Studies

• Learn

• Solve

Previous

Next

Related Posts

List

AI

Cloud

Tutorials

Learn

通过Genkit Go的Agent Skills实现按需专家能力

2026年7月31日

Case Studies

How-To Guides

在MaxText中复现Olmo 3 7B预训练：TPU大规模训练案例研究

2026年9月24日

Announcements

Problem-Solving

使用Google Cloud API Gateway将REST API转换为MCP工具

Gemini Enterprise Agent Platform的Agent和模型评估现已正式发布

导航点 /