Google称TPU上的稀疏视频注意力提速2.4倍

Google发布了一项案例研究，介绍如何加速视频扩散模型中的注意力计算——这是生成长视频、高分辨率视频片段时开销较大的一步。其JAX和Pallas内核采用稀疏注意力，并让注意力掩码与TPU处理数据时使用的计算分块对齐。

在单颗TPU v6e芯片上的测试中，最终版本耗时32.76毫秒，而稠密Splash注意力耗时78.70毫秒；据称速度提升了2.40倍。测试使用合成输入，未计入路由、令牌重排和设备间通信，因此不能证明端到端视频生成速度有所提升。