본 논문은 게이트가 있는 선형 RNN이 언어 모델링에서 Transformer와 비교하여 경쟁력 있는 성능을 보임에도 불구하고, 실제로는 최적화된 커스텀 커널이 필요하다는 점을 지적합니다. 기존의 Flash Linear Attention (FLA)는 입력 시퀀스의 청크(chunk)를 병렬 처리하여 선형 RNN 커널이 Flash Attention보다 빠르다는 것을 보였지만, 청크 크기의 제한으로 인해 GPU 메모리에 많은 중간 상태를 저장해야 하는 문제가 있습니다. 본 논문에서는 청크 내에서 추가적인 시퀀스 병렬화를 도입하여 임의의 큰 청크 크기와 높은 연산 강도를 가능하게 하는 새로운 커널 알고리즘인 Tiled Flash Linear Attention (TFLA)를 제시합니다. TFLA를 matrix memory를 사용하는 xLSTM인 mLSTM에 적용하고, 더 빠른 커널 실행 시간을 위해 시그모이드 입력 게이트와 계산량을 줄인 mLSTM 변형을 제안합니다. 실험 결과, TFLA 기반의 새로운 mLSTM 커널이 최적화된 Flash Attention, Linear Attention 및 Mamba 커널보다 우수한 성능을 보이며, 효율적인 장문맥 시퀀스 모델링 기본 요소에 대한 새로운 최첨단 기술을 제시합니다.
