优化传统RNN在现代硬件上的输入输出性能，实现50倍加速
FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
- 通过寄存器级内核优化和并行化处理小隐藏状态的RNN
- 相比原生PyTorch实现快50倍，隐藏层大小可扩大40倍
- 适用于需要状态追踪的时间序列建模与逻辑推理任务
尽管Transformer等可并行化架构在序列建模中表现优异，但缺乏状态追踪能力，而传统RNN如LSTM、GRU及sLSTM虽具备该能力，却受限于严格串行计算。本文提出FlashRNN，在Triton和CUDA上实现寄存器级内核优化，通过并行处理多个小隐藏状态的RNN（类似Transformer的头并行），显著提升效率。引入新的硬件缓存大小、内存与计算处理优化框架，基于类似多面体的约束建模，包含可除性概念，加速ConstrINT库中通用整数约束满足问题（integer CSP）求解。实验表明，该方法相较原生PyTorch实现提速50倍，隐藏层规模可扩大40倍。开源代码已发布，助力状态追踪型RNN与序列建模研究。
原文摘要 · Abstract (English)
While Transformers and other sequence-parallelizable neural network architectures seem like the current state of the art in sequence modeling, they specifically lack state-tracking capabilities. These are important for time-series tasks and logical reasoning. Traditional RNNs like LSTMs and GRUs, as well as modern variants like sLSTM do have these capabilities at the cost of strictly sequential processing. While this is often seen as a strong limitation, we show how fast these networks can get with our hardware-optimization FlashRNN in Triton and CUDA, optimizing kernels to the register level on modern GPUs. We extend traditional RNNs with a parallelization variant that processes multiple RNNs of smaller hidden state in parallel, similar to the head-wise processing in Transformers. To enable flexibility on different GPU variants, we introduce a new optimization framework for hardware-internal cache sizes, memory and compute handling. It models the hardware in a setting using polyhedral-like constraints, including the notion of divisibility. This speeds up the solution process in our ConstrINT library for general integer constraint satisfaction problems (integer CSPs). We show that our kernels can achieve 50x speed-ups over a vanilla PyTorch implementation and allow 40x larger hidden sizes compared to our Triton implementation. Our open-source kernels and the optimization library are released here to boost research in the direction of state-tracking enabled RNNs and sequence modeling: https://github.com/NX-AI/flashrnn
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。