본 논문은 시퀀스 병렬 처리가 가능한 Transformer와 같은 최신 신경망 아키텍처가 상태 추적 기능이 부족하다는 점에 주목합니다. 시간 시계열 작업 및 논리적 추론에는 상태 추적 기능이 중요하며, LSTM 및 GRU와 같은 기존 RNN은 이러한 기능을 갖지만 순차적 처리로 인해 속도가 느립니다. 본 논문에서는 Triton과 CUDA를 사용한 하드웨어 최적화 FlashRNN을 통해 기존 RNN의 속도를 높이는 방법을 제시합니다. GPU 레지스터 수준에서 커널을 최적화하고, Transformer의 헤드별 처리와 유사하게 작은 은닉 상태를 가진 여러 RNN을 병렬로 처리하는 병렬화 변형을 기존 RNN에 확장합니다. 다양한 GPU 변형에 대한 유연성을 확보하기 위해 하드웨어 내부 캐시 크기, 메모리 및 계산 처리를 위한 새로운 최적화 프레임워크를 도입하며, 일반 정수 제약 만족 문제(integer CSPs)를 위한 ConstrINT 라이브러리에서 솔루션 프로세스 속도를 높입니다. 실험 결과, 제안된 커널은 기본 PyTorch 구현보다 50배 빠르고, Triton 구현과 비교하여 40배 더 큰 은닉 크기를 허용함을 보여줍니다. 개발된 커널과 최적화 라이브러리는 오픈소스로 공개됩니다.
