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

通过 JAX 与 Pallas Splash Attention 内核，开发团队在 TPU v6e 上成功将视频扩散模型的稀疏时空注意力转化为硬件推理加速。针对 81 帧 1440p 视频下自注意力耗时占比达 88.2% 的瓶颈，方案结合 Sparse VideoGen 划分空间与时间头，仅保留约 38.87% 的 query-key 对，并通过硬件瓦片跳过机制大幅节省计算与带宽开销。