A kernel-centric path to real-time video generation on Trainium

以内核为中心的Trainium实时视频生成方法 - Amazon Science

AWS Neuron Science团队与Reactor合作，将实时视频生成功能引入Trainium。

机器学习

以内核为中心的Trainium实时视频生成方法

通过Neuron Kernel Interface，Reactor与AWS合作解决了动态形状、内存访问模式和缓存管理等挑战，这些挑战使得实时自回归扩散变得困难——开发了可跨模型通用的技术。

作者：Stephen Zorio 2026年9月25日 12分钟阅读

分享

• 复制链接
• 邮件
• X
• 领英
• Facebook
• Line
• Reddit
• QZone
• 新浪微博
• 微信
• WhatsApp

分享到微信 x

2018年，Jürgen Schmidhuber（他在1990年首次提出世界模型作为机器学习概念）与David Ha发表了一篇论文。他们讨论了“预测性世界模型”，该模型可以“提取空间和时间的有用表示”，并利用这些表示来训练代理执行包括驾驶在内的任务。自这篇开创性论文发表以来的八年里，世界模型的能力和实际应用的快速增长令人震惊。

大规模训练数据集的出现，结合变分自编码器，以及扩散模型和自回归变换器的出现，使当今的世界模型能够将文本、图像、音频和视频输入转换为潜在空间，用于构建和维护令人惊叹的沉浸式世界模型，这些模型可以预测状态变化并响应操作。这些模型在机器人技术、交通运输、气候建模、视频游戏开发、科学模拟和设计原型等领域发挥着至关重要的作用，并且应用范围不断扩大。

世界模型可以生成沉浸式、可探索的环境，能够实时响应用户输入。

世界模型的能力和潜力的显著提升，伴随着对能够支持它们的硬件和基础设施需求的增加。Bryce Schmidtchen（Reactor联合创始人兼首席技术官）对此非常熟悉。

Schmidtchen指出：“这关乎实时性，关乎低延迟，关乎在规模上尽可能高效地实现这些目标。”Reactor是一个平台，允许开发人员、设计师和研究人员部署、使用和扩展实时交互式AI模型。“效率涉及从如何在给定芯片（在我们的情况下是Trainium）上调度推理，到如何考虑将模型的每次前向传递最大化地打包到硬件中。”

Reactor的全球存在使得效率挑战更加严峻。“我们以区域为单位考虑GPU集群——我们全球有数百个集群。从GPU生成的潜在表示转换为像素并传输到网络时，必须能够直接传输而无需通过Kubernetes跳转。并且无论集群是得到良好支持的还是封闭空间中的裸金属，都必须以相同方式实现。”

Schmidtchen继续说道：“然后，你需要将这些与世界级的网络和媒体流技术相结合，支持不同的编解码器、不同的分辨率，并能够连接到跨不同语言的API和SDK，以支持各种类型的应用。”

Reactor的平台提供跨多种语言和编解码器的API和SDK，支持在AWS上运行实时交互式AI应用——从游戏和机器人技术到影视和研究。

AWS 全球基础设施已帮助 Reactor 实现许多目标。Schmidtchen 表示：“他们推理层拥有的所有能力，以及我们在不同地区实现的规模水平，正是让这种规模的平台成为可能并保持可靠性的关键。”

Reactor 与亚马逊 Neuron Science 团队最近的合作，为世界模型开辟了更高效的新路径。本文将探讨这两个团队所采用的优化策略，以及这些策略如何为未来模型在 Trainium 上实现实时能力奠定基础。

自回归扩散模型的兴起

Neuron 团队首席应用科学家 Jun Wu 解释道：“Neuron Science 团队探索生成式 AI 模型优化的新技术，这可能涉及新的模型架构、优化算法，或是运行这些模型时生成和优化代码的新方法。我们会识别出可以构建原型并实现生产流水线移植的机遇。”

团队在视频生成中使用扩散模型时发现了这样的机会。Wu 解释称：“我们注意到视频生成从较短固定长度向无限长度或动态长度演变。在这些长度上生成高质量视频帧，催生了自回归扩散模型的应用。” 这类结合大语言模型中顺序预测机制与扩散模型迭代优化特性的模型，对需要生成可交互环境视频的用户至关重要。

Wu 表示：“自回归扩散意味着用户输入的按键可以被模型吸收，模型可根据之前生成的视频帧，准确判断下一步动作或场景切换。” 他指出：“目前大多数交互式视频生成模型都在使用这种技术。”

Wu 与团队成员、高级应用科学家 Mason Fu 和 Lingfan Yu 看到了在 Trainium 上优化这些模型部署以实现实时视频生成的机会，并希望与 Reactor 合作推进。Schmidtchen 表示：“在那之前，我们已经在实时视频和交互式视频生成领域投入了超过 18 个月的研发。” 团队评估了多种视频生成技术，最终选择采用 Rolling Forcing 方法，因其既能稳定生成高质量 30 秒视频，又具有相对紧凑的模型规模。

实时视频生成模型的难点不在于质量，而在于必须实时运行。与传统视频生成技术不同，后者会离线渲染完整视频片段并稍后返回，而像Rolling Forcing这样的模型是流式处理的。每一帧生成后会立即被消费，展示给用户或作为下一帧的输入反馈。这正是开发者和客户实际使用的方式，任何延迟到达的帧都会破坏体验，因为生成速度必须始终领先于播放时间线。"高质量生成扩散模型面临非常长序列的挑战，这需要消耗大量内存，"吴解释道，"Rolling Forcing模型体积相对较小，但序列长度非常长。"这种小模型、超长序列和严格的帧率下限的组合，使得实时运行变得极具挑战性。Rolling Forcing每秒生成16帧的能力，达到了视频播放的标准，这也满足了延迟要求。这时Trainium的价值就体现出来了，它通过提供足够的性能和内存，在保持速度的同时持续处理长序列工作负载，实现每秒超过16帧的实时生成，而不仅仅是最终产出高质量帧。

开发者友好的方法

Neuron Science和Reactor团队采用了以内核为中心的自底向上方法，旨在简化开发者的工作。他们专注于实时视频生成对通用编译器提出的三个挑战——动态形状、非常规内存访问模式和繁重的缓存管理。这些挑战在每次前向传递中都会重复出现，因此在其他工作负载中可能微不足道的延迟，会在实时视频生成中累积成严重的延迟问题。

例如，动态形状带来的挑战部分源于实时视频生成场景的动态特性：想象一个镜头从特写投手单独在投手丘的画面，快速拉远展示其他球员，然后在几秒内展现成千上万的观众。实时视频生成还涉及通过连续滑动窗口过程生成和淘汰帧。"这意味着不同前向传递的注意力长度会变化，每个窗口会有两次独立的传递：去噪，然后缓存清理，"吴说道，"所有这些都使静态编译变得困难。"

除了滑动KV缓存复制，视频生成工作负载还包含一些操作——旋转位置嵌入（RoPE）和注意力转置——它们的内存访问模式具有工作负载特定性，这对任何通用编译器来说都难以完全优化。例如，在视频生成中，RoPE必须处理三个轴（高度、宽度和时间），而不仅仅是LLM中单个位置。"三维旋转嵌入在最内层维度上交错奇数和偶数元素，当通用编译时会产生大量微小的数据传输，"吴解释道。

最后，由于KV缓存发生在每一层——每一层都会读写滚动KV缓存——因此设备内复制需要高吞吐量。在每一层的每一步都需要读取旧缓存内容并写入新内容，这也成为内存性能的另一个显著瓶颈。

使用Neuron内核接口

为了解决这一额外复杂性，Neuron Science团队转向了Neuron内核接口（NKI）。"NKI使开发者能够编写直接在NeuronCore硬件上运行的计算内核，并对数据在内存与计算引擎之间的移动进行精确控制，"吴解释道。他指出，NKI为开发者提供了一条自助式路径，可直接修复热点问题，通过将特定瓶颈操作替换为经过硬件调优的实现（当性能分析显示仅编译模型已不足以解决问题时）。

"在我们所做的工作中，3D-RoPE内核的执行时间从5秒缩短至1.8毫秒，缓存复制操作从每层23毫秒降至1.9毫秒，注意力转置操作则通过将其融合到注意力内核中被完全消除，"吴指出。"对于实时工作负载而言，每个毫秒都至关重要，这种直接访问硬件的能力正是实现生产级性能的关键。"

此外，NKI-Dev-Suite——一个用于生成NKI内核的工具——在首次尝试时就成功生成了可用的3D-RoPE内核。"综合效果是：该流水线使用了11GB的高带宽内存，而标准的急切模式路径则因内存不足而无法运行，"吴指出。

混合分片策略

如前所述，视频扩散模型在实时要求下生成的标记序列远长于文本模型。这使得自注意力机制的挑战更加严峻。"此处的自注意力机制需要处理23,400个查询标记对32,760个上下文标记的注意力计算，占用了约70%的计算时间，"于指出。"没有将工作分布到多个核心上，单个核心无法高效处理这一任务。"

为了解决这一问题，团队采用了一种混合分片策略，包含序列并行（SP）——即按顺序划分数据——和张量并行（TP），后者沿特定维度对张量进行分片，以在多个设备间分配计算。对于模块中非自注意力部分，团队使用了序列并行。然而，对于自注意力部分，只有混合方法才能满足需求。

"LLM的注意力机制是因果且单向的——每个标记仅关注之前的标记，KV缓存单调增长，每层只有一个注意力类型和单一缓存策略，"吴表示。然而，Rolling Forcing每个块有两个注意力类型：用于跨帧时空一致性的自注意力，以及用于文本条件和活动窗口内双向注意力的交叉注意力，因为所有帧都是从噪声中联合优化的。这结合了双策略KV缓存（一个用于近期上下文的滑动窗口，一个用于全局上下文的永久注意力池）、每个窗口两次前向传递（去噪过程写入噪声KV条目，随后通过缓存更新传递用干净值覆盖它们，确保后续窗口始终关注干净的上下文），以及3D视频标记结构，这些因素共同限制了序列在核心间的分割方式。

于解释称，单独使用TP会带来数学上的难题。"WAN扩散变换器模型有12个注意力头，但每块芯片只有8个Neuron核心，因此无法整除。单独使用TP意味着必须进行填充，但填充会浪费计算资源。"

模型结构变化

Reactor 和 Neuron Science 团队还优化了 Rolling Forcing 模型代码的部分组件，使其在 Trainium 上运行效率更高。"基本上，该模型有两个阶段：一个是扩散阶段，另一个是缓存更新阶段，" Yu 说道。"这两个阶段是通过两次独立运行执行的，但问题是缓存更新阶段的计算量显著更少，因此如果单独运行，硬件利用率会大幅下降。这是因为我们还需要对其进行分片，所以高层原则是：提供给芯片的数据越少，硬件利用率越差。"

团队优化了模型代码，将这两个阶段共有的组件进行批量处理。

"当我们遇到略有差异的组件时，我们会再次拆分并单独处理不同部分，" Yu 解释道，"但对大多数流水线组件而言，它们都会被批量处理。这对于 Trainium 来说特别有用，因为每个实例有 16 个芯片，每个芯片有 8 个核心。如果你将计算分散到过多核心上，每次调用只会进行少量分片计算，这会导致容量利用率低下。"

结果

在应用这些优化和其他改进后，Reactor 和 Neuron Science 团队成功在首次端到端运行中生成了正确的视频，利用 Trainium 实现了实时模型。团队强调，这些成果具有通用性。

在单个 Trainium2 芯片上运行 Rolling Forcing 时，该合作项目实现了实时视频生成——团队表示这一成果可推广到所有自回归扩散模型。

"Rolling Forcing 是一个管道，它本身是一个非常小的模型，" 首席解决方案架构师 Yahav Biran 说道，"你可以快速对其进行迭代，但它仍然是一个端到端的稳健系统。它包含了稳健系统的所有复杂组件：编码器、DIT、VAE 和解码器。因此，如果你使用更稳健的系统，它本质上使用的是相同的构建模块。"

Yu 补充道："我们正在为采用自回归扩散的模型构建通用技术，这些模型还需要实时交互能力。我们不是在优化单一模型，而是在构建支持所有需要实时流式处理的模型的通用技术。"

未来展望

Schmidtchen 表示，他对于这类工作可能带来的未来感到兴奋。"在不太遥远的未来，每个像素都能实时、交互式地生成，" 他说道，"世界模型可以实时生成完整的故事：能够响应、可以参与互动，甚至能实时改变整个场景。"

他还指出，亚马逊已经开展并将继续推进的工作，将极大地推动这些愿景成为现实。

“Trainium 明显体现了 AWS 和亚马逊整体对人工智能领域的巨大投入，”Schmidtchen 表示。“我们看到了更清晰的性能提升路线图，包括更优的成本效益、更广泛的全球扩展能力。在未来的实时交互式 AI 场景中，当需要将 AI 服务大规模分发给消费者、物理 AI 等应用场景时，Trainium 在推理层的表现非常突出——无论是并行计算能力、内存性能还是软件栈设计，都能与 AWS 全球基础设施实现良好协同。我们非常期待与亚马逊继续深入合作，共同探索世界模型的未开发潜力。这仅仅是开始。”

研究领域

• 机器学习

标签

• 生成式 AI
• 亚马逊云服务（AWS）
• 硬件
• 扩散建模

关于作者

Stephen Zorio 是 Amazon Science 官方网站的撰稿人和编辑，拥有超过 15 年的编辑、写作和数字战略工作经验。