arXiv:2503.14376v2 宣告类型: replace-cross 摘要: 具有门控机制的线性RNN最近在语言建模方面展示了与Transformer相竞争的性能。尽管它们在序列长度上的线性计算扩展在理论上为Transformer提供了运行时间优势，但在实践中实现这些优势需要优化的定制内核，因为Transformer依赖于高度高效的Flash Attention内核（Dao, 2024）。利用线性RNN的分块并行形式，Flash Linear Attention (FLA) (Yang & Zhang, 2024) 显示出线性RNN内核比Flash Attention更快，通过在输入序列的块上并行化实现。然而，由于FLA的块大小有限，许多中间状态必须在GPU内存中实现。这导致了低算术强度和高内存消耗及IO成本，尤其是在长上下文预训练时。在本工作中，我们提出了一种新的线性RNN内核算法Tiled Flash Linear Attention (TFLA)，通过在每个块内引入额外的序列并行化层次，使得可以实现任意大的块大小和高算术强度。首先，我们将TFLA应用于具有矩阵记忆的xLSTM。其次，我们提出了一种具有Sigmoid输入门控和减少计算量的mLSTM变体，即使在相等的语言建模性能下也能实现更快的内核运行时间。在我们的速度基准测试中，我们展示了基于TFLA的新mLSTM内核超越了高度优化的Flash Attention、Linear Attention和Mamba内核，为高效长上下文序列建模提供了新的性能标准。