Visko 如何在 Google Cloud TPU 上流式传输实时世界模型

📌 One-Sentence Summary

Visko 与 Google Cloud 合作，在 30 天内将其 Orbis 实时世界模型部署在 TPU 上，通过 Ulysses 上下文并行和自定义 Pallas 内核等架构优化，实现了 1.07 倍的播放速度。

📝 Summary

本文技术深度深入探讨了 AI 初创公司 Visko 如何解决实时交互式世界模型的服务挑战。与生成静态片段的离线视频生成器不同，Visko 的“实时模型”能够流式传输响应用户指令连续视频和音频。文章详细介绍了如何使用 JAX、torchax 和 Pallas，在单个 TPU 切片上构建多模态流水线——结合视频生成、超分辨率放大和同步音频——从而在批次大小为 1 的情况下满足严格的延迟 SLO。

💡 Main Points

实时视频建模是硬性截止的延迟驱动问题

与用户等待队列的离线生成器不同，实时模型必须比实时播放速度更快地渲染片段（去噪、解码、放大），防止冻结。

结构限制包括批次大小为 1 和顺序依赖

每个流都是独立的会话，无法进行标准的基于批次的优化，且每个片段都依赖于前一个片段的潜变量，无法在时间上进行并行化。

Ulysses 上下文并行实现了跨芯片的高速推理

Ulysses 不是划分模型权重，而是将 Token 序列划分到各个芯片，允许每个芯片在持有完整序列的同时处理特定的注意力头。

自定义 Pallas Attention 优化了内存带宽使用

通过在 TPU 架构上使用类似 FlashAttention 的分块技术，团队在片上内存中计算注意力，而不是在高带宽内存 (HBM) 中实例化大型矩阵。

从 torchax 策略性迁移到原生 JAX 加速

Visko 使用 torchax 作为 PyTorch 模型的快速桥梁，但使用原生 JAX 重写了音频模型，从而实现了 5.8 倍的速度提升。

💬 Key Quotes

实时世界模型没有候室。

每一个片段都必须完成去噪、从潜变量解码、放大到 1664x936（18 fps）并与生成的音频配对，且速度快于实时播放速度。

启示不是‘跳过 torchax’ —— 而是将其用作快速启动的桥梁。

📊 Article Meta

AI Screening： 92

Featured： Yes

Source： Google Cloud Blog

Author： Will Yin, Ivan Nardini

Category： 人工智能

Language： 英文

Read Time： 7 min

Word Count： 1731

Tags：

AI 与智能应用 , 机器学习 , AI 硬件与芯片 , AI Agent , 世界模型