How Visko streams real-time world models on Google Cloud TPUs

Visko如何在Google Cloud TPUs上流式传输实时世界模型 | Google Cloud 博客

媒体与娱乐

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

2026年10月16日

##### Will Yin

Visko首席执行官

##### Ivan Nardini

Google Cloud开发人员关系工程师

##### 立即体验Gemini

职场人工智能的入口

立即体验

编辑注：生成式AI初创公司正在迅速超越静态片段生成，构建交互式实时世界。在这次技术深度解析中，我们将探讨AI初创公司Visko如何与Google Cloud TPU工程师合作，将Orbis——其持续流式传输的Live World Model——部署到Cloud TPUs上，仅用30天就实现了从初始启动到生产交接的全过程。

实时视频世界模型在生产环境中为何如此难以部署？瓶颈并非平均吞吐量，而是如何应对持续不断的实时时钟压力。与离线视频生成器用户提交提示并排队等待不同，实时世界模型需要持续流式传输。每一段新视频和同步音频必须在观众看完前一段之前完成渲染，否则世界会陷入冻结。

今天，我们很高兴分享Visko——这家开创交互式世界模型的AI初创公司——如何与Google Cloud合作解决这一挑战。双方工程师并肩协作，将Visko的端到端实时推理流水线（结合视频生成、超分辨率放大和同步音频）部署到单个Cloud TPU切片上。通过团队合作，端到端生成速度降至1.07倍播放速度。

以下是我们如何在一个月内构建该流水线实现这一里程碑的内部洞察，推动性能提升的架构优化，以及为其他在Google Cloud上构建超低延迟生成式AI的初创公司提供的实用经验。

用实时世界模型重新定义交互式媒体

虽然大多数生成式视频平台专注于渲染固定预烘焙片段，Visko正在构建其称为Live Models的实时模型：能够持续生成并实时响应的视频世界模型。

Visko的旗舰模型Orbis可以从单个文本或图像提示开始流式传输实时视频。当视频播放时，观众可以实时输入新的导航提示，动态演变场景而无需重启或破坏连续性。例如，在一次测试会话中，直播画面从一只在壁炉旁打盹的猫开始；当观众引导场景引入一碗食物时，碗无缝出现，猫伸展身体、站起来并发出喵叫。

自2026年9月1日公开演示发布以来，Orbis在互动娱乐、直播电商、AI伴侣、沉浸式教育和机器人模拟等领域开启了全新可能性。为可靠扩展这些交互会话，Visko于2026年8月与Google Cloud TPU工程师合作，将Orbis实时流水线迁移至Google的Tensor Processing Unit（TPU）。在30天内，完整的多模态流水线已在TPU上实现实时运行。

Visko 的目标是为新一代交互式世界提供动力，人们可以实时进入并塑造这些世界。在严格时钟下流式传输连续视频和音频，与离线剪辑生成相比，是完全不同的系统挑战。通过与 Google Cloud 的 TPU 工程师并肩合作，我们仅用 30 天就将完整的 Orbis 管道迁移至 TPU —— 这使我们能够实现所需的低延迟计算和内存带宽，从而在大规模上将实时模型带给创作者和开发者。

"Visko 正将一类激动人心的新实时可操控工作负载引入 Cloud TPU 生态系统。通过使用 JAX、torchax 和自定义 Pallas 内核，在单个 TPU 切片上联合开发其实时多模态管道，我们证明了精益创业团队从 PyTorch 原型到在 Google Cloud AI 超级计算机上以 1.07 倍播放速度渲染的转化速度。" —— Google Cloud AI 开发者倡导负责人 Dave Elliott

实时视频是硬截止期限问题

当你提供离线剪辑生成器时，延迟是可容忍的 —— 用户只需等待一次，而后台处理帧。但实时世界模型没有等待室。

每个视频块都必须在会话持续期间以比实时播放速度更快的速度完成去噪、从潜变量中解码、放大至 1664x936（18fps）并配以生成的音频。

两个结构性约束使实时世界模型比标准批量推理更具挑战性：

• 批量大小为 1：每个实时流代表单个用户的交互会话。你无法通过将数十个并发请求分组到大批量中来掩盖延迟 —— 这是保持 AI 加速器饱和的标准策略。

• 严格的顺序依赖性：每个块直接依赖于前一个块。由于视频模型维护着最近视频潜变量的滚动内存以确保块 N+1 从块 N 的动作平滑延续，你无法按任意顺序预渲染块或跨未来时间戳并行处理。

对于来自大语言模型服务的初创基础设施工程师，将实时世界建模视为每步延迟 SLO（服务等级协议）而非每秒令牌吞吐量基准会更有帮助。

在单个 TPU 切片上构建多模态管道

为在 Visko 扩展时保持基础设施精简和成本效益，团队设计了实时管道，将三个不同的生成模型链接在单个 Cloud TPU 切片上运行（精度为 BF16）。

当用户提交初始提示或实时操控更新时，核心视频扩散模型、超分辨率放大器和音频生成器这三个生成模型会共享同一芯片，轮流在硬件上处理每个块。

图 1：Visko 的端到端 Orbis 实时管道在单个 Cloud TPU 切片上按顺序链接视频扩散、潜变量解码、超分辨率放大和纯 JAX 音频生成，每个块的处理速度都快于播放速度。

为结合快速启动迭代与裸机加速器速度，软件栈基于三个开源支柱构建：

• JAX：Google 的高性能 Python 数组库，通过 XLA 编译器编译为优化的机器代码。

• torchax：通过 JAX 在 Cloud TPU 上直接运行 PyTorch 模型的互操作层，使初创团队无需从头重写即可迁移现有 PyTorch 组件。

• Pallas：JAX 中用于 TPU 的自定义内核编程语言——相当于在 GPU 上编写自定义 Triton 或 CUDA 内核的 TPU 对应方案。

实时性能分析：达到时钟频率目标

通过在 BF16 的单个切片上对每个阶段进行性能分析和优化，工程团队在保持约 7% 的余量以吸收实时流媒体中的变化的同时，实现了每块数据的总延迟低于目标值：

图 2：单个 TPU 切片上每块数据延迟预算图表。

由于视频去噪占总计算预算的近三分之二，该阶段成为我们内核和并行优化的主要关注点。

技术栈解析：实现实时流媒体的三大优化

#### 1. 所有 TPU 芯片上的 Ulysses 上下文并行

视频块在视频变压器中会扩展为大量的时空标记序列。为了在批量大小为 1 的情况下将工作负载分布到 TPU 芯片上，团队实现了 Ulysses 上下文并行。

与拆分模型权重不同，Ulysses 在芯片之间拆分标记序列。在注意力层之前，高速的全连接互连交换将张量重新分组，使得每个芯片持有完整的视频序列，但仅对应四分之一的注意力头。每个芯片在分配的头中对完整序列计算注意力后，第二次全连接交换会恢复前馈层的序列拆分布局。

图 3：批量大小为 1 时，TPU 芯片上的 Ulysses 上下文并行。高速全连接芯片间互连（ICI）交换在前馈层的序列拆分布局与 Pallas splash 注意力的头拆分布局之间转换张量。

#### 2. 来自 MaxDiffusion 的自定义 Pallas splash 注意力

为了加速注意力计算本身，团队集成了来自 MaxDiffusion（Google Cloud 的 Cloud TPU 扩散模型开源参考库）的自定义 Pallas splash 注意力内核。

Splash 注意力将 FlashAttention 风格的分块技术引入 TPU 架构：无需在高带宽内存（HBM）中生成完整的二次注意力矩阵，而是在超快片上内存（VMEM）中以紧凑的分块形式计算注意力。初创团队可通过 custom_splash_attention.py 直接将此端口使用的精确内核插入到自己的流水线中。

#### 3. 通过从 torchax 迁移到纯 JAX 实现 5.8 倍音频加速

我们首先使用 *torchax* 将所有三个 Orbis 模型部署到 Cloud TPU 上，*torchax* 通过 JAX 和 XLA 在 TPU 上运行 PyTorch 代码而无需重写。对于视频和超分辨率模型，torchax 配合 Pallas 注意力内核直接达到了延迟目标，最终流水线中这两个模型仍使用 torchax。

只有音频模型需要重写。将音频阶段从 torchax 迁移到 *原生 JAX* 使其速度提升了 *5.8 倍*，并使整个流水线在播放截止时间之前完成。

初创构建者提示：战略性迁移。关键点不是“跳过 torchax”——而是将其作为快速启动的桥梁。*torchax* 让我们的 PyTorch 模型在几天内而非几周内运行在 TPU 上，为我们提供了定位真实瓶颈所需的端到端遥测数据。先测量完整流水线，再将工程周期投入那些主导您延迟预算的组件重写中。

AI 构建者实地笔记：Cloud TPU 教训

在 Cloud TPU 上调整实时多模态推理的要点：

• * 单流视频的Ulysses上下文并行:* 在批量大小为1时，使用all-to-all互联交换将令牌序列跨芯片拆分的速度，优于张量并行和环形注意力两种方式。

• * 连续流的原生BF16支持:* FP8量化可以降低内存带宽压力，但连续流中各块之间的微小量化误差会逐渐累积。原生BF16在延迟预算范围内保持了视觉质量的完整性。

• * 选择性JAX端口:* torchax将视频和超分模型保留在PyTorch中，因此团队只需将音频瓶颈部分用原生JAX重写。

• *XLA编译缓存:* 首次启动时需要花费几分钟编译XLA图。通过缓存编译后的工件，可将工作节点重启时间控制在1分钟以内。

在Cloud TPU上构建实时生成模型？从这里开始

如果您的初创公司正在突破实时视频、世界模型或超低延迟多模态推理的边界，您可以使用与Visko端口相同的工具包进行开发：

• 开箱即用加速扩散注意力：获取MaxDiffusion合作中使用的开源splash注意力内核，通过Pallas文档编写自定义TPU内核进行探索。

• 快速迁移PyTorch模型并选择性优化：使用torchax快速将现有PyTorch模型迁移至TPU，通过Cloud TPU Profiler定位瓶颈，将关键延迟阶段移植到原生JAX。

• 按需启动TPU：使用TPU Flex-start虚拟机按需分配TPU资源，无需长期预留即可进行短期开发和基准测试。

• 与Google合作扩展：寻求云积分、技术指导以及访问Google Cloud AI超算堆栈的早期AI初创公司，可申请Google for Startups云计划。

特别感谢Visko.ai团队的紧密技术合作，以及领导该项目的Google Cloud TPU工程师团队，包括：Yixin Shi, James Huang和Elisa Tsai。

发布于

• 媒体与娱乐

• 人工智能与机器学习

• 客户案例