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

## Abstract

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 2 ⁡ L ) 𝑂 𝐿 superscript 2 𝐿 O(L\log^{2}L) italic_O ( italic_L roman_log start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_L ) 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 1.6 × 1.6\times 1.6 × end-to-end improvement over standard inference by improving 50 × 50\times 50 × within the position-mixing part.

## 1 Introduction

A lot of recent progress in deep learning, particularly in the form of large language models (LLMs) has been driven by the transformer architecture (Vaswani et al., ). While these models have great quality, it comes at a computation cost which scales quadratically in sequence length - both during training and inference. This can become prohibitive for very long contexts and as such a number of alternative architectures with better computational scaling in context length have been proposed (Gu and Dao, ; Poli et al., ; Fu et al., ). While most of these works have improved computational efficiency for training, some still scale quadratically in sequence length when it comes to inference, thus not improving asymptotically over transformers.

In this work, we propose a framework for optimizing inference efficiency for a general class of such models. As a case study, which inspired the method, we focus on long convolution sequence models (LCSMs) (Poli et al., ; Fu et al., ; Romero et al., ; Li et al., ; Karami and Ghodsi, ; Fu et al., ). However, our approach is not limited to LCSMs alone and we identify the properties that allow for such inference speedups in hope to guide the design of future architectures.

In the particular case of LCSMs (including Hyena), the building block of the architecture is that of convolving the input sequence with a sequence-length long, (potentially underparameterized) filter. If we let $L$ be the sequence length (e.g. number of tokens in the case of LLMs), then a naive implementation of convolution during training would take $\Omega(L^{2})$ FLOPs, but one can employ FFT to bring that down to $O(L\log L)$. The issue that occurs during inference is that FFT cannot be used directly since the whole input sequence is not known ahead of time, but rather incrementally computed. Because of this, the naive inference approach goes up to $\Omega(L^{2})$ - this is the apparent cost of moving from a *static* input to a *dynamic* one.

It turns out that in fact, at least in the case of “dynamic FFT”, one can obtain $O(L\log^{2}L)$ time complexity by using van der Hoeven ()’s result on relaxed polynomial convolution. Crucially, this result works by a direct reduction to several applications of FFT. This allows us to phrase a more general framework for turning training-efficient architectures to inference-efficient ones.

Our main contributions are:

- •
Proposing the first quasilinear-in-sequence-length $O(L\log^{2}L)$, *exact* inference algorithm for LCSMs.
- •
Identifying the main features of LCSMs that allow for faster inference to propose a more general framework.
- •
Highlighting and exploiting the potential for dealing with most workload of different layers in parallel, rather than having to wait for all computation of a layer to finish before moving on to subsequent layers.
- •
Saving up on data movement thanks to the tiling approach used.
We show how to save factors of $2$ in several places, including activation storage, when the convolutional filters are data-independent. However, our method extends to any causal, data-dependent filters.
- •
Our framework provides an $O(L\log^{2}L)$ inference algorithm which, when run empirically, yields up to $1.6\times$ end-to-end time-efficiency improvement; and up to $50\times$ on the long convolution part.

Figure: Figure 1: Cell ($i$, $j$) corresponds to the contribution of mixer-input $y_{i}$ to mixer-output $z_{j}$. To compute $z_{j}$, all its line of contributions should be be accounted for. Because of the autoregressive nature of inference, one only has access to $y_{i}$ after $z_{i-1}$ has been computed. *(Left Top)* represents the standard (lazy) approach, *(Left Bottom)* represents the eager approach, and *(Right)* represents our suggested method.
Refer to caption: x1.png

## 2 Setting and Related Work

Much recent effort has been put in designing more efficient alternatives to transformers (Tay et al., ). A couple of well-performing classes are state space models (SSMs) (Gu et al., ; Gu and Dao, ) and convolution-based models(Fu et al., ; Poli et al., ). In this work we focus on the latter class. To avoid a popular misconception, it is worth emphasizing that while any linear-time invariant (LTI) SSM has an equivalent convolutional filter, it is *not* true that any convolutional filter has an equivalent *low-dimensional* LTI SSM representation. It could be the case that the smallest LTI SSM equivalent to a length-$L$ convolution filter has a dimension of $L$. However, such a dimension would deem an LTI SSM approach have quadratic time complexity both during training as well as during inference.

### 2.1 Notations and Definitions

In the most general form, we consider architectures obtained by stacking several position-mixing layers potentially interleaved with element-wise, feature-mixing modules, partially following Arora et al. ()’s notation. Put formally, let $\mathrm{block}_{1},\ldots\mathrm{block}_{M}:\mathbb{R}^{D}\to\mathbb{R}^{D}$ be feature-mixing modules and $\mathrm{mixer}_{1}\ldots\mathrm{mixer}_{M}:\mathbb{R}^{L\times D}\to\mathbb{R}
^{L\times D}$ be position-mixing modules. One can think of the blocks as being combinations of MLPs, layernorms, skip connections and other modules of similar purpose and of mixers as playing the role of multi-head attention (MHA) in a transformer - the part where embeddings at different positions interact with each other. Assume these are all learnt from some well-defined parameterized classes. Here, $M$ is the number of layers, $D$ is the used embedding dimension and $L$ can be any sequence length (unless positional embeddings are used, case in which it can only take values up to some $L_{max}$). We define the activations $a_{0},\ldots,a_{M}:\mathbb{R}^{L\times D}\to\mathbb{R}^{L\times D}$ at level $1\leq\ell\leq M$ and position $1\leq i\leq L$ as $(a_{\ell}(x))_{i}\triangleq\mathrm{block}_{\ell}(\mathrm{mixer}_{\ell}(a_{\ell
-1}(x))_{i})$ where $a_{0}(x)=x$ represents the input embeddings ($L$ $D$-dimensional tokens) and the model output at position $i$ is read from the last layer $(a_{M}(x))_{i}$.

We focus on autoregressive sequence models where for an input sequence $x_{1},\ldots x_{p}$, one generates the $(p+1)^{\text{th}}$ token by sampling from a distribution given by $(a_{M}(x))_{p}$ (usually via a linear map). In order for this generation procedure to be sensible, we constrain (by construction of $\mathrm{mixer}$s) the models to be causal, that is, $(a_{\ell}(x))_{i}$ is a function of only $x_{[1,i]}$ or, equivalently, $\mathrm{mixer}_{\ell}(y)_{i}$ is a function of $y_{[1..i]}$.

We use $\odot$ to denote the element-wise (Hadamard) product: for two matrices $A,B\in\mathbb{R}^{N\times M}$, $A\odot B\in\mathbb{R}^{N\times M}$ is given by $(A\odot B)_{i,j}=A_{i,j}\cdot B_{i,j}$. We also use indexing by ranges to mean slicing a tensor in the respective dimension according to the range and indexing by a dot to mean leaving the dimension unchanged: for example, if $A:\mathbb{R}^{N\times M\times K}$, then $A_{\cdot,[l,r],5}\in\mathbb{R}^{N\times[r-l+1]}$ and $(A_{\cdot,[l,r],5})_{i,j}=A_{i,j+l-1,5}$.

We call a function $f:{\mathcal{X}}^{*}\to{\mathcal{X}}$ associative, if $f(x_{1},x_{2},\ldots x_{p})=f(f(x_{1},\ldots x_{i}),f(x_{i+1},\ldots x_{p}))$ for any $1\leq i<p$ and $x_{1}\ldots x_{p}\in{\mathcal{X}}$ - that is, it is invariant to bracketing. Finally, we define the time complexity of an algorithm as the number of floating-point operations (FLOPs) it executes - this is a notion independent of hardware.

### 2.2 Self-attention

Self-attention is the building block of transformers (Vaswani et al., ). In its simplest form, it implements a $\mathrm{mixer}$ by using three projection matrices $Q,K,V\in\mathbb{R}^{D\times D}$ to obtain three sequences of vectors $q,k,v\in\mathbb{R}^{L\times D}$ by projecting the input $y\in\mathbb{R}^{L\times D}$ via $q=yQ,k=yK,v=yV$ - these are called queries, keys and values, respectively. The (causal) self-attention operator $attention:\mathbb{R}^{L\times D}\to\mathbb{R}^{L\times D}$ is then given by:

$$ $\displaystyle attention(y)_{j}\triangleq\mathrm{mixer}(y)_{j}=\frac{\sum_{i=1} ^{j}{v_{j}\cdot e^{\langle{q_{j},k_{i}}\rangle}}}{\sum_{i=1}^{j}{e^{\langle{q_ {j},k_{i}}\rangle}}}=(v_{[1,j]})^{\top}\mathrm{softmax}(k_{[1,j]}q_{j})$ (1) $$

which represents, at position $j$, the average of the previous value vectors $v_{[1,j]}$ exponentially-weighted by how well the values’ corresponding keys match $j$^th query. Put differently, the bigger $\langle{q_{i},k_{j}}\rangle$ is, the more $j$^th output attends to $i$^th input. Note that both in the above notation, as well as in general, we assume any $1$-dimensional tensor to be a column vector (for example $q_{j}\in\mathbb{R}^{D\times 1})$. In the transformer architecture this is how one “head" operates and each embedding is normally split into several such heads - this is called multihead attention (MHA). If we take the causality away (that is, the $i\leq j$) for simplicity, one could think of attention as being $\mathrm{mixer}(y)=\mathrm{softmax}(qk^{\top})v$ where $\mathrm{softmax}$ is computed along the rows of what is called the attention matrix: $qk^{\top}$.

### 2.3 Long Convolution Sequence Models (LCSMs)

LCSMs (Poli et al., ; Li et al., ; Gaido et al., ) work by creating a SISO (single-input-single-output) primitive to map an input sequence $y\in\mathbb{R}^{L}$ to $z\in\mathbb{R}^{L}$ via $z_{t}\mapsto\sum_{i=1}^{t}{y_{i}\cdot\rho_{t-i}}$ where $\rho$ is a (possibly infinite, and of length at least $L$) convolution filter which is often parameterized by a smaller latent space $\Theta$: $\rho=f(\theta)$ where $\theta\in\Theta,\dim{\Theta}\ll L$ is learnt. These convolution primitives operate on independent axes of the hidden dimensions to create a positional mixer: $\mathrm{mixer}(y)_{t,c}=\sum_{i=1}^{t}{y_{i,c}\cdot\rho_{t-i,c}}$. This assumes the filters to be independent, but shared ones are possible as well (as is the case for multi-head Hyena (Massaroli et al., )).

For example, the Hyena architecture (Poli et al., ) maps back to our setup when mixers are defined as above and the $\mathrm{block}$ functions, depending on the layer, are either MLPs or gates - that is, element-wise multiplications with a projection of activations at same position, but lower level. We do not focus on the details of the $\mathrm{block}$s since they all involve some $D\times D$ matrix-vector multiplication that is performed once for every layer and position $1\leq i\leq L$ and thus scale as $\Theta(LD^{2})$ per layer - that is, linearly in context length.

#### 2.3.1 The Inference Problem

Whereas during training, the convolutions can be performed in $O(L\log L)$ time via FFT as the whole of $y$ is known beforehand, during inference this is non-trivial. Suppose there is a prompt of size $P$ and we want to generate $L-P$ more tokens. The bottleneck problem can be understood as computing, for every layer and every dimension, the below:

$$ $\displaystyle z_{t}\triangleq\sum\limits_{i=1}^{t}y_{i}\cdot\rho_{t-i}$ (2) $$

for all $1\leq t\leq L$ . To pre-fill the first $P$ values $z_{[1..p]}$ at all levels of activations, since the first $P$ inputs are known, one can perform FFT as they would normally do during train-time (incurring an $O(P\log P)$ time cost), but past $P$^th position it is important to note that $y_{i}$ is not available before computing $z_{i-1}$ - this is where *static* FFT breaks. Since dealing with the prompt can be done easily (Massaroli et al., , Lemma 2.1) by essentially filling in all contributions of $y_{[1..P]}$ to $z_{[1..L]}$ and then forgetting the prompt ever existed, we henceforth assume $P=0$.

#### 2.3.2 Previous Work on Efficient Inference

Speeding up Equation 2 from the naive $\Omega(L^{2})$ is the object of study of Massaroli et al. (): they optimize the process by finding a *low-dimensional* LTI SSM whose equivalent convolution filter closely resembles $\rho$, thus getting an RNN-like problem to simulate. If the learnt SSM has size $D^{\prime}$, then this yields an $\Theta(LDD^{\prime})$ algorithm for generating $L$ tokens. This has the significant added benefit of not needing to store the activations (or even inputs) thus far which makes it memory efficient and practically time-efficient.

The downside, however, is that by definition this will only be an approximation of the learnt filter. More importantly, this approximation represents a projection to a fundamentally smaller space of models, thus defeating the purpose of using LCSMs instead of LTI SSMs - it is practically a different training procedure for an LTI SSM model. Furthermore, this approach assumes the filter $\rho$ is data-independent - this is needed in order to undergo an expensive distillation procedure as a precomputation step. This can be qualitatively limiting as shown in Arora et al. (). While we do store all activations, our proposed framework exactly simulates the architecture and data-dependent filters can also be accommodated.

## 3 Fast LCSM inference

In this section we describe our method for the case of long convolution sequence models. That is, we still work with models as described in Section 2.1, but we further assume the mixers to be convolution-based, as described in Section 2.3. Hence, for every $\mathrm{mixer}_{\ell}\in\{\mathrm{mixer},\ldots\mathrm{mixer}_{M}\}$, there exists a filter $\rho_{\ell}\in R^{L\times D}$ such that

$$ $\displaystyle\mathrm{mixer}_{\ell}(y)_{t}=\sum_{i=1}^{t}{y_{i}\odot\rho_{\ell, t-i}}$ $$

We assume here, for simplicity, that $\rho$ is data-independent and part of the model, as is the case for Hyena. However, we discuss in Appendix B how our method can be extended to data-dependent filters. We also assume $L=2^{P}$ for some integer $P$ for simplicity - one can always round $L$ up to the closest power of $2$ without changing it by more than a factor of $2$ and thus keeping the asymptotics of all our analysis unchanged.

### 3.1 The Proposed Algorithm

We derive inspiration from the the work of van der Hoeven () regarding relaxed polynomial interpolation - they propose a fix to dealing with *dynamic* structure of Eq. 2 to achieve an overall $O(L\log^{2}L)$ time complexity.

#### 3.1.1 Relaxed Polynomial Interpolation

Consider Eq. 2:

$$ $\displaystyle z_{t}\triangleq\sum\limits_{i=1}^{t}y_{i}\cdot\rho_{t-i}$ $$

The problem of relaxed polynomial interpolation can be thought of as having to compute $z_{t}$ for every $1\leq t\leq L$ with the further constraint that $y_{t}$ and $\rho_{t}$ are only made available once $z_{t-1}$ has been outputted. While this general setting is covered in Appendix B, here we focus on the case that all of $\rho$ is known ahead of time and only $y$ is incrementally revealed. One could simply use the more general approach of (van der Hoeven, ) even when $\rho$ is known ahead of time, but we can take advantage of this particularity to get a twice-faster and slightly simpler algorithm. We present this both for brevity and because existing architectures have data-independent $\rho$’s so this is the more efficient choice for them (and thus, for practical purposes).

The main approaches to relaxed polynomial interpolation, as identified in (Van der Hoeven, ), are:

- •
The lazy approach: for every $t\in\{1\ldots L\}$, one computes $z_{t}$ by simply applying the formula. It takes $O(t)$ FLOPs to do so and the overall complexity is $O(L^{2})$. This approach can be thought of as the naive approach which only performs work when it is strictly needed.
- •
At the opposite end, there is the eager (or zealous) approach: as soon as a new value of $y$ or $\rho$ becomes available, one accounts for all its contributions to $z$. Since we assumed the entire $\rho$ is known beforehand, this translates to the following:
After computing $z_{t-1}$, $y_{t}$ becomes available, so for every $t\leq i\leq L$, we increase $z_{i}$ by $y_{t}\cdot\rho_{i-t}$. $z_{t}$ is now fully computed so we proceed to the next iteration.
Hence, the eager approach can be thought of as performing work as soon as it can be performed. Here, $t$^th iteration takes $O(L-t)$ so we still have an overall complexity of $O(L^{2})$.
- •
Relaxed approaches: these approaches perform some but not all available work ahead of time - their advantage stems from grouping contributions (i.e. accounting collectively for groups of pairs $y_{i}\cdot\rho_{j}$) in such a way that fast convolution methods (such as FFT) can be employed for a speedup.

The employment of FFT for speeding up accounting for groups of contributions is formalized in the following lemma:

###### Lemma 1 .

Let $1\leq l\leq r\leq l^{\prime}\leq r^{\prime}\leq L$ represent ranges of lengths $L_{1}=r-l+1$ and $L_{2}=r^{\prime}-l^{\prime}+1$ of $y$ and $z$, respectively.
There exists an FFT-based algorithm running in $O(L_{1}+L_{2})$ space and $O((L_{1}+L_{2})\log(L_{1}+L_{2}))$ time complexity that, given access to $y_{[l,r]}$, computes the aggregated contributions of all $y_{[l,r]}$ to all of $z_{[l^{\prime},r^{\prime}]}$.
That is, for each $l^{\prime}\leq t\leq r^{\prime}$, it computes:

$$ $\displaystyle\tau(y,[l,r],\rho,[l^{\prime},r^{\prime}])_{t}\triangleq\sum_{i=l }^{r}{y_{i}\cdot\rho_{t-i}}$ (3) $$

for all $l^{\prime}\leq t\leq r^{\prime}$.

What this lemma says is that we can compute efficiently the contribution of a range of inputs to a range of outputs, where the computational cost is asymptotically given by the larger of the ranges (this is because $O(L_{1}+L_{2})=O(\max(L_{1},L_{2}))$. Furthermore, observe that we assumed $r\leq l^{\prime}$: this is because in order to account for all these contributions at once, we need to have all of $y_{[l,r]}$ available which means that all of $z_{[1,r-1]}$ was computed, so by that time we should have already accounted for the contributions of $y_{[l,r]}$ to all of $z_{[1,r-1]}$; hence, there is no point to (re)computing any of these.

The lazy, eager, and our version of a relaxed approach are depicted in Figure 1 under the form of tilings of the contribution space. Note that the lazy and eager approaches can still be thought of as grouping contribution accounting, but do so along thin strips: the contribution of a long range of $y$’s to one element of $z$ in the lazy approach and that of a singe element of $y$ to a long range of $z$’s in the eager approach. The key to our relaxed tiling is using more balanced tiles since, as noted, the cost of a tile is given by the larger side (which is very expensive for thin tiles). The tiles still have to respect $r\leq l^{\prime}$ and one needs to ensure that the whole line of contribution to some $z_{t}$ has been dealt with by the time we return the value of $z_{t}$ - this is indeed the case for our tiling.

Figure: Algorithm 1 Fast Relaxed Polynomial Interpolation

While the lazy and eager algorithms have been described, one needs to make explicit the order of processing the tiles in the relaxed setting since we still cannot use $y_{t}$ before $z_{t-1}$ has been computed. This is shown in Algorithm 1: *return* marks a value of $z$ having been computed and *unlock* represents a value of $y$ becoming available. At the beginning of iteration $i$ it always holds that $z_{i}=\sum_{j=1}^{i-1}{y_{j}\cdot\rho_{i-j}}$. We start the iteration by completing this with the freshly unlocked $y_{i}\cdot\rho_{0}$ at line 4 - this corresponds to the red tiles of Figure 1. Following this step $z_{i}$ is ready to be returned. Then, we account for the gray tile that has just been unlocked (line 6) - this tile corresponds to the contribution of $y_{[i-U+1,i]}$ to $z_{[i+1,i+U]}$ where $U$ is the side of the tile, namely the largest power of $2$ dividing $i$. Finally we return $z_{i}$ as being in its final form and get access to $y_{i+1}$ - we could have done this before dealing with the gray tile as well. Note that the relative ordering of the red and gray lines does not matter within an iteration, but it is important to finish up the first red and gray tiles before second ones, the second ones before the third ones and so on.

###### Proposition 1 .

When $L=2^{P}$, Algorithm 1 performs $2^{P-1-q}$ calls to $\tau$ on ranges of length $2^{q}$ for every $q\in\{0\ldots P-1\}$. If we base $\tau$’s implementation on Lemma 1, this entails a total time complexity of $O(L\log^{2}L)$. Since the calls to $\tau$ are sequential and outputs are not stored, the peak memory usage is given by the largest call to $\tau$ and remains $O(L)$.

###### Lemma 1 .

###### Proposition 1 .

#### 3.1.2 Performing LCSM inference via fast relaxed polynomial interpolation

We can now use Algorithm 1 to get an efficient inference algorithm for LCSMs, but to do so we need to be careful about data-dependency. First, for brevity, in our previous notation of activations, disregard the dependency on input $x$ - that is, use $a_{\ell,i}$ instead of $(a_{\ell}(x))_{i}$ - since we are only going to work with one input $x$ (the one we are generating). Furthermore, denote the intermediate mixer computation by $b\in\mathbb{R}^{M\times L\times D}$ as follows:

$$ $\displaystyle b_{\ell,i}\triangleq\mathrm{mixer}_{\ell}(a_{\ell-1})_{i}=\sum_{ k=0}^{i}{a_{\ell-1,k}\odot\rho_{\ell,i-k}}$ (4) $\displaystyle a_{\ell,i}\triangleq\mathrm{block}_{\ell}(b_{\ell,i})$ (5) $$

for any $1\leq\ell\leq M$, where $a_{0,i}$ is the embedding of $i$^th token.
Whenever one has generated the tokens $x_{[1,i]}$ up to position $i$ or, equivalently, $a_{0,[1,i]}$, the whole of $a_{[0,M]\times[1,i-1]}$ will have been computed (since $a_{M,[1,i-1]}$ is needed to sample $x_{[1,i-1]}$). Caching these up is the equivalent of a KV-cache in transformers and is simply the natural way to not repeat work when autoregressively sampling. To generate the next token $x_{i+1}=a_{0,i+1}$, we want to fill in the values of $a_{[1,M],i}$. To do so, as per Eq. 4, we need to compute $b_{\ell,i}$ which is practically, for each level $\ell$, and within each dimension $1\leq c\leq D$, a relaxed interpolation between $a_{\ell-1,:,c}$ (playing the role of $y$) and $\rho_{\ell,:,c}$. Further applying a $\mathrm{block}_{\ell}$ to $b_{\ell}$, one gets $a_{\ell}$. Thus, we can pipeline several instances of Algorithm 1 sequentially across layers.

The resulting algorithm is illustrated in Algorithm 2. The outer loop (line 3) iterates through positions $i$: at each iteration, one computes the activations at position $i$ across all layers, does some (limited) eager work and samples next token. To do so, the inner loop (line 5) iterates through the several mixer-layers. It first accounts for the red cells in the diagram, that is, the direct dependency on previous layer’s activations at position $i$ (line 7), thus finalizing $b_{\ell,i}$. After computing the activation at position $i$ (line 8), it also accounts for the contribution of the gray tile that has just been unlocked (line 10) at the current layer. We reuse the notation from previous section for $\tau$ to be the same as before but within each of the $D$ dimensions - that is:

$$ $\displaystyle\tau(y,[l,r],\rho,[l^{\prime},r^{\prime}])_{t}\triangleq\sum_{i=l }^{r}{y_{i}\odot\rho_{t-i}}$ $$

for all $l^{\prime}\leq t\leq r^{\prime}$. The accounting works by considering (in-place) the influence of last $U$ positions of $a$ to next $U$ ones of $b$.

Figure: Algorithm 2 Flash Inference for LCSMs

As far as performance goes, Algorithm 2 practically calls Algorithm 1 $MD$ times (within each layer and within each of the $D$ dimensions of embeddings). Proposition 2 covers its complexity analysis:

###### Proposition 2 .

The overall time complexity of Algorithm 2 is $O(MDL\log^{2}L)$ for the mixer-part plus $LM$ $\mathrm{block}$ calls, to generate $L=2^{P}$ tokens. Furthermore, for each $0\leq q\leq P-1$, there are $MD2^{P-1-q}$ calls to $\tau$ on ranges of length $2^{q}$. The overall amount of memory it takes to store the activations is $O(MLD)$ (as is the case for the naive lazy approach) and the peak memory usage does not change the asymptotics.

###### Proposition 2 .

### 3.2 Across-layer Parallelization

One important feature of the proposed method is that it allows for higher parallelization across layers. In particular, the gray lines can be taken out of the inner loop and performed in parallel across all layers at once, just after the loop (Algorithm 3) - this is because all their inputs and outputs are disjoint. However, the red lines need to be performed sequentially since $a_{\ell,i}$ is a function of $b_{\ell,i}$ which is a function of $a_{\ell-1,i}$ for every $\ell$.

Figure: Algorithm 3 Flash Inference for LCSMs, parallel across layers

Note that this optimization can be applied to the eager and lazy approaches as well (provided that the red tiles are accounted for separately from the gray ones). However, the benefit of extra parallelization is practically relevant only in settings that are not memory-bandwidth bound - this is the case for smaller tiles which represent a significant amount of the tiles in our method, but an insignificant amount in the eager and lazy approaches. This is further discussed in Appendix E.

### 3.3 Memory Considerations

As follows from Proposition 1, the overall size of the inputs to (and outputs of) $\tau$ in the proposed algorithm is only $O(MDL\log L)$ - that is, we access on average activation values at $\log L$ positions per iteration as opposed to the naive implementations (such as lazy and eager approaches) that access an average of $\Omega(L)$ positions. This directly translates to data movement improvements by a factor of up to $L/\log L$.

A simple optimization of static memory is to never store $b$, but rather use $a_{\ell,i}$ to store $b_{\ell,i}$ until hitting line 8: this will work because from that point onwards $b_{\ell,i}$ is never used again and $a_{\ell,i}$ was not needed (or available) thus far. Hence, the static memory overhead is minimal: at least in the LCSM case, one directly aggregates contributions to the same activation tensor, thus not using any extra storage (on top of the inherent $MLD$ it takes to store the activations).

The only extra cost we pay is in the peak memory usage which is given by the largest tile - this is $O(MLD)$, but can be dropped to $O(LD)$ at essentially no cost by dropping parallelization for large tiles, as further discussed in Appendix E.

Lastly, if one is restricted to only using one GPU and the whole tensor of $M\times L\times D$ cannot fit it, it is possible to drop the requirement to only storing $M\times(L/2)\times D$ activations by dropping parallelization in the largest tile (which was alluded above). This, together with a more granular analysis of memory usage is further discussed in Appendices D and Appendix E.

## 4 The Flash Inference Framework

We propose a framework called Flash Inference to generalize our fast LCSM inference algorithm. To do so, we identify the main properties that are being exploited by Algorithm 2 and show how, granted these properties, one can get fast inference following the same kind of tiling.

Notice that in the case of LCSMs, it never mattered how one computes the contribution of a range of inputs to another range: only that it is well-defined (as soon as the inputs are available) and that it can be done more efficiently than the brute-force approach. Thus, the crux of Algorithm 2 is not convolution-specific, but rather the way we tile the space of contributions, while ensuring autoregressive generation can still take place.

### 4.1 Architectural Properties

In order for our framework to apply, we need the mixers involved to have certain properties:

- P.1
Contribution-based The used mixers work by aggregating contributions of each input position to each subsequent output position. That is, for any $\mathrm{mixer}_{\ell}$, there exist an *associative* aggregation function $\mathrm{agg}:{\mathcal{X}}^{*}\to{\mathcal{X}}$,
a read function $\mathrm{read}:{\mathcal{X}}\to\mathbb{R}^{D}$ and a contribution function $cont:\mathbb{R}^{D}\times\mathbb{N}\times\mathbb{N}\to{\mathcal{X}}$ such that:
$\displaystyle\mathrm{mixer}(y)_{i}=\mathrm{read}(\mathrm{agg}(\mathrm{cont}(y,
1,i),\mathrm{cont}(y,2,i),\ldots\mathrm{cont}(y,i,i)))$
(6)
where, ${\mathcal{X}}$ is a set of intermediate states and $\mathrm{read}$ is a function to map those back to embeddings. Recall that $\mathrm{agg}$ being associative means that $\mathrm{agg}(x_{1},x_{2},\ldots x_{\tau})=\mathrm{agg}(\mathrm{agg}(x_{1},
\ldots x_{i}),\mathrm{agg}(x_{i+1},\ldots x_{\tau}))$ for any $1\leq i<\tau$. For a sensible architecture, the size of ${\mathcal{X}}$ and cost of $\mathrm{agg}$ should be of order $D$.
In the case of self-attention this translates to having $\mathrm{read}\circ\mathrm{agg}$ simulate the $\mathrm{softmax}$, by letting ${\mathcal{X}}=\mathbb{R}^{D}\times\mathbb{R}$, $\mathrm{agg}=\sum$ and
$\displaystyle\mathrm{cont}(y,i,j)=(Vy_{i}\cdot e^{y_{j}^{\top}QK^{\top}y_{i}},
e^{y_{j}^{\top}QK^{\top}y_{i}})=(v_{i}\cdot e^{\langle{k_{i},q_{j}}\rangle},e^
{\langle{k_{i},q_{j}}\rangle})$
that is, the exponentially weighted value vector along with the exponential weight. Finally one can use $\mathrm{read}(v,w)=v/w$ to implemnent the $\mathrm{softmax}$ normalization step.
In the case of LCSMs, one can simply choose ${\mathcal{X}}=\mathbb{R}^{D}$, $\mathrm{read}$ be the identity function, $\mathrm{agg}$ be the sum again and $\mathrm{cont}(y,i,j)=y_{i}\odot\rho_{j-i}$.
- P.2
Query-independent The contribution function $\mathrm{cont}(y,i,j)$ does not depend on $y_{[i+1,L]}$. Note that this is the case in LCSMs since $y_{i}\odot\rho_{j-i}$ only depends on $y_{i}$. However, it is not the case for transformers, since $\mathrm{cont}(y,i,j)$ depends on $q_{j}$ which depends on $y_{j}$.

### 4.2 Setting and Definitions

Suppose P.1 holds and there exists an algorithm ${\mathcal{A}}$ that for any given input sequence $y$ and indices $l\leq r<l^{\prime}\leq r^{\prime}$, computes the contributions of $y_{[l,r]}$ to outputs at every position $p\in[l^{\prime},r^{\prime}]$:

$$ $\displaystyle{\mathcal{A}}(y,[l,r],[l^{\prime},r^{\prime}])_{l^{\prime}\leq p \leq r^{\prime}}=\mathrm{agg}(\mathrm{cont}(y,l,p),\mathrm{cont}(y,l+1,p), \ldots\mathrm{cont}(y,r,p)).$ $$

Furthermore, for this choice of ${\mathcal{A}}$, let ${\mathcal{T}}:\mathbb{N}\times\mathbb{N}\to\mathbb{N}$ such that for any $l\leq r<l^{\prime}\leq r^{\prime}$, evaluating $A(y,[l,r],[l^{\prime},r^{\prime}])$ takes at most ${\mathcal{T}}(r-l+1,r^{\prime}-l^{\prime}+1)$ FLOPs.

If we dropped the $r<l^{\prime}$ condition and set $l=l^{\prime}=1$ and $r=r^{\prime}=L$, this algorithm would actually represent the procedure one needs to perform a forward pass during training - which we refer to as the *static* setting. In the case of LCSMs, ${\mathcal{A}}$ is the FFT-based algorithm underlying Lemma 1, with an associated ${\mathcal{T}}(L_{1},L_{2})=D(L_{1}+L_{2})\log(L_{1}+L_{2})$, as opposed to the naive implementation that would have ${\mathcal{T}}_{\text{naive}}(L_{1},L_{2})=DL_{1}\cdot L_{2}$.

### 4.3 Main Result

Our main result can be phrased as follows:

###### Theorem 2 .

Under the assumptions P.1 and P.2, one can generate $L=2^{P}$ tokens autoregressively by performing, per layer, $L-1$ black-box calls to ${\mathcal{A}}$ as well as $L$ more calls to $\mathrm{cont}$, $\mathrm{agg}$, $\mathrm{read}$ and $\mathrm{block}$. Concretely, there are $L/2=2^{P-1}$ calls to ${\mathcal{A}}$ of length $1$ each, $L/4=2^{P-2}$ calls of length $2$ each and so on up to $1$ call of length $L/2$. Hence, neglecting the $\mathrm{cont}$, $\mathrm{agg}$, $\mathrm{read}$ and $\mathrm{block}$ part, the overall time complexity per mixer layer is:

$$ $\displaystyle\text{FLOPs}=\sum_{q=0}^{P-1}{2^{P-1-q}{\mathcal{T}}(2^{q},2^{q})}$ $$

Furthermore, during every token-generation iteration, the calls to ${\mathcal{A}}$ across different layers can be performed in parallel as there is no data-dependency between them.

Here, we are being ignorant of the $L$ calls to $\mathrm{cont}$, $\mathrm{agg}$, $\mathrm{read}$ and $\mathrm{block}$ because they only scale linearly with $L$ - however MLPs scale quadratically in $D$ and therefore these can be the limiting factor when $D$ is large in comparison to $L$.

The method underlying Theorem 2 is illustrated in Algorithm 4 and follows the same shape as its LCSM counterpart. The only difference is that ${\mathcal{X}}$ is not assumed to be $\mathbb{R}^{D}$ necessarily, so one cannot store the intermediate accumulated states in the same place they store the activations as per the discussion in Section 3.3 (though they could reuse some of the space since $a_{\ell,i}$ and $b_{\ell,i}$ never need to coexist). We use $b_{\ell,i}$ to store the inner part of Equation 6, namely $b_{\ell,i}$ incrementally computes $\mathrm{agg}(\mathrm{cont}(a_{\ell-1},1,i),\mathrm{cont}(a_{\ell-1},2,i),
\ldots\mathrm{cont}(a_{\ell-1},i,i))$. Finally, we explicitly moved the gray tile calls to happen after the inner loop and in parallel - this simply shows the last point of the theorem, namely that the calls to ${\mathcal{A}}$ can be done parallelly across layers.

Figure: Algorithm 4 Generic Flash Inference

###### Theorem 2 .

## 5 Experiments

We performed two kinds of experiment: one synthetic and one on the Hyena architecture (Poli et al., ). The synthetic setup defines all $\mathrm{block}$s to be MLPs with a hidden dimension of $2D$ and GELU activation; it also simply sets $a_{0,i+1}$ as $a_{M,i}$ plus some noise to avoid dependency on vocabulary size since that is out of the scope of our framework and will be negligible at scale. Note that this can be viewed as a sampler: a function from logits at the last layer and previous position to the next token’s embedding. In both settings the weights are initialized to random noise since it does not affect the runtime - this avoid unnecessarily training a new model for each hyperparameter choice.

In terms of notation, we introduce the batch dimension $B$ and keep $M,D,L$ refer to the number of layers, embedding dimension and number of tokens generated, respectively. We use $U$ to refer to the square-tile length, as per line 4. We sweep over $B\in\{1,2,4,8\},M\in\{18,36\}$(^2^22In the Hyena setting, M represents the effective number of mixers, but we fix the Hyena Order to 3 - so M=18 corresponds to 9 Hyena Operators and 18 mixers), $D\in\{256,768\}$ and generate tokens up to the greatest power of $2$ that fits the memory. All results are obtained by averaging over $4$ runs following $2$ runs of warm-up. We evaluate our approach on on the latest NVIDIA H100 and A100 GPUs.

For our baselines, (1) we consider the eager and lazy approaches described in Section 3.1, depicted in 1. We also consider their implementation exploiting our observation regarding parallelization across layers (3.2) (we denote the the parallel versions simply as lazy and eager).

Our Flash Inference framework explored various implementations of $\tau$, covered in Section 5.2. Our best method is a *Hybrid* that dynamically chooses the optimal implementation of $\tau$ depending on $(B,D,M,U)$.

### 5.1 Integrating Flash Inference in real world setting (Hyena Architecture)

Flash Inference significantly speeds-up the end-to-end inference by up to 1.6$\times$ and the convolution-based mixer component of the Hyena architecture by 50$\times$ compared to the baselines as shown in Figures‘2(a) and 2(b).

Figure 2(c) shows the per token response time of Hybrid and baselines. Hybrid shows low variance in per-token time except at the tokens positions where large tiles are computed. For a given sequence length of $L$, we have $L/2$ positions that use tile size 1, $L/4$ positions using tile size 2, and so on. That is, 93.75% of tokens use a tile size $U\leq 8$. Hence the spikes occur rarely.

Figure: (a)
Refer to caption: extracted/5930784/plots/hyena_time_breakdown_B4_N19_D768_L32768_H100.png

### 5.2 τ 𝜏 \tau italic_τ Implementations

Algorithm 2 assumes an implementation of primitive $\tau$ that accounts for the contribution of a range of inputs to a range of outputs of a convolution. We considered $7$ different implementations but, here, we only present results of the ones on the Pareto Frontier - that is, those optimal for at least some $(B,N,D,U)$ setting. There are $4$ such implementations, of two types:

- (1)
Depthwise-separable 1D Convolution: We use two types of implementations (a) *Conv1D* refers to the default PyTorch (Paszke et al., ) kernel, that requires explicit padding (b) *Flash Conv1D* refers to fused kernel by FlashFFTConv (Fu et al., ) - these are asymptotically quadratic in tile size $U$.
- (2)
FFT based convolution: we again use two implementations (a) PyTorch native, shown as *FFT*, that requires computing the FFTs of inputs, followed by pointwise multiplication and lastly performing the inverse FFT operation. (b) the second is provided by FlashFFTConv that does all the above in a single fused kernel shown as *FlashFFT*- both these scale as $O(U\log U)$.

### 5.3 Mixer Isolation study and Hybridization of τ 𝜏 \tau italic_τ implementations

We evaluate the different convolution implementations for all settings of $B,D,M,U$ and observe that each of these four implementation lies on the pareto frontier curve of tile size vs latency as shown in Figure 3(a). Our *Hybrid* approach dynamically chooses the best $\tau$ implementation for a given tile size $U$ based on the isolated empirically-measured efficiency of each implementation as shown in Figures 3(b). Figure 3(c) shows the cumulative token time breakdown for the mixer and non-mixer components of the synthetic setup for all our $\tau$ implementations. We observe an increase in the non-mixer components due to significant reduction in GPU kernel time that exposes the CPU kernel dispatch overhead for the MLP blocks in the non-mixer components.

Figure: (a)
Refer to caption: extracted/5930784/plots/conv_bench_B=1_M=36_D=256_maxT=65536_H100.png

### 5.4 Improvements justification

This significant speed-up can be attributed to:

- (1)
The significantly lower $O(L\log^{2}L)$ FLOPs required for our tile computation approach compared to the $\Omega(L^{2})$ FLOPs required by *Eager* and *Lazy* counterparts. This is shown in Figure 2(b) where methods based on our tiling outperform by a large constant the quadratic methods in terms of time spent on the mixer components.
- (2)
Drastically reduced activation memory access from $\Omega(L^{2})$ for *Eager* and *Lazy* to out $O(L\log L)$ tiling-based methods. This can be shown through the performance of *Flash Conv1D* (Figure 3(b)) which outperforms lazy and eager by a margin although it also performs $\Omega(L^{2})$ FLOPs - it does so in a more memory-friendly and kernel-optimizable way.
- (3)
The dynamic choice of best $\tau$ implementation for given tile $U$ - hybrid outperforming any method using a fixed implementation (Figure 3(b)).
- (4)
Our engineering contributions: first, the DFT for the convolutional kernel is pre-computed ahead of time for $\log_{2}(L)-1$ tile sizes. Second, Flash-FFT configurations are pre-initialized for these tile sizes to maximize hardware performance. Third, right padding is used instead of left padding to reduce computation time by half. Fourth, properties of circular convolution are exploited to halven FFT length. Finally, tile calculations are parallelized across layers to saturate memory bandwidth for small tile computations and optimize computation for large tile computations, resulting in improved performance (one can notice improvements of $10-20$% even in the Eager and Lazy implementations alone)

## 6 Conclusion and Further Work

We propose a framework for performing inference in certain autoregressive sequence models. Among such models, LCSMs such as Hyena, are noteworthy: there, our framework provides an $O(L\log^{2}L)$ inference algorithm which, when run empirically, yields up to $\times 1.6$ time-efficiency improvement. The framework exploits a causal, fractal tiling that helps save on data movement and share computation. Moreover, it allows for almost-complete across-layers parallelization of $\mathrm{mixer}$-related workload.

An interesting future direction to pursue is that of designing architectures that fit out framework requirements and thus get fast-inference by construction. Furthermore, in the class of LCSMs, we have noted that one can achieve the same theoretical complexity if filters are made data-dependent, and, while previous works Arora et al. (); Karami and Ghodsi () have shown the potential for these, they are not yet causal so looking into how to make filters data-dependent in a causal way is another promising direction.

## 7 Acknowledgments

This work was partially enabled by use of the Kempner AI cluster. Sham Kakade acknowledges: this work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. SK and CO also acknowledge support from the Office of Naval Research under award N00014-22-1-2377, and the National Science Foundation Grant under award #IIS 2229881.

## Appendix A Proofs

###### Lemma .

Let $1\leq l\leq r\leq l^{\prime}\leq r^{\prime}\leq L$ represent ranges of lengths $L_{1}=r-l+1$ and $L_{2}=r^{\prime}-l^{\prime}+1$ of $y$ and $z$, respectively.
There exists an FFT-based algorithm running in $O(L_{1}+L_{2})$ space and $O((L_{1}+L_{2})\log(L_{1}+L_{2}))$ time complexity that, given access to $y_{[l,r]}$, computes the aggregated contributions of all $y_{[l,r]}$ to all of $z_{[l^{\prime},r^{\prime}]}$.
That is, for each $l^{\prime}\leq t\leq r^{\prime}$, it computes:

$$ $\displaystyle\tau(y,[l,r],\rho,[l^{\prime},r^{\prime}])_{t}\triangleq\sum_{i=l }^{r}{y_{i}\cdot\rho_{t-i}}$ (7) $$

for all $l^{\prime}\leq t\leq r^{\prime}$.

###### Proof.

Consider performing a convolution between $y_{[l,r]}$ and $\rho_{[l^{\prime}-r,r^{\prime}-l]}$. It holds that for each $j\in[l^{\prime},r^{\prime}]$ and $i\in[l,r]$, $j-i$ is in the range of $\rho$. On the other hand, truncating the output of the convolution appropriately, one can keep only the corresponding outputs for each $j\in[l^{\prime},r^{\prime}]$. The overall size of the FFT is then given by the sum of the lengths, which is at most $r-l+1+(r^{\prime}-l)-(l^{\prime}-r)+1=2(r-l+1)+(r^{\prime}-l^{\prime}+1)-1=O((
r-l+1)+(r^{\prime}-l^{\prime}+1))$. Since an FFT of order $L$ runs in $O(L\log L)$ and takes up $O(L)$ memory, the conclusion follows.
∎

###### Theorem .

Under the assumptions P.1 and P.2, one can generate $L=2^{P}$ tokens autoregressively by performing, per layer, $L-1$ black-box calls to ${\mathcal{A}}$ as well as $L$ more calls to $\mathrm{cont}$, $\mathrm{agg}$, $\mathrm{read}$ and $\mathrm{block}$. Concretely, there are $L/2=2^{P-1}$ calls to ${\mathcal{A}}$ of length $1$ each, $L/4=2^{P-2}$ calls of length $2$ each and so on up to $1$ call of length $L/2$. Hence, neglecting the $\mathrm{cont}$, $\mathrm{agg}$, $\mathrm{read}$ and $\mathrm{block}$ part, the overall time complexity per mixer layer is:

$$ $\displaystyle\text{FLOPs}=\sum_{q=0}^{P-1}{2^{P-1-q}{\mathcal{T}}(2^{q},2^{q})}$ $$

Furthermore, during every token-generation iteration, the calls to ${\mathcal{A}}$ across different layers can be performed in parallel as there is no data-dependency between them.

###### Proof.

Figure 1 shows why the tiling used in Algorithm 2 covers every pair of contributions exactly once and in the correct order (remember, we only assumed $\mathrm{agg}$ to be associative, not necessarily commutative).

Only part remaining to be proved is the right counting of calls to ${\mathcal{A}}$ per layer: calls of length $2^{r}$ happen whenever $2^{r}$ divides $i$ but $2^{r+1}$ does not - that is, for $i\in\{2^{r},3\cdot 2^{r},5\cdot r^{2},\ldots(2^{P-r}-1)\cdot 2^{r}\}$, so $2^{P-r-1}$ calls.
∎

###### Lemma .

###### Proof.

###### Theorem .

###### Proof.

## Appendix B Extension to Data-Dependent Filters

The reason why data-dependent filters are harder to deal with is that $\mathrm{cont}(y,i,j)=y_{i}\cdot\rho_{j-i}$ depends on both $y_{i}$ and on what $\rho_{j-i}$ depends on. By only assuming causality, we can only access $\rho_{j-i}$ after $z_{j-i-1}$ has been computed. This stops Algorithm 2 from working out-of-the-box since when $i$ is a power of $2$, in order to account for the contribution of $a_{\ell-1,[1,i]}$ to $a_{\ell,[i+1,2i]}$ one will need access to $\rho_{\ell,[1,2i-1]}$ which is not guaranteed to be accessible (only $\rho_{\ell,[1,i]}$ can be assumed to be known at this stage).

The modified algorithm is shown in Algorithm 5 and precisely follows the tiling of van der Hoeven [] and, thus, its correctness transfers. Note that rather than using the ${\mathcal{A}}$ algorithm, it directly uses the untruncated convolution (implementable via FFT) - in the contribution space this corresponds to a parallelogram tile rather than rectangle - and this cannot be simulated via a rectangle, although the other way around is possible (and is what we did through Lemma 1). This implementation performs convolutions between two sequences of length $U$ each and uses the whole output and thus requires an order $2U$ FFT. Note that this happens twice for a given $i$ when $i+1$ is not a power of $2$ (almost always). Algorithm 2, on the other hand, only performs a convolution between a length-$U$ sequence and a length-$2U$ sequence. However, as noted in Appendix C, we can get away without padding and thus performing only one order $2U$ FFT. Thus, our proposed tiling improves FLOPs by a factor of $2$, but it requires the kernel to be data-independent.

Figure: Algorithm 5 Flash Inference for LCSMs with data-dependent filters

## Appendix C Implementation Improvements

In Algorithm 2, the calls to ${\mathcal{A}}$ all have a shape of ${\mathcal{A}}(y,i-U+1,i,i+1,i+U)$ where $U$ is a power of $2$ that divides $i$. Following the proof of Lemma 1, we note that to implement ${\mathcal{A}}$, we need to perform the convolution of a segment of length $U$ of $y$ and a prefix of length $2U$ of $\rho$. However, following the convolution, of the $3U-1$ outputs, indexed $[0,3U-2]$, we are only interested in the middle $U$ of them, namely indices $[U,2U-1]$. The canonical way of performing convolutions via FFT involves padding by enough $0s$ and using an order large enough to store the whole output which rounded up to the closest power of $2$ would mean an order $4U$ FFT. However, using a $2U$ FFT call, which will perform a cyclical convolution, is enough for our purposes, since the values of interest are not affected by the cyclicity - that is, when folding outputs at $[2U,3U-2]$ onto $[0,U-2]$, we do not intersect $[U,2U-1]$. Furthermore, there are only $\log L$ different lengths of prefixes of $\rho$ involved in these convolution so one could, in fact, cache the DFTs for these as a precomputation step, thus dropping the number of DFT’s per convolution from $3$ to $2$, speeding up by a further $\times 1.5$ factor.

## Appendix D Storing Only Half of the Activations

If the whole activation tensor ($M\times L\times D$) cannot fit the memory of one GPU, we can save a factor of $2$ along the $L$-axis. To do this, observe that after $(L/2)$^th iteration completes we never need to look back to any activation at position $i\leq L/2$. Hence, one can hope to reuse the space used for the first half to store the second half of activations. While doing this directly by having ${\mathcal{A}}$ work as an in-place operator requires storing $MLD$ intermediate values, one can choose to not apply ${\mathcal{A}}$ parallelly across all layers (or even dimensions), as discussed in Section 3.3. Processing the tiles sequentially allows for the reuse of the space for interim results; once the result of ${\mathcal{A}}$ is read, we place it back to where the input used to be - similarly in nature to gradient accumulation - always placing the output back into the input slot. This way, the peak extra memory required for the call to ${\mathcal{A}}$ can be $LD$ or even $L$.

As we get the output of ${\mathcal{A}}(a_{0},1,L/2,L/2+1,L)$, which has a shape of $(L/2)\times D$, we overwrite $a_{0}$ by it and then move on to the next layer. We can do this because the contribution of $a_{0,i}$ to $a_{1,i}$ has already been properly accounted for so the first half of $a_{0}$ truly becomes irrelevant for next layers and iterations to come. Hence, $a_{0}$ contains the intended second half of $a_{1}$. We then proceed to do the same thing for $a_{1}$ and so on. At the end, $a_{[0\ldots M-1]}$ will contain the second halves of $a_{[1\ldots M]}$, so we can shift these by $1$, emptying $a_{0}$ and now finally having $a$ represent the second halves of activations - the first halves have been discarded.

Although per layer, one will have a peak memory containing both $a_{\ell}$, the output, and whatever other necessary temporary values needed to perform ${\mathcal{A}}$, this does not scale with $M$ - because we essentially reuse the extra space needed by ${\mathcal{A}}$ across different layers. We seemingly pay the cost of not performing the calls to ${\mathcal{A}}$ in parallel across layers, but if storing the $M\times L\times D$ tensor of activations was an issue, then one would very likely have to make this sort of compromise to be able to run ${\mathcal{A}}$ in the first place.

## Appendix E Discussion on memory

In Section 3.3, we discussed the peak memory usage being given by the gray calls to ${\mathcal{A}}$. Since each call to ${\mathcal{A}}$ performs $D$ FFTs of length $L$ each (when dealing with the biggest tile of side $L/2$), the naive implementation would use up to $MD\cdot L$ extra memory to perform these FFT calls.
However, this assumes the biggest gray tile is being processed in parallel across all layers and across all dimensions at once. We do not need to maximize parallelization on all sizes of tiles: for the largest ones, we can process the tiles at different layers sequentially by reusing the space, thus dropping to an overhead of $O(LD)$, or even just $O(L)$ if one is to treat sequentially the different dimensions. Dropping parallelization could theoretically yield a slowdown, but that is only when we are not memory bandwidth-bound which is the case if one has to worry about allocating more than $O(LD)$ - in that case, the FFT calls are anyway not happening in parallel. Hence, one can easily drop the peak memory overhead to $O(LD)$ or even $O(L)$ without incurring an actual time cost.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Please see my question above.
soundness: 4
strengths: * As far as  I know, the interpolation perspective presented in this paper is original and inspiring. The writing is exceptionally clear, and Figure 1 has been extremely helpful in understanding the proposed method.

* The algorithm introduced in the paper largely solves the long-standing problem of quadratic inference complexity for long convolution models like SGConv and Hyena. This has been a significant bottleneck for the practical deployment of these architectures (note that there are also other long conv architectures that do not suffer from this, please see below).
summary: This paper introduces a novel linear-time inference algorithm for long convolution architectures like SGConv and Hyena. These models typically have a quadratic inference complexity with respect to the sequence length N, which can be prohibitive for large-scale applications.

The key insight behind the proposed approach is a clever partitioning and precomputation strategy for the contributions to the convolution of future outputs. This allows the inference complexity to be reduced to O(N log^2 N). The paper provides extensive experimental evidence demonstrating the acceleration of inference achieved by this new algorithm.
weaknesses: This is a good paper and there is no much weakness to say about its methodology. However, I find the significance of the work depends on a line of work on long convolution architectures that the authors unfortunately have not discussed or compared.

Long convolution kernels can be contructed from smaller convolutions in a tree style hierarchical dilations such as those in WaveNet [1]. Recently, people have shown that these architectures, with nonlinearities removed and weight sharing, can be interpreted as having a wavelet based state and are competitive with SSM and long convs on sequence modeling benchmarks [2].

Crucially, these dilated convolution architectures also support linear-time inference, as detailed in [3], by maintaining a cache per layer. This allows for efficient inference without the need for techniques like FFT that are required for standard long convolutions.

Given these relevant prior results, I believe a comparison of the proposed approach to these dilated convolution models and their linear-time inference capabilities would greatly strengthen the paper. These alternative architectures offer a potentially simpler approach. As a reader of this paper, it would be nice to know when it is preferable to have a less structured long conv combined with the inference acceleration presented here.

References:

[1] Van Den Oord, Aaron, et al. "Wavenet: A generative model for raw audio." arXiv preprint arXiv:1609.03499 12 (2016).

[2] Shi, Jiaxin, Ke Alexander Wang, and Emily Fox. "Sequence modeling with multiresolution convolutional memory." International Conference on Machine Learning. PMLR, 2023.

[3] Paine, Tom Le, et al. "Fast wavenet generation algorithm." arXiv preprint arXiv:1611.09482 (2016).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: I think the paper needs to position with respect to other types of LCSMs with more nuance. For example, it could benchmark how the proposed algorithm for general LCSMs compares to the inference speed of recurrent LCSMs; e.g. for a given sequence length $L$, at what recurrent state size does the speed of a recurrent SSM cross over the speed of the proposed algorithm? This would at least provide some more useful context for the reader interested in the pure algorithmic aspects.

However, overall, in order for this paper to be valuable to the machine learning community, it should be applied to actual models, and current understanding of LCSMs is that general LCSMs essentially do not benefit over low-dimensional LTI SSMs. The contributions are interesting from a purely algorithmic perspective, and the paper does suggest potential extensions beyond time-invariant convolutions. However, it is not clear whether this has any chance of being extended into performant models - no actual downstream model is proposed. Thus, I think that a conference such as ICLR is not the most appropriate venue for such a submission without anchoring the algorithmic contributions to real models. In order to increase my score substantially, I think the paper needs to show a practical utility, for example either
- Showing that there exist classes of models that the algorithm can be applied to that are not dominated by others (e.g. showing that some non-SSM LCSM model is actually faster or more performant than an equivalent SSM)
- Showing that there exist data-dependent extensions that outperform data-independent convolutions in modeling performance
soundness: 4
strengths: The paper's main strength is that it is technically sound and novel, and addresses the problem it aims to solve.
- The technical writing is clear and is helped by the inclusion of helpful graphics and rigorous algorithm boxes.
- Many considerations and variants of the core algorithm are proposed.
- An actual implementation is provided and all variants and baselines are benchmarked empirically.

There is a conception that long convolutions cannot be implemented efficiently in autoregressive inference settings, and so I do think that this paper presents an original algorithmic contribution.
summary: This paper proposes more efficient algorithms for auto-aggressive inference of long convolutional sequence models (LCSMs). The aggregate running time on a sequence of length $L$ is reduced from $O(L^2)$ to $O(L \log^2 L)$, and the actual wall clock time of the implemented algorithm reflects substantial increases in efficiency.
weaknesses: While the paper provides a technical contribution, the paper's main weakness is that of significance and direction with respect to the broader field; it aims to solve a problem that I believe does not need solving. Correspondingly, the papers writing (in terms of positioning and related works / baselines) could also use improvement.

- The paper's related work is sparse and I think it is important to present the lineage of these models more carefully. The original (depth separable) LCSMs were independently developed by two lines of work: the implicit convolution (CKConv and FlexConv) line of work, and the SSM line of work (LSSL, S4, DSS/S4D, H3, MEGA/Megalodon, and many more). Even though I understand why the paper deliberately puts emphasis on LCSMs that are not SSMs, because this is where its results are most applicable to, the positioning is at times misleading (e.g. in paragraph 2 of the introduction, where LCSMs are implicitly defined as non-SSM models, even though LTI SSMs are in fact the original models that popularized LCSMs).
- It is odd that the paper is heavily anchored around the Hyena architecture, even though it does not actually do anything Hyena-specific. The experiments don't use an actual trained model or look at empirical performance, only the speed of an architecture, for which any model with a LCSM convolution (e.g. H3, to which Hyena is equivalent except for the choice of convolution kernel) could equivalently be substituted into the writing without changing the algorithmic results.
- A distinction is made that low-dimensional LTI SSMs cannot represent general LCSMs, which is true and where the paper's potential benefits can come from. However, current understanding of LCSMs is that they need to be defined using certain priors (e.g. baking in exponential decay as in SGConv and Hyena) that essentially are similar to the priors imposed by low-dimensional SSMs. Empirically, it is known that there is little difference between these models (e.g. https://arxiv.org/abs/2312.00678v1 Fig 1, which claims that H3 in fact performs better than Hyena, where the only difference is the choice of convolution kernel parameterization). Thus a major weakness of this paper is that it is currently only applicable to models that have better alternatives empirically.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: ### Additional questions and observations

- Contribution 1. This is true, however, for clarity, I would recommend that the authors mention Laughing Hyena before making this claim to put in context that that method is not exact.

- Line 233. There’s a typo here that changes the whole meaning of the sentence. Please fix.

 ### Conclusion

While I believe that the core contributions of this paper are very valuable, I am not really convinced by the current presentation of the paper. I would recommend the authors to improve readability and make specific the different design decisions that make the paper. 

Do note that the contribution of this paper is rather an algorithmic one. Yet, the current very abstract depiction of the method makes it unnecessarily difficult to implement and reproduce. I am, therefore, only able to provide this paper a “weak acceptance” at this time. However, that if the authors were to improve over the weaknesses outlined here, I would be more than happy to increase my score.
soundness: 4
strengths: - The method is novel and offers interesting improvements in the inference speed of LCSMs.

- The paper offers interesting perspectives that could be used for the design of more efficient (causal, input-dependent) LCSMs in the future.
summary: This work introduces a method to speed-up the autoregressive inference of long convolution sequence models (LCSMs) to near linear time. The approach is based on findings from relaxed polynomial interpolation’s literature, which is gratefully adapted to convolutions.
The resulting algorithm results in important speeds-up both for the position-mixing components and end-do-end inference.
weaknesses: - The main weakness of the paper is that the presentation, design decisions and final implementation of the method remains quite abstract, even after reading the paper multiple times. Given that the paper presents an inference strategy, it should be feasible to have an stand-alone implementation (at least for one layer) incorporated in the Appendix of the paper. This would give clarity to the final, concrete version of the algorithm.

- Next, I feel that the presentation of the paper could be improved. For example, the function $\tau$ –which is crucial for the method– remains undefined through the whole body of the paper up to Sec. 4.2, where it is briefly and loosely defined. It is mentioned that 7 possible versions were tested, but only 4 were mention-worthy. Then, it is mentioned that some are used, but it is not specified in which ranges and under which parameters one is preferred over the other. Given that this is core to the method, this should definitely be improved.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. The paper suggests applicability beyond LCSMs. Could the authors provide additional empirical results or theoretical discussion on adapting the method to different architectures, such as transformers, beyond high-level discussion?

2. While data-independent filters optimize the framework's efficiency, data-dependent filters could require modifications. Could the authors elaborate on practical adaptations to support data-dependent filters effectively, and whether the framework maintains similar efficiency in such cases?

3. The work compares primarily with the eager and lazy approaches. How does the proposed framework perform against other efficient models or architectures, such as those using efficient transformer variants?

4. The paper assumes causality in model design. Are there potential modifications to the framework that could apply it to non-causal architectures? This could expand the impact and applicability of the proposed method.
soundness: 2
strengths: 1. The paper addresses a fundamental issue in sequence models, particularly long convolution sequence models (LCSMs) like Hyena, where inference time scales quadratically with sequence length. The proposed framework reduces this to quasilinear time, achieving significant improvements by leveraging a tiling-based approach inspired by relaxed polynomial interpolation.

2. Besides speed, the paper focuses on reducing memory movement, a bottleneck in handling large models. It suggests methods for activation storage optimization and discusses adjustments for memory-constrained hardware, making the work relevant for deployment.

3. The paper presents thorough empirical results that demonstrate the proposed method's effectiveness. It reports up to 1.6x end-to-end improvement in speed for Hyena and a 50x speedup within the convolutional mixer component, which provides concrete evidence of practical impact.
summary: This paper proposes a method for speeding up Long Convolution Sequence Models (LCSM)'s exact inference to quasilinear time, identifies the key properties that make this possible, and proposes a general framework that exploits these. The proposed approach is inspired by relaxed polynomial interpolation and uses a tiling method that minimizes memory movement and enhances computation sharing, allowing near-complete parallelization across layers in the architecture’s position-mixing component. Through a proof-of-concept implementation on Hyena, we demonstrate up to a 1.6× improvement in end-to-end inference time and a 50× speedup in the position-mixing part.
weaknesses: 1. While the study is well-documented, the experiments focus primarily on synthetic setups and Hyena. Additional tests on a broader range of sequence models or with real-world tasks could further validate the framework’s generalizability.

2. The framework's efficiency relies on data-independent filters for optimal performance. Although data-dependent filters can be supported, doing so may lead to higher complexity or additional constraints. This could limit the application in cases where data-dependent filters are essential.

3. The framework assumes autoregressive causality, which could restrict its application to non-causal architectures. This point is mentioned briefly, but it would benefit from further exploration of how the framework could adapt to different model constraints.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. How does the performance of Flash Inference compare to existing methods when dealing with very long sequences (e.g., $L>$32K) in non synthetic scenarios? Are there specific challenges (other than memory) or additional optimizations possible at these scales?
2. Could you elaborate on how the tiling strategy might be adapted for architectures with more complex dependencies between layers? This seems particularly relevant for hybrid architectures that combine convolution and attention mechanisms.
2. The main limitation of the proposed method is certainly the linear growth of cache size with sequence lenght. How does this affect the practical benefits of exact vs approximate inference methods? In the approximate method, the cost of distillation can be amortized offline for a subsequent constant memory autoregressive generation. Moreover, one can trade additional memory (number of states in the SSM) for additional precision in representing the filters. Using a pre-trained LCSM, could you verify the accuracy-memory trade-off of the approximate vs exact method for practical sequence lengths? 
3. Have you considered how this framework might be extended to handle structured sparsity patterns in the filters? This could potentially lead to additional efficiency gains while maintaining the theoretical guarantees.
soundness: 4
strengths: 1. Technical Innovation and Soundness:
- The paper presents a mathematically rigorous approach to improving inference efficiency, with clear proofs and careful analysis of complexity bounds
- The implementation details are thoroughly considered, including memory optimization and parallelization strategies
2. Practical Impact:
- The achieved speedups (1.6× end-to-end, 50× for position-mixing) are significant and well-documented
- The method is exact rather than approximate, maintaining model fidelity while improving performance
- The framework extends beyond just LCSMs, providing valuable insights for future architecture design
3. Experimental Validation:
- Comprehensive empirical evaluation across different hyper parameters parameters ($B$, $M$, $D$)
- Careful ablation studies of different $\tau$ implementations
- Clear breakdown of performance improvements and their sources
summary: This paper introduces "Flash Inference", a novel framework for efficient inference in long convolution sequence models (LCSMs) employing a tiling strategy inspired by relaxed polynomial interpolation . The authors propose an algorithm that achieves quasilinear $O(L \log^2_2 L)$ time complexity for exact inference, improving upon the quadratic complexity of standard approaches without resorting to any approximation techniques. The work specifically demonstrates the method's effectiveness on Hyena architectures, achieving significant end-to-end speedup and even more significant improvement within the position-mixing component of the model. Beyond LCSMs, the authors identify key properties that enable such speedups and propose a general framework to guide future architecture design. The proposed approach further reduces memory movement and enables parallel computation across layers.
weaknesses: 1. Technical Presentation:
- The notation could be improved for clarity, particularly in handling subscripts that simultaneously indicate sequence position, feature dimension, and layer. For example, you could consider using superscripts to indicate features (channels) and/or paranthesis around the subscripts/superscripts for the layers. 
- the use of "$\mapsto$" notation on line 159 is imprecise; it should be $y\mapsto z$.
- Algorithm 1 should reference equation (3) for the definition of $\tau$ for better clarity.

2. Contextual Discussion:
- The discussion of approximate inference methods could be more nuanced. The statement about "defeating the purpose of using LCSMs instead of LTI SSMs" is overly strong, given that Hyena filters are intentionally underparameterized and some *implicit regularization* (convergence of the filters to ones that can be effectively represented by a truncated basis of complex exponential functions) is to be expected. Te reviewer suggests that the authors provide a more balanced discussion of the trade-offs between exact and approximate methods, considering the intentional underparameterization of Hyena filters and the potential benefits of implicit regularization.
- The introduction would benefit from a clearer distinction between prefilling and autoregressive generation phases in the complexity analysis. Elaborating on this could strengthen the clarity of how the method affects the performance of the system.
- The claim about "popular misconception" (line 109) regarding the realization theory of LTI systems could be better supported with appropriate citations. It is the reviewer belief that the this is not a misconception in the control theory literature as the problem of realization of LTI system has bee studied for more than 60 years  (e.g., Kalman and Ho 1960). Perhaps the authors could clarify on that.

3. Technical Specifics:
- The asymptotic complexity expressions should consistently specify the base of the logarithm ($log_2$)
- The parallelization discussion could benefit from more detailed analysis of the memory bandwidth implications
- The interesting extension to data-dependent filters, while covered in the appendix, could deserve more attention in the main text if supported by quality experiments (do we need input-dependent filters in LCSMs?)

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Please kindly provide more details to the questions raised in the "Weaknesses" section. 
2. How about the case when the non-mixers’ runtime is dominant, e.g. the large MLP module is common in practice (such as Transformers in applications). Do the improvements shown in Fig. 2(a) & 3(c) become marginal? 
3. For Fig. 3(c), why is the non-mixers’ runtime less for Eager (NP) & Lazy (NP)? In addition, it seems that there is no reduction of mixers’ runtime of Hybrid compared to (Flash) FFT?
soundness: 3
strengths: 1. The paper is well-written and easy to follow. 
2. The complexity analysis regarding both the computation and memory is detailed. 
3. The experiments are also provided to numerically verify the effectiveness of the proposed method.
summary: This paper develops a computation framework for efficient inference of long convolution sequence models (LCSMs). The proposed approach is based on relaxed polynomial interpolation, and can speed up LCSMs’ exact inference to the quasilinear time, as is shown by both complexity analysis and numerical verifications.
weaknesses: 1. The main concern is about the novelty of this paper, particularly compared to the work by van der Hoeven and "dynamic FFT". It would be clearer if authors can provide more details about the differences or improvements of this work compared to former references. 
2. The applied setting seems somewhat limited, since it currently works only for Hyena-related architectures. Does the proposed method have the potential to inspire further extensions for general state-space models (SSMs; as with an equivalent convolutional filter) and Transformers (despite reasons stated in Sec. 3.4)? 
3. The numerical experiments part can be strengthened by e.g. adding tests for multiple configurations of hyper-parameters.

Meta-review: This paper proposes a method for speeding up long convolution sequence models exact inference to quasilinear time. The paper provides extensive experimental evidence demonstrating the inference benefits of this new algorithm. The reviews for the paper were mostly positive with reviewers highlighting the significance of the work and novelty. I recommend acceptance.

Decision: Accept (Poster)