对角批处理解锁循环记忆 Transformer 在长上下文中的并行性

Transformer模型由于其二次时间复杂度和线性内存复杂度，在长上下文推理方面面临挑战。循环记忆Transformer（RMT）通过将渐近成本降低到线性时间和常数内存使用量，提供了一种解决方案。然而，它们的内存更新机制导致顺序执行，从而造成性能瓶颈。我们引入了对角批处理（Diagonal Batching），这是一种调度方案，它在RMT中实现了跨段并行化，同时保持了精确的循环。这种方法消除了顺序约束，即使对于单个长上下文输入，无需复杂的批处理和流水线技术也能实现高效的GPU推理。因为该技术纯粹是一种运行时计算重排序，所以现有的RMT模型无需重新训练即可采用它。应用于LLaMA-1B ARMT模型，对角批处理在131,072-token序列上比标准的全注意力LLaMA-1B实现了3.3倍的加速，比顺序RMT实现加速了1.8倍。通过消除顺序瓶颈，对角批处理降低了推理成本和延迟，从而使RMT成为现实世界长上下文应用的实用解决方案。