2018年，于尔根·施密德胡伯（Jürgen Schmidhuber）——他早在1990年就首次提出了世界模型作为机器学习概念——与大卫·哈（David Ha）共同发表了一篇论文。他们在文中描述了一种"预测性世界模型"，该模型能够"提取空间和时间的有用表征"，并利用这些表征训练智能体完成驾驶等任务。自这篇开创性论文发表以来的八年间，世界模型在能力和实际应用方面的增长令人瞩目。
海量训练数据集的出现，加上变分自编码器，以及扩散和自回归Transformer的结合，使得今天的世界模型能够将文本、图像、音频和视频输入转换为潜在空间，进而构建和维护令人惊叹的沉浸式世界模型，这些模型可以预测状态变化并对动作做出响应。这些模型在机器人技术、交通运输、气候建模、视频游戏开发、科学模拟和设计原型制作中发挥着至关重要且不断扩大的作用。
世界模型能力和潜力的迅速提升，伴随着对能够支撑这些模型的硬件和基础设施需求的增长。这正是Reactor联合创始人兼首席技术官布莱斯·施密特钦（Bryce Schmidtchen）非常熟悉的挑战。
"这关乎实时性，关乎低延迟，也关乎如何以尽可能高效的方式实现规模化，"施密特钦观察到。Reactor是一个允许开发者、设计师和研究人员部署、使用和扩展实时交互式AI模型的平台。"效率意味着从如何在特定芯片（在我们的案例中是Trainium）上调度推理，到如何最大化地打包模型的每一次前向传播的方方面面。"
Reactor的全球业务布局使效率挑战变得更加严峻。"我们以区域为单位来思考GPU集群——我们在全球拥有数百个集群。从GPU输出的转换为像素的潜在数据，需要能够直接进入网络，而不必在Kubernetes中来回跳转。你需要以这样一种方式来实现这一点：无论是维护良好的集群还是机房里的裸机，都不会有区别。
"然后，"施密特钦继续说道，"你还需要将其与世界一流的网络和媒体流技术结合起来，支持不同的编解码器、不同的分辨率，并允许其连接到跨不同语言的API和SDK，从而支持所有不同类型的应用程序。"
AWS的全球覆盖范围已经帮助Reactor实现了许多目标。"他们在推理层拥有的一切，我们能够在不同区域实现的规模，正是这种规模的平台之所以可能且可靠的原因，"施密特钦观察到。
如今，Reactor与Amazon Neuron科学团队最近的合作，为世界模型开辟了一条实现更高效率的道路。以下是这些团队所追求的优化策略，以及这如何为未来模型在Trainium上实现实时能力奠定了基础。
自回归扩散模型的兴起
"Neuron科学团队探索用于生成式AI模型能力优化的新技术，"Neuron团队的首席应用科学家吴军（Jun Wu）解释道。"这可能是一种新的模型架构，或是一种新的优化算法，或是一种在Trainium上运行这些模型时生成和优化代码的新方法。我们识别出机会，构建原型，并使其能够可行地移植到生产管道中。"
该团队在视频生成中扩散模型的使用方面发现了这样一个机会。"我们注意到从生成较短的固定长度视频演变为生成无限长度或动态长度视频的趋势，"吴解释道。"要在这些长度下生成高质量的视频帧，就催生了自回归扩散模型的使用。"
这类模型结合了大语言模型中的顺序下一Token预测与扩散模型中的迭代和细化能力，对于希望在生成环境中自由导航的用户来说至关重要。
"自回归扩散意味着用户的按键操作可以被模型作为输入吸收，并以之前生成的视频帧为条件，模型可以正确判断下一步动作或下一个场景，"吴观察到。"我们今天看到的大多数交互式视频生成模型都在使用这种方法。"
吴和他的团队成员——两位高级应用科学家梅森·付（Mason Fu）和林凡·于（Lingfan Yu）——看到了优化这些模型在Trainium上实时视频生成部署方式的机会，并抓住了与Reactor合作的契机。"当时我们已经在实时视频和交互式视频生成方面工作了18个多月，"施密特钦指出。团队考虑了各种视频生成技术，并一致选择使用Rolling Forcing，理由是其能够持续生成高质量的30秒视频，以及其相对较小的规模。
这些模型的难点不在于质量，而在于必须实时运行。与传统视频生成不同——后者离线渲染完整片段然后返回——像Rolling Forcing这样的模型是流式的。每一帧都是生成后立即被消费，展示给用户或作为下一个输入反馈回去。这就是开发者和客户实际使用它们的方式，迟到的帧会破坏体验，因为生成过程必须始终领先于播放时间线。"高质量的生成扩散模型面临序列很长的挑战，这需要大量内存消耗，"吴解释道。"Rolling Forcing相对较小，但其序列长度非常大。"这种小模型、超长序列和硬性帧率下限的组合，正是实时性要求如此苛刻的原因所在。Rolling Forcing能够以每秒16帧的速度生成视频——这是视频播放的标准——这意味着它也满足了延迟要求。这正是Trainium能够增添价值之处，它提供了足以支撑这种长序列工作负载的性能和内存，能够以超过16帧/秒的速度实现实时生成，而不仅仅是最终生成出好的画面。因此，团队开始致力于在Trainium上实现Rolling Forcing，作为自回归扩散视频生成的代表案例。
对开发者友好的方法
Neuron科学团队和Reactor团队采用了以内核为中心、自下而上的方法，旨在让开发者的工作更加轻松。他们聚焦于三大挑战——动态形状、非常规内存访问模式和繁重的缓存管理——这些都是实时视频生成给通用编译器带来的难题。这些挑战在每一次前向传播中都会重复出现，因此在其他工作负载中可能微不足道的延迟，在实时视频生成中却会累积成延迟故障。
例如，动态形状带来的挑战部分源于实时视频生成中场景的变化本质：想象一个镜头，展示一位投手独自站在投手丘上，几秒钟内画面就扩展到展示其他球员，再到看台上成千上万的球迷。实时视频生成还涉及在持续滑动窗口过程中生成和淘汰的帧。"这意味着注意力长度在各次前向传播中会有所变化，并且每个窗口有两个不同的传播过程：去噪，然后是缓存清理，"吴表示。"这一切都使得静态编译变得困难。"
除了滑动KV缓存复制之外，视频生成工作负载还包含一些操作——旋转位置嵌入（RoPE）和注意力转置——它们的内存访问模式是特定于工作负载的，这使得任何通用编译器都难以对其进行充分优化。例如，在视频生成中，RoPE必须处理三个轴（高度、宽度和时间），而不是它在大语言模型中所处理的单一位置。"3D旋转嵌入沿最内层维度交错排列奇数和偶数元素，当以通用方式编译时会产生大量微小的数据传输，"吴解释道。
最后，由于KV缓存在每一层都会发生——每一层都要读写一个滚动的KV缓存——因此设备上的复制操作需要很高的吞吐量。每一层每一步都需要读取旧缓存内容并写入新内容，这对内存造成了另一个显著的拖累。
使用Neuron内核接口
为了解决这一额外的复杂性，Neuron科学团队转向了Neuron内核接口（NKI）。
"NKI让开发者能够编写直接在NeuronCore硬件上运行的计算内核，精确控制数据在内存和计算引擎之间的移动方式，"吴解释道。他指出，当性能分析显示仅仅编译模型是不够的时候，NKI为开发者提供了一条自助路径，可以直接修复热点问题，用针对硬件调优的实现来替代特定的瓶颈操作。
"在我们所做的工作中，3D-RoPE内核的耗时从5秒降至1.8毫秒，每层的缓存复制从23毫秒降至1.9毫秒，注意力转置则通过将其融合到注意力内核中而被完全消除，"吴指出。"对于每一毫秒都至关重要的实时工作负载而言，这种直接的硬件访问正是实现生产级性能的关键。"
此外，用于生成NKI内核的智能体NKI-Dev-Suite，在第一次尝试时就成功生成了一个可用的3D-RoPE内核。"综合效果是：该流水线使用了11GB的高带宽内存，而标准的即时模式路径则会耗尽内存，"吴指出。
混合分片策略
如前所述，在实时性要求下，视频扩散模型产生的Token序列远比文本模型更长。这使得自注意力机制的挑战更加突出。"这里的自注意力操作涉及23400个查询Token关注32760个上下文Token，占据约70%的计算时间，"于观察到。"没有单个核心能够高效地处理这项工作，除非将工作分布开来。"
为解决这一问题，团队采用了一种混合分片策略，包括序列并行（SP，即按顺序划分数据）和张量并行（TP，即沿特定维度分片张量，将计算分布到多个设备上）。对于模块中某些非自注意力的部分，团队采用了序列并行。然而，对于自注意力部分，只有混合方法才足够有效。
"大语言模型的注意力是因果且单向的——每个Token只关注之前的Token，KV缓存单调增长，每层只有一种注意力类型和单一的缓存策略，"吴表示。然而，Rolling Forcing在每个块中有两种注意力类型：用于跨帧时空一致性的自注意力，以及用于文本条件和活动窗口内双向注意力的交叉注意力，因为所有帧都是从噪声中联合细化的。
这一点，再加上双策略KV缓存（用于近期上下文的滑动窗口，加上用于全局上下文的永久注意力汇聚点）、每个窗口两次前向传播（去噪写入含噪声的KV条目，然后缓存更新过程用干净值覆盖它们，确保未来的窗口始终关注干净的上下文），以及3D视频Token结构，共同限制了序列如何在各核心之间进行拆分。
于解释道，仅使用TP会带来一个数学难题。"WAN扩散Transformer模型有12个注意力头，但我们每块芯片只有8个Neuron核心，无法整除。仅使用TP意味着你必须进行填充，但填充会浪费计算资源。"
另一方面，仅使用SP可能会破坏3D结构，正如于所指出的，"你无法保证序列分区恰好落在帧边界上。你的数据必须至少达到帧的粒度，但如果你在帧边界上进行分区，那些操作就无法正常工作。"这种混合方法将注意力头分散到4个核心上，将序列分散到2个核心上，从而保持了数学计算和数据布局的正确性。"VAE解码器使用了空间W轴分片，实现了8.25倍的超线性加速，"于表示。
模型结构变化
Reactor和Neuron科学团队还优化了Rolling Forcing模型代码的部分内容，以便在Trainium上更高效地运行。"基本上，该模型有两个阶段：一个是扩散，另一个是缓存更新，"于说道。"它们在两次独立的运行中执行，但问题在于缓存更新阶段的计算量明显更少，因此如果在单独的一轮中执行它，硬件利用率会低得多。这是因为我们也必须对其进行分片，所以从原理上讲，输入芯片的数据越少，硬件利用率就越差。"
团队优化了模型代码，将这两个阶段所共有的组件批处理在一起。
"当我们遇到略有不同的组件时，我们会再次进行拆分，然后分别处理不同的组件。但对于流水线的大部分内容，它们是批处理在一起的，"于解释道。"这对于Trainium来说尤其有用，因为每个实例有16块芯片，每块芯片有8个核心。如果你将计算分配到过多的核心上，每次调用只会获得少量的分区计算，这就造成了产能利用不足。"
成果
在采用了这些以及其他优化措施之后，Reactor和Neuron科学团队成功在首次端到端运行中就生成了正确的视频，利用Trainium实现了实时模型。团队还强调，这些成果是可推广的。
"Rolling Forcing是这条流水线的核心，它是一个非常小的模型，"首席解决方案架构师亚哈夫·比兰（Yahav Biran）表示。"你可以在它上面快速迭代，但它仍然是一个端到端的健壮系统。它具备健壮系统所拥有的全部复杂性：编码器、DIT、VAE、解码器。所以基本上，如果你采用一个更健壮的系统，它是在相同的构建模块上运行的。"
"我们正在为采用自回归扩散且同时具有实时交互需求的模型构建通用技术，"于补充道。"我们并不仅仅是在优化单一模型。我们正在构建通用技术，以支持所有具备实时流媒体需求的模型。"
未来展望
施密特钦表示，他对这类工作可能带来的未来感到兴奋。"在不远的将来，每一个像素都将实时、交互式地生成，"他说道。"世界模型将能够实时创作完整的故事：这些故事能够做出反应，你可以参与其中，甚至可以随时改变整个场景。"
他还指出，亚马逊正在进行的以及已经完成的工作，将在很大程度上帮助这些愿景成为现实。
"Trainium清楚地展现了AWS和亚马逊整体的巨大投入，"施密特钦指出。"未来的路线图更加清晰：更高的性能、更好的性价比、更大的全球规模。在这个需要将实时交互式AI大规模分发给消费者、物理AI等领域的未来，Trainium在推理层面——无论是并行化、内存还是软件栈方面——都处于非常有利的位置，并能够很好地与AWS的全球规模基础设施集成。我们很高兴能继续与亚马逊紧密合作，探索世界模型尚未开发的潜力。这仅仅是个开始。"
Q&A
Q1：Trainium芯片在实时视频生成中起到了什么作用？
A：Trainium为长序列扩散视频生成工作负载提供了足够的性能和内存支撑，使模型能够以超过每秒16帧的速度实现真正的实时生成，而不仅仅是最终产出高质量画面。
Q2：什么是自回归扩散模型？它有什么特点？
A：自回归扩散模型结合了大语言模型的顺序预测能力与扩散模型的迭代细化能力，能够根据用户的实时输入（如按键操作），基于之前生成的画面来正确判断并生成下一个场景，常用于交互式视频生成。
Q3：Neuron内核接口（NKI）为什么能大幅提升性能？
A：NKI允许开发者编写直接在硬件上运行的计算内核，精确控制数据流动。使用NKI后，3D-RoPE内核耗时从5秒降至1.8毫秒，缓存复制从23毫秒降至1.9毫秒，显著提升了实时视频生成的效率。
