Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond

While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length.Several subquatratic architectures have been proposed to address this computational issue. Some of them, including long convolution sequence models (LCSMs), such as Hyena, address this issue at training time but remain quadratic during inference. We propose a method for speeding up LCSMs' inference to quasilinear time, identify the key properties that make this possible, and propose a general framework that exploits these. Our approach, inspired by previous work on relaxed polynomial interpolation, is based on a tiling which helps decrease memory movement and share computation. It has the added benefit of allowing for almost complete parallelization across layers of the position-mixing part of the architecture. Empirically, we provide a proof of concept implementation for Hyena-like settings, which gets up to 4x improvement over standard inference.