A fast RG-LRU Kernel

TL;DR: I develop a IO-aware kernel for the RG-LRU in Triton. When training on sequences greater than 4K tokens, the kernel is faster than PyTorch's Flash Attention 2 implementation. Introduction Hawk is an RNN proposed by Google Deepmind in “ Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models”. 1 At its core is the Real-Gated Linear Recurrent Unit (RG-LRU)…