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

TL;DR

This paper introduces a novel method to accelerate inference in long convolution sequence models, achieving near-linear time complexity and significant speedups over standard approaches, enabling more efficient sequence processing.

Contribution

The authors develop a general framework for quasilinear inference in LCSMs, significantly reducing inference time and enabling high parallelization, with empirical validation on Hyena.

Findings

Up to 7.8× end-to-end speedup in inference

110× speedup within position-mixing component

Achieves O(L log^2 L) inference complexity

Abstract

While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length. Several subquadratic 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' exact inference to quasilinear $O(L\log^2L)$ 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, which gets up to $7.8\times$ end-to-end improvement over standard inference by improving $110\times$ within the position-mixing part.