FlashRNN：现代硬件上传统 RNN 的 I/O 感知优化
FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware
AI总结：
本文提出 FlashRNN，通过 Triton/CUDA 寄存器级 kernel、小隐藏状态多 RNN 并行及硬件约束优化框架，显著提升传统 RNN 在现代 GPU 上的速度与可扩展隐藏尺寸。
摘要：
尽管 Transformers 和其他可序列并行化的神经网络架构似乎代表了序列建模的当前最先进水平，但它们尤其缺乏状态跟踪能力。这些能力对于时间序列任务和逻辑推理非常重要。传统 RNN（如 LSTM 和 GRU）以及 sLSTM 等现代变体确实具备这些能力，但其代价是严格的顺序处理。虽然这通常被视为一个严重限制，但我们展示了借助使用 Triton 和 CUDA 实现的硬件优化 FlashRNN，这些网络能够达到多快的速度；该方法在现代 GPU 上将 kernel 优化到寄存器级别。我们通过一种并行化变体扩展传统 RNN，即并行处理多个具有较小隐藏状态的 RNN，类似于 Transformers 中按头处理的方式。为了在不同 GPU 变体上实现灵活性，我们提出了一个面向硬件内部缓存大小、内存和计算处理的新优化框架。该框架在一种使用类多面体约束的设置中对硬件建模，其中包括可整除性概念。这加快了我们用于一般整数约束满足问题（integer CSP）的 ConstrINT 库中的求解过程。我们表明，与原生 PyTorch 实现相比，我们的 kernel 可实现 50 倍加速，并且与我们的 Triton 实现相比，可支持 40 倍更大的隐藏尺寸。我们的开源 kernel 和优化库已在此发布，以推动具备状态跟踪能力的 RNN 与序列建模方向的研究：https://github.com/NX-AI/flashrnn
英文摘要：
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
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
