FlashRNN: Optimizing Traditional RNNs on Modern Hardware

虽然Transformer和其他可并行处理序列的神经网络架构似乎是当前序列建模的最新技术，但它们特别缺乏状态跟踪能力。这些能力对于时间序列任务和逻辑推理非常重要。传统的RNN（如LSTM和GRU）以及现代变体（如sLSTM）确实具有这些能力，但代价是必须进行严格的顺序处理。尽管这通常被视为一个严重的限制，但我们展示了通过我们的硬件优化FlashRNN在Triton和CUDA中优化到寄存器级别的内核，这些网络可以变得多快。我们通过一种并行化变体扩展了传统RNN，该变体并行处理多个隐藏状态较小的RNN，类似于Transformer中的分头处理。为了在不同的GPU变体上实现灵活性，我们引入了一个新的优化框架，用于硬件内部缓存大小、内存和计算处理。该框架使用类似多面体的约束来建模硬件，包括可除性的概念。这加快了解决过程，特别是在我们用于一般整数约束满足问题（整数CSP）的ConstrINT库中。我们展示了我们的内核可以比纯PyTorch实现提速50倍，并且与我们的Triton实现相比，允许隐藏状态大40倍。我们开源了这些内核和优化库，以促进状态跟踪功能RNN和序列建模的研究：[https://github.com/NX-AI/flashrnn](https://github.com/NX-AI/flashrnn)

该论文旨在解决传统RNN（如LSTM和GRU）在处理时间序列任务和逻辑推理时的性能瓶颈问题，特别是它们由于严格顺序处理而导致的速度较慢的问题。这是一个现有问题，但论文提出了新的解决方案。

论文的关键思路是通过硬件优化和并行化技术来加速RNN的处理速度。具体来说，作者提出了FlashRNN，这是一种在现代GPU上优化到寄存器级别的RNN实现方法，并引入了一种并行处理多个小型RNN的方法，类似于Transformer中的多头机制。此外，还开发了一个新的优化框架，用于建模硬件内部的缓存大小、内存和计算处理，以提高效率。

论文展示了FlashRNN在速度和模型规模上的显著提升，与PyTorch的默认实现相比，实现了50倍的速度提升，并支持40倍更大的隐藏层大小。此外，作者开源了优化后的内核和优化库，为未来的研究提供了工具。这些成果不仅提高了RNN的实用性，也为进一步探索状态跟踪能力在序列建模中的应用奠定了基础。

近期在这个领域中，相关的研究包括：1) 《Accelerating Recurrent Neural Networks with Custom Hardware》探讨了使用定制硬件加速RNN；2) 《Efficient Parallelization of RNNs on Modern GPUs》研究了RNN在现代GPU上的并行化策略；3) 《Optimizing RNNs for Time-Series Forecasting》专注于时间序列预测任务中的RNN优化。这些研究都从不同角度尝试解决RNN的性能问题，而本论文则通过硬件优化和并行化技术取得了显著进展。