加速 Roblox 上 3D 创作的人工智能推理

• Roblox 采用了 CUDA 图和键值对缓存技术，以加快 3D 网格生成速度，从而实现更流畅的迭代。
• 上线时，Cube 3D 模型可在 7.8 毫秒内生成令牌（较之前的 60.5 毫秒大幅缩短），并在 4 秒内生成完整对象（较之前的 31 秒大幅缩短）。 

今年早些时候，Roblox 分享了 Cube 3D 基础模型的首批功能。借助 Cube 3D，创作者可以直接根据文本提示生成 3D 模型和环境。 我们从一开始就将延迟优化作为首要任务，因为我们意识到，生成速度过慢会破坏本应具有迭代性的创作过程。在3月份Cube 3D正式发布之前，我们已经将推理步骤的速度提升了7.8倍，使开发者和用户都能获得更流畅的体验。 

自上线以来，多个知名体验中已生成超过57.8万个物体。开发者还表示希望让用户能在体验中通过“猫”、“汉堡”等文本提示生成3D物体。其中最引人注目的是《Mic Up》——这款热门的语音聊天聚会游戏利用Cube 3D，为玩家提供了生成物体的趣味互动方式。 在该游戏的实现中，玩家可以打开左侧菜单访问更多功能，其中包括一个 AI 图标。点击该图标后，玩家可以输入文本提示来生成 3D 对象。对于用户而言，过长的生成时间会造成体验摩擦，使他们无法实时见证自己的创意转化为 3D 模型的魔力。 

我们希望将3D生成体验从“停顿等待”的交互模式，转变为响应迅速且自然流畅的体验，从而支持快速实验。能够快速向场景中添加物体对开发者至关重要。为了加速Cube 3D，我们首先对推理管道进行了性能分析，以识别性能瓶颈。尽管使用了强大的GPU，但我们发现操作之间存在显著的空闲时间。

现代深度学习框架依赖于 CPU 来调度并启动 GPU 上的操作（或内核）。CPU 会准备每个操作，将其发送至 GPU，并在收到确认后才开始准备下一个操作。这种等待会形成调度瓶颈，导致在 CPU 准备下一批任务时，GPU 可能处于空闲状态。理想情况下，我们希望 CPU 能领先于 GPU 运行，预先准备并排队操作，从而确保 GPU 始终有任务可处理。

对于 Cube 3D 等变压器类模型中的自回归解码器而言，这一问题尤为突出，因为它们需要按顺序处理输入并生成令牌。这些模型单次生成就需要数千个独立操作，且计算开销会随着序列中的每一步累积增加。

“我们希望构建一种能够实现四维交互的技术，”工程副总裁阿努帕姆·辛格（Anupam Singh）在解释Roblox为何选择自回归方法时说道，“我们不仅想造一辆车，还希望能够打开车门并坐进车里。”

每次操作产生的：

• 为每个内核准备所需的CPU时间
• 启动内核产生的开销
• GPU执行时间（实际计算）
• 检查完成状态时的同步开销

对于在 GPU 上执行迅速的小型操作，这些开销可能占据推断时间的大部分。GPU 实际进行计算的时间可能仅占总推断时间的一小部分。

为解决这一瓶颈，我们利用了 CUDA Graphs——该功能允许在不需 CPU 干预的情况下记录和重放 GPU 操作序列。Cube 3D 架构中的自回归解码器组件处理文本提示，并通过固定长度向量生成形状令牌。 

虽然在功能上与传统的大型语言模型（LLM）相似，但我们的双流解码器架构有一个重要区别——它使用了两个并行注意力流。一个流专门处理条件令牌，另一个则处理形状令牌。现成的 LLM 推理引擎无法满足我们的需求，我们需要针对特定架构进行定制化实现。

可以将 CUDA 图（CUDA Graphs）理解为 GPU 的宏录制功能。它不再需要 CPU 逐条发出指令，而是将整个 GPU 操作序列（即图）记录下来，并通过一条 CPU 指令启动整个图。这种方法通过消除 CPU 在推理过程中逐个调度操作的必要性，极大地降低了内核启动的开销。一旦图被启动，GPU 就会自主执行整个序列，无需等待进一步的指令。

CUDA Graphs确实存在一些局限性。由于图结构需要预先确定，因此它要求固定的批处理大小和输入维度。这意味着针对每种批处理大小或输入形状都需要创建独立的图。对于我们Cube 3D的应用场景，这一限制是可以接受的，因为我们可以围绕常见的输入维度来标准化推理过程。

为了在 Cube 3D 模型中实现 CUDA 图，我们不得不调整方法。在传统的大型语言模型（LLMs）中，注意力操作始终针对相同的序列长度进行，从而提供了一个静态的处理形状。然而，在我们自定义的双流架构中，部分注意力层仅基于序列长度进行操作，而另一些则基于序列长度与条件长度的组合进行操作。

尽管面临这些挑战，但在实现 CUDA Graphs 后，我们获得了显著的成果。我们使用“每个输出令牌所需时间”（TPOT）来衡量推理过程中每个令牌的生成时间。在实现 CUDA Graphs 后，我们的 TPOT 从 60.5 毫秒提升至 20.5 毫秒，提升了 2.9 倍。总体生成时间下降了 66%，从 31 秒缩短至 10.5 秒。

为了进一步降低延迟，我们实现了键值对缓存（KV caching），这是大语言模型（LLM）推理中的一种标准做法，已在业界证明非常有效。 

在 Cube 3D 等基于 Transformer 的模型中，每次生成令牌都需要根据之前生成的所有令牌计算键 (K) 和值 (V) 矩阵。随着序列变长，为每个令牌重新计算这些矩阵的效率会越来越低。

KV缓存通过以下方式解决此问题：

1. 存储所有先前生成的令牌对应的 K 和 V 矩阵
2. 仅针对新令牌计算 K 和 V 矩阵
3. 将这些新矩阵追加到缓存值中

这种方法消除了冗余计算，减少了每个新令牌所需的工作量。随着生成的序列变长，这种优化效果尤为显著。

我们将 KV 缓存与 CUDA 图实现集成的方法与传统的 LLM 推理类似。添加 KV 缓存后，我们的 TPOT 缩短至仅 7.8 毫秒。总体生成时间减少了 87%，从最初的 31 秒缩短至仅 4 秒。这一显著的时间缩短，使该工具对创作者的使用效率大幅提升。

我们正在探索进一步降低延迟并提升用户体验的技术，包括优化内核、用于实现更快速推理的模型量化、针对特定硬件的优化以及并行令牌生成。

当扩展到完整的场景生成与理解时，这项工作变得更加复杂，因为在布局中，许多 3D 元素需要在上下文中协同工作。我们还希望创建的 3D 对象和世界能够完全正常运作，例如门可以开合、车轮可以转动等。为了实现这一目标，我们需要快速生成和迭代，以扩展到整个场景、完全功能性的对象和虚拟角色。 随着我们不断扩展 Cube 3D 基础模型，我们期待分享更多改进和新功能，并期待看到创作者社区利用这些模型构建出的沉浸式世界。